From 51c9cb9a1df6808715b231b29938feb2b4564b74 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 21 Oct 2024 14:15:22 +0800 Subject: [PATCH 001/324] feat: prewarm (#15) Signed-off-by: usamoi --- Cargo.toml | 2 ++ src/algorithm/rabitq.rs | 6 ++++++ src/gucs/mod.rs | 7 +++++++ src/gucs/prewarm.rs | 31 +++++++++++++++++++++++++++++++ 4 files changed, 46 insertions(+) create mode 100644 src/gucs/prewarm.rs diff --git a/Cargo.toml b/Cargo.toml index bef4e791..ab88f57a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,8 @@ path = "./src/bin/pgrx_embed.rs" [features] default = [] +pg12 = ["pgrx/pg12"] +pg13 = ["pgrx/pg13"] pg14 = ["pgrx/pg14"] pg15 = ["pgrx/pg15"] pg16 = ["pgrx/pg16"] diff --git a/src/algorithm/rabitq.rs b/src/algorithm/rabitq.rs index 45b512be..a0522bef 100644 --- a/src/algorithm/rabitq.rs +++ b/src/algorithm/rabitq.rs @@ -47,6 +47,12 @@ fn orthogonal_matrix(n: usize) -> Vec> { static MATRIXS: [OnceLock>>; 1 + 2000] = [const { OnceLock::new() }; 1 + 2000]; +pub fn prewarm(n: usize) { + if n <= 2000 { + MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); + } +} + pub fn project(vector: &[f32]) -> Vec { use base::scalar::ScalarLike; let n = vector.len(); diff --git a/src/gucs/mod.rs b/src/gucs/mod.rs index cc0145ef..e072844c 100644 --- a/src/gucs/mod.rs +++ b/src/gucs/mod.rs @@ -1,7 +1,14 @@ pub mod executing; +pub mod prewarm; pub unsafe fn init() { unsafe { executing::init(); + prewarm::init(); + prewarm::prewarm(); + #[cfg(any(feature = "pg12", feature = "pg13", feature = "pg14"))] + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"rabbithole".as_ptr()); + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] + pgrx::pg_sys::MarkGUCPrefixReserved(c"rabbithole".as_ptr()); } } diff --git a/src/gucs/prewarm.rs b/src/gucs/prewarm.rs new file mode 100644 index 00000000..8cd7456f --- /dev/null +++ b/src/gucs/prewarm.rs @@ -0,0 +1,31 @@ +use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; +use std::ffi::CStr; + +static PREWARM_DIM: GucSetting> = GucSetting::>::new(None); + +pub unsafe fn init() { + GucRegistry::define_string_guc( + "rabbithole.prewarm_dim", + "prewarm_dim when the extension is loading.", + "prewarm_dim when the extension is loading.", + &PREWARM_DIM, + GucContext::Userset, + GucFlags::default(), + ); +} + +pub fn prewarm() { + if let Some(prewarm_dim) = PREWARM_DIM.get() { + if let Ok(prewarm_dim) = prewarm_dim.to_str() { + for dim in prewarm_dim.split(',') { + if let Ok(dim) = dim.trim().parse::() { + crate::algorithm::rabitq::prewarm(dim as _); + } else { + pgrx::warning!("{dim:?} is not a valid integer"); + } + } + } else { + pgrx::warning!("rabbithole.prewarm_dim is not a valid UTF-8 string"); + } + } +} From 4723e2b19f7908c6776f560307485254caaec6f0 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 21 Oct 2024 15:59:39 +0800 Subject: [PATCH 002/324] feat: index progress report (#16) Signed-off-by: usamoi --- Cargo.lock | 206 +++++++++++++++++++++-------------------- Cargo.toml | 4 +- src/algorithm/build.rs | 16 +++- src/index/am.rs | 22 ++++- 4 files changed, 146 insertions(+), 102 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 97654851..067ba488 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.86" +version = "1.0.90" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da" +checksum = "37bf3594c4c988a53154954629820791dde498571819ae4ca50ca811e060cc95" [[package]] name = "approx" @@ -59,9 +59,9 @@ dependencies = [ [[package]] name = "autocfg" -version = "1.3.0" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0" +checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" [[package]] name = "base" @@ -88,7 +88,7 @@ source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole- dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] @@ -107,7 +107,7 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] @@ -152,9 +152,9 @@ dependencies = [ [[package]] name = "bytemuck" -version = "1.17.1" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "773d90827bc3feecfb67fab12e24de0749aad83c74b9504ecde46237b5cd24e2" +checksum = "8334215b81e418a0a7bdb8ef0849474f40bb10c8b71f1c4ed315cff49f32494d" [[package]] name = "byteorder" @@ -164,9 +164,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.7.1" +version = "1.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8318a53db07bb3f8dca91a600466bdb3f2eaadeedfdbcf02e1accbad9271ba50" +checksum = "428d9aa8fbc0670b7b8d6030a7fadd0f86151cae55e4dbbece15f3780a3dfaf3" [[package]] name = "cargo_toml" @@ -180,9 +180,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.1.15" +version = "1.1.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57b6a275aa2903740dc87da01c62040406b8812552e97129a63ea8850a17c6e6" +checksum = "c2e7962b54006dcfcc61cb72735f4d89bb97061dd6a7ed882ec6b8ee53714c6f" dependencies = [ "shlex", ] @@ -298,7 +298,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] @@ -309,7 +309,7 @@ checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" dependencies = [ "darling_core", "quote", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] @@ -327,7 +327,7 @@ source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole- dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] @@ -353,7 +353,7 @@ checksum = "f282cfdfe92516eb26c2af8589c274c7c17681f5ecc03c18255fe741c6aa64eb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] @@ -465,9 +465,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.14.5" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb" [[package]] name = "heapless" @@ -518,12 +518,12 @@ checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" [[package]] name = "indexmap" -version = "2.5.0" +version = "2.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68b900aa2f7301e21c36462b170ee99994de34dff39a4a6a528e80e7376d07e5" +checksum = "707907fe3c25f5424cce2cb7e1cbcafee6bdbe735ca90ef77c29e84591e5b9da" dependencies = [ "equivalent", - "hashbrown 0.14.5", + "hashbrown 0.15.0", ] [[package]] @@ -545,9 +545,9 @@ checksum = "7655c9839580ee829dfacba1d1278c2b7883e50a277ff7541299489d6bdfdc45" [[package]] name = "itertools" -version = "0.12.1" +version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" dependencies = [ "either", ] @@ -572,9 +572,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.158" +version = "0.2.161" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439" +checksum = "8e9489c2807c139ffd9c1794f4af0ebe86a828db53ecdc7fea2111d0fed085d1" [[package]] name = "libloading" @@ -622,9 +622,9 @@ checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" [[package]] name = "memmap2" -version = "0.9.4" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe751422e4a8caa417e13c3ea66452215d7d63e19e604f4980461212f3ae1322" +checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f" dependencies = [ "libc", ] @@ -637,9 +637,9 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" [[package]] name = "nalgebra" -version = "0.33.0" +version = "0.33.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c4b5f057b303842cf3262c27e465f4c303572e7f6b0648f60e16248ac3397f4" +checksum = "3bf139e93ad757869338ad85239cb1d6c067b23b94e5846e637ca6328ee4be60" dependencies = [ "approx", "matrixmultiply", @@ -659,7 +659,7 @@ checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] @@ -723,17 +723,18 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.19.0" +version = "1.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" [[package]] name = "owo-colors" -version = "4.0.0" +version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "caff54706df99d2a78a5a4e3455ff45448d81ef1bb63c22cd14052ca0e993a3f" +checksum = "fb37767f6569cd834a413442455e0f066d0d522de8630436e2a1761d9726ba56" dependencies = [ - "supports-color", + "supports-color 2.1.0", + "supports-color 3.0.1", ] [[package]] @@ -760,9 +761,9 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "pest" -version = "2.7.11" +version = "2.7.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd53dff83f26735fdc1ca837098ccf133605d794cdae66acfc2bfac3ec809d95" +checksum = "879952a81a83930934cbf1786752d6dedc3b1f29e8f8fb2ad1d0a36f377cf442" dependencies = [ "memchr", "thiserror", @@ -781,8 +782,8 @@ dependencies = [ [[package]] name = "pgrx" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.12.6" +source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" dependencies = [ "atomic-traits", "bitflags", @@ -805,8 +806,8 @@ dependencies = [ [[package]] name = "pgrx-bindgen" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.12.6" +source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" dependencies = [ "bindgen", "cc", @@ -816,25 +817,25 @@ dependencies = [ "proc-macro2", "quote", "shlex", - "syn 2.0.77", + "syn 2.0.82", "walkdir", ] [[package]] name = "pgrx-macros" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.12.6" +source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] name = "pgrx-pg-config" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.12.6" +source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" dependencies = [ "cargo_toml", "eyre", @@ -850,8 +851,8 @@ dependencies = [ [[package]] name = "pgrx-pg-sys" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.12.6" +source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" dependencies = [ "cee-scape", "libc", @@ -864,15 +865,15 @@ dependencies = [ [[package]] name = "pgrx-sql-entity-graph" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.12.6" +source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" dependencies = [ "convert_case", "eyre", "petgraph", "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", "thiserror", "unescape", ] @@ -912,9 +913,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.86" +version = "1.0.88" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e719e8df665df0d1c8fbfd238015744736151d4445ec0836b8e628aae103b77" +checksum = "7c3a7fc5db1e57d5a779a352c8cdb57b29aa4c40cc69c3a68a7fedc815fbf2f9" dependencies = [ "unicode-ident", ] @@ -1062,9 +1063,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.10.6" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4219d74c6b67a3654a9fbebc4b419e22126d13d2f3c4a07ee0cb61ff79a79619" +checksum = "38200e5ee88914975b69f657f0801b6f6dccafd44fd9326302a4aaeecfacb1d8" dependencies = [ "aho-corasick", "memchr", @@ -1074,9 +1075,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.7" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38caf58cc5ef2fed281f89292ef23f6365465ed9a41b7a7754eb4e26496c92df" +checksum = "368758f23274712b504848e9d5a6f010445cc8b87a7cdb4d7cbee666c1288da3" dependencies = [ "aho-corasick", "memchr", @@ -1085,9 +1086,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.4" +version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a66a03ae7c801facd77a29370b4faec201768915ac14a721ba36f20bc9c209b" +checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" [[package]] name = "rend" @@ -1144,9 +1145,9 @@ dependencies = [ [[package]] name = "rustix" -version = "0.38.35" +version = "0.38.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a85d50532239da68e9addb745ba38ff4612a242c1c7ceea689c4bc7c2f43c36f" +checksum = "8acb788b847c24f28525660c4d7758620a7210875711f79e7f663cc152726811" dependencies = [ "bitflags", "errno", @@ -1211,9 +1212,9 @@ checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" [[package]] name = "serde" -version = "1.0.209" +version = "1.0.210" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99fce0ffe7310761ca6bf9faf5115afbc19688edd00171d81b1bb1b116c63e09" +checksum = "c8e3592472072e6e22e0a54d5904d9febf8508f65fb8552499a1abc7d1078c3a" dependencies = [ "serde_derive", ] @@ -1230,20 +1231,20 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.209" +version = "1.0.210" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5831b979fd7b5439637af1752d535ff49f4860c0f341d1baeb6faf0f4242170" +checksum = "243902eda00fad750862fc144cea25caca5e20d615af0a81bee94ca738f1df1f" dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] name = "serde_json" -version = "1.0.127" +version = "1.0.132" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8043c06d9f82bd7271361ed64f415fe5e12a77fdb52e573e7f06a516dea329ad" +checksum = "d726bfaff4b320266d395898905d0eba0345aae23b54aee3a737e260fd46db03" dependencies = [ "itoa", "memchr", @@ -1253,9 +1254,9 @@ dependencies = [ [[package]] name = "serde_spanned" -version = "0.6.7" +version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb5b1b31579f3811bf615c144393417496f152e12ac8b7663bf664f4a815306d" +checksum = "87607cb1398ed59d48732e575a4c28a7a8ebf2454b964fe3f224f2afc07909e1" dependencies = [ "serde", ] @@ -1281,9 +1282,9 @@ dependencies = [ [[package]] name = "simdutf8" -version = "0.1.4" +version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f27f6278552951f1f2b8cf9da965d10969b2efdea95a6ec47987ab46edfe263a" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" [[package]] name = "smawk" @@ -1319,6 +1320,15 @@ dependencies = [ "is_ci", ] +[[package]] +name = "supports-color" +version = "3.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8775305acf21c96926c900ad056abeef436701108518cf890020387236ac5a77" +dependencies = [ + "is_ci", +] + [[package]] name = "syn" version = "1.0.109" @@ -1332,9 +1342,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.77" +version = "2.0.82" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f35bcdf61fd8e7be6caf75f429fdca8beb3ed76584befb503b1569faee373ed" +checksum = "83540f837a8afc019423a8edb95b52a8effe46957ee402287f4292fae35be021" dependencies = [ "proc-macro2", "quote", @@ -1349,22 +1359,22 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" [[package]] name = "thiserror" -version = "1.0.63" +version = "1.0.64" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0342370b38b6a11b6cc11d6a805569958d54cfa061a29969c3b5ce2ea405724" +checksum = "d50af8abc119fb8bb6dbabcfa89656f46f84aa0ac7688088608076ad2b459a84" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.63" +version = "1.0.64" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4558b58466b9ad7ca0f102865eccc95938dca1a74a856f2b57b6629050da261" +checksum = "08904e7672f5eb876eaaf87e0ce17857500934f4981c4a0ab2b4aa98baac7fc3" dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] @@ -1405,9 +1415,9 @@ dependencies = [ [[package]] name = "toml_edit" -version = "0.22.20" +version = "0.22.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "583c44c02ad26b0c3f3066fe629275e50627026c51ac2e595cca4c230ce1ce1d" +checksum = "4ae48d6208a266e853d946088ed816055e556cc6028c5e8e2b84d9fa5dd7c7f5" dependencies = [ "indexmap", "serde", @@ -1424,9 +1434,9 @@ checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" [[package]] name = "ucd-trie" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed646292ffc8188ef8ea4d1e0e0150fb15a5c2e12ad9b8fc191ae7a8a7f3c4b9" +checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" [[package]] name = "unescape" @@ -1436,30 +1446,30 @@ checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] name = "unicode-bidi" -version = "0.3.15" +version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08f95100a766bf4f8f28f90d77e0a5461bbdb219042e7679bebe79004fed8d75" +checksum = "5ab17db44d7388991a428b2ee655ce0c212e862eff1768a455c58f9aad6e7893" [[package]] name = "unicode-ident" -version = "1.0.12" +version = "1.0.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" [[package]] name = "unicode-normalization" -version = "0.1.23" +version = "0.1.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a56d1686db2308d901306f92a263857ef59ea39678a5458e7cb17f01415101f5" +checksum = "5033c97c4262335cded6d6fc3e5c18ab755e1a3dc96376350f3d8e9f009ad956" dependencies = [ "tinyvec", ] [[package]] name = "unicode-segmentation" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4c87d22b6e3f4a18d4d40ef354e97c90fcb14dd91d7dc0aa9d8a1172ebf7202" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" [[package]] name = "unicode-width" @@ -1480,9 +1490,9 @@ dependencies = [ [[package]] name = "uuid" -version = "1.10.0" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81dfa00651efa65069b0b6b651f4aaa31ba9e3c3ce0137aaad053604ee7e0314" +checksum = "f8c5f0a0af699448548ad1a2fbf920fb4bee257eae39953ba95cb84891a0446a" dependencies = [ "getrandom", ] @@ -1514,7 +1524,7 @@ dependencies = [ "proc-macro-error", "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", ] [[package]] @@ -1664,9 +1674,9 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "winnow" -version = "0.6.18" +version = "0.6.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68a9bda4691f099d435ad181000724da8e5899daa10713c2d432552b9ccd3a6f" +checksum = "36c1fec1a2bb5866f07c25f68c26e565c4c200aebb96d7e55710c19d3e8ac49b" dependencies = [ "memchr", ] @@ -1707,5 +1717,5 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn 2.0.82", ] diff --git a/Cargo.toml b/Cargo.toml index ab88f57a..3cafdf11 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,7 +21,7 @@ pg16 = ["pgrx/pg16"] pg17 = ["pgrx/pg17"] [dependencies] -pgrx = { version = "=0.12.5", default-features = false, features = [] } +pgrx = { version = "=0.12.6", default-features = false, features = [] } base = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } common = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } detect = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } @@ -39,7 +39,7 @@ seq-macro = "0.3.5" validator = "0.18.1" [patch.crates-io] -pgrx = { git = "https://github.com/tensorchord/pgrx.git", branch = "v0.12.5-patch" } +pgrx = { git = "https://github.com/tensorchord/pgrx.git", branch = "v0.12.6-patch" } [lints] rust.unsafe_op_in_unsafe_fn = "deny" diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index 4c7450af..417babdd 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -20,21 +20,29 @@ pub trait HeapRelation { F: FnMut((Pointer, Vec)); } -pub fn build( +pub trait Reporter { + fn tuples_total(&mut self, tuples_total: usize); + fn tuples_done(&mut self, tuples_done: usize); +} + +pub fn build( vector_options: VectorOptions, rabbithole_options: RabbitholeIndexingOptions, heap_relation: T, index_relation: Relation, + mut reporter: R, ) { let dims = vector_options.dims; let is_residual = rabbithole_options.residual_quantization && vector_options.d == DistanceKind::L2; + let mut tuples_total = 0_usize; let samples = { let mut rand = rand::thread_rng(); let max_number_of_samples = rabbithole_options.nlist.saturating_mul(256); let mut samples = Vec::new(); let mut number_of_samples = 0_u32; heap_relation.traverse(|(_, vector)| { + pgrx::check_for_interrupts!(); assert_eq!(dims as usize, vector.len(), "invalid vector dimensions",); let vector = rabitq::project(&vector); if number_of_samples < max_number_of_samples { @@ -46,9 +54,11 @@ pub fn build( let end = start + dims as usize; samples[start..end].copy_from_slice(&vector); } + tuples_total += 1; }); Vec2::from_vec((number_of_samples as _, dims as _), samples) }; + reporter.tuples_total(tuples_total); let structure = Structure::compute(vector_options.clone(), rabbithole_options.clone(), samples); let h2_len = structure.h2_len(); let h1_len = structure.h1_len(); @@ -156,7 +166,11 @@ pub fn build( for i in 0..structure.h1_len() { heads.push(h1_firsts[i as usize]); } + let mut tuples_done = 0; heap_relation.traverse(|(payload, vector)| { + pgrx::check_for_interrupts!(); + tuples_done += 1; + reporter.tuples_done(tuples_done); assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); let vector = rabitq::project(&vector); let h0_vector = vectors.push(&VectorTuple { diff --git a/src/index/am.rs b/src/index/am.rs index 329a257f..37d38181 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -1,5 +1,5 @@ use crate::algorithm; -use crate::algorithm::build::HeapRelation; +use crate::algorithm::build::{HeapRelation, Reporter}; use crate::index::am_options::{Opfamily, Reloption}; use crate::index::am_scan::Scanner; use crate::index::utils::{ctid_to_pointer, pointer_to_ctid}; @@ -200,6 +200,25 @@ pub unsafe extern "C" fn ambuild( } } } + pub struct PgReporter {} + impl Reporter for PgReporter { + fn tuples_total(&mut self, tuples_total: usize) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, + tuples_total as _, + ); + } + } + fn tuples_done(&mut self, tuples_done: usize) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, + tuples_done as _, + ); + } + } + } let (vector_options, rabbithole_options) = unsafe { am_options::options(index) }; let heap_relation = Heap { heap, @@ -213,6 +232,7 @@ pub unsafe extern "C" fn ambuild( rabbithole_options, heap_relation, index_relation, + PgReporter {}, ); unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } From cd149aa6e0b20f8609d40dd99284b8b133a0dc85 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 21 Oct 2024 16:29:05 +0800 Subject: [PATCH 003/324] fix: set default value of rabbithole.prewarm_dim, forbid being loaded in session process, lock nalgebra version (#17) * fix: set rabbithole.prewarm_dim to 64,128,256,384,512,768,1024,1536 Signed-off-by: usamoi * fix: lock nalgebra version Signed-off-by: usamoi --------- Signed-off-by: usamoi --- Cargo.lock | 4 ++-- Cargo.toml | 3 ++- src/algorithm/rabitq.rs | 16 ++++++++++++++++ src/gucs/prewarm.rs | 3 ++- src/lib.rs | 3 +++ 5 files changed, 25 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 067ba488..10aa9443 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -637,9 +637,9 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" [[package]] name = "nalgebra" -version = "0.33.1" +version = "0.33.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3bf139e93ad757869338ad85239cb1d6c067b23b94e5846e637ca6328ee4be60" +checksum = "3c4b5f057b303842cf3262c27e465f4c303572e7f6b0648f60e16248ac3397f4" dependencies = [ "approx", "matrixmultiply", diff --git a/Cargo.toml b/Cargo.toml index 3cafdf11..b1397e74 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +31,8 @@ paste = "1" serde = "1" toml = "0.8.19" rand = "0.8.5" -nalgebra = "0.33.0" +nalgebra = { version = "=0.33.0", default-features = false } +# lock algebra version forever so that the QR decomposition never changes for same input rand_distr = "0.4.3" rkyv = { version = "0.7.45", features = ["validation"] } rand_chacha = "0.3.1" diff --git a/src/algorithm/rabitq.rs b/src/algorithm/rabitq.rs index a0522bef..e3c5b002 100644 --- a/src/algorithm/rabitq.rs +++ b/src/algorithm/rabitq.rs @@ -32,6 +32,22 @@ fn check_all_matrixs_are_full_rank() { }); } +#[test] +fn check_matrices() { + assert_eq!( + orthogonal_matrix(2), + vec![vec![-0.5424608, -0.8400813], vec![0.8400813, -0.54246056]] + ); + assert_eq!( + orthogonal_matrix(3), + vec![ + vec![-0.5309615, -0.69094884, -0.49058124], + vec![0.8222731, -0.56002235, -0.10120347], + vec![0.20481002, 0.45712686, -0.86549866] + ] + ); +} + fn orthogonal_matrix(n: usize) -> Vec> { use nalgebra::QR; let matrix = random_matrix(n); diff --git a/src/gucs/prewarm.rs b/src/gucs/prewarm.rs index 8cd7456f..0638b385 100644 --- a/src/gucs/prewarm.rs +++ b/src/gucs/prewarm.rs @@ -1,7 +1,8 @@ use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; use std::ffi::CStr; -static PREWARM_DIM: GucSetting> = GucSetting::>::new(None); +static PREWARM_DIM: GucSetting> = + GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); pub unsafe fn init() { GucRegistry::define_string_guc( diff --git a/src/lib.rs b/src/lib.rs index 3de782c0..347aee74 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,6 +18,9 @@ pgrx::extension_sql_file!("./sql/finalize.sql", finalize); #[pgrx::pg_guard] unsafe extern "C" fn _PG_init() { + if unsafe { pgrx::pg_sys::IsUnderPostmaster } { + pgrx::error!("rabbithole must be loaded via shared_preload_libraries."); + } detect::init(); unsafe { index::init(); From 2dfd7de614f73d0c6f32026182038787c368df25 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 21 Oct 2024 18:04:15 +0800 Subject: [PATCH 004/324] feat: pgvector (#18) * feat: pgvector Signed-off-by: usamoi * feat: pick pgvecto.rs features Signed-off-by: usamoi --------- Signed-off-by: usamoi --- rabbithole.control | 2 +- ...ry_vecf32.rs => memory_pgvector_vector.rs} | 97 ++++++++++--------- src/datatype/mod.rs | 3 +- src/datatype/operators_pgvector_vector.rs | 24 +++++ src/index/am_options.rs | 8 +- src/sql/bootstrap.sql | 3 + src/sql/finalize.sql | 21 ++++ 7 files changed, 104 insertions(+), 54 deletions(-) rename src/datatype/{memory_vecf32.rs => memory_pgvector_vector.rs} (65%) create mode 100644 src/datatype/operators_pgvector_vector.rs diff --git a/rabbithole.control b/rabbithole.control index 74b4cc1d..a0f760d2 100644 --- a/rabbithole.control +++ b/rabbithole.control @@ -4,4 +4,4 @@ module_pathname = '$libdir/rabbithole' relocatable = false superuser = true schema = rabbithole -requires = 'vectors' +requires = 'vector' diff --git a/src/datatype/memory_vecf32.rs b/src/datatype/memory_pgvector_vector.rs similarity index 65% rename from src/datatype/memory_vecf32.rs rename to src/datatype/memory_pgvector_vector.rs index 2802134c..c3ac76ab 100644 --- a/src/datatype/memory_vecf32.rs +++ b/src/datatype/memory_pgvector_vector.rs @@ -15,20 +15,20 @@ use std::ptr::NonNull; pub const HEADER_MAGIC: u16 = 0; #[repr(C, align(8))] -pub struct Vecf32Header { +pub struct PgvectorVectorHeader { varlena: u32, dims: u16, magic: u16, phantom: [f32; 0], } -impl Vecf32Header { +impl PgvectorVectorHeader { fn varlena(size: usize) -> u32 { (size << 2) as u32 } fn layout(len: usize) -> Layout { u16::try_from(len).expect("Vector is too large."); - let layout_alpha = Layout::new::(); + let layout_alpha = Layout::new::(); let layout_beta = Layout::array::(len).unwrap(); let layout = layout_alpha.extend(layout_beta).unwrap().0; layout.pad_to_align() @@ -45,7 +45,7 @@ impl Vecf32Header { } } -impl Deref for Vecf32Header { +impl Deref for PgvectorVectorHeader { type Target = [f32]; fn deref(&self) -> &Self::Target { @@ -53,69 +53,70 @@ impl Deref for Vecf32Header { } } -pub enum Vecf32Input<'a> { - Owned(Vecf32Output), - Borrowed(&'a Vecf32Header), +pub enum PgvectorVectorInput<'a> { + Owned(PgvectorVectorOutput), + Borrowed(&'a PgvectorVectorHeader), } -impl<'a> Vecf32Input<'a> { - unsafe fn new(p: NonNull) -> Self { +impl<'a> PgvectorVectorInput<'a> { + unsafe fn new(p: NonNull) -> Self { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() }; if p != q { - Vecf32Input::Owned(Vecf32Output(q)) + PgvectorVectorInput::Owned(PgvectorVectorOutput(q)) } else { - unsafe { Vecf32Input::Borrowed(p.as_ref()) } + unsafe { PgvectorVectorInput::Borrowed(p.as_ref()) } } } } -impl Deref for Vecf32Input<'_> { - type Target = Vecf32Header; +impl Deref for PgvectorVectorInput<'_> { + type Target = PgvectorVectorHeader; fn deref(&self) -> &Self::Target { match self { - Vecf32Input::Owned(x) => x, - Vecf32Input::Borrowed(x) => x, + PgvectorVectorInput::Owned(x) => x, + PgvectorVectorInput::Borrowed(x) => x, } } } -pub struct Vecf32Output(NonNull); +pub struct PgvectorVectorOutput(NonNull); -impl Vecf32Output { - pub fn new(vector: VectBorrowed<'_, f32>) -> Vecf32Output { +impl PgvectorVectorOutput { + pub fn new(vector: VectBorrowed<'_, f32>) -> PgvectorVectorOutput { unsafe { let slice = vector.slice(); - let layout = Vecf32Header::layout(slice.len()); + let layout = PgvectorVectorHeader::layout(slice.len()); let dims = vector.dims(); let internal_dims = dims as u16; - let ptr = pgrx::pg_sys::palloc(layout.size()) as *mut Vecf32Header; + let ptr = pgrx::pg_sys::palloc(layout.size()) as *mut PgvectorVectorHeader; ptr.cast::().add(layout.size() - 8).write_bytes(0, 8); - std::ptr::addr_of_mut!((*ptr).varlena).write(Vecf32Header::varlena(layout.size())); + std::ptr::addr_of_mut!((*ptr).varlena) + .write(PgvectorVectorHeader::varlena(layout.size())); std::ptr::addr_of_mut!((*ptr).magic).write(HEADER_MAGIC); std::ptr::addr_of_mut!((*ptr).dims).write(internal_dims); std::ptr::copy_nonoverlapping(slice.as_ptr(), (*ptr).phantom.as_mut_ptr(), slice.len()); - Vecf32Output(NonNull::new(ptr).unwrap()) + PgvectorVectorOutput(NonNull::new(ptr).unwrap()) } } - pub fn into_raw(self) -> *mut Vecf32Header { + pub fn into_raw(self) -> *mut PgvectorVectorHeader { let result = self.0.as_ptr(); std::mem::forget(self); result } } -impl Deref for Vecf32Output { - type Target = Vecf32Header; +impl Deref for PgvectorVectorOutput { + type Target = PgvectorVectorHeader; fn deref(&self) -> &Self::Target { unsafe { self.0.as_ref() } } } -impl Drop for Vecf32Output { +impl Drop for PgvectorVectorOutput { fn drop(&mut self) { unsafe { pgrx::pg_sys::pfree(self.0.as_ptr() as _); @@ -123,72 +124,72 @@ impl Drop for Vecf32Output { } } -impl<'a> FromDatum for Vecf32Input<'a> { +impl<'a> FromDatum for PgvectorVectorInput<'a> { unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { if is_null { None } else { - let ptr = NonNull::new(datum.cast_mut_ptr::()).unwrap(); - unsafe { Some(Vecf32Input::new(ptr)) } + let ptr = NonNull::new(datum.cast_mut_ptr::()).unwrap(); + unsafe { Some(PgvectorVectorInput::new(ptr)) } } } } -impl IntoDatum for Vecf32Output { +impl IntoDatum for PgvectorVectorOutput { fn into_datum(self) -> Option { Some(Datum::from(self.into_raw() as *mut ())) } fn type_oid() -> Oid { - let namespace = pgrx::pg_catalog::PgNamespace::search_namespacename(c"vectors").unwrap(); - let namespace = namespace.get().expect("pgvecto.rs is not installed."); - let t = pgrx::pg_catalog::PgType::search_typenamensp(c"vector", namespace.oid()).unwrap(); - let t = t.get().expect("pgvecto.rs is not installed."); - t.oid() + panic!("calling `type_oid` is never expected") + } + + fn is_compatible_with(_: Oid) -> bool { + true } } -impl FromDatum for Vecf32Output { +impl FromDatum for PgvectorVectorOutput { unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { if is_null { None } else { - let p = NonNull::new(datum.cast_mut_ptr::())?; + let p = NonNull::new(datum.cast_mut_ptr::())?; let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast())? }; if p != q { - Some(Vecf32Output(q)) + Some(PgvectorVectorOutput(q)) } else { let header = p.as_ptr(); let vector = unsafe { (*header).as_borrowed() }; - Some(Vecf32Output::new(vector)) + Some(PgvectorVectorOutput::new(vector)) } } } } -unsafe impl pgrx::datum::UnboxDatum for Vecf32Output { - type As<'src> = Vecf32Output; +unsafe impl pgrx::datum::UnboxDatum for PgvectorVectorOutput { + type As<'src> = PgvectorVectorOutput; #[inline] unsafe fn unbox<'src>(d: pgrx::datum::Datum<'src>) -> Self::As<'src> where Self: 'src, { - let p = NonNull::new(d.sans_lifetime().cast_mut_ptr::()).unwrap(); + let p = NonNull::new(d.sans_lifetime().cast_mut_ptr::()).unwrap(); let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() }; if p != q { - Vecf32Output(q) + PgvectorVectorOutput(q) } else { let header = p.as_ptr(); let vector = unsafe { (*header).as_borrowed() }; - Vecf32Output::new(vector) + PgvectorVectorOutput::new(vector) } } } -unsafe impl SqlTranslatable for Vecf32Input<'_> { +unsafe impl SqlTranslatable for PgvectorVectorInput<'_> { fn argument_sql() -> Result { Ok(SqlMapping::As(String::from("vector"))) } @@ -197,7 +198,7 @@ unsafe impl SqlTranslatable for Vecf32Input<'_> { } } -unsafe impl SqlTranslatable for Vecf32Output { +unsafe impl SqlTranslatable for PgvectorVectorOutput { fn argument_sql() -> Result { Ok(SqlMapping::As(String::from("vector"))) } @@ -206,13 +207,13 @@ unsafe impl SqlTranslatable for Vecf32Output { } } -unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Vecf32Input<'fcx> { +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for PgvectorVectorInput<'fcx> { unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { unsafe { arg.unbox_arg_using_from_datum().unwrap() } } } -unsafe impl pgrx::callconv::BoxRet for Vecf32Output { +unsafe impl pgrx::callconv::BoxRet for PgvectorVectorOutput { unsafe fn box_into<'fcx>( self, fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 527ad4aa..9152f33a 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -1,2 +1,3 @@ -pub mod memory_vecf32; +pub mod memory_pgvector_vector; +pub mod operators_pgvector_vector; pub mod typmod; diff --git a/src/datatype/operators_pgvector_vector.rs b/src/datatype/operators_pgvector_vector.rs new file mode 100644 index 00000000..0a77d24d --- /dev/null +++ b/src/datatype/operators_pgvector_vector.rs @@ -0,0 +1,24 @@ +use crate::datatype::memory_pgvector_vector::*; +use base::scalar::ScalarLike; +use std::num::NonZero; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _rabbithole_pgvector_vector_sphere_l2_in( + lhs: PgvectorVectorInput<'_>, + rhs: pgrx::composite_type!("sphere_vector"), +) -> bool { + let center: PgvectorVectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + f32::reduce_sum_of_d2(lhs.slice(), center.slice()).to_f32() < radius +} diff --git a/src/index/am_options.rs b/src/index/am_options.rs index a3486a73..1a977e16 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -1,5 +1,5 @@ -use crate::datatype::memory_vecf32::Vecf32Input; -use crate::datatype::memory_vecf32::Vecf32Output; +use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; +use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use crate::datatype::typmod::Typmod; use crate::types::RabbitholeIndexingOptions; use base::distance::*; @@ -157,7 +157,7 @@ impl Opfamily { } let vector = match self.vector { VectorKind::Vecf32 => { - let vector = unsafe { Vecf32Input::from_datum(datum, false).unwrap() }; + let vector = unsafe { PgvectorVectorInput::from_datum(datum, false).unwrap() }; self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed())) } _ => unreachable!(), @@ -175,7 +175,7 @@ impl Opfamily { let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; let center = match self.vector { VectorKind::Vecf32 => tuple - .get_by_index::(NonZero::new(1).unwrap()) + .get_by_index::(NonZero::new(1).unwrap()) .unwrap() .map(|vector| self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed()))), _ => unreachable!(), diff --git a/src/sql/bootstrap.sql b/src/sql/bootstrap.sql index e69de29b..6cc59f1a 100644 --- a/src/sql/bootstrap.sql +++ b/src/sql/bootstrap.sql @@ -0,0 +1,3 @@ +-- List of shell types + +CREATE TYPE sphere_vector; diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 8fdb4add..96dfaaa6 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -1,3 +1,24 @@ +-- List of data types + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <<->> ( + PROCEDURE = _rabbithole_pgvector_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + -- List of access methods CREATE ACCESS METHOD rabbithole TYPE INDEX HANDLER _rabbithole_amhandler; From 5e05e9f0f1e645981f7578721e323a77bb6f823f Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Tue, 22 Oct 2024 11:45:51 +0800 Subject: [PATCH 005/324] feat: external centroids (#14) * feat: external centroids Signed-off-by: cutecutecat * fix by comments Signed-off-by: cutecutecat * load returns Structure Signed-off-by: cutecutecat --------- Signed-off-by: cutecutecat --- src/algorithm/build.rs | 117 ++++++++++++++++++++++++++++++++--------- src/index/utils.rs | 29 ++++++++++ src/types.rs | 12 +++++ 3 files changed, 134 insertions(+), 24 deletions(-) diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index 417babdd..e1be9a58 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -1,7 +1,9 @@ use crate::algorithm::rabitq; use crate::algorithm::tuples::*; +use crate::index::utils::load_proj_vectors; use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; +use crate::types::ExternalCentroids; use crate::types::RabbitholeIndexingOptions; use base::distance::DistanceKind; use base::index::VectorOptions; @@ -35,31 +37,36 @@ pub fn build( let dims = vector_options.dims; let is_residual = rabbithole_options.residual_quantization && vector_options.d == DistanceKind::L2; - let mut tuples_total = 0_usize; - let samples = { - let mut rand = rand::thread_rng(); - let max_number_of_samples = rabbithole_options.nlist.saturating_mul(256); - let mut samples = Vec::new(); - let mut number_of_samples = 0_u32; - heap_relation.traverse(|(_, vector)| { - pgrx::check_for_interrupts!(); - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions",); - let vector = rabitq::project(&vector); - if number_of_samples < max_number_of_samples { - samples.extend(vector); - number_of_samples += 1; - } else { - let index = rand.gen_range(0..max_number_of_samples) as usize; - let start = index * dims as usize; - let end = start + dims as usize; - samples[start..end].copy_from_slice(&vector); - } - tuples_total += 1; - }); - Vec2::from_vec((number_of_samples as _, dims as _), samples) + let structure = match &rabbithole_options.external_centroids { + Some(_) => Structure::load(vector_options.clone(), rabbithole_options.clone()), + None => { + let mut tuples_total = 0_usize; + let samples = { + let mut rand = rand::thread_rng(); + let max_number_of_samples = rabbithole_options.nlist.saturating_mul(256); + let mut samples = Vec::new(); + let mut number_of_samples = 0_u32; + heap_relation.traverse(|(_, vector)| { + pgrx::check_for_interrupts!(); + assert_eq!(dims as usize, vector.len(), "invalid vector dimensions",); + let vector = rabitq::project(&vector); + if number_of_samples < max_number_of_samples { + samples.extend(vector); + number_of_samples += 1; + } else { + let index = rand.gen_range(0..max_number_of_samples) as usize; + let start = index * dims as usize; + let end = start + dims as usize; + samples[start..end].copy_from_slice(&vector); + } + tuples_total += 1; + }); + Vec2::from_vec((number_of_samples as _, dims as _), samples) + }; + reporter.tuples_total(tuples_total); + Structure::compute(vector_options.clone(), rabbithole_options.clone(), samples) + } }; - reporter.tuples_total(tuples_total); - let structure = Structure::compute(vector_options.clone(), rabbithole_options.clone(), samples); let h2_len = structure.h2_len(); let h1_len = structure.h1_len(); let mut meta = Tape::create(&index_relation); @@ -303,6 +310,68 @@ impl Structure { h1_children: (0..rabbithole_options.nlist).map(|_| Vec::new()).collect(), } } + fn load(vector_options: VectorOptions, rabbithole_options: RabbitholeIndexingOptions) -> Self { + let dims = vector_options.dims; + let h1_means = match &rabbithole_options.external_centroids { + Some(ExternalCentroids { + table, + h1_means_column: h1, + .. + }) => load_proj_vectors(table, h1, rabbithole_options.nlist, vector_options.dims), + _ => unreachable!(), + }; + let h1_children = match &rabbithole_options.external_centroids { + Some(ExternalCentroids { + table, + h1_children_column: Some(h1), + .. + }) => load_proj_vectors(table, h1, 1, vector_options.dims) + .into_iter() + .map(|v| v.into_iter().map(|f| f as u32).collect()) + .collect(), + _ => (0..rabbithole_options.nlist).map(|_| Vec::new()).collect(), + }; + let h2_mean = match &rabbithole_options.external_centroids { + Some(ExternalCentroids { + table, + h2_mean_column: Some(h2), + .. + }) => load_proj_vectors(table, h2, 1, vector_options.dims) + .pop() + .expect("load h2_mean panic"), + _ => { + let mut centroid = vec![0.0; dims as _]; + for i in 0..rabbithole_options.nlist { + for j in 0..dims { + centroid[j as usize] += h1_means[i as usize][j as usize]; + } + } + for j in 0..dims { + centroid[j as usize] /= rabbithole_options.nlist as f32; + } + centroid + } + }; + let h2_children = match &rabbithole_options.external_centroids { + Some(ExternalCentroids { + table, + h2_children_column: Some(h2), + .. + }) => load_proj_vectors(table, h2, 1, vector_options.dims) + .pop() + .expect("load h2_children panic") + .into_iter() + .map(|f| f as u32) + .collect(), + _ => (0..rabbithole_options.nlist).collect(), + }; + Structure { + h2_mean, + h2_children, + h1_means, + h1_children, + } + } fn h2_len(&self) -> u32 { 1 } diff --git a/src/index/utils.rs b/src/index/utils.rs index a5d85a3f..d6bde6a9 100644 --- a/src/index/utils.rs +++ b/src/index/utils.rs @@ -1,4 +1,10 @@ use base::search::*; +use base::vector::VectorBorrowed; +use pgrx::pg_sys::panic::ErrorReportable; +use pgrx::{error, Spi}; + +use crate::algorithm::rabitq; +use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { let value = pointer.as_u64(); @@ -18,3 +24,26 @@ pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer { value |= ctid.ip_posid as u64; Pointer::new(value) } + +pub fn load_proj_vectors(table_name: &str, column_name: &str, rows: u32, dims: u32) -> Vec> { + let query = format!("SELECT {column_name} FROM {table_name};"); + let mut centroids = Vec::new(); + + Spi::connect(|client| { + let tup_table = client.select(&query, None, None).unwrap_or_report(); + assert_eq!(tup_table.len(), rows as usize); + + for row in tup_table { + let vector = row[column_name].value::(); + if let Ok(Some(v)) = vector { + let borrowed = v.as_borrowed(); + assert_eq!(borrowed.dims(), dims); + let projected_centroids = rabitq::project(borrowed.slice()); + centroids.push(projected_centroids); + } else { + error!("load vectors from column is not valid") + } + } + centroids + }) +} diff --git a/src/types.rs b/src/types.rs index 788504ab..b60fa149 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1,6 +1,16 @@ use serde::{Deserialize, Serialize}; use validator::Validate; +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ExternalCentroids { + pub table: String, + pub h1_means_column: String, + pub h1_children_column: Option, + pub h2_mean_column: Option, + pub h2_children_column: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] pub struct RabbitholeIndexingOptions { @@ -14,6 +24,7 @@ pub struct RabbitholeIndexingOptions { #[serde(default = "RabbitholeIndexingOptions::default_build_threads")] #[validate(range(min = 1, max = 255))] pub build_threads: u16, + pub external_centroids: Option, } impl RabbitholeIndexingOptions { @@ -38,6 +49,7 @@ impl Default for RabbitholeIndexingOptions { spherical_centroids: Self::default_spherical_centroids(), residual_quantization: Self::default_residual_quantization(), build_threads: Self::default_build_threads(), + external_centroids: None, } } } From fb0cf73eeb05ce05d3770ae09dd917a695363b49 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Tue, 22 Oct 2024 17:39:33 +0800 Subject: [PATCH 006/324] bench script (#19) Signed-off-by: cutecutecat --- bench/README.md | 68 ++++++++++++++++++++++++ bench/bench.py | 69 ++++++++++++++++++++++++ bench/internal.py | 95 +++++++++++++++++++++++++++++++++ bench/load.py | 114 ++++++++++++++++++++++++++++++++++++++++ bench/train.py | 61 +++++++++++++++++++++ docker/Dockerfile | 6 +++ tools/package.sh | 44 ++++++++++++++++ tools/schema-codegen.sh | 43 +++++++++++++++ tools/schema.sh | 47 +++++++++++++++++ 9 files changed, 547 insertions(+) create mode 100644 bench/README.md create mode 100644 bench/bench.py create mode 100644 bench/internal.py create mode 100644 bench/load.py create mode 100644 bench/train.py create mode 100644 docker/Dockerfile create mode 100755 tools/package.sh create mode 100755 tools/schema-codegen.sh create mode 100755 tools/schema.sh diff --git a/bench/README.md b/bench/README.md new file mode 100644 index 00000000..f36106f1 --- /dev/null +++ b/bench/README.md @@ -0,0 +1,68 @@ +## Build Docker + +```shell +cargo build --package rabbithole --lib --features pg16 --target x86_64-unknown-linux-gnu --release +./tools/schema.sh --features pg16 --target x86_64-unknown-linux-gnu --release | expand -t 4 > ./target/schema.sql + +export SEMVER="0.0.0" +export VERSION="16" +export ARCH="x86_64" +export PLATFORM="amd64" +./tools/package.sh + +docker build -t rabbithole:pg16-latest -f ./docker/Dockerfile . +``` + +Or you can use `starkind/rabbithole:pg16-latest` to run the bench. + +## Run Instance + +```shell +docker run --name rabbithole -e POSTGRES_PASSWORD=123 -p 5432:5432 -d starkind/rabbithole:pg16-latest + +PGPASSWORD=123 psql -h 127.0.0.1 -U postgres -c "CREATE USER bench WITH PASSWORD '123';" +PGPASSWORD=123 psql -h 127.0.0.1 -U postgres -c "ALTER ROLE bench SUPERUSER;" +``` + +## Download Data + +```shell +export AWS_ACCESS_KEY_ID=xxx +export AWS_SECRET_ACCESS_KEY=xxx + +aws s3 cp s3://pgvecto.rs-bench/sift_128_1m/sift.hdf5 $(whoami)/sift/sift.hdf5 +aws s3 cp s3://pgvecto.rs-bench/gist_960_1m/gist.hdf5 $(whoami)/gist/gist.hdf5 +aws s3 cp s3://pgvecto.rs-bench/glove_200_1m/glove.hdf5 $(whoami)/glove/glove.hdf5 +aws s3 cp s3://pgvecto.rs-bench/openai_1536_500k/openai.hdf5 $(whoami)/openai/openai.hdf5 +aws s3 cp s3://pgvecto.rs-bench/cohere-768-1m-2022/cohere.hdf5 $(whoami)/cohere_1m_22/cohere_1m_22.hdf5 +aws s3 cp s3://pgvecto.rs-bench/cohere-1024-1m-2023/cohere-1m-23.hdf5 $(whoami)/cohere_1m_23/cohere_1m_23.hdf5 +aws s3 cp s3://pgvecto.rs-bench/cohere-1024-10m-2023/cohere-10m-23.hdf5 $(whoami)/cohere_10m_23/cohere_10m_23.hdf5 +``` + +## Run Bench + +Options for `-n`: +- sift +- glove +- gist +- openai +- cohere_1m_22 +- cohere_1m_23 +- cohere_10m_23 + +```shell +# pip install pgvector numpy faiss-cpu psycopg h5py + +# If you want to use internal k-means(slow) +python internal.py -n sift -m l2 + +# Or external k-means +## K-means generate centroids +python train.py -n sift + +## Load centroids into PG and build index +python load.py -n sift -m l2 + +# Run bench +python bench.py -n sift -m l2 +``` \ No newline at end of file diff --git a/bench/bench.py b/bench/bench.py new file mode 100644 index 00000000..991f7b26 --- /dev/null +++ b/bench/bench.py @@ -0,0 +1,69 @@ +from os.path import join +import os +import time +import argparse +from pathlib import Path +from tqdm import tqdm + +import psycopg +import h5py +from pgvector.psycopg import register_vector +import numpy as np + +parser = argparse.ArgumentParser() +parser.add_argument( + "-m", "--metric", help="Metric to pick, in l2 or cos", required=True +) +parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) +args = parser.parse_args() + +HOME = Path.home() +INDEX_PATH = join(HOME, f"indexes/pg/{args.name}/{args.metric}") +DATA_PATH = join(HOME, f"{args.name}/{args.name}.hdf5") + +os.makedirs(join(HOME, f"indexes/pg/{args.name}"), exist_ok=True) +dataset = h5py.File(DATA_PATH, "r") + +test = dataset["test"][:] + +if args.metric == "l2": + metric_ops = "<->" +elif args.metric == "cos": + metric_ops = "<=>" +else: + raise ValueError + +answer = dataset["neighbors"][:] +n, dims = np.shape(test) +m = np.shape(test)[0] + +# reconnect for updated GUC variables to take effect +conn = psycopg.connect( + conninfo="postgres://bench:123@localhost:5432/postgres", + dbname="postgres", + autocommit=True, +) +conn.execute("SET search_path TO public, vectors, rabbithole") +conn.execute("CREATE EXTENSION IF NOT EXISTS vector") +conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") +conn.execute("SET rabbithole.nprobe=300") +register_vector(conn) + +Ks = [10, 100] +for k in Ks: + hits = 0 + delta = 0 + pbar = tqdm(enumerate(test), total=m) + for i, query in pbar: + start = time.perf_counter() + result = conn.execute( + f"SELECT id FROM {args.name} ORDER BY embedding {metric_ops} %s LIMIT {k}", + (query,), + ).fetchall() + end = time.perf_counter() + hits += len(set([p[0] for p in result[:k]]) & set(answer[i][:k].tolist())) + delta += end - start + pbar.set_description(f"recall: {hits / k / (i+1)} QPS: {(i+1) / delta} ") + recall = hits / k / m + qps = m / delta + print(f"Top: {k} recall: {recall:.4f} QPS: {qps:.2f}") \ No newline at end of file diff --git a/bench/internal.py b/bench/internal.py new file mode 100644 index 00000000..dd5806f1 --- /dev/null +++ b/bench/internal.py @@ -0,0 +1,95 @@ +import os +import time +import argparse +from pathlib import Path +import pickle + +import psycopg +import h5py +from pgvecto_rs.psycopg import register_vector +import numpy as np + +parser = argparse.ArgumentParser() +parser.add_argument( + "-m", "--metric", help="Metric to pick, in l2 or cos", required=True +) +parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) +parser.add_argument("-k", help="K-means centroids or nlist, default: 4096", required=False, default=4096) +args = parser.parse_args() + +HOME = Path.home() +DATA_PATH = join(HOME, f"{args.name}/{args.name}.hdf5") + +os.makedirs(join(HOME, f"indexes/pg/{args.name}"), exist_ok=True) +dataset = h5py.File(DATA_PATH, "r") + +train = dataset["train"][:] +test = dataset["test"][:] + +K = 4096 +if args.metric == "l2": + metric_ops = "vector_l2_ops" + ivf_config = f""" +nlist = {K} +residual_quantization = true +spherical_centroids = false +""" +elif args.metric == "cos": + metric_ops = "vector_cos_ops" + ivf_config = f""" +nlist = {K} +residual_quantization = false +spherical_centroids = true +""" +else: + raise ValueError + +answer = dataset["neighbors"][:] +n, dims = np.shape(train) +m = np.shape(test)[0] + +keepalive_kwargs = { + "keepalives": 1, + "keepalives_idle": 30, + "keepalives_interval": 5, + "keepalives_count": 5, +} + +start = time.perf_counter() + +with open(f"{args.name}.pickle", "rb") as f: + centroids = pickle.load(f) + +conn = psycopg.connect( + conninfo="postgres://bench:123@localhost:5432/postgres", + dbname="postgres", + autocommit=True, + **keepalive_kwargs, +) +conn.execute("SET search_path TO public, vectors, rabbithole") +conn.execute("CREATE EXTENSION IF NOT EXISTS vector") +conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") +register_vector(conn) + +conn.execute(f"DROP TABLE IF EXISTS public.centroids") +conn.execute(f"DROP TABLE IF EXISTS {args.name}") +conn.execute(f"CREATE TABLE {args.name} (id integer, embedding vector({dims}))") + +with conn.cursor().copy( + f"COPY {args.name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)" +) as copy: + copy.set_types(["integer", "vector"]) + + for i in range(n): + copy.write_row([i, train[i]]) + while conn.pgconn.flush() == 1: + pass + +start = time.perf_counter() +conn.execute( + f"CREATE INDEX ON {args.name} USING rabbithole (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" +) +end = time.perf_counter() + +delta = end - start +print(f"Index build time: {delta:.2f}s") \ No newline at end of file diff --git a/bench/load.py b/bench/load.py new file mode 100644 index 00000000..6b98b004 --- /dev/null +++ b/bench/load.py @@ -0,0 +1,114 @@ +from os.path import join +import os +import time +import argparse +from pathlib import Path +import pickle + +import psycopg +import h5py +from pgvector.psycopg import register_vector +import numpy as np + +parser = argparse.ArgumentParser() +parser.add_argument( + "-m", "--metric", help="Metric to pick, in l2 or cos", required=True +) +parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) +parser.add_argument("-k", help="K-means centroids or nlist, default: 4096", required=False, default=4096) +args = parser.parse_args() + +HOME = Path.home() +DATA_PATH = join(HOME, f"{args.name}/{args.name}.hdf5") + +os.makedirs(join(HOME, f"indexes/pg/{args.name}"), exist_ok=True) +dataset = h5py.File(DATA_PATH, "r") + +train = dataset["train"][:] +test = dataset["test"][:] + +K = args.k +if args.metric == "l2": + metric_ops = "vector_l2_ops" + ivf_config = f""" +nlist = {K} +residual_quantization = true +spherical_centroids = false +[external_centroids] +table = 'public.centroids' +h1_means_column = 'coordinate' +""" +elif args.metric == "cos": + metric_ops = "vector_cos_ops" + ivf_config = f""" +nlist = {K} +residual_quantization = false +spherical_centroids = true +[external_centroids] +table = 'public.centroids' +h1_means_column = 'coordinate' +""" +else: + raise ValueError + +answer = dataset["neighbors"][:] +n, dims = np.shape(train) +m = np.shape(test)[0] + +keepalive_kwargs = { + "keepalives": 1, + "keepalives_idle": 30, + "keepalives_interval": 5, + "keepalives_count": 5, +} + +start = time.perf_counter() + +with open(f"{args.name}.pickle", "rb") as f: + centroids = pickle.load(f) + +conn = psycopg.connect( + conninfo="postgres://bench:123@localhost:5432/postgres", + dbname="postgres", + autocommit=True, + **keepalive_kwargs, +) +conn.execute("SET search_path TO public, vectors, rabbithole") +conn.execute("CREATE EXTENSION IF NOT EXISTS vector") +conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") +register_vector(conn) + +conn.execute(f"DROP TABLE IF EXISTS public.centroids") +conn.execute(f"CREATE TABLE public.centroids (coordinate vector({dims}))") + +with conn.cursor().copy( + f"COPY public.centroids (coordinate) FROM STDIN WITH (FORMAT BINARY)" +) as copy: + copy.set_types(["vector"]) + + for i in range(K): + copy.write_row([centroids[i]]) + while conn.pgconn.flush() == 1: + pass + +conn.execute(f"DROP TABLE IF EXISTS {args.name}") +conn.execute(f"CREATE TABLE {args.name} (id integer, embedding vector({dims}))") + +with conn.cursor().copy( + f"COPY {args.name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)" +) as copy: + copy.set_types(["integer", "vector"]) + + for i in range(n): + copy.write_row([i, train[i]]) + while conn.pgconn.flush() == 1: + pass + +start = time.perf_counter() +conn.execute( + f"CREATE INDEX ON {args.name} USING rabbithole (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" +) +end = time.perf_counter() + +delta = end - start +print(f"Index build time: {delta:.2f}s") \ No newline at end of file diff --git a/bench/train.py b/bench/train.py new file mode 100644 index 00000000..802ae82b --- /dev/null +++ b/bench/train.py @@ -0,0 +1,61 @@ +from os.path import join +import time +import argparse +from pathlib import Path +import pickle + +import h5py +import faiss +import numpy as np + +K = 4096 +SEED = 42 + +parser = argparse.ArgumentParser() +parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) +parser.add_argument("-k", help="K-means centroids or nlist, default: 4096", required=False, default=4096) +args = parser.parse_args() + +K = args.k + +HOME = Path.home() +DATA_PATH = join(HOME, f"{args.name}/{args.name}.hdf5") + +dataset = h5py.File(DATA_PATH, "r") + +if len(dataset["train"]) > 256 * K: + rs = np.random.RandomState(SEED) + idx = rs.choice(len(dataset["train"]), size=256 * K, replace=False) + train = dataset["train"][np.sort(idx)] +else: + train = dataset["train"][:] + +test = dataset["test"][:] + +if np.shape(train)[0] > 256 * K: + rs = np.random.RandomState(SEED) + idx = rs.choice(np.shape(train)[0], size=256 * K, replace=False) + train = train[idx] + +answer = dataset["neighbors"][:] +n, dims = np.shape(train) +m = np.shape(test)[0] + +start = time.perf_counter() + +index = faiss.IndexFlatL2(dims) +clustering = faiss.Clustering(dims, K) +clustering.verbose = True +clustering.seed = 42 +clustering.niter = 10 +clustering.train(train, index) +centroids = faiss.vector_float_to_array(clustering.centroids) + +end = time.perf_counter() +delta = end - start +print(f"K-means time: {delta:.2f}s") + +centroids = centroids.reshape([K, -1]) + +with open(f"{args.name}.pickle", "wb") as f: + pickle.dump(centroids, f) \ No newline at end of file diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 00000000..0cbfeae5 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,6 @@ +FROM pgvector/pgvector:0.7.4-pg16 + +COPY ./build/rabbithole-pg16_0.0.0_amd64.deb /tmp/rabbithole.deb +RUN apt-get install -y /tmp/rabbithole.deb && rm -f /tmp/rabbithole.deb + +CMD ["postgres", "-c" ,"shared_preload_libraries=vector,rabbithole.so", "-c", "search_path=\"$user\", public, rabbithole", "-c", "logging_collector=on"] \ No newline at end of file diff --git a/tools/package.sh b/tools/package.sh new file mode 100755 index 00000000..68ac3a32 --- /dev/null +++ b/tools/package.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -e + +printf "SEMVER = ${SEMVER}\n" +printf "VERSION = ${VERSION}\n" +printf "ARCH = ${ARCH}\n" +printf "PLATFORM = ${PLATFORM}\n" + +rm -rf ./build/dir_zip +rm -rf ./build/rabbithole-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip +rm -rf ./build/dir_deb +rm -rf ./build/rabbithole-pg${VERSION}_${SEMVER}_${PLATFORM}.deb + +mkdir -p ./build/dir_zip +cp ./target/schema.sql ./build/dir_zip/rabbithole--$SEMVER.sql +sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./rabbithole.control > ./build/dir_zip/rabbithole.control +cp ./target/${ARCH}-unknown-linux-gnu/release/librabbithole.so ./build/dir_zip/rabbithole.so +zip ./build/rabbithole-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip -j ./build/dir_zip/* + +mkdir -p ./build/dir_deb +mkdir -p ./build/dir_deb/DEBIAN +mkdir -p ./build/dir_deb/usr/share/postgresql/$VERSION/extension/ +mkdir -p ./build/dir_deb/usr/lib/postgresql/$VERSION/lib/ +for file in $(ls ./build/dir_zip/*.sql | xargs -n 1 basename); do + cp ./build/dir_zip/$file ./build/dir_deb/usr/share/postgresql/$VERSION/extension/$file +done +for file in $(ls ./build/dir_zip/*.control | xargs -n 1 basename); do + cp ./build/dir_zip/$file ./build/dir_deb/usr/share/postgresql/$VERSION/extension/$file +done +for file in $(ls ./build/dir_zip/*.so | xargs -n 1 basename); do + cp ./build/dir_zip/$file ./build/dir_deb/usr/lib/postgresql/$VERSION/lib/$file +done +echo "Package: rabbithole-pg${VERSION} +Version: ${SEMVER} +Section: database +Priority: optional +Architecture: ${PLATFORM} +Maintainer: Tensorchord +Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM +Homepage: https://pgvecto.rs/ +License: apache2" \ +> ./build/dir_deb/DEBIAN/control +(cd ./build/dir_deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/dir_deb/DEBIAN/md5sums +dpkg-deb -Zxz --build ./build/dir_deb/ ./build/rabbithole-pg${VERSION}_${SEMVER}_${PLATFORM}.deb \ No newline at end of file diff --git a/tools/schema-codegen.sh b/tools/schema-codegen.sh new file mode 100755 index 00000000..63a7e272 --- /dev/null +++ b/tools/schema-codegen.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +set -e + +# CONTROL_FILEPATH +# SO_FILEPATH + +printf "fn main() {\n" + +cat << EOF + rabbithole::__pgrx_marker(); + + let mut entities = Vec::new(); + let control_file_path = std::path::PathBuf::from("$CONTROL_FILEPATH"); + let control_file = ::pgrx::pgrx_sql_entity_graph::ControlFile::try_from(control_file_path).expect(".control file should properly formatted"); + let control_file_entity = ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity::ExtensionRoot(control_file); + + entities.push(control_file_entity); +EOF + +while read -r name_ident; do +cat << EOF + extern "Rust" { + fn $name_ident() -> ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity; + } + let entity = unsafe { $name_ident() }; + entities.push(entity); +EOF +done <<< $(nm -D -g $SO_FILEPATH | grep "T __pgrx_internals_" | awk '{print $3}') + +cat << EOF + let pgrx_sql = ::pgrx::pgrx_sql_entity_graph::PgrxSql::build( + entities.into_iter(), + "vectors".to_string(), + false, + ) + .expect("SQL generation error"); + eprintln!("SQL entities to {}", "/dev/stdout",); + pgrx_sql + .write(&mut std::io::stdout()) + .expect("Could not write SQL to stdout"); +EOF + +printf "}\n" \ No newline at end of file diff --git a/tools/schema.sh b/tools/schema.sh new file mode 100755 index 00000000..1d868ad6 --- /dev/null +++ b/tools/schema.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +set -e +if [[ " $@ " =~ --target' '([^ ]+) ]]; then + TARGET="${BASH_REMATCH[1]}" + if [[ " $@ " =~ " --release " ]]; then + DIR="./target/$TARGET/release" + elif [[ " $@ " =~ " --profile opt " ]]; then + DIR="./target/$TARGET/opt" + else + DIR="./target/$TARGET/debug" + fi +else + if [[ " $@ " =~ " --release " ]]; then + DIR="./target/release" + elif [[ " $@ " =~ " --profile opt " ]]; then + DIR="./target/opt" + else + DIR="./target/debug" + fi +fi + +if [ "$TARGET" = "" ]; then + printf "Target: [not specified]\n" 1>&2 + RUNNER=() +elif [ "$TARGET" = $(rustc -vV | awk '/^host/ { print $2 }') ]; then + printf "Target: [host]\n" 1>&2 + RUNNER=() +elif [ "$TARGET" = "aarch64-unknown-linux-gnu" ]; then + printf "Target: $TARGET\n" 1>&2 + QEMU_LD_PREFIX="/usr/aarch64-linux-gnu" + RUNNER=("qemu-aarch64-static") +elif [ "$TARGET" = "riscv64gc-unknown-linux-gnu" ]; then + printf "Target: $TARGET\n" 1>&2 + QEMU_LD_PREFIX="/usr/riscv64-linux-gnu" + RUNNER=("qemu-riscv64-static") +else + printf "Unknown target: $TARGET\n" 1>&2 + exit 1 +fi + +code=$(mktemp) +chmod 700 $code +CONTROL_FILEPATH="./rabbithole.control" SO_FILEPATH="$DIR/librabbithole.so" $(dirname "$0")/schema-codegen.sh >> $code + +PGRX_EMBED=$code cargo rustc --package rabbithole --bin pgrx_embed_rabbithole "$@" -- --cfg pgrx_embed + +CARGO_PKG_VERSION="0.0.0" QEMU_LD_PREFIX=$QEMU_LD_PREFIX "${RUNNER[@]}" "$DIR/pgrx_embed_rabbithole" | expand -t 4 \ No newline at end of file From 86e8c7580c784855696a6362b3d450db80db1d0f Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 23 Oct 2024 17:03:20 +0800 Subject: [PATCH 007/324] feat: inline filtering (#20) Signed-off-by: usamoi --- src/algorithm/build.rs | 10 ++++-- src/algorithm/insert.rs | 6 +++- src/algorithm/scan.rs | 8 ++++- src/algorithm/tuples.rs | 3 ++ src/index/am.rs | 71 +++++++++++++++++++++++++++++++++-------- src/index/am_options.rs | 4 +-- src/index/am_scan.rs | 10 ++++-- src/sql/finalize.sql | 5 +++ 8 files changed, 95 insertions(+), 22 deletions(-) diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index e1be9a58..c7b67187 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -19,7 +19,7 @@ use std::sync::Arc; pub trait HeapRelation { fn traverse(&self, callback: F) where - F: FnMut((Pointer, Vec)); + F: FnMut((Pointer, Option, Vec)); } pub trait Reporter { @@ -46,7 +46,7 @@ pub fn build( let max_number_of_samples = rabbithole_options.nlist.saturating_mul(256); let mut samples = Vec::new(); let mut number_of_samples = 0_u32; - heap_relation.traverse(|(_, vector)| { + heap_relation.traverse(|(_, _, vector)| { pgrx::check_for_interrupts!(); assert_eq!(dims as usize, vector.len(), "invalid vector dimensions",); let vector = rabitq::project(&vector); @@ -174,7 +174,7 @@ pub fn build( heads.push(h1_firsts[i as usize]); } let mut tuples_done = 0; - heap_relation.traverse(|(payload, vector)| { + heap_relation.traverse(|(payload, extra, vector)| { pgrx::check_for_interrupts!(); tuples_done += 1; reporter.tuples_done(tuples_done); @@ -210,6 +210,7 @@ pub fn build( &code, h0_vector, h0_payload, + extra, ); if flag { return; @@ -224,6 +225,7 @@ pub fn build( factor_ip: [0.0; 32], factor_err: [0.0; 32], t: vec![0; (dims.div_ceil(4) * 16) as usize], + extra: [None; 32], }) .unwrap(); if let Some(i) = page.alloc(&tuple) { @@ -233,6 +235,7 @@ pub fn build( &code, h0_vector, h0_payload, + extra, ); assert!(flag, "a put fails even on a fresh tuple"); return; @@ -248,6 +251,7 @@ pub fn build( &code, h0_vector, h0_payload, + extra, ); assert!(flag, "a put fails even on a fresh tuple"); return; diff --git a/src/algorithm/insert.rs b/src/algorithm/insert.rs index d40da899..ebd2abbc 100644 --- a/src/algorithm/insert.rs +++ b/src/algorithm/insert.rs @@ -5,7 +5,7 @@ use base::distance::Distance; use base::scalar::ScalarLike; use base::search::Pointer; -pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { +pub fn insert(relation: Relation, payload: Pointer, extra: Option, vector: Vec) { let meta_guard = relation.read(0); let meta_tuple = meta_guard .get() @@ -111,6 +111,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { mask: [false; 32], mean: [(0, 0); 32], payload: [0; 32], + extra: [None; 32], dis_u_2: [0.0f32; 32], factor_ppc: [0.0f32; 32], factor_ip: [0.0f32; 32], @@ -153,6 +154,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { &code, h0_vector, h0_payload, + extra, ); if flag { return; @@ -165,6 +167,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { &code, h0_vector, h0_payload, + extra, ); assert!(flag, "a put fails even on a fresh tuple"); return; @@ -179,6 +182,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { &code, h0_vector, h0_payload, + extra, ); assert!(flag, "a put fails even on a fresh tuple"); return; diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs index fc1d3718..d39de635 100644 --- a/src/algorithm/scan.rs +++ b/src/algorithm/scan.rs @@ -12,6 +12,7 @@ pub fn scan( relation: Relation, vector: Vec, h1_nprobe: u32, + filters: Vec, ) -> impl Iterator { assert!(h1_nprobe >= 1); let meta_guard = relation.read(0); @@ -147,7 +148,12 @@ pub fn scan( ), ); for j in 0..32 { - if h0_tuple.mask[j] { + if h0_tuple.mask[j] + && filters + .iter() + .copied() + .all(|x| Some(x) == h0_tuple.extra[j]) + { results.push(( Reverse(lowerbounds[j]), AlwaysEqual(h0_tuple.mean[j]), diff --git a/src/algorithm/tuples.rs b/src/algorithm/tuples.rs index aaf3380c..9f6cdabb 100644 --- a/src/algorithm/tuples.rs +++ b/src/algorithm/tuples.rs @@ -45,6 +45,7 @@ pub struct Height0Tuple { pub mean: [(u32, u16); 32], // for height 0 tuple, it's pointers to heap relation pub payload: [u64; 32], + pub extra: [Option; 32], // RaBitQ algoithm pub dis_u_2: [f32; 32], pub factor_ppc: [f32; 32], @@ -59,6 +60,7 @@ pub fn put( code: &rabitq::Code, vector: (u32, u16), payload: u64, + extra: Option, ) -> bool { // todo: use mutable api let mut x = rkyv::from_bytes::(bytes).expect("data corruption"); @@ -66,6 +68,7 @@ pub fn put( if !x.mask[j] { x.mean[j] = vector; x.payload[j] = payload; + x.extra[j] = extra; x.mask[j] = true; x.dis_u_2[j] = code.dis_u_2; x.factor_ppc[j] = code.factor_ppc; diff --git a/src/index/am.rs b/src/index/am.rs index 37d38181..41c38cba 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -54,6 +54,7 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = { // for vector index scans without `ORDER BY` clauses a large number // and throw errors if someone really wants such a path. am_routine.amoptionalkey = true; + am_routine.amcanmulticol = true; am_routine.amvalidate = Some(amvalidate); am_routine.amoptions = Some(amoptions); @@ -141,7 +142,7 @@ pub unsafe extern "C" fn ambuild( impl HeapRelation for Heap { fn traverse(&self, callback: F) where - F: FnMut((Pointer, Vec)), + F: FnMut((Pointer, Option, Vec)), { pub struct State<'a, F> { pub this: &'a Heap, @@ -149,14 +150,14 @@ pub unsafe extern "C" fn ambuild( } #[pgrx::pg_guard] unsafe extern "C" fn call( - _index: pgrx::pg_sys::Relation, + index: pgrx::pg_sys::Relation, ctid: pgrx::pg_sys::ItemPointer, values: *mut Datum, is_null: *mut bool, _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((Pointer, Vec)), + F: FnMut((Pointer, Option, Vec)), { pgrx::check_for_interrupts!(); use base::vector::OwnedVector; @@ -167,15 +168,28 @@ pub unsafe extern "C" fn ambuild( .opfamily .datum_to_vector(*values.add(0), *is_null.add(0)) }; - let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { + let pointer = unsafe { ctid_to_pointer(ctid.read()) }; + let extra = unsafe { + match (*(*index).rd_att).natts { + 1 => None, + 2 => { + if !is_null.add(1).read() { + Some(values.add(1).read().value() as u32) + } else { + None + } + } + _ => unreachable!(), + } + }; let vector = match vector { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), OwnedVector::BVector(_) => unreachable!(), }; - (state.callback)((pointer, vector.into_vec())); + (state.callback)((pointer, extra, vector.into_vec())); } } let table_am = unsafe { &*(*self.heap).rd_tableam }; @@ -257,14 +271,32 @@ pub unsafe extern "C" fn aminsert( let opfamily = unsafe { am_options::opfamily(index) }; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { + let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); + let extra = unsafe { + match (*(*index).rd_att).natts { + 1 => None, + 2 => { + if !is_null.add(1).read() { + Some(values.add(1).read().value() as u32) + } else { + None + } + } + _ => unreachable!(), + } + }; let vector = match vector { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), OwnedVector::BVector(_) => unreachable!(), }; - let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - algorithm::insert::insert(unsafe { Relation::new(index) }, pointer, vector.into_vec()); + algorithm::insert::insert( + unsafe { Relation::new(index) }, + pointer, + extra, + vector.into_vec(), + ); } false } @@ -279,7 +311,7 @@ pub unsafe extern "C" fn ambeginscan( let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index, n_keys, n_orderbys) }; unsafe { - let scanner = am_scan::scan_make(None, None, false); + let scanner = am_scan::scan_make(None, None, Vec::new(), false); (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); } scan @@ -301,9 +333,10 @@ pub unsafe extern "C" fn amrescan( std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); } let opfamily = am_options::opfamily((*scan).indexRelation); - let (orderbys, spheres) = { + let (orderbys, spheres, filters) = { let mut orderbys = Vec::new(); let mut spheres = Vec::new(); + let mut filters = Vec::new(); if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { pgrx::error!( "vector search with no WHERE clause and no ORDER BY clause is not supported" @@ -324,14 +357,26 @@ pub unsafe extern "C" fn amrescan( let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; match (*data).sk_strategy { 2 => spheres.push(opfamily.datum_to_sphere(value, is_null)), + 11 => { + if is_null { + filters.push(0); + filters.push(1); + } else { + filters.push(value.value() as u32); + } + } _ => unreachable!(), } } - (orderbys, spheres) + (orderbys, spheres, filters) }; - let (vector, threshold, recheck) = am_scan::scan_build(orderbys, spheres, opfamily); + let (vector, threshold, filters, recheck) = + am_scan::scan_build(orderbys, spheres, filters, opfamily); let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace(scanner, am_scan::scan_make(vector, threshold, recheck)); + let scanner = std::mem::replace( + scanner, + am_scan::scan_make(vector, threshold, filters, recheck), + ); am_scan::scan_release(scanner); } } @@ -371,7 +416,7 @@ pub unsafe extern "C" fn amgettuple( pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { unsafe { let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace(scanner, am_scan::scan_make(None, None, false)); + let scanner = std::mem::replace(scanner, am_scan::scan_make(None, None, Vec::new(), false)); am_scan::scan_release(scanner); } } diff --git a/src/index/am_options.rs b/src/index/am_options.rs index 1a977e16..eb9aa88c 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -116,8 +116,8 @@ pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, Rabbitho if atts.is_empty() { pgrx::error!("indexing on no columns is not supported"); } - if atts.len() != 1 { - pgrx::error!("multicolumn index is not supported"); + if atts.len() > 2 { + pgrx::error!("indexing on many columns is not supported"); } // get dims let typmod = Typmod::parse_from_i32(atts[0].type_mod()).unwrap(); diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs index 9792729a..62f6a0ad 100644 --- a/src/index/am_scan.rs +++ b/src/index/am_scan.rs @@ -10,6 +10,7 @@ pub enum Scanner { Initial { vector: Option<(OwnedVector, Opfamily)>, threshold: Option, + filters: Vec, recheck: bool, }, Vbase { @@ -24,8 +25,9 @@ pub enum Scanner { pub fn scan_build( orderbys: Vec>, spheres: Vec<(Option, Option)>, + filters: Vec, opfamily: Opfamily, -) -> (Option<(OwnedVector, Opfamily)>, Option, bool) { +) -> (Option<(OwnedVector, Opfamily)>, Option, Vec, bool) { let mut pair = None; let mut threshold = None; let mut recheck = false; @@ -49,17 +51,19 @@ pub fn scan_build( break; } } - (pair.map(|x| (x, opfamily)), threshold, recheck) + (pair.map(|x| (x, opfamily)), threshold, filters, recheck) } pub fn scan_make( vector: Option<(OwnedVector, Opfamily)>, threshold: Option, + filters: Vec, recheck: bool, ) -> Scanner { Scanner::Initial { vector, threshold, + filters, recheck, } } @@ -68,6 +72,7 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, if let Scanner::Initial { vector, threshold, + filters, recheck, } = scanner { @@ -81,6 +86,7 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, OwnedVector::BVector(_) => unreachable!(), }, nprobe(), + filters.clone(), ); *scanner = Scanner::Vbase { vbase: Box::new(vbase), diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 96dfaaa6..0f659e29 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -27,6 +27,7 @@ COMMENT ON ACCESS METHOD rabbithole IS 'rabbithole index access method'; -- List of operator families CREATE OPERATOR FAMILY vector_l2_ops USING rabbithole; +CREATE OPERATOR FAMILY int4_eq_ops USING rabbithole; -- List of operator classes @@ -34,3 +35,7 @@ CREATE OPERATOR CLASS vector_l2_ops FOR TYPE vector USING rabbithole FAMILY vector_l2_ops AS OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH; + +CREATE OPERATOR CLASS int4_eq_ops + FOR TYPE int4 USING rabbithole FAMILY int4_eq_ops AS + OPERATOR 11 = (int4, int4) FOR SEARCH; From 9e57e043337b19ec963802e1557a20d5a1197a0f Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 24 Oct 2024 16:08:16 +0800 Subject: [PATCH 008/324] feat: insert by scan (#21) Signed-off-by: usamoi --- src/algorithm/build.rs | 450 ++++++++++++++++++++++++++-------------- src/algorithm/insert.rs | 122 +++++++---- src/index/utils.rs | 7 +- 3 files changed, 384 insertions(+), 195 deletions(-) diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index c7b67187..fb9264c2 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -5,6 +5,8 @@ use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; use crate::types::ExternalCentroids; use crate::types::RabbitholeIndexingOptions; +use base::always_equal::AlwaysEqual; +use base::distance::Distance; use base::distance::DistanceKind; use base::index::VectorOptions; use base::scalar::ScalarLike; @@ -12,6 +14,8 @@ use base::search::Pointer; use common::vec2::Vec2; use rand::Rng; use rkyv::ser::serializers::AllocSerializer; +use std::cmp::Reverse; +use std::collections::BinaryHeap; use std::marker::PhantomData; use std::sync::atomic::AtomicBool; use std::sync::Arc; @@ -31,7 +35,7 @@ pub fn build( vector_options: VectorOptions, rabbithole_options: RabbitholeIndexingOptions, heap_relation: T, - index_relation: Relation, + relation: Relation, mut reporter: R, ) { let dims = vector_options.dims; @@ -67,56 +71,92 @@ pub fn build( Structure::compute(vector_options.clone(), rabbithole_options.clone(), samples) } }; - let h2_len = structure.h2_len(); - let h1_len = structure.h1_len(); - let mut meta = Tape::create(&index_relation); - assert_eq!(meta.first(), 0); - let mut vectors = Tape::create(&index_relation); - assert_eq!(vectors.first(), 1); - let h2_means = (0..h2_len) - .map(|i| { - vectors.push(&VectorTuple { - payload: None, - vector: structure.h2_means(i).clone(), + let mut vectors_head; + { + let h2_len = structure.h2_len(); + let h1_len = structure.h1_len(); + let mut meta = Tape::create(&relation); + assert_eq!(meta.first(), 0); + let mut vectors = Tape::create(&relation); + assert_eq!(vectors.first(), 1); + let h2_means = (0..h2_len) + .map(|i| { + vectors.push(&VectorTuple { + payload: None, + vector: structure.h2_means(i).clone(), + }) }) - }) - .collect::>(); - let h1_means = (0..h1_len) - .map(|i| { - vectors.push(&VectorTuple { - payload: None, - vector: structure.h1_means(i).clone(), + .collect::>(); + let h1_means = (0..h1_len) + .map(|i| { + vectors.push(&VectorTuple { + payload: None, + vector: structure.h1_means(i).clone(), + }) }) - }) - .collect::>(); - let h1_firsts = (0..h1_len) - .map(|_| { - let tape = Tape::::create(&index_relation); - tape.first() - }) - .collect::>(); - let h2_firsts = (0..h2_len) - .map(|i| { - let mut tape = Tape::::create(&index_relation); - let mut cache = Vec::new(); - let h2_mean = structure.h2_means(i); - let children = structure.h2_children(i); - for child in children.iter().copied() { - let h1_mean = structure.h1_means(child); - let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(h1_mean, h2_mean)) - } else { - rabitq::code(dims, h1_mean) - }; - cache.push((child, code)); - if cache.len() == 32 { + .collect::>(); + let h1_firsts = (0..h1_len) + .map(|_| { + let tape = Tape::::create(&relation); + tape.first() + }) + .collect::>(); + let h2_firsts = (0..h2_len) + .map(|i| { + let mut tape = Tape::::create(&relation); + let mut cache = Vec::new(); + let h2_mean = structure.h2_means(i); + let children = structure.h2_children(i); + for child in children.iter().copied() { + let h1_mean = structure.h1_means(child); + let code = if is_residual { + rabitq::code(dims, &f32::vector_sub(h1_mean, h2_mean)) + } else { + rabitq::code(dims, h1_mean) + }; + cache.push((child, code)); + if cache.len() == 32 { + let group = std::mem::take(&mut cache); + let code = std::array::from_fn(|i| group[i].1.clone()); + let packed = rabitq::pack_codes(dims, code); + tape.push(&Height1Tuple { + mask: [true; 32], + mean: std::array::from_fn(|i| h1_means[group[i].0 as usize]), + first: std::array::from_fn(|i| h1_firsts[group[i].0 as usize]), + dis_u_2: packed.dis_u_2, + factor_ppc: packed.factor_ppc, + factor_ip: packed.factor_ip, + factor_err: packed.factor_err, + t: packed.t, + }); + } + } + if !cache.is_empty() { let group = std::mem::take(&mut cache); - let code = std::array::from_fn(|i| group[i].1.clone()); - let packed = rabitq::pack_codes(dims, code); + let codes = std::array::from_fn(|i| { + if i < group.len() { + group[i].1.clone() + } else { + rabitq::dummy_code(dims) + } + }); + let packed = rabitq::pack_codes(dims, codes); tape.push(&Height1Tuple { - mask: [true; 32], - mean: std::array::from_fn(|i| h1_means[group[i].0 as usize]), - first: std::array::from_fn(|i| h1_firsts[group[i].0 as usize]), + mask: std::array::from_fn(|i| i < group.len()), + mean: std::array::from_fn(|i| { + if i < group.len() { + h1_means[group[i].0 as usize] + } else { + Default::default() + } + }), + first: std::array::from_fn(|i| { + if i < group.len() { + h1_firsts[group[i].0 as usize] + } else { + Default::default() + } + }), dis_u_2: packed.dis_u_2, factor_ppc: packed.factor_ppc, factor_ip: packed.factor_ip, @@ -124,139 +164,235 @@ pub fn build( t: packed.t, }); } - } - if !cache.is_empty() { - let group = std::mem::take(&mut cache); - let codes = std::array::from_fn(|i| { - if i < group.len() { - group[i].1.clone() - } else { - rabitq::dummy_code(dims) - } - }); - let packed = rabitq::pack_codes(dims, codes); - tape.push(&Height1Tuple { - mask: std::array::from_fn(|i| i < group.len()), - mean: std::array::from_fn(|i| { - if i < group.len() { - h1_means[group[i].0 as usize] - } else { - Default::default() - } - }), - first: std::array::from_fn(|i| { - if i < group.len() { - h1_firsts[group[i].0 as usize] - } else { - Default::default() - } - }), - dis_u_2: packed.dis_u_2, - factor_ppc: packed.factor_ppc, - factor_ip: packed.factor_ip, - factor_err: packed.factor_err, - t: packed.t, - }); - } - tape.first() - }) - .collect::>(); - meta.push(&MetaTuple { - dims, - is_residual, - vectors_first: vectors.first(), - mean: h2_means[0], - first: h2_firsts[0], - }); - drop(meta); - let mut heads = Vec::new(); - for i in 0..structure.h1_len() { - heads.push(h1_firsts[i as usize]); + tape.first() + }) + .collect::>(); + meta.push(&MetaTuple { + dims, + is_residual, + vectors_first: vectors.first(), + mean: h2_means[0], + first: h2_firsts[0], + }); + vectors_head = vectors.head.id(); } let mut tuples_done = 0; heap_relation.traverse(|(payload, extra, vector)| { pgrx::check_for_interrupts!(); tuples_done += 1; reporter.tuples_done(tuples_done); + let meta_guard = relation.read(0); + let meta_tuple = meta_guard + .get() + .get(1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); let vector = rabitq::project(&vector); - let h0_vector = vectors.push(&VectorTuple { - vector: vector.clone(), - payload: Some(payload.as_u64()), - }); + let default_lut = if !is_residual { + Some(rabitq::fscan_preprocess(&vector)) + } else { + None + }; + let h0_vector = 'h0_vector: { + let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { + vector: vector.clone(), + payload: Some(payload.as_u64()), + }) + .unwrap(); + let mut write = relation.write(vectors_head); + if let Some(i) = write.get_mut().alloc(&tuple) { + break 'h0_vector (vectors_head, i); + } + let mut extend = relation.extend(); + write.get_mut().get_opaque_mut().next = extend.id(); + if let Some(i) = extend.get_mut().alloc(&tuple) { + vectors_head = extend.id(); + break 'h0_vector (extend.id(), i); + } else { + panic!("a tuple cannot even be fit in a fresh page"); + } + }; let h0_payload = payload.as_u64(); - let h2_id = 0_u32; - let h1_id = { - let mut target = (0_u32, f32::INFINITY); - for &i in structure.h2_children(h2_id) { - let dis = f32::reduce_sum_of_d2(&vector, structure.h1_means(i)); - if dis < target.1 { - target = (i, dis); + let list = ( + meta_tuple.first, + if is_residual { + let vector_guard = relation.read(meta_tuple.mean.0); + let vector_tuple = vector_guard + .get() + .get(meta_tuple.mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + Some(vector_tuple.vector.to_vec()) + } else { + None + }, + ); + let list = { + let mut results = Vec::new(); + { + let lut = if is_residual { + &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) + } else { + default_lut.as_ref().unwrap() + }; + let mut current = list.0; + while current != u32::MAX { + let h1_guard = relation.read(current); + for i in 1..=h1_guard.get().len() { + let h1_tuple = h1_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let lowerbounds = rabitq::fscan_process_lowerbound( + dims, + lut, + ( + &h1_tuple.dis_u_2, + &h1_tuple.factor_ppc, + &h1_tuple.factor_ip, + &h1_tuple.factor_err, + &h1_tuple.t, + ), + ); + for j in 0..32 { + if h1_tuple.mask[j] { + results.push(( + Reverse(lowerbounds[j]), + AlwaysEqual(h1_tuple.mean[j]), + AlwaysEqual(h1_tuple.first[j]), + )); + } + } + } + current = h1_guard.get().get_opaque().next; + } + } + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); + { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + let vector_guard = relation.read(mean.0); + let vector_tuple = vector_guard + .get() + .get(mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let dis_u = + Distance::from_f32(f32::reduce_sum_of_d2(&vector, &vector_tuple.vector)); + cache.push(( + Reverse(dis_u), + AlwaysEqual(first), + AlwaysEqual(if is_residual { + Some(vector_tuple.vector.to_vec()) + } else { + None + }), + )); } + let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop().unwrap(); + (first, mean) } - target.0 }; let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(&vector, structure.h1_means(h1_id))) + rabitq::code(dims, &f32::vector_sub(&vector, list.1.as_ref().unwrap())) } else { rabitq::code(dims, &vector) }; - let mut write = index_relation.write(heads[h1_id as usize]); - let page = write.get_mut(); - if page.len() != 0 { - let flag = put( - page.get_mut(page.len()).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - extra, - ); - if flag { - return; - } - } - let tuple = rkyv::to_bytes::<_, 8192>(&Height0Tuple { + let dummy = rkyv::to_bytes::<_, 8192>(&Height0Tuple { mask: [false; 32], mean: [(0, 0); 32], payload: [0; 32], - dis_u_2: [0.0; 32], - factor_ppc: [0.0; 32], - factor_ip: [0.0; 32], - factor_err: [0.0; 32], - t: vec![0; (dims.div_ceil(4) * 16) as usize], extra: [None; 32], + dis_u_2: [0.0f32; 32], + factor_ppc: [0.0f32; 32], + factor_ip: [0.0f32; 32], + factor_err: [0.0f32; 32], + t: vec![0; (dims.div_ceil(4) * 16) as usize], }) .unwrap(); - if let Some(i) = page.alloc(&tuple) { - let flag = put( - page.get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - extra, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; - } - let mut extend = index_relation.extend(); - heads[h1_id as usize] = extend.id(); - page.get_opaque_mut().next = extend.id(); - let page = extend.get_mut(); - if let Some(i) = page.alloc(&tuple) { - let flag = put( - page.get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - extra, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; + let first = list.0; + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = relation.read(current); + let flag = 'flag: { + for i in 1..=read.get().len() { + let h0_tuple = read + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + if h0_tuple.mask.iter().any(|x| *x) { + break 'flag true; + } + } + if read.get().freespace() as usize >= dummy.len() { + break 'flag true; + } + if read.get().get_opaque().next == u32::MAX { + break 'flag true; + } + false + }; + if flag { + drop(read); + let mut write = relation.write(current); + for i in 1..=write.get().len() { + let flag = put( + write.get_mut().get_mut(i).expect("data corruption"), + dims, + &code, + h0_vector, + h0_payload, + extra, + ); + if flag { + return; + } + } + if let Some(i) = write.get_mut().alloc(&dummy) { + let flag = put( + write.get_mut().get_mut(i).expect("data corruption"), + dims, + &code, + h0_vector, + h0_payload, + extra, + ); + assert!(flag, "a put fails even on a fresh tuple"); + return; + } + if write.get().get_opaque().next == u32::MAX { + let mut extend = relation.extend(); + write.get_mut().get_opaque_mut().next = extend.id(); + if let Some(i) = extend.get_mut().alloc(&dummy) { + let flag = put( + extend.get_mut().get_mut(i).expect("data corruption"), + dims, + &code, + h0_vector, + h0_payload, + extra, + ); + assert!(flag, "a put fails even on a fresh tuple"); + return; + } else { + panic!("a tuple cannot even be fit in a fresh page"); + } + } + current = write.get().get_opaque().next; + } else { + current = read.get().get_opaque().next; + } } - panic!("a tuple cannot even be fit in a fresh page") }); } diff --git a/src/algorithm/insert.rs b/src/algorithm/insert.rs index ebd2abbc..82dd544e 100644 --- a/src/algorithm/insert.rs +++ b/src/algorithm/insert.rs @@ -1,9 +1,12 @@ use crate::algorithm::rabitq; use crate::algorithm::tuples::*; use crate::postgres::Relation; +use base::always_equal::AlwaysEqual; use base::distance::Distance; use base::scalar::ScalarLike; use base::search::Pointer; +use std::cmp::Reverse; +use std::collections::BinaryHeap; pub fn insert(relation: Relation, payload: Pointer, extra: Option, vector: Vec) { let meta_guard = relation.read(0); @@ -17,6 +20,11 @@ pub fn insert(relation: Relation, payload: Pointer, extra: Option, vector: assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); let vector = rabitq::project(&vector); let is_residual = meta_tuple.is_residual; + let default_lut = if !is_residual { + Some(rabitq::fscan_preprocess(&vector)) + } else { + None + }; let h0_vector = { let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { vector: vector.clone(), @@ -57,50 +65,90 @@ pub fn insert(relation: Relation, payload: Pointer, extra: Option, vector: } }; let h0_payload = payload.as_u64(); - let list = (meta_tuple.first,); + let list = ( + meta_tuple.first, + if is_residual { + let vector_guard = relation.read(meta_tuple.mean.0); + let vector_tuple = vector_guard + .get() + .get(meta_tuple.mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + Some(vector_tuple.vector.to_vec()) + } else { + None + }, + ); let list = { - let mut result = None::<(Distance, u32, Option>)>; - let mut current = list.0; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { - let h1_tuple = h1_guard - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - for j in 0..32 { - if h1_tuple.mask[j] { - let vector_guard = relation.read(h1_tuple.mean[j].0); - let vector_tuple = vector_guard - .get() - .get(h1_tuple.mean[j].1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dis = Distance::from_f32(f32::reduce_sum_of_d2( - &vector, - &vector_tuple.vector, - )); - if result.is_none() || dis < result.as_ref().unwrap().0 { - result = Some(( - dis, - h1_tuple.first[j], - if is_residual { - Some(vector_tuple.vector.to_vec()) - } else { - None - }, + let mut results = Vec::new(); + { + let lut = if is_residual { + &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) + } else { + default_lut.as_ref().unwrap() + }; + let mut current = list.0; + while current != u32::MAX { + let h1_guard = relation.read(current); + for i in 1..=h1_guard.get().len() { + let h1_tuple = h1_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let lowerbounds = rabitq::fscan_process_lowerbound( + dims, + lut, + ( + &h1_tuple.dis_u_2, + &h1_tuple.factor_ppc, + &h1_tuple.factor_ip, + &h1_tuple.factor_err, + &h1_tuple.t, + ), + ); + for j in 0..32 { + if h1_tuple.mask[j] { + results.push(( + Reverse(lowerbounds[j]), + AlwaysEqual(h1_tuple.mean[j]), + AlwaysEqual(h1_tuple.first[j]), )); } } } + current = h1_guard.get().get_opaque().next; + } + } + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); + { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + let vector_guard = relation.read(mean.0); + let vector_tuple = vector_guard + .get() + .get(mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let dis_u = + Distance::from_f32(f32::reduce_sum_of_d2(&vector, &vector_tuple.vector)); + cache.push(( + Reverse(dis_u), + AlwaysEqual(first), + AlwaysEqual(if is_residual { + Some(vector_tuple.vector.to_vec()) + } else { + None + }), + )); } - current = h1_guard.get().get_opaque().next; + let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop().unwrap(); + (first, mean) } - let result = result.unwrap(); - (result.1, result.2) }; let code = if is_residual { rabitq::code(dims, &f32::vector_sub(&vector, list.1.as_ref().unwrap())) diff --git a/src/index/utils.rs b/src/index/utils.rs index d6bde6a9..4e500917 100644 --- a/src/index/utils.rs +++ b/src/index/utils.rs @@ -25,7 +25,12 @@ pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer { Pointer::new(value) } -pub fn load_proj_vectors(table_name: &str, column_name: &str, rows: u32, dims: u32) -> Vec> { +pub fn load_proj_vectors( + table_name: &str, + column_name: &str, + rows: u32, + dims: u32, +) -> Vec> { let query = format!("SELECT {column_name} FROM {table_name};"); let mut centroids = Vec::new(); From 3ae5c489b65ef4d76075076c07894d7a34add253 Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 25 Oct 2024 10:20:20 +0800 Subject: [PATCH 009/324] feat: dump db emb to local h5, sampling for kmeans (#22) * init Signed-off-by: Keming * add dump Signed-off-by: Keming * fix train Signed-off-by: Keming * w/wo centroids for indexing Signed-off-by: Keming --------- Signed-off-by: Keming --- bench/README.md | 22 ++++---- bench/bench.py | 118 +++++++++++++++++++++------------------ bench/dump.py | 69 +++++++++++++++++++++++ bench/index.py | 137 ++++++++++++++++++++++++++++++++++++++++++++++ bench/internal.py | 95 -------------------------------- bench/load.py | 114 -------------------------------------- bench/train.py | 133 ++++++++++++++++++++++++++++++-------------- 7 files changed, 375 insertions(+), 313 deletions(-) create mode 100644 bench/dump.py create mode 100644 bench/index.py delete mode 100644 bench/internal.py delete mode 100644 bench/load.py diff --git a/bench/README.md b/bench/README.md index f36106f1..8d17e28b 100644 --- a/bench/README.md +++ b/bench/README.md @@ -1,6 +1,9 @@ ## Build Docker ```shell +sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config clang +cargo install --locked cargo-pgrx +cargo pgrx init cargo build --package rabbithole --lib --features pg16 --target x86_64-unknown-linux-gnu --release ./tools/schema.sh --features pg16 --target x86_64-unknown-linux-gnu --release | expand -t 4 > ./target/schema.sql @@ -53,16 +56,15 @@ Options for `-n`: ```shell # pip install pgvector numpy faiss-cpu psycopg h5py -# If you want to use internal k-means(slow) -python internal.py -n sift -m l2 +# dump table embedding column to a local h5 file["train"] +python dump.py -n sift -o sift.h5 -c embedding -d 128 -# Or external k-means -## K-means generate centroids -python train.py -n sift +# external k-means +python train.py -i sift.hdf5 -o sift_centroids_4096 -## Load centroids into PG and build index -python load.py -n sift -m l2 +# build index (w/wo external centroids) +python index.py -n sift -c sift_centroids_4096.npy -i sift.hdf5 -# Run bench -python bench.py -n sift -m l2 -``` \ No newline at end of file +# bench +python bench.py -n sift -i sift.hdf5 +``` diff --git a/bench/bench.py b/bench/bench.py index 991f7b26..0271004d 100644 --- a/bench/bench.py +++ b/bench/bench.py @@ -1,5 +1,3 @@ -from os.path import join -import os import time import argparse from pathlib import Path @@ -8,62 +6,76 @@ import psycopg import h5py from pgvector.psycopg import register_vector -import numpy as np -parser = argparse.ArgumentParser() -parser.add_argument( - "-m", "--metric", help="Metric to pick, in l2 or cos", required=True -) -parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) -args = parser.parse_args() -HOME = Path.home() -INDEX_PATH = join(HOME, f"indexes/pg/{args.name}/{args.metric}") -DATA_PATH = join(HOME, f"{args.name}/{args.name}.hdf5") +def build_arg_parse(): + parser = argparse.ArgumentParser() + parser.add_argument( + "-m", + "--metric", + help="Metric to pick, in l2 or cos", + choices=["l2", "cos"], + default="l2", + ) + parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) + parser.add_argument("-i", "--input", help="input filepath", required=True) + parser.add_argument( + "-p", "--password", help="Database password", default="password" + ) + return parser -os.makedirs(join(HOME, f"indexes/pg/{args.name}"), exist_ok=True) -dataset = h5py.File(DATA_PATH, "r") -test = dataset["test"][:] +def create_connection(password): + keepalive_kwargs = { + "keepalives": 1, + "keepalives_idle": 30, + "keepalives_interval": 5, + "keepalives_count": 5, + } + conn = psycopg.connect( + conninfo=f"postgresql://postgres:{password}@localhost:5432/postgres", + dbname="postgres", + autocommit=True, + **keepalive_kwargs, + ) + conn.execute("SET search_path TO public, vectors, rabbithole") + conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") + register_vector(conn) + return conn -if args.metric == "l2": - metric_ops = "<->" -elif args.metric == "cos": - metric_ops = "<=>" -else: - raise ValueError -answer = dataset["neighbors"][:] -n, dims = np.shape(test) -m = np.shape(test)[0] +def bench(name, test, answer, metric_ops, conn): + m = test.shape[0] + for k in [10, 100]: + hits = 0 + delta = 0 + pbar = tqdm(enumerate(test), total=m) + for i, query in pbar: + start = time.perf_counter() + result = conn.execute( + f"SELECT id FROM {name} ORDER BY embedding {metric_ops} %s LIMIT {k}", + (query,), + ).fetchall() + end = time.perf_counter() + hits += len(set([p[0] for p in result[:k]]) & set(answer[i][:k].tolist())) + delta += end - start + pbar.set_description(f"recall: {hits / k / (i+1)} QPS: {(i+1) / delta} ") + recall = hits / k / m + qps = m / delta + print(f"Top: {k} recall: {recall:.4f} QPS: {qps:.2f}") -# reconnect for updated GUC variables to take effect -conn = psycopg.connect( - conninfo="postgres://bench:123@localhost:5432/postgres", - dbname="postgres", - autocommit=True, -) -conn.execute("SET search_path TO public, vectors, rabbithole") -conn.execute("CREATE EXTENSION IF NOT EXISTS vector") -conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") -conn.execute("SET rabbithole.nprobe=300") -register_vector(conn) -Ks = [10, 100] -for k in Ks: - hits = 0 - delta = 0 - pbar = tqdm(enumerate(test), total=m) - for i, query in pbar: - start = time.perf_counter() - result = conn.execute( - f"SELECT id FROM {args.name} ORDER BY embedding {metric_ops} %s LIMIT {k}", - (query,), - ).fetchall() - end = time.perf_counter() - hits += len(set([p[0] for p in result[:k]]) & set(answer[i][:k].tolist())) - delta += end - start - pbar.set_description(f"recall: {hits / k / (i+1)} QPS: {(i+1) / delta} ") - recall = hits / k / m - qps = m / delta - print(f"Top: {k} recall: {recall:.4f} QPS: {qps:.2f}") \ No newline at end of file +if __name__ == "__main__": + parser = build_arg_parse() + args = parser.parse_args() + print(args) + + dataset = h5py.File(Path(args.input), "r") + test = dataset["test"][:] + answer = dataset["neighbors"][:] + metric_ops = "<->" if args.metric == "l2" else "<=>" + conn = create_connection(args.password) + conn.execute("SET rabbithole.nprobe=300") + + bench(args.name, test, answer, metric_ops, conn) diff --git a/bench/dump.py b/bench/dump.py new file mode 100644 index 00000000..528ef056 --- /dev/null +++ b/bench/dump.py @@ -0,0 +1,69 @@ +import argparse + +import h5py +import numpy as np +import psycopg +from tqdm import tqdm +from pgvector.psycopg import register_vector + + +def build_arg_parse(): + parser = argparse.ArgumentParser(description="Dump embeddings to a local file") + parser.add_argument("-n", "--name", help="table name", required=True) + parser.add_argument( + "-p", "--password", help="Database password", default="password" + ) + parser.add_argument("-o", "--output", help="Output filepath", required=True) + parser.add_argument("-c", "--column", help="Column name", default="embedding") + parser.add_argument("-d", "--dim", help="Dimension", type=int, required=True) + return parser + + +def create_connection(password): + keepalive_kwargs = { + "keepalives": 1, + "keepalives_idle": 30, + "keepalives_interval": 5, + "keepalives_count": 5, + } + conn = psycopg.connect( + conninfo=f"postgresql://postgres:{password}@localhost:5432/postgres", + dbname="postgres", + autocommit=True, + **keepalive_kwargs, + ) + conn.execute("SET search_path TO public, vectors, rabbithole") + conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") + register_vector(conn) + return conn + + +def extract_vectors(conn, name, column): + with conn.execute(f"SELECT {column} FROM {name}") as cursor: + for row in cursor: + yield row[0] + + +def write_to_h5(filepath, vecs, dim): + with h5py.File(filepath, "w") as file: + dataset = file.create_dataset( + "train", (0, dim), maxshape=(None, dim), chunks=True, dtype=np.float32 + ) + current_size = 0 + for vec in tqdm(vecs): + if dataset.shape[0] == current_size: + dataset.resize(current_size + 64, axis=0) + dataset[current_size] = vec + current_size += 1 + dataset.resize(current_size, axis=0) + dataset.flush() + + +if __name__ == "__main__": + parser = build_arg_parse() + args = parser.parse_args() + print(args) + + conn = create_connection(args.password) + write_to_h5(args.output, extract_vectors(conn, args.name, args.column), args.dim) diff --git a/bench/index.py b/bench/index.py new file mode 100644 index 00000000..ac32ae7e --- /dev/null +++ b/bench/index.py @@ -0,0 +1,137 @@ +from sys import version_info +from time import perf_counter +import argparse +from pathlib import Path + +if version_info >= (3, 12): + raise RuntimeError("h5py doesn't support 3.12") + +import psycopg +import h5py +from pgvector.psycopg import register_vector +import numpy as np +from tqdm import tqdm + + +def build_arg_parse(): + parser = argparse.ArgumentParser(description="Build index with K-means centroids") + parser.add_argument( + "-m", "--metric", help="Distance metric", default="l2", choices=["l2", "cos"] + ) + parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) + parser.add_argument( + "-c", "--centroid", help="K-means centroids file", required=True + ) + parser.add_argument("-i", "--input", help="Input filepath", required=True) + parser.add_argument( + "-p", "--password", help="Database password", default="password" + ) + parser.add_argument("-k", help="Number of centroids", type=int, required=True) + parser.add_argument("-d", "--dim", help="Dimension", type=int, required=True) + return parser + + +def get_ivf_ops_config(metric, k, name=None): + external_centroids = """ + [external_centroids] + table = 'public.{name}_centroids' + h1_means_column = 'coordinate' + """ + if metric == "l2": + metric_ops = "vector_l2_ops" + ivf_config = f""" + nlist = {k} + residual_quantization = true + spherical_centroids = false + """ + elif metric == "cos": + metric_ops = "vector_cos_ops" + ivf_config = f""" + nlist = {k} + residual_quantization = false + spherical_centroids = true + """ + else: + raise ValueError + + if name: + ivf_config += external_centroids.format(name=name) + return metric_ops, ivf_config + + +def create_connection(password): + keepalive_kwargs = { + "keepalives": 1, + "keepalives_idle": 30, + "keepalives_interval": 5, + "keepalives_count": 5, + } + conn = psycopg.connect( + conninfo=f"postgresql://postgres:{password}@localhost:5432/postgres", + dbname="postgres", + autocommit=True, + **keepalive_kwargs, + ) + conn.execute("SET search_path TO public, vectors, rabbithole") + conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") + register_vector(conn) + return conn + + +def add_centroids(conn, name, centroids): + dim = centroids.shape[1] + conn.execute(f"DROP TABLE IF EXISTS public.{name}_centroids") + conn.execute(f"CREATE TABLE public.{name}_centroids (coordinate vector({dim}))") + with conn.cursor().copy( + f"COPY public.{name}_centroids (coordinate) FROM STDIN WITH (FORMAT BINARY)" + ) as copy: + copy.set_types(["vector"]) + for centroid in tqdm(centroids, desc="Adding centroids"): + copy.write_row((centroid,)) + while conn.pgconn.flush() == 1: + pass + + +def build_index(conn, name, metric_ops, ivf_config, dim, train): + conn.execute(f"DROP TABLE IF EXISTS {name}") + conn.execute(f"CREATE TABLE {name} (id integer, embedding vector({dim}))") + + with conn.cursor().copy( + f"COPY {name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)" + ) as copy: + copy.set_types(["integer", "vector"]) + + for i, vec in tqdm(enumerate(train), desc="Adding embeddings"): + copy.write_row((i, vec)) + while conn.pgconn.flush() == 1: + pass + + start_time = perf_counter() + conn.execute( + f"CREATE INDEX ON {name} USING rabbithole (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" + ) + print(f"Index build time: {perf_counter() - start_time:.2f}s") + + +if __name__ == "__main__": + parser = build_arg_parse() + args = parser.parse_args() + print(args) + + dataset = h5py.File(Path(args.input), "r") + conn = create_connection(args.password) + if args.centroids: + centroids = np.load(args.centroid, allow_pickle=False) + add_centroids(conn, args.name, centroids) + metric_ops, ivf_config = get_ivf_ops_config( + args.metric, args.k, args.name if args.centroids else None + ) + build_index( + conn, + args.name, + metric_ops, + ivf_config, + args.dim, + dataset["train"], + ) diff --git a/bench/internal.py b/bench/internal.py deleted file mode 100644 index dd5806f1..00000000 --- a/bench/internal.py +++ /dev/null @@ -1,95 +0,0 @@ -import os -import time -import argparse -from pathlib import Path -import pickle - -import psycopg -import h5py -from pgvecto_rs.psycopg import register_vector -import numpy as np - -parser = argparse.ArgumentParser() -parser.add_argument( - "-m", "--metric", help="Metric to pick, in l2 or cos", required=True -) -parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) -parser.add_argument("-k", help="K-means centroids or nlist, default: 4096", required=False, default=4096) -args = parser.parse_args() - -HOME = Path.home() -DATA_PATH = join(HOME, f"{args.name}/{args.name}.hdf5") - -os.makedirs(join(HOME, f"indexes/pg/{args.name}"), exist_ok=True) -dataset = h5py.File(DATA_PATH, "r") - -train = dataset["train"][:] -test = dataset["test"][:] - -K = 4096 -if args.metric == "l2": - metric_ops = "vector_l2_ops" - ivf_config = f""" -nlist = {K} -residual_quantization = true -spherical_centroids = false -""" -elif args.metric == "cos": - metric_ops = "vector_cos_ops" - ivf_config = f""" -nlist = {K} -residual_quantization = false -spherical_centroids = true -""" -else: - raise ValueError - -answer = dataset["neighbors"][:] -n, dims = np.shape(train) -m = np.shape(test)[0] - -keepalive_kwargs = { - "keepalives": 1, - "keepalives_idle": 30, - "keepalives_interval": 5, - "keepalives_count": 5, -} - -start = time.perf_counter() - -with open(f"{args.name}.pickle", "rb") as f: - centroids = pickle.load(f) - -conn = psycopg.connect( - conninfo="postgres://bench:123@localhost:5432/postgres", - dbname="postgres", - autocommit=True, - **keepalive_kwargs, -) -conn.execute("SET search_path TO public, vectors, rabbithole") -conn.execute("CREATE EXTENSION IF NOT EXISTS vector") -conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") -register_vector(conn) - -conn.execute(f"DROP TABLE IF EXISTS public.centroids") -conn.execute(f"DROP TABLE IF EXISTS {args.name}") -conn.execute(f"CREATE TABLE {args.name} (id integer, embedding vector({dims}))") - -with conn.cursor().copy( - f"COPY {args.name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)" -) as copy: - copy.set_types(["integer", "vector"]) - - for i in range(n): - copy.write_row([i, train[i]]) - while conn.pgconn.flush() == 1: - pass - -start = time.perf_counter() -conn.execute( - f"CREATE INDEX ON {args.name} USING rabbithole (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" -) -end = time.perf_counter() - -delta = end - start -print(f"Index build time: {delta:.2f}s") \ No newline at end of file diff --git a/bench/load.py b/bench/load.py deleted file mode 100644 index 6b98b004..00000000 --- a/bench/load.py +++ /dev/null @@ -1,114 +0,0 @@ -from os.path import join -import os -import time -import argparse -from pathlib import Path -import pickle - -import psycopg -import h5py -from pgvector.psycopg import register_vector -import numpy as np - -parser = argparse.ArgumentParser() -parser.add_argument( - "-m", "--metric", help="Metric to pick, in l2 or cos", required=True -) -parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) -parser.add_argument("-k", help="K-means centroids or nlist, default: 4096", required=False, default=4096) -args = parser.parse_args() - -HOME = Path.home() -DATA_PATH = join(HOME, f"{args.name}/{args.name}.hdf5") - -os.makedirs(join(HOME, f"indexes/pg/{args.name}"), exist_ok=True) -dataset = h5py.File(DATA_PATH, "r") - -train = dataset["train"][:] -test = dataset["test"][:] - -K = args.k -if args.metric == "l2": - metric_ops = "vector_l2_ops" - ivf_config = f""" -nlist = {K} -residual_quantization = true -spherical_centroids = false -[external_centroids] -table = 'public.centroids' -h1_means_column = 'coordinate' -""" -elif args.metric == "cos": - metric_ops = "vector_cos_ops" - ivf_config = f""" -nlist = {K} -residual_quantization = false -spherical_centroids = true -[external_centroids] -table = 'public.centroids' -h1_means_column = 'coordinate' -""" -else: - raise ValueError - -answer = dataset["neighbors"][:] -n, dims = np.shape(train) -m = np.shape(test)[0] - -keepalive_kwargs = { - "keepalives": 1, - "keepalives_idle": 30, - "keepalives_interval": 5, - "keepalives_count": 5, -} - -start = time.perf_counter() - -with open(f"{args.name}.pickle", "rb") as f: - centroids = pickle.load(f) - -conn = psycopg.connect( - conninfo="postgres://bench:123@localhost:5432/postgres", - dbname="postgres", - autocommit=True, - **keepalive_kwargs, -) -conn.execute("SET search_path TO public, vectors, rabbithole") -conn.execute("CREATE EXTENSION IF NOT EXISTS vector") -conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") -register_vector(conn) - -conn.execute(f"DROP TABLE IF EXISTS public.centroids") -conn.execute(f"CREATE TABLE public.centroids (coordinate vector({dims}))") - -with conn.cursor().copy( - f"COPY public.centroids (coordinate) FROM STDIN WITH (FORMAT BINARY)" -) as copy: - copy.set_types(["vector"]) - - for i in range(K): - copy.write_row([centroids[i]]) - while conn.pgconn.flush() == 1: - pass - -conn.execute(f"DROP TABLE IF EXISTS {args.name}") -conn.execute(f"CREATE TABLE {args.name} (id integer, embedding vector({dims}))") - -with conn.cursor().copy( - f"COPY {args.name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)" -) as copy: - copy.set_types(["integer", "vector"]) - - for i in range(n): - copy.write_row([i, train[i]]) - while conn.pgconn.flush() == 1: - pass - -start = time.perf_counter() -conn.execute( - f"CREATE INDEX ON {args.name} USING rabbithole (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" -) -end = time.perf_counter() - -delta = end - start -print(f"Index build time: {delta:.2f}s") \ No newline at end of file diff --git a/bench/train.py b/bench/train.py index 802ae82b..48b24a7b 100644 --- a/bench/train.py +++ b/bench/train.py @@ -1,61 +1,112 @@ -from os.path import join -import time +from time import perf_counter import argparse from pathlib import Path -import pickle +from sys import version_info + +if version_info >= (3, 12): + raise RuntimeError("h5py doesn't support 3.12") import h5py -import faiss +from faiss import Kmeans import numpy as np +from tqdm import tqdm -K = 4096 +DEFAULT_K = 4096 +N_ITER = 25 SEED = 42 +MAX_POINTS_PER_CLUSTER = 256 + + +def build_arg_parse(): + parser = argparse.ArgumentParser(description="Train K-means centroids") + parser.add_argument("-i", "--input", help="input filepath", required=True) + parser.add_argument("-o", "--output", help="output filepath", required=True) + parser.add_argument( + "-k", + help="K-means centroids or nlist", + type=int, + default=DEFAULT_K, + ) + parser.add_argument("--child-k", type=int, help="lower layer nlist (if enabled)") + parser.add_argument( + "--niter", help="number of iterations", type=int, default=N_ITER + ) + parser.add_argument("-m", "--metric", choices=["l2", "cos"], default="l2") + return parser -parser = argparse.ArgumentParser() -parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) -parser.add_argument("-k", help="K-means centroids or nlist, default: 4096", required=False, default=4096) -args = parser.parse_args() -K = args.k +def reservoir_sampling(iterator, k: int): + """Reservoir sampling from an iterator.""" + res = [] + while len(res) < k: + try: + res.append(next(iterator)) + except StopIteration: + return np.vstack(res) + for i, vec in enumerate(iterator, k + 1): + j = np.random.randint(0, i) + if j < k: + res[j] = vec + return np.vstack(res) -HOME = Path.home() -DATA_PATH = join(HOME, f"{args.name}/{args.name}.hdf5") -dataset = h5py.File(DATA_PATH, "r") +def filter_by_label(iter, labels, target): + for i, vec in enumerate(iter): + if labels[i] == target: + yield vec -if len(dataset["train"]) > 256 * K: - rs = np.random.RandomState(SEED) - idx = rs.choice(len(dataset["train"]), size=256 * K, replace=False) - train = dataset["train"][np.sort(idx)] -else: - train = dataset["train"][:] -test = dataset["test"][:] +def kmeans_cluster(data, k, child_k, niter, metric): + n, dim = data.shape + if n > MAX_POINTS_PER_CLUSTER * k: + train = reservoir_sampling(iter(data), MAX_POINTS_PER_CLUSTER * args.k) + else: + train = data[:] + kmeans = Kmeans( + dim, k, verbose=True, niter=niter, seed=SEED, spherical=metric == "cos" + ) + kmeans.train(train) + if not child_k: + return kmeans.centroids -if np.shape(train)[0] > 256 * K: - rs = np.random.RandomState(SEED) - idx = rs.choice(np.shape(train)[0], size=256 * K, replace=False) - train = train[idx] + # train the lower layer k-means + labels = np.zeros(n, dtype=np.uint32) + for i, vec in tqdm(enumerate(data), desc="Assigning labels"): + _, label = kmeans.assign(vec.reshape((1, -1))) + labels[i] = label[0] -answer = dataset["neighbors"][:] -n, dims = np.shape(train) -m = np.shape(test)[0] + centroids = [] + total_k = k * child_k + for i in tqdm(range(k), desc="training k-means for child layers"): + samples = np.sum(labels == i) / n * total_k * MAX_POINTS_PER_CLUSTER + child_train = reservoir_sampling( + filter_by_label(iter(data), labels, i), samples + ) + child_kmeans = Kmeans( + dim, + child_k, + verbose=True, + niter=niter, + seed=SEED, + spherical=metric == "cos", + ) + child_kmeans.train(child_train) + centroids.append(child_kmeans.centroids) + return np.vstack(centroids) -start = time.perf_counter() -index = faiss.IndexFlatL2(dims) -clustering = faiss.Clustering(dims, K) -clustering.verbose = True -clustering.seed = 42 -clustering.niter = 10 -clustering.train(train, index) -centroids = faiss.vector_float_to_array(clustering.centroids) +if __name__ == "__main__": + parser = build_arg_parse() + args = parser.parse_args() + print(args) -end = time.perf_counter() -delta = end - start -print(f"K-means time: {delta:.2f}s") + dataset = h5py.File(Path(args.input), "r") + n, dim = dataset["train"].shape -centroids = centroids.reshape([K, -1]) + start_time = perf_counter() + centroids = kmeans_cluster( + dataset["train"], args.k, args.child_k, args.niter, args.metric + ) + print(f"K-means (k=({args.k}, {args.child_k})): {perf_counter() - start_time:.2f}s") -with open(f"{args.name}.pickle", "wb") as f: - pickle.dump(centroids, f) \ No newline at end of file + np.save(Path(args.output), centroids, allow_pickle=False) From 669e056ba8cc270edced76a0d2a7cd472b7e430c Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 25 Oct 2024 20:01:11 +0800 Subject: [PATCH 010/324] fix: vacuum fails if vacuum runs before (#26) Signed-off-by: usamoi --- src/algorithm/vacuum.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/algorithm/vacuum.rs b/src/algorithm/vacuum.rs index 880a4bc6..0dd52093 100644 --- a/src/algorithm/vacuum.rs +++ b/src/algorithm/vacuum.rs @@ -100,11 +100,10 @@ pub fn vacuum(relation: Relation, callback: impl Fn(Pointer) -> bool) { let read = relation.read(current); let flag = 'flag: { for i in 1..=read.get().len() { - let vector_tuple = read - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") + let Some(vector_tuple) = read.get().get(i) else { + continue; + }; + let vector_tuple = rkyv::check_archived_root::(vector_tuple) .expect("data corruption"); if let Some(payload) = vector_tuple.payload.as_ref().copied() { if callback(Pointer::new(payload)) { From feccc1973714a5661fc60d8373cb6aa313385cdc Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 28 Oct 2024 12:42:07 +0800 Subject: [PATCH 011/324] Revert "feat: inline filtering (#20)" (#28) This reverts commit 845e2ed6b781e79e0d172a3ae954948e5db28326. --- src/algorithm/build.rs | 10 ++---- src/algorithm/insert.rs | 6 +--- src/algorithm/scan.rs | 8 +---- src/algorithm/tuples.rs | 3 -- src/index/am.rs | 71 ++++++++--------------------------------- src/index/am_options.rs | 4 +-- src/index/am_scan.rs | 10 ++---- src/sql/finalize.sql | 5 --- 8 files changed, 22 insertions(+), 95 deletions(-) diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index fb9264c2..782cd345 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -23,7 +23,7 @@ use std::sync::Arc; pub trait HeapRelation { fn traverse(&self, callback: F) where - F: FnMut((Pointer, Option, Vec)); + F: FnMut((Pointer, Vec)); } pub trait Reporter { @@ -50,7 +50,7 @@ pub fn build( let max_number_of_samples = rabbithole_options.nlist.saturating_mul(256); let mut samples = Vec::new(); let mut number_of_samples = 0_u32; - heap_relation.traverse(|(_, _, vector)| { + heap_relation.traverse(|(_, vector)| { pgrx::check_for_interrupts!(); assert_eq!(dims as usize, vector.len(), "invalid vector dimensions",); let vector = rabitq::project(&vector); @@ -177,7 +177,7 @@ pub fn build( vectors_head = vectors.head.id(); } let mut tuples_done = 0; - heap_relation.traverse(|(payload, extra, vector)| { + heap_relation.traverse(|(payload, vector)| { pgrx::check_for_interrupts!(); tuples_done += 1; reporter.tuples_done(tuples_done); @@ -309,7 +309,6 @@ pub fn build( mask: [false; 32], mean: [(0, 0); 32], payload: [0; 32], - extra: [None; 32], dis_u_2: [0.0f32; 32], factor_ppc: [0.0f32; 32], factor_ip: [0.0f32; 32], @@ -352,7 +351,6 @@ pub fn build( &code, h0_vector, h0_payload, - extra, ); if flag { return; @@ -365,7 +363,6 @@ pub fn build( &code, h0_vector, h0_payload, - extra, ); assert!(flag, "a put fails even on a fresh tuple"); return; @@ -380,7 +377,6 @@ pub fn build( &code, h0_vector, h0_payload, - extra, ); assert!(flag, "a put fails even on a fresh tuple"); return; diff --git a/src/algorithm/insert.rs b/src/algorithm/insert.rs index 82dd544e..fff7678e 100644 --- a/src/algorithm/insert.rs +++ b/src/algorithm/insert.rs @@ -8,7 +8,7 @@ use base::search::Pointer; use std::cmp::Reverse; use std::collections::BinaryHeap; -pub fn insert(relation: Relation, payload: Pointer, extra: Option, vector: Vec) { +pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { let meta_guard = relation.read(0); let meta_tuple = meta_guard .get() @@ -159,7 +159,6 @@ pub fn insert(relation: Relation, payload: Pointer, extra: Option, vector: mask: [false; 32], mean: [(0, 0); 32], payload: [0; 32], - extra: [None; 32], dis_u_2: [0.0f32; 32], factor_ppc: [0.0f32; 32], factor_ip: [0.0f32; 32], @@ -202,7 +201,6 @@ pub fn insert(relation: Relation, payload: Pointer, extra: Option, vector: &code, h0_vector, h0_payload, - extra, ); if flag { return; @@ -215,7 +213,6 @@ pub fn insert(relation: Relation, payload: Pointer, extra: Option, vector: &code, h0_vector, h0_payload, - extra, ); assert!(flag, "a put fails even on a fresh tuple"); return; @@ -230,7 +227,6 @@ pub fn insert(relation: Relation, payload: Pointer, extra: Option, vector: &code, h0_vector, h0_payload, - extra, ); assert!(flag, "a put fails even on a fresh tuple"); return; diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs index d39de635..fc1d3718 100644 --- a/src/algorithm/scan.rs +++ b/src/algorithm/scan.rs @@ -12,7 +12,6 @@ pub fn scan( relation: Relation, vector: Vec, h1_nprobe: u32, - filters: Vec, ) -> impl Iterator { assert!(h1_nprobe >= 1); let meta_guard = relation.read(0); @@ -148,12 +147,7 @@ pub fn scan( ), ); for j in 0..32 { - if h0_tuple.mask[j] - && filters - .iter() - .copied() - .all(|x| Some(x) == h0_tuple.extra[j]) - { + if h0_tuple.mask[j] { results.push(( Reverse(lowerbounds[j]), AlwaysEqual(h0_tuple.mean[j]), diff --git a/src/algorithm/tuples.rs b/src/algorithm/tuples.rs index 9f6cdabb..aaf3380c 100644 --- a/src/algorithm/tuples.rs +++ b/src/algorithm/tuples.rs @@ -45,7 +45,6 @@ pub struct Height0Tuple { pub mean: [(u32, u16); 32], // for height 0 tuple, it's pointers to heap relation pub payload: [u64; 32], - pub extra: [Option; 32], // RaBitQ algoithm pub dis_u_2: [f32; 32], pub factor_ppc: [f32; 32], @@ -60,7 +59,6 @@ pub fn put( code: &rabitq::Code, vector: (u32, u16), payload: u64, - extra: Option, ) -> bool { // todo: use mutable api let mut x = rkyv::from_bytes::(bytes).expect("data corruption"); @@ -68,7 +66,6 @@ pub fn put( if !x.mask[j] { x.mean[j] = vector; x.payload[j] = payload; - x.extra[j] = extra; x.mask[j] = true; x.dis_u_2[j] = code.dis_u_2; x.factor_ppc[j] = code.factor_ppc; diff --git a/src/index/am.rs b/src/index/am.rs index 41c38cba..37d38181 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -54,7 +54,6 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = { // for vector index scans without `ORDER BY` clauses a large number // and throw errors if someone really wants such a path. am_routine.amoptionalkey = true; - am_routine.amcanmulticol = true; am_routine.amvalidate = Some(amvalidate); am_routine.amoptions = Some(amoptions); @@ -142,7 +141,7 @@ pub unsafe extern "C" fn ambuild( impl HeapRelation for Heap { fn traverse(&self, callback: F) where - F: FnMut((Pointer, Option, Vec)), + F: FnMut((Pointer, Vec)), { pub struct State<'a, F> { pub this: &'a Heap, @@ -150,14 +149,14 @@ pub unsafe extern "C" fn ambuild( } #[pgrx::pg_guard] unsafe extern "C" fn call( - index: pgrx::pg_sys::Relation, + _index: pgrx::pg_sys::Relation, ctid: pgrx::pg_sys::ItemPointer, values: *mut Datum, is_null: *mut bool, _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((Pointer, Option, Vec)), + F: FnMut((Pointer, Vec)), { pgrx::check_for_interrupts!(); use base::vector::OwnedVector; @@ -168,28 +167,15 @@ pub unsafe extern "C" fn ambuild( .opfamily .datum_to_vector(*values.add(0), *is_null.add(0)) }; + let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { - let pointer = unsafe { ctid_to_pointer(ctid.read()) }; - let extra = unsafe { - match (*(*index).rd_att).natts { - 1 => None, - 2 => { - if !is_null.add(1).read() { - Some(values.add(1).read().value() as u32) - } else { - None - } - } - _ => unreachable!(), - } - }; let vector = match vector { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), OwnedVector::BVector(_) => unreachable!(), }; - (state.callback)((pointer, extra, vector.into_vec())); + (state.callback)((pointer, vector.into_vec())); } } let table_am = unsafe { &*(*self.heap).rd_tableam }; @@ -271,32 +257,14 @@ pub unsafe extern "C" fn aminsert( let opfamily = unsafe { am_options::opfamily(index) }; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { - let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - let extra = unsafe { - match (*(*index).rd_att).natts { - 1 => None, - 2 => { - if !is_null.add(1).read() { - Some(values.add(1).read().value() as u32) - } else { - None - } - } - _ => unreachable!(), - } - }; let vector = match vector { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), OwnedVector::BVector(_) => unreachable!(), }; - algorithm::insert::insert( - unsafe { Relation::new(index) }, - pointer, - extra, - vector.into_vec(), - ); + let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); + algorithm::insert::insert(unsafe { Relation::new(index) }, pointer, vector.into_vec()); } false } @@ -311,7 +279,7 @@ pub unsafe extern "C" fn ambeginscan( let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index, n_keys, n_orderbys) }; unsafe { - let scanner = am_scan::scan_make(None, None, Vec::new(), false); + let scanner = am_scan::scan_make(None, None, false); (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); } scan @@ -333,10 +301,9 @@ pub unsafe extern "C" fn amrescan( std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); } let opfamily = am_options::opfamily((*scan).indexRelation); - let (orderbys, spheres, filters) = { + let (orderbys, spheres) = { let mut orderbys = Vec::new(); let mut spheres = Vec::new(); - let mut filters = Vec::new(); if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { pgrx::error!( "vector search with no WHERE clause and no ORDER BY clause is not supported" @@ -357,26 +324,14 @@ pub unsafe extern "C" fn amrescan( let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; match (*data).sk_strategy { 2 => spheres.push(opfamily.datum_to_sphere(value, is_null)), - 11 => { - if is_null { - filters.push(0); - filters.push(1); - } else { - filters.push(value.value() as u32); - } - } _ => unreachable!(), } } - (orderbys, spheres, filters) + (orderbys, spheres) }; - let (vector, threshold, filters, recheck) = - am_scan::scan_build(orderbys, spheres, filters, opfamily); + let (vector, threshold, recheck) = am_scan::scan_build(orderbys, spheres, opfamily); let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace( - scanner, - am_scan::scan_make(vector, threshold, filters, recheck), - ); + let scanner = std::mem::replace(scanner, am_scan::scan_make(vector, threshold, recheck)); am_scan::scan_release(scanner); } } @@ -416,7 +371,7 @@ pub unsafe extern "C" fn amgettuple( pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { unsafe { let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace(scanner, am_scan::scan_make(None, None, Vec::new(), false)); + let scanner = std::mem::replace(scanner, am_scan::scan_make(None, None, false)); am_scan::scan_release(scanner); } } diff --git a/src/index/am_options.rs b/src/index/am_options.rs index eb9aa88c..1a977e16 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -116,8 +116,8 @@ pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, Rabbitho if atts.is_empty() { pgrx::error!("indexing on no columns is not supported"); } - if atts.len() > 2 { - pgrx::error!("indexing on many columns is not supported"); + if atts.len() != 1 { + pgrx::error!("multicolumn index is not supported"); } // get dims let typmod = Typmod::parse_from_i32(atts[0].type_mod()).unwrap(); diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs index 62f6a0ad..9792729a 100644 --- a/src/index/am_scan.rs +++ b/src/index/am_scan.rs @@ -10,7 +10,6 @@ pub enum Scanner { Initial { vector: Option<(OwnedVector, Opfamily)>, threshold: Option, - filters: Vec, recheck: bool, }, Vbase { @@ -25,9 +24,8 @@ pub enum Scanner { pub fn scan_build( orderbys: Vec>, spheres: Vec<(Option, Option)>, - filters: Vec, opfamily: Opfamily, -) -> (Option<(OwnedVector, Opfamily)>, Option, Vec, bool) { +) -> (Option<(OwnedVector, Opfamily)>, Option, bool) { let mut pair = None; let mut threshold = None; let mut recheck = false; @@ -51,19 +49,17 @@ pub fn scan_build( break; } } - (pair.map(|x| (x, opfamily)), threshold, filters, recheck) + (pair.map(|x| (x, opfamily)), threshold, recheck) } pub fn scan_make( vector: Option<(OwnedVector, Opfamily)>, threshold: Option, - filters: Vec, recheck: bool, ) -> Scanner { Scanner::Initial { vector, threshold, - filters, recheck, } } @@ -72,7 +68,6 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, if let Scanner::Initial { vector, threshold, - filters, recheck, } = scanner { @@ -86,7 +81,6 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, OwnedVector::BVector(_) => unreachable!(), }, nprobe(), - filters.clone(), ); *scanner = Scanner::Vbase { vbase: Box::new(vbase), diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 0f659e29..96dfaaa6 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -27,7 +27,6 @@ COMMENT ON ACCESS METHOD rabbithole IS 'rabbithole index access method'; -- List of operator families CREATE OPERATOR FAMILY vector_l2_ops USING rabbithole; -CREATE OPERATOR FAMILY int4_eq_ops USING rabbithole; -- List of operator classes @@ -35,7 +34,3 @@ CREATE OPERATOR CLASS vector_l2_ops FOR TYPE vector USING rabbithole FAMILY vector_l2_ops AS OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH; - -CREATE OPERATOR CLASS int4_eq_ops - FOR TYPE int4 USING rabbithole FAMILY int4_eq_ops AS - OPERATOR 11 = (int4, int4) FOR SEARCH; From f12170f7e214b432f305cf9c454df33744c5464d Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 28 Oct 2024 16:19:04 +0800 Subject: [PATCH 012/324] feat: parallel build (#30) Signed-off-by: usamoi --- Cargo.toml | 2 +- src/algorithm/build.rs | 402 ++++++++++------------------------------ src/algorithm/insert.rs | 17 +- src/algorithm/tuples.rs | 5 +- src/index/am.rs | 347 +++++++++++++++++++++++++++++++++- src/postgres.rs | 56 ++++-- 6 files changed, 499 insertions(+), 330 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index b1397e74..5b88d474 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,7 +21,7 @@ pg16 = ["pgrx/pg16"] pg17 = ["pgrx/pg17"] [dependencies] -pgrx = { version = "=0.12.6", default-features = false, features = [] } +pgrx = { version = "=0.12.6", default-features = false, features = ["cshim"] } base = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } common = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } detect = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index 782cd345..f58c4496 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -5,8 +5,6 @@ use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; use crate::types::ExternalCentroids; use crate::types::RabbitholeIndexingOptions; -use base::always_equal::AlwaysEqual; -use base::distance::Distance; use base::distance::DistanceKind; use base::index::VectorOptions; use base::scalar::ScalarLike; @@ -14,8 +12,6 @@ use base::search::Pointer; use common::vec2::Vec2; use rand::Rng; use rkyv::ser::serializers::AllocSerializer; -use std::cmp::Reverse; -use std::collections::BinaryHeap; use std::marker::PhantomData; use std::sync::atomic::AtomicBool; use std::sync::Arc; @@ -71,92 +67,58 @@ pub fn build( Structure::compute(vector_options.clone(), rabbithole_options.clone(), samples) } }; - let mut vectors_head; - { - let h2_len = structure.h2_len(); - let h1_len = structure.h1_len(); - let mut meta = Tape::create(&relation); - assert_eq!(meta.first(), 0); - let mut vectors = Tape::create(&relation); - assert_eq!(vectors.first(), 1); - let h2_means = (0..h2_len) - .map(|i| { - vectors.push(&VectorTuple { - payload: None, - vector: structure.h2_means(i).clone(), - }) + let h2_len = structure.h2_len(); + let h1_len = structure.h1_len(); + let mut meta = Tape::create(&relation, false); + assert_eq!(meta.first(), 0); + let mut forwards = Tape::::create(&relation, false); + assert_eq!(forwards.first(), 1); + let mut vectors = Tape::create(&relation, true); + assert_eq!(vectors.first(), 2); + let h2_means = (0..h2_len) + .map(|i| { + vectors.push(&VectorTuple { + payload: None, + vector: structure.h2_means(i).clone(), }) - .collect::>(); - let h1_means = (0..h1_len) - .map(|i| { - vectors.push(&VectorTuple { - payload: None, - vector: structure.h1_means(i).clone(), - }) - }) - .collect::>(); - let h1_firsts = (0..h1_len) - .map(|_| { - let tape = Tape::::create(&relation); - tape.first() + }) + .collect::>(); + let h1_means = (0..h1_len) + .map(|i| { + vectors.push(&VectorTuple { + payload: None, + vector: structure.h1_means(i).clone(), }) - .collect::>(); - let h2_firsts = (0..h2_len) - .map(|i| { - let mut tape = Tape::::create(&relation); - let mut cache = Vec::new(); - let h2_mean = structure.h2_means(i); - let children = structure.h2_children(i); - for child in children.iter().copied() { - let h1_mean = structure.h1_means(child); - let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(h1_mean, h2_mean)) - } else { - rabitq::code(dims, h1_mean) - }; - cache.push((child, code)); - if cache.len() == 32 { - let group = std::mem::take(&mut cache); - let code = std::array::from_fn(|i| group[i].1.clone()); - let packed = rabitq::pack_codes(dims, code); - tape.push(&Height1Tuple { - mask: [true; 32], - mean: std::array::from_fn(|i| h1_means[group[i].0 as usize]), - first: std::array::from_fn(|i| h1_firsts[group[i].0 as usize]), - dis_u_2: packed.dis_u_2, - factor_ppc: packed.factor_ppc, - factor_ip: packed.factor_ip, - factor_err: packed.factor_err, - t: packed.t, - }); - } - } - if !cache.is_empty() { + }) + .collect::>(); + let h1_firsts = (0..h1_len) + .map(|_| { + let tape = Tape::::create(&relation, false); + tape.first() + }) + .collect::>(); + let h2_firsts = (0..h2_len) + .map(|i| { + let mut tape = Tape::::create(&relation, false); + let mut cache = Vec::new(); + let h2_mean = structure.h2_means(i); + let children = structure.h2_children(i); + for child in children.iter().copied() { + let h1_mean = structure.h1_means(child); + let code = if is_residual { + rabitq::code(dims, &f32::vector_sub(h1_mean, h2_mean)) + } else { + rabitq::code(dims, h1_mean) + }; + cache.push((child, code)); + if cache.len() == 32 { let group = std::mem::take(&mut cache); - let codes = std::array::from_fn(|i| { - if i < group.len() { - group[i].1.clone() - } else { - rabitq::dummy_code(dims) - } - }); - let packed = rabitq::pack_codes(dims, codes); + let code = std::array::from_fn(|i| group[i].1.clone()); + let packed = rabitq::pack_codes(dims, code); tape.push(&Height1Tuple { - mask: std::array::from_fn(|i| i < group.len()), - mean: std::array::from_fn(|i| { - if i < group.len() { - h1_means[group[i].0 as usize] - } else { - Default::default() - } - }), - first: std::array::from_fn(|i| { - if i < group.len() { - h1_firsts[group[i].0 as usize] - } else { - Default::default() - } - }), + mask: [true; 32], + mean: std::array::from_fn(|i| h1_means[group[i].0 as usize]), + first: std::array::from_fn(|i| h1_firsts[group[i].0 as usize]), dis_u_2: packed.dis_u_2, factor_ppc: packed.factor_ppc, factor_ip: packed.factor_ip, @@ -164,231 +126,51 @@ pub fn build( t: packed.t, }); } - tape.first() - }) - .collect::>(); - meta.push(&MetaTuple { - dims, - is_residual, - vectors_first: vectors.first(), - mean: h2_means[0], - first: h2_firsts[0], - }); - vectors_head = vectors.head.id(); - } - let mut tuples_done = 0; - heap_relation.traverse(|(payload, vector)| { - pgrx::check_for_interrupts!(); - tuples_done += 1; - reporter.tuples_done(tuples_done); - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get() - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); - let vector = rabitq::project(&vector); - let default_lut = if !is_residual { - Some(rabitq::fscan_preprocess(&vector)) - } else { - None - }; - let h0_vector = 'h0_vector: { - let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { - vector: vector.clone(), - payload: Some(payload.as_u64()), - }) - .unwrap(); - let mut write = relation.write(vectors_head); - if let Some(i) = write.get_mut().alloc(&tuple) { - break 'h0_vector (vectors_head, i); } - let mut extend = relation.extend(); - write.get_mut().get_opaque_mut().next = extend.id(); - if let Some(i) = extend.get_mut().alloc(&tuple) { - vectors_head = extend.id(); - break 'h0_vector (extend.id(), i); - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - }; - let h0_payload = payload.as_u64(); - let list = ( - meta_tuple.first, - if is_residual { - let vector_guard = relation.read(meta_tuple.mean.0); - let vector_tuple = vector_guard - .get() - .get(meta_tuple.mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - Some(vector_tuple.vector.to_vec()) - } else { - None - }, - ); - let list = { - let mut results = Vec::new(); - { - let lut = if is_residual { - &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) - } else { - default_lut.as_ref().unwrap() - }; - let mut current = list.0; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { - let h1_tuple = h1_guard - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let lowerbounds = rabitq::fscan_process_lowerbound( - dims, - lut, - ( - &h1_tuple.dis_u_2, - &h1_tuple.factor_ppc, - &h1_tuple.factor_ip, - &h1_tuple.factor_err, - &h1_tuple.t, - ), - ); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(( - Reverse(lowerbounds[j]), - AlwaysEqual(h1_tuple.mean[j]), - AlwaysEqual(h1_tuple.first[j]), - )); - } - } + if !cache.is_empty() { + let group = std::mem::take(&mut cache); + let codes = std::array::from_fn(|i| { + if i < group.len() { + group[i].1.clone() + } else { + rabitq::dummy_code(dims) } - current = h1_guard.get().get_opaque().next; - } - } - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get() - .get(mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dis_u = - Distance::from_f32(f32::reduce_sum_of_d2(&vector, &vector_tuple.vector)); - cache.push(( - Reverse(dis_u), - AlwaysEqual(first), - AlwaysEqual(if is_residual { - Some(vector_tuple.vector.to_vec()) + }); + let packed = rabitq::pack_codes(dims, codes); + tape.push(&Height1Tuple { + mask: std::array::from_fn(|i| i < group.len()), + mean: std::array::from_fn(|i| { + if i < group.len() { + h1_means[group[i].0 as usize] } else { - None - }), - )); - } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop().unwrap(); - (first, mean) + Default::default() + } + }), + first: std::array::from_fn(|i| { + if i < group.len() { + h1_firsts[group[i].0 as usize] + } else { + Default::default() + } + }), + dis_u_2: packed.dis_u_2, + factor_ppc: packed.factor_ppc, + factor_ip: packed.factor_ip, + factor_err: packed.factor_err, + t: packed.t, + }); } - }; - let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(&vector, list.1.as_ref().unwrap())) - } else { - rabitq::code(dims, &vector) - }; - let dummy = rkyv::to_bytes::<_, 8192>(&Height0Tuple { - mask: [false; 32], - mean: [(0, 0); 32], - payload: [0; 32], - dis_u_2: [0.0f32; 32], - factor_ppc: [0.0f32; 32], - factor_ip: [0.0f32; 32], - factor_err: [0.0f32; 32], - t: vec![0; (dims.div_ceil(4) * 16) as usize], + tape.first() }) - .unwrap(); - let first = list.0; - assert!(first != u32::MAX); - let mut current = first; - loop { - let read = relation.read(current); - let flag = 'flag: { - for i in 1..=read.get().len() { - let h0_tuple = read - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - if h0_tuple.mask.iter().any(|x| *x) { - break 'flag true; - } - } - if read.get().freespace() as usize >= dummy.len() { - break 'flag true; - } - if read.get().get_opaque().next == u32::MAX { - break 'flag true; - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current); - for i in 1..=write.get().len() { - let flag = put( - write.get_mut().get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - if flag { - return; - } - } - if let Some(i) = write.get_mut().alloc(&dummy) { - let flag = put( - write.get_mut().get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; - } - if write.get().get_opaque().next == u32::MAX { - let mut extend = relation.extend(); - write.get_mut().get_opaque_mut().next = extend.id(); - if let Some(i) = extend.get_mut().alloc(&dummy) { - let flag = put( - extend.get_mut().get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - } - current = write.get().get_opaque().next; - } else { - current = read.get().get_opaque().next; - } - } + .collect::>(); + forwards.head.get_mut().get_opaque_mut().fast_forward = vectors.first(); + meta.push(&MetaTuple { + dims, + is_residual, + vectors_first: vectors.first(), + forwards_first: forwards.first(), + mean: h2_means[0], + first: h2_firsts[0], }); } @@ -535,17 +317,19 @@ struct Tape<'a, T> { relation: &'a Relation, head: BufferWriteGuard, first: u32, + tracking_freespace: bool, _phantom: PhantomData T>, } impl<'a, T> Tape<'a, T> { - fn create(relation: &'a Relation) -> Self { - let head = relation.extend(); + fn create(relation: &'a Relation, tracking_freespace: bool) -> Self { + let head = relation.extend(tracking_freespace); let first = head.id(); Self { relation, head, first, + tracking_freespace, _phantom: PhantomData, } } @@ -563,7 +347,7 @@ where if let Some(i) = self.head.get_mut().alloc(&bytes) { (self.head.id(), i) } else { - let next = self.relation.extend(); + let next = self.relation.extend(self.tracking_freespace); self.head.get_mut().get_opaque_mut().next = next.id(); self.head = next; if let Some(i) = self.head.get_mut().alloc(&bytes) { diff --git a/src/algorithm/insert.rs b/src/algorithm/insert.rs index fff7678e..4e3a2577 100644 --- a/src/algorithm/insert.rs +++ b/src/algorithm/insert.rs @@ -25,13 +25,18 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { } else { None }; - let h0_vector = { + let h0_vector = 'h0_vector: { let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { vector: vector.clone(), payload: Some(payload.as_u64()), }) .unwrap(); - let mut current = meta_tuple.vectors_first; + if let Some(mut write) = relation.search(tuple.len()) { + let i = write.get_mut().alloc(&tuple).unwrap(); + break 'h0_vector (write.id(), i); + } + let mut current = relation.read(1).get().get_opaque().fast_forward; + let mut changed = false; loop { let read = relation.read(current); let flag = 'flag: { @@ -50,7 +55,10 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { break (current, i); } if write.get().get_opaque().next == u32::MAX { - let mut extend = relation.extend(); + if changed { + relation.write(1).get_mut().get_opaque_mut().fast_forward = write.id(); + } + let mut extend = relation.extend(true); write.get_mut().get_opaque_mut().next = extend.id(); if let Some(i) = extend.get_mut().alloc(&tuple) { break (extend.id(), i); @@ -62,6 +70,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { } else { current = read.get().get_opaque().next; } + changed = true; } }; let h0_payload = payload.as_u64(); @@ -218,7 +227,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { return; } if write.get().get_opaque().next == u32::MAX { - let mut extend = relation.extend(); + let mut extend = relation.extend(false); write.get_mut().get_opaque_mut().next = extend.id(); if let Some(i) = extend.get_mut().alloc(&dummy) { let flag = put( diff --git a/src/algorithm/tuples.rs b/src/algorithm/tuples.rs index aaf3380c..331d0d22 100644 --- a/src/algorithm/tuples.rs +++ b/src/algorithm/tuples.rs @@ -7,6 +7,7 @@ pub struct MetaTuple { pub dims: u32, pub is_residual: bool, pub vectors_first: u32, + pub forwards_first: u32, // raw vector pub mean: (u32, u16), // for meta tuple, it's pointers to next level @@ -29,7 +30,7 @@ pub struct Height1Tuple { pub mean: [(u32, u16); 32], // for height 1 tuple, it's pointers to next level pub first: [u32; 32], - // RaBitQ algoithm + // RaBitQ algorithm pub dis_u_2: [f32; 32], pub factor_ppc: [f32; 32], pub factor_ip: [f32; 32], @@ -45,7 +46,7 @@ pub struct Height0Tuple { pub mean: [(u32, u16); 32], // for height 0 tuple, it's pointers to heap relation pub payload: [u64; 32], - // RaBitQ algoithm + // RaBitQ algorithm pub dis_u_2: [f32; 32], pub factor_ppc: [f32; 32], pub factor_ip: [f32; 32], diff --git a/src/index/am.rs b/src/index/am.rs index 37d38181..c2b0d400 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -46,6 +46,11 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = { am_routine.amcanorderbyop = true; + #[cfg(feature = "pg17")] + { + am_routine.amcanbuildparallel = true; + } + // Index access methods that set `amoptionalkey` to `false` // must index all tuples, even if the first column is `NULL`. // However, PostgreSQL does not generate a path if there is no @@ -132,6 +137,7 @@ pub unsafe extern "C" fn ambuild( index: pgrx::pg_sys::Relation, index_info: *mut pgrx::pg_sys::IndexInfo, ) -> *mut pgrx::pg_sys::IndexBuildResult { + #[derive(Debug, Clone)] pub struct Heap { heap: pgrx::pg_sys::Relation, index: pgrx::pg_sys::Relation, @@ -200,6 +206,7 @@ pub unsafe extern "C" fn ambuild( } } } + #[derive(Debug, Clone)] pub struct PgReporter {} impl Reporter for PgReporter { fn tuples_total(&mut self, tuples_total: usize) { @@ -226,17 +233,351 @@ pub unsafe extern "C" fn ambuild( index_info, opfamily: unsafe { am_options::opfamily(index) }, }; + let mut reporter = PgReporter {}; let index_relation = unsafe { Relation::new(index) }; algorithm::build::build( vector_options, rabbithole_options, - heap_relation, - index_relation, - PgReporter {}, + heap_relation.clone(), + index_relation.clone(), + reporter.clone(), ); + if let Some(leader) = + unsafe { RabbitholeLeader::enter(heap, index, (*index_info).ii_Concurrent) } + { + unsafe { + let nparticipanttuplesorts = leader.nparticipanttuplesorts; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.rabbitholeshared).mutex); + if (*leader.rabbitholeshared).nparticipantsdone == nparticipanttuplesorts { + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.rabbitholeshared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.rabbitholeshared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*leader.rabbitholeshared).workersdonecv, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + } + } else { + let mut tuples_done = 0; + reporter.tuples_done(tuples_done); + heap_relation.traverse(|(payload, vector)| { + pgrx::check_for_interrupts!(); + algorithm::insert::insert(index_relation.clone(), payload, vector); + tuples_done += 1; + reporter.tuples_done(tuples_done); + }); + } unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } +struct RabbitholeShared { + /* Immutable state */ + heaprelid: pgrx::pg_sys::Oid, + indexrelid: pgrx::pg_sys::Oid, + isconcurrent: bool, + + /* Worker progress */ + workersdonecv: pgrx::pg_sys::ConditionVariable, + + /* Mutex for mutable state */ + mutex: pgrx::pg_sys::slock_t, + + /* Mutable state */ + nparticipantsdone: i32, +} + +fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { + matches!( + unsafe { (*snapshot).snapshot_type }, + pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + | pgrx::pg_sys::SnapshotType::SNAPSHOT_HISTORIC_MVCC + ) +} + +struct RabbitholeLeader { + pcxt: *mut pgrx::pg_sys::ParallelContext, + nparticipanttuplesorts: i32, + rabbitholeshared: *mut RabbitholeShared, + snapshot: pgrx::pg_sys::Snapshot, +} + +impl RabbitholeLeader { + pub unsafe fn enter( + heap: pgrx::pg_sys::Relation, + index: pgrx::pg_sys::Relation, + isconcurrent: bool, + ) -> Option { + unsafe fn compute_parallel_workers( + heap: pgrx::pg_sys::Relation, + index: pgrx::pg_sys::Relation, + ) -> i32 { + unsafe { + if pgrx::pg_sys::plan_create_index_workers((*heap).rd_id, (*index).rd_id) == 0 { + return 0; + } + if !(*heap).rd_options.is_null() { + let std_options = (*heap).rd_options.cast::(); + std::cmp::min( + (*std_options).parallel_workers, + pgrx::pg_sys::max_parallel_maintenance_workers, + ) + } else { + pgrx::pg_sys::max_parallel_maintenance_workers + } + } + } + + let request = unsafe { compute_parallel_workers(heap, index) }; + if request <= 0 { + return None; + } + + unsafe { + pgrx::pg_sys::EnterParallelMode(); + } + let pcxt = unsafe { + pgrx::pg_sys::CreateParallelContext( + c"rabbithole".as_ptr(), + c"rabbithole_parallel_build_main".as_ptr(), + request, + ) + }; + + let snapshot = if isconcurrent { + unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } + } else { + &raw mut pgrx::pg_sys::SnapshotAnyData + }; + + fn estimate_chunk(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.space_for_chunks += x.next_multiple_of(pgrx::pg_sys::ALIGNOF_BUFFER as _); + } + fn estimate_keys(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.number_of_keys += x; + } + let est_tablescandesc = + unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap, snapshot) }; + unsafe { + estimate_chunk(&mut (*pcxt).estimator, size_of::()); + estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); + estimate_keys(&mut (*pcxt).estimator, 1); + } + + unsafe { + pgrx::pg_sys::InitializeParallelDSM(pcxt); + if (*pcxt).seg.is_null() { + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + let rabbitholeshared = unsafe { + let rabbitholeshared = + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) + .cast::(); + rabbitholeshared.write(RabbitholeShared { + heaprelid: (*heap).rd_id, + indexrelid: (*index).rd_id, + isconcurrent, + workersdonecv: std::mem::zeroed(), + mutex: std::mem::zeroed(), + nparticipantsdone: 0, + }); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*rabbitholeshared).workersdonecv); + pgrx::pg_sys::SpinLockInit(&raw mut (*rabbitholeshared).mutex); + rabbitholeshared + }; + + let tablescandesc = unsafe { + let tablescandesc = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, est_tablescandesc) + .cast::(); + pgrx::pg_sys::table_parallelscan_initialize(heap, tablescandesc, snapshot); + tablescandesc + }; + + unsafe { + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000001, rabbitholeshared.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); + } + + unsafe { + pgrx::pg_sys::LaunchParallelWorkers(pcxt); + } + + let nparticipanttuplesorts = unsafe { (*pcxt).nworkers_launched }; + + unsafe { + if nparticipanttuplesorts == 0 { + pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToAttach(pcxt); + } + Some(Self { + pcxt, + nparticipanttuplesorts, + rabbitholeshared, + snapshot, + }) + } +} + +impl Drop for RabbitholeLeader { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); + if is_mvcc_snapshot(self.snapshot) { + pgrx::pg_sys::UnregisterSnapshot(self.snapshot); + } + pgrx::pg_sys::DestroyParallelContext(self.pcxt); + pgrx::pg_sys::ExitParallelMode(); + } + } +} + +#[pgrx::pg_guard] +#[no_mangle] +pub unsafe extern "C" fn rabbithole_parallel_build_main( + _seg: *mut pgrx::pg_sys::dsm_segment, + toc: *mut pgrx::pg_sys::shm_toc, +) { + let rabbitholeshared = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() + }; + let tablescandesc = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) + .cast::() + }; + let heap_lockmode; + let index_lockmode; + if unsafe { !(*rabbitholeshared).isconcurrent } { + heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; + index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; + } else { + heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; + index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; + } + let heap = unsafe { pgrx::pg_sys::table_open((*rabbitholeshared).heaprelid, heap_lockmode) }; + let index = unsafe { pgrx::pg_sys::index_open((*rabbitholeshared).indexrelid, index_lockmode) }; + let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; + unsafe { + (*index_info).ii_Concurrent = (*rabbitholeshared).isconcurrent; + } + + #[derive(Debug, Clone)] + pub struct Heap { + heap: pgrx::pg_sys::Relation, + index: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + opfamily: Opfamily, + scan: *mut pgrx::pg_sys::TableScanDescData, + } + impl HeapRelation for Heap { + fn traverse(&self, callback: F) + where + F: FnMut((Pointer, Vec)), + { + pub struct State<'a, F> { + pub this: &'a Heap, + pub callback: F, + } + #[pgrx::pg_guard] + unsafe extern "C" fn call( + _index: pgrx::pg_sys::Relation, + ctid: pgrx::pg_sys::ItemPointer, + values: *mut Datum, + is_null: *mut bool, + _tuple_is_alive: bool, + state: *mut core::ffi::c_void, + ) where + F: FnMut((Pointer, Vec)), + { + pgrx::check_for_interrupts!(); + use base::vector::OwnedVector; + let state = unsafe { &mut *state.cast::>() }; + let vector = unsafe { + state + .this + .opfamily + .datum_to_vector(*values.add(0), *is_null.add(0)) + }; + let pointer = unsafe { ctid_to_pointer(ctid.read()) }; + if let Some(vector) = vector { + let vector = match vector { + OwnedVector::Vecf32(x) => x, + OwnedVector::Vecf16(_) => unreachable!(), + OwnedVector::SVecf32(_) => unreachable!(), + OwnedVector::BVector(_) => unreachable!(), + }; + (state.callback)((pointer, vector.into_vec())); + } + } + let table_am = unsafe { &*(*self.heap).rd_tableam }; + let mut state = State { + this: self, + callback, + }; + unsafe { + table_am.index_build_range_scan.unwrap()( + self.heap, + self.index, + self.index_info, + true, + false, + false, + 0, + pgrx::pg_sys::InvalidBlockNumber, + Some(call::), + (&mut state) as *mut State as *mut _, + self.scan, + ); + } + } + } + + let index_relation = unsafe { Relation::new(index) }; + let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap, tablescandesc) }; + let heap_relation = Heap { + heap, + index, + index_info, + opfamily: unsafe { am_options::opfamily(index) }, + scan, + }; + heap_relation.traverse(|(payload, vector)| { + pgrx::check_for_interrupts!(); + algorithm::insert::insert(index_relation.clone(), payload, vector); + }); + + unsafe { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*rabbitholeshared).mutex); + (*rabbitholeshared).nparticipantsdone += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*rabbitholeshared).mutex); + pgrx::pg_sys::ConditionVariableSignal(&raw mut (*rabbitholeshared).workersdonecv); + } + + unsafe { + pgrx::pg_sys::index_close(index, index_lockmode); + pgrx::pg_sys::table_close(heap, heap_lockmode); + } +} + #[pgrx::pg_guard] pub unsafe extern "C" fn ambuildempty(_index: pgrx::pg_sys::Relation) { pgrx::error!("Unlogged indexes are not supported."); diff --git a/src/postgres.rs b/src/postgres.rs index 4597931b..80d9e717 100644 --- a/src/postgres.rs +++ b/src/postgres.rs @@ -24,14 +24,18 @@ const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usi const _: () = assert!(size_of::() == pgrx::pg_sys::BLCKSZ as usize); impl Page { - pub fn init_mut(this: &mut MaybeUninit) -> &mut Self { + pub fn init_mut(this: &mut MaybeUninit, tracking_freespace: bool) -> &mut Self { unsafe { pgrx::pg_sys::PageInit( this.as_mut_ptr() as pgrx::pg_sys::Page, pgrx::pg_sys::BLCKSZ as usize, size_of::(), ); - (&raw mut (*this.as_mut_ptr()).opaque).write(Opaque::default()); + (&raw mut (*this.as_mut_ptr()).opaque).write(Opaque { + next: u32::MAX, + tracking_freespace, + fast_forward: u32::MAX, + }); } let this = unsafe { MaybeUninit::assume_init_mut(this) }; assert_eq!(offset_of!(Self, opaque), this.header.pd_special as usize); @@ -138,12 +142,8 @@ impl Page { #[repr(C, align(8))] pub struct Opaque { pub next: u32, -} - -impl Default for Opaque { - fn default() -> Self { - Self { next: u32::MAX } - } + pub tracking_freespace: bool, + pub fast_forward: u32, } const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); @@ -173,6 +173,7 @@ impl Drop for BufferReadGuard { } pub struct BufferWriteGuard { + raw: pgrx::pg_sys::Relation, buf: i32, page: NonNull, state: *mut pgrx::pg_sys::GenericXLogState, @@ -197,6 +198,14 @@ impl Drop for BufferWriteGuard { if std::thread::panicking() { pgrx::pg_sys::GenericXLogAbort(self.state); } else { + if self.get().get_opaque().tracking_freespace { + pgrx::pg_sys::RecordPageWithFreeSpace( + self.raw, + self.id, + self.get().freespace() as _, + ); + pgrx::pg_sys::FreeSpaceMapVacuumRange(self.raw, self.id, self.id + 1); + } pgrx::pg_sys::GenericXLogFinish(self.state); } pgrx::pg_sys::UnlockReleaseBuffer(self.buf); @@ -204,7 +213,7 @@ impl Drop for BufferWriteGuard { } } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct Relation { raw: pgrx::pg_sys::Relation, } @@ -253,6 +262,7 @@ impl Relation { ) .expect("failed to get page"); BufferWriteGuard { + raw: self.raw, buf, page: page.cast(), state, @@ -260,7 +270,7 @@ impl Relation { } } } - pub fn extend(&self) -> BufferWriteGuard { + pub fn extend(&self, tracking_freespace: bool) -> BufferWriteGuard { unsafe { use pgrx::pg_sys::{ ExclusiveLock, ForkNumber, GenericXLogRegisterBuffer, GenericXLogStart, LockBuffer, @@ -283,8 +293,9 @@ impl Relation { .cast::>(), ) .expect("failed to get page"); - Page::init_mut(page.as_mut()); + Page::init_mut(page.as_mut(), tracking_freespace); BufferWriteGuard { + raw: self.raw, buf, page: page.cast(), state, @@ -292,4 +303,27 @@ impl Relation { } } } + pub fn search(&self, freespace: usize) -> Option { + unsafe { + loop { + let id = pgrx::pg_sys::GetPageWithFreeSpace(self.raw, freespace); + if id == u32::MAX { + return None; + } + let write = self.write(id); + assert!(write.get().get_opaque().tracking_freespace); + if write.get().freespace() < freespace as _ { + // the free space is recorded incorrectly + pgrx::pg_sys::RecordPageWithFreeSpace( + self.raw, + id, + write.get().freespace() as _, + ); + pgrx::pg_sys::FreeSpaceMapVacuumRange(self.raw, id, id + 1); + continue; + } + return Some(write); + } + } + } } From 157d11b52607411b9b3814aa0768e30a84350ddd Mon Sep 17 00:00:00 2001 From: Keming Date: Mon, 28 Oct 2024 16:51:46 +0800 Subject: [PATCH 013/324] chore: use https to avoid git auth for downloading dependencies (#23) Signed-off-by: Keming Co-authored-by: usamoi --- Cargo.toml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 5b88d474..8b8c003b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,11 +22,11 @@ pg17 = ["pgrx/pg17"] [dependencies] pgrx = { version = "=0.12.6", default-features = false, features = ["cshim"] } -base = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -common = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -detect = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -k_means = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -quantization = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } +base = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } +common = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } +detect = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } +k_means = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } +quantization = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } paste = "1" serde = "1" toml = "0.8.19" From 245b68b708184b1879de7505c2c119ad92792b11 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 28 Oct 2024 18:23:29 +0800 Subject: [PATCH 014/324] fix: add vacuum_delay_point (#31) Signed-off-by: usamoi --- Cargo.lock | 14 +++++++------- src/algorithm/vacuum.rs | 4 +++- src/index/am.rs | 8 +++++++- 3 files changed, 17 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 10aa9443..8252a07b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -66,7 +66,7 @@ checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" [[package]] name = "base" version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" dependencies = [ "base_macros", "detect", @@ -84,7 +84,7 @@ dependencies = [ [[package]] name = "base_macros" version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" dependencies = [ "proc-macro2", "quote", @@ -226,7 +226,7 @@ dependencies = [ [[package]] name = "common" version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" dependencies = [ "base", "log", @@ -315,7 +315,7 @@ dependencies = [ [[package]] name = "detect" version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" dependencies = [ "detect_macros", ] @@ -323,7 +323,7 @@ dependencies = [ [[package]] name = "detect_macros" version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" dependencies = [ "proc-macro2", "quote", @@ -561,7 +561,7 @@ checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" [[package]] name = "k_means" version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" dependencies = [ "base", "common", @@ -943,7 +943,7 @@ dependencies = [ [[package]] name = "quantization" version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" dependencies = [ "base", "common", diff --git a/src/algorithm/vacuum.rs b/src/algorithm/vacuum.rs index 0dd52093..bddd6267 100644 --- a/src/algorithm/vacuum.rs +++ b/src/algorithm/vacuum.rs @@ -2,7 +2,7 @@ use crate::algorithm::tuples::*; use crate::postgres::Relation; use base::search::Pointer; -pub fn vacuum(relation: Relation, callback: impl Fn(Pointer) -> bool) { +pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) -> bool) { // step 1: vacuum height_0_tuple { let h1_firsts = { @@ -42,6 +42,7 @@ pub fn vacuum(relation: Relation, callback: impl Fn(Pointer) -> bool) { for first in h0_firsts { let mut current = first; while current != u32::MAX { + delay(); let mut h0_guard = relation.write(current); for i in 1..=h0_guard.get().len() { let h0_tuple = h0_guard @@ -97,6 +98,7 @@ pub fn vacuum(relation: Relation, callback: impl Fn(Pointer) -> bool) { meta_tuple.vectors_first }; while current != u32::MAX { + delay(); let read = relation.read(current); let flag = 'flag: { for i in 1..=read.get().len() { diff --git a/src/index/am.rs b/src/index/am.rs index c2b0d400..59d33ddc 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -732,7 +732,13 @@ pub unsafe extern "C" fn ambulkdelete( } let callback = callback.unwrap(); let callback = |p: Pointer| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; - algorithm::vacuum::vacuum(unsafe { Relation::new((*info).index) }, callback); + algorithm::vacuum::vacuum( + unsafe { Relation::new((*info).index) }, + || unsafe { + pgrx::pg_sys::vacuum_delay_point(); + }, + callback, + ); stats } From 7e4eea5f19fd37f202a301a4c120a2bdacd87c69 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Tue, 29 Oct 2024 11:42:51 +0800 Subject: [PATCH 015/324] progresss bar of index build (#32) Signed-off-by: cutecutecat --- bench/bench.py | 12 ++++- bench/index.py | 140 +++++++++++++++++++++++++++++++++++-------------- 2 files changed, 111 insertions(+), 41 deletions(-) diff --git a/bench/bench.py b/bench/bench.py index 0271004d..26423b5d 100644 --- a/bench/bench.py +++ b/bench/bench.py @@ -14,7 +14,7 @@ def build_arg_parse(): "-m", "--metric", help="Metric to pick, in l2 or cos", - choices=["l2", "cos"], + choices=["l2", "cos", "dot"], default="l2", ) parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) @@ -74,7 +74,15 @@ def bench(name, test, answer, metric_ops, conn): dataset = h5py.File(Path(args.input), "r") test = dataset["test"][:] answer = dataset["neighbors"][:] - metric_ops = "<->" if args.metric == "l2" else "<=>" + + if args.metric == "l2": + metric_ops = "<->" + elif args.metric == "cos": + metric_ops = "<=>" + elif args.metric == "dot": + metric_ops = "<#>" + else: + raise ValueError conn = create_connection(args.password) conn.execute("SET rabbithole.nprobe=300") diff --git a/bench/index.py b/bench/index.py index ac32ae7e..56a4527d 100644 --- a/bench/index.py +++ b/bench/index.py @@ -1,26 +1,39 @@ +import asyncio from sys import version_info from time import perf_counter import argparse from pathlib import Path +import multiprocessing if version_info >= (3, 12): raise RuntimeError("h5py doesn't support 3.12") import psycopg import h5py -from pgvector.psycopg import register_vector +from pgvector.psycopg import register_vector_async import numpy as np from tqdm import tqdm +KEEPALIVE_KWARGS = { + "keepalives": 1, + "keepalives_idle": 30, + "keepalives_interval": 5, + "keepalives_count": 5, +} + def build_arg_parse(): parser = argparse.ArgumentParser(description="Build index with K-means centroids") parser.add_argument( - "-m", "--metric", help="Distance metric", default="l2", choices=["l2", "cos"] + "-m", + "--metric", + help="Distance metric", + default="l2", + choices=["l2", "cos", "dot"], ) parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) parser.add_argument( - "-c", "--centroid", help="K-means centroids file", required=True + "-c", "--centroids", help="K-means centroids file", required=True ) parser.add_argument("-i", "--input", help="Input filepath", required=True) parser.add_argument( @@ -28,6 +41,14 @@ def build_arg_parse(): ) parser.add_argument("-k", help="Number of centroids", type=int, required=True) parser.add_argument("-d", "--dim", help="Dimension", type=int, required=True) + parser.add_argument( + "-w", + "--workers", + help="Workers to build index", + type=int, + required=False, + default=max(multiprocessing.cpu_count() - 2, 1), + ) return parser @@ -51,6 +72,13 @@ def get_ivf_ops_config(metric, k, name=None): residual_quantization = false spherical_centroids = true """ + elif metric == "dot": + metric_ops = "vector_dot_ops" + ivf_config = f""" + nlist = {k} + residual_quantization = false + spherical_centroids = true + """ else: raise ValueError @@ -59,79 +87,113 @@ def get_ivf_ops_config(metric, k, name=None): return metric_ops, ivf_config -def create_connection(password): - keepalive_kwargs = { - "keepalives": 1, - "keepalives_idle": 30, - "keepalives_interval": 5, - "keepalives_count": 5, - } - conn = psycopg.connect( +async def create_connection(password): + conn = await psycopg.AsyncConnection.connect( conninfo=f"postgresql://postgres:{password}@localhost:5432/postgres", dbname="postgres", autocommit=True, - **keepalive_kwargs, + **KEEPALIVE_KWARGS, ) - conn.execute("SET search_path TO public, vectors, rabbithole") - conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") - register_vector(conn) + await conn.execute("SET search_path TO public, vectors, rabbithole") + await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + await conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") + await register_vector_async(conn) return conn -def add_centroids(conn, name, centroids): +async def add_centroids(conn, name, centroids): dim = centroids.shape[1] - conn.execute(f"DROP TABLE IF EXISTS public.{name}_centroids") - conn.execute(f"CREATE TABLE public.{name}_centroids (coordinate vector({dim}))") - with conn.cursor().copy( + await conn.execute(f"DROP TABLE IF EXISTS public.{name}_centroids") + await conn.execute( + f"CREATE TABLE public.{name}_centroids (coordinate vector({dim}))" + ) + async with conn.cursor().copy( f"COPY public.{name}_centroids (coordinate) FROM STDIN WITH (FORMAT BINARY)" ) as copy: copy.set_types(["vector"]) for centroid in tqdm(centroids, desc="Adding centroids"): - copy.write_row((centroid,)) + await copy.write_row((centroid,)) while conn.pgconn.flush() == 1: pass -def build_index(conn, name, metric_ops, ivf_config, dim, train): - conn.execute(f"DROP TABLE IF EXISTS {name}") - conn.execute(f"CREATE TABLE {name} (id integer, embedding vector({dim}))") +async def add_embeddings(conn, name, dim, train): + await conn.execute(f"DROP TABLE IF EXISTS {name}") + await conn.execute(f"CREATE TABLE {name} (id integer, embedding vector({dim}))") - with conn.cursor().copy( + async with conn.cursor().copy( f"COPY {name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)" ) as copy: copy.set_types(["integer", "vector"]) - for i, vec in tqdm(enumerate(train), desc="Adding embeddings"): - copy.write_row((i, vec)) + for i, vec in tqdm( + enumerate(train), desc="Adding embeddings", total=len(train) + ): + await copy.write_row((i, vec)) while conn.pgconn.flush() == 1: pass + +async def build_index( + conn, name, workers, metric_ops, ivf_config, finish: asyncio.Event +): start_time = perf_counter() - conn.execute( + await conn.execute(f"SET max_parallel_maintenance_workers TO {workers}") + await conn.execute( f"CREATE INDEX ON {name} USING rabbithole (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" ) print(f"Index build time: {perf_counter() - start_time:.2f}s") + finish.set() -if __name__ == "__main__": - parser = build_arg_parse() - args = parser.parse_args() - print(args) - +async def monitor_index_build(password, finish: asyncio.Event): + conn = await psycopg.AsyncConnection.connect( + conninfo=f"postgresql://postgres:{password}@localhost:5432/postgres", + dbname="postgres", + autocommit=True, + **KEEPALIVE_KWARGS, + ) + pbar = tqdm(smoothing=0.0) + async with conn.cursor() as acur: + while True: + if finish.is_set(): + return + await acur.execute(f"SELECT tuples_done FROM pg_stat_progress_create_index") + tuples_done = await acur.fetchone() + update = 0 if tuples_done is None else tuples_done[0] + pbar.update(update - pbar.n) + await asyncio.sleep(1) + + +async def main(dataset): dataset = h5py.File(Path(args.input), "r") - conn = create_connection(args.password) + conn = await create_connection(args.password) if args.centroids: - centroids = np.load(args.centroid, allow_pickle=False) - add_centroids(conn, args.name, centroids) + centroids = np.load(args.centroids, allow_pickle=False) + await add_centroids(conn, args.name, centroids) metric_ops, ivf_config = get_ivf_ops_config( args.metric, args.k, args.name if args.centroids else None ) - build_index( + await add_embeddings(conn, args.name, args.dim, dataset["train"]) + + index_finish = asyncio.Event() + index_task = build_index( conn, args.name, + args.workers, metric_ops, ivf_config, - args.dim, - dataset["train"], + index_finish, + ) + monitor_task = monitor_index_build( + args.password, + index_finish, ) + await asyncio.gather(index_task, monitor_task) + + +if __name__ == "__main__": + parser = build_arg_parse() + args = parser.parse_args() + print(args) + asyncio.run(main(args)) From 8462aa5f1762352d2b971c8b67021c548db70891 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Thu, 31 Oct 2024 11:02:46 +0800 Subject: [PATCH 016/324] feat: cos metric (#24) * feat: cos metric Signed-off-by: cutecutecat * fix by comments Signed-off-by: cutecutecat * fix by comments Signed-off-by: cutecutecat --------- Signed-off-by: cutecutecat --- bench/README.md | 4 +-- docker/Dockerfile | 2 +- src/algorithm/build.rs | 43 +++++++++++++++++++---- src/algorithm/insert.rs | 11 +++--- src/algorithm/rabitq.rs | 29 ++++++++++----- src/algorithm/scan.rs | 18 ++++++---- src/datatype/operators_pgvector_vector.rs | 43 +++++++++++++++++++++++ src/index/am.rs | 27 +++++++++++--- src/index/am_options.rs | 23 +++++++++--- src/index/am_scan.rs | 1 + src/index/utils.rs | 25 +++++++++---- src/sql/finalize.sql | 26 ++++++++++++++ 12 files changed, 207 insertions(+), 45 deletions(-) diff --git a/bench/README.md b/bench/README.md index 8d17e28b..c0ae63c0 100644 --- a/bench/README.md +++ b/bench/README.md @@ -21,7 +21,7 @@ Or you can use `starkind/rabbithole:pg16-latest` to run the bench. ## Run Instance ```shell -docker run --name rabbithole -e POSTGRES_PASSWORD=123 -p 5432:5432 -d starkind/rabbithole:pg16-latest +docker run --name rabbithole -e POSTGRES_PASSWORD=123 -p 5432:5432 -d rabbithole:pg16-latest PGPASSWORD=123 psql -h 127.0.0.1 -U postgres -c "CREATE USER bench WITH PASSWORD '123';" PGPASSWORD=123 psql -h 127.0.0.1 -U postgres -c "ALTER ROLE bench SUPERUSER;" @@ -54,7 +54,7 @@ Options for `-n`: - cohere_10m_23 ```shell -# pip install pgvector numpy faiss-cpu psycopg h5py +# pip install pgvector numpy faiss-cpu psycopg h5py tqdm # dump table embedding column to a local h5 file["train"] python dump.py -n sift -o sift.h5 -c embedding -d 128 diff --git a/docker/Dockerfile b/docker/Dockerfile index 0cbfeae5..cf3d63f1 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -3,4 +3,4 @@ FROM pgvector/pgvector:0.7.4-pg16 COPY ./build/rabbithole-pg16_0.0.0_amd64.deb /tmp/rabbithole.deb RUN apt-get install -y /tmp/rabbithole.deb && rm -f /tmp/rabbithole.deb -CMD ["postgres", "-c" ,"shared_preload_libraries=vector,rabbithole.so", "-c", "search_path=\"$user\", public, rabbithole", "-c", "logging_collector=on"] \ No newline at end of file +CMD ["postgres", "-c" ,"shared_preload_libraries=vector,rabbithole.so", "-c", "search_path=\"$user\", public, rabbithole"] \ No newline at end of file diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index f58c4496..e175ee19 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -1,6 +1,7 @@ use crate::algorithm::rabitq; use crate::algorithm::tuples::*; -use crate::index::utils::load_proj_vectors; +use crate::index::am_options::PgDistanceKind; +use crate::index::utils::load_table_vectors; use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; use crate::types::ExternalCentroids; @@ -9,6 +10,8 @@ use base::distance::DistanceKind; use base::index::VectorOptions; use base::scalar::ScalarLike; use base::search::Pointer; +use base::vector::VectBorrowed; +use base::vector::VectorBorrowed; use common::vec2::Vec2; use rand::Rng; use rkyv::ser::serializers::AllocSerializer; @@ -32,13 +35,18 @@ pub fn build( rabbithole_options: RabbitholeIndexingOptions, heap_relation: T, relation: Relation, + pg_distance: PgDistanceKind, mut reporter: R, ) { let dims = vector_options.dims; let is_residual = rabbithole_options.residual_quantization && vector_options.d == DistanceKind::L2; let structure = match &rabbithole_options.external_centroids { - Some(_) => Structure::load(vector_options.clone(), rabbithole_options.clone()), + Some(_) => Structure::load( + vector_options.clone(), + rabbithole_options.clone(), + pg_distance, + ), None => { let mut tuples_total = 0_usize; let samples = { @@ -228,14 +236,35 @@ impl Structure { h1_children: (0..rabbithole_options.nlist).map(|_| Vec::new()).collect(), } } - fn load(vector_options: VectorOptions, rabbithole_options: RabbitholeIndexingOptions) -> Self { + fn load( + vector_options: VectorOptions, + rabbithole_options: RabbitholeIndexingOptions, + pg_distance: PgDistanceKind, + ) -> Self { let dims = vector_options.dims; + let preprocess_data = match pg_distance { + PgDistanceKind::L2 | PgDistanceKind::Dot => { + |b: VectBorrowed| rabitq::project(b.slice()) + } + PgDistanceKind::Cos => { + |b: VectBorrowed| rabitq::project(b.function_normalize().slice()) + } + }; + let preprocess_index = |b: VectBorrowed| b.slice().to_vec(); + let h1_means = match &rabbithole_options.external_centroids { Some(ExternalCentroids { table, h1_means_column: h1, .. - }) => load_proj_vectors(table, h1, rabbithole_options.nlist, vector_options.dims), + }) => load_table_vectors( + table, + h1, + rabbithole_options.nlist, + vector_options.dims, + preprocess_data, + ), + _ => unreachable!(), }; let h1_children = match &rabbithole_options.external_centroids { @@ -243,7 +272,7 @@ impl Structure { table, h1_children_column: Some(h1), .. - }) => load_proj_vectors(table, h1, 1, vector_options.dims) + }) => load_table_vectors(table, h1, 1, vector_options.dims, preprocess_index) .into_iter() .map(|v| v.into_iter().map(|f| f as u32).collect()) .collect(), @@ -254,7 +283,7 @@ impl Structure { table, h2_mean_column: Some(h2), .. - }) => load_proj_vectors(table, h2, 1, vector_options.dims) + }) => load_table_vectors(table, h2, 1, vector_options.dims, preprocess_data) .pop() .expect("load h2_mean panic"), _ => { @@ -275,7 +304,7 @@ impl Structure { table, h2_children_column: Some(h2), .. - }) => load_proj_vectors(table, h2, 1, vector_options.dims) + }) => load_table_vectors(table, h2, 1, vector_options.dims, preprocess_index) .pop() .expect("load h2_children panic") .into_iter() diff --git a/src/algorithm/insert.rs b/src/algorithm/insert.rs index 4e3a2577..ed9bdb6b 100644 --- a/src/algorithm/insert.rs +++ b/src/algorithm/insert.rs @@ -1,14 +1,17 @@ use crate::algorithm::rabitq; +use crate::algorithm::rabitq::fscan_process_lowerbound; use crate::algorithm::tuples::*; +use crate::index::utils::distance; use crate::postgres::Relation; use base::always_equal::AlwaysEqual; use base::distance::Distance; +use base::distance::DistanceKind; use base::scalar::ScalarLike; use base::search::Pointer; use std::cmp::Reverse; use std::collections::BinaryHeap; -pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { +pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_kind: DistanceKind) { let meta_guard = relation.read(0); let meta_tuple = meta_guard .get() @@ -107,7 +110,8 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - let lowerbounds = rabitq::fscan_process_lowerbound( + let lowerbounds = fscan_process_lowerbound( + distance_kind, dims, lut, ( @@ -143,8 +147,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - let dis_u = - Distance::from_f32(f32::reduce_sum_of_d2(&vector, &vector_tuple.vector)); + let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); cache.push(( Reverse(dis_u), AlwaysEqual(first), diff --git a/src/algorithm/rabitq.rs b/src/algorithm/rabitq.rs index e3c5b002..ead21846 100644 --- a/src/algorithm/rabitq.rs +++ b/src/algorithm/rabitq.rs @@ -1,4 +1,4 @@ -use base::distance::Distance; +use base::distance::{Distance, DistanceKind}; use base::scalar::ScalarLike; use nalgebra::DMatrix; use quantization::utils::InfiniteByteChunks; @@ -168,6 +168,7 @@ pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { } pub fn fscan_process_lowerbound( + distance_kind: DistanceKind, dims: u32, lut: &(f32, f32, f32, f32, Vec), (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( @@ -180,14 +181,24 @@ pub fn fscan_process_lowerbound( ) -> [Distance; 32] { let &(dis_v_2, b, k, qvector_sum, ref s) = lut; let r = quantization::fast_scan::b4::fast_scan_b4(dims.div_ceil(4), t, s); - std::array::from_fn(|i| { - let rough = dis_u_2[i] - + dis_v_2 - + b * factor_ppc[i] - + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - 1.9 * err) - }) + match distance_kind { + DistanceKind::L2 => std::array::from_fn(|i| { + let rough = dis_u_2[i] + + dis_v_2 + + b * factor_ppc[i] + + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; + let err = factor_err[i] * dis_v_2.sqrt(); + Distance::from_f32(rough - 1.9 * err) + }), + DistanceKind::Dot => std::array::from_fn(|i| { + let rough = 0.5 * b * factor_ppc[i] + + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; + let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); + Distance::from_f32(rough - 1.9 * err) + }), + DistanceKind::Hamming => unreachable!(), + DistanceKind::Jaccard => unreachable!(), + } } fn compress(mut qvector: Vec) -> Vec { diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs index fc1d3718..dec0cbeb 100644 --- a/src/algorithm/scan.rs +++ b/src/algorithm/scan.rs @@ -1,8 +1,11 @@ use crate::algorithm::rabitq; +use crate::algorithm::rabitq::fscan_process_lowerbound; use crate::algorithm::tuples::*; +use crate::index::utils::distance; use crate::postgres::Relation; use base::always_equal::AlwaysEqual; use base::distance::Distance; +use base::distance::DistanceKind; use base::scalar::ScalarLike; use base::search::Pointer; use std::cmp::Reverse; @@ -11,6 +14,7 @@ use std::collections::BinaryHeap; pub fn scan( relation: Relation, vector: Vec, + distance_kind: DistanceKind, h1_nprobe: u32, ) -> impl Iterator { assert!(h1_nprobe >= 1); @@ -24,7 +28,7 @@ pub fn scan( let dims = meta_tuple.dims; assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); let vector = rabitq::project(&vector); - let is_residual = meta_tuple.is_residual; + let is_residual = meta_tuple.is_residual && distance_kind == DistanceKind::L2; let default_lut = if !is_residual { Some(rabitq::fscan_preprocess(&vector)) } else { @@ -63,7 +67,8 @@ pub fn scan( .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - let lowerbounds = rabitq::fscan_process_lowerbound( + let lowerbounds = fscan_process_lowerbound( + distance_kind, dims, lut, ( @@ -99,8 +104,7 @@ pub fn scan( .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - let dis_u = - Distance::from_f32(f32::reduce_sum_of_d2(&vector, &vector_tuple.vector)); + let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); cache.push(( Reverse(dis_u), AlwaysEqual(first), @@ -135,7 +139,8 @@ pub fn scan( .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - let lowerbounds = rabitq::fscan_process_lowerbound( + let lowerbounds = fscan_process_lowerbound( + distance_kind, dims, lut, ( @@ -175,8 +180,7 @@ pub fn scan( // fails consistency check continue; } - let dis_u = - Distance::from_f32(f32::reduce_sum_of_d2(&vector, &vector_tuple.vector)); + let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); } let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; diff --git a/src/datatype/operators_pgvector_vector.rs b/src/datatype/operators_pgvector_vector.rs index 0a77d24d..44e17b80 100644 --- a/src/datatype/operators_pgvector_vector.rs +++ b/src/datatype/operators_pgvector_vector.rs @@ -1,5 +1,6 @@ use crate::datatype::memory_pgvector_vector::*; use base::scalar::ScalarLike; +use base::vector::{VectBorrowed, VectorBorrowed}; use std::num::NonZero; #[pgrx::pg_extern(immutable, strict, parallel_safe)] @@ -22,3 +23,45 @@ fn _rabbithole_pgvector_vector_sphere_l2_in( }; f32::reduce_sum_of_d2(lhs.slice(), center.slice()).to_f32() < radius } + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _rabbithole_pgvector_vector_sphere_dot_in( + lhs: PgvectorVectorInput<'_>, + rhs: pgrx::composite_type!("sphere_vector"), +) -> bool { + let center: PgvectorVectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + -f32::reduce_sum_of_xy(lhs.slice(), center.slice()) < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _rabbithole_pgvector_vector_sphere_cos_in( + lhs: PgvectorVectorInput<'_>, + rhs: pgrx::composite_type!("sphere_vector"), +) -> bool { + let center: PgvectorVectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + VectBorrowed::operator_cos(lhs.as_borrowed(), center.as_borrowed()).to_f32() < radius +} diff --git a/src/index/am.rs b/src/index/am.rs index 59d33ddc..d6367dc3 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -226,7 +226,7 @@ pub unsafe extern "C" fn ambuild( } } } - let (vector_options, rabbithole_options) = unsafe { am_options::options(index) }; + let (vector_options, rabbithole_options, pg_distance) = unsafe { am_options::options(index) }; let heap_relation = Heap { heap, index, @@ -240,6 +240,7 @@ pub unsafe extern "C" fn ambuild( rabbithole_options, heap_relation.clone(), index_relation.clone(), + pg_distance, reporter.clone(), ); if let Some(leader) = @@ -266,7 +267,12 @@ pub unsafe extern "C" fn ambuild( reporter.tuples_done(tuples_done); heap_relation.traverse(|(payload, vector)| { pgrx::check_for_interrupts!(); - algorithm::insert::insert(index_relation.clone(), payload, vector); + algorithm::insert::insert( + index_relation.clone(), + payload, + vector, + pg_distance.to_distance(), + ); tuples_done += 1; reporter.tuples_done(tuples_done); }); @@ -553,16 +559,22 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( let index_relation = unsafe { Relation::new(index) }; let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap, tablescandesc) }; + let opfamily = unsafe { am_options::opfamily(index) }; let heap_relation = Heap { heap, index, index_info, - opfamily: unsafe { am_options::opfamily(index) }, + opfamily, scan, }; heap_relation.traverse(|(payload, vector)| { pgrx::check_for_interrupts!(); - algorithm::insert::insert(index_relation.clone(), payload, vector); + algorithm::insert::insert( + index_relation.clone(), + payload, + vector, + opfamily.distance_kind(), + ); }); unsafe { @@ -605,7 +617,12 @@ pub unsafe extern "C" fn aminsert( OwnedVector::BVector(_) => unreachable!(), }; let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - algorithm::insert::insert(unsafe { Relation::new(index) }, pointer, vector.into_vec()); + algorithm::insert::insert( + unsafe { Relation::new(index) }, + pointer, + vector.into_vec(), + opfamily.distance_kind(), + ); } false } diff --git a/src/index/am_options.rs b/src/index/am_options.rs index 1a977e16..10d5da9d 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -38,12 +38,15 @@ impl Reloption { #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum PgDistanceKind { L2, + Dot, + Cos, } impl PgDistanceKind { pub fn to_distance(self) -> DistanceKind { match self { PgDistanceKind::L2 => DistanceKind::L2, + PgDistanceKind::Dot | PgDistanceKind::Cos => DistanceKind::Dot, } } } @@ -85,6 +88,8 @@ pub fn convert_opfamily_to_vd( fn convert_name_to_vd(name: &str) -> Option<(VectorKind, PgDistanceKind)> { match name.strip_suffix("_ops") { Some("vector_l2") => Some((VectorKind::Vecf32, PgDistanceKind::L2)), + Some("vector_dot") => Some((VectorKind::Vecf32, PgDistanceKind::Dot)), + Some("vector_cos") => Some((VectorKind::Vecf32, PgDistanceKind::Cos)), _ => None, } } @@ -109,7 +114,9 @@ unsafe fn convert_reloptions_to_options( } } -pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, RabbitholeIndexingOptions) { +pub unsafe fn options( + index: pgrx::pg_sys::Relation, +) -> (VectorOptions, RabbitholeIndexingOptions, PgDistanceKind) { let opfamily = unsafe { (*index).rd_opfamily.read() }; let att = unsafe { &mut *(*index).rd_att }; let atts = unsafe { att.attrs.as_slice(att.natts as _) }; @@ -137,7 +144,7 @@ pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, Rabbitho }; // get indexing, segment, optimizing let rabitq = unsafe { convert_reloptions_to_options((*index).rd_options) }; - (vector, rabitq) + (vector, rabitq, pg_d) } #[derive(Debug, Clone, Copy)] @@ -187,14 +194,22 @@ impl Opfamily { use BorrowedVector as B; use OwnedVector as O; match (vector, self.pg_distance) { - (B::Vecf32(x), _) => O::Vecf32(x.own()), + (B::Vecf32(x), PgDistanceKind::L2) => O::Vecf32(x.own()), + (B::Vecf32(x), PgDistanceKind::Dot) => O::Vecf32(x.own()), + (B::Vecf32(x), PgDistanceKind::Cos) => O::Vecf32(x.function_normalize()), (B::Vecf16(x), _) => O::Vecf16(x.own()), (B::SVecf32(x), _) => O::SVecf32(x.own()), (B::BVector(x), _) => O::BVector(x.own()), } } pub fn process(self, x: Distance) -> f32 { - f32::from(x) + match self.pg_distance { + PgDistanceKind::Cos => f32::from(x) + 1.0f32, + _ => f32::from(x), + } + } + pub fn distance_kind(self) -> DistanceKind { + self.pg_distance.to_distance() } } diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs index 9792729a..f248c008 100644 --- a/src/index/am_scan.rs +++ b/src/index/am_scan.rs @@ -80,6 +80,7 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, OwnedVector::SVecf32(_) => unreachable!(), OwnedVector::BVector(_) => unreachable!(), }, + opfamily.distance_kind(), nprobe(), ); *scanner = Scanner::Vbase { diff --git a/src/index/utils.rs b/src/index/utils.rs index 4e500917..68df340f 100644 --- a/src/index/utils.rs +++ b/src/index/utils.rs @@ -1,9 +1,10 @@ +use base::distance::{Distance, DistanceKind}; +use base::scalar::ScalarLike; use base::search::*; -use base::vector::VectorBorrowed; +use base::vector::{VectBorrowed, VectorBorrowed}; use pgrx::pg_sys::panic::ErrorReportable; use pgrx::{error, Spi}; -use crate::algorithm::rabitq; use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { @@ -25,12 +26,16 @@ pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer { Pointer::new(value) } -pub fn load_proj_vectors( +pub fn load_table_vectors( table_name: &str, column_name: &str, rows: u32, dims: u32, -) -> Vec> { + preprocess: F, +) -> Vec> +where + F: Fn(VectBorrowed) -> Vec, +{ let query = format!("SELECT {column_name} FROM {table_name};"); let mut centroids = Vec::new(); @@ -43,8 +48,7 @@ pub fn load_proj_vectors( if let Ok(Some(v)) = vector { let borrowed = v.as_borrowed(); assert_eq!(borrowed.dims(), dims); - let projected_centroids = rabitq::project(borrowed.slice()); - centroids.push(projected_centroids); + centroids.push(preprocess(borrowed)); } else { error!("load vectors from column is not valid") } @@ -52,3 +56,12 @@ pub fn load_proj_vectors( centroids }) } + +pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance { + match d { + DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)), + DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)), + DistanceKind::Hamming => unimplemented!(), + DistanceKind::Jaccard => unimplemented!(), + } +} diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 96dfaaa6..efcff726 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -14,6 +14,20 @@ CREATE OPERATOR <<->> ( COMMUTATOR = <<->> ); +CREATE OPERATOR <<=>> ( + PROCEDURE = _rabbithole_pgvector_vector_sphere_cos_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _rabbithole_pgvector_vector_sphere_dot_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + -- List of functions CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector @@ -27,6 +41,8 @@ COMMENT ON ACCESS METHOD rabbithole IS 'rabbithole index access method'; -- List of operator families CREATE OPERATOR FAMILY vector_l2_ops USING rabbithole; +CREATE OPERATOR FAMILY vector_dot_ops USING rabbithole; +CREATE OPERATOR FAMILY vector_cos_ops USING rabbithole; -- List of operator classes @@ -34,3 +50,13 @@ CREATE OPERATOR CLASS vector_l2_ops FOR TYPE vector USING rabbithole FAMILY vector_l2_ops AS OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH; + +CREATE OPERATOR CLASS vector_dot_ops + FOR TYPE vector USING rabbithole FAMILY vector_dot_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH; + +CREATE OPERATOR CLASS vector_cos_ops + FOR TYPE vector USING rabbithole FAMILY vector_cos_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH; \ No newline at end of file From f9ac8529285d06124d539650a3388f0c19e257e8 Mon Sep 17 00:00:00 2001 From: Jinjing Zhou Date: Thu, 31 Oct 2024 16:32:28 +0800 Subject: [PATCH 017/324] Add AGPL license (#13) Signed-off-by: Jinjing Zhou --- LICENSE | 661 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 661 insertions(+) create mode 100644 LICENSE diff --git a/LICENSE b/LICENSE new file mode 100644 index 00000000..0ad25db4 --- /dev/null +++ b/LICENSE @@ -0,0 +1,661 @@ + GNU AFFERO GENERAL PUBLIC LICENSE + Version 3, 19 November 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU Affero General Public License is a free, copyleft license for +software and other kinds of works, specifically designed to ensure +cooperation with the community in the case of network server software. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +our General Public Licenses are intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + Developers that use our General Public Licenses protect your rights +with two steps: (1) assert copyright on the software, and (2) offer +you this License which gives you legal permission to copy, distribute +and/or modify the software. + + A secondary benefit of defending all users' freedom is that +improvements made in alternate versions of the program, if they +receive widespread use, become available for other developers to +incorporate. Many developers of free software are heartened and +encouraged by the resulting cooperation. However, in the case of +software used on network servers, this result may fail to come about. +The GNU General Public License permits making a modified version and +letting the public access it on a server without ever releasing its +source code to the public. + + The GNU Affero General Public License is designed specifically to +ensure that, in such cases, the modified source code becomes available +to the community. It requires the operator of a network server to +provide the source code of the modified version running there to the +users of that server. Therefore, public use of a modified version, on +a publicly accessible server, gives the public access to the source +code of the modified version. + + An older license, called the Affero General Public License and +published by Affero, was designed to accomplish similar goals. This is +a different license, not a version of the Affero GPL, but Affero has +released a new version of the Affero GPL which permits relicensing under +this license. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU Affero General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Remote Network Interaction; Use with the GNU General Public License. + + Notwithstanding any other provision of this License, if you modify the +Program, your modified version must prominently offer all users +interacting with it remotely through a computer network (if your version +supports such interaction) an opportunity to receive the Corresponding +Source of your version by providing access to the Corresponding Source +from a network server at no charge, through some standard or customary +means of facilitating copying of software. This Corresponding Source +shall include the Corresponding Source for any work covered by version 3 +of the GNU General Public License that is incorporated pursuant to the +following paragraph. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the work with which it is combined will remain governed by version +3 of the GNU General Public License. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU Affero General Public License from time to time. Such new versions +will be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU Affero General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU Affero General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU Affero General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published + by the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Affero General Public License for more details. + + You should have received a copy of the GNU Affero General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If your software can interact with users remotely through a computer +network, you should also make sure that it provides a way for users to +get its source. For example, if your program is a web application, its +interface could display a "Source" link that leads users to an archive +of the code. There are many ways you could offer source, and different +solutions will be better for different programs; see section 13 for the +specific requirements. + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU AGPL, see +. From bd18a2096c5cd2722c3d024e3ce57b97fa1aa3c9 Mon Sep 17 00:00:00 2001 From: Jinjing Zhou Date: Fri, 1 Nov 2024 11:28:36 +0800 Subject: [PATCH 018/324] add index build progress report (#35) Signed-off-by: Jinjing Zhou --- bench/index.py | 60 +++++++++++++++++++++++++++----------------------- 1 file changed, 32 insertions(+), 28 deletions(-) diff --git a/bench/index.py b/bench/index.py index 56a4527d..4fde97f6 100644 --- a/bench/index.py +++ b/bench/index.py @@ -1,13 +1,9 @@ import asyncio -from sys import version_info from time import perf_counter import argparse from pathlib import Path import multiprocessing -if version_info >= (3, 12): - raise RuntimeError("h5py doesn't support 3.12") - import psycopg import h5py from pgvector.psycopg import register_vector_async @@ -47,7 +43,7 @@ def build_arg_parse(): help="Workers to build index", type=int, required=False, - default=max(multiprocessing.cpu_count() - 2, 1), + default=max(multiprocessing.cpu_count() - 1, 1), ) return parser @@ -87,9 +83,9 @@ def get_ivf_ops_config(metric, k, name=None): return metric_ops, ivf_config -async def create_connection(password): +async def create_connection(url): conn = await psycopg.AsyncConnection.connect( - conninfo=f"postgresql://postgres:{password}@localhost:5432/postgres", + conninfo=url, dbname="postgres", autocommit=True, **KEEPALIVE_KWARGS, @@ -113,8 +109,8 @@ async def add_centroids(conn, name, centroids): copy.set_types(["vector"]) for centroid in tqdm(centroids, desc="Adding centroids"): await copy.write_row((centroid,)) - while conn.pgconn.flush() == 1: - pass + while conn.pgconn.flush() == 1: + await asyncio.sleep(0) async def add_embeddings(conn, name, dim, train): @@ -130,8 +126,8 @@ async def add_embeddings(conn, name, dim, train): enumerate(train), desc="Adding embeddings", total=len(train) ): await copy.write_row((i, vec)) - while conn.pgconn.flush() == 1: - pass + while conn.pgconn.flush() == 1: + await asyncio.sleep(0) async def build_index( @@ -139,6 +135,7 @@ async def build_index( ): start_time = perf_counter() await conn.execute(f"SET max_parallel_maintenance_workers TO {workers}") + await conn.execute(f"SET max_parallel_workers TO {workers}") await conn.execute( f"CREATE INDEX ON {name} USING rabbithole (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" ) @@ -146,28 +143,33 @@ async def build_index( finish.set() -async def monitor_index_build(password, finish: asyncio.Event): - conn = await psycopg.AsyncConnection.connect( - conninfo=f"postgresql://postgres:{password}@localhost:5432/postgres", - dbname="postgres", - autocommit=True, - **KEEPALIVE_KWARGS, - ) - pbar = tqdm(smoothing=0.0) +async def monitor_index_build(conn, finish: asyncio.Event): async with conn.cursor() as acur: + blocks_total = None + while blocks_total is None: + await asyncio.sleep(1) + await acur.execute( + f"SELECT blocks_total FROM pg_stat_progress_create_index" + ) + blocks_total = await acur.fetchone() + total = 0 if blocks_total is None else blocks_total[0] + pbar = tqdm(smoothing=0.0, total=total, desc="Building index") while True: if finish.is_set(): + pbar.update(pbar.total - pbar.n) return - await acur.execute(f"SELECT tuples_done FROM pg_stat_progress_create_index") - tuples_done = await acur.fetchone() - update = 0 if tuples_done is None else tuples_done[0] - pbar.update(update - pbar.n) + await acur.execute(f"SELECT blocks_done FROM pg_stat_progress_create_index") + blocks_done = await acur.fetchone() + done = 0 if blocks_done is None else blocks_done[0] + pbar.update(done - pbar.n) await asyncio.sleep(1) + pbar.close() async def main(dataset): dataset = h5py.File(Path(args.input), "r") - conn = await create_connection(args.password) + url = f"postgresql://postgres:{args.password}@localhost:5432/postgres", + conn = await create_connection(url) if args.centroids: centroids = np.load(args.centroids, allow_pickle=False) await add_centroids(conn, args.name, centroids) @@ -177,6 +179,12 @@ async def main(dataset): await add_embeddings(conn, args.name, args.dim, dataset["train"]) index_finish = asyncio.Event() + # Need a seperate connection for monitor process + monitor_conn = await create_connection(url) + monitor_task = monitor_index_build( + monitor_conn, + index_finish, + ) index_task = build_index( conn, args.name, @@ -185,10 +193,6 @@ async def main(dataset): ivf_config, index_finish, ) - monitor_task = monitor_index_build( - args.password, - index_finish, - ) await asyncio.gather(index_task, monitor_task) From 5521e62bfd70d49e61f2bcc93f0a553ad9e8c976 Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 1 Nov 2024 11:29:00 +0800 Subject: [PATCH 019/324] feat: add gpu argparser for kmeans training (#36) Signed-off-by: Keming --- bench/train.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/bench/train.py b/bench/train.py index 48b24a7b..43339926 100644 --- a/bench/train.py +++ b/bench/train.py @@ -32,6 +32,9 @@ def build_arg_parse(): "--niter", help="number of iterations", type=int, default=N_ITER ) parser.add_argument("-m", "--metric", choices=["l2", "cos"], default="l2") + parser.add_argument( + "-g", "--gpu", help="enable GPU for KMeans", action="store_true" + ) return parser @@ -56,14 +59,14 @@ def filter_by_label(iter, labels, target): yield vec -def kmeans_cluster(data, k, child_k, niter, metric): +def kmeans_cluster(data, k, child_k, niter, metric, gpu=False): n, dim = data.shape if n > MAX_POINTS_PER_CLUSTER * k: train = reservoir_sampling(iter(data), MAX_POINTS_PER_CLUSTER * args.k) else: train = data[:] kmeans = Kmeans( - dim, k, verbose=True, niter=niter, seed=SEED, spherical=metric == "cos" + dim, k, gpu=gpu, verbose=True, niter=niter, seed=SEED, spherical=metric == "cos" ) kmeans.train(train) if not child_k: @@ -85,6 +88,7 @@ def kmeans_cluster(data, k, child_k, niter, metric): child_kmeans = Kmeans( dim, child_k, + gpu=gpu, verbose=True, niter=niter, seed=SEED, @@ -105,7 +109,7 @@ def kmeans_cluster(data, k, child_k, niter, metric): start_time = perf_counter() centroids = kmeans_cluster( - dataset["train"], args.k, args.child_k, args.niter, args.metric + dataset["train"], args.k, args.child_k, args.niter, args.metric, args.gpu ) print(f"K-means (k=({args.k}, {args.child_k})): {perf_counter() - start_time:.2f}s") From 4514e033879d6e5ea5bd255796439cffe19abc8f Mon Sep 17 00:00:00 2001 From: Jinjing Zhou Date: Fri, 1 Nov 2024 11:29:38 +0800 Subject: [PATCH 020/324] fix: enable block progress report in parallel build (#34) Signed-off-by: Jinjing Zhou --- src/index/am.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/index/am.rs b/src/index/am.rs index d6367dc3..375ba022 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -546,7 +546,7 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( self.index_info, true, false, - false, + true, 0, pgrx::pg_sys::InvalidBlockNumber, Some(call::), From 47e601d458dab8c3e523fce7c36a21fe2753eeb4 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 4 Nov 2024 14:27:33 +0800 Subject: [PATCH 021/324] chore: bump dependency version (#43) Signed-off-by: usamoi --- Cargo.lock | 125 +++++++++++++++++++++------------------- Cargo.toml | 28 ++++----- src/index/am_options.rs | 8 +-- 3 files changed, 85 insertions(+), 76 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8252a07b..6805b04a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.90" +version = "1.0.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37bf3594c4c988a53154954629820791dde498571819ae4ca50ca811e060cc95" +checksum = "74f37166d7d48a0284b99dd824694c26119c700b53bf0d1540cdb147dbdaaf13" [[package]] name = "approx" @@ -88,7 +88,7 @@ source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2# dependencies = [ "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -107,7 +107,7 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -164,9 +164,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.7.2" +version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "428d9aa8fbc0670b7b8d6030a7fadd0f86151cae55e4dbbece15f3780a3dfaf3" +checksum = "9ac0150caa2ae65ca5bd83f25c7de183dea78d4d366469f148435e2acfbad0da" [[package]] name = "cargo_toml" @@ -180,9 +180,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.1.31" +version = "1.1.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2e7962b54006dcfcc61cb72735f4d89bb97061dd6a7ed882ec6b8ee53714c6f" +checksum = "67b9470d453346108f93a59222a9a1a5724db32d0a4727b7ab7ace4b4d822dc9" dependencies = [ "shlex", ] @@ -298,7 +298,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -309,7 +309,7 @@ checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" dependencies = [ "darling_core", "quote", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -327,7 +327,7 @@ source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2# dependencies = [ "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -353,7 +353,7 @@ checksum = "f282cfdfe92516eb26c2af8589c274c7c17681f5ecc03c18255fe741c6aa64eb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -588,9 +588,9 @@ dependencies = [ [[package]] name = "libm" -version = "0.2.8" +version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" +checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" [[package]] name = "linux-raw-sys" @@ -659,7 +659,7 @@ checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" dependencies = [ "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -782,8 +782,9 @@ dependencies = [ [[package]] name = "pgrx" -version = "0.12.6" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" +version = "0.12.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6468e2a7c4085707209cf7f7e14a7b73c8b7f41fc716283024a09180b620a45" dependencies = [ "atomic-traits", "bitflags", @@ -792,7 +793,6 @@ dependencies = [ "heapless", "libc", "once_cell", - "paste", "pgrx-macros", "pgrx-pg-sys", "pgrx-sql-entity-graph", @@ -806,8 +806,9 @@ dependencies = [ [[package]] name = "pgrx-bindgen" -version = "0.12.6" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" +version = "0.12.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04d822012e4919882cb9b1851ae9e4b4b03703ddf2ef4af3522025c8944e9c7d" dependencies = [ "bindgen", "cc", @@ -817,25 +818,37 @@ dependencies = [ "proc-macro2", "quote", "shlex", - "syn 2.0.82", + "syn 2.0.87", "walkdir", ] +[[package]] +name = "pgrx-catalog" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38fc7bbd6334277565c322bf9d48c80bb8858cad493ec29e337df39dc259fc6a" +dependencies = [ + "paste", + "pgrx", +] + [[package]] name = "pgrx-macros" -version = "0.12.6" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" +version = "0.12.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ffdb879a3880d3034661b4d19f802029abfcde6b8232299f4d4d2c202680af" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] name = "pgrx-pg-config" -version = "0.12.6" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" +version = "0.12.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "945d664e710e4dcccd459628c6e58b3ccbd0495d2ed751b1d6f73996ceca01a2" dependencies = [ "cargo_toml", "eyre", @@ -851,8 +864,9 @@ dependencies = [ [[package]] name = "pgrx-pg-sys" -version = "0.12.6" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" +version = "0.12.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "489c96adc8af0b917165f47a11e25f77753e227144aa3790be7e241afab9886c" dependencies = [ "cee-scape", "libc", @@ -865,15 +879,16 @@ dependencies = [ [[package]] name = "pgrx-sql-entity-graph" -version = "0.12.6" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.6-patch#62fbc03707cd469d8189ce1e3e4489df88afb4db" +version = "0.12.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de77a74f3dc80283a97322ab4d65f9f81691eb54d0a5b7c37b93686d351c6096" dependencies = [ "convert_case", "eyre", "petgraph", "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", "thiserror", "unescape", ] @@ -913,9 +928,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.88" +version = "1.0.89" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c3a7fc5db1e57d5a779a352c8cdb57b29aa4c40cc69c3a68a7fedc815fbf2f9" +checksum = "f139b0662de085916d1fb67d2b4169d1addddda1919e696f3252b740b629986e" dependencies = [ "unicode-ident", ] @@ -978,12 +993,12 @@ dependencies = [ "nalgebra", "paste", "pgrx", + "pgrx-catalog", "quantization", "rand", "rand_chacha", "rand_distr", "rkyv", - "seq-macro", "serde", "toml", "validator", @@ -1063,9 +1078,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.11.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38200e5ee88914975b69f657f0801b6f6dccafd44fd9326302a4aaeecfacb1d8" +checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191" dependencies = [ "aho-corasick", "memchr", @@ -1145,9 +1160,9 @@ dependencies = [ [[package]] name = "rustix" -version = "0.38.37" +version = "0.38.38" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8acb788b847c24f28525660c4d7758620a7210875711f79e7f663cc152726811" +checksum = "aa260229e6538e52293eeb577aabd09945a09d6d9cc0fc550ed7529056c2e32a" dependencies = [ "bitflags", "errno", @@ -1204,17 +1219,11 @@ dependencies = [ "pest", ] -[[package]] -name = "seq-macro" -version = "0.3.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" - [[package]] name = "serde" -version = "1.0.210" +version = "1.0.214" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8e3592472072e6e22e0a54d5904d9febf8508f65fb8552499a1abc7d1078c3a" +checksum = "f55c3193aca71c12ad7890f1785d2b73e1b9f63a0bbc353c08ef26fe03fc56b5" dependencies = [ "serde_derive", ] @@ -1231,13 +1240,13 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.210" +version = "1.0.214" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "243902eda00fad750862fc144cea25caca5e20d615af0a81bee94ca738f1df1f" +checksum = "de523f781f095e28fa605cdce0f8307e451cc0fd14e2eb4cd2e98a355b147766" dependencies = [ "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -1342,9 +1351,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.82" +version = "2.0.87" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83540f837a8afc019423a8edb95b52a8effe46957ee402287f4292fae35be021" +checksum = "25aa4ce346d03a6dcd68dd8b4010bcb74e54e62c90c573f394c46eae99aba32d" dependencies = [ "proc-macro2", "quote", @@ -1359,22 +1368,22 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" [[package]] name = "thiserror" -version = "1.0.64" +version = "1.0.67" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d50af8abc119fb8bb6dbabcfa89656f46f84aa0ac7688088608076ad2b459a84" +checksum = "3b3c6efbfc763e64eb85c11c25320f0737cb7364c4b6336db90aa9ebe27a0bbd" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.64" +version = "1.0.67" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08904e7672f5eb876eaaf87e0ce17857500934f4981c4a0ab2b4aa98baac7fc3" +checksum = "b607164372e89797d78b8e23a6d67d5d1038c1c65efd52e1389ef8b77caba2a6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -1524,7 +1533,7 @@ dependencies = [ "proc-macro-error", "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", ] [[package]] @@ -1717,5 +1726,5 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.82", + "syn 2.0.87", ] diff --git a/Cargo.toml b/Cargo.toml index 8b8c003b..f47b7b4c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,15 +13,16 @@ path = "./src/bin/pgrx_embed.rs" [features] default = [] -pg12 = ["pgrx/pg12"] -pg13 = ["pgrx/pg13"] -pg14 = ["pgrx/pg14"] -pg15 = ["pgrx/pg15"] -pg16 = ["pgrx/pg16"] -pg17 = ["pgrx/pg17"] +pg12 = ["pgrx/pg12", "pgrx-catalog/pg12"] +pg13 = ["pgrx/pg13", "pgrx-catalog/pg13"] +pg14 = ["pgrx/pg14", "pgrx-catalog/pg14"] +pg15 = ["pgrx/pg15", "pgrx-catalog/pg15"] +pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] +pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] -pgrx = { version = "=0.12.6", default-features = false, features = ["cshim"] } +pgrx = { version = "=0.12.8", default-features = false, features = ["cshim"] } +pgrx-catalog = "0.1.0" base = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } common = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } detect = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } @@ -31,16 +32,15 @@ paste = "1" serde = "1" toml = "0.8.19" rand = "0.8.5" -nalgebra = { version = "=0.33.0", default-features = false } -# lock algebra version forever so that the QR decomposition never changes for same input -rand_distr = "0.4.3" -rkyv = { version = "0.7.45", features = ["validation"] } rand_chacha = "0.3.1" -seq-macro = "0.3.5" +rand_distr = "0.4.3" validator = "0.18.1" -[patch.crates-io] -pgrx = { git = "https://github.com/tensorchord/pgrx.git", branch = "v0.12.6-patch" } +# lock algebra version forever so that the QR decomposition never changes for same input +nalgebra = { version = "=0.33.0", default-features = false } + +# lock rkyv version forever so that data is always compatible +rkyv = { version = "=0.7.45", features = ["validation"] } [lints] rust.unsafe_op_in_unsafe_fn = "deny" diff --git a/src/index/am_options.rs b/src/index/am_options.rs index 10d5da9d..f9803e6d 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -54,9 +54,9 @@ impl PgDistanceKind { pub fn convert_opclass_to_vd( opclass_oid: pgrx::pg_sys::Oid, ) -> Option<(VectorKind, PgDistanceKind)> { - let namespace = pgrx::pg_catalog::PgNamespace::search_namespacename(c"rabbithole").unwrap(); + let namespace = pgrx_catalog::PgNamespace::search_namespacename(c"rabbithole").unwrap(); let namespace = namespace.get().expect("rabbithole is not installed."); - let opclass = pgrx::pg_catalog::PgOpclass::search_claoid(opclass_oid).unwrap(); + let opclass = pgrx_catalog::PgOpclass::search_claoid(opclass_oid).unwrap(); let opclass = opclass.get().expect("rabbithole is not installed."); if opclass.opcnamespace() == namespace.oid() { if let Ok(name) = opclass.opcname().to_str() { @@ -71,9 +71,9 @@ pub fn convert_opclass_to_vd( pub fn convert_opfamily_to_vd( opfamily_oid: pgrx::pg_sys::Oid, ) -> Option<(VectorKind, PgDistanceKind)> { - let namespace = pgrx::pg_catalog::PgNamespace::search_namespacename(c"rabbithole").unwrap(); + let namespace = pgrx_catalog::PgNamespace::search_namespacename(c"rabbithole").unwrap(); let namespace = namespace.get().expect("rabbithole is not installed."); - let opfamily = pgrx::pg_catalog::PgOpfamily::search_opfamilyoid(opfamily_oid).unwrap(); + let opfamily = pgrx_catalog::PgOpfamily::search_opfamilyoid(opfamily_oid).unwrap(); let opfamily = opfamily.get().expect("rabbithole is not installed."); if opfamily.opfnamespace() == namespace.oid() { if let Ok(name) = opfamily.opfname().to_str() { From 7fd752f1fd897dd2cea8c8112e1c08bd488f5f81 Mon Sep 17 00:00:00 2001 From: Keming Date: Tue, 5 Nov 2024 16:56:14 +0800 Subject: [PATCH 022/324] chore: add CI lint and build test (#41) * chore: add CI lint and build test Signed-off-by: Keming * fix typo Signed-off-by: Keming * fix pgrx build dependencies Signed-off-by: Keming * use pgrx docker image to accelerate the dev-build/clippy/test Signed-off-by: Keming * fix pgrx docker image ubuntu 22.04 aarch64 Signed-off-by: Keming * address commments Signed-off-by: Keming * Update .github/workflows/style.yml Co-authored-by: usamoi * enable sccache Signed-off-by: Keming --------- Signed-off-by: Keming Co-authored-by: usamoi --- .github/workflows/rust.yml | 69 +++++++++++++++++++++++++++++++++++++ .github/workflows/style.yml | 37 ++++++++++++++++++++ .taplo.toml | 10 ++++++ Cargo.toml | 19 +++++----- bench/index.py | 2 +- docker/pgrx.Dockerfile | 62 +++++++++++++++++++++++++++++++++ 6 files changed, 189 insertions(+), 10 deletions(-) create mode 100644 .github/workflows/rust.yml create mode 100644 .github/workflows/style.yml create mode 100644 .taplo.toml create mode 100644 docker/pgrx.Dockerfile diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml new file mode 100644 index 00000000..018139ab --- /dev/null +++ b/.github/workflows/rust.yml @@ -0,0 +1,69 @@ +name: Rust + +on: + pull_request: + paths: + - '.github/workflows/rust.yml' + - 'src/**' + - 'Cargo.lock' + - 'Cargo.toml' + - '*.control' + - 'rust-toolchain.toml' + push: + branches: + - main + paths: + - '.github/workflows/rust.yml' + - 'src/**' + - 'Cargo.lock' + - 'Cargo.toml' + - '*.control' + - 'rust-toolchain.toml' + merge_group: + workflow_dispatch: + +concurrency: + group: ${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + arch: ["x86_64", "aarch64"] + env: + PGRX_IMAGE: "kemingy/pgrx:0.12.8" + + steps: + - uses: actions/checkout@v4 + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + const url = process.env.ACTIONS_CACHE_URL || ''; + const token = process.env.ACTIONS_RUNTIME_TOKEN || ''; + core.exportVariable( + 'CACHE_ENVS', + `-e CARGO_INCREMENTAL=0 -e SCCACHE_GHA_ENABLED=true -e RUSTC_WRAPPER=sccache -e ACTIONS_CACHE_URL=${url} -e ACTIONS_RUNTIME_TOKEN=${token}`, + ); + - name: Set up docker images and permissions + run: | + docker pull $PGRX_IMAGE + echo "Default user: $(id -u):$(id -g)" + sudo chown -R 1000:1000 . + + - name: Clippy + run: | + for v in {14..17}; do + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE clippy --target ${{ matrix.arch }}-unknown-linux-gnu --features "pg$v" -- -D warnings + done + - name: Build + run: | + for v in {14..17}; do + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE build --lib --target ${{ matrix.arch }}-unknown-linux-gnu --features "pg$v" + done + - name: Test + run: | + # pg agnostic tests + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE test --no-fail-fast --target ${{ matrix.arch }}-unknown-linux-gnu --features pg17 diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml new file mode 100644 index 00000000..230d7d0c --- /dev/null +++ b/.github/workflows/style.yml @@ -0,0 +1,37 @@ +name: Style + +on: + pull_request: + branches: + - main + push: + branches: + - main + merge_group: + workflow_dispatch: + +concurrency: + group: ${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +env: + CARGO_TERM_COLOR: always + +jobs: + lint: + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v4 + + - name: Typos + uses: crate-ci/typos@master + + - name: TOML lint + run: | + curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-x86_64.gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo + taplo fmt --check + + - name: Cargo Lint + run: | + cargo fmt --check diff --git a/.taplo.toml b/.taplo.toml new file mode 100644 index 00000000..d9a9fdac --- /dev/null +++ b/.taplo.toml @@ -0,0 +1,10 @@ +[formatting] +indent_string = " " + +[[rule]] +keys = ["dependencies", "*-denpendencies", "lints", "patch.*", "profile.*"] + +[rule.formatting] +reorder_keys = true +reorder_arrays = true +align_comments = true diff --git a/Cargo.toml b/Cargo.toml index f47b7b4c..bd6643f8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,20 +21,21 @@ pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] +paste = "1" pgrx = { version = "=0.12.8", default-features = false, features = ["cshim"] } pgrx-catalog = "0.1.0" +rand = "0.8.5" +rand_chacha = "0.3.1" +rand_distr = "0.4.3" +serde = "1" +toml = "0.8.19" +validator = "0.18.1" + base = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } common = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } detect = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } k_means = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } quantization = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -paste = "1" -serde = "1" -toml = "0.8.19" -rand = "0.8.5" -rand_chacha = "0.3.1" -rand_distr = "0.4.3" -validator = "0.18.1" # lock algebra version forever so that the QR decomposition never changes for same input nalgebra = { version = "=0.33.0", default-features = false } @@ -48,12 +49,12 @@ rust.unused_lifetimes = "warn" rust.unused_qualifications = "warn" [profile.opt] +debug-assertions = false inherits = "dev" opt-level = 3 -debug-assertions = false overflow-checks = false [profile.release] -lto = "fat" codegen-units = 1 debug = true +lto = "fat" diff --git a/bench/index.py b/bench/index.py index 4fde97f6..97b67a01 100644 --- a/bench/index.py +++ b/bench/index.py @@ -179,7 +179,7 @@ async def main(dataset): await add_embeddings(conn, args.name, args.dim, dataset["train"]) index_finish = asyncio.Event() - # Need a seperate connection for monitor process + # Need a separate connection for monitor process monitor_conn = await create_connection(url) monitor_task = monitor_index_build( monitor_conn, diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile new file mode 100644 index 00000000..920817b5 --- /dev/null +++ b/docker/pgrx.Dockerfile @@ -0,0 +1,62 @@ +# CNPG only support Debian 12 (Bookworm) +FROM ubuntu:22.04 + +ARG PGRX_VERSION=0.12.8 +ARG SCCACHE_VERSION=0.8.2 + +ENV DEBIAN_FRONTEND=noninteractive \ + LANG=en_US.UTF-8 \ + LC_ALL=en_US.UTF-8 \ + RUSTFLAGS="-Dwarnings" \ + RUST_BACKTRACE=1 \ + CARGO_TERM_COLOR=always + +RUN set -eux; \ + apt update; \ + apt install -y --no-install-recommends \ + curl \ + ca-certificates \ + build-essential \ + postgresql-common gnupg \ + crossbuild-essential-arm64 \ + qemu-user-static \ + libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config \ + clang + +# set up sccache +RUN set -ex; \ + curl -fsSL -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v${SCCACHE_VERSION}/sccache-v${SCCACHE_VERSION}-x86_64-unknown-linux-musl.tar.gz; \ + tar -xzf sccache.tar.gz --strip-components=1; \ + rm sccache.tar.gz; \ + mv sccache /usr/local/bin/ + +RUN /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y +# install all the PostgresQL +RUN set -ex; \ + for v in $(seq 14 17); do \ + apt install -y --no-install-recommends postgresql-$v postgresql-server-dev-$v; \ + done; \ + rm -rf /var/lib/apt/lists/*; + +# create a non-root user (make it compatible with Ubuntu 24.04) +RUN useradd -u 1000 -U -m ubuntu +USER ubuntu +ENV PATH="$PATH:/home/ubuntu/.cargo/bin" +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- --default-toolchain=none -y + +WORKDIR /workspace +COPY rust-toolchain.toml /workspace/rust-toolchain.toml +# ref: https://github.com/pgcentralfoundation/pgrx/blob/develop/docs/src/extension/build/cross-compile.md +RUN set -ex; \ + echo 'target.aarch64-unknown-linux-gnu.linker = "aarch64-linux-gnu-gcc"' >> ~/.cargo/config.toml; \ + echo 'target.aarch64-unknown-linux-gnu.runner = ["qemu-aarch64-static", "-L", "/usr/aarch64-linux-gnu"]' >> ~/.cargo/config.toml +RUN rustup target add x86_64-unknown-linux-gnu aarch64-unknown-linux-gnu + +RUN cargo install cargo-pgrx --locked --version=${PGRX_VERSION} + +RUN set -ex; \ + for v in $(seq 14 17); do \ + cargo pgrx init --pg$v=/usr/lib/postgresql/$v/bin/pg_config; \ + done; + +ENTRYPOINT [ "cargo" ] From 39176bf1223cda3be20f4a0a8cd1fc960b278a15 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 7 Nov 2024 11:21:22 +0800 Subject: [PATCH 023/324] fix: do not release resources when an interrupt comes (#49) Signed-off-by: usamoi --- src/algorithm/build.rs | 1 - src/index/am.rs | 30 ++++++++++++++---------------- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index e175ee19..9f3eef83 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -55,7 +55,6 @@ pub fn build( let mut samples = Vec::new(); let mut number_of_samples = 0_u32; heap_relation.traverse(|(_, vector)| { - pgrx::check_for_interrupts!(); assert_eq!(dims as usize, vector.len(), "invalid vector dimensions",); let vector = rabitq::project(&vector); if number_of_samples < max_number_of_samples { diff --git a/src/index/am.rs b/src/index/am.rs index 375ba022..04f78e8f 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -164,7 +164,6 @@ pub unsafe extern "C" fn ambuild( ) where F: FnMut((Pointer, Vec)), { - pgrx::check_for_interrupts!(); use base::vector::OwnedVector; let state = unsafe { &mut *state.cast::>() }; let vector = unsafe { @@ -247,10 +246,10 @@ pub unsafe extern "C" fn ambuild( unsafe { RabbitholeLeader::enter(heap, index, (*index_info).ii_Concurrent) } { unsafe { - let nparticipanttuplesorts = leader.nparticipanttuplesorts; + let nparticipants = leader.nparticipants; loop { pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.rabbitholeshared).mutex); - if (*leader.rabbitholeshared).nparticipantsdone == nparticipanttuplesorts { + if (*leader.rabbitholeshared).nparticipantsdone == nparticipants { pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.rabbitholeshared).mutex); break; } @@ -266,7 +265,6 @@ pub unsafe extern "C" fn ambuild( let mut tuples_done = 0; reporter.tuples_done(tuples_done); heap_relation.traverse(|(payload, vector)| { - pgrx::check_for_interrupts!(); algorithm::insert::insert( index_relation.clone(), payload, @@ -306,7 +304,7 @@ fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { struct RabbitholeLeader { pcxt: *mut pgrx::pg_sys::ParallelContext, - nparticipanttuplesorts: i32, + nparticipants: i32, rabbitholeshared: *mut RabbitholeShared, snapshot: pgrx::pg_sys::Snapshot, } @@ -419,10 +417,10 @@ impl RabbitholeLeader { pgrx::pg_sys::LaunchParallelWorkers(pcxt); } - let nparticipanttuplesorts = unsafe { (*pcxt).nworkers_launched }; + let nparticipants = unsafe { (*pcxt).nworkers_launched }; unsafe { - if nparticipanttuplesorts == 0 { + if nparticipants == 0 { pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); if is_mvcc_snapshot(snapshot) { pgrx::pg_sys::UnregisterSnapshot(snapshot); @@ -437,7 +435,7 @@ impl RabbitholeLeader { } Some(Self { pcxt, - nparticipanttuplesorts, + nparticipants, rabbitholeshared, snapshot, }) @@ -446,13 +444,15 @@ impl RabbitholeLeader { impl Drop for RabbitholeLeader { fn drop(&mut self) { - unsafe { - pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); - if is_mvcc_snapshot(self.snapshot) { - pgrx::pg_sys::UnregisterSnapshot(self.snapshot); + if !std::thread::panicking() { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); + if is_mvcc_snapshot(self.snapshot) { + pgrx::pg_sys::UnregisterSnapshot(self.snapshot); + } + pgrx::pg_sys::DestroyParallelContext(self.pcxt); + pgrx::pg_sys::ExitParallelMode(); } - pgrx::pg_sys::DestroyParallelContext(self.pcxt); - pgrx::pg_sys::ExitParallelMode(); } } } @@ -514,7 +514,6 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( ) where F: FnMut((Pointer, Vec)), { - pgrx::check_for_interrupts!(); use base::vector::OwnedVector; let state = unsafe { &mut *state.cast::>() }; let vector = unsafe { @@ -568,7 +567,6 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( scan, }; heap_relation.traverse(|(payload, vector)| { - pgrx::check_for_interrupts!(); algorithm::insert::insert( index_relation.clone(), payload, From 3c94db357ced1ed36f6c07687c4c0178cb999518 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 7 Nov 2024 14:35:30 +0800 Subject: [PATCH 024/324] fix: cosine, l2 distance and external build for multilevel kmeans (#44) * fix: cosine distance and external build for multilevel kmeans Signed-off-by: usamoi * chore: use pgvector names Signed-off-by: usamoi --------- Signed-off-by: usamoi --- bench/index.py | 8 +- src/algorithm/build.rs | 255 +++++++++++++--------- src/algorithm/scan.rs | 8 +- src/datatype/operators_pgvector_vector.rs | 20 +- src/index/am.rs | 38 ++-- src/index/am_options.rs | 11 +- src/index/am_scan.rs | 2 +- src/index/utils.rs | 36 --- src/lib.rs | 5 +- src/sql/finalize.sql | 26 +-- src/types.rs | 74 +++++-- 11 files changed, 264 insertions(+), 219 deletions(-) diff --git a/bench/index.py b/bench/index.py index 97b67a01..bd0322d1 100644 --- a/bench/index.py +++ b/bench/index.py @@ -61,15 +61,15 @@ def get_ivf_ops_config(metric, k, name=None): residual_quantization = true spherical_centroids = false """ - elif metric == "cos": - metric_ops = "vector_cos_ops" + elif metric == "cosine": + metric_ops = "vector_cosine_ops" ivf_config = f""" nlist = {k} residual_quantization = false spherical_centroids = true """ - elif metric == "dot": - metric_ops = "vector_dot_ops" + elif metric == "ip": + metric_ops = "vector_ip_ops" ivf_config = f""" nlist = {k} residual_quantization = false diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index 9f3eef83..6c111d0e 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -1,17 +1,16 @@ use crate::algorithm::rabitq; use crate::algorithm::tuples::*; -use crate::index::am_options::PgDistanceKind; -use crate::index::utils::load_table_vectors; +use crate::index::am_options::Opfamily; use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; -use crate::types::ExternalCentroids; +use crate::types::RabbitholeBuildOptions; +use crate::types::RabbitholeExternalBuildOptions; use crate::types::RabbitholeIndexingOptions; +use crate::types::RabbitholeInternalBuildOptions; use base::distance::DistanceKind; use base::index::VectorOptions; use base::scalar::ScalarLike; use base::search::Pointer; -use base::vector::VectBorrowed; -use base::vector::VectorBorrowed; use common::vec2::Vec2; use rand::Rng; use rkyv::ser::serializers::AllocSerializer; @@ -23,6 +22,7 @@ pub trait HeapRelation { fn traverse(&self, callback: F) where F: FnMut((Pointer, Vec)); + fn opfamily(&self) -> Opfamily; } pub trait Reporter { @@ -35,28 +35,26 @@ pub fn build( rabbithole_options: RabbitholeIndexingOptions, heap_relation: T, relation: Relation, - pg_distance: PgDistanceKind, mut reporter: R, ) { let dims = vector_options.dims; let is_residual = rabbithole_options.residual_quantization && vector_options.d == DistanceKind::L2; - let structure = match &rabbithole_options.external_centroids { - Some(_) => Structure::load( + let structure = match rabbithole_options.build { + RabbitholeBuildOptions::External(external_build) => Structure::extern_build( vector_options.clone(), - rabbithole_options.clone(), - pg_distance, + heap_relation.opfamily(), + external_build.clone(), ), - None => { + RabbitholeBuildOptions::Internal(internal_build) => { let mut tuples_total = 0_usize; let samples = { let mut rand = rand::thread_rng(); - let max_number_of_samples = rabbithole_options.nlist.saturating_mul(256); + let max_number_of_samples = internal_build.nlist.saturating_mul(256); let mut samples = Vec::new(); let mut number_of_samples = 0_u32; heap_relation.traverse(|(_, vector)| { - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions",); - let vector = rabitq::project(&vector); + assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); if number_of_samples < max_number_of_samples { samples.extend(vector); number_of_samples += 1; @@ -71,7 +69,7 @@ pub fn build( Vec2::from_vec((number_of_samples as _, dims as _), samples) }; reporter.tuples_total(tuples_total); - Structure::compute(vector_options.clone(), rabbithole_options.clone(), samples) + Structure::internal_build(vector_options.clone(), internal_build.clone(), samples) } }; let h2_len = structure.h2_len(); @@ -182,33 +180,37 @@ pub fn build( } struct Structure { - h2_mean: Vec, - h2_children: Vec, + h2_means: Vec>, + h2_children: Vec>, h1_means: Vec>, h1_children: Vec>, } impl Structure { - fn compute( + fn internal_build( vector_options: VectorOptions, - rabbithole_options: RabbitholeIndexingOptions, - samples: Vec2, + internal_build: RabbitholeInternalBuildOptions, + mut samples: Vec2, ) -> Self { let dims = vector_options.dims; + for i in 0..samples.shape_0() { + let vector = &mut samples[(i,)]; + vector.copy_from_slice(&rabitq::project(vector)); + } let h1_means = base::parallelism::RayonParallelism::scoped( - rabbithole_options.build_threads as _, + internal_build.build_threads as _, Arc::new(AtomicBool::new(false)), |parallelism| { let raw = k_means::k_means( parallelism, - rabbithole_options.nlist as usize, + internal_build.nlist as usize, samples, - rabbithole_options.spherical_centroids, + internal_build.spherical_centroids, 10, false, ); let mut centroids = Vec::new(); - for i in 0..rabbithole_options.nlist { + for i in 0..internal_build.nlist { centroids.push(raw[(i as usize,)].to_vec()); } centroids @@ -218,116 +220,161 @@ impl Structure { .expect("k_means interrupted"); let h2_mean = { let mut centroid = vec![0.0; dims as _]; - for i in 0..rabbithole_options.nlist { + for i in 0..internal_build.nlist { for j in 0..dims { centroid[j as usize] += h1_means[i as usize][j as usize]; } } for j in 0..dims { - centroid[j as usize] /= rabbithole_options.nlist as f32; + centroid[j as usize] /= internal_build.nlist as f32; } centroid }; Structure { - h2_mean, - h2_children: (0..rabbithole_options.nlist).collect(), + h2_means: vec![h2_mean], + h2_children: vec![(0..internal_build.nlist).collect()], h1_means, - h1_children: (0..rabbithole_options.nlist).map(|_| Vec::new()).collect(), + h1_children: (0..internal_build.nlist).map(|_| Vec::new()).collect(), } } - fn load( + fn extern_build( vector_options: VectorOptions, - rabbithole_options: RabbitholeIndexingOptions, - pg_distance: PgDistanceKind, + _opfamily: Opfamily, + external_build: RabbitholeExternalBuildOptions, ) -> Self { - let dims = vector_options.dims; - let preprocess_data = match pg_distance { - PgDistanceKind::L2 | PgDistanceKind::Dot => { - |b: VectBorrowed| rabitq::project(b.slice()) + use std::collections::BTreeMap; + let RabbitholeExternalBuildOptions { table } = external_build; + let query = format!("SELECT id, parent, vector FROM {table};"); + let mut parents = BTreeMap::new(); + let mut vectors = BTreeMap::new(); + pgrx::spi::Spi::connect(|client| { + use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; + use base::vector::VectorBorrowed; + use pgrx::pg_sys::panic::ErrorReportable; + let table = client.select(&query, None, None).unwrap_or_report(); + for row in table { + let id: Option = row.get_by_name("id").unwrap(); + let parent: Option = row.get_by_name("parent").unwrap(); + let vector: Option = row.get_by_name("vector").unwrap(); + let id = id.expect("extern build: id could not be NULL"); + let vector = vector.expect("extern build: vector could not be NULL"); + let pop = parents.insert(id, parent); + if pop.is_some() { + pgrx::error!( + "external build: there are at least two lines have same id, id = {id}" + ); + } + if vector_options.dims != vector.as_borrowed().dims() { + pgrx::error!("extern build: incorrect dimension, id = {id}"); + } + vectors.insert(id, rabitq::project(vector.slice())); } - PgDistanceKind::Cos => { - |b: VectBorrowed| rabitq::project(b.function_normalize().slice()) + }); + let mut children = parents + .keys() + .map(|x| (*x, Vec::new())) + .collect::>(); + let mut root = None; + for (&id, &parent) in parents.iter() { + if let Some(parent) = parent { + if let Some(parent) = children.get_mut(&parent) { + parent.push(id); + } else { + pgrx::error!( + "external build: parent does not exist, id = {id}, parent = {parent}" + ); + } + } else { + if let Some(root) = root { + pgrx::error!("external build: two root, id = {root}, id = {id}"); + } else { + root = Some(id); + } } + } + let Some(root) = root else { + pgrx::error!("extern build: there are no root"); }; - let preprocess_index = |b: VectBorrowed| b.slice().to_vec(); - - let h1_means = match &rabbithole_options.external_centroids { - Some(ExternalCentroids { - table, - h1_means_column: h1, - .. - }) => load_table_vectors( - table, - h1, - rabbithole_options.nlist, - vector_options.dims, - preprocess_data, - ), - - _ => unreachable!(), - }; - let h1_children = match &rabbithole_options.external_centroids { - Some(ExternalCentroids { - table, - h1_children_column: Some(h1), - .. - }) => load_table_vectors(table, h1, 1, vector_options.dims, preprocess_index) - .into_iter() - .map(|v| v.into_iter().map(|f| f as u32).collect()) - .collect(), - _ => (0..rabbithole_options.nlist).map(|_| Vec::new()).collect(), - }; - let h2_mean = match &rabbithole_options.external_centroids { - Some(ExternalCentroids { - table, - h2_mean_column: Some(h2), - .. - }) => load_table_vectors(table, h2, 1, vector_options.dims, preprocess_data) - .pop() - .expect("load h2_mean panic"), - _ => { - let mut centroid = vec![0.0; dims as _]; - for i in 0..rabbithole_options.nlist { - for j in 0..dims { - centroid[j as usize] += h1_means[i as usize][j as usize]; + let mut heights = BTreeMap::<_, _>::new(); + fn dfs_for_heights( + heights: &mut BTreeMap>, + children: &BTreeMap>, + u: i32, + ) { + if heights.contains_key(&u) { + pgrx::error!("extern build: detect a cycle, id = {u}"); + } + heights.insert(u, None); + let mut height = None; + for &v in children[&u].iter() { + dfs_for_heights(heights, children, v); + let new = heights[&v].unwrap() + 1; + if let Some(height) = height { + if height != new { + pgrx::error!("extern build: two heights, id = {u}"); } + } else { + height = Some(new); } - for j in 0..dims { - centroid[j as usize] /= rabbithole_options.nlist as f32; - } - centroid } - }; - let h2_children = match &rabbithole_options.external_centroids { - Some(ExternalCentroids { - table, - h2_children_column: Some(h2), - .. - }) => load_table_vectors(table, h2, 1, vector_options.dims, preprocess_index) - .pop() - .expect("load h2_children panic") - .into_iter() - .map(|f| f as u32) - .collect(), - _ => (0..rabbithole_options.nlist).collect(), - }; - Structure { - h2_mean, + if height.is_none() { + height = Some(1); + } + heights.insert(u, height); + } + dfs_for_heights(&mut heights, &children, root); + let heights = heights + .into_iter() + .map(|(k, v)| (k, v.expect("not a connected graph"))) + .collect::>(); + if heights[&root] != 2 { + pgrx::error!( + "extern build: unexpected tree height, height = {}", + heights[&root] + ); + } + let mut cursors = vec![0_u32; 1 + heights[&root] as usize]; + let mut labels = BTreeMap::new(); + for id in parents.keys().copied() { + let height = heights[&id]; + let cursor = cursors[height as usize]; + labels.insert(id, (height, cursor)); + cursors[height as usize] += 1; + } + fn extract( + height: u32, + labels: &BTreeMap, + vectors: &BTreeMap>, + children: &BTreeMap>, + ) -> (Vec>, Vec>) { + labels + .iter() + .filter(|(_, &(h, _))| h == height) + .map(|(id, _)| { + ( + vectors[id].clone(), + children[id].iter().map(|id| labels[id].1).collect(), + ) + }) + .unzip() + } + let (h2_means, h2_children) = extract(2, &labels, &vectors, &children); + let (h1_means, h1_children) = extract(1, &labels, &vectors, &children); + Self { + h2_means, h2_children, h1_means, h1_children, } } fn h2_len(&self) -> u32 { - 1 + self.h2_means.len() as _ } fn h2_means(&self, i: u32) -> &Vec { - assert!(i == 0); - &self.h2_mean + &self.h2_means[i as usize] } fn h2_children(&self, i: u32) -> &Vec { - assert!(i == 0); - &self.h2_children + &self.h2_children[i as usize] } fn h1_len(&self) -> u32 { self.h1_means.len() as _ diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs index dec0cbeb..56950f13 100644 --- a/src/algorithm/scan.rs +++ b/src/algorithm/scan.rs @@ -15,9 +15,9 @@ pub fn scan( relation: Relation, vector: Vec, distance_kind: DistanceKind, - h1_nprobe: u32, + nprobe_1: u32, ) -> impl Iterator { - assert!(h1_nprobe >= 1); + assert!(nprobe_1 >= 1); let meta_guard = relation.read(0); let meta_tuple = meta_guard .get() @@ -28,7 +28,7 @@ pub fn scan( let dims = meta_tuple.dims; assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); let vector = rabitq::project(&vector); - let is_residual = meta_tuple.is_residual && distance_kind == DistanceKind::L2; + let is_residual = meta_tuple.is_residual; let default_lut = if !is_residual { Some(rabitq::fscan_preprocess(&vector)) } else { @@ -118,7 +118,7 @@ pub fn scan( let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; Some((first, mean)) }) - .take(h1_nprobe as usize) + .take(nprobe_1 as usize) .collect() }; { diff --git a/src/datatype/operators_pgvector_vector.rs b/src/datatype/operators_pgvector_vector.rs index 44e17b80..1d8ac1aa 100644 --- a/src/datatype/operators_pgvector_vector.rs +++ b/src/datatype/operators_pgvector_vector.rs @@ -1,5 +1,4 @@ use crate::datatype::memory_pgvector_vector::*; -use base::scalar::ScalarLike; use base::vector::{VectBorrowed, VectorBorrowed}; use std::num::NonZero; @@ -21,11 +20,14 @@ fn _rabbithole_pgvector_vector_sphere_l2_in( Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), Err(_) => unreachable!(), }; - f32::reduce_sum_of_d2(lhs.slice(), center.slice()).to_f32() < radius + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + let d = VectBorrowed::operator_l2(lhs, center).to_f32().sqrt(); + d < radius } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_pgvector_vector_sphere_dot_in( +fn _rabbithole_pgvector_vector_sphere_ip_in( lhs: PgvectorVectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { @@ -42,11 +44,14 @@ fn _rabbithole_pgvector_vector_sphere_dot_in( Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), Err(_) => unreachable!(), }; - -f32::reduce_sum_of_xy(lhs.slice(), center.slice()) < radius + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + let d = VectBorrowed::operator_dot(lhs, center).to_f32(); + d < radius } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_pgvector_vector_sphere_cos_in( +fn _rabbithole_pgvector_vector_sphere_cosine_in( lhs: PgvectorVectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { @@ -63,5 +68,8 @@ fn _rabbithole_pgvector_vector_sphere_cos_in( Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), Err(_) => unreachable!(), }; - VectBorrowed::operator_cos(lhs.as_borrowed(), center.as_borrowed()).to_f32() < radius + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + let d = VectBorrowed::operator_cos(lhs, center).to_f32(); + d < radius } diff --git a/src/index/am.rs b/src/index/am.rs index 04f78e8f..005a9ff6 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -166,15 +166,11 @@ pub unsafe extern "C" fn ambuild( { use base::vector::OwnedVector; let state = unsafe { &mut *state.cast::>() }; - let vector = unsafe { - state - .this - .opfamily - .datum_to_vector(*values.add(0), *is_null.add(0)) - }; + let opfamily = state.this.opfamily; + let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { - let vector = match vector { + let vector = match opfamily.preprocess(vector.as_borrowed()) { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), @@ -204,6 +200,10 @@ pub unsafe extern "C" fn ambuild( ); } } + + fn opfamily(&self) -> Opfamily { + self.opfamily + } } #[derive(Debug, Clone)] pub struct PgReporter {} @@ -225,12 +225,13 @@ pub unsafe extern "C" fn ambuild( } } } - let (vector_options, rabbithole_options, pg_distance) = unsafe { am_options::options(index) }; + let (vector_options, rabbithole_options) = unsafe { am_options::options(index) }; + let opfamily = unsafe { am_options::opfamily(index) }; let heap_relation = Heap { heap, index, index_info, - opfamily: unsafe { am_options::opfamily(index) }, + opfamily, }; let mut reporter = PgReporter {}; let index_relation = unsafe { Relation::new(index) }; @@ -239,7 +240,6 @@ pub unsafe extern "C" fn ambuild( rabbithole_options, heap_relation.clone(), index_relation.clone(), - pg_distance, reporter.clone(), ); if let Some(leader) = @@ -269,7 +269,7 @@ pub unsafe extern "C" fn ambuild( index_relation.clone(), payload, vector, - pg_distance.to_distance(), + opfamily.distance_kind(), ); tuples_done += 1; reporter.tuples_done(tuples_done); @@ -516,15 +516,11 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( { use base::vector::OwnedVector; let state = unsafe { &mut *state.cast::>() }; - let vector = unsafe { - state - .this - .opfamily - .datum_to_vector(*values.add(0), *is_null.add(0)) - }; + let opfamily = state.this.opfamily; + let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { - let vector = match vector { + let vector = match opfamily.preprocess(vector.as_borrowed()) { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), @@ -554,6 +550,10 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( ); } } + + fn opfamily(&self) -> Opfamily { + self.opfamily + } } let index_relation = unsafe { Relation::new(index) }; @@ -608,7 +608,7 @@ pub unsafe extern "C" fn aminsert( let opfamily = unsafe { am_options::opfamily(index) }; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { - let vector = match vector { + let vector = match opfamily.preprocess(vector.as_borrowed()) { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), diff --git a/src/index/am_options.rs b/src/index/am_options.rs index f9803e6d..19396536 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -88,8 +88,8 @@ pub fn convert_opfamily_to_vd( fn convert_name_to_vd(name: &str) -> Option<(VectorKind, PgDistanceKind)> { match name.strip_suffix("_ops") { Some("vector_l2") => Some((VectorKind::Vecf32, PgDistanceKind::L2)), - Some("vector_dot") => Some((VectorKind::Vecf32, PgDistanceKind::Dot)), - Some("vector_cos") => Some((VectorKind::Vecf32, PgDistanceKind::Cos)), + Some("vector_ip") => Some((VectorKind::Vecf32, PgDistanceKind::Dot)), + Some("vector_cosine") => Some((VectorKind::Vecf32, PgDistanceKind::Cos)), _ => None, } } @@ -114,9 +114,7 @@ unsafe fn convert_reloptions_to_options( } } -pub unsafe fn options( - index: pgrx::pg_sys::Relation, -) -> (VectorOptions, RabbitholeIndexingOptions, PgDistanceKind) { +pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, RabbitholeIndexingOptions) { let opfamily = unsafe { (*index).rd_opfamily.read() }; let att = unsafe { &mut *(*index).rd_att }; let atts = unsafe { att.attrs.as_slice(att.natts as _) }; @@ -144,7 +142,7 @@ pub unsafe fn options( }; // get indexing, segment, optimizing let rabitq = unsafe { convert_reloptions_to_options((*index).rd_options) }; - (vector, rabitq, pg_d) + (vector, rabitq) } #[derive(Debug, Clone, Copy)] @@ -205,6 +203,7 @@ impl Opfamily { pub fn process(self, x: Distance) -> f32 { match self.pg_distance { PgDistanceKind::Cos => f32::from(x) + 1.0f32, + PgDistanceKind::L2 => f32::from(x).sqrt(), _ => f32::from(x), } } diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs index f248c008..0c6406db 100644 --- a/src/index/am_scan.rs +++ b/src/index/am_scan.rs @@ -74,7 +74,7 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, if let Some((vector, opfamily)) = vector.as_ref() { let vbase = scan( relation, - match vector { + match opfamily.preprocess(vector.as_borrowed()) { OwnedVector::Vecf32(x) => x.slice().to_vec(), OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), diff --git a/src/index/utils.rs b/src/index/utils.rs index 68df340f..35697ab3 100644 --- a/src/index/utils.rs +++ b/src/index/utils.rs @@ -1,11 +1,6 @@ use base::distance::{Distance, DistanceKind}; use base::scalar::ScalarLike; use base::search::*; -use base::vector::{VectBorrowed, VectorBorrowed}; -use pgrx::pg_sys::panic::ErrorReportable; -use pgrx::{error, Spi}; - -use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { let value = pointer.as_u64(); @@ -26,37 +21,6 @@ pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer { Pointer::new(value) } -pub fn load_table_vectors( - table_name: &str, - column_name: &str, - rows: u32, - dims: u32, - preprocess: F, -) -> Vec> -where - F: Fn(VectBorrowed) -> Vec, -{ - let query = format!("SELECT {column_name} FROM {table_name};"); - let mut centroids = Vec::new(); - - Spi::connect(|client| { - let tup_table = client.select(&query, None, None).unwrap_or_report(); - assert_eq!(tup_table.len(), rows as usize); - - for row in tup_table { - let vector = row[column_name].value::(); - if let Ok(Some(v)) = vector { - let borrowed = v.as_borrowed(); - assert_eq!(borrowed.dims(), dims); - centroids.push(preprocess(borrowed)); - } else { - error!("load vectors from column is not valid") - } - } - centroids - }) -} - pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance { match d { DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)), diff --git a/src/lib.rs b/src/lib.rs index 347aee74..db5d53c7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,8 @@ -#![allow(clippy::too_many_arguments)] +#![allow(clippy::collapsible_else_if)] +#![allow(clippy::identity_op)] #![allow(clippy::needless_range_loop)] +#![allow(clippy::too_many_arguments)] #![allow(clippy::type_complexity)] -#![allow(clippy::identity_op)] mod algorithm; mod datatype; diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index efcff726..259b7797 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -14,18 +14,18 @@ CREATE OPERATOR <<->> ( COMMUTATOR = <<->> ); -CREATE OPERATOR <<=>> ( - PROCEDURE = _rabbithole_pgvector_vector_sphere_cos_in, +CREATE OPERATOR <<#>> ( + PROCEDURE = _rabbithole_pgvector_vector_sphere_ip_in, LEFTARG = vector, RIGHTARG = sphere_vector, - COMMUTATOR = <<=>> + COMMUTATOR = <<#>> ); -CREATE OPERATOR <<#>> ( - PROCEDURE = _rabbithole_pgvector_vector_sphere_dot_in, +CREATE OPERATOR <<=>> ( + PROCEDURE = _rabbithole_pgvector_vector_sphere_cosine_in, LEFTARG = vector, RIGHTARG = sphere_vector, - COMMUTATOR = <<#>> + COMMUTATOR = <<=>> ); -- List of functions @@ -41,8 +41,8 @@ COMMENT ON ACCESS METHOD rabbithole IS 'rabbithole index access method'; -- List of operator families CREATE OPERATOR FAMILY vector_l2_ops USING rabbithole; -CREATE OPERATOR FAMILY vector_dot_ops USING rabbithole; -CREATE OPERATOR FAMILY vector_cos_ops USING rabbithole; +CREATE OPERATOR FAMILY vector_ip_ops USING rabbithole; +CREATE OPERATOR FAMILY vector_cosine_ops USING rabbithole; -- List of operator classes @@ -51,12 +51,12 @@ CREATE OPERATOR CLASS vector_l2_ops OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH; -CREATE OPERATOR CLASS vector_dot_ops - FOR TYPE vector USING rabbithole FAMILY vector_dot_ops AS +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING rabbithole FAMILY vector_ip_ops AS OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH; -CREATE OPERATOR CLASS vector_cos_ops - FOR TYPE vector USING rabbithole FAMILY vector_cos_ops AS +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING rabbithole FAMILY vector_cosine_ops AS OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, - OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH; \ No newline at end of file + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH; diff --git a/src/types.rs b/src/types.rs index b60fa149..eb977b69 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1,55 +1,81 @@ use serde::{Deserialize, Serialize}; use validator::Validate; -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ExternalCentroids { - pub table: String, - pub h1_means_column: String, - pub h1_children_column: Option, - pub h2_mean_column: Option, - pub h2_children_column: Option, -} - #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] -pub struct RabbitholeIndexingOptions { - #[serde(default = "RabbitholeIndexingOptions::default_nlist")] +pub struct RabbitholeInternalBuildOptions { + #[serde(default = "RabbitholeInternalBuildOptions::default_nlist")] #[validate(range(min = 1, max = 1_000_000))] pub nlist: u32, - #[serde(default = "RabbitholeIndexingOptions::default_spherical_centroids")] + #[serde(default = "RabbitholeInternalBuildOptions::default_spherical_centroids")] pub spherical_centroids: bool, - #[serde(default = "RabbitholeIndexingOptions::default_residual_quantization")] - pub residual_quantization: bool, - #[serde(default = "RabbitholeIndexingOptions::default_build_threads")] + #[serde(default = "RabbitholeInternalBuildOptions::default_build_threads")] #[validate(range(min = 1, max = 255))] pub build_threads: u16, - pub external_centroids: Option, } -impl RabbitholeIndexingOptions { +impl RabbitholeInternalBuildOptions { fn default_nlist() -> u32 { 1000 } fn default_spherical_centroids() -> bool { false } - fn default_residual_quantization() -> bool { - false - } fn default_build_threads() -> u16 { 1 } } -impl Default for RabbitholeIndexingOptions { +impl Default for RabbitholeInternalBuildOptions { fn default() -> Self { Self { nlist: Self::default_nlist(), spherical_centroids: Self::default_spherical_centroids(), - residual_quantization: Self::default_residual_quantization(), build_threads: Self::default_build_threads(), - external_centroids: None, } } } + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct RabbitholeExternalBuildOptions { + pub table: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +#[serde(rename_all = "snake_case")] +pub enum RabbitholeBuildOptions { + Internal(RabbitholeInternalBuildOptions), + External(RabbitholeExternalBuildOptions), +} + +impl Default for RabbitholeBuildOptions { + fn default() -> Self { + Self::Internal(Default::default()) + } +} + +impl Validate for RabbitholeBuildOptions { + fn validate(&self) -> Result<(), validator::ValidationErrors> { + use RabbitholeBuildOptions::*; + match self { + Internal(internal_build) => internal_build.validate(), + External(external_build) => external_build.validate(), + } + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct RabbitholeIndexingOptions { + #[serde(default = "RabbitholeIndexingOptions::default_residual_quantization")] + pub residual_quantization: bool, + pub build: RabbitholeBuildOptions, +} + +impl RabbitholeIndexingOptions { + fn default_residual_quantization() -> bool { + false + } +} From 6f0e271176b0c6d7e6f884f6bcc61b51e99a6461 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 7 Nov 2024 16:42:45 +0800 Subject: [PATCH 025/324] feat: leader participation (#50) Signed-off-by: usamoi --- src/algorithm/build.rs | 4 +-- src/index/am.rs | 69 +++++++++++++++++++++++++++++++----------- 2 files changed, 54 insertions(+), 19 deletions(-) diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index 6c111d0e..ebd73e25 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -19,7 +19,7 @@ use std::sync::atomic::AtomicBool; use std::sync::Arc; pub trait HeapRelation { - fn traverse(&self, callback: F) + fn traverse(&self, progress: bool, callback: F) where F: FnMut((Pointer, Vec)); fn opfamily(&self) -> Opfamily; @@ -53,7 +53,7 @@ pub fn build( let max_number_of_samples = internal_build.nlist.saturating_mul(256); let mut samples = Vec::new(); let mut number_of_samples = 0_u32; - heap_relation.traverse(|(_, vector)| { + heap_relation.traverse(false, |(_, vector)| { assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); if number_of_samples < max_number_of_samples { samples.extend(vector); diff --git a/src/index/am.rs b/src/index/am.rs index 005a9ff6..545974a6 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -145,7 +145,7 @@ pub unsafe extern "C" fn ambuild( opfamily: Opfamily, } impl HeapRelation for Heap { - fn traverse(&self, callback: F) + fn traverse(&self, progress: bool, callback: F) where F: FnMut((Pointer, Vec)), { @@ -191,7 +191,7 @@ pub unsafe extern "C" fn ambuild( self.index_info, true, false, - true, + progress, 0, pgrx::pg_sys::InvalidBlockNumber, Some(call::), @@ -246,6 +246,15 @@ pub unsafe extern "C" fn ambuild( unsafe { RabbitholeLeader::enter(heap, index, (*index_info).ii_Concurrent) } { unsafe { + parallel_build( + index, + heap, + index_info, + leader.tablescandesc, + leader.rabbitholeshared, + true, + ); + leader.wait(); let nparticipants = leader.nparticipants; loop { pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.rabbitholeshared).mutex); @@ -264,7 +273,7 @@ pub unsafe extern "C" fn ambuild( } else { let mut tuples_done = 0; reporter.tuples_done(tuples_done); - heap_relation.traverse(|(payload, vector)| { + heap_relation.traverse(true, |(payload, vector)| { algorithm::insert::insert( index_relation.clone(), payload, @@ -306,6 +315,7 @@ struct RabbitholeLeader { pcxt: *mut pgrx::pg_sys::ParallelContext, nparticipants: i32, rabbitholeshared: *mut RabbitholeShared, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, snapshot: pgrx::pg_sys::Snapshot, } @@ -417,10 +427,10 @@ impl RabbitholeLeader { pgrx::pg_sys::LaunchParallelWorkers(pcxt); } - let nparticipants = unsafe { (*pcxt).nworkers_launched }; + let nworkers_launched = unsafe { (*pcxt).nworkers_launched }; unsafe { - if nparticipants == 0 { + if nworkers_launched == 0 { pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); if is_mvcc_snapshot(snapshot) { pgrx::pg_sys::UnregisterSnapshot(snapshot); @@ -430,16 +440,21 @@ impl RabbitholeLeader { return None; } } - unsafe { - pgrx::pg_sys::WaitForParallelWorkersToAttach(pcxt); - } + Some(Self { pcxt, - nparticipants, + nparticipants: nworkers_launched + 1, rabbitholeshared, + tablescandesc, snapshot, }) } + + pub fn wait(&self) { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToAttach(self.pcxt); + } + } } impl Drop for RabbitholeLeader { @@ -486,6 +501,31 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( (*index_info).ii_Concurrent = (*rabbitholeshared).isconcurrent; } + unsafe { + parallel_build( + index, + heap, + index_info, + tablescandesc, + rabbitholeshared, + false, + ); + } + + unsafe { + pgrx::pg_sys::index_close(index, index_lockmode); + pgrx::pg_sys::table_close(heap, heap_lockmode); + } +} + +unsafe fn parallel_build( + index: *mut pgrx::pg_sys::RelationData, + heap: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + rabbitholeshared: *mut RabbitholeShared, + progress: bool, +) { #[derive(Debug, Clone)] pub struct Heap { heap: pgrx::pg_sys::Relation, @@ -495,7 +535,7 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( scan: *mut pgrx::pg_sys::TableScanDescData, } impl HeapRelation for Heap { - fn traverse(&self, callback: F) + fn traverse(&self, progress: bool, callback: F) where F: FnMut((Pointer, Vec)), { @@ -541,7 +581,7 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( self.index_info, true, false, - true, + progress, 0, pgrx::pg_sys::InvalidBlockNumber, Some(call::), @@ -566,7 +606,7 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( opfamily, scan, }; - heap_relation.traverse(|(payload, vector)| { + heap_relation.traverse(progress, |(payload, vector)| { algorithm::insert::insert( index_relation.clone(), payload, @@ -581,11 +621,6 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( pgrx::pg_sys::SpinLockRelease(&raw mut (*rabbitholeshared).mutex); pgrx::pg_sys::ConditionVariableSignal(&raw mut (*rabbitholeshared).workersdonecv); } - - unsafe { - pgrx::pg_sys::index_close(index, index_lockmode); - pgrx::pg_sys::table_close(heap, heap_lockmode); - } } #[pgrx::pg_guard] From e434d121c8c28f4369eff868cd139fa5c6a3b519 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Fri, 8 Nov 2024 09:11:10 +0800 Subject: [PATCH 026/324] feat: nprob and parallel sample (#46) * feat: nprob and parallel sample Signed-off-by: cutecutecat * chunk insert at index and norm at train Signed-off-by: cutecutecat --------- Signed-off-by: cutecutecat --- bench/README.md | 27 +++++--------- bench/bench.py | 5 ++- bench/index.py | 42 ++++++++++++++------- bench/train.py | 97 +++++++++++++++++++++++++++++++++++++++++++++---- 4 files changed, 132 insertions(+), 39 deletions(-) diff --git a/bench/README.md b/bench/README.md index c0ae63c0..4f3643c3 100644 --- a/bench/README.md +++ b/bench/README.md @@ -27,21 +27,6 @@ PGPASSWORD=123 psql -h 127.0.0.1 -U postgres -c "CREATE USER bench WITH PASSWORD PGPASSWORD=123 psql -h 127.0.0.1 -U postgres -c "ALTER ROLE bench SUPERUSER;" ``` -## Download Data - -```shell -export AWS_ACCESS_KEY_ID=xxx -export AWS_SECRET_ACCESS_KEY=xxx - -aws s3 cp s3://pgvecto.rs-bench/sift_128_1m/sift.hdf5 $(whoami)/sift/sift.hdf5 -aws s3 cp s3://pgvecto.rs-bench/gist_960_1m/gist.hdf5 $(whoami)/gist/gist.hdf5 -aws s3 cp s3://pgvecto.rs-bench/glove_200_1m/glove.hdf5 $(whoami)/glove/glove.hdf5 -aws s3 cp s3://pgvecto.rs-bench/openai_1536_500k/openai.hdf5 $(whoami)/openai/openai.hdf5 -aws s3 cp s3://pgvecto.rs-bench/cohere-768-1m-2022/cohere.hdf5 $(whoami)/cohere_1m_22/cohere_1m_22.hdf5 -aws s3 cp s3://pgvecto.rs-bench/cohere-1024-1m-2023/cohere-1m-23.hdf5 $(whoami)/cohere_1m_23/cohere_1m_23.hdf5 -aws s3 cp s3://pgvecto.rs-bench/cohere-1024-10m-2023/cohere-10m-23.hdf5 $(whoami)/cohere_10m_23/cohere_10m_23.hdf5 -``` - ## Run Bench Options for `-n`: @@ -56,15 +41,21 @@ Options for `-n`: ```shell # pip install pgvector numpy faiss-cpu psycopg h5py tqdm +# If using GPU for train.py: +# conda install pytorch::faiss-gpu + # dump table embedding column to a local h5 file["train"] python dump.py -n sift -o sift.h5 -c embedding -d 128 # external k-means -python train.py -i sift.hdf5 -o sift_centroids_4096 +python train.py -i sift.hdf5 -o sift_centroids_4096 -m l2 # build index (w/wo external centroids) -python index.py -n sift -c sift_centroids_4096.npy -i sift.hdf5 +## with external centroids +python index.py -n sift -c sift_centroids_4096.npy -i sift.hdf5 -m l2 -p 123 -k 4096 -d 768 -w 4 +## without external centroids +## python index.py -n sift -i sift.hdf5 -m l2 -p 123 -k 4096 -d 768 -w 4 # bench -python bench.py -n sift -i sift.hdf5 +python bench.py -n sift -i sift.hdf5 --nprob 100 ``` diff --git a/bench/bench.py b/bench/bench.py index 26423b5d..ced1c53a 100644 --- a/bench/bench.py +++ b/bench/bench.py @@ -22,6 +22,9 @@ def build_arg_parse(): parser.add_argument( "-p", "--password", help="Database password", default="password" ) + parser.add_argument( + "--nprob", help="argument rabbithole.nprobe for query", default=300, type=int + ) return parser @@ -84,6 +87,6 @@ def bench(name, test, answer, metric_ops, conn): else: raise ValueError conn = create_connection(args.password) - conn.execute("SET rabbithole.nprobe=300") + conn.execute(f"SET rabbithole.nprobe={args.nprob}") bench(args.name, test, answer, metric_ops, conn) diff --git a/bench/index.py b/bench/index.py index bd0322d1..871a75cf 100644 --- a/bench/index.py +++ b/bench/index.py @@ -1,4 +1,5 @@ import asyncio +import math from time import perf_counter import argparse from pathlib import Path @@ -16,6 +17,7 @@ "keepalives_interval": 5, "keepalives_count": 5, } +CHUNKS = 10 def build_arg_parse(): @@ -45,6 +47,12 @@ def build_arg_parse(): required=False, default=max(multiprocessing.cpu_count() - 1, 1), ) + parser.add_argument( + "--chunks", + help="chunks for in-memory mode. If OOM, increase it", + type=int, + default=CHUNKS, + ) return parser @@ -113,21 +121,29 @@ async def add_centroids(conn, name, centroids): await asyncio.sleep(0) -async def add_embeddings(conn, name, dim, train): +async def add_embeddings(conn, name, dim, train, chunks): await conn.execute(f"DROP TABLE IF EXISTS {name}") await conn.execute(f"CREATE TABLE {name} (id integer, embedding vector({dim}))") - async with conn.cursor().copy( - f"COPY {name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)" - ) as copy: - copy.set_types(["integer", "vector"]) + n, dim = train.shape + chunk_size = math.ceil(n / chunks) + pbar = tqdm(desc="Adding embeddings", total=n) + for i in range(chunks): + chunk_start = i * chunk_size + chunk_len = min(chunk_size, n - i * chunk_size) + data = train[chunk_start : chunk_start + chunk_len] - for i, vec in tqdm( - enumerate(train), desc="Adding embeddings", total=len(train) - ): - await copy.write_row((i, vec)) - while conn.pgconn.flush() == 1: - await asyncio.sleep(0) + async with conn.cursor().copy( + f"COPY {name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)" + ) as copy: + copy.set_types(["integer", "vector"]) + + for i, vec in enumerate(data): + await copy.write_row((chunk_start + i, vec)) + while conn.pgconn.flush() == 1: + await asyncio.sleep(0) + pbar.update(chunk_len) + pbar.close() async def build_index( @@ -168,7 +184,7 @@ async def monitor_index_build(conn, finish: asyncio.Event): async def main(dataset): dataset = h5py.File(Path(args.input), "r") - url = f"postgresql://postgres:{args.password}@localhost:5432/postgres", + url = f"postgresql://postgres:{args.password}@localhost:5432/postgres" conn = await create_connection(url) if args.centroids: centroids = np.load(args.centroids, allow_pickle=False) @@ -176,7 +192,7 @@ async def main(dataset): metric_ops, ivf_config = get_ivf_ops_config( args.metric, args.k, args.name if args.centroids else None ) - await add_embeddings(conn, args.name, args.dim, dataset["train"]) + await add_embeddings(conn, args.name, args.dim, dataset["train"], args.chunks) index_finish = asyncio.Event() # Need a separate connection for monitor process diff --git a/bench/train.py b/bench/train.py index 43339926..26403000 100644 --- a/bench/train.py +++ b/bench/train.py @@ -1,7 +1,11 @@ +import itertools +from multiprocessing import Pool, cpu_count from time import perf_counter import argparse from pathlib import Path from sys import version_info +from tqdm import tqdm +from numpy import linalg as LA if version_info >= (3, 12): raise RuntimeError("h5py doesn't support 3.12") @@ -13,6 +17,7 @@ DEFAULT_K = 4096 N_ITER = 25 +CHUNKS = 10 SEED = 42 MAX_POINTS_PER_CLUSTER = 256 @@ -31,42 +36,112 @@ def build_arg_parse(): parser.add_argument( "--niter", help="number of iterations", type=int, default=N_ITER ) - parser.add_argument("-m", "--metric", choices=["l2", "cos"], default="l2") + parser.add_argument("-m", "--metric", choices=["l2", "cos", "dot"], default="l2") parser.add_argument( "-g", "--gpu", help="enable GPU for KMeans", action="store_true" ) + parser.add_argument( + "--in-memory", + help="use numpy in-memory mode sampling, quicker for large dataset", + action="store_true", + ) + parser.add_argument( + "--chunks", + help="chunks for in-memory mode. If OOM, increase it", + type=int, + default=CHUNKS, + ) return parser def reservoir_sampling(iterator, k: int): """Reservoir sampling from an iterator.""" res = [] - while len(res) < k: + for _ in tqdm(range(k), desc="Collect train subset"): try: res.append(next(iterator)) except StopIteration: return np.vstack(res) - for i, vec in enumerate(iterator, k + 1): + for i, vec in tqdm(enumerate(iterator, k + 1), total=k, desc="Random Pick"): j = np.random.randint(0, i) if j < k: res[j] = vec return np.vstack(res) +def _slice_chunk(args: tuple[int, str, np.ndarray]): + k, file_path, chunk, start_idx = args + dataset = h5py.File(Path(file_path), "r") + data = dataset["train"] + start, end = min(chunk), max(chunk) + indexes = [c - start for c in chunk] + source = data[start : end + 1] + select = source[indexes] + delta, dim = select.shape + + output = np.memmap("index.mmap", dtype=np.float32, mode="r+", shape=(k, dim)) + output[start_idx : start_idx + delta, :] = select + output.flush() + + +def reservoir_sampling_np(data, file_path, k: int, chunks: int): + """Reservoir sampling in memory by numpy.""" + index = np.random.permutation(len(data))[:k] + indices = np.sort(index) + num_processes = cpu_count() + # Split indices into chunks for parallel processing + index_chunks = np.array_split(indices, chunks) + _, dim = data.shape + np.memmap("index.mmap", dtype=np.float32, mode="w+", shape=(k, dim)) + # Create arguments for each chunk + start_idx_acu = [0] + start_idx_acu.extend( + list(itertools.accumulate([len(c) for c in index_chunks[:-1]])) + ) + chunk_args = [ + (k, file_path, chunk, start_idx_acu[i]) for i, chunk in enumerate(index_chunks) + ] + # Process chunks in parallel + with Pool(processes=num_processes) as pool: + list(pool.map(_slice_chunk, chunk_args)) + + def filter_by_label(iter, labels, target): for i, vec in enumerate(iter): if labels[i] == target: yield vec -def kmeans_cluster(data, k, child_k, niter, metric, gpu=False): +def kmeans_cluster( + data, + file_path, + k, + child_k, + niter, + metric, + gpu=False, + in_memory=False, + chunks=CHUNKS, +): n, dim = data.shape - if n > MAX_POINTS_PER_CLUSTER * k: + if n > MAX_POINTS_PER_CLUSTER * k and not in_memory: train = reservoir_sampling(iter(data), MAX_POINTS_PER_CLUSTER * args.k) + elif n > MAX_POINTS_PER_CLUSTER * k and in_memory: + reservoir_sampling_np(data, file_path, MAX_POINTS_PER_CLUSTER * args.k, chunks) + train = np.array( + np.memmap( + "index.mmap", + dtype=np.float32, + mode="r", + shape=(MAX_POINTS_PER_CLUSTER * k, dim), + ) + ) else: train = data[:] + if metric == "cos": + train = train / LA.norm(train, axis=1, keepdims=True) kmeans = Kmeans( - dim, k, gpu=gpu, verbose=True, niter=niter, seed=SEED, spherical=metric == "cos" + dim, k, gpu=gpu, verbose=True, niter=niter, seed=SEED, spherical=metric != "l2" ) kmeans.train(train) if not child_k: @@ -109,7 +184,15 @@ def kmeans_cluster(data, k, child_k, niter, metric, gpu=False): start_time = perf_counter() centroids = kmeans_cluster( - dataset["train"], args.k, args.child_k, args.niter, args.metric, args.gpu + dataset["train"], + args.input, + args.k, + args.child_k, + args.niter, + args.metric, + args.gpu, + args.in_memory, + args.chunks, ) print(f"K-means (k=({args.k}, {args.child_k})): {perf_counter() - start_time:.2f}s") From c7843fb0d140da14e37ddea3995b722d56714580 Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 8 Nov 2024 11:53:12 +0800 Subject: [PATCH 027/324] feat: add sqllogictest (#47) * add sqllogictest Signed-off-by: Keming * fix matrix variables Signed-off-by: Keming * chmod 777 Signed-off-by: Keming * write to ./target/schema.sql directly Signed-off-by: Keming * use profile in package Signed-off-by: Keming * fix pgvector image Signed-off-by: Keming * fix arg scope in dockerfile Signed-off-by: Keming * fix docker run --name Signed-off-by: Keming * add more tests Signed-off-by: Keming * fix partition test Signed-off-by: Keming * fix new config for index Signed-off-by: Keming --------- Signed-off-by: Keming Co-authored-by: Jinjing Zhou --- .github/workflows/psql.yml | 95 +++++++++++++++++++++ bench/README.md | 7 +- docker/Dockerfile | 11 ++- docker/pgrx.Dockerfile | 3 +- tests/README.md | 14 ++++ tests/logic/index.slt | 61 ++++++++++++++ tests/logic/issue427.slt | 19 +++++ tests/logic/null.fail | 39 +++++++++ tests/logic/partition.slt | 51 +++++++++++ tests/logic/pushdown_plan.fail | 144 ++++++++++++++++++++++++++++++++ tests/logic/pushdown_range.fail | 102 ++++++++++++++++++++++ tests/logic/reindex.slt | 36 ++++++++ tests/logic/vector.slt | 27 ++++++ tools/package.sh | 5 +- tools/schema.sh | 2 +- 15 files changed, 606 insertions(+), 10 deletions(-) create mode 100644 .github/workflows/psql.yml create mode 100644 tests/README.md create mode 100644 tests/logic/index.slt create mode 100644 tests/logic/issue427.slt create mode 100644 tests/logic/null.fail create mode 100644 tests/logic/partition.slt create mode 100644 tests/logic/pushdown_plan.fail create mode 100644 tests/logic/pushdown_range.fail create mode 100644 tests/logic/reindex.slt create mode 100644 tests/logic/vector.slt diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml new file mode 100644 index 00000000..56feedc8 --- /dev/null +++ b/.github/workflows/psql.yml @@ -0,0 +1,95 @@ +name: PostgresSQL + +on: + pull_request: + paths: + - '.github/workflows/rust.yml' + - 'src/**' + - 'Cargo.lock' + - 'Cargo.toml' + - '*.control' + - 'rust-toolchain.toml' + - 'tests/**' + - 'tools/**' + push: + branches: + - main + paths: + - '.github/workflows/rust.yml' + - 'src/**' + - 'Cargo.lock' + - 'Cargo.toml' + - '*.control' + - 'rust-toolchain.toml' + - 'tests/**' + - 'tools/**' + merge_group: + workflow_dispatch: + +concurrency: + group: ${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +env: + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 + RUSTFLAGS: "-Dwarnings" + CARGO_PROFILE_OPT_BUILD_OVERRIDE_DEBUG: true + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + version: ["14", "15", "16", "17"] + arch: ["x86_64"] + env: + PGRX_IMAGE: "kemingy/pgrx:0.12.8" + SQLLOGICTEST: "0.22.0" + + steps: + - uses: actions/checkout@v4 + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + const url = process.env.ACTIONS_CACHE_URL || ''; + const token = process.env.ACTIONS_RUNTIME_TOKEN || ''; + core.exportVariable( + 'CACHE_ENVS', + `-e CARGO_INCREMENTAL=0 -e SCCACHE_GHA_ENABLED=true -e RUSTC_WRAPPER=sccache -e ACTIONS_CACHE_URL=${url} -e ACTIONS_RUNTIME_TOKEN=${token}`, + ); + - name: Set up pgrx docker images and permissions + run: | + docker pull $PGRX_IMAGE + echo "Default user: $(id -u):$(id -g)" + sudo chmod -R 777 . + + - name: Build + env: + SEMVER: "0.0.0" + VERSION: ${{ matrix.version }} + ARCH: ${{ matrix.arch }} + PLATFORM: "amd64" + PROFILE: "opt" + run: | + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE build --lib --features pg${{ matrix.version }} --target ${{ matrix.arch }}-unknown-linux-gnu --profile opt + docker run --rm -v .:/workspace $CACHE_ENVS --entrypoint bash $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --target ${{ matrix.arch }}-unknown-linux-gnu --profile $PROFILE + ./tools/package.sh + docker build -t rabbithole:pg${{ matrix.version }} --build-arg PG_VERSION=${{ matrix.version }} -f ./docker/Dockerfile . + + - name: Setup SQL Logic Test + run: | + curl -fsSL -o sqllogictest.tar.gz https://github.com/risinglightdb/sqllogictest-rs/releases/download/v${SQLLOGICTEST}/sqllogictest-bin-v${SQLLOGICTEST}-${{ matrix.arch }}-unknown-linux-musl.tar.gz + tar -xzf sqllogictest.tar.gz + mv sqllogictest /usr/local/bin/ + + - name: SQL Test + env: + PGPASSWORD: postgres + run: | + docker run --rm --name test -d -e POSTGRES_PASSWORD=${PGPASSWORD} -p 5432:5432 rabbithole:pg${{ matrix.version }} + sleep 5 + psql -h localhost -U postgres -d postgres -c 'CREATE EXTENSION IF NOT EXISTS rabbithole CASCADE;' + sqllogictest './tests/**/*.slt' + docker stop test diff --git a/bench/README.md b/bench/README.md index 4f3643c3..07d02d50 100644 --- a/bench/README.md +++ b/bench/README.md @@ -4,16 +4,17 @@ sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config clang cargo install --locked cargo-pgrx cargo pgrx init -cargo build --package rabbithole --lib --features pg16 --target x86_64-unknown-linux-gnu --release -./tools/schema.sh --features pg16 --target x86_64-unknown-linux-gnu --release | expand -t 4 > ./target/schema.sql +cargo build --package rabbithole --lib --features pg16 --target x86_64-unknown-linux-gnu --profile opt +./tools/schema.sh --features pg16 --target x86_64-unknown-linux-gnu --profile opt export SEMVER="0.0.0" export VERSION="16" export ARCH="x86_64" export PLATFORM="amd64" +export PROFILE="opt" ./tools/package.sh -docker build -t rabbithole:pg16-latest -f ./docker/Dockerfile . +docker build -t rabbithole:pg16-latest --build-arg PG_VERSION=16 -f ./docker/Dockerfile . ``` Or you can use `starkind/rabbithole:pg16-latest` to run the bench. diff --git a/docker/Dockerfile b/docker/Dockerfile index cf3d63f1..15f3e7f9 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,6 +1,11 @@ -FROM pgvector/pgvector:0.7.4-pg16 +ARG PG_VERSION=17 +ARG PGVECTOR=0.8.0 -COPY ./build/rabbithole-pg16_0.0.0_amd64.deb /tmp/rabbithole.deb +FROM pgvector/pgvector:${PGVECTOR}-pg${PG_VERSION} + +ARG PG_VERSION +RUN echo ${PG_VERSION} +COPY ./build/rabbithole-pg${PG_VERSION}_0.0.0_amd64.deb /tmp/rabbithole.deb RUN apt-get install -y /tmp/rabbithole.deb && rm -f /tmp/rabbithole.deb -CMD ["postgres", "-c" ,"shared_preload_libraries=vector,rabbithole.so", "-c", "search_path=\"$user\", public, rabbithole"] \ No newline at end of file +CMD ["postgres", "-c" ,"shared_preload_libraries=rabbithole.so", "-c", "search_path=\"$user\", public, rabbithole"] diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile index 920817b5..02941aa5 100644 --- a/docker/pgrx.Dockerfile +++ b/docker/pgrx.Dockerfile @@ -34,12 +34,13 @@ RUN /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y # install all the PostgresQL RUN set -ex; \ for v in $(seq 14 17); do \ - apt install -y --no-install-recommends postgresql-$v postgresql-server-dev-$v; \ + apt install -y --no-install-recommends postgresql-$v postgresql-server-dev-$v postgresql-$v-pgvector; \ done; \ rm -rf /var/lib/apt/lists/*; # create a non-root user (make it compatible with Ubuntu 24.04) RUN useradd -u 1000 -U -m ubuntu +RUN chown -R ubuntu:ubuntu /usr/share/postgresql/ /usr/lib/postgresql/ USER ubuntu ENV PATH="$PATH:/home/ubuntu/.cargo/bin" RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- --default-toolchain=none -y diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 00000000..c3217e8f --- /dev/null +++ b/tests/README.md @@ -0,0 +1,14 @@ +# SQL logic test + +Requires `sqllogictest` to be installed: + +```bash +cargo install sqllogictest-bin +``` + +To run all the tests: + +```bash +PGPASSWORD=postgres psql -h localhost -U postgres -d postgres -c 'CREATE EXTENSION IF NOT EXISTS rabbithole CASCADE;' +sqllogictest './tests/**/*.slt' +``` diff --git a/tests/logic/index.slt b/tests/logic/index.slt new file mode 100644 index 00000000..cccb64c7 --- /dev/null +++ b/tests/logic/index.slt @@ -0,0 +1,61 @@ +statement ok +SET search_path TO public, rabbithole, pg_temp; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement error +CREATE INDEX ON t USING rabbithole (val vector_l2_ops) +WITH (options = $$ +unknown_options=true +$$); + +# multiple index on single column +statement ok +CREATE INDEX ON t USING rabbithole (val vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.internal] +nlist = 32 +spherical_centroids = false +$$); + +statement ok +CREATE INDEX ON t USING rabbithole (val vector_ip_ops) +WITH (options = $$ +residual_quantization = false +[build.internal] +nlist = 32 +spherical_centroids = true +$$); + +statement ok +CREATE INDEX ON t USING rabbithole (val vector_cosine_ops) +WITH (options = $$ +residual_quantization = false +[build.internal] +nlist = 32 +spherical_centroids = true +$$); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +---- +DROP TABLE t; diff --git a/tests/logic/issue427.slt b/tests/logic/issue427.slt new file mode 100644 index 00000000..f3569900 --- /dev/null +++ b/tests/logic/issue427.slt @@ -0,0 +1,19 @@ +# https://github.com/tensorchord/pgvecto.rs/issues/427 + +statement ok +SET search_path TO public, rabbithole, pg_temp; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT NULL::vector FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING rabbithole (val vector_l2_ops); + +statement ok +SELECT val FROM t ORDER BY val <-> (SELECT val FROM t LIMIT 1) limit 10; + +statement ok +DROP TABLE t; diff --git a/tests/logic/null.fail b/tests/logic/null.fail new file mode 100644 index 00000000..458382ce --- /dev/null +++ b/tests/logic/null.fail @@ -0,0 +1,39 @@ +statement ok +SET search_path TO public, rabbithole, pg_temp; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +INSERT INTO t (val) SELECT ('[NaN, Infinity, -Infinity]') FROM generate_series(1, 100); + +statement ok +INSERT INTO t (val) SELECT (NULL) FROM generate_series(1, 100); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +CREATE INDEX rabitq ON t USING rabbithole (val vector_l2_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +REINDEX INDEX rabitq; + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +---- +DROP TABLE t; diff --git a/tests/logic/partition.slt b/tests/logic/partition.slt new file mode 100644 index 00000000..e7811ec5 --- /dev/null +++ b/tests/logic/partition.slt @@ -0,0 +1,51 @@ +statement ok +SET search_path TO public, rabbithole, pg_temp; + +# partition table +statement ok +CREATE TABLE items (val vector(3), category_id int) PARTITION BY LIST(category_id); + +statement ok +CREATE TABLE id_123 PARTITION OF items FOR VALUES IN (1, 2, 3); + +statement ok +CREATE TABLE id_456 PARTITION OF items FOR VALUES IN (4, 5, 6); + +statement ok +CREATE TABLE id_789 PARTITION OF items FOR VALUES IN (7, 8, 9); + +statement ok +INSERT INTO items (val, category_id) +SELECT + ARRAY[random(), random(), random()]::real[], + (random() * 6 + 1)::int +FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON items USING rabbithole (val rabbithole.vector_l2_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +CREATE INDEX ON id_123 USING rabbithole (val vector_cosine_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +# partial index +statement ok +CREATE INDEX ON items USING rabbithole (val rabbithole.vector_ip_ops) WHERE (category_id = 1); + +query I +SELECT COUNT(1) FROM +(SELECT 1 FROM items WHERE (category_id = 1) ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE id_789, id_456, id_123, items; diff --git a/tests/logic/pushdown_plan.fail b/tests/logic/pushdown_plan.fail new file mode 100644 index 00000000..0cdabdef --- /dev/null +++ b/tests/logic/pushdown_plan.fail @@ -0,0 +1,144 @@ +statement ok +SET search_path TO public, rabbithole, pg_temp; + +statement ok +CREATE TABLE t (val0 vector(3), val1 halfvec(3), val2 sparsevec(3)); + +statement ok +INSERT INTO t (val0, val1, val2) +SELECT + ARRAY[random(), random(), random()]::real[]::vector, + ARRAY[random(), random(), random()]::real[]::vector::halfvec, + ARRAY[random(), random(), random()]::real[]::vector::sparsevec +FROM generate_series(1, 10000); + +statement ok +CREATE INDEX ind0 ON t USING rabbithole (val0 vector_l2_ops); + +# statement ok +# CREATE INDEX ind1 ON t USING rabbithole (val1 halfvec_dot_ops); + +# 1 vector key + 1 corresponding order_by key + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) ORDER BY val0 <-> '[0, 0, 0]'; +---- + Index Scan using ind0 on t + Index Cond: (val0 <<->> '("[0,0,0]",1)'::sphere_vector) + Order By: (val0 <-> '[0,0,0]'::vector) + +# 1 vector key + 0 order_by key + original style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +# 1 vector key + 0 order_by key + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1); +---- + Index Scan using ind0 on t + Index Cond: (val0 <<->> '("[0,0,0]",1)'::sphere_vector) + +# 0 vector key + 1 order_by key +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t ORDER BY val0 <-> '[0, 0, 0]'; +---- + Index Scan using ind0 on t + Order By: (val0 <-> '[0,0,0]'::vector) + +# 2 vector key(1 of them is corresponding) + 1 order_by key + original style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1 +AND val1 <#> '[0, 0, 0]' < 1 +ORDER BY val0 <-> '[0, 0, 0]'; +---- + Index Scan using ind0 on t + Order By: (val0 <-> '[0,0,0]'::vector) + Filter: (((val0 <-> '[0,0,0]'::vector) < '1'::double precision) AND ((val1 <#> '[0,0,0]'::halfvec) < '1'::double precision)) + +# 2 vector key(1 of them is corresponding) + 1 order_by key + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) +AND val1 <<#>> sphere('[0, 0, 0]'::halfvec, 1) +ORDER BY val0 <-> '[0, 0, 0]'; +---- + Index Scan using ind0 on t + Index Cond: (val0 <<->> '("[0, 0, 0]",1)'::sphere_vector) + Order By: (val0 <-> '[0, 0, 0]'::vector) + Filter: (val1 <<#>> '("[0, 0, 0]",1)'::sphere_vecf16) + +# 2 vector key(none of them is corresponding) + 1 order_by key + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val2 <<->> sphere('{}/3'::svector, 1) +AND val1 <<#>> sphere('[0, 0, 0]'::halfvec, 1) +ORDER BY val0 <-> '[0, 0, 0]'; +---- + Index Scan using ind0 on t + Order By: (val0 <-> '[0, 0, 0]'::vector) + Filter: ((val2 <<->> '({}/3,1)'::sphere_svector) AND (val1 <<#>> '("[0, 0, 0]",1)'::sphere_vecf16)) + +# 2 vector keys(both indexed) + 0 order_by key + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) +AND val1 <<#>> sphere('[0, 0, 0]'::halfvec, 1); +---- + Index Scan using ind1 on t + Index Cond: (val1 <<#>> '("[0, 0, 0]",1)'::sphere_vecf16) + Filter: (val0 <<->> '("[0, 0, 0]",1)'::sphere_vector) + +# 2 vector keys(both not indexed) + 0 order_by key + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<#>> sphere('[0, 0, 0]'::vector, 1) +AND val1 <<->> sphere('[0, 0, 0]'::halfvec, 1); +---- + Seq Scan on t + Filter: ((val0 <<#>> '("[0, 0, 0]",1)'::sphere_vector) AND (val1 <<->> '("[0, 0, 0]",1)'::sphere_vecf16)) + +# 2 vector key(1 indexed) + 0 order_by key + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) +AND val1 <<->> sphere('[0, 0, 0]'::halfvec, 1); +---- + Index Scan using ind0 on t + Index Cond: (val0 <<->> '("[0, 0, 0]",1)'::sphere_vector) + Filter: (val1 <<->> '("[0, 0, 0]",1)'::sphere_vecf16) + +# 1 vector key + 1 not corresponding order_by key(operator) + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<#>> sphere('[0, 0, 0]'::vector, 1) +ORDER BY val0 <-> '[0, 0, 0]'; +---- + Index Scan using ind0 on t + Order By: (val0 <-> '[0, 0, 0]'::vector) + Filter: (val0 <<#>> '("[0, 0, 0]",1)'::sphere_vector) + +# 1 vector key + 1 not corresponding order_by key(variable) + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) +ORDER BY val1 <#> '[1, 1, 1]'; +---- + Index Scan using ind1 on t + Order By: (val1 <#> '[1, 1, 1]'::halfvec) + Filter: (val0 <<->> '("[0, 0, 0]",1)'::sphere_vector) + +# 0 vector key + 0 order_by key(variable) +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +statement ok +DROP TABLE t; diff --git a/tests/logic/pushdown_range.fail b/tests/logic/pushdown_range.fail new file mode 100644 index 00000000..b13e6505 --- /dev/null +++ b/tests/logic/pushdown_range.fail @@ -0,0 +1,102 @@ +statement ok +SET search_path TO public, rabbithole, pg_temp; + +statement ok +CREATE TABLE t (val0 vector(3), val1 halfvec(3), val2 sparsevec(3), val3 bit(3)); + +statement ok +INSERT INTO t (val0, val1, val2, val3) VALUES + ('[0.1, 0.1, 0.1]', '[0.1, 0.1, 0.1]', '{1:-0.1, 2:0.1, 3:0.1}/3', '000'), + ('[0.2, 0.2, 0.2]', '[-0.2, 0.2, 0.2]', '{1:0.2, 2:-0.2, 3:0.2}/3', '001'), + ('[0.3, 0.3, 0.3]', '[0.3, 0.3, -0.3]', '{1:-0.3, 2:0.3, 3:-0.3}/3', '110'), + ('[0.4, 0.4, 0.4]', '[0.4, -0.4, 0.4]', '{1:-0.4, 2:-0.4, 3:-0.4}/3', '111'); + +statement ok +CREATE INDEX ON t USING rabbithole (val0 vector_l2_ops); + +# original style +query I +SELECT val0 FROM t WHERE val0 <-> '[0.24, 0.24, 0.24]' < 0.12 ORDER BY val0 <-> '[0.24, 0.24, 0.24]'; +---- +[0.2, 0.2, 0.2] +[0.3, 0.3, 0.3] + +# sphere style +query I +SELECT val0 FROM t WHERE val0 <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012) ORDER BY val0 <-> '[0.24, 0.24, 0.24]'; +---- +[0.2, 0.2, 0.2] +[0.3, 0.3, 0.3] + +statement ok +CREATE INDEX ON t USING rabbithole (val1 halfvec_dot_ops); + +# original style +query I +SELECT val1 FROM t WHERE val1 <#> '[0.24, -0.24, 0.24]' < 0 ORDER BY val1 <#> '[0.24, -0.24, 0.24]'; +---- +[0.39990234, -0.39990234, 0.39990234] +[0.099975586, 0.099975586, 0.099975586] + +# sphere style +query I +SELECT val1 FROM t WHERE val1 <<#>> sphere('[0.24, -0.24, 0.24]'::halfvec, 0) ORDER BY val1 <#> '[0.24, -0.24, 0.24]'; +---- +[0.39990234, -0.39990234, 0.39990234] +[0.099975586, 0.099975586, 0.099975586] + +statement ok +CREATE INDEX ON t USING rabbithole (val2 sparsevec_cos_ops) +WITH (options = "[indexing.hnsw]"); + +# original style +query I +SELECT val2 FROM t WHERE val2 <=> '{0:0.12, 1:0.24, 2:0.36}/3' < 1 ORDER BY val2 <=> '{0:0.12, 1:0.24, 2:0.36}/3'; +---- +{0:-0.1, 1:0.1, 2:0.1}/3 +{0:0.2, 1:-0.2, 2:0.2}/3 + +# sphere style +query I +SELECT val2 FROM t WHERE val2 <<=>> sphere('{0:0.12, 1:0.24, 2:0.36}/3'::sparsevec, 1) +ORDER BY val2 <=> '{0:0.12, 1:0.24, 2:0.36}/3'; +---- +{0:-0.1, 1:0.1, 2:0.1}/3 +{0:0.2, 1:-0.2, 2:0.2}/3 + +statement ok +CREATE INDEX ON t USING rabbithole (val3 bit_jaccard_ops) +WITH (options = "[indexing.hnsw]"); + +# original style +query I +SELECT val3 FROM t WHERE val3 <~> '[1, 1, 1]' <= 0.4 ORDER BY val3 <~> '[1, 1, 1]'; +---- +[1, 1, 1] +[1, 1, 0] + +# sphere style +query I +SELECT val3 FROM t WHERE val3 <<~>> sphere('[1, 1, 1]'::bit, 0.4) ORDER BY val3 <~> '[1, 1, 1]'; +---- +[1, 1, 1] +[1, 1, 0] + +# sphere style: multiple vector keys and no order-by key +query I +SELECT val0 FROM t WHERE val0 <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012) +AND val1 <<#>> sphere('[0.24, -0.24, 0.24]'::halfvec, 0.05) +ORDER BY val0 <-> '[0.24, 0.24, 0.24]'; +---- +[0.2, 0.2, 0.2] + +# sphere style: vectors in key and order-by key are different +query I +SELECT val0 FROM t WHERE val0 <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012) +ORDER BY val1 <#> '[1, 1, -1]'; +---- +[0.3, 0.3, 0.3] +[0.2, 0.2, 0.2] + +statement ok +DROP TABLE t; diff --git a/tests/logic/reindex.slt b/tests/logic/reindex.slt new file mode 100644 index 00000000..ac0922fd --- /dev/null +++ b/tests/logic/reindex.slt @@ -0,0 +1,36 @@ +statement ok +SET search_path TO public, rabbithole, pg_temp; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING rabbithole (val vector_l2_ops); + +statement ok +REINDEX INDEX t_val_idx; + +statement ok +INSERT INTO t (val) VALUES ('[0.6,0.6,0.6]'); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +REINDEX INDEX CONCURRENTLY t_val_idx; + +statement ok +INSERT INTO t (val) VALUES ('[0.6,0.6,0.6]'); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE t; diff --git a/tests/logic/vector.slt b/tests/logic/vector.slt new file mode 100644 index 00000000..0c507236 --- /dev/null +++ b/tests/logic/vector.slt @@ -0,0 +1,27 @@ +statement ok +SET search_path TO public, rabbithole, pg_temp; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +---- +DROP TABLE t; diff --git a/tools/package.sh b/tools/package.sh index 68ac3a32..82234cd9 100755 --- a/tools/package.sh +++ b/tools/package.sh @@ -5,6 +5,7 @@ printf "SEMVER = ${SEMVER}\n" printf "VERSION = ${VERSION}\n" printf "ARCH = ${ARCH}\n" printf "PLATFORM = ${PLATFORM}\n" +printf "PROFILE = ${PROFILE}\n" rm -rf ./build/dir_zip rm -rf ./build/rabbithole-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip @@ -12,9 +13,9 @@ rm -rf ./build/dir_deb rm -rf ./build/rabbithole-pg${VERSION}_${SEMVER}_${PLATFORM}.deb mkdir -p ./build/dir_zip -cp ./target/schema.sql ./build/dir_zip/rabbithole--$SEMVER.sql +cp ./target/${ARCH}-unknown-linux-gnu/${PROFILE}/schema.sql ./build/dir_zip/rabbithole--$SEMVER.sql sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./rabbithole.control > ./build/dir_zip/rabbithole.control -cp ./target/${ARCH}-unknown-linux-gnu/release/librabbithole.so ./build/dir_zip/rabbithole.so +cp ./target/${ARCH}-unknown-linux-gnu/${PROFILE}/librabbithole.so ./build/dir_zip/rabbithole.so zip ./build/rabbithole-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip -j ./build/dir_zip/* mkdir -p ./build/dir_deb diff --git a/tools/schema.sh b/tools/schema.sh index 1d868ad6..e50ee988 100755 --- a/tools/schema.sh +++ b/tools/schema.sh @@ -44,4 +44,4 @@ CONTROL_FILEPATH="./rabbithole.control" SO_FILEPATH="$DIR/librabbithole.so" $(di PGRX_EMBED=$code cargo rustc --package rabbithole --bin pgrx_embed_rabbithole "$@" -- --cfg pgrx_embed -CARGO_PKG_VERSION="0.0.0" QEMU_LD_PREFIX=$QEMU_LD_PREFIX "${RUNNER[@]}" "$DIR/pgrx_embed_rabbithole" | expand -t 4 \ No newline at end of file +CARGO_PKG_VERSION="0.0.0" QEMU_LD_PREFIX=$QEMU_LD_PREFIX "${RUNNER[@]}" "$DIR/pgrx_embed_rabbithole" | expand -t 4 > $DIR/schema.sql From de16fbce4154e9e49bb1bb4f0b34b6bddf2f8d99 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 8 Nov 2024 16:15:46 +0800 Subject: [PATCH 028/324] fix: support ctrl-c on k-means (#53) Signed-off-by: usamoi --- Cargo.lock | 340 ++++++++++++++++++++++++++++++++++- Cargo.toml | 6 +- src/algorithm/build.rs | 84 ++++----- src/algorithm/k_means.rs | 213 ++++++++++++++++++++++ src/algorithm/mod.rs | 2 + src/algorithm/parallelism.rs | 63 +++++++ src/algorithm/scan.rs | 1 - 7 files changed, 659 insertions(+), 50 deletions(-) create mode 100644 src/algorithm/k_means.rs create mode 100644 src/algorithm/parallelism.rs diff --git a/Cargo.lock b/Cargo.lock index 6805b04a..b79c1b30 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -78,7 +78,7 @@ dependencies = [ "serde", "thiserror", "toml", - "validator", + "validator 0.18.1", ] [[package]] @@ -330,6 +330,17 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "either" version = "1.13.0" @@ -494,6 +505,124 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "icu_collections" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db2fa452206ebee18c4b5c2274dbf1de17008e874b4dc4f0aea9d01ca79e4526" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locid" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13acbb8371917fc971be86fc8057c41a64b521c184808a698c02acc242dbf637" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_locid_transform" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01d11ac35de8e40fdeda00d9e1e9d92525f3f9d887cdd7aa81d727596788b54e" +dependencies = [ + "displaydoc", + "icu_locid", + "icu_locid_transform_data", + "icu_provider", + "tinystr", + "zerovec", +] + +[[package]] +name = "icu_locid_transform_data" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdc8ff3388f852bede6b579ad4e978ab004f139284d7b28715f773507b946f6e" + +[[package]] +name = "icu_normalizer" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19ce3e0da2ec68599d193c93d088142efd7f9c5d6fc9b803774855747dc6a84f" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "utf16_iter", + "utf8_iter", + "write16", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8cafbf7aa791e9b22bec55a167906f9e1215fd475cd22adfcf660e03e989516" + +[[package]] +name = "icu_properties" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93d6020766cfc6302c15dbbc9c8778c37e62c14427cb7f6e601d849e092aeef5" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_locid_transform", + "icu_properties_data", + "icu_provider", + "tinystr", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67a8effbc3dd3e4ba1afa8ad918d5684b8868b3b26500753effea8d2eed19569" + +[[package]] +name = "icu_provider" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ed421c8a8ef78d3e2dbc98a973be2f3770cb42b606e3ab18d6237c4dfde68d9" +dependencies = [ + "displaydoc", + "icu_locid", + "icu_provider_macros", + "stable_deref_trait", + "tinystr", + "writeable", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_provider_macros" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "ident_case" version = "1.0.1" @@ -510,6 +639,27 @@ dependencies = [ "unicode-normalization", ] +[[package]] +name = "idna" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daca1df1c957320b2cf139ac61e7bd64fed304c5040df000a745aa1de3b4ef71" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + [[package]] name = "indenter" version = "0.3.3" @@ -598,6 +748,12 @@ version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" +[[package]] +name = "litemap" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "643cb0b8d4fcc284004d5fd0d67ccf61dfffadb7f75e1e71bc420f4688a3a704" + [[package]] name = "log" version = "0.4.22" @@ -926,6 +1082,28 @@ dependencies = [ "version_check", ] +[[package]] +name = "proc-macro-error-attr2" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" +dependencies = [ + "proc-macro2", + "quote", +] + +[[package]] +name = "proc-macro-error2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" +dependencies = [ + "proc-macro-error-attr2", + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "proc-macro2" version = "1.0.89" @@ -989,7 +1167,8 @@ dependencies = [ "base", "common", "detect", - "k_means", + "half 2.4.1", + "log", "nalgebra", "paste", "pgrx", @@ -998,10 +1177,11 @@ dependencies = [ "rand", "rand_chacha", "rand_distr", + "rayon", "rkyv", "serde", "toml", - "validator", + "validator 0.19.0", ] [[package]] @@ -1295,6 +1475,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" +[[package]] +name = "smallvec" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" + [[package]] name = "smawk" version = "0.3.2" @@ -1360,6 +1546,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "synstructure" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "tap" version = "1.0.1" @@ -1386,6 +1583,16 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "tinystr" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9117f5d4db391c1cf6927e7bea3db74b9a1c1add8f7eda9ffd5364f40f57b82f" +dependencies = [ + "displaydoc", + "zerovec", +] + [[package]] name = "tinyvec" version = "1.8.0" @@ -1493,10 +1700,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22784dbdf76fdde8af1aeda5622b546b422b6fc585325248a2bf9f5e41e94d6c" dependencies = [ "form_urlencoded", - "idna", + "idna 0.5.0", "percent-encoding", ] +[[package]] +name = "utf16_iter" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + [[package]] name = "uuid" version = "1.11.0" @@ -1512,14 +1731,30 @@ version = "0.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db79c75af171630a3148bd3e6d7c4f42b6a9a014c2945bc5ed0020cbb8d9478e" dependencies = [ - "idna", + "idna 0.5.0", "once_cell", "regex", "serde", "serde_derive", "serde_json", "url", - "validator_derive", + "validator_derive 0.18.2", +] + +[[package]] +name = "validator" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0b4a29d8709210980a09379f27ee31549b73292c87ab9899beee1c0d3be6303" +dependencies = [ + "idna 1.0.3", + "once_cell", + "regex", + "serde", + "serde_derive", + "serde_json", + "url", + "validator_derive 0.19.0", ] [[package]] @@ -1536,6 +1771,20 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "validator_derive" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bac855a2ce6f843beb229757e6e570a42e837bcb15e5f449dd48d5747d41bf77" +dependencies = [ + "darling", + "once_cell", + "proc-macro-error2", + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "version_check" version = "0.9.5" @@ -1690,6 +1939,18 @@ dependencies = [ "memchr", ] +[[package]] +name = "write16" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936" + +[[package]] +name = "writeable" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e9df38ee2d2c3c5948ea468a8406ff0db0b29ae1ffde1bcf20ef305bcc95c51" + [[package]] name = "wyz" version = "0.5.1" @@ -1708,6 +1969,30 @@ dependencies = [ "winapi", ] +[[package]] +name = "yoke" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c5b1314b079b0930c31e3af543d8ee1757b1951ae1e1565ec704403a7240ca5" +dependencies = [ + "serde", + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28cc31741b18cb6f1d5ff12f5b7523e3d6eb0852bbbad19d73905511d9849b95" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", + "synstructure", +] + [[package]] name = "zerocopy" version = "0.7.35" @@ -1728,3 +2013,46 @@ dependencies = [ "quote", "syn 2.0.87", ] + +[[package]] +name = "zerofrom" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91ec111ce797d0e0784a1116d0ddcdbea84322cd79e5d5ad173daeba4f93ab55" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ea7b4a3637ea8669cedf0f1fd5c286a17f3de97b8dd5a70a6c167a1730e63a5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", + "synstructure", +] + +[[package]] +name = "zerovec" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa2b893d79df23bfb12d5461018d408ea19dfafe76c2c7ef6d4eba614f8ff079" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", +] diff --git a/Cargo.toml b/Cargo.toml index bd6643f8..808654df 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] +half = "2.4.1" paste = "1" pgrx = { version = "=0.12.8", default-features = false, features = ["cshim"] } pgrx-catalog = "0.1.0" @@ -29,18 +30,19 @@ rand_chacha = "0.3.1" rand_distr = "0.4.3" serde = "1" toml = "0.8.19" -validator = "0.18.1" +validator = { version = "0.19.0", features = ["derive"] } base = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } common = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } detect = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -k_means = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } quantization = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } # lock algebra version forever so that the QR decomposition never changes for same input nalgebra = { version = "=0.33.0", default-features = false } # lock rkyv version forever so that data is always compatible +log = "0.4.22" +rayon = "1.10.0" rkyv = { version = "=0.7.45", features = ["validation"] } [lints] diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index ebd73e25..3224af1f 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -1,3 +1,4 @@ +use crate::algorithm::k_means; use crate::algorithm::rabitq; use crate::algorithm::tuples::*; use crate::index::am_options::Opfamily; @@ -11,11 +12,9 @@ use base::distance::DistanceKind; use base::index::VectorOptions; use base::scalar::ScalarLike; use base::search::Pointer; -use common::vec2::Vec2; use rand::Rng; use rkyv::ser::serializers::AllocSerializer; use std::marker::PhantomData; -use std::sync::atomic::AtomicBool; use std::sync::Arc; pub trait HeapRelation { @@ -56,17 +55,15 @@ pub fn build( heap_relation.traverse(false, |(_, vector)| { assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); if number_of_samples < max_number_of_samples { - samples.extend(vector); + samples.push(vector); number_of_samples += 1; } else { let index = rand.gen_range(0..max_number_of_samples) as usize; - let start = index * dims as usize; - let end = start + dims as usize; - samples[start..end].copy_from_slice(&vector); + samples[index] = vector; } tuples_total += 1; }); - Vec2::from_vec((number_of_samples as _, dims as _), samples) + samples }; reporter.tuples_total(tuples_total); Structure::internal_build(vector_options.clone(), internal_build.clone(), samples) @@ -190,51 +187,56 @@ impl Structure { fn internal_build( vector_options: VectorOptions, internal_build: RabbitholeInternalBuildOptions, - mut samples: Vec2, + samples: Vec>, ) -> Self { - let dims = vector_options.dims; - for i in 0..samples.shape_0() { - let vector = &mut samples[(i,)]; - vector.copy_from_slice(&rabitq::project(vector)); - } - let h1_means = base::parallelism::RayonParallelism::scoped( + let h1_means = crate::algorithm::parallelism::RayonParallelism::scoped( internal_build.build_threads as _, - Arc::new(AtomicBool::new(false)), + Arc::new(|| { + pgrx::check_for_interrupts!(); + }), |parallelism| { - let raw = k_means::k_means( + k_means::k_means( parallelism, internal_build.nlist as usize, + vector_options.dims as usize, samples, internal_build.spherical_centroids, 10, - false, - ); - let mut centroids = Vec::new(); - for i in 0..internal_build.nlist { - centroids.push(raw[(i as usize,)].to_vec()); - } - centroids + ) }, ) - .expect("k_means panics") - .expect("k_means interrupted"); - let h2_mean = { - let mut centroid = vec![0.0; dims as _]; - for i in 0..internal_build.nlist { - for j in 0..dims { - centroid[j as usize] += h1_means[i as usize][j as usize]; - } - } - for j in 0..dims { - centroid[j as usize] /= internal_build.nlist as f32; - } - centroid - }; + .expect("failed to create thread pool"); + let h1_children = vec![Vec::new(); h1_means.len()]; + let h2_means = crate::algorithm::parallelism::RayonParallelism::scoped( + internal_build.build_threads as _, + Arc::new(|| { + pgrx::check_for_interrupts!(); + }), + |parallelism| { + k_means::k_means( + parallelism, + 1, + vector_options.dims as usize, + h1_means.clone(), + internal_build.spherical_centroids, + 10, + ) + }, + ) + .expect("failed to create thread pool"); + let mut h2_children = vec![Vec::new(); h2_means.len()]; + for i in 0..h1_means.len() as u32 { + let target: usize = k_means::k_means_lookup(&h1_means[i as usize], &h2_means); + h2_children[target].push(i); + } + let (h2_means, h2_children) = std::iter::zip(h2_means, h2_children) + .filter(|(_, x)| !x.is_empty()) + .unzip::<_, _, Vec<_>, Vec<_>>(); Structure { - h2_means: vec![h2_mean], - h2_children: vec![(0..internal_build.nlist).collect()], - h1_means, - h1_children: (0..internal_build.nlist).map(|_| Vec::new()).collect(), + h2_means: h2_means.into_iter().map(|x| rabitq::project(&x)).collect(), + h2_children, + h1_means: h1_means.into_iter().map(|x| rabitq::project(&x)).collect(), + h1_children, } } fn extern_build( diff --git a/src/algorithm/k_means.rs b/src/algorithm/k_means.rs new file mode 100644 index 00000000..dbb8b9f7 --- /dev/null +++ b/src/algorithm/k_means.rs @@ -0,0 +1,213 @@ +use crate::algorithm::parallelism::{ParallelIterator, Parallelism}; +use base::scalar::*; +use half::f16; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +pub fn k_means( + parallelism: &impl Parallelism, + c: usize, + dims: usize, + mut samples: Vec>, + is_spherical: bool, + iterations: usize, +) -> Vec> { + assert!(c > 0); + let n = samples.len(); + assert!(dims > 0); + if is_spherical { + for i in 0..n { + let sample = &mut samples[i]; + let l = S::reduce_sum_of_x2(sample).sqrt(); + S::vector_mul_scalar_inplace(sample, 1.0 / l); + } + } + if n <= c { + return quick_centers(c, dims, samples); + } + let mut lloyd_k_means = LloydKMeans::new(parallelism, c, dims, samples, is_spherical); + for _ in 0..iterations { + parallelism.check(); + if lloyd_k_means.iterate() { + break; + } + } + lloyd_k_means.finish() +} + +pub fn k_means_lookup(vector: &[S], centroids: &[Vec]) -> usize { + assert_ne!(centroids.len(), 0); + let mut result = (f32::INFINITY, 0); + for i in 0..centroids.len() { + let dis = S::reduce_sum_of_d2(vector, ¢roids[i]); + if dis <= result.0 { + result = (dis, i); + } + } + result.1 +} + +pub fn quick_centers( + c: usize, + dims: usize, + mut samples: Vec>, +) -> Vec> { + let n = samples.len(); + let mut rng = rand::thread_rng(); + assert!(c >= n); + for _ in n + 1..c { + let r = (0..dims) + .map(|_| S::from_f32(rng.gen_range(-1.0f32..1.0f32))) + .collect(); + samples.push(r); + } + samples +} + +pub struct LloydKMeans<'a, P, S> { + parallelism: &'a P, + dims: usize, + c: usize, + is_spherical: bool, + centroids: Vec>, + assign: Vec, + rng: StdRng, + samples: Vec>, +} + +const DELTA: f32 = f16::EPSILON.to_f32_const(); + +impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { + pub fn new( + parallelism: &'a P, + c: usize, + dims: usize, + samples: Vec>, + is_spherical: bool, + ) -> Self { + let n = samples.len(); + + let mut rng = StdRng::from_entropy(); + let mut centroids = Vec::with_capacity(c); + + for index in rand::seq::index::sample(&mut rng, n, c).into_iter() { + centroids.push(samples[index].clone()); + } + + let assign = parallelism + .into_par_iter(0..n) + .map(|i| { + let mut result = (f32::INFINITY, 0); + for j in 0..c { + let dis_2 = S::reduce_sum_of_d2(&samples[i], ¢roids[j]); + if dis_2 <= result.0 { + result = (dis_2, j); + } + } + result.1 + }) + .collect::>(); + + Self { + parallelism, + dims, + c, + is_spherical, + centroids, + assign, + rng, + samples, + } + } + + pub fn iterate(&mut self) -> bool { + let dims = self.dims; + let c = self.c; + let rand = &mut self.rng; + let samples = &self.samples; + let n = samples.len(); + + let (sum, mut count) = self + .parallelism + .into_par_iter(0..n) + .fold( + || (vec![vec![S::zero(); dims]; c], vec![0.0f32; c]), + |(mut sum, mut count), i| { + S::vector_add_inplace(&mut sum[self.assign[i]], &samples[i]); + count[self.assign[i]] += 1.0; + (sum, count) + }, + ) + .reduce( + || (vec![vec![S::zero(); dims]; c], vec![0.0f32; c]), + |(mut sum, mut count), (sum_1, count_1)| { + for i in 0..c { + S::vector_add_inplace(&mut sum[i], &sum_1[i]); + count[i] += count_1[i]; + } + (sum, count) + }, + ); + + let mut centroids = self + .parallelism + .into_par_iter(0..c) + .map(|i| S::vector_mul_scalar(&sum[i], 1.0 / count[i])) + .collect::>(); + + for i in 0..c { + if count[i] != 0.0f32 { + continue; + } + let mut o = 0; + loop { + let alpha = rand.gen_range(0.0..1.0f32); + let beta = (count[o] - 1.0) / (n - c) as f32; + if alpha < beta { + break; + } + o = (o + 1) % c; + } + centroids[i] = centroids[o].clone(); + S::kmeans_helper(&mut centroids[i], 1.0 + DELTA, 1.0 - DELTA); + S::kmeans_helper(&mut centroids[o], 1.0 - DELTA, 1.0 + DELTA); + count[i] = count[o] / 2.0; + count[o] -= count[i]; + } + + if self.is_spherical { + self.parallelism + .into_par_iter(&mut centroids) + .for_each(|centroid| { + let l = S::reduce_sum_of_x2(centroid).sqrt(); + S::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + + let assign = self + .parallelism + .into_par_iter(0..n) + .map(|i| { + let mut result = (f32::INFINITY, 0); + for j in 0..c { + let dis_2 = S::reduce_sum_of_d2(&samples[i], ¢roids[j]); + if dis_2 <= result.0 { + result = (dis_2, j); + } + } + result.1 + }) + .collect::>(); + + let result = (0..n).all(|i| assign[i] == self.assign[i]); + + self.centroids = centroids; + self.assign = assign; + + result + } + + pub fn finish(self) -> Vec> { + self.centroids + } +} diff --git a/src/algorithm/mod.rs b/src/algorithm/mod.rs index 6fcddd1e..ea0d0eb3 100644 --- a/src/algorithm/mod.rs +++ b/src/algorithm/mod.rs @@ -1,5 +1,7 @@ pub mod build; pub mod insert; +pub mod k_means; +pub mod parallelism; pub mod rabitq; pub mod scan; pub mod tuples; diff --git a/src/algorithm/parallelism.rs b/src/algorithm/parallelism.rs new file mode 100644 index 00000000..bd11191c --- /dev/null +++ b/src/algorithm/parallelism.rs @@ -0,0 +1,63 @@ +use std::any::Any; +use std::panic::AssertUnwindSafe; +use std::sync::Arc; + +pub use rayon::iter::ParallelIterator; + +pub trait Parallelism: Send + Sync { + fn check(&self); + + #[allow(clippy::wrong_self_convention)] + fn into_par_iter(&self, x: I) -> I::Iter; +} + +struct ParallelismCheckPanic(Box); + +pub struct RayonParallelism { + stop: Arc, +} + +impl RayonParallelism { + pub fn scoped( + num_threads: usize, + stop: Arc, + f: impl FnOnce(&Self) -> R, + ) -> Result { + match std::panic::catch_unwind(AssertUnwindSafe(|| { + rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .panic_handler(|e| { + if e.downcast_ref::().is_some() { + return; + } + log::error!("Asynchronous task panickied."); + }) + .build_scoped( + |thread| thread.run(), + |_| { + let pool = Self { stop: stop.clone() }; + f(&pool) + }, + ) + })) { + Ok(x) => x, + Err(e) => match e.downcast::() { + Ok(payload) => std::panic::resume_unwind((*payload).0), + Err(e) => std::panic::resume_unwind(e), + }, + } + } +} + +impl Parallelism for RayonParallelism { + fn check(&self) { + match std::panic::catch_unwind(AssertUnwindSafe(|| (self.stop)())) { + Ok(()) => (), + Err(payload) => std::panic::panic_any(ParallelismCheckPanic(payload)), + } + } + + fn into_par_iter(&self, x: I) -> I::Iter { + x.into_par_iter() + } +} diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs index 56950f13..cbf11365 100644 --- a/src/algorithm/scan.rs +++ b/src/algorithm/scan.rs @@ -17,7 +17,6 @@ pub fn scan( distance_kind: DistanceKind, nprobe_1: u32, ) -> impl Iterator { - assert!(nprobe_1 >= 1); let meta_guard = relation.read(0); let meta_tuple = meta_guard .get() From 5766d72298bb7c5e929a7f58cc364355ff7702b6 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 8 Nov 2024 19:23:34 +0800 Subject: [PATCH 029/324] fix: record indtuples in multiprogress environment (#54) Signed-off-by: usamoi --- src/algorithm/build.rs | 5 ++- src/index/am.rs | 75 ++++++++++++++++++++++++++---------------- 2 files changed, 49 insertions(+), 31 deletions(-) diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index 3224af1f..48dc02bc 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -25,8 +25,7 @@ pub trait HeapRelation { } pub trait Reporter { - fn tuples_total(&mut self, tuples_total: usize); - fn tuples_done(&mut self, tuples_done: usize); + fn tuples_total(&mut self, tuples_total: u64); } pub fn build( @@ -46,7 +45,7 @@ pub fn build( external_build.clone(), ), RabbitholeBuildOptions::Internal(internal_build) => { - let mut tuples_total = 0_usize; + let mut tuples_total = 0_u64; let samples = { let mut rand = rand::thread_rng(); let max_number_of_samples = internal_build.nlist.saturating_mul(256); diff --git a/src/index/am.rs b/src/index/am.rs index 545974a6..e3d29ece 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -131,6 +131,31 @@ pub unsafe extern "C" fn amcostestimate( } } +#[derive(Debug, Clone)] +struct PgReporter {} + +impl Reporter for PgReporter { + fn tuples_total(&mut self, tuples_total: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, + tuples_total as _, + ); + } + } +} + +impl PgReporter { + fn tuples_done(&mut self, tuples_done: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, + tuples_done as _, + ); + } + } +} + #[pgrx::pg_guard] pub unsafe extern "C" fn ambuild( heap: pgrx::pg_sys::Relation, @@ -205,26 +230,6 @@ pub unsafe extern "C" fn ambuild( self.opfamily } } - #[derive(Debug, Clone)] - pub struct PgReporter {} - impl Reporter for PgReporter { - fn tuples_total(&mut self, tuples_total: usize) { - unsafe { - pgrx::pg_sys::pgstat_progress_update_param( - pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, - tuples_total as _, - ); - } - } - fn tuples_done(&mut self, tuples_done: usize) { - unsafe { - pgrx::pg_sys::pgstat_progress_update_param( - pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, - tuples_done as _, - ); - } - } - } let (vector_options, rabbithole_options) = unsafe { am_options::options(index) }; let opfamily = unsafe { am_options::opfamily(index) }; let heap_relation = Heap { @@ -252,7 +257,7 @@ pub unsafe extern "C" fn ambuild( index_info, leader.tablescandesc, leader.rabbitholeshared, - true, + Some(reporter), ); leader.wait(); let nparticipants = leader.nparticipants; @@ -271,8 +276,8 @@ pub unsafe extern "C" fn ambuild( pgrx::pg_sys::ConditionVariableCancelSleep(); } } else { - let mut tuples_done = 0; - reporter.tuples_done(tuples_done); + let mut indtuples = 0; + reporter.tuples_done(indtuples); heap_relation.traverse(true, |(payload, vector)| { algorithm::insert::insert( index_relation.clone(), @@ -280,8 +285,8 @@ pub unsafe extern "C" fn ambuild( vector, opfamily.distance_kind(), ); - tuples_done += 1; - reporter.tuples_done(tuples_done); + indtuples += 1; + reporter.tuples_done(indtuples); }); } unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } @@ -301,6 +306,7 @@ struct RabbitholeShared { /* Mutable state */ nparticipantsdone: i32, + indtuples: u64, } fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { @@ -405,6 +411,7 @@ impl RabbitholeLeader { workersdonecv: std::mem::zeroed(), mutex: std::mem::zeroed(), nparticipantsdone: 0, + indtuples: 0, }); pgrx::pg_sys::ConditionVariableInit(&raw mut (*rabbitholeshared).workersdonecv); pgrx::pg_sys::SpinLockInit(&raw mut (*rabbitholeshared).mutex); @@ -508,7 +515,7 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( index_info, tablescandesc, rabbitholeshared, - false, + None, ); } @@ -524,7 +531,7 @@ unsafe fn parallel_build( index_info: *mut pgrx::pg_sys::IndexInfo, tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, rabbitholeshared: *mut RabbitholeShared, - progress: bool, + mut reporter: Option, ) { #[derive(Debug, Clone)] pub struct Heap { @@ -606,13 +613,25 @@ unsafe fn parallel_build( opfamily, scan, }; - heap_relation.traverse(progress, |(payload, vector)| { + heap_relation.traverse(reporter.is_some(), |(payload, vector)| { algorithm::insert::insert( index_relation.clone(), payload, vector, opfamily.distance_kind(), ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*rabbitholeshared).mutex); + (*rabbitholeshared).indtuples += 1; + indtuples = (*rabbitholeshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*rabbitholeshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } }); unsafe { From ec9cfcb65a4182a2468497095ebdaed08eea77a8 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 8 Nov 2024 19:39:17 +0800 Subject: [PATCH 030/324] fix: return Oid::INVALID to skip binary-coercible test (#55) Signed-off-by: usamoi --- src/datatype/memory_pgvector_vector.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/datatype/memory_pgvector_vector.rs b/src/datatype/memory_pgvector_vector.rs index c3ac76ab..fe0932bb 100644 --- a/src/datatype/memory_pgvector_vector.rs +++ b/src/datatype/memory_pgvector_vector.rs @@ -141,7 +141,7 @@ impl IntoDatum for PgvectorVectorOutput { } fn type_oid() -> Oid { - panic!("calling `type_oid` is never expected") + Oid::INVALID } fn is_compatible_with(_: Oid) -> bool { From 1b8bf47732fa2b461b0a06b3bf7a381b65d419ae Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 11 Nov 2024 10:45:02 +0800 Subject: [PATCH 031/324] fix: add error message (#57) Signed-off-by: usamoi --- Cargo.lock | 336 +----------------------------------------------- Cargo.toml | 2 +- src/index/am.rs | 10 ++ 3 files changed, 16 insertions(+), 332 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b79c1b30..0fce77b9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -78,7 +78,7 @@ dependencies = [ "serde", "thiserror", "toml", - "validator 0.18.1", + "validator", ] [[package]] @@ -330,17 +330,6 @@ dependencies = [ "syn 2.0.87", ] -[[package]] -name = "displaydoc" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - [[package]] name = "either" version = "1.13.0" @@ -505,124 +494,6 @@ dependencies = [ "windows-sys 0.52.0", ] -[[package]] -name = "icu_collections" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db2fa452206ebee18c4b5c2274dbf1de17008e874b4dc4f0aea9d01ca79e4526" -dependencies = [ - "displaydoc", - "yoke", - "zerofrom", - "zerovec", -] - -[[package]] -name = "icu_locid" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13acbb8371917fc971be86fc8057c41a64b521c184808a698c02acc242dbf637" -dependencies = [ - "displaydoc", - "litemap", - "tinystr", - "writeable", - "zerovec", -] - -[[package]] -name = "icu_locid_transform" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01d11ac35de8e40fdeda00d9e1e9d92525f3f9d887cdd7aa81d727596788b54e" -dependencies = [ - "displaydoc", - "icu_locid", - "icu_locid_transform_data", - "icu_provider", - "tinystr", - "zerovec", -] - -[[package]] -name = "icu_locid_transform_data" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdc8ff3388f852bede6b579ad4e978ab004f139284d7b28715f773507b946f6e" - -[[package]] -name = "icu_normalizer" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "19ce3e0da2ec68599d193c93d088142efd7f9c5d6fc9b803774855747dc6a84f" -dependencies = [ - "displaydoc", - "icu_collections", - "icu_normalizer_data", - "icu_properties", - "icu_provider", - "smallvec", - "utf16_iter", - "utf8_iter", - "write16", - "zerovec", -] - -[[package]] -name = "icu_normalizer_data" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8cafbf7aa791e9b22bec55a167906f9e1215fd475cd22adfcf660e03e989516" - -[[package]] -name = "icu_properties" -version = "1.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93d6020766cfc6302c15dbbc9c8778c37e62c14427cb7f6e601d849e092aeef5" -dependencies = [ - "displaydoc", - "icu_collections", - "icu_locid_transform", - "icu_properties_data", - "icu_provider", - "tinystr", - "zerovec", -] - -[[package]] -name = "icu_properties_data" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67a8effbc3dd3e4ba1afa8ad918d5684b8868b3b26500753effea8d2eed19569" - -[[package]] -name = "icu_provider" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ed421c8a8ef78d3e2dbc98a973be2f3770cb42b606e3ab18d6237c4dfde68d9" -dependencies = [ - "displaydoc", - "icu_locid", - "icu_provider_macros", - "stable_deref_trait", - "tinystr", - "writeable", - "yoke", - "zerofrom", - "zerovec", -] - -[[package]] -name = "icu_provider_macros" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - [[package]] name = "ident_case" version = "1.0.1" @@ -639,27 +510,6 @@ dependencies = [ "unicode-normalization", ] -[[package]] -name = "idna" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e" -dependencies = [ - "idna_adapter", - "smallvec", - "utf8_iter", -] - -[[package]] -name = "idna_adapter" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "daca1df1c957320b2cf139ac61e7bd64fed304c5040df000a745aa1de3b4ef71" -dependencies = [ - "icu_normalizer", - "icu_properties", -] - [[package]] name = "indenter" version = "0.3.3" @@ -748,12 +598,6 @@ version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" -[[package]] -name = "litemap" -version = "0.7.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "643cb0b8d4fcc284004d5fd0d67ccf61dfffadb7f75e1e71bc420f4688a3a704" - [[package]] name = "log" version = "0.4.22" @@ -1082,28 +926,6 @@ dependencies = [ "version_check", ] -[[package]] -name = "proc-macro-error-attr2" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" -dependencies = [ - "proc-macro2", - "quote", -] - -[[package]] -name = "proc-macro-error2" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" -dependencies = [ - "proc-macro-error-attr2", - "proc-macro2", - "quote", - "syn 2.0.87", -] - [[package]] name = "proc-macro2" version = "1.0.89" @@ -1181,7 +1003,7 @@ dependencies = [ "rkyv", "serde", "toml", - "validator 0.19.0", + "validator", ] [[package]] @@ -1475,12 +1297,6 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" -[[package]] -name = "smallvec" -version = "1.13.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" - [[package]] name = "smawk" version = "0.3.2" @@ -1546,17 +1362,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "synstructure" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - [[package]] name = "tap" version = "1.0.1" @@ -1583,16 +1388,6 @@ dependencies = [ "syn 2.0.87", ] -[[package]] -name = "tinystr" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9117f5d4db391c1cf6927e7bea3db74b9a1c1add8f7eda9ffd5364f40f57b82f" -dependencies = [ - "displaydoc", - "zerovec", -] - [[package]] name = "tinyvec" version = "1.8.0" @@ -1700,22 +1495,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22784dbdf76fdde8af1aeda5622b546b422b6fc585325248a2bf9f5e41e94d6c" dependencies = [ "form_urlencoded", - "idna 0.5.0", + "idna", "percent-encoding", ] -[[package]] -name = "utf16_iter" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" - -[[package]] -name = "utf8_iter" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" - [[package]] name = "uuid" version = "1.11.0" @@ -1731,30 +1514,14 @@ version = "0.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db79c75af171630a3148bd3e6d7c4f42b6a9a014c2945bc5ed0020cbb8d9478e" dependencies = [ - "idna 0.5.0", - "once_cell", - "regex", - "serde", - "serde_derive", - "serde_json", - "url", - "validator_derive 0.18.2", -] - -[[package]] -name = "validator" -version = "0.19.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0b4a29d8709210980a09379f27ee31549b73292c87ab9899beee1c0d3be6303" -dependencies = [ - "idna 1.0.3", + "idna", "once_cell", "regex", "serde", "serde_derive", "serde_json", "url", - "validator_derive 0.19.0", + "validator_derive", ] [[package]] @@ -1771,20 +1538,6 @@ dependencies = [ "syn 2.0.87", ] -[[package]] -name = "validator_derive" -version = "0.19.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bac855a2ce6f843beb229757e6e570a42e837bcb15e5f449dd48d5747d41bf77" -dependencies = [ - "darling", - "once_cell", - "proc-macro-error2", - "proc-macro2", - "quote", - "syn 2.0.87", -] - [[package]] name = "version_check" version = "0.9.5" @@ -1939,18 +1692,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "write16" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936" - -[[package]] -name = "writeable" -version = "0.5.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e9df38ee2d2c3c5948ea468a8406ff0db0b29ae1ffde1bcf20ef305bcc95c51" - [[package]] name = "wyz" version = "0.5.1" @@ -1969,30 +1710,6 @@ dependencies = [ "winapi", ] -[[package]] -name = "yoke" -version = "0.7.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c5b1314b079b0930c31e3af543d8ee1757b1951ae1e1565ec704403a7240ca5" -dependencies = [ - "serde", - "stable_deref_trait", - "yoke-derive", - "zerofrom", -] - -[[package]] -name = "yoke-derive" -version = "0.7.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28cc31741b18cb6f1d5ff12f5b7523e3d6eb0852bbbad19d73905511d9849b95" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", - "synstructure", -] - [[package]] name = "zerocopy" version = "0.7.35" @@ -2013,46 +1730,3 @@ dependencies = [ "quote", "syn 2.0.87", ] - -[[package]] -name = "zerofrom" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91ec111ce797d0e0784a1116d0ddcdbea84322cd79e5d5ad173daeba4f93ab55" -dependencies = [ - "zerofrom-derive", -] - -[[package]] -name = "zerofrom-derive" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ea7b4a3637ea8669cedf0f1fd5c286a17f3de97b8dd5a70a6c167a1730e63a5" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", - "synstructure", -] - -[[package]] -name = "zerovec" -version = "0.10.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa2b893d79df23bfb12d5461018d408ea19dfafe76c2c7ef6d4eba614f8ff079" -dependencies = [ - "yoke", - "zerofrom", - "zerovec-derive", -] - -[[package]] -name = "zerovec-derive" -version = "0.10.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] diff --git a/Cargo.toml b/Cargo.toml index 808654df..7c494572 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,7 +30,7 @@ rand_chacha = "0.3.1" rand_distr = "0.4.3" serde = "1" toml = "0.8.19" -validator = { version = "0.19.0", features = ["derive"] } +validator = { version = "0.18.1", features = ["derive"] } base = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } common = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } diff --git a/src/index/am.rs b/src/index/am.rs index e3d29ece..3144b05c 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -162,6 +162,7 @@ pub unsafe extern "C" fn ambuild( index: pgrx::pg_sys::Relation, index_info: *mut pgrx::pg_sys::IndexInfo, ) -> *mut pgrx::pg_sys::IndexBuildResult { + use validator::Validate; #[derive(Debug, Clone)] pub struct Heap { heap: pgrx::pg_sys::Relation, @@ -231,6 +232,15 @@ pub unsafe extern "C" fn ambuild( } } let (vector_options, rabbithole_options) = unsafe { am_options::options(index) }; + if let Err(errors) = Validate::validate(&vector_options) { + pgrx::error!("error while validating options: {}", errors); + } + if vector_options.dims > 2000 { + pgrx::error!("error while validating options: dimension is too large"); + } + if let Err(errors) = Validate::validate(&rabbithole_options) { + pgrx::error!("error while validating options: {}", errors); + } let opfamily = unsafe { am_options::opfamily(index) }; let heap_relation = Heap { heap, From ddd76ecdc9072ab1dd99dd782c339e0fd0cfd53e Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 11 Nov 2024 15:48:29 +0800 Subject: [PATCH 032/324] feat: allow users to install rabbithole on any schema (#58) Signed-off-by: usamoi --- Cargo.lock | 12 ------- Cargo.toml | 14 ++++---- docker/Dockerfile | 2 +- rabbithole.control | 3 +- src/index/am.rs | 10 ++---- src/index/am_options.rs | 74 ++++++++++++++++++--------------------- src/index/mod.rs | 1 + src/index/opclass.rs | 14 ++++++++ src/sql/finalize.sql | 9 +++-- tests/logic/partition.slt | 4 +-- 10 files changed, 68 insertions(+), 75 deletions(-) create mode 100644 src/index/opclass.rs diff --git a/Cargo.lock b/Cargo.lock index 0fce77b9..50fcd6ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -822,16 +822,6 @@ dependencies = [ "walkdir", ] -[[package]] -name = "pgrx-catalog" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38fc7bbd6334277565c322bf9d48c80bb8858cad493ec29e337df39dc259fc6a" -dependencies = [ - "paste", - "pgrx", -] - [[package]] name = "pgrx-macros" version = "0.12.8" @@ -987,14 +977,12 @@ name = "rabbithole" version = "0.0.0" dependencies = [ "base", - "common", "detect", "half 2.4.1", "log", "nalgebra", "paste", "pgrx", - "pgrx-catalog", "quantization", "rand", "rand_chacha", diff --git a/Cargo.toml b/Cargo.toml index 7c494572..3dd459b7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,18 +13,17 @@ path = "./src/bin/pgrx_embed.rs" [features] default = [] -pg12 = ["pgrx/pg12", "pgrx-catalog/pg12"] -pg13 = ["pgrx/pg13", "pgrx-catalog/pg13"] -pg14 = ["pgrx/pg14", "pgrx-catalog/pg14"] -pg15 = ["pgrx/pg15", "pgrx-catalog/pg15"] -pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] -pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] +pg12 = ["pgrx/pg12"] +pg13 = ["pgrx/pg13"] +pg14 = ["pgrx/pg14"] +pg15 = ["pgrx/pg15"] +pg16 = ["pgrx/pg16"] +pg17 = ["pgrx/pg17"] [dependencies] half = "2.4.1" paste = "1" pgrx = { version = "=0.12.8", default-features = false, features = ["cshim"] } -pgrx-catalog = "0.1.0" rand = "0.8.5" rand_chacha = "0.3.1" rand_distr = "0.4.3" @@ -33,7 +32,6 @@ toml = "0.8.19" validator = { version = "0.18.1", features = ["derive"] } base = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -common = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } detect = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } quantization = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } diff --git a/docker/Dockerfile b/docker/Dockerfile index 15f3e7f9..87bf94d6 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -8,4 +8,4 @@ RUN echo ${PG_VERSION} COPY ./build/rabbithole-pg${PG_VERSION}_0.0.0_amd64.deb /tmp/rabbithole.deb RUN apt-get install -y /tmp/rabbithole.deb && rm -f /tmp/rabbithole.deb -CMD ["postgres", "-c" ,"shared_preload_libraries=rabbithole.so", "-c", "search_path=\"$user\", public, rabbithole"] +CMD ["postgres", "-c" ,"shared_preload_libraries=rabbithole.so"] diff --git a/rabbithole.control b/rabbithole.control index a0f760d2..b8da9741 100644 --- a/rabbithole.control +++ b/rabbithole.control @@ -1,7 +1,6 @@ comment = 'rabbithole: Vector database plugin for Postgres, written in Rust, specifically designed for LLM' default_version = '@CARGO_VERSION@' module_pathname = '$libdir/rabbithole' -relocatable = false +relocatable = true superuser = true -schema = rabbithole requires = 'vector' diff --git a/src/index/am.rs b/src/index/am.rs index 3144b05c..9ad426e4 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -44,6 +44,7 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = { am_routine.type_ = pgrx::pg_sys::NodeTag::T_IndexAmRoutine; + am_routine.amsupport = 1; am_routine.amcanorderbyop = true; #[cfg(feature = "pg17")] @@ -79,13 +80,8 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = { }; #[pgrx::pg_guard] -pub unsafe extern "C" fn amvalidate(opclass_oid: pgrx::pg_sys::Oid) -> bool { - if am_options::convert_opclass_to_vd(opclass_oid).is_some() { - pgrx::info!("Vector indexes can only be built on built-in operator classes."); - true - } else { - false - } +pub unsafe extern "C" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { + true } #[pgrx::pg_guard] diff --git a/src/index/am_options.rs b/src/index/am_options.rs index 19396536..167c4aa2 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -51,40 +51,6 @@ impl PgDistanceKind { } } -pub fn convert_opclass_to_vd( - opclass_oid: pgrx::pg_sys::Oid, -) -> Option<(VectorKind, PgDistanceKind)> { - let namespace = pgrx_catalog::PgNamespace::search_namespacename(c"rabbithole").unwrap(); - let namespace = namespace.get().expect("rabbithole is not installed."); - let opclass = pgrx_catalog::PgOpclass::search_claoid(opclass_oid).unwrap(); - let opclass = opclass.get().expect("rabbithole is not installed."); - if opclass.opcnamespace() == namespace.oid() { - if let Ok(name) = opclass.opcname().to_str() { - if let Some(p) = convert_name_to_vd(name) { - return Some(p); - } - } - } - None -} - -pub fn convert_opfamily_to_vd( - opfamily_oid: pgrx::pg_sys::Oid, -) -> Option<(VectorKind, PgDistanceKind)> { - let namespace = pgrx_catalog::PgNamespace::search_namespacename(c"rabbithole").unwrap(); - let namespace = namespace.get().expect("rabbithole is not installed."); - let opfamily = pgrx_catalog::PgOpfamily::search_opfamilyoid(opfamily_oid).unwrap(); - let opfamily = opfamily.get().expect("rabbithole is not installed."); - if opfamily.opfnamespace() == namespace.oid() { - if let Ok(name) = opfamily.opfname().to_str() { - if let Some(p) = convert_name_to_vd(name) { - return Some(p); - } - } - } - None -} - fn convert_name_to_vd(name: &str) -> Option<(VectorKind, PgDistanceKind)> { match name.strip_suffix("_ops") { Some("vector_l2") => Some((VectorKind::Vecf32, PgDistanceKind::L2)), @@ -115,7 +81,6 @@ unsafe fn convert_reloptions_to_options( } pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, RabbitholeIndexingOptions) { - let opfamily = unsafe { (*index).rd_opfamily.read() }; let att = unsafe { &mut *(*index).rd_att }; let atts = unsafe { att.attrs.as_slice(att.natts as _) }; if atts.is_empty() { @@ -134,11 +99,11 @@ pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, Rabbitho ); }; // get v, d - let (v, pg_d) = convert_opfamily_to_vd(opfamily).unwrap(); + let opfamily = unsafe { opfamily(index) }; let vector = VectorOptions { dims, - v, - d: pg_d.to_distance(), + v: opfamily.vector, + d: opfamily.distance_kind(), }; // get indexing, segment, optimizing let rabitq = unsafe { convert_reloptions_to_options((*index).rd_options) }; @@ -213,8 +178,37 @@ impl Opfamily { } pub unsafe fn opfamily(index: pgrx::pg_sys::Relation) -> Opfamily { - let opfamily = unsafe { (*index).rd_opfamily.read() }; - let (vector, pg_distance) = convert_opfamily_to_vd(opfamily).unwrap(); + use pgrx::pg_sys::Oid; + + let proc = unsafe { pgrx::pg_sys::index_getprocid(index, 1, 1) }; + + if proc == Oid::INVALID { + pgrx::error!("support function 1 is not found"); + } + + let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); + unsafe { + pgrx::pg_sys::fmgr_info(proc, &mut flinfo); + } + + let fn_addr = flinfo.fn_addr.expect("null function pointer"); + + let mut fcinfo = unsafe { std::mem::zeroed::() }; + fcinfo.flinfo = &mut flinfo; + fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; + fcinfo.context = std::ptr::null_mut(); + fcinfo.resultinfo = std::ptr::null_mut(); + fcinfo.isnull = true; + fcinfo.nargs = 0; + + let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; + + let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; + + let result_string = result_option.expect("null string"); + + let (vector, pg_distance) = convert_name_to_vd(&result_string).unwrap(); + Opfamily { vector, pg_distance, diff --git a/src/index/mod.rs b/src/index/mod.rs index 489f18f7..6162f7ce 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -1,6 +1,7 @@ pub mod am; pub mod am_options; pub mod am_scan; +pub mod opclass; pub mod utils; pub unsafe fn init() { diff --git a/src/index/opclass.rs b/src/index/opclass.rs new file mode 100644 index 00000000..4a237991 --- /dev/null +++ b/src/index/opclass.rs @@ -0,0 +1,14 @@ +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _rabbithole_support_vector_l2_ops() -> String { + "vector_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _rabbithole_support_vector_ip_ops() -> String { + "vector_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _rabbithole_support_vector_cosine_ops() -> String { + "vector_cosine_ops".to_string() +} diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 259b7797..4badb2e4 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -49,14 +49,17 @@ CREATE OPERATOR FAMILY vector_cosine_ops USING rabbithole; CREATE OPERATOR CLASS vector_l2_ops FOR TYPE vector USING rabbithole FAMILY vector_l2_ops AS OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, - OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH; + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _rabbithole_support_vector_l2_ops(); CREATE OPERATOR CLASS vector_ip_ops FOR TYPE vector USING rabbithole FAMILY vector_ip_ops AS OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, - OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH; + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _rabbithole_support_vector_ip_ops(); CREATE OPERATOR CLASS vector_cosine_ops FOR TYPE vector USING rabbithole FAMILY vector_cosine_ops AS OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, - OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH; + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _rabbithole_support_vector_cosine_ops(); diff --git a/tests/logic/partition.slt b/tests/logic/partition.slt index e7811ec5..08d43c3b 100644 --- a/tests/logic/partition.slt +++ b/tests/logic/partition.slt @@ -22,7 +22,7 @@ SELECT FROM generate_series(1, 1000); statement ok -CREATE INDEX ON items USING rabbithole (val rabbithole.vector_l2_ops); +CREATE INDEX ON items USING rabbithole (val public.vector_l2_ops); query I SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; @@ -39,7 +39,7 @@ SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <=> '[0.5,0.5,0.5]' limit # partial index statement ok -CREATE INDEX ON items USING rabbithole (val rabbithole.vector_ip_ops) WHERE (category_id = 1); +CREATE INDEX ON items USING rabbithole (val public.vector_ip_ops) WHERE (category_id = 1); query I SELECT COUNT(1) FROM From 50ae49f4347ff101a590b377ee53b50b48cde26d Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 11 Nov 2024 16:01:22 +0800 Subject: [PATCH 033/324] feat: epsilon (#59) Signed-off-by: usamoi --- src/algorithm/insert.rs | 1 + src/algorithm/rabitq.rs | 5 +++-- src/algorithm/scan.rs | 3 +++ src/gucs/executing.rs | 15 +++++++++++++++ src/index/am_scan.rs | 2 ++ 5 files changed, 24 insertions(+), 2 deletions(-) diff --git a/src/algorithm/insert.rs b/src/algorithm/insert.rs index ed9bdb6b..a92a724c 100644 --- a/src/algorithm/insert.rs +++ b/src/algorithm/insert.rs @@ -121,6 +121,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k &h1_tuple.factor_err, &h1_tuple.t, ), + 1.9, ); for j in 0..32 { if h1_tuple.mask[j] { diff --git a/src/algorithm/rabitq.rs b/src/algorithm/rabitq.rs index ead21846..d05dfd0a 100644 --- a/src/algorithm/rabitq.rs +++ b/src/algorithm/rabitq.rs @@ -178,6 +178,7 @@ pub fn fscan_process_lowerbound( &[f32; 32], &[u8], ), + epsilon: f32, ) -> [Distance; 32] { let &(dis_v_2, b, k, qvector_sum, ref s) = lut; let r = quantization::fast_scan::b4::fast_scan_b4(dims.div_ceil(4), t, s); @@ -188,13 +189,13 @@ pub fn fscan_process_lowerbound( + b * factor_ppc[i] + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; let err = factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - 1.9 * err) + Distance::from_f32(rough - epsilon * err) }), DistanceKind::Dot => std::array::from_fn(|i| { let rough = 0.5 * b * factor_ppc[i] + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - 1.9 * err) + Distance::from_f32(rough - epsilon * err) }), DistanceKind::Hamming => unreachable!(), DistanceKind::Jaccard => unreachable!(), diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs index cbf11365..94acedcd 100644 --- a/src/algorithm/scan.rs +++ b/src/algorithm/scan.rs @@ -16,6 +16,7 @@ pub fn scan( vector: Vec, distance_kind: DistanceKind, nprobe_1: u32, + epsilon: f32, ) -> impl Iterator { let meta_guard = relation.read(0); let meta_tuple = meta_guard @@ -77,6 +78,7 @@ pub fn scan( &h1_tuple.factor_err, &h1_tuple.t, ), + epsilon, ); for j in 0..32 { if h1_tuple.mask[j] { @@ -149,6 +151,7 @@ pub fn scan( &h0_tuple.factor_err, &h0_tuple.t, ), + epsilon, ); for j in 0..32 { if h0_tuple.mask[j] { diff --git a/src/gucs/executing.rs b/src/gucs/executing.rs index 421d5a40..0e4e2161 100644 --- a/src/gucs/executing.rs +++ b/src/gucs/executing.rs @@ -1,6 +1,7 @@ use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; static NPROBE: GucSetting = GucSetting::::new(10); +static EPSILON: GucSetting = GucSetting::::new(1.9); pub unsafe fn init() { GucRegistry::define_int_guc( @@ -13,8 +14,22 @@ pub unsafe fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_float_guc( + "rabbithole.epsilon", + "`epsilon` argument of rabbithole.", + "`epsilon` argument of rabbithole.", + &EPSILON, + 1.0, + 4.0, + GucContext::Userset, + GucFlags::default(), + ); } pub fn nprobe() -> u32 { NPROBE.get() as u32 } + +pub fn epsilon() -> f32 { + EPSILON.get() as f32 +} diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs index 0c6406db..dd12e728 100644 --- a/src/index/am_scan.rs +++ b/src/index/am_scan.rs @@ -1,5 +1,6 @@ use super::am_options::Opfamily; use crate::algorithm::scan::scan; +use crate::gucs::executing::epsilon; use crate::gucs::executing::nprobe; use crate::postgres::Relation; use base::distance::Distance; @@ -82,6 +83,7 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, }, opfamily.distance_kind(), nprobe(), + epsilon(), ); *scanner = Scanner::Vbase { vbase: Box::new(vbase), From e66311340e7d190230f2dafc1b516f9b5024ef68 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 11 Nov 2024 16:07:28 +0800 Subject: [PATCH 034/324] fix: free leaked memory from fmgr (#60) Signed-off-by: usamoi --- src/index/am_options.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/index/am_options.rs b/src/index/am_options.rs index 167c4aa2..f351877c 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -209,6 +209,10 @@ pub unsafe fn opfamily(index: pgrx::pg_sys::Relation) -> Opfamily { let (vector, pg_distance) = convert_name_to_vd(&result_string).unwrap(); + unsafe { + pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); + } + Opfamily { vector, pg_distance, From fd4259679fcb4a59783f29218e940a9205b37001 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 11 Nov 2024 16:38:26 +0800 Subject: [PATCH 035/324] feat: pgvector ivfflat compatibility (#61) Signed-off-by: usamoi --- bench/bench.py | 4 ++-- bench/index.py | 6 +++--- bench/train.py | 4 ++-- src/algorithm/build.rs | 4 ++-- src/algorithm/scan.rs | 4 ++-- src/gucs/executing.rs | 14 +++++++------- src/index/am_scan.rs | 4 ++-- src/types.rs | 8 ++++---- tests/logic/index.slt | 6 +++--- 9 files changed, 27 insertions(+), 27 deletions(-) diff --git a/bench/bench.py b/bench/bench.py index ced1c53a..74a23c95 100644 --- a/bench/bench.py +++ b/bench/bench.py @@ -23,7 +23,7 @@ def build_arg_parse(): "-p", "--password", help="Database password", default="password" ) parser.add_argument( - "--nprob", help="argument rabbithole.nprobe for query", default=300, type=int + "--nprob", help="argument rabbithole.probes for query", default=300, type=int ) return parser @@ -87,6 +87,6 @@ def bench(name, test, answer, metric_ops, conn): else: raise ValueError conn = create_connection(args.password) - conn.execute(f"SET rabbithole.nprobe={args.nprob}") + conn.execute(f"SET rabbithole.probes={args.nprob}") bench(args.name, test, answer, metric_ops, conn) diff --git a/bench/index.py b/bench/index.py index 871a75cf..b0174ee4 100644 --- a/bench/index.py +++ b/bench/index.py @@ -65,21 +65,21 @@ def get_ivf_ops_config(metric, k, name=None): if metric == "l2": metric_ops = "vector_l2_ops" ivf_config = f""" - nlist = {k} + lists = {k} residual_quantization = true spherical_centroids = false """ elif metric == "cosine": metric_ops = "vector_cosine_ops" ivf_config = f""" - nlist = {k} + lists = {k} residual_quantization = false spherical_centroids = true """ elif metric == "ip": metric_ops = "vector_ip_ops" ivf_config = f""" - nlist = {k} + lists = {k} residual_quantization = false spherical_centroids = true """ diff --git a/bench/train.py b/bench/train.py index 26403000..de4c18e9 100644 --- a/bench/train.py +++ b/bench/train.py @@ -28,11 +28,11 @@ def build_arg_parse(): parser.add_argument("-o", "--output", help="output filepath", required=True) parser.add_argument( "-k", - help="K-means centroids or nlist", + help="K-means centroids or lists", type=int, default=DEFAULT_K, ) - parser.add_argument("--child-k", type=int, help="lower layer nlist (if enabled)") + parser.add_argument("--child-k", type=int, help="lower layer lists (if enabled)") parser.add_argument( "--niter", help="number of iterations", type=int, default=N_ITER ) diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index 48dc02bc..67353944 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -48,7 +48,7 @@ pub fn build( let mut tuples_total = 0_u64; let samples = { let mut rand = rand::thread_rng(); - let max_number_of_samples = internal_build.nlist.saturating_mul(256); + let max_number_of_samples = internal_build.lists.saturating_mul(256); let mut samples = Vec::new(); let mut number_of_samples = 0_u32; heap_relation.traverse(false, |(_, vector)| { @@ -196,7 +196,7 @@ impl Structure { |parallelism| { k_means::k_means( parallelism, - internal_build.nlist as usize, + internal_build.lists as usize, vector_options.dims as usize, samples, internal_build.spherical_centroids, diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs index 94acedcd..40046cb2 100644 --- a/src/algorithm/scan.rs +++ b/src/algorithm/scan.rs @@ -15,7 +15,7 @@ pub fn scan( relation: Relation, vector: Vec, distance_kind: DistanceKind, - nprobe_1: u32, + probes: u32, epsilon: f32, ) -> impl Iterator { let meta_guard = relation.read(0); @@ -119,7 +119,7 @@ pub fn scan( let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; Some((first, mean)) }) - .take(nprobe_1 as usize) + .take(probes as usize) .collect() }; { diff --git a/src/gucs/executing.rs b/src/gucs/executing.rs index 0e4e2161..c0e08d8d 100644 --- a/src/gucs/executing.rs +++ b/src/gucs/executing.rs @@ -1,14 +1,14 @@ use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; -static NPROBE: GucSetting = GucSetting::::new(10); +static PROBES: GucSetting = GucSetting::::new(10); static EPSILON: GucSetting = GucSetting::::new(1.9); pub unsafe fn init() { GucRegistry::define_int_guc( - "rabbithole.nprobe", - "`nprobe` argument of rabbithole.", - "`nprobe` argument of rabbithole.", - &NPROBE, + "rabbithole.probes", + "`probes` argument of rabbithole.", + "`probes` argument of rabbithole.", + &PROBES, 1, u16::MAX as _, GucContext::Userset, @@ -26,8 +26,8 @@ pub unsafe fn init() { ); } -pub fn nprobe() -> u32 { - NPROBE.get() as u32 +pub fn probes() -> u32 { + PROBES.get() as u32 } pub fn epsilon() -> f32 { diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs index dd12e728..3737d3de 100644 --- a/src/index/am_scan.rs +++ b/src/index/am_scan.rs @@ -1,7 +1,7 @@ use super::am_options::Opfamily; use crate::algorithm::scan::scan; use crate::gucs::executing::epsilon; -use crate::gucs::executing::nprobe; +use crate::gucs::executing::probes; use crate::postgres::Relation; use base::distance::Distance; use base::search::*; @@ -82,7 +82,7 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, OwnedVector::BVector(_) => unreachable!(), }, opfamily.distance_kind(), - nprobe(), + probes(), epsilon(), ); *scanner = Scanner::Vbase { diff --git a/src/types.rs b/src/types.rs index eb977b69..0741624c 100644 --- a/src/types.rs +++ b/src/types.rs @@ -4,9 +4,9 @@ use validator::Validate; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] pub struct RabbitholeInternalBuildOptions { - #[serde(default = "RabbitholeInternalBuildOptions::default_nlist")] + #[serde(default = "RabbitholeInternalBuildOptions::default_lists")] #[validate(range(min = 1, max = 1_000_000))] - pub nlist: u32, + pub lists: u32, #[serde(default = "RabbitholeInternalBuildOptions::default_spherical_centroids")] pub spherical_centroids: bool, #[serde(default = "RabbitholeInternalBuildOptions::default_build_threads")] @@ -15,7 +15,7 @@ pub struct RabbitholeInternalBuildOptions { } impl RabbitholeInternalBuildOptions { - fn default_nlist() -> u32 { + fn default_lists() -> u32 { 1000 } fn default_spherical_centroids() -> bool { @@ -29,7 +29,7 @@ impl RabbitholeInternalBuildOptions { impl Default for RabbitholeInternalBuildOptions { fn default() -> Self { Self { - nlist: Self::default_nlist(), + lists: Self::default_lists(), spherical_centroids: Self::default_spherical_centroids(), build_threads: Self::default_build_threads(), } diff --git a/tests/logic/index.slt b/tests/logic/index.slt index cccb64c7..ea9ee7ee 100644 --- a/tests/logic/index.slt +++ b/tests/logic/index.slt @@ -19,7 +19,7 @@ CREATE INDEX ON t USING rabbithole (val vector_l2_ops) WITH (options = $$ residual_quantization = true [build.internal] -nlist = 32 +lists = 32 spherical_centroids = false $$); @@ -28,7 +28,7 @@ CREATE INDEX ON t USING rabbithole (val vector_ip_ops) WITH (options = $$ residual_quantization = false [build.internal] -nlist = 32 +lists = 32 spherical_centroids = true $$); @@ -37,7 +37,7 @@ CREATE INDEX ON t USING rabbithole (val vector_cosine_ops) WITH (options = $$ residual_quantization = false [build.internal] -nlist = 32 +lists = 32 spherical_centroids = true $$); From 049b882695aeb1bd7992969013ec1b2e0a03962e Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 11 Nov 2024 19:30:24 +0800 Subject: [PATCH 036/324] fix: pg13 (#63) Signed-off-by: usamoi --- Cargo.toml | 1 - src/gucs/mod.rs | 2 +- src/index/am.rs | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 34 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 3dd459b7..e2fed9eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,6 @@ path = "./src/bin/pgrx_embed.rs" [features] default = [] -pg12 = ["pgrx/pg12"] pg13 = ["pgrx/pg13"] pg14 = ["pgrx/pg14"] pg15 = ["pgrx/pg15"] diff --git a/src/gucs/mod.rs b/src/gucs/mod.rs index e072844c..ab90704a 100644 --- a/src/gucs/mod.rs +++ b/src/gucs/mod.rs @@ -6,7 +6,7 @@ pub unsafe fn init() { executing::init(); prewarm::init(); prewarm::prewarm(); - #[cfg(any(feature = "pg12", feature = "pg13", feature = "pg14"))] + #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"rabbithole".as_ptr()); #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] pgrx::pg_sys::MarkGUCPrefixReserved(c"rabbithole".as_ptr()); diff --git a/src/index/am.rs b/src/index/am.rs index 9ad426e4..b4b834ae 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -653,6 +653,39 @@ pub unsafe extern "C" fn ambuildempty(_index: pgrx::pg_sys::Relation) { pgrx::error!("Unlogged indexes are not supported."); } +#[cfg(feature = "pg13")] +#[pgrx::pg_guard] +pub unsafe extern "C" fn aminsert( + index: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, + _heap: pgrx::pg_sys::Relation, + _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, + _index_info: *mut pgrx::pg_sys::IndexInfo, +) -> bool { + use base::vector::OwnedVector; + let opfamily = unsafe { am_options::opfamily(index) }; + let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; + if let Some(vector) = vector { + let vector = match opfamily.preprocess(vector.as_borrowed()) { + OwnedVector::Vecf32(x) => x, + OwnedVector::Vecf16(_) => unreachable!(), + OwnedVector::SVecf32(_) => unreachable!(), + OwnedVector::BVector(_) => unreachable!(), + }; + let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); + algorithm::insert::insert( + unsafe { Relation::new(index) }, + pointer, + vector.into_vec(), + opfamily.distance_kind(), + ); + } + false +} + +#[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] #[pgrx::pg_guard] pub unsafe extern "C" fn aminsert( index: pgrx::pg_sys::Relation, From 671edd7a586761f042bb5a3483ffdc6033e5ec64 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 11 Nov 2024 20:22:10 +0800 Subject: [PATCH 037/324] feat: prewarm (#64) Signed-off-by: usamoi --- Cargo.lock | 15 ++++++- Cargo.toml | 33 ++++++++-------- src/algorithm/mod.rs | 1 + src/algorithm/prewarm.rs | 85 ++++++++++++++++++++++++++++++++++++++++ src/index/functions.rs | 26 ++++++++++++ src/index/mod.rs | 1 + src/sql/finalize.sql | 3 ++ 7 files changed, 146 insertions(+), 18 deletions(-) create mode 100644 src/algorithm/prewarm.rs create mode 100644 src/index/functions.rs diff --git a/Cargo.lock b/Cargo.lock index 50fcd6ec..0c612a74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -822,6 +822,16 @@ dependencies = [ "walkdir", ] +[[package]] +name = "pgrx-catalog" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38fc7bbd6334277565c322bf9d48c80bb8858cad493ec29e337df39dc259fc6a" +dependencies = [ + "paste", + "pgrx", +] + [[package]] name = "pgrx-macros" version = "0.12.8" @@ -983,6 +993,7 @@ dependencies = [ "nalgebra", "paste", "pgrx", + "pgrx-catalog", "quantization", "rand", "rand_chacha", @@ -1150,9 +1161,9 @@ dependencies = [ [[package]] name = "rustix" -version = "0.38.38" +version = "0.38.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa260229e6538e52293eeb577aabd09945a09d6d9cc0fc550ed7529056c2e32a" +checksum = "99e4ea3e1cdc4b559b8e5650f9c8e5998e3e5c1343b4eaf034565f32318d63c0" dependencies = [ "bitflags", "errno", diff --git a/Cargo.toml b/Cargo.toml index e2fed9eb..6c49d385 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,23 +13,13 @@ path = "./src/bin/pgrx_embed.rs" [features] default = [] -pg13 = ["pgrx/pg13"] -pg14 = ["pgrx/pg14"] -pg15 = ["pgrx/pg15"] -pg16 = ["pgrx/pg16"] -pg17 = ["pgrx/pg17"] +pg13 = ["pgrx/pg13", "pgrx-catalog/pg13"] +pg14 = ["pgrx/pg14", "pgrx-catalog/pg14"] +pg15 = ["pgrx/pg15", "pgrx-catalog/pg15"] +pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] +pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] -half = "2.4.1" -paste = "1" -pgrx = { version = "=0.12.8", default-features = false, features = ["cshim"] } -rand = "0.8.5" -rand_chacha = "0.3.1" -rand_distr = "0.4.3" -serde = "1" -toml = "0.8.19" -validator = { version = "0.18.1", features = ["derive"] } - base = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } detect = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } quantization = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } @@ -38,9 +28,20 @@ quantization = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = nalgebra = { version = "=0.33.0", default-features = false } # lock rkyv version forever so that data is always compatible +rkyv = { version = "=0.7.45", features = ["validation"] } + +half = "2.4.1" log = "0.4.22" +paste = "1" +pgrx = { version = "=0.12.8", default-features = false, features = ["cshim"] } +pgrx-catalog = "0.1.0" +rand = "0.8.5" +rand_chacha = "0.3.1" +rand_distr = "0.4.3" rayon = "1.10.0" -rkyv = { version = "=0.7.45", features = ["validation"] } +serde = "1" +toml = "0.8.19" +validator = { version = "0.18.1", features = ["derive"] } [lints] rust.unsafe_op_in_unsafe_fn = "deny" diff --git a/src/algorithm/mod.rs b/src/algorithm/mod.rs index ea0d0eb3..8f1a2f07 100644 --- a/src/algorithm/mod.rs +++ b/src/algorithm/mod.rs @@ -2,6 +2,7 @@ pub mod build; pub mod insert; pub mod k_means; pub mod parallelism; +pub mod prewarm; pub mod rabitq; pub mod scan; pub mod tuples; diff --git a/src/algorithm/prewarm.rs b/src/algorithm/prewarm.rs new file mode 100644 index 00000000..9b6c1691 --- /dev/null +++ b/src/algorithm/prewarm.rs @@ -0,0 +1,85 @@ +use crate::algorithm::tuples::*; +use crate::postgres::Relation; +use std::fmt::Write; + +pub fn prewarm(relation: Relation) -> String { + let mut message = String::new(); + let meta_guard = relation.read(0); + let meta_tuple = meta_guard + .get() + .get(1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let lists = vec![meta_tuple.first]; + writeln!(message, "number of h2 tuples: {}", lists.len()).unwrap(); + writeln!(message, "number of h2 pages: {}", 1).unwrap(); + let mut counter = 0_usize; + let lists: Vec<_> = { + let mut results = Vec::new(); + for list in lists { + let mut current = list; + while current != u32::MAX { + counter += 1; + pgrx::check_for_interrupts!(); + let h1_guard = relation.read(current); + for i in 1..=h1_guard.get().len() { + let h1_tuple = h1_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + for j in 0..32 { + if h1_tuple.mask[j] { + results.push(h1_tuple.first[j]); + let mean = h1_tuple.mean[j]; + let vector_guard = relation.read(mean.0); + let vector_tuple = vector_guard + .get() + .get(mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let _ = vector_tuple; + } + } + } + current = h1_guard.get().get_opaque().next; + } + } + results + }; + writeln!(message, "number of h1 tuples: {}", lists.len()).unwrap(); + writeln!(message, "number of h1 pages: {}", counter).unwrap(); + let mut counter = 0_usize; + let lists: Vec<_> = { + let mut results = Vec::new(); + for list in lists { + let mut current = list; + while current != u32::MAX { + counter += 1; + pgrx::check_for_interrupts!(); + let h0_guard = relation.read(current); + for i in 1..=h0_guard.get().len() { + let h0_tuple = h0_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + for j in 0..32 { + if h0_tuple.mask[j] { + results.push(()); + } + } + } + current = h0_guard.get().get_opaque().next; + } + } + results + }; + writeln!(message, "number of h0 tuples: {}", lists.len()).unwrap(); + writeln!(message, "number of h0 pages: {}", counter).unwrap(); + message +} diff --git a/src/index/functions.rs b/src/index/functions.rs new file mode 100644 index 00000000..2e909b65 --- /dev/null +++ b/src/index/functions.rs @@ -0,0 +1,26 @@ +use crate::algorithm::prewarm::prewarm; +use crate::postgres::Relation; +use pgrx::pg_sys::Oid; +use pgrx_catalog::{PgAm, PgClass}; + +#[pgrx::pg_extern(strict)] +fn _rabbithole_prewarm(indexrelid: Oid) -> String { + let pg_am = PgAm::search_amname(c"rabbithole").unwrap(); + let Some(pg_am) = pg_am.get() else { + pgrx::error!("rabbithole is not installed"); + }; + let pg_class = PgClass::search_reloid(indexrelid).unwrap(); + let Some(pg_class) = pg_class.get() else { + pgrx::error!("there is no such index"); + }; + if pg_class.relam() != pg_am.oid() { + pgrx::error!("{:?} is not a rabbithole index", pg_class.relname()); + } + let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; + let relation = unsafe { Relation::new(index) }; + let message = prewarm(relation); + unsafe { + pgrx::pg_sys::index_close(index, pgrx::pg_sys::ShareLock as _); + } + message +} diff --git a/src/index/mod.rs b/src/index/mod.rs index 6162f7ce..5203e4fb 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -1,6 +1,7 @@ pub mod am; pub mod am_options; pub mod am_scan; +pub mod functions; pub mod opclass; pub mod utils; diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 4badb2e4..9430400f 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -33,6 +33,9 @@ CREATE OPERATOR <<=>> ( CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; +CREATE FUNCTION rabbithole_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_rabbithole_prewarm_wrapper'; + -- List of access methods CREATE ACCESS METHOD rabbithole TYPE INDEX HANDLER _rabbithole_amhandler; From 412c7f275585cee6aeef1a7e062bbd96216538a1 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 12 Nov 2024 13:27:24 +0800 Subject: [PATCH 038/324] feat: add a parameter of prewarm (#65) Signed-off-by: usamoi --- src/algorithm/prewarm.rs | 11 ++++++++++- src/datatype/operators_pgvector_vector.rs | 6 +++--- src/gucs/executing.rs | 2 +- src/index/am.rs | 4 +--- src/index/functions.rs | 6 +++--- src/sql/finalize.sql | 13 ++++++++----- 6 files changed, 26 insertions(+), 16 deletions(-) diff --git a/src/algorithm/prewarm.rs b/src/algorithm/prewarm.rs index 9b6c1691..3a2590ed 100644 --- a/src/algorithm/prewarm.rs +++ b/src/algorithm/prewarm.rs @@ -2,7 +2,7 @@ use crate::algorithm::tuples::*; use crate::postgres::Relation; use std::fmt::Write; -pub fn prewarm(relation: Relation) -> String { +pub fn prewarm(relation: Relation, height: i32) -> String { let mut message = String::new(); let meta_guard = relation.read(0); let meta_tuple = meta_guard @@ -11,9 +11,15 @@ pub fn prewarm(relation: Relation) -> String { .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); + if height > 2 { + return message; + } let lists = vec![meta_tuple.first]; writeln!(message, "number of h2 tuples: {}", lists.len()).unwrap(); writeln!(message, "number of h2 pages: {}", 1).unwrap(); + if height == 2 { + return message; + } let mut counter = 0_usize; let lists: Vec<_> = { let mut results = Vec::new(); @@ -52,6 +58,9 @@ pub fn prewarm(relation: Relation) -> String { }; writeln!(message, "number of h1 tuples: {}", lists.len()).unwrap(); writeln!(message, "number of h1 pages: {}", counter).unwrap(); + if height == 1 { + return message; + } let mut counter = 0_usize; let lists: Vec<_> = { let mut results = Vec::new(); diff --git a/src/datatype/operators_pgvector_vector.rs b/src/datatype/operators_pgvector_vector.rs index 1d8ac1aa..77d8fc4d 100644 --- a/src/datatype/operators_pgvector_vector.rs +++ b/src/datatype/operators_pgvector_vector.rs @@ -3,7 +3,7 @@ use base::vector::{VectBorrowed, VectorBorrowed}; use std::num::NonZero; #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_pgvector_vector_sphere_l2_in( +fn _rabbithole_vector_sphere_l2_in( lhs: PgvectorVectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { @@ -27,7 +27,7 @@ fn _rabbithole_pgvector_vector_sphere_l2_in( } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_pgvector_vector_sphere_ip_in( +fn _rabbithole_vector_sphere_ip_in( lhs: PgvectorVectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { @@ -51,7 +51,7 @@ fn _rabbithole_pgvector_vector_sphere_ip_in( } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_pgvector_vector_sphere_cosine_in( +fn _rabbithole_vector_sphere_cosine_in( lhs: PgvectorVectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { diff --git a/src/gucs/executing.rs b/src/gucs/executing.rs index c0e08d8d..af834043 100644 --- a/src/gucs/executing.rs +++ b/src/gucs/executing.rs @@ -19,7 +19,7 @@ pub unsafe fn init() { "`epsilon` argument of rabbithole.", "`epsilon` argument of rabbithole.", &EPSILON, - 1.0, + 0.0, 4.0, GucContext::Userset, GucFlags::default(), diff --git a/src/index/am.rs b/src/index/am.rs index b4b834ae..bec78dc3 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -26,9 +26,7 @@ pub unsafe fn init() { } } -#[pgrx::pg_extern(sql = "\ -CREATE FUNCTION _rabbithole_amhandler(internal) RETURNS index_am_handler -IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';")] +#[pgrx::pg_extern(sql = "")] fn _rabbithole_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { type T = pgrx::pg_sys::IndexAmRoutine; unsafe { diff --git a/src/index/functions.rs b/src/index/functions.rs index 2e909b65..4f661524 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -3,8 +3,8 @@ use crate::postgres::Relation; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass}; -#[pgrx::pg_extern(strict)] -fn _rabbithole_prewarm(indexrelid: Oid) -> String { +#[pgrx::pg_extern(sql = "")] +fn _rabbithole_prewarm(indexrelid: Oid, height: i32) -> String { let pg_am = PgAm::search_amname(c"rabbithole").unwrap(); let Some(pg_am) = pg_am.get() else { pgrx::error!("rabbithole is not installed"); @@ -18,7 +18,7 @@ fn _rabbithole_prewarm(indexrelid: Oid) -> String { } let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; let relation = unsafe { Relation::new(index) }; - let message = prewarm(relation); + let message = prewarm(relation, height); unsafe { pgrx::pg_sys::index_close(index, pgrx::pg_sys::ShareLock as _); } diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 9430400f..9ffc886c 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -8,21 +8,21 @@ CREATE TYPE sphere_vector AS ( -- List of operators CREATE OPERATOR <<->> ( - PROCEDURE = _rabbithole_pgvector_vector_sphere_l2_in, + PROCEDURE = _rabbithole_vector_sphere_l2_in, LEFTARG = vector, RIGHTARG = sphere_vector, COMMUTATOR = <<->> ); CREATE OPERATOR <<#>> ( - PROCEDURE = _rabbithole_pgvector_vector_sphere_ip_in, + PROCEDURE = _rabbithole_vector_sphere_ip_in, LEFTARG = vector, RIGHTARG = sphere_vector, COMMUTATOR = <<#>> ); CREATE OPERATOR <<=>> ( - PROCEDURE = _rabbithole_pgvector_vector_sphere_cosine_in, + PROCEDURE = _rabbithole_vector_sphere_cosine_in, LEFTARG = vector, RIGHTARG = sphere_vector, COMMUTATOR = <<=>> @@ -33,12 +33,15 @@ CREATE OPERATOR <<=>> ( CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; -CREATE FUNCTION rabbithole_prewarm(regclass) RETURNS TEXT +CREATE FUNCTION rabbithole_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_rabbithole_amhandler_wrapper'; + +CREATE FUNCTION rabbithole_prewarm(regclass, integer default 0) RETURNS TEXT STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_rabbithole_prewarm_wrapper'; -- List of access methods -CREATE ACCESS METHOD rabbithole TYPE INDEX HANDLER _rabbithole_amhandler; +CREATE ACCESS METHOD rabbithole TYPE INDEX HANDLER rabbithole_amhandler; COMMENT ON ACCESS METHOD rabbithole IS 'rabbithole index access method'; -- List of operator families From fd427cd4fee05d8c5ebca6b5e6dd4437b955b707 Mon Sep 17 00:00:00 2001 From: Jinjing Zhou Date: Wed, 13 Nov 2024 13:17:02 +0800 Subject: [PATCH 039/324] Add README (#66) * add README Signed-off-by: Jinjing.Zhou * align center Signed-off-by: Jinjing.Zhou * minor fix Signed-off-by: Jinjing.Zhou * address comment Signed-off-by: Jinjing.Zhou * add typos config for Rabit Signed-off-by: Jinjing.Zhou * add todo Signed-off-by: Jinjing.Zhou * toml format Signed-off-by: Jinjing.Zhou * RaBit instead of Rabit Signed-off-by: Jinjing.Zhou * fix todo Signed-off-by: Jinjing.Zhou --------- Signed-off-by: Jinjing.Zhou --- README.md | 131 ++++++++++++++++++++++++++++++++++++++++++++++++++++ _typos.toml | 2 + 2 files changed, 133 insertions(+) create mode 100644 README.md create mode 100644 _typos.toml diff --git a/README.md b/README.md new file mode 100644 index 00000000..abd6c8ea --- /dev/null +++ b/README.md @@ -0,0 +1,131 @@ +
+

VectorChord

+

Host 100M 768-dim vector (250GB+) on a $250/month machine (4 vcpu, 32GB) on AWS with VectorChord

+
+ +

+discord invitation link +trackgit-views + +all-contributors +

+ +VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search. It serves as the successor to the pgvecto.rs project. + +## Features +- **Blazing-Fast Queries**: Achieve up to 3x faster queries compared to pgvector's HNSW, maintaining the same recall level. +- **External Index Precomputation**: Built on IVF, VectorChord enables KMeans clustering to be performed externally (e.g., on a GPU) and seamlessly imported into the database. +- **Lightning-Fast Index Building**: Build index up to 20x faster than pgvector hnsw with precomputed centroids. (1.5 min for 1M 960-dim vectors) + +- **Advanced Quantization**: Uses cutting-edge RaBitQ to compress float vectors into compact bit representations with autonomous reranking. +- **Optimized SIMD Kernels**: Features a highly tuned computation kernel optimized for fast scans using SIMD and efficient register management. +- **Disk-Friendly Performance**: Query laion-100M 768-dim vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%. +- **Seamless Compatibility**: Compatible with pgvector data types while delivering faster indexing and querying. +- **Simple Configuration**: No need to tweak quantization or rerank parameters — best defaults are provided out of the box. + +## Quick Start + +Run the following SQL to ensure the extension is enabled. + +```SQL +CREATE EXTENSION IF NOT EXISTS vchord CASCADE; +``` + +To create the VectorChord RaBitQ(vchordrq) index, you can use the following SQL. + +```SQL +CREATE INDEX ON gist_train USING vchordrq (embedding vchordrq.vector_l2_ops) WITH (options = $$ +lists = 4096 +spherical_centroids = true +$$); +``` + +## Documentation + +### Query + +The query statement is exactly the same as pgvector. +```SQL +SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5; +``` +Supported distance functions are: +- <-> - L2 distance +- <#> - (negative) inner product +- <=> - cosine distance + + + +### Query Performance Tuning +You can fine-tune the search performance by adjusting the `probes` and `epsilon` parameters: + +```sql +-- Set probes to control the number of lists scanned. +-- Recommended range: 3%–10% of the total `lists` value. +SET vchordrq.probes = 100; + +-- Set epsilon to control the reranking precision. +-- Recommended range: 1.0–1.9. +SET vchordrq.epsilon = 1.0; +``` + +And for postgres's setting +```SQL +-- If using SSDs, set `effective_io_concurrency` to 200 for faster disk I/O. +SET effective_io_concurrency = 200; + +-- Disable JIT (Just-In-Time Compilation) as it offers minimal benefit (1–2%) +-- and adds overhead for single-query workloads. +SET jit = off; + +-- Allocate at least 25% of total memory to `shared_buffers`. +-- For disk-heavy workloads, you can increase this to up to 90% of total memory. +-- Note: A restart is required for this setting to take effect. +ALTER SYSTEM SET shared_buffers = '8GB'; +``` + + + + +### Index Build Time +Index building can parallelized, and with external centroid precomputation, the total time is primarily limited by disk speed. Optimize parallelism using the following settings: + +```SQL +-- Set this to the number of CPU cores available for parallel operations. +SET max_parallel_maintenance_workers = 8; +SET max_parallel_workers = 8; + +-- Adjust the total number of worker processes. +-- Note: A restart is required for this setting to take effect. +ALTER SYSTEM SET max_worker_processes = 8; +``` + +### Indexing Progress +You can check the indexing progress by querying the `pg_stat_progress_create_index` view. +```SQL +SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index; +``` + +### Installing From Source +TODO + +## Limitations +- Data Type Support: Currently, only the `f32` data type is supported for vectors, and the dimensionality is limited to 2000. (Dimension support improvements are planned for future updates.) +- Architecture Compatibility: The fast-scan kernel is optimized for x86_64 architectures. While it runs on aarch64, performance may be significantly lower. +- KMeans Clustering: The built-in KMeans clustering is not yet fully optimized and may require substantial memory. We strongly recommend using external centroid precomputation for efficient index construction. + + +## License +This project is licensed under the [GNU Affero General Public License v3.0](./LICENSE) and as commercial software. For commercial licensing, please contact us at support@tensorchord.ai. + diff --git a/_typos.toml b/_typos.toml new file mode 100644 index 00000000..e64dd008 --- /dev/null +++ b/_typos.toml @@ -0,0 +1,2 @@ +[default.extend-words] +RaBit = "RaBit" From 186d05c78bf443f2712dcfd08536ae975518ad41 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 13 Nov 2024 15:01:04 +0800 Subject: [PATCH 040/324] chore: rename rabbithole to vchord (#68) Signed-off-by: usamoi --- .github/workflows/psql.yml | 6 +- Cargo.lock | 46 ++++++------- Cargo.toml | 6 +- README.md | 1 + _typos.toml | 2 - bench/README.md | 8 +-- bench/bench.py | 7 +- bench/dump.py | 3 +- bench/index.py | 5 +- docker/Dockerfile | 6 +- rabbithole.control | 6 -- src/algorithm/build.rs | 8 +-- src/datatype/operators_pgvector_vector.rs | 6 +- src/gucs/executing.rs | 12 ++-- src/gucs/mod.rs | 8 ++- src/gucs/prewarm.rs | 4 +- src/index/am.rs | 81 +++++++++++------------ src/index/am_options.rs | 8 +-- src/index/functions.rs | 8 +-- src/index/opclass.rs | 6 +- src/lib.rs | 2 +- src/sql/finalize.sql | 35 +++++----- src/types.rs | 6 +- tests/README.md | 2 +- tests/logic/index.slt | 11 ++- tests/logic/issue427.slt | 5 +- tests/logic/null.fail | 5 +- tests/logic/partition.slt | 9 +-- tests/logic/pushdown_plan.fail | 7 +- tests/logic/pushdown_range.fail | 11 ++- tests/logic/reindex.slt | 5 +- tests/logic/vector.slt | 3 - tools/package.sh | 16 ++--- tools/schema-codegen.sh | 4 +- tools/schema.sh | 6 +- vchord.control | 6 ++ 36 files changed, 169 insertions(+), 201 deletions(-) delete mode 100644 _typos.toml delete mode 100644 rabbithole.control create mode 100644 vchord.control diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml index 56feedc8..1e2f9720 100644 --- a/.github/workflows/psql.yml +++ b/.github/workflows/psql.yml @@ -76,7 +76,7 @@ jobs: docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE build --lib --features pg${{ matrix.version }} --target ${{ matrix.arch }}-unknown-linux-gnu --profile opt docker run --rm -v .:/workspace $CACHE_ENVS --entrypoint bash $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --target ${{ matrix.arch }}-unknown-linux-gnu --profile $PROFILE ./tools/package.sh - docker build -t rabbithole:pg${{ matrix.version }} --build-arg PG_VERSION=${{ matrix.version }} -f ./docker/Dockerfile . + docker build -t vchord:pg${{ matrix.version }} --build-arg PG_VERSION=${{ matrix.version }} -f ./docker/Dockerfile . - name: Setup SQL Logic Test run: | @@ -88,8 +88,8 @@ jobs: env: PGPASSWORD: postgres run: | - docker run --rm --name test -d -e POSTGRES_PASSWORD=${PGPASSWORD} -p 5432:5432 rabbithole:pg${{ matrix.version }} + docker run --rm --name test -d -e POSTGRES_PASSWORD=${PGPASSWORD} -p 5432:5432 vchord:pg${{ matrix.version }} sleep 5 - psql -h localhost -U postgres -d postgres -c 'CREATE EXTENSION IF NOT EXISTS rabbithole CASCADE;' + psql -h localhost -U postgres -d postgres -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' sqllogictest './tests/**/*.slt' docker stop test diff --git a/Cargo.lock b/Cargo.lock index 0c612a74..0a95bb23 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -982,29 +982,6 @@ dependencies = [ "proc-macro2", ] -[[package]] -name = "rabbithole" -version = "0.0.0" -dependencies = [ - "base", - "detect", - "half 2.4.1", - "log", - "nalgebra", - "paste", - "pgrx", - "pgrx-catalog", - "quantization", - "rand", - "rand_chacha", - "rand_distr", - "rayon", - "rkyv", - "serde", - "toml", - "validator", -] - [[package]] name = "radium" version = "0.7.0" @@ -1537,6 +1514,29 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "vchord" +version = "0.0.0" +dependencies = [ + "base", + "detect", + "half 2.4.1", + "log", + "nalgebra", + "paste", + "pgrx", + "pgrx-catalog", + "quantization", + "rand", + "rand_chacha", + "rand_distr", + "rayon", + "rkyv", + "serde", + "toml", + "validator", +] + [[package]] name = "version_check" version = "0.9.5" diff --git a/Cargo.toml b/Cargo.toml index 6c49d385..ee8c3fe8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,14 +1,14 @@ [package] -name = "rabbithole" +name = "vchord" version = "0.0.0" edition = "2021" [lib] -name = "rabbithole" +name = "vchord" crate-type = ["cdylib", "lib"] [[bin]] -name = "pgrx_embed_rabbithole" +name = "pgrx_embed_vchord" path = "./src/bin/pgrx_embed.rs" [features] diff --git a/README.md b/README.md index abd6c8ea..6bca7b67 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,7 @@ To create the VectorChord RaBitQ(vchordrq) index, you can use the following SQL. ```SQL CREATE INDEX ON gist_train USING vchordrq (embedding vchordrq.vector_l2_ops) WITH (options = $$ +[build.internal] lists = 4096 spherical_centroids = true $$); diff --git a/_typos.toml b/_typos.toml deleted file mode 100644 index e64dd008..00000000 --- a/_typos.toml +++ /dev/null @@ -1,2 +0,0 @@ -[default.extend-words] -RaBit = "RaBit" diff --git a/bench/README.md b/bench/README.md index 07d02d50..57368e67 100644 --- a/bench/README.md +++ b/bench/README.md @@ -4,7 +4,7 @@ sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config clang cargo install --locked cargo-pgrx cargo pgrx init -cargo build --package rabbithole --lib --features pg16 --target x86_64-unknown-linux-gnu --profile opt +cargo build --package vchord --lib --features pg16 --target x86_64-unknown-linux-gnu --profile opt ./tools/schema.sh --features pg16 --target x86_64-unknown-linux-gnu --profile opt export SEMVER="0.0.0" @@ -14,15 +14,15 @@ export PLATFORM="amd64" export PROFILE="opt" ./tools/package.sh -docker build -t rabbithole:pg16-latest --build-arg PG_VERSION=16 -f ./docker/Dockerfile . +docker build -t vchord:pg16-latest --build-arg PG_VERSION=16 -f ./docker/Dockerfile . ``` -Or you can use `starkind/rabbithole:pg16-latest` to run the bench. +Or you can use `starkind/vchord:pg16-latest` to run the bench. ## Run Instance ```shell -docker run --name rabbithole -e POSTGRES_PASSWORD=123 -p 5432:5432 -d rabbithole:pg16-latest +docker run --name vchord -e POSTGRES_PASSWORD=123 -p 5432:5432 -d vchord:pg16-latest PGPASSWORD=123 psql -h 127.0.0.1 -U postgres -c "CREATE USER bench WITH PASSWORD '123';" PGPASSWORD=123 psql -h 127.0.0.1 -U postgres -c "ALTER ROLE bench SUPERUSER;" diff --git a/bench/bench.py b/bench/bench.py index 74a23c95..5acd7e4b 100644 --- a/bench/bench.py +++ b/bench/bench.py @@ -23,7 +23,7 @@ def build_arg_parse(): "-p", "--password", help="Database password", default="password" ) parser.add_argument( - "--nprob", help="argument rabbithole.probes for query", default=300, type=int + "--nprob", help="argument vchordrq.probes for query", default=300, type=int ) return parser @@ -41,9 +41,8 @@ def create_connection(password): autocommit=True, **keepalive_kwargs, ) - conn.execute("SET search_path TO public, vectors, rabbithole") conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") + conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") register_vector(conn) return conn @@ -87,6 +86,6 @@ def bench(name, test, answer, metric_ops, conn): else: raise ValueError conn = create_connection(args.password) - conn.execute(f"SET rabbithole.probes={args.nprob}") + conn.execute(f"SET vchordrq.probes={args.nprob}") bench(args.name, test, answer, metric_ops, conn) diff --git a/bench/dump.py b/bench/dump.py index 528ef056..fb5b1ccb 100644 --- a/bench/dump.py +++ b/bench/dump.py @@ -32,9 +32,8 @@ def create_connection(password): autocommit=True, **keepalive_kwargs, ) - conn.execute("SET search_path TO public, vectors, rabbithole") conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") + conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") register_vector(conn) return conn diff --git a/bench/index.py b/bench/index.py index b0174ee4..1cfb54c7 100644 --- a/bench/index.py +++ b/bench/index.py @@ -98,9 +98,8 @@ async def create_connection(url): autocommit=True, **KEEPALIVE_KWARGS, ) - await conn.execute("SET search_path TO public, vectors, rabbithole") await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - await conn.execute("CREATE EXTENSION IF NOT EXISTS rabbithole") + await conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") await register_vector_async(conn) return conn @@ -153,7 +152,7 @@ async def build_index( await conn.execute(f"SET max_parallel_maintenance_workers TO {workers}") await conn.execute(f"SET max_parallel_workers TO {workers}") await conn.execute( - f"CREATE INDEX ON {name} USING rabbithole (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" + f"CREATE INDEX ON {name} USING vchordrq (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" ) print(f"Index build time: {perf_counter() - start_time:.2f}s") finish.set() diff --git a/docker/Dockerfile b/docker/Dockerfile index 87bf94d6..79293c76 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -5,7 +5,7 @@ FROM pgvector/pgvector:${PGVECTOR}-pg${PG_VERSION} ARG PG_VERSION RUN echo ${PG_VERSION} -COPY ./build/rabbithole-pg${PG_VERSION}_0.0.0_amd64.deb /tmp/rabbithole.deb -RUN apt-get install -y /tmp/rabbithole.deb && rm -f /tmp/rabbithole.deb +COPY ./build/vchord-pg${PG_VERSION}_0.0.0_amd64.deb /tmp/vchord.deb +RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb -CMD ["postgres", "-c" ,"shared_preload_libraries=rabbithole.so"] +CMD ["postgres", "-c" ,"shared_preload_libraries=vchord.so"] diff --git a/rabbithole.control b/rabbithole.control deleted file mode 100644 index b8da9741..00000000 --- a/rabbithole.control +++ /dev/null @@ -1,6 +0,0 @@ -comment = 'rabbithole: Vector database plugin for Postgres, written in Rust, specifically designed for LLM' -default_version = '@CARGO_VERSION@' -module_pathname = '$libdir/rabbithole' -relocatable = true -superuser = true -requires = 'vector' diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index 67353944..b2c6c0b1 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -6,8 +6,8 @@ use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; use crate::types::RabbitholeBuildOptions; use crate::types::RabbitholeExternalBuildOptions; -use crate::types::RabbitholeIndexingOptions; use crate::types::RabbitholeInternalBuildOptions; +use crate::types::VchordrqIndexingOptions; use base::distance::DistanceKind; use base::index::VectorOptions; use base::scalar::ScalarLike; @@ -30,15 +30,15 @@ pub trait Reporter { pub fn build( vector_options: VectorOptions, - rabbithole_options: RabbitholeIndexingOptions, + vchordrq_options: VchordrqIndexingOptions, heap_relation: T, relation: Relation, mut reporter: R, ) { let dims = vector_options.dims; let is_residual = - rabbithole_options.residual_quantization && vector_options.d == DistanceKind::L2; - let structure = match rabbithole_options.build { + vchordrq_options.residual_quantization && vector_options.d == DistanceKind::L2; + let structure = match vchordrq_options.build { RabbitholeBuildOptions::External(external_build) => Structure::extern_build( vector_options.clone(), heap_relation.opfamily(), diff --git a/src/datatype/operators_pgvector_vector.rs b/src/datatype/operators_pgvector_vector.rs index 77d8fc4d..4a1f0554 100644 --- a/src/datatype/operators_pgvector_vector.rs +++ b/src/datatype/operators_pgvector_vector.rs @@ -3,7 +3,7 @@ use base::vector::{VectBorrowed, VectorBorrowed}; use std::num::NonZero; #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_vector_sphere_l2_in( +fn _vchord_vector_sphere_l2_in( lhs: PgvectorVectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { @@ -27,7 +27,7 @@ fn _rabbithole_vector_sphere_l2_in( } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_vector_sphere_ip_in( +fn _vchord_vector_sphere_ip_in( lhs: PgvectorVectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { @@ -51,7 +51,7 @@ fn _rabbithole_vector_sphere_ip_in( } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_vector_sphere_cosine_in( +fn _vchord_vector_sphere_cosine_in( lhs: PgvectorVectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { diff --git a/src/gucs/executing.rs b/src/gucs/executing.rs index af834043..31ca6dc2 100644 --- a/src/gucs/executing.rs +++ b/src/gucs/executing.rs @@ -5,9 +5,9 @@ static EPSILON: GucSetting = GucSetting::::new(1.9); pub unsafe fn init() { GucRegistry::define_int_guc( - "rabbithole.probes", - "`probes` argument of rabbithole.", - "`probes` argument of rabbithole.", + "vchordrq.probes", + "`probes` argument of vchordrq.", + "`probes` argument of vchordrq.", &PROBES, 1, u16::MAX as _, @@ -15,9 +15,9 @@ pub unsafe fn init() { GucFlags::default(), ); GucRegistry::define_float_guc( - "rabbithole.epsilon", - "`epsilon` argument of rabbithole.", - "`epsilon` argument of rabbithole.", + "vchordrq.epsilon", + "`epsilon` argument of vchordrq.", + "`epsilon` argument of vchordrq.", &EPSILON, 0.0, 4.0, diff --git a/src/gucs/mod.rs b/src/gucs/mod.rs index ab90704a..f4742589 100644 --- a/src/gucs/mod.rs +++ b/src/gucs/mod.rs @@ -7,8 +7,12 @@ pub unsafe fn init() { prewarm::init(); prewarm::prewarm(); #[cfg(any(feature = "pg13", feature = "pg14"))] - pgrx::pg_sys::EmitWarningsOnPlaceholders(c"rabbithole".as_ptr()); + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] - pgrx::pg_sys::MarkGUCPrefixReserved(c"rabbithole".as_ptr()); + pgrx::pg_sys::MarkGUCPrefixReserved(c"vchord".as_ptr()); + #[cfg(any(feature = "pg13", feature = "pg14"))] + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] + pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrq".as_ptr()); } } diff --git a/src/gucs/prewarm.rs b/src/gucs/prewarm.rs index 0638b385..6d358a14 100644 --- a/src/gucs/prewarm.rs +++ b/src/gucs/prewarm.rs @@ -6,7 +6,7 @@ static PREWARM_DIM: GucSetting> = pub unsafe fn init() { GucRegistry::define_string_guc( - "rabbithole.prewarm_dim", + "vchordrq.prewarm_dim", "prewarm_dim when the extension is loading.", "prewarm_dim when the extension is loading.", &PREWARM_DIM, @@ -26,7 +26,7 @@ pub fn prewarm() { } } } else { - pgrx::warning!("rabbithole.prewarm_dim is not a valid UTF-8 string"); + pgrx::warning!("vchordrq.prewarm_dim is not a valid UTF-8 string"); } } } diff --git a/src/index/am.rs b/src/index/am.rs index bec78dc3..7a3cdd94 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -27,7 +27,7 @@ pub unsafe fn init() { } #[pgrx::pg_extern(sql = "")] -fn _rabbithole_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { +fn _vchordrq_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { type T = pgrx::pg_sys::IndexAmRoutine; unsafe { let index_am_routine = pgrx::pg_sys::palloc0(size_of::()) as *mut T; @@ -225,14 +225,14 @@ pub unsafe extern "C" fn ambuild( self.opfamily } } - let (vector_options, rabbithole_options) = unsafe { am_options::options(index) }; + let (vector_options, vchordrq_options) = unsafe { am_options::options(index) }; if let Err(errors) = Validate::validate(&vector_options) { pgrx::error!("error while validating options: {}", errors); } if vector_options.dims > 2000 { pgrx::error!("error while validating options: dimension is too large"); } - if let Err(errors) = Validate::validate(&rabbithole_options) { + if let Err(errors) = Validate::validate(&vchordrq_options) { pgrx::error!("error while validating options: {}", errors); } let opfamily = unsafe { am_options::opfamily(index) }; @@ -246,7 +246,7 @@ pub unsafe extern "C" fn ambuild( let index_relation = unsafe { Relation::new(index) }; algorithm::build::build( vector_options, - rabbithole_options, + vchordrq_options, heap_relation.clone(), index_relation.clone(), reporter.clone(), @@ -260,20 +260,20 @@ pub unsafe extern "C" fn ambuild( heap, index_info, leader.tablescandesc, - leader.rabbitholeshared, + leader.vchordrqshared, Some(reporter), ); leader.wait(); let nparticipants = leader.nparticipants; loop { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.rabbitholeshared).mutex); - if (*leader.rabbitholeshared).nparticipantsdone == nparticipants { - pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.rabbitholeshared).mutex); + pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordrqshared).mutex); + if (*leader.vchordrqshared).nparticipantsdone == nparticipants { + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); break; } - pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.rabbitholeshared).mutex); + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); pgrx::pg_sys::ConditionVariableSleep( - &raw mut (*leader.rabbitholeshared).workersdonecv, + &raw mut (*leader.vchordrqshared).workersdonecv, pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN, ); } @@ -324,7 +324,7 @@ fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { struct RabbitholeLeader { pcxt: *mut pgrx::pg_sys::ParallelContext, nparticipants: i32, - rabbitholeshared: *mut RabbitholeShared, + vchordrqshared: *mut RabbitholeShared, tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, snapshot: pgrx::pg_sys::Snapshot, } @@ -365,8 +365,8 @@ impl RabbitholeLeader { } let pcxt = unsafe { pgrx::pg_sys::CreateParallelContext( - c"rabbithole".as_ptr(), - c"rabbithole_parallel_build_main".as_ptr(), + c"vchordrq".as_ptr(), + c"vchordrq_parallel_build_main".as_ptr(), request, ) }; @@ -404,11 +404,11 @@ impl RabbitholeLeader { } } - let rabbitholeshared = unsafe { - let rabbitholeshared = + let vchordrqshared = unsafe { + let vchordrqshared = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) .cast::(); - rabbitholeshared.write(RabbitholeShared { + vchordrqshared.write(RabbitholeShared { heaprelid: (*heap).rd_id, indexrelid: (*index).rd_id, isconcurrent, @@ -417,9 +417,9 @@ impl RabbitholeLeader { nparticipantsdone: 0, indtuples: 0, }); - pgrx::pg_sys::ConditionVariableInit(&raw mut (*rabbitholeshared).workersdonecv); - pgrx::pg_sys::SpinLockInit(&raw mut (*rabbitholeshared).mutex); - rabbitholeshared + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).workersdonecv); + pgrx::pg_sys::SpinLockInit(&raw mut (*vchordrqshared).mutex); + vchordrqshared }; let tablescandesc = unsafe { @@ -430,7 +430,7 @@ impl RabbitholeLeader { }; unsafe { - pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000001, rabbitholeshared.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000001, vchordrqshared.cast()); pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); } @@ -455,7 +455,7 @@ impl RabbitholeLeader { Some(Self { pcxt, nparticipants: nworkers_launched + 1, - rabbitholeshared, + vchordrqshared, tablescandesc, snapshot, }) @@ -485,11 +485,11 @@ impl Drop for RabbitholeLeader { #[pgrx::pg_guard] #[no_mangle] -pub unsafe extern "C" fn rabbithole_parallel_build_main( +pub unsafe extern "C" fn vchordrq_parallel_build_main( _seg: *mut pgrx::pg_sys::dsm_segment, toc: *mut pgrx::pg_sys::shm_toc, ) { - let rabbitholeshared = unsafe { + let vchordrqshared = unsafe { pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() }; let tablescandesc = unsafe { @@ -498,29 +498,22 @@ pub unsafe extern "C" fn rabbithole_parallel_build_main( }; let heap_lockmode; let index_lockmode; - if unsafe { !(*rabbitholeshared).isconcurrent } { + if unsafe { !(*vchordrqshared).isconcurrent } { heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; } else { heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; } - let heap = unsafe { pgrx::pg_sys::table_open((*rabbitholeshared).heaprelid, heap_lockmode) }; - let index = unsafe { pgrx::pg_sys::index_open((*rabbitholeshared).indexrelid, index_lockmode) }; + let heap = unsafe { pgrx::pg_sys::table_open((*vchordrqshared).heaprelid, heap_lockmode) }; + let index = unsafe { pgrx::pg_sys::index_open((*vchordrqshared).indexrelid, index_lockmode) }; let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; unsafe { - (*index_info).ii_Concurrent = (*rabbitholeshared).isconcurrent; + (*index_info).ii_Concurrent = (*vchordrqshared).isconcurrent; } unsafe { - parallel_build( - index, - heap, - index_info, - tablescandesc, - rabbitholeshared, - None, - ); + parallel_build(index, heap, index_info, tablescandesc, vchordrqshared, None); } unsafe { @@ -534,7 +527,7 @@ unsafe fn parallel_build( heap: pgrx::pg_sys::Relation, index_info: *mut pgrx::pg_sys::IndexInfo, tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, - rabbitholeshared: *mut RabbitholeShared, + vchordrqshared: *mut RabbitholeShared, mut reporter: Option, ) { #[derive(Debug, Clone)] @@ -627,10 +620,10 @@ unsafe fn parallel_build( unsafe { let indtuples; { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*rabbitholeshared).mutex); - (*rabbitholeshared).indtuples += 1; - indtuples = (*rabbitholeshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*rabbitholeshared).mutex); + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } if let Some(reporter) = reporter.as_mut() { reporter.tuples_done(indtuples); @@ -639,10 +632,10 @@ unsafe fn parallel_build( }); unsafe { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*rabbitholeshared).mutex); - (*rabbitholeshared).nparticipantsdone += 1; - pgrx::pg_sys::SpinLockRelease(&raw mut (*rabbitholeshared).mutex); - pgrx::pg_sys::ConditionVariableSignal(&raw mut (*rabbitholeshared).workersdonecv); + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).nparticipantsdone += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordrqshared).workersdonecv); } } diff --git a/src/index/am_options.rs b/src/index/am_options.rs index f351877c..1f271baf 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -1,7 +1,7 @@ use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use crate::datatype::typmod::Typmod; -use crate::types::RabbitholeIndexingOptions; +use crate::types::VchordrqIndexingOptions; use base::distance::*; use base::index::*; use base::vector::*; @@ -62,12 +62,12 @@ fn convert_name_to_vd(name: &str) -> Option<(VectorKind, PgDistanceKind)> { unsafe fn convert_reloptions_to_options( reloptions: *const pgrx::pg_sys::varlena, -) -> RabbitholeIndexingOptions { +) -> VchordrqIndexingOptions { #[derive(Debug, Clone, Deserialize, Default)] #[serde(deny_unknown_fields)] struct Parsed { #[serde(flatten)] - rabitq: RabbitholeIndexingOptions, + rabitq: VchordrqIndexingOptions, } let reloption = reloptions as *const Reloption; if reloption.is_null() || unsafe { (*reloption).options == 0 } { @@ -80,7 +80,7 @@ unsafe fn convert_reloptions_to_options( } } -pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, RabbitholeIndexingOptions) { +pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, VchordrqIndexingOptions) { let att = unsafe { &mut *(*index).rd_att }; let atts = unsafe { att.attrs.as_slice(att.natts as _) }; if atts.is_empty() { diff --git a/src/index/functions.rs b/src/index/functions.rs index 4f661524..c6f820dd 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -4,17 +4,17 @@ use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass}; #[pgrx::pg_extern(sql = "")] -fn _rabbithole_prewarm(indexrelid: Oid, height: i32) -> String { - let pg_am = PgAm::search_amname(c"rabbithole").unwrap(); +fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { + let pg_am = PgAm::search_amname(c"vchordrq").unwrap(); let Some(pg_am) = pg_am.get() else { - pgrx::error!("rabbithole is not installed"); + pgrx::error!("vchord is not installed"); }; let pg_class = PgClass::search_reloid(indexrelid).unwrap(); let Some(pg_class) = pg_class.get() else { pgrx::error!("there is no such index"); }; if pg_class.relam() != pg_am.oid() { - pgrx::error!("{:?} is not a rabbithole index", pg_class.relname()); + pgrx::error!("{:?} is not a vchordrq index", pg_class.relname()); } let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; let relation = unsafe { Relation::new(index) }; diff --git a/src/index/opclass.rs b/src/index/opclass.rs index 4a237991..e71da44a 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -1,14 +1,14 @@ #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_support_vector_l2_ops() -> String { +fn _vchordrq_support_vector_l2_ops() -> String { "vector_l2_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_support_vector_ip_ops() -> String { +fn _vchordrq_support_vector_ip_ops() -> String { "vector_ip_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _rabbithole_support_vector_cosine_ops() -> String { +fn _vchordrq_support_vector_cosine_ops() -> String { "vector_cosine_ops".to_string() } diff --git a/src/lib.rs b/src/lib.rs index db5d53c7..9d44ac73 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,7 +20,7 @@ pgrx::extension_sql_file!("./sql/finalize.sql", finalize); #[pgrx::pg_guard] unsafe extern "C" fn _PG_init() { if unsafe { pgrx::pg_sys::IsUnderPostmaster } { - pgrx::error!("rabbithole must be loaded via shared_preload_libraries."); + pgrx::error!("vchord must be loaded via shared_preload_libraries."); } detect::init(); unsafe { diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 9ffc886c..be75da37 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -8,21 +8,21 @@ CREATE TYPE sphere_vector AS ( -- List of operators CREATE OPERATOR <<->> ( - PROCEDURE = _rabbithole_vector_sphere_l2_in, + PROCEDURE = _vchord_vector_sphere_l2_in, LEFTARG = vector, RIGHTARG = sphere_vector, COMMUTATOR = <<->> ); CREATE OPERATOR <<#>> ( - PROCEDURE = _rabbithole_vector_sphere_ip_in, + PROCEDURE = _vchord_vector_sphere_ip_in, LEFTARG = vector, RIGHTARG = sphere_vector, COMMUTATOR = <<#>> ); CREATE OPERATOR <<=>> ( - PROCEDURE = _rabbithole_vector_sphere_cosine_in, + PROCEDURE = _vchord_vector_sphere_cosine_in, LEFTARG = vector, RIGHTARG = sphere_vector, COMMUTATOR = <<=>> @@ -33,39 +33,38 @@ CREATE OPERATOR <<=>> ( CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; -CREATE FUNCTION rabbithole_amhandler(internal) RETURNS index_am_handler -IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_rabbithole_amhandler_wrapper'; +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; -CREATE FUNCTION rabbithole_prewarm(regclass, integer default 0) RETURNS TEXT -STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_rabbithole_prewarm_wrapper'; +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; -- List of access methods -CREATE ACCESS METHOD rabbithole TYPE INDEX HANDLER rabbithole_amhandler; -COMMENT ON ACCESS METHOD rabbithole IS 'rabbithole index access method'; +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; -- List of operator families -CREATE OPERATOR FAMILY vector_l2_ops USING rabbithole; -CREATE OPERATOR FAMILY vector_ip_ops USING rabbithole; -CREATE OPERATOR FAMILY vector_cosine_ops USING rabbithole; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; -- List of operator classes CREATE OPERATOR CLASS vector_l2_ops - FOR TYPE vector USING rabbithole FAMILY vector_l2_ops AS + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, - FUNCTION 1 _rabbithole_support_vector_l2_ops(); + FUNCTION 1 _vchordrq_support_vector_l2_ops(); CREATE OPERATOR CLASS vector_ip_ops - FOR TYPE vector USING rabbithole FAMILY vector_ip_ops AS + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, - FUNCTION 1 _rabbithole_support_vector_ip_ops(); + FUNCTION 1 _vchordrq_support_vector_ip_ops(); CREATE OPERATOR CLASS vector_cosine_ops - FOR TYPE vector USING rabbithole FAMILY vector_cosine_ops AS + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, - FUNCTION 1 _rabbithole_support_vector_cosine_ops(); + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); diff --git a/src/types.rs b/src/types.rs index 0741624c..4ae331c1 100644 --- a/src/types.rs +++ b/src/types.rs @@ -68,13 +68,13 @@ impl Validate for RabbitholeBuildOptions { #[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] -pub struct RabbitholeIndexingOptions { - #[serde(default = "RabbitholeIndexingOptions::default_residual_quantization")] +pub struct VchordrqIndexingOptions { + #[serde(default = "VchordrqIndexingOptions::default_residual_quantization")] pub residual_quantization: bool, pub build: RabbitholeBuildOptions, } -impl RabbitholeIndexingOptions { +impl VchordrqIndexingOptions { fn default_residual_quantization() -> bool { false } diff --git a/tests/README.md b/tests/README.md index c3217e8f..d99133fb 100644 --- a/tests/README.md +++ b/tests/README.md @@ -9,6 +9,6 @@ cargo install sqllogictest-bin To run all the tests: ```bash -PGPASSWORD=postgres psql -h localhost -U postgres -d postgres -c 'CREATE EXTENSION IF NOT EXISTS rabbithole CASCADE;' +PGPASSWORD=postgres psql -h localhost -U postgres -d postgres -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' sqllogictest './tests/**/*.slt' ``` diff --git a/tests/logic/index.slt b/tests/logic/index.slt index ea9ee7ee..89449642 100644 --- a/tests/logic/index.slt +++ b/tests/logic/index.slt @@ -1,6 +1,3 @@ -statement ok -SET search_path TO public, rabbithole, pg_temp; - statement ok CREATE TABLE t (val vector(3)); @@ -8,14 +5,14 @@ statement ok INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); statement error -CREATE INDEX ON t USING rabbithole (val vector_l2_ops) +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) WITH (options = $$ unknown_options=true $$); # multiple index on single column statement ok -CREATE INDEX ON t USING rabbithole (val vector_l2_ops) +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) WITH (options = $$ residual_quantization = true [build.internal] @@ -24,7 +21,7 @@ spherical_centroids = false $$); statement ok -CREATE INDEX ON t USING rabbithole (val vector_ip_ops) +CREATE INDEX ON t USING vchordrq (val vector_ip_ops) WITH (options = $$ residual_quantization = false [build.internal] @@ -33,7 +30,7 @@ spherical_centroids = true $$); statement ok -CREATE INDEX ON t USING rabbithole (val vector_cosine_ops) +CREATE INDEX ON t USING vchordrq (val vector_cosine_ops) WITH (options = $$ residual_quantization = false [build.internal] diff --git a/tests/logic/issue427.slt b/tests/logic/issue427.slt index f3569900..de79cb87 100644 --- a/tests/logic/issue427.slt +++ b/tests/logic/issue427.slt @@ -1,8 +1,5 @@ # https://github.com/tensorchord/pgvecto.rs/issues/427 -statement ok -SET search_path TO public, rabbithole, pg_temp; - statement ok CREATE TABLE t (val vector(3)); @@ -10,7 +7,7 @@ statement ok INSERT INTO t (val) SELECT NULL::vector FROM generate_series(1, 1000); statement ok -CREATE INDEX ON t USING rabbithole (val vector_l2_ops); +CREATE INDEX ON t USING vchordrq (val vector_l2_ops); statement ok SELECT val FROM t ORDER BY val <-> (SELECT val FROM t LIMIT 1) limit 10; diff --git a/tests/logic/null.fail b/tests/logic/null.fail index 458382ce..3878d502 100644 --- a/tests/logic/null.fail +++ b/tests/logic/null.fail @@ -1,6 +1,3 @@ -statement ok -SET search_path TO public, rabbithole, pg_temp; - statement ok CREATE TABLE t (val vector(3)); @@ -19,7 +16,7 @@ SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) 10 statement ok -CREATE INDEX rabitq ON t USING rabbithole (val vector_l2_ops); +CREATE INDEX rabitq ON t USING vchordrq (val vector_l2_ops); query I SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; diff --git a/tests/logic/partition.slt b/tests/logic/partition.slt index 08d43c3b..68687722 100644 --- a/tests/logic/partition.slt +++ b/tests/logic/partition.slt @@ -1,6 +1,3 @@ -statement ok -SET search_path TO public, rabbithole, pg_temp; - # partition table statement ok CREATE TABLE items (val vector(3), category_id int) PARTITION BY LIST(category_id); @@ -22,7 +19,7 @@ SELECT FROM generate_series(1, 1000); statement ok -CREATE INDEX ON items USING rabbithole (val public.vector_l2_ops); +CREATE INDEX ON items USING vchordrq (val public.vector_l2_ops); query I SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; @@ -30,7 +27,7 @@ SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <-> '[0.5,0.5,0.5]' limit 10 statement ok -CREATE INDEX ON id_123 USING rabbithole (val vector_cosine_ops); +CREATE INDEX ON id_123 USING vchordrq (val vector_cosine_ops); query I SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; @@ -39,7 +36,7 @@ SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <=> '[0.5,0.5,0.5]' limit # partial index statement ok -CREATE INDEX ON items USING rabbithole (val public.vector_ip_ops) WHERE (category_id = 1); +CREATE INDEX ON items USING vchordrq (val public.vector_ip_ops) WHERE (category_id = 1); query I SELECT COUNT(1) FROM diff --git a/tests/logic/pushdown_plan.fail b/tests/logic/pushdown_plan.fail index 0cdabdef..392d5c0b 100644 --- a/tests/logic/pushdown_plan.fail +++ b/tests/logic/pushdown_plan.fail @@ -1,6 +1,3 @@ -statement ok -SET search_path TO public, rabbithole, pg_temp; - statement ok CREATE TABLE t (val0 vector(3), val1 halfvec(3), val2 sparsevec(3)); @@ -13,10 +10,10 @@ SELECT FROM generate_series(1, 10000); statement ok -CREATE INDEX ind0 ON t USING rabbithole (val0 vector_l2_ops); +CREATE INDEX ind0 ON t USING vchordrq (val0 vector_l2_ops); # statement ok -# CREATE INDEX ind1 ON t USING rabbithole (val1 halfvec_dot_ops); +# CREATE INDEX ind1 ON t USING vchordrq (val1 halfvec_dot_ops); # 1 vector key + 1 corresponding order_by key + sphere style query I diff --git a/tests/logic/pushdown_range.fail b/tests/logic/pushdown_range.fail index b13e6505..e58e63fa 100644 --- a/tests/logic/pushdown_range.fail +++ b/tests/logic/pushdown_range.fail @@ -1,6 +1,3 @@ -statement ok -SET search_path TO public, rabbithole, pg_temp; - statement ok CREATE TABLE t (val0 vector(3), val1 halfvec(3), val2 sparsevec(3), val3 bit(3)); @@ -12,7 +9,7 @@ INSERT INTO t (val0, val1, val2, val3) VALUES ('[0.4, 0.4, 0.4]', '[0.4, -0.4, 0.4]', '{1:-0.4, 2:-0.4, 3:-0.4}/3', '111'); statement ok -CREATE INDEX ON t USING rabbithole (val0 vector_l2_ops); +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops); # original style query I @@ -29,7 +26,7 @@ SELECT val0 FROM t WHERE val0 <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012) [0.3, 0.3, 0.3] statement ok -CREATE INDEX ON t USING rabbithole (val1 halfvec_dot_ops); +CREATE INDEX ON t USING vchordrq (val1 halfvec_dot_ops); # original style query I @@ -46,7 +43,7 @@ SELECT val1 FROM t WHERE val1 <<#>> sphere('[0.24, -0.24, 0.24]'::halfvec, 0) OR [0.099975586, 0.099975586, 0.099975586] statement ok -CREATE INDEX ON t USING rabbithole (val2 sparsevec_cos_ops) +CREATE INDEX ON t USING vchordrq (val2 sparsevec_cos_ops) WITH (options = "[indexing.hnsw]"); # original style @@ -65,7 +62,7 @@ ORDER BY val2 <=> '{0:0.12, 1:0.24, 2:0.36}/3'; {0:0.2, 1:-0.2, 2:0.2}/3 statement ok -CREATE INDEX ON t USING rabbithole (val3 bit_jaccard_ops) +CREATE INDEX ON t USING vchordrq (val3 bit_jaccard_ops) WITH (options = "[indexing.hnsw]"); # original style diff --git a/tests/logic/reindex.slt b/tests/logic/reindex.slt index ac0922fd..7d763433 100644 --- a/tests/logic/reindex.slt +++ b/tests/logic/reindex.slt @@ -1,6 +1,3 @@ -statement ok -SET search_path TO public, rabbithole, pg_temp; - statement ok CREATE TABLE t (val vector(3)); @@ -8,7 +5,7 @@ statement ok INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); statement ok -CREATE INDEX ON t USING rabbithole (val vector_l2_ops); +CREATE INDEX ON t USING vchordrq (val vector_l2_ops); statement ok REINDEX INDEX t_val_idx; diff --git a/tests/logic/vector.slt b/tests/logic/vector.slt index 0c507236..335266a5 100644 --- a/tests/logic/vector.slt +++ b/tests/logic/vector.slt @@ -1,6 +1,3 @@ -statement ok -SET search_path TO public, rabbithole, pg_temp; - statement ok CREATE TABLE t (val vector(3)); diff --git a/tools/package.sh b/tools/package.sh index 82234cd9..a1c2a010 100755 --- a/tools/package.sh +++ b/tools/package.sh @@ -8,15 +8,15 @@ printf "PLATFORM = ${PLATFORM}\n" printf "PROFILE = ${PROFILE}\n" rm -rf ./build/dir_zip -rm -rf ./build/rabbithole-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip +rm -rf ./build/vchord-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip rm -rf ./build/dir_deb -rm -rf ./build/rabbithole-pg${VERSION}_${SEMVER}_${PLATFORM}.deb +rm -rf ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb mkdir -p ./build/dir_zip -cp ./target/${ARCH}-unknown-linux-gnu/${PROFILE}/schema.sql ./build/dir_zip/rabbithole--$SEMVER.sql -sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./rabbithole.control > ./build/dir_zip/rabbithole.control -cp ./target/${ARCH}-unknown-linux-gnu/${PROFILE}/librabbithole.so ./build/dir_zip/rabbithole.so -zip ./build/rabbithole-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip -j ./build/dir_zip/* +cp ./target/${ARCH}-unknown-linux-gnu/${PROFILE}/schema.sql ./build/dir_zip/vchord--$SEMVER.sql +sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/dir_zip/vchord.control +cp ./target/${ARCH}-unknown-linux-gnu/${PROFILE}/libvchord.so ./build/dir_zip/vchord.so +zip ./build/vchord-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip -j ./build/dir_zip/* mkdir -p ./build/dir_deb mkdir -p ./build/dir_deb/DEBIAN @@ -31,7 +31,7 @@ done for file in $(ls ./build/dir_zip/*.so | xargs -n 1 basename); do cp ./build/dir_zip/$file ./build/dir_deb/usr/lib/postgresql/$VERSION/lib/$file done -echo "Package: rabbithole-pg${VERSION} +echo "Package: vchord-pg${VERSION} Version: ${SEMVER} Section: database Priority: optional @@ -42,4 +42,4 @@ Homepage: https://pgvecto.rs/ License: apache2" \ > ./build/dir_deb/DEBIAN/control (cd ./build/dir_deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/dir_deb/DEBIAN/md5sums -dpkg-deb -Zxz --build ./build/dir_deb/ ./build/rabbithole-pg${VERSION}_${SEMVER}_${PLATFORM}.deb \ No newline at end of file +dpkg-deb -Zxz --build ./build/dir_deb/ ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb \ No newline at end of file diff --git a/tools/schema-codegen.sh b/tools/schema-codegen.sh index 63a7e272..12f8289e 100755 --- a/tools/schema-codegen.sh +++ b/tools/schema-codegen.sh @@ -7,7 +7,7 @@ set -e printf "fn main() {\n" cat << EOF - rabbithole::__pgrx_marker(); + vchord::__pgrx_marker(); let mut entities = Vec::new(); let control_file_path = std::path::PathBuf::from("$CONTROL_FILEPATH"); @@ -30,7 +30,7 @@ done <<< $(nm -D -g $SO_FILEPATH | grep "T __pgrx_internals_" | awk '{print $3}' cat << EOF let pgrx_sql = ::pgrx::pgrx_sql_entity_graph::PgrxSql::build( entities.into_iter(), - "vectors".to_string(), + "vchord".to_string(), false, ) .expect("SQL generation error"); diff --git a/tools/schema.sh b/tools/schema.sh index e50ee988..2c430b8e 100755 --- a/tools/schema.sh +++ b/tools/schema.sh @@ -40,8 +40,8 @@ fi code=$(mktemp) chmod 700 $code -CONTROL_FILEPATH="./rabbithole.control" SO_FILEPATH="$DIR/librabbithole.so" $(dirname "$0")/schema-codegen.sh >> $code +CONTROL_FILEPATH="./vchord.control" SO_FILEPATH="$DIR/libvchord.so" $(dirname "$0")/schema-codegen.sh >> $code -PGRX_EMBED=$code cargo rustc --package rabbithole --bin pgrx_embed_rabbithole "$@" -- --cfg pgrx_embed +PGRX_EMBED=$code cargo rustc --package vchord --bin pgrx_embed_vchord "$@" -- --cfg pgrx_embed -CARGO_PKG_VERSION="0.0.0" QEMU_LD_PREFIX=$QEMU_LD_PREFIX "${RUNNER[@]}" "$DIR/pgrx_embed_rabbithole" | expand -t 4 > $DIR/schema.sql +CARGO_PKG_VERSION="0.0.0" QEMU_LD_PREFIX=$QEMU_LD_PREFIX "${RUNNER[@]}" "$DIR/pgrx_embed_vchord" | expand -t 4 > $DIR/schema.sql diff --git a/vchord.control b/vchord.control new file mode 100644 index 00000000..314ed5df --- /dev/null +++ b/vchord.control @@ -0,0 +1,6 @@ +comment = 'vchord: Vector database plugin for Postgres, written in Rust, specifically designed for LLM' +default_version = '@CARGO_VERSION@' +module_pathname = '$libdir/vchord' +relocatable = true +superuser = true +requires = 'vector' From 488f86460d2b7c7e223d83d3182657d4dc9badb7 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 13 Nov 2024 15:45:36 +0800 Subject: [PATCH 041/324] fix: off-by-one error in quick_centers (#69) Signed-off-by: usamoi --- src/algorithm/k_means.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/algorithm/k_means.rs b/src/algorithm/k_means.rs index dbb8b9f7..6cc1363a 100644 --- a/src/algorithm/k_means.rs +++ b/src/algorithm/k_means.rs @@ -55,7 +55,7 @@ pub fn quick_centers( let n = samples.len(); let mut rng = rand::thread_rng(); assert!(c >= n); - for _ in n + 1..c { + for _ in n..c { let r = (0..dims) .map(|_| S::from_f32(rng.gen_range(-1.0f32..1.0f32))) .collect(); From 66de953a5bba59b2cd984eca4f937055cc5ece36 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 13 Nov 2024 16:32:34 +0800 Subject: [PATCH 042/324] fix: spherical behavior (#70) Signed-off-by: usamoi --- src/algorithm/k_means.rs | 55 +++++++++++++++++++++------------------- 1 file changed, 29 insertions(+), 26 deletions(-) diff --git a/src/algorithm/k_means.rs b/src/algorithm/k_means.rs index 6cc1363a..a7e2f4fb 100644 --- a/src/algorithm/k_means.rs +++ b/src/algorithm/k_means.rs @@ -8,31 +8,25 @@ pub fn k_means( parallelism: &impl Parallelism, c: usize, dims: usize, - mut samples: Vec>, + samples: Vec>, is_spherical: bool, iterations: usize, ) -> Vec> { assert!(c > 0); - let n = samples.len(); assert!(dims > 0); - if is_spherical { - for i in 0..n { - let sample = &mut samples[i]; - let l = S::reduce_sum_of_x2(sample).sqrt(); - S::vector_mul_scalar_inplace(sample, 1.0 / l); - } - } + let n = samples.len(); if n <= c { - return quick_centers(c, dims, samples); - } - let mut lloyd_k_means = LloydKMeans::new(parallelism, c, dims, samples, is_spherical); - for _ in 0..iterations { - parallelism.check(); - if lloyd_k_means.iterate() { - break; + quick_centers(c, dims, samples, is_spherical) + } else { + let mut lloyd_k_means = LloydKMeans::new(parallelism, c, dims, samples, is_spherical); + for _ in 0..iterations { + parallelism.check(); + if lloyd_k_means.iterate() { + break; + } } + lloyd_k_means.finish() } - lloyd_k_means.finish() } pub fn k_means_lookup(vector: &[S], centroids: &[Vec]) -> usize { @@ -47,24 +41,33 @@ pub fn k_means_lookup(vector: &[S], centroids: &[Vec]) -> usiz result.1 } -pub fn quick_centers( +fn quick_centers( c: usize, dims: usize, - mut samples: Vec>, + samples: Vec>, + is_spherical: bool, ) -> Vec> { let n = samples.len(); - let mut rng = rand::thread_rng(); assert!(c >= n); + let mut rng = rand::thread_rng(); + let mut centroids = samples; for _ in n..c { let r = (0..dims) .map(|_| S::from_f32(rng.gen_range(-1.0f32..1.0f32))) .collect(); - samples.push(r); + centroids.push(r); + } + if is_spherical { + for i in 0..c { + let centroid = &mut centroids[i]; + let l = S::reduce_sum_of_x2(centroid).sqrt(); + S::vector_mul_scalar_inplace(centroid, 1.0 / l); + } } - samples + centroids } -pub struct LloydKMeans<'a, P, S> { +struct LloydKMeans<'a, P, S> { parallelism: &'a P, dims: usize, c: usize, @@ -78,7 +81,7 @@ pub struct LloydKMeans<'a, P, S> { const DELTA: f32 = f16::EPSILON.to_f32_const(); impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { - pub fn new( + fn new( parallelism: &'a P, c: usize, dims: usize, @@ -120,7 +123,7 @@ impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { } } - pub fn iterate(&mut self) -> bool { + fn iterate(&mut self) -> bool { let dims = self.dims; let c = self.c; let rand = &mut self.rng; @@ -207,7 +210,7 @@ impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { result } - pub fn finish(self) -> Vec> { + fn finish(self) -> Vec> { self.centroids } } From 009fab1f9242aeac83cca30ae79f02f7003e7d33 Mon Sep 17 00:00:00 2001 From: Jinjing Zhou Date: Wed, 13 Nov 2024 19:02:35 +0800 Subject: [PATCH 043/324] Update README (#71) * try enable sphere query Signed-off-by: Jinjing.Zhou * update readme Signed-off-by: Jinjing.Zhou * disable pushdown test Signed-off-by: Jinjing.Zhou --------- Signed-off-by: Jinjing.Zhou --- README.md | 33 ++++++++++++++--- tests/logic/pushdown_plan.fail | 4 +- tests/logic/pushdown_range.fail | 65 +++------------------------------ 3 files changed, 35 insertions(+), 67 deletions(-) diff --git a/README.md b/README.md index 6bca7b67..2a2a8cbd 100644 --- a/README.md +++ b/README.md @@ -10,18 +10,18 @@ all-contributors

-VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search. It serves as the successor to the pgvecto.rs project. +VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search. It serves as the successor to the pgvecto.rs project. VectorChord incorporates the insights and lessons learned from pgvecto.rs, providing faster query speeds, more flexible build options, significantly enhanced index build performance, and greater stability. It is entirely based on Postgres storage, allowing for physical replication and WAL incremental backups for index. This also enables effective ssd usage to reduce memory requirements and can easily handle vectors ranging from millions to billions of entries. ## Features - **Blazing-Fast Queries**: Achieve up to 3x faster queries compared to pgvector's HNSW, maintaining the same recall level. - **External Index Precomputation**: Built on IVF, VectorChord enables KMeans clustering to be performed externally (e.g., on a GPU) and seamlessly imported into the database. - **Lightning-Fast Index Building**: Build index up to 20x faster than pgvector hnsw with precomputed centroids. (1.5 min for 1M 960-dim vectors) - - **Advanced Quantization**: Uses cutting-edge RaBitQ to compress float vectors into compact bit representations with autonomous reranking. - **Optimized SIMD Kernels**: Features a highly tuned computation kernel optimized for fast scans using SIMD and efficient register management. - **Disk-Friendly Performance**: Query laion-100M 768-dim vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%. - **Seamless Compatibility**: Compatible with pgvector data types while delivering faster indexing and querying. - **Simple Configuration**: No need to tweak quantization or rerank parameters — best defaults are provided out of the box. + ## Quick Start + ### Query Performance Tuning You can fine-tune the search performance by adjusting the `probes` and `epsilon` parameters: @@ -93,6 +104,13 @@ SET jit = off; -- For disk-heavy workloads, you can increase this to up to 90% of total memory. -- Note: A restart is required for this setting to take effect. ALTER SYSTEM SET shared_buffers = '8GB'; + +-- vchordrq relies on a projection matrix to optimize performance. +-- Add your vector dimensions to the `prewarm_dim` list to reduce latency for the first query in each new connection. +-- If this is not configured, the first query will have higher latency as the matrix is generated on demand. +-- Default value: '64,128,256,384,512,768,1024,1536' +-- Note: This setting requires a database restart to take effect. +ALTER SYSTEM SET vchordrq.prewarm_dim = '64,128,256,384,512,768,1024,1536'; ```

VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search. It serves as the successor to the [pgvecto.rs](https://github.com/tensorchord/pgvecto.rs) project. @@ -16,6 +17,7 @@ VectorChord incorporates the insights and lessons learned from pgvecto.rs, provi ## Features - **Blazing-Fast Queries**: Achieve up to 3x faster queries compared to pgvector's HNSW, maintaining the same recall level. +- **High-throughput Update**: Achieve 16x faster insert throughput compared to pgvector's HNSW. - **External Index Precomputation**: Built on IVF, VectorChord enables KMeans clustering to be performed externally (e.g., on a GPU) and seamlessly imported into the database. - **Lightning-Fast Index Building**: Build index up to 20x faster than pgvector hnsw with precomputed centroids. (1.5 min for 1M 960-dim vectors) - **Advanced Quantization**: Uses cutting-edge RaBitQ to compress float vectors into compact bit representations with autonomous reranking. @@ -23,21 +25,20 @@ VectorChord incorporates the insights and lessons learned from pgvecto.rs, provi - **Disk-Friendly Performance**: Query laion-100M 768-dim vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%. - **Seamless Compatibility**: Compatible with pgvector data types while delivering faster indexing and querying. - **Simple Configuration**: No need to tweak quantization or rerank parameters — best defaults are provided out of the box. - ## Quick Start - +``` Run the following SQL to ensure the extension is enabled. ```SQL @@ -77,9 +78,15 @@ Supported distance functions are: ### Query Performance Tuning @@ -92,8 +99,9 @@ SET vchordrq.probes = 100; -- Set epsilon to control the reranking precision. -- Larger value means more rerank for higher recall rate. --- Recommended range: 1.0–1.9. -SET vchordrq.epsilon = 1.0; +-- Don't change it unless you only have limited memory. +-- Recommended range: 1.0–1.9. Default value is 1.9. +SET vchordrq.epsilon = 1.9; -- vchordrq relies on a projection matrix to optimize performance. -- Add your vector dimensions to the `prewarm_dim` list to reduce latency. @@ -118,8 +126,12 @@ SET jit = off; ALTER SYSTEM SET shared_buffers = '8GB'; ``` - +### Indexing prewarm +To prewarm the index, you can use the following SQL. It will significantly improve performance when using limited memory. +```SQL +-- vchordrq_prewarm(index_name::regclass) to prewarm the index into the shared buffer +SELECT vchordrq_prewarm('gist_train_embedding_idx'::regclass)" +``` ### Index Build Time @@ -182,7 +194,7 @@ cargo pgrx install --release --sudo # To install the extension into the system p ## Limitations - Data Type Support: Currently, only the `f32` data type is supported for vectors, and the dimensionality is limited to 1600. (Dimension support improvements are planned for future updates.) -- Architecture Compatibility: The fast-scan kernel is optimized for x86_64 architectures. While it runs on aarch64, performance may be significantly lower. +- Architecture Compatibility: The fast-scan kernel is optimized for x86_64 architectures. While it runs on aarch64, performance may be lower. - KMeans Clustering: The built-in KMeans clustering is not yet fully optimized and may require substantial memory. We strongly recommend using external centroid precomputation for efficient index construction. From 2f5236b121ab28961291bd478c1611e1db2e9145 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 20 Nov 2024 17:22:16 +0800 Subject: [PATCH 064/324] feat: support up to 60000 dimensions (#100) Signed-off-by: usamoi --- README.md | 3 +- src/algorithm/build.rs | 75 +++---------- src/algorithm/insert.rs | 235 ++++++++++++++++----------------------- src/algorithm/mod.rs | 1 + src/algorithm/prewarm.rs | 34 +----- src/algorithm/rabitq.rs | 177 +++++++++++++---------------- src/algorithm/scan.rs | 120 +++++++++----------- src/algorithm/tuples.rs | 118 +++----------------- src/algorithm/vacuum.rs | 38 +------ src/algorithm/vectors.rs | 76 +++++++++++++ src/index/am.rs | 13 ++- src/index/am_scan.rs | 2 +- src/postgres.rs | 17 +++ 13 files changed, 375 insertions(+), 534 deletions(-) create mode 100644 src/algorithm/vectors.rs diff --git a/README.md b/README.md index 54d8fc6b..5bac8114 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,6 @@ VectorChord incorporates the insights and lessons learned from pgvecto.rs, provi - **External Index Precomputation**: Built on IVF, VectorChord enables KMeans clustering to be performed externally (e.g., on a GPU) and seamlessly imported into the database. - **Lightning-Fast Index Building**: Build index up to 20x faster than pgvector hnsw with precomputed centroids. (1.5 min for 1M 960-dim vectors) - **Advanced Quantization**: Uses cutting-edge RaBitQ to compress float vectors into compact bit representations with autonomous reranking. -- **Optimized SIMD Kernels**: Features a highly tuned computation kernel optimized for fast scans using SIMD and efficient register management. - **Disk-Friendly Performance**: Query laion-100M 768-dim vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%. - **Seamless Compatibility**: Compatible with pgvector data types while delivering faster indexing and querying. - **Simple Configuration**: No need to tweak quantization or rerank parameters — best defaults are provided out of the box. @@ -193,7 +192,7 @@ cargo pgrx install --release --sudo # To install the extension into the system p ``` ## Limitations -- Data Type Support: Currently, only the `f32` data type is supported for vectors, and the dimensionality is limited to 1600. (Dimension support improvements are planned for future updates.) +- Data Type Support: Currently, only the `f32` data type is supported for vectors. - Architecture Compatibility: The fast-scan kernel is optimized for x86_64 architectures. While it runs on aarch64, performance may be lower. - KMeans Clustering: The built-in KMeans clustering is not yet fully optimized and may require substantial memory. We strongly recommend using external centroid precomputation for efficient index construction. diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index d8a38990..0b22b93d 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -1,6 +1,7 @@ use crate::algorithm::k_means; use crate::algorithm::rabitq; use crate::algorithm::tuples::*; +use crate::algorithm::vectors; use crate::index::am_options::Opfamily; use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; @@ -79,11 +80,16 @@ pub fn build( for i in 0..structures.len() { let mut level = Vec::new(); for j in 0..structures[i].len() { - let pointer = vectors.push(&VectorTuple { - payload: None, - vector: structures[i].means[j].clone(), - }); - level.push(pointer); + let slices = vectors::vector_split(&structures[i].means[j]); + let mut chain = None; + for i in (0..slices.len()).rev() { + chain = Some(vectors.push(&VectorTuple { + payload: None, + slice: slices[i].to_vec(), + chain, + })); + } + level.push(chain.unwrap()); } pointer_of_means.push(level); } @@ -96,7 +102,6 @@ pub fn build( level.push(tape.first()); } else { let mut tape = Tape::::create(&relation, false); - let mut cache = Vec::new(); let h2_mean = &structures[i].means[j]; let h2_children = &structures[i].children[j]; for child in h2_children.iter().copied() { @@ -106,58 +111,14 @@ pub fn build( } else { rabitq::code(dims, h1_mean) }; - cache.push((child, code)); - if cache.len() == 32 { - let group = std::mem::take(&mut cache); - let codes = std::array::from_fn(|k| group[k].1.clone()); - let packed = rabitq::pack_codes(dims, codes); - tape.push(&Height1Tuple { - mask: [true; 32], - mean: std::array::from_fn(|k| { - pointer_of_means[i - 1][group[k].0 as usize] - }), - first: std::array::from_fn(|k| { - pointer_of_firsts[i - 1][group[k].0 as usize] - }), - dis_u_2: packed.dis_u_2, - factor_ppc: packed.factor_ppc, - factor_ip: packed.factor_ip, - factor_err: packed.factor_err, - t: packed.t, - }); - } - } - if !cache.is_empty() { - let group = std::mem::take(&mut cache); - let codes = std::array::from_fn(|k| { - if k < group.len() { - group[k].1.clone() - } else { - rabitq::dummy_code(dims) - } - }); - let packed = rabitq::pack_codes(dims, codes); tape.push(&Height1Tuple { - mask: std::array::from_fn(|k| k < group.len()), - mean: std::array::from_fn(|k| { - if k < group.len() { - pointer_of_means[i - 1][group[k].0 as usize] - } else { - Default::default() - } - }), - first: std::array::from_fn(|k| { - if k < group.len() { - pointer_of_firsts[i - 1][group[k].0 as usize] - } else { - Default::default() - } - }), - dis_u_2: packed.dis_u_2, - factor_ppc: packed.factor_ppc, - factor_ip: packed.factor_ip, - factor_err: packed.factor_err, - t: packed.t, + mean: pointer_of_means[i - 1][child as usize], + first: pointer_of_firsts[i - 1][child as usize], + dis_u_2: code.dis_u_2, + factor_ppc: code.factor_ppc, + factor_ip: code.factor_ip, + factor_err: code.factor_err, + t: code.t(), }); } level.push(tape.first()); diff --git a/src/algorithm/insert.rs b/src/algorithm/insert.rs index 08019f08..859b1ac8 100644 --- a/src/algorithm/insert.rs +++ b/src/algorithm/insert.rs @@ -1,7 +1,7 @@ use crate::algorithm::rabitq; use crate::algorithm::rabitq::fscan_process_lowerbound; use crate::algorithm::tuples::*; -use crate::index::utils::distance; +use crate::algorithm::vectors; use crate::postgres::Relation; use base::always_equal::AlwaysEqual; use base::distance::Distance; @@ -28,70 +28,77 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k } else { None }; - let h0_vector = 'h0_vector: { - let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { - vector: vector.clone(), - payload: Some(payload.as_u64()), - }) - .unwrap(); - if let Some(mut write) = relation.search(tuple.len()) { - let i = write.get_mut().alloc(&tuple).unwrap(); - break 'h0_vector (write.id(), i); - } - let mut current = relation.read(1).get().get_opaque().fast_forward; - let mut changed = false; - loop { - let read = relation.read(current); - let flag = 'flag: { - if read.get().freespace() as usize >= tuple.len() { - break 'flag true; - } - if read.get().get_opaque().next == u32::MAX { - break 'flag true; - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current); - if let Some(i) = write.get_mut().alloc(&tuple) { - break (current, i); + let h0_vector = { + let slices = vectors::vector_split(&vector); + let mut chain = None; + for i in (0..slices.len()).rev() { + let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { + slice: slices[i].to_vec(), + payload: Some(payload.as_u64()), + chain, + }) + .unwrap(); + chain = Some('chain: { + if let Some(mut write) = relation.search(tuple.len()) { + let i = write.get_mut().alloc(&tuple).unwrap(); + break 'chain (write.id(), i); } - if write.get().get_opaque().next == u32::MAX { - if changed { - relation.write(1).get_mut().get_opaque_mut().fast_forward = write.id(); - } - let mut extend = relation.extend(true); - write.get_mut().get_opaque_mut().next = extend.id(); - if let Some(i) = extend.get_mut().alloc(&tuple) { - break (extend.id(), i); + let mut current = relation.read(1).get().get_opaque().fast_forward; + let mut changed = false; + loop { + let read = relation.read(current); + let flag = 'flag: { + if read.get().freespace() as usize >= tuple.len() { + break 'flag true; + } + if read.get().get_opaque().next == u32::MAX { + break 'flag true; + } + false + }; + if flag { + drop(read); + let mut write = relation.write(current); + if let Some(i) = write.get_mut().alloc(&tuple) { + break 'chain (current, i); + } + if write.get().get_opaque().next == u32::MAX { + if changed { + relation.write(1).get_mut().get_opaque_mut().fast_forward = + write.id(); + } + let mut extend = relation.extend(true); + write.get_mut().get_opaque_mut().next = extend.id(); + if let Some(i) = extend.get_mut().alloc(&tuple) { + break 'chain (extend.id(), i); + } else { + panic!("a tuple cannot even be fit in a fresh page"); + } + } + current = write.get().get_opaque().next; } else { - panic!("a tuple cannot even be fit in a fresh page"); + current = read.get().get_opaque().next; } + changed = true; } - current = write.get().get_opaque().next; - } else { - current = read.get().get_opaque().next; - } - changed = true; + }); } + chain.unwrap() }; let h0_payload = payload.as_u64(); - let mut list = ( - meta_tuple.first, - if is_residual { - let vector_guard = relation.read(meta_tuple.mean.0); - let vector_tuple = vector_guard - .get() - .get(meta_tuple.mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - Some(vector_tuple.vector.to_vec()) - } else { - None - }, - ); + let mut list = { + let Some((_, original)) = vectors::vector_dist( + relation.clone(), + &vector, + meta_tuple.mean, + None, + None, + is_residual, + ) else { + panic!("data corruption") + }; + (meta_tuple.first, original) + }; let make_list = |list: (u32, Option>)| { let mut results = Vec::new(); { @@ -115,23 +122,19 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k dims, lut, ( - &h1_tuple.dis_u_2, - &h1_tuple.factor_ppc, - &h1_tuple.factor_ip, - &h1_tuple.factor_err, + h1_tuple.dis_u_2, + h1_tuple.factor_ppc, + h1_tuple.factor_ip, + h1_tuple.factor_err, &h1_tuple.t, ), 1.9, ); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(( - Reverse(lowerbounds[j]), - AlwaysEqual(h1_tuple.mean[j]), - AlwaysEqual(h1_tuple.first[j]), - )); - } - } + results.push(( + Reverse(lowerbounds), + AlwaysEqual(h1_tuple.mean), + AlwaysEqual(h1_tuple.first), + )); } current = h1_guard.get().get_opaque().next; } @@ -141,23 +144,17 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get() - .get(mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); - cache.push(( - Reverse(dis_u), - AlwaysEqual(first), - AlwaysEqual(if is_residual { - Some(vector_tuple.vector.to_vec()) - } else { - None - }), - )); + let Some((Some(dis_u), original)) = vectors::vector_dist( + relation.clone(), + &vector, + mean, + None, + Some(distance_kind), + is_residual, + ) else { + panic!("data corruption") + }; + cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(original))); } let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop().unwrap(); (first, mean) @@ -171,15 +168,14 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k } else { rabitq::code(dims, &vector) }; - let dummy = rkyv::to_bytes::<_, 8192>(&Height0Tuple { - mask: [false; 32], - mean: [(0, 0); 32], - payload: [0; 32], - dis_u_2: [0.0f32; 32], - factor_ppc: [0.0f32; 32], - factor_ip: [0.0f32; 32], - factor_err: [0.0f32; 32], - t: vec![0; (dims.div_ceil(4) * 16) as usize], + let h0_tuple = rkyv::to_bytes::<_, 8192>(&Height0Tuple { + mean: h0_vector, + payload: h0_payload, + dis_u_2: code.dis_u_2, + factor_ppc: code.factor_ppc, + factor_ip: code.factor_ip, + factor_err: code.factor_err, + t: code.t(), }) .unwrap(); let first = list.0; @@ -188,18 +184,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k loop { let read = relation.read(current); let flag = 'flag: { - for i in 1..=read.get().len() { - let h0_tuple = read - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - if h0_tuple.mask.iter().any(|x| *x) { - break 'flag true; - } - } - if read.get().freespace() as usize >= dummy.len() { + if read.get().freespace() as usize >= h0_tuple.len() { break 'flag true; } if read.get().get_opaque().next == u32::MAX { @@ -210,41 +195,13 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k if flag { drop(read); let mut write = relation.write(current); - for i in 1..=write.get().len() { - let flag = put( - write.get_mut().get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - if flag { - return; - } - } - if let Some(i) = write.get_mut().alloc(&dummy) { - let flag = put( - write.get_mut().get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); + if write.get_mut().alloc(&h0_tuple).is_some() { return; } if write.get().get_opaque().next == u32::MAX { let mut extend = relation.extend(false); write.get_mut().get_opaque_mut().next = extend.id(); - if let Some(i) = extend.get_mut().alloc(&dummy) { - let flag = put( - extend.get_mut().get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); + if extend.get_mut().alloc(&h0_tuple).is_some() { return; } else { panic!("a tuple cannot even be fit in a fresh page"); diff --git a/src/algorithm/mod.rs b/src/algorithm/mod.rs index 8f1a2f07..a160118b 100644 --- a/src/algorithm/mod.rs +++ b/src/algorithm/mod.rs @@ -7,3 +7,4 @@ pub mod rabitq; pub mod scan; pub mod tuples; pub mod vacuum; +pub mod vectors; diff --git a/src/algorithm/prewarm.rs b/src/algorithm/prewarm.rs index d735852f..a274458e 100644 --- a/src/algorithm/prewarm.rs +++ b/src/algorithm/prewarm.rs @@ -1,4 +1,5 @@ use crate::algorithm::tuples::*; +use crate::algorithm::vectors; use crate::postgres::Relation; use std::fmt::Write; @@ -20,14 +21,7 @@ pub fn prewarm(relation: Relation, height: i32) -> String { let mut results = Vec::new(); let counter = 1_usize; { - let vector_guard = relation.read(meta_tuple.mean.0); - let vector_tuple = vector_guard - .get() - .get(meta_tuple.mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let _ = vector_tuple; + vectors::vector_warm(relation.clone(), meta_tuple.mean); results.push(meta_tuple.first); } writeln!(message, "number of tuples: {}", results.len()).unwrap(); @@ -50,20 +44,8 @@ pub fn prewarm(relation: Relation, height: i32) -> String { .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(h1_tuple.first[j]); - let mean = h1_tuple.mean[j]; - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get() - .get(mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let _ = vector_tuple; - } - } + vectors::vector_warm(relation.clone(), h1_tuple.mean); + results.push(h1_tuple.first); } current = h1_guard.get().get_opaque().next; } @@ -85,17 +67,13 @@ pub fn prewarm(relation: Relation, height: i32) -> String { pgrx::check_for_interrupts!(); let h0_guard = relation.read(current); for i in 1..=h0_guard.get().len() { - let h0_tuple = h0_guard + let _h0_tuple = h0_guard .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - for j in 0..32 { - if h0_tuple.mask[j] { - results.push(()); - } - } + results.push(()); } current = h0_guard.get().get_opaque().next; } diff --git a/src/algorithm/rabitq.rs b/src/algorithm/rabitq.rs index d05dfd0a..7c9da8d0 100644 --- a/src/algorithm/rabitq.rs +++ b/src/algorithm/rabitq.rs @@ -1,7 +1,6 @@ use base::distance::{Distance, DistanceKind}; use base::scalar::ScalarLike; use nalgebra::DMatrix; -use quantization::utils::InfiniteByteChunks; use std::sync::OnceLock; fn random_matrix(n: usize) -> DMatrix { @@ -20,7 +19,7 @@ fn check_all_matrixs_are_full_rank() { let mut threads = vec![]; for remainder in 0..parallelism { threads.push(scope.spawn(move || { - for n in (0..=2000).filter(|x| x % parallelism == remainder) { + for n in (0..=60000).filter(|x| x % parallelism == remainder) { let matrix = random_matrix(n); assert!(matrix.is_invertible()); } @@ -61,10 +60,10 @@ fn orthogonal_matrix(n: usize) -> Vec> { projection } -static MATRIXS: [OnceLock>>; 1 + 2000] = [const { OnceLock::new() }; 1 + 2000]; +static MATRIXS: [OnceLock>>; 1 + 60000] = [const { OnceLock::new() }; 1 + 60000]; pub fn prewarm(n: usize) { - if n <= 2000 { + if n <= 60000 { MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); } } @@ -87,6 +86,23 @@ pub struct Code { pub signs: Vec, } +impl Code { + pub fn t(&self) -> Vec { + use quantization::utils::InfiniteByteChunks; + let mut result = Vec::new(); + for x in InfiniteByteChunks::<_, 64>::new(self.signs.iter().copied()) + .take(self.signs.len().div_ceil(64)) + { + let mut r = 0_u64; + for i in 0..64 { + r |= (x[i] as u64) << i; + } + result.push(r); + } + result + } +} + pub fn code(dims: u32, vector: &[f32]) -> Code { let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); let sum_of_x_2 = f32::reduce_sum_of_x2(vector); @@ -119,43 +135,9 @@ pub fn code(dims: u32, vector: &[f32]) -> Code { } } -pub fn dummy_code(dims: u32) -> Code { - Code { - dis_u_2: 0.0, - factor_ppc: 0.0, - factor_ip: 0.0, - factor_err: 0.0, - signs: vec![0; dims as _], - } -} - -pub struct PackedCodes { - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, -} - -pub fn pack_codes(dims: u32, codes: [Code; 32]) -> PackedCodes { - PackedCodes { - dis_u_2: std::array::from_fn(|i| codes[i].dis_u_2), - factor_ppc: std::array::from_fn(|i| codes[i].factor_ppc), - factor_ip: std::array::from_fn(|i| codes[i].factor_ip), - factor_err: std::array::from_fn(|i| codes[i].factor_err), - t: { - let signs = codes.map(|code| { - InfiniteByteChunks::new(code.signs.into_iter()) - .map(|[b0, b1, b2, b3]| b0 | b1 << 1 | b2 << 2 | b3 << 3) - .take(dims.div_ceil(4) as usize) - .collect::>() - }); - quantization::fast_scan::b4::pack(dims.div_ceil(4), signs).collect() - }, - } -} +pub type Lut = (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)); -pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { +pub fn fscan_preprocess(vector: &[f32]) -> Lut { use quantization::quantize; let dis_v_2 = f32::reduce_sum_of_x2(vector); let (k, b, qvector) = quantize::quantize::<15>(vector); @@ -164,75 +146,74 @@ pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { } else { quantize::reduce_sum_of_x_as_u32(&qvector) as f32 }; - (dis_v_2, b, k, qvector_sum, compress(qvector)) + (dis_v_2, b, k, qvector_sum, binarize(&qvector)) } pub fn fscan_process_lowerbound( distance_kind: DistanceKind, - dims: u32, - lut: &(f32, f32, f32, f32, Vec), - (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[u8], - ), + _dims: u32, + lut: &Lut, + (dis_u_2, factor_ppc, factor_ip, factor_err, t): (f32, f32, f32, f32, &[u64]), epsilon: f32, -) -> [Distance; 32] { - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let r = quantization::fast_scan::b4::fast_scan_b4(dims.div_ceil(4), t, s); +) -> Distance { match distance_kind { - DistanceKind::L2 => std::array::from_fn(|i| { - let rough = dis_u_2[i] + DistanceKind::L2 => { + let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let value = asymmetric_binary_dot_product(t, s) as u16; + let rough = dis_u_2 + dis_v_2 - + b * factor_ppc[i] - + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = factor_err[i] * dis_v_2.sqrt(); + + b * factor_ppc + + ((2.0 * value as f32) - qvector_sum) * factor_ip * k; + let err = factor_err * dis_v_2.sqrt(); Distance::from_f32(rough - epsilon * err) - }), - DistanceKind::Dot => std::array::from_fn(|i| { - let rough = 0.5 * b * factor_ppc[i] - + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); + } + DistanceKind::Dot => { + let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let value = asymmetric_binary_dot_product(t, s) as u16; + let rough = + 0.5 * b * factor_ppc + 0.5 * ((2.0 * value as f32) - qvector_sum) * factor_ip * k; + let err = 0.5 * factor_err * dis_v_2.sqrt(); Distance::from_f32(rough - epsilon * err) - }), - DistanceKind::Hamming => unreachable!(), - DistanceKind::Jaccard => unreachable!(), + } + DistanceKind::Hamming => unimplemented!(), + DistanceKind::Jaccard => unimplemented!(), } } -fn compress(mut qvector: Vec) -> Vec { - let dims = qvector.len() as u32; - let width = dims.div_ceil(4); - qvector.resize(qvector.len().next_multiple_of(4), 0); - let mut t = vec![0u8; width as usize * 16]; - for i in 0..width as usize { - unsafe { - // this hint is used to skip bound checks - std::hint::assert_unchecked(4 * i + 3 < qvector.len()); - std::hint::assert_unchecked(16 * i + 15 < t.len()); - } - let t0 = qvector[4 * i + 0]; - let t1 = qvector[4 * i + 1]; - let t2 = qvector[4 * i + 2]; - let t3 = qvector[4 * i + 3]; - t[16 * i + 0b0000] = 0; - t[16 * i + 0b0001] = t0; - t[16 * i + 0b0010] = t1; - t[16 * i + 0b0011] = t1 + t0; - t[16 * i + 0b0100] = t2; - t[16 * i + 0b0101] = t2 + t0; - t[16 * i + 0b0110] = t2 + t1; - t[16 * i + 0b0111] = t2 + t1 + t0; - t[16 * i + 0b1000] = t3; - t[16 * i + 0b1001] = t3 + t0; - t[16 * i + 0b1010] = t3 + t1; - t[16 * i + 0b1011] = t3 + t1 + t0; - t[16 * i + 0b1100] = t3 + t2; - t[16 * i + 0b1101] = t3 + t2 + t0; - t[16 * i + 0b1110] = t3 + t2 + t1; - t[16 * i + 0b1111] = t3 + t2 + t1 + t0; +fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { + let n = vector.len(); + let mut t0 = vec![0u64; n.div_ceil(64)]; + let mut t1 = vec![0u64; n.div_ceil(64)]; + let mut t2 = vec![0u64; n.div_ceil(64)]; + let mut t3 = vec![0u64; n.div_ceil(64)]; + for i in 0..n { + t0[i / 64] |= (((vector[i] >> 0) & 1) as u64) << (i % 64); + t1[i / 64] |= (((vector[i] >> 1) & 1) as u64) << (i % 64); + t2[i / 64] |= (((vector[i] >> 2) & 1) as u64) << (i % 64); + t3[i / 64] |= (((vector[i] >> 3) & 1) as u64) << (i % 64); + } + (t0, t1, t2, t3) +} + +#[detect::multiversion(v2, fallback)] +fn asymmetric_binary_dot_product(x: &[u64], y: &(Vec, Vec, Vec, Vec)) -> u32 { + assert_eq!(x.len(), y.0.len()); + assert_eq!(x.len(), y.1.len()); + assert_eq!(x.len(), y.2.len()); + assert_eq!(x.len(), y.3.len()); + let n = x.len(); + let (mut t0, mut t1, mut t2, mut t3) = (0, 0, 0, 0); + for i in 0..n { + t0 += (x[i] & y.0[i]).count_ones(); + } + for i in 0..n { + t1 += (x[i] & y.1[i]).count_ones(); + } + for i in 0..n { + t2 += (x[i] & y.2[i]).count_ones(); + } + for i in 0..n { + t3 += (x[i] & y.3[i]).count_ones(); } - t + (t0 << 0) + (t1 << 1) + (t2 << 2) + (t3 << 3) } diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs index 124e7527..1f665b36 100644 --- a/src/algorithm/scan.rs +++ b/src/algorithm/scan.rs @@ -1,7 +1,7 @@ use crate::algorithm::rabitq; use crate::algorithm::rabitq::fscan_process_lowerbound; use crate::algorithm::tuples::*; -use crate::index::utils::distance; +use crate::algorithm::vectors; use crate::postgres::Relation; use base::always_equal::AlwaysEqual; use base::distance::Distance; @@ -36,21 +36,19 @@ pub fn scan( } else { None }; - let mut lists: Vec<_> = vec![( - meta_tuple.first, - if is_residual { - let vector_guard = relation.read(meta_tuple.mean.0); - let vector_tuple = vector_guard - .get() - .get(meta_tuple.mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - Some(vector_tuple.vector.to_vec()) - } else { - None - }, - )]; + let mut lists: Vec<_> = vec![{ + let Some((_, original)) = vectors::vector_dist( + relation.clone(), + &vector, + meta_tuple.mean, + None, + None, + is_residual, + ) else { + panic!("data corruption") + }; + (meta_tuple.first, original) + }]; let make_lists = |lists: Vec<(u32, Option>)>, probes| { let mut results = Vec::new(); for list in lists { @@ -74,23 +72,19 @@ pub fn scan( dims, lut, ( - &h1_tuple.dis_u_2, - &h1_tuple.factor_ppc, - &h1_tuple.factor_ip, - &h1_tuple.factor_err, + h1_tuple.dis_u_2, + h1_tuple.factor_ppc, + h1_tuple.factor_ip, + h1_tuple.factor_err, &h1_tuple.t, ), epsilon, ); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(( - Reverse(lowerbounds[j]), - AlwaysEqual(h1_tuple.mean[j]), - AlwaysEqual(h1_tuple.first[j]), - )); - } - } + results.push(( + Reverse(lowerbounds), + AlwaysEqual(h1_tuple.mean), + AlwaysEqual(h1_tuple.first), + )); } current = h1_guard.get().get_opaque().next; } @@ -100,23 +94,17 @@ pub fn scan( std::iter::from_fn(|| { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get() - .get(mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); - cache.push(( - Reverse(dis_u), - AlwaysEqual(first), - AlwaysEqual(if is_residual { - Some(vector_tuple.vector.to_vec()) - } else { - None - }), - )); + let Some((Some(dis_u), original)) = vectors::vector_dist( + relation.clone(), + &vector, + mean, + None, + Some(distance_kind), + is_residual, + ) else { + panic!("data corruption") + }; + cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(original))); } let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; Some((first, mean)) @@ -150,23 +138,19 @@ pub fn scan( dims, lut, ( - &h0_tuple.dis_u_2, - &h0_tuple.factor_ppc, - &h0_tuple.factor_ip, - &h0_tuple.factor_err, + h0_tuple.dis_u_2, + h0_tuple.factor_ppc, + h0_tuple.factor_ip, + h0_tuple.factor_err, &h0_tuple.t, ), epsilon, ); - for j in 0..32 { - if h0_tuple.mask[j] { - results.push(( - Reverse(lowerbounds[j]), - AlwaysEqual(h0_tuple.mean[j]), - AlwaysEqual(h0_tuple.payload[j]), - )); - } - } + results.push(( + Reverse(lowerbounds), + AlwaysEqual(h0_tuple.mean), + AlwaysEqual(h0_tuple.payload), + )); } current = h0_guard.get().get_opaque().next; } @@ -176,18 +160,16 @@ pub fn scan( std::iter::from_fn(move || { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); - let vector_guard = relation.read(mean.0); - let Some(vector_tuple) = vector_guard.get().get(mean.1) else { - // fails consistency check + let Some((Some(dis_u), _)) = vectors::vector_dist( + relation.clone(), + &vector, + mean, + Some(pay_u), + Some(distance_kind), + false, + ) else { continue; }; - let vector_tuple = rkyv::check_archived_root::(vector_tuple) - .expect("data corruption"); - if vector_tuple.payload != Some(pay_u) { - // fails consistency check - continue; - } - let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); } let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; diff --git a/src/algorithm/tuples.rs b/src/algorithm/tuples.rs index 66e9d649..afdc1cda 100644 --- a/src/algorithm/tuples.rs +++ b/src/algorithm/tuples.rs @@ -1,4 +1,3 @@ -use crate::algorithm::rabitq; use rkyv::{Archive, Deserialize, Serialize}; #[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] @@ -18,122 +17,37 @@ pub struct MetaTuple { #[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] #[archive(check_bytes)] pub struct VectorTuple { - pub vector: Vec, - // this field is saved only for vacuum + pub slice: Vec, pub payload: Option, + pub chain: Option<(u32, u16)>, } #[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] #[archive(check_bytes)] pub struct Height1Tuple { - pub mask: [bool; 32], // raw vector - pub mean: [(u32, u16); 32], + pub mean: (u32, u16), // for height 1 tuple, it's pointers to next level - pub first: [u32; 32], + pub first: u32, // RaBitQ algorithm - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub t: Vec, } #[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] #[archive(check_bytes)] pub struct Height0Tuple { - pub mask: [bool; 32], // raw vector - pub mean: [(u32, u16); 32], + pub mean: (u32, u16), // for height 0 tuple, it's pointers to heap relation - pub payload: [u64; 32], + pub payload: u64, // RaBitQ algorithm - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, -} - -pub fn put( - bytes: &mut [u8], - dims: u32, - code: &rabitq::Code, - vector: (u32, u16), - payload: u64, -) -> bool { - // todo: use mutable api - let mut x = rkyv::from_bytes::(bytes).expect("data corruption"); - for j in 0..32 { - if !x.mask[j] { - x.mean[j] = vector; - x.payload[j] = payload; - x.mask[j] = true; - x.dis_u_2[j] = code.dis_u_2; - x.factor_ppc[j] = code.factor_ppc; - x.factor_ip[j] = code.factor_ip; - x.factor_err[j] = code.factor_err; - let width = dims.div_ceil(4) as usize; - let table = [ - (0, 0), - (2, 0), - (4, 0), - (6, 0), - (8, 0), - (10, 0), - (12, 0), - (14, 0), - (1, 0), - (3, 0), - (5, 0), - (7, 0), - (9, 0), - (11, 0), - (13, 0), - (15, 0), - (0, 1), - (2, 1), - (4, 1), - (6, 1), - (8, 1), - (10, 1), - (12, 1), - (14, 1), - (1, 1), - (3, 1), - (5, 1), - (7, 1), - (9, 1), - (11, 1), - (13, 1), - (15, 1), - ]; - let pos = table[j].0; - let mask = match table[j].1 { - 0 => 0xf0, - 1 => 0x0f, - _ => unreachable!(), - }; - let shift = match table[j].1 { - 0 => 0, - 1 => 4, - _ => unreachable!(), - }; - let mut buffer = vec![0u8; width]; - for j in 0..width { - let b0 = code.signs.get(4 * j + 0).copied().unwrap_or_default(); - let b1 = code.signs.get(4 * j + 1).copied().unwrap_or_default(); - let b2 = code.signs.get(4 * j + 2).copied().unwrap_or_default(); - let b3 = code.signs.get(4 * j + 3).copied().unwrap_or_default(); - buffer[j] = b0 | b1 << 1 | b2 << 2 | b3 << 3; - } - for j in 0..width { - x.t[16 * j + pos] &= mask; - x.t[16 * j + pos] |= buffer[j] << shift; - } - bytes.copy_from_slice(&rkyv::to_bytes::<_, 8192>(&x).unwrap()); - return true; - } - } - false + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub t: Vec, } diff --git a/src/algorithm/vacuum.rs b/src/algorithm/vacuum.rs index 43f45021..493b6812 100644 --- a/src/algorithm/vacuum.rs +++ b/src/algorithm/vacuum.rs @@ -26,11 +26,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(h1_tuple.first[j]); - } - } + results.push(h1_tuple.first); } current = h1_guard.get().get_opaque().next; } @@ -45,6 +41,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - while current != u32::MAX { delay(); let mut h0_guard = relation.write(current); + let mut reconstruct_removes = Vec::new(); for i in 1..=h0_guard.get().len() { let h0_tuple = h0_guard .get() @@ -52,36 +49,11 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - let flag = 'flag: { - for j in 0..32 { - if h0_tuple.mask[j] && callback(Pointer::new(h0_tuple.payload[j])) { - break 'flag true; - } - } - false - }; - if flag { - // todo: use mutable API - let mut temp = h0_guard - .get() - .get(i) - .map(rkyv::from_bytes::) - .expect("data corruption") - .expect("data corruption"); - for j in 0..32 { - if temp.mask[j] && callback(Pointer::new(temp.payload[j])) { - temp.mask[j] = false; - } - } - let temp = rkyv::to_bytes::<_, 8192>(&temp).expect("failed to serialize"); - h0_guard - .get_mut() - .get_mut(i) - .expect("data corruption") - .copy_from_slice(&temp); + if callback(Pointer::new(h0_tuple.payload)) { + reconstruct_removes.push(i); } } - // todo: cross-tuple vacuum so that we can skip a tuple + h0_guard.get_mut().reconstruct(&reconstruct_removes); current = h0_guard.get().get_opaque().next; } } diff --git a/src/algorithm/vectors.rs b/src/algorithm/vectors.rs new file mode 100644 index 00000000..35b57d89 --- /dev/null +++ b/src/algorithm/vectors.rs @@ -0,0 +1,76 @@ +use crate::algorithm::tuples::VectorTuple; +use crate::index::utils::distance; +use crate::postgres::Relation; +use base::distance::Distance; +use base::distance::DistanceKind; + +pub fn vector_split(vector: &[f32]) -> Vec<&[f32]> { + match vector.len() { + 0..=960 => vec![vector], + 961..=1280 => vec![&vector[..640], &vector[640..]], + 1281.. => vector.chunks(1920).collect(), + } +} + +pub fn vector_dist( + relation: Relation, + vector: &[f32], + mean: (u32, u16), + payload: Option, + for_distance: Option, + for_original: bool, +) -> Option<(Option, Option>)> { + if for_distance.is_none() && !for_original && payload.is_none() { + return Some((None, None)); + } + let slices = vector_split(vector); + let mut cursor = Some(mean); + let mut result = 0.0f32; + let mut original = Vec::new(); + for i in 0..slices.len() { + let Some(mean) = cursor else { + // fails consistency check + return None; + }; + let vector_guard = relation.read(mean.0); + let Some(vector_tuple) = vector_guard.get().get(mean.1) else { + // fails consistency check + return None; + }; + let vector_tuple = + rkyv::check_archived_root::(vector_tuple).expect("data corruption"); + if vector_tuple.payload != payload { + // fails consistency check + return None; + } + if let Some(distance_kind) = for_distance { + result += distance(distance_kind, slices[i], &vector_tuple.slice).to_f32(); + } + if for_original { + original.extend_from_slice(&vector_tuple.slice); + } + cursor = vector_tuple.chain.as_ref().cloned(); + } + Some(( + for_distance.map(|_| Distance::from_f32(result)), + for_original.then_some(original), + )) +} + +pub fn vector_warm(relation: Relation, mean: (u32, u16)) { + let mut cursor = Some(mean); + while let Some(mean) = cursor { + let vector_guard = relation.read(mean.0); + let Some(vector_tuple) = vector_guard.get().get(mean.1) else { + // fails consistency check + return; + }; + let vector_tuple = + rkyv::check_archived_root::(vector_tuple).expect("data corruption"); + if vector_tuple.payload.is_some() { + // fails consistency check + return; + } + cursor = vector_tuple.chain.as_ref().cloned(); + } +} diff --git a/src/index/am.rs b/src/index/am.rs index 02f44983..514c26f8 100644 --- a/src/index/am.rs +++ b/src/index/am.rs @@ -190,7 +190,7 @@ pub unsafe extern "C" fn ambuild( let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { - let vector = match opfamily.preprocess(vector.as_borrowed()) { + let vector = match vector { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), @@ -229,7 +229,10 @@ pub unsafe extern "C" fn ambuild( if let Err(errors) = Validate::validate(&vector_options) { pgrx::error!("error while validating options: {}", errors); } - if vector_options.dims > 1600 { + if vector_options.dims == 0 { + pgrx::error!("error while validating options: dimension cannot be 0"); + } + if vector_options.dims > 60000 { pgrx::error!("error while validating options: dimension is too large"); } if let Err(errors) = Validate::validate(&vchordrq_options) { @@ -564,7 +567,7 @@ unsafe fn parallel_build( let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { - let vector = match opfamily.preprocess(vector.as_borrowed()) { + let vector = match vector { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), @@ -659,7 +662,7 @@ pub unsafe extern "C" fn aminsert( let opfamily = unsafe { am_options::opfamily(index) }; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { - let vector = match opfamily.preprocess(vector.as_borrowed()) { + let vector = match vector { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), @@ -692,7 +695,7 @@ pub unsafe extern "C" fn aminsert( let opfamily = unsafe { am_options::opfamily(index) }; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { - let vector = match opfamily.preprocess(vector.as_borrowed()) { + let vector = match vector { OwnedVector::Vecf32(x) => x, OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs index 8eb3a692..8f97061e 100644 --- a/src/index/am_scan.rs +++ b/src/index/am_scan.rs @@ -76,7 +76,7 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, if let Some((vector, opfamily)) = vector.as_ref() { let vbase = scan( relation, - match opfamily.preprocess(vector.as_borrowed()) { + match vector { OwnedVector::Vecf32(x) => x.slice().to_vec(), OwnedVector::Vecf16(_) => unreachable!(), OwnedVector::SVecf32(_) => unreachable!(), diff --git a/src/postgres.rs b/src/postgres.rs index 80d9e717..6dcd2430 100644 --- a/src/postgres.rs +++ b/src/postgres.rs @@ -91,6 +91,7 @@ impl Page { Some(std::slice::from_raw_parts(ptr, lp_len as _)) } } + #[allow(unused)] pub fn get_mut(&mut self, i: u16) -> Option<&mut [u8]> { use pgrx::pg_sys::{ItemIdData, PageHeaderData}; if i == 0 { @@ -134,6 +135,22 @@ impl Page { pgrx::pg_sys::PageIndexTupleDeleteNoCompact((self as *mut Self).cast(), i); } } + pub fn reconstruct(&mut self, removes: &[u16]) { + let mut removes = removes.to_vec(); + removes.sort(); + removes.dedup(); + let n = removes.len(); + if n > 0 { + assert!(removes[n - 1] <= self.len()); + unsafe { + pgrx::pg_sys::PageIndexMultiDelete( + (self as *mut Self).cast(), + removes.as_ptr().cast_mut(), + removes.len() as _, + ); + } + } + } pub fn freespace(&self) -> u16 { unsafe { pgrx::pg_sys::PageGetFreeSpace((self as *const Self).cast_mut().cast()) as u16 } } From cadb4822784d546588d11e7734dcde679e758360 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 20 Nov 2024 17:48:57 +0800 Subject: [PATCH 065/324] chore: directory structure (#104) Signed-off-by: usamoi --- src/lib.rs | 14 ++-- src/projection.rs | 76 +++++++++++++++++++++ src/types.rs | 30 ++++---- src/utils/cells.rs | 21 ------ src/utils/mod.rs | 1 - src/{ => vchordrq}/algorithm/build.rs | 32 ++++----- src/{ => vchordrq}/algorithm/insert.rs | 10 +-- src/{ => vchordrq}/algorithm/k_means.rs | 2 +- src/{ => vchordrq}/algorithm/mod.rs | 0 src/{ => vchordrq}/algorithm/parallelism.rs | 0 src/{ => vchordrq}/algorithm/prewarm.rs | 4 +- src/{ => vchordrq}/algorithm/rabitq.rs | 76 --------------------- src/{ => vchordrq}/algorithm/scan.rs | 10 +-- src/{ => vchordrq}/algorithm/tuples.rs | 0 src/{ => vchordrq}/algorithm/vacuum.rs | 2 +- src/{ => vchordrq}/algorithm/vectors.rs | 4 +- src/{ => vchordrq}/gucs/executing.rs | 0 src/{ => vchordrq}/gucs/mod.rs | 4 -- src/{ => vchordrq}/gucs/prewarm.rs | 2 +- src/{ => vchordrq}/index/am.rs | 46 ++++++------- src/{ => vchordrq}/index/am_options.rs | 0 src/{ => vchordrq}/index/am_scan.rs | 8 +-- src/{ => vchordrq}/index/functions.rs | 2 +- src/{ => vchordrq}/index/mod.rs | 0 src/{ => vchordrq}/index/opclass.rs | 0 src/{ => vchordrq}/index/utils.rs | 0 src/vchordrq/mod.rs | 10 +++ 27 files changed, 169 insertions(+), 185 deletions(-) create mode 100644 src/projection.rs delete mode 100644 src/utils/cells.rs delete mode 100644 src/utils/mod.rs rename src/{ => vchordrq}/algorithm/build.rs (93%) rename src/{ => vchordrq}/algorithm/insert.rs (97%) rename src/{ => vchordrq}/algorithm/k_means.rs (98%) rename src/{ => vchordrq}/algorithm/mod.rs (100%) rename src/{ => vchordrq}/algorithm/parallelism.rs (100%) rename src/{ => vchordrq}/algorithm/prewarm.rs (97%) rename src/{ => vchordrq}/algorithm/rabitq.rs (68%) rename src/{ => vchordrq}/algorithm/scan.rs (96%) rename src/{ => vchordrq}/algorithm/tuples.rs (100%) rename src/{ => vchordrq}/algorithm/vacuum.rs (99%) rename src/{ => vchordrq}/algorithm/vectors.rs (96%) rename src/{ => vchordrq}/gucs/executing.rs (100%) rename src/{ => vchordrq}/gucs/mod.rs (61%) rename src/{ => vchordrq}/gucs/prewarm.rs (93%) rename src/{ => vchordrq}/index/am.rs (96%) rename src/{ => vchordrq}/index/am_options.rs (100%) rename src/{ => vchordrq}/index/am_scan.rs (95%) rename src/{ => vchordrq}/index/functions.rs (94%) rename src/{ => vchordrq}/index/mod.rs (100%) rename src/{ => vchordrq}/index/opclass.rs (100%) rename src/{ => vchordrq}/index/utils.rs (100%) create mode 100644 src/vchordrq/mod.rs diff --git a/src/lib.rs b/src/lib.rs index 9d44ac73..61062f3d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,14 +4,12 @@ #![allow(clippy::too_many_arguments)] #![allow(clippy::type_complexity)] -mod algorithm; mod datatype; -mod gucs; -mod index; mod postgres; +mod projection; mod types; mod upgrade; -mod utils; +mod vchordrq; pgrx::pg_module_magic!(); pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap); @@ -24,8 +22,12 @@ unsafe extern "C" fn _PG_init() { } detect::init(); unsafe { - index::init(); - gucs::init(); + vchordrq::init(); + + #[cfg(any(feature = "pg13", feature = "pg14"))] + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] + pgrx::pg_sys::MarkGUCPrefixReserved(c"vchord".as_ptr()); } } diff --git a/src/projection.rs b/src/projection.rs new file mode 100644 index 00000000..68926ac9 --- /dev/null +++ b/src/projection.rs @@ -0,0 +1,76 @@ +use nalgebra::DMatrix; +use std::sync::OnceLock; + +fn random_matrix(n: usize) -> DMatrix { + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha12Rng; + use rand_distr::StandardNormal; + let mut rng = ChaCha12Rng::from_seed([7; 32]); + DMatrix::from_fn(n, n, |_, _| rng.sample(StandardNormal)) +} + +#[ignore] +#[test] +fn check_all_matrixs_are_full_rank() { + let parallelism = std::thread::available_parallelism().unwrap().get(); + std::thread::scope(|scope| { + let mut threads = vec![]; + for remainder in 0..parallelism { + threads.push(scope.spawn(move || { + for n in (0..=60000).filter(|x| x % parallelism == remainder) { + let matrix = random_matrix(n); + assert!(matrix.is_invertible()); + } + })); + } + for thread in threads { + thread.join().unwrap(); + } + }); +} + +#[test] +fn check_matrices() { + assert_eq!( + orthogonal_matrix(2), + vec![vec![-0.5424608, -0.8400813], vec![0.8400813, -0.54246056]] + ); + assert_eq!( + orthogonal_matrix(3), + vec![ + vec![-0.5309615, -0.69094884, -0.49058124], + vec![0.8222731, -0.56002235, -0.10120347], + vec![0.20481002, 0.45712686, -0.86549866] + ] + ); +} + +fn orthogonal_matrix(n: usize) -> Vec> { + use nalgebra::QR; + let matrix = random_matrix(n); + // QR decomposition is unique if the matrix is full rank + let qr = QR::new(matrix); + let q = qr.q(); + let mut projection = Vec::new(); + for row in q.row_iter() { + projection.push(row.iter().copied().collect::>()); + } + projection +} + +static MATRIXS: [OnceLock>>; 1 + 60000] = [const { OnceLock::new() }; 1 + 60000]; + +pub fn prewarm(n: usize) { + if n <= 60000 { + MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); + } +} + +pub fn project(vector: &[f32]) -> Vec { + use base::scalar::ScalarLike; + let n = vector.len(); + let matrix = MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); + (0..n) + .map(|i| f32::reduce_sum_of_xy(vector, &matrix[i])) + .collect() +} diff --git a/src/types.rs b/src/types.rs index 82349059..2002d7d6 100644 --- a/src/types.rs +++ b/src/types.rs @@ -3,18 +3,18 @@ use validator::{Validate, ValidationError, ValidationErrors}; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] -pub struct RabbitholeInternalBuildOptions { - #[serde(default = "RabbitholeInternalBuildOptions::default_lists")] - #[validate(length(min = 1, max = 8), custom(function = RabbitholeInternalBuildOptions::validate_lists))] +pub struct VchordrqInternalBuildOptions { + #[serde(default = "VchordrqInternalBuildOptions::default_lists")] + #[validate(length(min = 1, max = 8), custom(function = VchordrqInternalBuildOptions::validate_lists))] pub lists: Vec, - #[serde(default = "RabbitholeInternalBuildOptions::default_spherical_centroids")] + #[serde(default = "VchordrqInternalBuildOptions::default_spherical_centroids")] pub spherical_centroids: bool, - #[serde(default = "RabbitholeInternalBuildOptions::default_build_threads")] + #[serde(default = "VchordrqInternalBuildOptions::default_build_threads")] #[validate(range(min = 1, max = 255))] pub build_threads: u16, } -impl RabbitholeInternalBuildOptions { +impl VchordrqInternalBuildOptions { fn default_lists() -> Vec { vec![1000] } @@ -35,7 +35,7 @@ impl RabbitholeInternalBuildOptions { } } -impl Default for RabbitholeInternalBuildOptions { +impl Default for VchordrqInternalBuildOptions { fn default() -> Self { Self { lists: Self::default_lists(), @@ -47,27 +47,27 @@ impl Default for RabbitholeInternalBuildOptions { #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] -pub struct RabbitholeExternalBuildOptions { +pub struct VchordrqExternalBuildOptions { pub table: String, } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] #[serde(rename_all = "snake_case")] -pub enum RabbitholeBuildOptions { - Internal(RabbitholeInternalBuildOptions), - External(RabbitholeExternalBuildOptions), +pub enum VchordrqBuildOptions { + Internal(VchordrqInternalBuildOptions), + External(VchordrqExternalBuildOptions), } -impl Default for RabbitholeBuildOptions { +impl Default for VchordrqBuildOptions { fn default() -> Self { Self::Internal(Default::default()) } } -impl Validate for RabbitholeBuildOptions { +impl Validate for VchordrqBuildOptions { fn validate(&self) -> Result<(), ValidationErrors> { - use RabbitholeBuildOptions::*; + use VchordrqBuildOptions::*; match self { Internal(internal_build) => internal_build.validate(), External(external_build) => external_build.validate(), @@ -80,7 +80,7 @@ impl Validate for RabbitholeBuildOptions { pub struct VchordrqIndexingOptions { #[serde(default = "VchordrqIndexingOptions::default_residual_quantization")] pub residual_quantization: bool, - pub build: RabbitholeBuildOptions, + pub build: VchordrqBuildOptions, } impl VchordrqIndexingOptions { diff --git a/src/utils/cells.rs b/src/utils/cells.rs deleted file mode 100644 index ff4a31c4..00000000 --- a/src/utils/cells.rs +++ /dev/null @@ -1,21 +0,0 @@ -use std::cell::Cell; - -pub struct PgCell(Cell); - -unsafe impl Send for PgCell {} -unsafe impl Sync for PgCell {} - -impl PgCell { - pub const unsafe fn new(x: T) -> Self { - Self(Cell::new(x)) - } -} - -impl PgCell { - pub fn get(&self) -> T { - self.0.get() - } - pub fn set(&self, value: T) { - self.0.set(value); - } -} diff --git a/src/utils/mod.rs b/src/utils/mod.rs deleted file mode 100644 index 742cee90..00000000 --- a/src/utils/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod cells; diff --git a/src/algorithm/build.rs b/src/vchordrq/algorithm/build.rs similarity index 93% rename from src/algorithm/build.rs rename to src/vchordrq/algorithm/build.rs index 0b22b93d..5e446748 100644 --- a/src/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -1,14 +1,14 @@ -use crate::algorithm::k_means; -use crate::algorithm::rabitq; -use crate::algorithm::tuples::*; -use crate::algorithm::vectors; -use crate::index::am_options::Opfamily; use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; -use crate::types::RabbitholeBuildOptions; -use crate::types::RabbitholeExternalBuildOptions; -use crate::types::RabbitholeInternalBuildOptions; +use crate::types::VchordrqBuildOptions; +use crate::types::VchordrqExternalBuildOptions; use crate::types::VchordrqIndexingOptions; +use crate::types::VchordrqInternalBuildOptions; +use crate::vchordrq::algorithm::k_means; +use crate::vchordrq::algorithm::rabitq; +use crate::vchordrq::algorithm::tuples::*; +use crate::vchordrq::algorithm::vectors; +use crate::vchordrq::index::am_options::Opfamily; use base::distance::DistanceKind; use base::index::VectorOptions; use base::scalar::ScalarLike; @@ -40,12 +40,12 @@ pub fn build( let is_residual = vchordrq_options.residual_quantization && vector_options.d == DistanceKind::L2; let structures = match vchordrq_options.build { - RabbitholeBuildOptions::External(external_build) => Structure::extern_build( + VchordrqBuildOptions::External(external_build) => Structure::extern_build( vector_options.clone(), heap_relation.opfamily(), external_build.clone(), ), - RabbitholeBuildOptions::Internal(internal_build) => { + VchordrqBuildOptions::Internal(internal_build) => { let mut tuples_total = 0_u64; let samples = { let mut rand = rand::thread_rng(); @@ -149,16 +149,16 @@ impl Structure { } fn internal_build( vector_options: VectorOptions, - internal_build: RabbitholeInternalBuildOptions, + internal_build: VchordrqInternalBuildOptions, mut samples: Vec>, ) -> Vec { use std::iter::once; for sample in samples.iter_mut() { - *sample = rabitq::project(sample); + *sample = crate::projection::project(sample); } let mut result = Vec::::new(); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { - let means = crate::algorithm::parallelism::RayonParallelism::scoped( + let means = crate::vchordrq::algorithm::parallelism::RayonParallelism::scoped( internal_build.build_threads as _, Arc::new(|| { pgrx::check_for_interrupts!(); @@ -199,10 +199,10 @@ impl Structure { fn extern_build( vector_options: VectorOptions, _opfamily: Opfamily, - external_build: RabbitholeExternalBuildOptions, + external_build: VchordrqExternalBuildOptions, ) -> Vec { use std::collections::BTreeMap; - let RabbitholeExternalBuildOptions { table } = external_build; + let VchordrqExternalBuildOptions { table } = external_build; let query = format!("SELECT id, parent, vector FROM {table};"); let mut parents = BTreeMap::new(); let mut vectors = BTreeMap::new(); @@ -226,7 +226,7 @@ impl Structure { if vector_options.dims != vector.as_borrowed().dims() { pgrx::error!("extern build: incorrect dimension, id = {id}"); } - vectors.insert(id, rabitq::project(vector.slice())); + vectors.insert(id, crate::projection::project(vector.slice())); } }); let mut children = parents diff --git a/src/algorithm/insert.rs b/src/vchordrq/algorithm/insert.rs similarity index 97% rename from src/algorithm/insert.rs rename to src/vchordrq/algorithm/insert.rs index 859b1ac8..e9318c3a 100644 --- a/src/algorithm/insert.rs +++ b/src/vchordrq/algorithm/insert.rs @@ -1,8 +1,8 @@ -use crate::algorithm::rabitq; -use crate::algorithm::rabitq::fscan_process_lowerbound; -use crate::algorithm::tuples::*; -use crate::algorithm::vectors; use crate::postgres::Relation; +use crate::vchordrq::algorithm::rabitq; +use crate::vchordrq::algorithm::rabitq::fscan_process_lowerbound; +use crate::vchordrq::algorithm::tuples::*; +use crate::vchordrq::algorithm::vectors; use base::always_equal::AlwaysEqual; use base::distance::Distance; use base::distance::DistanceKind; @@ -21,7 +21,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k .expect("data corruption"); let dims = meta_tuple.dims; assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); - let vector = rabitq::project(&vector); + let vector = crate::projection::project(&vector); let is_residual = meta_tuple.is_residual; let default_lut = if !is_residual { Some(rabitq::fscan_preprocess(&vector)) diff --git a/src/algorithm/k_means.rs b/src/vchordrq/algorithm/k_means.rs similarity index 98% rename from src/algorithm/k_means.rs rename to src/vchordrq/algorithm/k_means.rs index 4d708112..c9e72295 100644 --- a/src/algorithm/k_means.rs +++ b/src/vchordrq/algorithm/k_means.rs @@ -1,4 +1,4 @@ -use crate::algorithm::parallelism::{ParallelIterator, Parallelism}; +use crate::vchordrq::algorithm::parallelism::{ParallelIterator, Parallelism}; use base::scalar::*; use half::f16; use rand::rngs::StdRng; diff --git a/src/algorithm/mod.rs b/src/vchordrq/algorithm/mod.rs similarity index 100% rename from src/algorithm/mod.rs rename to src/vchordrq/algorithm/mod.rs diff --git a/src/algorithm/parallelism.rs b/src/vchordrq/algorithm/parallelism.rs similarity index 100% rename from src/algorithm/parallelism.rs rename to src/vchordrq/algorithm/parallelism.rs diff --git a/src/algorithm/prewarm.rs b/src/vchordrq/algorithm/prewarm.rs similarity index 97% rename from src/algorithm/prewarm.rs rename to src/vchordrq/algorithm/prewarm.rs index a274458e..92441383 100644 --- a/src/algorithm/prewarm.rs +++ b/src/vchordrq/algorithm/prewarm.rs @@ -1,6 +1,6 @@ -use crate::algorithm::tuples::*; -use crate::algorithm::vectors; use crate::postgres::Relation; +use crate::vchordrq::algorithm::tuples::*; +use crate::vchordrq::algorithm::vectors; use std::fmt::Write; pub fn prewarm(relation: Relation, height: i32) -> String { diff --git a/src/algorithm/rabitq.rs b/src/vchordrq/algorithm/rabitq.rs similarity index 68% rename from src/algorithm/rabitq.rs rename to src/vchordrq/algorithm/rabitq.rs index 7c9da8d0..e7bec93c 100644 --- a/src/algorithm/rabitq.rs +++ b/src/vchordrq/algorithm/rabitq.rs @@ -1,81 +1,5 @@ use base::distance::{Distance, DistanceKind}; use base::scalar::ScalarLike; -use nalgebra::DMatrix; -use std::sync::OnceLock; - -fn random_matrix(n: usize) -> DMatrix { - use rand::{Rng, SeedableRng}; - use rand_chacha::ChaCha12Rng; - use rand_distr::StandardNormal; - let mut rng = ChaCha12Rng::from_seed([7; 32]); - DMatrix::from_fn(n, n, |_, _| rng.sample(StandardNormal)) -} - -#[ignore] -#[test] -fn check_all_matrixs_are_full_rank() { - let parallelism = std::thread::available_parallelism().unwrap().get(); - std::thread::scope(|scope| { - let mut threads = vec![]; - for remainder in 0..parallelism { - threads.push(scope.spawn(move || { - for n in (0..=60000).filter(|x| x % parallelism == remainder) { - let matrix = random_matrix(n); - assert!(matrix.is_invertible()); - } - })); - } - for thread in threads { - thread.join().unwrap(); - } - }); -} - -#[test] -fn check_matrices() { - assert_eq!( - orthogonal_matrix(2), - vec![vec![-0.5424608, -0.8400813], vec![0.8400813, -0.54246056]] - ); - assert_eq!( - orthogonal_matrix(3), - vec![ - vec![-0.5309615, -0.69094884, -0.49058124], - vec![0.8222731, -0.56002235, -0.10120347], - vec![0.20481002, 0.45712686, -0.86549866] - ] - ); -} - -fn orthogonal_matrix(n: usize) -> Vec> { - use nalgebra::QR; - let matrix = random_matrix(n); - // QR decomposition is unique if the matrix is full rank - let qr = QR::new(matrix); - let q = qr.q(); - let mut projection = Vec::new(); - for row in q.row_iter() { - projection.push(row.iter().copied().collect::>()); - } - projection -} - -static MATRIXS: [OnceLock>>; 1 + 60000] = [const { OnceLock::new() }; 1 + 60000]; - -pub fn prewarm(n: usize) { - if n <= 60000 { - MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); - } -} - -pub fn project(vector: &[f32]) -> Vec { - use base::scalar::ScalarLike; - let n = vector.len(); - let matrix = MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); - (0..n) - .map(|i| f32::reduce_sum_of_xy(vector, &matrix[i])) - .collect() -} #[derive(Debug, Clone)] pub struct Code { diff --git a/src/algorithm/scan.rs b/src/vchordrq/algorithm/scan.rs similarity index 96% rename from src/algorithm/scan.rs rename to src/vchordrq/algorithm/scan.rs index 1f665b36..14ecba1c 100644 --- a/src/algorithm/scan.rs +++ b/src/vchordrq/algorithm/scan.rs @@ -1,8 +1,8 @@ -use crate::algorithm::rabitq; -use crate::algorithm::rabitq::fscan_process_lowerbound; -use crate::algorithm::tuples::*; -use crate::algorithm::vectors; use crate::postgres::Relation; +use crate::vchordrq::algorithm::rabitq; +use crate::vchordrq::algorithm::rabitq::fscan_process_lowerbound; +use crate::vchordrq::algorithm::tuples::*; +use crate::vchordrq::algorithm::vectors; use base::always_equal::AlwaysEqual; use base::distance::Distance; use base::distance::DistanceKind; @@ -29,7 +29,7 @@ pub fn scan( let height_of_root = meta_tuple.height_of_root; assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); assert_eq!(height_of_root as usize, 1 + probes.len(), "invalid probes"); - let vector = rabitq::project(&vector); + let vector = crate::projection::project(&vector); let is_residual = meta_tuple.is_residual; let default_lut = if !is_residual { Some(rabitq::fscan_preprocess(&vector)) diff --git a/src/algorithm/tuples.rs b/src/vchordrq/algorithm/tuples.rs similarity index 100% rename from src/algorithm/tuples.rs rename to src/vchordrq/algorithm/tuples.rs diff --git a/src/algorithm/vacuum.rs b/src/vchordrq/algorithm/vacuum.rs similarity index 99% rename from src/algorithm/vacuum.rs rename to src/vchordrq/algorithm/vacuum.rs index 493b6812..aec95a7d 100644 --- a/src/algorithm/vacuum.rs +++ b/src/vchordrq/algorithm/vacuum.rs @@ -1,5 +1,5 @@ -use crate::algorithm::tuples::*; use crate::postgres::Relation; +use crate::vchordrq::algorithm::tuples::*; use base::search::Pointer; pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) -> bool) { diff --git a/src/algorithm/vectors.rs b/src/vchordrq/algorithm/vectors.rs similarity index 96% rename from src/algorithm/vectors.rs rename to src/vchordrq/algorithm/vectors.rs index 35b57d89..bf4d7828 100644 --- a/src/algorithm/vectors.rs +++ b/src/vchordrq/algorithm/vectors.rs @@ -1,6 +1,6 @@ -use crate::algorithm::tuples::VectorTuple; -use crate::index::utils::distance; use crate::postgres::Relation; +use crate::vchordrq::algorithm::tuples::VectorTuple; +use crate::vchordrq::index::utils::distance; use base::distance::Distance; use base::distance::DistanceKind; diff --git a/src/gucs/executing.rs b/src/vchordrq/gucs/executing.rs similarity index 100% rename from src/gucs/executing.rs rename to src/vchordrq/gucs/executing.rs diff --git a/src/gucs/mod.rs b/src/vchordrq/gucs/mod.rs similarity index 61% rename from src/gucs/mod.rs rename to src/vchordrq/gucs/mod.rs index f4742589..2fb489e1 100644 --- a/src/gucs/mod.rs +++ b/src/vchordrq/gucs/mod.rs @@ -7,10 +7,6 @@ pub unsafe fn init() { prewarm::init(); prewarm::prewarm(); #[cfg(any(feature = "pg13", feature = "pg14"))] - pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); - #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] - pgrx::pg_sys::MarkGUCPrefixReserved(c"vchord".as_ptr()); - #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrq".as_ptr()); diff --git a/src/gucs/prewarm.rs b/src/vchordrq/gucs/prewarm.rs similarity index 93% rename from src/gucs/prewarm.rs rename to src/vchordrq/gucs/prewarm.rs index 6d358a14..bc484367 100644 --- a/src/gucs/prewarm.rs +++ b/src/vchordrq/gucs/prewarm.rs @@ -20,7 +20,7 @@ pub fn prewarm() { if let Ok(prewarm_dim) = prewarm_dim.to_str() { for dim in prewarm_dim.split(',') { if let Ok(dim) = dim.trim().parse::() { - crate::algorithm::rabitq::prewarm(dim as _); + crate::projection::prewarm(dim as _); } else { pgrx::warning!("{dim:?} is not a valid integer"); } diff --git a/src/index/am.rs b/src/vchordrq/index/am.rs similarity index 96% rename from src/index/am.rs rename to src/vchordrq/index/am.rs index 514c26f8..b702abe2 100644 --- a/src/index/am.rs +++ b/src/vchordrq/index/am.rs @@ -1,22 +1,21 @@ -use crate::algorithm; -use crate::algorithm::build::{HeapRelation, Reporter}; -use crate::index::am_options::{Opfamily, Reloption}; -use crate::index::am_scan::Scanner; -use crate::index::utils::{ctid_to_pointer, pointer_to_ctid}; -use crate::index::{am_options, am_scan}; use crate::postgres::Relation; -use crate::utils::cells::PgCell; +use crate::vchordrq::algorithm; +use crate::vchordrq::algorithm::build::{HeapRelation, Reporter}; +use crate::vchordrq::index::am_options::{Opfamily, Reloption}; +use crate::vchordrq::index::am_scan::Scanner; +use crate::vchordrq::index::utils::{ctid_to_pointer, pointer_to_ctid}; +use crate::vchordrq::index::{am_options, am_scan}; use base::search::Pointer; use pgrx::datum::Internal; use pgrx::pg_sys::Datum; -static RELOPT_KIND_RABBITHOLE: PgCell = unsafe { PgCell::new(0) }; +static mut RELOPT_KIND_VCHORDRQ: pgrx::pg_sys::relopt_kind::Type = 0; pub unsafe fn init() { unsafe { - RELOPT_KIND_RABBITHOLE.set(pgrx::pg_sys::add_reloption_kind()); + (&raw mut RELOPT_KIND_VCHORDRQ).write(pgrx::pg_sys::add_reloption_kind()); pgrx::pg_sys::add_string_reloption( - RELOPT_KIND_RABBITHOLE.get(), + (&raw const RELOPT_KIND_VCHORDRQ).read(), c"options".as_ptr(), c"Vector index options, represented as a TOML string.".as_ptr(), c"".as_ptr(), @@ -88,7 +87,7 @@ pub unsafe extern "C" fn amoptions(reloptions: Datum, validate: bool) -> *mut pg pgrx::pg_sys::build_reloptions( reloptions, validate, - RELOPT_KIND_RABBITHOLE.get(), + (&raw const RELOPT_KIND_VCHORDRQ).read(), size_of::(), Reloption::TAB.as_ptr(), Reloption::TAB.len() as _, @@ -254,8 +253,7 @@ pub unsafe extern "C" fn ambuild( index_relation.clone(), reporter.clone(), ); - if let Some(leader) = - unsafe { RabbitholeLeader::enter(heap, index, (*index_info).ii_Concurrent) } + if let Some(leader) = unsafe { VchordrqLeader::enter(heap, index, (*index_info).ii_Concurrent) } { unsafe { parallel_build( @@ -299,7 +297,7 @@ pub unsafe extern "C" fn ambuild( unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } -struct RabbitholeShared { +struct VchordrqShared { /* Immutable state */ heaprelid: pgrx::pg_sys::Oid, indexrelid: pgrx::pg_sys::Oid, @@ -324,15 +322,15 @@ fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { ) } -struct RabbitholeLeader { +struct VchordrqLeader { pcxt: *mut pgrx::pg_sys::ParallelContext, nparticipants: i32, - vchordrqshared: *mut RabbitholeShared, + vchordrqshared: *mut VchordrqShared, tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, snapshot: pgrx::pg_sys::Snapshot, } -impl RabbitholeLeader { +impl VchordrqLeader { pub unsafe fn enter( heap: pgrx::pg_sys::Relation, index: pgrx::pg_sys::Relation, @@ -389,7 +387,7 @@ impl RabbitholeLeader { let est_tablescandesc = unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap, snapshot) }; unsafe { - estimate_chunk(&mut (*pcxt).estimator, size_of::()); + estimate_chunk(&mut (*pcxt).estimator, size_of::()); estimate_keys(&mut (*pcxt).estimator, 1); estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); estimate_keys(&mut (*pcxt).estimator, 1); @@ -409,9 +407,9 @@ impl RabbitholeLeader { let vchordrqshared = unsafe { let vchordrqshared = - pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) - .cast::(); - vchordrqshared.write(RabbitholeShared { + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) + .cast::(); + vchordrqshared.write(VchordrqShared { heaprelid: (*heap).rd_id, indexrelid: (*index).rd_id, isconcurrent, @@ -471,7 +469,7 @@ impl RabbitholeLeader { } } -impl Drop for RabbitholeLeader { +impl Drop for VchordrqLeader { fn drop(&mut self) { if !std::thread::panicking() { unsafe { @@ -493,7 +491,7 @@ pub unsafe extern "C" fn vchordrq_parallel_build_main( toc: *mut pgrx::pg_sys::shm_toc, ) { let vchordrqshared = unsafe { - pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() }; let tablescandesc = unsafe { pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) @@ -530,7 +528,7 @@ unsafe fn parallel_build( heap: pgrx::pg_sys::Relation, index_info: *mut pgrx::pg_sys::IndexInfo, tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, - vchordrqshared: *mut RabbitholeShared, + vchordrqshared: *mut VchordrqShared, mut reporter: Option, ) { #[derive(Debug, Clone)] diff --git a/src/index/am_options.rs b/src/vchordrq/index/am_options.rs similarity index 100% rename from src/index/am_options.rs rename to src/vchordrq/index/am_options.rs diff --git a/src/index/am_scan.rs b/src/vchordrq/index/am_scan.rs similarity index 95% rename from src/index/am_scan.rs rename to src/vchordrq/index/am_scan.rs index 8f97061e..e97d352d 100644 --- a/src/index/am_scan.rs +++ b/src/vchordrq/index/am_scan.rs @@ -1,9 +1,9 @@ use super::am_options::Opfamily; -use crate::algorithm::scan::scan; -use crate::gucs::executing::epsilon; -use crate::gucs::executing::max_scan_tuples; -use crate::gucs::executing::probes; use crate::postgres::Relation; +use crate::vchordrq::algorithm::scan::scan; +use crate::vchordrq::gucs::executing::epsilon; +use crate::vchordrq::gucs::executing::max_scan_tuples; +use crate::vchordrq::gucs::executing::probes; use base::distance::Distance; use base::search::*; use base::vector::*; diff --git a/src/index/functions.rs b/src/vchordrq/index/functions.rs similarity index 94% rename from src/index/functions.rs rename to src/vchordrq/index/functions.rs index c6f820dd..40df8b53 100644 --- a/src/index/functions.rs +++ b/src/vchordrq/index/functions.rs @@ -1,5 +1,5 @@ -use crate::algorithm::prewarm::prewarm; use crate::postgres::Relation; +use crate::vchordrq::algorithm::prewarm::prewarm; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass}; diff --git a/src/index/mod.rs b/src/vchordrq/index/mod.rs similarity index 100% rename from src/index/mod.rs rename to src/vchordrq/index/mod.rs diff --git a/src/index/opclass.rs b/src/vchordrq/index/opclass.rs similarity index 100% rename from src/index/opclass.rs rename to src/vchordrq/index/opclass.rs diff --git a/src/index/utils.rs b/src/vchordrq/index/utils.rs similarity index 100% rename from src/index/utils.rs rename to src/vchordrq/index/utils.rs diff --git a/src/vchordrq/mod.rs b/src/vchordrq/mod.rs new file mode 100644 index 00000000..d7ddff43 --- /dev/null +++ b/src/vchordrq/mod.rs @@ -0,0 +1,10 @@ +mod algorithm; +mod gucs; +mod index; + +pub unsafe fn init() { + unsafe { + index::init(); + gucs::init(); + } +} From 92ff15f2068a7945510d119a12c174fd6e286cb5 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 20 Nov 2024 18:56:06 +0800 Subject: [PATCH 066/324] feat: vchordrqfscan (#105) Signed-off-by: usamoi --- src/lib.rs | 4 +- src/sql/finalize.sql | 30 + src/{vchordrq/algorithm => utils}/k_means.rs | 2 +- src/utils/mod.rs | 2 + .../algorithm => utils}/parallelism.rs | 0 src/vchordrq/algorithm/build.rs | 16 +- src/vchordrq/algorithm/mod.rs | 2 - src/vchordrq/algorithm/rabitq.rs | 9 + src/vchordrq/algorithm/vectors.rs | 2 +- src/vchordrq/index/am_options.rs | 2 +- src/vchordrq/index/utils.rs | 11 - src/vchordrq/mod.rs | 1 + src/{ => vchordrq}/types.rs | 0 src/vchordrqfscan/algorithm/build.rs | 412 +++++++++ src/vchordrqfscan/algorithm/insert.rs | 258 ++++++ src/vchordrqfscan/algorithm/mod.rs | 7 + src/vchordrqfscan/algorithm/prewarm.rs | 107 +++ src/vchordrqfscan/algorithm/rabitq.rs | 171 ++++ src/vchordrqfscan/algorithm/scan.rs | 197 ++++ src/vchordrqfscan/algorithm/tuples.rs | 139 +++ src/vchordrqfscan/algorithm/vacuum.rs | 141 +++ src/vchordrqfscan/gucs/executing.rs | 76 ++ src/vchordrqfscan/gucs/mod.rs | 14 + src/vchordrqfscan/gucs/prewarm.rs | 32 + src/vchordrqfscan/index/am.rs | 864 ++++++++++++++++++ src/vchordrqfscan/index/am_options.rs | 222 +++++ src/vchordrqfscan/index/am_scan.rs | 132 +++ src/vchordrqfscan/index/functions.rs | 26 + src/vchordrqfscan/index/mod.rs | 12 + src/vchordrqfscan/index/opclass.rs | 14 + src/vchordrqfscan/index/utils.rs | 20 + src/vchordrqfscan/mod.rs | 11 + src/vchordrqfscan/types.rs | 90 ++ 33 files changed, 3001 insertions(+), 25 deletions(-) rename src/{vchordrq/algorithm => utils}/k_means.rs (98%) create mode 100644 src/utils/mod.rs rename src/{vchordrq/algorithm => utils}/parallelism.rs (100%) rename src/{ => vchordrq}/types.rs (100%) create mode 100644 src/vchordrqfscan/algorithm/build.rs create mode 100644 src/vchordrqfscan/algorithm/insert.rs create mode 100644 src/vchordrqfscan/algorithm/mod.rs create mode 100644 src/vchordrqfscan/algorithm/prewarm.rs create mode 100644 src/vchordrqfscan/algorithm/rabitq.rs create mode 100644 src/vchordrqfscan/algorithm/scan.rs create mode 100644 src/vchordrqfscan/algorithm/tuples.rs create mode 100644 src/vchordrqfscan/algorithm/vacuum.rs create mode 100644 src/vchordrqfscan/gucs/executing.rs create mode 100644 src/vchordrqfscan/gucs/mod.rs create mode 100644 src/vchordrqfscan/gucs/prewarm.rs create mode 100644 src/vchordrqfscan/index/am.rs create mode 100644 src/vchordrqfscan/index/am_options.rs create mode 100644 src/vchordrqfscan/index/am_scan.rs create mode 100644 src/vchordrqfscan/index/functions.rs create mode 100644 src/vchordrqfscan/index/mod.rs create mode 100644 src/vchordrqfscan/index/opclass.rs create mode 100644 src/vchordrqfscan/index/utils.rs create mode 100644 src/vchordrqfscan/mod.rs create mode 100644 src/vchordrqfscan/types.rs diff --git a/src/lib.rs b/src/lib.rs index 61062f3d..94741f56 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,9 +7,10 @@ mod datatype; mod postgres; mod projection; -mod types; mod upgrade; +mod utils; mod vchordrq; +mod vchordrqfscan; pgrx::pg_module_magic!(); pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap); @@ -23,6 +24,7 @@ unsafe extern "C" fn _PG_init() { detect::init(); unsafe { vchordrq::init(); + vchordrqfscan::init(); #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index be75da37..32b356c4 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -39,16 +39,28 @@ IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhan CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; +CREATE FUNCTION vchordrqfscan_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrqfscan_amhandler_wrapper'; + +CREATE FUNCTION vchordrqfscan_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrqfscan_prewarm_wrapper'; + -- List of access methods CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD Vchordrqfscan TYPE INDEX HANDLER Vchordrqfscan_amhandler; + -- List of operator families CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING Vchordrqfscan; +CREATE OPERATOR FAMILY vector_ip_ops USING Vchordrqfscan; +CREATE OPERATOR FAMILY vector_cosine_ops USING Vchordrqfscan; + -- List of operator classes CREATE OPERATOR CLASS vector_l2_ops @@ -68,3 +80,21 @@ CREATE OPERATOR CLASS vector_cosine_ops OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING Vchordrqfscan FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _Vchordrqfscan_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING Vchordrqfscan FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _Vchordrqfscan_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING Vchordrqfscan FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _Vchordrqfscan_support_vector_cosine_ops(); diff --git a/src/vchordrq/algorithm/k_means.rs b/src/utils/k_means.rs similarity index 98% rename from src/vchordrq/algorithm/k_means.rs rename to src/utils/k_means.rs index c9e72295..8b4fd9d2 100644 --- a/src/vchordrq/algorithm/k_means.rs +++ b/src/utils/k_means.rs @@ -1,4 +1,4 @@ -use crate::vchordrq::algorithm::parallelism::{ParallelIterator, Parallelism}; +use super::parallelism::{ParallelIterator, Parallelism}; use base::scalar::*; use half::f16; use rand::rngs::StdRng; diff --git a/src/utils/mod.rs b/src/utils/mod.rs new file mode 100644 index 00000000..1b07dc6c --- /dev/null +++ b/src/utils/mod.rs @@ -0,0 +1,2 @@ +pub mod k_means; +pub mod parallelism; diff --git a/src/vchordrq/algorithm/parallelism.rs b/src/utils/parallelism.rs similarity index 100% rename from src/vchordrq/algorithm/parallelism.rs rename to src/utils/parallelism.rs diff --git a/src/vchordrq/algorithm/build.rs b/src/vchordrq/algorithm/build.rs index 5e446748..3b96f0c4 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -1,14 +1,13 @@ use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; -use crate::types::VchordrqBuildOptions; -use crate::types::VchordrqExternalBuildOptions; -use crate::types::VchordrqIndexingOptions; -use crate::types::VchordrqInternalBuildOptions; -use crate::vchordrq::algorithm::k_means; use crate::vchordrq::algorithm::rabitq; use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; use crate::vchordrq::index::am_options::Opfamily; +use crate::vchordrq::types::VchordrqBuildOptions; +use crate::vchordrq::types::VchordrqExternalBuildOptions; +use crate::vchordrq::types::VchordrqIndexingOptions; +use crate::vchordrq::types::VchordrqInternalBuildOptions; use base::distance::DistanceKind; use base::index::VectorOptions; use base::scalar::ScalarLike; @@ -158,13 +157,13 @@ impl Structure { } let mut result = Vec::::new(); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { - let means = crate::vchordrq::algorithm::parallelism::RayonParallelism::scoped( + let means = crate::utils::parallelism::RayonParallelism::scoped( internal_build.build_threads as _, Arc::new(|| { pgrx::check_for_interrupts!(); }), |parallelism| { - k_means::k_means( + crate::utils::k_means::k_means( parallelism, w as usize, vector_options.dims as usize, @@ -182,7 +181,8 @@ impl Structure { if let Some(structure) = result.last() { let mut children = vec![Vec::new(); means.len()]; for i in 0..structure.len() as u32 { - let target = k_means::k_means_lookup(&structure.means[i as usize], &means); + let target = + crate::utils::k_means::k_means_lookup(&structure.means[i as usize], &means); children[target].push(i); } let (means, children) = std::iter::zip(means, children) diff --git a/src/vchordrq/algorithm/mod.rs b/src/vchordrq/algorithm/mod.rs index a160118b..88239a8e 100644 --- a/src/vchordrq/algorithm/mod.rs +++ b/src/vchordrq/algorithm/mod.rs @@ -1,7 +1,5 @@ pub mod build; pub mod insert; -pub mod k_means; -pub mod parallelism; pub mod prewarm; pub mod rabitq; pub mod scan; diff --git a/src/vchordrq/algorithm/rabitq.rs b/src/vchordrq/algorithm/rabitq.rs index e7bec93c..b3746b80 100644 --- a/src/vchordrq/algorithm/rabitq.rs +++ b/src/vchordrq/algorithm/rabitq.rs @@ -141,3 +141,12 @@ fn asymmetric_binary_dot_product(x: &[u64], y: &(Vec, Vec, Vec, V } (t0 << 0) + (t1 << 1) + (t2 << 2) + (t3 << 3) } + +pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance { + match d { + DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)), + DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)), + DistanceKind::Hamming => unimplemented!(), + DistanceKind::Jaccard => unimplemented!(), + } +} diff --git a/src/vchordrq/algorithm/vectors.rs b/src/vchordrq/algorithm/vectors.rs index bf4d7828..c1f86273 100644 --- a/src/vchordrq/algorithm/vectors.rs +++ b/src/vchordrq/algorithm/vectors.rs @@ -1,6 +1,6 @@ use crate::postgres::Relation; +use crate::vchordrq::algorithm::rabitq::distance; use crate::vchordrq::algorithm::tuples::VectorTuple; -use crate::vchordrq::index::utils::distance; use base::distance::Distance; use base::distance::DistanceKind; diff --git a/src/vchordrq/index/am_options.rs b/src/vchordrq/index/am_options.rs index 1f271baf..971273fd 100644 --- a/src/vchordrq/index/am_options.rs +++ b/src/vchordrq/index/am_options.rs @@ -1,7 +1,7 @@ use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use crate::datatype::typmod::Typmod; -use crate::types::VchordrqIndexingOptions; +use crate::vchordrq::types::VchordrqIndexingOptions; use base::distance::*; use base::index::*; use base::vector::*; diff --git a/src/vchordrq/index/utils.rs b/src/vchordrq/index/utils.rs index 35697ab3..a5d85a3f 100644 --- a/src/vchordrq/index/utils.rs +++ b/src/vchordrq/index/utils.rs @@ -1,5 +1,3 @@ -use base::distance::{Distance, DistanceKind}; -use base::scalar::ScalarLike; use base::search::*; pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { @@ -20,12 +18,3 @@ pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer { value |= ctid.ip_posid as u64; Pointer::new(value) } - -pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance { - match d { - DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)), - DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)), - DistanceKind::Hamming => unimplemented!(), - DistanceKind::Jaccard => unimplemented!(), - } -} diff --git a/src/vchordrq/mod.rs b/src/vchordrq/mod.rs index d7ddff43..c2ae9456 100644 --- a/src/vchordrq/mod.rs +++ b/src/vchordrq/mod.rs @@ -1,6 +1,7 @@ mod algorithm; mod gucs; mod index; +mod types; pub unsafe fn init() { unsafe { diff --git a/src/types.rs b/src/vchordrq/types.rs similarity index 100% rename from src/types.rs rename to src/vchordrq/types.rs diff --git a/src/vchordrqfscan/algorithm/build.rs b/src/vchordrqfscan/algorithm/build.rs new file mode 100644 index 00000000..3215e017 --- /dev/null +++ b/src/vchordrqfscan/algorithm/build.rs @@ -0,0 +1,412 @@ +use crate::postgres::BufferWriteGuard; +use crate::postgres::Relation; +use crate::vchordrqfscan::algorithm::rabitq; +use crate::vchordrqfscan::algorithm::tuples::*; +use crate::vchordrqfscan::index::am_options::Opfamily; +use crate::vchordrqfscan::types::VchordrqfscanBuildOptions; +use crate::vchordrqfscan::types::VchordrqfscanExternalBuildOptions; +use crate::vchordrqfscan::types::VchordrqfscanIndexingOptions; +use crate::vchordrqfscan::types::VchordrqfscanInternalBuildOptions; +use base::distance::DistanceKind; +use base::index::VectorOptions; +use base::scalar::ScalarLike; +use base::search::Pointer; +use rand::Rng; +use rkyv::ser::serializers::AllocSerializer; +use std::marker::PhantomData; +use std::sync::Arc; + +pub trait HeapRelation { + fn traverse(&self, progress: bool, callback: F) + where + F: FnMut((Pointer, Vec)); + fn opfamily(&self) -> Opfamily; +} + +pub trait Reporter { + fn tuples_total(&mut self, tuples_total: u64); +} + +pub fn build( + vector_options: VectorOptions, + vchordrqfscan_options: VchordrqfscanIndexingOptions, + heap_relation: T, + relation: Relation, + mut reporter: R, +) { + let dims = vector_options.dims; + let is_residual = + vchordrqfscan_options.residual_quantization && vector_options.d == DistanceKind::L2; + let structures = match vchordrqfscan_options.build { + VchordrqfscanBuildOptions::External(external_build) => Structure::extern_build( + vector_options.clone(), + heap_relation.opfamily(), + external_build.clone(), + ), + VchordrqfscanBuildOptions::Internal(internal_build) => { + let mut tuples_total = 0_u64; + let samples = { + let mut rand = rand::thread_rng(); + let max_number_of_samples = + internal_build.lists.last().unwrap().saturating_mul(256); + let mut samples = Vec::new(); + let mut number_of_samples = 0_u32; + heap_relation.traverse(false, |(_, vector)| { + assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); + if number_of_samples < max_number_of_samples { + samples.push(vector); + number_of_samples += 1; + } else { + let index = rand.gen_range(0..max_number_of_samples) as usize; + samples[index] = vector; + } + tuples_total += 1; + }); + samples + }; + reporter.tuples_total(tuples_total); + Structure::internal_build(vector_options.clone(), internal_build.clone(), samples) + } + }; + let mut meta = Tape::create(&relation, false); + assert_eq!(meta.first(), 0); + let mut forwards = Tape::::create(&relation, false); + assert_eq!(forwards.first(), 1); + let mut vectors = Tape::create(&relation, true); + assert_eq!(vectors.first(), 2); + let mut pointer_of_means = Vec::>::new(); + for i in 0..structures.len() { + let mut level = Vec::new(); + for j in 0..structures[i].len() { + let pointer = vectors.push(&VectorTuple { + payload: None, + vector: structures[i].means[j].clone(), + }); + level.push(pointer); + } + pointer_of_means.push(level); + } + let mut pointer_of_firsts = Vec::>::new(); + for i in 0..structures.len() { + let mut level = Vec::new(); + for j in 0..structures[i].len() { + if i == 0 { + let tape = Tape::::create(&relation, false); + level.push(tape.first()); + } else { + let mut tape = Tape::::create(&relation, false); + let mut cache = Vec::new(); + let h2_mean = &structures[i].means[j]; + let h2_children = &structures[i].children[j]; + for child in h2_children.iter().copied() { + let h1_mean = &structures[i - 1].means[child as usize]; + let code = if is_residual { + rabitq::code(dims, &f32::vector_sub(h1_mean, h2_mean)) + } else { + rabitq::code(dims, h1_mean) + }; + cache.push((child, code)); + if cache.len() == 32 { + let group = std::mem::take(&mut cache); + let codes = std::array::from_fn(|k| group[k].1.clone()); + let packed = rabitq::pack_codes(dims, codes); + tape.push(&Height1Tuple { + mask: [true; 32], + mean: std::array::from_fn(|k| { + pointer_of_means[i - 1][group[k].0 as usize] + }), + first: std::array::from_fn(|k| { + pointer_of_firsts[i - 1][group[k].0 as usize] + }), + dis_u_2: packed.dis_u_2, + factor_ppc: packed.factor_ppc, + factor_ip: packed.factor_ip, + factor_err: packed.factor_err, + t: packed.t, + }); + } + } + if !cache.is_empty() { + let group = std::mem::take(&mut cache); + let codes = std::array::from_fn(|k| { + if k < group.len() { + group[k].1.clone() + } else { + rabitq::dummy_code(dims) + } + }); + let packed = rabitq::pack_codes(dims, codes); + tape.push(&Height1Tuple { + mask: std::array::from_fn(|k| k < group.len()), + mean: std::array::from_fn(|k| { + if k < group.len() { + pointer_of_means[i - 1][group[k].0 as usize] + } else { + Default::default() + } + }), + first: std::array::from_fn(|k| { + if k < group.len() { + pointer_of_firsts[i - 1][group[k].0 as usize] + } else { + Default::default() + } + }), + dis_u_2: packed.dis_u_2, + factor_ppc: packed.factor_ppc, + factor_ip: packed.factor_ip, + factor_err: packed.factor_err, + t: packed.t, + }); + } + level.push(tape.first()); + } + } + pointer_of_firsts.push(level); + } + forwards.head.get_mut().get_opaque_mut().fast_forward = vectors.first(); + meta.push(&MetaTuple { + dims, + height_of_root: structures.len() as u32, + is_residual, + vectors_first: vectors.first(), + forwards_first: forwards.first(), + mean: pointer_of_means.last().unwrap()[0], + first: pointer_of_firsts.last().unwrap()[0], + }); +} + +struct Structure { + means: Vec>, + children: Vec>, +} + +impl Structure { + fn len(&self) -> usize { + self.children.len() + } + fn internal_build( + vector_options: VectorOptions, + internal_build: VchordrqfscanInternalBuildOptions, + mut samples: Vec>, + ) -> Vec { + use std::iter::once; + for sample in samples.iter_mut() { + *sample = crate::projection::project(sample); + } + let mut result = Vec::::new(); + for w in internal_build.lists.iter().rev().copied().chain(once(1)) { + let means = crate::utils::parallelism::RayonParallelism::scoped( + internal_build.build_threads as _, + Arc::new(|| { + pgrx::check_for_interrupts!(); + }), + |parallelism| { + crate::utils::k_means::k_means( + parallelism, + w as usize, + vector_options.dims as usize, + if let Some(structure) = result.last() { + &structure.means + } else { + &samples + }, + internal_build.spherical_centroids, + 10, + ) + }, + ) + .expect("failed to create thread pool"); + if let Some(structure) = result.last() { + let mut children = vec![Vec::new(); means.len()]; + for i in 0..structure.len() as u32 { + let target = + crate::utils::k_means::k_means_lookup(&structure.means[i as usize], &means); + children[target].push(i); + } + let (means, children) = std::iter::zip(means, children) + .filter(|(_, x)| !x.is_empty()) + .unzip::<_, _, Vec<_>, Vec<_>>(); + result.push(Structure { means, children }); + } else { + let children = vec![Vec::new(); means.len()]; + result.push(Structure { means, children }); + } + } + result + } + fn extern_build( + vector_options: VectorOptions, + _opfamily: Opfamily, + external_build: VchordrqfscanExternalBuildOptions, + ) -> Vec { + use std::collections::BTreeMap; + let VchordrqfscanExternalBuildOptions { table } = external_build; + let query = format!("SELECT id, parent, vector FROM {table};"); + let mut parents = BTreeMap::new(); + let mut vectors = BTreeMap::new(); + pgrx::spi::Spi::connect(|client| { + use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; + use base::vector::VectorBorrowed; + use pgrx::pg_sys::panic::ErrorReportable; + let table = client.select(&query, None, None).unwrap_or_report(); + for row in table { + let id: Option = row.get_by_name("id").unwrap(); + let parent: Option = row.get_by_name("parent").unwrap(); + let vector: Option = row.get_by_name("vector").unwrap(); + let id = id.expect("extern build: id could not be NULL"); + let vector = vector.expect("extern build: vector could not be NULL"); + let pop = parents.insert(id, parent); + if pop.is_some() { + pgrx::error!( + "external build: there are at least two lines have same id, id = {id}" + ); + } + if vector_options.dims != vector.as_borrowed().dims() { + pgrx::error!("extern build: incorrect dimension, id = {id}"); + } + vectors.insert(id, crate::projection::project(vector.slice())); + } + }); + let mut children = parents + .keys() + .map(|x| (*x, Vec::new())) + .collect::>(); + let mut root = None; + for (&id, &parent) in parents.iter() { + if let Some(parent) = parent { + if let Some(parent) = children.get_mut(&parent) { + parent.push(id); + } else { + pgrx::error!( + "external build: parent does not exist, id = {id}, parent = {parent}" + ); + } + } else { + if let Some(root) = root { + pgrx::error!("external build: two root, id = {root}, id = {id}"); + } else { + root = Some(id); + } + } + } + let Some(root) = root else { + pgrx::error!("extern build: there are no root"); + }; + let mut heights = BTreeMap::<_, _>::new(); + fn dfs_for_heights( + heights: &mut BTreeMap>, + children: &BTreeMap>, + u: i32, + ) { + if heights.contains_key(&u) { + pgrx::error!("extern build: detect a cycle, id = {u}"); + } + heights.insert(u, None); + let mut height = None; + for &v in children[&u].iter() { + dfs_for_heights(heights, children, v); + let new = heights[&v].unwrap() + 1; + if let Some(height) = height { + if height != new { + pgrx::error!("extern build: two heights, id = {u}"); + } + } else { + height = Some(new); + } + } + if height.is_none() { + height = Some(1); + } + heights.insert(u, height); + } + dfs_for_heights(&mut heights, &children, root); + let heights = heights + .into_iter() + .map(|(k, v)| (k, v.expect("not a connected graph"))) + .collect::>(); + if !(1..=8).contains(&(heights[&root] - 1)) { + pgrx::error!( + "extern build: unexpected tree height, height = {}", + heights[&root] + ); + } + let mut cursors = vec![0_u32; 1 + heights[&root] as usize]; + let mut labels = BTreeMap::new(); + for id in parents.keys().copied() { + let height = heights[&id]; + let cursor = cursors[height as usize]; + labels.insert(id, (height, cursor)); + cursors[height as usize] += 1; + } + fn extract( + height: u32, + labels: &BTreeMap, + vectors: &BTreeMap>, + children: &BTreeMap>, + ) -> (Vec>, Vec>) { + labels + .iter() + .filter(|(_, &(h, _))| h == height) + .map(|(id, _)| { + ( + vectors[id].clone(), + children[id].iter().map(|id| labels[id].1).collect(), + ) + }) + .unzip() + } + let mut result = Vec::new(); + for height in 1..=heights[&root] { + let (means, children) = extract(height, &labels, &vectors, &children); + result.push(Structure { means, children }); + } + result + } +} + +struct Tape<'a, T> { + relation: &'a Relation, + head: BufferWriteGuard, + first: u32, + tracking_freespace: bool, + _phantom: PhantomData T>, +} + +impl<'a, T> Tape<'a, T> { + fn create(relation: &'a Relation, tracking_freespace: bool) -> Self { + let head = relation.extend(tracking_freespace); + let first = head.id(); + Self { + relation, + head, + first, + tracking_freespace, + _phantom: PhantomData, + } + } + fn first(&self) -> u32 { + self.first + } +} + +impl<'a, T> Tape<'a, T> +where + T: rkyv::Serialize>, +{ + fn push(&mut self, x: &T) -> (u32, u16) { + let bytes = rkyv::to_bytes(x).expect("failed to serialize"); + if let Some(i) = self.head.get_mut().alloc(&bytes) { + (self.head.id(), i) + } else { + let next = self.relation.extend(self.tracking_freespace); + self.head.get_mut().get_opaque_mut().next = next.id(); + self.head = next; + if let Some(i) = self.head.get_mut().alloc(&bytes) { + (self.head.id(), i) + } else { + panic!("tuple is too large to fit in a fresh page") + } + } + } +} diff --git a/src/vchordrqfscan/algorithm/insert.rs b/src/vchordrqfscan/algorithm/insert.rs new file mode 100644 index 00000000..d17c22ea --- /dev/null +++ b/src/vchordrqfscan/algorithm/insert.rs @@ -0,0 +1,258 @@ +use crate::postgres::Relation; +use crate::vchordrqfscan::algorithm::rabitq; +use crate::vchordrqfscan::algorithm::rabitq::distance; +use crate::vchordrqfscan::algorithm::rabitq::fscan_process_lowerbound; +use crate::vchordrqfscan::algorithm::tuples::*; +use base::always_equal::AlwaysEqual; +use base::distance::Distance; +use base::distance::DistanceKind; +use base::scalar::ScalarLike; +use base::search::Pointer; +use std::cmp::Reverse; +use std::collections::BinaryHeap; + +pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_kind: DistanceKind) { + let meta_guard = relation.read(0); + let meta_tuple = meta_guard + .get() + .get(1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let dims = meta_tuple.dims; + assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); + let vector = crate::projection::project(&vector); + let is_residual = meta_tuple.is_residual; + let default_lut = if !is_residual { + Some(rabitq::fscan_preprocess(&vector)) + } else { + None + }; + let h0_vector = 'h0_vector: { + let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { + vector: vector.clone(), + payload: Some(payload.as_u64()), + }) + .unwrap(); + if let Some(mut write) = relation.search(tuple.len()) { + let i = write.get_mut().alloc(&tuple).unwrap(); + break 'h0_vector (write.id(), i); + } + let mut current = relation.read(1).get().get_opaque().fast_forward; + let mut changed = false; + loop { + let read = relation.read(current); + let flag = 'flag: { + if read.get().freespace() as usize >= tuple.len() { + break 'flag true; + } + if read.get().get_opaque().next == u32::MAX { + break 'flag true; + } + false + }; + if flag { + drop(read); + let mut write = relation.write(current); + if let Some(i) = write.get_mut().alloc(&tuple) { + break (current, i); + } + if write.get().get_opaque().next == u32::MAX { + if changed { + relation.write(1).get_mut().get_opaque_mut().fast_forward = write.id(); + } + let mut extend = relation.extend(true); + write.get_mut().get_opaque_mut().next = extend.id(); + if let Some(i) = extend.get_mut().alloc(&tuple) { + break (extend.id(), i); + } else { + panic!("a tuple cannot even be fit in a fresh page"); + } + } + current = write.get().get_opaque().next; + } else { + current = read.get().get_opaque().next; + } + changed = true; + } + }; + let h0_payload = payload.as_u64(); + let mut list = ( + meta_tuple.first, + if is_residual { + let vector_guard = relation.read(meta_tuple.mean.0); + let vector_tuple = vector_guard + .get() + .get(meta_tuple.mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + Some(vector_tuple.vector.to_vec()) + } else { + None + }, + ); + let make_list = |list: (u32, Option>)| { + let mut results = Vec::new(); + { + let lut = if is_residual { + &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) + } else { + default_lut.as_ref().unwrap() + }; + let mut current = list.0; + while current != u32::MAX { + let h1_guard = relation.read(current); + for i in 1..=h1_guard.get().len() { + let h1_tuple = h1_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let lowerbounds = fscan_process_lowerbound( + distance_kind, + dims, + lut, + ( + &h1_tuple.dis_u_2, + &h1_tuple.factor_ppc, + &h1_tuple.factor_ip, + &h1_tuple.factor_err, + &h1_tuple.t, + ), + 1.9, + ); + for j in 0..32 { + if h1_tuple.mask[j] { + results.push(( + Reverse(lowerbounds[j]), + AlwaysEqual(h1_tuple.mean[j]), + AlwaysEqual(h1_tuple.first[j]), + )); + } + } + } + current = h1_guard.get().get_opaque().next; + } + } + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); + { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + let vector_guard = relation.read(mean.0); + let vector_tuple = vector_guard + .get() + .get(mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); + cache.push(( + Reverse(dis_u), + AlwaysEqual(first), + AlwaysEqual(if is_residual { + Some(vector_tuple.vector.to_vec()) + } else { + None + }), + )); + } + let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop().unwrap(); + (first, mean) + } + }; + for _ in (1..meta_tuple.height_of_root).rev() { + list = make_list(list); + } + let code = if is_residual { + rabitq::code(dims, &f32::vector_sub(&vector, list.1.as_ref().unwrap())) + } else { + rabitq::code(dims, &vector) + }; + let dummy = rkyv::to_bytes::<_, 8192>(&Height0Tuple { + mask: [false; 32], + mean: [(0, 0); 32], + payload: [0; 32], + dis_u_2: [0.0f32; 32], + factor_ppc: [0.0f32; 32], + factor_ip: [0.0f32; 32], + factor_err: [0.0f32; 32], + t: vec![0; (dims.div_ceil(4) * 16) as usize], + }) + .unwrap(); + let first = list.0; + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = relation.read(current); + let flag = 'flag: { + for i in 1..=read.get().len() { + let h0_tuple = read + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + if h0_tuple.mask.iter().any(|x| *x) { + break 'flag true; + } + } + if read.get().freespace() as usize >= dummy.len() { + break 'flag true; + } + if read.get().get_opaque().next == u32::MAX { + break 'flag true; + } + false + }; + if flag { + drop(read); + let mut write = relation.write(current); + for i in 1..=write.get().len() { + let flag = put( + write.get_mut().get_mut(i).expect("data corruption"), + dims, + &code, + h0_vector, + h0_payload, + ); + if flag { + return; + } + } + if let Some(i) = write.get_mut().alloc(&dummy) { + let flag = put( + write.get_mut().get_mut(i).expect("data corruption"), + dims, + &code, + h0_vector, + h0_payload, + ); + assert!(flag, "a put fails even on a fresh tuple"); + return; + } + if write.get().get_opaque().next == u32::MAX { + let mut extend = relation.extend(false); + write.get_mut().get_opaque_mut().next = extend.id(); + if let Some(i) = extend.get_mut().alloc(&dummy) { + let flag = put( + extend.get_mut().get_mut(i).expect("data corruption"), + dims, + &code, + h0_vector, + h0_payload, + ); + assert!(flag, "a put fails even on a fresh tuple"); + return; + } else { + panic!("a tuple cannot even be fit in a fresh page"); + } + } + current = write.get().get_opaque().next; + } else { + current = read.get().get_opaque().next; + } + } +} diff --git a/src/vchordrqfscan/algorithm/mod.rs b/src/vchordrqfscan/algorithm/mod.rs new file mode 100644 index 00000000..448d9194 --- /dev/null +++ b/src/vchordrqfscan/algorithm/mod.rs @@ -0,0 +1,7 @@ +pub mod build; +pub mod insert; +pub mod prewarm; +pub mod rabitq; +pub mod scan; +pub mod tuples; +pub mod vacuum; diff --git a/src/vchordrqfscan/algorithm/prewarm.rs b/src/vchordrqfscan/algorithm/prewarm.rs new file mode 100644 index 00000000..ec7642a8 --- /dev/null +++ b/src/vchordrqfscan/algorithm/prewarm.rs @@ -0,0 +1,107 @@ +use crate::postgres::Relation; +use crate::vchordrqfscan::algorithm::tuples::*; +use std::fmt::Write; + +pub fn prewarm(relation: Relation, height: i32) -> String { + let mut message = String::new(); + let meta_guard = relation.read(0); + let meta_tuple = meta_guard + .get() + .get(1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + writeln!(message, "height of root: {}", meta_tuple.height_of_root).unwrap(); + let prewarm_max_height = if height < 0 { 0 } else { height as u32 }; + if prewarm_max_height > meta_tuple.height_of_root { + return message; + } + let mut lists = { + let mut results = Vec::new(); + let counter = 1_usize; + { + let vector_guard = relation.read(meta_tuple.mean.0); + let vector_tuple = vector_guard + .get() + .get(meta_tuple.mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let _ = vector_tuple; + results.push(meta_tuple.first); + } + writeln!(message, "number of tuples: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", counter).unwrap(); + results + }; + let mut make_lists = |lists| { + let mut counter = 0_usize; + let mut results = Vec::new(); + for list in lists { + let mut current = list; + while current != u32::MAX { + counter += 1; + pgrx::check_for_interrupts!(); + let h1_guard = relation.read(current); + for i in 1..=h1_guard.get().len() { + let h1_tuple = h1_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + for j in 0..32 { + if h1_tuple.mask[j] { + results.push(h1_tuple.first[j]); + let mean = h1_tuple.mean[j]; + let vector_guard = relation.read(mean.0); + let vector_tuple = vector_guard + .get() + .get(mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let _ = vector_tuple; + } + } + } + current = h1_guard.get().get_opaque().next; + } + } + writeln!(message, "number of tuples: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", counter).unwrap(); + results + }; + for _ in (std::cmp::max(1, prewarm_max_height)..meta_tuple.height_of_root).rev() { + lists = make_lists(lists); + } + if prewarm_max_height == 0 { + let mut counter = 0_usize; + let mut results = Vec::new(); + for list in lists { + let mut current = list; + while current != u32::MAX { + counter += 1; + pgrx::check_for_interrupts!(); + let h0_guard = relation.read(current); + for i in 1..=h0_guard.get().len() { + let h0_tuple = h0_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + for j in 0..32 { + if h0_tuple.mask[j] { + results.push(()); + } + } + } + current = h0_guard.get().get_opaque().next; + } + } + writeln!(message, "number of tuples: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", counter).unwrap(); + } + message +} diff --git a/src/vchordrqfscan/algorithm/rabitq.rs b/src/vchordrqfscan/algorithm/rabitq.rs new file mode 100644 index 00000000..cf72ca54 --- /dev/null +++ b/src/vchordrqfscan/algorithm/rabitq.rs @@ -0,0 +1,171 @@ +use base::distance::{Distance, DistanceKind}; +use base::scalar::ScalarLike; +use quantization::utils::InfiniteByteChunks; + +#[derive(Debug, Clone)] +pub struct Code { + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub signs: Vec, +} + +pub fn code(dims: u32, vector: &[f32]) -> Code { + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x_2 = f32::reduce_sum_of_x2(vector); + let dis_u = sum_of_x_2.sqrt(); + let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); + let x_x0 = dis_u / x0; + let fac_norm = (dims as f32).sqrt(); + let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); + let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); + let factor_ip = -2.0f32 / fac_norm * x_x0; + let cnt_pos = vector + .iter() + .map(|x| x.is_sign_positive() as i32) + .sum::(); + let cnt_neg = vector + .iter() + .map(|x| x.is_sign_negative() as i32) + .sum::(); + let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; + let mut signs = Vec::new(); + for i in 0..dims { + signs.push(vector[i as usize].is_sign_positive() as u8); + } + Code { + dis_u_2: sum_of_x_2, + factor_ppc, + factor_ip, + factor_err, + signs, + } +} + +pub fn dummy_code(dims: u32) -> Code { + Code { + dis_u_2: 0.0, + factor_ppc: 0.0, + factor_ip: 0.0, + factor_err: 0.0, + signs: vec![0; dims as _], + } +} + +pub struct PackedCodes { + pub dis_u_2: [f32; 32], + pub factor_ppc: [f32; 32], + pub factor_ip: [f32; 32], + pub factor_err: [f32; 32], + pub t: Vec, +} + +pub fn pack_codes(dims: u32, codes: [Code; 32]) -> PackedCodes { + PackedCodes { + dis_u_2: std::array::from_fn(|i| codes[i].dis_u_2), + factor_ppc: std::array::from_fn(|i| codes[i].factor_ppc), + factor_ip: std::array::from_fn(|i| codes[i].factor_ip), + factor_err: std::array::from_fn(|i| codes[i].factor_err), + t: { + let signs = codes.map(|code| { + InfiniteByteChunks::new(code.signs.into_iter()) + .map(|[b0, b1, b2, b3]| b0 | b1 << 1 | b2 << 2 | b3 << 3) + .take(dims.div_ceil(4) as usize) + .collect::>() + }); + quantization::fast_scan::b4::pack(dims.div_ceil(4), signs).collect() + }, + } +} + +pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { + use quantization::quantize; + let dis_v_2 = f32::reduce_sum_of_x2(vector); + let (k, b, qvector) = quantize::quantize::<15>(vector); + let qvector_sum = if vector.len() <= 4369 { + quantize::reduce_sum_of_x_as_u16(&qvector) as f32 + } else { + quantize::reduce_sum_of_x_as_u32(&qvector) as f32 + }; + (dis_v_2, b, k, qvector_sum, compress(qvector)) +} + +pub fn fscan_process_lowerbound( + distance_kind: DistanceKind, + dims: u32, + lut: &(f32, f32, f32, f32, Vec), + (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( + &[f32; 32], + &[f32; 32], + &[f32; 32], + &[f32; 32], + &[u8], + ), + epsilon: f32, +) -> [Distance; 32] { + let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let r = quantization::fast_scan::b4::fast_scan_b4(dims.div_ceil(4), t, s); + match distance_kind { + DistanceKind::L2 => std::array::from_fn(|i| { + let rough = dis_u_2[i] + + dis_v_2 + + b * factor_ppc[i] + + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; + let err = factor_err[i] * dis_v_2.sqrt(); + Distance::from_f32(rough - epsilon * err) + }), + DistanceKind::Dot => std::array::from_fn(|i| { + let rough = 0.5 * b * factor_ppc[i] + + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; + let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); + Distance::from_f32(rough - epsilon * err) + }), + DistanceKind::Hamming => unreachable!(), + DistanceKind::Jaccard => unreachable!(), + } +} + +fn compress(mut qvector: Vec) -> Vec { + let dims = qvector.len() as u32; + let width = dims.div_ceil(4); + qvector.resize(qvector.len().next_multiple_of(4), 0); + let mut t = vec![0u8; width as usize * 16]; + for i in 0..width as usize { + unsafe { + // this hint is used to skip bound checks + std::hint::assert_unchecked(4 * i + 3 < qvector.len()); + std::hint::assert_unchecked(16 * i + 15 < t.len()); + } + let t0 = qvector[4 * i + 0]; + let t1 = qvector[4 * i + 1]; + let t2 = qvector[4 * i + 2]; + let t3 = qvector[4 * i + 3]; + t[16 * i + 0b0000] = 0; + t[16 * i + 0b0001] = t0; + t[16 * i + 0b0010] = t1; + t[16 * i + 0b0011] = t1 + t0; + t[16 * i + 0b0100] = t2; + t[16 * i + 0b0101] = t2 + t0; + t[16 * i + 0b0110] = t2 + t1; + t[16 * i + 0b0111] = t2 + t1 + t0; + t[16 * i + 0b1000] = t3; + t[16 * i + 0b1001] = t3 + t0; + t[16 * i + 0b1010] = t3 + t1; + t[16 * i + 0b1011] = t3 + t1 + t0; + t[16 * i + 0b1100] = t3 + t2; + t[16 * i + 0b1101] = t3 + t2 + t0; + t[16 * i + 0b1110] = t3 + t2 + t1; + t[16 * i + 0b1111] = t3 + t2 + t1 + t0; + } + t +} + +pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance { + match d { + DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)), + DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)), + DistanceKind::Hamming => unimplemented!(), + DistanceKind::Jaccard => unimplemented!(), + } +} diff --git a/src/vchordrqfscan/algorithm/scan.rs b/src/vchordrqfscan/algorithm/scan.rs new file mode 100644 index 00000000..63264f1f --- /dev/null +++ b/src/vchordrqfscan/algorithm/scan.rs @@ -0,0 +1,197 @@ +use crate::postgres::Relation; +use crate::vchordrqfscan::algorithm::rabitq; +use crate::vchordrqfscan::algorithm::rabitq::distance; +use crate::vchordrqfscan::algorithm::rabitq::fscan_process_lowerbound; +use crate::vchordrqfscan::algorithm::tuples::*; +use base::always_equal::AlwaysEqual; +use base::distance::Distance; +use base::distance::DistanceKind; +use base::scalar::ScalarLike; +use base::search::Pointer; +use std::cmp::Reverse; +use std::collections::BinaryHeap; + +pub fn scan( + relation: Relation, + vector: Vec, + distance_kind: DistanceKind, + probes: Vec, + epsilon: f32, +) -> impl Iterator { + let meta_guard = relation.read(0); + let meta_tuple = meta_guard + .get() + .get(1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let dims = meta_tuple.dims; + let height_of_root = meta_tuple.height_of_root; + assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); + assert_eq!(height_of_root as usize, 1 + probes.len(), "invalid probes"); + let vector = crate::projection::project(&vector); + let is_residual = meta_tuple.is_residual; + let default_lut = if !is_residual { + Some(rabitq::fscan_preprocess(&vector)) + } else { + None + }; + let mut lists: Vec<_> = vec![( + meta_tuple.first, + if is_residual { + let vector_guard = relation.read(meta_tuple.mean.0); + let vector_tuple = vector_guard + .get() + .get(meta_tuple.mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + Some(vector_tuple.vector.to_vec()) + } else { + None + }, + )]; + let make_lists = |lists: Vec<(u32, Option>)>, probes| { + let mut results = Vec::new(); + for list in lists { + let lut = if is_residual { + &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) + } else { + default_lut.as_ref().unwrap() + }; + let mut current = list.0; + while current != u32::MAX { + let h1_guard = relation.read(current); + for i in 1..=h1_guard.get().len() { + let h1_tuple = h1_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let lowerbounds = fscan_process_lowerbound( + distance_kind, + dims, + lut, + ( + &h1_tuple.dis_u_2, + &h1_tuple.factor_ppc, + &h1_tuple.factor_ip, + &h1_tuple.factor_err, + &h1_tuple.t, + ), + epsilon, + ); + for j in 0..32 { + if h1_tuple.mask[j] { + results.push(( + Reverse(lowerbounds[j]), + AlwaysEqual(h1_tuple.mean[j]), + AlwaysEqual(h1_tuple.first[j]), + )); + } + } + } + current = h1_guard.get().get_opaque().next; + } + } + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); + std::iter::from_fn(|| { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + let vector_guard = relation.read(mean.0); + let vector_tuple = vector_guard + .get() + .get(mean.1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); + cache.push(( + Reverse(dis_u), + AlwaysEqual(first), + AlwaysEqual(if is_residual { + Some(vector_tuple.vector.to_vec()) + } else { + None + }), + )); + } + let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; + Some((first, mean)) + }) + .take(probes as usize) + .collect() + }; + for i in (1..meta_tuple.height_of_root).rev() { + lists = make_lists(lists, probes[i as usize - 1]); + } + { + let mut results = Vec::new(); + for list in lists { + let lut = if is_residual { + &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) + } else { + default_lut.as_ref().unwrap() + }; + let mut current = list.0; + while current != u32::MAX { + let h0_guard = relation.read(current); + for i in 1..=h0_guard.get().len() { + let h0_tuple = h0_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let lowerbounds = fscan_process_lowerbound( + distance_kind, + dims, + lut, + ( + &h0_tuple.dis_u_2, + &h0_tuple.factor_ppc, + &h0_tuple.factor_ip, + &h0_tuple.factor_err, + &h0_tuple.t, + ), + epsilon, + ); + for j in 0..32 { + if h0_tuple.mask[j] { + results.push(( + Reverse(lowerbounds[j]), + AlwaysEqual(h0_tuple.mean[j]), + AlwaysEqual(h0_tuple.payload[j]), + )); + } + } + } + current = h0_guard.get().get_opaque().next; + } + } + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _)>::new(); + std::iter::from_fn(move || { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); + let vector_guard = relation.read(mean.0); + let Some(vector_tuple) = vector_guard.get().get(mean.1) else { + // fails consistency check + continue; + }; + let vector_tuple = rkyv::check_archived_root::(vector_tuple) + .expect("data corruption"); + if vector_tuple.payload != Some(pay_u) { + // fails consistency check + continue; + } + let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); + cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); + } + let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; + Some((dis_u, Pointer::new(pay_u))) + }) + } +} diff --git a/src/vchordrqfscan/algorithm/tuples.rs b/src/vchordrqfscan/algorithm/tuples.rs new file mode 100644 index 00000000..3b43dac8 --- /dev/null +++ b/src/vchordrqfscan/algorithm/tuples.rs @@ -0,0 +1,139 @@ +use crate::vchordrqfscan::algorithm::rabitq; +use rkyv::{Archive, Deserialize, Serialize}; + +#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] +#[archive(check_bytes)] +pub struct MetaTuple { + pub dims: u32, + pub height_of_root: u32, + pub is_residual: bool, + pub vectors_first: u32, + pub forwards_first: u32, + // raw vector + pub mean: (u32, u16), + // for meta tuple, it's pointers to next level + pub first: u32, +} + +#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] +#[archive(check_bytes)] +pub struct VectorTuple { + pub vector: Vec, + // this field is saved only for vacuum + pub payload: Option, +} + +#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] +#[archive(check_bytes)] +pub struct Height1Tuple { + pub mask: [bool; 32], + // raw vector + pub mean: [(u32, u16); 32], + // for height 1 tuple, it's pointers to next level + pub first: [u32; 32], + // RaBitQ algorithm + pub dis_u_2: [f32; 32], + pub factor_ppc: [f32; 32], + pub factor_ip: [f32; 32], + pub factor_err: [f32; 32], + pub t: Vec, +} + +#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] +#[archive(check_bytes)] +pub struct Height0Tuple { + pub mask: [bool; 32], + // raw vector + pub mean: [(u32, u16); 32], + // for height 0 tuple, it's pointers to heap relation + pub payload: [u64; 32], + // RaBitQ algorithm + pub dis_u_2: [f32; 32], + pub factor_ppc: [f32; 32], + pub factor_ip: [f32; 32], + pub factor_err: [f32; 32], + pub t: Vec, +} + +pub fn put( + bytes: &mut [u8], + dims: u32, + code: &rabitq::Code, + vector: (u32, u16), + payload: u64, +) -> bool { + // todo: use mutable api + let mut x = rkyv::from_bytes::(bytes).expect("data corruption"); + for j in 0..32 { + if !x.mask[j] { + x.mean[j] = vector; + x.payload[j] = payload; + x.mask[j] = true; + x.dis_u_2[j] = code.dis_u_2; + x.factor_ppc[j] = code.factor_ppc; + x.factor_ip[j] = code.factor_ip; + x.factor_err[j] = code.factor_err; + let width = dims.div_ceil(4) as usize; + let table = [ + (0, 0), + (2, 0), + (4, 0), + (6, 0), + (8, 0), + (10, 0), + (12, 0), + (14, 0), + (1, 0), + (3, 0), + (5, 0), + (7, 0), + (9, 0), + (11, 0), + (13, 0), + (15, 0), + (0, 1), + (2, 1), + (4, 1), + (6, 1), + (8, 1), + (10, 1), + (12, 1), + (14, 1), + (1, 1), + (3, 1), + (5, 1), + (7, 1), + (9, 1), + (11, 1), + (13, 1), + (15, 1), + ]; + let pos = table[j].0; + let mask = match table[j].1 { + 0 => 0xf0, + 1 => 0x0f, + _ => unreachable!(), + }; + let shift = match table[j].1 { + 0 => 0, + 1 => 4, + _ => unreachable!(), + }; + let mut buffer = vec![0u8; width]; + for j in 0..width { + let b0 = code.signs.get(4 * j + 0).copied().unwrap_or_default(); + let b1 = code.signs.get(4 * j + 1).copied().unwrap_or_default(); + let b2 = code.signs.get(4 * j + 2).copied().unwrap_or_default(); + let b3 = code.signs.get(4 * j + 3).copied().unwrap_or_default(); + buffer[j] = b0 | b1 << 1 | b2 << 2 | b3 << 3; + } + for j in 0..width { + x.t[16 * j + pos] &= mask; + x.t[16 * j + pos] |= buffer[j] << shift; + } + bytes.copy_from_slice(&rkyv::to_bytes::<_, 8192>(&x).unwrap()); + return true; + } + } + false +} diff --git a/src/vchordrqfscan/algorithm/vacuum.rs b/src/vchordrqfscan/algorithm/vacuum.rs new file mode 100644 index 00000000..1a77b980 --- /dev/null +++ b/src/vchordrqfscan/algorithm/vacuum.rs @@ -0,0 +1,141 @@ +use crate::postgres::Relation; +use crate::vchordrqfscan::algorithm::tuples::VectorTuple; +use crate::vchordrqfscan::algorithm::tuples::*; +use base::search::Pointer; + +pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) -> bool) { + // step 1: vacuum height_0_tuple + { + let meta_guard = relation.read(0); + let meta_tuple = meta_guard + .get() + .get(1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let mut firsts = vec![meta_tuple.first]; + let make_firsts = |firsts| { + let mut results = Vec::new(); + for first in firsts { + let mut current = first; + while current != u32::MAX { + let h1_guard = relation.read(current); + for i in 1..=h1_guard.get().len() { + let h1_tuple = h1_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + for j in 0..32 { + if h1_tuple.mask[j] { + results.push(h1_tuple.first[j]); + } + } + } + current = h1_guard.get().get_opaque().next; + } + } + results + }; + for _ in (1..meta_tuple.height_of_root).rev() { + firsts = make_firsts(firsts); + } + for first in firsts { + let mut current = first; + while current != u32::MAX { + delay(); + let mut h0_guard = relation.write(current); + for i in 1..=h0_guard.get().len() { + let h0_tuple = h0_guard + .get() + .get(i) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + let flag = 'flag: { + for j in 0..32 { + if h0_tuple.mask[j] && callback(Pointer::new(h0_tuple.payload[j])) { + break 'flag true; + } + } + false + }; + if flag { + // todo: use mutable API + let mut temp = h0_guard + .get() + .get(i) + .map(rkyv::from_bytes::) + .expect("data corruption") + .expect("data corruption"); + for j in 0..32 { + if temp.mask[j] && callback(Pointer::new(temp.payload[j])) { + temp.mask[j] = false; + } + } + let temp = rkyv::to_bytes::<_, 8192>(&temp).expect("failed to serialize"); + h0_guard + .get_mut() + .get_mut(i) + .expect("data corruption") + .copy_from_slice(&temp); + } + } + // todo: cross-tuple vacuum so that we can skip a tuple + current = h0_guard.get().get_opaque().next; + } + } + } + // step 2: vacuum vector_tuple + { + let mut current = { + let meta_guard = relation.read(0); + let meta_tuple = meta_guard + .get() + .get(1) + .map(rkyv::check_archived_root::) + .expect("data corruption") + .expect("data corruption"); + meta_tuple.vectors_first + }; + while current != u32::MAX { + delay(); + let read = relation.read(current); + let flag = 'flag: { + for i in 1..=read.get().len() { + let Some(vector_tuple) = read.get().get(i) else { + continue; + }; + let vector_tuple = rkyv::check_archived_root::(vector_tuple) + .expect("data corruption"); + if let Some(payload) = vector_tuple.payload.as_ref().copied() { + if callback(Pointer::new(payload)) { + break 'flag true; + } + } + } + false + }; + if flag { + drop(read); + let mut write = relation.write(current); + for i in 1..=write.get().len() { + let Some(vector_tuple) = write.get().get(i) else { + continue; + }; + let vector_tuple = rkyv::check_archived_root::(vector_tuple) + .expect("data corruption"); + if let Some(payload) = vector_tuple.payload.as_ref().copied() { + if callback(Pointer::new(payload)) { + write.get_mut().free(i); + } + } + } + current = write.get().get_opaque().next; + } else { + current = read.get().get_opaque().next; + } + } + } +} diff --git a/src/vchordrqfscan/gucs/executing.rs b/src/vchordrqfscan/gucs/executing.rs new file mode 100644 index 00000000..6ec186ec --- /dev/null +++ b/src/vchordrqfscan/gucs/executing.rs @@ -0,0 +1,76 @@ +use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; +use std::ffi::CStr; + +static PROBES: GucSetting> = GucSetting::>::new(Some(c"10")); +static EPSILON: GucSetting = GucSetting::::new(1.9); +static MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); + +pub unsafe fn init() { + GucRegistry::define_string_guc( + "vchordrqfscan.probes", + "`probes` argument of vchordrqfscan.", + "`probes` argument of vchordrqfscan.", + &PROBES, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_float_guc( + "vchordrqfscan.epsilon", + "`epsilon` argument of vchordrqfscan.", + "`epsilon` argument of vchordrqfscan.", + &EPSILON, + 0.0, + 4.0, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + "vchordrqfscan.max_scan_tuples", + "`max_scan_tuples` argument of vchordrqfscan.", + "`max_scan_tuples` argument of vchordrqfscan.", + &MAX_SCAN_TUPLES, + -1, + u16::MAX as _, + GucContext::Userset, + GucFlags::default(), + ); +} + +pub fn probes() -> Vec { + match PROBES.get() { + None => Vec::new(), + Some(probes) => { + let mut result = Vec::new(); + let mut current = None; + for &c in probes.to_bytes() { + match c { + b' ' => continue, + b',' => result.push(current.take().expect("empty probes")), + b'0'..=b'9' => { + if let Some(x) = current.as_mut() { + *x = *x * 10 + (c - b'0') as u32; + } else { + current = Some((c - b'0') as u32); + } + } + c => pgrx::error!("unknown character in probes: ASCII = {c}"), + } + } + result.push(current.take().expect("empty probes")); + result + } + } +} + +pub fn epsilon() -> f32 { + EPSILON.get() as f32 +} + +pub fn max_scan_tuples() -> Option { + let x = MAX_SCAN_TUPLES.get(); + if x < 0 { + None + } else { + Some(x as u32) + } +} diff --git a/src/vchordrqfscan/gucs/mod.rs b/src/vchordrqfscan/gucs/mod.rs new file mode 100644 index 00000000..48cc060d --- /dev/null +++ b/src/vchordrqfscan/gucs/mod.rs @@ -0,0 +1,14 @@ +pub mod executing; +pub mod prewarm; + +pub unsafe fn init() { + unsafe { + executing::init(); + prewarm::init(); + prewarm::prewarm(); + #[cfg(any(feature = "pg13", feature = "pg14"))] + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrqfscan".as_ptr()); + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] + pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrqfscan".as_ptr()); + } +} diff --git a/src/vchordrqfscan/gucs/prewarm.rs b/src/vchordrqfscan/gucs/prewarm.rs new file mode 100644 index 00000000..ae9180a5 --- /dev/null +++ b/src/vchordrqfscan/gucs/prewarm.rs @@ -0,0 +1,32 @@ +use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; +use std::ffi::CStr; + +static PREWARM_DIM: GucSetting> = + GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); + +pub unsafe fn init() { + GucRegistry::define_string_guc( + "vchordrqfscan.prewarm_dim", + "prewarm_dim when the extension is loading.", + "prewarm_dim when the extension is loading.", + &PREWARM_DIM, + GucContext::Userset, + GucFlags::default(), + ); +} + +pub fn prewarm() { + if let Some(prewarm_dim) = PREWARM_DIM.get() { + if let Ok(prewarm_dim) = prewarm_dim.to_str() { + for dim in prewarm_dim.split(',') { + if let Ok(dim) = dim.trim().parse::() { + crate::projection::prewarm(dim as _); + } else { + pgrx::warning!("{dim:?} is not a valid integer"); + } + } + } else { + pgrx::warning!("vchordrqfscan.prewarm_dim is not a valid UTF-8 string"); + } + } +} diff --git a/src/vchordrqfscan/index/am.rs b/src/vchordrqfscan/index/am.rs new file mode 100644 index 00000000..e2344148 --- /dev/null +++ b/src/vchordrqfscan/index/am.rs @@ -0,0 +1,864 @@ +use crate::postgres::Relation; +use crate::vchordrqfscan::algorithm; +use crate::vchordrqfscan::algorithm::build::{HeapRelation, Reporter}; +use crate::vchordrqfscan::index::am_options::{Opfamily, Reloption}; +use crate::vchordrqfscan::index::am_scan::Scanner; +use crate::vchordrqfscan::index::utils::{ctid_to_pointer, pointer_to_ctid}; +use crate::vchordrqfscan::index::{am_options, am_scan}; +use base::search::Pointer; +use pgrx::datum::Internal; +use pgrx::pg_sys::Datum; + +static mut RELOPT_KIND_VCHORDRQFSCAN: pgrx::pg_sys::relopt_kind::Type = 0; + +pub unsafe fn init() { + unsafe { + (&raw mut RELOPT_KIND_VCHORDRQFSCAN).write(pgrx::pg_sys::add_reloption_kind()); + pgrx::pg_sys::add_string_reloption( + (&raw const RELOPT_KIND_VCHORDRQFSCAN).read(), + c"options".as_ptr(), + c"Vector index options, represented as a TOML string.".as_ptr(), + c"".as_ptr(), + None, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + } +} + +#[pgrx::pg_extern(sql = "")] +fn _vchordrqfscan_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { + type T = pgrx::pg_sys::IndexAmRoutine; + unsafe { + let index_am_routine = pgrx::pg_sys::palloc0(size_of::()) as *mut T; + index_am_routine.write(AM_HANDLER); + Internal::from(Some(Datum::from(index_am_routine))) + } +} + +const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = { + let mut am_routine = + unsafe { std::mem::MaybeUninit::::zeroed().assume_init() }; + + am_routine.type_ = pgrx::pg_sys::NodeTag::T_IndexAmRoutine; + + am_routine.amsupport = 1; + am_routine.amcanorderbyop = true; + + #[cfg(feature = "pg17")] + { + am_routine.amcanbuildparallel = true; + } + + // Index access methods that set `amoptionalkey` to `false` + // must index all tuples, even if the first column is `NULL`. + // However, PostgreSQL does not generate a path if there is no + // index clauses, even if there is a `ORDER BY` clause. + // So we have to set it to `true` and set costs of every path + // for vector index scans without `ORDER BY` clauses a large number + // and throw errors if someone really wants such a path. + am_routine.amoptionalkey = true; + + am_routine.amvalidate = Some(amvalidate); + am_routine.amoptions = Some(amoptions); + am_routine.amcostestimate = Some(amcostestimate); + + am_routine.ambuild = Some(ambuild); + am_routine.ambuildempty = Some(ambuildempty); + am_routine.aminsert = Some(aminsert); + am_routine.ambulkdelete = Some(ambulkdelete); + am_routine.amvacuumcleanup = Some(amvacuumcleanup); + + am_routine.ambeginscan = Some(ambeginscan); + am_routine.amrescan = Some(amrescan); + am_routine.amgettuple = Some(amgettuple); + am_routine.amendscan = Some(amendscan); + + am_routine +}; + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { + true +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amoptions(reloptions: Datum, validate: bool) -> *mut pgrx::pg_sys::bytea { + let rdopts = unsafe { + pgrx::pg_sys::build_reloptions( + reloptions, + validate, + (&raw const RELOPT_KIND_VCHORDRQFSCAN).read(), + size_of::(), + Reloption::TAB.as_ptr(), + Reloption::TAB.len() as _, + ) + }; + rdopts as *mut pgrx::pg_sys::bytea +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amcostestimate( + _root: *mut pgrx::pg_sys::PlannerInfo, + path: *mut pgrx::pg_sys::IndexPath, + _loop_count: f64, + index_startup_cost: *mut pgrx::pg_sys::Cost, + index_total_cost: *mut pgrx::pg_sys::Cost, + index_selectivity: *mut pgrx::pg_sys::Selectivity, + index_correlation: *mut f64, + index_pages: *mut f64, +) { + unsafe { + if (*path).indexorderbys.is_null() && (*path).indexclauses.is_null() { + *index_startup_cost = f64::MAX; + *index_total_cost = f64::MAX; + *index_selectivity = 0.0; + *index_correlation = 0.0; + *index_pages = 0.0; + return; + } + *index_startup_cost = 0.0; + *index_total_cost = 0.0; + *index_selectivity = 1.0; + *index_correlation = 1.0; + *index_pages = 0.0; + } +} + +#[derive(Debug, Clone)] +struct PgReporter {} + +impl Reporter for PgReporter { + fn tuples_total(&mut self, tuples_total: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, + tuples_total as _, + ); + } + } +} + +impl PgReporter { + fn tuples_done(&mut self, tuples_done: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, + tuples_done as _, + ); + } + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn ambuild( + heap: pgrx::pg_sys::Relation, + index: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, +) -> *mut pgrx::pg_sys::IndexBuildResult { + use validator::Validate; + #[derive(Debug, Clone)] + pub struct Heap { + heap: pgrx::pg_sys::Relation, + index: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + opfamily: Opfamily, + } + impl HeapRelation for Heap { + fn traverse(&self, progress: bool, callback: F) + where + F: FnMut((Pointer, Vec)), + { + pub struct State<'a, F> { + pub this: &'a Heap, + pub callback: F, + } + #[pgrx::pg_guard] + unsafe extern "C" fn call( + _index: pgrx::pg_sys::Relation, + ctid: pgrx::pg_sys::ItemPointer, + values: *mut Datum, + is_null: *mut bool, + _tuple_is_alive: bool, + state: *mut core::ffi::c_void, + ) where + F: FnMut((Pointer, Vec)), + { + use base::vector::OwnedVector; + let state = unsafe { &mut *state.cast::>() }; + let opfamily = state.this.opfamily; + let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; + let pointer = unsafe { ctid_to_pointer(ctid.read()) }; + if let Some(vector) = vector { + let vector = match vector { + OwnedVector::Vecf32(x) => x, + OwnedVector::Vecf16(_) => unreachable!(), + OwnedVector::SVecf32(_) => unreachable!(), + OwnedVector::BVector(_) => unreachable!(), + }; + (state.callback)((pointer, vector.into_vec())); + } + } + let table_am = unsafe { &*(*self.heap).rd_tableam }; + let mut state = State { + this: self, + callback, + }; + unsafe { + table_am.index_build_range_scan.unwrap()( + self.heap, + self.index, + self.index_info, + true, + false, + progress, + 0, + pgrx::pg_sys::InvalidBlockNumber, + Some(call::), + (&mut state) as *mut State as *mut _, + std::ptr::null_mut(), + ); + } + } + + fn opfamily(&self) -> Opfamily { + self.opfamily + } + } + let (vector_options, vchordrqfscan_options) = unsafe { am_options::options(index) }; + if let Err(errors) = Validate::validate(&vector_options) { + pgrx::error!("error while validating options: {}", errors); + } + if vector_options.dims == 0 { + pgrx::error!("error while validating options: dimension cannot be 0"); + } + if vector_options.dims > 1600 { + pgrx::error!("error while validating options: dimension is too large"); + } + if let Err(errors) = Validate::validate(&vchordrqfscan_options) { + pgrx::error!("error while validating options: {}", errors); + } + let opfamily = unsafe { am_options::opfamily(index) }; + let heap_relation = Heap { + heap, + index, + index_info, + opfamily, + }; + let mut reporter = PgReporter {}; + let index_relation = unsafe { Relation::new(index) }; + algorithm::build::build( + vector_options, + vchordrqfscan_options, + heap_relation.clone(), + index_relation.clone(), + reporter.clone(), + ); + if let Some(leader) = + unsafe { VchordrqfscanLeader::enter(heap, index, (*index_info).ii_Concurrent) } + { + unsafe { + parallel_build( + index, + heap, + index_info, + leader.tablescandesc, + leader.vchordrqfscanshared, + Some(reporter), + ); + leader.wait(); + let nparticipants = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordrqfscanshared).mutex); + if (*leader.vchordrqfscanshared).nparticipantsdone == nparticipants { + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqfscanshared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqfscanshared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*leader.vchordrqfscanshared).workersdonecv, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + } + } else { + let mut indtuples = 0; + reporter.tuples_done(indtuples); + heap_relation.traverse(true, |(payload, vector)| { + algorithm::insert::insert( + index_relation.clone(), + payload, + vector, + opfamily.distance_kind(), + ); + indtuples += 1; + reporter.tuples_done(indtuples); + }); + } + unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } +} + +struct VchordrqfscanShared { + /* Immutable state */ + heaprelid: pgrx::pg_sys::Oid, + indexrelid: pgrx::pg_sys::Oid, + isconcurrent: bool, + + /* Worker progress */ + workersdonecv: pgrx::pg_sys::ConditionVariable, + + /* Mutex for mutable state */ + mutex: pgrx::pg_sys::slock_t, + + /* Mutable state */ + nparticipantsdone: i32, + indtuples: u64, +} + +fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { + matches!( + unsafe { (*snapshot).snapshot_type }, + pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + | pgrx::pg_sys::SnapshotType::SNAPSHOT_HISTORIC_MVCC + ) +} + +struct VchordrqfscanLeader { + pcxt: *mut pgrx::pg_sys::ParallelContext, + nparticipants: i32, + vchordrqfscanshared: *mut VchordrqfscanShared, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + snapshot: pgrx::pg_sys::Snapshot, +} + +impl VchordrqfscanLeader { + pub unsafe fn enter( + heap: pgrx::pg_sys::Relation, + index: pgrx::pg_sys::Relation, + isconcurrent: bool, + ) -> Option { + unsafe fn compute_parallel_workers( + heap: pgrx::pg_sys::Relation, + index: pgrx::pg_sys::Relation, + ) -> i32 { + unsafe { + if pgrx::pg_sys::plan_create_index_workers((*heap).rd_id, (*index).rd_id) == 0 { + return 0; + } + if !(*heap).rd_options.is_null() { + let std_options = (*heap).rd_options.cast::(); + std::cmp::min( + (*std_options).parallel_workers, + pgrx::pg_sys::max_parallel_maintenance_workers, + ) + } else { + pgrx::pg_sys::max_parallel_maintenance_workers + } + } + } + + let request = unsafe { compute_parallel_workers(heap, index) }; + if request <= 0 { + return None; + } + + unsafe { + pgrx::pg_sys::EnterParallelMode(); + } + let pcxt = unsafe { + pgrx::pg_sys::CreateParallelContext( + c"vchord".as_ptr(), + c"vchordrqfscan_parallel_build_main".as_ptr(), + request, + ) + }; + + let snapshot = if isconcurrent { + unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } + } else { + &raw mut pgrx::pg_sys::SnapshotAnyData + }; + + fn estimate_chunk(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.space_for_chunks += x.next_multiple_of(pgrx::pg_sys::ALIGNOF_BUFFER as _); + } + fn estimate_keys(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.number_of_keys += x; + } + let est_tablescandesc = + unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap, snapshot) }; + unsafe { + estimate_chunk(&mut (*pcxt).estimator, size_of::()); + estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); + estimate_keys(&mut (*pcxt).estimator, 1); + } + + unsafe { + pgrx::pg_sys::InitializeParallelDSM(pcxt); + if (*pcxt).seg.is_null() { + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + let vchordrqfscanshared = unsafe { + let vchordrqfscanshared = + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) + .cast::(); + vchordrqfscanshared.write(VchordrqfscanShared { + heaprelid: (*heap).rd_id, + indexrelid: (*index).rd_id, + isconcurrent, + workersdonecv: std::mem::zeroed(), + mutex: std::mem::zeroed(), + nparticipantsdone: 0, + indtuples: 0, + }); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqfscanshared).workersdonecv); + pgrx::pg_sys::SpinLockInit(&raw mut (*vchordrqfscanshared).mutex); + vchordrqfscanshared + }; + + let tablescandesc = unsafe { + let tablescandesc = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, est_tablescandesc) + .cast::(); + pgrx::pg_sys::table_parallelscan_initialize(heap, tablescandesc, snapshot); + tablescandesc + }; + + unsafe { + pgrx::pg_sys::shm_toc_insert( + (*pcxt).toc, + 0xA000000000000001, + vchordrqfscanshared.cast(), + ); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); + } + + unsafe { + pgrx::pg_sys::LaunchParallelWorkers(pcxt); + } + + let nworkers_launched = unsafe { (*pcxt).nworkers_launched }; + + unsafe { + if nworkers_launched == 0 { + pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + Some(Self { + pcxt, + nparticipants: nworkers_launched + 1, + vchordrqfscanshared, + tablescandesc, + snapshot, + }) + } + + pub fn wait(&self) { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToAttach(self.pcxt); + } + } +} + +impl Drop for VchordrqfscanLeader { + fn drop(&mut self) { + if !std::thread::panicking() { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); + if is_mvcc_snapshot(self.snapshot) { + pgrx::pg_sys::UnregisterSnapshot(self.snapshot); + } + pgrx::pg_sys::DestroyParallelContext(self.pcxt); + pgrx::pg_sys::ExitParallelMode(); + } + } + } +} + +#[pgrx::pg_guard] +#[no_mangle] +pub unsafe extern "C" fn vchordrqfscan_parallel_build_main( + _seg: *mut pgrx::pg_sys::dsm_segment, + toc: *mut pgrx::pg_sys::shm_toc, +) { + let vchordrqfscanshared = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() + }; + let tablescandesc = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) + .cast::() + }; + let heap_lockmode; + let index_lockmode; + if unsafe { !(*vchordrqfscanshared).isconcurrent } { + heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; + index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; + } else { + heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; + index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; + } + let heap = unsafe { pgrx::pg_sys::table_open((*vchordrqfscanshared).heaprelid, heap_lockmode) }; + let index = + unsafe { pgrx::pg_sys::index_open((*vchordrqfscanshared).indexrelid, index_lockmode) }; + let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; + unsafe { + (*index_info).ii_Concurrent = (*vchordrqfscanshared).isconcurrent; + } + + unsafe { + parallel_build( + index, + heap, + index_info, + tablescandesc, + vchordrqfscanshared, + None, + ); + } + + unsafe { + pgrx::pg_sys::index_close(index, index_lockmode); + pgrx::pg_sys::table_close(heap, heap_lockmode); + } +} + +unsafe fn parallel_build( + index: *mut pgrx::pg_sys::RelationData, + heap: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + vchordrqfscanshared: *mut VchordrqfscanShared, + mut reporter: Option, +) { + #[derive(Debug, Clone)] + pub struct Heap { + heap: pgrx::pg_sys::Relation, + index: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + opfamily: Opfamily, + scan: *mut pgrx::pg_sys::TableScanDescData, + } + impl HeapRelation for Heap { + fn traverse(&self, progress: bool, callback: F) + where + F: FnMut((Pointer, Vec)), + { + pub struct State<'a, F> { + pub this: &'a Heap, + pub callback: F, + } + #[pgrx::pg_guard] + unsafe extern "C" fn call( + _index: pgrx::pg_sys::Relation, + ctid: pgrx::pg_sys::ItemPointer, + values: *mut Datum, + is_null: *mut bool, + _tuple_is_alive: bool, + state: *mut core::ffi::c_void, + ) where + F: FnMut((Pointer, Vec)), + { + use base::vector::OwnedVector; + let state = unsafe { &mut *state.cast::>() }; + let opfamily = state.this.opfamily; + let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; + let pointer = unsafe { ctid_to_pointer(ctid.read()) }; + if let Some(vector) = vector { + let vector = match vector { + OwnedVector::Vecf32(x) => x, + OwnedVector::Vecf16(_) => unreachable!(), + OwnedVector::SVecf32(_) => unreachable!(), + OwnedVector::BVector(_) => unreachable!(), + }; + (state.callback)((pointer, vector.into_vec())); + } + } + let table_am = unsafe { &*(*self.heap).rd_tableam }; + let mut state = State { + this: self, + callback, + }; + unsafe { + table_am.index_build_range_scan.unwrap()( + self.heap, + self.index, + self.index_info, + true, + false, + progress, + 0, + pgrx::pg_sys::InvalidBlockNumber, + Some(call::), + (&mut state) as *mut State as *mut _, + self.scan, + ); + } + } + + fn opfamily(&self) -> Opfamily { + self.opfamily + } + } + + let index_relation = unsafe { Relation::new(index) }; + let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap, tablescandesc) }; + let opfamily = unsafe { am_options::opfamily(index) }; + let heap_relation = Heap { + heap, + index, + index_info, + opfamily, + scan, + }; + heap_relation.traverse(reporter.is_some(), |(payload, vector)| { + algorithm::insert::insert( + index_relation.clone(), + payload, + vector, + opfamily.distance_kind(), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqfscanshared).mutex); + (*vchordrqfscanshared).indtuples += 1; + indtuples = (*vchordrqfscanshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqfscanshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); + + unsafe { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqfscanshared).mutex); + (*vchordrqfscanshared).nparticipantsdone += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqfscanshared).mutex); + pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordrqfscanshared).workersdonecv); + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn ambuildempty(_index: pgrx::pg_sys::Relation) { + pgrx::error!("Unlogged indexes are not supported."); +} + +#[cfg(feature = "pg13")] +#[pgrx::pg_guard] +pub unsafe extern "C" fn aminsert( + index: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, + _heap: pgrx::pg_sys::Relation, + _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, + _index_info: *mut pgrx::pg_sys::IndexInfo, +) -> bool { + use base::vector::OwnedVector; + let opfamily = unsafe { am_options::opfamily(index) }; + let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; + if let Some(vector) = vector { + let vector = match vector { + OwnedVector::Vecf32(x) => x, + OwnedVector::Vecf16(_) => unreachable!(), + OwnedVector::SVecf32(_) => unreachable!(), + OwnedVector::BVector(_) => unreachable!(), + }; + let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); + algorithm::insert::insert( + unsafe { Relation::new(index) }, + pointer, + vector.into_vec(), + opfamily.distance_kind(), + ); + } + false +} + +#[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] +#[pgrx::pg_guard] +pub unsafe extern "C" fn aminsert( + index: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, + _heap: pgrx::pg_sys::Relation, + _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, + _index_unchanged: bool, + _index_info: *mut pgrx::pg_sys::IndexInfo, +) -> bool { + use base::vector::OwnedVector; + let opfamily = unsafe { am_options::opfamily(index) }; + let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; + if let Some(vector) = vector { + let vector = match vector { + OwnedVector::Vecf32(x) => x, + OwnedVector::Vecf16(_) => unreachable!(), + OwnedVector::SVecf32(_) => unreachable!(), + OwnedVector::BVector(_) => unreachable!(), + }; + let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); + algorithm::insert::insert( + unsafe { Relation::new(index) }, + pointer, + vector.into_vec(), + opfamily.distance_kind(), + ); + } + false +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn ambeginscan( + index: pgrx::pg_sys::Relation, + n_keys: std::os::raw::c_int, + n_orderbys: std::os::raw::c_int, +) -> pgrx::pg_sys::IndexScanDesc { + use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; + + let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index, n_keys, n_orderbys) }; + unsafe { + let scanner = am_scan::scan_make(None, None, false); + (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); + } + scan +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amrescan( + scan: pgrx::pg_sys::IndexScanDesc, + keys: pgrx::pg_sys::ScanKey, + _n_keys: std::os::raw::c_int, + orderbys: pgrx::pg_sys::ScanKey, + _n_orderbys: std::os::raw::c_int, +) { + unsafe { + if !keys.is_null() && (*scan).numberOfKeys > 0 { + std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); + } + if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { + std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); + } + let opfamily = am_options::opfamily((*scan).indexRelation); + let (orderbys, spheres) = { + let mut orderbys = Vec::new(); + let mut spheres = Vec::new(); + if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { + pgrx::error!( + "vector search with no WHERE clause and no ORDER BY clause is not supported" + ); + } + for i in 0..(*scan).numberOfOrderBys { + let data = (*scan).orderByData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + match (*data).sk_strategy { + 1 => orderbys.push(opfamily.datum_to_vector(value, is_null)), + _ => unreachable!(), + } + } + for i in 0..(*scan).numberOfKeys { + let data = (*scan).keyData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + match (*data).sk_strategy { + 2 => spheres.push(opfamily.datum_to_sphere(value, is_null)), + _ => unreachable!(), + } + } + (orderbys, spheres) + }; + let (vector, threshold, recheck) = am_scan::scan_build(orderbys, spheres, opfamily); + let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); + let scanner = std::mem::replace(scanner, am_scan::scan_make(vector, threshold, recheck)); + am_scan::scan_release(scanner); + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amgettuple( + scan: pgrx::pg_sys::IndexScanDesc, + direction: pgrx::pg_sys::ScanDirection::Type, +) -> bool { + if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { + pgrx::error!("vector search without a forward scan direction is not supported"); + } + // https://www.postgresql.org/docs/current/index-locking.html + // If heap entries referenced physical pointers are deleted before + // they are consumed by PostgreSQL, PostgreSQL will received wrong + // physical pointers: no rows or irreverent rows are referenced. + if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + { + pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); + } + let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; + let relation = unsafe { Relation::new((*scan).indexRelation) }; + if let Some((pointer, recheck)) = am_scan::scan_next(scanner, relation) { + let ctid = pointer_to_ctid(pointer); + unsafe { + (*scan).xs_heaptid = ctid; + (*scan).xs_recheckorderby = false; + (*scan).xs_recheck = recheck; + } + true + } else { + false + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { + unsafe { + let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); + let scanner = std::mem::replace(scanner, am_scan::scan_make(None, None, false)); + am_scan::scan_release(scanner); + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn ambulkdelete( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, + callback: pgrx::pg_sys::IndexBulkDeleteCallback, + callback_state: *mut std::os::raw::c_void, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let callback = callback.unwrap(); + let callback = |p: Pointer| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; + algorithm::vacuum::vacuum( + unsafe { Relation::new((*info).index) }, + || unsafe { + pgrx::pg_sys::vacuum_delay_point(); + }, + callback, + ); + stats +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amvacuumcleanup( + _info: *mut pgrx::pg_sys::IndexVacuumInfo, + _stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + std::ptr::null_mut() +} diff --git a/src/vchordrqfscan/index/am_options.rs b/src/vchordrqfscan/index/am_options.rs new file mode 100644 index 00000000..51a1009b --- /dev/null +++ b/src/vchordrqfscan/index/am_options.rs @@ -0,0 +1,222 @@ +use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; +use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; +use crate::datatype::typmod::Typmod; +use crate::vchordrqfscan::types::VchordrqfscanIndexingOptions; +use base::distance::*; +use base::index::*; +use base::vector::*; +use pgrx::datum::FromDatum; +use pgrx::heap_tuple::PgHeapTuple; +use serde::Deserialize; +use std::ffi::CStr; +use std::num::NonZero; + +#[derive(Copy, Clone, Debug, Default)] +#[repr(C)] +pub struct Reloption { + vl_len_: i32, + pub options: i32, +} + +impl Reloption { + pub const TAB: &'static [pgrx::pg_sys::relopt_parse_elt] = &[pgrx::pg_sys::relopt_parse_elt { + optname: c"options".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, + offset: std::mem::offset_of!(Reloption, options) as i32, + }]; + unsafe fn options(&self) -> &CStr { + unsafe { + let ptr = std::ptr::addr_of!(*self) + .cast::() + .offset(self.options as _); + CStr::from_ptr(ptr) + } + } +} + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum PgDistanceKind { + L2, + Dot, + Cos, +} + +impl PgDistanceKind { + pub fn to_distance(self) -> DistanceKind { + match self { + PgDistanceKind::L2 => DistanceKind::L2, + PgDistanceKind::Dot | PgDistanceKind::Cos => DistanceKind::Dot, + } + } +} + +fn convert_name_to_vd(name: &str) -> Option<(VectorKind, PgDistanceKind)> { + match name.strip_suffix("_ops") { + Some("vector_l2") => Some((VectorKind::Vecf32, PgDistanceKind::L2)), + Some("vector_ip") => Some((VectorKind::Vecf32, PgDistanceKind::Dot)), + Some("vector_cosine") => Some((VectorKind::Vecf32, PgDistanceKind::Cos)), + _ => None, + } +} + +unsafe fn convert_reloptions_to_options( + reloptions: *const pgrx::pg_sys::varlena, +) -> VchordrqfscanIndexingOptions { + #[derive(Debug, Clone, Deserialize, Default)] + #[serde(deny_unknown_fields)] + struct Parsed { + #[serde(flatten)] + rabitq: VchordrqfscanIndexingOptions, + } + let reloption = reloptions as *const Reloption; + if reloption.is_null() || unsafe { (*reloption).options == 0 } { + return Default::default(); + } + let s = unsafe { (*reloption).options() }.to_string_lossy(); + match toml::from_str::(&s) { + Ok(p) => p.rabitq, + Err(e) => pgrx::error!("failed to parse options: {}", e), + } +} + +pub unsafe fn options( + index: pgrx::pg_sys::Relation, +) -> (VectorOptions, VchordrqfscanIndexingOptions) { + let att = unsafe { &mut *(*index).rd_att }; + let atts = unsafe { att.attrs.as_slice(att.natts as _) }; + if atts.is_empty() { + pgrx::error!("indexing on no columns is not supported"); + } + if atts.len() != 1 { + pgrx::error!("multicolumn index is not supported"); + } + // get dims + let typmod = Typmod::parse_from_i32(atts[0].type_mod()).unwrap(); + let dims = if let Some(dims) = typmod.dims() { + dims.get() + } else { + pgrx::error!( + "Dimensions type modifier of a vector column is needed for building the index." + ); + }; + // get v, d + let opfamily = unsafe { opfamily(index) }; + let vector = VectorOptions { + dims, + v: opfamily.vector, + d: opfamily.distance_kind(), + }; + // get indexing, segment, optimizing + let rabitq = unsafe { convert_reloptions_to_options((*index).rd_options) }; + (vector, rabitq) +} + +#[derive(Debug, Clone, Copy)] +pub struct Opfamily { + vector: VectorKind, + pg_distance: PgDistanceKind, +} + +impl Opfamily { + pub unsafe fn datum_to_vector( + self, + datum: pgrx::pg_sys::Datum, + is_null: bool, + ) -> Option { + if is_null || datum.is_null() { + return None; + } + let vector = match self.vector { + VectorKind::Vecf32 => { + let vector = unsafe { PgvectorVectorInput::from_datum(datum, false).unwrap() }; + self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed())) + } + _ => unreachable!(), + }; + Some(vector) + } + pub unsafe fn datum_to_sphere( + self, + datum: pgrx::pg_sys::Datum, + is_null: bool, + ) -> (Option, Option) { + if is_null || datum.is_null() { + return (None, None); + } + let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; + let center = match self.vector { + VectorKind::Vecf32 => tuple + .get_by_index::(NonZero::new(1).unwrap()) + .unwrap() + .map(|vector| self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed()))), + _ => unreachable!(), + }; + let radius = tuple.get_by_index::(NonZero::new(2).unwrap()).unwrap(); + (center, radius) + } + pub fn preprocess(self, vector: BorrowedVector<'_>) -> OwnedVector { + use BorrowedVector as B; + use OwnedVector as O; + match (vector, self.pg_distance) { + (B::Vecf32(x), PgDistanceKind::L2) => O::Vecf32(x.own()), + (B::Vecf32(x), PgDistanceKind::Dot) => O::Vecf32(x.own()), + (B::Vecf32(x), PgDistanceKind::Cos) => O::Vecf32(x.function_normalize()), + (B::Vecf16(x), _) => O::Vecf16(x.own()), + (B::SVecf32(x), _) => O::SVecf32(x.own()), + (B::BVector(x), _) => O::BVector(x.own()), + } + } + pub fn process(self, x: Distance) -> f32 { + match self.pg_distance { + PgDistanceKind::Cos => f32::from(x) + 1.0f32, + PgDistanceKind::L2 => f32::from(x).sqrt(), + _ => f32::from(x), + } + } + pub fn distance_kind(self) -> DistanceKind { + self.pg_distance.to_distance() + } +} + +pub unsafe fn opfamily(index: pgrx::pg_sys::Relation) -> Opfamily { + use pgrx::pg_sys::Oid; + + let proc = unsafe { pgrx::pg_sys::index_getprocid(index, 1, 1) }; + + if proc == Oid::INVALID { + pgrx::error!("support function 1 is not found"); + } + + let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); + unsafe { + pgrx::pg_sys::fmgr_info(proc, &mut flinfo); + } + + let fn_addr = flinfo.fn_addr.expect("null function pointer"); + + let mut fcinfo = unsafe { std::mem::zeroed::() }; + fcinfo.flinfo = &mut flinfo; + fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; + fcinfo.context = std::ptr::null_mut(); + fcinfo.resultinfo = std::ptr::null_mut(); + fcinfo.isnull = true; + fcinfo.nargs = 0; + + let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; + + let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; + + let result_string = result_option.expect("null string"); + + let (vector, pg_distance) = convert_name_to_vd(&result_string).unwrap(); + + unsafe { + pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); + } + + Opfamily { + vector, + pg_distance, + } +} diff --git a/src/vchordrqfscan/index/am_scan.rs b/src/vchordrqfscan/index/am_scan.rs new file mode 100644 index 00000000..7396bd1b --- /dev/null +++ b/src/vchordrqfscan/index/am_scan.rs @@ -0,0 +1,132 @@ +use super::am_options::Opfamily; +use crate::postgres::Relation; +use crate::vchordrqfscan::algorithm::scan::scan; +use crate::vchordrqfscan::gucs::executing::epsilon; +use crate::vchordrqfscan::gucs::executing::max_scan_tuples; +use crate::vchordrqfscan::gucs::executing::probes; +use base::distance::Distance; +use base::search::*; +use base::vector::*; + +pub enum Scanner { + Initial { + vector: Option<(OwnedVector, Opfamily)>, + threshold: Option, + recheck: bool, + }, + Vbase { + vbase: Box>, + threshold: Option, + recheck: bool, + opfamily: Opfamily, + }, + Empty {}, +} + +pub fn scan_build( + orderbys: Vec>, + spheres: Vec<(Option, Option)>, + opfamily: Opfamily, +) -> (Option<(OwnedVector, Opfamily)>, Option, bool) { + let mut pair = None; + let mut threshold = None; + let mut recheck = false; + for orderby_vector in orderbys { + if pair.is_none() { + pair = orderby_vector; + } else if orderby_vector.is_some() && pair != orderby_vector { + pgrx::error!("vector search with multiple vectors is not supported"); + } + } + for (sphere_vector, sphere_threshold) in spheres { + if pair.is_none() { + pair = sphere_vector; + threshold = sphere_threshold; + } else if pair == sphere_vector { + if threshold.is_none() || sphere_threshold < threshold { + threshold = sphere_threshold; + } + } else { + recheck = true; + break; + } + } + (pair.map(|x| (x, opfamily)), threshold, recheck) +} + +pub fn scan_make( + vector: Option<(OwnedVector, Opfamily)>, + threshold: Option, + recheck: bool, +) -> Scanner { + Scanner::Initial { + vector, + threshold, + recheck, + } +} + +pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, bool)> { + if let Scanner::Initial { + vector, + threshold, + recheck, + } = scanner + { + if let Some((vector, opfamily)) = vector.as_ref() { + let vbase = scan( + relation, + match vector { + OwnedVector::Vecf32(x) => x.slice().to_vec(), + OwnedVector::Vecf16(_) => unreachable!(), + OwnedVector::SVecf32(_) => unreachable!(), + OwnedVector::BVector(_) => unreachable!(), + }, + opfamily.distance_kind(), + probes(), + epsilon(), + ); + *scanner = Scanner::Vbase { + vbase: if let Some(max_scan_tuples) = max_scan_tuples() { + Box::new(vbase.take(max_scan_tuples as usize)) + } else { + Box::new(vbase) + }, + threshold: *threshold, + recheck: *recheck, + opfamily: *opfamily, + }; + } else { + *scanner = Scanner::Empty {}; + } + } + match scanner { + Scanner::Initial { .. } => unreachable!(), + Scanner::Vbase { + vbase, + threshold, + recheck, + opfamily, + } => match ( + vbase.next().map(|(d, p)| (opfamily.process(d), p)), + threshold, + ) { + (Some((_, ptr)), None) => Some((ptr, *recheck)), + (Some((distance, ptr)), Some(t)) if distance < *t => Some((ptr, *recheck)), + _ => { + let scanner = std::mem::replace(scanner, Scanner::Empty {}); + scan_release(scanner); + None + } + }, + Scanner::Empty {} => None, + } +} + +pub fn scan_release(scanner: Scanner) { + match scanner { + Scanner::Initial { .. } => {} + Scanner::Vbase { .. } => {} + Scanner::Empty {} => {} + } +} diff --git a/src/vchordrqfscan/index/functions.rs b/src/vchordrqfscan/index/functions.rs new file mode 100644 index 00000000..98f0f251 --- /dev/null +++ b/src/vchordrqfscan/index/functions.rs @@ -0,0 +1,26 @@ +use crate::postgres::Relation; +use crate::vchordrqfscan::algorithm::prewarm::prewarm; +use pgrx::pg_sys::Oid; +use pgrx_catalog::{PgAm, PgClass}; + +#[pgrx::pg_extern(sql = "")] +fn _vchordrqfscan_prewarm(indexrelid: Oid, height: i32) -> String { + let pg_am = PgAm::search_amname(c"vchordrqfscan").unwrap(); + let Some(pg_am) = pg_am.get() else { + pgrx::error!("vchord is not installed"); + }; + let pg_class = PgClass::search_reloid(indexrelid).unwrap(); + let Some(pg_class) = pg_class.get() else { + pgrx::error!("there is no such index"); + }; + if pg_class.relam() != pg_am.oid() { + pgrx::error!("{:?} is not a vchordrqfscan index", pg_class.relname()); + } + let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; + let relation = unsafe { Relation::new(index) }; + let message = prewarm(relation, height); + unsafe { + pgrx::pg_sys::index_close(index, pgrx::pg_sys::ShareLock as _); + } + message +} diff --git a/src/vchordrqfscan/index/mod.rs b/src/vchordrqfscan/index/mod.rs new file mode 100644 index 00000000..5203e4fb --- /dev/null +++ b/src/vchordrqfscan/index/mod.rs @@ -0,0 +1,12 @@ +pub mod am; +pub mod am_options; +pub mod am_scan; +pub mod functions; +pub mod opclass; +pub mod utils; + +pub unsafe fn init() { + unsafe { + am::init(); + } +} diff --git a/src/vchordrqfscan/index/opclass.rs b/src/vchordrqfscan/index/opclass.rs new file mode 100644 index 00000000..d095b9a6 --- /dev/null +++ b/src/vchordrqfscan/index/opclass.rs @@ -0,0 +1,14 @@ +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrqfscan_support_vector_l2_ops() -> String { + "vector_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrqfscan_support_vector_ip_ops() -> String { + "vector_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrqfscan_support_vector_cosine_ops() -> String { + "vector_cosine_ops".to_string() +} diff --git a/src/vchordrqfscan/index/utils.rs b/src/vchordrqfscan/index/utils.rs new file mode 100644 index 00000000..a5d85a3f --- /dev/null +++ b/src/vchordrqfscan/index/utils.rs @@ -0,0 +1,20 @@ +use base::search::*; + +pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { + let value = pointer.as_u64(); + pgrx::pg_sys::ItemPointerData { + ip_blkid: pgrx::pg_sys::BlockIdData { + bi_hi: ((value >> 32) & 0xffff) as u16, + bi_lo: ((value >> 16) & 0xffff) as u16, + }, + ip_posid: (value & 0xffff) as u16, + } +} + +pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer { + let mut value = 0; + value |= (ctid.ip_blkid.bi_hi as u64) << 32; + value |= (ctid.ip_blkid.bi_lo as u64) << 16; + value |= ctid.ip_posid as u64; + Pointer::new(value) +} diff --git a/src/vchordrqfscan/mod.rs b/src/vchordrqfscan/mod.rs new file mode 100644 index 00000000..c2ae9456 --- /dev/null +++ b/src/vchordrqfscan/mod.rs @@ -0,0 +1,11 @@ +mod algorithm; +mod gucs; +mod index; +mod types; + +pub unsafe fn init() { + unsafe { + index::init(); + gucs::init(); + } +} diff --git a/src/vchordrqfscan/types.rs b/src/vchordrqfscan/types.rs new file mode 100644 index 00000000..b3a1067f --- /dev/null +++ b/src/vchordrqfscan/types.rs @@ -0,0 +1,90 @@ +use serde::{Deserialize, Serialize}; +use validator::{Validate, ValidationError, ValidationErrors}; + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqfscanInternalBuildOptions { + #[serde(default = "VchordrqfscanInternalBuildOptions::default_lists")] + #[validate(length(min = 1, max = 8), custom(function = VchordrqfscanInternalBuildOptions::validate_lists))] + pub lists: Vec, + #[serde(default = "VchordrqfscanInternalBuildOptions::default_spherical_centroids")] + pub spherical_centroids: bool, + #[serde(default = "VchordrqfscanInternalBuildOptions::default_build_threads")] + #[validate(range(min = 1, max = 255))] + pub build_threads: u16, +} + +impl VchordrqfscanInternalBuildOptions { + fn default_lists() -> Vec { + vec![1000] + } + fn validate_lists(lists: &[u32]) -> Result<(), ValidationError> { + if !lists.is_sorted() { + return Err(ValidationError::new("`lists` should be in ascending order")); + } + if !lists.iter().all(|x| (1..=1 << 24).contains(x)) { + return Err(ValidationError::new("list is too long or too short")); + } + Ok(()) + } + fn default_spherical_centroids() -> bool { + false + } + fn default_build_threads() -> u16 { + 1 + } +} + +impl Default for VchordrqfscanInternalBuildOptions { + fn default() -> Self { + Self { + lists: Self::default_lists(), + spherical_centroids: Self::default_spherical_centroids(), + build_threads: Self::default_build_threads(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqfscanExternalBuildOptions { + pub table: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +#[serde(rename_all = "snake_case")] +pub enum VchordrqfscanBuildOptions { + Internal(VchordrqfscanInternalBuildOptions), + External(VchordrqfscanExternalBuildOptions), +} + +impl Default for VchordrqfscanBuildOptions { + fn default() -> Self { + Self::Internal(Default::default()) + } +} + +impl Validate for VchordrqfscanBuildOptions { + fn validate(&self) -> Result<(), ValidationErrors> { + use VchordrqfscanBuildOptions::*; + match self { + Internal(internal_build) => internal_build.validate(), + External(external_build) => external_build.validate(), + } + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqfscanIndexingOptions { + #[serde(default = "VchordrqfscanIndexingOptions::default_residual_quantization")] + pub residual_quantization: bool, + pub build: VchordrqfscanBuildOptions, +} + +impl VchordrqfscanIndexingOptions { + fn default_residual_quantization() -> bool { + false + } +} From b95a18c25348eb2d9a5281e4d9475001d70d570d Mon Sep 17 00:00:00 2001 From: Keming Date: Thu, 21 Nov 2024 20:11:52 +0800 Subject: [PATCH 067/324] chore: add CI to build the pgrx image (#107) Signed-off-by: Keming --- .github/workflows/pgrx.yml | 43 ++++++++++++++++++++++++++++++++++++++ rust-toolchain.toml | 2 +- 2 files changed, 44 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/pgrx.yml diff --git a/.github/workflows/pgrx.yml b/.github/workflows/pgrx.yml new file mode 100644 index 00000000..7c90a345 --- /dev/null +++ b/.github/workflows/pgrx.yml @@ -0,0 +1,43 @@ +name: Build pgrx Image + +on: + workflow_dispatch: + inputs: + version: + description: 'pgrx version' + required: true + type: string + +concurrency: + group: ${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +env: + IMAGE_NAME: "ghcr.io/${{ github.repository }}-pgrx" + +permissions: + contents: write + packages: write + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + - name: Login to ghcr.io + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + - name: Build and push + uses: docker/build-push-action@v6 + with: + context: . + file: ./docker/pgrx.Dockerfile + push: true + build-args: | + PGRX_VERSION=${{ github.event.inputs.version }} + tags: ${{ env.IMAGE_NAME }}:${{ github.event.inputs.version }} diff --git a/rust-toolchain.toml b/rust-toolchain.toml index edfb9578..0d943e1a 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2024-09-14" +channel = "nightly-2024-11-19" From ba445e47569c05dcbc04a43e6bf1a5c30fc0fc76 Mon Sep 17 00:00:00 2001 From: Keming Date: Thu, 21 Nov 2024 20:38:38 +0800 Subject: [PATCH 068/324] chore: fix the pgrx ci (#108) Signed-off-by: Keming --- .github/workflows/pgrx.yml | 2 +- .github/workflows/psql.yml | 2 +- .github/workflows/release.yml | 2 +- .github/workflows/rust.yml | 2 +- src/datatype/memory_pgvector_vector.rs | 4 ++-- src/vchordrq/algorithm/build.rs | 2 +- src/vchordrqfscan/algorithm/build.rs | 2 +- 7 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/pgrx.yml b/.github/workflows/pgrx.yml index 7c90a345..bf96d954 100644 --- a/.github/workflows/pgrx.yml +++ b/.github/workflows/pgrx.yml @@ -13,7 +13,7 @@ concurrency: cancel-in-progress: true env: - IMAGE_NAME: "ghcr.io/${{ github.repository }}-pgrx" + IMAGE_NAME: "ghcr.io/tensorchord/vectorchord-pgrx" permissions: contents: write diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml index c5e9dd5e..55587859 100644 --- a/.github/workflows/psql.yml +++ b/.github/workflows/psql.yml @@ -44,7 +44,7 @@ jobs: version: ["14", "15", "16", "17"] arch: ["x86_64"] env: - PGRX_IMAGE: "kemingy/pgrx:0.12.8" + PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.8" SQLLOGICTEST: "0.22.0" steps: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 24136a86..be1f8e2c 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -40,7 +40,7 @@ jobs: version: ["14", "15", "16", "17"] arch: ["x86_64", "aarch64"] env: - PGRX_IMAGE: "kemingy/pgrx:0.12.8" + PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.8" SEMVER: ${{ needs.semver.outputs.SEMVER }} steps: diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 018139ab..7f6ffe04 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -33,7 +33,7 @@ jobs: matrix: arch: ["x86_64", "aarch64"] env: - PGRX_IMAGE: "kemingy/pgrx:0.12.8" + PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.8" steps: - uses: actions/checkout@v4 diff --git a/src/datatype/memory_pgvector_vector.rs b/src/datatype/memory_pgvector_vector.rs index fe0932bb..7166ace0 100644 --- a/src/datatype/memory_pgvector_vector.rs +++ b/src/datatype/memory_pgvector_vector.rs @@ -58,7 +58,7 @@ pub enum PgvectorVectorInput<'a> { Borrowed(&'a PgvectorVectorHeader), } -impl<'a> PgvectorVectorInput<'a> { +impl PgvectorVectorInput<'_> { unsafe fn new(p: NonNull) -> Self { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() @@ -124,7 +124,7 @@ impl Drop for PgvectorVectorOutput { } } -impl<'a> FromDatum for PgvectorVectorInput<'a> { +impl FromDatum for PgvectorVectorInput<'_> { unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { if is_null { None diff --git a/src/vchordrq/algorithm/build.rs b/src/vchordrq/algorithm/build.rs index 3b96f0c4..3cb83e7e 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -351,7 +351,7 @@ impl<'a, T> Tape<'a, T> { } } -impl<'a, T> Tape<'a, T> +impl Tape<'_, T> where T: rkyv::Serialize>, { diff --git a/src/vchordrqfscan/algorithm/build.rs b/src/vchordrqfscan/algorithm/build.rs index 3215e017..49820c9e 100644 --- a/src/vchordrqfscan/algorithm/build.rs +++ b/src/vchordrqfscan/algorithm/build.rs @@ -390,7 +390,7 @@ impl<'a, T> Tape<'a, T> { } } -impl<'a, T> Tape<'a, T> +impl Tape<'_, T> where T: rkyv::Serialize>, { From 0d8e3f3f6f57da526265989875592df3d7d95796 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 21 Nov 2024 20:52:10 +0800 Subject: [PATCH 069/324] fix: optimize insertions in building when lists = 1 (#106) Signed-off-by: usamoi Co-authored-by: Keming --- src/postgres.rs | 22 ++--- src/vchordrq/algorithm/build.rs | 8 +- src/vchordrq/algorithm/insert.rs | 126 +++++++++++++------------- src/vchordrq/algorithm/tuples.rs | 1 - src/vchordrq/algorithm/vacuum.rs | 4 +- src/vchordrq/index/am.rs | 3 + src/vchordrqfscan/algorithm/build.rs | 2 +- src/vchordrqfscan/algorithm/insert.rs | 16 +++- src/vchordrqfscan/algorithm/vacuum.rs | 4 +- 9 files changed, 94 insertions(+), 92 deletions(-) diff --git a/src/postgres.rs b/src/postgres.rs index 6dcd2430..06b402d9 100644 --- a/src/postgres.rs +++ b/src/postgres.rs @@ -24,7 +24,7 @@ const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usi const _: () = assert!(size_of::() == pgrx::pg_sys::BLCKSZ as usize); impl Page { - pub fn init_mut(this: &mut MaybeUninit, tracking_freespace: bool) -> &mut Self { + pub fn init_mut(this: &mut MaybeUninit) -> &mut Self { unsafe { pgrx::pg_sys::PageInit( this.as_mut_ptr() as pgrx::pg_sys::Page, @@ -33,8 +33,7 @@ impl Page { ); (&raw mut (*this.as_mut_ptr()).opaque).write(Opaque { next: u32::MAX, - tracking_freespace, - fast_forward: u32::MAX, + skip: u32::MAX, }); } let this = unsafe { MaybeUninit::assume_init_mut(this) }; @@ -59,7 +58,7 @@ impl Page { assert!(self.header.pd_upper as usize <= size_of::()); let lower = self.header.pd_lower as usize; let upper = self.header.pd_upper as usize; - assert!(lower < upper); + assert!(lower <= upper); ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16 } pub fn get(&self, i: u16) -> Option<&[u8]> { @@ -159,8 +158,7 @@ impl Page { #[repr(C, align(8))] pub struct Opaque { pub next: u32, - pub tracking_freespace: bool, - pub fast_forward: u32, + pub skip: u32, } const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); @@ -195,6 +193,7 @@ pub struct BufferWriteGuard { page: NonNull, state: *mut pgrx::pg_sys::GenericXLogState, id: u32, + tracking_freespace: bool, } impl BufferWriteGuard { @@ -215,7 +214,7 @@ impl Drop for BufferWriteGuard { if std::thread::panicking() { pgrx::pg_sys::GenericXLogAbort(self.state); } else { - if self.get().get_opaque().tracking_freespace { + if self.tracking_freespace { pgrx::pg_sys::RecordPageWithFreeSpace( self.raw, self.id, @@ -257,7 +256,7 @@ impl Relation { BufferReadGuard { buf, page, id } } } - pub fn write(&self, id: u32) -> BufferWriteGuard { + pub fn write(&self, id: u32, tracking_freespace: bool) -> BufferWriteGuard { assert!(id != u32::MAX, "no such page"); unsafe { use pgrx::pg_sys::{ @@ -284,6 +283,7 @@ impl Relation { page: page.cast(), state, id, + tracking_freespace, } } } @@ -310,13 +310,14 @@ impl Relation { .cast::>(), ) .expect("failed to get page"); - Page::init_mut(page.as_mut(), tracking_freespace); + Page::init_mut(page.as_mut()); BufferWriteGuard { raw: self.raw, buf, page: page.cast(), state, id: pgrx::pg_sys::BufferGetBlockNumber(buf), + tracking_freespace, } } } @@ -327,8 +328,7 @@ impl Relation { if id == u32::MAX { return None; } - let write = self.write(id); - assert!(write.get().get_opaque().tracking_freespace); + let write = self.write(id, true); if write.get().freespace() < freespace as _ { // the free space is recorded incorrectly pgrx::pg_sys::RecordPageWithFreeSpace( diff --git a/src/vchordrq/algorithm/build.rs b/src/vchordrq/algorithm/build.rs index 3cb83e7e..8df19052 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -71,10 +71,7 @@ pub fn build( }; let mut meta = Tape::create(&relation, false); assert_eq!(meta.first(), 0); - let mut forwards = Tape::::create(&relation, false); - assert_eq!(forwards.first(), 1); let mut vectors = Tape::create(&relation, true); - assert_eq!(vectors.first(), 2); let mut pointer_of_means = Vec::>::new(); for i in 0..structures.len() { let mut level = Vec::new(); @@ -125,13 +122,11 @@ pub fn build( } pointer_of_firsts.push(level); } - forwards.head.get_mut().get_opaque_mut().fast_forward = vectors.first(); meta.push(&MetaTuple { dims, height_of_root: structures.len() as u32, is_residual, vectors_first: vectors.first(), - forwards_first: forwards.first(), mean: pointer_of_means.last().unwrap()[0], first: pointer_of_firsts.last().unwrap()[0], }); @@ -336,7 +331,8 @@ struct Tape<'a, T> { impl<'a, T> Tape<'a, T> { fn create(relation: &'a Relation, tracking_freespace: bool) -> Self { - let head = relation.extend(tracking_freespace); + let mut head = relation.extend(tracking_freespace); + head.get_mut().get_opaque_mut().skip = head.id(); let first = head.id(); Self { relation, diff --git a/src/vchordrq/algorithm/insert.rs b/src/vchordrq/algorithm/insert.rs index e9318c3a..c37f7c38 100644 --- a/src/vchordrq/algorithm/insert.rs +++ b/src/vchordrq/algorithm/insert.rs @@ -11,7 +11,13 @@ use base::search::Pointer; use std::cmp::Reverse; use std::collections::BinaryHeap; -pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_kind: DistanceKind) { +pub fn insert( + relation: Relation, + payload: Pointer, + vector: Vec, + distance_kind: DistanceKind, + in_building: bool, +) { let meta_guard = relation.read(0); let meta_tuple = meta_guard .get() @@ -38,50 +44,14 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k chain, }) .unwrap(); - chain = Some('chain: { - if let Some(mut write) = relation.search(tuple.len()) { - let i = write.get_mut().alloc(&tuple).unwrap(); - break 'chain (write.id(), i); - } - let mut current = relation.read(1).get().get_opaque().fast_forward; - let mut changed = false; - loop { - let read = relation.read(current); - let flag = 'flag: { - if read.get().freespace() as usize >= tuple.len() { - break 'flag true; - } - if read.get().get_opaque().next == u32::MAX { - break 'flag true; - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current); - if let Some(i) = write.get_mut().alloc(&tuple) { - break 'chain (current, i); - } - if write.get().get_opaque().next == u32::MAX { - if changed { - relation.write(1).get_mut().get_opaque_mut().fast_forward = - write.id(); - } - let mut extend = relation.extend(true); - write.get_mut().get_opaque_mut().next = extend.id(); - if let Some(i) = extend.get_mut().alloc(&tuple) { - break 'chain (extend.id(), i); - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - } - current = write.get().get_opaque().next; - } else { - current = read.get().get_opaque().next; - } - changed = true; - } - }); + chain = Some(append( + relation.clone(), + meta_tuple.vectors_first, + &tuple, + true, + true, + true, + )); } chain.unwrap() }; @@ -168,7 +138,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k } else { rabitq::code(dims, &vector) }; - let h0_tuple = rkyv::to_bytes::<_, 8192>(&Height0Tuple { + let tuple = rkyv::to_bytes::<_, 8192>(&Height0Tuple { mean: h0_vector, payload: h0_payload, dis_u_2: code.dis_u_2, @@ -178,38 +148,64 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k t: code.t(), }) .unwrap(); - let first = list.0; + append(relation, list.0, &tuple, false, in_building, in_building); +} + +fn append( + relation: Relation, + first: u32, + tuple: &[u8], + tracking_freespace: bool, + skipping_traversal: bool, + updating_skip: bool, +) -> (u32, u16) { + if tracking_freespace { + if let Some(mut write) = relation.search(tuple.len()) { + let i = write.get_mut().alloc(tuple).unwrap(); + return (write.id(), i); + } + } assert!(first != u32::MAX); let mut current = first; loop { let read = relation.read(current); - let flag = 'flag: { - if read.get().freespace() as usize >= h0_tuple.len() { - break 'flag true; - } - if read.get().get_opaque().next == u32::MAX { - break 'flag true; - } - false - }; - if flag { + if read.get().freespace() as usize >= tuple.len() + || read.get().get_opaque().next == u32::MAX + { drop(read); - let mut write = relation.write(current); - if write.get_mut().alloc(&h0_tuple).is_some() { - return; + let mut write = relation.write(current, tracking_freespace); + if let Some(i) = write.get_mut().alloc(tuple) { + return (current, i); } if write.get().get_opaque().next == u32::MAX { - let mut extend = relation.extend(false); + let mut extend = relation.extend(tracking_freespace); write.get_mut().get_opaque_mut().next = extend.id(); - if extend.get_mut().alloc(&h0_tuple).is_some() { - return; + drop(write); + if let Some(i) = extend.get_mut().alloc(tuple) { + let result = (extend.id(), i); + drop(extend); + if updating_skip { + let mut past = relation.write(first, tracking_freespace); + let skip = &mut past.get_mut().get_opaque_mut().skip; + assert!(*skip != u32::MAX); + *skip = std::cmp::max(*skip, result.0); + } + return result; } else { panic!("a tuple cannot even be fit in a fresh page"); } } - current = write.get().get_opaque().next; + if skipping_traversal && current == first && write.get().get_opaque().skip != first { + current = write.get().get_opaque().skip; + } else { + current = write.get().get_opaque().next; + } } else { - current = read.get().get_opaque().next; + if skipping_traversal && current == first && read.get().get_opaque().skip != first { + current = read.get().get_opaque().skip; + } else { + current = read.get().get_opaque().next; + } } } } diff --git a/src/vchordrq/algorithm/tuples.rs b/src/vchordrq/algorithm/tuples.rs index afdc1cda..cf6236f0 100644 --- a/src/vchordrq/algorithm/tuples.rs +++ b/src/vchordrq/algorithm/tuples.rs @@ -7,7 +7,6 @@ pub struct MetaTuple { pub height_of_root: u32, pub is_residual: bool, pub vectors_first: u32, - pub forwards_first: u32, // raw vector pub mean: (u32, u16), // for meta tuple, it's pointers to next level diff --git a/src/vchordrq/algorithm/vacuum.rs b/src/vchordrq/algorithm/vacuum.rs index aec95a7d..27377020 100644 --- a/src/vchordrq/algorithm/vacuum.rs +++ b/src/vchordrq/algorithm/vacuum.rs @@ -40,7 +40,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - let mut current = first; while current != u32::MAX { delay(); - let mut h0_guard = relation.write(current); + let mut h0_guard = relation.write(current, false); let mut reconstruct_removes = Vec::new(); for i in 1..=h0_guard.get().len() { let h0_tuple = h0_guard @@ -90,7 +90,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - }; if flag { drop(read); - let mut write = relation.write(current); + let mut write = relation.write(current, true); for i in 1..=write.get().len() { let Some(vector_tuple) = write.get().get(i) else { continue; diff --git a/src/vchordrq/index/am.rs b/src/vchordrq/index/am.rs index b702abe2..d3c3e8b2 100644 --- a/src/vchordrq/index/am.rs +++ b/src/vchordrq/index/am.rs @@ -289,6 +289,7 @@ pub unsafe extern "C" fn ambuild( payload, vector, opfamily.distance_kind(), + true, ); indtuples += 1; reporter.tuples_done(indtuples); @@ -617,6 +618,7 @@ unsafe fn parallel_build( payload, vector, opfamily.distance_kind(), + true, ); unsafe { let indtuples; @@ -705,6 +707,7 @@ pub unsafe extern "C" fn aminsert( pointer, vector.into_vec(), opfamily.distance_kind(), + false, ); } false diff --git a/src/vchordrqfscan/algorithm/build.rs b/src/vchordrqfscan/algorithm/build.rs index 49820c9e..9bdc3b92 100644 --- a/src/vchordrqfscan/algorithm/build.rs +++ b/src/vchordrqfscan/algorithm/build.rs @@ -164,7 +164,7 @@ pub fn build( } pointer_of_firsts.push(level); } - forwards.head.get_mut().get_opaque_mut().fast_forward = vectors.first(); + forwards.head.get_mut().get_opaque_mut().skip = vectors.first(); meta.push(&MetaTuple { dims, height_of_root: structures.len() as u32, diff --git a/src/vchordrqfscan/algorithm/insert.rs b/src/vchordrqfscan/algorithm/insert.rs index d17c22ea..e89b1101 100644 --- a/src/vchordrqfscan/algorithm/insert.rs +++ b/src/vchordrqfscan/algorithm/insert.rs @@ -38,7 +38,11 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k let i = write.get_mut().alloc(&tuple).unwrap(); break 'h0_vector (write.id(), i); } - let mut current = relation.read(1).get().get_opaque().fast_forward; + let mut current = relation + .read(meta_tuple.forwards_first) + .get() + .get_opaque() + .skip; let mut changed = false; loop { let read = relation.read(current); @@ -53,13 +57,17 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k }; if flag { drop(read); - let mut write = relation.write(current); + let mut write = relation.write(current, true); if let Some(i) = write.get_mut().alloc(&tuple) { break (current, i); } if write.get().get_opaque().next == u32::MAX { if changed { - relation.write(1).get_mut().get_opaque_mut().fast_forward = write.id(); + relation + .write(meta_tuple.forwards_first, false) + .get_mut() + .get_opaque_mut() + .skip = write.id(); } let mut extend = relation.extend(true); write.get_mut().get_opaque_mut().next = extend.id(); @@ -209,7 +217,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k }; if flag { drop(read); - let mut write = relation.write(current); + let mut write = relation.write(current, false); for i in 1..=write.get().len() { let flag = put( write.get_mut().get_mut(i).expect("data corruption"), diff --git a/src/vchordrqfscan/algorithm/vacuum.rs b/src/vchordrqfscan/algorithm/vacuum.rs index 1a77b980..7bb51796 100644 --- a/src/vchordrqfscan/algorithm/vacuum.rs +++ b/src/vchordrqfscan/algorithm/vacuum.rs @@ -45,7 +45,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - let mut current = first; while current != u32::MAX { delay(); - let mut h0_guard = relation.write(current); + let mut h0_guard = relation.write(current, false); for i in 1..=h0_guard.get().len() { let h0_tuple = h0_guard .get() @@ -119,7 +119,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - }; if flag { drop(read); - let mut write = relation.write(current); + let mut write = relation.write(current, true); for i in 1..=write.get().len() { let Some(vector_tuple) = write.get().get(i) else { continue; From 45a95325d7788ee94934aca137a7a3243d33400b Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 22 Nov 2024 17:25:37 +0800 Subject: [PATCH 070/324] chore: build multiarch docker images (#112) - close #102 Tested https://github.com/tensorchord/VectorChord/actions/runs/11968955634 Images can be found in https://hub.docker.com/r/tensorchord/vchord-postgres/tags --------- Signed-off-by: Keming --- .github/workflows/release.yml | 18 ++++++++++-------- README.md | 2 +- docker/Dockerfile | 4 ++-- docker/binary.Dockerfile | 6 +++++- 4 files changed, 18 insertions(+), 12 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index be1f8e2c..9c2e04d7 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -82,7 +82,6 @@ jobs: strategy: matrix: version: ["14", "15", "16", "17"] - platform: ["amd64", "arm64"] env: SEMVER: ${{ needs.semver.outputs.SEMVER }} steps: @@ -92,9 +91,10 @@ jobs: env: GH_TOKEN: ${{ github.token }} run: | - gh release download $SEMVER --pattern "vchord-pg${{ matrix.version }}_${SEMVER}_${{ matrix.platform }}.deb" --output vchord-binary-release.deb mkdir -p build - cp vchord-binary-release.deb ./build/vchord-pg${{ matrix.version }}_${SEMVER}_${{ matrix.platform }}.deb + for arch in amd64 arm64; do + gh release download $SEMVER --pattern "vchord-pg${{ matrix.version }}_${SEMVER}_${arch}.deb" --output ./build/vchord-pg${{ matrix.version }}_${SEMVER}_${arch}.deb + done - name: Set up QEMU uses: docker/setup-qemu-action@v3 - name: Set up Docker Buildx @@ -109,19 +109,21 @@ jobs: with: context: . push: true - platforms: "linux/${{ matrix.platform }}" + platforms: "linux/amd64,linux/arm64" file: ./docker/binary.Dockerfile - tags: tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ matrix.platform }} + tags: tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }} + build-args: | + PG_VERSION=${{ matrix.version }} + SEMVER=${{ env.SEMVER }} - name: Push PostgreSQL release to Docker Registry uses: docker/build-push-action@v6 with: context: . push: true - platforms: "linux/${{ matrix.platform }}" + platforms: "linux/amd64,linux/arm64" file: ./docker/Dockerfile - tags: tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ matrix.platform }} + tags: tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} build-args: | PG_VERSION=${{ matrix.version }} SEMVER=${{ env.SEMVER }} PGVECTOR=0.8.0 - PLATFORM=${{ matrix.platform }} diff --git a/README.md b/README.md index 5bac8114..c8957cd0 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d tensorchord/vchord-postgres:pg17-v0.1.0-amd64 + -d tensorchord/vchord-postgres:pg17-v0.1.0 Then you can connect to the database using the `psql` command line tool. The default username is `postgres`, and the default password is `mysecretpassword`. ```bash diff --git a/docker/Dockerfile b/docker/Dockerfile index 7fec75b0..ac5b1aad 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -5,10 +5,10 @@ FROM pgvector/pgvector:${PGVECTOR}-pg${PG_VERSION} ARG PG_VERSION ARG SEMVER=0.0.0 -ARG PLATFORM=amd64 +ARG TARGETARCH RUN echo ${PG_VERSION} -COPY ./build/vchord-pg${PG_VERSION}_${SEMVER}_${PLATFORM}.deb /tmp/vchord.deb +COPY ./build/vchord-pg${PG_VERSION}_${SEMVER}_${TARGETARCH}.deb /tmp/vchord.deb RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb CMD ["postgres", "-c" ,"shared_preload_libraries=vchord.so"] diff --git a/docker/binary.Dockerfile b/docker/binary.Dockerfile index 71e54ef6..001bcc99 100644 --- a/docker/binary.Dockerfile +++ b/docker/binary.Dockerfile @@ -1,4 +1,8 @@ FROM scratch +ARG SEMVER +ARG PG_VERSION +ARG TARGETARCH + WORKDIR /workspace -COPY ./vchord-binary-release.deb /workspace/ +COPY ./build/vchord-pg${PG_VERSION}_${SEMVER}_${TARGETARCH}.deb /workspace/ From 3b80d621d92ff20bd0aead8b6359b1f5a95daf03 Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 22 Nov 2024 19:29:07 +0800 Subject: [PATCH 071/324] docs(readme): fix markdown style for docker run (#113) Signed-off-by: Keming --- README.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index c8957cd0..ff5e70fc 100644 --- a/README.md +++ b/README.md @@ -33,8 +33,10 @@ docker run \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ -d tensorchord/vchord-postgres:pg17-v0.1.0 +``` Then you can connect to the database using the `psql` command line tool. The default username is `postgres`, and the default password is `mysecretpassword`. + ```bash psql -h localhost -p 5432 -U postgres ``` @@ -44,7 +46,7 @@ Run the following SQL to ensure the extension is enabled. CREATE EXTENSION IF NOT EXISTS vchord CASCADE; ``` -And make sure to add vchord.so to the shared_preload_libraries in postgresql.conf. +And make sure to add `vchord.so` to the `shared_preload_libraries` in `postgresql.conf`. ```SQL -- Add vchord and pgvector to shared_preload_libraries -- From a9d19b3786e311021176c60012a063c13ca835bd Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 26 Nov 2024 17:58:55 +0800 Subject: [PATCH 072/324] feat: improve internal build (#115) closes #99 Signed-off-by: usamoi --- src/utils/k_means.rs | 280 +++++++++++++++++++++++---- src/vchordrq/algorithm/build.rs | 7 +- src/vchordrq/types.rs | 7 + src/vchordrqfscan/algorithm/build.rs | 7 +- src/vchordrqfscan/types.rs | 7 + 5 files changed, 261 insertions(+), 47 deletions(-) diff --git a/src/utils/k_means.rs b/src/utils/k_means.rs index 8b4fd9d2..ac58a088 100644 --- a/src/utils/k_means.rs +++ b/src/utils/k_means.rs @@ -1,24 +1,34 @@ +#![allow(clippy::ptr_arg)] + use super::parallelism::{ParallelIterator, Parallelism}; use base::scalar::*; use half::f16; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; -pub fn k_means( - parallelism: &impl Parallelism, +pub fn k_means( + parallelism: &P, c: usize, dims: usize, - samples: &Vec>, + samples: &Vec>, is_spherical: bool, iterations: usize, -) -> Vec> { +) -> Vec> { assert!(c > 0); assert!(dims > 0); let n = samples.len(); if n <= c { quick_centers(c, dims, samples.clone(), is_spherical) } else { - let mut lloyd_k_means = LloydKMeans::new(parallelism, c, dims, samples, is_spherical); + let compute = |parallelism: &P, centroids: &Vec>| { + if n >= 1000 && c >= 1000 { + rabitq_index(parallelism, dims, n, c, samples, centroids) + } else { + flat_index(parallelism, dims, n, c, samples, centroids) + } + }; + let mut lloyd_k_means = + LloydKMeans::new(parallelism, c, dims, samples, is_spherical, compute); for _ in 0..iterations { parallelism.check(); if lloyd_k_means.iterate() { @@ -29,11 +39,11 @@ pub fn k_means( } } -pub fn k_means_lookup(vector: &[S], centroids: &[Vec]) -> usize { +pub fn k_means_lookup(vector: &[f32], centroids: &[Vec]) -> usize { assert_ne!(centroids.len(), 0); let mut result = (f32::INFINITY, 0); for i in 0..centroids.len() { - let dis = S::reduce_sum_of_d2(vector, ¢roids[i]); + let dis = f32::reduce_sum_of_d2(vector, ¢roids[i]); if dis <= result.0 { result = (dis, i); } @@ -41,52 +51,248 @@ pub fn k_means_lookup(vector: &[S], centroids: &[Vec]) -> usiz result.1 } -fn quick_centers( +fn quick_centers( c: usize, dims: usize, - samples: Vec>, + samples: Vec>, is_spherical: bool, -) -> Vec> { +) -> Vec> { let n = samples.len(); assert!(c >= n); let mut rng = rand::thread_rng(); let mut centroids = samples; for _ in n..c { let r = (0..dims) - .map(|_| S::from_f32(rng.gen_range(-1.0f32..1.0f32))) + .map(|_| f32::from_f32(rng.gen_range(-1.0f32..1.0f32))) .collect(); centroids.push(r); } if is_spherical { for i in 0..c { let centroid = &mut centroids[i]; - let l = S::reduce_sum_of_x2(centroid).sqrt(); - S::vector_mul_scalar_inplace(centroid, 1.0 / l); + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); } } centroids } -struct LloydKMeans<'a, P, S> { +fn rabitq_index( + parallelism: &P, + dims: usize, + n: usize, + c: usize, + samples: &Vec>, + centroids: &Vec>, +) -> Vec { + fn code_alpha(vector: &[f32]) -> (f32, f32, f32, f32) { + let dims = vector.len(); + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x2 = f32::reduce_sum_of_x2(vector); + let dis_u = sum_of_x2.sqrt(); + let x0 = sum_of_abs_x / (sum_of_x2 * (dims as f32)).sqrt(); + let x_x0 = dis_u / x0; + let fac_norm = (dims as f32).sqrt(); + let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); + let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); + let factor_ip = -2.0f32 / fac_norm * x_x0; + let cnt_pos = vector + .iter() + .map(|x| x.scalar_is_sign_positive() as i32) + .sum::(); + let cnt_neg = vector + .iter() + .map(|x| x.scalar_is_sign_negative() as i32) + .sum::(); + let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; + (sum_of_x2, factor_ppc, factor_ip, factor_err) + } + fn code_beta(vector: &[f32]) -> Vec { + let dims = vector.len(); + let mut code = Vec::new(); + for i in 0..dims { + code.push(vector[i].scalar_is_sign_positive() as u8); + } + code + } + let mut a0 = Vec::new(); + let mut a1 = Vec::new(); + let mut a2 = Vec::new(); + let mut a3 = Vec::new(); + let mut a4 = Vec::new(); + for vectors in centroids.chunks(32) { + use quantization::fast_scan::b4::pack; + let code_alphas = std::array::from_fn::<_, 32, _>(|i| { + if let Some(vector) = vectors.get(i) { + code_alpha(vector) + } else { + (0.0, 0.0, 0.0, 0.0) + } + }); + let code_betas = std::array::from_fn::<_, 32, _>(|i| { + let mut result = vec![0_u8; dims.div_ceil(4)]; + if let Some(vector) = vectors.get(i) { + let mut c = code_beta(vector); + c.resize(dims.next_multiple_of(4), 0); + for i in 0..dims.div_ceil(4) { + for j in 0..4 { + result[i] |= c[i * 4 + j] << j; + } + } + } + result + }); + a0.push(code_alphas.map(|x| x.0)); + a1.push(code_alphas.map(|x| x.1)); + a2.push(code_alphas.map(|x| x.2)); + a3.push(code_alphas.map(|x| x.3)); + a4.push(pack(dims.div_ceil(4) as _, code_betas).collect::>()); + } + parallelism + .into_par_iter(0..n) + .map(|i| { + fn gen(mut qvector: Vec) -> Vec { + let dims = qvector.len() as u32; + let t = dims.div_ceil(4); + qvector.resize(qvector.len().next_multiple_of(4), 0); + let mut lut = vec![0u8; t as usize * 16]; + for i in 0..t as usize { + unsafe { + // this hint is used to skip bound checks + std::hint::assert_unchecked(4 * i + 3 < qvector.len()); + std::hint::assert_unchecked(16 * i + 15 < lut.len()); + } + let t0 = qvector[4 * i + 0]; + let t1 = qvector[4 * i + 1]; + let t2 = qvector[4 * i + 2]; + let t3 = qvector[4 * i + 3]; + lut[16 * i + 0b0000] = 0; + lut[16 * i + 0b0001] = t0; + lut[16 * i + 0b0010] = t1; + lut[16 * i + 0b0011] = t1 + t0; + lut[16 * i + 0b0100] = t2; + lut[16 * i + 0b0101] = t2 + t0; + lut[16 * i + 0b0110] = t2 + t1; + lut[16 * i + 0b0111] = t2 + t1 + t0; + lut[16 * i + 0b1000] = t3; + lut[16 * i + 0b1001] = t3 + t0; + lut[16 * i + 0b1010] = t3 + t1; + lut[16 * i + 0b1011] = t3 + t1 + t0; + lut[16 * i + 0b1100] = t3 + t2; + lut[16 * i + 0b1101] = t3 + t2 + t0; + lut[16 * i + 0b1110] = t3 + t2 + t1; + lut[16 * i + 0b1111] = t3 + t2 + t1 + t0; + } + lut + } + fn fscan_process_lowerbound( + dims: u32, + lut: &(f32, f32, f32, f32, Vec), + (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( + &[f32; 32], + &[f32; 32], + &[f32; 32], + &[f32; 32], + &[u8], + ), + epsilon: f32, + ) -> [Distance; 32] { + use quantization::fast_scan::b4::fast_scan_b4; + let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let r = fast_scan_b4(dims.div_ceil(4), t, s); + std::array::from_fn(|i| { + let rough = dis_u_2[i] + + dis_v_2 + + b * factor_ppc[i] + + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; + let err = factor_err[i] * dis_v_2.sqrt(); + Distance::from_f32(rough - epsilon * err) + }) + } + use base::distance::Distance; + use quantization::quantize; + + let lut = { + let vector = &samples[i]; + let dis_v_2 = f32::reduce_sum_of_x2(vector); + let (k, b, qvector) = + quantize::quantize::<15>(f32::vector_to_f32_borrowed(vector).as_ref()); + let qvector_sum = if vector.len() <= 4369 { + quantize::reduce_sum_of_x_as_u16(&qvector) as f32 + } else { + quantize::reduce_sum_of_x_as_u32(&qvector) as f32 + }; + (dis_v_2, b, k, qvector_sum, gen(qvector)) + }; + + let mut result = (Distance::INFINITY, 0); + for block in 0..c.div_ceil(32) { + let lowerbound = fscan_process_lowerbound( + dims as _, + &lut, + (&a0[block], &a1[block], &a2[block], &a3[block], &a4[block]), + 1.9, + ); + for j in block * 32..std::cmp::min(block * 32 + 32, c) { + if lowerbound[j - block * 32] < result.0 { + let dis = + Distance::from_f32(f32::reduce_sum_of_d2(&samples[i], ¢roids[j])); + if dis <= result.0 { + result = (dis, j); + } + } + } + } + result.1 + }) + .collect::>() +} + +fn flat_index( + parallelism: &P, + _dims: usize, + n: usize, + c: usize, + samples: &Vec>, + centroids: &Vec>, +) -> Vec { + parallelism + .into_par_iter(0..n) + .map(|i| { + let mut result = (f32::INFINITY, 0); + for j in 0..c { + let dis_2 = f32::reduce_sum_of_d2(&samples[i], ¢roids[j]); + if dis_2 <= result.0 { + result = (dis_2, j); + } + } + result.1 + }) + .collect::>() +} + +struct LloydKMeans<'a, P, F> { parallelism: &'a P, dims: usize, c: usize, is_spherical: bool, - centroids: Vec>, + centroids: Vec>, assign: Vec, rng: StdRng, - samples: &'a Vec>, + samples: &'a Vec>, + compute: F, } const DELTA: f32 = f16::EPSILON.to_f32_const(); -impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { +impl<'a, P: Parallelism, F: Fn(&P, &Vec>) -> Vec> LloydKMeans<'a, P, F> { fn new( parallelism: &'a P, c: usize, dims: usize, - samples: &'a Vec>, + samples: &'a Vec>, is_spherical: bool, + compute: F, ) -> Self { let n = samples.len(); @@ -102,7 +308,7 @@ impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { .map(|i| { let mut result = (f32::INFINITY, 0); for j in 0..c { - let dis_2 = S::reduce_sum_of_d2(&samples[i], ¢roids[j]); + let dis_2 = f32::reduce_sum_of_d2(&samples[i], ¢roids[j]); if dis_2 <= result.0 { result = (dis_2, j); } @@ -120,6 +326,7 @@ impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { assign, rng, samples, + compute, } } @@ -134,18 +341,18 @@ impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { .parallelism .into_par_iter(0..n) .fold( - || (vec![vec![S::zero(); dims]; c], vec![0.0f32; c]), + || (vec![vec![f32::zero(); dims]; c], vec![0.0f32; c]), |(mut sum, mut count), i| { - S::vector_add_inplace(&mut sum[self.assign[i]], &samples[i]); + f32::vector_add_inplace(&mut sum[self.assign[i]], &samples[i]); count[self.assign[i]] += 1.0; (sum, count) }, ) .reduce( - || (vec![vec![S::zero(); dims]; c], vec![0.0f32; c]), + || (vec![vec![f32::zero(); dims]; c], vec![0.0f32; c]), |(mut sum, mut count), (sum_1, count_1)| { for i in 0..c { - S::vector_add_inplace(&mut sum[i], &sum_1[i]); + f32::vector_add_inplace(&mut sum[i], &sum_1[i]); count[i] += count_1[i]; } (sum, count) @@ -155,7 +362,7 @@ impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { let mut centroids = self .parallelism .into_par_iter(0..c) - .map(|i| S::vector_mul_scalar(&sum[i], 1.0 / count[i])) + .map(|i| f32::vector_mul_scalar(&sum[i], 1.0 / count[i])) .collect::>(); for i in 0..c { @@ -172,8 +379,8 @@ impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { o = (o + 1) % c; } centroids[i] = centroids[o].clone(); - S::kmeans_helper(&mut centroids[i], 1.0 + DELTA, 1.0 - DELTA); - S::kmeans_helper(&mut centroids[o], 1.0 - DELTA, 1.0 + DELTA); + f32::kmeans_helper(&mut centroids[i], 1.0 + DELTA, 1.0 - DELTA); + f32::kmeans_helper(&mut centroids[o], 1.0 - DELTA, 1.0 + DELTA); count[i] = count[o] / 2.0; count[o] -= count[i]; } @@ -182,25 +389,12 @@ impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { self.parallelism .into_par_iter(&mut centroids) .for_each(|centroid| { - let l = S::reduce_sum_of_x2(centroid).sqrt(); - S::vector_mul_scalar_inplace(centroid, 1.0 / l); + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); }); } - let assign = self - .parallelism - .into_par_iter(0..n) - .map(|i| { - let mut result = (f32::INFINITY, 0); - for j in 0..c { - let dis_2 = S::reduce_sum_of_d2(&samples[i], ¢roids[j]); - if dis_2 <= result.0 { - result = (dis_2, j); - } - } - result.1 - }) - .collect::>(); + let assign = (self.compute)(self.parallelism, ¢roids); let result = (0..n).all(|i| assign[i] == self.assign[i]); @@ -210,7 +404,7 @@ impl<'a, P: Parallelism, S: ScalarLike> LloydKMeans<'a, P, S> { result } - fn finish(self) -> Vec> { + fn finish(self) -> Vec> { self.centroids } } diff --git a/src/vchordrq/algorithm/build.rs b/src/vchordrq/algorithm/build.rs index 8df19052..c65bd851 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -48,8 +48,11 @@ pub fn build( let mut tuples_total = 0_u64; let samples = { let mut rand = rand::thread_rng(); - let max_number_of_samples = - internal_build.lists.last().unwrap().saturating_mul(256); + let max_number_of_samples = internal_build + .lists + .last() + .unwrap() + .saturating_mul(internal_build.sampling_factor); let mut samples = Vec::new(); let mut number_of_samples = 0_u32; heap_relation.traverse(false, |(_, vector)| { diff --git a/src/vchordrq/types.rs b/src/vchordrq/types.rs index 2002d7d6..23011509 100644 --- a/src/vchordrq/types.rs +++ b/src/vchordrq/types.rs @@ -9,6 +9,9 @@ pub struct VchordrqInternalBuildOptions { pub lists: Vec, #[serde(default = "VchordrqInternalBuildOptions::default_spherical_centroids")] pub spherical_centroids: bool, + #[serde(default = "VchordrqInternalBuildOptions::default_sampling_factor")] + #[validate(range(min = 1, max = 1024))] + pub sampling_factor: u32, #[serde(default = "VchordrqInternalBuildOptions::default_build_threads")] #[validate(range(min = 1, max = 255))] pub build_threads: u16, @@ -30,6 +33,9 @@ impl VchordrqInternalBuildOptions { fn default_spherical_centroids() -> bool { false } + fn default_sampling_factor() -> u32 { + 256 + } fn default_build_threads() -> u16 { 1 } @@ -40,6 +46,7 @@ impl Default for VchordrqInternalBuildOptions { Self { lists: Self::default_lists(), spherical_centroids: Self::default_spherical_centroids(), + sampling_factor: Self::default_sampling_factor(), build_threads: Self::default_build_threads(), } } diff --git a/src/vchordrqfscan/algorithm/build.rs b/src/vchordrqfscan/algorithm/build.rs index 9bdc3b92..0528a5c4 100644 --- a/src/vchordrqfscan/algorithm/build.rs +++ b/src/vchordrqfscan/algorithm/build.rs @@ -47,8 +47,11 @@ pub fn build( let mut tuples_total = 0_u64; let samples = { let mut rand = rand::thread_rng(); - let max_number_of_samples = - internal_build.lists.last().unwrap().saturating_mul(256); + let max_number_of_samples = internal_build + .lists + .last() + .unwrap() + .saturating_mul(internal_build.sampling_factor); let mut samples = Vec::new(); let mut number_of_samples = 0_u32; heap_relation.traverse(false, |(_, vector)| { diff --git a/src/vchordrqfscan/types.rs b/src/vchordrqfscan/types.rs index b3a1067f..0fbe82e9 100644 --- a/src/vchordrqfscan/types.rs +++ b/src/vchordrqfscan/types.rs @@ -9,6 +9,9 @@ pub struct VchordrqfscanInternalBuildOptions { pub lists: Vec, #[serde(default = "VchordrqfscanInternalBuildOptions::default_spherical_centroids")] pub spherical_centroids: bool, + #[serde(default = "VchordrqfscanInternalBuildOptions::default_sampling_factor")] + #[validate(range(min = 1, max = 1024))] + pub sampling_factor: u32, #[serde(default = "VchordrqfscanInternalBuildOptions::default_build_threads")] #[validate(range(min = 1, max = 255))] pub build_threads: u16, @@ -30,6 +33,9 @@ impl VchordrqfscanInternalBuildOptions { fn default_spherical_centroids() -> bool { false } + fn default_sampling_factor() -> u32 { + 256 + } fn default_build_threads() -> u16 { 1 } @@ -40,6 +46,7 @@ impl Default for VchordrqfscanInternalBuildOptions { Self { lists: Self::default_lists(), spherical_centroids: Self::default_spherical_centroids(), + sampling_factor: Self::default_sampling_factor(), build_threads: Self::default_build_threads(), } } From 5fb0052a5b70b13fcf3a2e85889218a59adac486 Mon Sep 17 00:00:00 2001 From: xieydd Date: Wed, 27 Nov 2024 17:35:02 +0800 Subject: [PATCH 073/324] add enterprise image build step to ci (#114) Signed-off-by: xieydd --- .github/workflows/release.yml | 56 +++ .github/workflows/release_pg_slim.yml | 36 ++ docker/pg-cnpg/Dockerfile | 190 +++++++++ docker/pg-cnpg/requirements.txt | 513 +++++++++++++++++++++++++ docker/pg-cnpg/trunk-install.sh | 80 ++++ docker/pg-slim/Dockerfile | 137 +++++++ docker/pg-slim/docker-ensure-initdb.sh | 71 ++++ docker/pg-slim/docker-entrypoint.sh | 356 +++++++++++++++++ 8 files changed, 1439 insertions(+) create mode 100644 .github/workflows/release_pg_slim.yml create mode 100644 docker/pg-cnpg/Dockerfile create mode 100644 docker/pg-cnpg/requirements.txt create mode 100755 docker/pg-cnpg/trunk-install.sh create mode 100644 docker/pg-slim/Dockerfile create mode 100755 docker/pg-slim/docker-ensure-initdb.sh create mode 100755 docker/pg-slim/docker-entrypoint.sh diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 9c2e04d7..700dc858 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -127,3 +127,59 @@ jobs: PG_VERSION=${{ matrix.version }} SEMVER=${{ env.SEMVER }} PGVECTOR=0.8.0 + - name: Login to modelzai Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} + password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} + - name: Build and push Enterprise image to Docker Registry + uses: docker/build-push-action@v6 + if: ${{ matrix.version != '17' }} + with: + context: ./docker/pg-cnpg + push: true + platforms: "linux/amd64" + file: ./docker/pg-cnpg/Dockerfile + build-args: | + PG_MAJOR=${{ matrix.version }} + SEMVER=${{ env.SEMVER }} + LIB_DIR=/usr/lib/x86_64-linux-gnu + TARGETARCH=amd64 + PGVECTOR=0.8.0 + tags: modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }} + + test: + name: Run tests + runs-on: + - ubuntu-latest + needs: ["semver", "build", "docker"] + strategy: + matrix: + version: [14, 15, 16] + platform: ["amd64"] + container: + image: modelzai/vchord-cnpg:${{ matrix.version }}-v${{ needs.semver.outputs.SEMVER }} + options: --user root + credentials: + username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} + password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} + env: + PGHOST: "localhost" + PGPORT: "5432" + PGDATABASE: "postgres" + PGUSER: "postgres" + PGPASSWORD: "postgres" + POSTGRES_PASSWORD: "password" + PGDATA: "/var/lib/postgresql/data2" + + steps: + - name: Install all extensions in registry + # Entrypoint is overwritten by GitHub Action. We need to execute it manually in order to start Postgres. + # More information here https://github.com/actions/runner/issues/1964 + run: | + bash /usr/local/bin/docker-entrypoint.sh postgres & + sleep 5 + curl https://registry.pgtrunk.io/extensions/all | jq -r ".[] | .name" > /tmp/extensions.txt + trunk-install.sh | tee /tmp/output.txt + cat /tmp/output.txt + diff --git a/.github/workflows/release_pg_slim.yml b/.github/workflows/release_pg_slim.yml new file mode 100644 index 00000000..779b35db --- /dev/null +++ b/.github/workflows/release_pg_slim.yml @@ -0,0 +1,36 @@ +name: Release for Postgres slim + +on: + workflow_dispatch: + +jobs: + pg-slim: + strategy: + matrix: + version: [14, 15, 16, 17] + platform: ["amd64"] + runs-on: ubuntu-latest + env: + PG_MAJOR: ${{ matrix.version }} + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + - name: Login to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} + password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} + - name: Push binary release to Docker Registry + uses: docker/build-push-action@v4 + with: + context: ./docker/pg-slim + push: true + platforms: "linux/${{ matrix.platform }}" + file: ./docker/pg-slim/Dockerfile + build-args: | + PG_MAJOR=${{ matrix.version }} + tags: modelzai/pg-slim:${{ matrix.version }}-${{ matrix.platform }} \ No newline at end of file diff --git a/docker/pg-cnpg/Dockerfile b/docker/pg-cnpg/Dockerfile new file mode 100644 index 00000000..99b04915 --- /dev/null +++ b/docker/pg-cnpg/Dockerfile @@ -0,0 +1,190 @@ +ARG PG_MAJOR +ARG SEMVER +ARG TARGETARCH + +FROM tensorchord/vchord-binary:pg${PG_MAJOR}-v${SEMVER} as binary + +FROM rust:1.78-bookworm as builder +ARG TRUNK_VER=0.12.25 +ENV CARGO_REGISTRIES_CRATES_IO_PROTOCOL sparse +RUN cargo install --version $TRUNK_VER pg-trunk + +FROM modelzai/pg-slim:${PG_MAJOR}-${TARGETARCH} +ARG PG_MAJOR +ARG SEMVER +ARG TARGETARCH +ARG LIB_DIR +ARG PGVECTOR +ARG ALTDIR=/var/lib/postgresql/data/tensorchord + +USER root + +COPY --from=binary /workspace/vchord-pg${PG_MAJOR}_${SEMVER}_${TARGETARCH}.deb /tmp/vchord.deb +RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb + +# PGDATA is set in pg-slim and used by dependents on this image. +RUN if [ -z "${PGDATA}" ]; then echo "PGDATA is not set"; exit 1; fi + +# Install trunk +COPY --from=builder /usr/local/cargo/bin/trunk /usr/bin/trunk +COPY requirements.txt . + +# Install barman-cloud +RUN set -xe; \ + apt-get update; \ + apt-get install -y --no-install-recommends \ + python3-pip \ + python3-psycopg2 \ + python3-setuptools \ + ; \ + pip3 install --upgrade pip; \ + # TODO: Remove --no-deps once https://github.com/pypa/pip/issues/9644 is solved + pip3 install --no-deps -r requirements.txt; \ + apt-get autoremove -y; \ + apt-get clean; \ + rm -rf /var/lib/apt/lists/*; + +RUN chown -R postgres:postgres ${ALTDIR}/${PG_MAJOR} && \ + chmod -R 0700 ${ALTDIR}/${PG_MAJOR} +RUN chown postgres /usr/share/postgresql/${PG_MAJOR}/extension + +RUN apt-get update && apt-get install -y \ + jq \ + curl \ + wget \ + && rm -rf /var/lib/apt/lists/* + +# Install extension dependencies +RUN apt-get update && apt-get install -y \ + libmysqlclient-dev \ + libtcl8.6 \ + libgeos-dev \ + libproj-dev \ + libjson-c-dev \ + libprotobuf-c-dev \ + libxml2-dev \ + libboost-serialization1.74-dev \ + libhiredis-dev \ + libsybdb5 \ + libpython3.10-dev \ + r-base-core \ + openssl \ + liblz4-1 \ + libpcre2-8-0 \ + libuuid1 \ + libgroonga0 \ + libopenblas0-pthread \ + libcurl4 \ + libjson-c5 \ + libsodium23 \ + libgcc-s1 \ + libselinux1 \ + librdkafka1 \ + libgdal30 \ + libcrypt1 \ + liburiparser1 \ + libfreetype6 \ + libzstd1 \ + zlib1g \ + libperl5.34 \ + libgomp1 \ + libssl3 \ + libsfcgal1 \ + openjdk-11-jdk \ + libaio1 \ + libbson-dev \ + libgsl-dev \ + && rm -rf /var/lib/apt/lists/* +RUN ln -s /usr/lib/jvm/java-11-openjdk-amd64/lib/server/libjvm.so ${LIB_DIR}/libjvm.so +RUN wget https://download.oracle.com/otn_software/linux/instantclient/1920000/instantclient-basiclite-linux.x64-19.20.0.0.0dbru.zip && \ + unzip instantclient-basiclite-linux.x64-19.20.0.0.0dbru.zip && \ + cp instantclient_19_20/libclntsh.so.19.1 ${LIB_DIR}/ && \ + cp instantclient_19_20/libnnz19.so ${LIB_DIR}/ && \ + cp instantclient_19_20/libclntshcore.so.19.1 ${LIB_DIR}/ && \ + rm -rf instantclient_19_20 && \ + rm instantclient-basiclite-linux.x64-19.20.0.0.0dbru.zip + +# Install zhparser dependency +RUN wget http://www.xunsearch.com/scws/down/scws-1.2.3.tar.bz2 && \ + tar xvf scws-1.2.3.tar.bz2 && \ + cd scws-1.2.3 && \ + ./configure && \ + make install && \ + cd .. && \ + rm -rf scws-1.2.3.tar.bz2 scws-1.2.3 && \ + ln -s /usr/local/lib/libscws.so ${LIB_DIR}/libscws.so + +# Install duckdb libs +RUN wget https://github.com/duckdb/duckdb/releases/download/v0.8.1/libduckdb-linux-amd64.zip && \ + unzip libduckdb-linux-amd64.zip && \ + cp libduckdb.so ${LIB_DIR}/ && \ + rm -rf libduckdb-linux-amd64.zip libduckdb.so + +# Install pg_stat_statements +RUN trunk install pg_stat_statements + +# Install auto_explain +RUN trunk install auto_explain + +# Install plpython3u +RUN trunk install plpython3u + +# Install pgvector +RUN trunk install pgvector --version ${PGVECTOR} + +# Clone and build AWS SDK for C++ +RUN git clone https://github.com/aws/aws-sdk-cpp.git && \ + cd aws-sdk-cpp && \ + git checkout 1.9.263 && \ + git submodule update --init --recursive && \ + mkdir build && cd build && \ + cmake -DBUILD_ONLY="s3;core;config;sts;cognito-identity;transfer;identity-management" -DAUTORUN_UNIT_TESTS=OFF -DCMAKE_CXX_FLAGS=-Wno-error=deprecated-declarations .. && \ + make -j$(nproc) && \ + make install && \ + cd ../../../ && rm -rf aws-sdk-cpp + +# Clone and build Apache Arrow +RUN git clone https://github.com/apache/arrow.git && \ + cd arrow && \ + git checkout apache-arrow-7.0.1 && \ + cd cpp && \ + mkdir build && cd build && \ + cmake -DARROW_PARQUET=ON -DARROW_S3=ON -DARROW_WITH_SNAPPY=ON .. && \ + make -j$(nproc) && \ + make install && \ + cd ../../../ && rm -rf arrow + +# Clone and build pgaudit +RUN git clone https://github.com/pgaudit/pgaudit.git && \ + cd pgaudit && \ + git checkout REL_${PG_MAJOR}_STABLE && \ + make install USE_PGXS=1 PG_CONFIG=/usr/lib/postgresql/${PG_MAJOR}/bin/pg_config && \ + cd ../ && rm -rf pgaudit + +# Clone and build pg_failover_slots +RUN git clone https://github.com/EnterpriseDB/pg_failover_slots.git && \ + cd pg_failover_slots && \ + make install PG_CONFIG=/usr/lib/postgresql/${PG_MAJOR}/bin/pg_config && \ + cd ../ && rm -rf pg_failover_slots + +# cache all extensions +ENV LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH + +# Test trunk +COPY trunk-install.sh /usr/local/bin/ + +# Change the uid of postgres to 26 +RUN usermod -u 26 postgres +RUN chown -R postgres:postgres ${ALTDIR} +RUN cp /usr/share/postgresql/${PG_MAJOR}/extension/* ${ALTDIR}/extension/ +RUN cp /usr/lib/postgresql/${PG_MAJOR}/lib/* ${ALTDIR}/${PG_MAJOR}/lib/ + +RUN set -eux; \ + mkdir /tmp/pg_pkglibdir; \ + mkdir /tmp/pg_sharedir; \ + cp -r $(pg_config --pkglibdir)/* /tmp/pg_pkglibdir; \ + cp -r $(pg_config --sharedir)/* /tmp/pg_sharedir + +RUN chown -R postgres:postgres /tmp +USER 26 +ENV PATH $PATH:/usr/lib/postgresql/${PG_MAJOR}/bin diff --git a/docker/pg-cnpg/requirements.txt b/docker/pg-cnpg/requirements.txt new file mode 100644 index 00000000..f7c47313 --- /dev/null +++ b/docker/pg-cnpg/requirements.txt @@ -0,0 +1,513 @@ +# +# This file is autogenerated by pip-compile with Python 3.11 +# by the following command: +# +# pip-compile --generate-hashes +# +azure-core==1.32.0 \ + --hash=sha256:22b3c35d6b2dae14990f6c1be2912bf23ffe50b220e708a28ab1bb92b1c730e5 \ + --hash=sha256:eac191a0efb23bfa83fddf321b27b122b4ec847befa3091fa736a5c32c50d7b4 + # via + # azure-identity + # azure-storage-blob +azure-identity==1.19.0 \ + --hash=sha256:500144dc18197d7019b81501165d4fa92225f03778f17d7ca8a2a180129a9c83 \ + --hash=sha256:e3f6558c181692d7509f09de10cca527c7dce426776454fb97df512a46527e81 +azure-storage-blob==12.23.1 \ + --hash=sha256:1c2238aa841d1545f42714a5017c010366137a44a0605da2d45f770174bfc6b4 \ + --hash=sha256:a587e54d4e39d2a27bd75109db164ffa2058fe194061e5446c5a89bca918272f +barman[azure,cloud,google,snappy]==3.11.1 \ + --hash=sha256:295b9b7e058e064338f66ca0d10e4892e784a2347f06e4a225164995f6114498 \ + --hash=sha256:4f424f3327cb24fb82d6a29dc1cdf02222b950c447c78273273d6eb76d7ce8d7 + # via -r requirements.in +boto3==1.35.54 \ + --hash=sha256:2d5e160b614db55fbee7981001c54476cb827c441cef65b2fcb2c52a62019909 \ + --hash=sha256:7d9c359bbbc858a60b51c86328db813353c8bd1940212cdbd0a7da835291c2e1 +botocore==1.35.54 \ + --hash=sha256:131bb59ce59c8a939b31e8e647242d70cf11d32d4529fa4dca01feea1e891a76 \ + --hash=sha256:9cca1811094b6cdc144c2c063a3ec2db6d7c88194b04d4277cd34fc8e3473aff + # via + # boto3 + # s3transfer +cachetools==5.5.0 \ + --hash=sha256:02134e8439cdc2ffb62023ce1debca2944c3f289d66bb17ead3ab3dede74b292 \ + --hash=sha256:2cc24fb4cbe39633fb7badd9db9ca6295d766d9c2995f245725a46715d050f2a + # via google-auth +certifi==2024.8.30 \ + --hash=sha256:922820b53db7a7257ffbda3f597266d435245903d80737e34f8a45ff3e3230d8 \ + --hash=sha256:bec941d2aa8195e248a60b31ff9f0558284cf01a52591ceda73ea9afffd69fd9 + # via requests +cffi==1.17.1 \ + --hash=sha256:045d61c734659cc045141be4bae381a41d89b741f795af1dd018bfb532fd0df8 \ + --hash=sha256:0984a4925a435b1da406122d4d7968dd861c1385afe3b45ba82b750f229811e2 \ + --hash=sha256:0e2b1fac190ae3ebfe37b979cc1ce69c81f4e4fe5746bb401dca63a9062cdaf1 \ + --hash=sha256:0f048dcf80db46f0098ccac01132761580d28e28bc0f78ae0d58048063317e15 \ + --hash=sha256:1257bdabf294dceb59f5e70c64a3e2f462c30c7ad68092d01bbbfb1c16b1ba36 \ + --hash=sha256:1c39c6016c32bc48dd54561950ebd6836e1670f2ae46128f67cf49e789c52824 \ + --hash=sha256:1d599671f396c4723d016dbddb72fe8e0397082b0a77a4fab8028923bec050e8 \ + --hash=sha256:28b16024becceed8c6dfbc75629e27788d8a3f9030691a1dbf9821a128b22c36 \ + --hash=sha256:2bb1a08b8008b281856e5971307cc386a8e9c5b625ac297e853d36da6efe9c17 \ + --hash=sha256:30c5e0cb5ae493c04c8b42916e52ca38079f1b235c2f8ae5f4527b963c401caf \ + --hash=sha256:31000ec67d4221a71bd3f67df918b1f88f676f1c3b535a7eb473255fdc0b83fc \ + --hash=sha256:386c8bf53c502fff58903061338ce4f4950cbdcb23e2902d86c0f722b786bbe3 \ + --hash=sha256:3edc8d958eb099c634dace3c7e16560ae474aa3803a5df240542b305d14e14ed \ + --hash=sha256:45398b671ac6d70e67da8e4224a065cec6a93541bb7aebe1b198a61b58c7b702 \ + --hash=sha256:46bf43160c1a35f7ec506d254e5c890f3c03648a4dbac12d624e4490a7046cd1 \ + --hash=sha256:4ceb10419a9adf4460ea14cfd6bc43d08701f0835e979bf821052f1805850fe8 \ + --hash=sha256:51392eae71afec0d0c8fb1a53b204dbb3bcabcb3c9b807eedf3e1e6ccf2de903 \ + --hash=sha256:5da5719280082ac6bd9aa7becb3938dc9f9cbd57fac7d2871717b1feb0902ab6 \ + --hash=sha256:610faea79c43e44c71e1ec53a554553fa22321b65fae24889706c0a84d4ad86d \ + --hash=sha256:636062ea65bd0195bc012fea9321aca499c0504409f413dc88af450b57ffd03b \ + --hash=sha256:6883e737d7d9e4899a8a695e00ec36bd4e5e4f18fabe0aca0efe0a4b44cdb13e \ + --hash=sha256:6b8b4a92e1c65048ff98cfe1f735ef8f1ceb72e3d5f0c25fdb12087a23da22be \ + --hash=sha256:6f17be4345073b0a7b8ea599688f692ac3ef23ce28e5df79c04de519dbc4912c \ + --hash=sha256:706510fe141c86a69c8ddc029c7910003a17353970cff3b904ff0686a5927683 \ + --hash=sha256:72e72408cad3d5419375fc87d289076ee319835bdfa2caad331e377589aebba9 \ + --hash=sha256:733e99bc2df47476e3848417c5a4540522f234dfd4ef3ab7fafdf555b082ec0c \ + --hash=sha256:7596d6620d3fa590f677e9ee430df2958d2d6d6de2feeae5b20e82c00b76fbf8 \ + --hash=sha256:78122be759c3f8a014ce010908ae03364d00a1f81ab5c7f4a7a5120607ea56e1 \ + --hash=sha256:805b4371bf7197c329fcb3ead37e710d1bca9da5d583f5073b799d5c5bd1eee4 \ + --hash=sha256:85a950a4ac9c359340d5963966e3e0a94a676bd6245a4b55bc43949eee26a655 \ + --hash=sha256:8f2cdc858323644ab277e9bb925ad72ae0e67f69e804f4898c070998d50b1a67 \ + --hash=sha256:9755e4345d1ec879e3849e62222a18c7174d65a6a92d5b346b1863912168b595 \ + --hash=sha256:98e3969bcff97cae1b2def8ba499ea3d6f31ddfdb7635374834cf89a1a08ecf0 \ + --hash=sha256:a08d7e755f8ed21095a310a693525137cfe756ce62d066e53f502a83dc550f65 \ + --hash=sha256:a1ed2dd2972641495a3ec98445e09766f077aee98a1c896dcb4ad0d303628e41 \ + --hash=sha256:a24ed04c8ffd54b0729c07cee15a81d964e6fee0e3d4d342a27b020d22959dc6 \ + --hash=sha256:a45e3c6913c5b87b3ff120dcdc03f6131fa0065027d0ed7ee6190736a74cd401 \ + --hash=sha256:a9b15d491f3ad5d692e11f6b71f7857e7835eb677955c00cc0aefcd0669adaf6 \ + --hash=sha256:ad9413ccdeda48c5afdae7e4fa2192157e991ff761e7ab8fdd8926f40b160cc3 \ + --hash=sha256:b2ab587605f4ba0bf81dc0cb08a41bd1c0a5906bd59243d56bad7668a6fc6c16 \ + --hash=sha256:b62ce867176a75d03a665bad002af8e6d54644fad99a3c70905c543130e39d93 \ + --hash=sha256:c03e868a0b3bc35839ba98e74211ed2b05d2119be4e8a0f224fba9384f1fe02e \ + --hash=sha256:c59d6e989d07460165cc5ad3c61f9fd8f1b4796eacbd81cee78957842b834af4 \ + --hash=sha256:c7eac2ef9b63c79431bc4b25f1cd649d7f061a28808cbc6c47b534bd789ef964 \ + --hash=sha256:c9c3d058ebabb74db66e431095118094d06abf53284d9c81f27300d0e0d8bc7c \ + --hash=sha256:ca74b8dbe6e8e8263c0ffd60277de77dcee6c837a3d0881d8c1ead7268c9e576 \ + --hash=sha256:caaf0640ef5f5517f49bc275eca1406b0ffa6aa184892812030f04c2abf589a0 \ + --hash=sha256:cdf5ce3acdfd1661132f2a9c19cac174758dc2352bfe37d98aa7512c6b7178b3 \ + --hash=sha256:d016c76bdd850f3c626af19b0542c9677ba156e4ee4fccfdd7848803533ef662 \ + --hash=sha256:d01b12eeeb4427d3110de311e1774046ad344f5b1a7403101878976ecd7a10f3 \ + --hash=sha256:d63afe322132c194cf832bfec0dc69a99fb9bb6bbd550f161a49e9e855cc78ff \ + --hash=sha256:da95af8214998d77a98cc14e3a3bd00aa191526343078b530ceb0bd710fb48a5 \ + --hash=sha256:dd398dbc6773384a17fe0d3e7eeb8d1a21c2200473ee6806bb5e6a8e62bb73dd \ + --hash=sha256:de2ea4b5833625383e464549fec1bc395c1bdeeb5f25c4a3a82b5a8c756ec22f \ + --hash=sha256:de55b766c7aa2e2a3092c51e0483d700341182f08e67c63630d5b6f200bb28e5 \ + --hash=sha256:df8b1c11f177bc2313ec4b2d46baec87a5f3e71fc8b45dab2ee7cae86d9aba14 \ + --hash=sha256:e03eab0a8677fa80d646b5ddece1cbeaf556c313dcfac435ba11f107ba117b5d \ + --hash=sha256:e221cf152cff04059d011ee126477f0d9588303eb57e88923578ace7baad17f9 \ + --hash=sha256:e31ae45bc2e29f6b2abd0de1cc3b9d5205aa847cafaecb8af1476a609a2f6eb7 \ + --hash=sha256:edae79245293e15384b51f88b00613ba9f7198016a5948b5dddf4917d4d26382 \ + --hash=sha256:f1e22e8c4419538cb197e4dd60acc919d7696e5ef98ee4da4e01d3f8cfa4cc5a \ + --hash=sha256:f3a2b4222ce6b60e2e8b337bb9596923045681d71e5a082783484d845390938e \ + --hash=sha256:f6a16c31041f09ead72d69f583767292f750d24913dadacf5756b966aacb3f1a \ + --hash=sha256:f75c7ab1f9e4aca5414ed4d8e5c0e303a34f4421f8a0d47a4d019ceff0ab6af4 \ + --hash=sha256:f79fc4fc25f1c8698ff97788206bb3c2598949bfe0fef03d299eb1b5356ada99 \ + --hash=sha256:f7f5baafcc48261359e14bcd6d9bff6d4b28d9103847c9e136694cb0501aef87 \ + --hash=sha256:fc48c783f9c87e60831201f2cce7f3b2e4846bf4d8728eabe54d60700b318a0b + # via cryptography +charset-normalizer==3.4.0 \ + --hash=sha256:0099d79bdfcf5c1f0c2c72f91516702ebf8b0b8ddd8905f97a8aecf49712c621 \ + --hash=sha256:0713f3adb9d03d49d365b70b84775d0a0d18e4ab08d12bc46baa6132ba78aaf6 \ + --hash=sha256:07afec21bbbbf8a5cc3651aa96b980afe2526e7f048fdfb7f1014d84acc8b6d8 \ + --hash=sha256:0b309d1747110feb25d7ed6b01afdec269c647d382c857ef4663bbe6ad95a912 \ + --hash=sha256:0d99dd8ff461990f12d6e42c7347fd9ab2532fb70e9621ba520f9e8637161d7c \ + --hash=sha256:0de7b687289d3c1b3e8660d0741874abe7888100efe14bd0f9fd7141bcbda92b \ + --hash=sha256:1110e22af8ca26b90bd6364fe4c763329b0ebf1ee213ba32b68c73de5752323d \ + --hash=sha256:130272c698667a982a5d0e626851ceff662565379baf0ff2cc58067b81d4f11d \ + --hash=sha256:136815f06a3ae311fae551c3df1f998a1ebd01ddd424aa5603a4336997629e95 \ + --hash=sha256:14215b71a762336254351b00ec720a8e85cada43b987da5a042e4ce3e82bd68e \ + --hash=sha256:1db4e7fefefd0f548d73e2e2e041f9df5c59e178b4c72fbac4cc6f535cfb1565 \ + --hash=sha256:1ffd9493de4c922f2a38c2bf62b831dcec90ac673ed1ca182fe11b4d8e9f2a64 \ + --hash=sha256:2006769bd1640bdf4d5641c69a3d63b71b81445473cac5ded39740a226fa88ab \ + --hash=sha256:20587d20f557fe189b7947d8e7ec5afa110ccf72a3128d61a2a387c3313f46be \ + --hash=sha256:223217c3d4f82c3ac5e29032b3f1c2eb0fb591b72161f86d93f5719079dae93e \ + --hash=sha256:27623ba66c183eca01bf9ff833875b459cad267aeeb044477fedac35e19ba907 \ + --hash=sha256:285e96d9d53422efc0d7a17c60e59f37fbf3dfa942073f666db4ac71e8d726d0 \ + --hash=sha256:2de62e8801ddfff069cd5c504ce3bc9672b23266597d4e4f50eda28846c322f2 \ + --hash=sha256:2f6c34da58ea9c1a9515621f4d9ac379871a8f21168ba1b5e09d74250de5ad62 \ + --hash=sha256:309a7de0a0ff3040acaebb35ec45d18db4b28232f21998851cfa709eeff49d62 \ + --hash=sha256:35c404d74c2926d0287fbd63ed5d27eb911eb9e4a3bb2c6d294f3cfd4a9e0c23 \ + --hash=sha256:3710a9751938947e6327ea9f3ea6332a09bf0ba0c09cae9cb1f250bd1f1549bc \ + --hash=sha256:3d59d125ffbd6d552765510e3f31ed75ebac2c7470c7274195b9161a32350284 \ + --hash=sha256:40d3ff7fc90b98c637bda91c89d51264a3dcf210cade3a2c6f838c7268d7a4ca \ + --hash=sha256:425c5f215d0eecee9a56cdb703203dda90423247421bf0d67125add85d0c4455 \ + --hash=sha256:43193c5cda5d612f247172016c4bb71251c784d7a4d9314677186a838ad34858 \ + --hash=sha256:44aeb140295a2f0659e113b31cfe92c9061622cadbc9e2a2f7b8ef6b1e29ef4b \ + --hash=sha256:47334db71978b23ebcf3c0f9f5ee98b8d65992b65c9c4f2d34c2eaf5bcaf0594 \ + --hash=sha256:4796efc4faf6b53a18e3d46343535caed491776a22af773f366534056c4e1fbc \ + --hash=sha256:4a51b48f42d9358460b78725283f04bddaf44a9358197b889657deba38f329db \ + --hash=sha256:4b67fdab07fdd3c10bb21edab3cbfe8cf5696f453afce75d815d9d7223fbe88b \ + --hash=sha256:4ec9dd88a5b71abfc74e9df5ebe7921c35cbb3b641181a531ca65cdb5e8e4dea \ + --hash=sha256:4f9fc98dad6c2eaa32fc3af1417d95b5e3d08aff968df0cd320066def971f9a6 \ + --hash=sha256:54b6a92d009cbe2fb11054ba694bc9e284dad30a26757b1e372a1fdddaf21920 \ + --hash=sha256:55f56e2ebd4e3bc50442fbc0888c9d8c94e4e06a933804e2af3e89e2f9c1c749 \ + --hash=sha256:5726cf76c982532c1863fb64d8c6dd0e4c90b6ece9feb06c9f202417a31f7dd7 \ + --hash=sha256:5d447056e2ca60382d460a604b6302d8db69476fd2015c81e7c35417cfabe4cd \ + --hash=sha256:5ed2e36c3e9b4f21dd9422f6893dec0abf2cca553af509b10cd630f878d3eb99 \ + --hash=sha256:5ff2ed8194587faf56555927b3aa10e6fb69d931e33953943bc4f837dfee2242 \ + --hash=sha256:62f60aebecfc7f4b82e3f639a7d1433a20ec32824db2199a11ad4f5e146ef5ee \ + --hash=sha256:63bc5c4ae26e4bc6be6469943b8253c0fd4e4186c43ad46e713ea61a0ba49129 \ + --hash=sha256:6b40e8d38afe634559e398cc32b1472f376a4099c75fe6299ae607e404c033b2 \ + --hash=sha256:6b493a043635eb376e50eedf7818f2f322eabbaa974e948bd8bdd29eb7ef2a51 \ + --hash=sha256:6dba5d19c4dfab08e58d5b36304b3f92f3bd5d42c1a3fa37b5ba5cdf6dfcbcee \ + --hash=sha256:6fd30dc99682dc2c603c2b315bded2799019cea829f8bf57dc6b61efde6611c8 \ + --hash=sha256:707b82d19e65c9bd28b81dde95249b07bf9f5b90ebe1ef17d9b57473f8a64b7b \ + --hash=sha256:7706f5850360ac01d80c89bcef1640683cc12ed87f42579dab6c5d3ed6888613 \ + --hash=sha256:7782afc9b6b42200f7362858f9e73b1f8316afb276d316336c0ec3bd73312742 \ + --hash=sha256:79983512b108e4a164b9c8d34de3992f76d48cadc9554c9e60b43f308988aabe \ + --hash=sha256:7f683ddc7eedd742e2889d2bfb96d69573fde1d92fcb811979cdb7165bb9c7d3 \ + --hash=sha256:82357d85de703176b5587dbe6ade8ff67f9f69a41c0733cf2425378b49954de5 \ + --hash=sha256:84450ba661fb96e9fd67629b93d2941c871ca86fc38d835d19d4225ff946a631 \ + --hash=sha256:86f4e8cca779080f66ff4f191a685ced73d2f72d50216f7112185dc02b90b9b7 \ + --hash=sha256:8cda06946eac330cbe6598f77bb54e690b4ca93f593dee1568ad22b04f347c15 \ + --hash=sha256:8ce7fd6767a1cc5a92a639b391891bf1c268b03ec7e021c7d6d902285259685c \ + --hash=sha256:8ff4e7cdfdb1ab5698e675ca622e72d58a6fa2a8aa58195de0c0061288e6e3ea \ + --hash=sha256:9289fd5dddcf57bab41d044f1756550f9e7cf0c8e373b8cdf0ce8773dc4bd417 \ + --hash=sha256:92a7e36b000bf022ef3dbb9c46bfe2d52c047d5e3f3343f43204263c5addc250 \ + --hash=sha256:92db3c28b5b2a273346bebb24857fda45601aef6ae1c011c0a997106581e8a88 \ + --hash=sha256:95c3c157765b031331dd4db3c775e58deaee050a3042fcad72cbc4189d7c8dca \ + --hash=sha256:980b4f289d1d90ca5efcf07958d3eb38ed9c0b7676bf2831a54d4f66f9c27dfa \ + --hash=sha256:9ae4ef0b3f6b41bad6366fb0ea4fc1d7ed051528e113a60fa2a65a9abb5b1d99 \ + --hash=sha256:9c98230f5042f4945f957d006edccc2af1e03ed5e37ce7c373f00a5a4daa6149 \ + --hash=sha256:9fa2566ca27d67c86569e8c85297aaf413ffab85a8960500f12ea34ff98e4c41 \ + --hash=sha256:a14969b8691f7998e74663b77b4c36c0337cb1df552da83d5c9004a93afdb574 \ + --hash=sha256:a8aacce6e2e1edcb6ac625fb0f8c3a9570ccc7bfba1f63419b3769ccf6a00ed0 \ + --hash=sha256:a8e538f46104c815be19c975572d74afb53f29650ea2025bbfaef359d2de2f7f \ + --hash=sha256:aa41e526a5d4a9dfcfbab0716c7e8a1b215abd3f3df5a45cf18a12721d31cb5d \ + --hash=sha256:aa693779a8b50cd97570e5a0f343538a8dbd3e496fa5dcb87e29406ad0299654 \ + --hash=sha256:ab22fbd9765e6954bc0bcff24c25ff71dcbfdb185fcdaca49e81bac68fe724d3 \ + --hash=sha256:ab2e5bef076f5a235c3774b4f4028a680432cded7cad37bba0fd90d64b187d19 \ + --hash=sha256:ab973df98fc99ab39080bfb0eb3a925181454d7c3ac8a1e695fddfae696d9e90 \ + --hash=sha256:af73657b7a68211996527dbfeffbb0864e043d270580c5aef06dc4b659a4b578 \ + --hash=sha256:b197e7094f232959f8f20541ead1d9862ac5ebea1d58e9849c1bf979255dfac9 \ + --hash=sha256:b295729485b06c1a0683af02a9e42d2caa9db04a373dc38a6a58cdd1e8abddf1 \ + --hash=sha256:b8831399554b92b72af5932cdbbd4ddc55c55f631bb13ff8fe4e6536a06c5c51 \ + --hash=sha256:b8dcd239c743aa2f9c22ce674a145e0a25cb1566c495928440a181ca1ccf6719 \ + --hash=sha256:bcb4f8ea87d03bc51ad04add8ceaf9b0f085ac045ab4d74e73bbc2dc033f0236 \ + --hash=sha256:bd7af3717683bea4c87acd8c0d3d5b44d56120b26fd3f8a692bdd2d5260c620a \ + --hash=sha256:bf4475b82be41b07cc5e5ff94810e6a01f276e37c2d55571e3fe175e467a1a1c \ + --hash=sha256:c3e446d253bd88f6377260d07c895816ebf33ffffd56c1c792b13bff9c3e1ade \ + --hash=sha256:c57516e58fd17d03ebe67e181a4e4e2ccab1168f8c2976c6a334d4f819fe5944 \ + --hash=sha256:c94057af19bc953643a33581844649a7fdab902624d2eb739738a30e2b3e60fc \ + --hash=sha256:cab5d0b79d987c67f3b9e9c53f54a61360422a5a0bc075f43cab5621d530c3b6 \ + --hash=sha256:ce031db0408e487fd2775d745ce30a7cd2923667cf3b69d48d219f1d8f5ddeb6 \ + --hash=sha256:cee4373f4d3ad28f1ab6290684d8e2ebdb9e7a1b74fdc39e4c211995f77bec27 \ + --hash=sha256:d5b054862739d276e09928de37c79ddeec42a6e1bfc55863be96a36ba22926f6 \ + --hash=sha256:dbe03226baf438ac4fda9e2d0715022fd579cb641c4cf639fa40d53b2fe6f3e2 \ + --hash=sha256:dc15e99b2d8a656f8e666854404f1ba54765871104e50c8e9813af8a7db07f12 \ + --hash=sha256:dcaf7c1524c0542ee2fc82cc8ec337f7a9f7edee2532421ab200d2b920fc97cf \ + --hash=sha256:dd4eda173a9fcccb5f2e2bd2a9f423d180194b1bf17cf59e3269899235b2a114 \ + --hash=sha256:dd9a8bd8900e65504a305bf8ae6fa9fbc66de94178c420791d0293702fce2df7 \ + --hash=sha256:de7376c29d95d6719048c194a9cf1a1b0393fbe8488a22008610b0361d834ecf \ + --hash=sha256:e7fdd52961feb4c96507aa649550ec2a0d527c086d284749b2f582f2d40a2e0d \ + --hash=sha256:e91f541a85298cf35433bf66f3fab2a4a2cff05c127eeca4af174f6d497f0d4b \ + --hash=sha256:e9e3c4c9e1ed40ea53acf11e2a386383c3304212c965773704e4603d589343ed \ + --hash=sha256:ee803480535c44e7f5ad00788526da7d85525cfefaf8acf8ab9a310000be4b03 \ + --hash=sha256:f09cb5a7bbe1ecae6e87901a2eb23e0256bb524a79ccc53eb0b7629fbe7677c4 \ + --hash=sha256:f19c1585933c82098c2a520f8ec1227f20e339e33aca8fa6f956f6691b784e67 \ + --hash=sha256:f1a2f519ae173b5b6a2c9d5fa3116ce16e48b3462c8b96dfdded11055e3d6365 \ + --hash=sha256:f28f891ccd15c514a0981f3b9db9aa23d62fe1a99997512b0491d2ed323d229a \ + --hash=sha256:f3e73a4255342d4eb26ef6df01e3962e73aa29baa3124a8e824c5d3364a65748 \ + --hash=sha256:f606a1881d2663630ea5b8ce2efe2111740df4b687bd78b34a8131baa007f79b \ + --hash=sha256:fe9f97feb71aa9896b81973a7bbada8c49501dc73e58a10fcef6663af95e5079 \ + --hash=sha256:ffc519621dce0c767e96b9c53f09c5d215578e10b02c285809f76509a3931482 + # via requests +cramjam==2.9.0 \ + --hash=sha256:00d96f798bc980b29f8e1c3ed7d554050e05d4cde23d1633ffed4cd63110024a \ + --hash=sha256:013cb872205641c6e5269f530ed40aaaa5640d84e0d8f33b89f5a1bf7f655527 \ + --hash=sha256:017b7066f18b7b676068f51b1dbdecc02d76d9af10092252b22dcbd03a78ed33 \ + --hash=sha256:05970fb640f236767003e62c256a085754536169bac863f4a3502ecb59cbf197 \ + --hash=sha256:0711c776750e243ae347d6609c975f0ff4be9ae65b2764d29e4bbdad8e574c3a \ + --hash=sha256:0b1410e68c464666473a89cade17483b94bb4639d9161c440ee54ee1e0eca583 \ + --hash=sha256:0b4f1b5e33915ed591c0c19b8c3bbdd7aa0f6a9bfe2b7246b475d497bda15f18 \ + --hash=sha256:0ed2fef010d1caca9ea63814e9cb5b1d47d907b80302b8cc0b3a1e116ea241e2 \ + --hash=sha256:0ed6362cb6c964f8d0c6e7f790e8961b9242cd3acd87c56169ca14d642653707 \ + --hash=sha256:0fb5ea631dbf998f667766a9e485e757817d66ed559916ba553a0ec2f902d788 \ + --hash=sha256:132db7d3346ea21ba44e7ee23ec73bd6fa9eb1e77133ca6dfe1f7449a69999af \ + --hash=sha256:170a50407f9400073621cc1d5f3200ca3ad9de3000831e3e86f5561ca8048a08 \ + --hash=sha256:1b4ca30c9f27e3b88bc082d4637e7648f93da5cb69a2dbe0c0300bc51353c820 \ + --hash=sha256:1ca28a8f6ab5fca35f163fd7d7a970880ce4fc1a0bead1249ecdaa96ec9ac1f4 \ + --hash=sha256:1d5b5512dc61ea78f32e021e88a5fd5b46a821409479e6657d33614fc9e45677 \ + --hash=sha256:21569f19d5848606b85ac0dde0dc3639319d26fed8522c7103515df875bcb300 \ + --hash=sha256:24afad3ba62774abbb150dc25aab21b047ab999c4143c7a8d96577848baf7af6 \ + --hash=sha256:2addf801c88bead21256ccd87dc97cffead03758c4a4947fad8e454f4abfda0a \ + --hash=sha256:3121e2fbec58907fa70636adaeaf30c27614c867e08a7a5bd2887b33786ff790 \ + --hash=sha256:34578e4c1518b10dad5e0ba40c721e529ef13e7742a528843b40e1f20dd6078c \ + --hash=sha256:35cad507eb02c775e6c5444312f98b28dd8bf122425677ae199484996e838673 \ + --hash=sha256:36426e3f1920f6aa4c644d007bf9cfad06dd9f1a30cd0a921d72b010492d8447 \ + --hash=sha256:37054c73704a3183b60869e7fec1614648752c31d89f44de1ffe1f01ad4d20d5 \ + --hash=sha256:399baf80fea574e3870f233e12e6a12f02c53b054e13d792348b272b0614370a \ + --hash=sha256:441d3875cdffe5df9294b93ef570058837732dd727cd9d18efa0f089f1c2687a \ + --hash=sha256:4714e1ea0c3329368b83fe5ad6e831d5ca11fb794ca7cf491622eb6b2d420d2f \ + --hash=sha256:47d7253b5a10c201cc65aecfb517dfa1c0b5831b2524ac32dd2964fceafc0dc4 \ + --hash=sha256:4a63c4e63319bf7dfc3ab46c06afb76d3d9cc1c94369b609dde480e5cc78e4de \ + --hash=sha256:4adbf4366f8dc29b7c5c731c800cf633be76c9911e928daeb606827d6ae7c599 \ + --hash=sha256:4bd76b654275736fd4f55521981b73751c34dacf70a1dbce96e454a39d43201f \ + --hash=sha256:4d5a39118008bb9f2fba36a0ceea6c41fbd0b55d2647b043ba51a868e5f6de92 \ + --hash=sha256:598eac1713ddbe69c3b30dcc890d69b206ce08903fc3aed58149aae87c61973a \ + --hash=sha256:604c16052cf29d0c796927ed7e107f65429d2036c82c9a8009bd453c94e5e4f0 \ + --hash=sha256:65a097ea765dd4ef2fb868b5b0959d7c93a64c250b2c52f462898c823ae4b950 \ + --hash=sha256:65bded20fd2cef17b22246c336ddd67fac842341ee311042b4a70e65dc745aa7 \ + --hash=sha256:6ebee5f5d7e2b9277895ea4fd94646b72075fe9cfc0e8f4770b65c9e72b1fec1 \ + --hash=sha256:72e9ebc27c557706a3c9964c1d1b4522857760dbd60c105a4f5421f3b66e31a2 \ + --hash=sha256:78e7349f945a83bc48855fb042873092a69b155a088b8c11942eb76418b32705 \ + --hash=sha256:7de19a382bcab93cd4d028d51f6f581920a3b79659a384775188135b7fc64f15 \ + --hash=sha256:7f33a83969fa94ee8e0c1f0aef8eb303ead3e9142338dc543abeb7e1a28734ab \ + --hash=sha256:7f6ef35eba883927af2678b561cc4407e0b3b0d58a251c863bec4b3d8258cc2f \ + --hash=sha256:8982925d179b940efa860513a31b839bb06343501077cca3e67f7a2f7360d355 \ + --hash=sha256:8e33ebe4d709b21bc15e7ddf485ac6b30d7fdc7ed7c3c65130654c007f50c183 \ + --hash=sha256:904be92e3bc25e78343ee52aa0fd5fba3a31d11d474e8af4623a9d00baa84bc2 \ + --hash=sha256:912c94781c8ff318a4d3f3306f8d94d41ae5aa7b9760c4bb0476b01142084845 \ + --hash=sha256:9221297c547d702e1431e96705fce26c6a87df34a681a6b97fe63b536d09c1d8 \ + --hash=sha256:97a6311bd32f301ff1b922bc9de62ace3d9fd845e20efc0f71b4d0239a45b8d2 \ + --hash=sha256:9862ca8ead80857ecfb9b07f02f577733261e981346f31585fe118975eabb738 \ + --hash=sha256:9de33ef3bc006c11fbad1dc8b15341dcc78430df2c5ce1e790dfb729b11ab593 \ + --hash=sha256:9f685fe4e49b2f3e233548e3397b3f9189d71a265718ec631d13eca3d5718ddb \ + --hash=sha256:a0f654c739a6bc4a69a2aaf31463328a208757ed780ff886234532f78e06a864 \ + --hash=sha256:a4156fcefa1dfaa65d35ff82c252d1e32be12820f26d04748be6cd3b461cf85f \ + --hash=sha256:a41b4b10a381be1d42a1a7dd07b8c3faccd3d12c7e98e973a6ec558fd040a607 \ + --hash=sha256:ab17a429a92db90bf40115efb97d10e71b94b0dcacf30cf724552df2794a58fb \ + --hash=sha256:abd8bf9a94e3866215ac181a7dbcfa1ddbedca4f8048494a79934febe88537df \ + --hash=sha256:ad301801afa0eecdacabf353a2802df5e6770f9bfb0a559d6c069813d83cfd42 \ + --hash=sha256:b0078727fe8c28ef1695e5d04aae5c41ac697eb087cba387c6a02b825f9071c0 \ + --hash=sha256:b21e55b5cfdaff96eae1f323ae9a0d36e86852cdf62fe23b60a2481d2fed5571 \ + --hash=sha256:b4a3104022129d7463100dfaf12efd398ebfa4b7e4e50832ccc596754f7c26df \ + --hash=sha256:b4b8d8160685c11ffb4e8e6daaab79cb351a1c54ceec41cc18a0a62c89309fe0 \ + --hash=sha256:b8f8b1117b4e697d39950ecab01700ce0aef66541e4478eb4d7b3ade8703347b \ + --hash=sha256:b99efaf81be8e381de1cde6574e2c89030ed53994e73b0e75b62d6e232f491c5 \ + --hash=sha256:ba7e2d33e1d092dffd0a3ff4bd1b86177594aa3c2901fd478e78e1fb2aee8ed3 \ + --hash=sha256:bafc32f01d4ab64f83fdbc29bc5bd25a920b59c751c12e06e6f4b1e379be7600 \ + --hash=sha256:bd04205b2a87087ffc2257c3ad33f11daabc053956f64ac1ec7bae299cac3f2f \ + --hash=sha256:bd26d71939de5dcf169d479fbc7fcfed21e6675bab33e7f7e9f8405f19711c71 \ + --hash=sha256:c3464d0042a03e8ef38a2b774ef23163cf3c0cdc41b8dfbf7c4aadf93e40b459 \ + --hash=sha256:c48da60a5eb481b412e5e462b81ad307fb2203178a2840a743f0a7c5fc1718c9 \ + --hash=sha256:c4fa6c23e56d48df18f534af921ec936c812743a8972ecdd5e5ff47b464fea00 \ + --hash=sha256:c902e56e60c48f5f15e55257aaa1c2678323df5f18a1b839e8d05cac1107576c \ + --hash=sha256:ca880f555c8db40942acc8a50722c33e229b6be90e598acc1a201f36487b917d \ + --hash=sha256:cb1e86bfea656b51f2e75f2cedb17fc08b552d105b814d19b595294ecbe94d8d \ + --hash=sha256:cd4d4ab9deb5846af0ac6cf1fa139cfa40291ad14d073efa8b8e20c8d1aa90bd \ + --hash=sha256:dbbd6fba677e1cbc9d6bd4ebbe3e8b3667d0295f1731489db2a971c95f0ceca0 \ + --hash=sha256:dc07376aa33b6004ea372ac9b0ba0ed3455aa2fc4e18727414142ecb46b176b8 \ + --hash=sha256:dd70ea5d7b2c5e479e04ac3a00d8bc3deca146d2b5dbfbe3d7b42ed136e19de4 \ + --hash=sha256:ddb9c4db36188a8f08c2303100a83100f26a8572803ae35eadff359bebd3d204 \ + --hash=sha256:df089639983a03070be6eabc60317aa1ffbf2c5409023b57a5fc2e4975163bc4 \ + --hash=sha256:e0b062d261fa3fac00146cf801896c8cfafe1e41332eb047aa0a36558299daa6 \ + --hash=sha256:e248510f8e2dbc71fa99f86238c9023365dbe1a4520eb40e33d73416527349f2 \ + --hash=sha256:e94021c541eb2a199b5a2ffae0ea84fb8b99863dab99a5b154b00bc7a44b5c48 \ + --hash=sha256:e98a18c22a85f321091cc8db6694af1d713a369c2d60ec611c10ccfe24ab103a \ + --hash=sha256:ea9bcaff298f5d35ef67346d474fca388c5cf6d4edab1d06b84868800f88bd36 \ + --hash=sha256:eb16d995e454b0155b166f6e6da7df4ac812d44e0f3b6dc0f344a934609fd5bc \ + --hash=sha256:ed486e57a79ccc7aebaa2ec12517d891fdc5d2fde16915e3db705b8a47570981 \ + --hash=sha256:ed7fd7bc2b86ec3161fe0cc49f5f392e6efa55c91a95397d5047820c38117660 \ + --hash=sha256:ef553d4080368006817c1a935ed619c71987cf10417a32386acc00c5418a2934 \ + --hash=sha256:f103e648aa3ebe9b8e2c1a3a92719288d8f3f41007c319ad298cdce2d0c28641 \ + --hash=sha256:fc49b6575e3cb15da3180c5a3926ec81db33b109e48530708da76614b306904b \ + --hash=sha256:fe9af350dfbdc7ed4c93a8016a8ad7b5492fc116e7197cad7cbce99b434d3fe1 + # via + # barman + # python-snappy +cryptography==43.0.3 \ + --hash=sha256:0c580952eef9bf68c4747774cde7ec1d85a6e61de97281f2dba83c7d2c806362 \ + --hash=sha256:0f996e7268af62598f2fc1204afa98a3b5712313a55c4c9d434aef49cadc91d4 \ + --hash=sha256:1ec0bcf7e17c0c5669d881b1cd38c4972fade441b27bda1051665faaa89bdcaa \ + --hash=sha256:281c945d0e28c92ca5e5930664c1cefd85efe80e5c0d2bc58dd63383fda29f83 \ + --hash=sha256:2ce6fae5bdad59577b44e4dfed356944fbf1d925269114c28be377692643b4ff \ + --hash=sha256:315b9001266a492a6ff443b61238f956b214dbec9910a081ba5b6646a055a805 \ + --hash=sha256:443c4a81bb10daed9a8f334365fe52542771f25aedaf889fd323a853ce7377d6 \ + --hash=sha256:4a02ded6cd4f0a5562a8887df8b3bd14e822a90f97ac5e544c162899bc467664 \ + --hash=sha256:53a583b6637ab4c4e3591a15bc9db855b8d9dee9a669b550f311480acab6eb08 \ + --hash=sha256:63efa177ff54aec6e1c0aefaa1a241232dcd37413835a9b674b6e3f0ae2bfd3e \ + --hash=sha256:74f57f24754fe349223792466a709f8e0c093205ff0dca557af51072ff47ab18 \ + --hash=sha256:7e1ce50266f4f70bf41a2c6dc4358afadae90e2a1e5342d3c08883df1675374f \ + --hash=sha256:81ef806b1fef6b06dcebad789f988d3b37ccaee225695cf3e07648eee0fc6b73 \ + --hash=sha256:846da004a5804145a5f441b8530b4bf35afbf7da70f82409f151695b127213d5 \ + --hash=sha256:8ac43ae87929a5982f5948ceda07001ee5e83227fd69cf55b109144938d96984 \ + --hash=sha256:9762ea51a8fc2a88b70cf2995e5675b38d93bf36bd67d91721c309df184f49bd \ + --hash=sha256:a2a431ee15799d6db9fe80c82b055bae5a752bef645bba795e8e52687c69efe3 \ + --hash=sha256:bf7a1932ac4176486eab36a19ed4c0492da5d97123f1406cf15e41b05e787d2e \ + --hash=sha256:c2e6fc39c4ab499049df3bdf567f768a723a5e8464816e8f009f121a5a9f4405 \ + --hash=sha256:cbeb489927bd7af4aa98d4b261af9a5bc025bd87f0e3547e11584be9e9427be2 \ + --hash=sha256:d03b5621a135bffecad2c73e9f4deb1a0f977b9a8ffe6f8e002bf6c9d07b918c \ + --hash=sha256:d56e96520b1020449bbace2b78b603442e7e378a9b3bd68de65c782db1507995 \ + --hash=sha256:df6b6c6d742395dd77a23ea3728ab62f98379eff8fb61be2744d4679ab678f73 \ + --hash=sha256:e1be4655c7ef6e1bbe6b5d0403526601323420bcf414598955968c9ef3eb7d16 \ + --hash=sha256:f18c716be16bc1fea8e95def49edf46b82fccaa88587a45f8dc0ff6ab5d8e0a7 \ + --hash=sha256:f46304d6f0c6ab8e52770addfa2fc41e6629495548862279641972b6215451cd \ + --hash=sha256:f7b178f11ed3664fd0e995a47ed2b5ff0a12d893e41dd0494f406d1cf555cab7 + # via + # azure-identity + # azure-storage-blob + # msal + # pyjwt +google-api-core==2.22.0 \ + --hash=sha256:26f8d76b96477db42b55fd02a33aae4a42ec8b86b98b94969b7333a2c828bf35 \ + --hash=sha256:a6652b6bd51303902494998626653671703c420f6f4c88cfd3f50ed723e9d021 + # via + # google-cloud-core + # google-cloud-storage +google-auth==2.35.0 \ + --hash=sha256:25df55f327ef021de8be50bad0dfd4a916ad0de96da86cd05661c9297723ad3f \ + --hash=sha256:f4c64ed4e01e8e8b646ef34c018f8bf3338df0c8e37d8b3bba40e7f574a3278a + # via + # google-api-core + # google-cloud-core + # google-cloud-storage +google-cloud-core==2.4.1 \ + --hash=sha256:9b7749272a812bde58fff28868d0c5e2f585b82f37e09a1f6ed2d4d10f134073 \ + --hash=sha256:a9e6a4422b9ac5c29f79a0ede9485473338e2ce78d91f2370c01e730eab22e61 + # via google-cloud-storage +google-cloud-storage==2.18.2 \ + --hash=sha256:97a4d45c368b7d401ed48c4fdfe86e1e1cb96401c9e199e419d289e2c0370166 \ + --hash=sha256:aaf7acd70cdad9f274d29332673fcab98708d0e1f4dceb5a5356aaef06af4d99 +google-crc32c==1.6.0 \ + --hash=sha256:05e2d8c9a2f853ff116db9706b4a27350587f341eda835f46db3c0a8c8ce2f24 \ + --hash=sha256:18e311c64008f1f1379158158bb3f0c8d72635b9eb4f9545f8cf990c5668e59d \ + --hash=sha256:236c87a46cdf06384f614e9092b82c05f81bd34b80248021f729396a78e55d7e \ + --hash=sha256:35834855408429cecf495cac67ccbab802de269e948e27478b1e47dfb6465e57 \ + --hash=sha256:386122eeaaa76951a8196310432c5b0ef3b53590ef4c317ec7588ec554fec5d2 \ + --hash=sha256:40b05ab32a5067525670880eb5d169529089a26fe35dce8891127aeddc1950e8 \ + --hash=sha256:48abd62ca76a2cbe034542ed1b6aee851b6f28aaca4e6551b5599b6f3ef175cc \ + --hash=sha256:50cf2a96da226dcbff8671233ecf37bf6e95de98b2a2ebadbfdf455e6d05df42 \ + --hash=sha256:51c4f54dd8c6dfeb58d1df5e4f7f97df8abf17a36626a217f169893d1d7f3e9f \ + --hash=sha256:5bcc90b34df28a4b38653c36bb5ada35671ad105c99cfe915fb5bed7ad6924aa \ + --hash=sha256:62f6d4a29fea082ac4a3c9be5e415218255cf11684ac6ef5488eea0c9132689b \ + --hash=sha256:6eceb6ad197656a1ff49ebfbbfa870678c75be4344feb35ac1edf694309413dc \ + --hash=sha256:7aec8e88a3583515f9e0957fe4f5f6d8d4997e36d0f61624e70469771584c760 \ + --hash=sha256:91ca8145b060679ec9176e6de4f89b07363d6805bd4760631ef254905503598d \ + --hash=sha256:a184243544811e4a50d345838a883733461e67578959ac59964e43cca2c791e7 \ + --hash=sha256:a9e4b426c3702f3cd23b933436487eb34e01e00327fac20c9aebb68ccf34117d \ + --hash=sha256:bb0966e1c50d0ef5bc743312cc730b533491d60585a9a08f897274e57c3f70e0 \ + --hash=sha256:bb8b3c75bd157010459b15222c3fd30577042a7060e29d42dabce449c087f2b3 \ + --hash=sha256:bd5e7d2445d1a958c266bfa5d04c39932dc54093fa391736dbfdb0f1929c1fb3 \ + --hash=sha256:c87d98c7c4a69066fd31701c4e10d178a648c2cac3452e62c6b24dc51f9fcc00 \ + --hash=sha256:d2952396dc604544ea7476b33fe87faedc24d666fb0c2d5ac971a2b9576ab871 \ + --hash=sha256:d8797406499f28b5ef791f339594b0b5fdedf54e203b5066675c406ba69d705c \ + --hash=sha256:d9e9913f7bd69e093b81da4535ce27af842e7bf371cde42d1ae9e9bd382dc0e9 \ + --hash=sha256:e2806553238cd076f0a55bddab37a532b53580e699ed8e5606d0de1f856b5205 \ + --hash=sha256:ebab974b1687509e5c973b5c4b8b146683e101e102e17a86bd196ecaa4d099fc \ + --hash=sha256:ed767bf4ba90104c1216b68111613f0d5926fb3780660ea1198fc469af410e9d \ + --hash=sha256:f7a1fc29803712f80879b0806cb83ab24ce62fc8daf0569f2204a0cfd7f68ed4 + # via + # google-cloud-storage + # google-resumable-media +google-resumable-media==2.7.2 \ + --hash=sha256:3ce7551e9fe6d99e9a126101d2536612bb73486721951e9562fee0f90c6ababa \ + --hash=sha256:5280aed4629f2b60b847b0d42f9857fd4935c11af266744df33d8074cae92fe0 + # via google-cloud-storage +googleapis-common-protos==1.65.0 \ + --hash=sha256:2972e6c496f435b92590fd54045060867f3fe9be2c82ab148fc8885035479a63 \ + --hash=sha256:334a29d07cddc3aa01dee4988f9afd9b2916ee2ff49d6b757155dc0d197852c0 + # via google-api-core +idna==3.10 \ + --hash=sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9 \ + --hash=sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3 + # via requests +isodate==0.7.2 \ + --hash=sha256:28009937d8031054830160fce6d409ed342816b543597cece116d966c6d99e15 \ + --hash=sha256:4cd1aa0f43ca76f4a6c6c0292a85f40b35ec2e43e315b59f06e6d32171a953e6 + # via azure-storage-blob +jmespath==1.0.1 \ + --hash=sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980 \ + --hash=sha256:90261b206d6defd58fdd5e85f478bf633a2901798906be2ad389150c5c60edbe + # via + # boto3 + # botocore +msal==1.31.0 \ + --hash=sha256:2c4f189cf9cc8f00c80045f66d39b7c0f3ed45873fd3d1f2af9f22db2e12ff4b \ + --hash=sha256:96bc37cff82ebe4b160d5fc0f1196f6ca8b50e274ecd0ec5bf69c438514086e7 + # via + # azure-identity + # msal-extensions +msal-extensions==1.2.0 \ + --hash=sha256:6f41b320bfd2933d631a215c91ca0dd3e67d84bd1a2f50ce917d5874ec646bef \ + --hash=sha256:cf5ba83a2113fa6dc011a254a72f1c223c88d7dfad74cc30617c4679a417704d + # via azure-identity +portalocker==2.10.1 \ + --hash=sha256:53a5984ebc86a025552264b459b46a2086e269b21823cb572f8f28ee759e45bf \ + --hash=sha256:ef1bf844e878ab08aee7e40184156e1151f228f103aa5c6bd0724cc330960f8f + # via msal-extensions +proto-plus==1.25.0 \ + --hash=sha256:c91fc4a65074ade8e458e95ef8bac34d4008daa7cce4a12d6707066fca648961 \ + --hash=sha256:fbb17f57f7bd05a68b7707e745e26528b0b3c34e378db91eef93912c54982d91 + # via google-api-core +protobuf==5.28.3 \ + --hash=sha256:0c4eec6f987338617072592b97943fdbe30d019c56126493111cf24344c1cc24 \ + --hash=sha256:135658402f71bbd49500322c0f736145731b16fc79dc8f367ab544a17eab4535 \ + --hash=sha256:27b246b3723692bf1068d5734ddaf2fccc2cdd6e0c9b47fe099244d80200593b \ + --hash=sha256:3e6101d095dfd119513cde7259aa703d16c6bbdfae2554dfe5cfdbe94e32d548 \ + --hash=sha256:3fa2de6b8b29d12c61911505d893afe7320ce7ccba4df913e2971461fa36d584 \ + --hash=sha256:64badbc49180a5e401f373f9ce7ab1d18b63f7dd4a9cdc43c92b9f0b481cef7b \ + --hash=sha256:70585a70fc2dd4818c51287ceef5bdba6387f88a578c86d47bb34669b5552c36 \ + --hash=sha256:712319fbdddb46f21abb66cd33cb9e491a5763b2febd8f228251add221981135 \ + --hash=sha256:91fba8f445723fcf400fdbe9ca796b19d3b1242cd873907979b9ed71e4afe868 \ + --hash=sha256:a3f6857551e53ce35e60b403b8a27b0295f7d6eb63d10484f12bc6879c715687 \ + --hash=sha256:cee1757663fa32a1ee673434fcf3bf24dd54763c79690201208bafec62f19eed + # via + # google-api-core + # googleapis-common-protos + # proto-plus +pyasn1==0.6.1 \ + --hash=sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629 \ + --hash=sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034 + # via + # pyasn1-modules + # rsa +pyasn1-modules==0.4.1 \ + --hash=sha256:49bfa96b45a292b711e986f222502c1c9a5e1f4e568fc30e2574a6c7d07838fd \ + --hash=sha256:c28e2dbf9c06ad61c71a075c7e0f9fd0f1b0bb2d2ad4377f240d33ac2ab60a7c + # via google-auth +pycparser==2.22 \ + --hash=sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6 \ + --hash=sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc + # via cffi +pyjwt[crypto]==2.9.0 \ + --hash=sha256:3b02fb0f44517787776cf48f2ae25d8e14f300e6d7545a4315cee571a415e850 \ + --hash=sha256:7e1e5b56cc735432a7369cbfa0efe50fa113ebecdc04ae6922deba8b84582d0c + # via + # msal + # pyjwt +python-dateutil==2.9.0.post0 \ + --hash=sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3 \ + --hash=sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427 + # via + # barman + # botocore +python-snappy==0.7.3 \ + --hash=sha256:074c0636cfcd97e7251330f428064050ac81a52c62ed884fc2ddebbb60ed7f50 \ + --hash=sha256:40216c1badfb2d38ac781ecb162a1d0ec40f8ee9747e610bcfefdfa79486cee3 +requests==2.32.3 \ + --hash=sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760 \ + --hash=sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6 + # via + # azure-core + # google-api-core + # google-cloud-storage + # msal +rsa==4.9 \ + --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \ + --hash=sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21 + # via google-auth +s3transfer==0.10.3 \ + --hash=sha256:263ed587a5803c6c708d3ce44dc4dfedaab4c1a32e8329bab818933d79ddcf5d \ + --hash=sha256:4f50ed74ab84d474ce614475e0b8d5047ff080810aac5d01ea25231cfc944b0c + # via boto3 +six==1.16.0 \ + --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \ + --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254 + # via + # azure-core + # python-dateutil +typing-extensions==4.12.2 \ + --hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \ + --hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8 + # via + # azure-core + # azure-identity + # azure-storage-blob +urllib3==2.2.3 \ + --hash=sha256:ca899ca043dcb1bafa3e262d73aa25c465bfb49e0bd9dd5d59f1d0acba2f8fac \ + --hash=sha256:e7d814a81dad81e6caf2ec9fdedb284ecc9c73076b62654547cc64ccdcae26e9 + # via + # botocore + # requests diff --git a/docker/pg-cnpg/trunk-install.sh b/docker/pg-cnpg/trunk-install.sh new file mode 100755 index 00000000..d64bd495 --- /dev/null +++ b/docker/pg-cnpg/trunk-install.sh @@ -0,0 +1,80 @@ +#!/bin/bash +trunk_install_failed_extensions=() +need_load_shared_preload_libraries_extensions=() +version_not_found_extensions=() +file='/tmp/extensions.txt' +extension_count=$(<$file wc -l) +lines=$(cat $file) +for line in $lines +do + output=$(trunk install $line 2>&1) + + if [ $? -ne 0 ]; then + if [[ $output == *"Failed to find an archive for"* || $output == *"Failed to fetch Trunk archive from"* ]]; then + version_not_found_extensions+=("$line") + else + echo "trunk install command failed" + trunk_install_failed_extensions+=("$line") + fi + fi + echo $output + printf "\n\n" +done +IFS=$'\n' extensions=(`psql postgres://postgres:postgres@localhost:5432 -tA postgres -c 'select name from pg_available_extensions;'`) +for ext in "${extensions[@]}" +do + # drop schema columnar if ext name is columnar + if [ "$ext" == "columnar" ]; then + psql postgres://postgres:postgres@localhost:5432 -c "drop extension if exists citus_columnar cascade;" + fi + # drop type semver if ext name is semver + if [ "$ext" == "semver" ]; then + psql postgres://postgres:postgres@localhost:5432 -c "drop extension if exists pg_text_semver cascade;" + fi + # if extension name is meta_triggers, create extension meta first and create extension meta_triggers + if [ "$ext" == "meta_triggers" ]; then + psql postgres://postgres:postgres@localhost:5432 -c "create extension if not exists hstore cascade;" + psql postgres://postgres:postgres@localhost:5432 -c "create extension if not exists meta cascade;" + psql postgres://postgres:postgres@localhost:5432 -c "create extension if not exists meta_triggers cascade;" + fi + output=$(psql postgres://postgres:postgres@localhost:5432 -c "create extension if not exists \"$ext\" cascade;" 2>&1) + if [ $? -ne 0 ]; then + if [[ $output == *"shared_preload_libraries"* ]]; then + need_load_shared_preload_libraries_extensions+=("$ext") + elif [[ $output == *"already exists"* ]]; then + echo "extension \"$ext\" already exists" + else + echo "CREATE EXTENSION command failed" + failed_extensions+=("$ext") + fi + fi + echo $output + printf "\n\n" +done +available_extensions_count=${#extensions[@]} +failure_count=${#failed_extensions[@]} +need_load_shared_preload_libraries_count=${#need_load_shared_preload_libraries_extensions[@]} +not_found_count=${#version_not_found_extensions[@]} +success=$(($available_extensions_count-$failure_count)) +success_percent=$(awk "BEGIN { pc=100*${success}/${extension_count}; i=int(pc); print (pc-i<0.5)?i:i+1 }") +failure_percent=$(awk "BEGIN { pc=100*${failure_count}/${extension_count}; i=int(pc); print (pc-i<0.5)?i:i+1 }") + +printf "\n\n***TRUNK INSTALL EXTENSIONS THAT VERSION NOT FOUND RATE***\n" +echo "$not_found_count / $extension_count" +printf "***CREATE EXTENSIONS SUCCESS RATE***\n" +echo "$success / $extension_count ($success_percent%)" +printf "\n\n***CREATE EXTENSION FAILURE RATE***\n" +echo "$failure_count / $extension_count ($failure_percent%)" +printf "\n\n***CREATE EXTENSIONS THAT NEED TO BE LOADED IN shared_preload_libraries***\n" +echo "$need_load_shared_preload_libraries_count / $extension_count" +printf "\n\n***NEED TO LOAD shared_preload_libraries EXTENSIONS***\n" +for need in "${need_load_shared_preload_libraries_extensions[@]}" +do + echo $need +done + +printf "\n\n***FAILED EXTENSIONS***\n" +for failed in "${failed_extensions[@]}" +do + echo $failed +done diff --git a/docker/pg-slim/Dockerfile b/docker/pg-slim/Dockerfile new file mode 100644 index 00000000..06049c9e --- /dev/null +++ b/docker/pg-slim/Dockerfile @@ -0,0 +1,137 @@ +FROM ubuntu:22.04 + +# 14, 15, 16, 17 +ARG PG_MAJOR +ARG ALTDIR=/var/lib/postgresql/data/tensorchord +ENV DEBIAN_FRONTEND=noninteractive \ + TZ=Etc/UTC \ + PGDATA=/var/lib/postgresql/data + +# Get latest package updates +RUN set -eux; \ + apt-get update; \ + apt-get upgrade -y + +# explicitly set user/group IDs +RUN set -eux; \ + groupadd -r postgres --gid=999; \ +# https://salsa.debian.org/postgresql/postgresql-common/blob/997d842ee744687d99a2b2d95c1083a2615c79e8/debian/postgresql-common.postinst#L32-35 + useradd -r -g postgres --uid=999 --home-dir=/var/lib/postgresql --shell=/bin/bash postgres; \ +# also create the postgres user's home directory with appropriate permissions +# see https://github.com/docker-library/postgres/issues/274 + install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql + +RUN set -eux; \ + apt-get update; \ + apt-get install -y --no-install-recommends \ + locales curl ca-certificates gnupg lsb-release lbzip2 git cmake \ + ; \ + rm -rf /var/lib/apt/lists/*; \ + localedef -i en_US -c -f UTF-8 -A /usr/share/locale/locale.alias en_US.UTF-8 +ENV LANG en_US.utf8 + +RUN mkdir /docker-entrypoint-initdb.d + +RUN git clone https://github.com/postgres/postgres.git && \ + cd postgres && \ + PG_RELEASE=$(git tag | grep '^REL_'${PG_MAJOR}'_' | grep -vE 'RC|BETA' | sort -V | tail -1 | sed 's/REL_//') && \ + echo $PG_RELEASE && \ + git checkout REL_$PG_RELEASE + +RUN set -eux; \ + apt-get update && apt-get install -y \ + libreadline-dev \ + zlib1g-dev \ + libpq-dev \ + build-essential \ + python3-dev \ + tcl-dev \ + libxslt1-dev \ + libperl-dev \ + libpam0g-dev \ + libreadline-dev \ + libssl-dev \ + xz-utils \ + libnss-wrapper \ + llvm \ + clang \ + icu-devtools \ + pkg-config \ + libgss-dev \ + libkrb5-dev \ + uuid-dev \ + gettext \ + liblz4-dev \ + libsystemd-dev \ + libselinux1-dev \ + libzstd-dev \ + vim \ + flex \ + bison; \ + apt-get autoremove -y; \ + apt-get clean -y; \ + rm -rf /var/lib/apt/lists/* + +WORKDIR postgres + +ENV CFLAGS "-g -O2 -fstack-protector-strong -Wformat -Werror=format-security -fno-omit-frame-pointer" +ENV LDFLAGS "-Wl,-z,relro -Wl,-z,now" +RUN ./configure --prefix=/usr/lib/postgresql/${PG_MAJOR} \ + --datarootdir=${ALTDIR} \ + --libdir=${ALTDIR}/${PG_MAJOR}/lib \ + --with-perl \ + --with-python \ + --with-tcl \ + --with-pam \ + --with-libxml \ + --with-libxslt \ + --with-openssl \ + --enable-nls \ + --enable-thread-safety \ + --enable-debug \ + --disable-rpath \ + --with-uuid=e2fs \ + --with-gnu-ld \ + --with-gssapi \ + --with-pgport=5432 \ + --with-system-tzdata=/usr/share/zoneinfo \ + --with-icu \ + --with-llvm \ + --with-lz4 \ + --with-zstd \ + --with-systemd \ + --with-selinux + +RUN make -j$(nproc) +RUN make install + +WORKDIR / +RUN rm -rf /postgres + +RUN mkdir -p /var/run/postgresql && chmod 775 /var/run/postgresql +RUN mkdir -p /usr/share/postgresql/${PG_MAJOR}/extension && chmod 775 /usr/share/postgresql/${PG_MAJOR}/extension + +COPY --from=tianon/gosu /gosu /usr/local/bin/ + +# make the sample config easier to munge (and "correct by default") +RUN set -eux; \ + echo "unix_socket_directories = '/var/run/postgresql'" >> ${ALTDIR}/postgresql.conf.sample; \ + sed -ri "s!^#?(listen_addresses)\s*=\s*\S+.*!\1 = '*'!" ${ALTDIR}/postgresql.conf.sample; \ + grep -F "listen_addresses = '*'" ${ALTDIR}/postgresql.conf.sample + +RUN install --verbose --directory --owner postgres --group postgres --mode 3777 /var/run/postgresql + +# this 1777 will be replaced by 0700 at runtime (allows semi-arbitrary "--user" values) +RUN install --verbose --directory --owner postgres --group postgres --mode 1777 ${PGDATA} + +ENV PATH $PATH:/usr/lib/postgresql/$PG_MAJOR/bin:/usr/local/bin +COPY docker-entrypoint.sh docker-ensure-initdb.sh /usr/local/bin/ +RUN ln -sT docker-ensure-initdb.sh /usr/local/bin/docker-enforce-initdb.sh +ENTRYPOINT ["docker-entrypoint.sh"] + +# Remove pre-installed pg_config +RUN rm /usr/bin/pg_config + +STOPSIGNAL SIGINT +EXPOSE 5432 +CMD ["postgres"] \ No newline at end of file diff --git a/docker/pg-slim/docker-ensure-initdb.sh b/docker/pg-slim/docker-ensure-initdb.sh new file mode 100755 index 00000000..ae1f6b6b --- /dev/null +++ b/docker/pg-slim/docker-ensure-initdb.sh @@ -0,0 +1,71 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +# +# This script is intended for three main use cases: +# +# 1. (most importantly) as an example of how to use "docker-entrypoint.sh" to extend/reuse the initialization behavior +# +# 2. ("docker-ensure-initdb.sh") as a Kubernetes "init container" to ensure the provided database directory is initialized; see also "startup probes" for an alternative solution +# (no-op if database is already initialized) +# +# 3. ("docker-enforce-initdb.sh") as part of CI to ensure the database is fully initialized before use +# (error if database is already initialized) +# + +source /usr/local/bin/docker-entrypoint.sh + +# arguments to this script are assumed to be arguments to the "postgres" server (same as "docker-entrypoint.sh"), and most "docker-entrypoint.sh" functions assume "postgres" is the first argument (see "_main" over there) +if [ "$#" -eq 0 ] || [ "$1" != 'postgres' ]; then + set -- postgres "$@" +fi + +# see also "_main" in "docker-entrypoint.sh" + +docker_setup_env +# setup data directories and permissions (when run as root) +docker_create_db_directories +if [ "$(id -u)" = '0' ]; then + # then restart script as postgres user + exec gosu postgres "$BASH_SOURCE" "$@" +fi + +# only run initialization on an empty data directory +if [ -z "$DATABASE_ALREADY_EXISTS" ]; then + docker_verify_minimum_env + + # check dir permissions to reduce likelihood of half-initialized database + ls /docker-entrypoint-initdb.d/ > /dev/null + + docker_init_database_dir + pg_setup_hba_conf "$@" + + # PGPASSWORD is required for psql when authentication is required for 'local' connections via pg_hba.conf and is otherwise harmless + # e.g. when '--auth=md5' or '--auth-local=md5' is used in POSTGRES_INITDB_ARGS + export PGPASSWORD="${PGPASSWORD:-$POSTGRES_PASSWORD}" + docker_temp_server_start "$@" + + docker_setup_db + docker_process_init_files /docker-entrypoint-initdb.d/* + + docker_temp_server_stop + unset PGPASSWORD +else + self="$(basename "$0")" + case "$self" in + docker-ensure-initdb.sh) + echo >&2 "$self: note: database already initialized in '$PGDATA'!" + exit 0 + ;; + + docker-enforce-initdb.sh) + echo >&2 "$self: error: (unexpected) database found in '$PGDATA'!" + exit 1 + ;; + + *) + echo >&2 "$self: error: unknown file name: $self" + exit 99 + ;; + esac +fi diff --git a/docker/pg-slim/docker-entrypoint.sh b/docker/pg-slim/docker-entrypoint.sh new file mode 100755 index 00000000..6f59993e --- /dev/null +++ b/docker/pg-slim/docker-entrypoint.sh @@ -0,0 +1,356 @@ +#!/usr/bin/env bash +set -Eeo pipefail +# TODO swap to -Eeuo pipefail above (after handling all potentially-unset variables) + +# usage: file_env VAR [DEFAULT] +# ie: file_env 'XYZ_DB_PASSWORD' 'example' +# (will allow for "$XYZ_DB_PASSWORD_FILE" to fill in the value of +# "$XYZ_DB_PASSWORD" from a file, especially for Docker's secrets feature) +file_env() { + local var="$1" + local fileVar="${var}_FILE" + local def="${2:-}" + if [ "${!var:-}" ] && [ "${!fileVar:-}" ]; then + printf >&2 'error: both %s and %s are set (but are exclusive)\n' "$var" "$fileVar" + exit 1 + fi + local val="$def" + if [ "${!var:-}" ]; then + val="${!var}" + elif [ "${!fileVar:-}" ]; then + val="$(< "${!fileVar}")" + fi + export "$var"="$val" + unset "$fileVar" +} + +# check to see if this file is being run or sourced from another script +_is_sourced() { + # https://unix.stackexchange.com/a/215279 + [ "${#FUNCNAME[@]}" -ge 2 ] \ + && [ "${FUNCNAME[0]}" = '_is_sourced' ] \ + && [ "${FUNCNAME[1]}" = 'source' ] +} + +# used to create initial postgres directories and if run as root, ensure ownership to the "postgres" user +docker_create_db_directories() { + local user; user="$(id -u)" + + mkdir -p "$PGDATA" + # ignore failure since there are cases where we can't chmod (and PostgreSQL might fail later anyhow - it's picky about permissions of this directory) + chmod 00700 "$PGDATA" || : + + # ignore failure since it will be fine when using the image provided directory; see also https://github.com/docker-library/postgres/pull/289 + mkdir -p /var/run/postgresql || : + chmod 03775 /var/run/postgresql || : + + # Create the transaction log directory before initdb is run so the directory is owned by the correct user + if [ -n "${POSTGRES_INITDB_WALDIR:-}" ]; then + mkdir -p "$POSTGRES_INITDB_WALDIR" + if [ "$user" = '0' ]; then + find "$POSTGRES_INITDB_WALDIR" \! -user postgres -exec chown postgres '{}' + + fi + chmod 700 "$POSTGRES_INITDB_WALDIR" + fi + + # allow the container to be started with `--user` + if [ "$user" = '0' ]; then + find "$PGDATA" \! -user postgres -exec chown postgres '{}' + + find /var/run/postgresql \! -user postgres -exec chown postgres '{}' + + fi +} + +# initialize empty PGDATA directory with new database via 'initdb' +# arguments to `initdb` can be passed via POSTGRES_INITDB_ARGS or as arguments to this function +# `initdb` automatically creates the "postgres", "template0", and "template1" dbnames +# this is also where the database user is created, specified by `POSTGRES_USER` env +docker_init_database_dir() { + # "initdb" is particular about the current user existing in "/etc/passwd", so we use "nss_wrapper" to fake that if necessary + # see https://github.com/docker-library/postgres/pull/253, https://github.com/docker-library/postgres/issues/359, https://cwrap.org/nss_wrapper.html + local uid; uid="$(id -u)" + if ! getent passwd "$uid" &> /dev/null; then + # see if we can find a suitable "libnss_wrapper.so" (https://salsa.debian.org/sssd-team/nss-wrapper/-/commit/b9925a653a54e24d09d9b498a2d913729f7abb15) + local wrapper + for wrapper in {/usr,}/lib{/*,}/libnss_wrapper.so; do + if [ -s "$wrapper" ]; then + NSS_WRAPPER_PASSWD="$(mktemp)" + NSS_WRAPPER_GROUP="$(mktemp)" + export LD_PRELOAD="$wrapper" NSS_WRAPPER_PASSWD NSS_WRAPPER_GROUP + local gid; gid="$(id -g)" + printf 'postgres:x:%s:%s:PostgreSQL:%s:/bin/false\n' "$uid" "$gid" "$PGDATA" > "$NSS_WRAPPER_PASSWD" + printf 'postgres:x:%s:\n' "$gid" > "$NSS_WRAPPER_GROUP" + break + fi + done + fi + + if [ -n "${POSTGRES_INITDB_WALDIR:-}" ]; then + set -- --waldir "$POSTGRES_INITDB_WALDIR" "$@" + fi + + # --pwfile refuses to handle a properly-empty file (hence the "\n"): https://github.com/docker-library/postgres/issues/1025 + eval 'initdb --username="$POSTGRES_USER" --pwfile=<(printf "%s\n" "$POSTGRES_PASSWORD") '"$POSTGRES_INITDB_ARGS"' "$@"' + + # unset/cleanup "nss_wrapper" bits + if [[ "${LD_PRELOAD:-}" == */libnss_wrapper.so ]]; then + rm -f "$NSS_WRAPPER_PASSWD" "$NSS_WRAPPER_GROUP" + unset LD_PRELOAD NSS_WRAPPER_PASSWD NSS_WRAPPER_GROUP + fi +} + +# print large warning if POSTGRES_PASSWORD is long +# error if both POSTGRES_PASSWORD is empty and POSTGRES_HOST_AUTH_METHOD is not 'trust' +# print large warning if POSTGRES_HOST_AUTH_METHOD is set to 'trust' +# assumes database is not set up, ie: [ -z "$DATABASE_ALREADY_EXISTS" ] +docker_verify_minimum_env() { + case "${PG_MAJOR:-}" in + 12 | 13) # https://github.com/postgres/postgres/commit/67a472d71c98c3d2fa322a1b4013080b20720b98 + # check password first so we can output the warning before postgres + # messes it up + if [ "${#POSTGRES_PASSWORD}" -ge 100 ]; then + cat >&2 <<-'EOWARN' + + WARNING: The supplied POSTGRES_PASSWORD is 100+ characters. + + This will not work if used via PGPASSWORD with "psql". + + https://www.postgresql.org/message-id/flat/E1Rqxp2-0004Qt-PL%40wrigleys.postgresql.org (BUG #6412) + https://github.com/docker-library/postgres/issues/507 + + EOWARN + fi + ;; + esac + if [ -z "$POSTGRES_PASSWORD" ] && [ 'trust' != "$POSTGRES_HOST_AUTH_METHOD" ]; then + # The - option suppresses leading tabs but *not* spaces. :) + cat >&2 <<-'EOE' + Error: Database is uninitialized and superuser password is not specified. + You must specify POSTGRES_PASSWORD to a non-empty value for the + superuser. For example, "-e POSTGRES_PASSWORD=password" on "docker run". + + You may also use "POSTGRES_HOST_AUTH_METHOD=trust" to allow all + connections without a password. This is *not* recommended. + + See PostgreSQL documentation about "trust": + https://www.postgresql.org/docs/current/auth-trust.html + EOE + exit 1 + fi + if [ 'trust' = "$POSTGRES_HOST_AUTH_METHOD" ]; then + cat >&2 <<-'EOWARN' + ******************************************************************************** + WARNING: POSTGRES_HOST_AUTH_METHOD has been set to "trust". This will allow + anyone with access to the Postgres port to access your database without + a password, even if POSTGRES_PASSWORD is set. See PostgreSQL + documentation about "trust": + https://www.postgresql.org/docs/current/auth-trust.html + In Docker's default configuration, this is effectively any other + container on the same system. + + It is not recommended to use POSTGRES_HOST_AUTH_METHOD=trust. Replace + it with "-e POSTGRES_PASSWORD=password" instead to set a password in + "docker run". + ******************************************************************************** + EOWARN + fi +} + +# usage: docker_process_init_files [file [file [...]]] +# ie: docker_process_init_files /always-initdb.d/* +# process initializer files, based on file extensions and permissions +docker_process_init_files() { + # psql here for backwards compatibility "${psql[@]}" + psql=( docker_process_sql ) + + printf '\n' + local f + for f; do + case "$f" in + *.sh) + # https://github.com/docker-library/postgres/issues/450#issuecomment-393167936 + # https://github.com/docker-library/postgres/pull/452 + if [ -x "$f" ]; then + printf '%s: running %s\n' "$0" "$f" + "$f" + else + printf '%s: sourcing %s\n' "$0" "$f" + . "$f" + fi + ;; + *.sql) printf '%s: running %s\n' "$0" "$f"; docker_process_sql -f "$f"; printf '\n' ;; + *.sql.gz) printf '%s: running %s\n' "$0" "$f"; gunzip -c "$f" | docker_process_sql; printf '\n' ;; + *.sql.xz) printf '%s: running %s\n' "$0" "$f"; xzcat "$f" | docker_process_sql; printf '\n' ;; + *.sql.zst) printf '%s: running %s\n' "$0" "$f"; zstd -dc "$f" | docker_process_sql; printf '\n' ;; + *) printf '%s: ignoring %s\n' "$0" "$f" ;; + esac + printf '\n' + done +} + +# Execute sql script, passed via stdin (or -f flag of pqsl) +# usage: docker_process_sql [psql-cli-args] +# ie: docker_process_sql --dbname=mydb <<<'INSERT ...' +# ie: docker_process_sql -f my-file.sql +# ie: docker_process_sql > "$PGDATA/pg_hba.conf" +} + +# start socket-only postgresql server for setting up or running scripts +# all arguments will be passed along as arguments to `postgres` (via pg_ctl) +docker_temp_server_start() { + if [ "$1" = 'postgres' ]; then + shift + fi + + # internal start of server in order to allow setup using psql client + # does not listen on external TCP/IP and waits until start finishes + set -- "$@" -c listen_addresses='' -p "${PGPORT:-5432}" + + PGUSER="${PGUSER:-$POSTGRES_USER}" \ + pg_ctl -D "$PGDATA" \ + -o "$(printf '%q ' "$@")" \ + -w start +} + +# stop postgresql server after done setting up user and running scripts +docker_temp_server_stop() { + PGUSER="${PGUSER:-postgres}" \ + pg_ctl -D "$PGDATA" -m fast -w stop +} + +# check arguments for an option that would cause postgres to stop +# return true if there is one +_pg_want_help() { + local arg + for arg; do + case "$arg" in + # postgres --help | grep 'then exit' + # leaving out -C on purpose since it always fails and is unhelpful: + # postgres: could not access the server configuration file "/var/lib/postgresql/data/postgresql.conf": No such file or directory + -'?'|--help|--describe-config|-V|--version) + return 0 + ;; + esac + done + return 1 +} + +_main() { + # if first arg looks like a flag, assume we want to run postgres server + if [ "${1:0:1}" = '-' ]; then + set -- postgres "$@" + fi + + if [ "$1" = 'postgres' ] && ! _pg_want_help "$@"; then + docker_setup_env + # setup data directories and permissions (when run as root) + docker_create_db_directories + if [ "$(id -u)" = '0' ]; then + # then restart script as postgres user + exec gosu postgres "$BASH_SOURCE" "$@" + fi + + # only run initialization on an empty data directory + if [ -z "$DATABASE_ALREADY_EXISTS" ]; then + docker_verify_minimum_env + + # check dir permissions to reduce likelihood of half-initialized database + ls /docker-entrypoint-initdb.d/ > /dev/null + + docker_init_database_dir + pg_setup_hba_conf "$@" + + # PGPASSWORD is required for psql when authentication is required for 'local' connections via pg_hba.conf and is otherwise harmless + # e.g. when '--auth=md5' or '--auth-local=md5' is used in POSTGRES_INITDB_ARGS + export PGPASSWORD="${PGPASSWORD:-$POSTGRES_PASSWORD}" + docker_temp_server_start "$@" + + docker_setup_db + docker_process_init_files /docker-entrypoint-initdb.d/* + + docker_temp_server_stop + unset PGPASSWORD + + cat <<-'EOM' + + PostgreSQL init process complete; ready for start up. + + EOM + else + cat <<-'EOM' + + PostgreSQL Database directory appears to contain a database; Skipping initialization + + EOM + fi + fi + + exec "$@" +} + +if ! _is_sourced; then + _main "$@" +fi From e74ce9c0eb61b6052f6a661d94f241b5117321c3 Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Sat, 7 Dec 2024 17:59:58 +0800 Subject: [PATCH 074/324] chore(README): Add some benchmark data (#126) Ref #79 Preview: https://github.com/gaocegege/VectorChord/tree/refine --------- Signed-off-by: Ce Gao --- README.md | 38 ++++++++++++++++++++++---------------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index ff5e70fc..a1e9b255 100644 --- a/README.md +++ b/README.md @@ -1,29 +1,36 @@

VectorChord

-

Host 100M 768-dim vector (250GB+) on an AWS i4i.xlarge machine ($250/month) with 4 vCPUs and 32GB of RAM using VectorChord.

+

Effortlessly host 100 million 768-dimensional vectors (250GB+) on an AWS i4i.xlarge instance ($250/month), featuring 4 vCPUs and 32GB of RAM with VectorChord.

discord invitation link -trackgit-views - -

Prior release:

- +Twitter +Docker pulls +

Prior release: Previous Docker pulls

-VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search. It serves as the successor to the [pgvecto.rs](https://github.com/tensorchord/pgvecto.rs) project. +VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search, and serves as the successor to [pgvecto.rs](https://github.com/tensorchord/pgvecto.rs). -VectorChord incorporates the insights and lessons learned from pgvecto.rs, providing faster query speeds, more flexible build options, significantly enhanced index build performance, and greater stability. It is entirely based on Postgres storage, allowing for physical replication and WAL incremental backups for index. This also enables effective ssd usage to reduce memory requirements and can easily handle vectors ranging from hundreds of millions to multi billions of entries. +With VectorChord, you can store 400,000 vectors for just $1, enabling significant savings: 6x more vectors compared to Pinecone's optimized storage and 26x more than pgvector/pgvecto.rs for the same price[^1]. For further insights, check out our [launch blog post](https://blog.pgvecto.rs/vectorchord-store-400k-vectors-for-1-in-postgresql). + +[^1]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/#/) with 768-dimensional vectors and 95% recall. ## Features -- **Blazing-Fast Queries**: Achieve up to 3x faster queries compared to pgvector's HNSW, maintaining the same recall level. -- **High-throughput Update**: Achieve 16x faster insert throughput compared to pgvector's HNSW. -- **External Index Precomputation**: Built on IVF, VectorChord enables KMeans clustering to be performed externally (e.g., on a GPU) and seamlessly imported into the database. -- **Lightning-Fast Index Building**: Build index up to 20x faster than pgvector hnsw with precomputed centroids. (1.5 min for 1M 960-dim vectors) -- **Advanced Quantization**: Uses cutting-edge RaBitQ to compress float vectors into compact bit representations with autonomous reranking. -- **Disk-Friendly Performance**: Query laion-100M 768-dim vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%. -- **Seamless Compatibility**: Compatible with pgvector data types while delivering faster indexing and querying. -- **Simple Configuration**: No need to tweak quantization or rerank parameters — best defaults are provided out of the box. + +VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: + +**⚡ Enhanced Performance**: Delivering optimized operations with up to 5x faster queries, 16x higher insert throughput, and 16x quicker[^3] index building compared to pgvector's HNSW implementation. + +[^3]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/#/) with 768-dimensional vectors. Please checkout our [blog post](https://blog.pgvecto.rs/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. + +**💰 Affordable Vector Search**: Query 100M 768-dimensional vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%, helping you keep infrastructure costs down while maintaining high search quality. + +**🔌 Seamless Integration**: Fully compatible with pgvector data types and syntax while providing optimal defaults out of the box - no manual parameter tuning needed. Just drop in VectorChord for enhanced performance. + +**🔧 External Index Build**: Leverage IVF to build indexes externally (e.g., on GPU) for faster KMeans clustering, combined with RaBitQ[^2] compression to efficiently store vectors while maintaining search quality through autonomous reranking. + +[^2]: Gao, Jianyang, and Cheng Long. "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search." Proceedings of the ACM on Management of Data 2.3 (2024): 1-27. ## Quick Start For new users, we recommend using the Docker image to get started quickly. @@ -201,4 +208,3 @@ cargo pgrx install --release --sudo # To install the extension into the system p ## License This project is licensed under the [GNU Affero General Public License v3.0](./LICENSE) and as commercial software. For commercial licensing, please contact us at support@tensorchord.ai. - From 67356198ca51840a545aca47524a97823f5468e8 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 9 Dec 2024 09:31:25 +0800 Subject: [PATCH 075/324] fix: use stable toolchain by default (#128) Signed-off-by: usamoi --- docker/pgrx.Dockerfile | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile index 02941aa5..3a01627b 100644 --- a/docker/pgrx.Dockerfile +++ b/docker/pgrx.Dockerfile @@ -43,15 +43,17 @@ RUN useradd -u 1000 -U -m ubuntu RUN chown -R ubuntu:ubuntu /usr/share/postgresql/ /usr/lib/postgresql/ USER ubuntu ENV PATH="$PATH:/home/ubuntu/.cargo/bin" -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- --default-toolchain=none -y +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y +RUN rustup target add x86_64-unknown-linux-gnu aarch64-unknown-linux-gnu WORKDIR /workspace COPY rust-toolchain.toml /workspace/rust-toolchain.toml +RUN rustup target add x86_64-unknown-linux-gnu aarch64-unknown-linux-gnu + # ref: https://github.com/pgcentralfoundation/pgrx/blob/develop/docs/src/extension/build/cross-compile.md RUN set -ex; \ echo 'target.aarch64-unknown-linux-gnu.linker = "aarch64-linux-gnu-gcc"' >> ~/.cargo/config.toml; \ echo 'target.aarch64-unknown-linux-gnu.runner = ["qemu-aarch64-static", "-L", "/usr/aarch64-linux-gnu"]' >> ~/.cargo/config.toml -RUN rustup target add x86_64-unknown-linux-gnu aarch64-unknown-linux-gnu RUN cargo install cargo-pgrx --locked --version=${PGRX_VERSION} From 340a56701bb704940a210e6ffaa75005be60e108 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 10 Dec 2024 14:24:48 +0800 Subject: [PATCH 076/324] feat: scalar8 & indexing on halfvec (#131) closes #91 (indexing on halfvec) closes #118 (scalar8) --------- Signed-off-by: usamoi --- Cargo.lock | 116 +-------- Cargo.toml | 13 +- src/datatype/binary_scalar8.rs | 75 ++++++ src/datatype/functions_scalar8.rs | 26 ++ src/datatype/memory_pgvector_halfvec.rs | 204 +++++++++++++++ src/datatype/memory_pgvector_vector.rs | 56 ++-- src/datatype/memory_scalar8.rs | 215 ++++++++++++++++ src/datatype/mod.rs | 7 + src/datatype/operators_pgvector_halfvec.rs | 75 ++++++ src/datatype/operators_pgvector_vector.rs | 18 +- src/datatype/operators_scalar8.rs | 106 ++++++++ src/datatype/text_scalar8.rs | 138 ++++++++++ src/datatype/typmod.rs | 30 ++- src/lib.rs | 6 +- src/projection.rs | 2 +- src/sql/bootstrap.sql | 3 + src/sql/finalize.sql | 118 +++++++++ src/types/mod.rs | 1 + src/types/scalar8.rs | 286 +++++++++++++++++++++ src/utils/infinite_byte_chunks.rs | 20 ++ src/utils/k_means.rs | 14 +- src/utils/mod.rs | 1 + src/vchordrq/algorithm/build.rs | 32 +-- src/vchordrq/algorithm/insert.rs | 47 ++-- src/vchordrq/algorithm/prewarm.rs | 6 +- src/vchordrq/algorithm/rabitq.rs | 44 +--- src/vchordrq/algorithm/scan.rs | 44 ++-- src/vchordrq/algorithm/tuples.rs | 237 ++++++++++++++++- src/vchordrq/algorithm/vacuum.rs | 10 +- src/vchordrq/algorithm/vectors.rs | 57 ++-- src/vchordrq/index/am.rs | 263 ++++++++++++------- src/vchordrq/index/am_options.rs | 27 +- src/vchordrq/index/am_scan.rs | 74 ++++-- src/vchordrq/index/functions.rs | 10 +- src/vchordrq/index/opclass.rs | 15 ++ src/vchordrq/types.rs | 47 ++++ src/vchordrqfscan/algorithm/build.rs | 6 +- src/vchordrqfscan/algorithm/insert.rs | 2 +- src/vchordrqfscan/algorithm/rabitq.rs | 16 +- src/vchordrqfscan/algorithm/scan.rs | 2 +- src/vchordrqfscan/index/am.rs | 20 +- src/vchordrqfscan/index/am_options.rs | 12 +- src/vchordrqfscan/index/am_scan.rs | 11 +- src/vchordrqfscan/types.rs | 41 +++ tests/logic/distance.slt | 59 +++++ 45 files changed, 2137 insertions(+), 475 deletions(-) create mode 100644 src/datatype/binary_scalar8.rs create mode 100644 src/datatype/functions_scalar8.rs create mode 100644 src/datatype/memory_pgvector_halfvec.rs create mode 100644 src/datatype/memory_scalar8.rs create mode 100644 src/datatype/operators_pgvector_halfvec.rs create mode 100644 src/datatype/operators_scalar8.rs create mode 100644 src/datatype/text_scalar8.rs create mode 100644 src/types/mod.rs create mode 100644 src/types/scalar8.rs create mode 100644 src/utils/infinite_byte_chunks.rs create mode 100644 tests/logic/distance.slt diff --git a/Cargo.lock b/Cargo.lock index 5fe8ddb5..235b7822 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "ahash" @@ -66,15 +66,13 @@ checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" [[package]] name = "base" version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=c911e8a476effaf05cd4b1a037826b100177bdad#c911e8a476effaf05cd4b1a037826b100177bdad" dependencies = [ "base_macros", "detect", "half 2.4.1", "libc", - "log", "rand", - "rayon", "serde", "thiserror", "toml", @@ -84,7 +82,7 @@ dependencies = [ [[package]] name = "base_macros" version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=c911e8a476effaf05cd4b1a037826b100177bdad#c911e8a476effaf05cd4b1a037826b100177bdad" dependencies = [ "proc-macro2", "quote", @@ -152,9 +150,9 @@ dependencies = [ [[package]] name = "bytemuck" -version = "1.19.0" +version = "1.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8334215b81e418a0a7bdb8ef0849474f40bb10c8b71f1c4ed315cff49f32494d" +checksum = "8b37c88a63ffd85d15b406896cc343916d7cf57838a847b3a6f2ca5d39a5695a" [[package]] name = "byteorder" @@ -223,20 +221,6 @@ dependencies = [ "libloading", ] -[[package]] -name = "common" -version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" -dependencies = [ - "base", - "log", - "memmap2", - "rand", - "rustix", - "serde", - "serde_json", -] - [[package]] name = "convert_case" version = "0.6.0" @@ -315,7 +299,7 @@ dependencies = [ [[package]] name = "detect" version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=c911e8a476effaf05cd4b1a037826b100177bdad#c911e8a476effaf05cd4b1a037826b100177bdad" dependencies = [ "detect_macros", ] @@ -323,7 +307,7 @@ dependencies = [ [[package]] name = "detect_macros" version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=c911e8a476effaf05cd4b1a037826b100177bdad#c911e8a476effaf05cd4b1a037826b100177bdad" dependencies = [ "proc-macro2", "quote", @@ -362,16 +346,6 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" -[[package]] -name = "errno" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "534c5cf6194dfab3db3242765c03bbe257cf92f22b38f6bc0c58d59108a820ba" -dependencies = [ - "libc", - "windows-sys 0.52.0", -] - [[package]] name = "eyre" version = "0.6.12" @@ -435,13 +409,13 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403" [[package]] name = "half" version = "2.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +source = "git+https://github.com/tensorchord/half-rs.git#5b7fedc636c0eb1624763a40840c5cbf54cffd02" dependencies = [ "cfg-if", "crunchy", "rand", "rand_distr", + "rkyv", "serde", ] @@ -558,18 +532,6 @@ version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" -[[package]] -name = "k_means" -version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" -dependencies = [ - "base", - "common", - "half 2.4.1", - "rand", - "smawk", -] - [[package]] name = "libc" version = "0.2.161" @@ -592,12 +554,6 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" -[[package]] -name = "linux-raw-sys" -version = "0.4.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" - [[package]] name = "log" version = "0.4.22" @@ -620,15 +576,6 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" -[[package]] -name = "memmap2" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f" -dependencies = [ - "libc", -] - [[package]] name = "minimal-lexical" version = "0.2.1" @@ -955,24 +902,6 @@ dependencies = [ "syn 1.0.109", ] -[[package]] -name = "quantization" -version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" -dependencies = [ - "base", - "common", - "detect", - "k_means", - "log", - "nalgebra", - "rand", - "rand_chacha", - "rand_distr", - "serde", - "serde_json", -] - [[package]] name = "quote" version = "1.0.37" @@ -1136,19 +1065,6 @@ dependencies = [ "semver", ] -[[package]] -name = "rustix" -version = "0.38.40" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99e4ea3e1cdc4b559b8e5650f9c8e5998e3e5c1343b4eaf034565f32318d63c0" -dependencies = [ - "bitflags", - "errno", - "libc", - "linux-raw-sys", - "windows-sys 0.52.0", -] - [[package]] name = "ryu" version = "1.0.18" @@ -1273,12 +1189,6 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" -[[package]] -name = "smawk" -version = "0.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7c388c1b5e93756d0c740965c41e8822f866621d41acbdf6336a6a168f8840c" - [[package]] name = "sptr" version = "0.3.2" @@ -1516,17 +1426,15 @@ dependencies = [ [[package]] name = "vchord" -version = "0.1.0" +version = "0.0.0" dependencies = [ "base", - "detect", "half 2.4.1", "log", "nalgebra", "paste", "pgrx", "pgrx-catalog", - "quantization", "rand", "rand_chacha", "rand_distr", @@ -1561,9 +1469,9 @@ checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" [[package]] name = "wide" -version = "0.7.28" +version = "0.7.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b828f995bf1e9622031f8009f8481a85406ce1f4d4588ff746d872043e855690" +checksum = "58e6db2670d2be78525979e9a5f9c69d296fd7d670549fe9ebf70f8708cb5019" dependencies = [ "bytemuck", "safe_arch", diff --git a/Cargo.toml b/Cargo.toml index 000def6c..66af418d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vchord" -version = "0.1.0" +version = "0.0.0" edition = "2021" [lib] @@ -20,17 +20,15 @@ pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] -base = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -detect = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -quantization = { git = "https://github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } +base = { git = "https://github.com/tensorchord/pgvecto.rs.git", rev = "c911e8a476effaf05cd4b1a037826b100177bdad" } # lock algebra version forever so that the QR decomposition never changes for same input -nalgebra = { version = "=0.33.0", default-features = false } +nalgebra = "=0.33.0" # lock rkyv version forever so that data is always compatible rkyv = { version = "=0.7.45", features = ["validation"] } -half = "2.4.1" +half = { version = "2.4.1", features = ["rkyv"] } log = "0.4.22" paste = "1" pgrx = { version = "=0.12.8", default-features = false, features = ["cshim"] } @@ -43,6 +41,9 @@ serde = "1" toml = "0.8.19" validator = { version = "0.18.1", features = ["derive"] } +[patch.crates-io] +half = { git = "https://github.com/tensorchord/half-rs.git" } + [lints] rust.unsafe_op_in_unsafe_fn = "deny" rust.unused_lifetimes = "warn" diff --git a/src/datatype/binary_scalar8.rs b/src/datatype/binary_scalar8.rs new file mode 100644 index 00000000..4dfe844b --- /dev/null +++ b/src/datatype/binary_scalar8.rs @@ -0,0 +1,75 @@ +use super::memory_scalar8::{Scalar8Input, Scalar8Output}; +use crate::types::scalar8::Scalar8Borrowed; +use base::vector::VectorBorrowed; +use pgrx::datum::Internal; +use pgrx::pg_sys::Oid; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_send(vector: Scalar8Input<'_>) -> Vec { + let vector = vector.as_borrowed(); + let mut stream = Vec::::new(); + stream.extend(vector.dims().to_be_bytes()); + stream.extend(vector.sum_of_x2().to_be_bytes()); + stream.extend(vector.k().to_be_bytes()); + stream.extend(vector.b().to_be_bytes()); + stream.extend(vector.sum_of_code().to_be_bytes()); + for &c in vector.code() { + stream.extend(c.to_be_bytes()); + } + stream +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_recv(internal: Internal, oid: Oid, typmod: i32) -> Scalar8Output { + let _ = (oid, typmod); + let buf = unsafe { internal.get_mut::().unwrap() }; + + let dims = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + u32::from_be_bytes(raw) + }; + let sum_of_x2 = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let k = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let b = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let sum_of_code = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let code = { + let mut result = Vec::with_capacity(dims as _); + for _ in 0..dims { + result.push({ + assert!(buf.cursor < i32::MAX - 1 && buf.cursor + 1 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 1]>().read() }; + buf.cursor += 1; + u8::from_be_bytes(raw) + }); + } + result + }; + + if let Some(x) = Scalar8Borrowed::new_checked(sum_of_x2, k, b, sum_of_code, &code) { + Scalar8Output::new(x) + } else { + pgrx::error!("detect data corruption"); + } +} diff --git a/src/datatype/functions_scalar8.rs b/src/datatype/functions_scalar8.rs new file mode 100644 index 00000000..fc1c221c --- /dev/null +++ b/src/datatype/functions_scalar8.rs @@ -0,0 +1,26 @@ +use crate::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput; +use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; +use crate::datatype::memory_scalar8::Scalar8Output; +use crate::types::scalar8::Scalar8Borrowed; +use base::simd::ScalarLike; +use half::f16; + +#[pgrx::pg_extern(sql = "")] +fn _vchord_vector_quantize_to_scalar8(vector: PgvectorVectorInput) -> Scalar8Output { + let vector = vector.as_borrowed(); + let sum_of_x2 = f32::reduce_sum_of_x2(vector.slice()); + let (k, b, code) = + base::simd::quantize::quantize(f32::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); + let sum_of_code = base::simd::u8::reduce_sum_of_x_as_u32(&code) as f32; + Scalar8Output::new(Scalar8Borrowed::new(sum_of_x2, k, b, sum_of_code, &code)) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_halfvec_quantize_to_scalar8(vector: PgvectorHalfvecInput) -> Scalar8Output { + let vector = vector.as_borrowed(); + let sum_of_x2 = f16::reduce_sum_of_x2(vector.slice()); + let (k, b, code) = + base::simd::quantize::quantize(f16::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); + let sum_of_code = base::simd::u8::reduce_sum_of_x_as_u32(&code) as f32; + Scalar8Output::new(Scalar8Borrowed::new(sum_of_x2, k, b, sum_of_code, &code)) +} diff --git a/src/datatype/memory_pgvector_halfvec.rs b/src/datatype/memory_pgvector_halfvec.rs new file mode 100644 index 00000000..7265a1be --- /dev/null +++ b/src/datatype/memory_pgvector_halfvec.rs @@ -0,0 +1,204 @@ +use base::vector::*; +use half::f16; +use pgrx::datum::FromDatum; +use pgrx::datum::IntoDatum; +use pgrx::pg_sys::Datum; +use pgrx::pg_sys::Oid; +use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; +use pgrx::pgrx_sql_entity_graph::metadata::Returns; +use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; +use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; +use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; +use std::ops::Deref; +use std::ptr::NonNull; + +#[repr(C, align(8))] +pub struct PgvectorHalfvecHeader { + varlena: u32, + dims: u16, + unused: u16, + phantom: [f16; 0], +} + +impl PgvectorHalfvecHeader { + fn size_of(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + (size_of::() + size_of::() * len).next_multiple_of(8) + } + pub fn as_borrowed(&self) -> VectBorrowed<'_, f16> { + unsafe { + VectBorrowed::new_unchecked(std::slice::from_raw_parts( + self.phantom.as_ptr(), + self.dims as usize, + )) + } + } +} + +pub enum PgvectorHalfvecInput<'a> { + Owned(PgvectorHalfvecOutput), + Borrowed(&'a PgvectorHalfvecHeader), +} + +impl PgvectorHalfvecInput<'_> { + unsafe fn new(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() + }; + if p != q { + PgvectorHalfvecInput::Owned(PgvectorHalfvecOutput(q)) + } else { + unsafe { PgvectorHalfvecInput::Borrowed(p.as_ref()) } + } + } +} + +impl Deref for PgvectorHalfvecInput<'_> { + type Target = PgvectorHalfvecHeader; + + fn deref(&self) -> &Self::Target { + match self { + PgvectorHalfvecInput::Owned(x) => x, + PgvectorHalfvecInput::Borrowed(x) => x, + } + } +} + +pub struct PgvectorHalfvecOutput(NonNull); + +impl PgvectorHalfvecOutput { + pub fn new(vector: VectBorrowed<'_, f16>) -> PgvectorHalfvecOutput { + unsafe { + let slice = vector.slice(); + let size = PgvectorHalfvecHeader::size_of(slice.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut PgvectorHalfvecHeader; + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dims).write(vector.dims() as _); + (&raw mut (*ptr).unused).write(0); + std::ptr::copy_nonoverlapping(slice.as_ptr(), (*ptr).phantom.as_mut_ptr(), slice.len()); + PgvectorHalfvecOutput(NonNull::new(ptr).unwrap()) + } + } + pub fn into_raw(self) -> *mut PgvectorHalfvecHeader { + let result = self.0.as_ptr(); + std::mem::forget(self); + result + } +} + +impl Deref for PgvectorHalfvecOutput { + type Target = PgvectorHalfvecHeader; + + fn deref(&self) -> &Self::Target { + unsafe { self.0.as_ref() } + } +} + +impl Drop for PgvectorHalfvecOutput { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr() as _); + } + } +} + +impl FromDatum for PgvectorHalfvecInput<'_> { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr::()).unwrap(); + unsafe { Some(PgvectorHalfvecInput::new(ptr)) } + } + } +} + +impl IntoDatum for PgvectorHalfvecOutput { + fn into_datum(self) -> Option { + Some(Datum::from(self.into_raw() as *mut ())) + } + + fn type_oid() -> Oid { + Oid::INVALID + } + + fn is_compatible_with(_: Oid) -> bool { + true + } +} + +impl FromDatum for PgvectorHalfvecOutput { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let p = NonNull::new(datum.cast_mut_ptr::())?; + let q = + unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast())? }; + if p != q { + Some(PgvectorHalfvecOutput(q)) + } else { + let header = p.as_ptr(); + let vector = unsafe { (*header).as_borrowed() }; + Some(PgvectorHalfvecOutput::new(vector)) + } + } + } +} + +unsafe impl pgrx::datum::UnboxDatum for PgvectorHalfvecOutput { + type As<'src> = PgvectorHalfvecOutput; + #[inline] + unsafe fn unbox<'src>(d: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let p = NonNull::new(d.sans_lifetime().cast_mut_ptr::()).unwrap(); + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() + }; + if p != q { + PgvectorHalfvecOutput(q) + } else { + let header = p.as_ptr(); + let vector = unsafe { (*header).as_borrowed() }; + PgvectorHalfvecOutput::new(vector) + } + } +} + +unsafe impl SqlTranslatable for PgvectorHalfvecInput<'_> { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("halfvec"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("halfvec")))) + } +} + +unsafe impl SqlTranslatable for PgvectorHalfvecOutput { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("halfvec"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("halfvec")))) + } +} + +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for PgvectorHalfvecInput<'fcx> { + unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { + unsafe { arg.unbox_arg_using_from_datum().unwrap() } + } +} + +unsafe impl pgrx::callconv::BoxRet for PgvectorHalfvecOutput { + unsafe fn box_into<'fcx>( + self, + fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, + ) -> pgrx::datum::Datum<'fcx> { + unsafe { fcinfo.return_raw_datum(Datum::from(self.into_raw() as *mut ())) } + } +} diff --git a/src/datatype/memory_pgvector_vector.rs b/src/datatype/memory_pgvector_vector.rs index 7166ace0..d81492cf 100644 --- a/src/datatype/memory_pgvector_vector.rs +++ b/src/datatype/memory_pgvector_vector.rs @@ -8,48 +8,31 @@ use pgrx::pgrx_sql_entity_graph::metadata::Returns; use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; -use std::alloc::Layout; use std::ops::Deref; use std::ptr::NonNull; -pub const HEADER_MAGIC: u16 = 0; - #[repr(C, align(8))] pub struct PgvectorVectorHeader { varlena: u32, dims: u16, - magic: u16, + unused: u16, phantom: [f32; 0], } impl PgvectorVectorHeader { - fn varlena(size: usize) -> u32 { - (size << 2) as u32 - } - fn layout(len: usize) -> Layout { - u16::try_from(len).expect("Vector is too large."); - let layout_alpha = Layout::new::(); - let layout_beta = Layout::array::(len).unwrap(); - let layout = layout_alpha.extend(layout_beta).unwrap().0; - layout.pad_to_align() - } - #[allow(dead_code)] - pub fn dims(&self) -> u32 { - self.dims as u32 - } - pub fn slice(&self) -> &[f32] { - unsafe { std::slice::from_raw_parts(self.phantom.as_ptr(), self.dims as usize) } + fn size_of(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + (size_of::() + size_of::() * len).next_multiple_of(8) } pub fn as_borrowed(&self) -> VectBorrowed<'_, f32> { - unsafe { VectBorrowed::new_unchecked(self.slice()) } - } -} - -impl Deref for PgvectorVectorHeader { - type Target = [f32]; - - fn deref(&self) -> &Self::Target { - self.slice() + unsafe { + VectBorrowed::new_unchecked(std::slice::from_raw_parts( + self.phantom.as_ptr(), + self.dims as usize, + )) + } } } @@ -88,15 +71,12 @@ impl PgvectorVectorOutput { pub fn new(vector: VectBorrowed<'_, f32>) -> PgvectorVectorOutput { unsafe { let slice = vector.slice(); - let layout = PgvectorVectorHeader::layout(slice.len()); - let dims = vector.dims(); - let internal_dims = dims as u16; - let ptr = pgrx::pg_sys::palloc(layout.size()) as *mut PgvectorVectorHeader; - ptr.cast::().add(layout.size() - 8).write_bytes(0, 8); - std::ptr::addr_of_mut!((*ptr).varlena) - .write(PgvectorVectorHeader::varlena(layout.size())); - std::ptr::addr_of_mut!((*ptr).magic).write(HEADER_MAGIC); - std::ptr::addr_of_mut!((*ptr).dims).write(internal_dims); + let size = PgvectorVectorHeader::size_of(slice.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut PgvectorVectorHeader; + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dims).write(vector.dims() as _); + (&raw mut (*ptr).unused).write(0); std::ptr::copy_nonoverlapping(slice.as_ptr(), (*ptr).phantom.as_mut_ptr(), slice.len()); PgvectorVectorOutput(NonNull::new(ptr).unwrap()) } diff --git a/src/datatype/memory_scalar8.rs b/src/datatype/memory_scalar8.rs new file mode 100644 index 00000000..3a7dab4a --- /dev/null +++ b/src/datatype/memory_scalar8.rs @@ -0,0 +1,215 @@ +use crate::types::scalar8::Scalar8Borrowed; +use base::vector::*; +use pgrx::datum::FromDatum; +use pgrx::datum::IntoDatum; +use pgrx::pg_sys::Datum; +use pgrx::pg_sys::Oid; +use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; +use pgrx::pgrx_sql_entity_graph::metadata::Returns; +use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; +use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; +use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; +use std::ops::Deref; +use std::ptr::NonNull; + +#[repr(C, align(8))] +pub struct Scalar8Header { + varlena: u32, + dims: u16, + unused: u16, + sum_of_x2: f32, + k: f32, + b: f32, + sum_of_code: f32, + phantom: [u8; 0], +} + +impl Scalar8Header { + fn size_of(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + (size_of::() + size_of::() * len).next_multiple_of(8) + } + pub fn as_borrowed(&self) -> Scalar8Borrowed<'_> { + unsafe { + Scalar8Borrowed::new_unchecked( + self.sum_of_x2, + self.k, + self.b, + self.sum_of_code, + std::slice::from_raw_parts(self.phantom.as_ptr(), self.dims as usize), + ) + } + } +} + +pub enum Scalar8Input<'a> { + Owned(Scalar8Output), + Borrowed(&'a Scalar8Header), +} + +impl Scalar8Input<'_> { + unsafe fn new(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() + }; + if p != q { + Scalar8Input::Owned(Scalar8Output(q)) + } else { + unsafe { Scalar8Input::Borrowed(p.as_ref()) } + } + } +} + +impl Deref for Scalar8Input<'_> { + type Target = Scalar8Header; + + fn deref(&self) -> &Self::Target { + match self { + Scalar8Input::Owned(x) => x, + Scalar8Input::Borrowed(x) => x, + } + } +} + +pub struct Scalar8Output(NonNull); + +impl Scalar8Output { + pub fn new(vector: Scalar8Borrowed<'_>) -> Scalar8Output { + unsafe { + let code = vector.code(); + let size = Scalar8Header::size_of(code.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut Scalar8Header; + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dims).write(vector.dims() as _); + (&raw mut (*ptr).unused).write(0); + (&raw mut (*ptr).sum_of_x2).write(vector.sum_of_x2()); + (&raw mut (*ptr).k).write(vector.k()); + (&raw mut (*ptr).b).write(vector.b()); + (&raw mut (*ptr).sum_of_code).write(vector.sum_of_code()); + std::ptr::copy_nonoverlapping(code.as_ptr(), (*ptr).phantom.as_mut_ptr(), code.len()); + Scalar8Output(NonNull::new(ptr).unwrap()) + } + } + pub fn into_raw(self) -> *mut Scalar8Header { + let result = self.0.as_ptr(); + std::mem::forget(self); + result + } +} + +impl Deref for Scalar8Output { + type Target = Scalar8Header; + + fn deref(&self) -> &Self::Target { + unsafe { self.0.as_ref() } + } +} + +impl Drop for Scalar8Output { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr() as _); + } + } +} + +impl FromDatum for Scalar8Input<'_> { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr::()).unwrap(); + unsafe { Some(Scalar8Input::new(ptr)) } + } + } +} + +impl IntoDatum for Scalar8Output { + fn into_datum(self) -> Option { + Some(Datum::from(self.into_raw() as *mut ())) + } + + fn type_oid() -> Oid { + Oid::INVALID + } + + fn is_compatible_with(_: Oid) -> bool { + true + } +} + +impl FromDatum for Scalar8Output { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let p = NonNull::new(datum.cast_mut_ptr::())?; + let q = + unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast())? }; + if p != q { + Some(Scalar8Output(q)) + } else { + let header = p.as_ptr(); + let vector = unsafe { (*header).as_borrowed() }; + Some(Scalar8Output::new(vector)) + } + } + } +} + +unsafe impl pgrx::datum::UnboxDatum for Scalar8Output { + type As<'src> = Scalar8Output; + #[inline] + unsafe fn unbox<'src>(d: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let p = NonNull::new(d.sans_lifetime().cast_mut_ptr::()).unwrap(); + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() + }; + if p != q { + Scalar8Output(q) + } else { + let header = p.as_ptr(); + let vector = unsafe { (*header).as_borrowed() }; + Scalar8Output::new(vector) + } + } +} + +unsafe impl SqlTranslatable for Scalar8Input<'_> { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("scalar8"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("scalar8")))) + } +} + +unsafe impl SqlTranslatable for Scalar8Output { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("scalar8"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("scalar8")))) + } +} + +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Scalar8Input<'fcx> { + unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { + unsafe { arg.unbox_arg_using_from_datum().unwrap() } + } +} + +unsafe impl pgrx::callconv::BoxRet for Scalar8Output { + unsafe fn box_into<'fcx>( + self, + fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, + ) -> pgrx::datum::Datum<'fcx> { + unsafe { fcinfo.return_raw_datum(Datum::from(self.into_raw() as *mut ())) } + } +} diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 9152f33a..98b6650e 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -1,3 +1,10 @@ +pub mod binary_scalar8; +pub mod functions_scalar8; +pub mod memory_pgvector_halfvec; pub mod memory_pgvector_vector; +pub mod memory_scalar8; +pub mod operators_pgvector_halfvec; pub mod operators_pgvector_vector; +pub mod operators_scalar8; +pub mod text_scalar8; pub mod typmod; diff --git a/src/datatype/operators_pgvector_halfvec.rs b/src/datatype/operators_pgvector_halfvec.rs new file mode 100644 index 00000000..2e707ebf --- /dev/null +++ b/src/datatype/operators_pgvector_halfvec.rs @@ -0,0 +1,75 @@ +use crate::datatype::memory_pgvector_halfvec::*; +use base::vector::{VectBorrowed, VectorBorrowed}; +use std::num::NonZero; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_halfvec_sphere_l2_in( + lhs: PgvectorHalfvecInput<'_>, + rhs: pgrx::composite_type!("sphere_halfvec"), +) -> bool { + let center: PgvectorHalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } + let d = VectBorrowed::operator_l2(lhs, center).to_f32().sqrt(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_halfvec_sphere_ip_in( + lhs: PgvectorHalfvecInput<'_>, + rhs: pgrx::composite_type!("sphere_halfvec"), +) -> bool { + let center: PgvectorHalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } + let d = VectBorrowed::operator_dot(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_halfvec_sphere_cosine_in( + lhs: PgvectorHalfvecInput<'_>, + rhs: pgrx::composite_type!("sphere_halfvec"), +) -> bool { + let center: PgvectorHalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } + let d = VectBorrowed::operator_cos(lhs, center).to_f32(); + d < radius +} diff --git a/src/datatype/operators_pgvector_vector.rs b/src/datatype/operators_pgvector_vector.rs index 4a1f0554..2308ab9d 100644 --- a/src/datatype/operators_pgvector_vector.rs +++ b/src/datatype/operators_pgvector_vector.rs @@ -12,9 +12,6 @@ fn _vchord_vector_sphere_l2_in( Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), }; - if lhs.dims() != center.dims() { - pgrx::error!("dimension is not matched"); - } let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), @@ -22,6 +19,9 @@ fn _vchord_vector_sphere_l2_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } let d = VectBorrowed::operator_l2(lhs, center).to_f32().sqrt(); d < radius } @@ -36,9 +36,6 @@ fn _vchord_vector_sphere_ip_in( Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), }; - if lhs.dims() != center.dims() { - pgrx::error!("dimension is not matched"); - } let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), @@ -46,6 +43,9 @@ fn _vchord_vector_sphere_ip_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } let d = VectBorrowed::operator_dot(lhs, center).to_f32(); d < radius } @@ -60,9 +60,6 @@ fn _vchord_vector_sphere_cosine_in( Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), }; - if lhs.dims() != center.dims() { - pgrx::error!("dimension is not matched"); - } let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), @@ -70,6 +67,9 @@ fn _vchord_vector_sphere_cosine_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } let d = VectBorrowed::operator_cos(lhs, center).to_f32(); d < radius } diff --git a/src/datatype/operators_scalar8.rs b/src/datatype/operators_scalar8.rs new file mode 100644 index 00000000..db6a3727 --- /dev/null +++ b/src/datatype/operators_scalar8.rs @@ -0,0 +1,106 @@ +use crate::datatype::memory_scalar8::{Scalar8Input, Scalar8Output}; +use crate::types::scalar8::Scalar8Borrowed; +use base::vector::*; +use std::num::NonZero; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_operator_ip(lhs: Scalar8Input<'_>, rhs: Scalar8Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dims() != rhs.dims() { + pgrx::error!("dimension is not matched"); + } + Scalar8Borrowed::operator_dot(lhs, rhs).to_f32() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_operator_l2(lhs: Scalar8Input<'_>, rhs: Scalar8Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dims() != rhs.dims() { + pgrx::error!("dimension is not matched"); + } + Scalar8Borrowed::operator_l2(lhs, rhs).to_f32().sqrt() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_operator_cosine(lhs: Scalar8Input<'_>, rhs: Scalar8Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dims() != rhs.dims() { + pgrx::error!("dimension is not matched"); + } + Scalar8Borrowed::operator_cos(lhs, rhs).to_f32() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_sphere_ip_in( + lhs: Scalar8Input<'_>, + rhs: pgrx::composite_type!("sphere_scalar8"), +) -> bool { + let center: Scalar8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } + let d = Scalar8Borrowed::operator_dot(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_sphere_l2_in( + lhs: Scalar8Input<'_>, + rhs: pgrx::composite_type!("sphere_scalar8"), +) -> bool { + let center: Scalar8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } + let d = Scalar8Borrowed::operator_l2(lhs, center).to_f32().sqrt(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_sphere_cosine_in( + lhs: Scalar8Input<'_>, + rhs: pgrx::composite_type!("sphere_scalar8"), +) -> bool { + let center: Scalar8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dims() != center.dims() { + pgrx::error!("dimension is not matched"); + } + let d = Scalar8Borrowed::operator_cos(lhs, center).to_f32(); + d < radius +} diff --git a/src/datatype/text_scalar8.rs b/src/datatype/text_scalar8.rs new file mode 100644 index 00000000..5de82da6 --- /dev/null +++ b/src/datatype/text_scalar8.rs @@ -0,0 +1,138 @@ +use super::memory_scalar8::Scalar8Output; +use crate::datatype::memory_scalar8::Scalar8Input; +use crate::types::scalar8::Scalar8Borrowed; +use pgrx::pg_sys::Oid; +use std::ffi::{CStr, CString}; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_in(input: &CStr, oid: Oid, typmod: i32) -> Scalar8Output { + let _ = (oid, typmod); + let mut input = input.to_bytes().iter(); + let mut p0 = Vec::::new(); + let mut p1 = Vec::::new(); + { + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + match c { + b' ' => (), + b'(' => break, + _ => pgrx::error!("incorrect vector"), + } + } + } + { + let mut s = Option::::None; + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + s = match (s, c) { + (s, b' ') => s, + (None, c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + Some(String::from(c as char)) + } + (Some(s), c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + let mut x = s; + x.push(c as char); + Some(x) + } + (Some(s), b',') => { + p0.push(s.parse().expect("failed to parse number")); + None + } + (None, b',') => { + pgrx::error!("incorrect vector") + } + (Some(s), b')') => { + p0.push(s.parse().expect("failed to parse number")); + break; + } + (None, b')') => break, + _ => pgrx::error!("incorrect vector"), + }; + } + } + { + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + match c { + b' ' => (), + b'[' => break, + _ => pgrx::error!("incorrect vector"), + } + } + } + { + let mut s = Option::::None; + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + s = match (s, c) { + (s, b' ') => s, + (None, c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + Some(String::from(c as char)) + } + (Some(s), c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + let mut x = s; + x.push(c as char); + Some(x) + } + (Some(s), b',') => { + p1.push(s.parse().expect("failed to parse number")); + None + } + (None, b',') => { + pgrx::error!("incorrect vector") + } + (Some(s), b']') => { + p1.push(s.parse().expect("failed to parse number")); + break; + } + (None, b']') => break, + _ => pgrx::error!("incorrect vector"), + }; + } + } + if p0.len() != 4 { + pgrx::error!("incorrect vector"); + } + if p1.is_empty() { + pgrx::error!("vector must have at least 1 dimension"); + } + let sum_of_x2 = p0[0]; + let k = p0[1]; + let b = p0[2]; + let sum_of_code = p0[3]; + let code = p1; + if let Some(x) = Scalar8Borrowed::new_checked(sum_of_x2, k, b, sum_of_code, &code) { + Scalar8Output::new(x) + } else { + pgrx::error!("incorrect vector"); + } +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_scalar8_out(vector: Scalar8Input<'_>) -> CString { + let vector = vector.as_borrowed(); + let mut buffer = String::new(); + buffer.push('('); + buffer.push_str(format!("{}", vector.sum_of_x2()).as_str()); + buffer.push_str(format!(", {}", vector.k()).as_str()); + buffer.push_str(format!(", {}", vector.b()).as_str()); + buffer.push_str(format!(", {}", vector.sum_of_code()).as_str()); + buffer.push(')'); + buffer.push('['); + if let Some(&x) = vector.code().first() { + buffer.push_str(format!("{}", x).as_str()); + } + for &x in vector.code().iter().skip(1) { + buffer.push_str(format!(", {}", x).as_str()); + } + buffer.push(']'); + CString::new(buffer).unwrap() +} diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index d67bb0be..fe90a6d2 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -1,5 +1,6 @@ use serde::{Deserialize, Serialize}; -use std::num::NonZeroU32; +use std::ffi::{CStr, CString}; +use std::num::{NonZero, NonZeroU32}; #[derive(Debug, Clone, Copy, Serialize, Deserialize)] pub enum Typmod { @@ -42,3 +43,30 @@ impl Typmod { } } } + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_typmod_in_65535(list: pgrx::datum::Array<&CStr>) -> i32 { + if list.is_empty() { + -1 + } else if list.len() == 1 { + let s = list.get(0).unwrap().unwrap().to_str().unwrap(); + let d = s.parse::().ok(); + if let Some(d @ 1..=65535) = d { + let typmod = Typmod::Dims(NonZero::new(d).unwrap()); + typmod.into_i32() + } else { + pgrx::error!("Modifier of the type is invalid.") + } + } else { + pgrx::error!("Modifier of the type is invalid.") + } +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_typmod_out(typmod: i32) -> CString { + let typmod = Typmod::parse_from_i32(typmod).unwrap(); + match typmod.into_option_string() { + Some(s) => CString::new(format!("({})", s)).unwrap(), + None => CString::new("()").unwrap(), + } +} diff --git a/src/lib.rs b/src/lib.rs index 94741f56..4c55ac98 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,10 +3,14 @@ #![allow(clippy::needless_range_loop)] #![allow(clippy::too_many_arguments)] #![allow(clippy::type_complexity)] +#![allow(clippy::int_plus_one)] +#![allow(clippy::unused_unit)] +#![allow(clippy::infallible_destructuring_match)] mod datatype; mod postgres; mod projection; +mod types; mod upgrade; mod utils; mod vchordrq; @@ -21,7 +25,7 @@ unsafe extern "C" fn _PG_init() { if unsafe { pgrx::pg_sys::IsUnderPostmaster } { pgrx::error!("vchord must be loaded via shared_preload_libraries."); } - detect::init(); + base::simd::enable(); unsafe { vchordrq::init(); vchordrqfscan::init(); diff --git a/src/projection.rs b/src/projection.rs index 68926ac9..4273180a 100644 --- a/src/projection.rs +++ b/src/projection.rs @@ -67,7 +67,7 @@ pub fn prewarm(n: usize) { } pub fn project(vector: &[f32]) -> Vec { - use base::scalar::ScalarLike; + use base::simd::ScalarLike; let n = vector.len(); let matrix = MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); (0..n) diff --git a/src/sql/bootstrap.sql b/src/sql/bootstrap.sql index 6cc59f1a..acc6a551 100644 --- a/src/sql/bootstrap.sql +++ b/src/sql/bootstrap.sql @@ -1,3 +1,6 @@ -- List of shell types +CREATE TYPE scalar8; CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 32b356c4..c00aab1a 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -1,12 +1,55 @@ -- List of data types +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + CREATE TYPE sphere_vector AS ( center vector, radius REAL ); +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + -- List of operators +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + CREATE OPERATOR <<->> ( PROCEDURE = _vchord_vector_sphere_l2_in, LEFTARG = vector, @@ -14,6 +57,20 @@ CREATE OPERATOR <<->> ( COMMUTATOR = <<->> ); +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + CREATE OPERATOR <<#>> ( PROCEDURE = _vchord_vector_sphere_ip_in, LEFTARG = vector, @@ -21,6 +78,20 @@ CREATE OPERATOR <<#>> ( COMMUTATOR = <<#>> ); +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + CREATE OPERATOR <<=>> ( PROCEDURE = _vchord_vector_sphere_cosine_in, LEFTARG = vector, @@ -28,11 +99,37 @@ CREATE OPERATOR <<=>> ( COMMUTATOR = <<=>> ); +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + -- List of functions CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; @@ -56,6 +153,9 @@ CREATE ACCESS METHOD Vchordrqfscan TYPE INDEX HANDLER Vchordrqfscan_amhandler; CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; CREATE OPERATOR FAMILY vector_l2_ops USING Vchordrqfscan; CREATE OPERATOR FAMILY vector_ip_ops USING Vchordrqfscan; @@ -81,6 +181,24 @@ CREATE OPERATOR CLASS vector_cosine_ops OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, FUNCTION 1 _vchordrq_support_vector_cosine_ops(); +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + CREATE OPERATOR CLASS vector_l2_ops FOR TYPE vector USING Vchordrqfscan FAMILY vector_l2_ops AS OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, diff --git a/src/types/mod.rs b/src/types/mod.rs new file mode 100644 index 00000000..af08ee7f --- /dev/null +++ b/src/types/mod.rs @@ -0,0 +1 @@ +pub mod scalar8; diff --git a/src/types/scalar8.rs b/src/types/scalar8.rs new file mode 100644 index 00000000..55c0074a --- /dev/null +++ b/src/types/scalar8.rs @@ -0,0 +1,286 @@ +use base::distance::Distance; +use base::vector::{VectorBorrowed, VectorOwned}; +use serde::{Deserialize, Serialize}; +use std::ops::RangeBounds; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Scalar8Owned { + sum_of_x2: f32, + k: f32, + b: f32, + sum_of_code: f32, + code: Vec, +} + +impl Scalar8Owned { + #[allow(dead_code)] + #[inline(always)] + pub fn new(sum_of_x2: f32, k: f32, b: f32, sum_of_code: f32, code: Vec) -> Self { + Self::new_checked(sum_of_x2, k, b, sum_of_code, code).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + sum_of_x2: f32, + k: f32, + b: f32, + sum_of_code: f32, + code: Vec, + ) -> Option { + if !(1..=65535).contains(&code.len()) { + return None; + } + Some(unsafe { Self::new_unchecked(sum_of_x2, k, b, sum_of_code, code) }) + } + + /// # Safety + /// + /// * `code.len()` must not be zero. + #[inline(always)] + pub unsafe fn new_unchecked( + sum_of_x2: f32, + k: f32, + b: f32, + sum_of_code: f32, + code: Vec, + ) -> Self { + Self { + sum_of_x2, + k, + b, + sum_of_code, + code, + } + } +} + +impl VectorOwned for Scalar8Owned { + type Borrowed<'a> = Scalar8Borrowed<'a>; + + #[inline(always)] + fn as_borrowed(&self) -> Scalar8Borrowed<'_> { + Scalar8Borrowed { + sum_of_x2: self.sum_of_x2, + k: self.k, + b: self.b, + sum_of_code: self.sum_of_code, + code: self.code.as_slice(), + } + } + + #[inline(always)] + fn zero(dims: u32) -> Self { + Self { + sum_of_x2: 0.0, + k: 0.0, + b: 0.0, + sum_of_code: 0.0, + code: vec![0; dims as usize], + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct Scalar8Borrowed<'a> { + sum_of_x2: f32, + k: f32, + b: f32, + sum_of_code: f32, + code: &'a [u8], +} + +impl<'a> Scalar8Borrowed<'a> { + #[inline(always)] + pub fn new(sum_of_x2: f32, k: f32, b: f32, sum_of_code: f32, code: &'a [u8]) -> Self { + Self::new_checked(sum_of_x2, k, b, sum_of_code, code).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + sum_of_x2: f32, + k: f32, + b: f32, + sum_of_code: f32, + code: &'a [u8], + ) -> Option { + if !(1..=65535).contains(&code.len()) { + return None; + } + Some(unsafe { Self::new_unchecked(sum_of_x2, k, b, sum_of_code, code) }) + } + + /// # Safety + /// + /// * `code.len()` must not be zero. + #[inline(always)] + pub unsafe fn new_unchecked( + sum_of_x2: f32, + k: f32, + b: f32, + sum_of_code: f32, + code: &'a [u8], + ) -> Self { + Self { + sum_of_x2, + k, + b, + sum_of_code, + code, + } + } + + #[inline(always)] + pub fn sum_of_x2(&self) -> f32 { + self.sum_of_x2 + } + + #[inline(always)] + pub fn k(&self) -> f32 { + self.k + } + + #[inline(always)] + pub fn b(&self) -> f32 { + self.b + } + + #[inline(always)] + pub fn sum_of_code(&self) -> f32 { + self.sum_of_code + } + + #[inline(always)] + pub fn code(&self) -> &'a [u8] { + self.code + } +} + +impl VectorBorrowed for Scalar8Borrowed<'_> { + type Owned = Scalar8Owned; + + #[inline(always)] + fn dims(&self) -> u32 { + self.code.len() as u32 + } + + #[inline(always)] + fn own(&self) -> Scalar8Owned { + Scalar8Owned { + sum_of_x2: self.sum_of_x2, + k: self.k, + b: self.b, + sum_of_code: self.sum_of_code, + code: self.code.to_owned(), + } + } + + #[inline(always)] + fn norm(&self) -> f32 { + self.sum_of_x2.sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + assert_eq!(self.code.len(), rhs.code.len()); + let xy = self.k * rhs.k * base::simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 + + self.b * rhs.b * self.code.len() as f32 + + self.k * rhs.b * self.sum_of_code + + self.b * rhs.k * rhs.sum_of_code; + Distance::from(-xy) + } + + #[inline(always)] + fn operator_l2(self, rhs: Self) -> Distance { + assert_eq!(self.code.len(), rhs.code.len()); + let xy = self.k * rhs.k * base::simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 + + self.b * rhs.b * self.code.len() as f32 + + self.k * rhs.b * self.sum_of_code + + self.b * rhs.k * rhs.sum_of_code; + let x2 = self.sum_of_x2; + let y2 = rhs.sum_of_x2; + Distance::from(x2 + y2 - 2.0 * xy) + } + + #[inline(always)] + fn operator_cos(self, rhs: Self) -> Distance { + assert_eq!(self.code.len(), rhs.code.len()); + let xy = self.k * rhs.k * base::simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 + + self.b * rhs.b * self.code.len() as f32 + + self.k * rhs.b * self.sum_of_code + + self.b * rhs.k * rhs.sum_of_code; + let x2 = self.sum_of_x2; + let y2 = rhs.sum_of_x2; + Distance::from(1.0 - xy / (x2 * y2).sqrt()) + } + + #[inline(always)] + fn operator_hamming(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_jaccard(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn function_normalize(&self) -> Scalar8Owned { + let l = self.sum_of_x2.sqrt(); + Scalar8Owned { + sum_of_x2: 1.0, + k: self.k / l, + b: self.b / l, + sum_of_code: self.sum_of_code, + code: self.code.to_owned(), + } + } + + fn operator_add(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_sub(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_mul(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_and(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_or(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_xor(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + #[inline(always)] + fn subvector(&self, bounds: impl RangeBounds) -> Option { + let start_bound = bounds.start_bound().map(|x| *x as usize); + let end_bound = bounds.end_bound().map(|x| *x as usize); + let code = self.code.get((start_bound, end_bound))?; + if code.is_empty() { + return None; + } + Self::Owned::new_checked( + { + // recover it as much as possible + let mut result = 0.0; + for &x in code { + let y = self.k * (x as f32) + self.b; + result += y * y; + } + result + }, + self.k, + self.b, + base::simd::u8::reduce_sum_of_x_as_u32(code) as f32, + code.to_owned(), + ) + } +} diff --git a/src/utils/infinite_byte_chunks.rs b/src/utils/infinite_byte_chunks.rs new file mode 100644 index 00000000..c61b87e6 --- /dev/null +++ b/src/utils/infinite_byte_chunks.rs @@ -0,0 +1,20 @@ +#[derive(Debug, Clone)] +pub struct InfiniteByteChunks { + iter: I, +} + +impl InfiniteByteChunks { + pub fn new(iter: I) -> Self { + Self { iter } + } +} + +impl, const N: usize> Iterator for InfiniteByteChunks { + type Item = [u8; N]; + + fn next(&mut self) -> Option { + Some(std::array::from_fn::(|_| { + self.iter.next().unwrap_or(0) + })) + } +} diff --git a/src/utils/k_means.rs b/src/utils/k_means.rs index ac58a088..97a810f8 100644 --- a/src/utils/k_means.rs +++ b/src/utils/k_means.rs @@ -1,7 +1,7 @@ #![allow(clippy::ptr_arg)] use super::parallelism::{ParallelIterator, Parallelism}; -use base::scalar::*; +use base::simd::*; use half::f16; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; @@ -121,7 +121,7 @@ fn rabitq_index( let mut a3 = Vec::new(); let mut a4 = Vec::new(); for vectors in centroids.chunks(32) { - use quantization::fast_scan::b4::pack; + use base::simd::fast_scan::b4::pack; let code_alphas = std::array::from_fn::<_, 32, _>(|i| { if let Some(vector) = vectors.get(i) { code_alpha(vector) @@ -197,7 +197,7 @@ fn rabitq_index( ), epsilon: f32, ) -> [Distance; 32] { - use quantization::fast_scan::b4::fast_scan_b4; + use base::simd::fast_scan::b4::fast_scan_b4; let &(dis_v_2, b, k, qvector_sum, ref s) = lut; let r = fast_scan_b4(dims.div_ceil(4), t, s); std::array::from_fn(|i| { @@ -210,17 +210,17 @@ fn rabitq_index( }) } use base::distance::Distance; - use quantization::quantize; + use base::simd::quantize; let lut = { let vector = &samples[i]; let dis_v_2 = f32::reduce_sum_of_x2(vector); let (k, b, qvector) = - quantize::quantize::<15>(f32::vector_to_f32_borrowed(vector).as_ref()); + quantize::quantize(f32::vector_to_f32_borrowed(vector).as_ref(), 15.0); let qvector_sum = if vector.len() <= 4369 { - quantize::reduce_sum_of_x_as_u16(&qvector) as f32 + u8::reduce_sum_of_x_as_u16(&qvector) as f32 } else { - quantize::reduce_sum_of_x_as_u32(&qvector) as f32 + u8::reduce_sum_of_x_as_u32(&qvector) as f32 }; (dis_v_2, b, k, qvector_sum, gen(qvector)) }; diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 1b07dc6c..2d9a3b7a 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,2 +1,3 @@ +pub mod infinite_byte_chunks; pub mod k_means; pub mod parallelism; diff --git a/src/vchordrq/algorithm/build.rs b/src/vchordrq/algorithm/build.rs index c65bd851..d292628a 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -2,25 +2,25 @@ use crate::postgres::BufferWriteGuard; use crate::postgres::Relation; use crate::vchordrq::algorithm::rabitq; use crate::vchordrq::algorithm::tuples::*; -use crate::vchordrq::algorithm::vectors; use crate::vchordrq::index::am_options::Opfamily; use crate::vchordrq::types::VchordrqBuildOptions; use crate::vchordrq::types::VchordrqExternalBuildOptions; use crate::vchordrq::types::VchordrqIndexingOptions; use crate::vchordrq::types::VchordrqInternalBuildOptions; +use crate::vchordrq::types::VectorOptions; use base::distance::DistanceKind; -use base::index::VectorOptions; -use base::scalar::ScalarLike; use base::search::Pointer; +use base::simd::ScalarLike; +use base::vector::VectorBorrowed; use rand::Rng; use rkyv::ser::serializers::AllocSerializer; use std::marker::PhantomData; use std::sync::Arc; -pub trait HeapRelation { +pub trait HeapRelation { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((Pointer, Vec)); + F: FnMut((Pointer, V)); fn opfamily(&self) -> Opfamily; } @@ -28,7 +28,7 @@ pub trait Reporter { fn tuples_total(&mut self, tuples_total: u64); } -pub fn build( +pub fn build, R: Reporter>( vector_options: VectorOptions, vchordrq_options: VchordrqIndexingOptions, heap_relation: T, @@ -56,13 +56,14 @@ pub fn build( let mut samples = Vec::new(); let mut number_of_samples = 0_u32; heap_relation.traverse(false, |(_, vector)| { - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); + let vector = vector.as_borrowed(); + assert_eq!(dims, vector.dims(), "invalid vector dimensions"); if number_of_samples < max_number_of_samples { - samples.push(vector); + samples.push(V::build_to_vecf32(vector)); number_of_samples += 1; } else { let index = rand.gen_range(0..max_number_of_samples) as usize; - samples[index] = vector; + samples[index] = V::build_to_vecf32(vector); } tuples_total += 1; }); @@ -74,21 +75,22 @@ pub fn build( }; let mut meta = Tape::create(&relation, false); assert_eq!(meta.first(), 0); - let mut vectors = Tape::create(&relation, true); + let mut vectors = Tape::>::create(&relation, true); let mut pointer_of_means = Vec::>::new(); for i in 0..structures.len() { let mut level = Vec::new(); for j in 0..structures[i].len() { - let slices = vectors::vector_split(&structures[i].means[j]); - let mut chain = None; + let vector = V::build_from_vecf32(&structures[i].means[j]); + let (metadata, slices) = V::vector_split(vector.as_borrowed()); + let mut chain = Err(metadata); for i in (0..slices.len()).rev() { - chain = Some(vectors.push(&VectorTuple { + chain = Ok(vectors.push(&VectorTuple { payload: None, slice: slices[i].to_vec(), chain, })); } - level.push(chain.unwrap()); + level.push(chain.ok().unwrap()); } pointer_of_means.push(level); } @@ -224,7 +226,7 @@ impl Structure { if vector_options.dims != vector.as_borrowed().dims() { pgrx::error!("extern build: incorrect dimension, id = {id}"); } - vectors.insert(id, crate::projection::project(vector.slice())); + vectors.insert(id, crate::projection::project(vector.as_borrowed().slice())); } }); let mut children = parents diff --git a/src/vchordrq/algorithm/insert.rs b/src/vchordrq/algorithm/insert.rs index c37f7c38..ee47e9a7 100644 --- a/src/vchordrq/algorithm/insert.rs +++ b/src/vchordrq/algorithm/insert.rs @@ -1,23 +1,23 @@ use crate::postgres::Relation; -use crate::vchordrq::algorithm::rabitq; use crate::vchordrq::algorithm::rabitq::fscan_process_lowerbound; use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; use base::always_equal::AlwaysEqual; use base::distance::Distance; use base::distance::DistanceKind; -use base::scalar::ScalarLike; use base::search::Pointer; +use base::vector::VectorBorrowed; use std::cmp::Reverse; use std::collections::BinaryHeap; -pub fn insert( +pub fn insert( relation: Relation, payload: Pointer, - vector: Vec, + vector: V, distance_kind: DistanceKind, in_building: bool, ) { + let vector = vector.as_borrowed(); let meta_guard = relation.read(0); let meta_tuple = meta_guard .get() @@ -26,25 +26,26 @@ pub fn insert( .expect("data corruption") .expect("data corruption"); let dims = meta_tuple.dims; - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); - let vector = crate::projection::project(&vector); + assert_eq!(dims, vector.dims(), "invalid vector dimensions"); + let vector = V::random_projection(vector); + let vector = vector.as_borrowed(); let is_residual = meta_tuple.is_residual; let default_lut = if !is_residual { - Some(rabitq::fscan_preprocess(&vector)) + Some(V::rabitq_fscan_preprocess(vector)) } else { None }; let h0_vector = { - let slices = vectors::vector_split(&vector); - let mut chain = None; + let (metadata, slices) = V::vector_split(vector); + let mut chain = Err(metadata); for i in (0..slices.len()).rev() { - let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { + let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple:: { slice: slices[i].to_vec(), payload: Some(payload.as_u64()), chain, }) .unwrap(); - chain = Some(append( + chain = Ok(append( relation.clone(), meta_tuple.vectors_first, &tuple, @@ -53,13 +54,13 @@ pub fn insert( true, )); } - chain.unwrap() + chain.ok().unwrap() }; let h0_payload = payload.as_u64(); let mut list = { - let Some((_, original)) = vectors::vector_dist( + let Some((_, original)) = vectors::vector_dist::( relation.clone(), - &vector, + vector, meta_tuple.mean, None, None, @@ -69,11 +70,14 @@ pub fn insert( }; (meta_tuple.first, original) }; - let make_list = |list: (u32, Option>)| { + let make_list = |list: (u32, Option)| { let mut results = Vec::new(); { let lut = if is_residual { - &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) + &V::rabitq_fscan_preprocess( + V::residual(vector, list.1.as_ref().map(|x| x.as_borrowed()).unwrap()) + .as_borrowed(), + ) } else { default_lut.as_ref().unwrap() }; @@ -114,9 +118,9 @@ pub fn insert( { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let Some((Some(dis_u), original)) = vectors::vector_dist( + let Some((Some(dis_u), original)) = vectors::vector_dist::( relation.clone(), - &vector, + vector, mean, None, Some(distance_kind), @@ -134,9 +138,12 @@ pub fn insert( list = make_list(list); } let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(&vector, list.1.as_ref().unwrap())) + V::rabitq_code( + dims, + V::residual(vector, list.1.as_ref().map(|x| x.as_borrowed()).unwrap()).as_borrowed(), + ) } else { - rabitq::code(dims, &vector) + V::rabitq_code(dims, vector) }; let tuple = rkyv::to_bytes::<_, 8192>(&Height0Tuple { mean: h0_vector, diff --git a/src/vchordrq/algorithm/prewarm.rs b/src/vchordrq/algorithm/prewarm.rs index 92441383..26bb9866 100644 --- a/src/vchordrq/algorithm/prewarm.rs +++ b/src/vchordrq/algorithm/prewarm.rs @@ -3,7 +3,7 @@ use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; use std::fmt::Write; -pub fn prewarm(relation: Relation, height: i32) -> String { +pub fn prewarm(relation: Relation, height: i32) -> String { let mut message = String::new(); let meta_guard = relation.read(0); let meta_tuple = meta_guard @@ -21,7 +21,7 @@ pub fn prewarm(relation: Relation, height: i32) -> String { let mut results = Vec::new(); let counter = 1_usize; { - vectors::vector_warm(relation.clone(), meta_tuple.mean); + vectors::vector_warm::(relation.clone(), meta_tuple.mean); results.push(meta_tuple.first); } writeln!(message, "number of tuples: {}", results.len()).unwrap(); @@ -44,7 +44,7 @@ pub fn prewarm(relation: Relation, height: i32) -> String { .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - vectors::vector_warm(relation.clone(), h1_tuple.mean); + vectors::vector_warm::(relation.clone(), h1_tuple.mean); results.push(h1_tuple.first); } current = h1_guard.get().get_opaque().next; diff --git a/src/vchordrq/algorithm/rabitq.rs b/src/vchordrq/algorithm/rabitq.rs index b3746b80..b7b3858d 100644 --- a/src/vchordrq/algorithm/rabitq.rs +++ b/src/vchordrq/algorithm/rabitq.rs @@ -1,5 +1,5 @@ use base::distance::{Distance, DistanceKind}; -use base::scalar::ScalarLike; +use base::simd::ScalarLike; #[derive(Debug, Clone)] pub struct Code { @@ -12,7 +12,7 @@ pub struct Code { impl Code { pub fn t(&self) -> Vec { - use quantization::utils::InfiniteByteChunks; + use crate::utils::infinite_byte_chunks::InfiniteByteChunks; let mut result = Vec::new(); for x in InfiniteByteChunks::<_, 64>::new(self.signs.iter().copied()) .take(self.signs.len().div_ceil(64)) @@ -62,13 +62,13 @@ pub fn code(dims: u32, vector: &[f32]) -> Code { pub type Lut = (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)); pub fn fscan_preprocess(vector: &[f32]) -> Lut { - use quantization::quantize; + use base::simd::quantize; let dis_v_2 = f32::reduce_sum_of_x2(vector); - let (k, b, qvector) = quantize::quantize::<15>(vector); + let (k, b, qvector) = quantize::quantize(vector, 15.0); let qvector_sum = if vector.len() <= 4369 { - quantize::reduce_sum_of_x_as_u16(&qvector) as f32 + base::simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 } else { - quantize::reduce_sum_of_x_as_u32(&qvector) as f32 + base::simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 }; (dis_v_2, b, k, qvector_sum, binarize(&qvector)) } @@ -119,34 +119,10 @@ fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { (t0, t1, t2, t3) } -#[detect::multiversion(v2, fallback)] fn asymmetric_binary_dot_product(x: &[u64], y: &(Vec, Vec, Vec, Vec)) -> u32 { - assert_eq!(x.len(), y.0.len()); - assert_eq!(x.len(), y.1.len()); - assert_eq!(x.len(), y.2.len()); - assert_eq!(x.len(), y.3.len()); - let n = x.len(); - let (mut t0, mut t1, mut t2, mut t3) = (0, 0, 0, 0); - for i in 0..n { - t0 += (x[i] & y.0[i]).count_ones(); - } - for i in 0..n { - t1 += (x[i] & y.1[i]).count_ones(); - } - for i in 0..n { - t2 += (x[i] & y.2[i]).count_ones(); - } - for i in 0..n { - t3 += (x[i] & y.3[i]).count_ones(); - } + let t0 = base::simd::bit::sum_of_and(x, &y.0); + let t1 = base::simd::bit::sum_of_and(x, &y.1); + let t2 = base::simd::bit::sum_of_and(x, &y.2); + let t3 = base::simd::bit::sum_of_and(x, &y.3); (t0 << 0) + (t1 << 1) + (t2 << 2) + (t3 << 3) } - -pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance { - match d { - DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)), - DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)), - DistanceKind::Hamming => unimplemented!(), - DistanceKind::Jaccard => unimplemented!(), - } -} diff --git a/src/vchordrq/algorithm/scan.rs b/src/vchordrq/algorithm/scan.rs index 14ecba1c..df4d93de 100644 --- a/src/vchordrq/algorithm/scan.rs +++ b/src/vchordrq/algorithm/scan.rs @@ -1,23 +1,23 @@ use crate::postgres::Relation; -use crate::vchordrq::algorithm::rabitq; use crate::vchordrq::algorithm::rabitq::fscan_process_lowerbound; use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; use base::always_equal::AlwaysEqual; use base::distance::Distance; use base::distance::DistanceKind; -use base::scalar::ScalarLike; use base::search::Pointer; +use base::vector::VectorBorrowed; use std::cmp::Reverse; use std::collections::BinaryHeap; -pub fn scan( +pub fn scan( relation: Relation, - vector: Vec, + vector: V, distance_kind: DistanceKind, probes: Vec, epsilon: f32, ) -> impl Iterator { + let vector = vector.as_borrowed(); let meta_guard = relation.read(0); let meta_tuple = meta_guard .get() @@ -27,19 +27,19 @@ pub fn scan( .expect("data corruption"); let dims = meta_tuple.dims; let height_of_root = meta_tuple.height_of_root; - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); + assert_eq!(dims, vector.dims(), "invalid vector dimensions"); assert_eq!(height_of_root as usize, 1 + probes.len(), "invalid probes"); - let vector = crate::projection::project(&vector); + let vector = V::random_projection(vector); let is_residual = meta_tuple.is_residual; let default_lut = if !is_residual { - Some(rabitq::fscan_preprocess(&vector)) + Some(V::rabitq_fscan_preprocess(vector.as_borrowed())) } else { None }; let mut lists: Vec<_> = vec![{ - let Some((_, original)) = vectors::vector_dist( + let Some((_, original)) = vectors::vector_dist::( relation.clone(), - &vector, + vector.as_borrowed(), meta_tuple.mean, None, None, @@ -49,11 +49,17 @@ pub fn scan( }; (meta_tuple.first, original) }]; - let make_lists = |lists: Vec<(u32, Option>)>, probes| { + let make_lists = |lists: Vec<(u32, Option)>, probes| { let mut results = Vec::new(); for list in lists { let lut = if is_residual { - &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) + &V::rabitq_fscan_preprocess( + V::residual( + vector.as_borrowed(), + list.1.as_ref().map(|x| x.as_borrowed()).unwrap(), + ) + .as_borrowed(), + ) } else { default_lut.as_ref().unwrap() }; @@ -94,9 +100,9 @@ pub fn scan( std::iter::from_fn(|| { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let Some((Some(dis_u), original)) = vectors::vector_dist( + let Some((Some(dis_u), original)) = vectors::vector_dist::( relation.clone(), - &vector, + vector.as_borrowed(), mean, None, Some(distance_kind), @@ -119,7 +125,13 @@ pub fn scan( let mut results = Vec::new(); for list in lists { let lut = if is_residual { - &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) + &V::rabitq_fscan_preprocess( + V::residual( + vector.as_borrowed(), + list.1.as_ref().map(|x| x.as_borrowed()).unwrap(), + ) + .as_borrowed(), + ) } else { default_lut.as_ref().unwrap() }; @@ -160,9 +172,9 @@ pub fn scan( std::iter::from_fn(move || { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); - let Some((Some(dis_u), _)) = vectors::vector_dist( + let Some((Some(dis_u), _)) = vectors::vector_dist::( relation.clone(), - &vector, + vector.as_borrowed(), mean, Some(pay_u), Some(distance_kind), diff --git a/src/vchordrq/algorithm/tuples.rs b/src/vchordrq/algorithm/tuples.rs index cf6236f0..40a795ce 100644 --- a/src/vchordrq/algorithm/tuples.rs +++ b/src/vchordrq/algorithm/tuples.rs @@ -1,4 +1,233 @@ -use rkyv::{Archive, Deserialize, Serialize}; +use super::rabitq::{self, Code, Lut}; +use crate::vchordrq::types::OwnedVector; +use base::distance::DistanceKind; +use base::simd::ScalarLike; +use base::vector::{VectOwned, VectorOwned}; +use half::f16; +use rkyv::{Archive, ArchiveUnsized, CheckBytes, Deserialize, Serialize}; + +pub trait Vector: VectorOwned { + type Metadata: Copy + + Serialize< + rkyv::ser::serializers::CompositeSerializer< + rkyv::ser::serializers::AlignedSerializer, + rkyv::ser::serializers::FallbackScratch< + rkyv::ser::serializers::HeapScratch<8192>, + rkyv::ser::serializers::AllocScratch, + >, + rkyv::ser::serializers::SharedSerializeMap, + >, + > + for<'a> CheckBytes>; + type Element: Copy + + Serialize< + rkyv::ser::serializers::CompositeSerializer< + rkyv::ser::serializers::AlignedSerializer, + rkyv::ser::serializers::FallbackScratch< + rkyv::ser::serializers::HeapScratch<8192>, + rkyv::ser::serializers::AllocScratch, + >, + rkyv::ser::serializers::SharedSerializeMap, + >, + > + for<'a> CheckBytes> + + Archive; + + fn metadata_from_archived( + archived: &::Archived, + ) -> Self::Metadata; + + fn vector_split(vector: Self::Borrowed<'_>) -> (Self::Metadata, Vec<&[Self::Element]>); + fn vector_merge(metadata: Self::Metadata, slice: &[Self::Element]) -> Self; + fn from_owned(vector: OwnedVector) -> Self; + + type DistanceAccumulator; + fn distance_begin(distance_kind: DistanceKind) -> Self::DistanceAccumulator; + fn distance_next( + accumulator: &mut Self::DistanceAccumulator, + left: &[Self::Element], + right: &[Self::Element], + ); + fn distance_end( + accumulator: Self::DistanceAccumulator, + left: Self::Metadata, + right: Self::Metadata, + ) -> f32; + + fn random_projection(vector: Self::Borrowed<'_>) -> Self; + + fn residual(vector: Self::Borrowed<'_>, center: Self::Borrowed<'_>) -> Self; + + fn rabitq_fscan_preprocess(vector: Self::Borrowed<'_>) -> Lut; + + fn rabitq_code(dims: u32, vector: Self::Borrowed<'_>) -> Code; + + fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec; + + fn build_from_vecf32(x: &[f32]) -> Self; +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f32; + + fn metadata_from_archived(_: &::Archived) -> Self::Metadata { + () + } + + fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f32]>) { + let vector = vector.slice(); + ( + (), + match vector.len() { + 0..=960 => vec![vector], + 961..=1280 => vec![&vector[..640], &vector[640..]], + 1281.. => vector.chunks(1920).collect(), + }, + ) + } + + fn vector_merge((): Self::Metadata, slice: &[Self::Element]) -> Self { + VectOwned::new(slice.to_vec()) + } + + fn from_owned(vector: OwnedVector) -> Self { + match vector { + OwnedVector::Vecf32(x) => x, + _ => unreachable!(), + } + } + + type DistanceAccumulator = (DistanceKind, f32); + fn distance_begin(distance_kind: DistanceKind) -> Self::DistanceAccumulator { + (distance_kind, 0.0) + } + fn distance_next( + accumulator: &mut Self::DistanceAccumulator, + left: &[Self::Element], + right: &[Self::Element], + ) { + match accumulator.0 { + DistanceKind::L2 => accumulator.1 += f32::reduce_sum_of_d2(left, right), + DistanceKind::Dot => accumulator.1 += -f32::reduce_sum_of_xy(left, right), + DistanceKind::Hamming => unreachable!(), + DistanceKind::Jaccard => unreachable!(), + } + } + fn distance_end( + accumulator: Self::DistanceAccumulator, + (): Self::Metadata, + (): Self::Metadata, + ) -> f32 { + accumulator.1 + } + + fn random_projection(vector: Self::Borrowed<'_>) -> Self { + Self::new(crate::projection::project(vector.slice())) + } + + fn residual(vector: Self::Borrowed<'_>, center: Self::Borrowed<'_>) -> Self { + Self::new(ScalarLike::vector_sub(vector.slice(), center.slice())) + } + + fn rabitq_fscan_preprocess(vector: Self::Borrowed<'_>) -> Lut { + rabitq::fscan_preprocess(vector.slice()) + } + + fn rabitq_code(dims: u32, vector: Self::Borrowed<'_>) -> Code { + rabitq::code(dims, vector.slice()) + } + + fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { + vector.slice().to_vec() + } + + fn build_from_vecf32(x: &[f32]) -> Self { + Self::new(x.to_vec()) + } +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f16; + + fn metadata_from_archived(_: &::Archived) -> Self::Metadata { + () + } + + fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f16]>) { + let vector = vector.slice(); + ( + (), + match vector.len() { + 0..=1920 => vec![vector], + 1921..=2560 => vec![&vector[..1280], &vector[1280..]], + 2561.. => vector.chunks(3840).collect(), + }, + ) + } + + fn vector_merge((): Self::Metadata, slice: &[Self::Element]) -> Self { + VectOwned::new(slice.to_vec()) + } + + fn from_owned(vector: OwnedVector) -> Self { + match vector { + OwnedVector::Vecf16(x) => x, + _ => unreachable!(), + } + } + + type DistanceAccumulator = (DistanceKind, f32); + fn distance_begin(distance_kind: DistanceKind) -> Self::DistanceAccumulator { + (distance_kind, 0.0) + } + fn distance_next( + accumulator: &mut Self::DistanceAccumulator, + left: &[Self::Element], + right: &[Self::Element], + ) { + match accumulator.0 { + DistanceKind::L2 => accumulator.1 += f16::reduce_sum_of_d2(left, right), + DistanceKind::Dot => accumulator.1 += -f16::reduce_sum_of_xy(left, right), + DistanceKind::Hamming => unreachable!(), + DistanceKind::Jaccard => unreachable!(), + } + } + fn distance_end( + accumulator: Self::DistanceAccumulator, + (): Self::Metadata, + (): Self::Metadata, + ) -> f32 { + accumulator.1 + } + + fn random_projection(vector: Self::Borrowed<'_>) -> Self { + Self::new(f16::vector_from_f32(&crate::projection::project( + &f16::vector_to_f32(vector.slice()), + ))) + } + + fn residual(vector: Self::Borrowed<'_>, center: Self::Borrowed<'_>) -> Self { + Self::new(ScalarLike::vector_sub(vector.slice(), center.slice())) + } + + fn rabitq_fscan_preprocess(vector: Self::Borrowed<'_>) -> Lut { + rabitq::fscan_preprocess(&f16::vector_to_f32(vector.slice())) + } + + fn rabitq_code(dims: u32, vector: Self::Borrowed<'_>) -> Code { + rabitq::code(dims, &f16::vector_to_f32(vector.slice())) + } + + fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { + f16::vector_to_f32(vector.slice()) + } + + fn build_from_vecf32(x: &[f32]) -> Self { + Self::new(f16::vector_from_f32(x)) + } +} #[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] #[archive(check_bytes)] @@ -15,10 +244,10 @@ pub struct MetaTuple { #[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] #[archive(check_bytes)] -pub struct VectorTuple { - pub slice: Vec, +pub struct VectorTuple { + pub slice: Vec, pub payload: Option, - pub chain: Option<(u32, u16)>, + pub chain: Result<(u32, u16), V::Metadata>, } #[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] diff --git a/src/vchordrq/algorithm/vacuum.rs b/src/vchordrq/algorithm/vacuum.rs index 27377020..2b219c49 100644 --- a/src/vchordrq/algorithm/vacuum.rs +++ b/src/vchordrq/algorithm/vacuum.rs @@ -2,7 +2,7 @@ use crate::postgres::Relation; use crate::vchordrq::algorithm::tuples::*; use base::search::Pointer; -pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) -> bool) { +pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) -> bool) { // step 1: vacuum height_0_tuple { let meta_guard = relation.read(0); @@ -78,8 +78,8 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - let Some(vector_tuple) = read.get().get(i) else { continue; }; - let vector_tuple = rkyv::check_archived_root::(vector_tuple) - .expect("data corruption"); + let vector_tuple = + unsafe { rkyv::archived_root::>(vector_tuple) }; if let Some(payload) = vector_tuple.payload.as_ref().copied() { if callback(Pointer::new(payload)) { break 'flag true; @@ -95,8 +95,8 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - let Some(vector_tuple) = write.get().get(i) else { continue; }; - let vector_tuple = rkyv::check_archived_root::(vector_tuple) - .expect("data corruption"); + let vector_tuple = + unsafe { rkyv::archived_root::>(vector_tuple) }; if let Some(payload) = vector_tuple.payload.as_ref().copied() { if callback(Pointer::new(payload)) { write.get_mut().free(i); diff --git a/src/vchordrq/algorithm/vectors.rs b/src/vchordrq/algorithm/vectors.rs index c1f86273..6a23f746 100644 --- a/src/vchordrq/algorithm/vectors.rs +++ b/src/vchordrq/algorithm/vectors.rs @@ -1,34 +1,26 @@ +use super::tuples::Vector; use crate::postgres::Relation; -use crate::vchordrq::algorithm::rabitq::distance; use crate::vchordrq::algorithm::tuples::VectorTuple; use base::distance::Distance; use base::distance::DistanceKind; -pub fn vector_split(vector: &[f32]) -> Vec<&[f32]> { - match vector.len() { - 0..=960 => vec![vector], - 961..=1280 => vec![&vector[..640], &vector[640..]], - 1281.. => vector.chunks(1920).collect(), - } -} - -pub fn vector_dist( +pub fn vector_dist( relation: Relation, - vector: &[f32], + vector: V::Borrowed<'_>, mean: (u32, u16), payload: Option, for_distance: Option, for_original: bool, -) -> Option<(Option, Option>)> { +) -> Option<(Option, Option)> { if for_distance.is_none() && !for_original && payload.is_none() { return Some((None, None)); } - let slices = vector_split(vector); - let mut cursor = Some(mean); - let mut result = 0.0f32; + let (left_metadata, slices) = V::vector_split(vector); + let mut cursor = Ok(mean); + let mut result = for_distance.map(|x| V::distance_begin(x)); let mut original = Vec::new(); for i in 0..slices.len() { - let Some(mean) = cursor else { + let Ok(mean) = cursor else { // fails consistency check return None; }; @@ -37,40 +29,47 @@ pub fn vector_dist( // fails consistency check return None; }; - let vector_tuple = - rkyv::check_archived_root::(vector_tuple).expect("data corruption"); + let vector_tuple = unsafe { rkyv::archived_root::>(vector_tuple) }; if vector_tuple.payload != payload { // fails consistency check return None; } - if let Some(distance_kind) = for_distance { - result += distance(distance_kind, slices[i], &vector_tuple.slice).to_f32(); + if let Some(result) = result.as_mut() { + V::distance_next(result, slices[i], &vector_tuple.slice); } if for_original { original.extend_from_slice(&vector_tuple.slice); } - cursor = vector_tuple.chain.as_ref().cloned(); + cursor = match &vector_tuple.chain { + rkyv::result::ArchivedResult::Ok(x) => Ok(*x), + rkyv::result::ArchivedResult::Err(x) => Err(V::metadata_from_archived(x)), + }; } + let Err(right_metadata) = cursor else { + panic!("data corruption") + }; Some(( - for_distance.map(|_| Distance::from_f32(result)), - for_original.then_some(original), + result.map(|r| Distance::from_f32(V::distance_end(r, left_metadata, right_metadata))), + for_original.then(|| V::vector_merge(right_metadata, &original)), )) } -pub fn vector_warm(relation: Relation, mean: (u32, u16)) { - let mut cursor = Some(mean); - while let Some(mean) = cursor { +pub fn vector_warm(relation: Relation, mean: (u32, u16)) { + let mut cursor = Ok(mean); + while let Ok(mean) = cursor { let vector_guard = relation.read(mean.0); let Some(vector_tuple) = vector_guard.get().get(mean.1) else { // fails consistency check return; }; - let vector_tuple = - rkyv::check_archived_root::(vector_tuple).expect("data corruption"); + let vector_tuple = unsafe { rkyv::archived_root::>(vector_tuple) }; if vector_tuple.payload.is_some() { // fails consistency check return; } - cursor = vector_tuple.chain.as_ref().cloned(); + cursor = match &vector_tuple.chain { + rkyv::result::ArchivedResult::Ok(x) => Ok(*x), + rkyv::result::ArchivedResult::Err(x) => Err(V::metadata_from_archived(x)), + }; } } diff --git a/src/vchordrq/index/am.rs b/src/vchordrq/index/am.rs index d3c3e8b2..e9182c7a 100644 --- a/src/vchordrq/index/am.rs +++ b/src/vchordrq/index/am.rs @@ -1,11 +1,15 @@ use crate::postgres::Relation; use crate::vchordrq::algorithm; use crate::vchordrq::algorithm::build::{HeapRelation, Reporter}; +use crate::vchordrq::algorithm::tuples::Vector; use crate::vchordrq::index::am_options::{Opfamily, Reloption}; use crate::vchordrq::index::am_scan::Scanner; use crate::vchordrq::index::utils::{ctid_to_pointer, pointer_to_ctid}; use crate::vchordrq::index::{am_options, am_scan}; +use crate::vchordrq::types::VectorKind; use base::search::Pointer; +use base::vector::VectOwned; +use half::f16; use pgrx::datum::Internal; use pgrx::pg_sys::Datum; @@ -163,17 +167,17 @@ pub unsafe extern "C" fn ambuild( index_info: *mut pgrx::pg_sys::IndexInfo, opfamily: Opfamily, } - impl HeapRelation for Heap { + impl HeapRelation for Heap { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((Pointer, Vec)), + F: FnMut((Pointer, V)), { pub struct State<'a, F> { pub this: &'a Heap, pub callback: F, } #[pgrx::pg_guard] - unsafe extern "C" fn call( + unsafe extern "C" fn call( _index: pgrx::pg_sys::Relation, ctid: pgrx::pg_sys::ItemPointer, values: *mut Datum, @@ -181,21 +185,14 @@ pub unsafe extern "C" fn ambuild( _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((Pointer, Vec)), + F: FnMut((Pointer, V)), { - use base::vector::OwnedVector; let state = unsafe { &mut *state.cast::>() }; let opfamily = state.this.opfamily; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), - }; - (state.callback)((pointer, vector.into_vec())); + (state.callback)((pointer, V::from_owned(vector))); } } let table_am = unsafe { &*(*self.heap).rd_tableam }; @@ -213,7 +210,7 @@ pub unsafe extern "C" fn ambuild( progress, 0, pgrx::pg_sys::InvalidBlockNumber, - Some(call::), + Some(call::), (&mut state) as *mut State as *mut _, std::ptr::null_mut(), ); @@ -246,13 +243,22 @@ pub unsafe extern "C" fn ambuild( }; let mut reporter = PgReporter {}; let index_relation = unsafe { Relation::new(index) }; - algorithm::build::build( - vector_options, - vchordrq_options, - heap_relation.clone(), - index_relation.clone(), - reporter.clone(), - ); + match opfamily.vector_kind() { + VectorKind::Vecf32 => algorithm::build::build::, Heap, _>( + vector_options, + vchordrq_options, + heap_relation.clone(), + index_relation.clone(), + reporter.clone(), + ), + VectorKind::Vecf16 => algorithm::build::build::, Heap, _>( + vector_options, + vchordrq_options, + heap_relation.clone(), + index_relation.clone(), + reporter.clone(), + ), + } if let Some(leader) = unsafe { VchordrqLeader::enter(heap, index, (*index_info).ii_Concurrent) } { unsafe { @@ -283,17 +289,42 @@ pub unsafe extern "C" fn ambuild( } else { let mut indtuples = 0; reporter.tuples_done(indtuples); - heap_relation.traverse(true, |(payload, vector)| { - algorithm::insert::insert( - index_relation.clone(), - payload, - vector, - opfamily.distance_kind(), - true, - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }); + match opfamily.vector_kind() { + VectorKind::Vecf32 => { + HeapRelation::>::traverse( + &heap_relation, + true, + |(pointer, vector)| { + algorithm::insert::insert::>( + unsafe { Relation::new(index) }, + pointer, + vector, + opfamily.distance_kind(), + true, + ); + indtuples += 1; + reporter.tuples_done(indtuples); + }, + ); + } + VectorKind::Vecf16 => { + HeapRelation::>::traverse( + &heap_relation, + true, + |(pointer, vector)| { + algorithm::insert::insert::>( + unsafe { Relation::new(index) }, + pointer, + vector, + opfamily.distance_kind(), + true, + ); + indtuples += 1; + reporter.tuples_done(indtuples); + }, + ); + } + } } unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -540,17 +571,17 @@ unsafe fn parallel_build( opfamily: Opfamily, scan: *mut pgrx::pg_sys::TableScanDescData, } - impl HeapRelation for Heap { + impl HeapRelation for Heap { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((Pointer, Vec)), + F: FnMut((Pointer, V)), { pub struct State<'a, F> { pub this: &'a Heap, pub callback: F, } #[pgrx::pg_guard] - unsafe extern "C" fn call( + unsafe extern "C" fn call( _index: pgrx::pg_sys::Relation, ctid: pgrx::pg_sys::ItemPointer, values: *mut Datum, @@ -558,21 +589,14 @@ unsafe fn parallel_build( _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((Pointer, Vec)), + F: FnMut((Pointer, V)), { - use base::vector::OwnedVector; let state = unsafe { &mut *state.cast::>() }; let opfamily = state.this.opfamily; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), - }; - (state.callback)((pointer, vector.into_vec())); + (state.callback)((pointer, V::from_owned(vector))); } } let table_am = unsafe { &*(*self.heap).rd_tableam }; @@ -590,7 +614,7 @@ unsafe fn parallel_build( progress, 0, pgrx::pg_sys::InvalidBlockNumber, - Some(call::), + Some(call::), (&mut state) as *mut State as *mut _, self.scan, ); @@ -612,28 +636,54 @@ unsafe fn parallel_build( opfamily, scan, }; - heap_relation.traverse(reporter.is_some(), |(payload, vector)| { - algorithm::insert::insert( - index_relation.clone(), - payload, - vector, - opfamily.distance_kind(), - true, - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } + match opfamily.vector_kind() { + VectorKind::Vecf32 => { + HeapRelation::>::traverse(&heap_relation, true, |(pointer, vector)| { + algorithm::insert::insert::>( + index_relation.clone(), + pointer, + vector, + opfamily.distance_kind(), + true, + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); } - }); - + VectorKind::Vecf16 => { + HeapRelation::>::traverse(&heap_relation, true, |(pointer, vector)| { + algorithm::insert::insert::>( + index_relation.clone(), + pointer, + vector, + opfamily.distance_kind(), + true, + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); + } + } unsafe { pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); (*vchordrqshared).nparticipantsdone += 1; @@ -658,23 +708,26 @@ pub unsafe extern "C" fn aminsert( _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, _index_info: *mut pgrx::pg_sys::IndexInfo, ) -> bool { - use base::vector::OwnedVector; let opfamily = unsafe { am_options::opfamily(index) }; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), - }; let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - algorithm::insert::insert( - unsafe { Relation::new(index) }, - pointer, - vector.into_vec(), - opfamily.distance_kind(), - ); + match opfamily.vector_kind() { + VectorKind::Vecf32 => algorithm::insert::insert::>( + unsafe { Relation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + opfamily.distance_kind(), + false, + ), + VectorKind::Vecf16 => algorithm::insert::insert::>( + unsafe { Relation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + opfamily.distance_kind(), + false, + ), + } } false } @@ -691,24 +744,26 @@ pub unsafe extern "C" fn aminsert( _index_unchanged: bool, _index_info: *mut pgrx::pg_sys::IndexInfo, ) -> bool { - use base::vector::OwnedVector; let opfamily = unsafe { am_options::opfamily(index) }; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), - }; let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - algorithm::insert::insert( - unsafe { Relation::new(index) }, - pointer, - vector.into_vec(), - opfamily.distance_kind(), - false, - ); + match opfamily.vector_kind() { + VectorKind::Vecf32 => algorithm::insert::insert::>( + unsafe { Relation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + opfamily.distance_kind(), + false, + ), + VectorKind::Vecf16 => algorithm::insert::insert::>( + unsafe { Relation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + opfamily.distance_kind(), + false, + ), + } } false } @@ -833,15 +888,25 @@ pub unsafe extern "C" fn ambulkdelete( pgrx::pg_sys::palloc0(size_of::()).cast() }; } + let opfamily = unsafe { am_options::opfamily((*info).index) }; let callback = callback.unwrap(); let callback = |p: Pointer| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; - algorithm::vacuum::vacuum( - unsafe { Relation::new((*info).index) }, - || unsafe { - pgrx::pg_sys::vacuum_delay_point(); - }, - callback, - ); + match opfamily.vector_kind() { + VectorKind::Vecf32 => algorithm::vacuum::vacuum::>( + unsafe { Relation::new((*info).index) }, + || unsafe { + pgrx::pg_sys::vacuum_delay_point(); + }, + callback, + ), + VectorKind::Vecf16 => algorithm::vacuum::vacuum::>( + unsafe { Relation::new((*info).index) }, + || unsafe { + pgrx::pg_sys::vacuum_delay_point(); + }, + callback, + ), + } stats } diff --git a/src/vchordrq/index/am_options.rs b/src/vchordrq/index/am_options.rs index 971273fd..a357da40 100644 --- a/src/vchordrq/index/am_options.rs +++ b/src/vchordrq/index/am_options.rs @@ -1,10 +1,13 @@ +use crate::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput; +use crate::datatype::memory_pgvector_halfvec::PgvectorHalfvecOutput; use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use crate::datatype::typmod::Typmod; use crate::vchordrq::types::VchordrqIndexingOptions; +use crate::vchordrq::types::VectorOptions; +use crate::vchordrq::types::{BorrowedVector, OwnedVector, VectorKind}; use base::distance::*; -use base::index::*; -use base::vector::*; +use base::vector::VectorBorrowed; use pgrx::datum::FromDatum; use pgrx::heap_tuple::PgHeapTuple; use serde::Deserialize; @@ -26,7 +29,7 @@ impl Reloption { }]; unsafe fn options(&self) -> &CStr { unsafe { - let ptr = std::ptr::addr_of!(*self) + let ptr = (&raw const *self) .cast::() .offset(self.options as _); CStr::from_ptr(ptr) @@ -56,6 +59,9 @@ fn convert_name_to_vd(name: &str) -> Option<(VectorKind, PgDistanceKind)> { Some("vector_l2") => Some((VectorKind::Vecf32, PgDistanceKind::L2)), Some("vector_ip") => Some((VectorKind::Vecf32, PgDistanceKind::Dot)), Some("vector_cosine") => Some((VectorKind::Vecf32, PgDistanceKind::Cos)), + Some("halfvec_l2") => Some((VectorKind::Vecf16, PgDistanceKind::L2)), + Some("halfvec_ip") => Some((VectorKind::Vecf16, PgDistanceKind::Dot)), + Some("halfvec_cosine") => Some((VectorKind::Vecf16, PgDistanceKind::Cos)), _ => None, } } @@ -130,7 +136,10 @@ impl Opfamily { let vector = unsafe { PgvectorVectorInput::from_datum(datum, false).unwrap() }; self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed())) } - _ => unreachable!(), + VectorKind::Vecf16 => { + let vector = unsafe { PgvectorHalfvecInput::from_datum(datum, false).unwrap() }; + self.preprocess(BorrowedVector::Vecf16(vector.as_borrowed())) + } }; Some(vector) } @@ -148,7 +157,10 @@ impl Opfamily { .get_by_index::(NonZero::new(1).unwrap()) .unwrap() .map(|vector| self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed()))), - _ => unreachable!(), + VectorKind::Vecf16 => tuple + .get_by_index::(NonZero::new(1).unwrap()) + .unwrap() + .map(|vector| self.preprocess(BorrowedVector::Vecf16(vector.as_borrowed()))), }; let radius = tuple.get_by_index::(NonZero::new(2).unwrap()).unwrap(); (center, radius) @@ -161,8 +173,6 @@ impl Opfamily { (B::Vecf32(x), PgDistanceKind::Dot) => O::Vecf32(x.own()), (B::Vecf32(x), PgDistanceKind::Cos) => O::Vecf32(x.function_normalize()), (B::Vecf16(x), _) => O::Vecf16(x.own()), - (B::SVecf32(x), _) => O::SVecf32(x.own()), - (B::BVector(x), _) => O::BVector(x.own()), } } pub fn process(self, x: Distance) -> f32 { @@ -175,6 +185,9 @@ impl Opfamily { pub fn distance_kind(self) -> DistanceKind { self.pg_distance.to_distance() } + pub fn vector_kind(self) -> VectorKind { + self.vector + } } pub unsafe fn opfamily(index: pgrx::pg_sys::Relation) -> Opfamily { diff --git a/src/vchordrq/index/am_scan.rs b/src/vchordrq/index/am_scan.rs index e97d352d..1b78ff08 100644 --- a/src/vchordrq/index/am_scan.rs +++ b/src/vchordrq/index/am_scan.rs @@ -1,12 +1,16 @@ use super::am_options::Opfamily; use crate::postgres::Relation; use crate::vchordrq::algorithm::scan::scan; +use crate::vchordrq::algorithm::tuples::Vector; use crate::vchordrq::gucs::executing::epsilon; use crate::vchordrq::gucs::executing::max_scan_tuples; use crate::vchordrq::gucs::executing::probes; +use crate::vchordrq::types::OwnedVector; +use crate::vchordrq::types::VectorKind; use base::distance::Distance; use base::search::*; -use base::vector::*; +use base::vector::VectOwned; +use half::f16; pub enum Scanner { Initial { @@ -34,7 +38,7 @@ pub fn scan_build( for orderby_vector in orderbys { if pair.is_none() { pair = orderby_vector; - } else if orderby_vector.is_some() && pair != orderby_vector { + } else if orderby_vector.is_some() { pgrx::error!("vector search with multiple vectors is not supported"); } } @@ -42,10 +46,6 @@ pub fn scan_build( if pair.is_none() { pair = sphere_vector; threshold = sphere_threshold; - } else if pair == sphere_vector { - if threshold.is_none() || sphere_threshold < threshold { - threshold = sphere_threshold; - } } else { recheck = true; break; @@ -74,28 +74,46 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, } = scanner { if let Some((vector, opfamily)) = vector.as_ref() { - let vbase = scan( - relation, - match vector { - OwnedVector::Vecf32(x) => x.slice().to_vec(), - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), - }, - opfamily.distance_kind(), - probes(), - epsilon(), - ); - *scanner = Scanner::Vbase { - vbase: if let Some(max_scan_tuples) = max_scan_tuples() { - Box::new(vbase.take(max_scan_tuples as usize)) - } else { - Box::new(vbase) - }, - threshold: *threshold, - recheck: *recheck, - opfamily: *opfamily, - }; + match opfamily.vector_kind() { + VectorKind::Vecf32 => { + let vbase = scan::>( + relation, + VectOwned::::from_owned(vector.clone()), + opfamily.distance_kind(), + probes(), + epsilon(), + ); + *scanner = Scanner::Vbase { + vbase: if let Some(max_scan_tuples) = max_scan_tuples() { + Box::new(vbase.take(max_scan_tuples as usize)) + } else { + Box::new(vbase) + }, + threshold: *threshold, + recheck: *recheck, + opfamily: *opfamily, + }; + } + VectorKind::Vecf16 => { + let vbase = scan::>( + relation, + VectOwned::::from_owned(vector.clone()), + opfamily.distance_kind(), + probes(), + epsilon(), + ); + *scanner = Scanner::Vbase { + vbase: if let Some(max_scan_tuples) = max_scan_tuples() { + Box::new(vbase.take(max_scan_tuples as usize)) + } else { + Box::new(vbase) + }, + threshold: *threshold, + recheck: *recheck, + opfamily: *opfamily, + }; + } + } } else { *scanner = Scanner::Empty {}; } diff --git a/src/vchordrq/index/functions.rs b/src/vchordrq/index/functions.rs index 40df8b53..05f348f5 100644 --- a/src/vchordrq/index/functions.rs +++ b/src/vchordrq/index/functions.rs @@ -1,5 +1,9 @@ +use super::am_options; use crate::postgres::Relation; use crate::vchordrq::algorithm::prewarm::prewarm; +use crate::vchordrq::types::VectorKind; +use base::vector::VectOwned; +use half::f16; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass}; @@ -18,7 +22,11 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { } let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; let relation = unsafe { Relation::new(index) }; - let message = prewarm(relation, height); + let opfamily = unsafe { am_options::opfamily(index) }; + let message = match opfamily.vector_kind() { + VectorKind::Vecf32 => prewarm::>(relation, height), + VectorKind::Vecf16 => prewarm::>(relation, height), + }; unsafe { pgrx::pg_sys::index_close(index, pgrx::pg_sys::ShareLock as _); } diff --git a/src/vchordrq/index/opclass.rs b/src/vchordrq/index/opclass.rs index e71da44a..a2dc8618 100644 --- a/src/vchordrq/index/opclass.rs +++ b/src/vchordrq/index/opclass.rs @@ -12,3 +12,18 @@ fn _vchordrq_support_vector_ip_ops() -> String { fn _vchordrq_support_vector_cosine_ops() -> String { "vector_cosine_ops".to_string() } + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_halfvec_l2_ops() -> String { + "halfvec_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_halfvec_ip_ops() -> String { + "halfvec_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_halfvec_cosine_ops() -> String { + "halfvec_cosine_ops".to_string() +} diff --git a/src/vchordrq/types.rs b/src/vchordrq/types.rs index 23011509..0e1bdc07 100644 --- a/src/vchordrq/types.rs +++ b/src/vchordrq/types.rs @@ -1,3 +1,6 @@ +use base::distance::DistanceKind; +use base::vector::{VectBorrowed, VectOwned}; +use half::f16; use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError, ValidationErrors}; @@ -95,3 +98,47 @@ impl VchordrqIndexingOptions { false } } + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum OwnedVector { + Vecf32(VectOwned), + Vecf16(VectOwned), +} + +#[derive(Debug, Clone, Copy)] +pub enum BorrowedVector<'a> { + Vecf32(VectBorrowed<'a, f32>), + Vecf16(VectBorrowed<'a, f16>), +} + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +pub enum VectorKind { + Vecf32, + Vecf16, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +#[validate(schema(function = "Self::validate_self"))] +pub struct VectorOptions { + #[validate(range(min = 1, max = 1_048_575))] + #[serde(rename = "dimensions")] + pub dims: u32, + #[serde(rename = "vector")] + pub v: VectorKind, + #[serde(rename = "distance")] + pub d: DistanceKind, +} + +impl VectorOptions { + pub fn validate_self(&self) -> Result<(), ValidationError> { + match (self.v, self.d, self.dims) { + (VectorKind::Vecf32, DistanceKind::L2, 1..65536) => Ok(()), + (VectorKind::Vecf32, DistanceKind::Dot, 1..65536) => Ok(()), + (VectorKind::Vecf16, DistanceKind::L2, 1..65536) => Ok(()), + (VectorKind::Vecf16, DistanceKind::Dot, 1..65536) => Ok(()), + _ => Err(ValidationError::new("not valid vector options")), + } + } +} diff --git a/src/vchordrqfscan/algorithm/build.rs b/src/vchordrqfscan/algorithm/build.rs index 0528a5c4..25a66119 100644 --- a/src/vchordrqfscan/algorithm/build.rs +++ b/src/vchordrqfscan/algorithm/build.rs @@ -7,10 +7,10 @@ use crate::vchordrqfscan::types::VchordrqfscanBuildOptions; use crate::vchordrqfscan::types::VchordrqfscanExternalBuildOptions; use crate::vchordrqfscan::types::VchordrqfscanIndexingOptions; use crate::vchordrqfscan::types::VchordrqfscanInternalBuildOptions; +use crate::vchordrqfscan::types::VectorOptions; use base::distance::DistanceKind; -use base::index::VectorOptions; -use base::scalar::ScalarLike; use base::search::Pointer; +use base::simd::ScalarLike; use rand::Rng; use rkyv::ser::serializers::AllocSerializer; use std::marker::PhantomData; @@ -268,7 +268,7 @@ impl Structure { if vector_options.dims != vector.as_borrowed().dims() { pgrx::error!("extern build: incorrect dimension, id = {id}"); } - vectors.insert(id, crate::projection::project(vector.slice())); + vectors.insert(id, crate::projection::project(vector.as_borrowed().slice())); } }); let mut children = parents diff --git a/src/vchordrqfscan/algorithm/insert.rs b/src/vchordrqfscan/algorithm/insert.rs index e89b1101..4dfd432f 100644 --- a/src/vchordrqfscan/algorithm/insert.rs +++ b/src/vchordrqfscan/algorithm/insert.rs @@ -6,8 +6,8 @@ use crate::vchordrqfscan::algorithm::tuples::*; use base::always_equal::AlwaysEqual; use base::distance::Distance; use base::distance::DistanceKind; -use base::scalar::ScalarLike; use base::search::Pointer; +use base::simd::ScalarLike; use std::cmp::Reverse; use std::collections::BinaryHeap; diff --git a/src/vchordrqfscan/algorithm/rabitq.rs b/src/vchordrqfscan/algorithm/rabitq.rs index cf72ca54..65c4996e 100644 --- a/src/vchordrqfscan/algorithm/rabitq.rs +++ b/src/vchordrqfscan/algorithm/rabitq.rs @@ -1,6 +1,6 @@ +use crate::utils::infinite_byte_chunks::InfiniteByteChunks; use base::distance::{Distance, DistanceKind}; -use base::scalar::ScalarLike; -use quantization::utils::InfiniteByteChunks; +use base::simd::ScalarLike; #[derive(Debug, Clone)] pub struct Code { @@ -74,19 +74,19 @@ pub fn pack_codes(dims: u32, codes: [Code; 32]) -> PackedCodes { .take(dims.div_ceil(4) as usize) .collect::>() }); - quantization::fast_scan::b4::pack(dims.div_ceil(4), signs).collect() + base::simd::fast_scan::b4::pack(dims.div_ceil(4), signs).collect() }, } } pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { - use quantization::quantize; + use base::simd::quantize; let dis_v_2 = f32::reduce_sum_of_x2(vector); - let (k, b, qvector) = quantize::quantize::<15>(vector); + let (k, b, qvector) = quantize::quantize(vector, 15.0); let qvector_sum = if vector.len() <= 4369 { - quantize::reduce_sum_of_x_as_u16(&qvector) as f32 + base::simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 } else { - quantize::reduce_sum_of_x_as_u32(&qvector) as f32 + base::simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 }; (dis_v_2, b, k, qvector_sum, compress(qvector)) } @@ -105,7 +105,7 @@ pub fn fscan_process_lowerbound( epsilon: f32, ) -> [Distance; 32] { let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let r = quantization::fast_scan::b4::fast_scan_b4(dims.div_ceil(4), t, s); + let r = base::simd::fast_scan::b4::fast_scan_b4(dims.div_ceil(4), t, s); match distance_kind { DistanceKind::L2 => std::array::from_fn(|i| { let rough = dis_u_2[i] diff --git a/src/vchordrqfscan/algorithm/scan.rs b/src/vchordrqfscan/algorithm/scan.rs index 63264f1f..202949ef 100644 --- a/src/vchordrqfscan/algorithm/scan.rs +++ b/src/vchordrqfscan/algorithm/scan.rs @@ -6,8 +6,8 @@ use crate::vchordrqfscan::algorithm::tuples::*; use base::always_equal::AlwaysEqual; use base::distance::Distance; use base::distance::DistanceKind; -use base::scalar::ScalarLike; use base::search::Pointer; +use base::simd::ScalarLike; use std::cmp::Reverse; use std::collections::BinaryHeap; diff --git a/src/vchordrqfscan/index/am.rs b/src/vchordrqfscan/index/am.rs index e2344148..84a4c8ba 100644 --- a/src/vchordrqfscan/index/am.rs +++ b/src/vchordrqfscan/index/am.rs @@ -183,7 +183,7 @@ pub unsafe extern "C" fn ambuild( ) where F: FnMut((Pointer, Vec)), { - use base::vector::OwnedVector; + use crate::vchordrqfscan::types::OwnedVector; let state = unsafe { &mut *state.cast::>() }; let opfamily = state.this.opfamily; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; @@ -191,9 +191,6 @@ pub unsafe extern "C" fn ambuild( if let Some(vector) = vector { let vector = match vector { OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), }; (state.callback)((pointer, vector.into_vec())); } @@ -572,7 +569,7 @@ unsafe fn parallel_build( ) where F: FnMut((Pointer, Vec)), { - use base::vector::OwnedVector; + use crate::vchordrqfscan::types::OwnedVector; let state = unsafe { &mut *state.cast::>() }; let opfamily = state.this.opfamily; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; @@ -580,9 +577,6 @@ unsafe fn parallel_build( if let Some(vector) = vector { let vector = match vector { OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), }; (state.callback)((pointer, vector.into_vec())); } @@ -669,15 +663,12 @@ pub unsafe extern "C" fn aminsert( _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, _index_info: *mut pgrx::pg_sys::IndexInfo, ) -> bool { - use base::vector::OwnedVector; + use crate::vchordrqfscan::types::OwnedVector; let opfamily = unsafe { am_options::opfamily(index) }; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { let vector = match vector { OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), }; let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); algorithm::insert::insert( @@ -702,15 +693,12 @@ pub unsafe extern "C" fn aminsert( _index_unchanged: bool, _index_info: *mut pgrx::pg_sys::IndexInfo, ) -> bool { - use base::vector::OwnedVector; + use crate::vchordrqfscan::types::OwnedVector; let opfamily = unsafe { am_options::opfamily(index) }; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { let vector = match vector { OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), }; let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); algorithm::insert::insert( diff --git a/src/vchordrqfscan/index/am_options.rs b/src/vchordrqfscan/index/am_options.rs index 51a1009b..b49b7a21 100644 --- a/src/vchordrqfscan/index/am_options.rs +++ b/src/vchordrqfscan/index/am_options.rs @@ -1,10 +1,9 @@ use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use crate::datatype::typmod::Typmod; -use crate::vchordrqfscan::types::VchordrqfscanIndexingOptions; +use crate::vchordrqfscan::types::*; use base::distance::*; -use base::index::*; -use base::vector::*; +use base::vector::VectorBorrowed; use pgrx::datum::FromDatum; use pgrx::heap_tuple::PgHeapTuple; use serde::Deserialize; @@ -26,7 +25,7 @@ impl Reloption { }]; unsafe fn options(&self) -> &CStr { unsafe { - let ptr = std::ptr::addr_of!(*self) + let ptr = (&raw const *self) .cast::() .offset(self.options as _); CStr::from_ptr(ptr) @@ -132,7 +131,6 @@ impl Opfamily { let vector = unsafe { PgvectorVectorInput::from_datum(datum, false).unwrap() }; self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed())) } - _ => unreachable!(), }; Some(vector) } @@ -150,7 +148,6 @@ impl Opfamily { .get_by_index::(NonZero::new(1).unwrap()) .unwrap() .map(|vector| self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed()))), - _ => unreachable!(), }; let radius = tuple.get_by_index::(NonZero::new(2).unwrap()).unwrap(); (center, radius) @@ -162,9 +159,6 @@ impl Opfamily { (B::Vecf32(x), PgDistanceKind::L2) => O::Vecf32(x.own()), (B::Vecf32(x), PgDistanceKind::Dot) => O::Vecf32(x.own()), (B::Vecf32(x), PgDistanceKind::Cos) => O::Vecf32(x.function_normalize()), - (B::Vecf16(x), _) => O::Vecf16(x.own()), - (B::SVecf32(x), _) => O::SVecf32(x.own()), - (B::BVector(x), _) => O::BVector(x.own()), } } pub fn process(self, x: Distance) -> f32 { diff --git a/src/vchordrqfscan/index/am_scan.rs b/src/vchordrqfscan/index/am_scan.rs index 7396bd1b..b07edb7e 100644 --- a/src/vchordrqfscan/index/am_scan.rs +++ b/src/vchordrqfscan/index/am_scan.rs @@ -4,9 +4,9 @@ use crate::vchordrqfscan::algorithm::scan::scan; use crate::vchordrqfscan::gucs::executing::epsilon; use crate::vchordrqfscan::gucs::executing::max_scan_tuples; use crate::vchordrqfscan::gucs::executing::probes; +use crate::vchordrqfscan::types::OwnedVector; use base::distance::Distance; use base::search::*; -use base::vector::*; pub enum Scanner { Initial { @@ -34,7 +34,7 @@ pub fn scan_build( for orderby_vector in orderbys { if pair.is_none() { pair = orderby_vector; - } else if orderby_vector.is_some() && pair != orderby_vector { + } else if orderby_vector.is_some() { pgrx::error!("vector search with multiple vectors is not supported"); } } @@ -42,10 +42,6 @@ pub fn scan_build( if pair.is_none() { pair = sphere_vector; threshold = sphere_threshold; - } else if pair == sphere_vector { - if threshold.is_none() || sphere_threshold < threshold { - threshold = sphere_threshold; - } } else { recheck = true; break; @@ -78,9 +74,6 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, relation, match vector { OwnedVector::Vecf32(x) => x.slice().to_vec(), - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), }, opfamily.distance_kind(), probes(), diff --git a/src/vchordrqfscan/types.rs b/src/vchordrqfscan/types.rs index 0fbe82e9..1180e649 100644 --- a/src/vchordrqfscan/types.rs +++ b/src/vchordrqfscan/types.rs @@ -1,3 +1,5 @@ +use base::distance::DistanceKind; +use base::vector::{VectBorrowed, VectOwned}; use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError, ValidationErrors}; @@ -95,3 +97,42 @@ impl VchordrqfscanIndexingOptions { false } } + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum OwnedVector { + Vecf32(VectOwned), +} + +#[derive(Debug, Clone, Copy)] +pub enum BorrowedVector<'a> { + Vecf32(VectBorrowed<'a, f32>), +} + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +pub enum VectorKind { + Vecf32, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +#[validate(schema(function = "Self::validate_self"))] +pub struct VectorOptions { + #[validate(range(min = 1, max = 1_048_575))] + #[serde(rename = "dimensions")] + pub dims: u32, + #[serde(rename = "vector")] + pub v: VectorKind, + #[serde(rename = "distance")] + pub d: DistanceKind, +} + +impl VectorOptions { + pub fn validate_self(&self) -> Result<(), ValidationError> { + match (self.v, self.d, self.dims) { + (VectorKind::Vecf32, DistanceKind::L2, 1..65536) => Ok(()), + (VectorKind::Vecf32, DistanceKind::Dot, 1..65536) => Ok(()), + _ => Err(ValidationError::new("not valid vector options")), + } + } +} diff --git a/tests/logic/distance.slt b/tests/logic/distance.slt new file mode 100644 index 00000000..e03f32f6 --- /dev/null +++ b/tests/logic/distance.slt @@ -0,0 +1,59 @@ +query I +SELECT round(('[1,2,3]'::vector <-> '[2,3,4]'::vector):: numeric, 3); +---- +1.732 + +query I +SELECT round(('[1,2,3]'::vector <#> '[2,3,4]'::vector):: numeric, 3); +---- +-20.000 + +query I +SELECT round(('[1,2,3]'::vector <=> '[2,3,4]'::vector):: numeric, 3); +---- +0.007 + +query I +SELECT round(('[1,2,3]'::halfvec <-> '[2,3,4]'::halfvec):: numeric, 3); +---- +1.732 + +query I +SELECT round(('[1,2,3]'::halfvec <#> '[2,3,4]'::halfvec):: numeric, 3); +---- +-20.000 + +query I +SELECT round(('[1,2,3]'::halfvec <=> '[2,3,4]'::halfvec):: numeric, 3); +---- +0.007 + +query I +SELECT round((quantize_to_scalar8('[1,2,3]'::vector) <-> quantize_to_scalar8('[2,3,4]'::vector)):: numeric, 1); +---- +1.7 + +query I +SELECT round((quantize_to_scalar8('[1,2,3]'::vector) <#> quantize_to_scalar8('[2,3,4]'::vector)):: numeric, 1); +---- +-20.0 + +query I +SELECT round((quantize_to_scalar8('[1,2,3]'::vector) <=> quantize_to_scalar8('[2,3,4]'::vector)):: numeric, 2); +---- +0.01 + +query I +SELECT round((quantize_to_scalar8('[1,2,3]'::halfvec) <-> quantize_to_scalar8('[2,3,4]'::halfvec)):: numeric, 1); +---- +1.7 + +query I +SELECT round((quantize_to_scalar8('[1,2,3]'::halfvec) <#> quantize_to_scalar8('[2,3,4]'::halfvec)):: numeric, 1); +---- +-20.0 + +query I +SELECT round((quantize_to_scalar8('[1,2,3]'::halfvec) <=> quantize_to_scalar8('[2,3,4]'::halfvec)):: numeric, 2); +---- +0.01 From da0bf0053203918258c9d1cd0b492e939efb2542 Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Tue, 10 Dec 2024 22:38:36 +0800 Subject: [PATCH 077/324] feat: Use dual license (AGPLv3 and ELv2) (#130) Close #123 --------- Signed-off-by: Ce Gao --- LICENSE | 664 +--------------------------------------- README.md | 9 +- licenses/LICENSE.AGPLv3 | 661 +++++++++++++++++++++++++++++++++++++++ licenses/LICENSE.ELv2 | 91 ++++++ 4 files changed, 773 insertions(+), 652 deletions(-) create mode 100644 licenses/LICENSE.AGPLv3 create mode 100644 licenses/LICENSE.ELv2 diff --git a/LICENSE b/LICENSE index 0ad25db4..e6d43143 100644 --- a/LICENSE +++ b/LICENSE @@ -1,661 +1,23 @@ - GNU AFFERO GENERAL PUBLIC LICENSE - Version 3, 19 November 2007 +# Dual License Notice - Copyright (C) 2007 Free Software Foundation, Inc. - Everyone is permitted to copy and distribute verbatim copies - of this license document, but changing it is not allowed. +This software is licensed under a dual license model. You may choose to use this software under one of the following licenses: - Preamble +1. **GNU Affero General Public License v3 (AGPLv3)** - The GNU Affero General Public License is a free, copyleft license for -software and other kinds of works, specifically designed to ensure -cooperation with the community in the case of network server software. + This program is free software: you can redistribute it and/or modify it under the terms of the AGPLv3. + You should have received a copy of the AGPLv3 along with this program. If not, see . - The licenses for most software and other practical works are designed -to take away your freedom to share and change the works. By contrast, -our General Public Licenses are intended to guarantee your freedom to -share and change all versions of a program--to make sure it remains free -software for all its users. + The AGPLv3 license allows you to use, modify, and distribute the software, but requires you to share your modifications under the same license. - When we speak of free software, we are referring to freedom, not -price. Our General Public Licenses are designed to make sure that you -have the freedom to distribute copies of free software (and charge for -them if you wish), that you receive source code or can get it if you -want it, that you can change the software or use pieces of it in new -free programs, and that you know you can do these things. +2. **Elastic License v2 (ELv2)** - Developers that use our General Public Licenses protect your rights -with two steps: (1) assert copyright on the software, and (2) offer -you this License which gives you legal permission to copy, distribute -and/or modify the software. + This software is also licensed under the Elastic License v2. + For details, see . - A secondary benefit of defending all users' freedom is that -improvements made in alternate versions of the program, if they -receive widespread use, become available for other developers to -incorporate. Many developers of free software are heartened and -encouraged by the resulting cooperation. However, in the case of -software used on network servers, this result may fail to come about. -The GNU General Public License permits making a modified version and -letting the public access it on a server without ever releasing its -source code to the public. + The Elastic License v2 allows you to use, modify, and distribute the software, but imposes certain restrictions, particularly regarding the provision of the software as a service. - The GNU Affero General Public License is designed specifically to -ensure that, in such cases, the modified source code becomes available -to the community. It requires the operator of a network server to -provide the source code of the modified version running there to the -users of that server. Therefore, public use of a modified version, on -a publicly accessible server, gives the public access to the source -code of the modified version. +## Choosing a License - An older license, called the Affero General Public License and -published by Affero, was designed to accomplish similar goals. This is -a different license, not a version of the Affero GPL, but Affero has -released a new version of the Affero GPL which permits relicensing under -this license. +You may choose to use this software under either the AGPLv3 or the Elastic License v2. If you choose the AGPLv3, you must comply with its terms. If you choose the Elastic License v2, you must comply with its terms. - The precise terms and conditions for copying, distribution and -modification follow. - - TERMS AND CONDITIONS - - 0. Definitions. - - "This License" refers to version 3 of the GNU Affero General Public License. - - "Copyright" also means copyright-like laws that apply to other kinds of -works, such as semiconductor masks. - - "The Program" refers to any copyrightable work licensed under this -License. Each licensee is addressed as "you". "Licensees" and -"recipients" may be individuals or organizations. - - To "modify" a work means to copy from or adapt all or part of the work -in a fashion requiring copyright permission, other than the making of an -exact copy. The resulting work is called a "modified version" of the -earlier work or a work "based on" the earlier work. - - A "covered work" means either the unmodified Program or a work based -on the Program. - - To "propagate" a work means to do anything with it that, without -permission, would make you directly or secondarily liable for -infringement under applicable copyright law, except executing it on a -computer or modifying a private copy. Propagation includes copying, -distribution (with or without modification), making available to the -public, and in some countries other activities as well. - - To "convey" a work means any kind of propagation that enables other -parties to make or receive copies. Mere interaction with a user through -a computer network, with no transfer of a copy, is not conveying. - - An interactive user interface displays "Appropriate Legal Notices" -to the extent that it includes a convenient and prominently visible -feature that (1) displays an appropriate copyright notice, and (2) -tells the user that there is no warranty for the work (except to the -extent that warranties are provided), that licensees may convey the -work under this License, and how to view a copy of this License. If -the interface presents a list of user commands or options, such as a -menu, a prominent item in the list meets this criterion. - - 1. Source Code. - - The "source code" for a work means the preferred form of the work -for making modifications to it. "Object code" means any non-source -form of a work. - - A "Standard Interface" means an interface that either is an official -standard defined by a recognized standards body, or, in the case of -interfaces specified for a particular programming language, one that -is widely used among developers working in that language. - - The "System Libraries" of an executable work include anything, other -than the work as a whole, that (a) is included in the normal form of -packaging a Major Component, but which is not part of that Major -Component, and (b) serves only to enable use of the work with that -Major Component, or to implement a Standard Interface for which an -implementation is available to the public in source code form. A -"Major Component", in this context, means a major essential component -(kernel, window system, and so on) of the specific operating system -(if any) on which the executable work runs, or a compiler used to -produce the work, or an object code interpreter used to run it. - - The "Corresponding Source" for a work in object code form means all -the source code needed to generate, install, and (for an executable -work) run the object code and to modify the work, including scripts to -control those activities. However, it does not include the work's -System Libraries, or general-purpose tools or generally available free -programs which are used unmodified in performing those activities but -which are not part of the work. For example, Corresponding Source -includes interface definition files associated with source files for -the work, and the source code for shared libraries and dynamically -linked subprograms that the work is specifically designed to require, -such as by intimate data communication or control flow between those -subprograms and other parts of the work. - - The Corresponding Source need not include anything that users -can regenerate automatically from other parts of the Corresponding -Source. - - The Corresponding Source for a work in source code form is that -same work. - - 2. Basic Permissions. - - All rights granted under this License are granted for the term of -copyright on the Program, and are irrevocable provided the stated -conditions are met. This License explicitly affirms your unlimited -permission to run the unmodified Program. The output from running a -covered work is covered by this License only if the output, given its -content, constitutes a covered work. This License acknowledges your -rights of fair use or other equivalent, as provided by copyright law. - - You may make, run and propagate covered works that you do not -convey, without conditions so long as your license otherwise remains -in force. You may convey covered works to others for the sole purpose -of having them make modifications exclusively for you, or provide you -with facilities for running those works, provided that you comply with -the terms of this License in conveying all material for which you do -not control copyright. Those thus making or running the covered works -for you must do so exclusively on your behalf, under your direction -and control, on terms that prohibit them from making any copies of -your copyrighted material outside their relationship with you. - - Conveying under any other circumstances is permitted solely under -the conditions stated below. Sublicensing is not allowed; section 10 -makes it unnecessary. - - 3. Protecting Users' Legal Rights From Anti-Circumvention Law. - - No covered work shall be deemed part of an effective technological -measure under any applicable law fulfilling obligations under article -11 of the WIPO copyright treaty adopted on 20 December 1996, or -similar laws prohibiting or restricting circumvention of such -measures. - - When you convey a covered work, you waive any legal power to forbid -circumvention of technological measures to the extent such circumvention -is effected by exercising rights under this License with respect to -the covered work, and you disclaim any intention to limit operation or -modification of the work as a means of enforcing, against the work's -users, your or third parties' legal rights to forbid circumvention of -technological measures. - - 4. Conveying Verbatim Copies. - - You may convey verbatim copies of the Program's source code as you -receive it, in any medium, provided that you conspicuously and -appropriately publish on each copy an appropriate copyright notice; -keep intact all notices stating that this License and any -non-permissive terms added in accord with section 7 apply to the code; -keep intact all notices of the absence of any warranty; and give all -recipients a copy of this License along with the Program. - - You may charge any price or no price for each copy that you convey, -and you may offer support or warranty protection for a fee. - - 5. Conveying Modified Source Versions. - - You may convey a work based on the Program, or the modifications to -produce it from the Program, in the form of source code under the -terms of section 4, provided that you also meet all of these conditions: - - a) The work must carry prominent notices stating that you modified - it, and giving a relevant date. - - b) The work must carry prominent notices stating that it is - released under this License and any conditions added under section - 7. This requirement modifies the requirement in section 4 to - "keep intact all notices". - - c) You must license the entire work, as a whole, under this - License to anyone who comes into possession of a copy. This - License will therefore apply, along with any applicable section 7 - additional terms, to the whole of the work, and all its parts, - regardless of how they are packaged. This License gives no - permission to license the work in any other way, but it does not - invalidate such permission if you have separately received it. - - d) If the work has interactive user interfaces, each must display - Appropriate Legal Notices; however, if the Program has interactive - interfaces that do not display Appropriate Legal Notices, your - work need not make them do so. - - A compilation of a covered work with other separate and independent -works, which are not by their nature extensions of the covered work, -and which are not combined with it such as to form a larger program, -in or on a volume of a storage or distribution medium, is called an -"aggregate" if the compilation and its resulting copyright are not -used to limit the access or legal rights of the compilation's users -beyond what the individual works permit. Inclusion of a covered work -in an aggregate does not cause this License to apply to the other -parts of the aggregate. - - 6. Conveying Non-Source Forms. - - You may convey a covered work in object code form under the terms -of sections 4 and 5, provided that you also convey the -machine-readable Corresponding Source under the terms of this License, -in one of these ways: - - a) Convey the object code in, or embodied in, a physical product - (including a physical distribution medium), accompanied by the - Corresponding Source fixed on a durable physical medium - customarily used for software interchange. - - b) Convey the object code in, or embodied in, a physical product - (including a physical distribution medium), accompanied by a - written offer, valid for at least three years and valid for as - long as you offer spare parts or customer support for that product - model, to give anyone who possesses the object code either (1) a - copy of the Corresponding Source for all the software in the - product that is covered by this License, on a durable physical - medium customarily used for software interchange, for a price no - more than your reasonable cost of physically performing this - conveying of source, or (2) access to copy the - Corresponding Source from a network server at no charge. - - c) Convey individual copies of the object code with a copy of the - written offer to provide the Corresponding Source. This - alternative is allowed only occasionally and noncommercially, and - only if you received the object code with such an offer, in accord - with subsection 6b. - - d) Convey the object code by offering access from a designated - place (gratis or for a charge), and offer equivalent access to the - Corresponding Source in the same way through the same place at no - further charge. You need not require recipients to copy the - Corresponding Source along with the object code. If the place to - copy the object code is a network server, the Corresponding Source - may be on a different server (operated by you or a third party) - that supports equivalent copying facilities, provided you maintain - clear directions next to the object code saying where to find the - Corresponding Source. Regardless of what server hosts the - Corresponding Source, you remain obligated to ensure that it is - available for as long as needed to satisfy these requirements. - - e) Convey the object code using peer-to-peer transmission, provided - you inform other peers where the object code and Corresponding - Source of the work are being offered to the general public at no - charge under subsection 6d. - - A separable portion of the object code, whose source code is excluded -from the Corresponding Source as a System Library, need not be -included in conveying the object code work. - - A "User Product" is either (1) a "consumer product", which means any -tangible personal property which is normally used for personal, family, -or household purposes, or (2) anything designed or sold for incorporation -into a dwelling. In determining whether a product is a consumer product, -doubtful cases shall be resolved in favor of coverage. For a particular -product received by a particular user, "normally used" refers to a -typical or common use of that class of product, regardless of the status -of the particular user or of the way in which the particular user -actually uses, or expects or is expected to use, the product. A product -is a consumer product regardless of whether the product has substantial -commercial, industrial or non-consumer uses, unless such uses represent -the only significant mode of use of the product. - - "Installation Information" for a User Product means any methods, -procedures, authorization keys, or other information required to install -and execute modified versions of a covered work in that User Product from -a modified version of its Corresponding Source. The information must -suffice to ensure that the continued functioning of the modified object -code is in no case prevented or interfered with solely because -modification has been made. - - If you convey an object code work under this section in, or with, or -specifically for use in, a User Product, and the conveying occurs as -part of a transaction in which the right of possession and use of the -User Product is transferred to the recipient in perpetuity or for a -fixed term (regardless of how the transaction is characterized), the -Corresponding Source conveyed under this section must be accompanied -by the Installation Information. But this requirement does not apply -if neither you nor any third party retains the ability to install -modified object code on the User Product (for example, the work has -been installed in ROM). - - The requirement to provide Installation Information does not include a -requirement to continue to provide support service, warranty, or updates -for a work that has been modified or installed by the recipient, or for -the User Product in which it has been modified or installed. Access to a -network may be denied when the modification itself materially and -adversely affects the operation of the network or violates the rules and -protocols for communication across the network. - - Corresponding Source conveyed, and Installation Information provided, -in accord with this section must be in a format that is publicly -documented (and with an implementation available to the public in -source code form), and must require no special password or key for -unpacking, reading or copying. - - 7. Additional Terms. - - "Additional permissions" are terms that supplement the terms of this -License by making exceptions from one or more of its conditions. -Additional permissions that are applicable to the entire Program shall -be treated as though they were included in this License, to the extent -that they are valid under applicable law. If additional permissions -apply only to part of the Program, that part may be used separately -under those permissions, but the entire Program remains governed by -this License without regard to the additional permissions. - - When you convey a copy of a covered work, you may at your option -remove any additional permissions from that copy, or from any part of -it. (Additional permissions may be written to require their own -removal in certain cases when you modify the work.) You may place -additional permissions on material, added by you to a covered work, -for which you have or can give appropriate copyright permission. - - Notwithstanding any other provision of this License, for material you -add to a covered work, you may (if authorized by the copyright holders of -that material) supplement the terms of this License with terms: - - a) Disclaiming warranty or limiting liability differently from the - terms of sections 15 and 16 of this License; or - - b) Requiring preservation of specified reasonable legal notices or - author attributions in that material or in the Appropriate Legal - Notices displayed by works containing it; or - - c) Prohibiting misrepresentation of the origin of that material, or - requiring that modified versions of such material be marked in - reasonable ways as different from the original version; or - - d) Limiting the use for publicity purposes of names of licensors or - authors of the material; or - - e) Declining to grant rights under trademark law for use of some - trade names, trademarks, or service marks; or - - f) Requiring indemnification of licensors and authors of that - material by anyone who conveys the material (or modified versions of - it) with contractual assumptions of liability to the recipient, for - any liability that these contractual assumptions directly impose on - those licensors and authors. - - All other non-permissive additional terms are considered "further -restrictions" within the meaning of section 10. If the Program as you -received it, or any part of it, contains a notice stating that it is -governed by this License along with a term that is a further -restriction, you may remove that term. If a license document contains -a further restriction but permits relicensing or conveying under this -License, you may add to a covered work material governed by the terms -of that license document, provided that the further restriction does -not survive such relicensing or conveying. - - If you add terms to a covered work in accord with this section, you -must place, in the relevant source files, a statement of the -additional terms that apply to those files, or a notice indicating -where to find the applicable terms. - - Additional terms, permissive or non-permissive, may be stated in the -form of a separately written license, or stated as exceptions; -the above requirements apply either way. - - 8. Termination. - - You may not propagate or modify a covered work except as expressly -provided under this License. Any attempt otherwise to propagate or -modify it is void, and will automatically terminate your rights under -this License (including any patent licenses granted under the third -paragraph of section 11). - - However, if you cease all violation of this License, then your -license from a particular copyright holder is reinstated (a) -provisionally, unless and until the copyright holder explicitly and -finally terminates your license, and (b) permanently, if the copyright -holder fails to notify you of the violation by some reasonable means -prior to 60 days after the cessation. - - Moreover, your license from a particular copyright holder is -reinstated permanently if the copyright holder notifies you of the -violation by some reasonable means, this is the first time you have -received notice of violation of this License (for any work) from that -copyright holder, and you cure the violation prior to 30 days after -your receipt of the notice. - - Termination of your rights under this section does not terminate the -licenses of parties who have received copies or rights from you under -this License. If your rights have been terminated and not permanently -reinstated, you do not qualify to receive new licenses for the same -material under section 10. - - 9. Acceptance Not Required for Having Copies. - - You are not required to accept this License in order to receive or -run a copy of the Program. Ancillary propagation of a covered work -occurring solely as a consequence of using peer-to-peer transmission -to receive a copy likewise does not require acceptance. However, -nothing other than this License grants you permission to propagate or -modify any covered work. These actions infringe copyright if you do -not accept this License. Therefore, by modifying or propagating a -covered work, you indicate your acceptance of this License to do so. - - 10. Automatic Licensing of Downstream Recipients. - - Each time you convey a covered work, the recipient automatically -receives a license from the original licensors, to run, modify and -propagate that work, subject to this License. You are not responsible -for enforcing compliance by third parties with this License. - - An "entity transaction" is a transaction transferring control of an -organization, or substantially all assets of one, or subdividing an -organization, or merging organizations. If propagation of a covered -work results from an entity transaction, each party to that -transaction who receives a copy of the work also receives whatever -licenses to the work the party's predecessor in interest had or could -give under the previous paragraph, plus a right to possession of the -Corresponding Source of the work from the predecessor in interest, if -the predecessor has it or can get it with reasonable efforts. - - You may not impose any further restrictions on the exercise of the -rights granted or affirmed under this License. For example, you may -not impose a license fee, royalty, or other charge for exercise of -rights granted under this License, and you may not initiate litigation -(including a cross-claim or counterclaim in a lawsuit) alleging that -any patent claim is infringed by making, using, selling, offering for -sale, or importing the Program or any portion of it. - - 11. Patents. - - A "contributor" is a copyright holder who authorizes use under this -License of the Program or a work on which the Program is based. The -work thus licensed is called the contributor's "contributor version". - - A contributor's "essential patent claims" are all patent claims -owned or controlled by the contributor, whether already acquired or -hereafter acquired, that would be infringed by some manner, permitted -by this License, of making, using, or selling its contributor version, -but do not include claims that would be infringed only as a -consequence of further modification of the contributor version. For -purposes of this definition, "control" includes the right to grant -patent sublicenses in a manner consistent with the requirements of -this License. - - Each contributor grants you a non-exclusive, worldwide, royalty-free -patent license under the contributor's essential patent claims, to -make, use, sell, offer for sale, import and otherwise run, modify and -propagate the contents of its contributor version. - - In the following three paragraphs, a "patent license" is any express -agreement or commitment, however denominated, not to enforce a patent -(such as an express permission to practice a patent or covenant not to -sue for patent infringement). To "grant" such a patent license to a -party means to make such an agreement or commitment not to enforce a -patent against the party. - - If you convey a covered work, knowingly relying on a patent license, -and the Corresponding Source of the work is not available for anyone -to copy, free of charge and under the terms of this License, through a -publicly available network server or other readily accessible means, -then you must either (1) cause the Corresponding Source to be so -available, or (2) arrange to deprive yourself of the benefit of the -patent license for this particular work, or (3) arrange, in a manner -consistent with the requirements of this License, to extend the patent -license to downstream recipients. "Knowingly relying" means you have -actual knowledge that, but for the patent license, your conveying the -covered work in a country, or your recipient's use of the covered work -in a country, would infringe one or more identifiable patents in that -country that you have reason to believe are valid. - - If, pursuant to or in connection with a single transaction or -arrangement, you convey, or propagate by procuring conveyance of, a -covered work, and grant a patent license to some of the parties -receiving the covered work authorizing them to use, propagate, modify -or convey a specific copy of the covered work, then the patent license -you grant is automatically extended to all recipients of the covered -work and works based on it. - - A patent license is "discriminatory" if it does not include within -the scope of its coverage, prohibits the exercise of, or is -conditioned on the non-exercise of one or more of the rights that are -specifically granted under this License. You may not convey a covered -work if you are a party to an arrangement with a third party that is -in the business of distributing software, under which you make payment -to the third party based on the extent of your activity of conveying -the work, and under which the third party grants, to any of the -parties who would receive the covered work from you, a discriminatory -patent license (a) in connection with copies of the covered work -conveyed by you (or copies made from those copies), or (b) primarily -for and in connection with specific products or compilations that -contain the covered work, unless you entered into that arrangement, -or that patent license was granted, prior to 28 March 2007. - - Nothing in this License shall be construed as excluding or limiting -any implied license or other defenses to infringement that may -otherwise be available to you under applicable patent law. - - 12. No Surrender of Others' Freedom. - - If conditions are imposed on you (whether by court order, agreement or -otherwise) that contradict the conditions of this License, they do not -excuse you from the conditions of this License. If you cannot convey a -covered work so as to satisfy simultaneously your obligations under this -License and any other pertinent obligations, then as a consequence you may -not convey it at all. For example, if you agree to terms that obligate you -to collect a royalty for further conveying from those to whom you convey -the Program, the only way you could satisfy both those terms and this -License would be to refrain entirely from conveying the Program. - - 13. Remote Network Interaction; Use with the GNU General Public License. - - Notwithstanding any other provision of this License, if you modify the -Program, your modified version must prominently offer all users -interacting with it remotely through a computer network (if your version -supports such interaction) an opportunity to receive the Corresponding -Source of your version by providing access to the Corresponding Source -from a network server at no charge, through some standard or customary -means of facilitating copying of software. This Corresponding Source -shall include the Corresponding Source for any work covered by version 3 -of the GNU General Public License that is incorporated pursuant to the -following paragraph. - - Notwithstanding any other provision of this License, you have -permission to link or combine any covered work with a work licensed -under version 3 of the GNU General Public License into a single -combined work, and to convey the resulting work. The terms of this -License will continue to apply to the part which is the covered work, -but the work with which it is combined will remain governed by version -3 of the GNU General Public License. - - 14. Revised Versions of this License. - - The Free Software Foundation may publish revised and/or new versions of -the GNU Affero General Public License from time to time. Such new versions -will be similar in spirit to the present version, but may differ in detail to -address new problems or concerns. - - Each version is given a distinguishing version number. If the -Program specifies that a certain numbered version of the GNU Affero General -Public License "or any later version" applies to it, you have the -option of following the terms and conditions either of that numbered -version or of any later version published by the Free Software -Foundation. If the Program does not specify a version number of the -GNU Affero General Public License, you may choose any version ever published -by the Free Software Foundation. - - If the Program specifies that a proxy can decide which future -versions of the GNU Affero General Public License can be used, that proxy's -public statement of acceptance of a version permanently authorizes you -to choose that version for the Program. - - Later license versions may give you additional or different -permissions. However, no additional obligations are imposed on any -author or copyright holder as a result of your choosing to follow a -later version. - - 15. Disclaimer of Warranty. - - THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY -APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT -HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY -OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, -THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM -IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF -ALL NECESSARY SERVICING, REPAIR OR CORRECTION. - - 16. Limitation of Liability. - - IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING -WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS -THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY -GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE -USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF -DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD -PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), -EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF -SUCH DAMAGES. - - 17. Interpretation of Sections 15 and 16. - - If the disclaimer of warranty and limitation of liability provided -above cannot be given local legal effect according to their terms, -reviewing courts shall apply local law that most closely approximates -an absolute waiver of all civil liability in connection with the -Program, unless a warranty or assumption of liability accompanies a -copy of the Program in return for a fee. - - END OF TERMS AND CONDITIONS - - How to Apply These Terms to Your New Programs - - If you develop a new program, and you want it to be of the greatest -possible use to the public, the best way to achieve this is to make it -free software which everyone can redistribute and change under these terms. - - To do so, attach the following notices to the program. It is safest -to attach them to the start of each source file to most effectively -state the exclusion of warranty; and each file should have at least -the "copyright" line and a pointer to where the full notice is found. - - - Copyright (C) - - This program is free software: you can redistribute it and/or modify - it under the terms of the GNU Affero General Public License as published - by the Free Software Foundation, either version 3 of the License, or - (at your option) any later version. - - This program is distributed in the hope that it will be useful, - but WITHOUT ANY WARRANTY; without even the implied warranty of - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - GNU Affero General Public License for more details. - - You should have received a copy of the GNU Affero General Public License - along with this program. If not, see . - -Also add information on how to contact you by electronic and paper mail. - - If your software can interact with users remotely through a computer -network, you should also make sure that it provides a way for users to -get its source. For example, if your program is a web application, its -interface could display a "Source" link that leads users to an archive -of the code. There are many ways you could offer source, and different -solutions will be better for different programs; see section 13 for the -specific requirements. - - You should also get your employer (if you work as a programmer) or school, -if any, to sign a "copyright disclaimer" for the program, if necessary. -For more information on this, and how to apply and follow the GNU AGPL, see -. +By using this software, you agree to adhere to the terms of the license you select. diff --git a/README.md b/README.md index a1e9b255..3f55e0c3 100644 --- a/README.md +++ b/README.md @@ -207,4 +207,11 @@ cargo pgrx install --release --sudo # To install the extension into the system p ## License -This project is licensed under the [GNU Affero General Public License v3.0](./LICENSE) and as commercial software. For commercial licensing, please contact us at support@tensorchord.ai. + +This software is licensed under a dual license model: + +1. **GNU Affero General Public License v3 (AGPLv3)**: You may use, modify, and distribute this software under the terms of the AGPLv3. + +2. **Elastic License v2 (ELv2)**: You may also use, modify, and distribute this software under the Elastic License v2, which has specific restrictions. + +You may choose either license based on your needs. We welcome any commercial collaboration or support, so please email us with any questions or requests regarding the licenses. diff --git a/licenses/LICENSE.AGPLv3 b/licenses/LICENSE.AGPLv3 new file mode 100644 index 00000000..0ad25db4 --- /dev/null +++ b/licenses/LICENSE.AGPLv3 @@ -0,0 +1,661 @@ + GNU AFFERO GENERAL PUBLIC LICENSE + Version 3, 19 November 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU Affero General Public License is a free, copyleft license for +software and other kinds of works, specifically designed to ensure +cooperation with the community in the case of network server software. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +our General Public Licenses are intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + Developers that use our General Public Licenses protect your rights +with two steps: (1) assert copyright on the software, and (2) offer +you this License which gives you legal permission to copy, distribute +and/or modify the software. + + A secondary benefit of defending all users' freedom is that +improvements made in alternate versions of the program, if they +receive widespread use, become available for other developers to +incorporate. Many developers of free software are heartened and +encouraged by the resulting cooperation. However, in the case of +software used on network servers, this result may fail to come about. +The GNU General Public License permits making a modified version and +letting the public access it on a server without ever releasing its +source code to the public. + + The GNU Affero General Public License is designed specifically to +ensure that, in such cases, the modified source code becomes available +to the community. It requires the operator of a network server to +provide the source code of the modified version running there to the +users of that server. Therefore, public use of a modified version, on +a publicly accessible server, gives the public access to the source +code of the modified version. + + An older license, called the Affero General Public License and +published by Affero, was designed to accomplish similar goals. This is +a different license, not a version of the Affero GPL, but Affero has +released a new version of the Affero GPL which permits relicensing under +this license. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU Affero General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Remote Network Interaction; Use with the GNU General Public License. + + Notwithstanding any other provision of this License, if you modify the +Program, your modified version must prominently offer all users +interacting with it remotely through a computer network (if your version +supports such interaction) an opportunity to receive the Corresponding +Source of your version by providing access to the Corresponding Source +from a network server at no charge, through some standard or customary +means of facilitating copying of software. This Corresponding Source +shall include the Corresponding Source for any work covered by version 3 +of the GNU General Public License that is incorporated pursuant to the +following paragraph. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the work with which it is combined will remain governed by version +3 of the GNU General Public License. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU Affero General Public License from time to time. Such new versions +will be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU Affero General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU Affero General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU Affero General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published + by the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Affero General Public License for more details. + + You should have received a copy of the GNU Affero General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If your software can interact with users remotely through a computer +network, you should also make sure that it provides a way for users to +get its source. For example, if your program is a web application, its +interface could display a "Source" link that leads users to an archive +of the code. There are many ways you could offer source, and different +solutions will be better for different programs; see section 13 for the +specific requirements. + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU AGPL, see +. diff --git a/licenses/LICENSE.ELv2 b/licenses/LICENSE.ELv2 new file mode 100644 index 00000000..efa8dfbe --- /dev/null +++ b/licenses/LICENSE.ELv2 @@ -0,0 +1,91 @@ +Elastic License 2.0 + +## Acceptance + +By using the software, you agree to all of the terms and conditions below. + +## Copyright License + +The licensor grants you a non-exclusive, royalty-free, worldwide, +non-sublicensable, non-transferable license to use, copy, distribute, make +available, and prepare derivative works of the software, in each case subject to +the limitations and conditions below. + +## Limitations + +You may not provide the software to third parties as a hosted or managed +service, where the service provides users with access to any substantial set of +the features or functionality of the software. + +You may not move, change, disable, or circumvent the license key functionality +in the software, and you may not remove or obscure any functionality in the +software that is protected by the license key. + +You may not alter, remove, or obscure any licensing, copyright, or other notices +of the licensor in the software. Any use of the licensor’s trademarks is subject +to applicable law. + +## Patents + +The licensor grants you a license, under any patent claims the licensor can +license, or becomes able to license, to make, have made, use, sell, offer for +sale, import and have imported the software, in each case subject to the +limitations and conditions in this license. This license does not cover any +patent claims that you cause to be infringed by modifications or additions to +the software. If you or your company make any written claim that the software +infringes or contributes to infringement of any patent, your patent license for +the software granted under these terms ends immediately. If your company makes +such a claim, your patent license ends immediately for work on behalf of your +company. + +## Notices + +You must ensure that anyone who gets a copy of any part of the software from you +also gets a copy of these terms. + +If you modify the software, you must include in any modified copies of the +software prominent notices stating that you have modified the software. + +## No Other Rights + +These terms do not imply any licenses other than those expressly granted in +these terms. + +## Termination + +If you use the software in violation of these terms, such use is not licensed, +and your licenses will automatically terminate. If the licensor provides you +with a notice of your violation, and you cease all violation of this license no +later than 30 days after you receive that notice, your licenses will be +reinstated retroactively. However, if you violate these terms after such +reinstatement, any additional violation of these terms will cause your licenses +to terminate automatically and permanently. + +## No Liability + +*As far as the law allows, the software comes as is, without any warranty or +condition, and the licensor will not be liable to you for any damages arising +out of these terms or the use or nature of the software, under any kind of +legal claim.* + +## Definitions + +The **licensor** is the entity offering these terms, and the **software** is the +software the licensor makes available under these terms, including any portion +of it. + +**you** refers to the individual or entity agreeing to these terms. + +**your company** is any legal entity, sole proprietorship, or other kind of +organization that you work for, plus all organizations that have control over, +are under the control of, or are under common control with that +organization. **control** means ownership of substantially all the assets of an +entity, or the power to direct its management and policies by vote, contract, or +otherwise. Control can be direct or indirect. + +**your licenses** are all the licenses granted to you for the software under +these terms. + +**use** means anything you do with the software requiring one of your licenses. + +**trademark** means trademarks, service marks, and similar rights. From 6f310a27c6444b5e93447b96bac16759e051aafe Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 13 Dec 2024 19:21:06 +0800 Subject: [PATCH 078/324] chore: update base (#137) Signed-off-by: usamoi --- Cargo.lock | 444 +++++++++++++++++++++++++++++++++-------- Cargo.toml | 2 +- docker/pgrx.Dockerfile | 6 +- src/lib.rs | 1 - 4 files changed, 365 insertions(+), 88 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 235b7822..982074bd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.92" +version = "1.0.94" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74f37166d7d48a0284b99dd824694c26119c700b53bf0d1540cdb147dbdaaf13" +checksum = "c1fd03a028ef38ba2276dce7e33fcd6369c158a1bca17946c4b1b701891c1ff7" [[package]] name = "approx" @@ -66,15 +66,15 @@ checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" [[package]] name = "base" version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=c911e8a476effaf05cd4b1a037826b100177bdad#c911e8a476effaf05cd4b1a037826b100177bdad" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=0374439c8593ead1f61fc2477a4b79a49fcdceb8#0374439c8593ead1f61fc2477a4b79a49fcdceb8" dependencies = [ "base_macros", - "detect", + "cc", "half 2.4.1", "libc", "rand", "serde", - "thiserror", + "thiserror 1.0.69", "toml", "validator", ] @@ -82,11 +82,11 @@ dependencies = [ [[package]] name = "base_macros" version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=c911e8a476effaf05cd4b1a037826b100177bdad#c911e8a476effaf05cd4b1a037826b100177bdad" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=0374439c8593ead1f61fc2477a4b79a49fcdceb8#0374439c8593ead1f61fc2477a4b79a49fcdceb8" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] @@ -105,7 +105,7 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] @@ -162,9 +162,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.8.0" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ac0150caa2ae65ca5bd83f25c7de183dea78d4d366469f148435e2acfbad0da" +checksum = "325918d6fe32f23b19878fe4b34794ae41fc19ddbe53b10571a4874d44ffd39b" [[package]] name = "cargo_toml" @@ -178,9 +178,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.1.34" +version = "1.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b9470d453346108f93a59222a9a1a5724db32d0a4727b7ab7ace4b4d822dc9" +checksum = "27f657647bcff5394bf56c7317665bbf790a137a50eaaa5c6bfbb9e27a518f2d" dependencies = [ "shlex", ] @@ -282,7 +282,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] @@ -293,25 +293,18 @@ checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" dependencies = [ "darling_core", "quote", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] -name = "detect" -version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=c911e8a476effaf05cd4b1a037826b100177bdad#c911e8a476effaf05cd4b1a037826b100177bdad" -dependencies = [ - "detect_macros", -] - -[[package]] -name = "detect_macros" -version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=c911e8a476effaf05cd4b1a037826b100177bdad#c911e8a476effaf05cd4b1a037826b100177bdad" +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] @@ -337,7 +330,7 @@ checksum = "f282cfdfe92516eb26c2af8589c274c7c17681f5ecc03c18255fe741c6aa64eb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] @@ -439,9 +432,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.15.0" +version = "0.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb" +checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" [[package]] name = "heapless" @@ -468,6 +461,124 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "icu_collections" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db2fa452206ebee18c4b5c2274dbf1de17008e874b4dc4f0aea9d01ca79e4526" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locid" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13acbb8371917fc971be86fc8057c41a64b521c184808a698c02acc242dbf637" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_locid_transform" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01d11ac35de8e40fdeda00d9e1e9d92525f3f9d887cdd7aa81d727596788b54e" +dependencies = [ + "displaydoc", + "icu_locid", + "icu_locid_transform_data", + "icu_provider", + "tinystr", + "zerovec", +] + +[[package]] +name = "icu_locid_transform_data" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdc8ff3388f852bede6b579ad4e978ab004f139284d7b28715f773507b946f6e" + +[[package]] +name = "icu_normalizer" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19ce3e0da2ec68599d193c93d088142efd7f9c5d6fc9b803774855747dc6a84f" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "utf16_iter", + "utf8_iter", + "write16", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8cafbf7aa791e9b22bec55a167906f9e1215fd475cd22adfcf660e03e989516" + +[[package]] +name = "icu_properties" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93d6020766cfc6302c15dbbc9c8778c37e62c14427cb7f6e601d849e092aeef5" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_locid_transform", + "icu_properties_data", + "icu_provider", + "tinystr", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67a8effbc3dd3e4ba1afa8ad918d5684b8868b3b26500753effea8d2eed19569" + +[[package]] +name = "icu_provider" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ed421c8a8ef78d3e2dbc98a973be2f3770cb42b606e3ab18d6237c4dfde68d9" +dependencies = [ + "displaydoc", + "icu_locid", + "icu_provider_macros", + "stable_deref_trait", + "tinystr", + "writeable", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_provider_macros" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.90", +] + [[package]] name = "ident_case" version = "1.0.1" @@ -484,6 +595,27 @@ dependencies = [ "unicode-normalization", ] +[[package]] +name = "idna" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daca1df1c957320b2cf139ac61e7bd64fed304c5040df000a745aa1de3b4ef71" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + [[package]] name = "indenter" version = "0.3.3" @@ -492,12 +624,12 @@ checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" [[package]] name = "indexmap" -version = "2.6.0" +version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "707907fe3c25f5424cce2cb7e1cbcafee6bdbe735ca90ef77c29e84591e5b9da" +checksum = "62f822373a4fe84d4bb149bf54e584a7f4abec90e072ed49cda0edea5b95471f" dependencies = [ "equivalent", - "hashbrown 0.15.0", + "hashbrown 0.15.2", ] [[package]] @@ -528,21 +660,21 @@ dependencies = [ [[package]] name = "itoa" -version = "1.0.11" +version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674" [[package]] name = "libc" -version = "0.2.161" +version = "0.2.168" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e9489c2807c139ffd9c1794f4af0ebe86a828db53ecdc7fea2111d0fed085d1" +checksum = "5aaeb2981e0606ca11d79718f8bb01164f1d6ed75080182d3abf017e6d244b6d" [[package]] name = "libloading" -version = "0.8.5" +version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4979f22fdb869068da03c9f7528f8297c6fd2606bc3a4affe42e6a823fdb8da4" +checksum = "fc2f4eb4bc735547cfed7c0a4922cbd04a4655978c09b54f1f7b228750664c34" dependencies = [ "cfg-if", "windows-targets", @@ -554,6 +686,12 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" +[[package]] +name = "litemap" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ee93343901ab17bd981295f2cf0026d4ad018c7c31ba84549a4ddbb47a45104" + [[package]] name = "log" version = "0.4.22" @@ -606,7 +744,7 @@ checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] @@ -681,7 +819,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fb37767f6569cd834a413442455e0f066d0d522de8630436e2a1761d9726ba56" dependencies = [ "supports-color 2.1.0", - "supports-color 3.0.1", + "supports-color 3.0.2", ] [[package]] @@ -708,12 +846,12 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "pest" -version = "2.7.14" +version = "2.7.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "879952a81a83930934cbf1786752d6dedc3b1f29e8f8fb2ad1d0a36f377cf442" +checksum = "8b7cafe60d6cf8e62e1b9b2ea516a089c008945bb5a275416789e7db0bc199dc" dependencies = [ "memchr", - "thiserror", + "thiserror 2.0.6", "ucd-trie", ] @@ -747,15 +885,15 @@ dependencies = [ "serde", "serde_cbor", "serde_json", - "thiserror", + "thiserror 1.0.69", "uuid", ] [[package]] name = "pgrx-bindgen" -version = "0.12.8" +version = "0.12.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04d822012e4919882cb9b1851ae9e4b4b03703ddf2ef4af3522025c8944e9c7d" +checksum = "81cbcd956c2da35baaf0a116e6f6a49a6c2fbc8f6b332f66d6fd060bfd00615f" dependencies = [ "bindgen", "cc", @@ -765,7 +903,7 @@ dependencies = [ "proc-macro2", "quote", "shlex", - "syn 2.0.87", + "syn 2.0.90", "walkdir", ] @@ -788,14 +926,14 @@ dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] name = "pgrx-pg-config" -version = "0.12.8" +version = "0.12.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "945d664e710e4dcccd459628c6e58b3ccbd0495d2ed751b1d6f73996ceca01a2" +checksum = "86a64a4c6e4e43e73cf8d3379d9533df98ded45c920e1ba8131c979633d74132" dependencies = [ "cargo_toml", "eyre", @@ -804,7 +942,7 @@ dependencies = [ "pathsearch", "serde", "serde_json", - "thiserror", + "thiserror 1.0.69", "toml", "url", ] @@ -835,8 +973,8 @@ dependencies = [ "petgraph", "proc-macro2", "quote", - "syn 2.0.87", - "thiserror", + "syn 2.0.90", + "thiserror 1.0.69", "unescape", ] @@ -875,9 +1013,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.89" +version = "1.0.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f139b0662de085916d1fb67d2b4169d1addddda1919e696f3252b740b629986e" +checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0" dependencies = [ "unicode-ident", ] @@ -997,9 +1135,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.8" +version = "0.4.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "368758f23274712b504848e9d5a6f010445cc8b87a7cdb4d7cbee666c1288da3" +checksum = "809e8dc61f6de73b46c85f4c96486310fe304c434cfa43669d7b40f711150908" dependencies = [ "aho-corasick", "memchr", @@ -1106,18 +1244,18 @@ dependencies = [ [[package]] name = "semver-parser" -version = "0.10.2" +version = "0.10.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00b0bef5b7f9e0df16536d3961cfb6e84331c065b4066afb39768d0e319411f7" +checksum = "9900206b54a3527fdc7b8a938bffd94a568bac4f4aa8113b209df75a09c0dec2" dependencies = [ "pest", ] [[package]] name = "serde" -version = "1.0.214" +version = "1.0.216" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f55c3193aca71c12ad7890f1785d2b73e1b9f63a0bbc353c08ef26fe03fc56b5" +checksum = "0b9781016e935a97e8beecf0c933758c97a5520d32930e460142b4cd80c6338e" dependencies = [ "serde_derive", ] @@ -1134,20 +1272,20 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.214" +version = "1.0.216" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de523f781f095e28fa605cdce0f8307e451cc0fd14e2eb4cd2e98a355b147766" +checksum = "46f859dbbf73865c6627ed570e78961cd3ac92407a2d117204c49232485da55e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] name = "serde_json" -version = "1.0.132" +version = "1.0.133" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d726bfaff4b320266d395898905d0eba0345aae23b54aee3a737e260fd46db03" +checksum = "c7fceb2473b9166b2294ef05efcb65a3db80803f0b03ef86a5fc88a2b85ee377" dependencies = [ "itoa", "memchr", @@ -1189,6 +1327,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" +[[package]] +name = "smallvec" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" + [[package]] name = "sptr" version = "0.3.2" @@ -1219,9 +1363,9 @@ dependencies = [ [[package]] name = "supports-color" -version = "3.0.1" +version = "3.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8775305acf21c96926c900ad056abeef436701108518cf890020387236ac5a77" +checksum = "c64fc7232dd8d2e4ac5ce4ef302b1d81e0b80d055b9d77c7c4f51f6aa4c867d6" dependencies = [ "is_ci", ] @@ -1239,15 +1383,26 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.87" +version = "2.0.90" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25aa4ce346d03a6dcd68dd8b4010bcb74e54e62c90c573f394c46eae99aba32d" +checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31" dependencies = [ "proc-macro2", "quote", "unicode-ident", ] +[[package]] +name = "synstructure" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.90", +] + [[package]] name = "tap" version = "1.0.1" @@ -1256,22 +1411,52 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" [[package]] name = "thiserror" -version = "1.0.67" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec2a1820ebd077e2b90c4df007bebf344cd394098a13c563957d0afc83ea47" +dependencies = [ + "thiserror-impl 2.0.6", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b3c6efbfc763e64eb85c11c25320f0737cb7364c4b6336db90aa9ebe27a0bbd" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ - "thiserror-impl", + "proc-macro2", + "quote", + "syn 2.0.90", ] [[package]] name = "thiserror-impl" -version = "1.0.67" +version = "2.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b607164372e89797d78b8e23a6d67d5d1038c1c65efd52e1389ef8b77caba2a6" +checksum = "d65750cab40f4ff1929fb1ba509e9914eb756131cef4210da8d5d700d26f6312" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.90", +] + +[[package]] +name = "tinystr" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9117f5d4db391c1cf6927e7bea3db74b9a1c1add8f7eda9ffd5364f40f57b82f" +dependencies = [ + "displaydoc", + "zerovec", ] [[package]] @@ -1349,9 +1534,9 @@ checksum = "5ab17db44d7388991a428b2ee655ce0c212e862eff1768a455c58f9aad6e7893" [[package]] name = "unicode-ident" -version = "1.0.13" +version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" +checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" [[package]] name = "unicode-normalization" @@ -1376,15 +1561,27 @@ checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" [[package]] name = "url" -version = "2.5.2" +version = "2.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22784dbdf76fdde8af1aeda5622b546b422b6fc585325248a2bf9f5e41e94d6c" +checksum = "32f8b686cadd1473f4bd0117a5d28d36b1ade384ea9b5069a1c40aefed7fda60" dependencies = [ "form_urlencoded", - "idna", + "idna 1.0.3", "percent-encoding", ] +[[package]] +name = "utf16_iter" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + [[package]] name = "uuid" version = "1.11.0" @@ -1400,7 +1597,7 @@ version = "0.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db79c75af171630a3148bd3e6d7c4f42b6a9a014c2945bc5ed0020cbb8d9478e" dependencies = [ - "idna", + "idna 0.5.0", "once_cell", "regex", "serde", @@ -1421,7 +1618,7 @@ dependencies = [ "proc-macro-error", "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.90", ] [[package]] @@ -1599,6 +1796,18 @@ dependencies = [ "memchr", ] +[[package]] +name = "write16" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936" + +[[package]] +name = "writeable" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e9df38ee2d2c3c5948ea468a8406ff0db0b29ae1ffde1bcf20ef305bcc95c51" + [[package]] name = "wyz" version = "0.5.1" @@ -1617,6 +1826,30 @@ dependencies = [ "winapi", ] +[[package]] +name = "yoke" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "120e6aef9aa629e3d4f52dc8cc43a015c7724194c97dfaf45180d2daf2b77f40" +dependencies = [ + "serde", + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.90", + "synstructure", +] + [[package]] name = "zerocopy" version = "0.7.35" @@ -1635,5 +1868,48 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.90", +] + +[[package]] +name = "zerofrom" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cff3ee08c995dee1859d998dea82f7374f2826091dd9cd47def953cae446cd2e" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.90", + "synstructure", +] + +[[package]] +name = "zerovec" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa2b893d79df23bfb12d5461018d408ea19dfafe76c2c7ef6d4eba614f8ff079" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.90", ] diff --git a/Cargo.toml b/Cargo.toml index 66af418d..2996cc6c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,7 +20,7 @@ pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] -base = { git = "https://github.com/tensorchord/pgvecto.rs.git", rev = "c911e8a476effaf05cd4b1a037826b100177bdad" } +base = { git = "https://github.com/tensorchord/pgvecto.rs.git", rev = "0374439c8593ead1f61fc2477a4b79a49fcdceb8" } # lock algebra version forever so that the QR decomposition never changes for same input nalgebra = "=0.33.0" diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile index 3a01627b..306527fb 100644 --- a/docker/pgrx.Dockerfile +++ b/docker/pgrx.Dockerfile @@ -20,8 +20,10 @@ RUN set -eux; \ postgresql-common gnupg \ crossbuild-essential-arm64 \ qemu-user-static \ - libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config \ - clang + libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config + +RUN curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18; \ + sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 # set up sccache RUN set -ex; \ diff --git a/src/lib.rs b/src/lib.rs index 4c55ac98..2af441b7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,7 +25,6 @@ unsafe extern "C" fn _PG_init() { if unsafe { pgrx::pg_sys::IsUnderPostmaster } { pgrx::error!("vchord must be loaded via shared_preload_libraries."); } - base::simd::enable(); unsafe { vchordrq::init(); vchordrqfscan::init(); From 7578655d86a5e7f22571740f4b07c122ae5bab3f Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 13 Dec 2024 19:25:04 +0800 Subject: [PATCH 079/324] fix: remove sudo in dockerfile (#138) Signed-off-by: usamoi --- docker/pgrx.Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile index 306527fb..94e1ff0c 100644 --- a/docker/pgrx.Dockerfile +++ b/docker/pgrx.Dockerfile @@ -22,8 +22,8 @@ RUN set -eux; \ qemu-user-static \ libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config -RUN curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18; \ - sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 +RUN curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | bash -s -- 18; \ + update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 # set up sccache RUN set -ex; \ From ae8dab4c4b8f6e900bec5f4cad9455f7a2d0eecb Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 13 Dec 2024 19:33:10 +0800 Subject: [PATCH 080/324] fix: ci (#139) Signed-off-by: usamoi --- docker/pgrx.Dockerfile | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile index 94e1ff0c..f4d3032b 100644 --- a/docker/pgrx.Dockerfile +++ b/docker/pgrx.Dockerfile @@ -22,7 +22,8 @@ RUN set -eux; \ qemu-user-static \ libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config -RUN curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | bash -s -- 18; \ +RUN apt -y install lsb-release wget software-properties-common gnupg; \ + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | bash -s -- 18; \ update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 # set up sccache From ab89bf8397c5ded1ec91d1760f6ded5f916fa2a3 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 16 Dec 2024 15:46:37 +0800 Subject: [PATCH 081/324] fix: preprocess for halfvec (#140) Signed-off-by: usamoi --- src/vchordrq/index/am_options.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/vchordrq/index/am_options.rs b/src/vchordrq/index/am_options.rs index a357da40..25fcd0c3 100644 --- a/src/vchordrq/index/am_options.rs +++ b/src/vchordrq/index/am_options.rs @@ -172,14 +172,16 @@ impl Opfamily { (B::Vecf32(x), PgDistanceKind::L2) => O::Vecf32(x.own()), (B::Vecf32(x), PgDistanceKind::Dot) => O::Vecf32(x.own()), (B::Vecf32(x), PgDistanceKind::Cos) => O::Vecf32(x.function_normalize()), - (B::Vecf16(x), _) => O::Vecf16(x.own()), + (B::Vecf16(x), PgDistanceKind::L2) => O::Vecf16(x.own()), + (B::Vecf16(x), PgDistanceKind::Dot) => O::Vecf16(x.own()), + (B::Vecf16(x), PgDistanceKind::Cos) => O::Vecf16(x.function_normalize()), } } pub fn process(self, x: Distance) -> f32 { match self.pg_distance { PgDistanceKind::Cos => f32::from(x) + 1.0f32, PgDistanceKind::L2 => f32::from(x).sqrt(), - _ => f32::from(x), + PgDistanceKind::Dot => x.into(), } } pub fn distance_kind(self) -> DistanceKind { From 01e6a6e46e358ff3efb0abf42f5a3f903b87f216 Mon Sep 17 00:00:00 2001 From: Jinjing Zhou Date: Thu, 19 Dec 2024 10:12:47 +0800 Subject: [PATCH 082/324] docs: update README for clarity and new features (#142) Signed-off-by: Jinjing.Zhou Signed-off-by: Jinjing.Zhou --- README.md | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 3f55e0c3..3d53e542 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ discord invitation link Twitter Docker pulls -

Prior release: Previous Docker pulls

+

Docker pull for pgvecto.rs: Previous Docker pulls

VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search, and serves as the successor to [pgvecto.rs](https://github.com/tensorchord/pgvecto.rs). @@ -30,6 +30,8 @@ VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: **🔧 External Index Build**: Leverage IVF to build indexes externally (e.g., on GPU) for faster KMeans clustering, combined with RaBitQ[^2] compression to efficiently store vectors while maintaining search quality through autonomous reranking. +**📏 Long Vector Support**: Store and search vectors up to 65,535 dimensions, enabling the use of the best high-dimensional models like text-embedding-3-large with ease. + [^2]: Gao, Jianyang, and Cheng Long. "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search." Proceedings of the ACM on Management of Data 2.3 (2024): 1-27. ## Quick Start @@ -63,12 +65,22 @@ ALTER SYSTEM SET shared_preload_libraries = 'vchord.so'; To create the VectorChord RaBitQ(vchordrq) index, you can use the following SQL. ```SQL +-- Set residual_quantization to true and spherical_centroids to false for L2 distance -- CREATE INDEX ON gist_train USING vchordrq (embedding vector_l2_ops) WITH (options = $$ residual_quantization = true [build.internal] lists = [4096] spherical_centroids = false $$); + + +-- Set residual_quantization to false and spherical_centroids to true for cos/dot distance -- +CREATE INDEX ON laion USING vchordrq (embedding vector_cos_ops) WITH (options = $$ +residual_quantization = false +[build.internal] +lists = [4096] +spherical_centroids = true +$$); ``` ## Documentation @@ -201,9 +213,7 @@ cargo pgrx install --release --sudo # To install the extension into the system p ``` ## Limitations -- Data Type Support: Currently, only the `f32` data type is supported for vectors. -- Architecture Compatibility: The fast-scan kernel is optimized for x86_64 architectures. While it runs on aarch64, performance may be lower. -- KMeans Clustering: The built-in KMeans clustering is not yet fully optimized and may require substantial memory. We strongly recommend using external centroid precomputation for efficient index construction. +- KMeans Clustering: The built-in KMeans clustering depends on multi-thread in-memory build and may require substantial memory. We strongly recommend using external centroid precomputation for efficient index construction. ## License From 32f402ca744bca492b6ed5286adf9b79b817683a Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 23 Dec 2024 17:28:07 +0800 Subject: [PATCH 083/324] fix: set an implicit root in external build if parents are not set (#147) Signed-off-by: usamoi --- src/vchordrq/algorithm/build.rs | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/vchordrq/algorithm/build.rs b/src/vchordrq/algorithm/build.rs index d292628a..7c139353 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -229,6 +229,29 @@ impl Structure { vectors.insert(id, crate::projection::project(vector.as_borrowed().slice())); } }); + if parents.len() >= 2 && parents.values().all(|x| x.is_none()) { + // if there are more than one vertexs and no edges, + // assume there is an implicit root + let n = parents.len(); + let mut result = Vec::new(); + result.push(Structure { + means: vectors.values().cloned().collect::>(), + children: vec![Vec::new(); n], + }); + result.push(Structure { + means: vec![{ + // compute the vector on root, without normalizing it + let mut sum = vec![0.0f32; vector_options.dims as _]; + for vector in vectors.values() { + f32::vector_add_inplace(&mut sum, vector); + } + f32::vector_mul_scalar_inplace(&mut sum, 1.0 / n as f32); + sum + }], + children: vec![(0..n as u32).collect()], + }); + return result; + } let mut children = parents .keys() .map(|x| (*x, Vec::new())) From 23d8505674c98e156496f0a4ada501ec5896b63d Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Wed, 25 Dec 2024 09:51:34 +0800 Subject: [PATCH 084/324] chore: type check and external test (#149) Close #144 --------- Signed-off-by: cutecutecat --- scripts/train.py | 2 +- src/vchordrq/algorithm/build.rs | 32 ++++++--- tests/logic/external_build.slt | 117 ++++++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+), 11 deletions(-) create mode 100644 tests/logic/external_build.slt diff --git a/scripts/train.py b/scripts/train.py index 99b952c4..6d3c44b9 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -170,7 +170,7 @@ def kmeans_cluster( verbose=True, niter=niter, seed=SEED, - spherical=metric == "cos", + spherical=metric != "l2", ) child_kmeans.train(child_train) centroids.append(child_kmeans.centroids) diff --git a/src/vchordrq/algorithm/build.rs b/src/vchordrq/algorithm/build.rs index 7c139353..06f835ee 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -203,20 +203,32 @@ impl Structure { ) -> Vec { use std::collections::BTreeMap; let VchordrqExternalBuildOptions { table } = external_build; - let query = format!("SELECT id, parent, vector FROM {table};"); let mut parents = BTreeMap::new(); let mut vectors = BTreeMap::new(); pgrx::spi::Spi::connect(|client| { use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use base::vector::VectorBorrowed; use pgrx::pg_sys::panic::ErrorReportable; - let table = client.select(&query, None, None).unwrap_or_report(); - for row in table { + let schema_query = "SELECT n.nspname::TEXT + FROM pg_catalog.pg_extension e + LEFT JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vector';"; + let pgvector_schema: String = client + .select(schema_query, None, None) + .unwrap_or_report() + .first() + .get_by_name("nspname") + .expect("external build: cannot get schema of pgvector") + .expect("external build: cannot get schema of pgvector"); + let dump_query = + format!("SELECT id, parent, vector::{pgvector_schema}.vector FROM {table};"); + let centroids = client.select(&dump_query, None, None).unwrap_or_report(); + for row in centroids { let id: Option = row.get_by_name("id").unwrap(); let parent: Option = row.get_by_name("parent").unwrap(); let vector: Option = row.get_by_name("vector").unwrap(); - let id = id.expect("extern build: id could not be NULL"); - let vector = vector.expect("extern build: vector could not be NULL"); + let id = id.expect("external build: id could not be NULL"); + let vector = vector.expect("external build: vector could not be NULL"); let pop = parents.insert(id, parent); if pop.is_some() { pgrx::error!( @@ -224,7 +236,7 @@ impl Structure { ); } if vector_options.dims != vector.as_borrowed().dims() { - pgrx::error!("extern build: incorrect dimension, id = {id}"); + pgrx::error!("external build: incorrect dimension, id = {id}"); } vectors.insert(id, crate::projection::project(vector.as_borrowed().slice())); } @@ -275,7 +287,7 @@ impl Structure { } } let Some(root) = root else { - pgrx::error!("extern build: there are no root"); + pgrx::error!("external build: there are no root"); }; let mut heights = BTreeMap::<_, _>::new(); fn dfs_for_heights( @@ -284,7 +296,7 @@ impl Structure { u: i32, ) { if heights.contains_key(&u) { - pgrx::error!("extern build: detect a cycle, id = {u}"); + pgrx::error!("external build: detect a cycle, id = {u}"); } heights.insert(u, None); let mut height = None; @@ -293,7 +305,7 @@ impl Structure { let new = heights[&v].unwrap() + 1; if let Some(height) = height { if height != new { - pgrx::error!("extern build: two heights, id = {u}"); + pgrx::error!("external build: two heights, id = {u}"); } } else { height = Some(new); @@ -311,7 +323,7 @@ impl Structure { .collect::>(); if !(1..=8).contains(&(heights[&root] - 1)) { pgrx::error!( - "extern build: unexpected tree height, height = {}", + "external build: unexpected tree height, height = {}", heights[&root] ); } diff --git a/tests/logic/external_build.slt b/tests/logic/external_build.slt new file mode 100644 index 00000000..81dc35d8 --- /dev/null +++ b/tests/logic/external_build.slt @@ -0,0 +1,117 @@ +statement ok +CREATE TABLE t (val0 vector(3), val1 halfvec(3)); + +statement ok +INSERT INTO t (val0, val1) +SELECT + ARRAY[random(), random(), random()]::real[]::vector, + ARRAY[random(), random(), random()]::real[]::halfvec +FROM generate_series(1, 100); + +statement ok +CREATE TABLE vector_centroid (id integer, parent integer, vector vector(3)); + +statement ok +INSERT INTO vector_centroid (id, vector) VALUES + (0, '[1.0, 0.0, 0.0]'), + (1, '[0.0, 1.0, 0.0]'), + (2, '[0.0, 0.0, 1.0]'); + +statement ok +CREATE TABLE halfvec_centroid (id integer, parent integer, vector halfvec(3)); + +statement ok +INSERT INTO halfvec_centroid (id, vector) VALUES + (0, '[1.0, 0.0, 0.0]'), + (1, '[0.0, 1.0, 0.0]'), + (2, '[0.0, 0.0, 1.0]'); + +statement ok +CREATE TABLE real_centroid (id integer, parent integer, vector real[]); + +statement ok +INSERT INTO real_centroid (id, vector) VALUES + (0, '{1.0, 0.0, 0.0}'), + (1, '{0.0, 1.0, 0.0}'), + (2, '{0.0, 0.0, 1.0}'); + +statement ok +CREATE TABLE bad_type_centroid (id integer, parent integer, vector integer); + +statement ok +INSERT INTO bad_type_centroid (id, vector) VALUES + (0, 0), + (1, 0), + (2, 0); + +statement ok +CREATE TABLE bad_duplicate_id (id integer, parent integer, vector vector(3)); + +statement ok +INSERT INTO bad_duplicate_id (id, vector) VALUES + (1, '[1.0, 0.0, 0.0]'), + (1, '[0.0, 1.0, 0.0]'), + (2, '[0.0, 0.0, 1.0]'); + +# external build for vector column + +statement ok +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.vector_centroid' +$$); + +# external build for halfvec column + +statement ok +CREATE INDEX ON t USING vchordrq (val1 halfvec_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.vector_centroid' +$$); + +# external build for halfvec column by a halfvec table + +statement ok +CREATE INDEX ON t USING vchordrq (val1 halfvec_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.halfvec_centroid' +$$); + +# external build for halfvec column by a real[] table + +statement ok +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.real_centroid' +$$); + +# failed: bad vector data type + +statement error cannot cast type integer to (.*)vector +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.bad_type_centroid' +$$); + +# failed: duplicate id + +statement error external build: there are at least two lines have same id, id = 1 +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.bad_duplicate_id' +$$); + +statement ok +DROP TABLE t, vector_centroid, halfvec_centroid, real_centroid, bad_type_centroid, bad_duplicate_id; \ No newline at end of file From d855fac65fac99d07423e350467a821533fbe975 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 25 Dec 2024 14:09:52 +0800 Subject: [PATCH 085/324] chore: update dependencies (#151) Signed-off-by: usamoi --- Cargo.lock | 177 ++++++++++++++-------------------- Cargo.toml | 11 ++- rust-toolchain.toml | 2 +- src/bin/pgrx_embed.rs | 1 + src/lib.rs | 1 + src/upgrade/symbols.rs | 4 +- src/utils/k_means.rs | 4 +- src/vchordrq/index/am.rs | 2 +- src/vchordrqfscan/index/am.rs | 2 +- 9 files changed, 92 insertions(+), 112 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 982074bd..9439996e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.94" +version = "1.0.95" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1fd03a028ef38ba2276dce7e33fcd6369c158a1bca17946c4b1b701891c1ff7" +checksum = "34ac096ce696dc2fcabef30516bb13c0a68a11d30131d3df6f04711467681b04" [[package]] name = "approx" @@ -66,7 +66,7 @@ checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" [[package]] name = "base" version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=0374439c8593ead1f61fc2477a4b79a49fcdceb8#0374439c8593ead1f61fc2477a4b79a49fcdceb8" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74#9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74" dependencies = [ "base_macros", "cc", @@ -74,7 +74,7 @@ dependencies = [ "libc", "rand", "serde", - "thiserror 1.0.69", + "thiserror 2.0.9", "toml", "validator", ] @@ -82,11 +82,11 @@ dependencies = [ [[package]] name = "base_macros" version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=0374439c8593ead1f61fc2477a4b79a49fcdceb8#0374439c8593ead1f61fc2477a4b79a49fcdceb8" +source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74#9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -105,7 +105,7 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -150,9 +150,9 @@ dependencies = [ [[package]] name = "bytemuck" -version = "1.20.0" +version = "1.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b37c88a63ffd85d15b406896cc343916d7cf57838a847b3a6f2ca5d39a5695a" +checksum = "ef657dfab802224e671f5818e9a4935f9b1957ed18e58292690cc39e7a4092a3" [[package]] name = "byteorder" @@ -178,9 +178,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.3" +version = "1.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27f657647bcff5394bf56c7317665bbf790a137a50eaaa5c6bfbb9e27a518f2d" +checksum = "c31a0499c1dc64f458ad13872de75c0eb7e3fdb0e67964610c914b034fc5956e" dependencies = [ "shlex", ] @@ -232,9 +232,9 @@ dependencies = [ [[package]] name = "crossbeam-deque" -version = "0.8.5" +version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" dependencies = [ "crossbeam-epoch", "crossbeam-utils", @@ -251,9 +251,9 @@ dependencies = [ [[package]] name = "crossbeam-utils" -version = "0.8.20" +version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] name = "crunchy" @@ -282,7 +282,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -293,7 +293,7 @@ checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" dependencies = [ "darling_core", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -304,7 +304,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -330,7 +330,7 @@ checksum = "f282cfdfe92516eb26c2af8589c274c7c17681f5ecc03c18255fe741c6aa64eb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -454,11 +454,11 @@ checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc" [[package]] name = "home" -version = "0.5.9" +version = "0.5.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5" +checksum = "589533453244b0995c858700322199b2becb13b627df2851f64a2775d024abcf" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -576,7 +576,7 @@ checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -585,16 +585,6 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" -[[package]] -name = "idna" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "634d9b1461af396cad843f47fdba5597a4f9e6ddd4bfb6ff5d85028c25cb12f6" -dependencies = [ - "unicode-bidi", - "unicode-normalization", -] - [[package]] name = "idna" version = "1.0.3" @@ -666,9 +656,9 @@ checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674" [[package]] name = "libc" -version = "0.2.168" +version = "0.2.169" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5aaeb2981e0606ca11d79718f8bb01164f1d6ed75080182d3abf017e6d244b6d" +checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a" [[package]] name = "libloading" @@ -744,7 +734,7 @@ checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -851,7 +841,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b7cafe60d6cf8e62e1b9b2ea516a089c008945bb5a275416789e7db0bc199dc" dependencies = [ "memchr", - "thiserror 2.0.6", + "thiserror 2.0.9", "ucd-trie", ] @@ -867,9 +857,9 @@ dependencies = [ [[package]] name = "pgrx" -version = "0.12.8" +version = "0.12.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6468e2a7c4085707209cf7f7e14a7b73c8b7f41fc716283024a09180b620a45" +checksum = "227bf7e162ce710994306a97bc56bb3fe305f21120ab6692e2151c48416f5c0d" dependencies = [ "atomic-traits", "bitflags", @@ -903,7 +893,7 @@ dependencies = [ "proc-macro2", "quote", "shlex", - "syn 2.0.90", + "syn 2.0.91", "walkdir", ] @@ -919,14 +909,14 @@ dependencies = [ [[package]] name = "pgrx-macros" -version = "0.12.8" +version = "0.12.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52ffdb879a3880d3034661b4d19f802029abfcde6b8232299f4d4d2c202680af" +checksum = "e2f4291450d65e4deb770ce57ea93e22353d97950566222429cd166ebdf6f938" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -949,9 +939,9 @@ dependencies = [ [[package]] name = "pgrx-pg-sys" -version = "0.12.8" +version = "0.12.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "489c96adc8af0b917165f47a11e25f77753e227144aa3790be7e241afab9886c" +checksum = "63a5dc64f2a8226434118aa2c4700450fa42b04f29488ad98268848b21c1a4ec" dependencies = [ "cee-scape", "libc", @@ -964,16 +954,16 @@ dependencies = [ [[package]] name = "pgrx-sql-entity-graph" -version = "0.12.8" +version = "0.12.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de77a74f3dc80283a97322ab4d65f9f81691eb54d0a5b7c37b93686d351c6096" +checksum = "d81cc2e851c7e36b2f47c03e22d64d56c1d0e762fbde0039ba2cd490cfef3615" dependencies = [ "convert_case", "eyre", "petgraph", "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", "thiserror 1.0.69", "unescape", ] @@ -988,27 +978,25 @@ dependencies = [ ] [[package]] -name = "proc-macro-error" -version = "1.0.4" +name = "proc-macro-error-attr2" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" dependencies = [ - "proc-macro-error-attr", "proc-macro2", "quote", - "syn 1.0.109", - "version_check", ] [[package]] -name = "proc-macro-error-attr" -version = "1.0.4" +name = "proc-macro-error2" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" dependencies = [ + "proc-macro-error-attr2", "proc-macro2", "quote", - "version_check", + "syn 2.0.91", ] [[package]] @@ -1211,9 +1199,9 @@ checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" [[package]] name = "safe_arch" -version = "0.7.2" +version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3460605018fdc9612bce72735cba0d27efbcd9904780d44c7e3a9948f96148a" +checksum = "96b02de82ddbe1b636e6170c21be622223aea188ef2e139be0a5b219ec215323" dependencies = [ "bytemuck", ] @@ -1278,14 +1266,14 @@ checksum = "46f859dbbf73865c6627ed570e78961cd3ac92407a2d117204c49232485da55e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] name = "serde_json" -version = "1.0.133" +version = "1.0.134" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7fceb2473b9166b2294ef05efcb65a3db80803f0b03ef86a5fc88a2b85ee377" +checksum = "d00f4175c42ee48b15416f6193a959ba3a0d67fc699a0db9ad12df9f83991c7d" dependencies = [ "itoa", "memchr", @@ -1383,9 +1371,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.90" +version = "2.0.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31" +checksum = "d53cbcb5a243bd33b7858b1d7f4aca2153490815872d86d955d6ea29f743c035" dependencies = [ "proc-macro2", "quote", @@ -1400,7 +1388,7 @@ checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -1420,11 +1408,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.6" +version = "2.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fec2a1820ebd077e2b90c4df007bebf344cd394098a13c563957d0afc83ea47" +checksum = "f072643fd0190df67a8bab670c20ef5d8737177d6ac6b2e9a236cb096206b2cc" dependencies = [ - "thiserror-impl 2.0.6", + "thiserror-impl 2.0.9", ] [[package]] @@ -1435,18 +1423,18 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] name = "thiserror-impl" -version = "2.0.6" +version = "2.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d65750cab40f4ff1929fb1ba509e9914eb756131cef4210da8d5d700d26f6312" +checksum = "7b50fa271071aae2e6ee85f842e2e28ba8cd2c5fb67f11fcb1fd70b276f9e7d4" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -1461,9 +1449,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.8.0" +version = "1.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "445e881f4f6d382d5f27c034e25eb92edd7c784ceab92a0937db7f2e9471b938" +checksum = "022db8904dfa342efe721985167e9fcd16c29b226db4397ed752a761cfce81e8" dependencies = [ "tinyvec_macros", ] @@ -1526,27 +1514,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" -[[package]] -name = "unicode-bidi" -version = "0.3.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ab17db44d7388991a428b2ee655ce0c212e862eff1768a455c58f9aad6e7893" - [[package]] name = "unicode-ident" version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" -[[package]] -name = "unicode-normalization" -version = "0.1.24" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5033c97c4262335cded6d6fc3e5c18ab755e1a3dc96376350f3d8e9f009ad956" -dependencies = [ - "tinyvec", -] - [[package]] name = "unicode-segmentation" version = "1.12.0" @@ -1566,7 +1539,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32f8b686cadd1473f4bd0117a5d28d36b1ade384ea9b5069a1c40aefed7fda60" dependencies = [ "form_urlencoded", - "idna 1.0.3", + "idna", "percent-encoding", ] @@ -1593,11 +1566,11 @@ dependencies = [ [[package]] name = "validator" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db79c75af171630a3148bd3e6d7c4f42b6a9a014c2945bc5ed0020cbb8d9478e" +checksum = "d0b4a29d8709210980a09379f27ee31549b73292c87ab9899beee1c0d3be6303" dependencies = [ - "idna 0.5.0", + "idna", "once_cell", "regex", "serde", @@ -1609,16 +1582,16 @@ dependencies = [ [[package]] name = "validator_derive" -version = "0.18.2" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df0bcf92720c40105ac4b2dda2a4ea3aa717d4d6a862cc217da653a4bd5c6b10" +checksum = "bac855a2ce6f843beb229757e6e570a42e837bcb15e5f449dd48d5747d41bf77" dependencies = [ "darling", "once_cell", - "proc-macro-error", + "proc-macro-error2", "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -1846,7 +1819,7 @@ checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", "synstructure", ] @@ -1868,7 +1841,7 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] [[package]] @@ -1888,7 +1861,7 @@ checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", "synstructure", ] @@ -1911,5 +1884,5 @@ checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.90", + "syn 2.0.91", ] diff --git a/Cargo.toml b/Cargo.toml index 2996cc6c..2697ad16 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,7 +20,7 @@ pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] -base = { git = "https://github.com/tensorchord/pgvecto.rs.git", rev = "0374439c8593ead1f61fc2477a4b79a49fcdceb8" } +base = { git = "https://github.com/tensorchord/pgvecto.rs.git", rev = "9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74" } # lock algebra version forever so that the QR decomposition never changes for same input nalgebra = "=0.33.0" @@ -31,7 +31,7 @@ rkyv = { version = "=0.7.45", features = ["validation"] } half = { version = "2.4.1", features = ["rkyv"] } log = "0.4.22" paste = "1" -pgrx = { version = "=0.12.8", default-features = false, features = ["cshim"] } +pgrx = { version = "=0.12.9", default-features = false, features = ["cshim"] } pgrx-catalog = "0.1.0" rand = "0.8.5" rand_chacha = "0.3.1" @@ -39,12 +39,17 @@ rand_distr = "0.4.3" rayon = "1.10.0" serde = "1" toml = "0.8.19" -validator = { version = "0.18.1", features = ["derive"] } +validator = { version = "0.19.0", features = ["derive"] } [patch.crates-io] half = { git = "https://github.com/tensorchord/half-rs.git" } [lints] +rust.fuzzy_provenance_casts = "deny" +rust.unexpected_cfgs = { level = "warn", check-cfg = [ + 'cfg(feature, values("pg12"))', + 'cfg(pgrx_embed)', +] } rust.unsafe_op_in_unsafe_fn = "deny" rust.unused_lifetimes = "warn" rust.unused_qualifications = "warn" diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 0d943e1a..eb254600 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2024-11-19" +channel = "nightly-2024-12-25" diff --git a/src/bin/pgrx_embed.rs b/src/bin/pgrx_embed.rs index 5f5c4d85..20a006e8 100644 --- a/src/bin/pgrx_embed.rs +++ b/src/bin/pgrx_embed.rs @@ -1 +1,2 @@ +#![feature(strict_provenance_lints)] ::pgrx::pgrx_embed!(); diff --git a/src/lib.rs b/src/lib.rs index 2af441b7..26d80ed0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,7 @@ #![allow(clippy::int_plus_one)] #![allow(clippy::unused_unit)] #![allow(clippy::infallible_destructuring_match)] +#![feature(strict_provenance_lints)] mod datatype; mod postgres; diff --git a/src/upgrade/symbols.rs b/src/upgrade/symbols.rs index c275bb10..5998c1cb 100644 --- a/src/upgrade/symbols.rs +++ b/src/upgrade/symbols.rs @@ -8,7 +8,7 @@ macro_rules! symbol { ($t:ident) => { paste::paste! { - #[no_mangle] + #[unsafe(no_mangle)] #[doc(hidden)] #[pgrx::pg_guard] extern "C" fn [<$t _wrapper>](_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> pgrx::pg_sys::Datum { @@ -17,7 +17,7 @@ macro_rules! symbol { stringify!($t), ); } - #[no_mangle] + #[unsafe(no_mangle)] #[doc(hidden)] pub extern "C" fn []() -> &'static ::pgrx::pg_sys::Pg_finfo_record { const V1_API: ::pgrx::pg_sys::Pg_finfo_record = ::pgrx::pg_sys::Pg_finfo_record { diff --git a/src/utils/k_means.rs b/src/utils/k_means.rs index 97a810f8..8aeac72a 100644 --- a/src/utils/k_means.rs +++ b/src/utils/k_means.rs @@ -151,7 +151,7 @@ fn rabitq_index( parallelism .into_par_iter(0..n) .map(|i| { - fn gen(mut qvector: Vec) -> Vec { + fn generate(mut qvector: Vec) -> Vec { let dims = qvector.len() as u32; let t = dims.div_ceil(4); qvector.resize(qvector.len().next_multiple_of(4), 0); @@ -222,7 +222,7 @@ fn rabitq_index( } else { u8::reduce_sum_of_x_as_u32(&qvector) as f32 }; - (dis_v_2, b, k, qvector_sum, gen(qvector)) + (dis_v_2, b, k, qvector_sum, generate(qvector)) }; let mut result = (Distance::INFINITY, 0); diff --git a/src/vchordrq/index/am.rs b/src/vchordrq/index/am.rs index e9182c7a..8f22db6f 100644 --- a/src/vchordrq/index/am.rs +++ b/src/vchordrq/index/am.rs @@ -517,7 +517,7 @@ impl Drop for VchordrqLeader { } #[pgrx::pg_guard] -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn vchordrq_parallel_build_main( _seg: *mut pgrx::pg_sys::dsm_segment, toc: *mut pgrx::pg_sys::shm_toc, diff --git a/src/vchordrqfscan/index/am.rs b/src/vchordrqfscan/index/am.rs index 84a4c8ba..2b834fbc 100644 --- a/src/vchordrqfscan/index/am.rs +++ b/src/vchordrqfscan/index/am.rs @@ -487,7 +487,7 @@ impl Drop for VchordrqfscanLeader { } #[pgrx::pg_guard] -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn vchordrqfscan_parallel_build_main( _seg: *mut pgrx::pg_sys::dsm_segment, toc: *mut pgrx::pg_sys::shm_toc, From 6ae6548d72b3ed83a5d49d188f9ced56cfe174dc Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 25 Dec 2024 14:27:49 +0800 Subject: [PATCH 086/324] fix: update docker image in ci (#152) Signed-off-by: usamoi --- .github/workflows/psql.yml | 2 +- .github/workflows/release.yml | 2 +- .github/workflows/rust.yml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml index 55587859..18daba00 100644 --- a/.github/workflows/psql.yml +++ b/.github/workflows/psql.yml @@ -44,7 +44,7 @@ jobs: version: ["14", "15", "16", "17"] arch: ["x86_64"] env: - PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.8" + PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9" SQLLOGICTEST: "0.22.0" steps: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 700dc858..88f2e04b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -40,7 +40,7 @@ jobs: version: ["14", "15", "16", "17"] arch: ["x86_64", "aarch64"] env: - PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.8" + PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9" SEMVER: ${{ needs.semver.outputs.SEMVER }} steps: diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 7f6ffe04..fb1dbe6a 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -33,7 +33,7 @@ jobs: matrix: arch: ["x86_64", "aarch64"] env: - PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.8" + PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9" steps: - uses: actions/checkout@v4 From 6b31cd56c96f96d10bf32571dfbaf95db7137997 Mon Sep 17 00:00:00 2001 From: xieydd Date: Fri, 27 Dec 2024 11:28:04 +0800 Subject: [PATCH 087/324] Update enterprise dockerfile (#148) Add pg_later and pgmq. --------- Signed-off-by: xieydd --- docker/pg-cnpg/Dockerfile | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/docker/pg-cnpg/Dockerfile b/docker/pg-cnpg/Dockerfile index 99b04915..7f7c621c 100644 --- a/docker/pg-cnpg/Dockerfile +++ b/docker/pg-cnpg/Dockerfile @@ -4,9 +4,9 @@ ARG TARGETARCH FROM tensorchord/vchord-binary:pg${PG_MAJOR}-v${SEMVER} as binary -FROM rust:1.78-bookworm as builder -ARG TRUNK_VER=0.12.25 -ENV CARGO_REGISTRIES_CRATES_IO_PROTOCOL sparse +FROM rust:1.82-bookworm as builder +ARG TRUNK_VER=0.15.6 +ENV CARGO_REGISTRIES_CRATES_IO_PROTOCOL sparse RUN cargo install --version $TRUNK_VER pg-trunk FROM modelzai/pg-slim:${PG_MAJOR}-${TARGETARCH} @@ -132,6 +132,12 @@ RUN trunk install plpython3u # Install pgvector RUN trunk install pgvector --version ${PGVECTOR} +# Install pgmq +RUN trunk install pgmq --version 1.4.5 + +# Install pg_later +RUN trunk install pg_later --version 0.3.0 + # Clone and build AWS SDK for C++ RUN git clone https://github.com/aws/aws-sdk-cpp.git && \ cd aws-sdk-cpp && \ From c1b3b9d597289a4a7df54291670649cdf22f2b0c Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 27 Dec 2024 12:05:57 +0800 Subject: [PATCH 088/324] feat: build multiarch pgrx image (#153) - `sqllogictest-rs ` doesn't have Linux ARM64 release, tracked [here](https://github.com/risinglightdb/sqllogictest-rs/issues/240) - add rust-toolchain channel to image tag - use arm64 runner for clippy/build & release --------- Signed-off-by: Keming --- .github/workflows/pgrx.yml | 14 ++++++++++++-- .github/workflows/psql.yml | 16 ++++++++-------- .github/workflows/release.yml | 14 +++++++------- .github/workflows/rust.yml | 16 ++++++++++------ docker/pgrx.Dockerfile | 27 +++++++++++---------------- 5 files changed, 48 insertions(+), 39 deletions(-) diff --git a/.github/workflows/pgrx.yml b/.github/workflows/pgrx.yml index bf96d954..212de577 100644 --- a/.github/workflows/pgrx.yml +++ b/.github/workflows/pgrx.yml @@ -7,6 +7,10 @@ on: description: 'pgrx version' required: true type: string + toolchain: + description: 'additional rust toolchain' + required: true + type: string concurrency: group: ${{ github.ref }}-${{ github.workflow }} @@ -21,9 +25,11 @@ permissions: jobs: build: - runs-on: ubuntu-latest + runs-on: ubicloud-standard-8 steps: - uses: actions/checkout@v4 + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Login to ghcr.io @@ -38,6 +44,10 @@ jobs: context: . file: ./docker/pgrx.Dockerfile push: true + cache-from: type=gha + cache-to: type=gha,mode=max + platforms: "linux/amd64,linux/arm64" build-args: | PGRX_VERSION=${{ github.event.inputs.version }} - tags: ${{ env.IMAGE_NAME }}:${{ github.event.inputs.version }} + RUST_TOOLCHAIN=${{ github.event.inputs.toolchain }} + tags: ${{ env.IMAGE_NAME }}:${{ github.event.inputs.version }}-${{ github.event.inputs.toolchain }} diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml index 18daba00..0400ac56 100644 --- a/.github/workflows/psql.yml +++ b/.github/workflows/psql.yml @@ -38,14 +38,16 @@ env: jobs: test: - runs-on: ubuntu-latest + runs-on: ${{ matrix.runner }} strategy: matrix: version: ["14", "15", "16", "17"] - arch: ["x86_64"] + runner: ["ubuntu-latest"] env: - PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9" + PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" SQLLOGICTEST: "0.22.0" + ARCH: "x86_64" + PLATFORM: "amd64" steps: - uses: actions/checkout@v4 @@ -69,18 +71,16 @@ jobs: env: SEMVER: "0.0.0" VERSION: ${{ matrix.version }} - ARCH: ${{ matrix.arch }} - PLATFORM: "amd64" PROFILE: "opt" run: | - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE build --lib --features pg${{ matrix.version }} --target ${{ matrix.arch }}-unknown-linux-gnu --profile $PROFILE - docker run --rm -v .:/workspace $CACHE_ENVS --entrypoint bash $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --target ${{ matrix.arch }}-unknown-linux-gnu --profile $PROFILE + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg${{ matrix.version }} --target $ARCH-unknown-linux-gnu --profile $PROFILE + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --target $ARCH-unknown-linux-gnu --profile $PROFILE ./tools/package.sh docker build -t vchord:pg${{ matrix.version }} --build-arg PG_VERSION=${{ matrix.version }} -f ./docker/Dockerfile . - name: Setup SQL Logic Test run: | - curl -fsSL -o sqllogictest.tar.gz https://github.com/risinglightdb/sqllogictest-rs/releases/download/v${SQLLOGICTEST}/sqllogictest-bin-v${SQLLOGICTEST}-${{ matrix.arch }}-unknown-linux-musl.tar.gz + curl -fsSL -o sqllogictest.tar.gz https://github.com/risinglightdb/sqllogictest-rs/releases/download/v${SQLLOGICTEST}/sqllogictest-bin-v${SQLLOGICTEST}-$ARCH-unknown-linux-musl.tar.gz tar -xzf sqllogictest.tar.gz mv sqllogictest /usr/local/bin/ diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 88f2e04b..d7b17588 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -33,15 +33,17 @@ jobs: core.setOutput('SEMVER', tag); build: - runs-on: ubuntu-latest + runs-on: ${{ matrix.runner }} needs: ["semver"] strategy: matrix: version: ["14", "15", "16", "17"] - arch: ["x86_64", "aarch64"] + runner: ["ubicloud-standard-4", "ubicloud-standard-4-arm"] env: - PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9" + PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" SEMVER: ${{ needs.semver.outputs.SEMVER }} + ARCH: ${{ matrix.runner == 'ubicloud-standard-4' && 'x86_64' || 'aarch64' }} + PLATFORM: ${{ matrix.runner == 'ubicloud-standard-4' && 'amd64' || 'arm64' }} steps: - uses: actions/checkout@v4 @@ -64,13 +66,11 @@ jobs: - name: Build env: VERSION: ${{ matrix.version }} - ARCH: ${{ matrix.arch }} - PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} PROFILE: "release" GH_TOKEN: ${{ github.token }} run: | - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE build --lib --features pg${{ matrix.version }} --target ${{ matrix.arch }}-unknown-linux-gnu --profile $PROFILE - docker run --rm -v .:/workspace $CACHE_ENVS --entrypoint bash -e SEMVER=${SEMVER} $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --target ${{ matrix.arch }}-unknown-linux-gnu --profile $PROFILE + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg$VERSION --profile $PROFILE + docker run --rm -v .:/workspace $CACHE_ENVS -e SEMVER=${SEMVER} $PGRX_IMAGE ./tools/schema.sh --features pg$VERSION --profile $PROFILE ./tools/package.sh ls ./build gh release upload --clobber $SEMVER ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index fb1dbe6a..df5547a4 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -28,12 +28,16 @@ concurrency: jobs: test: - runs-on: ubuntu-latest strategy: matrix: - arch: ["x86_64", "aarch64"] + include: + - runner: "ubicloud-standard-4" + arch: "x86_64" + - runner: "ubicloud-standard-4-arm" + arch: "aarch64" + runs-on: ${{ matrix.runner }} env: - PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9" + PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" steps: - uses: actions/checkout@v4 @@ -56,14 +60,14 @@ jobs: - name: Clippy run: | for v in {14..17}; do - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE clippy --target ${{ matrix.arch }}-unknown-linux-gnu --features "pg$v" -- -D warnings + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo clippy --target ${{ matrix.arch }}-unknown-linux-gnu --features "pg$v" -- -D warnings done - name: Build run: | for v in {14..17}; do - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE build --lib --target ${{ matrix.arch }}-unknown-linux-gnu --features "pg$v" + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --target ${{ matrix.arch }}-unknown-linux-gnu --features "pg$v" done - name: Test run: | # pg agnostic tests - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE test --no-fail-fast --target ${{ matrix.arch }}-unknown-linux-gnu --features pg17 + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo test --no-fail-fast --target ${{ matrix.arch }}-unknown-linux-gnu --features pg17 diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile index f4d3032b..60b6225b 100644 --- a/docker/pgrx.Dockerfile +++ b/docker/pgrx.Dockerfile @@ -1,15 +1,17 @@ # CNPG only support Debian 12 (Bookworm) FROM ubuntu:22.04 -ARG PGRX_VERSION=0.12.8 -ARG SCCACHE_VERSION=0.8.2 +ARG PGRX_VERSION +ARG RUST_TOOLCHAIN +ARG TARGETARCH ENV DEBIAN_FRONTEND=noninteractive \ LANG=en_US.UTF-8 \ LC_ALL=en_US.UTF-8 \ RUSTFLAGS="-Dwarnings" \ RUST_BACKTRACE=1 \ - CARGO_TERM_COLOR=always + CARGO_TERM_COLOR=always \ + SCCACHE_VERSION=0.9.0 RUN set -eux; \ apt update; \ @@ -18,17 +20,16 @@ RUN set -eux; \ ca-certificates \ build-essential \ postgresql-common gnupg \ - crossbuild-essential-arm64 \ - qemu-user-static \ libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config -RUN apt -y install lsb-release wget software-properties-common gnupg; \ +RUN set -eux; \ + apt -y install lsb-release wget software-properties-common gnupg; \ curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | bash -s -- 18; \ update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 # set up sccache RUN set -ex; \ - curl -fsSL -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v${SCCACHE_VERSION}/sccache-v${SCCACHE_VERSION}-x86_64-unknown-linux-musl.tar.gz; \ + curl -fsSL -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v${SCCACHE_VERSION}/sccache-v${SCCACHE_VERSION}-$(uname -m)-unknown-linux-musl.tar.gz; \ tar -xzf sccache.tar.gz --strip-components=1; \ rm sccache.tar.gz; \ mv sccache /usr/local/bin/ @@ -47,16 +48,10 @@ RUN chown -R ubuntu:ubuntu /usr/share/postgresql/ /usr/lib/postgresql/ USER ubuntu ENV PATH="$PATH:/home/ubuntu/.cargo/bin" RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y -RUN rustup target add x86_64-unknown-linux-gnu aarch64-unknown-linux-gnu WORKDIR /workspace -COPY rust-toolchain.toml /workspace/rust-toolchain.toml -RUN rustup target add x86_64-unknown-linux-gnu aarch64-unknown-linux-gnu - -# ref: https://github.com/pgcentralfoundation/pgrx/blob/develop/docs/src/extension/build/cross-compile.md -RUN set -ex; \ - echo 'target.aarch64-unknown-linux-gnu.linker = "aarch64-linux-gnu-gcc"' >> ~/.cargo/config.toml; \ - echo 'target.aarch64-unknown-linux-gnu.runner = ["qemu-aarch64-static", "-L", "/usr/aarch64-linux-gnu"]' >> ~/.cargo/config.toml +RUN rustup toolchain install ${RUST_TOOLCHAIN} +RUN rustup target add $(uname -m)-unknown-linux-gnu RUN cargo install cargo-pgrx --locked --version=${PGRX_VERSION} @@ -65,4 +60,4 @@ RUN set -ex; \ cargo pgrx init --pg$v=/usr/lib/postgresql/$v/bin/pg_config; \ done; -ENTRYPOINT [ "cargo" ] +CMD [ "/bin/bash" ] From d35c8752b07a19c9553ea72c0673b929515dc47f Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 27 Dec 2024 14:40:42 +0800 Subject: [PATCH 089/324] chore: impl dereference traits for page guards (#156) Signed-off-by: usamoi --- src/postgres.rs | 61 ++++++++++++++++++-------- src/vchordrq/algorithm/build.rs | 32 +++++++------- src/vchordrq/algorithm/insert.rs | 50 +++++++++++---------- src/vchordrq/algorithm/mod.rs | 59 +++++++++++++++++++++++++ src/vchordrq/algorithm/prewarm.rs | 15 +++---- src/vchordrq/algorithm/scan.rs | 16 +++---- src/vchordrq/algorithm/vacuum.rs | 36 +++++++-------- src/vchordrq/algorithm/vectors.rs | 10 ++--- src/vchordrq/index/am.rs | 6 ++- src/vchordrqfscan/algorithm/build.rs | 8 ++-- src/vchordrqfscan/algorithm/insert.rs | 60 +++++++++++-------------- src/vchordrqfscan/algorithm/prewarm.rs | 13 ++---- src/vchordrqfscan/algorithm/scan.rs | 15 +++---- src/vchordrqfscan/algorithm/vacuum.rs | 28 +++++------- 14 files changed, 233 insertions(+), 176 deletions(-) diff --git a/src/postgres.rs b/src/postgres.rs index 06b402d9..c24d669e 100644 --- a/src/postgres.rs +++ b/src/postgres.rs @@ -1,4 +1,5 @@ use std::mem::{offset_of, MaybeUninit}; +use std::ops::{Deref, DerefMut}; use std::ptr::NonNull; const _: () = assert!( @@ -46,6 +47,14 @@ impl Page { assert_eq!(offset_of!(Self, opaque), this.header.pd_special as usize); this } + #[allow(dead_code)] + pub fn clone_into_boxed(&self) -> Box { + let mut result = Box::new_uninit(); + unsafe { + std::ptr::copy(self as *const Self, result.as_mut_ptr(), 1); + result.assume_init() + } + } pub fn get_opaque(&self) -> &Opaque { &self.opaque } @@ -170,15 +179,20 @@ pub struct BufferReadGuard { } impl BufferReadGuard { - pub fn get(&self) -> &Page { - unsafe { self.page.as_ref() } - } #[allow(dead_code)] pub fn id(&self) -> u32 { self.id } } +impl Deref for BufferReadGuard { + type Target = Page; + + fn deref(&self) -> &Page { + unsafe { self.page.as_ref() } + } +} + impl Drop for BufferReadGuard { fn drop(&mut self) { unsafe { @@ -197,15 +211,23 @@ pub struct BufferWriteGuard { } impl BufferWriteGuard { - pub fn get(&self) -> &Page { + pub fn id(&self) -> u32 { + self.id + } +} + +impl Deref for BufferWriteGuard { + type Target = Page; + + fn deref(&self) -> &Page { unsafe { self.page.as_ref() } } - pub fn get_mut(&mut self) -> &mut Page { +} + +impl DerefMut for BufferWriteGuard { + fn deref_mut(&mut self) -> &mut Page { unsafe { self.page.as_mut() } } - pub fn id(&self) -> u32 { - self.id - } } impl Drop for BufferWriteGuard { @@ -215,11 +237,7 @@ impl Drop for BufferWriteGuard { pgrx::pg_sys::GenericXLogAbort(self.state); } else { if self.tracking_freespace { - pgrx::pg_sys::RecordPageWithFreeSpace( - self.raw, - self.id, - self.get().freespace() as _, - ); + pgrx::pg_sys::RecordPageWithFreeSpace(self.raw, self.id, self.freespace() as _); pgrx::pg_sys::FreeSpaceMapVacuumRange(self.raw, self.id, self.id + 1); } pgrx::pg_sys::GenericXLogFinish(self.state); @@ -329,13 +347,9 @@ impl Relation { return None; } let write = self.write(id, true); - if write.get().freespace() < freespace as _ { + if write.freespace() < freespace as _ { // the free space is recorded incorrectly - pgrx::pg_sys::RecordPageWithFreeSpace( - self.raw, - id, - write.get().freespace() as _, - ); + pgrx::pg_sys::RecordPageWithFreeSpace(self.raw, id, write.freespace() as _); pgrx::pg_sys::FreeSpaceMapVacuumRange(self.raw, id, id + 1); continue; } @@ -343,4 +357,13 @@ impl Relation { } } } + #[allow(dead_code)] + pub fn len(&self) -> u32 { + unsafe { + pgrx::pg_sys::RelationGetNumberOfBlocksInFork( + self.raw, + pgrx::pg_sys::ForkNumber::MAIN_FORKNUM, + ) + } + } } diff --git a/src/vchordrq/algorithm/build.rs b/src/vchordrq/algorithm/build.rs index 06f835ee..6554bfc3 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -1,7 +1,7 @@ -use crate::postgres::BufferWriteGuard; -use crate::postgres::Relation; +use super::RelationWrite; use crate::vchordrq::algorithm::rabitq; use crate::vchordrq::algorithm::tuples::*; +use crate::vchordrq::algorithm::PageGuard; use crate::vchordrq::index::am_options::Opfamily; use crate::vchordrq::types::VchordrqBuildOptions; use crate::vchordrq::types::VchordrqExternalBuildOptions; @@ -32,7 +32,7 @@ pub fn build, R: Reporter>( vector_options: VectorOptions, vchordrq_options: VchordrqIndexingOptions, heap_relation: T, - relation: Relation, + relation: impl RelationWrite, mut reporter: R, ) { let dims = vector_options.dims; @@ -75,7 +75,7 @@ pub fn build, R: Reporter>( }; let mut meta = Tape::create(&relation, false); assert_eq!(meta.first(), 0); - let mut vectors = Tape::>::create(&relation, true); + let mut vectors = Tape::, _>::create(&relation, true); let mut pointer_of_means = Vec::>::new(); for i in 0..structures.len() { let mut level = Vec::new(); @@ -99,10 +99,10 @@ pub fn build, R: Reporter>( let mut level = Vec::new(); for j in 0..structures[i].len() { if i == 0 { - let tape = Tape::::create(&relation, false); + let tape = Tape::::create(&relation, false); level.push(tape.first()); } else { - let mut tape = Tape::::create(&relation, false); + let mut tape = Tape::::create(&relation, false); let h2_mean = &structures[i].means[j]; let h2_children = &structures[i].children[j]; for child in h2_children.iter().copied() { @@ -361,18 +361,18 @@ impl Structure { } } -struct Tape<'a, T> { - relation: &'a Relation, - head: BufferWriteGuard, +struct Tape<'a: 'b, 'b, T, R: 'b + RelationWrite> { + relation: &'a R, + head: R::WriteGuard<'b>, first: u32, tracking_freespace: bool, _phantom: PhantomData T>, } -impl<'a, T> Tape<'a, T> { - fn create(relation: &'a Relation, tracking_freespace: bool) -> Self { +impl<'a: 'b, 'b, T, R: 'b + RelationWrite> Tape<'a, 'b, T, R> { + fn create(relation: &'a R, tracking_freespace: bool) -> Self { let mut head = relation.extend(tracking_freespace); - head.get_mut().get_opaque_mut().skip = head.id(); + head.get_opaque_mut().skip = head.id(); let first = head.id(); Self { relation, @@ -387,19 +387,19 @@ impl<'a, T> Tape<'a, T> { } } -impl Tape<'_, T> +impl<'a: 'b, 'b, T, R: 'b + RelationWrite> Tape<'a, 'b, T, R> where T: rkyv::Serialize>, { fn push(&mut self, x: &T) -> (u32, u16) { let bytes = rkyv::to_bytes(x).expect("failed to serialize"); - if let Some(i) = self.head.get_mut().alloc(&bytes) { + if let Some(i) = self.head.alloc(&bytes) { (self.head.id(), i) } else { let next = self.relation.extend(self.tracking_freespace); - self.head.get_mut().get_opaque_mut().next = next.id(); + self.head.get_opaque_mut().next = next.id(); self.head = next; - if let Some(i) = self.head.get_mut().alloc(&bytes) { + if let Some(i) = self.head.alloc(&bytes) { (self.head.id(), i) } else { panic!("tuple is too large to fit in a fresh page") diff --git a/src/vchordrq/algorithm/insert.rs b/src/vchordrq/algorithm/insert.rs index ee47e9a7..323fa24f 100644 --- a/src/vchordrq/algorithm/insert.rs +++ b/src/vchordrq/algorithm/insert.rs @@ -1,7 +1,8 @@ -use crate::postgres::Relation; +use super::RelationWrite; use crate::vchordrq::algorithm::rabitq::fscan_process_lowerbound; use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; +use crate::vchordrq::algorithm::PageGuard; use base::always_equal::AlwaysEqual; use base::distance::Distance; use base::distance::DistanceKind; @@ -11,7 +12,7 @@ use std::cmp::Reverse; use std::collections::BinaryHeap; pub fn insert( - relation: Relation, + relation: impl RelationWrite + Clone, payload: Pointer, vector: V, distance_kind: DistanceKind, @@ -20,7 +21,6 @@ pub fn insert( let vector = vector.as_borrowed(); let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -84,9 +84,8 @@ pub fn insert( let mut current = list.0; while current != u32::MAX { let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { + for i in 1..=h1_guard.len() { let h1_tuple = h1_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -110,7 +109,7 @@ pub fn insert( AlwaysEqual(h1_tuple.first), )); } - current = h1_guard.get().get_opaque().next; + current = h1_guard.get_opaque().next; } } let mut heap = BinaryHeap::from(results); @@ -155,11 +154,18 @@ pub fn insert( t: code.t(), }) .unwrap(); - append(relation, list.0, &tuple, false, in_building, in_building); + append( + relation.clone(), + list.0, + &tuple, + false, + in_building, + in_building, + ); } fn append( - relation: Relation, + relation: impl RelationWrite, first: u32, tuple: &[u8], tracking_freespace: bool, @@ -168,7 +174,7 @@ fn append( ) -> (u32, u16) { if tracking_freespace { if let Some(mut write) = relation.search(tuple.len()) { - let i = write.get_mut().alloc(tuple).unwrap(); + let i = write.alloc(tuple).unwrap(); return (write.id(), i); } } @@ -176,24 +182,22 @@ fn append( let mut current = first; loop { let read = relation.read(current); - if read.get().freespace() as usize >= tuple.len() - || read.get().get_opaque().next == u32::MAX - { + if read.freespace() as usize >= tuple.len() || read.get_opaque().next == u32::MAX { drop(read); let mut write = relation.write(current, tracking_freespace); - if let Some(i) = write.get_mut().alloc(tuple) { + if let Some(i) = write.alloc(tuple) { return (current, i); } - if write.get().get_opaque().next == u32::MAX { + if write.get_opaque().next == u32::MAX { let mut extend = relation.extend(tracking_freespace); - write.get_mut().get_opaque_mut().next = extend.id(); + write.get_opaque_mut().next = extend.id(); drop(write); - if let Some(i) = extend.get_mut().alloc(tuple) { + if let Some(i) = extend.alloc(tuple) { let result = (extend.id(), i); drop(extend); if updating_skip { let mut past = relation.write(first, tracking_freespace); - let skip = &mut past.get_mut().get_opaque_mut().skip; + let skip = &mut past.get_opaque_mut().skip; assert!(*skip != u32::MAX); *skip = std::cmp::max(*skip, result.0); } @@ -202,16 +206,16 @@ fn append( panic!("a tuple cannot even be fit in a fresh page"); } } - if skipping_traversal && current == first && write.get().get_opaque().skip != first { - current = write.get().get_opaque().skip; + if skipping_traversal && current == first && write.get_opaque().skip != first { + current = write.get_opaque().skip; } else { - current = write.get().get_opaque().next; + current = write.get_opaque().next; } } else { - if skipping_traversal && current == first && read.get().get_opaque().skip != first { - current = read.get().get_opaque().skip; + if skipping_traversal && current == first && read.get_opaque().skip != first { + current = read.get_opaque().skip; } else { - current = read.get().get_opaque().next; + current = read.get_opaque().next; } } } diff --git a/src/vchordrq/algorithm/mod.rs b/src/vchordrq/algorithm/mod.rs index 88239a8e..41744d7d 100644 --- a/src/vchordrq/algorithm/mod.rs +++ b/src/vchordrq/algorithm/mod.rs @@ -6,3 +6,62 @@ pub mod scan; pub mod tuples; pub mod vacuum; pub mod vectors; + +use crate::postgres::Page; +use std::ops::{Deref, DerefMut}; + +pub trait PageGuard { + fn id(&self) -> u32; +} + +pub trait RelationRead { + type ReadGuard<'a>: PageGuard + Deref + where + Self: 'a; + fn read(&self, id: u32) -> Self::ReadGuard<'_>; +} + +pub trait RelationWrite: RelationRead { + type WriteGuard<'a>: PageGuard + DerefMut + where + Self: 'a; + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_>; + fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_>; + fn search(&self, freespace: usize) -> Option>; +} + +impl PageGuard for crate::postgres::BufferReadGuard { + fn id(&self) -> u32 { + self.id() + } +} + +impl PageGuard for crate::postgres::BufferWriteGuard { + fn id(&self) -> u32 { + self.id() + } +} + +impl RelationRead for crate::postgres::Relation { + type ReadGuard<'a> = crate::postgres::BufferReadGuard; + + fn read(&self, id: u32) -> Self::ReadGuard<'_> { + self.read(id) + } +} + +impl RelationWrite for crate::postgres::Relation { + type WriteGuard<'a> = crate::postgres::BufferWriteGuard; + + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.write(id, tracking_freespace) + } + + fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.extend(tracking_freespace) + } + + fn search(&self, freespace: usize) -> Option> { + self.search(freespace) + } +} diff --git a/src/vchordrq/algorithm/prewarm.rs b/src/vchordrq/algorithm/prewarm.rs index 26bb9866..6d01f7c1 100644 --- a/src/vchordrq/algorithm/prewarm.rs +++ b/src/vchordrq/algorithm/prewarm.rs @@ -1,13 +1,12 @@ -use crate::postgres::Relation; +use super::RelationRead; use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; use std::fmt::Write; -pub fn prewarm(relation: Relation, height: i32) -> String { +pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> String { let mut message = String::new(); let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -37,9 +36,8 @@ pub fn prewarm(relation: Relation, height: i32) -> String { counter += 1; pgrx::check_for_interrupts!(); let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { + for i in 1..=h1_guard.len() { let h1_tuple = h1_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -47,7 +45,7 @@ pub fn prewarm(relation: Relation, height: i32) -> String { vectors::vector_warm::(relation.clone(), h1_tuple.mean); results.push(h1_tuple.first); } - current = h1_guard.get().get_opaque().next; + current = h1_guard.get_opaque().next; } } writeln!(message, "number of tuples: {}", results.len()).unwrap(); @@ -66,16 +64,15 @@ pub fn prewarm(relation: Relation, height: i32) -> String { counter += 1; pgrx::check_for_interrupts!(); let h0_guard = relation.read(current); - for i in 1..=h0_guard.get().len() { + for i in 1..=h0_guard.len() { let _h0_tuple = h0_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); results.push(()); } - current = h0_guard.get().get_opaque().next; + current = h0_guard.get_opaque().next; } } writeln!(message, "number of tuples: {}", results.len()).unwrap(); diff --git a/src/vchordrq/algorithm/scan.rs b/src/vchordrq/algorithm/scan.rs index df4d93de..42357c06 100644 --- a/src/vchordrq/algorithm/scan.rs +++ b/src/vchordrq/algorithm/scan.rs @@ -1,4 +1,4 @@ -use crate::postgres::Relation; +use super::RelationRead; use crate::vchordrq::algorithm::rabitq::fscan_process_lowerbound; use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; @@ -11,7 +11,7 @@ use std::cmp::Reverse; use std::collections::BinaryHeap; pub fn scan( - relation: Relation, + relation: impl RelationRead + Clone, vector: V, distance_kind: DistanceKind, probes: Vec, @@ -20,7 +20,6 @@ pub fn scan( let vector = vector.as_borrowed(); let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -66,9 +65,8 @@ pub fn scan( let mut current = list.0; while current != u32::MAX { let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { + for i in 1..=h1_guard.len() { let h1_tuple = h1_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -92,7 +90,7 @@ pub fn scan( AlwaysEqual(h1_tuple.first), )); } - current = h1_guard.get().get_opaque().next; + current = h1_guard.get_opaque().next; } } let mut heap = BinaryHeap::from(results); @@ -121,6 +119,7 @@ pub fn scan( for i in (1..meta_tuple.height_of_root).rev() { lists = make_lists(lists, probes[i as usize - 1]); } + drop(meta_guard); { let mut results = Vec::new(); for list in lists { @@ -138,9 +137,8 @@ pub fn scan( let mut current = list.0; while current != u32::MAX { let h0_guard = relation.read(current); - for i in 1..=h0_guard.get().len() { + for i in 1..=h0_guard.len() { let h0_tuple = h0_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -164,7 +162,7 @@ pub fn scan( AlwaysEqual(h0_tuple.payload), )); } - current = h0_guard.get().get_opaque().next; + current = h0_guard.get_opaque().next; } } let mut heap = BinaryHeap::from(results); diff --git a/src/vchordrq/algorithm/vacuum.rs b/src/vchordrq/algorithm/vacuum.rs index 2b219c49..c77bd216 100644 --- a/src/vchordrq/algorithm/vacuum.rs +++ b/src/vchordrq/algorithm/vacuum.rs @@ -1,13 +1,16 @@ -use crate::postgres::Relation; +use super::RelationWrite; use crate::vchordrq::algorithm::tuples::*; use base::search::Pointer; -pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) -> bool) { +pub fn vacuum( + relation: impl RelationWrite, + delay: impl Fn(), + callback: impl Fn(Pointer) -> bool, +) { // step 1: vacuum height_0_tuple { let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -19,16 +22,15 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn let mut current = first; while current != u32::MAX { let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { + for i in 1..=h1_guard.len() { let h1_tuple = h1_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); results.push(h1_tuple.first); } - current = h1_guard.get().get_opaque().next; + current = h1_guard.get_opaque().next; } } results @@ -42,9 +44,8 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn delay(); let mut h0_guard = relation.write(current, false); let mut reconstruct_removes = Vec::new(); - for i in 1..=h0_guard.get().len() { + for i in 1..=h0_guard.len() { let h0_tuple = h0_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -53,8 +54,8 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn reconstruct_removes.push(i); } } - h0_guard.get_mut().reconstruct(&reconstruct_removes); - current = h0_guard.get().get_opaque().next; + h0_guard.reconstruct(&reconstruct_removes); + current = h0_guard.get_opaque().next; } } } @@ -63,7 +64,6 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn let mut current = { let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -74,8 +74,8 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn delay(); let read = relation.read(current); let flag = 'flag: { - for i in 1..=read.get().len() { - let Some(vector_tuple) = read.get().get(i) else { + for i in 1..=read.len() { + let Some(vector_tuple) = read.get(i) else { continue; }; let vector_tuple = @@ -91,21 +91,21 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn if flag { drop(read); let mut write = relation.write(current, true); - for i in 1..=write.get().len() { - let Some(vector_tuple) = write.get().get(i) else { + for i in 1..=write.len() { + let Some(vector_tuple) = write.get(i) else { continue; }; let vector_tuple = unsafe { rkyv::archived_root::>(vector_tuple) }; if let Some(payload) = vector_tuple.payload.as_ref().copied() { if callback(Pointer::new(payload)) { - write.get_mut().free(i); + write.free(i); } } } - current = write.get().get_opaque().next; + current = write.get_opaque().next; } else { - current = read.get().get_opaque().next; + current = read.get_opaque().next; } } } diff --git a/src/vchordrq/algorithm/vectors.rs b/src/vchordrq/algorithm/vectors.rs index 6a23f746..06075d37 100644 --- a/src/vchordrq/algorithm/vectors.rs +++ b/src/vchordrq/algorithm/vectors.rs @@ -1,11 +1,11 @@ use super::tuples::Vector; -use crate::postgres::Relation; +use super::RelationRead; use crate::vchordrq::algorithm::tuples::VectorTuple; use base::distance::Distance; use base::distance::DistanceKind; pub fn vector_dist( - relation: Relation, + relation: impl RelationRead, vector: V::Borrowed<'_>, mean: (u32, u16), payload: Option, @@ -25,7 +25,7 @@ pub fn vector_dist( return None; }; let vector_guard = relation.read(mean.0); - let Some(vector_tuple) = vector_guard.get().get(mean.1) else { + let Some(vector_tuple) = vector_guard.get(mean.1) else { // fails consistency check return None; }; @@ -54,11 +54,11 @@ pub fn vector_dist( )) } -pub fn vector_warm(relation: Relation, mean: (u32, u16)) { +pub fn vector_warm(relation: impl RelationRead, mean: (u32, u16)) { let mut cursor = Ok(mean); while let Ok(mean) = cursor { let vector_guard = relation.read(mean.0); - let Some(vector_tuple) = vector_guard.get().get(mean.1) else { + let Some(vector_tuple) = vector_guard.get(mean.1) else { // fails consistency check return; }; diff --git a/src/vchordrq/index/am.rs b/src/vchordrq/index/am.rs index 8f22db6f..1dc7cdcf 100644 --- a/src/vchordrq/index/am.rs +++ b/src/vchordrq/index/am.rs @@ -289,6 +289,7 @@ pub unsafe extern "C" fn ambuild( } else { let mut indtuples = 0; reporter.tuples_done(indtuples); + let relation = unsafe { Relation::new(index) }; match opfamily.vector_kind() { VectorKind::Vecf32 => { HeapRelation::>::traverse( @@ -296,7 +297,7 @@ pub unsafe extern "C" fn ambuild( true, |(pointer, vector)| { algorithm::insert::insert::>( - unsafe { Relation::new(index) }, + relation.clone(), pointer, vector, opfamily.distance_kind(), @@ -313,7 +314,7 @@ pub unsafe extern "C" fn ambuild( true, |(pointer, vector)| { algorithm::insert::insert::>( - unsafe { Relation::new(index) }, + relation.clone(), pointer, vector, opfamily.distance_kind(), @@ -627,6 +628,7 @@ unsafe fn parallel_build( } let index_relation = unsafe { Relation::new(index) }; + let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap, tablescandesc) }; let opfamily = unsafe { am_options::opfamily(index) }; let heap_relation = Heap { diff --git a/src/vchordrqfscan/algorithm/build.rs b/src/vchordrqfscan/algorithm/build.rs index 25a66119..9a0daa6e 100644 --- a/src/vchordrqfscan/algorithm/build.rs +++ b/src/vchordrqfscan/algorithm/build.rs @@ -167,7 +167,7 @@ pub fn build( } pointer_of_firsts.push(level); } - forwards.head.get_mut().get_opaque_mut().skip = vectors.first(); + forwards.head.get_opaque_mut().skip = vectors.first(); meta.push(&MetaTuple { dims, height_of_root: structures.len() as u32, @@ -399,13 +399,13 @@ where { fn push(&mut self, x: &T) -> (u32, u16) { let bytes = rkyv::to_bytes(x).expect("failed to serialize"); - if let Some(i) = self.head.get_mut().alloc(&bytes) { + if let Some(i) = self.head.alloc(&bytes) { (self.head.id(), i) } else { let next = self.relation.extend(self.tracking_freespace); - self.head.get_mut().get_opaque_mut().next = next.id(); + self.head.get_opaque_mut().next = next.id(); self.head = next; - if let Some(i) = self.head.get_mut().alloc(&bytes) { + if let Some(i) = self.head.alloc(&bytes) { (self.head.id(), i) } else { panic!("tuple is too large to fit in a fresh page") diff --git a/src/vchordrqfscan/algorithm/insert.rs b/src/vchordrqfscan/algorithm/insert.rs index 4dfd432f..1a0ab778 100644 --- a/src/vchordrqfscan/algorithm/insert.rs +++ b/src/vchordrqfscan/algorithm/insert.rs @@ -14,7 +14,6 @@ use std::collections::BinaryHeap; pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_kind: DistanceKind) { let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -35,22 +34,18 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k }) .unwrap(); if let Some(mut write) = relation.search(tuple.len()) { - let i = write.get_mut().alloc(&tuple).unwrap(); + let i = write.alloc(&tuple).unwrap(); break 'h0_vector (write.id(), i); } - let mut current = relation - .read(meta_tuple.forwards_first) - .get() - .get_opaque() - .skip; + let mut current = relation.read(meta_tuple.forwards_first).get_opaque().skip; let mut changed = false; loop { let read = relation.read(current); let flag = 'flag: { - if read.get().freespace() as usize >= tuple.len() { + if read.freespace() as usize >= tuple.len() { break 'flag true; } - if read.get().get_opaque().next == u32::MAX { + if read.get_opaque().next == u32::MAX { break 'flag true; } false @@ -58,28 +53,27 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k if flag { drop(read); let mut write = relation.write(current, true); - if let Some(i) = write.get_mut().alloc(&tuple) { + if let Some(i) = write.alloc(&tuple) { break (current, i); } - if write.get().get_opaque().next == u32::MAX { + if write.get_opaque().next == u32::MAX { if changed { relation .write(meta_tuple.forwards_first, false) - .get_mut() .get_opaque_mut() .skip = write.id(); } let mut extend = relation.extend(true); - write.get_mut().get_opaque_mut().next = extend.id(); - if let Some(i) = extend.get_mut().alloc(&tuple) { + write.get_opaque_mut().next = extend.id(); + if let Some(i) = extend.alloc(&tuple) { break (extend.id(), i); } else { panic!("a tuple cannot even be fit in a fresh page"); } } - current = write.get().get_opaque().next; + current = write.get_opaque().next; } else { - current = read.get().get_opaque().next; + current = read.get_opaque().next; } changed = true; } @@ -90,7 +84,6 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k if is_residual { let vector_guard = relation.read(meta_tuple.mean.0); let vector_tuple = vector_guard - .get() .get(meta_tuple.mean.1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -111,9 +104,8 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k let mut current = list.0; while current != u32::MAX { let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { + for i in 1..=h1_guard.len() { let h1_tuple = h1_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -141,7 +133,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k } } } - current = h1_guard.get().get_opaque().next; + current = h1_guard.get_opaque().next; } } let mut heap = BinaryHeap::from(results); @@ -151,7 +143,6 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); let vector_guard = relation.read(mean.0); let vector_tuple = vector_guard - .get() .get(mean.1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -196,9 +187,8 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k loop { let read = relation.read(current); let flag = 'flag: { - for i in 1..=read.get().len() { + for i in 1..=read.len() { let h0_tuple = read - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -207,10 +197,10 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k break 'flag true; } } - if read.get().freespace() as usize >= dummy.len() { + if read.freespace() as usize >= dummy.len() { break 'flag true; } - if read.get().get_opaque().next == u32::MAX { + if read.get_opaque().next == u32::MAX { break 'flag true; } false @@ -218,9 +208,9 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k if flag { drop(read); let mut write = relation.write(current, false); - for i in 1..=write.get().len() { + for i in 1..=write.len() { let flag = put( - write.get_mut().get_mut(i).expect("data corruption"), + write.get_mut(i).expect("data corruption"), dims, &code, h0_vector, @@ -230,9 +220,9 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k return; } } - if let Some(i) = write.get_mut().alloc(&dummy) { + if let Some(i) = write.alloc(&dummy) { let flag = put( - write.get_mut().get_mut(i).expect("data corruption"), + write.get_mut(i).expect("data corruption"), dims, &code, h0_vector, @@ -241,12 +231,12 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k assert!(flag, "a put fails even on a fresh tuple"); return; } - if write.get().get_opaque().next == u32::MAX { + if write.get_opaque().next == u32::MAX { let mut extend = relation.extend(false); - write.get_mut().get_opaque_mut().next = extend.id(); - if let Some(i) = extend.get_mut().alloc(&dummy) { + write.get_opaque_mut().next = extend.id(); + if let Some(i) = extend.alloc(&dummy) { let flag = put( - extend.get_mut().get_mut(i).expect("data corruption"), + extend.get_mut(i).expect("data corruption"), dims, &code, h0_vector, @@ -258,9 +248,9 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k panic!("a tuple cannot even be fit in a fresh page"); } } - current = write.get().get_opaque().next; + current = write.get_opaque().next; } else { - current = read.get().get_opaque().next; + current = read.get_opaque().next; } } } diff --git a/src/vchordrqfscan/algorithm/prewarm.rs b/src/vchordrqfscan/algorithm/prewarm.rs index ec7642a8..ec8976ab 100644 --- a/src/vchordrqfscan/algorithm/prewarm.rs +++ b/src/vchordrqfscan/algorithm/prewarm.rs @@ -6,7 +6,6 @@ pub fn prewarm(relation: Relation, height: i32) -> String { let mut message = String::new(); let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -22,7 +21,6 @@ pub fn prewarm(relation: Relation, height: i32) -> String { { let vector_guard = relation.read(meta_tuple.mean.0); let vector_tuple = vector_guard - .get() .get(meta_tuple.mean.1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -43,9 +41,8 @@ pub fn prewarm(relation: Relation, height: i32) -> String { counter += 1; pgrx::check_for_interrupts!(); let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { + for i in 1..=h1_guard.len() { let h1_tuple = h1_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -56,7 +53,6 @@ pub fn prewarm(relation: Relation, height: i32) -> String { let mean = h1_tuple.mean[j]; let vector_guard = relation.read(mean.0); let vector_tuple = vector_guard - .get() .get(mean.1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -65,7 +61,7 @@ pub fn prewarm(relation: Relation, height: i32) -> String { } } } - current = h1_guard.get().get_opaque().next; + current = h1_guard.get_opaque().next; } } writeln!(message, "number of tuples: {}", results.len()).unwrap(); @@ -84,9 +80,8 @@ pub fn prewarm(relation: Relation, height: i32) -> String { counter += 1; pgrx::check_for_interrupts!(); let h0_guard = relation.read(current); - for i in 1..=h0_guard.get().len() { + for i in 1..=h0_guard.len() { let h0_tuple = h0_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -97,7 +92,7 @@ pub fn prewarm(relation: Relation, height: i32) -> String { } } } - current = h0_guard.get().get_opaque().next; + current = h0_guard.get_opaque().next; } } writeln!(message, "number of tuples: {}", results.len()).unwrap(); diff --git a/src/vchordrqfscan/algorithm/scan.rs b/src/vchordrqfscan/algorithm/scan.rs index 202949ef..a691b6c4 100644 --- a/src/vchordrqfscan/algorithm/scan.rs +++ b/src/vchordrqfscan/algorithm/scan.rs @@ -20,7 +20,6 @@ pub fn scan( ) -> impl Iterator { let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -41,7 +40,6 @@ pub fn scan( if is_residual { let vector_guard = relation.read(meta_tuple.mean.0); let vector_tuple = vector_guard - .get() .get(meta_tuple.mean.1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -62,9 +60,8 @@ pub fn scan( let mut current = list.0; while current != u32::MAX { let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { + for i in 1..=h1_guard.len() { let h1_tuple = h1_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -92,7 +89,7 @@ pub fn scan( } } } - current = h1_guard.get().get_opaque().next; + current = h1_guard.get_opaque().next; } } let mut heap = BinaryHeap::from(results); @@ -102,7 +99,6 @@ pub fn scan( let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); let vector_guard = relation.read(mean.0); let vector_tuple = vector_guard - .get() .get(mean.1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -138,9 +134,8 @@ pub fn scan( let mut current = list.0; while current != u32::MAX { let h0_guard = relation.read(current); - for i in 1..=h0_guard.get().len() { + for i in 1..=h0_guard.len() { let h0_tuple = h0_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -168,7 +163,7 @@ pub fn scan( } } } - current = h0_guard.get().get_opaque().next; + current = h0_guard.get_opaque().next; } } let mut heap = BinaryHeap::from(results); @@ -177,7 +172,7 @@ pub fn scan( while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); let vector_guard = relation.read(mean.0); - let Some(vector_tuple) = vector_guard.get().get(mean.1) else { + let Some(vector_tuple) = vector_guard.get(mean.1) else { // fails consistency check continue; }; diff --git a/src/vchordrqfscan/algorithm/vacuum.rs b/src/vchordrqfscan/algorithm/vacuum.rs index 7bb51796..8ede4f64 100644 --- a/src/vchordrqfscan/algorithm/vacuum.rs +++ b/src/vchordrqfscan/algorithm/vacuum.rs @@ -8,7 +8,6 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - { let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -20,9 +19,8 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - let mut current = first; while current != u32::MAX { let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { + for i in 1..=h1_guard.len() { let h1_tuple = h1_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -33,7 +31,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - } } } - current = h1_guard.get().get_opaque().next; + current = h1_guard.get_opaque().next; } } results @@ -46,9 +44,8 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - while current != u32::MAX { delay(); let mut h0_guard = relation.write(current, false); - for i in 1..=h0_guard.get().len() { + for i in 1..=h0_guard.len() { let h0_tuple = h0_guard - .get() .get(i) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -64,7 +61,6 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - if flag { // todo: use mutable API let mut temp = h0_guard - .get() .get(i) .map(rkyv::from_bytes::) .expect("data corruption") @@ -76,14 +72,13 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - } let temp = rkyv::to_bytes::<_, 8192>(&temp).expect("failed to serialize"); h0_guard - .get_mut() .get_mut(i) .expect("data corruption") .copy_from_slice(&temp); } } // todo: cross-tuple vacuum so that we can skip a tuple - current = h0_guard.get().get_opaque().next; + current = h0_guard.get_opaque().next; } } } @@ -92,7 +87,6 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - let mut current = { let meta_guard = relation.read(0); let meta_tuple = meta_guard - .get() .get(1) .map(rkyv::check_archived_root::) .expect("data corruption") @@ -103,8 +97,8 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - delay(); let read = relation.read(current); let flag = 'flag: { - for i in 1..=read.get().len() { - let Some(vector_tuple) = read.get().get(i) else { + for i in 1..=read.len() { + let Some(vector_tuple) = read.get(i) else { continue; }; let vector_tuple = rkyv::check_archived_root::(vector_tuple) @@ -120,21 +114,21 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - if flag { drop(read); let mut write = relation.write(current, true); - for i in 1..=write.get().len() { - let Some(vector_tuple) = write.get().get(i) else { + for i in 1..=write.len() { + let Some(vector_tuple) = write.get(i) else { continue; }; let vector_tuple = rkyv::check_archived_root::(vector_tuple) .expect("data corruption"); if let Some(payload) = vector_tuple.payload.as_ref().copied() { if callback(Pointer::new(payload)) { - write.get_mut().free(i); + write.free(i); } } } - current = write.get().get_opaque().next; + current = write.get_opaque().next; } else { - current = read.get().get_opaque().next; + current = read.get_opaque().next; } } } From 5237a0a11a9186da24a308b280003505d56ba0fd Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 27 Dec 2024 16:14:30 +0800 Subject: [PATCH 090/324] chore: fix the psql & release CI target, update readme (#158) - clean the target option - fix release CI - update readme about how to use the vectorchord-pgrx image --------- Signed-off-by: Keming --- .github/workflows/psql.yml | 4 ++-- README.md | 5 +++++ scripts/README.md | 28 ++++++++++++++++++++++++++-- tools/package.sh | 4 ++-- 4 files changed, 35 insertions(+), 6 deletions(-) diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml index 0400ac56..8a7f7984 100644 --- a/.github/workflows/psql.yml +++ b/.github/workflows/psql.yml @@ -73,8 +73,8 @@ jobs: VERSION: ${{ matrix.version }} PROFILE: "opt" run: | - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg${{ matrix.version }} --target $ARCH-unknown-linux-gnu --profile $PROFILE - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --target $ARCH-unknown-linux-gnu --profile $PROFILE + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg${{ matrix.version }} --profile $PROFILE + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --profile $PROFILE ./tools/package.sh docker build -t vchord:pg${{ matrix.version }} --build-arg PG_VERSION=${{ matrix.version }} -f ./docker/Dockerfile . diff --git a/README.md b/README.md index 3d53e542..9cd6cf2c 100644 --- a/README.md +++ b/README.md @@ -204,7 +204,12 @@ $$); To simplify the workflow, we provide end-to-end scripts for external index pre-computation, see [scripts](./scripts/README.md#run-external-index-precomputation-toolkit). +### Build the Postgres Docker Image with VectorChord extension + +Follow the steps in [Dev Guidance](./scripts/README.md#build-docker). + ### Installing From Source + Install pgrx according to [pgrx's instruction](https://github.com/pgcentralfoundation/pgrx?tab=readme-ov-file#getting-started). ```bash cargo install --locked cargo-pgrx diff --git a/scripts/README.md b/scripts/README.md index 1c282ea2..938ae900 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -1,19 +1,43 @@ ## Build Docker +Users can choose to build the package with the provided docker image or create the development environment by themselves. + +- (option 1) With `vectorchord-pgrx` Image + +```shell +# use the required version of `pgrx` and `rust` +export PGRX_VERSION=$(awk -F'version = "=|"' '/^pgrx\s*=.*version/ {print $2}' Cargo.toml) +export RUST_TOOLCHAIN=$(awk -F'"' '/^\s*channel\s*=/ {print $2}' rust-toolchain.toml) +export PGRX_IMAGE=ghcr.io/tensorchord/vectorchord-pgrx:$PGRX_VERSION-$RUST_TOOLCHAIN + +docker run --rm -v .:/workspace $PGRX_IMAGE cargo build --lib --features pg16 --profile opt +docker run --rm -v .:/workspace $PGRX_IMAGE ./tools/schema.sh --features pg16 --profile opt +``` + +- (option 2) With Local Development Environment + ```shell sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config clang cargo install --locked cargo-pgrx cargo pgrx init -cargo build --package vchord --lib --features pg16 --target x86_64-unknown-linux-gnu --profile opt -./tools/schema.sh --features pg16 --target x86_64-unknown-linux-gnu --profile opt +cargo build --package vchord --lib --features pg16 --profile opt +./tools/schema.sh --features pg16 --profile opt +``` + +- build the debian package +```shell export SEMVER="0.0.0" export VERSION="16" export ARCH="x86_64" export PLATFORM="amd64" export PROFILE="opt" ./tools/package.sh +``` + +- build the docker image +```shell docker build -t vchord:pg16-latest --build-arg PG_VERSION=16 -f ./docker/Dockerfile . ``` diff --git a/tools/package.sh b/tools/package.sh index 55f70369..3f6172ee 100755 --- a/tools/package.sh +++ b/tools/package.sh @@ -13,9 +13,9 @@ rm -rf ./build/dir_deb rm -rf ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb mkdir -p ./build/dir_zip -cp ./target/${ARCH}-unknown-linux-gnu/${PROFILE}/schema.sql ./build/dir_zip/vchord--$SEMVER.sql +cp ./target/${PROFILE}/schema.sql ./build/dir_zip/vchord--$SEMVER.sql sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/dir_zip/vchord.control -cp ./target/${ARCH}-unknown-linux-gnu/${PROFILE}/libvchord.so ./build/dir_zip/vchord.so +cp ./target/${PROFILE}/libvchord.so ./build/dir_zip/vchord.so zip ./build/vchord-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip -j ./build/dir_zip/* mkdir -p ./build/dir_deb From 4878ab8f9c4da34a9d751f7f93a576485b76e74b Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 27 Dec 2024 16:26:03 +0800 Subject: [PATCH 091/324] chore: add postgres sqllogicaltest for arm (#159) Signed-off-by: Keming --- .github/workflows/psql.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml index 8a7f7984..17f4bacb 100644 --- a/.github/workflows/psql.yml +++ b/.github/workflows/psql.yml @@ -3,7 +3,7 @@ name: PostgresSQL on: pull_request: paths: - - '.github/workflows/rust.yml' + - '.github/workflows/psql.yml' - 'src/**' - 'Cargo.lock' - 'Cargo.toml' @@ -15,7 +15,7 @@ on: branches: - main paths: - - '.github/workflows/rust.yml' + - '.github/workflows/psql.yml' - 'src/**' - 'Cargo.lock' - 'Cargo.toml' @@ -42,12 +42,12 @@ jobs: strategy: matrix: version: ["14", "15", "16", "17"] - runner: ["ubuntu-latest"] + runner: ["ubicloud-standard-4", "ubicloud-standard-4-arm"] env: PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" - SQLLOGICTEST: "0.22.0" - ARCH: "x86_64" - PLATFORM: "amd64" + SQLLOGICTEST: "0.25.0" + ARCH: ${{ matrix.runner == 'ubicloud-standard-4' && 'x86_64' || 'aarch64' }} + PLATFORM: ${{ matrix.runner == 'ubicloud-standard-4' && 'amd64' || 'arm64' }} steps: - uses: actions/checkout@v4 From a8059410ed0929744c94caa5c422ec8a49a1168e Mon Sep 17 00:00:00 2001 From: Jinjing Zhou Date: Mon, 30 Dec 2024 16:34:43 +0800 Subject: [PATCH 092/324] chore: update link in readme (#160) Signed-off-by: Jinjing.Zhou --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 9cd6cf2c..0cff5bb9 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search, and serves as the successor to [pgvecto.rs](https://github.com/tensorchord/pgvecto.rs). -With VectorChord, you can store 400,000 vectors for just $1, enabling significant savings: 6x more vectors compared to Pinecone's optimized storage and 26x more than pgvector/pgvecto.rs for the same price[^1]. For further insights, check out our [launch blog post](https://blog.pgvecto.rs/vectorchord-store-400k-vectors-for-1-in-postgresql). +With VectorChord, you can store 400,000 vectors for just $1, enabling significant savings: 6x more vectors compared to Pinecone's optimized storage and 26x more than pgvector/pgvecto.rs for the same price[^1]. For further insights, check out our [launch blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql). [^1]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/#/) with 768-dimensional vectors and 95% recall. @@ -22,7 +22,7 @@ VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: **⚡ Enhanced Performance**: Delivering optimized operations with up to 5x faster queries, 16x higher insert throughput, and 16x quicker[^3] index building compared to pgvector's HNSW implementation. -[^3]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/#/) with 768-dimensional vectors. Please checkout our [blog post](https://blog.pgvecto.rs/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. +[^3]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/#/) with 768-dimensional vectors. Please checkout our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. **💰 Affordable Vector Search**: Query 100M 768-dimensional vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%, helping you keep infrastructure costs down while maintaining high search quality. From ee123ada9b9dc100ee45d1804dfec724e3e48a26 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 7 Jan 2025 14:59:10 +0800 Subject: [PATCH 093/324] refactor: move pgvecto.rs base to this repo (#161) Signed-off-by: usamoi --- .github/workflows/psql.yml | 6 +- .github/workflows/release.yml | 5 +- Cargo.lock | 169 +- Cargo.toml | 54 +- build.rs | 4 + crates/algorithm/Cargo.toml | 7 + crates/algorithm/src/lib.rs | 41 + crates/always_equal/Cargo.toml | 7 + crates/always_equal/src/lib.rs | 30 + crates/distance/Cargo.toml | 7 + crates/distance/src/lib.rs | 79 + crates/rabitq/Cargo.toml | 11 + crates/rabitq/src/binary.rs | 123 + crates/rabitq/src/block.rs | 172 ++ crates/rabitq/src/lib.rs | 5 + .../rabitq/src/utils.rs | 0 crates/random_orthogonal_matrix/Cargo.toml | 15 + crates/random_orthogonal_matrix/src/lib.rs | 55 + crates/simd/Cargo.toml | 18 + crates/simd/build.rs | 12 + crates/simd/cshim.c | 223 ++ crates/simd/src/aligned.rs | 4 + crates/simd/src/bit.rs | 399 +++ crates/simd/src/emulate.rs | 208 ++ crates/simd/src/f16.rs | 1185 +++++++++ crates/simd/src/f32.rs | 2277 +++++++++++++++++ crates/simd/src/fast_scan/mod.rs | 496 ++++ crates/simd/src/lib.rs | 107 + crates/simd/src/packed_u4.rs | 18 + crates/simd/src/quantize.rs | 291 +++ crates/simd/src/u8.rs | 343 +++ crates/simd_macros/Cargo.toml | 18 + crates/simd_macros/src/lib.rs | 206 ++ crates/simd_macros/src/target.rs | 83 + crates/vector/Cargo.toml | 13 + crates/vector/src/bvect.rs | 274 ++ crates/vector/src/lib.rs | 48 + {src/types => crates/vector/src}/scalar8.rs | 13 +- crates/vector/src/svect.rs | 445 ++++ crates/vector/src/vect.rs | 216 ++ rustfmt.toml | 1 + scripts/README.md | 9 +- src/bin/pgrx_embed.rs | 1 - src/datatype/binary_scalar8.rs | 4 +- src/datatype/functions_scalar8.rs | 12 +- src/datatype/memory_pgvector_halfvec.rs | 3 +- src/datatype/memory_pgvector_vector.rs | 3 +- src/datatype/memory_scalar8.rs | 4 +- src/datatype/operators_pgvector_halfvec.rs | 5 +- src/datatype/operators_pgvector_vector.rs | 5 +- src/datatype/operators_scalar8.rs | 4 +- src/datatype/text_scalar8.rs | 2 +- src/datatype/typmod.rs | 2 - src/lib.rs | 8 +- src/postgres.rs | 158 +- src/projection.rs | 72 +- src/types/mod.rs | 1 - src/utils/k_means.rs | 194 +- src/utils/mod.rs | 1 - src/utils/parallelism.rs | 5 +- src/vchordrq/algorithm/build.rs | 15 +- src/vchordrq/algorithm/insert.rs | 23 +- src/vchordrq/algorithm/mod.rs | 59 - src/vchordrq/algorithm/prewarm.rs | 2 +- src/vchordrq/algorithm/rabitq.rs | 131 +- src/vchordrq/algorithm/scan.rs | 22 +- src/vchordrq/algorithm/tuples.rs | 59 +- src/vchordrq/algorithm/vacuum.rs | 12 +- src/vchordrq/algorithm/vectors.rs | 9 +- src/vchordrq/gucs/executing.rs | 6 +- src/vchordrq/index/am.rs | 36 +- src/vchordrq/index/am_options.rs | 10 +- src/vchordrq/index/am_scan.rs | 12 +- src/vchordrq/index/functions.rs | 6 +- src/vchordrq/index/utils.rs | 10 +- src/vchordrq/types.rs | 10 +- src/vchordrqfscan/algorithm/build.rs | 33 +- src/vchordrqfscan/algorithm/insert.rs | 27 +- src/vchordrqfscan/algorithm/prewarm.rs | 4 +- src/vchordrqfscan/algorithm/rabitq.rs | 171 +- src/vchordrqfscan/algorithm/scan.rs | 21 +- src/vchordrqfscan/algorithm/tuples.rs | 8 +- src/vchordrqfscan/algorithm/vacuum.rs | 18 +- src/vchordrqfscan/gucs/executing.rs | 6 +- src/vchordrqfscan/index/am.rs | 26 +- src/vchordrqfscan/index/am_options.rs | 8 +- src/vchordrqfscan/index/am_scan.rs | 10 +- src/vchordrqfscan/index/functions.rs | 4 +- src/vchordrqfscan/index/utils.rs | 10 +- src/vchordrqfscan/types.rs | 19 +- tests/logic/reindex.slt | 2 +- tools/package.sh | 5 +- tools/schema.sh | 8 - 93 files changed, 7995 insertions(+), 988 deletions(-) create mode 100644 build.rs create mode 100644 crates/algorithm/Cargo.toml create mode 100644 crates/algorithm/src/lib.rs create mode 100644 crates/always_equal/Cargo.toml create mode 100644 crates/always_equal/src/lib.rs create mode 100644 crates/distance/Cargo.toml create mode 100644 crates/distance/src/lib.rs create mode 100644 crates/rabitq/Cargo.toml create mode 100644 crates/rabitq/src/binary.rs create mode 100644 crates/rabitq/src/block.rs create mode 100644 crates/rabitq/src/lib.rs rename src/utils/infinite_byte_chunks.rs => crates/rabitq/src/utils.rs (100%) create mode 100644 crates/random_orthogonal_matrix/Cargo.toml create mode 100644 crates/random_orthogonal_matrix/src/lib.rs create mode 100644 crates/simd/Cargo.toml create mode 100644 crates/simd/build.rs create mode 100644 crates/simd/cshim.c create mode 100644 crates/simd/src/aligned.rs create mode 100644 crates/simd/src/bit.rs create mode 100644 crates/simd/src/emulate.rs create mode 100644 crates/simd/src/f16.rs create mode 100644 crates/simd/src/f32.rs create mode 100644 crates/simd/src/fast_scan/mod.rs create mode 100644 crates/simd/src/lib.rs create mode 100644 crates/simd/src/packed_u4.rs create mode 100644 crates/simd/src/quantize.rs create mode 100644 crates/simd/src/u8.rs create mode 100644 crates/simd_macros/Cargo.toml create mode 100644 crates/simd_macros/src/lib.rs create mode 100644 crates/simd_macros/src/target.rs create mode 100644 crates/vector/Cargo.toml create mode 100644 crates/vector/src/bvect.rs create mode 100644 crates/vector/src/lib.rs rename {src/types => crates/vector/src}/scalar8.rs (93%) create mode 100644 crates/vector/src/svect.rs create mode 100644 crates/vector/src/vect.rs create mode 100644 rustfmt.toml delete mode 100644 src/types/mod.rs diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml index 17f4bacb..a6311893 100644 --- a/.github/workflows/psql.yml +++ b/.github/workflows/psql.yml @@ -34,7 +34,6 @@ env: CARGO_TERM_COLOR: always RUST_BACKTRACE: 1 RUSTFLAGS: "-Dwarnings" - CARGO_PROFILE_OPT_BUILD_OVERRIDE_DEBUG: true jobs: test: @@ -71,10 +70,9 @@ jobs: env: SEMVER: "0.0.0" VERSION: ${{ matrix.version }} - PROFILE: "opt" run: | - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg${{ matrix.version }} --profile $PROFILE - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --profile $PROFILE + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg${{ matrix.version }} --release + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --release ./tools/package.sh docker build -t vchord:pg${{ matrix.version }} --build-arg PG_VERSION=${{ matrix.version }} -f ./docker/Dockerfile . diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index d7b17588..96088248 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -66,11 +66,10 @@ jobs: - name: Build env: VERSION: ${{ matrix.version }} - PROFILE: "release" GH_TOKEN: ${{ github.token }} run: | - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg$VERSION --profile $PROFILE - docker run --rm -v .:/workspace $CACHE_ENVS -e SEMVER=${SEMVER} $PGRX_IMAGE ./tools/schema.sh --features pg$VERSION --profile $PROFILE + docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg$VERSION --release + docker run --rm -v .:/workspace $CACHE_ENVS -e SEMVER=${SEMVER} $PGRX_IMAGE ./tools/schema.sh --features pg$VERSION --release ./tools/package.sh ls ./build gh release upload --clobber $SEMVER ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb diff --git a/Cargo.lock b/Cargo.lock index 9439996e..5d17e1a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -22,6 +22,14 @@ dependencies = [ "memchr", ] +[[package]] +name = "algorithm" +version = "0.0.0" + +[[package]] +name = "always_equal" +version = "0.0.0" + [[package]] name = "annotate-snippets" version = "0.9.2" @@ -63,32 +71,6 @@ version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" -[[package]] -name = "base" -version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74#9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74" -dependencies = [ - "base_macros", - "cc", - "half 2.4.1", - "libc", - "rand", - "serde", - "thiserror 2.0.9", - "toml", - "validator", -] - -[[package]] -name = "base_macros" -version = "0.0.0" -source = "git+https://github.com/tensorchord/pgvecto.rs.git?rev=9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74#9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.91", -] - [[package]] name = "bindgen" version = "0.70.1" @@ -105,7 +87,7 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -178,9 +160,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.5" +version = "1.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c31a0499c1dc64f458ad13872de75c0eb7e3fdb0e67964610c914b034fc5956e" +checksum = "8d6dbb628b8f8555f86d0323c2eb39e3ec81901f4b83e091db8a6a76d316a333" dependencies = [ "shlex", ] @@ -282,7 +264,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -293,7 +275,7 @@ checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" dependencies = [ "darling_core", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -304,9 +286,13 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] +[[package]] +name = "distance" +version = "0.0.0" + [[package]] name = "either" version = "1.13.0" @@ -330,7 +316,7 @@ checksum = "f282cfdfe92516eb26c2af8589c274c7c17681f5ecc03c18255fe741c6aa64eb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -389,9 +375,9 @@ dependencies = [ [[package]] name = "glob" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" +checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2" [[package]] name = "half" @@ -402,12 +388,10 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403" [[package]] name = "half" version = "2.4.1" -source = "git+https://github.com/tensorchord/half-rs.git#5b7fedc636c0eb1624763a40840c5cbf54cffd02" +source = "git+https://github.com/tensorchord/half-rs.git#2d8a66092bee436aebb26ea7ac47d11150cda31d" dependencies = [ "cfg-if", "crunchy", - "rand", - "rand_distr", "rkyv", "serde", ] @@ -576,7 +560,7 @@ checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -734,7 +718,7 @@ checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -893,7 +877,7 @@ dependencies = [ "proc-macro2", "quote", "shlex", - "syn 2.0.91", + "syn 2.0.94", "walkdir", ] @@ -916,7 +900,7 @@ dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -963,7 +947,7 @@ dependencies = [ "petgraph", "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", "thiserror 1.0.69", "unescape", ] @@ -996,7 +980,7 @@ dependencies = [ "proc-macro-error-attr2", "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -1030,13 +1014,21 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.37" +version = "1.0.38" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" +checksum = "0e4dccaaaf89514f546c693ddc140f729f958c247918a13380cccc6078391acc" dependencies = [ "proc-macro2", ] +[[package]] +name = "rabitq" +version = "0.0.0" +dependencies = [ + "distance", + "simd", +] + [[package]] name = "radium" version = "0.7.0" @@ -1083,6 +1075,16 @@ dependencies = [ "rand", ] +[[package]] +name = "random_orthogonal_matrix" +version = "0.0.0" +dependencies = [ + "nalgebra", + "rand", + "rand_chacha", + "rand_distr", +] + [[package]] name = "rawpointer" version = "0.2.1" @@ -1241,9 +1243,9 @@ dependencies = [ [[package]] name = "serde" -version = "1.0.216" +version = "1.0.217" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b9781016e935a97e8beecf0c933758c97a5520d32930e460142b4cd80c6338e" +checksum = "02fc4265df13d6fa1d00ecff087228cc0a2b5f3c0e87e258d8b94a156e984c70" dependencies = [ "serde_derive", ] @@ -1260,13 +1262,13 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.216" +version = "1.0.217" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46f859dbbf73865c6627ed570e78961cd3ac92407a2d117204c49232485da55e" +checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -1309,6 +1311,26 @@ dependencies = [ "wide", ] +[[package]] +name = "simd" +version = "0.0.0" +dependencies = [ + "cc", + "half 2.4.1", + "rand", + "serde", + "simd_macros", +] + +[[package]] +name = "simd_macros" +version = "0.0.0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.94", +] + [[package]] name = "simdutf8" version = "0.1.5" @@ -1371,9 +1393,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.91" +version = "2.0.94" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d53cbcb5a243bd33b7858b1d7f4aca2153490815872d86d955d6ea29f743c035" +checksum = "987bc0be1cdea8b10216bd06e2ca407d40b9543468fafd3ddfb02f36e77f71f3" dependencies = [ "proc-macro2", "quote", @@ -1388,7 +1410,7 @@ checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -1423,7 +1445,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -1434,7 +1456,7 @@ checksum = "7b50fa271071aae2e6ee85f842e2e28ba8cd2c5fb67f11fcb1fd70b276f9e7d4" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -1591,28 +1613,41 @@ dependencies = [ "proc-macro-error2", "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] name = "vchord" version = "0.0.0" dependencies = [ - "base", + "algorithm", + "always_equal", + "distance", "half 2.4.1", "log", - "nalgebra", "paste", "pgrx", "pgrx-catalog", + "rabitq", "rand", - "rand_chacha", - "rand_distr", + "random_orthogonal_matrix", "rayon", "rkyv", "serde", + "simd", "toml", "validator", + "vector", +] + +[[package]] +name = "vector" +version = "0.0.0" +dependencies = [ + "distance", + "half 2.4.1", + "serde", + "simd", ] [[package]] @@ -1762,9 +1797,9 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "winnow" -version = "0.6.20" +version = "0.6.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36c1fec1a2bb5866f07c25f68c26e565c4c200aebb96d7e55710c19d3e8ac49b" +checksum = "e6f5bb5257f2407a5425c6e749bfd9692192a73e70a6060516ac04f889087d68" dependencies = [ "memchr", ] @@ -1819,7 +1854,7 @@ checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", "synstructure", ] @@ -1841,7 +1876,7 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] [[package]] @@ -1861,7 +1896,7 @@ checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", "synstructure", ] @@ -1884,5 +1919,5 @@ checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.91", + "syn 2.0.94", ] diff --git a/Cargo.toml b/Cargo.toml index 2697ad16..13272c2d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "vchord" -version = "0.0.0" -edition = "2021" +version.workspace = true +edition.workspace = true [lib] name = "vchord" @@ -20,24 +20,25 @@ pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] -base = { git = "https://github.com/tensorchord/pgvecto.rs.git", rev = "9d87afd75ca3dd6819da2a0a38d9fefdfb5b1c74" } - -# lock algebra version forever so that the QR decomposition never changes for same input -nalgebra = "=0.33.0" +algorithm = { path = "./crates/algorithm" } +always_equal = { path = "./crates/always_equal" } +distance = { path = "./crates/distance" } +rabitq = { path = "./crates/rabitq" } +random_orthogonal_matrix = { path = "./crates/random_orthogonal_matrix" } +simd = { path = "./crates/simd" } +vector = { path = "./crates/vector" } # lock rkyv version forever so that data is always compatible rkyv = { version = "=0.7.45", features = ["validation"] } -half = { version = "2.4.1", features = ["rkyv"] } +half.workspace = true log = "0.4.22" paste = "1" pgrx = { version = "=0.12.9", default-features = false, features = ["cshim"] } pgrx-catalog = "0.1.0" -rand = "0.8.5" -rand_chacha = "0.3.1" -rand_distr = "0.4.3" +rand.workspace = true rayon = "1.10.0" -serde = "1" +serde.workspace = true toml = "0.8.19" validator = { version = "0.19.0", features = ["derive"] } @@ -45,21 +46,30 @@ validator = { version = "0.19.0", features = ["derive"] } half = { git = "https://github.com/tensorchord/half-rs.git" } [lints] -rust.fuzzy_provenance_casts = "deny" -rust.unexpected_cfgs = { level = "warn", check-cfg = [ - 'cfg(feature, values("pg12"))', - 'cfg(pgrx_embed)', -] } +workspace = true + +[workspace] +resolver = "2" +members = ["crates/*"] + +[workspace.package] +version = "0.0.0" +edition = "2021" + +[workspace.dependencies] +half = { version = "2.4.1", features = ["rkyv", "serde"] } +rand = "0.8.5" +serde = "1" + +[workspace.lints] +clippy.identity_op = "allow" +clippy.int_plus_one = "allow" +clippy.needless_range_loop = "allow" +clippy.nonminimal_bool = "allow" rust.unsafe_op_in_unsafe_fn = "deny" rust.unused_lifetimes = "warn" rust.unused_qualifications = "warn" -[profile.opt] -debug-assertions = false -inherits = "dev" -opt-level = 3 -overflow-checks = false - [profile.release] codegen-units = 1 debug = true diff --git a/build.rs b/build.rs new file mode 100644 index 00000000..1fadc37c --- /dev/null +++ b/build.rs @@ -0,0 +1,4 @@ +fn main() { + println!(r#"cargo::rustc-check-cfg=cfg(pgrx_embed)"#); + println!(r#"cargo::rustc-check-cfg=cfg(feature, values("pg12"))"#); +} diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml new file mode 100644 index 00000000..56da6cec --- /dev/null +++ b/crates/algorithm/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "algorithm" +version.workspace = true +edition.workspace = true + +[lints] +workspace = true diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs new file mode 100644 index 00000000..8bee0167 --- /dev/null +++ b/crates/algorithm/src/lib.rs @@ -0,0 +1,41 @@ +use std::ops::{Deref, DerefMut}; + +#[repr(C, align(8))] +pub struct Opaque { + pub next: u32, + pub skip: u32, +} + +#[allow(clippy::len_without_is_empty)] +pub trait Page: Sized { + fn get_opaque(&self) -> &Opaque; + fn get_opaque_mut(&mut self) -> &mut Opaque; + fn len(&self) -> u16; + fn get(&self, i: u16) -> Option<&[u8]>; + fn get_mut(&mut self, i: u16) -> Option<&mut [u8]>; + fn alloc(&mut self, data: &[u8]) -> Option; + fn free(&mut self, i: u16); + fn reconstruct(&mut self, removes: &[u16]); + fn freespace(&self) -> u16; +} + +pub trait PageGuard { + fn id(&self) -> u32; +} + +pub trait RelationRead { + type Page: Page; + type ReadGuard<'a>: PageGuard + Deref + where + Self: 'a; + fn read(&self, id: u32) -> Self::ReadGuard<'_>; +} + +pub trait RelationWrite: RelationRead { + type WriteGuard<'a>: PageGuard + DerefMut + where + Self: 'a; + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_>; + fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_>; + fn search(&self, freespace: usize) -> Option>; +} diff --git a/crates/always_equal/Cargo.toml b/crates/always_equal/Cargo.toml new file mode 100644 index 00000000..923fa8e1 --- /dev/null +++ b/crates/always_equal/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "always_equal" +version.workspace = true +edition.workspace = true + +[lints] +workspace = true diff --git a/crates/always_equal/src/lib.rs b/crates/always_equal/src/lib.rs new file mode 100644 index 00000000..8cb2cc1c --- /dev/null +++ b/crates/always_equal/src/lib.rs @@ -0,0 +1,30 @@ +use std::cmp::Ordering; +use std::hash::Hash; + +#[derive(Debug, Clone, Copy, Default)] +#[repr(transparent)] +pub struct AlwaysEqual(pub T); + +impl PartialEq for AlwaysEqual { + fn eq(&self, _: &Self) -> bool { + true + } +} + +impl Eq for AlwaysEqual {} + +impl PartialOrd for AlwaysEqual { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for AlwaysEqual { + fn cmp(&self, _: &Self) -> Ordering { + Ordering::Equal + } +} + +impl Hash for AlwaysEqual { + fn hash(&self, _: &mut H) {} +} diff --git a/crates/distance/Cargo.toml b/crates/distance/Cargo.toml new file mode 100644 index 00000000..6c578193 --- /dev/null +++ b/crates/distance/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "distance" +version.workspace = true +edition.workspace = true + +[lints] +workspace = true diff --git a/crates/distance/src/lib.rs b/crates/distance/src/lib.rs new file mode 100644 index 00000000..5f21aace --- /dev/null +++ b/crates/distance/src/lib.rs @@ -0,0 +1,79 @@ +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +#[repr(transparent)] +pub struct Distance(i32); + +impl Distance { + pub const ZERO: Self = Distance::from_f32(0.0f32); + pub const INFINITY: Self = Distance::from_f32(f32::INFINITY); + pub const NEG_INFINITY: Self = Distance::from_f32(f32::NEG_INFINITY); + + #[inline(always)] + pub const fn from_f32(value: f32) -> Self { + let bits = value.to_bits() as i32; + let mask = ((bits >> 31) as u32) >> 1; + let res = bits ^ (mask as i32); + Self(res) + } + + #[inline(always)] + pub const fn to_f32(self) -> f32 { + let bits = self.0; + let mask = ((bits >> 31) as u32) >> 1; + let res = bits ^ (mask as i32); + f32::from_bits(res as u32) + } + + #[inline(always)] + pub const fn to_i32(self) -> i32 { + self.0 + } +} + +impl From for Distance { + #[inline(always)] + fn from(value: f32) -> Self { + Distance::from_f32(value) + } +} + +impl From for f32 { + #[inline(always)] + fn from(value: Distance) -> Self { + Distance::to_f32(value) + } +} + +#[test] +fn distance_conversions() { + assert_eq!(Distance::from(0.0f32), Distance::ZERO); + assert_eq!(Distance::from(f32::INFINITY), Distance::INFINITY); + assert_eq!(Distance::from(f32::NEG_INFINITY), Distance::NEG_INFINITY); + for i in -100..100 { + let val = (i as f32) * 0.1; + assert_eq!(f32::from(Distance::from(val)).to_bits(), val.to_bits()); + } + assert_eq!( + f32::from(Distance::from(0.0f32)).to_bits(), + 0.0f32.to_bits() + ); + assert_eq!( + f32::from(Distance::from(-0.0f32)).to_bits(), + (-0.0f32).to_bits() + ); + assert_eq!( + f32::from(Distance::from(f32::NAN)).to_bits(), + f32::NAN.to_bits() + ); + assert_eq!( + f32::from(Distance::from(-f32::NAN)).to_bits(), + (-f32::NAN).to_bits() + ); + assert_eq!( + f32::from(Distance::from(f32::INFINITY)).to_bits(), + f32::INFINITY.to_bits() + ); + assert_eq!( + f32::from(Distance::from(-f32::INFINITY)).to_bits(), + (-f32::INFINITY).to_bits() + ); +} diff --git a/crates/rabitq/Cargo.toml b/crates/rabitq/Cargo.toml new file mode 100644 index 00000000..5ca300c4 --- /dev/null +++ b/crates/rabitq/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "rabitq" +version.workspace = true +edition.workspace = true + +[dependencies] +distance = { path = "../distance" } +simd = { path = "../simd" } + +[lints] +workspace = true diff --git a/crates/rabitq/src/binary.rs b/crates/rabitq/src/binary.rs new file mode 100644 index 00000000..8916ffd9 --- /dev/null +++ b/crates/rabitq/src/binary.rs @@ -0,0 +1,123 @@ +use distance::Distance; +use simd::Floating; + +#[derive(Debug, Clone)] +pub struct Code { + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub signs: Vec, +} + +impl Code { + pub fn t(&self) -> Vec { + use crate::utils::InfiniteByteChunks; + let mut result = Vec::new(); + for x in InfiniteByteChunks::<_, 64>::new(self.signs.iter().copied()) + .take(self.signs.len().div_ceil(64)) + { + let mut r = 0_u64; + for i in 0..64 { + r |= (x[i] as u64) << i; + } + result.push(r); + } + result + } +} + +pub fn code(dims: u32, vector: &[f32]) -> Code { + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x_2 = f32::reduce_sum_of_x2(vector); + let dis_u = sum_of_x_2.sqrt(); + let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); + let x_x0 = dis_u / x0; + let fac_norm = (dims as f32).sqrt(); + let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); + let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); + let factor_ip = -2.0f32 / fac_norm * x_x0; + let cnt_pos = vector + .iter() + .map(|x| x.is_sign_positive() as i32) + .sum::(); + let cnt_neg = vector + .iter() + .map(|x| x.is_sign_negative() as i32) + .sum::(); + let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; + let mut signs = Vec::new(); + for i in 0..dims { + signs.push(vector[i as usize].is_sign_positive() as u8); + } + Code { + dis_u_2: sum_of_x_2, + factor_ppc, + factor_ip, + factor_err, + signs, + } +} + +pub type Lut = (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)); + +pub fn preprocess(vector: &[f32]) -> Lut { + let dis_v_2 = f32::reduce_sum_of_x2(vector); + let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); + let qvector_sum = if vector.len() <= 4369 { + simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 + } else { + simd::u8::reduce_sum_of_x(&qvector) as f32 + }; + (dis_v_2, b, k, qvector_sum, binarize(&qvector)) +} + +pub fn process_lowerbound_l2( + _: u32, + lut: &Lut, + (dis_u_2, factor_ppc, factor_ip, factor_err, t): (f32, f32, f32, f32, &[u64]), + epsilon: f32, +) -> Distance { + let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let value = asymmetric_binary_dot_product(t, s) as u16; + let rough = + dis_u_2 + dis_v_2 + b * factor_ppc + ((2.0 * value as f32) - qvector_sum) * factor_ip * k; + let err = factor_err * dis_v_2.sqrt(); + Distance::from_f32(rough - epsilon * err) +} + +pub fn process_lowerbound_dot( + _: u32, + lut: &Lut, + (_, factor_ppc, factor_ip, factor_err, t): (f32, f32, f32, f32, &[u64]), + epsilon: f32, +) -> Distance { + let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let value = asymmetric_binary_dot_product(t, s) as u16; + let rough = 0.5 * b * factor_ppc + 0.5 * ((2.0 * value as f32) - qvector_sum) * factor_ip * k; + let err = 0.5 * factor_err * dis_v_2.sqrt(); + Distance::from_f32(rough - epsilon * err) +} + +fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { + let n = vector.len(); + let mut t0 = vec![0u64; n.div_ceil(64)]; + let mut t1 = vec![0u64; n.div_ceil(64)]; + let mut t2 = vec![0u64; n.div_ceil(64)]; + let mut t3 = vec![0u64; n.div_ceil(64)]; + for i in 0..n { + t0[i / 64] |= (((vector[i] >> 0) & 1) as u64) << (i % 64); + t1[i / 64] |= (((vector[i] >> 1) & 1) as u64) << (i % 64); + t2[i / 64] |= (((vector[i] >> 2) & 1) as u64) << (i % 64); + t3[i / 64] |= (((vector[i] >> 3) & 1) as u64) << (i % 64); + } + (t0, t1, t2, t3) +} + +fn asymmetric_binary_dot_product(x: &[u64], y: &(Vec, Vec, Vec, Vec)) -> u32 { + let t0 = simd::bit::sum_of_and(x, &y.0); + let t1 = simd::bit::sum_of_and(x, &y.1); + let t2 = simd::bit::sum_of_and(x, &y.2); + let t3 = simd::bit::sum_of_and(x, &y.3); + (t0 << 0) + (t1 << 1) + (t2 << 2) + (t3 << 3) +} diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs new file mode 100644 index 00000000..0a8dab70 --- /dev/null +++ b/crates/rabitq/src/block.rs @@ -0,0 +1,172 @@ +use distance::Distance; +use simd::Floating; + +#[derive(Debug, Clone)] +pub struct Code { + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub signs: Vec, +} + +pub fn code(dims: u32, vector: &[f32]) -> Code { + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x_2 = f32::reduce_sum_of_x2(vector); + let dis_u = sum_of_x_2.sqrt(); + let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); + let x_x0 = dis_u / x0; + let fac_norm = (dims as f32).sqrt(); + let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); + let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); + let factor_ip = -2.0f32 / fac_norm * x_x0; + let cnt_pos = vector + .iter() + .map(|x| x.is_sign_positive() as i32) + .sum::(); + let cnt_neg = vector + .iter() + .map(|x| x.is_sign_negative() as i32) + .sum::(); + let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; + let mut signs = Vec::new(); + for i in 0..dims { + signs.push(vector[i as usize].is_sign_positive() as u8); + } + Code { + dis_u_2: sum_of_x_2, + factor_ppc, + factor_ip, + factor_err, + signs, + } +} + +pub fn dummy_code(dims: u32) -> Code { + Code { + dis_u_2: 0.0, + factor_ppc: 0.0, + factor_ip: 0.0, + factor_err: 0.0, + signs: vec![0; dims as _], + } +} + +pub struct PackedCodes { + pub dis_u_2: [f32; 32], + pub factor_ppc: [f32; 32], + pub factor_ip: [f32; 32], + pub factor_err: [f32; 32], + pub t: Vec, +} + +pub fn pack_codes(dims: u32, codes: [Code; 32]) -> PackedCodes { + use crate::utils::InfiniteByteChunks; + PackedCodes { + dis_u_2: std::array::from_fn(|i| codes[i].dis_u_2), + factor_ppc: std::array::from_fn(|i| codes[i].factor_ppc), + factor_ip: std::array::from_fn(|i| codes[i].factor_ip), + factor_err: std::array::from_fn(|i| codes[i].factor_err), + t: { + let signs = codes.map(|code| { + InfiniteByteChunks::new(code.signs.into_iter()) + .map(|[b0, b1, b2, b3]| b0 | b1 << 1 | b2 << 2 | b3 << 3) + .take(dims.div_ceil(4) as usize) + .collect::>() + }); + simd::fast_scan::pack(dims.div_ceil(4), signs).collect() + }, + } +} + +pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { + let dis_v_2 = f32::reduce_sum_of_x2(vector); + let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); + let qvector_sum = if vector.len() <= 4369 { + simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 + } else { + simd::u8::reduce_sum_of_x(&qvector) as f32 + }; + (dis_v_2, b, k, qvector_sum, compress(qvector)) +} + +pub fn fscan_process_lowerbound_l2( + dims: u32, + lut: &(f32, f32, f32, f32, Vec), + (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( + &[f32; 32], + &[f32; 32], + &[f32; 32], + &[f32; 32], + &[u8], + ), + epsilon: f32, +) -> [Distance; 32] { + let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let r = simd::fast_scan::fast_scan(dims.div_ceil(4), t, s); + std::array::from_fn(|i| { + let rough = dis_u_2[i] + + dis_v_2 + + b * factor_ppc[i] + + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; + let err = factor_err[i] * dis_v_2.sqrt(); + Distance::from_f32(rough - epsilon * err) + }) +} + +pub fn fscan_process_lowerbound_dot( + dims: u32, + lut: &(f32, f32, f32, f32, Vec), + (_, factor_ppc, factor_ip, factor_err, t): ( + &[f32; 32], + &[f32; 32], + &[f32; 32], + &[f32; 32], + &[u8], + ), + epsilon: f32, +) -> [Distance; 32] { + let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let r = simd::fast_scan::fast_scan(dims.div_ceil(4), t, s); + std::array::from_fn(|i| { + let rough = + 0.5 * b * factor_ppc[i] + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; + let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); + Distance::from_f32(rough - epsilon * err) + }) +} + +fn compress(mut qvector: Vec) -> Vec { + let dims = qvector.len() as u32; + let width = dims.div_ceil(4); + qvector.resize(qvector.len().next_multiple_of(4), 0); + let mut t = vec![0u8; width as usize * 16]; + for i in 0..width as usize { + unsafe { + // this hint is used to skip bound checks + std::hint::assert_unchecked(4 * i + 3 < qvector.len()); + std::hint::assert_unchecked(16 * i + 15 < t.len()); + } + let t0 = qvector[4 * i + 0]; + let t1 = qvector[4 * i + 1]; + let t2 = qvector[4 * i + 2]; + let t3 = qvector[4 * i + 3]; + t[16 * i + 0b0000] = 0; + t[16 * i + 0b0001] = t0; + t[16 * i + 0b0010] = t1; + t[16 * i + 0b0011] = t1 + t0; + t[16 * i + 0b0100] = t2; + t[16 * i + 0b0101] = t2 + t0; + t[16 * i + 0b0110] = t2 + t1; + t[16 * i + 0b0111] = t2 + t1 + t0; + t[16 * i + 0b1000] = t3; + t[16 * i + 0b1001] = t3 + t0; + t[16 * i + 0b1010] = t3 + t1; + t[16 * i + 0b1011] = t3 + t1 + t0; + t[16 * i + 0b1100] = t3 + t2; + t[16 * i + 0b1101] = t3 + t2 + t0; + t[16 * i + 0b1110] = t3 + t2 + t1; + t[16 * i + 0b1111] = t3 + t2 + t1 + t0; + } + t +} diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs new file mode 100644 index 00000000..0796b7ac --- /dev/null +++ b/crates/rabitq/src/lib.rs @@ -0,0 +1,5 @@ +#![allow(clippy::type_complexity)] + +pub mod binary; +pub mod block; +mod utils; diff --git a/src/utils/infinite_byte_chunks.rs b/crates/rabitq/src/utils.rs similarity index 100% rename from src/utils/infinite_byte_chunks.rs rename to crates/rabitq/src/utils.rs diff --git a/crates/random_orthogonal_matrix/Cargo.toml b/crates/random_orthogonal_matrix/Cargo.toml new file mode 100644 index 00000000..16c3a39a --- /dev/null +++ b/crates/random_orthogonal_matrix/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "random_orthogonal_matrix" +version.workspace = true +edition.workspace = true + +[dependencies] +# lock algebra version forever so that the QR decomposition never changes for same input +nalgebra = "=0.33.0" + +rand.workspace = true +rand_chacha = "0.3.1" +rand_distr = "0.4.3" + +[lints] +workspace = true diff --git a/crates/random_orthogonal_matrix/src/lib.rs b/crates/random_orthogonal_matrix/src/lib.rs new file mode 100644 index 00000000..e2711504 --- /dev/null +++ b/crates/random_orthogonal_matrix/src/lib.rs @@ -0,0 +1,55 @@ +use nalgebra::DMatrix; + +#[ignore] +#[test] +fn check_full_rank_matrix() { + let parallelism = std::thread::available_parallelism().unwrap().get(); + std::thread::scope(|scope| { + let mut threads = vec![]; + for remainder in 0..parallelism { + threads.push(scope.spawn(move || { + for n in (0..=60000).filter(|x| x % parallelism == remainder) { + let matrix = random_full_rank_matrix(n); + assert!(matrix.is_invertible()); + } + })); + } + for thread in threads { + thread.join().unwrap(); + } + }); +} + +fn random_full_rank_matrix(n: usize) -> DMatrix { + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha12Rng; + use rand_distr::StandardNormal; + let mut rng = ChaCha12Rng::from_seed([7; 32]); + DMatrix::from_fn(n, n, |_, _| rng.sample(StandardNormal)) +} + +#[test] +fn check_random_orthogonal_matrix() { + assert_eq!(random_orthogonal_matrix(2), vec![ + vec![-0.5424608, -0.8400813], + vec![0.8400813, -0.54246056] + ]); + assert_eq!(random_orthogonal_matrix(3), vec![ + vec![-0.5309615, -0.69094884, -0.49058124], + vec![0.8222731, -0.56002235, -0.10120347], + vec![0.20481002, 0.45712686, -0.86549866] + ]); +} + +pub fn random_orthogonal_matrix(n: usize) -> Vec> { + use nalgebra::QR; + let matrix = random_full_rank_matrix(n); + // QR decomposition is unique if the matrix is full rank + let qr = QR::new(matrix); + let q = qr.q(); + let mut projection = Vec::new(); + for row in q.row_iter() { + projection.push(row.iter().copied().collect::>()); + } + projection +} diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml new file mode 100644 index 00000000..0905fb14 --- /dev/null +++ b/crates/simd/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "simd" +version.workspace = true +edition.workspace = true + +[dependencies] +half.workspace = true +serde.workspace = true +simd_macros = { path = "../simd_macros" } + +[dev-dependencies] +rand.workspace = true + +[build-dependencies] +cc = "1.2.6" + +[lints] +workspace = true diff --git a/crates/simd/build.rs b/crates/simd/build.rs new file mode 100644 index 00000000..22ebf935 --- /dev/null +++ b/crates/simd/build.rs @@ -0,0 +1,12 @@ +fn main() { + println!("cargo::rerun-if-changed=cshim.c"); + cc::Build::new() + .compiler("clang") + .file("cshim.c") + .opt_level(3) + .flag("-fassociative-math") + .flag("-ffp-contract=fast") + .flag("-freciprocal-math") + .flag("-fno-signed-zeros") + .compile("base_cshim"); +} diff --git a/crates/simd/cshim.c b/crates/simd/cshim.c new file mode 100644 index 00000000..1374bbfe --- /dev/null +++ b/crates/simd/cshim.c @@ -0,0 +1,223 @@ +#if !(__clang_major__ >= 16) +#error "clang version must be >= 16" +#endif + +#include +#include + +#ifdef __aarch64__ + +#include +#include + +__attribute__((target("v8.3a,fp16"))) float +fp16_reduce_sum_of_xy_v8_3a_fp16_unroll(__fp16 *__restrict a, + __fp16 *__restrict b, size_t n) { + float16x8_t xy_0 = vdupq_n_f16(0.0); + float16x8_t xy_1 = vdupq_n_f16(0.0); + float16x8_t xy_2 = vdupq_n_f16(0.0); + float16x8_t xy_3 = vdupq_n_f16(0.0); + while (n >= 32) { + float16x8_t x_0 = vld1q_f16(a + 0); + float16x8_t x_1 = vld1q_f16(a + 8); + float16x8_t x_2 = vld1q_f16(a + 16); + float16x8_t x_3 = vld1q_f16(a + 24); + float16x8_t y_0 = vld1q_f16(b + 0); + float16x8_t y_1 = vld1q_f16(b + 8); + float16x8_t y_2 = vld1q_f16(b + 16); + float16x8_t y_3 = vld1q_f16(b + 24); + a += 32; + b += 32; + n -= 32; + xy_0 = vfmaq_f16(xy_0, x_0, y_0); + xy_1 = vfmaq_f16(xy_1, x_1, y_1); + xy_2 = vfmaq_f16(xy_2, x_2, y_2); + xy_3 = vfmaq_f16(xy_3, x_3, y_3); + } + if (n > 0) { + __fp16 A[32] = {}; + __fp16 B[32] = {}; + for (size_t i = 0; i < n; i += 1) { + A[i] = a[i]; + B[i] = b[i]; + } + float16x8_t x_0 = vld1q_f16(A + 0); + float16x8_t x_1 = vld1q_f16(A + 8); + float16x8_t x_2 = vld1q_f16(A + 16); + float16x8_t x_3 = vld1q_f16(A + 24); + float16x8_t y_0 = vld1q_f16(B + 0); + float16x8_t y_1 = vld1q_f16(B + 8); + float16x8_t y_2 = vld1q_f16(B + 16); + float16x8_t y_3 = vld1q_f16(B + 24); + xy_0 = vfmaq_f16(xy_0, x_0, y_0); + xy_1 = vfmaq_f16(xy_1, x_1, y_1); + xy_2 = vfmaq_f16(xy_2, x_2, y_2); + xy_3 = vfmaq_f16(xy_3, x_3, y_3); + } + float16x8_t xy = vaddq_f16(vaddq_f16(xy_0, xy_1), vaddq_f16(xy_2, xy_3)); + return vgetq_lane_f16(xy, 0) + vgetq_lane_f16(xy, 1) + vgetq_lane_f16(xy, 2) + + vgetq_lane_f16(xy, 3) + vgetq_lane_f16(xy, 4) + vgetq_lane_f16(xy, 5) + + vgetq_lane_f16(xy, 6) + vgetq_lane_f16(xy, 7); +} + +__attribute__((target("v8.3a,sve"))) float +fp16_reduce_sum_of_xy_v8_3a_sve(__fp16 *__restrict a, __fp16 *__restrict b, + size_t n) { + svfloat16_t xy = svdup_f16(0.0); + for (size_t i = 0; i < n; i += svcnth()) { + svbool_t mask = svwhilelt_b16(i, n); + svfloat16_t x = svld1_f16(mask, a + i); + svfloat16_t y = svld1_f16(mask, b + i); + xy = svmla_f16_x(mask, xy, x, y); + } + return svaddv_f16(svptrue_b16(), xy); +} + +__attribute__((target("v8.3a,fp16"))) float +fp16_reduce_sum_of_d2_v8_3a_fp16_unroll(__fp16 *__restrict a, + __fp16 *__restrict b, size_t n) { + float16x8_t d2_0 = vdupq_n_f16(0.0); + float16x8_t d2_1 = vdupq_n_f16(0.0); + float16x8_t d2_2 = vdupq_n_f16(0.0); + float16x8_t d2_3 = vdupq_n_f16(0.0); + while (n >= 32) { + float16x8_t x_0 = vld1q_f16(a + 0); + float16x8_t x_1 = vld1q_f16(a + 8); + float16x8_t x_2 = vld1q_f16(a + 16); + float16x8_t x_3 = vld1q_f16(a + 24); + float16x8_t y_0 = vld1q_f16(b + 0); + float16x8_t y_1 = vld1q_f16(b + 8); + float16x8_t y_2 = vld1q_f16(b + 16); + float16x8_t y_3 = vld1q_f16(b + 24); + a += 32; + b += 32; + n -= 32; + float16x8_t d_0 = vsubq_f16(x_0, y_0); + float16x8_t d_1 = vsubq_f16(x_1, y_1); + float16x8_t d_2 = vsubq_f16(x_2, y_2); + float16x8_t d_3 = vsubq_f16(x_3, y_3); + d2_0 = vfmaq_f16(d2_0, d_0, d_0); + d2_1 = vfmaq_f16(d2_1, d_1, d_1); + d2_2 = vfmaq_f16(d2_2, d_2, d_2); + d2_3 = vfmaq_f16(d2_3, d_3, d_3); + } + if (n > 0) { + __fp16 A[32] = {}; + __fp16 B[32] = {}; + for (size_t i = 0; i < n; i += 1) { + A[i] = a[i]; + B[i] = b[i]; + } + float16x8_t x_0 = vld1q_f16(A + 0); + float16x8_t x_1 = vld1q_f16(A + 8); + float16x8_t x_2 = vld1q_f16(A + 16); + float16x8_t x_3 = vld1q_f16(A + 24); + float16x8_t y_0 = vld1q_f16(B + 0); + float16x8_t y_1 = vld1q_f16(B + 8); + float16x8_t y_2 = vld1q_f16(B + 16); + float16x8_t y_3 = vld1q_f16(B + 24); + float16x8_t d_0 = vsubq_f16(x_0, y_0); + float16x8_t d_1 = vsubq_f16(x_1, y_1); + float16x8_t d_2 = vsubq_f16(x_2, y_2); + float16x8_t d_3 = vsubq_f16(x_3, y_3); + d2_0 = vfmaq_f16(d2_0, d_0, d_0); + d2_1 = vfmaq_f16(d2_1, d_1, d_1); + d2_2 = vfmaq_f16(d2_2, d_2, d_2); + d2_3 = vfmaq_f16(d2_3, d_3, d_3); + } + float16x8_t d2 = vaddq_f16(vaddq_f16(d2_0, d2_1), vaddq_f16(d2_2, d2_3)); + return vgetq_lane_f16(d2, 0) + vgetq_lane_f16(d2, 1) + vgetq_lane_f16(d2, 2) + + vgetq_lane_f16(d2, 3) + vgetq_lane_f16(d2, 4) + vgetq_lane_f16(d2, 5) + + vgetq_lane_f16(d2, 6) + vgetq_lane_f16(d2, 7); +} + +__attribute__((target("v8.3a,sve"))) float +fp16_reduce_sum_of_d2_v8_3a_sve(__fp16 *__restrict a, __fp16 *__restrict b, + size_t n) { + svfloat16_t d2 = svdup_f16(0.0); + for (size_t i = 0; i < n; i += svcnth()) { + svbool_t mask = svwhilelt_b16(i, n); + svfloat16_t x = svld1_f16(mask, a + i); + svfloat16_t y = svld1_f16(mask, b + i); + svfloat16_t d = svsub_f16_x(mask, x, y); + d2 = svmla_f16_x(mask, d2, d, d); + } + return svaddv_f16(svptrue_b16(), d2); +} + +__attribute__((target("v8.3a,sve"))) float +fp32_reduce_sum_of_x_v8_3a_sve(float *__restrict this, size_t n) { + svfloat32_t sum = svdup_f32(0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32(i, n); + svfloat32_t x = svld1_f32(mask, this + i); + sum = svadd_f32_x(mask, sum, x); + } + return svaddv_f32(svptrue_b32(), sum); +} + +__attribute__((target("v8.3a,sve"))) float +fp32_reduce_sum_of_abs_x_v8_3a_sve(float *__restrict this, size_t n) { + svfloat32_t sum = svdup_f32(0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32(i, n); + svfloat32_t x = svld1_f32(mask, this + i); + sum = svadd_f32_x(mask, sum, svabs_f32_x(mask, x)); + } + return svaddv_f32(svptrue_b32(), sum); +} + +__attribute__((target("v8.3a,sve"))) float +fp32_reduce_sum_of_x2_v8_3a_sve(float *__restrict this, size_t n) { + svfloat32_t sum = svdup_f32(0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32(i, n); + svfloat32_t x = svld1_f32(mask, this + i); + sum = svmla_f32_x(mask, sum, x, x); + } + return svaddv_f32(svptrue_b32(), sum); +} + +__attribute__((target("v8.3a,sve"))) void +fp32_reduce_min_max_of_x_v8_3a_sve(float *__restrict this, size_t n, + float *out_min, float *out_max) { + svfloat32_t min = svdup_f32(1.0 / 0.0); + svfloat32_t max = svdup_f32(-1.0 / 0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32(i, n); + svfloat32_t x = svld1_f32(mask, this + i); + min = svmin_f32_x(mask, min, x); + max = svmax_f32_x(mask, max, x); + } + *out_min = svminv_f32(svptrue_b32(), min); + *out_max = svmaxv_f32(svptrue_b32(), max); +} + +__attribute__((target("v8.3a,sve"))) float +fp32_reduce_sum_of_xy_v8_3a_sve(float *__restrict lhs, float *__restrict rhs, + size_t n) { + svfloat32_t sum = svdup_f32(0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32(i, n); + svfloat32_t x = svld1_f32(mask, lhs + i); + svfloat32_t y = svld1_f32(mask, rhs + i); + sum = svmla_f32_x(mask, sum, x, y); + } + return svaddv_f32(svptrue_b32(), sum); +} + +__attribute__((target("v8.3a,sve"))) float +fp32_reduce_sum_of_d2_v8_3a_sve(float *__restrict lhs, float *__restrict rhs, + size_t n) { + svfloat32_t sum = svdup_f32(0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32(i, n); + svfloat32_t x = svld1_f32(mask, lhs + i); + svfloat32_t y = svld1_f32(mask, rhs + i); + svfloat32_t d = svsub_f32_x(mask, x, y); + sum = svmla_f32_x(mask, sum, d, d); + } + return svaddv_f32(svptrue_b32(), sum); +} + +#endif diff --git a/crates/simd/src/aligned.rs b/crates/simd/src/aligned.rs new file mode 100644 index 00000000..18db9d85 --- /dev/null +++ b/crates/simd/src/aligned.rs @@ -0,0 +1,4 @@ +#[allow(dead_code)] +#[derive(Debug, Clone, Copy)] +#[repr(C, align(16))] +pub struct Aligned16(pub T); diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs new file mode 100644 index 00000000..dc417bda --- /dev/null +++ b/crates/simd/src/bit.rs @@ -0,0 +1,399 @@ +#[inline(always)] +pub fn sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { + sum_of_and::sum_of_and(lhs, rhs) +} + +mod sum_of_and { + // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn sum_of_and_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut and = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = _mm512_loadu_si512(a.cast()); + let y = _mm512_loadu_si512(b.cast()); + a = a.add(8); + b = b.add(8); + n -= 8; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = _mm512_maskz_loadu_epi64(mask, a.cast()); + let y = _mm512_maskz_loadu_epi64(mask, b.cast()); + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + } + _mm512_reduce_add_epi64(and) as u32 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn sum_of_and_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { sum_of_and_v4_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = sum_of_and_fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut and = 0; + for i in 0..n { + and += (lhs[i] & rhs[i]).count_ones(); + } + and + } +} + +#[inline(always)] +pub fn sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { + sum_of_or::sum_of_or(lhs, rhs) +} + +mod sum_of_or { + // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn sum_of_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut or = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = _mm512_loadu_si512(a.cast()); + let y = _mm512_loadu_si512(b.cast()); + a = a.add(8); + b = b.add(8); + n -= 8; + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = _mm512_maskz_loadu_epi64(mask, a.cast()); + let y = _mm512_maskz_loadu_epi64(mask, b.cast()); + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + } + _mm512_reduce_add_epi64(or) as u32 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn sum_of_or_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { sum_of_or_v4_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = sum_of_or_fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut or = 0; + for i in 0..n { + or += (lhs[i] | rhs[i]).count_ones(); + } + or + } +} + +#[inline(always)] +pub fn sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { + sum_of_xor::sum_of_xor(lhs, rhs) +} + +mod sum_of_xor { + // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn sum_of_xor_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut xor = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = _mm512_loadu_si512(a.cast()); + let y = _mm512_loadu_si512(b.cast()); + a = a.add(8); + b = b.add(8); + n -= 8; + xor = _mm512_add_epi64(xor, _mm512_popcnt_epi64(_mm512_xor_si512(x, y))); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = _mm512_maskz_loadu_epi64(mask, a.cast()); + let y = _mm512_maskz_loadu_epi64(mask, b.cast()); + xor = _mm512_add_epi64(xor, _mm512_popcnt_epi64(_mm512_xor_si512(x, y))); + } + _mm512_reduce_add_epi64(xor) as u32 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn sum_of_xor_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { sum_of_xor_v4_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = sum_of_xor_fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut xor = 0; + for i in 0..n { + xor += (lhs[i] ^ rhs[i]).count_ones(); + } + xor + } +} + +#[inline(always)] +pub fn sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + sum_of_and_or::sum_of_and_or(lhs, rhs) +} + +mod sum_of_and_or { + // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn sum_of_and_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut and = _mm512_setzero_si512(); + let mut or = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = _mm512_loadu_si512(a.cast()); + let y = _mm512_loadu_si512(b.cast()); + a = a.add(8); + b = b.add(8); + n -= 8; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = _mm512_maskz_loadu_epi64(mask, a.cast()); + let y = _mm512_maskz_loadu_epi64(mask, b.cast()); + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + } + ( + _mm512_reduce_add_epi64(and) as u32, + _mm512_reduce_add_epi64(or) as u32, + ) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn sum_of_xor_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { sum_of_and_or_v4_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = sum_of_and_or_fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut and = 0; + let mut or = 0; + for i in 0..n { + and += (lhs[i] & rhs[i]).count_ones(); + or += (lhs[i] | rhs[i]).count_ones(); + } + (and, or) + } +} + +#[inline(always)] +pub fn sum_of_x(this: &[u64]) -> u32 { + sum_of_x::sum_of_x(this) +} + +mod sum_of_x { + // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn sum_of_x_v4_avx512vpopcntdq(this: &[u64]) -> u32 { + unsafe { + use std::arch::x86_64::*; + let mut and = _mm512_setzero_si512(); + let mut a = this.as_ptr(); + let mut n = this.len(); + while n >= 8 { + let x = _mm512_loadu_si512(a.cast()); + a = a.add(8); + n -= 8; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(x)); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = _mm512_maskz_loadu_epi64(mask, a.cast()); + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(x)); + } + _mm512_reduce_add_epi64(and) as u32 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn sum_of_x_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let this = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { sum_of_x_v4_avx512vpopcntdq(&this) }; + let fallback = sum_of_x_fallback(&this); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn sum_of_x(this: &[u64]) -> u32 { + let n = this.len(); + let mut and = 0; + for i in 0..n { + and += this[i].count_ones(); + } + and + } +} + +#[inline(always)] +pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { + vector_and::vector_and(lhs, rhs) +} + +mod vector_and { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] & rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +#[inline(always)] +pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { + vector_or::vector_or(lhs, rhs) +} + +mod vector_or { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] | rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +#[inline(always)] +pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { + vector_xor::vector_xor(lhs, rhs) +} + +mod vector_xor { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] ^ rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs new file mode 100644 index 00000000..520b6a1d --- /dev/null +++ b/crates/simd/src/emulate.rs @@ -0,0 +1,208 @@ +// VP2INTERSECT emulation. +// Díez-Cañas, G. (2021). Faster-Than-Native Alternatives for x86 VP2INTERSECT +// Instructions. arXiv preprint arXiv:2112.06342. +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v4")] +pub fn emulate_mm512_2intersect_epi32( + a: std::arch::x86_64::__m512i, + b: std::arch::x86_64::__m512i, +) -> (std::arch::x86_64::__mmask16, std::arch::x86_64::__mmask16) { + unsafe { + use std::arch::x86_64::*; + + let a1 = _mm512_alignr_epi32(a, a, 4); + let a2 = _mm512_alignr_epi32(a, a, 8); + let a3 = _mm512_alignr_epi32(a, a, 12); + let b1 = _mm512_shuffle_epi32(b, _MM_PERM_ADCB); + let b2 = _mm512_shuffle_epi32(b, _MM_PERM_BADC); + let b3 = _mm512_shuffle_epi32(b, _MM_PERM_CBAD); + let m00 = _mm512_cmpeq_epi32_mask(a, b); + let m01 = _mm512_cmpeq_epi32_mask(a, b1); + let m02 = _mm512_cmpeq_epi32_mask(a, b2); + let m03 = _mm512_cmpeq_epi32_mask(a, b3); + let m10 = _mm512_cmpeq_epi32_mask(a1, b); + let m11 = _mm512_cmpeq_epi32_mask(a1, b1); + let m12 = _mm512_cmpeq_epi32_mask(a1, b2); + let m13 = _mm512_cmpeq_epi32_mask(a1, b3); + let m20 = _mm512_cmpeq_epi32_mask(a2, b); + let m21 = _mm512_cmpeq_epi32_mask(a2, b1); + let m22 = _mm512_cmpeq_epi32_mask(a2, b2); + let m23 = _mm512_cmpeq_epi32_mask(a2, b3); + let m30 = _mm512_cmpeq_epi32_mask(a3, b); + let m31 = _mm512_cmpeq_epi32_mask(a3, b1); + let m32 = _mm512_cmpeq_epi32_mask(a3, b2); + let m33 = _mm512_cmpeq_epi32_mask(a3, b3); + + let m0 = m00 | m10 | m20 | m30; + let m1 = m01 | m11 | m21 | m31; + let m2 = m02 | m12 | m22 | m32; + let m3 = m03 | m13 | m23 | m33; + + let res_a = m00 + | m01 + | m02 + | m03 + | (m10 | m11 | m12 | m13).rotate_left(4) + | (m20 | m21 | m22 | m23).rotate_left(8) + | (m30 | m31 | m32 | m33).rotate_right(4); + + let res_b = m0 + | ((0x7777 & m1) << 1) + | ((m1 >> 3) & 0x1111) + | ((0x3333 & m2) << 2) + | ((m2 >> 2) & 0x3333) + | ((0x1111 & m3) << 3) + | ((m3 >> 1) & 0x7777); + (res_a, res_b) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v3")] +pub fn emulate_mm256_reduce_add_ps(mut x: std::arch::x86_64::__m256) -> f32 { + unsafe { + use std::arch::x86_64::*; + x = _mm256_add_ps(x, _mm256_permute2f128_ps(x, x, 1)); + x = _mm256_hadd_ps(x, x); + x = _mm256_hadd_ps(x, x); + _mm256_cvtss_f32(x) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v2")] +pub fn emulate_mm_reduce_add_ps(mut x: std::arch::x86_64::__m128) -> f32 { + unsafe { + use std::arch::x86_64::*; + x = _mm_hadd_ps(x, x); + x = _mm_hadd_ps(x, x); + _mm_cvtss_f32(x) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v4")] +pub fn emulate_mm512_reduce_add_epi16(x: std::arch::x86_64::__m512i) -> i16 { + unsafe { + use std::arch::x86_64::*; + _mm256_reduce_add_epi16(_mm512_castsi512_si256(x)) + + _mm256_reduce_add_epi16(_mm512_extracti32x8_epi32(x, 1)) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v3")] +pub fn emulate_mm256_reduce_add_epi16(mut x: std::arch::x86_64::__m256i) -> i16 { + unsafe { + use std::arch::x86_64::*; + x = _mm256_add_epi16(x, _mm256_permute2f128_si256(x, x, 1)); + x = _mm256_hadd_epi16(x, x); + x = _mm256_hadd_epi16(x, x); + let x = _mm256_cvtsi256_si32(x); + (x as i16) + ((x >> 16) as i16) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v2")] +pub fn emulate_mm_reduce_add_epi16(mut x: std::arch::x86_64::__m128i) -> i16 { + unsafe { + use std::arch::x86_64::*; + x = _mm_hadd_epi16(x, x); + x = _mm_hadd_epi16(x, x); + let x = _mm_cvtsi128_si32(x); + (x as i16) + ((x >> 16) as i16) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v3")] +pub fn emulate_mm256_reduce_add_epi32(mut x: std::arch::x86_64::__m256i) -> i32 { + unsafe { + use std::arch::x86_64::*; + x = _mm256_add_epi32(x, _mm256_permute2f128_si256(x, x, 1)); + x = _mm256_hadd_epi32(x, x); + x = _mm256_hadd_epi32(x, x); + _mm256_cvtsi256_si32(x) + } +} + +#[expect(dead_code)] +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v2")] +pub fn emulate_mm_reduce_add_epi32(mut x: std::arch::x86_64::__m128i) -> i32 { + unsafe { + use std::arch::x86_64::*; + x = _mm_hadd_epi32(x, x); + x = _mm_hadd_epi32(x, x); + _mm_cvtsi128_si32(x) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v3")] +pub fn emulate_mm256_reduce_min_ps(x: std::arch::x86_64::__m256) -> f32 { + use crate::aligned::Aligned16; + unsafe { + use std::arch::x86_64::*; + let lo = _mm256_castps256_ps128(x); + let hi = _mm256_extractf128_ps(x, 1); + let min = _mm_min_ps(lo, hi); + let mut x = Aligned16([0.0f32; 4]); + _mm_store_ps(x.0.as_mut_ptr(), min); + f32::min(f32::min(x.0[0], x.0[1]), f32::min(x.0[2], x.0[3])) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v2")] +pub fn emulate_mm_reduce_min_ps(x: std::arch::x86_64::__m128) -> f32 { + use crate::aligned::Aligned16; + unsafe { + use std::arch::x86_64::*; + let min = x; + let mut x = Aligned16([0.0f32; 4]); + _mm_store_ps(x.0.as_mut_ptr(), min); + f32::min(f32::min(x.0[0], x.0[1]), f32::min(x.0[2], x.0[3])) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v3")] +pub fn emulate_mm256_reduce_max_ps(x: std::arch::x86_64::__m256) -> f32 { + use crate::aligned::Aligned16; + unsafe { + use std::arch::x86_64::*; + let lo = _mm256_castps256_ps128(x); + let hi = _mm256_extractf128_ps(x, 1); + let max = _mm_max_ps(lo, hi); + let mut x = Aligned16([0.0f32; 4]); + _mm_store_ps(x.0.as_mut_ptr(), max); + f32::max(f32::max(x.0[0], x.0[1]), f32::max(x.0[2], x.0[3])) + } +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v2")] +pub fn emulate_mm_reduce_max_ps(x: std::arch::x86_64::__m128) -> f32 { + use crate::aligned::Aligned16; + unsafe { + use std::arch::x86_64::*; + let max = x; + let mut x = Aligned16([0.0f32; 4]); + _mm_store_ps(x.0.as_mut_ptr(), max); + f32::max(f32::max(x.0[0], x.0[1]), f32::max(x.0[2], x.0[3])) + } +} diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs new file mode 100644 index 00000000..53bddb78 --- /dev/null +++ b/crates/simd/src/f16.rs @@ -0,0 +1,1185 @@ +use crate::{Floating, f32}; +use half::f16; + +impl Floating for f16 { + #[inline(always)] + fn zero() -> Self { + f16::ZERO + } + + #[inline(always)] + fn infinity() -> Self { + f16::INFINITY + } + + #[inline(always)] + fn mask(self, m: bool) -> Self { + f16::from_bits(self.to_bits() & (m as u16).wrapping_neg()) + } + + #[inline(always)] + fn scalar_neg(this: Self) -> Self { + -this + } + + #[inline(always)] + fn scalar_add(lhs: Self, rhs: Self) -> Self { + lhs + rhs + } + + #[inline(always)] + fn scalar_sub(lhs: Self, rhs: Self) -> Self { + lhs - rhs + } + + #[inline(always)] + fn scalar_mul(lhs: Self, rhs: Self) -> Self { + lhs * rhs + } + + #[inline(always)] + fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { + reduce_or_of_is_zero_x::reduce_or_of_is_zero_x(this) + } + + #[inline(always)] + fn reduce_sum_of_x(this: &[f16]) -> f32 { + reduce_sum_of_x::reduce_sum_of_x(this) + } + + #[inline(always)] + fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { + reduce_sum_of_abs_x::reduce_sum_of_abs_x(this) + } + + #[inline(always)] + fn reduce_sum_of_x2(this: &[f16]) -> f32 { + reduce_sum_of_x2::reduce_sum_of_x2(this) + } + + #[inline(always)] + fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { + reduce_min_max_of_x::reduce_min_max_of_x(this) + } + + #[inline(always)] + fn reduce_sum_of_xy(lhs: &[Self], rhs: &[Self]) -> f32 { + reduce_sum_of_xy::reduce_sum_of_xy(lhs, rhs) + } + + #[inline(always)] + fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { + reduce_sum_of_d2::reduce_sum_of_d2(lhs, rhs) + } + + #[inline(always)] + fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { + reduce_sum_of_xy_sparse::reduce_sum_of_xy_sparse(lidx, lval, ridx, rval) + } + + #[inline(always)] + fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { + reduce_sum_of_d2_sparse::reduce_sum_of_d2_sparse(lidx, lval, ridx, rval) + } + + #[inline(always)] + fn vector_from_f32(this: &[f32]) -> Vec { + vector_from_f32::vector_from_f32(this) + } + + #[inline(always)] + fn vector_to_f32(this: &[Self]) -> Vec { + vector_to_f32::vector_to_f32(this) + } + + #[inline(always)] + fn vector_add(lhs: &[Self], rhs: &[Self]) -> Vec { + vector_add::vector_add(lhs, rhs) + } + + #[inline(always)] + fn vector_add_inplace(lhs: &mut [Self], rhs: &[Self]) { + vector_add_inplace::vector_add_inplace(lhs, rhs) + } + + #[inline(always)] + fn vector_sub(lhs: &[Self], rhs: &[Self]) -> Vec { + vector_sub::vector_sub(lhs, rhs) + } + + #[inline(always)] + fn vector_mul(lhs: &[Self], rhs: &[Self]) -> Vec { + vector_mul::vector_mul(lhs, rhs) + } + + #[inline(always)] + fn vector_mul_scalar(lhs: &[Self], rhs: f32) -> Vec { + vector_mul_scalar::vector_mul_scalar(lhs, rhs) + } + + #[inline(always)] + fn vector_mul_scalar_inplace(lhs: &mut [Self], rhs: f32) { + vector_mul_scalar_inplace::vector_mul_scalar_inplace(lhs, rhs) + } + + #[inline(always)] + fn vector_to_f32_borrowed(this: &[Self]) -> impl AsRef<[f32]> { + Self::vector_to_f32(this) + } +} + +mod reduce_or_of_is_zero_x { + // FIXME: add manually-implemented SIMD version + + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { + for &x in this { + if x == f16::ZERO { + return true; + } + } + false + } +} + +mod reduce_sum_of_x { + // FIXME: add manually-implemented SIMD version + + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_x(this: &[f16]) -> f32 { + let n = this.len(); + let mut x = 0.0f32; + for i in 0..n { + x += this[i].to_f32(); + } + x + } +} + +mod reduce_sum_of_abs_x { + // FIXME: add manually-implemented SIMD version + + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { + let n = this.len(); + let mut x = 0.0f32; + for i in 0..n { + x += this[i].to_f32().abs(); + } + x + } +} + +mod reduce_sum_of_x2 { + // FIXME: add manually-implemented SIMD version + + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_x2(this: &[f16]) -> f32 { + let n = this.len(); + let mut x2 = 0.0f32; + for i in 0..n { + x2 += this[i].to_f32() * this[i].to_f32(); + } + x2 + } +} + +mod reduce_min_max_of_x { + // FIXME: add manually-implemented SIMD version + + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { + let mut min = f32::INFINITY; + let mut max = f32::NEG_INFINITY; + let n = this.len(); + for i in 0..n { + min = min.min(this[i].to_f32()); + max = max.max(this[i].to_f32()); + } + (min, max) + } +} + +mod reduce_sum_of_xy { + use half::f16; + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + pub fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm512_setzero_ph(); + while n >= 32 { + let x = _mm512_loadu_ph(a.cast()); + let y = _mm512_loadu_ph(b.cast()); + a = a.add(32); + b = b.add(32); + n -= 32; + xy = _mm512_fmadd_ph(x, y, xy); + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())); + let y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())); + xy = _mm512_fmadd_ph(x, y, xy); + } + _mm512_reduce_add_ph(xy) as f32 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v4_avx512fp16_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4_avx512fp16(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + pub fn reduce_sum_of_xy_v4(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm512_setzero_ps(); + while n >= 16 { + let x = _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())); + let y = _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())); + a = a.add(16); + b = b.add(16); + n -= 16; + xy = _mm512_fmadd_ps(x, y, xy); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())); + let y = _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())); + xy = _mm512_fmadd_ps(x, y, xy); + } + _mm512_reduce_add_ps(xy) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v4_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let specialized = unsafe { reduce_sum_of_xy_v4(&lhs, &rhs) }; + let fallback = reduce_sum_of_xy_fallback(&lhs, &rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + pub fn reduce_sum_of_xy_v3(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::emulate_mm256_reduce_add_ps; + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm256_setzero_ps(); + while n >= 8 { + let x = _mm256_cvtph_ps(_mm_loadu_si128(a.cast())); + let y = _mm256_cvtph_ps(_mm_loadu_si128(b.cast())); + a = a.add(8); + b = b.add(8); + n -= 8; + xy = _mm256_fmadd_ps(x, y, xy); + } + let mut xy = emulate_mm256_reduce_add_ps(xy); + while n > 0 { + let x = a.read().to_f32(); + let y = b.read().to_f32(); + a = a.add(1); + b = b.add(1); + n -= 1; + xy += x * y; + } + xy + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v3_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "f16c")] + #[target_feature(enable = "fma")] + pub fn reduce_sum_of_xy_v2_f16c_fma(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::emulate_mm_reduce_add_ps; + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm_setzero_ps(); + while n >= 4 { + let x = _mm_cvtph_ps(_mm_loadu_si128(a.cast())); + let y = _mm_cvtph_ps(_mm_loadu_si128(b.cast())); + a = a.add(4); + b = b.add(4); + n -= 4; + xy = _mm_fmadd_ps(x, y, xy); + } + let mut xy = emulate_mm_reduce_add_ps(xy); + while n > 0 { + let x = a.read().to_f32(); + let y = b.read().to_f32(); + a = a.add(1); + b = b.add(1); + n -= 1; + xy += x * y; + } + xy + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v2_f16c_fma_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v2") + || !crate::is_feature_detected!("f16c") + || !crate::is_feature_detected!("fma") + { + println!("test {} ... skipped (v2:f16c:fma)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v2_f16c_fma(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "fp16")] + pub fn reduce_sum_of_xy_v8_3a_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + extern "C" { + fn fp16_reduce_sum_of_xy_v8_3a_fp16_unroll( + a: *const (), + b: *const (), + n: usize, + ) -> f32; + } + fp16_reduce_sum_of_xy_v8_3a_fp16_unroll( + lhs.as_ptr().cast(), + rhs.as_ptr().cast(), + lhs.len(), + ) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v8_3a_fp16_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (v8.3a:fp16)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v8_3a_fp16(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + // temporarily disables this for uncertain precision + #[expect(dead_code)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "sve")] + pub fn reduce_sum_of_xy_v8_3a_sve(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + extern "C" { + fn fp16_reduce_sum_of_xy_v8_3a_sve(a: *const (), b: *const (), n: usize) -> f32; + } + fp16_reduce_sum_of_xy_v8_3a_sve(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v8_3a_sve_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { + println!("test {} ... skipped (v8.3a:sve)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v8_3a_sve(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"v2:f16c:fma", @"v8.3a:fp16")] + pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let mut xy = 0.0f32; + for i in 0..n { + xy += lhs[i].to_f32() * rhs[i].to_f32(); + } + xy + } +} + +mod reduce_sum_of_d2 { + use half::f16; + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + pub fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len() as u32; + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm512_setzero_ph(); + while n >= 32 { + let x = _mm512_loadu_ph(a.cast()); + let y = _mm512_loadu_ph(b.cast()); + a = a.add(32); + b = b.add(32); + n -= 32; + let d = _mm512_sub_ph(x, y); + d2 = _mm512_fmadd_ph(d, d, d2); + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n); + let x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())); + let y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())); + let d = _mm512_sub_ph(x, y); + d2 = _mm512_fmadd_ph(d, d, d2); + } + _mm512_reduce_add_ph(d2) as f32 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_v4_avx512fp16_test() { + use rand::Rng; + const EPSILON: f32 = 6.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v4_avx512fp16(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + pub fn reduce_sum_of_d2_v4(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len() as u32; + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm512_setzero_ps(); + while n >= 16 { + let x = _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())); + let y = _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())); + a = a.add(16); + b = b.add(16); + n -= 16; + let d = _mm512_sub_ps(x, y); + d2 = _mm512_fmadd_ps(d, d, d2); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n) as u16; + let x = _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())); + let y = _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())); + let d = _mm512_sub_ps(x, y); + d2 = _mm512_fmadd_ps(d, d, d2); + } + _mm512_reduce_add_ps(d2) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_v4_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v4(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + pub fn reduce_sum_of_d2_v3(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::emulate_mm256_reduce_add_ps; + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len() as u32; + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm256_setzero_ps(); + while n >= 8 { + let x = _mm256_cvtph_ps(_mm_loadu_si128(a.cast())); + let y = _mm256_cvtph_ps(_mm_loadu_si128(b.cast())); + a = a.add(8); + b = b.add(8); + n -= 8; + let d = _mm256_sub_ps(x, y); + d2 = _mm256_fmadd_ps(d, d, d2); + } + let mut d2 = emulate_mm256_reduce_add_ps(d2); + while n > 0 { + let x = a.read().to_f32(); + let y = b.read().to_f32(); + a = a.add(1); + b = b.add(1); + n -= 1; + let d = x - y; + d2 += d * d; + } + d2 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_v3_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v3(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "f16c")] + #[target_feature(enable = "fma")] + pub fn reduce_sum_of_d2_v2_f16c_fma(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::emulate_mm_reduce_add_ps; + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len() as u32; + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm_setzero_ps(); + while n >= 4 { + let x = _mm_cvtph_ps(_mm_loadu_si128(a.cast())); + let y = _mm_cvtph_ps(_mm_loadu_si128(b.cast())); + a = a.add(4); + b = b.add(4); + n -= 4; + let d = _mm_sub_ps(x, y); + d2 = _mm_fmadd_ps(d, d, d2); + } + let mut d2 = emulate_mm_reduce_add_ps(d2); + while n > 0 { + let x = a.read().to_f32(); + let y = b.read().to_f32(); + a = a.add(1); + b = b.add(1); + n -= 1; + let d = x - y; + d2 += d * d; + } + d2 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_v2_f16c_fma_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v2") + || !crate::is_feature_detected!("f16c") + || !crate::is_feature_detected!("fma") + { + println!("test {} ... skipped (v2:f16c:fma)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v2_f16c_fma(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "fp16")] + pub fn reduce_sum_of_d2_v8_3a_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + extern "C" { + fn fp16_reduce_sum_of_d2_v8_3a_fp16_unroll( + a: *const (), + b: *const (), + n: usize, + ) -> f32; + } + fp16_reduce_sum_of_d2_v8_3a_fp16_unroll( + lhs.as_ptr().cast(), + rhs.as_ptr().cast(), + lhs.len(), + ) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_v8_3a_fp16_test() { + use rand::Rng; + const EPSILON: f32 = 6.0; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (v8.3a:fp16)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v8_3a_fp16(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + // temporarily disables this for uncertain precision + #[expect(dead_code)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "sve")] + pub fn reduce_sum_of_d2_v8_3a_sve(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + extern "C" { + fn fp16_reduce_sum_of_d2_v8_3a_sve(a: *const (), b: *const (), n: usize) -> f32; + } + fp16_reduce_sum_of_d2_v8_3a_sve(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_v8_3a_sve_test() { + use rand::Rng; + const EPSILON: f32 = 6.0; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { + println!("test {} ... skipped (v8.3a:sve)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v8_3a_sve(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"v2:f16c:fma", @"v8.3a:fp16")] + pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let mut d2 = 0.0; + for i in 0..n { + let d = lhs[i].to_f32() - rhs[i].to_f32(); + d2 += d * d; + } + d2 + } +} + +mod reduce_sum_of_xy_sparse { + // There is no manually-implemented SIMD version. + // Add it if `svecf16` is supported. + + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { + use std::cmp::Ordering; + assert_eq!(lidx.len(), lval.len()); + assert_eq!(ridx.len(), rval.len()); + let (mut lp, ln) = (0, lidx.len()); + let (mut rp, rn) = (0, ridx.len()); + let mut xy = 0.0f32; + while lp < ln && rp < rn { + match Ord::cmp(&lidx[lp], &ridx[rp]) { + Ordering::Equal => { + xy += lval[lp].to_f32() * rval[rp].to_f32(); + lp += 1; + rp += 1; + } + Ordering::Less => { + lp += 1; + } + Ordering::Greater => { + rp += 1; + } + } + } + xy + } +} + +mod reduce_sum_of_d2_sparse { + // There is no manually-implemented SIMD version. + // Add it if `svecf16` is supported. + + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { + use std::cmp::Ordering; + assert_eq!(lidx.len(), lval.len()); + assert_eq!(ridx.len(), rval.len()); + let (mut lp, ln) = (0, lidx.len()); + let (mut rp, rn) = (0, ridx.len()); + let mut d2 = 0.0f32; + while lp < ln && rp < rn { + match Ord::cmp(&lidx[lp], &ridx[rp]) { + Ordering::Equal => { + let d = lval[lp].to_f32() - rval[rp].to_f32(); + d2 += d * d; + lp += 1; + rp += 1; + } + Ordering::Less => { + d2 += lval[lp].to_f32() * lval[lp].to_f32(); + lp += 1; + } + Ordering::Greater => { + d2 += rval[rp].to_f32() * rval[rp].to_f32(); + rp += 1; + } + } + } + for i in lp..ln { + d2 += lval[i].to_f32() * lval[i].to_f32(); + } + for i in rp..rn { + d2 += rval[i].to_f32() * rval[i].to_f32(); + } + d2 + } +} + +mod vector_add { + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_add(lhs: &[f16], rhs: &[f16]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] + rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_add_inplace { + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_add_inplace(lhs: &mut [f16], rhs: &[f16]) { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + for i in 0..n { + lhs[i] += rhs[i]; + } + } +} + +mod vector_sub { + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_sub(lhs: &[f16], rhs: &[f16]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] - rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul { + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_mul(lhs: &[f16], rhs: &[f16]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] * rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul_scalar { + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_mul_scalar(lhs: &[f16], rhs: f32) -> Vec { + let rhs = f16::from_f32(rhs); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] * rhs); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul_scalar_inplace { + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_mul_scalar_inplace(lhs: &mut [f16], rhs: f32) { + let rhs = f16::from_f32(rhs); + let n = lhs.len(); + for i in 0..n { + lhs[i] *= rhs; + } + } +} + +mod vector_abs_inplace { + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_abs_inplace(this: &mut [f16]) { + let n = this.len(); + for i in 0..n { + this[i] = f16::from_f32(this[i].to_f32().abs()); + } + } +} + +mod vector_from_f32 { + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_from_f32(this: &[f32]) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(f16::from_f32(this[i])); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_to_f32 { + use half::f16; + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_to_f32(this: &[f16]) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(this[i].to_f32()); + } + } + unsafe { + r.set_len(n); + } + r + } +} diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs new file mode 100644 index 00000000..555da5c7 --- /dev/null +++ b/crates/simd/src/f32.rs @@ -0,0 +1,2277 @@ +use crate::Floating; + +impl Floating for f32 { + #[inline(always)] + fn zero() -> Self { + 0.0f32 + } + + #[inline(always)] + fn infinity() -> Self { + f32::INFINITY + } + + #[inline(always)] + fn mask(self, m: bool) -> Self { + f32::from_bits(self.to_bits() & (m as u32).wrapping_neg()) + } + + #[inline(always)] + fn scalar_neg(this: Self) -> Self { + -this + } + + #[inline(always)] + fn scalar_add(lhs: Self, rhs: Self) -> Self { + lhs + rhs + } + + #[inline(always)] + fn scalar_sub(lhs: Self, rhs: Self) -> Self { + lhs - rhs + } + + #[inline(always)] + fn scalar_mul(lhs: Self, rhs: Self) -> Self { + lhs * rhs + } + + #[inline(always)] + fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { + reduce_or_of_is_zero_x::reduce_or_of_is_zero_x(this) + } + + #[inline(always)] + fn reduce_sum_of_x(this: &[f32]) -> f32 { + reduce_sum_of_x::reduce_sum_of_x(this) + } + + #[inline(always)] + fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { + reduce_sum_of_abs_x::reduce_sum_of_abs_x(this) + } + + #[inline(always)] + fn reduce_sum_of_x2(this: &[f32]) -> f32 { + reduce_sum_of_x2::reduce_sum_of_x2(this) + } + + #[inline(always)] + fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { + reduce_min_max_of_x::reduce_min_max_of_x(this) + } + + #[inline(always)] + fn reduce_sum_of_xy(lhs: &[Self], rhs: &[Self]) -> f32 { + reduce_sum_of_xy::reduce_sum_of_xy(lhs, rhs) + } + + #[inline(always)] + fn reduce_sum_of_d2(lhs: &[Self], rhs: &[Self]) -> f32 { + reduce_sum_of_d2::reduce_sum_of_d2(lhs, rhs) + } + + #[inline(always)] + fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { + reduce_sum_of_xy_sparse::reduce_sum_of_xy_sparse(lidx, lval, ridx, rval) + } + + #[inline(always)] + fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { + reduce_sum_of_d2_sparse::reduce_sum_of_d2_sparse(lidx, lval, ridx, rval) + } + + fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { + vector_add::vector_add(lhs, rhs) + } + + fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { + vector_add_inplace::vector_add_inplace(lhs, rhs) + } + + fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { + vector_sub::vector_sub(lhs, rhs) + } + + fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { + vector_mul::vector_mul(lhs, rhs) + } + + fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { + vector_mul_scalar::vector_mul_scalar(lhs, rhs) + } + + fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { + vector_mul_scalar_inplace::vector_mul_scalar_inplace(lhs, rhs); + } + + fn vector_from_f32(this: &[f32]) -> Vec { + this.to_vec() + } + + fn vector_to_f32(this: &[f32]) -> Vec { + this.to_vec() + } + + fn vector_to_f32_borrowed(this: &[Self]) -> impl AsRef<[f32]> { + this + } +} + +mod reduce_or_of_is_zero_x { + // FIXME: add manually-implemented SIMD version + + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { + for &x in this { + if x == 0.0f32 { + return true; + } + } + false + } +} + +mod reduce_sum_of_x { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[f32]) -> f32 { + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = _mm512_loadu_ps(a); + a = a.add(16); + n -= 16; + sum = _mm512_add_ps(x, sum); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = _mm512_maskz_loadu_ps(mask, a); + sum = _mm512_add_ps(x, sum); + } + _mm512_reduce_add_ps(sum) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_v4_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v4(this) }; + let fallback = reduce_sum_of_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_v3(this: &[f32]) -> f32 { + use crate::emulate::emulate_mm256_reduce_add_ps; + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = _mm256_loadu_ps(a); + a = a.add(8); + n -= 8; + sum = _mm256_add_ps(x, sum); + } + if n >= 4 { + let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); + a = a.add(4); + n -= 4; + sum = _mm256_add_ps(x, sum); + } + let mut sum = emulate_mm256_reduce_add_ps(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + sum += x; + } + sum + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v3_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v3(this) }; + let fallback = reduce_sum_of_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_x_v2(this: &[f32]) -> f32 { + use crate::emulate::emulate_mm_reduce_add_ps; + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_ps(); + while n >= 4 { + let x = _mm_loadu_ps(a); + a = a.add(4); + n -= 4; + sum = _mm_add_ps(x, sum); + } + let mut sum = emulate_mm_reduce_add_ps(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + sum += x; + } + sum + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v2_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v2(this) }; + let fallback = reduce_sum_of_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + fn reduce_sum_of_x_v8_3a(this: &[f32]) -> f32 { + unsafe { + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = vld1q_f32(a); + a = a.add(4); + n -= 4; + sum = vaddq_f32(x, sum); + } + let mut sum = vaddvq_f32(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + sum += x; + } + sum + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_v8_3a_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v8.3a") { + println!("test {} ... skipped (v8_3a)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v8_3a(this) }; + let fallback = reduce_sum_of_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "sve")] + fn reduce_sum_of_x_v8_3a_sve(this: &[f32]) -> f32 { + unsafe { + extern "C" { + fn fp32_reduce_sum_of_x_v8_3a_sve(this: *const f32, n: usize) -> f32; + } + fp32_reduce_sum_of_x_v8_3a_sve(this.as_ptr(), this.len()) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_v8_3a_sve_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { + println!("test {} ... skipped (v8_3a:sve)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v8_3a_sve(this) }; + let fallback = reduce_sum_of_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a:sve", @"v8.3a")] + pub fn reduce_sum_of_x(this: &[f32]) -> f32 { + let n = this.len(); + let mut sum = 0.0f32; + for i in 0..n { + sum += this[i]; + } + sum + } +} + +mod reduce_sum_of_abs_x { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_abs_x_v4(this: &[f32]) -> f32 { + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = _mm512_loadu_ps(a); + let abs_x = _mm512_abs_ps(x); + a = a.add(16); + n -= 16; + sum = _mm512_add_ps(abs_x, sum); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = _mm512_maskz_loadu_ps(mask, a); + let abs_x = _mm512_abs_ps(x); + sum = _mm512_add_ps(abs_x, sum); + } + _mm512_reduce_add_ps(sum) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_abs_x_v4_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v4(this) }; + let fallback = reduce_sum_of_abs_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_abs_x_v3(this: &[f32]) -> f32 { + use crate::emulate::emulate_mm256_reduce_add_ps; + unsafe { + use std::arch::x86_64::*; + let abs = _mm256_castsi256_ps(_mm256_srli_epi32(_mm256_set1_epi32(-1), 1)); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = _mm256_loadu_ps(a); + let abs_x = _mm256_and_ps(abs, x); + a = a.add(8); + n -= 8; + sum = _mm256_add_ps(abs_x, sum); + } + if n >= 4 { + let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); + let abs_x = _mm256_and_ps(abs, x); + a = a.add(4); + n -= 4; + sum = _mm256_add_ps(abs_x, sum); + } + let mut sum = emulate_mm256_reduce_add_ps(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let abs_x = x.abs(); + a = a.add(1); + n -= 1; + sum += abs_x; + } + sum + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_abs_x_v3_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v3(this) }; + let fallback = reduce_sum_of_abs_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_abs_x_v2(this: &[f32]) -> f32 { + use crate::emulate::emulate_mm_reduce_add_ps; + unsafe { + use std::arch::x86_64::*; + let abs = _mm_castsi128_ps(_mm_srli_epi32(_mm_set1_epi32(-1), 1)); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_ps(); + while n >= 4 { + let x = _mm_loadu_ps(a); + let abs_x = _mm_and_ps(abs, x); + a = a.add(4); + n -= 4; + sum = _mm_add_ps(abs_x, sum); + } + let mut sum = emulate_mm_reduce_add_ps(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let abs_x = x.abs(); + a = a.add(1); + n -= 1; + sum += abs_x; + } + sum + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_abs_x_v2_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v2(this) }; + let fallback = reduce_sum_of_abs_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + fn reduce_sum_of_abs_x_v8_3a(this: &[f32]) -> f32 { + unsafe { + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = vld1q_f32(a); + let abs_x = vabsq_f32(x); + a = a.add(4); + n -= 4; + sum = vaddq_f32(abs_x, sum); + } + let mut sum = vaddvq_f32(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let abs_x = x.abs(); + a = a.add(1); + n -= 1; + sum += abs_x; + } + sum + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_abs_x_v8_3a_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v8.3a") { + println!("test {} ... skipped (v8.3a)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v8_3a(this) }; + let fallback = reduce_sum_of_abs_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "sve")] + fn reduce_sum_of_abs_x_v8_3a_sve(this: &[f32]) -> f32 { + unsafe { + extern "C" { + fn fp32_reduce_sum_of_abs_x_v8_3a_sve(this: *const f32, n: usize) -> f32; + } + fp32_reduce_sum_of_abs_x_v8_3a_sve(this.as_ptr(), this.len()) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_abs_x_v8_3a_sve_test() { + use rand::Rng; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { + println!("test {} ... skipped (v8.3a:sve)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v8_3a_sve(this) }; + let fallback = reduce_sum_of_abs_x_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a:sve", @"v8.3a")] + pub fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { + let n = this.len(); + let mut sum = 0.0f32; + for i in 0..n { + sum += this[i].abs(); + } + sum + } +} + +mod reduce_sum_of_x2 { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x2_v4(this: &[f32]) -> f32 { + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut x2 = _mm512_setzero_ps(); + while n >= 16 { + let x = _mm512_loadu_ps(a); + a = a.add(16); + n -= 16; + x2 = _mm512_fmadd_ps(x, x, x2); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = _mm512_maskz_loadu_ps(mask, a); + x2 = _mm512_fmadd_ps(x, x, x2); + } + _mm512_reduce_add_ps(x2) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_x2_v4_test() { + use rand::Rng; + const EPSILON: f32 = 0.006; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v4(this) }; + let fallback = reduce_sum_of_x2_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x2_v3(this: &[f32]) -> f32 { + use crate::emulate::emulate_mm256_reduce_add_ps; + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut x2 = _mm256_setzero_ps(); + while n >= 8 { + let x = _mm256_loadu_ps(a); + a = a.add(8); + n -= 8; + x2 = _mm256_fmadd_ps(x, x, x2); + } + if n >= 4 { + let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); + a = a.add(4); + n -= 4; + x2 = _mm256_fmadd_ps(x, x, x2); + } + let mut x2 = emulate_mm256_reduce_add_ps(x2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + x2 += x * x; + } + x2 + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x2_v3_test() { + use rand::Rng; + const EPSILON: f32 = 0.006; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v3(this) }; + let fallback = reduce_sum_of_x2_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "fma")] + fn reduce_sum_of_x2_v2_fma(this: &[f32]) -> f32 { + use crate::emulate::emulate_mm_reduce_add_ps; + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut x2 = _mm_setzero_ps(); + while n >= 4 { + let x = _mm_loadu_ps(a); + a = a.add(4); + n -= 4; + x2 = _mm_fmadd_ps(x, x, x2); + } + let mut x2 = emulate_mm_reduce_add_ps(x2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + x2 += x * x; + } + x2 + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x2_v2_fma_test() { + use rand::Rng; + const EPSILON: f32 = 0.006; + if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { + println!("test {} ... skipped (v2:fma)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v2_fma(this) }; + let fallback = reduce_sum_of_x2_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + fn reduce_sum_of_x2_v8_3a(this: &[f32]) -> f32 { + unsafe { + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut x2 = vdupq_n_f32(0.0); + while n >= 4 { + let x = vld1q_f32(a); + a = a.add(4); + n -= 4; + x2 = vfmaq_f32(x2, x, x); + } + let mut x2 = vaddvq_f32(x2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + x2 += x * x; + } + x2 + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_x2_v8_3a_test() { + use rand::Rng; + const EPSILON: f32 = 0.006; + if !crate::is_cpu_detected!("v8.3a") { + println!("test {} ... skipped (v8.3a)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v8_3a(this) }; + let fallback = reduce_sum_of_x2_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "sve")] + fn reduce_sum_of_x2_v8_3a_sve(this: &[f32]) -> f32 { + unsafe { + extern "C" { + fn fp32_reduce_sum_of_x2_v8_3a_sve(this: *const f32, n: usize) -> f32; + } + fp32_reduce_sum_of_x2_v8_3a_sve(this.as_ptr(), this.len()) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_x2_v8_3a_sve_test() { + use rand::Rng; + const EPSILON: f32 = 0.006; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { + println!("test {} ... skipped (v8.3a:sve)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v8_3a_sve(this) }; + let fallback = reduce_sum_of_x2_fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"v8.3a:sve", @"v8.3a")] + pub fn reduce_sum_of_x2(this: &[f32]) -> f32 { + let n = this.len(); + let mut x2 = 0.0f32; + for i in 0..n { + x2 += this[i] * this[i]; + } + x2 + } +} + +mod reduce_min_max_of_x { + // Semanctics of `f32::min` is different from `_mm256_min_ps`, + // which may lead to issues... + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_min_max_of_x_v4(this: &[f32]) -> (f32, f32) { + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm512_set1_ps(f32::INFINITY); + let mut max = _mm512_set1_ps(f32::NEG_INFINITY); + while n >= 16 { + let x = _mm512_loadu_ps(a); + a = a.add(16); + n -= 16; + min = _mm512_min_ps(x, min); + max = _mm512_max_ps(x, max); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = _mm512_maskz_loadu_ps(mask, a); + min = _mm512_mask_min_ps(min, mask, x, min); + max = _mm512_mask_max_ps(max, mask, x, max); + } + let min = _mm512_reduce_min_ps(min); + let max = _mm512_reduce_max_ps(max); + (min, max) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_min_max_of_x_v4_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let x = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v4(x) }; + let fallback = reduce_min_max_of_x_fallback(x); + assert_eq!(specialized.0, fallback.0); + assert_eq!(specialized.1, fallback.1); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_min_max_of_x_v3(this: &[f32]) -> (f32, f32) { + use crate::emulate::emulate_mm256_reduce_max_ps; + use crate::emulate::emulate_mm256_reduce_min_ps; + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm256_set1_ps(f32::INFINITY); + let mut max = _mm256_set1_ps(f32::NEG_INFINITY); + while n >= 8 { + let x = _mm256_loadu_ps(a); + a = a.add(8); + n -= 8; + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + } + let mut min = emulate_mm256_reduce_min_ps(min); + let mut max = emulate_mm256_reduce_max_ps(max); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + min = x.min(min); + max = x.max(max); + } + (min, max) + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_min_max_of_x_v3_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let x = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v3(x) }; + let fallback = reduce_min_max_of_x_fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_min_max_of_x_v2(this: &[f32]) -> (f32, f32) { + use crate::emulate::emulate_mm_reduce_max_ps; + use crate::emulate::emulate_mm_reduce_min_ps; + unsafe { + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm_set1_ps(f32::INFINITY); + let mut max = _mm_set1_ps(f32::NEG_INFINITY); + while n >= 4 { + let x = _mm_loadu_ps(a); + a = a.add(4); + n -= 4; + min = _mm_min_ps(x, min); + max = _mm_max_ps(x, max); + } + let mut min = emulate_mm_reduce_min_ps(min); + let mut max = emulate_mm_reduce_max_ps(max); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + min = x.min(min); + max = x.max(max); + } + (min, max) + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_min_max_of_x_v2_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let x = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v2(x) }; + let fallback = reduce_min_max_of_x_fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + fn reduce_min_max_of_x_v8_3a(this: &[f32]) -> (f32, f32) { + unsafe { + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = vdupq_n_f32(f32::INFINITY); + let mut max = vdupq_n_f32(f32::NEG_INFINITY); + while n >= 4 { + let x = vld1q_f32(a); + a = a.add(4); + n -= 4; + min = vminq_f32(x, min); + max = vmaxq_f32(x, max); + } + let mut min = vminvq_f32(min); + let mut max = vmaxvq_f32(max); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + min = x.min(min); + max = x.max(max); + } + (min, max) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_min_max_of_x_v8_3a_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v8.3a") { + println!("test {} ... skipped (v8.3a)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let x = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v8_3a(x) }; + let fallback = reduce_min_max_of_x_fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "sve")] + fn reduce_min_max_of_x_v8_3a_sve(this: &[f32]) -> (f32, f32) { + let mut min = f32::INFINITY; + let mut max = -f32::INFINITY; + unsafe { + extern "C" { + fn fp32_reduce_min_max_of_x_v8_3a_sve( + this: *const f32, + n: usize, + out_min: &mut f32, + out_max: &mut f32, + ); + } + fp32_reduce_min_max_of_x_v8_3a_sve(this.as_ptr(), this.len(), &mut min, &mut max); + } + (min, max) + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_min_max_of_x_v8_3a_sve_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { + println!("test {} ... skipped (v8.3a:sve)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let x = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v8_3a_sve(x) }; + let fallback = reduce_min_max_of_x_fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a:sve", @"v8.3a")] + pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { + let mut min = f32::INFINITY; + let mut max = f32::NEG_INFINITY; + let n = this.len(); + for i in 0..n { + min = min.min(this[i]); + max = max.max(this[i]); + } + (min, max) + } +} + +mod reduce_sum_of_xy { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_v4(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm512_setzero_ps(); + while n >= 16 { + let x = _mm512_loadu_ps(a); + let y = _mm512_loadu_ps(b); + a = a.add(16); + b = b.add(16); + n -= 16; + xy = _mm512_fmadd_ps(x, y, xy); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = _mm512_maskz_loadu_ps(mask, a); + let y = _mm512_maskz_loadu_ps(mask, b); + xy = _mm512_fmadd_ps(x, y, xy); + } + _mm512_reduce_add_ps(xy) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v4_test() { + use rand::Rng; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_xy_v3(lhs: &[f32], rhs: &[f32]) -> f32 { + use crate::emulate::emulate_mm256_reduce_add_ps; + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm256_setzero_ps(); + while n >= 8 { + let x = _mm256_loadu_ps(a); + let y = _mm256_loadu_ps(b); + a = a.add(8); + b = b.add(8); + n -= 8; + xy = _mm256_fmadd_ps(x, y, xy); + } + if n >= 4 { + let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); + let y = _mm256_zextps128_ps256(_mm_loadu_ps(b)); + a = a.add(4); + b = b.add(4); + n -= 4; + xy = _mm256_fmadd_ps(x, y, xy); + } + let mut xy = emulate_mm256_reduce_add_ps(xy); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let y = b.read(); + a = a.add(1); + b = b.add(1); + n -= 1; + xy += x * y; + } + xy + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v3_test() { + use rand::Rng; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "fma")] + fn reduce_sum_of_xy_v2_fma(lhs: &[f32], rhs: &[f32]) -> f32 { + use crate::emulate::emulate_mm_reduce_add_ps; + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm_setzero_ps(); + while n >= 4 { + let x = _mm_loadu_ps(a); + let y = _mm_loadu_ps(b); + a = a.add(4); + b = b.add(4); + n -= 4; + xy = _mm_fmadd_ps(x, y, xy); + } + let mut xy = emulate_mm_reduce_add_ps(xy); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let y = b.read(); + a = a.add(1); + b = b.add(1); + n -= 1; + xy += x * y; + } + xy + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v2_fma_test() { + use rand::Rng; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { + println!("test {} ... skipped (v2:fma)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v2_fma(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + fn reduce_sum_of_xy_v8_3a(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::aarch64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = vdupq_n_f32(0.0); + while n >= 4 { + let x = vld1q_f32(a); + let y = vld1q_f32(b); + a = a.add(4); + b = b.add(4); + n -= 4; + xy = vfmaq_f32(xy, x, y); + } + let mut xy = vaddvq_f32(xy); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let y = b.read(); + a = a.add(1); + b = b.add(1); + n -= 1; + xy += x * y; + } + xy + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v8_3a_test() { + use rand::Rng; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("v8.3a") { + println!("test {} ... skipped (v8.3a)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v8_3a(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "sve")] + fn reduce_sum_of_xy_v8_3a_sve(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + extern "C" { + fn fp32_reduce_sum_of_xy_v8_3a_sve(a: *const f32, b: *const f32, n: usize) -> f32; + } + fp32_reduce_sum_of_xy_v8_3a_sve(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v8_3a_sve_test() { + use rand::Rng; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { + println!("test {} ... skipped (v8.3a:sve)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v8_3a_sve(lhs, rhs) }; + let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"v8.3a:sve", @"v8.3a")] + pub fn reduce_sum_of_xy(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let mut xy = 0.0f32; + for i in 0..n { + xy += lhs[i] * rhs[i]; + } + xy + } +} + +mod reduce_sum_of_d2 { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_d2_v4(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len() as u32; + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm512_setzero_ps(); + while n >= 16 { + let x = _mm512_loadu_ps(a); + let y = _mm512_loadu_ps(b); + a = a.add(16); + b = b.add(16); + n -= 16; + let d = _mm512_sub_ps(x, y); + d2 = _mm512_fmadd_ps(d, d, d2); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n) as u16; + let x = _mm512_maskz_loadu_ps(mask, a); + let y = _mm512_maskz_loadu_ps(mask, b); + let d = _mm512_sub_ps(x, y); + d2 = _mm512_fmadd_ps(d, d, d2); + } + _mm512_reduce_add_ps(d2) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_v4_test() { + use rand::Rng; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v4(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_d2_v3(lhs: &[f32], rhs: &[f32]) -> f32 { + use crate::emulate::emulate_mm256_reduce_add_ps; + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm256_setzero_ps(); + while n >= 8 { + let x = _mm256_loadu_ps(a); + let y = _mm256_loadu_ps(b); + a = a.add(8); + b = b.add(8); + n -= 8; + let d = _mm256_sub_ps(x, y); + d2 = _mm256_fmadd_ps(d, d, d2); + } + if n >= 4 { + let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); + let y = _mm256_zextps128_ps256(_mm_loadu_ps(b)); + a = a.add(4); + b = b.add(4); + n -= 4; + let d = _mm256_sub_ps(x, y); + d2 = _mm256_fmadd_ps(d, d, d2); + } + let mut d2 = emulate_mm256_reduce_add_ps(d2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let y = b.read(); + a = a.add(1); + b = b.add(1); + n -= 1; + let d = x - y; + d2 += d * d; + } + d2 + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_d2_v3_test() { + use rand::Rng; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v3(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "fma")] + fn reduce_sum_of_d2_v2_fma(lhs: &[f32], rhs: &[f32]) -> f32 { + use crate::emulate::emulate_mm_reduce_add_ps; + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm_setzero_ps(); + while n >= 4 { + let x = _mm_loadu_ps(a); + let y = _mm_loadu_ps(b); + a = a.add(4); + b = b.add(4); + n -= 4; + let d = _mm_sub_ps(x, y); + d2 = _mm_fmadd_ps(d, d, d2); + } + let mut d2 = emulate_mm_reduce_add_ps(d2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let y = b.read(); + a = a.add(1); + b = b.add(1); + n -= 1; + let d = x - y; + d2 += d * d; + } + d2 + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_d2_v2_fma_test() { + use rand::Rng; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { + println!("test {} ... skipped (v2:fma)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v2_fma(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + fn reduce_sum_of_d2_v8_3a(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + use std::arch::aarch64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = vdupq_n_f32(0.0); + while n >= 4 { + let x = vld1q_f32(a); + let y = vld1q_f32(b); + a = a.add(4); + b = b.add(4); + n -= 4; + let d = vsubq_f32(x, y); + d2 = vfmaq_f32(d2, d, d); + } + let mut d2 = vaddvq_f32(d2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let y = b.read(); + a = a.add(1); + b = b.add(1); + n -= 1; + let d = x - y; + d2 += d * d; + } + d2 + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_v8_3a_test() { + use rand::Rng; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("v8.3a") { + println!("test {} ... skipped (v8.3a)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v8_3a(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + #[target_feature(enable = "sve")] + fn reduce_sum_of_d2_v8_3a_sve(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + unsafe { + extern "C" { + fn fp32_reduce_sum_of_d2_v8_3a_sve(a: *const f32, b: *const f32, n: usize) -> f32; + } + fp32_reduce_sum_of_d2_v8_3a_sve(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_v8_3a_sve_test() { + use rand::Rng; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { + println!("test {} ... skipped (v8.3a:sve)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v8_3a_sve(lhs, rhs) }; + let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"v8.3a:sve", @"v8.3a")] + pub fn reduce_sum_of_d2(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let mut d2 = 0.0f32; + for i in 0..n { + let d = lhs[i] - rhs[i]; + d2 += d * d; + } + d2 + } +} + +mod reduce_sum_of_xy_sparse { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_sparse_v4(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { + use crate::emulate::emulate_mm512_2intersect_epi32; + assert_eq!(li.len(), lv.len()); + assert_eq!(ri.len(), rv.len()); + let (mut lp, ln) = (0, li.len()); + let (mut rp, rn) = (0, ri.len()); + let (li, lv) = (li.as_ptr(), lv.as_ptr()); + let (ri, rv) = (ri.as_ptr(), rv.as_ptr()); + unsafe { + use std::arch::x86_64::*; + let mut xy = _mm512_setzero_ps(); + while lp + 16 <= ln && rp + 16 <= rn { + let lx = _mm512_loadu_epi32(li.add(lp).cast()); + let rx = _mm512_loadu_epi32(ri.add(rp).cast()); + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))); + let rv = _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))); + xy = _mm512_fmadd_ps(lv, rv, xy); + let lt = li.add(lp + 16 - 1).read(); + let rt = ri.add(rp + 16 - 1).read(); + lp += (lt <= rt) as usize * 16; + rp += (lt >= rt) as usize * 16; + } + while lp < ln && rp < rn { + let lw = 16.min(ln - lp); + let rw = 16.min(rn - rp); + let lm = _bzhi_u32(0xffff, lw as _) as u16; + let rm = _bzhi_u32(0xffff, rw as _) as u16; + let lx = _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), lm, li.add(lp).cast()); + let rx = _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), rm, ri.add(rp).cast()); + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))); + let rv = _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))); + xy = _mm512_fmadd_ps(lv, rv, xy); + let lt = li.add(lp + lw - 1).read(); + let rt = ri.add(rp + rw - 1).read(); + lp += (lt <= rt) as usize * lw; + rp += (lt >= rt) as usize * rw; + } + _mm512_reduce_add_ps(xy) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_sparse_v4_test() { + use rand::Rng; + const EPSILON: f32 = 0.000001; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + pub fn sample_u32_sorted( + rng: &mut (impl Rng + ?Sized), + length: u32, + amount: u32, + ) -> Vec { + let mut x = match rand::seq::index::sample(rng, length as usize, amount as usize) { + rand::seq::index::IndexVec::U32(x) => x, + _ => unreachable!(), + }; + x.sort(); + x + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lm = 300; + let lidx = sample_u32_sorted(&mut rng, 10000, lm); + let lval = (0..lm) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rm = 350; + let ridx = sample_u32_sorted(&mut rng, 10000, rm); + let rval = (0..rm) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let specialized = unsafe { reduce_sum_of_xy_sparse_v4(&lidx, &lval, &ridx, &rval) }; + let fallback = reduce_sum_of_xy_sparse_fallback(&lidx, &lval, &ridx, &rval); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[crate::multiversion(@"v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { + use std::cmp::Ordering; + assert_eq!(lidx.len(), lval.len()); + assert_eq!(ridx.len(), rval.len()); + let (mut lp, ln) = (0, lidx.len()); + let (mut rp, rn) = (0, ridx.len()); + let mut xy = 0.0f32; + while lp < ln && rp < rn { + match Ord::cmp(&lidx[lp], &ridx[rp]) { + Ordering::Equal => { + xy += lval[lp] * rval[rp]; + lp += 1; + rp += 1; + } + Ordering::Less => { + lp += 1; + } + Ordering::Greater => { + rp += 1; + } + } + } + xy + } +} + +mod reduce_sum_of_d2_sparse { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_d2_sparse_v4(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { + use crate::emulate::emulate_mm512_2intersect_epi32; + assert_eq!(li.len(), lv.len()); + assert_eq!(ri.len(), rv.len()); + let (mut lp, ln) = (0, li.len()); + let (mut rp, rn) = (0, ri.len()); + let (li, lv) = (li.as_ptr(), lv.as_ptr()); + let (ri, rv) = (ri.as_ptr(), rv.as_ptr()); + unsafe { + use std::arch::x86_64::*; + let mut d2 = _mm512_setzero_ps(); + while lp + 16 <= ln && rp + 16 <= rn { + let lx = _mm512_loadu_epi32(li.add(lp).cast()); + let rx = _mm512_loadu_epi32(ri.add(rp).cast()); + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))); + let rv = _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))); + let d = _mm512_sub_ps(lv, rv); + d2 = _mm512_fmadd_ps(d, d, d2); + d2 = _mm512_sub_ps(d2, _mm512_mul_ps(lv, lv)); + d2 = _mm512_sub_ps(d2, _mm512_mul_ps(rv, rv)); + let lt = li.add(lp + 16 - 1).read(); + let rt = ri.add(rp + 16 - 1).read(); + lp += (lt <= rt) as usize * 16; + rp += (lt >= rt) as usize * 16; + } + while lp < ln && rp < rn { + let lw = 16.min(ln - lp); + let rw = 16.min(rn - rp); + let lm = _bzhi_u32(0xffff, lw as _) as u16; + let rm = _bzhi_u32(0xffff, rw as _) as u16; + let lx = _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), lm, li.add(lp).cast()); + let rx = _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), rm, ri.add(rp).cast()); + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))); + let rv = _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))); + let d = _mm512_sub_ps(lv, rv); + d2 = _mm512_fmadd_ps(d, d, d2); + d2 = _mm512_sub_ps(d2, _mm512_mul_ps(lv, lv)); + d2 = _mm512_sub_ps(d2, _mm512_mul_ps(rv, rv)); + let lt = li.add(lp + lw - 1).read(); + let rt = ri.add(rp + rw - 1).read(); + lp += (lt <= rt) as usize * lw; + rp += (lt >= rt) as usize * rw; + } + { + let mut lp = 0; + while lp + 16 <= ln { + let d = _mm512_loadu_ps(lv.add(lp)); + d2 = _mm512_fmadd_ps(d, d, d2); + lp += 16; + } + if lp < ln { + let lw = ln - lp; + let lm = _bzhi_u32(0xffff, lw as _) as u16; + let d = _mm512_maskz_loadu_ps(lm, lv.add(lp)); + d2 = _mm512_fmadd_ps(d, d, d2); + } + } + { + let mut rp = 0; + while rp + 16 <= rn { + let d = _mm512_loadu_ps(rv.add(rp)); + d2 = _mm512_fmadd_ps(d, d, d2); + rp += 16; + } + if rp < rn { + let rw = rn - rp; + let rm = _bzhi_u32(0xffff, rw as _) as u16; + let d = _mm512_maskz_loadu_ps(rm, rv.add(rp)); + d2 = _mm512_fmadd_ps(d, d, d2); + } + } + _mm512_reduce_add_ps(d2) + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_d2_sparse_v4_test() { + use rand::Rng; + const EPSILON: f32 = 0.0004; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + pub fn sample_u32_sorted( + rng: &mut (impl Rng + ?Sized), + length: u32, + amount: u32, + ) -> Vec { + let mut x = match rand::seq::index::sample(rng, length as usize, amount as usize) { + rand::seq::index::IndexVec::U32(x) => x, + _ => unreachable!(), + }; + x.sort(); + x + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lm = 300; + let lidx = sample_u32_sorted(&mut rng, 10000, lm); + let lval = (0..lm) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let rm = 350; + let ridx = sample_u32_sorted(&mut rng, 10000, rm); + let rval = (0..rm) + .map(|_| rng.gen_range(-1.0..=1.0)) + .collect::>(); + let specialized = unsafe { reduce_sum_of_d2_sparse_v4(&lidx, &lval, &ridx, &rval) }; + let fallback = reduce_sum_of_d2_sparse_fallback(&lidx, &lval, &ridx, &rval); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[crate::multiversion(@"v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { + use std::cmp::Ordering; + assert_eq!(lidx.len(), lval.len()); + assert_eq!(ridx.len(), rval.len()); + let (mut lp, ln) = (0, lidx.len()); + let (mut rp, rn) = (0, ridx.len()); + let mut d2 = 0.0f32; + while lp < ln && rp < rn { + match Ord::cmp(&lidx[lp], &ridx[rp]) { + Ordering::Equal => { + let d = lval[lp] - rval[rp]; + d2 += d * d; + lp += 1; + rp += 1; + } + Ordering::Less => { + d2 += lval[lp] * lval[lp]; + lp += 1; + } + Ordering::Greater => { + d2 += rval[rp] * rval[rp]; + rp += 1; + } + } + } + for i in lp..ln { + d2 += lval[i] * lval[i]; + } + for i in rp..rn { + d2 += rval[i] * rval[i]; + } + d2 + } +} + +mod vector_add { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] + rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_add_inplace { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + for i in 0..n { + lhs[i] += rhs[i]; + } + } +} + +mod vector_sub { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] - rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] * rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul_scalar { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] * rhs); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul_scalar_inplace { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { + let n = lhs.len(); + for i in 0..n { + lhs[i] *= rhs; + } + } +} + +mod vector_abs_inplace { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn vector_abs_inplace(this: &mut [f32]) { + let n = this.len(); + for i in 0..n { + this[i] = this[i].abs(); + } + } +} diff --git a/crates/simd/src/fast_scan/mod.rs b/crates/simd/src/fast_scan/mod.rs new file mode 100644 index 00000000..c2579af1 --- /dev/null +++ b/crates/simd/src/fast_scan/mod.rs @@ -0,0 +1,496 @@ +/* + +## codes layout for 4-bit quantizer + +group i = | vector i | (total bytes = width/2) + +byte: | 0 | 1 | 2 | ... | width/2 - 1 | +bits 0..3: | code 0 | code 2 | code 4 | ... | code width-2 | +bits 4..7: | code 1 | code 3 | code 5 | ... | code width-1 | + +## packed_codes layout for 4-bit quantizer + +group i = | vector 32i | vector 32i+1 | vector 32i+2 | ... | vector 32i+31 | (total bytes = width * 16) + +byte | 0 | 1 | 2 | ... | 14 | 15 | +bits 0..3 | code 0,vector 0 | code 0,vector 8 | code 0,vector 1 | ... | code 0,vector 14 | code 0,vector 15 | +bits 4..7 | code 0,vector 16 | code 0,vector 24 | code 0,vector 17 | ... | code 0,vector 30 | code 0,vector 31 | + +byte | 16 | 17 | 18 | ... | 30 | 31 | +bits 0..3 | code 1,vector 0 | code 1,vector 8 | code 1,vector 1 | ... | code 1,vector 14 | code 1,vector 15 | +bits 4..7 | code 1,vector 16 | code 1,vector 24 | code 1,vector 17 | ... | code 1,vector 30 | code 1,vector 31 | + +byte | 32 | 33 | 34 | ... | 46 | 47 | +bits 0..3 | code 2,vector 0 | code 2,vector 8 | code 2,vector 1 | ... | code 2,vector 14 | code 2,vector 15 | +bits 4..7 | code 2,vector 16 | code 2,vector 24 | code 2,vector 17 | ... | code 2,vector 30 | code 2,vector 31 | + +... + +byte | width*32-32 | width*32-31 | ... | width*32-1 | +bits 0..3 | code (width-1),vector 0 | code (width-1),vector 8 | ... | code (width-1),vector 15 | +bits 4..7 | code (width-1),vector 16 | code (width-1),vector 24 | ... | code (width-1),vector 31 | + +*/ + +pub fn pack(width: u32, r: [Vec; 32]) -> impl Iterator { + (0..width as usize).flat_map(move |i| { + [ + r[0][i] | (r[16][i] << 4), + r[8][i] | (r[24][i] << 4), + r[1][i] | (r[17][i] << 4), + r[9][i] | (r[25][i] << 4), + r[2][i] | (r[18][i] << 4), + r[10][i] | (r[26][i] << 4), + r[3][i] | (r[19][i] << 4), + r[11][i] | (r[27][i] << 4), + r[4][i] | (r[20][i] << 4), + r[12][i] | (r[28][i] << 4), + r[5][i] | (r[21][i] << 4), + r[13][i] | (r[29][i] << 4), + r[6][i] | (r[22][i] << 4), + r[14][i] | (r[30][i] << 4), + r[7][i] | (r[23][i] << 4), + r[15][i] | (r[31][i] << 4), + ] + .into_iter() + }) +} + +#[allow(clippy::module_inception)] +mod fast_scan { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn fast_scan_v4(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(codes.len(), width as usize * 16); + assert_eq!(lut.len(), width as usize * 16); + + unsafe { + use std::arch::x86_64::*; + + #[inline] + #[crate::target_cpu(enable = "v4")] + fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { + unsafe { + let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); + let x0y1 = _mm256_blend_epi32(x0x1, y0y1, 0xf0); + _mm256_add_epi16(x1y0, x0y1) + } + } + + #[inline] + #[crate::target_cpu(enable = "v4")] + fn combine4x2(x0x1x2x3: __m512i, y0y1y2y3: __m512i) -> __m256i { + unsafe { + let x0x1 = _mm512_castsi512_si256(x0x1x2x3); + let x2x3 = _mm512_extracti64x4_epi64(x0x1x2x3, 1); + let y0y1 = _mm512_castsi512_si256(y0y1y2y3); + let y2y3 = _mm512_extracti64x4_epi64(y0y1y2y3, 1); + let x01y01 = combine2x2(x0x1, y0y1); + let x23y23 = combine2x2(x2x3, y2y3); + _mm256_add_epi16(x01y01, x23y23) + } + } + + let mut accu_0 = _mm512_setzero_si512(); + let mut accu_1 = _mm512_setzero_si512(); + let mut accu_2 = _mm512_setzero_si512(); + let mut accu_3 = _mm512_setzero_si512(); + + let mut i = 0_usize; + while i + 4 <= width as usize { + let c = _mm512_loadu_si512(codes.as_ptr().add(i * 16).cast()); + + let mask = _mm512_set1_epi8(0xf); + let clo = _mm512_and_si512(c, mask); + let chi = _mm512_and_si512(_mm512_srli_epi16(c, 4), mask); + + let lut = _mm512_loadu_si512(lut.as_ptr().add(i * 16).cast()); + let res_lo = _mm512_shuffle_epi8(lut, clo); + accu_0 = _mm512_add_epi16(accu_0, res_lo); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_shuffle_epi8(lut, chi); + accu_2 = _mm512_add_epi16(accu_2, res_hi); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + + i += 4; + } + if i + 2 <= width as usize { + let c = _mm256_loadu_si256(codes.as_ptr().add(i * 16).cast()); + + let mask = _mm256_set1_epi8(0xf); + let clo = _mm256_and_si256(c, mask); + let chi = _mm256_and_si256(_mm256_srli_epi16(c, 4), mask); + + let lut = _mm256_loadu_si256(lut.as_ptr().add(i * 16).cast()); + let res_lo = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, clo)); + accu_0 = _mm512_add_epi16(accu_0, res_lo); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, chi)); + accu_2 = _mm512_add_epi16(accu_2, res_hi); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + + i += 2; + } + if i < width as usize { + let c = _mm_loadu_si128(codes.as_ptr().add(i * 16).cast()); + + let mask = _mm_set1_epi8(0xf); + let clo = _mm_and_si128(c, mask); + let chi = _mm_and_si128(_mm_srli_epi16(c, 4), mask); + + let lut = _mm_loadu_si128(lut.as_ptr().add(i * 16).cast()); + let res_lo = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, clo)); + accu_0 = _mm512_add_epi16(accu_0, res_lo); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, chi)); + accu_2 = _mm512_add_epi16(accu_2, res_hi); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, width as usize); + + let mut result = [0_u16; 32]; + + accu_0 = _mm512_sub_epi16(accu_0, _mm512_slli_epi16(accu_1, 8)); + _mm256_storeu_si256( + result.as_mut_ptr().add(0).cast(), + combine4x2(accu_0, accu_1), + ); + + accu_2 = _mm512_sub_epi16(accu_2, _mm512_slli_epi16(accu_3, 8)); + _mm256_storeu_si256( + result.as_mut_ptr().add(16).cast(), + combine4x2(accu_2, accu_3), + ); + + result + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn fast_scan_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + for width in 90..110 { + let codes = (0..16 * width).map(|_| rand::random()).collect::>(); + let lut = (0..16 * width).map(|_| rand::random()).collect::>(); + unsafe { + assert_eq!( + fast_scan_v4(width, &codes, &lut), + fast_scan_fallback(width, &codes, &lut) + ); + } + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn fast_scan_v3(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(codes.len(), width as usize * 16); + assert_eq!(lut.len(), width as usize * 16); + + unsafe { + use std::arch::x86_64::*; + + #[inline] + #[crate::target_cpu(enable = "v3")] + fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { + unsafe { + let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); + let x0y1 = _mm256_blend_epi32(x0x1, y0y1, 0xf0); + _mm256_add_epi16(x1y0, x0y1) + } + } + + let mut accu_0 = _mm256_setzero_si256(); + let mut accu_1 = _mm256_setzero_si256(); + let mut accu_2 = _mm256_setzero_si256(); + let mut accu_3 = _mm256_setzero_si256(); + + let mut i = 0_usize; + while i + 2 <= width as usize { + let c = _mm256_loadu_si256(codes.as_ptr().add(i * 16).cast()); + + let mask = _mm256_set1_epi8(0xf); + let clo = _mm256_and_si256(c, mask); + let chi = _mm256_and_si256(_mm256_srli_epi16(c, 4), mask); + + let lut = _mm256_loadu_si256(lut.as_ptr().add(i * 16).cast()); + let res_lo = _mm256_shuffle_epi8(lut, clo); + accu_0 = _mm256_add_epi16(accu_0, res_lo); + accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); + let res_hi = _mm256_shuffle_epi8(lut, chi); + accu_2 = _mm256_add_epi16(accu_2, res_hi); + accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); + + i += 2; + } + if i < width as usize { + let c = _mm_loadu_si128(codes.as_ptr().add(i * 16).cast()); + + let mask = _mm_set1_epi8(0xf); + let clo = _mm_and_si128(c, mask); + let chi = _mm_and_si128(_mm_srli_epi16(c, 4), mask); + + let lut = _mm_loadu_si128(lut.as_ptr().add(i * 16).cast()); + let res_lo = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, clo)); + accu_0 = _mm256_add_epi16(accu_0, res_lo); + accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); + let res_hi = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, chi)); + accu_2 = _mm256_add_epi16(accu_2, res_hi); + accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, width as usize); + + let mut result = [0_u16; 32]; + + accu_0 = _mm256_sub_epi16(accu_0, _mm256_slli_epi16(accu_1, 8)); + _mm256_storeu_si256( + result.as_mut_ptr().add(0).cast(), + combine2x2(accu_0, accu_1), + ); + + accu_2 = _mm256_sub_epi16(accu_2, _mm256_slli_epi16(accu_3, 8)); + _mm256_storeu_si256( + result.as_mut_ptr().add(16).cast(), + combine2x2(accu_2, accu_3), + ); + + result + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn fast_scan_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + for width in 90..110 { + let codes = (0..16 * width).map(|_| rand::random()).collect::>(); + let lut = (0..16 * width).map(|_| rand::random()).collect::>(); + unsafe { + assert_eq!( + fast_scan_v3(width, &codes, &lut), + fast_scan_fallback(width, &codes, &lut) + ); + } + } + } + } + + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn fast_scan_v2(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(codes.len(), width as usize * 16); + assert_eq!(lut.len(), width as usize * 16); + + unsafe { + use std::arch::x86_64::*; + + let mut accu_0 = _mm_setzero_si128(); + let mut accu_1 = _mm_setzero_si128(); + let mut accu_2 = _mm_setzero_si128(); + let mut accu_3 = _mm_setzero_si128(); + + let mut i = 0_usize; + while i < width as usize { + let c = _mm_loadu_si128(codes.as_ptr().add(i * 16).cast()); + + let mask = _mm_set1_epi8(0xf); + let clo = _mm_and_si128(c, mask); + let chi = _mm_and_si128(_mm_srli_epi16(c, 4), mask); + + let lut = _mm_loadu_si128(lut.as_ptr().add(i * 16).cast()); + let res_lo = _mm_shuffle_epi8(lut, clo); + accu_0 = _mm_add_epi16(accu_0, res_lo); + accu_1 = _mm_add_epi16(accu_1, _mm_srli_epi16(res_lo, 8)); + let res_hi = _mm_shuffle_epi8(lut, chi); + accu_2 = _mm_add_epi16(accu_2, res_hi); + accu_3 = _mm_add_epi16(accu_3, _mm_srli_epi16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, width as usize); + + let mut result = [0_u16; 32]; + + accu_0 = _mm_sub_epi16(accu_0, _mm_slli_epi16(accu_1, 8)); + _mm_storeu_si128(result.as_mut_ptr().add(0).cast(), accu_0); + _mm_storeu_si128(result.as_mut_ptr().add(8).cast(), accu_1); + + accu_2 = _mm_sub_epi16(accu_2, _mm_slli_epi16(accu_3, 8)); + _mm_storeu_si128(result.as_mut_ptr().add(16).cast(), accu_2); + _mm_storeu_si128(result.as_mut_ptr().add(24).cast(), accu_3); + + result + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn fast_scan_v2_test() { + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + for width in 90..110 { + let codes = (0..16 * width).map(|_| rand::random()).collect::>(); + let lut = (0..16 * width).map(|_| rand::random()).collect::>(); + unsafe { + assert_eq!( + fast_scan_v2(width, &codes, &lut), + fast_scan_fallback(width, &codes, &lut) + ); + } + } + } + } + + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + fn fast_scan_v8_3a(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(codes.len(), width as usize * 16); + assert_eq!(lut.len(), width as usize * 16); + + unsafe { + use std::arch::aarch64::*; + + let mut accu_0 = vdupq_n_u16(0); + let mut accu_1 = vdupq_n_u16(0); + let mut accu_2 = vdupq_n_u16(0); + let mut accu_3 = vdupq_n_u16(0); + + let mut i = 0_usize; + while i < width as usize { + let c = vld1q_u8(codes.as_ptr().add(i * 16).cast()); + + let mask = vdupq_n_u8(0xf); + let clo = vandq_u8(c, mask); + let chi = vandq_u8(vshrq_n_u8(c, 4), mask); + + let lut = vld1q_u8(lut.as_ptr().add(i * 16).cast()); + let res_lo = vreinterpretq_u16_u8(vqtbl1q_u8(lut, clo)); + accu_0 = vaddq_u16(accu_0, res_lo); + accu_1 = vaddq_u16(accu_1, vshrq_n_u16(res_lo, 8)); + let res_hi = vreinterpretq_u16_u8(vqtbl1q_u8(lut, chi)); + accu_2 = vaddq_u16(accu_2, res_hi); + accu_3 = vaddq_u16(accu_3, vshrq_n_u16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, width as usize); + + let mut result = [0_u16; 32]; + + accu_0 = vsubq_u16(accu_0, vshlq_n_u16(accu_1, 8)); + vst1q_u16(result.as_mut_ptr().add(0).cast(), accu_0); + vst1q_u16(result.as_mut_ptr().add(8).cast(), accu_1); + + accu_2 = vsubq_u16(accu_2, vshlq_n_u16(accu_3, 8)); + vst1q_u16(result.as_mut_ptr().add(16).cast(), accu_2); + vst1q_u16(result.as_mut_ptr().add(24).cast(), accu_3); + + result + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn fast_scan_v8_3a_test() { + if !crate::is_cpu_detected!("v8.3a") { + println!("test {} ... skipped (v8.3a)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + for width in 90..110 { + let codes = (0..16 * width).map(|_| rand::random()).collect::>(); + let lut = (0..16 * width).map(|_| rand::random()).collect::>(); + unsafe { + assert_eq!( + fast_scan_v8_3a(width, &codes, &lut), + fast_scan_fallback(width, &codes, &lut) + ); + } + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a")] + pub fn fast_scan(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + let width = width as usize; + + assert_eq!(codes.len(), width * 16); + assert_eq!(lut.len(), width * 16); + + use std::array::from_fn; + use std::ops::BitAnd; + + fn load(slice: &[T]) -> [T; N] { + from_fn(|i| slice[i]) + } + fn unary(op: impl Fn(T) -> U, a: [T; N]) -> [U; N] { + from_fn(|i| op(a[i])) + } + fn binary(op: impl Fn(T, T) -> T, a: [T; N], b: [T; N]) -> [T; N] { + from_fn(|i| op(a[i], b[i])) + } + fn shuffle(a: [T; N], b: [u8; N]) -> [T; N] { + from_fn(|i| a[b[i] as usize]) + } + fn cast(x: [u8; 16]) -> [u16; 8] { + from_fn(|i| u16::from_le_bytes([x[i << 1 | 0], x[i << 1 | 1]])) + } + fn setr(x: [[T; 8]; 4]) -> [T; 32] { + from_fn(|i| x[i >> 3][i & 7]) + } + + let mut a_0 = [0u16; 8]; + let mut a_1 = [0u16; 8]; + let mut a_2 = [0u16; 8]; + let mut a_3 = [0u16; 8]; + + for i in 0..width { + let c = load(&codes[16 * i..]); + + let mask = [0xfu8; 16]; + let clo = binary(u8::bitand, c, mask); + let chi = binary(u8::bitand, unary(|x| x >> 4, c), mask); + + let lut = load(&lut[16 * i..]); + let res_lo = cast(shuffle(lut, clo)); + a_0 = binary(u16::wrapping_add, a_0, res_lo); + a_1 = binary(u16::wrapping_add, a_1, unary(|x| x >> 8, res_lo)); + let res_hi = cast(shuffle(lut, chi)); + a_2 = binary(u16::wrapping_add, a_2, res_hi); + a_3 = binary(u16::wrapping_add, a_3, unary(|x| x >> 8, res_hi)); + } + + a_0 = binary(u16::wrapping_sub, a_0, unary(|x| x.wrapping_shl(8), a_1)); + a_2 = binary(u16::wrapping_sub, a_2, unary(|x| x.wrapping_shl(8), a_3)); + + setr([a_0, a_1, a_2, a_3]) + } +} + +#[inline(always)] +pub fn fast_scan(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + fast_scan::fast_scan(width, codes, lut) +} diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs new file mode 100644 index 00000000..aec5c523 --- /dev/null +++ b/crates/simd/src/lib.rs @@ -0,0 +1,107 @@ +#![feature(target_feature_11)] +#![feature(avx512_target_feature)] +#![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] +#![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512_f16))] + +mod aligned; +mod emulate; +mod f16; +mod f32; + +pub mod bit; +pub mod fast_scan; +pub mod packed_u4; +pub mod quantize; +pub mod u8; + +pub trait Floating: + Copy + + Send + + Sync + + std::fmt::Debug + + serde::Serialize + + for<'a> serde::Deserialize<'a> + + Default + + 'static + + PartialEq + + PartialOrd +{ + fn zero() -> Self; + fn infinity() -> Self; + fn mask(self, m: bool) -> Self; + + fn scalar_neg(this: Self) -> Self; + fn scalar_add(lhs: Self, rhs: Self) -> Self; + fn scalar_sub(lhs: Self, rhs: Self) -> Self; + fn scalar_mul(lhs: Self, rhs: Self) -> Self; + + fn reduce_or_of_is_zero_x(this: &[Self]) -> bool; + fn reduce_sum_of_x(this: &[Self]) -> f32; + fn reduce_sum_of_abs_x(this: &[Self]) -> f32; + fn reduce_sum_of_x2(this: &[Self]) -> f32; + fn reduce_min_max_of_x(this: &[Self]) -> (f32, f32); + fn reduce_sum_of_xy(lhs: &[Self], rhs: &[Self]) -> f32; + fn reduce_sum_of_d2(lhs: &[Self], rhs: &[Self]) -> f32; + fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[Self], ridx: &[u32], rval: &[Self]) -> f32; + fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[Self], ridx: &[u32], rval: &[Self]) -> f32; + + fn vector_from_f32(this: &[f32]) -> Vec; + fn vector_to_f32(this: &[Self]) -> Vec; + fn vector_to_f32_borrowed(this: &[Self]) -> impl AsRef<[f32]>; + fn vector_add(lhs: &[Self], rhs: &[Self]) -> Vec; + fn vector_add_inplace(lhs: &mut [Self], rhs: &[Self]); + fn vector_sub(lhs: &[Self], rhs: &[Self]) -> Vec; + fn vector_mul(lhs: &[Self], rhs: &[Self]) -> Vec; + fn vector_mul_scalar(lhs: &[Self], rhs: f32) -> Vec; + fn vector_mul_scalar_inplace(lhs: &mut [Self], rhs: f32); +} + +mod internal { + #[cfg(target_arch = "x86_64")] + simd_macros::define_is_cpu_detected!("x86_64"); + + #[cfg(target_arch = "aarch64")] + simd_macros::define_is_cpu_detected!("aarch64"); + + #[cfg(target_arch = "riscv64")] + simd_macros::define_is_cpu_detected!("riscv64"); + + #[cfg(target_arch = "x86_64")] + #[allow(unused_imports)] + pub use is_x86_64_cpu_detected; + + #[cfg(target_arch = "aarch64")] + #[allow(unused_imports)] + pub use is_aarch64_cpu_detected; + + #[cfg(target_arch = "riscv64")] + #[allow(unused_imports)] + pub use is_riscv64_cpu_detected; +} + +pub use simd_macros::multiversion; +pub use simd_macros::target_cpu; + +#[cfg(target_arch = "x86_64")] +#[allow(unused_imports)] +pub use std::arch::is_x86_feature_detected as is_feature_detected; + +#[cfg(target_arch = "aarch64")] +#[allow(unused_imports)] +pub use std::arch::is_aarch64_feature_detected as is_feature_detected; + +#[cfg(target_arch = "riscv64")] +#[allow(unused_imports)] +pub use std::arch::is_riscv_feature_detected as is_feature_detected; + +#[cfg(target_arch = "x86_64")] +#[allow(unused_imports)] +pub use internal::is_x86_64_cpu_detected as is_cpu_detected; + +#[cfg(target_arch = "aarch64")] +#[allow(unused_imports)] +pub use internal::is_aarch64_cpu_detected as is_cpu_detected; + +#[cfg(target_arch = "riscv64")] +#[allow(unused_imports)] +pub use internal::is_riscv64_cpu_detected as is_cpu_detected; diff --git a/crates/simd/src/packed_u4.rs b/crates/simd/src/packed_u4.rs new file mode 100644 index 00000000..d68333c7 --- /dev/null +++ b/crates/simd/src/packed_u4.rs @@ -0,0 +1,18 @@ +pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + reduce_sum_of_xy::reduce_sum_of_xy(s, t) +} + +mod reduce_sum_of_xy { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + assert_eq!(s.len(), t.len()); + let n = s.len(); + let mut result = 0; + for i in 0..n { + let (s, t) = (s[i], t[i]); + result += ((s & 15) as u32) * ((t & 15) as u32); + result += ((s >> 4) as u32) * ((t >> 4) as u32); + } + result + } +} diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs new file mode 100644 index 00000000..baaf594a --- /dev/null +++ b/crates/simd/src/quantize.rs @@ -0,0 +1,291 @@ +mod mul_add_round { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn mul_add_round_v4(this: &[f32], k: f32, b: f32) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + unsafe { + use std::arch::x86_64::*; + let lk = _mm512_set1_ps(k); + let lb = _mm512_set1_ps(b); + let mut n = n; + let mut a = this.as_ptr(); + let mut r = r.as_mut_ptr(); + while n >= 16 { + let x = _mm512_loadu_ps(a); + let v = + _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + let v = _mm512_cvtps_epi32(v); + let vfl = _mm512_cvtepi32_epi8(v); + _mm_storeu_si128(r.cast(), vfl); + n -= 16; + a = a.add(16); + r = r.add(16); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = _mm512_maskz_loadu_ps(mask, a); + let v = + _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + let v = _mm512_cvtps_epi32(v); + let vfl = _mm512_cvtepi32_epi8(v); + _mm_mask_storeu_epi8(r.cast(), mask, vfl); + } + } + unsafe { + r.set_len(n); + } + r + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn mul_add_round_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4010; + let x = (0..n).map(|_| rand::random::<_>()).collect::>(); + for z in 3990..4010 { + let x = &x[..z]; + let k = 20.0; + let b = 20.0; + let specialized = unsafe { mul_add_round_v4(x, k, b) }; + let fallback = mul_add_round_fallback(x, k, b); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn mul_add_round_v3(this: &[f32], k: f32, b: f32) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + unsafe { + use std::arch::x86_64::*; + let cons = _mm256_setr_epi8( + 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 + -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 + 0, 4, 8, 12, -1, -1, -1, -1, // 16..24 + -1, -1, -1, -1, -1, -1, -1, -1, // 24..32 + ); + let lk = _mm256_set1_ps(k); + let lb = _mm256_set1_ps(b); + let mut n = n; + let mut a = this.as_ptr(); + let mut r = r.as_mut_ptr(); + while n >= 8 { + let x = _mm256_loadu_ps(a); + let v = _mm256_fmadd_ps(x, lk, lb); + let v = _mm256_cvtps_epi32(_mm256_round_ps(v, 0x00)); + let vs = _mm256_shuffle_epi8(v, cons); + let vlo = _mm256_extract_epi32::<0>(vs) as u32; + let vhi = _mm256_extract_epi32::<4>(vs) as u32; + let vfl = vlo as u64 | ((vhi as u64) << 32); + r.cast::().write_unaligned(vfl); + n -= 8; + a = a.add(8); + r = r.add(8); + } + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let v = x.mul_add(k, b).round_ties_even() as u8; + r.write(v); + n -= 1; + a = a.add(1); + r = r.add(1); + } + } + unsafe { + r.set_len(n); + } + r + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn mul_add_round_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4010; + let x = (0..n).map(|_| rand::random::<_>()).collect::>(); + for z in 3990..4010 { + let x = &x[..z]; + let k = 20.0; + let b = 20.0; + let specialized = unsafe { mul_add_round_v3(x, k, b) }; + let fallback = mul_add_round_fallback(x, k, b); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "fma")] + fn mul_add_round_v2_fma(this: &[f32], k: f32, b: f32) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + unsafe { + use std::arch::x86_64::*; + let cons = _mm_setr_epi8( + 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 + -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 + ); + let lk = _mm_set1_ps(k); + let lb = _mm_set1_ps(b); + let mut n = n; + let mut a = this.as_ptr(); + let mut r = r.as_mut_ptr(); + while n >= 4 { + let x = _mm_loadu_ps(a); + let v = _mm_fmadd_ps(x, lk, lb); + let v = _mm_cvtps_epi32(_mm_round_ps(v, 0x00)); + let vs = _mm_shuffle_epi8(v, cons); + let vfl = _mm_extract_epi32::<0>(vs) as u32; + r.cast::().write_unaligned(vfl); + n -= 4; + a = a.add(4); + r = r.add(4); + } + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let v = x.mul_add(k, b).round_ties_even() as u8; + r.write(v); + n -= 1; + a = a.add(1); + r = r.add(1); + } + } + unsafe { + r.set_len(n); + } + r + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn mul_add_round_v2_fma_test() { + if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { + println!("test {} ... skipped (v2:fma)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4010; + let x = (0..n).map(|_| rand::random::<_>()).collect::>(); + for z in 3990..4010 { + let x = &x[..z]; + let k = 20.0; + let b = 20.0; + let specialized = unsafe { mul_add_round_v2_fma(x, k, b) }; + let fallback = mul_add_round_fallback(x, k, b); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + fn mul_add_round_v8_3a(this: &[f32], k: f32, b: f32) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + unsafe { + use std::arch::aarch64::*; + let cons = vld1q_u8( + [ + 0, 4, 8, 12, 0xff, 0xff, 0xff, 0xff, // 0..8 + 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // 8..15 + ] + .as_ptr(), + ); + let lk = vdupq_n_f32(k); + let lb = vdupq_n_f32(b); + let mut n = n; + let mut a = this.as_ptr(); + let mut r = r.as_mut_ptr(); + while n >= 4 { + let x = vld1q_f32(a); + let v = vfmaq_f32(lb, x, lk); + let v = vcvtnq_u32_f32(v); + let vs = vqtbl1q_u8(vreinterpretq_u8_u32(v), cons); + let vfl = vgetq_lane_u32::<0>(vreinterpretq_u32_u8(vs)); + r.cast::().write_unaligned(vfl); + n -= 4; + a = a.add(4); + r = r.add(4); + } + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + let v = x.mul_add(k, b).round_ties_even() as u8; + r.write(v); + n -= 1; + a = a.add(1); + r = r.add(1); + } + } + unsafe { + r.set_len(n); + } + r + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn mul_add_round_v8_3a_test() { + if !crate::is_cpu_detected!("v8.3a") { + println!("test {} ... skipped (v8.3a)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4010; + let x = (0..n).map(|_| rand::random::<_>()).collect::>(); + for z in 3990..4010 { + let x = &x[..z]; + let k = 20.0; + let b = 20.0; + let specialized = unsafe { mul_add_round_v8_3a(x, k, b) }; + let fallback = mul_add_round_fallback(x, k, b); + assert_eq!(specialized, fallback); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"v8.3a")] + pub fn mul_add_round(this: &[f32], k: f32, b: f32) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + let x = this[i]; + let v = x.mul_add(k, b).round_ties_even() as u8; + unsafe { + r.as_mut_ptr().add(i).write(v); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +#[inline(always)] +pub fn quantize(lut: &[f32], n: f32) -> (f32, f32, Vec) { + use crate::Floating; + let (min, max) = f32::reduce_min_max_of_x(lut); + let k = 0.0f32.max((max - min) / n); + let b = min; + (k, b, mul_add_round::mul_add_round(lut, 1.0 / k, -b / k)) +} diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs new file mode 100644 index 00000000..bc845332 --- /dev/null +++ b/crates/simd/src/u8.rs @@ -0,0 +1,343 @@ +mod reduce_sum_of_xy { + #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + assert_eq!(s.len(), t.len()); + let n = s.len(); + let mut result = 0; + for i in 0..n { + result += (s[i] as u32) * (t[i] as u32); + } + result + } +} + +#[inline(always)] +pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + reduce_sum_of_xy::reduce_sum_of_xy(s, t) +} + +mod reduce_sum_of_x_as_u16 { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_as_u16_v4(this: &[u8]) -> u16 { + use crate::emulate::emulate_mm512_reduce_add_epi16; + unsafe { + use std::arch::x86_64::*; + let us = _mm512_set1_epi16(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_si512(); + while n >= 32 { + let x = _mm256_loadu_si256(a.cast()); + a = a.add(32); + n -= 32; + sum = _mm512_add_epi16(_mm512_and_si512(us, _mm512_cvtepi8_epi16(x)), sum); + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x = _mm256_maskz_loadu_epi8(mask, a.cast()); + sum = _mm512_add_epi16(_mm512_and_si512(us, _mm512_cvtepi8_epi16(x)), sum); + } + emulate_mm512_reduce_add_epi16(sum) as u16 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_as_u16_v4_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_as_u16_v4(this) }; + let fallback = reduce_sum_of_x_as_u16_fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_as_u16_v3(this: &[u8]) -> u16 { + use crate::emulate::emulate_mm256_reduce_add_epi16; + unsafe { + use std::arch::x86_64::*; + let us = _mm256_set1_epi16(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_si256(); + while n >= 16 { + let x = _mm_loadu_si128(a.cast()); + a = a.add(16); + n -= 16; + sum = _mm256_add_epi16(_mm256_and_si256(us, _mm256_cvtepi8_epi16(x)), sum); + } + let mut sum = emulate_mm256_reduce_add_epi16(sum) as u16; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + sum += x as u16; + } + sum + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_as_u16_v3_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_as_u16_v3(this) }; + let fallback = reduce_sum_of_x_as_u16_fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_x_as_u16_v2(this: &[u8]) -> u16 { + use crate::emulate::emulate_mm_reduce_add_epi16; + unsafe { + use std::arch::x86_64::*; + let us = _mm_set1_epi16(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_si128(); + while n >= 8 { + let x = _mm_loadu_si64(a.cast()); + a = a.add(8); + n -= 8; + sum = _mm_add_epi16(_mm_and_si128(us, _mm_cvtepi8_epi16(x)), sum); + } + let mut sum = emulate_mm_reduce_add_epi16(sum) as u16; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + sum += x as u16; + } + sum + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_as_u16_v2_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_as_u16_v2(this) }; + let fallback = reduce_sum_of_x_as_u16_fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "v8.3a")] + fn reduce_sum_of_x_as_u16_v8_3a(this: &[u8]) -> u16 { + unsafe { + use std::arch::aarch64::*; + let us = vdupq_n_u16(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_u16(0); + while n >= 8 { + let x = vld1_u8(a); + a = a.add(8); + n -= 8; + sum = vaddq_u16(vandq_u16(us, vmovl_u8(x)), sum); + } + let mut sum = vaddvq_u16(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + sum += x as u16; + } + sum + } + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_as_u16_v8_3a_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v8.3a") { + println!("test {} ... skipped (v8.3a)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_as_u16_v8_3a(this) }; + let fallback = reduce_sum_of_x_as_u16_fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a")] + pub fn reduce_sum_of_x_as_u16(this: &[u8]) -> u16 { + let n = this.len(); + let mut sum = 0; + for i in 0..n { + sum += this[i] as u16; + } + sum + } +} + +#[inline(always)] +pub fn reduce_sum_of_x_as_u16(vector: &[u8]) -> u16 { + reduce_sum_of_x_as_u16::reduce_sum_of_x_as_u16(vector) +} + +mod reduce_sum_of_x { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[u8]) -> u32 { + unsafe { + use std::arch::x86_64::*; + let us = _mm512_set1_epi32(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_si512(); + while n >= 16 { + let x = _mm_loadu_epi8(a.cast()); + a = a.add(16); + n -= 16; + sum = _mm512_add_epi32(_mm512_and_si512(us, _mm512_cvtepi8_epi32(x)), sum); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = _mm_maskz_loadu_epi8(mask, a.cast()); + sum = _mm512_add_epi32(_mm512_and_si512(us, _mm512_cvtepi8_epi32(x)), sum); + } + _mm512_reduce_add_epi32(sum) as u32 + } + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_v4_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v4(this) }; + let fallback = reduce_sum_of_x_fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_v3(this: &[u8]) -> u32 { + use crate::emulate::emulate_mm256_reduce_add_epi32; + unsafe { + use std::arch::x86_64::*; + let us = _mm256_set1_epi32(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_si256(); + while n >= 8 { + let x = _mm_loadl_epi64(a.cast()); + a = a.add(8); + n -= 8; + sum = _mm256_add_epi32(_mm256_and_si256(us, _mm256_cvtepi8_epi32(x)), sum); + } + let mut sum = emulate_mm256_reduce_add_epi32(sum) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + sum += x as u32; + } + sum + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_as_u16_v3_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::thread_rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v3(this) }; + let fallback = reduce_sum_of_x_fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[crate::multiversion(@"v4", @"v3", "v2", "v8.3a:sve", "v8.3a")] + pub fn reduce_sum_of_x(this: &[u8]) -> u32 { + let n = this.len(); + let mut sum = 0; + for i in 0..n { + sum += this[i] as u32; + } + sum + } +} + +#[inline(always)] +pub fn reduce_sum_of_x(vector: &[u8]) -> u32 { + reduce_sum_of_x::reduce_sum_of_x(vector) +} diff --git a/crates/simd_macros/Cargo.toml b/crates/simd_macros/Cargo.toml new file mode 100644 index 00000000..799db570 --- /dev/null +++ b/crates/simd_macros/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "simd_macros" +version.workspace = true +edition.workspace = true + +[lib] +proc-macro = true + +[dependencies] +proc-macro2 = { version = "1.0.79", features = ["proc-macro"] } +quote = "1.0.35" +syn = { version = "2.0.53", default-features = false, features = [ + "clone-impls", + "full", + "parsing", + "printing", + "proc-macro", +] } diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs new file mode 100644 index 00000000..fe455646 --- /dev/null +++ b/crates/simd_macros/src/lib.rs @@ -0,0 +1,206 @@ +mod target; + +struct MultiversionVersion { + target: String, + import: bool, +} + +impl syn::parse::Parse for MultiversionVersion { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + let lookahead1 = input.lookahead1(); + if lookahead1.peek(syn::Token![@]) { + let _: syn::Token![@] = input.parse()?; + let target: syn::LitStr = input.parse()?; + Ok(Self { + target: target.value(), + import: true, + }) + } else { + let target: syn::LitStr = input.parse()?; + Ok(Self { + target: target.value(), + import: false, + }) + } + } +} + +struct Multiversion { + versions: syn::punctuated::Punctuated, +} + +impl syn::parse::Parse for Multiversion { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + Ok(Multiversion { + versions: syn::punctuated::Punctuated::parse_terminated(input)?, + }) + } +} + +#[proc_macro_attribute] +pub fn multiversion( + attr: proc_macro::TokenStream, + item: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + let attr = syn::parse_macro_input!(attr as Multiversion); + let item_fn = syn::parse::(item).expect("not a function item"); + let syn::ItemFn { + attrs, + vis, + sig, + block, + } = item_fn; + let name = sig.ident.to_string(); + if sig.constness.is_some() { + panic!("const functions are not supported"); + } + if sig.asyncness.is_some() { + panic!("async functions are not supported"); + } + let generics_params = sig.generics.params.clone(); + for generic_param in generics_params.iter() { + if !matches!(generic_param, syn::GenericParam::Lifetime(_)) { + panic!("generic parameters are not supported"); + } + } + let generics_where = sig.generics.where_clause.clone(); + let inputs = sig.inputs.clone(); + let arguments = { + let mut list = vec![]; + for x in sig.inputs.iter() { + if let syn::FnArg::Typed(y) = x { + if let syn::Pat::Ident(ident) = *y.pat.clone() { + list.push(ident); + } else { + panic!("patterns on parameters are not supported") + } + } else { + panic!("receiver parameters are not supported") + } + } + list + }; + if sig.variadic.is_some() { + panic!("variadic parameters are not supported"); + } + let output = sig.output.clone(); + let mut versions = quote::quote! {}; + let mut branches = quote::quote! {}; + for version in attr.versions { + let target = version.target.clone(); + let name = syn::Ident::new( + &format!("{name}_{}", target.replace(":", "_").replace(".", "_")), + proc_macro2::Span::mixed_site(), + ); + let s = target.split(":").collect::>(); + let target_cpu = target::TARGET_CPUS + .iter() + .find(|target_cpu| target_cpu.target_cpu == s[0]) + .expect("unknown target_cpu"); + let additional_target_features = s[1..].to_vec(); + let target_arch = target_cpu.target_arch; + let target_cpu = target_cpu.target_cpu; + if !version.import { + versions.extend(quote::quote! { + #[inline] + #[cfg(any(target_arch = #target_arch))] + #[crate::target_cpu(enable = #target_cpu)] + #(#[target_feature(enable = #additional_target_features)])* + fn #name < #generics_params > (#inputs) #output #generics_where { #block } + }); + } + branches.extend(quote::quote! { + #[cfg(target_arch = #target_arch)] + if crate::is_cpu_detected!(#target_cpu) #(&& crate::is_feature_detected!(#additional_target_features))* { + let _multiversion_internal: unsafe fn(#inputs) #output = #name; + CACHE.store(_multiversion_internal as *mut (), core::sync::atomic::Ordering::Relaxed); + return unsafe { _multiversion_internal(#(#arguments,)*) }; + } + }); + } + let fallback_name = + syn::Ident::new(&format!("{name}_fallback"), proc_macro2::Span::mixed_site()); + quote::quote! { + #versions + fn #fallback_name < #generics_params > (#inputs) #output #generics_where { #block } + #[inline(always)] + #(#attrs)* #vis #sig { + static CACHE: core::sync::atomic::AtomicPtr<()> = core::sync::atomic::AtomicPtr::new(core::ptr::null_mut()); + let cache = CACHE.load(core::sync::atomic::Ordering::Relaxed); + if !cache.is_null() { + let f = unsafe { core::mem::transmute::<*mut (), unsafe fn(#inputs) #output>(cache as _) }; + return unsafe { f(#(#arguments,)*) }; + } + #branches + let _multiversion_internal: unsafe fn(#inputs) #output = #fallback_name; + CACHE.store(_multiversion_internal as *mut (), core::sync::atomic::Ordering::Relaxed); + unsafe { _multiversion_internal(#(#arguments,)*) } + } + } + .into() +} + +struct TargetCpu { + enable: String, +} + +impl syn::parse::Parse for TargetCpu { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + let _: syn::Ident = input.parse()?; + let _: syn::Token![=] = input.parse()?; + let enable: syn::LitStr = input.parse()?; + Ok(Self { + enable: enable.value(), + }) + } +} + +#[proc_macro_attribute] +pub fn target_cpu( + attr: proc_macro::TokenStream, + item: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + let attr = syn::parse_macro_input!(attr as TargetCpu); + let mut result = quote::quote! {}; + for s in attr.enable.split(',') { + let target_cpu = target::TARGET_CPUS + .iter() + .find(|target_cpu| target_cpu.target_cpu == s) + .expect("unknown target_cpu"); + let target_features = target_cpu.target_features; + result.extend(quote::quote!( + #(#[target_feature(enable = #target_features)])* + )); + } + result.extend(proc_macro2::TokenStream::from(item)); + result.into() +} + +#[proc_macro] +pub fn define_is_cpu_detected(input: proc_macro::TokenStream) -> proc_macro::TokenStream { + let target_arch = syn::parse_macro_input!(input as syn::LitStr).value(); + let mut arms = quote::quote! {}; + for target_cpu in target::TARGET_CPUS { + if target_cpu.target_arch != target_arch { + continue; + } + let target_features = target_cpu.target_features; + let target_cpu = target_cpu.target_cpu; + arms.extend(quote::quote! { + (#target_cpu) => { + true #(&& $crate::is_feature_detected!(#target_features))* + }; + }); + } + let name = syn::Ident::new( + &format!("is_{target_arch}_cpu_detected"), + proc_macro2::Span::mixed_site(), + ); + quote::quote! { + #[macro_export] + macro_rules! #name { + #arms + } + } + .into() +} diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs new file mode 100644 index 00000000..c2584013 --- /dev/null +++ b/crates/simd_macros/src/target.rs @@ -0,0 +1,83 @@ +pub struct TargetCpu { + pub target_cpu: &'static str, + pub target_arch: &'static str, + pub target_features: &'static [&'static str], +} + +pub const TARGET_CPUS: &[TargetCpu] = &[ + TargetCpu { + target_cpu: "v4", + target_arch: "x86_64", + target_features: &[ + "avx", + "avx2", + "avx512bw", + "avx512cd", + "avx512dq", + "avx512f", + "avx512vl", + "bmi1", + "bmi2", + "cmpxchg16b", + "f16c", + "fma", + "fxsr", + "lzcnt", + "movbe", + "popcnt", + "sse", + "sse2", + "sse3", + "sse4.1", + "sse4.2", + "ssse3", + "xsave", + ], + }, + TargetCpu { + target_cpu: "v3", + target_arch: "x86_64", + target_features: &[ + "avx", + "avx2", + "bmi1", + "bmi2", + "cmpxchg16b", + "f16c", + "fma", + "fxsr", + "lzcnt", + "movbe", + "popcnt", + "sse", + "sse2", + "sse3", + "sse4.1", + "sse4.2", + "ssse3", + "xsave", + ], + }, + TargetCpu { + target_cpu: "v2", + target_arch: "x86_64", + target_features: &[ + "cmpxchg16b", + "fxsr", + "popcnt", + "sse", + "sse2", + "sse3", + "sse4.1", + "sse4.2", + "ssse3", + ], + }, + TargetCpu { + target_cpu: "v8.3a", + target_arch: "aarch64", + target_features: &[ + "crc", "dpb", "fcma", "jsconv", "lse", "neon", "paca", "pacg", "rcpc", "rdm", + ], + }, +]; diff --git a/crates/vector/Cargo.toml b/crates/vector/Cargo.toml new file mode 100644 index 00000000..b910d361 --- /dev/null +++ b/crates/vector/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "vector" +version.workspace = true +edition.workspace = true + +[dependencies] +distance = { path = "../distance" } +half.workspace = true +serde.workspace = true +simd = { path = "../simd" } + +[lints] +workspace = true diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs new file mode 100644 index 00000000..fe80164c --- /dev/null +++ b/crates/vector/src/bvect.rs @@ -0,0 +1,274 @@ +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; +use serde::{Deserialize, Serialize}; +use std::ops::{Bound, RangeBounds}; + +pub const BVECTOR_WIDTH: u32 = u64::BITS; + +// When using binary vector, please ensure that the padding bits are always zero. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BVectOwned { + dims: u32, + data: Vec, +} + +impl BVectOwned { + #[inline(always)] + pub fn new(dims: u32, data: Vec) -> Self { + Self::new_checked(dims, data).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(dims: u32, data: Vec) -> Option { + if !(1..=65535).contains(&dims) { + return None; + } + if data.len() != dims.div_ceil(BVECTOR_WIDTH) as usize { + return None; + } + if dims % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dims % BVECTOR_WIDTH) != 0 { + return None; + } + unsafe { Some(Self::new_unchecked(dims, data)) } + } + + /// # Safety + /// + /// * `dims` must be in `1..=65535`. + /// * `data` must be of the correct length. + /// * The padding bits must be zero. + #[inline(always)] + pub unsafe fn new_unchecked(dims: u32, data: Vec) -> Self { + Self { dims, data } + } +} + +impl VectorOwned for BVectOwned { + type Borrowed<'a> = BVectBorrowed<'a>; + + #[inline(always)] + fn as_borrowed(&self) -> BVectBorrowed<'_> { + BVectBorrowed { + dims: self.dims, + data: &self.data, + } + } + + #[inline(always)] + fn zero(dims: u32) -> Self { + Self::new(dims, vec![0; dims.div_ceil(BVECTOR_WIDTH) as usize]) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct BVectBorrowed<'a> { + dims: u32, + data: &'a [u64], +} + +impl<'a> BVectBorrowed<'a> { + #[inline(always)] + pub fn new(dims: u32, data: &'a [u64]) -> Self { + Self::new_checked(dims, data).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(dims: u32, data: &'a [u64]) -> Option { + if !(1..=65535).contains(&dims) { + return None; + } + if data.len() != dims.div_ceil(BVECTOR_WIDTH) as usize { + return None; + } + if dims % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dims % BVECTOR_WIDTH) != 0 { + return None; + } + unsafe { Some(Self::new_unchecked(dims, data)) } + } + + /// # Safety + /// + /// * `dims` must be in `1..=65535`. + /// * `data` must be of the correct length. + /// * The padding bits must be zero. + #[inline(always)] + pub unsafe fn new_unchecked(dims: u32, data: &'a [u64]) -> Self { + Self { dims, data } + } + + #[inline(always)] + pub fn data(&self) -> &'a [u64] { + self.data + } + + #[inline(always)] + pub fn get(&self, index: u32) -> bool { + assert!(index < self.dims); + self.data[(index / BVECTOR_WIDTH) as usize] & (1 << (index % BVECTOR_WIDTH)) != 0 + } + + #[inline(always)] + pub fn iter(self) -> impl Iterator + 'a { + let mut index = 0_u32; + std::iter::from_fn(move || { + if index < self.dims { + let result = self.data[(index / BVECTOR_WIDTH) as usize] + & (1 << (index % BVECTOR_WIDTH)) + != 0; + index += 1; + Some(result) + } else { + None + } + }) + } +} + +impl VectorBorrowed for BVectBorrowed<'_> { + type Owned = BVectOwned; + + #[inline(always)] + fn dims(&self) -> u32 { + self.dims + } + + fn own(&self) -> BVectOwned { + BVectOwned { + dims: self.dims, + data: self.data.to_vec(), + } + } + + #[inline(always)] + fn norm(&self) -> f32 { + (simd::bit::sum_of_x(self.data) as f32).sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + Distance::from(-(simd::bit::sum_of_and(self.data, rhs.data) as f32)) + } + + #[inline(always)] + fn operator_l2(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_cos(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_hamming(self, rhs: Self) -> Distance { + Distance::from(simd::bit::sum_of_xor(self.data, rhs.data) as f32) + } + + #[inline(always)] + fn operator_jaccard(self, rhs: Self) -> Distance { + let (and, or) = simd::bit::sum_of_and_or(self.data, rhs.data); + Distance::from(1.0 - (and as f32 / or as f32)) + } + + #[inline(always)] + fn function_normalize(&self) -> BVectOwned { + unimplemented!() + } + + fn operator_add(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_sub(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_mul(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_and(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dims, rhs.dims); + let data = simd::bit::vector_and(self.data, self.data); + BVectOwned::new(self.dims, data) + } + + fn operator_or(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dims, rhs.dims); + let data = simd::bit::vector_or(self.data, rhs.data); + BVectOwned::new(self.dims, data) + } + + fn operator_xor(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dims, rhs.dims); + let data = simd::bit::vector_xor(self.data, rhs.data); + BVectOwned::new(self.dims, data) + } + + #[inline(always)] + fn subvector(&self, bounds: impl RangeBounds) -> Option { + let start = match bounds.start_bound().cloned() { + Bound::Included(x) => x, + Bound::Excluded(u32::MAX) => return None, + Bound::Excluded(x) => x + 1, + Bound::Unbounded => 0, + }; + let end = match bounds.end_bound().cloned() { + Bound::Included(u32::MAX) => return None, + Bound::Included(x) => x + 1, + Bound::Excluded(x) => x, + Bound::Unbounded => self.dims, + }; + if start >= end || end > self.dims { + return None; + } + let dims = end - start; + let mut data = vec![0_u64; dims.div_ceil(BVECTOR_WIDTH) as _]; + { + let mut i = 0; + let mut j = start; + while j < end { + if self.data[(j / BVECTOR_WIDTH) as usize] & (1 << (j % BVECTOR_WIDTH)) != 0 { + data[(i / BVECTOR_WIDTH) as usize] |= 1 << (i % BVECTOR_WIDTH); + } + i += 1; + j += 1; + } + } + Self::Owned::new_checked(dims, data) + } +} + +impl PartialEq for BVectBorrowed<'_> { + fn eq(&self, other: &Self) -> bool { + if self.dims != other.dims { + return false; + } + for (&l, &r) in self.data.iter().zip(other.data.iter()) { + let l = l.reverse_bits(); + let r = r.reverse_bits(); + if l != r { + return false; + } + } + true + } +} + +impl PartialOrd for BVectBorrowed<'_> { + fn partial_cmp(&self, other: &Self) -> Option { + use std::cmp::Ordering; + if self.dims != other.dims { + return None; + } + for (&l, &r) in self.data.iter().zip(other.data.iter()) { + let l = l.reverse_bits(); + let r = r.reverse_bits(); + match l.cmp(&r) { + Ordering::Equal => (), + x => return Some(x), + } + } + Some(Ordering::Equal) + } +} diff --git a/crates/vector/src/lib.rs b/crates/vector/src/lib.rs new file mode 100644 index 00000000..e82e4a6c --- /dev/null +++ b/crates/vector/src/lib.rs @@ -0,0 +1,48 @@ +pub mod bvect; +pub mod scalar8; +pub mod svect; +pub mod vect; + +pub trait VectorOwned: Clone + serde::Serialize + for<'a> serde::Deserialize<'a> + 'static { + type Borrowed<'a>: VectorBorrowed; + + fn as_borrowed(&self) -> Self::Borrowed<'_>; + + fn zero(dims: u32) -> Self; +} + +pub trait VectorBorrowed: Copy { + type Owned: VectorOwned; + + fn own(&self) -> Self::Owned; + + fn dims(&self) -> u32; + + fn norm(&self) -> f32; + + fn operator_dot(self, rhs: Self) -> distance::Distance; + + fn operator_l2(self, rhs: Self) -> distance::Distance; + + fn operator_cos(self, rhs: Self) -> distance::Distance; + + fn operator_hamming(self, rhs: Self) -> distance::Distance; + + fn operator_jaccard(self, rhs: Self) -> distance::Distance; + + fn function_normalize(&self) -> Self::Owned; + + fn operator_add(&self, rhs: Self) -> Self::Owned; + + fn operator_sub(&self, rhs: Self) -> Self::Owned; + + fn operator_mul(&self, rhs: Self) -> Self::Owned; + + fn operator_and(&self, rhs: Self) -> Self::Owned; + + fn operator_or(&self, rhs: Self) -> Self::Owned; + + fn operator_xor(&self, rhs: Self) -> Self::Owned; + + fn subvector(&self, bounds: impl std::ops::RangeBounds) -> Option; +} diff --git a/src/types/scalar8.rs b/crates/vector/src/scalar8.rs similarity index 93% rename from src/types/scalar8.rs rename to crates/vector/src/scalar8.rs index 55c0074a..ff9095aa 100644 --- a/src/types/scalar8.rs +++ b/crates/vector/src/scalar8.rs @@ -1,5 +1,5 @@ -use base::distance::Distance; -use base::vector::{VectorBorrowed, VectorOwned}; +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; use serde::{Deserialize, Serialize}; use std::ops::RangeBounds; @@ -13,7 +13,6 @@ pub struct Scalar8Owned { } impl Scalar8Owned { - #[allow(dead_code)] #[inline(always)] pub fn new(sum_of_x2: f32, k: f32, b: f32, sum_of_code: f32, code: Vec) -> Self { Self::new_checked(sum_of_x2, k, b, sum_of_code, code).expect("invalid data") @@ -182,7 +181,7 @@ impl VectorBorrowed for Scalar8Borrowed<'_> { #[inline(always)] fn operator_dot(self, rhs: Self) -> Distance { assert_eq!(self.code.len(), rhs.code.len()); - let xy = self.k * rhs.k * base::simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 + let xy = self.k * rhs.k * simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 + self.b * rhs.b * self.code.len() as f32 + self.k * rhs.b * self.sum_of_code + self.b * rhs.k * rhs.sum_of_code; @@ -192,7 +191,7 @@ impl VectorBorrowed for Scalar8Borrowed<'_> { #[inline(always)] fn operator_l2(self, rhs: Self) -> Distance { assert_eq!(self.code.len(), rhs.code.len()); - let xy = self.k * rhs.k * base::simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 + let xy = self.k * rhs.k * simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 + self.b * rhs.b * self.code.len() as f32 + self.k * rhs.b * self.sum_of_code + self.b * rhs.k * rhs.sum_of_code; @@ -204,7 +203,7 @@ impl VectorBorrowed for Scalar8Borrowed<'_> { #[inline(always)] fn operator_cos(self, rhs: Self) -> Distance { assert_eq!(self.code.len(), rhs.code.len()); - let xy = self.k * rhs.k * base::simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 + let xy = self.k * rhs.k * simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 + self.b * rhs.b * self.code.len() as f32 + self.k * rhs.b * self.sum_of_code + self.b * rhs.k * rhs.sum_of_code; @@ -279,7 +278,7 @@ impl VectorBorrowed for Scalar8Borrowed<'_> { }, self.k, self.b, - base::simd::u8::reduce_sum_of_x_as_u32(code) as f32, + simd::u8::reduce_sum_of_x(code) as f32, code.to_owned(), ) } diff --git a/crates/vector/src/svect.rs b/crates/vector/src/svect.rs new file mode 100644 index 00000000..08f678d3 --- /dev/null +++ b/crates/vector/src/svect.rs @@ -0,0 +1,445 @@ +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; +use serde::{Deserialize, Serialize}; +use simd::Floating; +use std::ops::{Bound, RangeBounds}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SVectOwned { + dims: u32, + indexes: Vec, + values: Vec, +} + +impl SVectOwned { + #[inline(always)] + pub fn new(dims: u32, indexes: Vec, values: Vec) -> Self { + Self::new_checked(dims, indexes, values).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(dims: u32, indexes: Vec, values: Vec) -> Option { + if !(1..=1_048_575).contains(&dims) { + return None; + } + if indexes.len() != values.len() { + return None; + } + let len = indexes.len(); + for i in 1..len { + if !(indexes[i - 1] < indexes[i]) { + return None; + } + } + if len != 0 && !(indexes[len - 1] < dims) { + return None; + } + if S::reduce_or_of_is_zero_x(&values) { + return None; + } + unsafe { Some(Self::new_unchecked(dims, indexes, values)) } + } + + /// # Safety + /// + /// * `dims` must be in `1..=1_048_575`. + /// * `indexes.len()` must be equal to `values.len()`. + /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dims`. + /// * A floating number in `values` must not be positive zero or negative zero. + #[inline(always)] + pub unsafe fn new_unchecked(dims: u32, indexes: Vec, values: Vec) -> Self { + Self { + dims, + indexes, + values, + } + } + + #[inline(always)] + pub fn indexes(&self) -> &[u32] { + &self.indexes + } + + #[inline(always)] + pub fn values(&self) -> &[S] { + &self.values + } +} + +impl VectorOwned for SVectOwned { + type Borrowed<'a> = SVectBorrowed<'a, S>; + + #[inline(always)] + fn as_borrowed(&self) -> SVectBorrowed<'_, S> { + SVectBorrowed { + dims: self.dims, + indexes: &self.indexes, + values: &self.values, + } + } + + #[inline(always)] + fn zero(dims: u32) -> Self { + Self::new(dims, vec![], vec![]) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct SVectBorrowed<'a, S> { + dims: u32, + indexes: &'a [u32], + values: &'a [S], +} + +impl<'a, S: Floating> SVectBorrowed<'a, S> { + #[inline(always)] + pub fn new(dims: u32, indexes: &'a [u32], values: &'a [S]) -> Self { + Self::new_checked(dims, indexes, values).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(dims: u32, indexes: &'a [u32], values: &'a [S]) -> Option { + if !(1..=1_048_575).contains(&dims) { + return None; + } + if indexes.len() != values.len() { + return None; + } + let len = indexes.len(); + for i in 1..len { + if !(indexes[i - 1] < indexes[i]) { + return None; + } + } + if len != 0 && !(indexes[len - 1] < dims) { + return None; + } + for i in 0..len { + if values[i] == S::zero() { + return None; + } + } + unsafe { Some(Self::new_unchecked(dims, indexes, values)) } + } + + /// # Safety + /// + /// * `dims` must be in `1..=1_048_575`. + /// * `indexes.len()` must be equal to `values.len()`. + /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dims`. + /// * A floating number in `values` must not be positive zero or negative zero. + #[inline(always)] + pub unsafe fn new_unchecked(dims: u32, indexes: &'a [u32], values: &'a [S]) -> Self { + Self { + dims, + indexes, + values, + } + } + + #[inline(always)] + pub fn indexes(&self) -> &'a [u32] { + self.indexes + } + + #[inline(always)] + pub fn values(&self) -> &'a [S] { + self.values + } + + #[inline(always)] + #[allow(clippy::len_without_is_empty)] + pub fn len(&self) -> u32 { + self.indexes.len() as u32 + } +} + +impl VectorBorrowed for SVectBorrowed<'_, S> { + type Owned = SVectOwned; + + #[inline(always)] + fn dims(&self) -> u32 { + self.dims + } + + #[inline(always)] + fn own(&self) -> SVectOwned { + SVectOwned { + dims: self.dims, + indexes: self.indexes.to_vec(), + values: self.values.to_vec(), + } + } + + #[inline(always)] + fn norm(&self) -> f32 { + S::reduce_sum_of_x2(self.values).sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + let xy = S::reduce_sum_of_xy_sparse(self.indexes, self.values, rhs.indexes, rhs.values); + Distance::from(-xy) + } + + #[inline(always)] + fn operator_l2(self, rhs: Self) -> Distance { + let d2 = S::reduce_sum_of_d2_sparse(self.indexes, self.values, rhs.indexes, rhs.values); + Distance::from(d2) + } + + #[inline(always)] + fn operator_cos(self, rhs: Self) -> Distance { + let xy = S::reduce_sum_of_xy_sparse(self.indexes, self.values, rhs.indexes, rhs.values); + let x2 = S::reduce_sum_of_x2(self.values); + let y2 = S::reduce_sum_of_x2(rhs.values); + Distance::from(1.0 - xy / (x2 * y2).sqrt()) + } + + #[inline(always)] + fn operator_hamming(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_jaccard(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn function_normalize(&self) -> SVectOwned { + let l = S::reduce_sum_of_x2(self.values).sqrt(); + let mut indexes = self.indexes.to_vec(); + let mut values = self.values.to_vec(); + let n = indexes.len(); + S::vector_mul_scalar_inplace(&mut values, 1.0 / l); + let mut j = 0_usize; + for i in 0..n { + if values[i] != S::zero() { + indexes[j] = indexes[i]; + values[j] = values[i]; + j += 1; + } + } + indexes.truncate(j); + values.truncate(j); + SVectOwned::new(self.dims, indexes, values) + } + + fn operator_add(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dims, rhs.dims); + let size1 = self.len(); + let size2 = rhs.len(); + let mut pos1 = 0; + let mut pos2 = 0; + let mut pos = 0; + let mut indexes = vec![0; (size1 + size2) as _]; + let mut values = vec![S::zero(); (size1 + size2) as _]; + while pos1 < size1 && pos2 < size2 { + let lhs_index = self.indexes[pos1 as usize]; + let rhs_index = rhs.indexes[pos2 as usize]; + let lhs_value = self.values[pos1 as usize]; + let rhs_value = rhs.values[pos2 as usize]; + indexes[pos] = lhs_index.min(rhs_index); + values[pos] = S::scalar_add( + lhs_value.mask(lhs_index <= rhs_index), + rhs_value.mask(lhs_index >= rhs_index), + ); + pos1 += (lhs_index <= rhs_index) as u32; + pos2 += (lhs_index >= rhs_index) as u32; + pos += (values[pos] != S::zero()) as usize; + } + for i in pos1..size1 { + indexes[pos] = self.indexes[i as usize]; + values[pos] = self.values[i as usize]; + pos += 1; + } + for i in pos2..size2 { + indexes[pos] = rhs.indexes[i as usize]; + values[pos] = rhs.values[i as usize]; + pos += 1; + } + indexes.truncate(pos); + values.truncate(pos); + SVectOwned::new(self.dims, indexes, values) + } + + fn operator_sub(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dims, rhs.dims); + let size1 = self.len(); + let size2 = rhs.len(); + let mut pos1 = 0; + let mut pos2 = 0; + let mut pos = 0; + let mut indexes = vec![0; (size1 + size2) as _]; + let mut values = vec![S::zero(); (size1 + size2) as _]; + while pos1 < size1 && pos2 < size2 { + let lhs_index = self.indexes[pos1 as usize]; + let rhs_index = rhs.indexes[pos2 as usize]; + let lhs_value = self.values[pos1 as usize]; + let rhs_value = rhs.values[pos2 as usize]; + indexes[pos] = lhs_index.min(rhs_index); + values[pos] = S::scalar_sub( + lhs_value.mask(lhs_index <= rhs_index), + rhs_value.mask(lhs_index >= rhs_index), + ); + pos1 += (lhs_index <= rhs_index) as u32; + pos2 += (lhs_index >= rhs_index) as u32; + pos += (values[pos] != S::zero()) as usize; + } + for i in pos1..size1 { + indexes[pos] = self.indexes[i as usize]; + values[pos] = self.values[i as usize]; + pos += 1; + } + for i in pos2..size2 { + indexes[pos] = rhs.indexes[i as usize]; + values[pos] = S::scalar_neg(rhs.values[i as usize]); + pos += 1; + } + indexes.truncate(pos); + values.truncate(pos); + SVectOwned::new(self.dims, indexes, values) + } + + fn operator_mul(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dims, rhs.dims); + let size1 = self.len(); + let size2 = rhs.len(); + let mut pos1 = 0; + let mut pos2 = 0; + let mut pos = 0; + let mut indexes = vec![0; std::cmp::min(size1, size2) as _]; + let mut values = vec![S::zero(); std::cmp::min(size1, size2) as _]; + while pos1 < size1 && pos2 < size2 { + let lhs_index = self.indexes[pos1 as usize]; + let rhs_index = rhs.indexes[pos2 as usize]; + match lhs_index.cmp(&rhs_index) { + std::cmp::Ordering::Less => { + pos1 += 1; + } + std::cmp::Ordering::Equal => { + // only both indexes are not zero, values are multiplied + let lhs_value = self.values[pos1 as usize]; + let rhs_value = rhs.values[pos2 as usize]; + indexes[pos] = lhs_index; + values[pos] = S::scalar_mul(lhs_value, rhs_value); + pos1 += 1; + pos2 += 1; + // only increment pos if the value is not zero + pos += (values[pos] != S::zero()) as usize; + } + std::cmp::Ordering::Greater => { + pos2 += 1; + } + } + } + indexes.truncate(pos); + values.truncate(pos); + SVectOwned::new(self.dims, indexes, values) + } + + fn operator_and(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_or(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_xor(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + #[inline(always)] + fn subvector(&self, bounds: impl RangeBounds) -> Option { + let start = match bounds.start_bound().cloned() { + Bound::Included(x) => x, + Bound::Excluded(u32::MAX) => return None, + Bound::Excluded(x) => x + 1, + Bound::Unbounded => 0, + }; + let end = match bounds.end_bound().cloned() { + Bound::Included(u32::MAX) => return None, + Bound::Included(x) => x + 1, + Bound::Excluded(x) => x, + Bound::Unbounded => self.dims, + }; + if start >= end || end > self.dims { + return None; + } + let dims = end - start; + let s = self.indexes.partition_point(|&x| x < start); + let e = self.indexes.partition_point(|&x| x < end); + let indexes = self.indexes[s..e] + .iter() + .map(|x| x - start) + .collect::>(); + let values = self.values[s..e].to_vec(); + Self::Owned::new_checked(dims, indexes, values) + } +} + +impl PartialEq for SVectBorrowed<'_, S> { + fn eq(&self, other: &Self) -> bool { + if self.dims != other.dims { + return false; + } + if self.indexes.len() != other.indexes.len() { + return false; + } + for (&l, &r) in self.indexes.iter().zip(other.indexes.iter()) { + if l != r { + return false; + } + } + for (&l, &r) in self.values.iter().zip(other.values.iter()) { + if l != r { + return false; + } + } + true + } +} + +impl PartialOrd for SVectBorrowed<'_, S> { + fn partial_cmp(&self, other: &Self) -> Option { + use std::cmp::Ordering; + if self.dims != other.dims { + return None; + } + let mut lhs = self + .indexes + .iter() + .copied() + .zip(self.values.iter().copied()); + let mut rhs = other + .indexes + .iter() + .copied() + .zip(other.values.iter().copied()); + loop { + return match (lhs.next(), rhs.next()) { + (Some(lh), Some(rh)) => match lh.0.cmp(&rh.0) { + Ordering::Equal => match lh.1.partial_cmp(&rh.1)? { + Ordering::Equal => continue, + x => Some(x), + }, + Ordering::Less => Some(if lh.1 < S::zero() { + Ordering::Less + } else { + Ordering::Greater + }), + Ordering::Greater => Some(if S::zero() < rh.1 { + Ordering::Less + } else { + Ordering::Greater + }), + }, + (Some((_, x)), None) => Some(PartialOrd::partial_cmp(&x, &S::zero())?), + (None, Some((_, y))) => Some(PartialOrd::partial_cmp(&S::zero(), &y)?), + (None, None) => Some(Ordering::Equal), + }; + } + } +} diff --git a/crates/vector/src/vect.rs b/crates/vector/src/vect.rs new file mode 100644 index 00000000..34b186f7 --- /dev/null +++ b/crates/vector/src/vect.rs @@ -0,0 +1,216 @@ +use super::{VectorBorrowed, VectorOwned}; +use distance::Distance; +use serde::{Deserialize, Serialize}; +use simd::Floating; +use std::cmp::Ordering; +use std::ops::RangeBounds; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[repr(transparent)] +pub struct VectOwned(Vec); + +impl VectOwned { + #[inline(always)] + pub fn new(slice: Vec) -> Self { + Self::new_checked(slice).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(slice: Vec) -> Option { + if !(1..=65535).contains(&slice.len()) { + return None; + } + Some(unsafe { Self::new_unchecked(slice) }) + } + + /// # Safety + /// + /// * `slice.len()` must not be zero. + #[inline(always)] + pub unsafe fn new_unchecked(slice: Vec) -> Self { + Self(slice) + } + + #[inline(always)] + pub fn slice(&self) -> &[S] { + self.0.as_slice() + } + + #[inline(always)] + pub fn slice_mut(&mut self) -> &mut [S] { + self.0.as_mut_slice() + } + + #[inline(always)] + pub fn into_vec(self) -> Vec { + self.0 + } +} + +impl VectorOwned for VectOwned { + type Borrowed<'a> = VectBorrowed<'a, S>; + + #[inline(always)] + fn as_borrowed(&self) -> VectBorrowed<'_, S> { + VectBorrowed(self.0.as_slice()) + } + + #[inline(always)] + fn zero(dims: u32) -> Self { + Self::new(vec![S::zero(); dims as usize]) + } +} + +#[derive(Debug, Clone, Copy)] +#[repr(transparent)] +pub struct VectBorrowed<'a, S>(&'a [S]); + +impl<'a, S: Floating> VectBorrowed<'a, S> { + #[inline(always)] + pub fn new(slice: &'a [S]) -> Self { + Self::new_checked(slice).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(slice: &'a [S]) -> Option { + if !(1..=65535).contains(&slice.len()) { + return None; + } + Some(unsafe { Self::new_unchecked(slice) }) + } + + /// # Safety + /// + /// * `slice.len()` must not be zero. + #[inline(always)] + pub unsafe fn new_unchecked(slice: &'a [S]) -> Self { + Self(slice) + } + + #[inline(always)] + pub fn slice(&self) -> &'a [S] { + self.0 + } +} + +impl VectorBorrowed for VectBorrowed<'_, S> { + type Owned = VectOwned; + + #[inline(always)] + fn dims(&self) -> u32 { + self.0.len() as u32 + } + + #[inline(always)] + fn own(&self) -> VectOwned { + VectOwned(self.0.to_vec()) + } + + #[inline(always)] + fn norm(&self) -> f32 { + S::reduce_sum_of_x2(self.0).sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + Distance::from(-S::reduce_sum_of_xy(self.slice(), rhs.slice())) + } + + #[inline(always)] + fn operator_l2(self, rhs: Self) -> Distance { + Distance::from(S::reduce_sum_of_d2(self.slice(), rhs.slice())) + } + + #[inline(always)] + fn operator_cos(self, rhs: Self) -> Distance { + let xy = S::reduce_sum_of_xy(self.slice(), rhs.slice()); + let x2 = S::reduce_sum_of_x2(self.0); + let y2 = S::reduce_sum_of_x2(rhs.0); + Distance::from(1.0 - xy / (x2 * y2).sqrt()) + } + + #[inline(always)] + fn operator_hamming(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_jaccard(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn function_normalize(&self) -> VectOwned { + let mut data = self.0.to_vec(); + let l = S::reduce_sum_of_x2(&data).sqrt(); + S::vector_mul_scalar_inplace(&mut data, 1.0 / l); + VectOwned(data) + } + + fn operator_add(&self, rhs: Self) -> Self::Owned { + VectOwned::new(S::vector_add(self.slice(), rhs.slice())) + } + + fn operator_sub(&self, rhs: Self) -> Self::Owned { + VectOwned::new(S::vector_sub(self.slice(), rhs.slice())) + } + + fn operator_mul(&self, rhs: Self) -> Self::Owned { + VectOwned::new(S::vector_mul(self.slice(), rhs.slice())) + } + + fn operator_and(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_or(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_xor(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + #[inline(always)] + fn subvector(&self, bounds: impl RangeBounds) -> Option { + let start_bound = bounds.start_bound().map(|x| *x as usize); + let end_bound = bounds.end_bound().map(|x| *x as usize); + let slice = self.0.get((start_bound, end_bound))?; + if slice.is_empty() { + return None; + } + Self::Owned::new_checked(slice.to_vec()) + } +} + +impl PartialEq for VectBorrowed<'_, S> { + fn eq(&self, other: &Self) -> bool { + if self.0.len() != other.0.len() { + return false; + } + let n = self.0.len(); + for i in 0..n { + if self.0[i] != other.0[i] { + return false; + } + } + true + } +} + +impl PartialOrd for VectBorrowed<'_, S> { + fn partial_cmp(&self, other: &Self) -> Option { + if self.0.len() != other.0.len() { + return None; + } + let n = self.0.len(); + for i in 0..n { + match PartialOrd::partial_cmp(&self.0[i], &other.0[i])? { + Ordering::Less => return Some(Ordering::Less), + Ordering::Equal => continue, + Ordering::Greater => return Some(Ordering::Greater), + } + } + Some(Ordering::Equal) + } +} diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 00000000..35011368 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1 @@ +style_edition = "2024" diff --git a/scripts/README.md b/scripts/README.md index 938ae900..9ed71720 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -10,8 +10,8 @@ export PGRX_VERSION=$(awk -F'version = "=|"' '/^pgrx\s*=.*version/ {print $2}' C export RUST_TOOLCHAIN=$(awk -F'"' '/^\s*channel\s*=/ {print $2}' rust-toolchain.toml) export PGRX_IMAGE=ghcr.io/tensorchord/vectorchord-pgrx:$PGRX_VERSION-$RUST_TOOLCHAIN -docker run --rm -v .:/workspace $PGRX_IMAGE cargo build --lib --features pg16 --profile opt -docker run --rm -v .:/workspace $PGRX_IMAGE ./tools/schema.sh --features pg16 --profile opt +docker run --rm -v .:/workspace $PGRX_IMAGE cargo build --lib --features pg16 --release +docker run --rm -v .:/workspace $PGRX_IMAGE ./tools/schema.sh --features pg16 --release ``` - (option 2) With Local Development Environment @@ -20,8 +20,8 @@ docker run --rm -v .:/workspace $PGRX_IMAGE ./tools/schema.sh --features pg16 -- sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config clang cargo install --locked cargo-pgrx cargo pgrx init -cargo build --package vchord --lib --features pg16 --profile opt -./tools/schema.sh --features pg16 --profile opt +cargo build --package vchord --lib --features pg16 --release +./tools/schema.sh --features pg16 --release ``` - build the debian package @@ -31,7 +31,6 @@ export SEMVER="0.0.0" export VERSION="16" export ARCH="x86_64" export PLATFORM="amd64" -export PROFILE="opt" ./tools/package.sh ``` diff --git a/src/bin/pgrx_embed.rs b/src/bin/pgrx_embed.rs index 20a006e8..5f5c4d85 100644 --- a/src/bin/pgrx_embed.rs +++ b/src/bin/pgrx_embed.rs @@ -1,2 +1 @@ -#![feature(strict_provenance_lints)] ::pgrx::pgrx_embed!(); diff --git a/src/datatype/binary_scalar8.rs b/src/datatype/binary_scalar8.rs index 4dfe844b..efb8758d 100644 --- a/src/datatype/binary_scalar8.rs +++ b/src/datatype/binary_scalar8.rs @@ -1,8 +1,8 @@ use super::memory_scalar8::{Scalar8Input, Scalar8Output}; -use crate::types::scalar8::Scalar8Borrowed; -use base::vector::VectorBorrowed; use pgrx::datum::Internal; use pgrx::pg_sys::Oid; +use vector::VectorBorrowed; +use vector::scalar8::Scalar8Borrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_scalar8_send(vector: Scalar8Input<'_>) -> Vec { diff --git a/src/datatype/functions_scalar8.rs b/src/datatype/functions_scalar8.rs index fc1c221c..adcdb30c 100644 --- a/src/datatype/functions_scalar8.rs +++ b/src/datatype/functions_scalar8.rs @@ -1,17 +1,17 @@ use crate::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput; use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; use crate::datatype::memory_scalar8::Scalar8Output; -use crate::types::scalar8::Scalar8Borrowed; -use base::simd::ScalarLike; use half::f16; +use simd::Floating; +use vector::scalar8::Scalar8Borrowed; #[pgrx::pg_extern(sql = "")] fn _vchord_vector_quantize_to_scalar8(vector: PgvectorVectorInput) -> Scalar8Output { let vector = vector.as_borrowed(); let sum_of_x2 = f32::reduce_sum_of_x2(vector.slice()); let (k, b, code) = - base::simd::quantize::quantize(f32::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); - let sum_of_code = base::simd::u8::reduce_sum_of_x_as_u32(&code) as f32; + simd::quantize::quantize(f32::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); + let sum_of_code = simd::u8::reduce_sum_of_x(&code) as f32; Scalar8Output::new(Scalar8Borrowed::new(sum_of_x2, k, b, sum_of_code, &code)) } @@ -20,7 +20,7 @@ fn _vchord_halfvec_quantize_to_scalar8(vector: PgvectorHalfvecInput) -> Scalar8O let vector = vector.as_borrowed(); let sum_of_x2 = f16::reduce_sum_of_x2(vector.slice()); let (k, b, code) = - base::simd::quantize::quantize(f16::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); - let sum_of_code = base::simd::u8::reduce_sum_of_x_as_u32(&code) as f32; + simd::quantize::quantize(f16::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); + let sum_of_code = simd::u8::reduce_sum_of_x(&code) as f32; Scalar8Output::new(Scalar8Borrowed::new(sum_of_x2, k, b, sum_of_code, &code)) } diff --git a/src/datatype/memory_pgvector_halfvec.rs b/src/datatype/memory_pgvector_halfvec.rs index 7265a1be..1c065d94 100644 --- a/src/datatype/memory_pgvector_halfvec.rs +++ b/src/datatype/memory_pgvector_halfvec.rs @@ -1,4 +1,3 @@ -use base::vector::*; use half::f16; use pgrx::datum::FromDatum; use pgrx::datum::IntoDatum; @@ -11,6 +10,8 @@ use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; use std::ops::Deref; use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; #[repr(C, align(8))] pub struct PgvectorHalfvecHeader { diff --git a/src/datatype/memory_pgvector_vector.rs b/src/datatype/memory_pgvector_vector.rs index d81492cf..e3ab9f90 100644 --- a/src/datatype/memory_pgvector_vector.rs +++ b/src/datatype/memory_pgvector_vector.rs @@ -1,4 +1,3 @@ -use base::vector::*; use pgrx::datum::FromDatum; use pgrx::datum::IntoDatum; use pgrx::pg_sys::Datum; @@ -10,6 +9,8 @@ use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; use std::ops::Deref; use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; #[repr(C, align(8))] pub struct PgvectorVectorHeader { diff --git a/src/datatype/memory_scalar8.rs b/src/datatype/memory_scalar8.rs index 3a7dab4a..1641c63d 100644 --- a/src/datatype/memory_scalar8.rs +++ b/src/datatype/memory_scalar8.rs @@ -1,5 +1,3 @@ -use crate::types::scalar8::Scalar8Borrowed; -use base::vector::*; use pgrx::datum::FromDatum; use pgrx::datum::IntoDatum; use pgrx::pg_sys::Datum; @@ -11,6 +9,8 @@ use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; use std::ops::Deref; use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::scalar8::Scalar8Borrowed; #[repr(C, align(8))] pub struct Scalar8Header { diff --git a/src/datatype/operators_pgvector_halfvec.rs b/src/datatype/operators_pgvector_halfvec.rs index 2e707ebf..fb0492ac 100644 --- a/src/datatype/operators_pgvector_halfvec.rs +++ b/src/datatype/operators_pgvector_halfvec.rs @@ -1,6 +1,7 @@ -use crate::datatype::memory_pgvector_halfvec::*; -use base::vector::{VectBorrowed, VectorBorrowed}; +use crate::datatype::memory_pgvector_halfvec::{PgvectorHalfvecInput, PgvectorHalfvecOutput}; use std::num::NonZero; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_halfvec_sphere_l2_in( diff --git a/src/datatype/operators_pgvector_vector.rs b/src/datatype/operators_pgvector_vector.rs index 2308ab9d..e6b4be10 100644 --- a/src/datatype/operators_pgvector_vector.rs +++ b/src/datatype/operators_pgvector_vector.rs @@ -1,6 +1,7 @@ -use crate::datatype::memory_pgvector_vector::*; -use base::vector::{VectBorrowed, VectorBorrowed}; +use crate::datatype::memory_pgvector_vector::{PgvectorVectorInput, PgvectorVectorOutput}; use std::num::NonZero; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_vector_sphere_l2_in( diff --git a/src/datatype/operators_scalar8.rs b/src/datatype/operators_scalar8.rs index db6a3727..8b75e592 100644 --- a/src/datatype/operators_scalar8.rs +++ b/src/datatype/operators_scalar8.rs @@ -1,7 +1,7 @@ use crate::datatype::memory_scalar8::{Scalar8Input, Scalar8Output}; -use crate::types::scalar8::Scalar8Borrowed; -use base::vector::*; use std::num::NonZero; +use vector::VectorBorrowed; +use vector::scalar8::Scalar8Borrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_scalar8_operator_ip(lhs: Scalar8Input<'_>, rhs: Scalar8Input<'_>) -> f32 { diff --git a/src/datatype/text_scalar8.rs b/src/datatype/text_scalar8.rs index 5de82da6..8669a015 100644 --- a/src/datatype/text_scalar8.rs +++ b/src/datatype/text_scalar8.rs @@ -1,8 +1,8 @@ use super::memory_scalar8::Scalar8Output; use crate::datatype::memory_scalar8::Scalar8Input; -use crate::types::scalar8::Scalar8Borrowed; use pgrx::pg_sys::Oid; use std::ffi::{CStr, CString}; +use vector::scalar8::Scalar8Borrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_scalar8_in(input: &CStr, oid: Oid, typmod: i32) -> Scalar8Output { diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index fe90a6d2..08ac3754 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -19,7 +19,6 @@ impl Typmod { None } } - #[allow(dead_code)] pub fn into_option_string(self) -> Option { use Typmod::*; match self { @@ -27,7 +26,6 @@ impl Typmod { Dims(x) => Some(x.get().to_string()), } } - #[allow(dead_code)] pub fn into_i32(self) -> i32 { use Typmod::*; match self { diff --git a/src/lib.rs b/src/lib.rs index 26d80ed0..018388d8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,17 +1,11 @@ #![allow(clippy::collapsible_else_if)] -#![allow(clippy::identity_op)] -#![allow(clippy::needless_range_loop)] +#![allow(clippy::infallible_destructuring_match)] #![allow(clippy::too_many_arguments)] #![allow(clippy::type_complexity)] -#![allow(clippy::int_plus_one)] -#![allow(clippy::unused_unit)] -#![allow(clippy::infallible_destructuring_match)] -#![feature(strict_provenance_lints)] mod datatype; mod postgres; mod projection; -mod types; mod upgrade; mod utils; mod vchordrq; diff --git a/src/postgres.rs b/src/postgres.rs index c24d669e..6ff91e23 100644 --- a/src/postgres.rs +++ b/src/postgres.rs @@ -1,4 +1,5 @@ -use std::mem::{offset_of, MaybeUninit}; +use algorithm::{Opaque, Page, PageGuard, RelationRead, RelationWrite}; +use std::mem::{MaybeUninit, offset_of}; use std::ops::{Deref, DerefMut}; use std::ptr::NonNull; @@ -7,7 +8,7 @@ const _: () = assert!( ); const fn size_of_contents() -> usize { - use pgrx::pg_sys::{PageHeaderData, BLCKSZ}; + use pgrx::pg_sys::{BLCKSZ, PageHeaderData}; let size_of_page = BLCKSZ as usize; let size_of_header = offset_of!(PageHeaderData, pd_linp); let size_of_opaque = size_of::(); @@ -15,17 +16,17 @@ const fn size_of_contents() -> usize { } #[repr(C, align(8))] -pub struct Page { +pub struct PostgresPage { header: pgrx::pg_sys::PageHeaderData, content: [u8; size_of_contents()], opaque: Opaque, } -const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); -const _: () = assert!(size_of::() == pgrx::pg_sys::BLCKSZ as usize); +const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); +const _: () = assert!(size_of::() == pgrx::pg_sys::BLCKSZ as usize); -impl Page { - pub fn init_mut(this: &mut MaybeUninit) -> &mut Self { +impl PostgresPage { + fn init_mut(this: &mut MaybeUninit) -> &mut Self { unsafe { pgrx::pg_sys::PageInit( this.as_mut_ptr() as pgrx::pg_sys::Page, @@ -42,26 +43,29 @@ impl Page { this } #[allow(dead_code)] - pub unsafe fn assume_init_mut(this: &mut MaybeUninit) -> &mut Self { + unsafe fn assume_init_mut(this: &mut MaybeUninit) -> &mut Self { let this = unsafe { MaybeUninit::assume_init_mut(this) }; assert_eq!(offset_of!(Self, opaque), this.header.pd_special as usize); this } #[allow(dead_code)] - pub fn clone_into_boxed(&self) -> Box { + fn clone_into_boxed(&self) -> Box { let mut result = Box::new_uninit(); unsafe { std::ptr::copy(self as *const Self, result.as_mut_ptr(), 1); result.assume_init() } } - pub fn get_opaque(&self) -> &Opaque { +} + +impl Page for PostgresPage { + fn get_opaque(&self) -> &Opaque { &self.opaque } - pub fn get_opaque_mut(&mut self) -> &mut Opaque { + fn get_opaque_mut(&mut self) -> &mut Opaque { &mut self.opaque } - pub fn len(&self) -> u16 { + fn len(&self) -> u16 { use pgrx::pg_sys::{ItemIdData, PageHeaderData}; assert!(self.header.pd_lower as usize <= size_of::()); assert!(self.header.pd_upper as usize <= size_of::()); @@ -70,7 +74,7 @@ impl Page { assert!(lower <= upper); ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16 } - pub fn get(&self, i: u16) -> Option<&[u8]> { + fn get(&self, i: u16) -> Option<&[u8]> { use pgrx::pg_sys::{ItemIdData, PageHeaderData}; if i == 0 { return None; @@ -99,8 +103,7 @@ impl Page { Some(std::slice::from_raw_parts(ptr, lp_len as _)) } } - #[allow(unused)] - pub fn get_mut(&mut self, i: u16) -> Option<&mut [u8]> { + fn get_mut(&mut self, i: u16) -> Option<&mut [u8]> { use pgrx::pg_sys::{ItemIdData, PageHeaderData}; if i == 0 { return None; @@ -122,7 +125,7 @@ impl Page { Some(std::slice::from_raw_parts_mut(ptr, lp_len as _)) } } - pub fn alloc(&mut self, data: &[u8]) -> Option { + fn alloc(&mut self, data: &[u8]) -> Option { unsafe { let i = pgrx::pg_sys::PageAddItemExtended( (self as *const Self).cast_mut().cast(), @@ -131,19 +134,15 @@ impl Page { 0, 0, ); - if i == 0 { - None - } else { - Some(i) - } + if i == 0 { None } else { Some(i) } } } - pub fn free(&mut self, i: u16) { + fn free(&mut self, i: u16) { unsafe { pgrx::pg_sys::PageIndexTupleDeleteNoCompact((self as *mut Self).cast(), i); } } - pub fn reconstruct(&mut self, removes: &[u16]) { + fn reconstruct(&mut self, removes: &[u16]) { let mut removes = removes.to_vec(); removes.sort(); removes.dedup(); @@ -159,41 +158,34 @@ impl Page { } } } - pub fn freespace(&self) -> u16 { + fn freespace(&self) -> u16 { unsafe { pgrx::pg_sys::PageGetFreeSpace((self as *const Self).cast_mut().cast()) as u16 } } } -#[repr(C, align(8))] -pub struct Opaque { - pub next: u32, - pub skip: u32, -} - const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); -pub struct BufferReadGuard { +pub struct PostgresBufferReadGuard { buf: i32, - page: NonNull, + page: NonNull, id: u32, } -impl BufferReadGuard { - #[allow(dead_code)] - pub fn id(&self) -> u32 { +impl PageGuard for PostgresBufferReadGuard { + fn id(&self) -> u32 { self.id } } -impl Deref for BufferReadGuard { - type Target = Page; +impl Deref for PostgresBufferReadGuard { + type Target = PostgresPage; - fn deref(&self) -> &Page { + fn deref(&self) -> &PostgresPage { unsafe { self.page.as_ref() } } } -impl Drop for BufferReadGuard { +impl Drop for PostgresBufferReadGuard { fn drop(&mut self) { unsafe { pgrx::pg_sys::UnlockReleaseBuffer(self.buf); @@ -201,36 +193,36 @@ impl Drop for BufferReadGuard { } } -pub struct BufferWriteGuard { +pub struct PostgresBufferWriteGuard { raw: pgrx::pg_sys::Relation, buf: i32, - page: NonNull, + page: NonNull, state: *mut pgrx::pg_sys::GenericXLogState, id: u32, tracking_freespace: bool, } -impl BufferWriteGuard { - pub fn id(&self) -> u32 { +impl PageGuard for PostgresBufferWriteGuard { + fn id(&self) -> u32 { self.id } } -impl Deref for BufferWriteGuard { - type Target = Page; +impl Deref for PostgresBufferWriteGuard { + type Target = PostgresPage; - fn deref(&self) -> &Page { + fn deref(&self) -> &PostgresPage { unsafe { self.page.as_ref() } } } -impl DerefMut for BufferWriteGuard { - fn deref_mut(&mut self) -> &mut Page { +impl DerefMut for PostgresBufferWriteGuard { + fn deref_mut(&mut self) -> &mut PostgresPage { unsafe { self.page.as_mut() } } } -impl Drop for BufferWriteGuard { +impl Drop for PostgresBufferWriteGuard { fn drop(&mut self) { unsafe { if std::thread::panicking() { @@ -248,19 +240,36 @@ impl Drop for BufferWriteGuard { } #[derive(Debug, Clone)] -pub struct Relation { +pub struct PostgresRelation { raw: pgrx::pg_sys::Relation, } -impl Relation { +impl PostgresRelation { pub unsafe fn new(raw: pgrx::pg_sys::Relation) -> Self { Self { raw } } - pub fn read(&self, id: u32) -> BufferReadGuard { + + #[allow(dead_code)] + pub fn len(&self) -> u32 { + unsafe { + pgrx::pg_sys::RelationGetNumberOfBlocksInFork( + self.raw, + pgrx::pg_sys::ForkNumber::MAIN_FORKNUM, + ) + } + } +} + +impl RelationRead for PostgresRelation { + type Page = PostgresPage; + + type ReadGuard<'a> = PostgresBufferReadGuard; + + fn read(&self, id: u32) -> Self::ReadGuard<'_> { assert!(id != u32::MAX, "no such page"); unsafe { use pgrx::pg_sys::{ - BufferGetPage, LockBuffer, ReadBufferExtended, ReadBufferMode, BUFFER_LOCK_SHARE, + BUFFER_LOCK_SHARE, BufferGetPage, LockBuffer, ReadBufferExtended, ReadBufferMode, }; let buf = ReadBufferExtended( self.raw, @@ -271,15 +280,21 @@ impl Relation { ); LockBuffer(buf, BUFFER_LOCK_SHARE as _); let page = NonNull::new(BufferGetPage(buf).cast()).expect("failed to get page"); - BufferReadGuard { buf, page, id } + PostgresBufferReadGuard { buf, page, id } } } - pub fn write(&self, id: u32, tracking_freespace: bool) -> BufferWriteGuard { +} + +impl RelationWrite for PostgresRelation { + type WriteGuard<'a> = PostgresBufferWriteGuard; + + fn write(&self, id: u32, tracking_freespace: bool) -> PostgresBufferWriteGuard { assert!(id != u32::MAX, "no such page"); unsafe { use pgrx::pg_sys::{ - ForkNumber, GenericXLogRegisterBuffer, GenericXLogStart, LockBuffer, - ReadBufferExtended, ReadBufferMode, BUFFER_LOCK_EXCLUSIVE, GENERIC_XLOG_FULL_IMAGE, + BUFFER_LOCK_EXCLUSIVE, ForkNumber, GENERIC_XLOG_FULL_IMAGE, + GenericXLogRegisterBuffer, GenericXLogStart, LockBuffer, ReadBufferExtended, + ReadBufferMode, }; let buf = ReadBufferExtended( self.raw, @@ -292,10 +307,10 @@ impl Relation { let state = GenericXLogStart(self.raw); let page = NonNull::new( GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) - .cast::>(), + .cast::>(), ) .expect("failed to get page"); - BufferWriteGuard { + PostgresBufferWriteGuard { raw: self.raw, buf, page: page.cast(), @@ -305,12 +320,12 @@ impl Relation { } } } - pub fn extend(&self, tracking_freespace: bool) -> BufferWriteGuard { + fn extend(&self, tracking_freespace: bool) -> PostgresBufferWriteGuard { unsafe { use pgrx::pg_sys::{ - ExclusiveLock, ForkNumber, GenericXLogRegisterBuffer, GenericXLogStart, LockBuffer, - LockRelationForExtension, ReadBufferExtended, ReadBufferMode, - UnlockRelationForExtension, BUFFER_LOCK_EXCLUSIVE, GENERIC_XLOG_FULL_IMAGE, + BUFFER_LOCK_EXCLUSIVE, ExclusiveLock, ForkNumber, GENERIC_XLOG_FULL_IMAGE, + GenericXLogRegisterBuffer, GenericXLogStart, LockBuffer, LockRelationForExtension, + ReadBufferExtended, ReadBufferMode, UnlockRelationForExtension, }; LockRelationForExtension(self.raw, ExclusiveLock as _); let buf = ReadBufferExtended( @@ -325,11 +340,11 @@ impl Relation { let state = GenericXLogStart(self.raw); let mut page = NonNull::new( GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) - .cast::>(), + .cast::>(), ) .expect("failed to get page"); - Page::init_mut(page.as_mut()); - BufferWriteGuard { + PostgresPage::init_mut(page.as_mut()); + PostgresBufferWriteGuard { raw: self.raw, buf, page: page.cast(), @@ -339,7 +354,7 @@ impl Relation { } } } - pub fn search(&self, freespace: usize) -> Option { + fn search(&self, freespace: usize) -> Option { unsafe { loop { let id = pgrx::pg_sys::GetPageWithFreeSpace(self.raw, freespace); @@ -357,13 +372,4 @@ impl Relation { } } } - #[allow(dead_code)] - pub fn len(&self) -> u32 { - unsafe { - pgrx::pg_sys::RelationGetNumberOfBlocksInFork( - self.raw, - pgrx::pg_sys::ForkNumber::MAIN_FORKNUM, - ) - } - } } diff --git a/src/projection.rs b/src/projection.rs index 4273180a..fbcaeffa 100644 --- a/src/projection.rs +++ b/src/projection.rs @@ -1,75 +1,21 @@ -use nalgebra::DMatrix; +use random_orthogonal_matrix::random_orthogonal_matrix; use std::sync::OnceLock; -fn random_matrix(n: usize) -> DMatrix { - use rand::{Rng, SeedableRng}; - use rand_chacha::ChaCha12Rng; - use rand_distr::StandardNormal; - let mut rng = ChaCha12Rng::from_seed([7; 32]); - DMatrix::from_fn(n, n, |_, _| rng.sample(StandardNormal)) +fn matrix(n: usize) -> Option<&'static Vec>> { + static MATRIXS: [OnceLock>>; 1 + 60000] = [const { OnceLock::new() }; 1 + 60000]; + MATRIXS + .get(n) + .map(|x| x.get_or_init(|| random_orthogonal_matrix(n))) } -#[ignore] -#[test] -fn check_all_matrixs_are_full_rank() { - let parallelism = std::thread::available_parallelism().unwrap().get(); - std::thread::scope(|scope| { - let mut threads = vec![]; - for remainder in 0..parallelism { - threads.push(scope.spawn(move || { - for n in (0..=60000).filter(|x| x % parallelism == remainder) { - let matrix = random_matrix(n); - assert!(matrix.is_invertible()); - } - })); - } - for thread in threads { - thread.join().unwrap(); - } - }); -} - -#[test] -fn check_matrices() { - assert_eq!( - orthogonal_matrix(2), - vec![vec![-0.5424608, -0.8400813], vec![0.8400813, -0.54246056]] - ); - assert_eq!( - orthogonal_matrix(3), - vec![ - vec![-0.5309615, -0.69094884, -0.49058124], - vec![0.8222731, -0.56002235, -0.10120347], - vec![0.20481002, 0.45712686, -0.86549866] - ] - ); -} - -fn orthogonal_matrix(n: usize) -> Vec> { - use nalgebra::QR; - let matrix = random_matrix(n); - // QR decomposition is unique if the matrix is full rank - let qr = QR::new(matrix); - let q = qr.q(); - let mut projection = Vec::new(); - for row in q.row_iter() { - projection.push(row.iter().copied().collect::>()); - } - projection -} - -static MATRIXS: [OnceLock>>; 1 + 60000] = [const { OnceLock::new() }; 1 + 60000]; - pub fn prewarm(n: usize) { - if n <= 60000 { - MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); - } + let _ = matrix(n); } pub fn project(vector: &[f32]) -> Vec { - use base::simd::ScalarLike; + use simd::Floating; let n = vector.len(); - let matrix = MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); + let matrix = matrix(n).expect("dimension too large"); (0..n) .map(|i| f32::reduce_sum_of_xy(vector, &matrix[i])) .collect() diff --git a/src/types/mod.rs b/src/types/mod.rs deleted file mode 100644 index af08ee7f..00000000 --- a/src/types/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod scalar8; diff --git a/src/utils/k_means.rs b/src/utils/k_means.rs index 8aeac72a..7b44a24c 100644 --- a/src/utils/k_means.rs +++ b/src/utils/k_means.rs @@ -1,16 +1,14 @@ -#![allow(clippy::ptr_arg)] - use super::parallelism::{ParallelIterator, Parallelism}; -use base::simd::*; use half::f16; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; +use simd::Floating; pub fn k_means( parallelism: &P, c: usize, dims: usize, - samples: &Vec>, + samples: &[Vec], is_spherical: bool, iterations: usize, ) -> Vec> { @@ -18,9 +16,9 @@ pub fn k_means( assert!(dims > 0); let n = samples.len(); if n <= c { - quick_centers(c, dims, samples.clone(), is_spherical) + quick_centers(c, dims, samples.to_vec(), is_spherical) } else { - let compute = |parallelism: &P, centroids: &Vec>| { + let compute = |parallelism: &P, centroids: &[Vec]| { if n >= 1000 && c >= 1000 { rabitq_index(parallelism, dims, n, c, samples, centroids) } else { @@ -62,9 +60,7 @@ fn quick_centers( let mut rng = rand::thread_rng(); let mut centroids = samples; for _ in n..c { - let r = (0..dims) - .map(|_| f32::from_f32(rng.gen_range(-1.0f32..1.0f32))) - .collect(); + let r = (0..dims).map(|_| rng.gen_range(-1.0f32..1.0f32)).collect(); centroids.push(r); } if is_spherical { @@ -82,152 +78,37 @@ fn rabitq_index( dims: usize, n: usize, c: usize, - samples: &Vec>, - centroids: &Vec>, + samples: &[Vec], + centroids: &[Vec], ) -> Vec { - fn code_alpha(vector: &[f32]) -> (f32, f32, f32, f32) { - let dims = vector.len(); - let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); - let sum_of_x2 = f32::reduce_sum_of_x2(vector); - let dis_u = sum_of_x2.sqrt(); - let x0 = sum_of_abs_x / (sum_of_x2 * (dims as f32)).sqrt(); - let x_x0 = dis_u / x0; - let fac_norm = (dims as f32).sqrt(); - let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); - let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); - let factor_ip = -2.0f32 / fac_norm * x_x0; - let cnt_pos = vector - .iter() - .map(|x| x.scalar_is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.scalar_is_sign_negative() as i32) - .sum::(); - let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; - (sum_of_x2, factor_ppc, factor_ip, factor_err) - } - fn code_beta(vector: &[f32]) -> Vec { - let dims = vector.len(); - let mut code = Vec::new(); - for i in 0..dims { - code.push(vector[i].scalar_is_sign_positive() as u8); - } - code - } let mut a0 = Vec::new(); let mut a1 = Vec::new(); let mut a2 = Vec::new(); let mut a3 = Vec::new(); let mut a4 = Vec::new(); for vectors in centroids.chunks(32) { - use base::simd::fast_scan::b4::pack; - let code_alphas = std::array::from_fn::<_, 32, _>(|i| { + use simd::fast_scan::pack; + let x = std::array::from_fn::<_, 32, _>(|i| { if let Some(vector) = vectors.get(i) { - code_alpha(vector) + rabitq::block::code(dims as _, vector) } else { - (0.0, 0.0, 0.0, 0.0) - } - }); - let code_betas = std::array::from_fn::<_, 32, _>(|i| { - let mut result = vec![0_u8; dims.div_ceil(4)]; - if let Some(vector) = vectors.get(i) { - let mut c = code_beta(vector); - c.resize(dims.next_multiple_of(4), 0); - for i in 0..dims.div_ceil(4) { - for j in 0..4 { - result[i] |= c[i * 4 + j] << j; - } - } + rabitq::block::dummy_code(dims as _) } - result }); - a0.push(code_alphas.map(|x| x.0)); - a1.push(code_alphas.map(|x| x.1)); - a2.push(code_alphas.map(|x| x.2)); - a3.push(code_alphas.map(|x| x.3)); - a4.push(pack(dims.div_ceil(4) as _, code_betas).collect::>()); + a0.push(x.each_ref().map(|x| x.dis_u_2)); + a1.push(x.each_ref().map(|x| x.factor_ppc)); + a2.push(x.each_ref().map(|x| x.factor_ip)); + a3.push(x.each_ref().map(|x| x.factor_err)); + a4.push(pack(dims.div_ceil(4) as _, x.map(|x| x.signs)).collect::>()); } parallelism - .into_par_iter(0..n) + .rayon_into_par_iter(0..n) .map(|i| { - fn generate(mut qvector: Vec) -> Vec { - let dims = qvector.len() as u32; - let t = dims.div_ceil(4); - qvector.resize(qvector.len().next_multiple_of(4), 0); - let mut lut = vec![0u8; t as usize * 16]; - for i in 0..t as usize { - unsafe { - // this hint is used to skip bound checks - std::hint::assert_unchecked(4 * i + 3 < qvector.len()); - std::hint::assert_unchecked(16 * i + 15 < lut.len()); - } - let t0 = qvector[4 * i + 0]; - let t1 = qvector[4 * i + 1]; - let t2 = qvector[4 * i + 2]; - let t3 = qvector[4 * i + 3]; - lut[16 * i + 0b0000] = 0; - lut[16 * i + 0b0001] = t0; - lut[16 * i + 0b0010] = t1; - lut[16 * i + 0b0011] = t1 + t0; - lut[16 * i + 0b0100] = t2; - lut[16 * i + 0b0101] = t2 + t0; - lut[16 * i + 0b0110] = t2 + t1; - lut[16 * i + 0b0111] = t2 + t1 + t0; - lut[16 * i + 0b1000] = t3; - lut[16 * i + 0b1001] = t3 + t0; - lut[16 * i + 0b1010] = t3 + t1; - lut[16 * i + 0b1011] = t3 + t1 + t0; - lut[16 * i + 0b1100] = t3 + t2; - lut[16 * i + 0b1101] = t3 + t2 + t0; - lut[16 * i + 0b1110] = t3 + t2 + t1; - lut[16 * i + 0b1111] = t3 + t2 + t1 + t0; - } - lut - } - fn fscan_process_lowerbound( - dims: u32, - lut: &(f32, f32, f32, f32, Vec), - (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[u8], - ), - epsilon: f32, - ) -> [Distance; 32] { - use base::simd::fast_scan::b4::fast_scan_b4; - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let r = fast_scan_b4(dims.div_ceil(4), t, s); - std::array::from_fn(|i| { - let rough = dis_u_2[i] - + dis_v_2 - + b * factor_ppc[i] - + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) - }) - } - use base::distance::Distance; - use base::simd::quantize; - - let lut = { - let vector = &samples[i]; - let dis_v_2 = f32::reduce_sum_of_x2(vector); - let (k, b, qvector) = - quantize::quantize(f32::vector_to_f32_borrowed(vector).as_ref(), 15.0); - let qvector_sum = if vector.len() <= 4369 { - u8::reduce_sum_of_x_as_u16(&qvector) as f32 - } else { - u8::reduce_sum_of_x_as_u32(&qvector) as f32 - }; - (dis_v_2, b, k, qvector_sum, generate(qvector)) - }; - + use distance::Distance; + let lut = rabitq::block::fscan_preprocess(&samples[i]); let mut result = (Distance::INFINITY, 0); for block in 0..c.div_ceil(32) { - let lowerbound = fscan_process_lowerbound( + let lowerbound = rabitq::block::fscan_process_lowerbound_l2( dims as _, &lut, (&a0[block], &a1[block], &a2[block], &a3[block], &a4[block]), @@ -253,11 +134,11 @@ fn flat_index( _dims: usize, n: usize, c: usize, - samples: &Vec>, - centroids: &Vec>, + samples: &[Vec], + centroids: &[Vec], ) -> Vec { parallelism - .into_par_iter(0..n) + .rayon_into_par_iter(0..n) .map(|i| { let mut result = (f32::INFINITY, 0); for j in 0..c { @@ -279,18 +160,18 @@ struct LloydKMeans<'a, P, F> { centroids: Vec>, assign: Vec, rng: StdRng, - samples: &'a Vec>, + samples: &'a [Vec], compute: F, } const DELTA: f32 = f16::EPSILON.to_f32_const(); -impl<'a, P: Parallelism, F: Fn(&P, &Vec>) -> Vec> LloydKMeans<'a, P, F> { +impl<'a, P: Parallelism, F: Fn(&P, &[Vec]) -> Vec> LloydKMeans<'a, P, F> { fn new( parallelism: &'a P, c: usize, dims: usize, - samples: &'a Vec>, + samples: &'a [Vec], is_spherical: bool, compute: F, ) -> Self { @@ -304,7 +185,7 @@ impl<'a, P: Parallelism, F: Fn(&P, &Vec>) -> Vec> LloydKMeans<'a } let assign = parallelism - .into_par_iter(0..n) + .rayon_into_par_iter(0..n) .map(|i| { let mut result = (f32::INFINITY, 0); for j in 0..c { @@ -339,7 +220,7 @@ impl<'a, P: Parallelism, F: Fn(&P, &Vec>) -> Vec> LloydKMeans<'a let (sum, mut count) = self .parallelism - .into_par_iter(0..n) + .rayon_into_par_iter(0..n) .fold( || (vec![vec![f32::zero(); dims]; c], vec![0.0f32; c]), |(mut sum, mut count), i| { @@ -361,7 +242,7 @@ impl<'a, P: Parallelism, F: Fn(&P, &Vec>) -> Vec> LloydKMeans<'a let mut centroids = self .parallelism - .into_par_iter(0..c) + .rayon_into_par_iter(0..c) .map(|i| f32::vector_mul_scalar(&sum[i], 1.0 / count[i])) .collect::>(); @@ -379,15 +260,15 @@ impl<'a, P: Parallelism, F: Fn(&P, &Vec>) -> Vec> LloydKMeans<'a o = (o + 1) % c; } centroids[i] = centroids[o].clone(); - f32::kmeans_helper(&mut centroids[i], 1.0 + DELTA, 1.0 - DELTA); - f32::kmeans_helper(&mut centroids[o], 1.0 - DELTA, 1.0 + DELTA); + vector_mul_scalars_inplace(&mut centroids[i], [1.0 + DELTA, 1.0 - DELTA]); + vector_mul_scalars_inplace(&mut centroids[o], [1.0 - DELTA, 1.0 + DELTA]); count[i] = count[o] / 2.0; count[o] -= count[i]; } if self.is_spherical { self.parallelism - .into_par_iter(&mut centroids) + .rayon_into_par_iter(&mut centroids) .for_each(|centroid| { let l = f32::reduce_sum_of_x2(centroid).sqrt(); f32::vector_mul_scalar_inplace(centroid, 1.0 / l); @@ -408,3 +289,14 @@ impl<'a, P: Parallelism, F: Fn(&P, &Vec>) -> Vec> LloydKMeans<'a self.centroids } } + +fn vector_mul_scalars_inplace(this: &mut [f32], scalars: [f32; 2]) { + let n: usize = this.len(); + for i in 0..n { + if i % 2 == 0 { + this[i] *= scalars[0]; + } else { + this[i] *= scalars[1]; + } + } +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 2d9a3b7a..1b07dc6c 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,3 +1,2 @@ -pub mod infinite_byte_chunks; pub mod k_means; pub mod parallelism; diff --git a/src/utils/parallelism.rs b/src/utils/parallelism.rs index bd11191c..b960b568 100644 --- a/src/utils/parallelism.rs +++ b/src/utils/parallelism.rs @@ -7,8 +7,7 @@ pub use rayon::iter::ParallelIterator; pub trait Parallelism: Send + Sync { fn check(&self); - #[allow(clippy::wrong_self_convention)] - fn into_par_iter(&self, x: I) -> I::Iter; + fn rayon_into_par_iter(&self, x: I) -> I::Iter; } struct ParallelismCheckPanic(Box); @@ -57,7 +56,7 @@ impl Parallelism for RayonParallelism { } } - fn into_par_iter(&self, x: I) -> I::Iter { + fn rayon_into_par_iter(&self, x: I) -> I::Iter { x.into_par_iter() } } diff --git a/src/vchordrq/algorithm/build.rs b/src/vchordrq/algorithm/build.rs index 6554bfc3..4213e59d 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/vchordrq/algorithm/build.rs @@ -1,26 +1,25 @@ -use super::RelationWrite; use crate::vchordrq::algorithm::rabitq; use crate::vchordrq::algorithm::tuples::*; -use crate::vchordrq::algorithm::PageGuard; use crate::vchordrq::index::am_options::Opfamily; +use crate::vchordrq::types::DistanceKind; use crate::vchordrq::types::VchordrqBuildOptions; use crate::vchordrq::types::VchordrqExternalBuildOptions; use crate::vchordrq::types::VchordrqIndexingOptions; use crate::vchordrq::types::VchordrqInternalBuildOptions; use crate::vchordrq::types::VectorOptions; -use base::distance::DistanceKind; -use base::search::Pointer; -use base::simd::ScalarLike; -use base::vector::VectorBorrowed; +use algorithm::{Page, PageGuard, RelationWrite}; use rand::Rng; use rkyv::ser::serializers::AllocSerializer; +use simd::Floating; use std::marker::PhantomData; +use std::num::NonZeroU64; use std::sync::Arc; +use vector::VectorBorrowed; pub trait HeapRelation { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((Pointer, V)); + F: FnMut((NonZeroU64, V)); fn opfamily(&self) -> Opfamily; } @@ -207,8 +206,8 @@ impl Structure { let mut vectors = BTreeMap::new(); pgrx::spi::Spi::connect(|client| { use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; - use base::vector::VectorBorrowed; use pgrx::pg_sys::panic::ErrorReportable; + use vector::VectorBorrowed; let schema_query = "SELECT n.nspname::TEXT FROM pg_catalog.pg_extension e LEFT JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace diff --git a/src/vchordrq/algorithm/insert.rs b/src/vchordrq/algorithm/insert.rs index 323fa24f..f625dcac 100644 --- a/src/vchordrq/algorithm/insert.rs +++ b/src/vchordrq/algorithm/insert.rs @@ -1,19 +1,18 @@ -use super::RelationWrite; -use crate::vchordrq::algorithm::rabitq::fscan_process_lowerbound; +use crate::vchordrq::algorithm::rabitq; use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; -use crate::vchordrq::algorithm::PageGuard; -use base::always_equal::AlwaysEqual; -use base::distance::Distance; -use base::distance::DistanceKind; -use base::search::Pointer; -use base::vector::VectorBorrowed; +use crate::vchordrq::types::DistanceKind; +use algorithm::{Page, PageGuard, RelationWrite}; +use always_equal::AlwaysEqual; +use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; +use std::num::NonZeroU64; +use vector::VectorBorrowed; pub fn insert( relation: impl RelationWrite + Clone, - payload: Pointer, + payload: NonZeroU64, vector: V, distance_kind: DistanceKind, in_building: bool, @@ -41,7 +40,7 @@ pub fn insert( for i in (0..slices.len()).rev() { let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple:: { slice: slices[i].to_vec(), - payload: Some(payload.as_u64()), + payload: Some(payload), chain, }) .unwrap(); @@ -56,7 +55,7 @@ pub fn insert( } chain.ok().unwrap() }; - let h0_payload = payload.as_u64(); + let h0_payload = payload; let mut list = { let Some((_, original)) = vectors::vector_dist::( relation.clone(), @@ -90,7 +89,7 @@ pub fn insert( .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - let lowerbounds = fscan_process_lowerbound( + let lowerbounds = rabitq::process_lowerbound( distance_kind, dims, lut, diff --git a/src/vchordrq/algorithm/mod.rs b/src/vchordrq/algorithm/mod.rs index 41744d7d..88239a8e 100644 --- a/src/vchordrq/algorithm/mod.rs +++ b/src/vchordrq/algorithm/mod.rs @@ -6,62 +6,3 @@ pub mod scan; pub mod tuples; pub mod vacuum; pub mod vectors; - -use crate::postgres::Page; -use std::ops::{Deref, DerefMut}; - -pub trait PageGuard { - fn id(&self) -> u32; -} - -pub trait RelationRead { - type ReadGuard<'a>: PageGuard + Deref - where - Self: 'a; - fn read(&self, id: u32) -> Self::ReadGuard<'_>; -} - -pub trait RelationWrite: RelationRead { - type WriteGuard<'a>: PageGuard + DerefMut - where - Self: 'a; - fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_>; - fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_>; - fn search(&self, freespace: usize) -> Option>; -} - -impl PageGuard for crate::postgres::BufferReadGuard { - fn id(&self) -> u32 { - self.id() - } -} - -impl PageGuard for crate::postgres::BufferWriteGuard { - fn id(&self) -> u32 { - self.id() - } -} - -impl RelationRead for crate::postgres::Relation { - type ReadGuard<'a> = crate::postgres::BufferReadGuard; - - fn read(&self, id: u32) -> Self::ReadGuard<'_> { - self.read(id) - } -} - -impl RelationWrite for crate::postgres::Relation { - type WriteGuard<'a> = crate::postgres::BufferWriteGuard; - - fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { - self.write(id, tracking_freespace) - } - - fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_> { - self.extend(tracking_freespace) - } - - fn search(&self, freespace: usize) -> Option> { - self.search(freespace) - } -} diff --git a/src/vchordrq/algorithm/prewarm.rs b/src/vchordrq/algorithm/prewarm.rs index 6d01f7c1..6a7dc252 100644 --- a/src/vchordrq/algorithm/prewarm.rs +++ b/src/vchordrq/algorithm/prewarm.rs @@ -1,6 +1,6 @@ -use super::RelationRead; use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; +use algorithm::{Page, RelationRead}; use std::fmt::Write; pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> String { diff --git a/src/vchordrq/algorithm/rabitq.rs b/src/vchordrq/algorithm/rabitq.rs index b7b3858d..4f406e18 100644 --- a/src/vchordrq/algorithm/rabitq.rs +++ b/src/vchordrq/algorithm/rabitq.rs @@ -1,128 +1,21 @@ -use base::distance::{Distance, DistanceKind}; -use base::simd::ScalarLike; +use crate::vchordrq::types::DistanceKind; +use distance::Distance; -#[derive(Debug, Clone)] -pub struct Code { - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub signs: Vec, -} - -impl Code { - pub fn t(&self) -> Vec { - use crate::utils::infinite_byte_chunks::InfiniteByteChunks; - let mut result = Vec::new(); - for x in InfiniteByteChunks::<_, 64>::new(self.signs.iter().copied()) - .take(self.signs.len().div_ceil(64)) - { - let mut r = 0_u64; - for i in 0..64 { - r |= (x[i] as u64) << i; - } - result.push(r); - } - result - } -} - -pub fn code(dims: u32, vector: &[f32]) -> Code { - let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); - let sum_of_x_2 = f32::reduce_sum_of_x2(vector); - let dis_u = sum_of_x_2.sqrt(); - let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); - let x_x0 = dis_u / x0; - let fac_norm = (dims as f32).sqrt(); - let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); - let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); - let factor_ip = -2.0f32 / fac_norm * x_x0; - let cnt_pos = vector - .iter() - .map(|x| x.is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.is_sign_negative() as i32) - .sum::(); - let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; - let mut signs = Vec::new(); - for i in 0..dims { - signs.push(vector[i as usize].is_sign_positive() as u8); - } - Code { - dis_u_2: sum_of_x_2, - factor_ppc, - factor_ip, - factor_err, - signs, - } -} +pub use rabitq::binary::Code; +pub use rabitq::binary::Lut; +pub use rabitq::binary::code; +pub use rabitq::binary::preprocess; +pub use rabitq::binary::{process_lowerbound_dot, process_lowerbound_l2}; -pub type Lut = (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)); - -pub fn fscan_preprocess(vector: &[f32]) -> Lut { - use base::simd::quantize; - let dis_v_2 = f32::reduce_sum_of_x2(vector); - let (k, b, qvector) = quantize::quantize(vector, 15.0); - let qvector_sum = if vector.len() <= 4369 { - base::simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 - } else { - base::simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 - }; - (dis_v_2, b, k, qvector_sum, binarize(&qvector)) -} - -pub fn fscan_process_lowerbound( +pub fn process_lowerbound( distance_kind: DistanceKind, - _dims: u32, + dims: u32, lut: &Lut, - (dis_u_2, factor_ppc, factor_ip, factor_err, t): (f32, f32, f32, f32, &[u64]), + code: (f32, f32, f32, f32, &[u64]), epsilon: f32, ) -> Distance { match distance_kind { - DistanceKind::L2 => { - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let value = asymmetric_binary_dot_product(t, s) as u16; - let rough = dis_u_2 - + dis_v_2 - + b * factor_ppc - + ((2.0 * value as f32) - qvector_sum) * factor_ip * k; - let err = factor_err * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) - } - DistanceKind::Dot => { - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let value = asymmetric_binary_dot_product(t, s) as u16; - let rough = - 0.5 * b * factor_ppc + 0.5 * ((2.0 * value as f32) - qvector_sum) * factor_ip * k; - let err = 0.5 * factor_err * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) - } - DistanceKind::Hamming => unimplemented!(), - DistanceKind::Jaccard => unimplemented!(), - } -} - -fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { - let n = vector.len(); - let mut t0 = vec![0u64; n.div_ceil(64)]; - let mut t1 = vec![0u64; n.div_ceil(64)]; - let mut t2 = vec![0u64; n.div_ceil(64)]; - let mut t3 = vec![0u64; n.div_ceil(64)]; - for i in 0..n { - t0[i / 64] |= (((vector[i] >> 0) & 1) as u64) << (i % 64); - t1[i / 64] |= (((vector[i] >> 1) & 1) as u64) << (i % 64); - t2[i / 64] |= (((vector[i] >> 2) & 1) as u64) << (i % 64); - t3[i / 64] |= (((vector[i] >> 3) & 1) as u64) << (i % 64); + DistanceKind::L2 => process_lowerbound_l2(dims, lut, code, epsilon), + DistanceKind::Dot => process_lowerbound_dot(dims, lut, code, epsilon), } - (t0, t1, t2, t3) -} - -fn asymmetric_binary_dot_product(x: &[u64], y: &(Vec, Vec, Vec, Vec)) -> u32 { - let t0 = base::simd::bit::sum_of_and(x, &y.0); - let t1 = base::simd::bit::sum_of_and(x, &y.1); - let t2 = base::simd::bit::sum_of_and(x, &y.2); - let t3 = base::simd::bit::sum_of_and(x, &y.3); - (t0 << 0) + (t1 << 1) + (t2 << 2) + (t3 << 3) } diff --git a/src/vchordrq/algorithm/scan.rs b/src/vchordrq/algorithm/scan.rs index 42357c06..e6915f0a 100644 --- a/src/vchordrq/algorithm/scan.rs +++ b/src/vchordrq/algorithm/scan.rs @@ -1,14 +1,14 @@ -use super::RelationRead; -use crate::vchordrq::algorithm::rabitq::fscan_process_lowerbound; +use crate::vchordrq::algorithm::rabitq; use crate::vchordrq::algorithm::tuples::*; use crate::vchordrq::algorithm::vectors; -use base::always_equal::AlwaysEqual; -use base::distance::Distance; -use base::distance::DistanceKind; -use base::search::Pointer; -use base::vector::VectorBorrowed; +use crate::vchordrq::types::DistanceKind; +use algorithm::{Page, RelationRead}; +use always_equal::AlwaysEqual; +use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; +use std::num::NonZeroU64; +use vector::VectorBorrowed; pub fn scan( relation: impl RelationRead + Clone, @@ -16,7 +16,7 @@ pub fn scan( distance_kind: DistanceKind, probes: Vec, epsilon: f32, -) -> impl Iterator { +) -> impl Iterator { let vector = vector.as_borrowed(); let meta_guard = relation.read(0); let meta_tuple = meta_guard @@ -71,7 +71,7 @@ pub fn scan( .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - let lowerbounds = fscan_process_lowerbound( + let lowerbounds = rabitq::process_lowerbound( distance_kind, dims, lut, @@ -143,7 +143,7 @@ pub fn scan( .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - let lowerbounds = fscan_process_lowerbound( + let lowerbounds = rabitq::process_lowerbound( distance_kind, dims, lut, @@ -183,7 +183,7 @@ pub fn scan( cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); } let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; - Some((dis_u, Pointer::new(pay_u))) + Some((dis_u, pay_u)) }) } } diff --git a/src/vchordrq/algorithm/tuples.rs b/src/vchordrq/algorithm/tuples.rs index 40a795ce..23fe664e 100644 --- a/src/vchordrq/algorithm/tuples.rs +++ b/src/vchordrq/algorithm/tuples.rs @@ -1,10 +1,13 @@ +use std::num::NonZeroU64; + use super::rabitq::{self, Code, Lut}; +use crate::vchordrq::types::DistanceKind; use crate::vchordrq::types::OwnedVector; -use base::distance::DistanceKind; -use base::simd::ScalarLike; -use base::vector::{VectOwned, VectorOwned}; use half::f16; use rkyv::{Archive, ArchiveUnsized, CheckBytes, Deserialize, Serialize}; +use simd::Floating; +use vector::VectorOwned; +use vector::vect::VectOwned; pub trait Vector: VectorOwned { type Metadata: Copy @@ -70,20 +73,15 @@ impl Vector for VectOwned { type Element = f32; - fn metadata_from_archived(_: &::Archived) -> Self::Metadata { - () - } + fn metadata_from_archived(_: &::Archived) -> Self::Metadata {} fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f32]>) { let vector = vector.slice(); - ( - (), - match vector.len() { - 0..=960 => vec![vector], - 961..=1280 => vec![&vector[..640], &vector[640..]], - 1281.. => vector.chunks(1920).collect(), - }, - ) + ((), match vector.len() { + 0..=960 => vec![vector], + 961..=1280 => vec![&vector[..640], &vector[640..]], + 1281.. => vector.chunks(1920).collect(), + }) } fn vector_merge((): Self::Metadata, slice: &[Self::Element]) -> Self { @@ -109,8 +107,6 @@ impl Vector for VectOwned { match accumulator.0 { DistanceKind::L2 => accumulator.1 += f32::reduce_sum_of_d2(left, right), DistanceKind::Dot => accumulator.1 += -f32::reduce_sum_of_xy(left, right), - DistanceKind::Hamming => unreachable!(), - DistanceKind::Jaccard => unreachable!(), } } fn distance_end( @@ -126,11 +122,11 @@ impl Vector for VectOwned { } fn residual(vector: Self::Borrowed<'_>, center: Self::Borrowed<'_>) -> Self { - Self::new(ScalarLike::vector_sub(vector.slice(), center.slice())) + Self::new(Floating::vector_sub(vector.slice(), center.slice())) } fn rabitq_fscan_preprocess(vector: Self::Borrowed<'_>) -> Lut { - rabitq::fscan_preprocess(vector.slice()) + rabitq::preprocess(vector.slice()) } fn rabitq_code(dims: u32, vector: Self::Borrowed<'_>) -> Code { @@ -151,20 +147,15 @@ impl Vector for VectOwned { type Element = f16; - fn metadata_from_archived(_: &::Archived) -> Self::Metadata { - () - } + fn metadata_from_archived(_: &::Archived) -> Self::Metadata {} fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f16]>) { let vector = vector.slice(); - ( - (), - match vector.len() { - 0..=1920 => vec![vector], - 1921..=2560 => vec![&vector[..1280], &vector[1280..]], - 2561.. => vector.chunks(3840).collect(), - }, - ) + ((), match vector.len() { + 0..=1920 => vec![vector], + 1921..=2560 => vec![&vector[..1280], &vector[1280..]], + 2561.. => vector.chunks(3840).collect(), + }) } fn vector_merge((): Self::Metadata, slice: &[Self::Element]) -> Self { @@ -190,8 +181,6 @@ impl Vector for VectOwned { match accumulator.0 { DistanceKind::L2 => accumulator.1 += f16::reduce_sum_of_d2(left, right), DistanceKind::Dot => accumulator.1 += -f16::reduce_sum_of_xy(left, right), - DistanceKind::Hamming => unreachable!(), - DistanceKind::Jaccard => unreachable!(), } } fn distance_end( @@ -209,11 +198,11 @@ impl Vector for VectOwned { } fn residual(vector: Self::Borrowed<'_>, center: Self::Borrowed<'_>) -> Self { - Self::new(ScalarLike::vector_sub(vector.slice(), center.slice())) + Self::new(Floating::vector_sub(vector.slice(), center.slice())) } fn rabitq_fscan_preprocess(vector: Self::Borrowed<'_>) -> Lut { - rabitq::fscan_preprocess(&f16::vector_to_f32(vector.slice())) + rabitq::preprocess(&f16::vector_to_f32(vector.slice())) } fn rabitq_code(dims: u32, vector: Self::Borrowed<'_>) -> Code { @@ -246,7 +235,7 @@ pub struct MetaTuple { #[archive(check_bytes)] pub struct VectorTuple { pub slice: Vec, - pub payload: Option, + pub payload: Option, pub chain: Result<(u32, u16), V::Metadata>, } @@ -271,7 +260,7 @@ pub struct Height0Tuple { // raw vector pub mean: (u32, u16), // for height 0 tuple, it's pointers to heap relation - pub payload: u64, + pub payload: NonZeroU64, // RaBitQ algorithm pub dis_u_2: f32, pub factor_ppc: f32, diff --git a/src/vchordrq/algorithm/vacuum.rs b/src/vchordrq/algorithm/vacuum.rs index c77bd216..ee97ca6c 100644 --- a/src/vchordrq/algorithm/vacuum.rs +++ b/src/vchordrq/algorithm/vacuum.rs @@ -1,11 +1,11 @@ -use super::RelationWrite; use crate::vchordrq::algorithm::tuples::*; -use base::search::Pointer; +use algorithm::{Page, RelationWrite}; +use std::num::NonZeroU64; pub fn vacuum( relation: impl RelationWrite, delay: impl Fn(), - callback: impl Fn(Pointer) -> bool, + callback: impl Fn(NonZeroU64) -> bool, ) { // step 1: vacuum height_0_tuple { @@ -50,7 +50,7 @@ pub fn vacuum( .map(rkyv::check_archived_root::) .expect("data corruption") .expect("data corruption"); - if callback(Pointer::new(h0_tuple.payload)) { + if callback(h0_tuple.payload) { reconstruct_removes.push(i); } } @@ -81,7 +81,7 @@ pub fn vacuum( let vector_tuple = unsafe { rkyv::archived_root::>(vector_tuple) }; if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(Pointer::new(payload)) { + if callback(payload) { break 'flag true; } } @@ -98,7 +98,7 @@ pub fn vacuum( let vector_tuple = unsafe { rkyv::archived_root::>(vector_tuple) }; if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(Pointer::new(payload)) { + if callback(payload) { write.free(i); } } diff --git a/src/vchordrq/algorithm/vectors.rs b/src/vchordrq/algorithm/vectors.rs index 06075d37..72e448a1 100644 --- a/src/vchordrq/algorithm/vectors.rs +++ b/src/vchordrq/algorithm/vectors.rs @@ -1,14 +1,15 @@ use super::tuples::Vector; -use super::RelationRead; use crate::vchordrq::algorithm::tuples::VectorTuple; -use base::distance::Distance; -use base::distance::DistanceKind; +use crate::vchordrq::types::DistanceKind; +use algorithm::{Page, RelationRead}; +use distance::Distance; +use std::num::NonZeroU64; pub fn vector_dist( relation: impl RelationRead, vector: V::Borrowed<'_>, mean: (u32, u16), - payload: Option, + payload: Option, for_distance: Option, for_original: bool, ) -> Option<(Option, Option)> { diff --git a/src/vchordrq/gucs/executing.rs b/src/vchordrq/gucs/executing.rs index 6b20f749..af6cce7d 100644 --- a/src/vchordrq/gucs/executing.rs +++ b/src/vchordrq/gucs/executing.rs @@ -68,9 +68,5 @@ pub fn epsilon() -> f32 { pub fn max_scan_tuples() -> Option { let x = MAX_SCAN_TUPLES.get(); - if x < 0 { - None - } else { - Some(x as u32) - } + if x < 0 { None } else { Some(x as u32) } } diff --git a/src/vchordrq/index/am.rs b/src/vchordrq/index/am.rs index 1dc7cdcf..3d19a490 100644 --- a/src/vchordrq/index/am.rs +++ b/src/vchordrq/index/am.rs @@ -1,4 +1,4 @@ -use crate::postgres::Relation; +use crate::postgres::PostgresRelation; use crate::vchordrq::algorithm; use crate::vchordrq::algorithm::build::{HeapRelation, Reporter}; use crate::vchordrq::algorithm::tuples::Vector; @@ -7,11 +7,11 @@ use crate::vchordrq::index::am_scan::Scanner; use crate::vchordrq::index::utils::{ctid_to_pointer, pointer_to_ctid}; use crate::vchordrq::index::{am_options, am_scan}; use crate::vchordrq::types::VectorKind; -use base::search::Pointer; -use base::vector::VectOwned; use half::f16; use pgrx::datum::Internal; use pgrx::pg_sys::Datum; +use std::num::NonZeroU64; +use vector::vect::VectOwned; static mut RELOPT_KIND_VCHORDRQ: pgrx::pg_sys::relopt_kind::Type = 0; @@ -170,7 +170,7 @@ pub unsafe extern "C" fn ambuild( impl HeapRelation for Heap { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((Pointer, V)), + F: FnMut((NonZeroU64, V)), { pub struct State<'a, F> { pub this: &'a Heap, @@ -185,7 +185,7 @@ pub unsafe extern "C" fn ambuild( _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((Pointer, V)), + F: FnMut((NonZeroU64, V)), { let state = unsafe { &mut *state.cast::>() }; let opfamily = state.this.opfamily; @@ -242,7 +242,7 @@ pub unsafe extern "C" fn ambuild( opfamily, }; let mut reporter = PgReporter {}; - let index_relation = unsafe { Relation::new(index) }; + let index_relation = unsafe { PostgresRelation::new(index) }; match opfamily.vector_kind() { VectorKind::Vecf32 => algorithm::build::build::, Heap, _>( vector_options, @@ -289,7 +289,7 @@ pub unsafe extern "C" fn ambuild( } else { let mut indtuples = 0; reporter.tuples_done(indtuples); - let relation = unsafe { Relation::new(index) }; + let relation = unsafe { PostgresRelation::new(index) }; match opfamily.vector_kind() { VectorKind::Vecf32 => { HeapRelation::>::traverse( @@ -575,7 +575,7 @@ unsafe fn parallel_build( impl HeapRelation for Heap { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((Pointer, V)), + F: FnMut((NonZeroU64, V)), { pub struct State<'a, F> { pub this: &'a Heap, @@ -590,7 +590,7 @@ unsafe fn parallel_build( _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((Pointer, V)), + F: FnMut((NonZeroU64, V)), { let state = unsafe { &mut *state.cast::>() }; let opfamily = state.this.opfamily; @@ -627,7 +627,7 @@ unsafe fn parallel_build( } } - let index_relation = unsafe { Relation::new(index) }; + let index_relation = unsafe { PostgresRelation::new(index) }; let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap, tablescandesc) }; let opfamily = unsafe { am_options::opfamily(index) }; @@ -716,14 +716,14 @@ pub unsafe extern "C" fn aminsert( let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); match opfamily.vector_kind() { VectorKind::Vecf32 => algorithm::insert::insert::>( - unsafe { Relation::new(index) }, + unsafe { PostgresRelation::new(index) }, pointer, VectOwned::::from_owned(vector), opfamily.distance_kind(), false, ), VectorKind::Vecf16 => algorithm::insert::insert::>( - unsafe { Relation::new(index) }, + unsafe { PostgresRelation::new(index) }, pointer, VectOwned::::from_owned(vector), opfamily.distance_kind(), @@ -752,14 +752,14 @@ pub unsafe extern "C" fn aminsert( let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); match opfamily.vector_kind() { VectorKind::Vecf32 => algorithm::insert::insert::>( - unsafe { Relation::new(index) }, + unsafe { PostgresRelation::new(index) }, pointer, VectOwned::::from_owned(vector), opfamily.distance_kind(), false, ), VectorKind::Vecf16 => algorithm::insert::insert::>( - unsafe { Relation::new(index) }, + unsafe { PostgresRelation::new(index) }, pointer, VectOwned::::from_owned(vector), opfamily.distance_kind(), @@ -854,7 +854,7 @@ pub unsafe extern "C" fn amgettuple( pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); } let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; - let relation = unsafe { Relation::new((*scan).indexRelation) }; + let relation = unsafe { PostgresRelation::new((*scan).indexRelation) }; if let Some((pointer, recheck)) = am_scan::scan_next(scanner, relation) { let ctid = pointer_to_ctid(pointer); unsafe { @@ -892,17 +892,17 @@ pub unsafe extern "C" fn ambulkdelete( } let opfamily = unsafe { am_options::opfamily((*info).index) }; let callback = callback.unwrap(); - let callback = |p: Pointer| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; + let callback = |p: NonZeroU64| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; match opfamily.vector_kind() { VectorKind::Vecf32 => algorithm::vacuum::vacuum::>( - unsafe { Relation::new((*info).index) }, + unsafe { PostgresRelation::new((*info).index) }, || unsafe { pgrx::pg_sys::vacuum_delay_point(); }, callback, ), VectorKind::Vecf16 => algorithm::vacuum::vacuum::>( - unsafe { Relation::new((*info).index) }, + unsafe { PostgresRelation::new((*info).index) }, || unsafe { pgrx::pg_sys::vacuum_delay_point(); }, diff --git a/src/vchordrq/index/am_options.rs b/src/vchordrq/index/am_options.rs index 25fcd0c3..5c730ed7 100644 --- a/src/vchordrq/index/am_options.rs +++ b/src/vchordrq/index/am_options.rs @@ -3,16 +3,16 @@ use crate::datatype::memory_pgvector_halfvec::PgvectorHalfvecOutput; use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use crate::datatype::typmod::Typmod; -use crate::vchordrq::types::VchordrqIndexingOptions; -use crate::vchordrq::types::VectorOptions; -use crate::vchordrq::types::{BorrowedVector, OwnedVector, VectorKind}; -use base::distance::*; -use base::vector::VectorBorrowed; +use crate::vchordrq::types::{BorrowedVector, OwnedVector}; +use crate::vchordrq::types::{DistanceKind, VectorKind}; +use crate::vchordrq::types::{VchordrqIndexingOptions, VectorOptions}; +use distance::Distance; use pgrx::datum::FromDatum; use pgrx::heap_tuple::PgHeapTuple; use serde::Deserialize; use std::ffi::CStr; use std::num::NonZero; +use vector::VectorBorrowed; #[derive(Copy, Clone, Debug, Default)] #[repr(C)] diff --git a/src/vchordrq/index/am_scan.rs b/src/vchordrq/index/am_scan.rs index 1b78ff08..6e2d30d5 100644 --- a/src/vchordrq/index/am_scan.rs +++ b/src/vchordrq/index/am_scan.rs @@ -1,5 +1,5 @@ use super::am_options::Opfamily; -use crate::postgres::Relation; +use crate::postgres::PostgresRelation; use crate::vchordrq::algorithm::scan::scan; use crate::vchordrq::algorithm::tuples::Vector; use crate::vchordrq::gucs::executing::epsilon; @@ -7,10 +7,10 @@ use crate::vchordrq::gucs::executing::max_scan_tuples; use crate::vchordrq::gucs::executing::probes; use crate::vchordrq::types::OwnedVector; use crate::vchordrq::types::VectorKind; -use base::distance::Distance; -use base::search::*; -use base::vector::VectOwned; +use distance::Distance; use half::f16; +use std::num::NonZeroU64; +use vector::vect::VectOwned; pub enum Scanner { Initial { @@ -19,7 +19,7 @@ pub enum Scanner { recheck: bool, }, Vbase { - vbase: Box>, + vbase: Box>, threshold: Option, recheck: bool, opfamily: Opfamily, @@ -66,7 +66,7 @@ pub fn scan_make( } } -pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, bool)> { +pub fn scan_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(NonZeroU64, bool)> { if let Scanner::Initial { vector, threshold, diff --git a/src/vchordrq/index/functions.rs b/src/vchordrq/index/functions.rs index 05f348f5..32e6d03d 100644 --- a/src/vchordrq/index/functions.rs +++ b/src/vchordrq/index/functions.rs @@ -1,11 +1,11 @@ use super::am_options; -use crate::postgres::Relation; +use crate::postgres::PostgresRelation; use crate::vchordrq::algorithm::prewarm::prewarm; use crate::vchordrq::types::VectorKind; -use base::vector::VectOwned; use half::f16; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass}; +use vector::vect::VectOwned; #[pgrx::pg_extern(sql = "")] fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { @@ -21,7 +21,7 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { pgrx::error!("{:?} is not a vchordrq index", pg_class.relname()); } let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; - let relation = unsafe { Relation::new(index) }; + let relation = unsafe { PostgresRelation::new(index) }; let opfamily = unsafe { am_options::opfamily(index) }; let message = match opfamily.vector_kind() { VectorKind::Vecf32 => prewarm::>(relation, height), diff --git a/src/vchordrq/index/utils.rs b/src/vchordrq/index/utils.rs index a5d85a3f..726a5979 100644 --- a/src/vchordrq/index/utils.rs +++ b/src/vchordrq/index/utils.rs @@ -1,7 +1,7 @@ -use base::search::*; +use std::num::NonZeroU64; -pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { - let value = pointer.as_u64(); +pub fn pointer_to_ctid(pointer: NonZeroU64) -> pgrx::pg_sys::ItemPointerData { + let value = pointer.get(); pgrx::pg_sys::ItemPointerData { ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi: ((value >> 32) & 0xffff) as u16, @@ -11,10 +11,10 @@ pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { } } -pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer { +pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> NonZeroU64 { let mut value = 0; value |= (ctid.ip_blkid.bi_hi as u64) << 32; value |= (ctid.ip_blkid.bi_lo as u64) << 16; value |= ctid.ip_posid as u64; - Pointer::new(value) + NonZeroU64::new(value).expect("invalid pointer") } diff --git a/src/vchordrq/types.rs b/src/vchordrq/types.rs index 0e1bdc07..4ef2171d 100644 --- a/src/vchordrq/types.rs +++ b/src/vchordrq/types.rs @@ -1,8 +1,7 @@ -use base::distance::DistanceKind; -use base::vector::{VectBorrowed, VectOwned}; use half::f16; use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError, ValidationErrors}; +use vector::vect::{VectBorrowed, VectOwned}; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] @@ -111,6 +110,13 @@ pub enum BorrowedVector<'a> { Vecf16(VectBorrowed<'a, f16>), } +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +pub enum DistanceKind { + L2, + Dot, +} + #[repr(u8)] #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] pub enum VectorKind { diff --git a/src/vchordrqfscan/algorithm/build.rs b/src/vchordrqfscan/algorithm/build.rs index 9a0daa6e..a245d776 100644 --- a/src/vchordrqfscan/algorithm/build.rs +++ b/src/vchordrqfscan/algorithm/build.rs @@ -1,25 +1,24 @@ -use crate::postgres::BufferWriteGuard; -use crate::postgres::Relation; use crate::vchordrqfscan::algorithm::rabitq; use crate::vchordrqfscan::algorithm::tuples::*; use crate::vchordrqfscan::index::am_options::Opfamily; +use crate::vchordrqfscan::types::DistanceKind; use crate::vchordrqfscan::types::VchordrqfscanBuildOptions; use crate::vchordrqfscan::types::VchordrqfscanExternalBuildOptions; use crate::vchordrqfscan::types::VchordrqfscanIndexingOptions; use crate::vchordrqfscan::types::VchordrqfscanInternalBuildOptions; use crate::vchordrqfscan::types::VectorOptions; -use base::distance::DistanceKind; -use base::search::Pointer; -use base::simd::ScalarLike; +use algorithm::{Page, PageGuard, RelationWrite}; use rand::Rng; use rkyv::ser::serializers::AllocSerializer; +use simd::Floating; use std::marker::PhantomData; +use std::num::NonZeroU64; use std::sync::Arc; pub trait HeapRelation { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((Pointer, Vec)); + F: FnMut((NonZeroU64, Vec)); fn opfamily(&self) -> Opfamily; } @@ -31,7 +30,7 @@ pub fn build( vector_options: VectorOptions, vchordrqfscan_options: VchordrqfscanIndexingOptions, heap_relation: T, - relation: Relation, + relation: impl RelationWrite, mut reporter: R, ) { let dims = vector_options.dims; @@ -73,7 +72,7 @@ pub fn build( }; let mut meta = Tape::create(&relation, false); assert_eq!(meta.first(), 0); - let mut forwards = Tape::::create(&relation, false); + let mut forwards = Tape::::create(&relation, false); assert_eq!(forwards.first(), 1); let mut vectors = Tape::create(&relation, true); assert_eq!(vectors.first(), 2); @@ -94,10 +93,10 @@ pub fn build( let mut level = Vec::new(); for j in 0..structures[i].len() { if i == 0 { - let tape = Tape::::create(&relation, false); + let tape = Tape::::create(&relation, false); level.push(tape.first()); } else { - let mut tape = Tape::::create(&relation, false); + let mut tape = Tape::::create(&relation, false); let mut cache = Vec::new(); let h2_mean = &structures[i].means[j]; let h2_children = &structures[i].children[j]; @@ -250,8 +249,8 @@ impl Structure { let mut vectors = BTreeMap::new(); pgrx::spi::Spi::connect(|client| { use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; - use base::vector::VectorBorrowed; use pgrx::pg_sys::panic::ErrorReportable; + use vector::VectorBorrowed; let table = client.select(&query, None, None).unwrap_or_report(); for row in table { let id: Option = row.get_by_name("id").unwrap(); @@ -368,16 +367,16 @@ impl Structure { } } -struct Tape<'a, T> { - relation: &'a Relation, - head: BufferWriteGuard, +struct Tape<'a, 'b, T, R: 'b + RelationWrite> { + relation: &'a R, + head: R::WriteGuard<'b>, first: u32, tracking_freespace: bool, _phantom: PhantomData T>, } -impl<'a, T> Tape<'a, T> { - fn create(relation: &'a Relation, tracking_freespace: bool) -> Self { +impl<'a: 'b, 'b, T, R: 'b + RelationWrite> Tape<'a, 'b, T, R> { + fn create(relation: &'a R, tracking_freespace: bool) -> Self { let head = relation.extend(tracking_freespace); let first = head.id(); Self { @@ -393,7 +392,7 @@ impl<'a, T> Tape<'a, T> { } } -impl Tape<'_, T> +impl<'a: 'b, 'b, T, R: 'b + RelationWrite> Tape<'a, 'b, T, R> where T: rkyv::Serialize>, { diff --git a/src/vchordrqfscan/algorithm/insert.rs b/src/vchordrqfscan/algorithm/insert.rs index 1a0ab778..ee447601 100644 --- a/src/vchordrqfscan/algorithm/insert.rs +++ b/src/vchordrqfscan/algorithm/insert.rs @@ -1,17 +1,22 @@ -use crate::postgres::Relation; use crate::vchordrqfscan::algorithm::rabitq; -use crate::vchordrqfscan::algorithm::rabitq::distance; use crate::vchordrqfscan::algorithm::rabitq::fscan_process_lowerbound; use crate::vchordrqfscan::algorithm::tuples::*; -use base::always_equal::AlwaysEqual; -use base::distance::Distance; -use base::distance::DistanceKind; -use base::search::Pointer; -use base::simd::ScalarLike; +use crate::vchordrqfscan::types::DistanceKind; +use crate::vchordrqfscan::types::distance; +use algorithm::{Page, PageGuard, RelationWrite}; +use always_equal::AlwaysEqual; +use distance::Distance; +use simd::Floating; use std::cmp::Reverse; use std::collections::BinaryHeap; +use std::num::NonZeroU64; -pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_kind: DistanceKind) { +pub fn insert( + relation: impl RelationWrite, + payload: NonZeroU64, + vector: Vec, + distance_kind: DistanceKind, +) { let meta_guard = relation.read(0); let meta_tuple = meta_guard .get(1) @@ -30,7 +35,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k let h0_vector = 'h0_vector: { let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { vector: vector.clone(), - payload: Some(payload.as_u64()), + payload: Some(payload), }) .unwrap(); if let Some(mut write) = relation.search(tuple.len()) { @@ -78,7 +83,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k changed = true; } }; - let h0_payload = payload.as_u64(); + let h0_payload = payload; let mut list = ( meta_tuple.first, if is_residual { @@ -173,7 +178,7 @@ pub fn insert(relation: Relation, payload: Pointer, vector: Vec, distance_k let dummy = rkyv::to_bytes::<_, 8192>(&Height0Tuple { mask: [false; 32], mean: [(0, 0); 32], - payload: [0; 32], + payload: [NonZeroU64::MIN; 32], dis_u_2: [0.0f32; 32], factor_ppc: [0.0f32; 32], factor_ip: [0.0f32; 32], diff --git a/src/vchordrqfscan/algorithm/prewarm.rs b/src/vchordrqfscan/algorithm/prewarm.rs index ec8976ab..c8500d43 100644 --- a/src/vchordrqfscan/algorithm/prewarm.rs +++ b/src/vchordrqfscan/algorithm/prewarm.rs @@ -1,8 +1,8 @@ -use crate::postgres::Relation; use crate::vchordrqfscan::algorithm::tuples::*; +use algorithm::{Page, RelationRead}; use std::fmt::Write; -pub fn prewarm(relation: Relation, height: i32) -> String { +pub fn prewarm(relation: impl RelationRead, height: i32) -> String { let mut message = String::new(); let meta_guard = relation.read(0); let meta_tuple = meta_guard diff --git a/src/vchordrqfscan/algorithm/rabitq.rs b/src/vchordrqfscan/algorithm/rabitq.rs index 65c4996e..707d81c6 100644 --- a/src/vchordrqfscan/algorithm/rabitq.rs +++ b/src/vchordrqfscan/algorithm/rabitq.rs @@ -1,171 +1,22 @@ -use crate::utils::infinite_byte_chunks::InfiniteByteChunks; -use base::distance::{Distance, DistanceKind}; -use base::simd::ScalarLike; +use crate::vchordrqfscan::types::DistanceKind; +use distance::Distance; -#[derive(Debug, Clone)] -pub struct Code { - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub signs: Vec, -} - -pub fn code(dims: u32, vector: &[f32]) -> Code { - let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); - let sum_of_x_2 = f32::reduce_sum_of_x2(vector); - let dis_u = sum_of_x_2.sqrt(); - let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); - let x_x0 = dis_u / x0; - let fac_norm = (dims as f32).sqrt(); - let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); - let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); - let factor_ip = -2.0f32 / fac_norm * x_x0; - let cnt_pos = vector - .iter() - .map(|x| x.is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.is_sign_negative() as i32) - .sum::(); - let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; - let mut signs = Vec::new(); - for i in 0..dims { - signs.push(vector[i as usize].is_sign_positive() as u8); - } - Code { - dis_u_2: sum_of_x_2, - factor_ppc, - factor_ip, - factor_err, - signs, - } -} - -pub fn dummy_code(dims: u32) -> Code { - Code { - dis_u_2: 0.0, - factor_ppc: 0.0, - factor_ip: 0.0, - factor_err: 0.0, - signs: vec![0; dims as _], - } -} - -pub struct PackedCodes { - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, -} - -pub fn pack_codes(dims: u32, codes: [Code; 32]) -> PackedCodes { - PackedCodes { - dis_u_2: std::array::from_fn(|i| codes[i].dis_u_2), - factor_ppc: std::array::from_fn(|i| codes[i].factor_ppc), - factor_ip: std::array::from_fn(|i| codes[i].factor_ip), - factor_err: std::array::from_fn(|i| codes[i].factor_err), - t: { - let signs = codes.map(|code| { - InfiniteByteChunks::new(code.signs.into_iter()) - .map(|[b0, b1, b2, b3]| b0 | b1 << 1 | b2 << 2 | b3 << 3) - .take(dims.div_ceil(4) as usize) - .collect::>() - }); - base::simd::fast_scan::b4::pack(dims.div_ceil(4), signs).collect() - }, - } -} - -pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { - use base::simd::quantize; - let dis_v_2 = f32::reduce_sum_of_x2(vector); - let (k, b, qvector) = quantize::quantize(vector, 15.0); - let qvector_sum = if vector.len() <= 4369 { - base::simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 - } else { - base::simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 - }; - (dis_v_2, b, k, qvector_sum, compress(qvector)) -} +pub use rabitq::block::Code; +pub use rabitq::block::code; +pub use rabitq::block::dummy_code; +pub use rabitq::block::fscan_preprocess; +pub use rabitq::block::pack_codes; +pub use rabitq::block::{fscan_process_lowerbound_dot, fscan_process_lowerbound_l2}; pub fn fscan_process_lowerbound( distance_kind: DistanceKind, dims: u32, lut: &(f32, f32, f32, f32, Vec), - (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[u8], - ), + code: (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], &[u8]), epsilon: f32, ) -> [Distance; 32] { - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let r = base::simd::fast_scan::b4::fast_scan_b4(dims.div_ceil(4), t, s); match distance_kind { - DistanceKind::L2 => std::array::from_fn(|i| { - let rough = dis_u_2[i] - + dis_v_2 - + b * factor_ppc[i] - + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) - }), - DistanceKind::Dot => std::array::from_fn(|i| { - let rough = 0.5 * b * factor_ppc[i] - + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) - }), - DistanceKind::Hamming => unreachable!(), - DistanceKind::Jaccard => unreachable!(), - } -} - -fn compress(mut qvector: Vec) -> Vec { - let dims = qvector.len() as u32; - let width = dims.div_ceil(4); - qvector.resize(qvector.len().next_multiple_of(4), 0); - let mut t = vec![0u8; width as usize * 16]; - for i in 0..width as usize { - unsafe { - // this hint is used to skip bound checks - std::hint::assert_unchecked(4 * i + 3 < qvector.len()); - std::hint::assert_unchecked(16 * i + 15 < t.len()); - } - let t0 = qvector[4 * i + 0]; - let t1 = qvector[4 * i + 1]; - let t2 = qvector[4 * i + 2]; - let t3 = qvector[4 * i + 3]; - t[16 * i + 0b0000] = 0; - t[16 * i + 0b0001] = t0; - t[16 * i + 0b0010] = t1; - t[16 * i + 0b0011] = t1 + t0; - t[16 * i + 0b0100] = t2; - t[16 * i + 0b0101] = t2 + t0; - t[16 * i + 0b0110] = t2 + t1; - t[16 * i + 0b0111] = t2 + t1 + t0; - t[16 * i + 0b1000] = t3; - t[16 * i + 0b1001] = t3 + t0; - t[16 * i + 0b1010] = t3 + t1; - t[16 * i + 0b1011] = t3 + t1 + t0; - t[16 * i + 0b1100] = t3 + t2; - t[16 * i + 0b1101] = t3 + t2 + t0; - t[16 * i + 0b1110] = t3 + t2 + t1; - t[16 * i + 0b1111] = t3 + t2 + t1 + t0; - } - t -} - -pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance { - match d { - DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)), - DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)), - DistanceKind::Hamming => unimplemented!(), - DistanceKind::Jaccard => unimplemented!(), + DistanceKind::L2 => fscan_process_lowerbound_l2(dims, lut, code, epsilon), + DistanceKind::Dot => fscan_process_lowerbound_dot(dims, lut, code, epsilon), } } diff --git a/src/vchordrqfscan/algorithm/scan.rs b/src/vchordrqfscan/algorithm/scan.rs index a691b6c4..5546af93 100644 --- a/src/vchordrqfscan/algorithm/scan.rs +++ b/src/vchordrqfscan/algorithm/scan.rs @@ -1,23 +1,23 @@ -use crate::postgres::Relation; use crate::vchordrqfscan::algorithm::rabitq; -use crate::vchordrqfscan::algorithm::rabitq::distance; use crate::vchordrqfscan::algorithm::rabitq::fscan_process_lowerbound; use crate::vchordrqfscan::algorithm::tuples::*; -use base::always_equal::AlwaysEqual; -use base::distance::Distance; -use base::distance::DistanceKind; -use base::search::Pointer; -use base::simd::ScalarLike; +use crate::vchordrqfscan::types::DistanceKind; +use crate::vchordrqfscan::types::distance; +use algorithm::{Page, RelationWrite}; +use always_equal::AlwaysEqual; +use distance::Distance; +use simd::Floating; use std::cmp::Reverse; use std::collections::BinaryHeap; +use std::num::NonZeroU64; pub fn scan( - relation: Relation, + relation: impl RelationWrite + Clone, vector: Vec, distance_kind: DistanceKind, probes: Vec, epsilon: f32, -) -> impl Iterator { +) -> impl Iterator { let meta_guard = relation.read(0); let meta_tuple = meta_guard .get(1) @@ -123,6 +123,7 @@ pub fn scan( for i in (1..meta_tuple.height_of_root).rev() { lists = make_lists(lists, probes[i as usize - 1]); } + drop(meta_guard); { let mut results = Vec::new(); for list in lists { @@ -186,7 +187,7 @@ pub fn scan( cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); } let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; - Some((dis_u, Pointer::new(pay_u))) + Some((dis_u, pay_u)) }) } } diff --git a/src/vchordrqfscan/algorithm/tuples.rs b/src/vchordrqfscan/algorithm/tuples.rs index 3b43dac8..ff947131 100644 --- a/src/vchordrqfscan/algorithm/tuples.rs +++ b/src/vchordrqfscan/algorithm/tuples.rs @@ -1,3 +1,5 @@ +use std::num::NonZeroU64; + use crate::vchordrqfscan::algorithm::rabitq; use rkyv::{Archive, Deserialize, Serialize}; @@ -20,7 +22,7 @@ pub struct MetaTuple { pub struct VectorTuple { pub vector: Vec, // this field is saved only for vacuum - pub payload: Option, + pub payload: Option, } #[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] @@ -46,7 +48,7 @@ pub struct Height0Tuple { // raw vector pub mean: [(u32, u16); 32], // for height 0 tuple, it's pointers to heap relation - pub payload: [u64; 32], + pub payload: [NonZeroU64; 32], // RaBitQ algorithm pub dis_u_2: [f32; 32], pub factor_ppc: [f32; 32], @@ -60,7 +62,7 @@ pub fn put( dims: u32, code: &rabitq::Code, vector: (u32, u16), - payload: u64, + payload: NonZeroU64, ) -> bool { // todo: use mutable api let mut x = rkyv::from_bytes::(bytes).expect("data corruption"); diff --git a/src/vchordrqfscan/algorithm/vacuum.rs b/src/vchordrqfscan/algorithm/vacuum.rs index 8ede4f64..7773ed27 100644 --- a/src/vchordrqfscan/algorithm/vacuum.rs +++ b/src/vchordrqfscan/algorithm/vacuum.rs @@ -1,9 +1,13 @@ -use crate::postgres::Relation; use crate::vchordrqfscan::algorithm::tuples::VectorTuple; use crate::vchordrqfscan::algorithm::tuples::*; -use base::search::Pointer; +use algorithm::{Page, RelationWrite}; +use std::num::NonZeroU64; -pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) -> bool) { +pub fn vacuum( + relation: impl RelationWrite, + delay: impl Fn(), + callback: impl Fn(NonZeroU64) -> bool, +) { // step 1: vacuum height_0_tuple { let meta_guard = relation.read(0); @@ -52,7 +56,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - .expect("data corruption"); let flag = 'flag: { for j in 0..32 { - if h0_tuple.mask[j] && callback(Pointer::new(h0_tuple.payload[j])) { + if h0_tuple.mask[j] && callback(h0_tuple.payload[j]) { break 'flag true; } } @@ -66,7 +70,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - .expect("data corruption") .expect("data corruption"); for j in 0..32 { - if temp.mask[j] && callback(Pointer::new(temp.payload[j])) { + if temp.mask[j] && callback(temp.payload[j]) { temp.mask[j] = false; } } @@ -104,7 +108,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - let vector_tuple = rkyv::check_archived_root::(vector_tuple) .expect("data corruption"); if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(Pointer::new(payload)) { + if callback(payload) { break 'flag true; } } @@ -121,7 +125,7 @@ pub fn vacuum(relation: Relation, delay: impl Fn(), callback: impl Fn(Pointer) - let vector_tuple = rkyv::check_archived_root::(vector_tuple) .expect("data corruption"); if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(Pointer::new(payload)) { + if callback(payload) { write.free(i); } } diff --git a/src/vchordrqfscan/gucs/executing.rs b/src/vchordrqfscan/gucs/executing.rs index 6ec186ec..ba204b9f 100644 --- a/src/vchordrqfscan/gucs/executing.rs +++ b/src/vchordrqfscan/gucs/executing.rs @@ -68,9 +68,5 @@ pub fn epsilon() -> f32 { pub fn max_scan_tuples() -> Option { let x = MAX_SCAN_TUPLES.get(); - if x < 0 { - None - } else { - Some(x as u32) - } + if x < 0 { None } else { Some(x as u32) } } diff --git a/src/vchordrqfscan/index/am.rs b/src/vchordrqfscan/index/am.rs index 2b834fbc..7011e26f 100644 --- a/src/vchordrqfscan/index/am.rs +++ b/src/vchordrqfscan/index/am.rs @@ -1,13 +1,13 @@ -use crate::postgres::Relation; +use crate::postgres::PostgresRelation; use crate::vchordrqfscan::algorithm; use crate::vchordrqfscan::algorithm::build::{HeapRelation, Reporter}; use crate::vchordrqfscan::index::am_options::{Opfamily, Reloption}; use crate::vchordrqfscan::index::am_scan::Scanner; use crate::vchordrqfscan::index::utils::{ctid_to_pointer, pointer_to_ctid}; use crate::vchordrqfscan::index::{am_options, am_scan}; -use base::search::Pointer; use pgrx::datum::Internal; use pgrx::pg_sys::Datum; +use std::num::NonZeroU64; static mut RELOPT_KIND_VCHORDRQFSCAN: pgrx::pg_sys::relopt_kind::Type = 0; @@ -166,7 +166,7 @@ pub unsafe extern "C" fn ambuild( impl HeapRelation for Heap { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((Pointer, Vec)), + F: FnMut((NonZeroU64, Vec)), { pub struct State<'a, F> { pub this: &'a Heap, @@ -181,7 +181,7 @@ pub unsafe extern "C" fn ambuild( _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((Pointer, Vec)), + F: FnMut((NonZeroU64, Vec)), { use crate::vchordrqfscan::types::OwnedVector; let state = unsafe { &mut *state.cast::>() }; @@ -242,7 +242,7 @@ pub unsafe extern "C" fn ambuild( opfamily, }; let mut reporter = PgReporter {}; - let index_relation = unsafe { Relation::new(index) }; + let index_relation = unsafe { PostgresRelation::new(index) }; algorithm::build::build( vector_options, vchordrqfscan_options, @@ -552,7 +552,7 @@ unsafe fn parallel_build( impl HeapRelation for Heap { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((Pointer, Vec)), + F: FnMut((NonZeroU64, Vec)), { pub struct State<'a, F> { pub this: &'a Heap, @@ -567,7 +567,7 @@ unsafe fn parallel_build( _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((Pointer, Vec)), + F: FnMut((NonZeroU64, Vec)), { use crate::vchordrqfscan::types::OwnedVector; let state = unsafe { &mut *state.cast::>() }; @@ -608,7 +608,7 @@ unsafe fn parallel_build( } } - let index_relation = unsafe { Relation::new(index) }; + let index_relation = unsafe { PostgresRelation::new(index) }; let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap, tablescandesc) }; let opfamily = unsafe { am_options::opfamily(index) }; let heap_relation = Heap { @@ -672,7 +672,7 @@ pub unsafe extern "C" fn aminsert( }; let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); algorithm::insert::insert( - unsafe { Relation::new(index) }, + unsafe { PostgresRelation::new(index) }, pointer, vector.into_vec(), opfamily.distance_kind(), @@ -702,7 +702,7 @@ pub unsafe extern "C" fn aminsert( }; let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); algorithm::insert::insert( - unsafe { Relation::new(index) }, + unsafe { PostgresRelation::new(index) }, pointer, vector.into_vec(), opfamily.distance_kind(), @@ -795,7 +795,7 @@ pub unsafe extern "C" fn amgettuple( pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); } let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; - let relation = unsafe { Relation::new((*scan).indexRelation) }; + let relation = unsafe { PostgresRelation::new((*scan).indexRelation) }; if let Some((pointer, recheck)) = am_scan::scan_next(scanner, relation) { let ctid = pointer_to_ctid(pointer); unsafe { @@ -832,9 +832,9 @@ pub unsafe extern "C" fn ambulkdelete( }; } let callback = callback.unwrap(); - let callback = |p: Pointer| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; + let callback = |p: NonZeroU64| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; algorithm::vacuum::vacuum( - unsafe { Relation::new((*info).index) }, + unsafe { PostgresRelation::new((*info).index) }, || unsafe { pgrx::pg_sys::vacuum_delay_point(); }, diff --git a/src/vchordrqfscan/index/am_options.rs b/src/vchordrqfscan/index/am_options.rs index b49b7a21..be4154ef 100644 --- a/src/vchordrqfscan/index/am_options.rs +++ b/src/vchordrqfscan/index/am_options.rs @@ -1,14 +1,16 @@ use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use crate::datatype::typmod::Typmod; -use crate::vchordrqfscan::types::*; -use base::distance::*; -use base::vector::VectorBorrowed; +use crate::vchordrqfscan::types::{BorrowedVector, OwnedVector}; +use crate::vchordrqfscan::types::{DistanceKind, VectorKind}; +use crate::vchordrqfscan::types::{VchordrqfscanIndexingOptions, VectorOptions}; +use distance::Distance; use pgrx::datum::FromDatum; use pgrx::heap_tuple::PgHeapTuple; use serde::Deserialize; use std::ffi::CStr; use std::num::NonZero; +use vector::VectorBorrowed; #[derive(Copy, Clone, Debug, Default)] #[repr(C)] diff --git a/src/vchordrqfscan/index/am_scan.rs b/src/vchordrqfscan/index/am_scan.rs index b07edb7e..da049ab4 100644 --- a/src/vchordrqfscan/index/am_scan.rs +++ b/src/vchordrqfscan/index/am_scan.rs @@ -1,12 +1,12 @@ use super::am_options::Opfamily; -use crate::postgres::Relation; +use crate::postgres::PostgresRelation; use crate::vchordrqfscan::algorithm::scan::scan; use crate::vchordrqfscan::gucs::executing::epsilon; use crate::vchordrqfscan::gucs::executing::max_scan_tuples; use crate::vchordrqfscan::gucs::executing::probes; use crate::vchordrqfscan::types::OwnedVector; -use base::distance::Distance; -use base::search::*; +use distance::Distance; +use std::num::NonZeroU64; pub enum Scanner { Initial { @@ -15,7 +15,7 @@ pub enum Scanner { recheck: bool, }, Vbase { - vbase: Box>, + vbase: Box>, threshold: Option, recheck: bool, opfamily: Opfamily, @@ -62,7 +62,7 @@ pub fn scan_make( } } -pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, bool)> { +pub fn scan_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(NonZeroU64, bool)> { if let Scanner::Initial { vector, threshold, diff --git a/src/vchordrqfscan/index/functions.rs b/src/vchordrqfscan/index/functions.rs index 98f0f251..27bd9ac2 100644 --- a/src/vchordrqfscan/index/functions.rs +++ b/src/vchordrqfscan/index/functions.rs @@ -1,4 +1,4 @@ -use crate::postgres::Relation; +use crate::postgres::PostgresRelation; use crate::vchordrqfscan::algorithm::prewarm::prewarm; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass}; @@ -17,7 +17,7 @@ fn _vchordrqfscan_prewarm(indexrelid: Oid, height: i32) -> String { pgrx::error!("{:?} is not a vchordrqfscan index", pg_class.relname()); } let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; - let relation = unsafe { Relation::new(index) }; + let relation = unsafe { PostgresRelation::new(index) }; let message = prewarm(relation, height); unsafe { pgrx::pg_sys::index_close(index, pgrx::pg_sys::ShareLock as _); diff --git a/src/vchordrqfscan/index/utils.rs b/src/vchordrqfscan/index/utils.rs index a5d85a3f..726a5979 100644 --- a/src/vchordrqfscan/index/utils.rs +++ b/src/vchordrqfscan/index/utils.rs @@ -1,7 +1,7 @@ -use base::search::*; +use std::num::NonZeroU64; -pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { - let value = pointer.as_u64(); +pub fn pointer_to_ctid(pointer: NonZeroU64) -> pgrx::pg_sys::ItemPointerData { + let value = pointer.get(); pgrx::pg_sys::ItemPointerData { ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi: ((value >> 32) & 0xffff) as u16, @@ -11,10 +11,10 @@ pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { } } -pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer { +pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> NonZeroU64 { let mut value = 0; value |= (ctid.ip_blkid.bi_hi as u64) << 32; value |= (ctid.ip_blkid.bi_lo as u64) << 16; value |= ctid.ip_posid as u64; - Pointer::new(value) + NonZeroU64::new(value).expect("invalid pointer") } diff --git a/src/vchordrqfscan/types.rs b/src/vchordrqfscan/types.rs index 1180e649..91ca7690 100644 --- a/src/vchordrqfscan/types.rs +++ b/src/vchordrqfscan/types.rs @@ -1,7 +1,7 @@ -use base::distance::DistanceKind; -use base::vector::{VectBorrowed, VectOwned}; +use distance::Distance; use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError, ValidationErrors}; +use vector::vect::{VectBorrowed, VectOwned}; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] @@ -108,6 +108,13 @@ pub enum BorrowedVector<'a> { Vecf32(VectBorrowed<'a, f32>), } +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +pub enum DistanceKind { + L2, + Dot, +} + #[repr(u8)] #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] pub enum VectorKind { @@ -136,3 +143,11 @@ impl VectorOptions { } } } + +pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance { + use simd::Floating; + match d { + DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)), + DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)), + } +} diff --git a/tests/logic/reindex.slt b/tests/logic/reindex.slt index 7d763433..9a58049b 100644 --- a/tests/logic/reindex.slt +++ b/tests/logic/reindex.slt @@ -2,7 +2,7 @@ statement ok CREATE TABLE t (val vector(3)); statement ok -INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 10000); statement ok CREATE INDEX ON t USING vchordrq (val vector_l2_ops); diff --git a/tools/package.sh b/tools/package.sh index 3f6172ee..c6955ae2 100755 --- a/tools/package.sh +++ b/tools/package.sh @@ -5,7 +5,6 @@ printf "SEMVER = ${SEMVER}\n" printf "VERSION = ${VERSION}\n" printf "ARCH = ${ARCH}\n" printf "PLATFORM = ${PLATFORM}\n" -printf "PROFILE = ${PROFILE}\n" rm -rf ./build/dir_zip rm -rf ./build/vchord-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip @@ -13,9 +12,9 @@ rm -rf ./build/dir_deb rm -rf ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb mkdir -p ./build/dir_zip -cp ./target/${PROFILE}/schema.sql ./build/dir_zip/vchord--$SEMVER.sql +cp ./target/release/schema.sql ./build/dir_zip/vchord--$SEMVER.sql sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/dir_zip/vchord.control -cp ./target/${PROFILE}/libvchord.so ./build/dir_zip/vchord.so +cp ./target/release/libvchord.so ./build/dir_zip/vchord.so zip ./build/vchord-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip -j ./build/dir_zip/* mkdir -p ./build/dir_deb diff --git a/tools/schema.sh b/tools/schema.sh index 24a26d0a..c43f7f80 100755 --- a/tools/schema.sh +++ b/tools/schema.sh @@ -4,20 +4,12 @@ if [[ " $@ " =~ --target' '([^ ]+) ]]; then TARGET="${BASH_REMATCH[1]}" if [[ " $@ " =~ " --release " ]]; then DIR="./target/$TARGET/release" - elif [[ " $@ " =~ " --profile opt " ]]; then - DIR="./target/$TARGET/opt" - elif [[ " $@ " =~ " --profile release " ]]; then - DIR="./target/$TARGET/release" else DIR="./target/$TARGET/debug" fi else if [[ " $@ " =~ " --release " ]]; then DIR="./target/release" - elif [[ " $@ " =~ " --profile opt " ]]; then - DIR="./target/opt" - elif [[ " $@ " =~ " --profile release " ]]; then - DIR="./target/release" else DIR="./target/debug" fi From 89b49a23ba70660d246e1745b5884fbc74f2b5fa Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 8 Jan 2025 17:08:39 +0800 Subject: [PATCH 094/324] fix: pick feat back to vchordrqfscan (#162) pick #106, #147 and #149 to vchordrqfscan Signed-off-by: usamoi --- src/vchordrqfscan/algorithm/build.rs | 65 +++++--- src/vchordrqfscan/algorithm/insert.rs | 210 +++++++++++++++----------- src/vchordrqfscan/algorithm/tuples.rs | 1 - src/vchordrqfscan/index/am.rs | 4 + 4 files changed, 176 insertions(+), 104 deletions(-) diff --git a/src/vchordrqfscan/algorithm/build.rs b/src/vchordrqfscan/algorithm/build.rs index a245d776..e46d6a7d 100644 --- a/src/vchordrqfscan/algorithm/build.rs +++ b/src/vchordrqfscan/algorithm/build.rs @@ -72,10 +72,7 @@ pub fn build( }; let mut meta = Tape::create(&relation, false); assert_eq!(meta.first(), 0); - let mut forwards = Tape::::create(&relation, false); - assert_eq!(forwards.first(), 1); let mut vectors = Tape::create(&relation, true); - assert_eq!(vectors.first(), 2); let mut pointer_of_means = Vec::>::new(); for i in 0..structures.len() { let mut level = Vec::new(); @@ -166,13 +163,11 @@ pub fn build( } pointer_of_firsts.push(level); } - forwards.head.get_opaque_mut().skip = vectors.first(); meta.push(&MetaTuple { dims, height_of_root: structures.len() as u32, is_residual, vectors_first: vectors.first(), - forwards_first: forwards.first(), mean: pointer_of_means.last().unwrap()[0], first: pointer_of_firsts.last().unwrap()[0], }); @@ -244,20 +239,32 @@ impl Structure { ) -> Vec { use std::collections::BTreeMap; let VchordrqfscanExternalBuildOptions { table } = external_build; - let query = format!("SELECT id, parent, vector FROM {table};"); let mut parents = BTreeMap::new(); let mut vectors = BTreeMap::new(); pgrx::spi::Spi::connect(|client| { use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use pgrx::pg_sys::panic::ErrorReportable; use vector::VectorBorrowed; - let table = client.select(&query, None, None).unwrap_or_report(); - for row in table { + let schema_query = "SELECT n.nspname::TEXT + FROM pg_catalog.pg_extension e + LEFT JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vector';"; + let pgvector_schema: String = client + .select(schema_query, None, None) + .unwrap_or_report() + .first() + .get_by_name("nspname") + .expect("external build: cannot get schema of pgvector") + .expect("external build: cannot get schema of pgvector"); + let dump_query = + format!("SELECT id, parent, vector::{pgvector_schema}.vector FROM {table};"); + let centroids = client.select(&dump_query, None, None).unwrap_or_report(); + for row in centroids { let id: Option = row.get_by_name("id").unwrap(); let parent: Option = row.get_by_name("parent").unwrap(); let vector: Option = row.get_by_name("vector").unwrap(); - let id = id.expect("extern build: id could not be NULL"); - let vector = vector.expect("extern build: vector could not be NULL"); + let id = id.expect("external build: id could not be NULL"); + let vector = vector.expect("external build: vector could not be NULL"); let pop = parents.insert(id, parent); if pop.is_some() { pgrx::error!( @@ -265,11 +272,34 @@ impl Structure { ); } if vector_options.dims != vector.as_borrowed().dims() { - pgrx::error!("extern build: incorrect dimension, id = {id}"); + pgrx::error!("external build: incorrect dimension, id = {id}"); } vectors.insert(id, crate::projection::project(vector.as_borrowed().slice())); } }); + if parents.len() >= 2 && parents.values().all(|x| x.is_none()) { + // if there are more than one vertexs and no edges, + // assume there is an implicit root + let n = parents.len(); + let mut result = Vec::new(); + result.push(Structure { + means: vectors.values().cloned().collect::>(), + children: vec![Vec::new(); n], + }); + result.push(Structure { + means: vec![{ + // compute the vector on root, without normalizing it + let mut sum = vec![0.0f32; vector_options.dims as _]; + for vector in vectors.values() { + f32::vector_add_inplace(&mut sum, vector); + } + f32::vector_mul_scalar_inplace(&mut sum, 1.0 / n as f32); + sum + }], + children: vec![(0..n as u32).collect()], + }); + return result; + } let mut children = parents .keys() .map(|x| (*x, Vec::new())) @@ -293,7 +323,7 @@ impl Structure { } } let Some(root) = root else { - pgrx::error!("extern build: there are no root"); + pgrx::error!("external build: there are no root"); }; let mut heights = BTreeMap::<_, _>::new(); fn dfs_for_heights( @@ -302,7 +332,7 @@ impl Structure { u: i32, ) { if heights.contains_key(&u) { - pgrx::error!("extern build: detect a cycle, id = {u}"); + pgrx::error!("external build: detect a cycle, id = {u}"); } heights.insert(u, None); let mut height = None; @@ -311,7 +341,7 @@ impl Structure { let new = heights[&v].unwrap() + 1; if let Some(height) = height { if height != new { - pgrx::error!("extern build: two heights, id = {u}"); + pgrx::error!("external build: two heights, id = {u}"); } } else { height = Some(new); @@ -329,7 +359,7 @@ impl Structure { .collect::>(); if !(1..=8).contains(&(heights[&root] - 1)) { pgrx::error!( - "extern build: unexpected tree height, height = {}", + "external build: unexpected tree height, height = {}", heights[&root] ); } @@ -367,7 +397,7 @@ impl Structure { } } -struct Tape<'a, 'b, T, R: 'b + RelationWrite> { +struct Tape<'a: 'b, 'b, T, R: 'b + RelationWrite> { relation: &'a R, head: R::WriteGuard<'b>, first: u32, @@ -377,7 +407,8 @@ struct Tape<'a, 'b, T, R: 'b + RelationWrite> { impl<'a: 'b, 'b, T, R: 'b + RelationWrite> Tape<'a, 'b, T, R> { fn create(relation: &'a R, tracking_freespace: bool) -> Self { - let head = relation.extend(tracking_freespace); + let mut head = relation.extend(tracking_freespace); + head.get_opaque_mut().skip = head.id(); let first = head.id(); Self { relation, diff --git a/src/vchordrqfscan/algorithm/insert.rs b/src/vchordrqfscan/algorithm/insert.rs index ee447601..cb30577f 100644 --- a/src/vchordrqfscan/algorithm/insert.rs +++ b/src/vchordrqfscan/algorithm/insert.rs @@ -12,10 +12,11 @@ use std::collections::BinaryHeap; use std::num::NonZeroU64; pub fn insert( - relation: impl RelationWrite, + relation: impl RelationWrite + Clone, payload: NonZeroU64, vector: Vec, distance_kind: DistanceKind, + in_building: bool, ) { let meta_guard = relation.read(0); let meta_tuple = meta_guard @@ -32,56 +33,20 @@ pub fn insert( } else { None }; - let h0_vector = 'h0_vector: { + let h0_vector = { let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { vector: vector.clone(), payload: Some(payload), }) .unwrap(); - if let Some(mut write) = relation.search(tuple.len()) { - let i = write.alloc(&tuple).unwrap(); - break 'h0_vector (write.id(), i); - } - let mut current = relation.read(meta_tuple.forwards_first).get_opaque().skip; - let mut changed = false; - loop { - let read = relation.read(current); - let flag = 'flag: { - if read.freespace() as usize >= tuple.len() { - break 'flag true; - } - if read.get_opaque().next == u32::MAX { - break 'flag true; - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current, true); - if let Some(i) = write.alloc(&tuple) { - break (current, i); - } - if write.get_opaque().next == u32::MAX { - if changed { - relation - .write(meta_tuple.forwards_first, false) - .get_opaque_mut() - .skip = write.id(); - } - let mut extend = relation.extend(true); - write.get_opaque_mut().next = extend.id(); - if let Some(i) = extend.alloc(&tuple) { - break (extend.id(), i); - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; - } - changed = true; - } + append( + relation.clone(), + meta_tuple.vectors_first, + &tuple, + true, + true, + true, + ) }; let h0_payload = payload; let mut list = ( @@ -186,23 +151,97 @@ pub fn insert( t: vec![0; (dims.div_ceil(4) * 16) as usize], }) .unwrap(); - let first = list.0; + append_by_update( + relation.clone(), + list.0, + &dummy, + in_building, + in_building, + |bytes| { + let t = rkyv::check_archived_root::(bytes).expect("data corruption"); + t.mask.iter().any(|x| *x) + }, + |bytes| put(bytes, dims, &code, h0_vector, h0_payload), + ); +} + +fn append( + relation: impl RelationWrite, + first: u32, + tuple: &[u8], + tracking_freespace: bool, + skipping_traversal: bool, + updating_skip: bool, +) -> (u32, u16) { + if tracking_freespace { + if let Some(mut write) = relation.search(tuple.len()) { + let i = write.alloc(tuple).unwrap(); + return (write.id(), i); + } + } + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = relation.read(current); + if read.freespace() as usize >= tuple.len() || read.get_opaque().next == u32::MAX { + drop(read); + let mut write = relation.write(current, tracking_freespace); + if let Some(i) = write.alloc(tuple) { + return (current, i); + } + if write.get_opaque().next == u32::MAX { + let mut extend = relation.extend(tracking_freespace); + write.get_opaque_mut().next = extend.id(); + drop(write); + if let Some(i) = extend.alloc(tuple) { + let result = (extend.id(), i); + drop(extend); + if updating_skip { + let mut past = relation.write(first, tracking_freespace); + let skip = &mut past.get_opaque_mut().skip; + assert!(*skip != u32::MAX); + *skip = std::cmp::max(*skip, result.0); + } + return result; + } else { + panic!("a tuple cannot even be fit in a fresh page"); + } + } + if skipping_traversal && current == first && write.get_opaque().skip != first { + current = write.get_opaque().skip; + } else { + current = write.get_opaque().next; + } + } else { + if skipping_traversal && current == first && read.get_opaque().skip != first { + current = read.get_opaque().skip; + } else { + current = read.get_opaque().next; + } + } + } +} + +fn append_by_update( + relation: impl RelationWrite, + first: u32, + tuple: &[u8], + skipping_traversal: bool, + updating_skip: bool, + can_update: impl Fn(&[u8]) -> bool, + mut update: impl FnMut(&mut [u8]) -> bool, +) { assert!(first != u32::MAX); let mut current = first; loop { let read = relation.read(current); let flag = 'flag: { for i in 1..=read.len() { - let h0_tuple = read - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - if h0_tuple.mask.iter().any(|x| *x) { + if can_update(read.get(i).expect("data corruption")) { break 'flag true; } } - if read.freespace() as usize >= dummy.len() { + if read.freespace() as usize >= tuple.len() { break 'flag true; } if read.get_opaque().next == u32::MAX { @@ -214,48 +253,47 @@ pub fn insert( drop(read); let mut write = relation.write(current, false); for i in 1..=write.len() { - let flag = put( - write.get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - if flag { + if update(write.get_mut(i).expect("data corruption")) { return; } } - if let Some(i) = write.alloc(&dummy) { - let flag = put( - write.get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; + if let Some(i) = write.alloc(tuple) { + if update(write.get_mut(i).expect("data corruption")) { + return; + } + panic!("an update fails on a fresh tuple"); } if write.get_opaque().next == u32::MAX { let mut extend = relation.extend(false); write.get_opaque_mut().next = extend.id(); - if let Some(i) = extend.alloc(&dummy) { - let flag = put( - extend.get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; - } else { - panic!("a tuple cannot even be fit in a fresh page"); + drop(write); + if let Some(i) = extend.alloc(tuple) { + if update(extend.get_mut(i).expect("data corruption")) { + let id = extend.id(); + drop(extend); + if updating_skip { + let mut past = relation.write(first, false); + let skip = &mut past.get_opaque_mut().skip; + assert!(*skip != u32::MAX); + *skip = std::cmp::max(*skip, id); + } + return; + } + panic!("an update fails on a fresh tuple"); } + panic!("a tuple cannot even be fit in a fresh page"); + } + if skipping_traversal && current == first && write.get_opaque().skip != first { + current = write.get_opaque().skip; + } else { + current = write.get_opaque().next; } - current = write.get_opaque().next; } else { - current = read.get_opaque().next; + if skipping_traversal && current == first && read.get_opaque().skip != first { + current = read.get_opaque().skip; + } else { + current = read.get_opaque().next; + } } } } diff --git a/src/vchordrqfscan/algorithm/tuples.rs b/src/vchordrqfscan/algorithm/tuples.rs index ff947131..7d1c97af 100644 --- a/src/vchordrqfscan/algorithm/tuples.rs +++ b/src/vchordrqfscan/algorithm/tuples.rs @@ -10,7 +10,6 @@ pub struct MetaTuple { pub height_of_root: u32, pub is_residual: bool, pub vectors_first: u32, - pub forwards_first: u32, // raw vector pub mean: (u32, u16), // for meta tuple, it's pointers to next level diff --git a/src/vchordrqfscan/index/am.rs b/src/vchordrqfscan/index/am.rs index 7011e26f..339b315c 100644 --- a/src/vchordrqfscan/index/am.rs +++ b/src/vchordrqfscan/index/am.rs @@ -287,6 +287,7 @@ pub unsafe extern "C" fn ambuild( payload, vector, opfamily.distance_kind(), + true, ); indtuples += 1; reporter.tuples_done(indtuples); @@ -624,6 +625,7 @@ unsafe fn parallel_build( payload, vector, opfamily.distance_kind(), + true, ); unsafe { let indtuples; @@ -676,6 +678,7 @@ pub unsafe extern "C" fn aminsert( pointer, vector.into_vec(), opfamily.distance_kind(), + false, ); } false @@ -706,6 +709,7 @@ pub unsafe extern "C" fn aminsert( pointer, vector.into_vec(), opfamily.distance_kind(), + false, ); } false From cb5108f6de6de5969c17c5b4ede7100577ab9d4b Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 17 Jan 2025 20:12:01 +0800 Subject: [PATCH 095/324] chore: fix discord and x badge (#170) Signed-off-by: Keming --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 0cff5bb9..8b2f92c2 100644 --- a/README.md +++ b/README.md @@ -4,8 +4,8 @@

-discord invitation link -Twitter +discord invitation link +Twitter Docker pulls

Docker pull for pgvecto.rs: Previous Docker pulls

From dce791f696642405c64e2ae41605dce8f1b2319b Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 20 Jan 2025 15:13:09 +0800 Subject: [PATCH 096/324] feat: unify vchordrq and vchordrqfscan (#167) closes #164 closes #163 Signed-off-by: usamoi --- Cargo.lock | 273 +--- Cargo.toml | 12 +- crates/algorithm/Cargo.toml | 7 - crates/rabitq/src/binary.rs | 72 +- crates/rabitq/src/block.rs | 160 +-- crates/rabitq/src/lib.rs | 94 +- crates/rabitq/src/utils.rs | 20 - crates/simd/src/f16.rs | 4 +- crates/simd/src/fast_scan/mod.rs | 306 +++-- src/{vchordrq => }/algorithm/build.rs | 139 +- src/algorithm/freepages.rs | 61 + src/algorithm/insert.rs | 190 +++ .../src/lib.rs => src/algorithm/mod.rs | 13 +- src/algorithm/operator.rs | 628 +++++++++ src/algorithm/prewarm.rs | 99 ++ src/algorithm/scan.rs | 171 +++ src/algorithm/tape.rs | 349 +++++ src/algorithm/tuples.rs | 1205 +++++++++++++++++ src/algorithm/vacuum.rs | 311 +++++ src/algorithm/vectors.rs | 133 ++ src/{vchordrq => }/gucs/executing.rs | 0 src/{vchordrq => }/gucs/mod.rs | 0 src/{vchordrq => }/gucs/prewarm.rs | 0 src/{vchordrq => }/index/am.rs | 448 ++++-- src/{vchordrq => }/index/am_options.rs | 6 +- src/{vchordrq => }/index/am_scan.rs | 64 +- src/{vchordrq => }/index/functions.rs | 22 +- src/{vchordrq => }/index/mod.rs | 0 src/{vchordrq => }/index/opclass.rs | 0 src/index/utils.rs | 34 + src/lib.rs | 11 +- src/postgres.rs | 55 +- src/sql/finalize.sql | 30 - src/{vchordrq => }/types.rs | 0 src/utils/k_means.rs | 76 +- src/utils/mod.rs | 1 + src/utils/pipe.rs | 14 + src/vchordrq/algorithm/insert.rs | 221 --- src/vchordrq/algorithm/mod.rs | 8 - src/vchordrq/algorithm/prewarm.rs | 82 -- src/vchordrq/algorithm/rabitq.rs | 21 - src/vchordrq/algorithm/scan.rs | 189 --- src/vchordrq/algorithm/tuples.rs | 270 ---- src/vchordrq/algorithm/vacuum.rs | 112 -- src/vchordrq/algorithm/vectors.rs | 76 -- src/vchordrq/index/utils.rs | 20 - src/vchordrq/mod.rs | 11 - src/vchordrqfscan/algorithm/build.rs | 445 ------ src/vchordrqfscan/algorithm/insert.rs | 299 ---- src/vchordrqfscan/algorithm/mod.rs | 7 - src/vchordrqfscan/algorithm/prewarm.rs | 102 -- src/vchordrqfscan/algorithm/rabitq.rs | 22 - src/vchordrqfscan/algorithm/scan.rs | 193 --- src/vchordrqfscan/algorithm/tuples.rs | 140 -- src/vchordrqfscan/algorithm/vacuum.rs | 139 -- src/vchordrqfscan/gucs/executing.rs | 72 - src/vchordrqfscan/gucs/mod.rs | 14 - src/vchordrqfscan/gucs/prewarm.rs | 32 - src/vchordrqfscan/index/am.rs | 856 ------------ src/vchordrqfscan/index/am_options.rs | 218 --- src/vchordrqfscan/index/am_scan.rs | 125 -- src/vchordrqfscan/index/functions.rs | 26 - src/vchordrqfscan/index/mod.rs | 12 - src/vchordrqfscan/index/opclass.rs | 14 - src/vchordrqfscan/index/utils.rs | 20 - src/vchordrqfscan/mod.rs | 11 - src/vchordrqfscan/types.rs | 153 --- 67 files changed, 4144 insertions(+), 4774 deletions(-) delete mode 100644 crates/algorithm/Cargo.toml delete mode 100644 crates/rabitq/src/utils.rs rename src/{vchordrq => }/algorithm/build.rs (78%) create mode 100644 src/algorithm/freepages.rs create mode 100644 src/algorithm/insert.rs rename crates/algorithm/src/lib.rs => src/algorithm/mod.rs (85%) create mode 100644 src/algorithm/operator.rs create mode 100644 src/algorithm/prewarm.rs create mode 100644 src/algorithm/scan.rs create mode 100644 src/algorithm/tape.rs create mode 100644 src/algorithm/tuples.rs create mode 100644 src/algorithm/vacuum.rs create mode 100644 src/algorithm/vectors.rs rename src/{vchordrq => }/gucs/executing.rs (100%) rename src/{vchordrq => }/gucs/mod.rs (100%) rename src/{vchordrq => }/gucs/prewarm.rs (100%) rename src/{vchordrq => }/index/am.rs (67%) rename src/{vchordrq => }/index/am_options.rs (97%) rename src/{vchordrq => }/index/am_scan.rs (64%) rename src/{vchordrq => }/index/functions.rs (58%) rename src/{vchordrq => }/index/mod.rs (100%) rename src/{vchordrq => }/index/opclass.rs (100%) create mode 100644 src/index/utils.rs rename src/{vchordrq => }/types.rs (100%) create mode 100644 src/utils/pipe.rs delete mode 100644 src/vchordrq/algorithm/insert.rs delete mode 100644 src/vchordrq/algorithm/mod.rs delete mode 100644 src/vchordrq/algorithm/prewarm.rs delete mode 100644 src/vchordrq/algorithm/rabitq.rs delete mode 100644 src/vchordrq/algorithm/scan.rs delete mode 100644 src/vchordrq/algorithm/tuples.rs delete mode 100644 src/vchordrq/algorithm/vacuum.rs delete mode 100644 src/vchordrq/algorithm/vectors.rs delete mode 100644 src/vchordrq/index/utils.rs delete mode 100644 src/vchordrq/mod.rs delete mode 100644 src/vchordrqfscan/algorithm/build.rs delete mode 100644 src/vchordrqfscan/algorithm/insert.rs delete mode 100644 src/vchordrqfscan/algorithm/mod.rs delete mode 100644 src/vchordrqfscan/algorithm/prewarm.rs delete mode 100644 src/vchordrqfscan/algorithm/rabitq.rs delete mode 100644 src/vchordrqfscan/algorithm/scan.rs delete mode 100644 src/vchordrqfscan/algorithm/tuples.rs delete mode 100644 src/vchordrqfscan/algorithm/vacuum.rs delete mode 100644 src/vchordrqfscan/gucs/executing.rs delete mode 100644 src/vchordrqfscan/gucs/mod.rs delete mode 100644 src/vchordrqfscan/gucs/prewarm.rs delete mode 100644 src/vchordrqfscan/index/am.rs delete mode 100644 src/vchordrqfscan/index/am_options.rs delete mode 100644 src/vchordrqfscan/index/am_scan.rs delete mode 100644 src/vchordrqfscan/index/functions.rs delete mode 100644 src/vchordrqfscan/index/mod.rs delete mode 100644 src/vchordrqfscan/index/opclass.rs delete mode 100644 src/vchordrqfscan/index/utils.rs delete mode 100644 src/vchordrqfscan/mod.rs delete mode 100644 src/vchordrqfscan/types.rs diff --git a/Cargo.lock b/Cargo.lock index 5d17e1a4..f5214d8e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,17 +2,6 @@ # It is not intended for manual editing. version = 4 -[[package]] -name = "ahash" -version = "0.7.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "891477e0c6a8957309ee5c45a6368af3ae14bb510732d2684ffa19af310920f9" -dependencies = [ - "getrandom", - "once_cell", - "version_check", -] - [[package]] name = "aho-corasick" version = "1.1.3" @@ -22,10 +11,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "algorithm" -version = "0.0.0" - [[package]] name = "always_equal" version = "0.0.0" @@ -87,14 +72,14 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.94", + "syn", ] [[package]] name = "bitflags" -version = "2.6.0" +version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" +checksum = "1be3f42a67d6d345ecd59f675f3f012d6974981560836e938c22b424b85ce1be" [[package]] name = "bitvec" @@ -108,28 +93,6 @@ dependencies = [ "wyz", ] -[[package]] -name = "bytecheck" -version = "0.6.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23cdc57ce23ac53c931e88a43d06d070a6fd142f2617be5855eb75efc9beb1c2" -dependencies = [ - "bytecheck_derive", - "ptr_meta", - "simdutf8", -] - -[[package]] -name = "bytecheck_derive" -version = "0.6.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3db406d29fbcd95542e92559bed4d8ad92636d1ca8b3b72ede10b4bcc010e659" -dependencies = [ - "proc-macro2", - "quote", - "syn 1.0.109", -] - [[package]] name = "bytemuck" version = "1.21.0" @@ -142,12 +105,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" -[[package]] -name = "bytes" -version = "1.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "325918d6fe32f23b19878fe4b34794ae41fc19ddbe53b10571a4874d44ffd39b" - [[package]] name = "cargo_toml" version = "0.19.2" @@ -160,9 +117,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.6" +version = "1.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d6dbb628b8f8555f86d0323c2eb39e3ec81901f4b83e091db8a6a76d316a333" +checksum = "c8293772165d9345bdaaa39b45b2109591e63fe5e6fbc23c6ff930a048aa310b" dependencies = [ "shlex", ] @@ -264,7 +221,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.94", + "syn", ] [[package]] @@ -275,7 +232,7 @@ checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" dependencies = [ "darling_core", "quote", - "syn 2.0.94", + "syn", ] [[package]] @@ -286,7 +243,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] @@ -316,7 +273,7 @@ checksum = "f282cfdfe92516eb26c2af8589c274c7c17681f5ecc03c18255fe741c6aa64eb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] @@ -388,12 +345,13 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403" [[package]] name = "half" version = "2.4.1" -source = "git+https://github.com/tensorchord/half-rs.git#2d8a66092bee436aebb26ea7ac47d11150cda31d" +source = "git+https://github.com/tensorchord/half-rs.git?rev=3f9a8843d6722bd1833de2289347640ad8770146#3f9a8843d6722bd1833de2289347640ad8770146" dependencies = [ "cfg-if", "crunchy", - "rkyv", "serde", + "zerocopy 0.8.14", + "zerocopy-derive 0.8.14", ] [[package]] @@ -405,15 +363,6 @@ dependencies = [ "byteorder", ] -[[package]] -name = "hashbrown" -version = "0.12.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" -dependencies = [ - "ahash", -] - [[package]] name = "hashbrown" version = "0.15.2" @@ -560,7 +509,7 @@ checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] @@ -603,7 +552,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "62f822373a4fe84d4bb149bf54e584a7f4abec90e072ed49cda0edea5b95471f" dependencies = [ "equivalent", - "hashbrown 0.15.2", + "hashbrown", ] [[package]] @@ -668,9 +617,9 @@ checksum = "4ee93343901ab17bd981295f2cf0026d4ad018c7c31ba84549a4ddbb47a45104" [[package]] name = "log" -version = "0.4.22" +version = "0.4.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" +checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f" [[package]] name = "matrixmultiply" @@ -718,7 +667,7 @@ checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] @@ -825,7 +774,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b7cafe60d6cf8e62e1b9b2ea516a089c008945bb5a275416789e7db0bc199dc" dependencies = [ "memchr", - "thiserror 2.0.9", + "thiserror 2.0.11", "ucd-trie", ] @@ -877,7 +826,7 @@ dependencies = [ "proc-macro2", "quote", "shlex", - "syn 2.0.94", + "syn", "walkdir", ] @@ -900,7 +849,7 @@ dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] @@ -947,7 +896,7 @@ dependencies = [ "petgraph", "proc-macro2", "quote", - "syn 2.0.94", + "syn", "thiserror 1.0.69", "unescape", ] @@ -958,7 +907,7 @@ version = "0.2.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" dependencies = [ - "zerocopy", + "zerocopy 0.7.35", ] [[package]] @@ -980,38 +929,18 @@ dependencies = [ "proc-macro-error-attr2", "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] name = "proc-macro2" -version = "1.0.92" +version = "1.0.93" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0" +checksum = "60946a68e5f9d28b0dc1c21bb8a97ee7d018a8b322fa57838ba31cc878e22d99" dependencies = [ "unicode-ident", ] -[[package]] -name = "ptr_meta" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0738ccf7ea06b608c10564b31debd4f5bc5e197fc8bfe088f68ae5ce81e7a4f1" -dependencies = [ - "ptr_meta_derive", -] - -[[package]] -name = "ptr_meta_derive" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16b845dbfca988fa33db069c0e230574d15a3088f147a87b64c7589eb662c9ac" -dependencies = [ - "proc-macro2", - "quote", - "syn 1.0.109", -] - [[package]] name = "quote" version = "1.0.38" @@ -1140,44 +1069,6 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" -[[package]] -name = "rend" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71fe3824f5629716b1589be05dacd749f6aa084c87e00e016714a8cdfccc997c" -dependencies = [ - "bytecheck", -] - -[[package]] -name = "rkyv" -version = "0.7.45" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9008cd6385b9e161d8229e1f6549dd23c3d022f132a2ea37ac3a10ac4935779b" -dependencies = [ - "bitvec", - "bytecheck", - "bytes", - "hashbrown 0.12.3", - "ptr_meta", - "rend", - "rkyv_derive", - "seahash", - "tinyvec", - "uuid", -] - -[[package]] -name = "rkyv_derive" -version = "0.7.45" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "503d1d27590a2b0a3a4ca4c94755aa2875657196ecbf401a42eff41d7de532c0" -dependencies = [ - "proc-macro2", - "quote", - "syn 1.0.109", -] - [[package]] name = "rustc-hash" version = "1.1.0" @@ -1268,14 +1159,14 @@ checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] name = "serde_json" -version = "1.0.134" +version = "1.0.135" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d00f4175c42ee48b15416f6193a959ba3a0d67fc699a0db9ad12df9f83991c7d" +checksum = "2b0d7ba2887406110130a978386c4e1befb98c674b4fba677954e4db976630d9" dependencies = [ "itoa", "memchr", @@ -1328,15 +1219,9 @@ version = "0.0.0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] -[[package]] -name = "simdutf8" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" - [[package]] name = "smallvec" version = "1.13.2" @@ -1382,20 +1267,9 @@ dependencies = [ [[package]] name = "syn" -version = "1.0.109" +version = "2.0.96" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "syn" -version = "2.0.94" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "987bc0be1cdea8b10216bd06e2ca407d40b9543468fafd3ddfb02f36e77f71f3" +checksum = "d5d0adab1ae378d7f53bdebc67a39f1f151407ef230f0ce2883572f5d8985c80" dependencies = [ "proc-macro2", "quote", @@ -1410,7 +1284,7 @@ checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] @@ -1430,11 +1304,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.9" +version = "2.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f072643fd0190df67a8bab670c20ef5d8737177d6ac6b2e9a236cb096206b2cc" +checksum = "d452f284b73e6d76dd36758a0c8684b1d5be31f92b89d07fd5822175732206fc" dependencies = [ - "thiserror-impl 2.0.9", + "thiserror-impl 2.0.11", ] [[package]] @@ -1445,18 +1319,18 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] name = "thiserror-impl" -version = "2.0.9" +version = "2.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b50fa271071aae2e6ee85f842e2e28ba8cd2c5fb67f11fcb1fd70b276f9e7d4" +checksum = "26afc1baea8a989337eeb52b6e72a039780ce45c3edfcc9c5b9d112feeb173c2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] @@ -1469,21 +1343,6 @@ dependencies = [ "zerovec", ] -[[package]] -name = "tinyvec" -version = "1.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "022db8904dfa342efe721985167e9fcd16c29b226db4397ed752a761cfce81e8" -dependencies = [ - "tinyvec_macros", -] - -[[package]] -name = "tinyvec_macros" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" - [[package]] name = "toml" version = "0.8.19" @@ -1579,9 +1438,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.11.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8c5f0a0af699448548ad1a2fbf920fb4bee257eae39953ba95cb84891a0446a" +checksum = "b913a3b5fe84142e269d63cc62b64319ccaf89b748fc31fe025177f767a756c4" dependencies = [ "getrandom", ] @@ -1613,14 +1472,13 @@ dependencies = [ "proc-macro-error2", "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] [[package]] name = "vchord" version = "0.0.0" dependencies = [ - "algorithm", "always_equal", "distance", "half 2.4.1", @@ -1632,12 +1490,13 @@ dependencies = [ "rand", "random_orthogonal_matrix", "rayon", - "rkyv", "serde", "simd", "toml", "validator", "vector", + "zerocopy 0.8.14", + "zerocopy-derive 0.8.14", ] [[package]] @@ -1650,12 +1509,6 @@ dependencies = [ "simd", ] -[[package]] -name = "version_check" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" - [[package]] name = "walkdir" version = "2.5.0" @@ -1674,9 +1527,9 @@ checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" [[package]] name = "wide" -version = "0.7.30" +version = "0.7.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "58e6db2670d2be78525979e9a5f9c69d296fd7d670549fe9ebf70f8708cb5019" +checksum = "41b5576b9a81633f3e8df296ce0063042a73507636cbe956c61133dd7034ab22" dependencies = [ "bytemuck", "safe_arch", @@ -1797,9 +1650,9 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "winnow" -version = "0.6.21" +version = "0.6.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6f5bb5257f2407a5425c6e749bfd9692192a73e70a6060516ac04f889087d68" +checksum = "c8d71a593cc5c42ad7876e2c1fda56f314f3754c084128833e64f1345ff8a03a" dependencies = [ "memchr", ] @@ -1854,7 +1707,7 @@ checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", "synstructure", ] @@ -1865,7 +1718,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" dependencies = [ "byteorder", - "zerocopy-derive", + "zerocopy-derive 0.7.35", +] + +[[package]] +name = "zerocopy" +version = "0.8.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a367f292d93d4eab890745e75a778da40909cab4d6ff8173693812f79c4a2468" +dependencies = [ + "zerocopy-derive 0.8.14", ] [[package]] @@ -1876,7 +1738,18 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3931cb58c62c13adec22e38686b559c86a30565e16ad6e8510a337cedc611e1" +dependencies = [ + "proc-macro2", + "quote", + "syn", ] [[package]] @@ -1896,7 +1769,7 @@ checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", "synstructure", ] @@ -1919,5 +1792,5 @@ checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.94", + "syn", ] diff --git a/Cargo.toml b/Cargo.toml index 13272c2d..3583ae26 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,7 +20,6 @@ pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] -algorithm = { path = "./crates/algorithm" } always_equal = { path = "./crates/always_equal" } distance = { path = "./crates/distance" } rabitq = { path = "./crates/rabitq" } @@ -28,11 +27,8 @@ random_orthogonal_matrix = { path = "./crates/random_orthogonal_matrix" } simd = { path = "./crates/simd" } vector = { path = "./crates/vector" } -# lock rkyv version forever so that data is always compatible -rkyv = { version = "=0.7.45", features = ["validation"] } - half.workspace = true -log = "0.4.22" +log = "0.4.25" paste = "1" pgrx = { version = "=0.12.9", default-features = false, features = ["cshim"] } pgrx-catalog = "0.1.0" @@ -41,9 +37,11 @@ rayon = "1.10.0" serde.workspace = true toml = "0.8.19" validator = { version = "0.19.0", features = ["derive"] } +zerocopy = "0.8.14" +zerocopy-derive = "0.8.14" [patch.crates-io] -half = { git = "https://github.com/tensorchord/half-rs.git" } +half = { git = "https://github.com/tensorchord/half-rs.git", rev = "3f9a8843d6722bd1833de2289347640ad8770146" } [lints] workspace = true @@ -57,7 +55,7 @@ version = "0.0.0" edition = "2021" [workspace.dependencies] -half = { version = "2.4.1", features = ["rkyv", "serde"] } +half = { version = "2.4.1", features = ["serde", "zerocopy"] } rand = "0.8.5" serde = "1" diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml deleted file mode 100644 index 56da6cec..00000000 --- a/crates/algorithm/Cargo.toml +++ /dev/null @@ -1,7 +0,0 @@ -[package] -name = "algorithm" -version.workspace = true -edition.workspace = true - -[lints] -workspace = true diff --git a/crates/rabitq/src/binary.rs b/crates/rabitq/src/binary.rs index 8916ffd9..04bcb6f8 100644 --- a/crates/rabitq/src/binary.rs +++ b/crates/rabitq/src/binary.rs @@ -1,67 +1,9 @@ use distance::Distance; use simd::Floating; -#[derive(Debug, Clone)] -pub struct Code { - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub signs: Vec, -} - -impl Code { - pub fn t(&self) -> Vec { - use crate::utils::InfiniteByteChunks; - let mut result = Vec::new(); - for x in InfiniteByteChunks::<_, 64>::new(self.signs.iter().copied()) - .take(self.signs.len().div_ceil(64)) - { - let mut r = 0_u64; - for i in 0..64 { - r |= (x[i] as u64) << i; - } - result.push(r); - } - result - } -} - -pub fn code(dims: u32, vector: &[f32]) -> Code { - let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); - let sum_of_x_2 = f32::reduce_sum_of_x2(vector); - let dis_u = sum_of_x_2.sqrt(); - let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); - let x_x0 = dis_u / x0; - let fac_norm = (dims as f32).sqrt(); - let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); - let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); - let factor_ip = -2.0f32 / fac_norm * x_x0; - let cnt_pos = vector - .iter() - .map(|x| x.is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.is_sign_negative() as i32) - .sum::(); - let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; - let mut signs = Vec::new(); - for i in 0..dims { - signs.push(vector[i as usize].is_sign_positive() as u8); - } - Code { - dis_u_2: sum_of_x_2, - factor_ppc, - factor_ip, - factor_err, - signs, - } -} - -pub type Lut = (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)); - -pub fn preprocess(vector: &[f32]) -> Lut { +pub fn preprocess( + vector: &[f32], +) -> (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)) { let dis_v_2 = f32::reduce_sum_of_x2(vector); let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); let qvector_sum = if vector.len() <= 4369 { @@ -73,8 +15,7 @@ pub fn preprocess(vector: &[f32]) -> Lut { } pub fn process_lowerbound_l2( - _: u32, - lut: &Lut, + lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), (dis_u_2, factor_ppc, factor_ip, factor_err, t): (f32, f32, f32, f32, &[u64]), epsilon: f32, ) -> Distance { @@ -87,8 +28,7 @@ pub fn process_lowerbound_l2( } pub fn process_lowerbound_dot( - _: u32, - lut: &Lut, + lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), (_, factor_ppc, factor_ip, factor_err, t): (f32, f32, f32, f32, &[u64]), epsilon: f32, ) -> Distance { @@ -99,7 +39,7 @@ pub fn process_lowerbound_dot( Distance::from_f32(rough - epsilon * err) } -fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { +pub fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { let n = vector.len(); let mut t0 = vec![0u64; n.div_ceil(64)]; let mut t1 = vec![0u64; n.div_ceil(64)]; diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs index 0a8dab70..9f26fce1 100644 --- a/crates/rabitq/src/block.rs +++ b/crates/rabitq/src/block.rs @@ -1,85 +1,7 @@ use distance::Distance; use simd::Floating; -#[derive(Debug, Clone)] -pub struct Code { - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub signs: Vec, -} - -pub fn code(dims: u32, vector: &[f32]) -> Code { - let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); - let sum_of_x_2 = f32::reduce_sum_of_x2(vector); - let dis_u = sum_of_x_2.sqrt(); - let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); - let x_x0 = dis_u / x0; - let fac_norm = (dims as f32).sqrt(); - let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); - let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); - let factor_ip = -2.0f32 / fac_norm * x_x0; - let cnt_pos = vector - .iter() - .map(|x| x.is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.is_sign_negative() as i32) - .sum::(); - let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; - let mut signs = Vec::new(); - for i in 0..dims { - signs.push(vector[i as usize].is_sign_positive() as u8); - } - Code { - dis_u_2: sum_of_x_2, - factor_ppc, - factor_ip, - factor_err, - signs, - } -} - -pub fn dummy_code(dims: u32) -> Code { - Code { - dis_u_2: 0.0, - factor_ppc: 0.0, - factor_ip: 0.0, - factor_err: 0.0, - signs: vec![0; dims as _], - } -} - -pub struct PackedCodes { - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, -} - -pub fn pack_codes(dims: u32, codes: [Code; 32]) -> PackedCodes { - use crate::utils::InfiniteByteChunks; - PackedCodes { - dis_u_2: std::array::from_fn(|i| codes[i].dis_u_2), - factor_ppc: std::array::from_fn(|i| codes[i].factor_ppc), - factor_ip: std::array::from_fn(|i| codes[i].factor_ip), - factor_err: std::array::from_fn(|i| codes[i].factor_err), - t: { - let signs = codes.map(|code| { - InfiniteByteChunks::new(code.signs.into_iter()) - .map(|[b0, b1, b2, b3]| b0 | b1 << 1 | b2 << 2 | b3 << 3) - .take(dims.div_ceil(4) as usize) - .collect::>() - }); - simd::fast_scan::pack(dims.div_ceil(4), signs).collect() - }, - } -} - -pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { +pub fn preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec<[u64; 2]>) { let dis_v_2 = f32::reduce_sum_of_x2(vector); let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); let qvector_sum = if vector.len() <= 4369 { @@ -90,20 +12,19 @@ pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { (dis_v_2, b, k, qvector_sum, compress(qvector)) } -pub fn fscan_process_lowerbound_l2( - dims: u32, - lut: &(f32, f32, f32, f32, Vec), +pub fn process_lowerbound_l2( + lut: &(f32, f32, f32, f32, Vec<[u64; 2]>), (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( &[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], - &[u8], + &[[u64; 2]], ), epsilon: f32, ) -> [Distance; 32] { let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let r = simd::fast_scan::fast_scan(dims.div_ceil(4), t, s); + let r = simd::fast_scan::fast_scan(t, s); std::array::from_fn(|i| { let rough = dis_u_2[i] + dis_v_2 @@ -114,20 +35,19 @@ pub fn fscan_process_lowerbound_l2( }) } -pub fn fscan_process_lowerbound_dot( - dims: u32, - lut: &(f32, f32, f32, f32, Vec), +pub fn process_lowerbound_dot( + lut: &(f32, f32, f32, f32, Vec<[u64; 2]>), (_, factor_ppc, factor_ip, factor_err, t): ( &[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], - &[u8], + &[[u64; 2]], ), epsilon: f32, ) -> [Distance; 32] { let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let r = simd::fast_scan::fast_scan(dims.div_ceil(4), t, s); + let r = simd::fast_scan::fast_scan(t, s); std::array::from_fn(|i| { let rough = 0.5 * b * factor_ppc[i] + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; @@ -136,37 +56,41 @@ pub fn fscan_process_lowerbound_dot( }) } -fn compress(mut qvector: Vec) -> Vec { - let dims = qvector.len() as u32; - let width = dims.div_ceil(4); - qvector.resize(qvector.len().next_multiple_of(4), 0); - let mut t = vec![0u8; width as usize * 16]; - for i in 0..width as usize { +pub fn compress(mut vector: Vec) -> Vec<[u64; 2]> { + let width = vector.len().div_ceil(4); + vector.resize(width * 4, 0); + let mut result = vec![[0u64, 0u64]; width]; + for i in 0..width { unsafe { // this hint is used to skip bound checks - std::hint::assert_unchecked(4 * i + 3 < qvector.len()); - std::hint::assert_unchecked(16 * i + 15 < t.len()); + std::hint::assert_unchecked(4 * i + 3 < vector.len()); } - let t0 = qvector[4 * i + 0]; - let t1 = qvector[4 * i + 1]; - let t2 = qvector[4 * i + 2]; - let t3 = qvector[4 * i + 3]; - t[16 * i + 0b0000] = 0; - t[16 * i + 0b0001] = t0; - t[16 * i + 0b0010] = t1; - t[16 * i + 0b0011] = t1 + t0; - t[16 * i + 0b0100] = t2; - t[16 * i + 0b0101] = t2 + t0; - t[16 * i + 0b0110] = t2 + t1; - t[16 * i + 0b0111] = t2 + t1 + t0; - t[16 * i + 0b1000] = t3; - t[16 * i + 0b1001] = t3 + t0; - t[16 * i + 0b1010] = t3 + t1; - t[16 * i + 0b1011] = t3 + t1 + t0; - t[16 * i + 0b1100] = t3 + t2; - t[16 * i + 0b1101] = t3 + t2 + t0; - t[16 * i + 0b1110] = t3 + t2 + t1; - t[16 * i + 0b1111] = t3 + t2 + t1 + t0; + let t_0 = vector[4 * i + 0]; + let t_1 = vector[4 * i + 1]; + let t_2 = vector[4 * i + 2]; + let t_3 = vector[4 * i + 3]; + result[i] = [ + u64::from_le_bytes([ + 0, + t_0, + t_1, + t_1 + t_0, + t_2, + t_2 + t_0, + t_2 + t_1, + t_2 + t_1 + t_0, + ]), + u64::from_le_bytes([ + t_3, + t_3 + t_0, + t_3 + t_1, + t_3 + t_1 + t_0, + t_3 + t_2, + t_3 + t_2 + t_0, + t_3 + t_2 + t_1, + t_3 + t_2 + t_1 + t_0, + ]), + ]; } - t + result } diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index 0796b7ac..1e379a30 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -2,4 +2,96 @@ pub mod binary; pub mod block; -mod utils; + +use simd::Floating; + +#[derive(Debug, Clone)] +pub struct Code { + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub signs: Vec, +} + +pub fn pack_to_u4(signs: &[bool]) -> Vec { + fn f(x: [bool; 4]) -> u8 { + x[0] as u8 | (x[1] as u8) << 1 | (x[2] as u8) << 2 | (x[3] as u8) << 3 + } + let mut result = Vec::with_capacity(signs.len().div_ceil(4)); + for i in 0..signs.len().div_ceil(4) { + let x = std::array::from_fn(|j| signs.get(i * 4 + j).copied().unwrap_or_default()); + result.push(f(x)); + } + result +} + +pub fn pack_to_u64(signs: &[bool]) -> Vec { + fn f(x: [bool; 64]) -> u64 { + let mut result = 0_u64; + for i in 0..64 { + result |= (x[i] as u64) << i; + } + result + } + let mut result = Vec::with_capacity(signs.len().div_ceil(64)); + for i in 0..signs.len().div_ceil(64) { + let x = std::array::from_fn(|j| signs.get(i * 64 + j).copied().unwrap_or_default()); + result.push(f(x)); + } + result +} + +pub fn code(dims: u32, vector: &[f32]) -> Code { + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x_2 = f32::reduce_sum_of_x2(vector); + let dis_u = sum_of_x_2.sqrt(); + let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); + let x_x0 = dis_u / x0; + let fac_norm = (dims as f32).sqrt(); + let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); + let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); + let factor_ip = -2.0f32 / fac_norm * x_x0; + let cnt_pos = vector + .iter() + .map(|x| x.is_sign_positive() as i32) + .sum::(); + let cnt_neg = vector + .iter() + .map(|x| x.is_sign_negative() as i32) + .sum::(); + let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; + let mut signs = Vec::new(); + for i in 0..dims { + signs.push(vector[i as usize].is_sign_positive()); + } + Code { + dis_u_2: sum_of_x_2, + factor_ppc, + factor_ip, + factor_err, + signs, + } +} + +pub fn compute_lut( + vector: &[f32], +) -> ( + (f32, f32, f32, f32, Vec<[u64; 2]>), + (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), +) { + use simd::Floating; + let dis_v_2 = f32::reduce_sum_of_x2(vector); + let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); + let qvector_sum = if vector.len() <= 4369 { + simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 + } else { + simd::u8::reduce_sum_of_x(&qvector) as f32 + }; + let binary = binary::binarize(&qvector); + let block = block::compress(qvector); + ( + (dis_v_2, b, k, qvector_sum, block), + (dis_v_2, b, k, qvector_sum, binary), + ) +} diff --git a/crates/rabitq/src/utils.rs b/crates/rabitq/src/utils.rs deleted file mode 100644 index c61b87e6..00000000 --- a/crates/rabitq/src/utils.rs +++ /dev/null @@ -1,20 +0,0 @@ -#[derive(Debug, Clone)] -pub struct InfiniteByteChunks { - iter: I, -} - -impl InfiniteByteChunks { - pub fn new(iter: I) -> Self { - Self { iter } - } -} - -impl, const N: usize> Iterator for InfiniteByteChunks { - type Item = [u8; N]; - - fn next(&mut self) -> Option { - Some(std::array::from_fn::(|_| { - self.iter.next().unwrap_or(0) - })) - } -} diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 53bddb78..ce906187 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -517,7 +517,7 @@ mod reduce_sum_of_xy { } // temporarily disables this for uncertain precision - #[expect(dead_code)] + #[cfg_attr(not(test), expect(dead_code))] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "v8.3a")] @@ -894,7 +894,7 @@ mod reduce_sum_of_d2 { } // temporarily disables this for uncertain precision - #[expect(dead_code)] + #[cfg_attr(not(test), expect(dead_code))] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "v8.3a")] diff --git a/crates/simd/src/fast_scan/mod.rs b/crates/simd/src/fast_scan/mod.rs index c2579af1..84ad9901 100644 --- a/crates/simd/src/fast_scan/mod.rs +++ b/crates/simd/src/fast_scan/mod.rs @@ -1,16 +1,16 @@ /* -## codes layout for 4-bit quantizer +## code layout for 4-bit quantizer -group i = | vector i | (total bytes = width/2) +group i = | vector i | (total bytes = n/2) -byte: | 0 | 1 | 2 | ... | width/2 - 1 | -bits 0..3: | code 0 | code 2 | code 4 | ... | code width-2 | -bits 4..7: | code 1 | code 3 | code 5 | ... | code width-1 | +byte: | 0 | 1 | 2 | ... | n/2 - 1 | +bits 0..3: | code 0 | code 2 | code 4 | ... | code n-2 | +bits 4..7: | code 1 | code 3 | code 5 | ... | code n-1 | -## packed_codes layout for 4-bit quantizer +## packed_code layout for 4-bit quantizer -group i = | vector 32i | vector 32i+1 | vector 32i+2 | ... | vector 32i+31 | (total bytes = width * 16) +group i = | vector 32i | vector 32i+1 | vector 32i+2 | ... | vector 32i+31 | (total bytes = n * 16) byte | 0 | 1 | 2 | ... | 14 | 15 | bits 0..3 | code 0,vector 0 | code 0,vector 8 | code 0,vector 1 | ... | code 0,vector 14 | code 0,vector 15 | @@ -26,34 +26,105 @@ bits 4..7 | code 2,vector 16 | code 2,vector 24 | code 2,vector 17 | ... | code ... -byte | width*32-32 | width*32-31 | ... | width*32-1 | -bits 0..3 | code (width-1),vector 0 | code (width-1),vector 8 | ... | code (width-1),vector 15 | -bits 4..7 | code (width-1),vector 16 | code (width-1),vector 24 | ... | code (width-1),vector 31 | +byte | n*32-32 | n*32-31 | ... | n*32-1 | +bits 0..3 | code (n-1),vector 0 | code (n-1),vector 8 | ... | code (n-1),vector 15 | +bits 4..7 | code (n-1),vector 16 | code (n-1),vector 24 | ... | code (n-1),vector 31 | */ -pub fn pack(width: u32, r: [Vec; 32]) -> impl Iterator { - (0..width as usize).flat_map(move |i| { - [ - r[0][i] | (r[16][i] << 4), - r[8][i] | (r[24][i] << 4), - r[1][i] | (r[17][i] << 4), - r[9][i] | (r[25][i] << 4), - r[2][i] | (r[18][i] << 4), - r[10][i] | (r[26][i] << 4), - r[3][i] | (r[19][i] << 4), - r[11][i] | (r[27][i] << 4), - r[4][i] | (r[20][i] << 4), - r[12][i] | (r[28][i] << 4), - r[5][i] | (r[21][i] << 4), - r[13][i] | (r[29][i] << 4), - r[6][i] | (r[22][i] << 4), - r[14][i] | (r[30][i] << 4), - r[7][i] | (r[23][i] << 4), - r[15][i] | (r[31][i] << 4), - ] - .into_iter() - }) +pub fn pack(x: [&[u8]; 32]) -> Vec<[u64; 2]> { + let n = { + let l = x.each_ref().map(|i| i.len()); + for i in 1..32 { + assert!(l[0] == l[i]); + } + l[0] + }; + let mut result = Vec::with_capacity(n); + for i in 0..n { + result.push([ + u64::from_le_bytes([ + x[0][i] | (x[16][i] << 4), + x[8][i] | (x[24][i] << 4), + x[1][i] | (x[17][i] << 4), + x[9][i] | (x[25][i] << 4), + x[2][i] | (x[18][i] << 4), + x[10][i] | (x[26][i] << 4), + x[3][i] | (x[19][i] << 4), + x[11][i] | (x[27][i] << 4), + ]), + u64::from_le_bytes([ + x[4][i] | (x[20][i] << 4), + x[12][i] | (x[28][i] << 4), + x[5][i] | (x[21][i] << 4), + x[13][i] | (x[29][i] << 4), + x[6][i] | (x[22][i] << 4), + x[14][i] | (x[30][i] << 4), + x[7][i] | (x[23][i] << 4), + x[15][i] | (x[31][i] << 4), + ]), + ]); + } + result +} + +pub fn unpack(x: &[[u64; 2]]) -> [Vec; 32] { + let n = x.len(); + let mut result = std::array::from_fn(|_| Vec::with_capacity(n)); + for i in 0..n { + let a = x[i][0].to_le_bytes(); + let b = x[i][1].to_le_bytes(); + result[0].push(a[0] & 0xf); + result[1].push(a[2] & 0xf); + result[2].push(a[4] & 0xf); + result[3].push(a[6] & 0xf); + result[4].push(b[0] & 0xf); + result[5].push(b[2] & 0xf); + result[6].push(b[4] & 0xf); + result[7].push(b[6] & 0xf); + result[8].push(a[1] & 0xf); + result[9].push(a[3] & 0xf); + result[10].push(a[5] & 0xf); + result[11].push(a[7] & 0xf); + result[12].push(b[1] & 0xf); + result[13].push(b[3] & 0xf); + result[14].push(b[5] & 0xf); + result[15].push(b[7] & 0xf); + result[16].push(a[0] >> 4); + result[17].push(a[2] >> 4); + result[18].push(a[4] >> 4); + result[19].push(a[6] >> 4); + result[20].push(b[0] >> 4); + result[21].push(b[2] >> 4); + result[22].push(b[4] >> 4); + result[23].push(b[6] >> 4); + result[24].push(a[1] >> 4); + result[25].push(a[3] >> 4); + result[26].push(a[5] >> 4); + result[27].push(a[7] >> 4); + result[28].push(b[1] >> 4); + result[29].push(b[3] >> 4); + result[30].push(b[5] >> 4); + result[31].push(b[7] >> 4); + } + result +} + +pub fn padding_pack(x: impl IntoIterator>) -> Vec<[u64; 2]> { + let x = x.into_iter().collect::>(); + let x = x.iter().map(|x| x.as_ref()).collect::>(); + if x.is_empty() || x.len() > 32 { + panic!("too few or too many slices"); + } + let n = x[0].len(); + let t = vec![0; n]; + pack(std::array::from_fn(|i| { + if i < x.len() { x[i] } else { t.as_slice() } + })) +} + +pub fn any_pack(mut x: impl Iterator) -> [T; 32] { + std::array::from_fn(|_| x.next()).map(|x| x.unwrap_or_default()) } #[allow(clippy::module_inception)] @@ -61,10 +132,10 @@ mod fast_scan { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - fn fast_scan_v4(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + fn fast_scan_v4(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually - assert_eq!(codes.len(), width as usize * 16); - assert_eq!(lut.len(), width as usize * 16); + assert_eq!(code.len(), lut.len()); + let n = code.len(); unsafe { use std::arch::x86_64::*; @@ -99,14 +170,14 @@ mod fast_scan { let mut accu_3 = _mm512_setzero_si512(); let mut i = 0_usize; - while i + 4 <= width as usize { - let c = _mm512_loadu_si512(codes.as_ptr().add(i * 16).cast()); + while i + 4 <= n { + let c = _mm512_loadu_si512(code.as_ptr().add(i).cast()); let mask = _mm512_set1_epi8(0xf); let clo = _mm512_and_si512(c, mask); let chi = _mm512_and_si512(_mm512_srli_epi16(c, 4), mask); - let lut = _mm512_loadu_si512(lut.as_ptr().add(i * 16).cast()); + let lut = _mm512_loadu_si512(lut.as_ptr().add(i).cast()); let res_lo = _mm512_shuffle_epi8(lut, clo); accu_0 = _mm512_add_epi16(accu_0, res_lo); accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); @@ -116,14 +187,14 @@ mod fast_scan { i += 4; } - if i + 2 <= width as usize { - let c = _mm256_loadu_si256(codes.as_ptr().add(i * 16).cast()); + if i + 2 <= n { + let c = _mm256_loadu_si256(code.as_ptr().add(i).cast()); let mask = _mm256_set1_epi8(0xf); let clo = _mm256_and_si256(c, mask); let chi = _mm256_and_si256(_mm256_srli_epi16(c, 4), mask); - let lut = _mm256_loadu_si256(lut.as_ptr().add(i * 16).cast()); + let lut = _mm256_loadu_si256(lut.as_ptr().add(i).cast()); let res_lo = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, clo)); accu_0 = _mm512_add_epi16(accu_0, res_lo); accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); @@ -133,14 +204,14 @@ mod fast_scan { i += 2; } - if i < width as usize { - let c = _mm_loadu_si128(codes.as_ptr().add(i * 16).cast()); + if i < n { + let c = _mm_loadu_si128(code.as_ptr().add(i).cast()); let mask = _mm_set1_epi8(0xf); let clo = _mm_and_si128(c, mask); let chi = _mm_and_si128(_mm_srli_epi16(c, 4), mask); - let lut = _mm_loadu_si128(lut.as_ptr().add(i * 16).cast()); + let lut = _mm_loadu_si128(lut.as_ptr().add(i).cast()); let res_lo = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, clo)); accu_0 = _mm512_add_epi16(accu_0, res_lo); accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); @@ -150,7 +221,7 @@ mod fast_scan { i += 1; } - debug_assert_eq!(i, width as usize); + debug_assert_eq!(i, n); let mut result = [0_u16; 32]; @@ -178,14 +249,15 @@ mod fast_scan { return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - for width in 90..110 { - let codes = (0..16 * width).map(|_| rand::random()).collect::>(); - let lut = (0..16 * width).map(|_| rand::random()).collect::>(); + for n in 90..110 { + let code = (0..n) + .map(|_| [rand::random(), rand::random()]) + .collect::>(); + let lut = (0..n) + .map(|_| [rand::random(), rand::random()]) + .collect::>(); unsafe { - assert_eq!( - fast_scan_v4(width, &codes, &lut), - fast_scan_fallback(width, &codes, &lut) - ); + assert_eq!(fast_scan_v4(&code, &lut), fast_scan_fallback(&code, &lut)); } } } @@ -194,10 +266,10 @@ mod fast_scan { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - fn fast_scan_v3(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + fn fast_scan_v3(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually - assert_eq!(codes.len(), width as usize * 16); - assert_eq!(lut.len(), width as usize * 16); + assert_eq!(code.len(), lut.len()); + let n = code.len(); unsafe { use std::arch::x86_64::*; @@ -218,14 +290,14 @@ mod fast_scan { let mut accu_3 = _mm256_setzero_si256(); let mut i = 0_usize; - while i + 2 <= width as usize { - let c = _mm256_loadu_si256(codes.as_ptr().add(i * 16).cast()); + while i + 2 <= n { + let c = _mm256_loadu_si256(code.as_ptr().add(i).cast()); let mask = _mm256_set1_epi8(0xf); let clo = _mm256_and_si256(c, mask); let chi = _mm256_and_si256(_mm256_srli_epi16(c, 4), mask); - let lut = _mm256_loadu_si256(lut.as_ptr().add(i * 16).cast()); + let lut = _mm256_loadu_si256(lut.as_ptr().add(i).cast()); let res_lo = _mm256_shuffle_epi8(lut, clo); accu_0 = _mm256_add_epi16(accu_0, res_lo); accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); @@ -235,14 +307,14 @@ mod fast_scan { i += 2; } - if i < width as usize { - let c = _mm_loadu_si128(codes.as_ptr().add(i * 16).cast()); + if i < n { + let c = _mm_loadu_si128(code.as_ptr().add(i).cast()); let mask = _mm_set1_epi8(0xf); let clo = _mm_and_si128(c, mask); let chi = _mm_and_si128(_mm_srli_epi16(c, 4), mask); - let lut = _mm_loadu_si128(lut.as_ptr().add(i * 16).cast()); + let lut = _mm_loadu_si128(lut.as_ptr().add(i).cast()); let res_lo = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, clo)); accu_0 = _mm256_add_epi16(accu_0, res_lo); accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); @@ -252,7 +324,7 @@ mod fast_scan { i += 1; } - debug_assert_eq!(i, width as usize); + debug_assert_eq!(i, n); let mut result = [0_u16; 32]; @@ -280,14 +352,15 @@ mod fast_scan { return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - for width in 90..110 { - let codes = (0..16 * width).map(|_| rand::random()).collect::>(); - let lut = (0..16 * width).map(|_| rand::random()).collect::>(); + for n in 90..110 { + let code = (0..n) + .map(|_| [rand::random(), rand::random()]) + .collect::>(); + let lut = (0..n) + .map(|_| [rand::random(), rand::random()]) + .collect::>(); unsafe { - assert_eq!( - fast_scan_v3(width, &codes, &lut), - fast_scan_fallback(width, &codes, &lut) - ); + assert_eq!(fast_scan_v3(&code, &lut), fast_scan_fallback(&code, &lut)); } } } @@ -295,10 +368,10 @@ mod fast_scan { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - fn fast_scan_v2(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + fn fast_scan_v2(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually - assert_eq!(codes.len(), width as usize * 16); - assert_eq!(lut.len(), width as usize * 16); + assert_eq!(code.len(), lut.len()); + let n = code.len(); unsafe { use std::arch::x86_64::*; @@ -309,14 +382,14 @@ mod fast_scan { let mut accu_3 = _mm_setzero_si128(); let mut i = 0_usize; - while i < width as usize { - let c = _mm_loadu_si128(codes.as_ptr().add(i * 16).cast()); + while i < n { + let c = _mm_loadu_si128(code.as_ptr().add(i).cast()); let mask = _mm_set1_epi8(0xf); let clo = _mm_and_si128(c, mask); let chi = _mm_and_si128(_mm_srli_epi16(c, 4), mask); - let lut = _mm_loadu_si128(lut.as_ptr().add(i * 16).cast()); + let lut = _mm_loadu_si128(lut.as_ptr().add(i).cast()); let res_lo = _mm_shuffle_epi8(lut, clo); accu_0 = _mm_add_epi16(accu_0, res_lo); accu_1 = _mm_add_epi16(accu_1, _mm_srli_epi16(res_lo, 8)); @@ -326,7 +399,7 @@ mod fast_scan { i += 1; } - debug_assert_eq!(i, width as usize); + debug_assert_eq!(i, n); let mut result = [0_u16; 32]; @@ -350,14 +423,15 @@ mod fast_scan { return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - for width in 90..110 { - let codes = (0..16 * width).map(|_| rand::random()).collect::>(); - let lut = (0..16 * width).map(|_| rand::random()).collect::>(); + for n in 90..110 { + let code = (0..n) + .map(|_| [rand::random(), rand::random()]) + .collect::>(); + let lut = (0..n) + .map(|_| [rand::random(), rand::random()]) + .collect::>(); unsafe { - assert_eq!( - fast_scan_v2(width, &codes, &lut), - fast_scan_fallback(width, &codes, &lut) - ); + assert_eq!(fast_scan_v2(&code, &lut), fast_scan_fallback(&code, &lut)); } } } @@ -365,10 +439,10 @@ mod fast_scan { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "v8.3a")] - fn fast_scan_v8_3a(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { + fn fast_scan_v8_3a(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually - assert_eq!(codes.len(), width as usize * 16); - assert_eq!(lut.len(), width as usize * 16); + assert_eq!(code.len(), lut.len()); + let n = code.len(); unsafe { use std::arch::aarch64::*; @@ -379,14 +453,14 @@ mod fast_scan { let mut accu_3 = vdupq_n_u16(0); let mut i = 0_usize; - while i < width as usize { - let c = vld1q_u8(codes.as_ptr().add(i * 16).cast()); + while i < n { + let c = vld1q_u8(code.as_ptr().add(i).cast()); let mask = vdupq_n_u8(0xf); let clo = vandq_u8(c, mask); let chi = vandq_u8(vshrq_n_u8(c, 4), mask); - let lut = vld1q_u8(lut.as_ptr().add(i * 16).cast()); + let lut = vld1q_u8(lut.as_ptr().add(i).cast()); let res_lo = vreinterpretq_u16_u8(vqtbl1q_u8(lut, clo)); accu_0 = vaddq_u16(accu_0, res_lo); accu_1 = vaddq_u16(accu_1, vshrq_n_u16(res_lo, 8)); @@ -396,7 +470,7 @@ mod fast_scan { i += 1; } - debug_assert_eq!(i, width as usize); + debug_assert_eq!(i, n); let mut result = [0_u16; 32]; @@ -420,13 +494,17 @@ mod fast_scan { return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - for width in 90..110 { - let codes = (0..16 * width).map(|_| rand::random()).collect::>(); - let lut = (0..16 * width).map(|_| rand::random()).collect::>(); + for n in 90..110 { + let code = (0..n) + .map(|_| [rand::random(), rand::random()]) + .collect::>(); + let lut = (0..n) + .map(|_| [rand::random(), rand::random()]) + .collect::>(); unsafe { assert_eq!( - fast_scan_v8_3a(width, &codes, &lut), - fast_scan_fallback(width, &codes, &lut) + fast_scan_v8_3a(&code, &lut), + fast_scan_fallback(&code, &lut) ); } } @@ -434,32 +512,24 @@ mod fast_scan { } #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a")] - pub fn fast_scan(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { - let width = width as usize; - - assert_eq!(codes.len(), width * 16); - assert_eq!(lut.len(), width * 16); + pub fn fast_scan(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { + assert_eq!(code.len(), lut.len()); + let n = code.len(); - use std::array::from_fn; - use std::ops::BitAnd; - - fn load(slice: &[T]) -> [T; N] { - from_fn(|i| slice[i]) - } fn unary(op: impl Fn(T) -> U, a: [T; N]) -> [U; N] { - from_fn(|i| op(a[i])) + std::array::from_fn(|i| op(a[i])) } fn binary(op: impl Fn(T, T) -> T, a: [T; N], b: [T; N]) -> [T; N] { - from_fn(|i| op(a[i], b[i])) + std::array::from_fn(|i| op(a[i], b[i])) } fn shuffle(a: [T; N], b: [u8; N]) -> [T; N] { - from_fn(|i| a[b[i] as usize]) + std::array::from_fn(|i| a[b[i] as usize]) } fn cast(x: [u8; 16]) -> [u16; 8] { - from_fn(|i| u16::from_le_bytes([x[i << 1 | 0], x[i << 1 | 1]])) + std::array::from_fn(|i| u16::from_le_bytes([x[i << 1 | 0], x[i << 1 | 1]])) } fn setr(x: [[T; 8]; 4]) -> [T; 32] { - from_fn(|i| x[i >> 3][i & 7]) + std::array::from_fn(|i| x[i >> 3][i & 7]) } let mut a_0 = [0u16; 8]; @@ -467,14 +537,14 @@ mod fast_scan { let mut a_2 = [0u16; 8]; let mut a_3 = [0u16; 8]; - for i in 0..width { - let c = load(&codes[16 * i..]); + for i in 0..n { + let c = unsafe { std::mem::transmute::<[u64; 2], [u8; 16]>(code[i]) }; let mask = [0xfu8; 16]; - let clo = binary(u8::bitand, c, mask); - let chi = binary(u8::bitand, unary(|x| x >> 4, c), mask); + let clo = binary(std::ops::BitAnd::bitand, c, mask); + let chi = binary(std::ops::BitAnd::bitand, unary(|x| x >> 4, c), mask); - let lut = load(&lut[16 * i..]); + let lut = unsafe { std::mem::transmute::<[u64; 2], [u8; 16]>(lut[i]) }; let res_lo = cast(shuffle(lut, clo)); a_0 = binary(u16::wrapping_add, a_0, res_lo); a_1 = binary(u16::wrapping_add, a_1, unary(|x| x >> 8, res_lo)); @@ -491,6 +561,6 @@ mod fast_scan { } #[inline(always)] -pub fn fast_scan(width: u32, codes: &[u8], lut: &[u8]) -> [u16; 32] { - fast_scan::fast_scan(width, codes, lut) +pub fn fast_scan(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { + fast_scan::fast_scan(code, lut) } diff --git a/src/vchordrq/algorithm/build.rs b/src/algorithm/build.rs similarity index 78% rename from src/vchordrq/algorithm/build.rs rename to src/algorithm/build.rs index 4213e59d..a0810a71 100644 --- a/src/vchordrq/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -1,25 +1,24 @@ -use crate::vchordrq::algorithm::rabitq; -use crate::vchordrq::algorithm::tuples::*; -use crate::vchordrq::index::am_options::Opfamily; -use crate::vchordrq::types::DistanceKind; -use crate::vchordrq::types::VchordrqBuildOptions; -use crate::vchordrq::types::VchordrqExternalBuildOptions; -use crate::vchordrq::types::VchordrqIndexingOptions; -use crate::vchordrq::types::VchordrqInternalBuildOptions; -use crate::vchordrq::types::VectorOptions; -use algorithm::{Page, PageGuard, RelationWrite}; +use crate::algorithm::RelationWrite; +use crate::algorithm::operator::{Operator, Vector}; +use crate::algorithm::tape::*; +use crate::algorithm::tuples::*; +use crate::index::am_options::Opfamily; +use crate::types::VchordrqBuildOptions; +use crate::types::VchordrqExternalBuildOptions; +use crate::types::VchordrqIndexingOptions; +use crate::types::VchordrqInternalBuildOptions; +use crate::types::VectorOptions; use rand::Rng; -use rkyv::ser::serializers::AllocSerializer; use simd::Floating; -use std::marker::PhantomData; use std::num::NonZeroU64; use std::sync::Arc; use vector::VectorBorrowed; +use vector::VectorOwned; -pub trait HeapRelation { +pub trait HeapRelation { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((NonZeroU64, V)); + F: FnMut((NonZeroU64, O::Vector)); fn opfamily(&self) -> Opfamily; } @@ -27,7 +26,7 @@ pub trait Reporter { fn tuples_total(&mut self, tuples_total: u64); } -pub fn build, R: Reporter>( +pub fn build, R: Reporter>( vector_options: VectorOptions, vchordrq_options: VchordrqIndexingOptions, heap_relation: T, @@ -35,8 +34,7 @@ pub fn build, R: Reporter>( mut reporter: R, ) { let dims = vector_options.dims; - let is_residual = - vchordrq_options.residual_quantization && vector_options.d == DistanceKind::L2; + let is_residual = vchordrq_options.residual_quantization && O::SUPPORTS_RESIDUAL; let structures = match vchordrq_options.build { VchordrqBuildOptions::External(external_build) => Structure::extern_build( vector_options.clone(), @@ -58,11 +56,11 @@ pub fn build, R: Reporter>( let vector = vector.as_borrowed(); assert_eq!(dims, vector.dims(), "invalid vector dimensions"); if number_of_samples < max_number_of_samples { - samples.push(V::build_to_vecf32(vector)); + samples.push(O::Vector::build_to_vecf32(vector)); number_of_samples += 1; } else { let index = rand.gen_range(0..max_number_of_samples) as usize; - samples[index] = V::build_to_vecf32(vector); + samples[index] = O::Vector::build_to_vecf32(vector); } tuples_total += 1; }); @@ -72,24 +70,32 @@ pub fn build, R: Reporter>( Structure::internal_build(vector_options.clone(), internal_build.clone(), samples) } }; - let mut meta = Tape::create(&relation, false); + let mut meta = TapeWriter::<_, _, MetaTuple>::create(|| relation.extend(false)); assert_eq!(meta.first(), 0); - let mut vectors = Tape::, _>::create(&relation, true); - let mut pointer_of_means = Vec::>::new(); + let freepage = TapeWriter::<_, _, FreepageTuple>::create(|| relation.extend(false)); + let mut vectors = TapeWriter::<_, _, VectorTuple>::create(|| relation.extend(true)); + let mut pointer_of_means = Vec::>::new(); for i in 0..structures.len() { let mut level = Vec::new(); for j in 0..structures[i].len() { - let vector = V::build_from_vecf32(&structures[i].means[j]); - let (metadata, slices) = V::vector_split(vector.as_borrowed()); - let mut chain = Err(metadata); + let vector = O::Vector::build_from_vecf32(&structures[i].means[j]); + let (metadata, slices) = O::Vector::vector_split(vector.as_borrowed()); + let mut chain = Ok(metadata); for i in (0..slices.len()).rev() { - chain = Ok(vectors.push(&VectorTuple { - payload: None, - slice: slices[i].to_vec(), - chain, + chain = Err(vectors.push(match chain { + Ok(metadata) => VectorTuple::_0 { + payload: None, + elements: slices[i].to_vec(), + metadata, + }, + Err(pointer) => VectorTuple::_1 { + payload: None, + elements: slices[i].to_vec(), + pointer, + }, })); } - level.push(chain.ok().unwrap()); + level.push(chain.err().unwrap()); } pointer_of_means.push(level); } @@ -98,10 +104,14 @@ pub fn build, R: Reporter>( let mut level = Vec::new(); for j in 0..structures[i].len() { if i == 0 { - let tape = Tape::::create(&relation, false); - level.push(tape.first()); + let tape = TapeWriter::<_, _, H0Tuple>::create(|| relation.extend(false)); + let mut jump = TapeWriter::<_, _, JumpTuple>::create(|| relation.extend(false)); + jump.push(JumpTuple { + first: tape.first(), + }); + level.push(jump.first()); } else { - let mut tape = Tape::::create(&relation, false); + let mut tape = H1TapeWriter::<_, _>::create(|| relation.extend(false)); let h2_mean = &structures[i].means[j]; let h2_children = &structures[i].children[j]; for child in h2_children.iter().copied() { @@ -111,28 +121,30 @@ pub fn build, R: Reporter>( } else { rabitq::code(dims, h1_mean) }; - tape.push(&Height1Tuple { + tape.push(H1Branch { mean: pointer_of_means[i - 1][child as usize], - first: pointer_of_firsts[i - 1][child as usize], dis_u_2: code.dis_u_2, factor_ppc: code.factor_ppc, factor_ip: code.factor_ip, factor_err: code.factor_err, - t: code.t(), + signs: code.signs, + first: pointer_of_firsts[i - 1][child as usize], }); } + let tape = tape.into_inner(); level.push(tape.first()); } } pointer_of_firsts.push(level); } - meta.push(&MetaTuple { + meta.push(MetaTuple { dims, height_of_root: structures.len() as u32, is_residual, vectors_first: vectors.first(), - mean: pointer_of_means.last().unwrap()[0], - first: pointer_of_firsts.last().unwrap()[0], + root_mean: pointer_of_means.last().unwrap()[0], + root_first: pointer_of_firsts.last().unwrap()[0], + freepage_first: freepage.first(), }); } @@ -342,7 +354,7 @@ impl Structure { ) -> (Vec>, Vec>) { labels .iter() - .filter(|(_, &(h, _))| h == height) + .filter(|(_, (h, _))| *h == height) .map(|(id, _)| { ( vectors[id].clone(), @@ -359,50 +371,3 @@ impl Structure { result } } - -struct Tape<'a: 'b, 'b, T, R: 'b + RelationWrite> { - relation: &'a R, - head: R::WriteGuard<'b>, - first: u32, - tracking_freespace: bool, - _phantom: PhantomData T>, -} - -impl<'a: 'b, 'b, T, R: 'b + RelationWrite> Tape<'a, 'b, T, R> { - fn create(relation: &'a R, tracking_freespace: bool) -> Self { - let mut head = relation.extend(tracking_freespace); - head.get_opaque_mut().skip = head.id(); - let first = head.id(); - Self { - relation, - head, - first, - tracking_freespace, - _phantom: PhantomData, - } - } - fn first(&self) -> u32 { - self.first - } -} - -impl<'a: 'b, 'b, T, R: 'b + RelationWrite> Tape<'a, 'b, T, R> -where - T: rkyv::Serialize>, -{ - fn push(&mut self, x: &T) -> (u32, u16) { - let bytes = rkyv::to_bytes(x).expect("failed to serialize"); - if let Some(i) = self.head.alloc(&bytes) { - (self.head.id(), i) - } else { - let next = self.relation.extend(self.tracking_freespace); - self.head.get_opaque_mut().next = next.id(); - self.head = next; - if let Some(i) = self.head.alloc(&bytes) { - (self.head.id(), i) - } else { - panic!("tuple is too large to fit in a fresh page") - } - } - } -} diff --git a/src/algorithm/freepages.rs b/src/algorithm/freepages.rs new file mode 100644 index 00000000..8984d3ce --- /dev/null +++ b/src/algorithm/freepages.rs @@ -0,0 +1,61 @@ +use crate::algorithm::tuples::*; +use crate::algorithm::*; +use crate::utils::pipe::Pipe; +use std::cmp::Reverse; + +pub fn mark(relation: impl RelationWrite, freepage_first: u32, pages: &[u32]) { + let mut pages = pages.to_vec(); + pages.sort_by_key(|x| Reverse(*x)); + pages.dedup(); + let first = freepage_first; + assert!(first != u32::MAX); + let (mut current, mut offset) = (first, 0_u32); + while pages.is_empty() { + let locals = { + let mut local = Vec::new(); + while let Some(target) = pages.pop_if(|x| (offset..offset + 32768).contains(x)) { + local.push(target - offset); + } + local + }; + let mut freespace_guard = relation.write(current, false); + if freespace_guard.len() == 0 { + freespace_guard.alloc(&serialize(&FreepageTuple {})); + } + let mut freespace_tuple = freespace_guard + .get_mut(1) + .expect("data corruption") + .pipe(write_tuple::); + for local in locals { + freespace_tuple.mark(local as _); + } + if freespace_guard.get_opaque().next == u32::MAX { + let extend = relation.extend(false); + freespace_guard.get_opaque_mut().next = extend.id(); + } + (current, offset) = (freespace_guard.get_opaque().next, offset + 32768); + } +} + +pub fn fetch(relation: impl RelationWrite, freepage_first: u32) -> Option { + let first = freepage_first; + assert!(first != u32::MAX); + let (mut current, mut offset) = (first, 0_u32); + loop { + let mut freespace_guard = relation.write(current, false); + if freespace_guard.len() == 0 { + return None; + } + let mut freespace_tuple = freespace_guard + .get_mut(1) + .expect("data corruption") + .pipe(write_tuple::); + if let Some(local) = freespace_tuple.fetch() { + return Some(local as u32 + offset); + } + if freespace_guard.get_opaque().next == u32::MAX { + return None; + } + (current, offset) = (freespace_guard.get_opaque().next, offset + 32768); + } +} diff --git a/src/algorithm/insert.rs b/src/algorithm/insert.rs new file mode 100644 index 00000000..f2b8cfbc --- /dev/null +++ b/src/algorithm/insert.rs @@ -0,0 +1,190 @@ +use crate::algorithm::operator::*; +use crate::algorithm::tape::read_h1_tape; +use crate::algorithm::tuples::*; +use crate::algorithm::vectors::{self}; +use crate::algorithm::{Page, PageGuard, RelationWrite}; +use crate::utils::pipe::Pipe; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::num::NonZeroU64; +use vector::VectorBorrowed; +use vector::VectorOwned; + +pub fn insert( + relation: impl RelationWrite + Clone, + payload: NonZeroU64, + vector: O::Vector, +) { + let vector = O::Vector::random_projection(vector.as_borrowed()); + let meta_guard = relation.read(0); + let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let dims = meta_tuple.dims(); + let is_residual = meta_tuple.is_residual(); + let height_of_root = meta_tuple.height_of_root(); + assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); + let root_mean = meta_tuple.root_mean(); + let root_first = meta_tuple.root_first(); + let vectors_first = meta_tuple.vectors_first(); + drop(meta_guard); + + let default_lut_block = if !is_residual { + Some(O::Vector::compute_lut_block(vector.as_borrowed())) + } else { + None + }; + + let mean = vectors::vector_append::( + relation.clone(), + vectors_first, + vector.as_borrowed(), + payload, + ); + + type State = (u32, Option<::Vector>); + let mut state: State = { + let mean = root_mean; + if is_residual { + let residual_u = vectors::vector_access_1::( + relation.clone(), + mean, + LAccess::new( + O::Vector::elements_and_metadata(vector.as_borrowed()), + O::ResidualAccessor::default(), + ), + ); + (root_first, Some(residual_u)) + } else { + (root_first, None) + } + }; + let step = |state: State| { + let mut results = Vec::new(); + { + let (first, residual) = state; + let lut = if let Some(residual) = residual { + &O::Vector::compute_lut_block(residual.as_borrowed()) + } else { + default_lut_block.as_ref().unwrap() + }; + read_h1_tape( + relation.clone(), + first, + || { + RAccess::new( + (&lut.4, (lut.0, lut.1, lut.2, lut.3, 1.9f32)), + O::Distance::block_accessor(), + ) + }, + |lowerbound, mean, first| { + results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(first))); + }, + ); + } + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); + { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + if is_residual { + let (dis_u, residual_u) = vectors::vector_access_1::( + relation.clone(), + mean, + LAccess::new( + O::Vector::elements_and_metadata(vector.as_borrowed()), + ( + O::DistanceAccessor::default(), + O::ResidualAccessor::default(), + ), + ), + ); + cache.push(( + Reverse(dis_u), + AlwaysEqual(first), + AlwaysEqual(Some(residual_u)), + )); + } else { + let dis_u = vectors::vector_access_1::( + relation.clone(), + mean, + LAccess::new( + O::Vector::elements_and_metadata(vector.as_borrowed()), + O::DistanceAccessor::default(), + ), + ); + cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); + } + } + let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop().unwrap(); + (first, mean) + } + }; + for _ in (1..height_of_root).rev() { + state = step(state); + } + + let (first, residual) = state; + let code = if let Some(residual) = residual { + O::Vector::code(residual.as_borrowed()) + } else { + O::Vector::code(vector.as_borrowed()) + }; + let bytes = serialize(&H0Tuple::_0 { + mean, + dis_u_2: code.dis_u_2, + factor_ppc: code.factor_ppc, + factor_ip: code.factor_ip, + factor_err: code.factor_err, + payload: Some(payload), + elements: rabitq::pack_to_u64(&code.signs), + }); + + let jump_guard = relation.read(first); + let jump_tuple = jump_guard + .get(1) + .expect("data corruption") + .pipe(read_tuple::); + + let first = jump_tuple.first(); + + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = relation.read(current); + if read.get_opaque().next == u32::MAX { + drop(read); + let mut write = relation.write(current, false); + if write.get_opaque().next == u32::MAX { + if write.alloc(&bytes).is_some() { + return; + } + let mut extend = relation.extend(false); + write.get_opaque_mut().next = extend.id(); + drop(write); + let fresh = extend.id(); + if extend.alloc(&bytes).is_some() { + drop(extend); + let mut past = relation.write(first, false); + past.get_opaque_mut().skip = std::cmp::max(past.get_opaque_mut().skip, fresh); + drop(past); + return; + } else { + panic!("a tuple cannot even be fit in a fresh page"); + } + } else { + if current == first && write.get_opaque().skip != first { + current = write.get_opaque().skip; + } else { + current = write.get_opaque().next; + } + } + } else { + if current == first && read.get_opaque().skip != first { + current = read.get_opaque().skip; + } else { + current = read.get_opaque().next; + } + } + } +} diff --git a/crates/algorithm/src/lib.rs b/src/algorithm/mod.rs similarity index 85% rename from crates/algorithm/src/lib.rs rename to src/algorithm/mod.rs index 8bee0167..7a9c3ff2 100644 --- a/crates/algorithm/src/lib.rs +++ b/src/algorithm/mod.rs @@ -1,3 +1,14 @@ +pub mod build; +pub mod freepages; +pub mod insert; +pub mod operator; +pub mod prewarm; +pub mod scan; +pub mod tape; +pub mod tuples; +pub mod vacuum; +pub mod vectors; + use std::ops::{Deref, DerefMut}; #[repr(C, align(8))] @@ -15,8 +26,8 @@ pub trait Page: Sized { fn get_mut(&mut self, i: u16) -> Option<&mut [u8]>; fn alloc(&mut self, data: &[u8]) -> Option; fn free(&mut self, i: u16); - fn reconstruct(&mut self, removes: &[u16]); fn freespace(&self) -> u16; + fn clear(&mut self); } pub trait PageGuard { diff --git a/src/algorithm/operator.rs b/src/algorithm/operator.rs new file mode 100644 index 00000000..9506d7a2 --- /dev/null +++ b/src/algorithm/operator.rs @@ -0,0 +1,628 @@ +use crate::types::{DistanceKind, OwnedVector}; +use distance::Distance; +use half::f16; +use simd::Floating; +use std::fmt::Debug; +use std::marker::PhantomData; +use vector::vect::VectOwned; +use vector::{VectorBorrowed, VectorOwned}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +pub trait Accessor2 { + type Output; + fn push(&mut self, input: &[E0], target: &[E1]); + fn finish(self, input: M0, target: M1) -> Self::Output; +} + +impl Accessor2 for () { + type Output = (); + + fn push(&mut self, _: &[E0], _: &[E1]) {} + + fn finish(self, _: M0, _: M1) -> Self::Output {} +} + +impl> Accessor2 for (A,) { + type Output = (A::Output,); + + fn push(&mut self, input: &[E0], target: &[E1]) { + self.0.push(input, target); + } + + fn finish(self, input: M0, target: M1) -> Self::Output { + (self.0.finish(input, target),) + } +} + +impl, B: Accessor2> + Accessor2 for (A, B) +{ + type Output = (A::Output, B::Output); + + fn push(&mut self, input: &[E0], target: &[E1]) { + self.0.push(input, target); + self.1.push(input, target); + } + + fn finish(self, input: M0, target: M1) -> Self::Output { + (self.0.finish(input, target), self.1.finish(input, target)) + } +} + +#[derive(Debug)] +pub struct Sum(f32, PhantomData O>); + +impl Default for Sum { + fn default() -> Self { + Self(0.0, PhantomData) + } +} + +impl Accessor2 for Sum, L2>> { + type Output = Distance; + + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0 += f32::reduce_sum_of_d2(target, input) + } + + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(self.0) + } +} + +impl Accessor2 for Sum, Dot>> { + type Output = Distance; + + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0 += f32::reduce_sum_of_xy(target, input) + } + + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(-self.0) + } +} + +impl Accessor2 for Sum, L2>> { + type Output = Distance; + + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0 += f16::reduce_sum_of_d2(target, input) + } + + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(self.0) + } +} + +impl Accessor2 for Sum, Dot>> { + type Output = Distance; + + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0 += f16::reduce_sum_of_xy(target, input) + } + + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(-self.0) + } +} + +#[derive(Debug, Clone)] +pub struct Diff(Vec<::Element>); + +impl Default for Diff { + fn default() -> Self { + Self(Vec::new()) + } +} + +impl Accessor2 for Diff, L2>> { + type Output = VectOwned; + + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0.extend(f32::vector_sub(target, input)); + } + + fn finish(self, (): (), (): ()) -> Self::Output { + VectOwned::new(self.0) + } +} + +impl Accessor2 for Diff, Dot>> { + type Output = VectOwned; + + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0.extend(f32::vector_sub(target, input)); + } + + fn finish(self, (): (), (): ()) -> Self::Output { + VectOwned::new(self.0) + } +} + +impl Accessor2 for Diff, L2>> { + type Output = VectOwned; + + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0.extend(f16::vector_sub(target, input)); + } + + fn finish(self, (): (), (): ()) -> Self::Output { + VectOwned::new(self.0) + } +} + +impl Accessor2 for Diff, Dot>> { + type Output = VectOwned; + + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0.extend(f16::vector_sub(target, input)); + } + + fn finish(self, (): (), (): ()) -> Self::Output { + VectOwned::new(self.0) + } +} + +#[derive(Debug)] +pub struct Block([u16; 32], PhantomData D>); + +impl Default for Block { + fn default() -> Self { + Self([0u16; 32], PhantomData) + } +} + +impl + Accessor2< + [u64; 2], + [u64; 2], + (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), + (f32, f32, f32, f32, f32), + > for Block +{ + type Output = [Distance; 32]; + + fn push(&mut self, input: &[[u64; 2]], target: &[[u64; 2]]) { + let t = simd::fast_scan::fast_scan(input, target); + for i in 0..32 { + self.0[i] += t[i]; + } + } + + fn finish( + self, + (dis_u_2, factor_ppc, factor_ip, factor_err): ( + &[f32; 32], + &[f32; 32], + &[f32; 32], + &[f32; 32], + ), + (dis_v_2, b, k, qvector_sum, epsilon): (f32, f32, f32, f32, f32), + ) -> Self::Output { + std::array::from_fn(|i| { + let rough = dis_u_2[i] + + dis_v_2 + + b * factor_ppc[i] + + ((2.0 * self.0[i] as f32) - qvector_sum) * factor_ip[i] * k; + let err = factor_err[i] * dis_v_2.sqrt(); + Distance::from_f32(rough - epsilon * err) + }) + } +} + +impl + Accessor2< + [u64; 2], + [u64; 2], + (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), + (f32, f32, f32, f32, f32), + > for Block +{ + type Output = [Distance; 32]; + + fn push(&mut self, input: &[[u64; 2]], target: &[[u64; 2]]) { + let t = simd::fast_scan::fast_scan(input, target); + for i in 0..32 { + self.0[i] += t[i]; + } + } + + fn finish( + self, + (_, factor_ppc, factor_ip, factor_err): (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), + (dis_v_2, b, k, qvector_sum, epsilon): (f32, f32, f32, f32, f32), + ) -> Self::Output { + std::array::from_fn(|i| { + let rough = 0.5 * b * factor_ppc[i] + + 0.5 * ((2.0 * self.0[i] as f32) - qvector_sum) * factor_ip[i] * k; + let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); + Distance::from_f32(rough - epsilon * err) + }) + } +} + +pub trait Accessor1 { + type Output; + fn push(&mut self, input: &[E]); + fn finish(self, input: M) -> Self::Output; +} + +impl Accessor1 for () { + type Output = (); + + fn push(&mut self, _: &[E]) {} + + fn finish(self, _: M) -> Self::Output {} +} + +impl Accessor1 for (A,) +where + A: Accessor1, +{ + type Output = (A::Output,); + + fn push(&mut self, input: &[E]) { + self.0.push(input); + } + + fn finish(self, input: M) -> Self::Output { + (self.0.finish(input),) + } +} + +impl Accessor1 for (A, B) +where + A: Accessor1, + B: Accessor1, +{ + type Output = (A::Output, B::Output); + + fn push(&mut self, input: &[E]) { + self.0.push(input); + self.1.push(input); + } + + fn finish(self, input: M) -> Self::Output { + (self.0.finish(input), self.1.finish(input)) + } +} + +pub struct LAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> LAccess<'a, E, M, A> { + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } + } +} + +impl> Accessor1 for LAccess<'_, E0, M0, A> { + type Output = A::Output; + + fn push(&mut self, rhs: &[E1]) { + let (lhs, elements) = self.elements.split_at(rhs.len()); + self.accessor.push(lhs, rhs); + self.elements = elements; + } + + fn finish(self, rhs: M1) -> Self::Output { + self.accessor.finish(self.metadata, rhs) + } +} + +pub struct RAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> RAccess<'a, E, M, A> { + #[allow(dead_code)] + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } + } +} + +impl> Accessor1 for RAccess<'_, E1, M1, A> { + type Output = A::Output; + + fn push(&mut self, lhs: &[E0]) { + let (rhs, elements) = self.elements.split_at(lhs.len()); + self.accessor.push(lhs, rhs); + self.elements = elements; + } + + fn finish(self, lhs: M0) -> Self::Output { + self.accessor.finish(lhs, self.metadata) + } +} + +pub trait Vector: VectorOwned { + type Element: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; + type Metadata: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; + + fn vector_split(vector: Self::Borrowed<'_>) -> (Self::Metadata, Vec<&[Self::Element]>); + fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); + fn from_owned(vector: OwnedVector) -> Self; + + fn random_projection(vector: Self::Borrowed<'_>) -> Self; + + fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u64; 2]>); + + fn compute_lut( + vector: Self::Borrowed<'_>, + ) -> ( + (f32, f32, f32, f32, Vec<[u64; 2]>), + (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), + ); + + fn code(vector: Self::Borrowed<'_>) -> rabitq::Code; + + fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec; + + fn build_from_vecf32(x: &[f32]) -> Self; +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f32; + + fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f32]>) { + let vector = vector.slice(); + ((), match vector.len() { + 0..=960 => vec![vector], + 961..=1280 => vec![&vector[..640], &vector[640..]], + 1281.. => vector.chunks(1920).collect(), + }) + } + + fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + (vector.slice(), ()) + } + + fn from_owned(vector: OwnedVector) -> Self { + match vector { + OwnedVector::Vecf32(x) => x, + _ => unreachable!(), + } + } + + fn random_projection(vector: Self::Borrowed<'_>) -> Self { + Self::new(crate::projection::project(vector.slice())) + } + + fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u64; 2]>) { + rabitq::block::preprocess(vector.slice()) + } + + fn compute_lut( + vector: Self::Borrowed<'_>, + ) -> ( + (f32, f32, f32, f32, Vec<[u64; 2]>), + (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), + ) { + rabitq::compute_lut(vector.slice()) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { + rabitq::code(vector.dims(), vector.slice()) + } + + fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { + vector.slice().to_vec() + } + + fn build_from_vecf32(x: &[f32]) -> Self { + Self::new(x.to_vec()) + } +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f16; + + fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f16]>) { + let vector = vector.slice(); + ((), match vector.len() { + 0..=1920 => vec![vector], + 1921..=2560 => vec![&vector[..1280], &vector[1280..]], + 2561.. => vector.chunks(3840).collect(), + }) + } + + fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + (vector.slice(), ()) + } + + fn from_owned(vector: OwnedVector) -> Self { + match vector { + OwnedVector::Vecf16(x) => x, + _ => unreachable!(), + } + } + + fn random_projection(vector: Self::Borrowed<'_>) -> Self { + Self::new(f16::vector_from_f32(&crate::projection::project( + &f16::vector_to_f32(vector.slice()), + ))) + } + + fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u64; 2]>) { + rabitq::block::preprocess(&f16::vector_to_f32(vector.slice())) + } + + fn compute_lut( + vector: Self::Borrowed<'_>, + ) -> ( + (f32, f32, f32, f32, Vec<[u64; 2]>), + (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), + ) { + rabitq::compute_lut(&f16::vector_to_f32(vector.slice())) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { + rabitq::code(vector.dims(), &f16::vector_to_f32(vector.slice())) + } + + fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { + f16::vector_to_f32(vector.slice()) + } + + fn build_from_vecf32(x: &[f32]) -> Self { + Self::new(f16::vector_from_f32(x)) + } +} + +pub trait OperatorDistance: 'static + Debug + Copy { + const KIND: DistanceKind; + + fn compute_lowerbound_binary( + lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), + code: (f32, f32, f32, f32, &[u64]), + epsilon: f32, + ) -> Distance; + + type BlockAccessor: for<'a> Accessor2< + [u64; 2], + [u64; 2], + (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), + (f32, f32, f32, f32, f32), + Output = [Distance; 32], + > + Default; + + fn block_accessor() -> Self::BlockAccessor { + Default::default() + } +} + +#[derive(Debug, Clone, Copy)] +pub struct L2; + +impl OperatorDistance for L2 { + const KIND: DistanceKind = DistanceKind::L2; + + fn compute_lowerbound_binary( + lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), + code: (f32, f32, f32, f32, &[u64]), + epsilon: f32, + ) -> Distance { + rabitq::binary::process_lowerbound_l2(lut, code, epsilon) + } + + type BlockAccessor = Block; +} + +#[derive(Debug, Clone, Copy)] +pub struct Dot; + +impl OperatorDistance for Dot { + const KIND: DistanceKind = DistanceKind::Dot; + + fn compute_lowerbound_binary( + lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), + code: (f32, f32, f32, f32, &[u64]), + epsilon: f32, + ) -> Distance { + rabitq::binary::process_lowerbound_dot(lut, code, epsilon) + } + + type BlockAccessor = Block; +} + +pub trait Operator: 'static + Debug + Copy { + type Vector: Vector; + + type Distance: OperatorDistance; + + type DistanceAccessor: Default + + Accessor2< + ::Element, + ::Element, + ::Metadata, + ::Metadata, + Output = Distance, + >; + + const SUPPORTS_RESIDUAL: bool; + + type ResidualAccessor: Default + + Accessor2< + ::Element, + ::Element, + ::Metadata, + ::Metadata, + Output = Self::Vector, + >; +} + +#[derive(Debug)] +pub struct Op(PhantomData<(fn(V) -> V, fn(D) -> D)>); + +impl Clone for Op { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for Op {} + +impl Operator for Op, L2> { + type Vector = VectOwned; + + type Distance = L2; + + type DistanceAccessor = Sum, L2>>; + + const SUPPORTS_RESIDUAL: bool = true; + + type ResidualAccessor = Diff, L2>>; +} + +impl Operator for Op, Dot> { + type Vector = VectOwned; + + type Distance = Dot; + + type DistanceAccessor = Sum, Dot>>; + + const SUPPORTS_RESIDUAL: bool = false; + + type ResidualAccessor = Diff, Dot>>; +} + +impl Operator for Op, L2> { + type Vector = VectOwned; + + type Distance = L2; + + type DistanceAccessor = Sum, L2>>; + + const SUPPORTS_RESIDUAL: bool = true; + + type ResidualAccessor = Diff, L2>>; +} + +impl Operator for Op, Dot> { + type Vector = VectOwned; + + type Distance = Dot; + + type DistanceAccessor = Sum, Dot>>; + + const SUPPORTS_RESIDUAL: bool = false; + + type ResidualAccessor = Diff, Dot>>; +} diff --git a/src/algorithm/prewarm.rs b/src/algorithm/prewarm.rs new file mode 100644 index 00000000..9373f4a9 --- /dev/null +++ b/src/algorithm/prewarm.rs @@ -0,0 +1,99 @@ +use crate::algorithm::operator::Operator; +use crate::algorithm::tuples::*; +use crate::algorithm::vectors; +use crate::algorithm::{Page, RelationRead}; +use crate::utils::pipe::Pipe; +use std::fmt::Write; + +pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> String { + let meta_guard = relation.read(0); + let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let height_of_root = meta_tuple.height_of_root(); + let root_mean = meta_tuple.root_mean(); + let root_first = meta_tuple.root_first(); + drop(meta_guard); + + let mut message = String::new(); + writeln!(message, "height of root: {}", height_of_root).unwrap(); + let prewarm_max_height = if height < 0 { 0 } else { height as u32 }; + if prewarm_max_height > height_of_root { + return message; + } + type State = Vec; + let mut state: State = { + let mut results = Vec::new(); + let counter = 1_usize; + { + vectors::vector_access_1::(relation.clone(), root_mean, ()); + results.push(root_first); + } + writeln!(message, "number of tuples: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", counter).unwrap(); + results + }; + let mut step = |state: State| { + let mut counter = 0_usize; + let mut results = Vec::new(); + for list in state { + let mut current = list; + while current != u32::MAX { + counter += 1; + pgrx::check_for_interrupts!(); + let h1_guard = relation.read(current); + for i in 1..=h1_guard.len() { + let h1_tuple = h1_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h1_tuple { + H1TupleReader::_0(h1_tuple) => { + for mean in h1_tuple.mean().iter().copied() { + vectors::vector_access_1::(relation.clone(), mean, ()); + } + for first in h1_tuple.first().iter().copied() { + results.push(first); + } + } + H1TupleReader::_1(_) => (), + } + } + current = h1_guard.get_opaque().next; + } + } + writeln!(message, "number of tuples: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", counter).unwrap(); + results + }; + for _ in (std::cmp::max(1, prewarm_max_height)..height_of_root).rev() { + state = step(state); + } + if prewarm_max_height == 0 { + let mut counter = 0_usize; + let mut results = Vec::new(); + for list in state { + let jump_guard = relation.read(list); + let jump_tuple = jump_guard + .get(1) + .expect("data corruption") + .pipe(read_tuple::); + let first = jump_tuple.first(); + let mut current = first; + while current != u32::MAX { + counter += 1; + pgrx::check_for_interrupts!(); + let h0_guard = relation.read(current); + for i in 1..=h0_guard.len() { + let _h0_tuple = h0_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + results.push(()); + } + current = h0_guard.get_opaque().next; + } + } + writeln!(message, "number of tuples: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", counter).unwrap(); + } + message +} diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs new file mode 100644 index 00000000..fee6df3b --- /dev/null +++ b/src/algorithm/scan.rs @@ -0,0 +1,171 @@ +use crate::algorithm::operator::*; +use crate::algorithm::tape::read_h0_tape; +use crate::algorithm::tape::read_h1_tape; +use crate::algorithm::tuples::*; +use crate::algorithm::vectors; +use crate::algorithm::{Page, RelationRead}; +use crate::utils::pipe::Pipe; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::num::NonZeroU64; +use vector::VectorBorrowed; +use vector::VectorOwned; + +pub fn scan( + relation: impl RelationRead + Clone, + vector: O::Vector, + probes: Vec, + epsilon: f32, +) -> impl Iterator { + let vector = O::Vector::random_projection(vector.as_borrowed()); + let meta_guard = relation.read(0); + let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let dims = meta_tuple.dims(); + let is_residual = meta_tuple.is_residual(); + let height_of_root = meta_tuple.height_of_root(); + assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); + assert_eq!(height_of_root as usize, 1 + probes.len(), "invalid probes"); + let root_mean = meta_tuple.root_mean(); + let root_first = meta_tuple.root_first(); + drop(meta_guard); + + let default_lut = if !is_residual { + Some(O::Vector::compute_lut(vector.as_borrowed())) + } else { + None + }; + + type State = Vec<(u32, Option<::Vector>)>; + let mut state: State = vec![{ + let mean = root_mean; + if is_residual { + let residual_u = vectors::vector_access_1::( + relation.clone(), + mean, + LAccess::new( + O::Vector::elements_and_metadata(vector.as_borrowed()), + O::ResidualAccessor::default(), + ), + ); + (root_first, Some(residual_u)) + } else { + (root_first, None) + } + }]; + let step = |state: State, probes| { + let mut results = Vec::new(); + for (first, residual) in state { + let lut = if let Some(residual) = residual { + &O::Vector::compute_lut_block(residual.as_borrowed()) + } else { + default_lut.as_ref().map(|x| &x.0).unwrap() + }; + read_h1_tape( + relation.clone(), + first, + || { + RAccess::new( + (&lut.4, (lut.0, lut.1, lut.2, lut.3, epsilon)), + O::Distance::block_accessor(), + ) + }, + |lowerbound, mean, first| { + results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(first))); + }, + ); + } + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); + std::iter::from_fn(|| { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + if is_residual { + let (dis_u, residual_u) = vectors::vector_access_1::( + relation.clone(), + mean, + LAccess::new( + O::Vector::elements_and_metadata(vector.as_borrowed()), + ( + O::DistanceAccessor::default(), + O::ResidualAccessor::default(), + ), + ), + ); + cache.push(( + Reverse(dis_u), + AlwaysEqual(first), + AlwaysEqual(Some(residual_u)), + )); + } else { + let dis_u = vectors::vector_access_1::( + relation.clone(), + mean, + LAccess::new( + O::Vector::elements_and_metadata(vector.as_borrowed()), + O::DistanceAccessor::default(), + ), + ); + cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); + } + } + let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; + Some((first, mean)) + }) + .take(probes as usize) + .collect() + }; + for i in (1..height_of_root).rev() { + state = step(state, probes[i as usize - 1]); + } + + let mut results = Vec::new(); + for (first, residual) in state { + let lut = if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { + &O::Vector::compute_lut(residual) + } else { + default_lut.as_ref().unwrap() + }; + let jump_guard = relation.read(first); + let jump_tuple = jump_guard + .get(1) + .expect("data corruption") + .pipe(read_tuple::); + let first = jump_tuple.first(); + read_h0_tape( + relation.clone(), + first, + || { + RAccess::new( + (&lut.0.4, (lut.0.0, lut.0.1, lut.0.2, lut.0.3, epsilon)), + O::Distance::block_accessor(), + ) + }, + |code| O::Distance::compute_lowerbound_binary(&lut.1, code, epsilon), + |lowerbound, mean, payload| { + results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(payload))); + }, + ); + } + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _)>::new(); + std::iter::from_fn(move || { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); + if let Some(dis_u) = vectors::vector_access_0::( + relation.clone(), + mean, + pay_u, + LAccess::new( + O::Vector::elements_and_metadata(vector.as_borrowed()), + O::DistanceAccessor::default(), + ), + ) { + cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); + }; + } + let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; + Some((dis_u, pay_u)) + }) +} diff --git a/src/algorithm/tape.rs b/src/algorithm/tape.rs new file mode 100644 index 00000000..4ee722a5 --- /dev/null +++ b/src/algorithm/tape.rs @@ -0,0 +1,349 @@ +use super::RelationRead; +use super::operator::Accessor1; +use crate::algorithm::Page; +use crate::algorithm::PageGuard; +use crate::algorithm::tuples::*; +use crate::utils::pipe::Pipe; +use distance::Distance; +use simd::fast_scan::any_pack; +use simd::fast_scan::padding_pack; +use std::marker::PhantomData; +use std::num::NonZeroU64; +use std::ops::DerefMut; + +pub struct TapeWriter { + head: G, + first: u32, + extend: E, + _phantom: PhantomData T>, +} + +impl TapeWriter +where + G: PageGuard + DerefMut, + G::Target: Page, + E: Fn() -> G, +{ + pub fn create(extend: E) -> Self { + let mut head = extend(); + head.get_opaque_mut().skip = head.id(); + let first = head.id(); + Self { + head, + first, + extend, + _phantom: PhantomData, + } + } + pub fn first(&self) -> u32 { + self.first + } + fn freespace(&self) -> u16 { + self.head.freespace() + } + fn tape_move(&mut self) { + if self.head.len() == 0 { + panic!("tuple is too large to fit in a fresh page"); + } + let next = (self.extend)(); + self.head.get_opaque_mut().next = next.id(); + self.head = next; + } +} + +impl TapeWriter +where + G: PageGuard + DerefMut, + G::Target: Page, + E: Fn() -> G, + T: Tuple, +{ + pub fn push(&mut self, x: T) -> IndexPointer { + let bytes = serialize(&x); + if let Some(i) = self.head.alloc(&bytes) { + pair_to_pointer((self.head.id(), i)) + } else { + let next = (self.extend)(); + self.head.get_opaque_mut().next = next.id(); + self.head = next; + if let Some(i) = self.head.alloc(&bytes) { + pair_to_pointer((self.head.id(), i)) + } else { + panic!("tuple is too large to fit in a fresh page") + } + } + } + fn tape_put(&mut self, x: T) -> IndexPointer { + let bytes = serialize(&x); + if let Some(i) = self.head.alloc(&bytes) { + pair_to_pointer((self.head.id(), i)) + } else { + panic!("tuple is too large to fit in this page") + } + } +} + +pub struct H1Branch { + pub mean: IndexPointer, + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub signs: Vec, + pub first: u32, +} + +pub struct H1TapeWriter { + tape: TapeWriter, + branches: Vec, +} + +impl H1TapeWriter +where + G: PageGuard + DerefMut, + G::Target: Page, + E: Fn() -> G, +{ + pub fn create(extend: E) -> Self { + Self { + tape: TapeWriter::create(extend), + branches: Vec::new(), + } + } + pub fn push(&mut self, branch: H1Branch) { + self.branches.push(branch); + if self.branches.len() == 32 { + let chunk = std::array::from_fn::<_, 32, _>(|_| self.branches.pop().unwrap()); + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + loop { + let freespace = self.tape.freespace(); + match (H1Tuple::fit_0(freespace), H1Tuple::fit_1(freespace)) { + (Some(w), _) if w >= remain.len() => { + self.tape.tape_put(H1Tuple::_0 { + mean: chunk.each_ref().map(|x| x.mean), + dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), + factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), + factor_ip: chunk.each_ref().map(|x| x.factor_ip), + factor_err: chunk.each_ref().map(|x| x.factor_err), + first: chunk.each_ref().map(|x| x.first), + len: chunk.len() as _, + elements: remain, + }); + break; + } + (_, Some(w)) => { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(H1Tuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } + (_, None) => self.tape.tape_move(), + } + } + } + } + pub fn into_inner(mut self) -> TapeWriter { + let chunk = self.branches; + if !chunk.is_empty() { + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + loop { + let freespace = self.tape.freespace(); + match (H1Tuple::fit_0(freespace), H1Tuple::fit_1(freespace)) { + (Some(w), _) if w >= remain.len() => { + self.tape.push(H1Tuple::_0 { + mean: any_pack(chunk.iter().map(|x| x.mean)), + dis_u_2: any_pack(chunk.iter().map(|x| x.dis_u_2)), + factor_ppc: any_pack(chunk.iter().map(|x| x.factor_ppc)), + factor_ip: any_pack(chunk.iter().map(|x| x.factor_ip)), + factor_err: any_pack(chunk.iter().map(|x| x.factor_err)), + first: any_pack(chunk.iter().map(|x| x.first)), + len: chunk.len() as _, + elements: remain, + }); + break; + } + (_, Some(w)) => { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(H1Tuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } + (_, None) => self.tape.tape_move(), + } + } + } + self.tape + } +} + +pub struct H0BranchWriter { + pub mean: IndexPointer, + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub signs: Vec, + pub payload: NonZeroU64, +} + +pub struct H0Tape { + tape: TapeWriter, + branches: Vec, +} + +impl H0Tape +where + G: PageGuard + DerefMut, + G::Target: Page, + E: Fn() -> G, +{ + pub fn create(extend: E) -> Self { + Self { + tape: TapeWriter::create(extend), + branches: Vec::new(), + } + } + pub fn push(&mut self, branch: H0BranchWriter) { + self.branches.push(branch); + if self.branches.len() == 32 { + let chunk = std::array::from_fn::<_, 32, _>(|_| self.branches.pop().unwrap()); + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + loop { + let freespace = self.tape.freespace(); + match (H0Tuple::fit_1(freespace), H0Tuple::fit_2(freespace)) { + (Some(w), _) if w >= remain.len() => { + self.tape.push(H0Tuple::_1 { + mean: chunk.each_ref().map(|x| x.mean), + dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), + factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), + factor_ip: chunk.each_ref().map(|x| x.factor_ip), + factor_err: chunk.each_ref().map(|x| x.factor_err), + payload: chunk.each_ref().map(|x| Some(x.payload)), + elements: remain, + }); + break; + } + (_, Some(w)) => { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(H0Tuple::_2 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } + (_, None) => self.tape.tape_move(), + } + } + } + } + pub fn into_inner(mut self) -> TapeWriter { + for x in self.branches { + self.tape.push(H0Tuple::_0 { + mean: x.mean, + dis_u_2: x.dis_u_2, + factor_ppc: x.factor_ppc, + factor_ip: x.factor_ip, + factor_err: x.factor_err, + payload: Some(x.payload), + elements: rabitq::pack_to_u64(&x.signs), + }); + } + self.tape + } +} + +pub fn read_h1_tape( + relation: impl RelationRead, + first: u32, + compute_block: impl Fn() -> A + Copy, + mut callback: impl FnMut(Distance, IndexPointer, u32), +) where + A: for<'a> Accessor1< + [u64; 2], + (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), + Output = [Distance; 32], + >, +{ + assert!(first != u32::MAX); + let mut current = first; + let mut computing = None; + while current != u32::MAX { + let h1_guard = relation.read(current); + for i in 1..=h1_guard.len() { + let h1_tuple = h1_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h1_tuple { + H1TupleReader::_0(h1_tuple) => { + let mut compute = computing.take().unwrap_or_else(compute_block); + compute.push(h1_tuple.elements()); + let lowerbounds = compute.finish(h1_tuple.metadata()); + for i in 0..h1_tuple.len() { + callback( + lowerbounds[i as usize], + h1_tuple.mean()[i as usize], + h1_tuple.first()[i as usize], + ); + } + } + H1TupleReader::_1(h1_tuple) => { + let computing = computing.get_or_insert_with(compute_block); + computing.push(h1_tuple.elements()); + } + } + } + current = h1_guard.get_opaque().next; + } +} + +pub fn read_h0_tape( + relation: impl RelationRead, + first: u32, + compute_block: impl Fn() -> A + Copy, + compute_binary: impl Fn((f32, f32, f32, f32, &[u64])) -> Distance, + mut callback: impl FnMut(Distance, IndexPointer, NonZeroU64), +) where + A: for<'a> Accessor1< + [u64; 2], + (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), + Output = [Distance; 32], + >, +{ + assert!(first != u32::MAX); + let mut current = first; + let mut computing = None; + while current != u32::MAX { + let h0_guard = relation.read(current); + for i in 1..=h0_guard.len() { + let h0_tuple = h0_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h0_tuple { + H0TupleReader::_0(h0_tuple) => { + let lowerbound = compute_binary(h0_tuple.code()); + if let Some(payload) = h0_tuple.payload() { + callback(lowerbound, h0_tuple.mean(), payload); + } + } + H0TupleReader::_1(h0_tuple) => { + let mut compute = computing.take().unwrap_or_else(compute_block); + compute.push(h0_tuple.elements()); + let lowerbounds = compute.finish(h0_tuple.metadata()); + for j in 0..32 { + if let Some(payload) = h0_tuple.payload()[j] { + callback(lowerbounds[j], h0_tuple.mean()[j], payload); + } + } + } + H0TupleReader::_2(h0_tuple) => { + let computing = computing.get_or_insert_with(compute_block); + computing.push(h0_tuple.elements()); + } + } + } + current = h0_guard.get_opaque().next; + } +} diff --git a/src/algorithm/tuples.rs b/src/algorithm/tuples.rs new file mode 100644 index 00000000..6ecef02e --- /dev/null +++ b/src/algorithm/tuples.rs @@ -0,0 +1,1205 @@ +use crate::algorithm::operator::Vector; +use std::num::{NonZeroU8, NonZeroU64}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; +use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +pub const ALIGN: usize = 8; +pub type Tag = u64; + +pub trait Tuple: 'static { + type Reader<'a>: TupleReader<'a, Tuple = Self>; + fn serialize(&self) -> Vec; +} + +pub trait WithWriter: Tuple { + type Writer<'a>: TupleWriter<'a, Tuple = Self>; +} + +pub trait TupleReader<'a>: Copy { + type Tuple: Tuple; + fn deserialize_ref(source: &'a [u8]) -> Self; +} + +pub trait TupleWriter<'a> { + type Tuple: Tuple; + fn deserialize_mut(source: &'a mut [u8]) -> Self; +} + +pub fn serialize(tuple: &T) -> Vec { + Tuple::serialize(tuple) +} + +pub fn read_tuple(source: &[u8]) -> T::Reader<'_> { + TupleReader::deserialize_ref(source) +} + +pub fn write_tuple(source: &mut [u8]) -> T::Writer<'_> { + TupleWriter::deserialize_mut(source) +} + +// meta tuple + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct MetaTupleHeader { + dims: u32, + height_of_root: u32, + is_residual: Bool, + _padding_0: [ZeroU8; 3], + vectors_first: u32, + // raw vector + root_mean: IndexPointer, + // for meta tuple, it's pointers to next level + root_first: u32, + freepage_first: u32, +} + +pub struct MetaTuple { + pub dims: u32, + pub height_of_root: u32, + pub is_residual: bool, + pub vectors_first: u32, + pub root_mean: IndexPointer, + pub root_first: u32, + pub freepage_first: u32, +} + +impl Tuple for MetaTuple { + type Reader<'a> = MetaTupleReader<'a>; + + fn serialize(&self) -> Vec { + MetaTupleHeader { + dims: self.dims, + height_of_root: self.height_of_root, + is_residual: self.is_residual.into(), + _padding_0: Default::default(), + vectors_first: self.vectors_first, + root_mean: self.root_mean, + root_first: self.root_first, + freepage_first: self.freepage_first, + } + .as_bytes() + .to_vec() + } +} + +#[derive(Debug, Clone, Copy)] +pub struct MetaTupleReader<'a> { + header: &'a MetaTupleHeader, +} + +impl<'a> TupleReader<'a> for MetaTupleReader<'a> { + type Tuple = MetaTuple; + fn deserialize_ref(source: &'a [u8]) -> Self { + let checker = RefChecker::new(source); + let header = checker.prefix(0); + Self { header } + } +} + +impl MetaTupleReader<'_> { + pub fn dims(self) -> u32 { + self.header.dims + } + pub fn height_of_root(self) -> u32 { + self.header.height_of_root + } + pub fn is_residual(self) -> bool { + self.header.is_residual.into() + } + pub fn vectors_first(self) -> u32 { + self.header.vectors_first + } + pub fn root_mean(self) -> IndexPointer { + self.header.root_mean + } + pub fn root_first(self) -> u32 { + self.header.root_first + } + pub fn freepage_first(self) -> u32 { + self.header.freepage_first + } +} + +// freepage tuple + +#[repr(C, align(8))] +#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct FreepageTupleHeader { + a: [u32; 1], + b: [u32; 32], + c: [u32; 32 * 32], + _padding_0: [ZeroU8; 4], +} + +const _: () = assert!(size_of::() == 4232); + +#[derive(Debug, Clone, PartialEq)] +pub struct FreepageTuple {} + +impl Tuple for FreepageTuple { + type Reader<'a> = FreepageTupleReader<'a>; + + fn serialize(&self) -> Vec { + FreepageTupleHeader { + a: std::array::from_fn(|_| 0), + b: std::array::from_fn(|_| 0), + c: std::array::from_fn(|_| 0), + _padding_0: Default::default(), + } + .as_bytes() + .to_vec() + } +} + +impl WithWriter for FreepageTuple { + type Writer<'a> = FreepageTupleWriter<'a>; +} + +#[derive(Debug, Clone, Copy)] +pub struct FreepageTupleReader<'a> { + #[allow(dead_code)] + header: &'a FreepageTupleHeader, +} + +impl<'a> TupleReader<'a> for FreepageTupleReader<'a> { + type Tuple = FreepageTuple; + + fn deserialize_ref(source: &'a [u8]) -> Self { + let checker = RefChecker::new(source); + let header = checker.prefix(0); + Self { header } + } +} + +pub struct FreepageTupleWriter<'a> { + header: &'a mut FreepageTupleHeader, +} + +impl<'a> TupleWriter<'a> for FreepageTupleWriter<'a> { + type Tuple = FreepageTuple; + + fn deserialize_mut(source: &'a mut [u8]) -> Self { + let mut checker = MutChecker::new(source); + let header = checker.prefix(0); + Self { header } + } +} + +impl FreepageTupleWriter<'_> { + pub fn mark(&mut self, i: usize) { + let c_i = i; + self.header.c[c_i / 32] |= 1 << (c_i % 32); + let b_i = i / 32; + self.header.b[b_i / 32] |= 1 << (b_i % 32); + let a_i = i / 32 / 32; + self.header.a[a_i / 32] |= 1 << (a_i % 32); + } + pub fn fetch(&mut self) -> Option { + if self.header.a[0].trailing_ones() == 32 { + return None; + } + let a_i = self.header.a[0].trailing_zeros() as usize; + let b_i = self.header.b[a_i].trailing_zeros() as usize + a_i * 32; + let c_i = self.header.c[b_i].trailing_zeros() as usize + b_i * 32; + self.header.c[c_i / 32] &= !(1 << (c_i % 32)); + if self.header.c[b_i] == 0 { + self.header.b[b_i / 32] &= !(1 << (b_i % 32)); + if self.header.b[a_i] == 0 { + self.header.a[a_i / 32] &= !(1 << (a_i % 32)); + } + } + Some(c_i) + } +} + +// vector tuple + +#[repr(C, align(8))] +#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VectorTupleHeader0 { + payload: Option, + metadata_s: usize, + elements_s: usize, + elements_e: usize, +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VectorTupleHeader1 { + payload: Option, + pointer: IndexPointer, + elements_s: usize, + elements_e: usize, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum VectorTuple { + _0 { + payload: Option, + metadata: V::Metadata, + elements: Vec, + }, + _1 { + payload: Option, + pointer: IndexPointer, + elements: Vec, + }, +} + +impl Tuple for VectorTuple { + type Reader<'a> = VectorTupleReader<'a, V>; + + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + VectorTuple::_0 { + payload, + metadata, + elements, + } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat(0).take(size_of::())); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let metadata_s = buffer.len(); + buffer.extend(metadata.as_bytes()); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len(); + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len(); + buffer[size_of::()..][..size_of::()].copy_from_slice( + VectorTupleHeader0 { + payload: *payload, + metadata_s, + elements_s, + elements_e, + } + .as_bytes(), + ); + } + VectorTuple::_1 { + payload, + pointer, + elements, + } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat(0).take(size_of::())); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len(); + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len(); + buffer[size_of::()..][..size_of::()].copy_from_slice( + VectorTupleHeader1 { + payload: *payload, + pointer: *pointer, + elements_s, + elements_e, + } + .as_bytes(), + ); + } + } + buffer + } +} + +#[derive(Clone)] +pub struct VectorTupleReader0<'a, V: Vector> { + header: &'a VectorTupleHeader0, + metadata: &'a V::Metadata, + elements: &'a [V::Element], +} + +impl Copy for VectorTupleReader0<'_, V> {} + +#[derive(Clone)] +pub struct VectorTupleReader1<'a, V: Vector> { + header: &'a VectorTupleHeader1, + elements: &'a [V::Element], +} + +impl Copy for VectorTupleReader1<'_, V> {} + +#[derive(Clone)] +pub enum VectorTupleReader<'a, V: Vector> { + _0(VectorTupleReader0<'a, V>), + _1(VectorTupleReader1<'a, V>), +} + +impl Copy for VectorTupleReader<'_, V> {} + +impl<'a, V: Vector> TupleReader<'a> for VectorTupleReader<'a, V> { + type Tuple = VectorTuple; + + fn deserialize_ref(source: &'a [u8]) -> Self { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader0 = checker.prefix(size_of::()); + let metadata = checker.prefix(header.metadata_s); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_0(VectorTupleReader0 { + header, + elements, + metadata, + }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_1(VectorTupleReader1 { header, elements }) + } + _ => panic!("bad bytes"), + } + } +} + +impl<'a, V: Vector> VectorTupleReader<'a, V> { + pub fn payload(self) -> Option { + match self { + VectorTupleReader::_0(this) => this.header.payload, + VectorTupleReader::_1(this) => this.header.payload, + } + } + pub fn elements(self) -> &'a [::Element] { + match self { + VectorTupleReader::_0(this) => this.elements, + VectorTupleReader::_1(this) => this.elements, + } + } + pub fn metadata_or_pointer(self) -> Result { + match self { + VectorTupleReader::_0(this) => Ok(*this.metadata), + VectorTupleReader::_1(this) => Err(this.header.pointer), + } + } +} + +// height1tuple + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct H1TupleHeader0 { + mean: [IndexPointer; 32], + dis_u_2: [f32; 32], + factor_ppc: [f32; 32], + factor_ip: [f32; 32], + factor_err: [f32; 32], + first: [u32; 32], + len: u32, + _padding_0: [ZeroU8; 4], + elements_s: usize, + elements_e: usize, +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct H1TupleHeader1 { + elements_s: usize, + elements_e: usize, +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug, Clone, PartialEq)] +pub enum H1Tuple { + _0 { + mean: [IndexPointer; 32], + dis_u_2: [f32; 32], + factor_ppc: [f32; 32], + factor_ip: [f32; 32], + factor_err: [f32; 32], + first: [u32; 32], + len: u32, + elements: Vec<[u64; 2]>, + }, + _1 { + elements: Vec<[u64; 2]>, + }, +} + +impl H1Tuple { + pub fn fit_0(freespace: u16) -> Option { + let mut freespace = freespace as isize; + freespace -= size_of::() as isize; + freespace -= size_of::() as isize; + if freespace >= 0 { + Some(freespace as usize / size_of::<[u64; 2]>()) + } else { + None + } + } + pub fn fit_1(freespace: u16) -> Option { + let mut freespace = freespace as isize; + freespace -= size_of::() as isize; + freespace -= size_of::() as isize; + if freespace >= 0 { + Some(freespace as usize / size_of::<[u64; 2]>()) + } else { + None + } + } +} + +impl Tuple for H1Tuple { + type Reader<'a> = H1TupleReader<'a>; + + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + Self::_0 { + mean, + dis_u_2, + factor_ppc, + factor_ip, + factor_err, + first, + len, + elements, + } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat(0).take(size_of::())); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len(); + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len(); + buffer[size_of::()..][..size_of::()].copy_from_slice( + H1TupleHeader0 { + mean: *mean, + dis_u_2: *dis_u_2, + factor_ppc: *factor_ppc, + factor_ip: *factor_ip, + factor_err: *factor_err, + first: *first, + len: *len, + _padding_0: Default::default(), + elements_s, + elements_e, + } + .as_bytes(), + ); + } + Self::_1 { elements } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat(0).take(size_of::())); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len(); + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len(); + buffer[size_of::()..][..size_of::()].copy_from_slice( + H1TupleHeader1 { + elements_s, + elements_e, + } + .as_bytes(), + ); + } + } + buffer + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum H1TupleReader<'a> { + _0(H1TupleReader0<'a>), + _1(H1TupleReader1<'a>), +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct H1TupleReader0<'a> { + header: &'a H1TupleHeader0, + elements: &'a [[u64; 2]], +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct H1TupleReader1<'a> { + header: &'a H1TupleHeader1, + elements: &'a [[u64; 2]], +} + +impl<'a> TupleReader<'a> for H1TupleReader<'a> { + type Tuple = H1Tuple; + + fn deserialize_ref(source: &'a [u8]) -> Self { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &H1TupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_0(H1TupleReader0 { header, elements }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &H1TupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_1(H1TupleReader1 { header, elements }) + } + _ => panic!("bad bytes"), + } + } +} + +impl<'a> H1TupleReader0<'a> { + pub fn len(self) -> u32 { + self.header.len + } + pub fn mean(self) -> &'a [IndexPointer] { + &self.header.mean[..self.header.len as usize] + } + pub fn metadata(self) -> (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]) { + ( + &self.header.dis_u_2, + &self.header.factor_ppc, + &self.header.factor_ip, + &self.header.factor_err, + ) + } + pub fn first(self) -> &'a [u32] { + &self.header.first[..self.header.len as usize] + } + pub fn elements(&self) -> &'a [[u64; 2]] { + self.elements + } +} + +impl<'a> H1TupleReader1<'a> { + pub fn elements(&self) -> &'a [[u64; 2]] { + self.elements + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct JumpTupleHeader { + first: u32, + _padding_0: [ZeroU8; 4], +} + +#[derive(Debug, Clone, PartialEq)] +pub struct JumpTuple { + pub first: u32, +} + +impl Tuple for JumpTuple { + type Reader<'a> = JumpTupleReader<'a>; + + fn serialize(&self) -> Vec { + JumpTupleHeader { + first: self.first, + _padding_0: Default::default(), + } + .as_bytes() + .to_vec() + } +} + +impl WithWriter for JumpTuple { + type Writer<'a> = JumpTupleWriter<'a>; +} + +#[derive(Debug, Clone, Copy)] +pub struct JumpTupleReader<'a> { + header: &'a JumpTupleHeader, +} + +impl<'a> TupleReader<'a> for JumpTupleReader<'a> { + type Tuple = JumpTuple; + + fn deserialize_ref(source: &'a [u8]) -> Self { + let checker = RefChecker::new(source); + let header: &JumpTupleHeader = checker.prefix(0); + Self { header } + } +} + +impl JumpTupleReader<'_> { + pub fn first(self) -> u32 { + self.header.first + } +} + +#[derive(Debug)] +pub struct JumpTupleWriter<'a> { + header: &'a mut JumpTupleHeader, +} + +impl<'a> TupleWriter<'a> for JumpTupleWriter<'a> { + type Tuple = JumpTuple; + + fn deserialize_mut(source: &'a mut [u8]) -> Self { + let mut checker = MutChecker::new(source); + let header: &mut JumpTupleHeader = checker.prefix(0); + Self { header } + } +} + +impl JumpTupleWriter<'_> { + pub fn first(&mut self) -> &mut u32 { + &mut self.header.first + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct H0TupleHeader0 { + mean: IndexPointer, + dis_u_2: f32, + factor_ppc: f32, + factor_ip: f32, + factor_err: f32, + payload: Option, + elements_s: usize, + elements_e: usize, +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct H0TupleHeader1 { + mean: [IndexPointer; 32], + dis_u_2: [f32; 32], + factor_ppc: [f32; 32], + factor_ip: [f32; 32], + factor_err: [f32; 32], + payload: [Option; 32], + elements_s: usize, + elements_e: usize, +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct H0TupleHeader2 { + elements_s: usize, + elements_e: usize, +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug, Clone, PartialEq)] +pub enum H0Tuple { + _0 { + mean: IndexPointer, + dis_u_2: f32, + factor_ppc: f32, + factor_ip: f32, + factor_err: f32, + payload: Option, + elements: Vec, + }, + _1 { + mean: [IndexPointer; 32], + dis_u_2: [f32; 32], + factor_ppc: [f32; 32], + factor_ip: [f32; 32], + factor_err: [f32; 32], + payload: [Option; 32], + elements: Vec<[u64; 2]>, + }, + _2 { + elements: Vec<[u64; 2]>, + }, +} + +impl H0Tuple { + pub fn fit_1(freespace: u16) -> Option { + let mut freespace = freespace as isize; + freespace -= size_of::() as isize; + freespace -= size_of::() as isize; + if freespace >= 0 { + Some(freespace as usize / size_of::<[u64; 2]>()) + } else { + None + } + } + pub fn fit_2(freespace: u16) -> Option { + let mut freespace = freespace as isize; + freespace -= size_of::() as isize; + freespace -= size_of::() as isize; + if freespace >= 0 { + Some(freespace as usize / size_of::<[u64; 2]>()) + } else { + None + } + } +} + +impl Tuple for H0Tuple { + type Reader<'a> = H0TupleReader<'a>; + + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + H0Tuple::_0 { + mean, + dis_u_2, + factor_ppc, + factor_ip, + factor_err, + payload, + elements, + } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat(0).take(size_of::())); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len(); + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len(); + buffer[size_of::()..][..size_of::()].copy_from_slice( + H0TupleHeader0 { + mean: *mean, + dis_u_2: *dis_u_2, + factor_ppc: *factor_ppc, + factor_ip: *factor_ip, + factor_err: *factor_err, + payload: *payload, + elements_s, + elements_e, + } + .as_bytes(), + ); + } + H0Tuple::_1 { + mean, + dis_u_2, + factor_ppc, + factor_ip, + factor_err, + payload, + elements, + } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat(0).take(size_of::())); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len(); + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len(); + buffer[size_of::()..][..size_of::()].copy_from_slice( + H0TupleHeader1 { + mean: *mean, + dis_u_2: *dis_u_2, + factor_ppc: *factor_ppc, + factor_ip: *factor_ip, + factor_err: *factor_err, + payload: *payload, + elements_s, + elements_e, + } + .as_bytes(), + ); + } + Self::_2 { elements } => { + buffer.extend((2 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat(0).take(size_of::())); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len(); + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len(); + buffer[size_of::()..][..size_of::()].copy_from_slice( + H0TupleHeader2 { + elements_s, + elements_e, + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithWriter for H0Tuple { + type Writer<'a> = H0TupleWriter<'a>; +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum H0TupleReader<'a> { + _0(H0TupleReader0<'a>), + _1(H0TupleReader1<'a>), + _2(H0TupleReader2<'a>), +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct H0TupleReader0<'a> { + header: &'a H0TupleHeader0, + elements: &'a [u64], +} + +impl<'a> H0TupleReader0<'a> { + pub fn mean(self) -> IndexPointer { + self.header.mean + } + pub fn code(self) -> (f32, f32, f32, f32, &'a [u64]) { + ( + self.header.dis_u_2, + self.header.factor_ppc, + self.header.factor_ip, + self.header.factor_err, + self.elements, + ) + } + pub fn payload(self) -> Option { + self.header.payload + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct H0TupleReader1<'a> { + header: &'a H0TupleHeader1, + elements: &'a [[u64; 2]], +} + +impl<'a> H0TupleReader1<'a> { + pub fn mean(self) -> &'a [IndexPointer; 32] { + &self.header.mean + } + pub fn metadata(self) -> (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]) { + ( + &self.header.dis_u_2, + &self.header.factor_ppc, + &self.header.factor_ip, + &self.header.factor_err, + ) + } + pub fn elements(self) -> &'a [[u64; 2]] { + self.elements + } + pub fn payload(self) -> &'a [Option; 32] { + &self.header.payload + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct H0TupleReader2<'a> { + header: &'a H0TupleHeader2, + elements: &'a [[u64; 2]], +} + +impl<'a> H0TupleReader2<'a> { + pub fn elements(self) -> &'a [[u64; 2]] { + self.elements + } +} + +impl<'a> TupleReader<'a> for H0TupleReader<'a> { + type Tuple = H0Tuple; + + fn deserialize_ref(source: &'a [u8]) -> Self { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &H0TupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_0(H0TupleReader0 { header, elements }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &H0TupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_1(H0TupleReader1 { header, elements }) + } + 2 => { + let checker = RefChecker::new(source); + let header: &H0TupleHeader2 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_2(H0TupleReader2 { header, elements }) + } + _ => panic!("bad bytes"), + } + } +} + +#[derive(Debug)] +pub enum H0TupleWriter<'a> { + _0(H0TupleWriter0<'a>), + _1(H0TupleWriter1<'a>), + #[allow(dead_code)] + _2(H0TupleWriter2<'a>), +} + +#[derive(Debug)] +pub struct H0TupleWriter0<'a> { + header: &'a mut H0TupleHeader0, + #[allow(dead_code)] + elements: &'a mut [u64], +} + +#[derive(Debug)] +pub struct H0TupleWriter1<'a> { + header: &'a mut H0TupleHeader1, + #[allow(dead_code)] + elements: &'a mut [[u64; 2]], +} + +#[derive(Debug)] +pub struct H0TupleWriter2<'a> { + #[allow(dead_code)] + header: &'a mut H0TupleHeader2, + #[allow(dead_code)] + elements: &'a mut [[u64; 2]], +} + +impl<'a> TupleWriter<'a> for H0TupleWriter<'a> { + type Tuple = H0Tuple; + + fn deserialize_mut(source: &'a mut [u8]) -> Self { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let mut checker = MutChecker::new(source); + let header: &mut H0TupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_0(H0TupleWriter0 { header, elements }) + } + 1 => { + let mut checker = MutChecker::new(source); + let header: &mut H0TupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_1(H0TupleWriter1 { header, elements }) + } + 2 => { + let mut checker = MutChecker::new(source); + let header: &mut H0TupleHeader2 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + Self::_2(H0TupleWriter2 { header, elements }) + } + _ => panic!("bad bytes"), + } + } +} + +impl H0TupleWriter0<'_> { + pub fn payload(&mut self) -> &mut Option { + &mut self.header.payload + } +} + +impl H0TupleWriter1<'_> { + pub fn payload(&mut self) -> &mut [Option; 32] { + &mut self.header.payload + } +} + +pub const fn pointer_to_pair(pointer: IndexPointer) -> (u32, u16) { + let value = pointer.0; + (((value >> 16) & 0xffffffff) as u32, (value & 0xffff) as u16) +} + +pub const fn pair_to_pointer(pair: (u32, u16)) -> IndexPointer { + let mut value = 0; + value |= (pair.0 as u64) << 16; + value |= pair.1 as u64; + IndexPointer(value) +} + +#[allow(dead_code)] +const fn soundness_check(a: (u32, u16)) { + let b = pair_to_pointer(a); + let c = pointer_to_pair(b); + assert!(a.0 == c.0); + assert!(a.1 == c.1); +} + +const _: () = soundness_check((111, 222)); + +#[repr(transparent)] +#[derive( + Debug, + Default, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct ZeroU8(Option); + +#[repr(transparent)] +#[derive( + Debug, Default, Clone, Copy, PartialEq, Eq, Hash, IntoBytes, FromBytes, Immutable, KnownLayout, +)] +pub struct IndexPointer(pub u64); + +#[repr(transparent)] +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct Bool(u8); + +impl Bool { + pub const FALSE: Self = Self(0); + pub const TRUE: Self = Self(1); +} + +impl From for bool { + fn from(value: Bool) -> Self { + value != Bool::FALSE + } +} + +impl From for Bool { + fn from(value: bool) -> Self { + if value { Self::TRUE } else { Self::FALSE } + } +} + +pub struct RefChecker<'a> { + bytes: &'a [u8], +} + +impl<'a> RefChecker<'a> { + pub fn new(bytes: &'a [u8]) -> Self { + Self { bytes } + } + pub fn prefix( + &self, + s: usize, + ) -> &'a T { + let start = s; + let end = s + size_of::(); + let bytes = &self.bytes[start..end]; + FromBytes::ref_from_bytes(bytes).expect("bad bytes") + } + pub fn bytes( + &self, + s: usize, + e: usize, + ) -> &'a T { + let start = s; + let end = e; + let bytes = &self.bytes[start..end]; + FromBytes::ref_from_bytes(bytes).expect("bad bytes") + } +} + +pub struct MutChecker<'a> { + flag: Vec, + bytes: &'a mut [u8], +} + +impl<'a> MutChecker<'a> { + pub fn new(bytes: &'a mut [u8]) -> Self { + Self { + flag: vec![0u64; bytes.len().div_ceil(64)], + bytes, + } + } + pub fn prefix( + &mut self, + s: usize, + ) -> &'a mut T { + let start = s; + let end = s + size_of::(); + if !(start <= end && end <= self.bytes.len()) { + panic!("bad bytes"); + } + for i in start..end { + if (self.flag[i / 64] & (1 << (i % 64))) != 0 { + panic!("bad bytes"); + } else { + self.flag[i / 64] |= 1 << (i % 64); + } + } + let bytes = unsafe { + std::slice::from_raw_parts_mut(self.bytes.as_mut_ptr().add(start), end - start) + }; + FromBytes::mut_from_bytes(bytes).expect("bad bytes") + } + pub fn bytes( + &mut self, + s: usize, + e: usize, + ) -> &'a mut T { + let start = s; + let end = e; + if !(start <= end && end <= self.bytes.len()) { + panic!("bad bytes"); + } + for i in start..end { + if (self.flag[i / 64] & (1 << (i % 64))) != 0 { + panic!("bad bytes"); + } else { + self.flag[i / 64] |= 1 << (i % 64); + } + } + let bytes = unsafe { + std::slice::from_raw_parts_mut(self.bytes.as_mut_ptr().add(start), end - start) + }; + FromBytes::mut_from_bytes(bytes).expect("bad bytes") + } +} + +// this test only passes if `MIRIFLAGS="-Zmiri-tree-borrows"` is set +#[test] +fn aliasing_test() { + #[repr(C, align(8))] + #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct ExampleHeader { + elements_s: usize, + elements_e: usize, + } + let serialized = { + let elements = (0u32..1111).collect::>(); + let mut buffer = Vec::::new(); + buffer.extend(std::iter::repeat(0).take(size_of::())); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len(); + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len(); + buffer[..size_of::()].copy_from_slice( + ExampleHeader { + elements_s, + elements_e, + } + .as_bytes(), + ); + buffer + }; + let mut source = vec![0u64; serialized.len().next_multiple_of(8)]; + source.as_mut_bytes()[..serialized.len()].copy_from_slice(&serialized); + let deserialized = { + let mut checker = MutChecker::new(source.as_mut_bytes()); + let header: &mut ExampleHeader = checker.prefix(0); + let elements: &mut [u32] = checker.bytes(header.elements_s, header.elements_e); + (header, elements) + }; + assert_eq!( + deserialized.1, + (0u32..1111).collect::>().as_slice() + ); +} diff --git a/src/algorithm/vacuum.rs b/src/algorithm/vacuum.rs new file mode 100644 index 00000000..17366256 --- /dev/null +++ b/src/algorithm/vacuum.rs @@ -0,0 +1,311 @@ +use crate::algorithm::freepages; +use crate::algorithm::operator::Operator; +use crate::algorithm::tape::*; +use crate::algorithm::tuples::*; +use crate::algorithm::{Page, RelationWrite}; +use crate::utils::pipe::Pipe; +use simd::fast_scan::unpack; +use std::num::NonZeroU64; + +pub fn bulkdelete( + relation: impl RelationWrite, + delay: impl Fn(), + callback: impl Fn(NonZeroU64) -> bool, +) { + let meta_guard = relation.read(0); + let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let height_of_root = meta_tuple.height_of_root(); + let root_first = meta_tuple.root_first(); + let vectors_first = meta_tuple.vectors_first(); + drop(meta_guard); + { + type State = Vec; + let mut state: State = vec![root_first]; + let step = |state: State| { + let mut results = Vec::new(); + for first in state { + let mut current = first; + while current != u32::MAX { + let h1_guard = relation.read(current); + for i in 1..=h1_guard.len() { + let h1_tuple = h1_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h1_tuple { + H1TupleReader::_0(h1_tuple) => { + for first in h1_tuple.first().iter().copied() { + results.push(first); + } + } + H1TupleReader::_1(_) => (), + } + } + current = h1_guard.get_opaque().next; + } + } + results + }; + for _ in (1..height_of_root).rev() { + state = step(state); + } + for first in state { + let jump_guard = relation.read(first); + let jump_tuple = jump_guard + .get(1) + .expect("data corruption") + .pipe(read_tuple::); + let first = jump_tuple.first(); + let mut current = first; + while current != u32::MAX { + delay(); + let read = relation.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + let h0_tuple = read + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h0_tuple { + H0TupleReader::_0(h0_tuple) => { + let p = h0_tuple.payload(); + if let Some(payload) = p { + if callback(payload) { + break 'flag true; + } + } + } + H0TupleReader::_1(h0_tuple) => { + let p = h0_tuple.payload(); + for j in 0..32 { + if let Some(payload) = p[j] { + if callback(payload) { + break 'flag true; + } + } + } + } + H0TupleReader::_2(_) => (), + } + } + false + }; + if flag { + drop(read); + let mut write = relation.write(current, false); + for i in 1..=write.len() { + let h0_tuple = write + .get_mut(i) + .expect("data corruption") + .pipe(write_tuple::); + match h0_tuple { + H0TupleWriter::_0(mut h0_tuple) => { + let p = h0_tuple.payload(); + if let Some(payload) = *p { + if callback(payload) { + *p = None; + } + } + } + H0TupleWriter::_1(mut h0_tuple) => { + let p = h0_tuple.payload(); + for j in 0..32 { + if let Some(payload) = p[j] { + if callback(payload) { + p[j] = None; + } + } + } + } + H0TupleWriter::_2(_) => (), + } + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; + } + } + } + } + { + let first = vectors_first; + let mut current = first; + while current != u32::MAX { + delay(); + let read = relation.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + if let Some(vector_bytes) = read.get(i) { + let vector_tuple = vector_bytes.pipe(read_tuple::>); + let p = vector_tuple.payload(); + if let Some(payload) = p { + if callback(payload) { + break 'flag true; + } + } + } + } + false + }; + if flag { + drop(read); + let mut write = relation.write(current, true); + for i in 1..=write.len() { + if let Some(vector_bytes) = write.get(i) { + let vector_tuple = vector_bytes.pipe(read_tuple::>); + let p = vector_tuple.payload(); + if let Some(payload) = p { + if callback(payload) { + write.free(i); + } + } + }; + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; + } + } + } +} + +pub fn maintain(relation: impl RelationWrite + Clone, delay: impl Fn()) { + let meta_guard = relation.read(0); + let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let dims = meta_tuple.dims(); + let height_of_root = meta_tuple.height_of_root(); + let root_first = meta_tuple.root_first(); + let freepage_first = meta_tuple.freepage_first(); + drop(meta_guard); + + let firsts = { + type State = Vec; + let mut state: State = vec![root_first]; + let step = |state: State| { + let mut results = Vec::new(); + for first in state { + let mut current = first; + while current != u32::MAX { + delay(); + let h1_guard = relation.read(current); + for i in 1..=h1_guard.len() { + let h1_tuple = h1_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h1_tuple { + H1TupleReader::_0(h1_tuple) => { + for first in h1_tuple.first().iter().copied() { + results.push(first); + } + } + H1TupleReader::_1(_) => (), + } + } + current = h1_guard.get_opaque().next; + } + } + results + }; + for _ in (1..height_of_root).rev() { + state = step(state); + } + state + }; + + for first in firsts { + let mut jump_guard = relation.write(first, false); + let mut jump_tuple = jump_guard + .get_mut(1) + .expect("data corruption") + .pipe(write_tuple::); + + let mut tape = H0Tape::<_, _>::create(|| { + if let Some(id) = freepages::fetch(relation.clone(), freepage_first) { + let mut write = relation.write(id, false); + write.clear(); + write + } else { + relation.extend(false) + } + }); + + let mut trace = Vec::new(); + + let first = *jump_tuple.first(); + let mut current = first; + let mut computing = None; + while current != u32::MAX { + delay(); + trace.push(current); + let h0_guard = relation.read(current); + for i in 1..=h0_guard.len() { + let h0_tuple = h0_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h0_tuple { + H0TupleReader::_0(h0_tuple) => { + if let Some(payload) = h0_tuple.payload() { + tape.push(H0BranchWriter { + mean: h0_tuple.mean(), + dis_u_2: h0_tuple.code().0, + factor_ppc: h0_tuple.code().1, + factor_ip: h0_tuple.code().2, + factor_err: h0_tuple.code().3, + signs: h0_tuple + .code() + .4 + .iter() + .flat_map(|x| { + std::array::from_fn::<_, 64, _>(|i| *x & (1 << i) != 0) + }) + .take(dims as _) + .collect::>(), + payload, + }); + } + } + H0TupleReader::_1(h0_tuple) => { + let computing = &mut computing.take().unwrap_or_else(Vec::new); + computing.extend_from_slice(h0_tuple.elements()); + let unpacked = unpack(computing); + for j in 0..32 { + if let Some(payload) = h0_tuple.payload()[j] { + tape.push(H0BranchWriter { + mean: h0_tuple.mean()[j], + dis_u_2: h0_tuple.metadata().0[j], + factor_ppc: h0_tuple.metadata().1[j], + factor_ip: h0_tuple.metadata().2[j], + factor_err: h0_tuple.metadata().3[j], + signs: unpacked[j] + .iter() + .flat_map(|&x| { + [x & 1 != 0, x & 2 != 0, x & 4 != 0, x & 8 != 0] + }) + .collect(), + payload, + }); + } + } + } + H0TupleReader::_2(h0_tuple) => { + let computing = computing.get_or_insert_with(Vec::new); + computing.extend_from_slice(h0_tuple.elements()); + } + } + } + current = h0_guard.get_opaque().next; + drop(h0_guard); + } + + let tape = tape.into_inner(); + let new = tape.first(); + drop(tape); + + *jump_tuple.first() = new; + drop(jump_guard); + + freepages::mark(relation.clone(), freepage_first, &trace); + } +} diff --git a/src/algorithm/vectors.rs b/src/algorithm/vectors.rs new file mode 100644 index 00000000..d71499bb --- /dev/null +++ b/src/algorithm/vectors.rs @@ -0,0 +1,133 @@ +use crate::algorithm::operator::*; +use crate::algorithm::tuples::*; +use crate::algorithm::{Page, PageGuard, RelationRead, RelationWrite}; +use crate::utils::pipe::Pipe; +use std::num::NonZeroU64; +use vector::VectorOwned; + +pub fn vector_access_1< + O: Operator, + A: Accessor1<::Element, ::Metadata>, +>( + relation: impl RelationRead, + mean: IndexPointer, + accessor: A, +) -> A::Output { + let mut cursor = Err(mean); + let mut result = accessor; + while let Err(mean) = cursor.map_err(pointer_to_pair) { + let vector_guard = relation.read(mean.0); + let vector_tuple = vector_guard + .get(mean.1) + .expect("data corruption") + .pipe(read_tuple::>); + if vector_tuple.payload().is_some() { + panic!("data corruption"); + } + result.push(vector_tuple.elements()); + cursor = vector_tuple.metadata_or_pointer(); + } + result.finish(cursor.expect("data corruption")) +} + +pub fn vector_access_0< + O: Operator, + A: Accessor1<::Element, ::Metadata>, +>( + relation: impl RelationRead, + mean: IndexPointer, + payload: NonZeroU64, + accessor: A, +) -> Option { + let mut cursor = Err(mean); + let mut result = accessor; + while let Err(mean) = cursor.map_err(pointer_to_pair) { + let vector_guard = relation.read(mean.0); + let vector_tuple = vector_guard + .get(mean.1)? + .pipe(read_tuple::>); + if vector_tuple.payload().is_none() { + panic!("data corruption"); + } + if vector_tuple.payload() != Some(payload) { + return None; + } + result.push(vector_tuple.elements()); + cursor = vector_tuple.metadata_or_pointer(); + } + Some(result.finish(cursor.ok()?)) +} + +pub fn vector_append( + relation: impl RelationWrite + Clone, + vectors_first: u32, + vector: ::Borrowed<'_>, + payload: NonZeroU64, +) -> IndexPointer { + fn append(relation: impl RelationWrite, first: u32, bytes: &[u8]) -> IndexPointer { + if let Some(mut write) = relation.search(bytes.len()) { + let i = write.alloc(bytes).unwrap(); + return pair_to_pointer((write.id(), i)); + } + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = relation.read(current); + if read.freespace() as usize >= bytes.len() || read.get_opaque().next == u32::MAX { + drop(read); + let mut write = relation.write(current, true); + if let Some(i) = write.alloc(bytes) { + return pair_to_pointer((current, i)); + } + if write.get_opaque().next == u32::MAX { + let mut extend = relation.extend(true); + write.get_opaque_mut().next = extend.id(); + drop(write); + if let Some(i) = extend.alloc(bytes) { + let result = (extend.id(), i); + drop(extend); + let mut past = relation.write(first, true); + let skip = &mut past.get_opaque_mut().skip; + assert!(*skip != u32::MAX); + *skip = std::cmp::max(*skip, result.0); + return pair_to_pointer(result); + } else { + panic!("a tuple cannot even be fit in a fresh page"); + } + } + if current == first && write.get_opaque().skip != first { + current = write.get_opaque().skip; + } else { + current = write.get_opaque().next; + } + } else { + if current == first && read.get_opaque().skip != first { + current = read.get_opaque().skip; + } else { + current = read.get_opaque().next; + } + } + } + } + let (metadata, slices) = O::Vector::vector_split(vector); + let mut chain = Ok(metadata); + for i in (0..slices.len()).rev() { + chain = Err(append( + relation.clone(), + vectors_first, + &serialize::>(&match chain { + Ok(metadata) => VectorTuple::_0 { + elements: slices[i].to_vec(), + payload: Some(payload), + metadata, + }, + Err(pointer) => VectorTuple::_1 { + elements: slices[i].to_vec(), + payload: Some(payload), + pointer, + }, + }), + )); + } + chain.err().unwrap() +} diff --git a/src/vchordrq/gucs/executing.rs b/src/gucs/executing.rs similarity index 100% rename from src/vchordrq/gucs/executing.rs rename to src/gucs/executing.rs diff --git a/src/vchordrq/gucs/mod.rs b/src/gucs/mod.rs similarity index 100% rename from src/vchordrq/gucs/mod.rs rename to src/gucs/mod.rs diff --git a/src/vchordrq/gucs/prewarm.rs b/src/gucs/prewarm.rs similarity index 100% rename from src/vchordrq/gucs/prewarm.rs rename to src/gucs/prewarm.rs diff --git a/src/vchordrq/index/am.rs b/src/index/am.rs similarity index 67% rename from src/vchordrq/index/am.rs rename to src/index/am.rs index 3d19a490..5db07d24 100644 --- a/src/vchordrq/index/am.rs +++ b/src/index/am.rs @@ -1,12 +1,13 @@ +use crate::algorithm; +use crate::algorithm::build::{HeapRelation, Reporter}; +use crate::algorithm::operator::{Dot, L2, Op}; +use crate::algorithm::operator::{Operator, Vector}; +use crate::index::am_options::{Opfamily, Reloption}; +use crate::index::am_scan::Scanner; +use crate::index::utils::{ctid_to_pointer, pointer_to_ctid}; +use crate::index::{am_options, am_scan}; use crate::postgres::PostgresRelation; -use crate::vchordrq::algorithm; -use crate::vchordrq::algorithm::build::{HeapRelation, Reporter}; -use crate::vchordrq::algorithm::tuples::Vector; -use crate::vchordrq::index::am_options::{Opfamily, Reloption}; -use crate::vchordrq::index::am_scan::Scanner; -use crate::vchordrq::index::utils::{ctid_to_pointer, pointer_to_ctid}; -use crate::vchordrq::index::{am_options, am_scan}; -use crate::vchordrq::types::VectorKind; +use crate::types::{DistanceKind, VectorKind}; use half::f16; use pgrx::datum::Internal; use pgrx::pg_sys::Datum; @@ -167,17 +168,17 @@ pub unsafe extern "C" fn ambuild( index_info: *mut pgrx::pg_sys::IndexInfo, opfamily: Opfamily, } - impl HeapRelation for Heap { + impl HeapRelation for Heap { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((NonZeroU64, V)), + F: FnMut((NonZeroU64, O::Vector)), { pub struct State<'a, F> { pub this: &'a Heap, pub callback: F, } #[pgrx::pg_guard] - unsafe extern "C" fn call( + unsafe extern "C" fn call( _index: pgrx::pg_sys::Relation, ctid: pgrx::pg_sys::ItemPointer, values: *mut Datum, @@ -185,14 +186,14 @@ pub unsafe extern "C" fn ambuild( _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((NonZeroU64, V)), + F: FnMut((NonZeroU64, O::Vector)), { let state = unsafe { &mut *state.cast::>() }; let opfamily = state.this.opfamily; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { - (state.callback)((pointer, V::from_owned(vector))); + (state.callback)((pointer, O::Vector::from_owned(vector))); } } let table_am = unsafe { &*(*self.heap).rd_tableam }; @@ -210,7 +211,7 @@ pub unsafe extern "C" fn ambuild( progress, 0, pgrx::pg_sys::InvalidBlockNumber, - Some(call::), + Some(call::), (&mut state) as *mut State as *mut _, std::ptr::null_mut(), ); @@ -243,21 +244,43 @@ pub unsafe extern "C" fn ambuild( }; let mut reporter = PgReporter {}; let index_relation = unsafe { PostgresRelation::new(index) }; - match opfamily.vector_kind() { - VectorKind::Vecf32 => algorithm::build::build::, Heap, _>( - vector_options, - vchordrq_options, - heap_relation.clone(), - index_relation.clone(), - reporter.clone(), - ), - VectorKind::Vecf16 => algorithm::build::build::, Heap, _>( - vector_options, - vchordrq_options, - heap_relation.clone(), - index_relation.clone(), - reporter.clone(), - ), + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + algorithm::build::build::, L2>, Heap, _>( + vector_options, + vchordrq_options, + heap_relation.clone(), + index_relation.clone(), + reporter.clone(), + ) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + algorithm::build::build::, Dot>, Heap, _>( + vector_options, + vchordrq_options, + heap_relation.clone(), + index_relation.clone(), + reporter.clone(), + ) + } + (VectorKind::Vecf16, DistanceKind::L2) => { + algorithm::build::build::, L2>, Heap, _>( + vector_options, + vchordrq_options, + heap_relation.clone(), + index_relation.clone(), + reporter.clone(), + ) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + algorithm::build::build::, Dot>, Heap, _>( + vector_options, + vchordrq_options, + heap_relation.clone(), + index_relation.clone(), + reporter.clone(), + ) + } } if let Some(leader) = unsafe { VchordrqLeader::enter(heap, index, (*index_info).ii_Concurrent) } { @@ -290,35 +313,61 @@ pub unsafe extern "C" fn ambuild( let mut indtuples = 0; reporter.tuples_done(indtuples); let relation = unsafe { PostgresRelation::new(index) }; - match opfamily.vector_kind() { - VectorKind::Vecf32 => { - HeapRelation::>::traverse( + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + HeapRelation::, L2>>::traverse( + &heap_relation, + true, + |(pointer, vector)| { + algorithm::insert::insert::, L2>>( + relation.clone(), + pointer, + vector, + ); + indtuples += 1; + reporter.tuples_done(indtuples); + }, + ); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + HeapRelation::, Dot>>::traverse( &heap_relation, true, |(pointer, vector)| { - algorithm::insert::insert::>( + algorithm::insert::insert::, Dot>>( relation.clone(), pointer, vector, - opfamily.distance_kind(), - true, ); indtuples += 1; reporter.tuples_done(indtuples); }, ); } - VectorKind::Vecf16 => { - HeapRelation::>::traverse( + (VectorKind::Vecf16, DistanceKind::L2) => { + HeapRelation::, L2>>::traverse( &heap_relation, true, |(pointer, vector)| { - algorithm::insert::insert::>( + algorithm::insert::insert::, L2>>( + relation.clone(), + pointer, + vector, + ); + indtuples += 1; + reporter.tuples_done(indtuples); + }, + ); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + HeapRelation::, Dot>>::traverse( + &heap_relation, + true, + |(pointer, vector)| { + algorithm::insert::insert::, Dot>>( relation.clone(), pointer, vector, - opfamily.distance_kind(), - true, ); indtuples += 1; reporter.tuples_done(indtuples); @@ -327,6 +376,28 @@ pub unsafe extern "C" fn ambuild( } } } + let relation = unsafe { PostgresRelation::new(index) }; + let delay = || { + pgrx::check_for_interrupts!(); + }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + type O = Op, L2>; + algorithm::vacuum::maintain::(relation, delay); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + type O = Op, Dot>; + algorithm::vacuum::maintain::(relation, delay); + } + (VectorKind::Vecf16, DistanceKind::L2) => { + type O = Op, L2>; + algorithm::vacuum::maintain::(relation, delay); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + type O = Op, Dot>; + algorithm::vacuum::maintain::(relation, delay); + } + } unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -572,17 +643,17 @@ unsafe fn parallel_build( opfamily: Opfamily, scan: *mut pgrx::pg_sys::TableScanDescData, } - impl HeapRelation for Heap { + impl HeapRelation for Heap { fn traverse(&self, progress: bool, callback: F) where - F: FnMut((NonZeroU64, V)), + F: FnMut((NonZeroU64, O::Vector)), { pub struct State<'a, F> { pub this: &'a Heap, pub callback: F, } #[pgrx::pg_guard] - unsafe extern "C" fn call( + unsafe extern "C" fn call( _index: pgrx::pg_sys::Relation, ctid: pgrx::pg_sys::ItemPointer, values: *mut Datum, @@ -590,14 +661,14 @@ unsafe fn parallel_build( _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((NonZeroU64, V)), + F: FnMut((NonZeroU64, O::Vector)), { let state = unsafe { &mut *state.cast::>() }; let opfamily = state.this.opfamily; let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; let pointer = unsafe { ctid_to_pointer(ctid.read()) }; if let Some(vector) = vector { - (state.callback)((pointer, V::from_owned(vector))); + (state.callback)((pointer, O::Vector::from_owned(vector))); } } let table_am = unsafe { &*(*self.heap).rd_tableam }; @@ -615,7 +686,7 @@ unsafe fn parallel_build( progress, 0, pgrx::pg_sys::InvalidBlockNumber, - Some(call::), + Some(call::), (&mut state) as *mut State as *mut _, self.scan, ); @@ -638,52 +709,106 @@ unsafe fn parallel_build( opfamily, scan, }; - match opfamily.vector_kind() { - VectorKind::Vecf32 => { - HeapRelation::>::traverse(&heap_relation, true, |(pointer, vector)| { - algorithm::insert::insert::>( - index_relation.clone(), - pointer, - vector, - opfamily.distance_kind(), - true, - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + HeapRelation::, L2>>::traverse( + &heap_relation, + true, + |(pointer, vector)| { + algorithm::insert::insert::, L2>>( + index_relation.clone(), + pointer, + vector, + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); + }, + ); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + HeapRelation::, Dot>>::traverse( + &heap_relation, + true, + |(pointer, vector)| { + algorithm::insert::insert::, Dot>>( + index_relation.clone(), + pointer, + vector, + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } } - } - }); + }, + ); } - VectorKind::Vecf16 => { - HeapRelation::>::traverse(&heap_relation, true, |(pointer, vector)| { - algorithm::insert::insert::>( - index_relation.clone(), - pointer, - vector, - opfamily.distance_kind(), - true, - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + (VectorKind::Vecf16, DistanceKind::L2) => { + HeapRelation::, L2>>::traverse( + &heap_relation, + true, + |(pointer, vector)| { + algorithm::insert::insert::, L2>>( + index_relation.clone(), + pointer, + vector, + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); + }, + ); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + HeapRelation::, Dot>>::traverse( + &heap_relation, + true, + |(pointer, vector)| { + algorithm::insert::insert::, Dot>>( + index_relation.clone(), + pointer, + vector, + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } } - } - }); + }, + ); } } unsafe { @@ -714,21 +839,35 @@ pub unsafe extern "C" fn aminsert( let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - match opfamily.vector_kind() { - VectorKind::Vecf32 => algorithm::insert::insert::>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - opfamily.distance_kind(), - false, - ), - VectorKind::Vecf16 => algorithm::insert::insert::>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - opfamily.distance_kind(), - false, - ), + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + algorithm::insert::insert::, L2>>( + unsafe { PostgresRelation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + ) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + algorithm::insert::insert::, Dot>>( + unsafe { PostgresRelation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + ) + } + (VectorKind::Vecf16, DistanceKind::L2) => { + algorithm::insert::insert::, L2>>( + unsafe { PostgresRelation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + ) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + algorithm::insert::insert::, Dot>>( + unsafe { PostgresRelation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + ) + } } } false @@ -750,21 +889,35 @@ pub unsafe extern "C" fn aminsert( let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; if let Some(vector) = vector { let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - match opfamily.vector_kind() { - VectorKind::Vecf32 => algorithm::insert::insert::>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - opfamily.distance_kind(), - false, - ), - VectorKind::Vecf16 => algorithm::insert::insert::>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - opfamily.distance_kind(), - false, - ), + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + algorithm::insert::insert::, L2>>( + unsafe { PostgresRelation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + ) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + algorithm::insert::insert::, Dot>>( + unsafe { PostgresRelation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + ) + } + (VectorKind::Vecf16, DistanceKind::L2) => { + algorithm::insert::insert::, L2>>( + unsafe { PostgresRelation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + ) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + algorithm::insert::insert::, Dot>>( + unsafe { PostgresRelation::new(index) }, + pointer, + VectOwned::::from_owned(vector), + ) + } } } false @@ -893,29 +1046,58 @@ pub unsafe extern "C" fn ambulkdelete( let opfamily = unsafe { am_options::opfamily((*info).index) }; let callback = callback.unwrap(); let callback = |p: NonZeroU64| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; - match opfamily.vector_kind() { - VectorKind::Vecf32 => algorithm::vacuum::vacuum::>( - unsafe { PostgresRelation::new((*info).index) }, - || unsafe { - pgrx::pg_sys::vacuum_delay_point(); - }, - callback, - ), - VectorKind::Vecf16 => algorithm::vacuum::vacuum::>( - unsafe { PostgresRelation::new((*info).index) }, - || unsafe { - pgrx::pg_sys::vacuum_delay_point(); - }, - callback, - ), + let index = unsafe { PostgresRelation::new((*info).index) }; + let delay = || unsafe { + pgrx::pg_sys::vacuum_delay_point(); + }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + type O = Op, L2>; + algorithm::vacuum::bulkdelete::(index.clone(), delay, callback); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + type O = Op, Dot>; + algorithm::vacuum::bulkdelete::(index.clone(), delay, callback); + } + (VectorKind::Vecf16, DistanceKind::L2) => { + type O = Op, L2>; + algorithm::vacuum::bulkdelete::(index.clone(), delay, callback); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + type O = Op, Dot>; + algorithm::vacuum::bulkdelete::(index.clone(), delay, callback); + } } stats } #[pgrx::pg_guard] pub unsafe extern "C" fn amvacuumcleanup( - _info: *mut pgrx::pg_sys::IndexVacuumInfo, + info: *mut pgrx::pg_sys::IndexVacuumInfo, _stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, ) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + let opfamily = unsafe { am_options::opfamily((*info).index) }; + let index = unsafe { PostgresRelation::new((*info).index) }; + let delay = || unsafe { + pgrx::pg_sys::vacuum_delay_point(); + }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + type O = Op, L2>; + algorithm::vacuum::maintain::(index, delay); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + type O = Op, Dot>; + algorithm::vacuum::maintain::(index, delay); + } + (VectorKind::Vecf16, DistanceKind::L2) => { + type O = Op, L2>; + algorithm::vacuum::maintain::(index, delay); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + type O = Op, Dot>; + algorithm::vacuum::maintain::(index, delay); + } + } std::ptr::null_mut() } diff --git a/src/vchordrq/index/am_options.rs b/src/index/am_options.rs similarity index 97% rename from src/vchordrq/index/am_options.rs rename to src/index/am_options.rs index 5c730ed7..34c76a9e 100644 --- a/src/vchordrq/index/am_options.rs +++ b/src/index/am_options.rs @@ -3,9 +3,9 @@ use crate::datatype::memory_pgvector_halfvec::PgvectorHalfvecOutput; use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; use crate::datatype::typmod::Typmod; -use crate::vchordrq::types::{BorrowedVector, OwnedVector}; -use crate::vchordrq::types::{DistanceKind, VectorKind}; -use crate::vchordrq::types::{VchordrqIndexingOptions, VectorOptions}; +use crate::types::{BorrowedVector, OwnedVector}; +use crate::types::{DistanceKind, VectorKind}; +use crate::types::{VchordrqIndexingOptions, VectorOptions}; use distance::Distance; use pgrx::datum::FromDatum; use pgrx::heap_tuple::PgHeapTuple; diff --git a/src/vchordrq/index/am_scan.rs b/src/index/am_scan.rs similarity index 64% rename from src/vchordrq/index/am_scan.rs rename to src/index/am_scan.rs index 6e2d30d5..83e62f33 100644 --- a/src/vchordrq/index/am_scan.rs +++ b/src/index/am_scan.rs @@ -1,12 +1,14 @@ use super::am_options::Opfamily; +use crate::algorithm::operator::Vector; +use crate::algorithm::operator::{Dot, L2, Op}; +use crate::algorithm::scan::scan; +use crate::gucs::executing::epsilon; +use crate::gucs::executing::max_scan_tuples; +use crate::gucs::executing::probes; use crate::postgres::PostgresRelation; -use crate::vchordrq::algorithm::scan::scan; -use crate::vchordrq::algorithm::tuples::Vector; -use crate::vchordrq::gucs::executing::epsilon; -use crate::vchordrq::gucs::executing::max_scan_tuples; -use crate::vchordrq::gucs::executing::probes; -use crate::vchordrq::types::OwnedVector; -use crate::vchordrq::types::VectorKind; +use crate::types::DistanceKind; +use crate::types::OwnedVector; +use crate::types::VectorKind; use distance::Distance; use half::f16; use std::num::NonZeroU64; @@ -74,12 +76,11 @@ pub fn scan_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(N } = scanner { if let Some((vector, opfamily)) = vector.as_ref() { - match opfamily.vector_kind() { - VectorKind::Vecf32 => { - let vbase = scan::>( + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + let vbase = scan::, L2>>( relation, VectOwned::::from_owned(vector.clone()), - opfamily.distance_kind(), probes(), epsilon(), ); @@ -94,11 +95,46 @@ pub fn scan_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(N opfamily: *opfamily, }; } - VectorKind::Vecf16 => { - let vbase = scan::>( + (VectorKind::Vecf32, DistanceKind::Dot) => { + let vbase = scan::, Dot>>( + relation, + VectOwned::::from_owned(vector.clone()), + probes(), + epsilon(), + ); + *scanner = Scanner::Vbase { + vbase: if let Some(max_scan_tuples) = max_scan_tuples() { + Box::new(vbase.take(max_scan_tuples as usize)) + } else { + Box::new(vbase) + }, + threshold: *threshold, + recheck: *recheck, + opfamily: *opfamily, + }; + } + (VectorKind::Vecf16, DistanceKind::L2) => { + let vbase = scan::, L2>>( + relation, + VectOwned::::from_owned(vector.clone()), + probes(), + epsilon(), + ); + *scanner = Scanner::Vbase { + vbase: if let Some(max_scan_tuples) = max_scan_tuples() { + Box::new(vbase.take(max_scan_tuples as usize)) + } else { + Box::new(vbase) + }, + threshold: *threshold, + recheck: *recheck, + opfamily: *opfamily, + }; + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + let vbase = scan::, Dot>>( relation, VectOwned::::from_owned(vector.clone()), - opfamily.distance_kind(), probes(), epsilon(), ); diff --git a/src/vchordrq/index/functions.rs b/src/index/functions.rs similarity index 58% rename from src/vchordrq/index/functions.rs rename to src/index/functions.rs index 32e6d03d..1f3b4e28 100644 --- a/src/vchordrq/index/functions.rs +++ b/src/index/functions.rs @@ -1,7 +1,9 @@ use super::am_options; +use crate::algorithm::operator::{Dot, L2, Op}; +use crate::algorithm::prewarm::prewarm; use crate::postgres::PostgresRelation; -use crate::vchordrq::algorithm::prewarm::prewarm; -use crate::vchordrq::types::VectorKind; +use crate::types::DistanceKind; +use crate::types::VectorKind; use half::f16; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass}; @@ -23,9 +25,19 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; let relation = unsafe { PostgresRelation::new(index) }; let opfamily = unsafe { am_options::opfamily(index) }; - let message = match opfamily.vector_kind() { - VectorKind::Vecf32 => prewarm::>(relation, height), - VectorKind::Vecf16 => prewarm::>(relation, height), + let message = match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + prewarm::, L2>>(relation, height) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + prewarm::, Dot>>(relation, height) + } + (VectorKind::Vecf16, DistanceKind::L2) => { + prewarm::, L2>>(relation, height) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + prewarm::, Dot>>(relation, height) + } }; unsafe { pgrx::pg_sys::index_close(index, pgrx::pg_sys::ShareLock as _); diff --git a/src/vchordrq/index/mod.rs b/src/index/mod.rs similarity index 100% rename from src/vchordrq/index/mod.rs rename to src/index/mod.rs diff --git a/src/vchordrq/index/opclass.rs b/src/index/opclass.rs similarity index 100% rename from src/vchordrq/index/opclass.rs rename to src/index/opclass.rs diff --git a/src/index/utils.rs b/src/index/utils.rs new file mode 100644 index 00000000..18234ac0 --- /dev/null +++ b/src/index/utils.rs @@ -0,0 +1,34 @@ +use std::num::NonZeroU64; + +pub const fn pointer_to_ctid(pointer: NonZeroU64) -> pgrx::pg_sys::ItemPointerData { + let value = pointer.get(); + pgrx::pg_sys::ItemPointerData { + ip_blkid: pgrx::pg_sys::BlockIdData { + bi_hi: ((value >> 32) & 0xffff) as u16, + bi_lo: ((value >> 16) & 0xffff) as u16, + }, + ip_posid: (value & 0xffff) as u16, + } +} + +pub const fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> NonZeroU64 { + let mut value = 0; + value |= (ctid.ip_blkid.bi_hi as u64) << 32; + value |= (ctid.ip_blkid.bi_lo as u64) << 16; + value |= ctid.ip_posid as u64; + NonZeroU64::new(value).expect("invalid pointer") +} + +#[allow(dead_code)] +const fn soundness_check(a: pgrx::pg_sys::ItemPointerData) { + let b = ctid_to_pointer(a); + let c = pointer_to_ctid(b); + assert!(a.ip_blkid.bi_hi == c.ip_blkid.bi_hi); + assert!(a.ip_blkid.bi_lo == c.ip_blkid.bi_lo); + assert!(a.ip_posid == c.ip_posid); +} + +const _: () = soundness_check(pgrx::pg_sys::ItemPointerData { + ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi: 1, bi_lo: 2 }, + ip_posid: 3, +}); diff --git a/src/lib.rs b/src/lib.rs index 018388d8..187e3cda 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,15 +1,18 @@ +#![feature(vec_pop_if)] #![allow(clippy::collapsible_else_if)] #![allow(clippy::infallible_destructuring_match)] #![allow(clippy::too_many_arguments)] #![allow(clippy::type_complexity)] +mod algorithm; mod datatype; +mod gucs; +mod index; mod postgres; mod projection; +mod types; mod upgrade; mod utils; -mod vchordrq; -mod vchordrqfscan; pgrx::pg_module_magic!(); pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap); @@ -21,8 +24,8 @@ unsafe extern "C" fn _PG_init() { pgrx::error!("vchord must be loaded via shared_preload_libraries."); } unsafe { - vchordrq::init(); - vchordrqfscan::init(); + index::init(); + gucs::init(); #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); diff --git a/src/postgres.rs b/src/postgres.rs index 6ff91e23..f68d0fa7 100644 --- a/src/postgres.rs +++ b/src/postgres.rs @@ -1,4 +1,4 @@ -use algorithm::{Opaque, Page, PageGuard, RelationRead, RelationWrite}; +use crate::algorithm::{Opaque, Page, PageGuard, RelationRead, RelationWrite}; use std::mem::{MaybeUninit, offset_of}; use std::ops::{Deref, DerefMut}; use std::ptr::NonNull; @@ -43,12 +43,6 @@ impl PostgresPage { this } #[allow(dead_code)] - unsafe fn assume_init_mut(this: &mut MaybeUninit) -> &mut Self { - let this = unsafe { MaybeUninit::assume_init_mut(this) }; - assert_eq!(offset_of!(Self, opaque), this.header.pd_special as usize); - this - } - #[allow(dead_code)] fn clone_into_boxed(&self) -> Box { let mut result = Box::new_uninit(); unsafe { @@ -56,6 +50,23 @@ impl PostgresPage { result.assume_init() } } + #[allow(dead_code)] + fn reconstruct(&mut self, removes: &[u16]) { + let mut removes = removes.to_vec(); + removes.sort(); + removes.dedup(); + let n = removes.len(); + if n > 0 { + assert!(removes[n - 1] <= self.len()); + unsafe { + pgrx::pg_sys::PageIndexMultiDelete( + (self as *mut Self).cast(), + removes.as_ptr().cast_mut(), + removes.len() as _, + ); + } + } + } } impl Page for PostgresPage { @@ -142,25 +153,23 @@ impl Page for PostgresPage { pgrx::pg_sys::PageIndexTupleDeleteNoCompact((self as *mut Self).cast(), i); } } - fn reconstruct(&mut self, removes: &[u16]) { - let mut removes = removes.to_vec(); - removes.sort(); - removes.dedup(); - let n = removes.len(); - if n > 0 { - assert!(removes[n - 1] <= self.len()); - unsafe { - pgrx::pg_sys::PageIndexMultiDelete( - (self as *mut Self).cast(), - removes.as_ptr().cast_mut(), - removes.len() as _, - ); - } - } - } fn freespace(&self) -> u16 { unsafe { pgrx::pg_sys::PageGetFreeSpace((self as *const Self).cast_mut().cast()) as u16 } } + fn clear(&mut self) { + unsafe { + pgrx::pg_sys::PageInit( + (self as *mut PostgresPage as pgrx::pg_sys::Page).cast(), + pgrx::pg_sys::BLCKSZ as usize, + size_of::(), + ); + (&raw mut self.opaque).write(Opaque { + next: u32::MAX, + skip: u32::MAX, + }); + } + assert_eq!(offset_of!(Self, opaque), self.header.pd_special as usize); + } } const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index c00aab1a..7bc36b62 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -136,18 +136,10 @@ IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhan CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; -CREATE FUNCTION vchordrqfscan_amhandler(internal) RETURNS index_am_handler -IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrqfscan_amhandler_wrapper'; - -CREATE FUNCTION vchordrqfscan_prewarm(regclass, integer default 0) RETURNS TEXT -STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrqfscan_prewarm_wrapper'; - -- List of access methods CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; -CREATE ACCESS METHOD Vchordrqfscan TYPE INDEX HANDLER Vchordrqfscan_amhandler; - -- List of operator families CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; @@ -157,10 +149,6 @@ CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; -CREATE OPERATOR FAMILY vector_l2_ops USING Vchordrqfscan; -CREATE OPERATOR FAMILY vector_ip_ops USING Vchordrqfscan; -CREATE OPERATOR FAMILY vector_cosine_ops USING Vchordrqfscan; - -- List of operator classes CREATE OPERATOR CLASS vector_l2_ops @@ -198,21 +186,3 @@ CREATE OPERATOR CLASS halfvec_cosine_ops OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); - -CREATE OPERATOR CLASS vector_l2_ops - FOR TYPE vector USING Vchordrqfscan FAMILY vector_l2_ops AS - OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, - OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, - FUNCTION 1 _Vchordrqfscan_support_vector_l2_ops(); - -CREATE OPERATOR CLASS vector_ip_ops - FOR TYPE vector USING Vchordrqfscan FAMILY vector_ip_ops AS - OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, - OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, - FUNCTION 1 _Vchordrqfscan_support_vector_ip_ops(); - -CREATE OPERATOR CLASS vector_cosine_ops - FOR TYPE vector USING Vchordrqfscan FAMILY vector_cosine_ops AS - OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, - OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, - FUNCTION 1 _Vchordrqfscan_support_vector_cosine_ops(); diff --git a/src/vchordrq/types.rs b/src/types.rs similarity index 100% rename from src/vchordrq/types.rs rename to src/types.rs diff --git a/src/utils/k_means.rs b/src/utils/k_means.rs index 7b44a24c..b1808c90 100644 --- a/src/utils/k_means.rs +++ b/src/utils/k_means.rs @@ -3,6 +3,7 @@ use half::f16; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use simd::Floating; +use simd::fast_scan::{any_pack, padding_pack}; pub fn k_means( parallelism: &P, @@ -81,39 +82,64 @@ fn rabitq_index( samples: &[Vec], centroids: &[Vec], ) -> Vec { - let mut a0 = Vec::new(); - let mut a1 = Vec::new(); - let mut a2 = Vec::new(); - let mut a3 = Vec::new(); - let mut a4 = Vec::new(); - for vectors in centroids.chunks(32) { - use simd::fast_scan::pack; - let x = std::array::from_fn::<_, 32, _>(|i| { - if let Some(vector) = vectors.get(i) { - rabitq::block::code(dims as _, vector) - } else { - rabitq::block::dummy_code(dims as _) - } + struct Branch { + dis_u_2: f32, + factor_ppc: f32, + factor_ip: f32, + factor_err: f32, + signs: Vec, + } + let branches = { + let mut branches = Vec::new(); + for centroid in centroids { + let code = rabitq::code(dims as _, centroid); + branches.push(Branch { + dis_u_2: code.dis_u_2, + factor_ppc: code.factor_ppc, + factor_ip: code.factor_ip, + factor_err: code.factor_err, + signs: code.signs, + }); + } + branches + }; + struct Block { + dis_u_2: [f32; 32], + factor_ppc: [f32; 32], + factor_ip: [f32; 32], + factor_err: [f32; 32], + elements: Vec<[u64; 2]>, + } + impl Block { + fn code(&self) -> (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], &[[u64; 2]]) { + ( + &self.dis_u_2, + &self.factor_ppc, + &self.factor_ip, + &self.factor_err, + &self.elements, + ) + } + } + let mut blocks = Vec::new(); + for chunk in branches.chunks(32) { + blocks.push(Block { + dis_u_2: any_pack(chunk.iter().map(|x| x.dis_u_2)), + factor_ppc: any_pack(chunk.iter().map(|x| x.factor_ppc)), + factor_ip: any_pack(chunk.iter().map(|x| x.factor_ip)), + factor_err: any_pack(chunk.iter().map(|x| x.factor_err)), + elements: padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))), }); - a0.push(x.each_ref().map(|x| x.dis_u_2)); - a1.push(x.each_ref().map(|x| x.factor_ppc)); - a2.push(x.each_ref().map(|x| x.factor_ip)); - a3.push(x.each_ref().map(|x| x.factor_err)); - a4.push(pack(dims.div_ceil(4) as _, x.map(|x| x.signs)).collect::>()); } parallelism .rayon_into_par_iter(0..n) .map(|i| { use distance::Distance; - let lut = rabitq::block::fscan_preprocess(&samples[i]); + let lut = rabitq::block::preprocess(&samples[i]); let mut result = (Distance::INFINITY, 0); for block in 0..c.div_ceil(32) { - let lowerbound = rabitq::block::fscan_process_lowerbound_l2( - dims as _, - &lut, - (&a0[block], &a1[block], &a2[block], &a3[block], &a4[block]), - 1.9, - ); + let lowerbound = + rabitq::block::process_lowerbound_l2(&lut, blocks[block].code(), 1.9); for j in block * 32..std::cmp::min(block * 32 + 32, c) { if lowerbound[j - block * 32] < result.0 { let dis = diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 1b07dc6c..85a84e0e 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,2 +1,3 @@ pub mod k_means; pub mod parallelism; +pub mod pipe; diff --git a/src/utils/pipe.rs b/src/utils/pipe.rs new file mode 100644 index 00000000..18cfc376 --- /dev/null +++ b/src/utils/pipe.rs @@ -0,0 +1,14 @@ +pub trait Pipe { + fn pipe(self, f: impl FnOnce(Self) -> T) -> T + where + Self: Sized; +} + +impl Pipe for S { + fn pipe(self, f: impl FnOnce(Self) -> T) -> T + where + Self: Sized, + { + f(self) + } +} diff --git a/src/vchordrq/algorithm/insert.rs b/src/vchordrq/algorithm/insert.rs deleted file mode 100644 index f625dcac..00000000 --- a/src/vchordrq/algorithm/insert.rs +++ /dev/null @@ -1,221 +0,0 @@ -use crate::vchordrq::algorithm::rabitq; -use crate::vchordrq::algorithm::tuples::*; -use crate::vchordrq::algorithm::vectors; -use crate::vchordrq::types::DistanceKind; -use algorithm::{Page, PageGuard, RelationWrite}; -use always_equal::AlwaysEqual; -use distance::Distance; -use std::cmp::Reverse; -use std::collections::BinaryHeap; -use std::num::NonZeroU64; -use vector::VectorBorrowed; - -pub fn insert( - relation: impl RelationWrite + Clone, - payload: NonZeroU64, - vector: V, - distance_kind: DistanceKind, - in_building: bool, -) { - let vector = vector.as_borrowed(); - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dims = meta_tuple.dims; - assert_eq!(dims, vector.dims(), "invalid vector dimensions"); - let vector = V::random_projection(vector); - let vector = vector.as_borrowed(); - let is_residual = meta_tuple.is_residual; - let default_lut = if !is_residual { - Some(V::rabitq_fscan_preprocess(vector)) - } else { - None - }; - let h0_vector = { - let (metadata, slices) = V::vector_split(vector); - let mut chain = Err(metadata); - for i in (0..slices.len()).rev() { - let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple:: { - slice: slices[i].to_vec(), - payload: Some(payload), - chain, - }) - .unwrap(); - chain = Ok(append( - relation.clone(), - meta_tuple.vectors_first, - &tuple, - true, - true, - true, - )); - } - chain.ok().unwrap() - }; - let h0_payload = payload; - let mut list = { - let Some((_, original)) = vectors::vector_dist::( - relation.clone(), - vector, - meta_tuple.mean, - None, - None, - is_residual, - ) else { - panic!("data corruption") - }; - (meta_tuple.first, original) - }; - let make_list = |list: (u32, Option)| { - let mut results = Vec::new(); - { - let lut = if is_residual { - &V::rabitq_fscan_preprocess( - V::residual(vector, list.1.as_ref().map(|x| x.as_borrowed()).unwrap()) - .as_borrowed(), - ) - } else { - default_lut.as_ref().unwrap() - }; - let mut current = list.0; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let lowerbounds = rabitq::process_lowerbound( - distance_kind, - dims, - lut, - ( - h1_tuple.dis_u_2, - h1_tuple.factor_ppc, - h1_tuple.factor_ip, - h1_tuple.factor_err, - &h1_tuple.t, - ), - 1.9, - ); - results.push(( - Reverse(lowerbounds), - AlwaysEqual(h1_tuple.mean), - AlwaysEqual(h1_tuple.first), - )); - } - current = h1_guard.get_opaque().next; - } - } - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let Some((Some(dis_u), original)) = vectors::vector_dist::( - relation.clone(), - vector, - mean, - None, - Some(distance_kind), - is_residual, - ) else { - panic!("data corruption") - }; - cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(original))); - } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop().unwrap(); - (first, mean) - } - }; - for _ in (1..meta_tuple.height_of_root).rev() { - list = make_list(list); - } - let code = if is_residual { - V::rabitq_code( - dims, - V::residual(vector, list.1.as_ref().map(|x| x.as_borrowed()).unwrap()).as_borrowed(), - ) - } else { - V::rabitq_code(dims, vector) - }; - let tuple = rkyv::to_bytes::<_, 8192>(&Height0Tuple { - mean: h0_vector, - payload: h0_payload, - dis_u_2: code.dis_u_2, - factor_ppc: code.factor_ppc, - factor_ip: code.factor_ip, - factor_err: code.factor_err, - t: code.t(), - }) - .unwrap(); - append( - relation.clone(), - list.0, - &tuple, - false, - in_building, - in_building, - ); -} - -fn append( - relation: impl RelationWrite, - first: u32, - tuple: &[u8], - tracking_freespace: bool, - skipping_traversal: bool, - updating_skip: bool, -) -> (u32, u16) { - if tracking_freespace { - if let Some(mut write) = relation.search(tuple.len()) { - let i = write.alloc(tuple).unwrap(); - return (write.id(), i); - } - } - assert!(first != u32::MAX); - let mut current = first; - loop { - let read = relation.read(current); - if read.freespace() as usize >= tuple.len() || read.get_opaque().next == u32::MAX { - drop(read); - let mut write = relation.write(current, tracking_freespace); - if let Some(i) = write.alloc(tuple) { - return (current, i); - } - if write.get_opaque().next == u32::MAX { - let mut extend = relation.extend(tracking_freespace); - write.get_opaque_mut().next = extend.id(); - drop(write); - if let Some(i) = extend.alloc(tuple) { - let result = (extend.id(), i); - drop(extend); - if updating_skip { - let mut past = relation.write(first, tracking_freespace); - let skip = &mut past.get_opaque_mut().skip; - assert!(*skip != u32::MAX); - *skip = std::cmp::max(*skip, result.0); - } - return result; - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - } - if skipping_traversal && current == first && write.get_opaque().skip != first { - current = write.get_opaque().skip; - } else { - current = write.get_opaque().next; - } - } else { - if skipping_traversal && current == first && read.get_opaque().skip != first { - current = read.get_opaque().skip; - } else { - current = read.get_opaque().next; - } - } - } -} diff --git a/src/vchordrq/algorithm/mod.rs b/src/vchordrq/algorithm/mod.rs deleted file mode 100644 index 88239a8e..00000000 --- a/src/vchordrq/algorithm/mod.rs +++ /dev/null @@ -1,8 +0,0 @@ -pub mod build; -pub mod insert; -pub mod prewarm; -pub mod rabitq; -pub mod scan; -pub mod tuples; -pub mod vacuum; -pub mod vectors; diff --git a/src/vchordrq/algorithm/prewarm.rs b/src/vchordrq/algorithm/prewarm.rs deleted file mode 100644 index 6a7dc252..00000000 --- a/src/vchordrq/algorithm/prewarm.rs +++ /dev/null @@ -1,82 +0,0 @@ -use crate::vchordrq::algorithm::tuples::*; -use crate::vchordrq::algorithm::vectors; -use algorithm::{Page, RelationRead}; -use std::fmt::Write; - -pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> String { - let mut message = String::new(); - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - writeln!(message, "height of root: {}", meta_tuple.height_of_root).unwrap(); - let prewarm_max_height = if height < 0 { 0 } else { height as u32 }; - if prewarm_max_height > meta_tuple.height_of_root { - return message; - } - let mut lists = { - let mut results = Vec::new(); - let counter = 1_usize; - { - vectors::vector_warm::(relation.clone(), meta_tuple.mean); - results.push(meta_tuple.first); - } - writeln!(message, "number of tuples: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); - results - }; - let mut make_lists = |lists| { - let mut counter = 0_usize; - let mut results = Vec::new(); - for list in lists { - let mut current = list; - while current != u32::MAX { - counter += 1; - pgrx::check_for_interrupts!(); - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - vectors::vector_warm::(relation.clone(), h1_tuple.mean); - results.push(h1_tuple.first); - } - current = h1_guard.get_opaque().next; - } - } - writeln!(message, "number of tuples: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); - results - }; - for _ in (std::cmp::max(1, prewarm_max_height)..meta_tuple.height_of_root).rev() { - lists = make_lists(lists); - } - if prewarm_max_height == 0 { - let mut counter = 0_usize; - let mut results = Vec::new(); - for list in lists { - let mut current = list; - while current != u32::MAX { - counter += 1; - pgrx::check_for_interrupts!(); - let h0_guard = relation.read(current); - for i in 1..=h0_guard.len() { - let _h0_tuple = h0_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - results.push(()); - } - current = h0_guard.get_opaque().next; - } - } - writeln!(message, "number of tuples: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); - } - message -} diff --git a/src/vchordrq/algorithm/rabitq.rs b/src/vchordrq/algorithm/rabitq.rs deleted file mode 100644 index 4f406e18..00000000 --- a/src/vchordrq/algorithm/rabitq.rs +++ /dev/null @@ -1,21 +0,0 @@ -use crate::vchordrq::types::DistanceKind; -use distance::Distance; - -pub use rabitq::binary::Code; -pub use rabitq::binary::Lut; -pub use rabitq::binary::code; -pub use rabitq::binary::preprocess; -pub use rabitq::binary::{process_lowerbound_dot, process_lowerbound_l2}; - -pub fn process_lowerbound( - distance_kind: DistanceKind, - dims: u32, - lut: &Lut, - code: (f32, f32, f32, f32, &[u64]), - epsilon: f32, -) -> Distance { - match distance_kind { - DistanceKind::L2 => process_lowerbound_l2(dims, lut, code, epsilon), - DistanceKind::Dot => process_lowerbound_dot(dims, lut, code, epsilon), - } -} diff --git a/src/vchordrq/algorithm/scan.rs b/src/vchordrq/algorithm/scan.rs deleted file mode 100644 index e6915f0a..00000000 --- a/src/vchordrq/algorithm/scan.rs +++ /dev/null @@ -1,189 +0,0 @@ -use crate::vchordrq::algorithm::rabitq; -use crate::vchordrq::algorithm::tuples::*; -use crate::vchordrq::algorithm::vectors; -use crate::vchordrq::types::DistanceKind; -use algorithm::{Page, RelationRead}; -use always_equal::AlwaysEqual; -use distance::Distance; -use std::cmp::Reverse; -use std::collections::BinaryHeap; -use std::num::NonZeroU64; -use vector::VectorBorrowed; - -pub fn scan( - relation: impl RelationRead + Clone, - vector: V, - distance_kind: DistanceKind, - probes: Vec, - epsilon: f32, -) -> impl Iterator { - let vector = vector.as_borrowed(); - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dims = meta_tuple.dims; - let height_of_root = meta_tuple.height_of_root; - assert_eq!(dims, vector.dims(), "invalid vector dimensions"); - assert_eq!(height_of_root as usize, 1 + probes.len(), "invalid probes"); - let vector = V::random_projection(vector); - let is_residual = meta_tuple.is_residual; - let default_lut = if !is_residual { - Some(V::rabitq_fscan_preprocess(vector.as_borrowed())) - } else { - None - }; - let mut lists: Vec<_> = vec![{ - let Some((_, original)) = vectors::vector_dist::( - relation.clone(), - vector.as_borrowed(), - meta_tuple.mean, - None, - None, - is_residual, - ) else { - panic!("data corruption") - }; - (meta_tuple.first, original) - }]; - let make_lists = |lists: Vec<(u32, Option)>, probes| { - let mut results = Vec::new(); - for list in lists { - let lut = if is_residual { - &V::rabitq_fscan_preprocess( - V::residual( - vector.as_borrowed(), - list.1.as_ref().map(|x| x.as_borrowed()).unwrap(), - ) - .as_borrowed(), - ) - } else { - default_lut.as_ref().unwrap() - }; - let mut current = list.0; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let lowerbounds = rabitq::process_lowerbound( - distance_kind, - dims, - lut, - ( - h1_tuple.dis_u_2, - h1_tuple.factor_ppc, - h1_tuple.factor_ip, - h1_tuple.factor_err, - &h1_tuple.t, - ), - epsilon, - ); - results.push(( - Reverse(lowerbounds), - AlwaysEqual(h1_tuple.mean), - AlwaysEqual(h1_tuple.first), - )); - } - current = h1_guard.get_opaque().next; - } - } - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - std::iter::from_fn(|| { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let Some((Some(dis_u), original)) = vectors::vector_dist::( - relation.clone(), - vector.as_borrowed(), - mean, - None, - Some(distance_kind), - is_residual, - ) else { - panic!("data corruption") - }; - cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(original))); - } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; - Some((first, mean)) - }) - .take(probes as usize) - .collect() - }; - for i in (1..meta_tuple.height_of_root).rev() { - lists = make_lists(lists, probes[i as usize - 1]); - } - drop(meta_guard); - { - let mut results = Vec::new(); - for list in lists { - let lut = if is_residual { - &V::rabitq_fscan_preprocess( - V::residual( - vector.as_borrowed(), - list.1.as_ref().map(|x| x.as_borrowed()).unwrap(), - ) - .as_borrowed(), - ) - } else { - default_lut.as_ref().unwrap() - }; - let mut current = list.0; - while current != u32::MAX { - let h0_guard = relation.read(current); - for i in 1..=h0_guard.len() { - let h0_tuple = h0_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let lowerbounds = rabitq::process_lowerbound( - distance_kind, - dims, - lut, - ( - h0_tuple.dis_u_2, - h0_tuple.factor_ppc, - h0_tuple.factor_ip, - h0_tuple.factor_err, - &h0_tuple.t, - ), - epsilon, - ); - results.push(( - Reverse(lowerbounds), - AlwaysEqual(h0_tuple.mean), - AlwaysEqual(h0_tuple.payload), - )); - } - current = h0_guard.get_opaque().next; - } - } - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _)>::new(); - std::iter::from_fn(move || { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); - let Some((Some(dis_u), _)) = vectors::vector_dist::( - relation.clone(), - vector.as_borrowed(), - mean, - Some(pay_u), - Some(distance_kind), - false, - ) else { - continue; - }; - cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); - } - let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; - Some((dis_u, pay_u)) - }) - } -} diff --git a/src/vchordrq/algorithm/tuples.rs b/src/vchordrq/algorithm/tuples.rs deleted file mode 100644 index 23fe664e..00000000 --- a/src/vchordrq/algorithm/tuples.rs +++ /dev/null @@ -1,270 +0,0 @@ -use std::num::NonZeroU64; - -use super::rabitq::{self, Code, Lut}; -use crate::vchordrq::types::DistanceKind; -use crate::vchordrq::types::OwnedVector; -use half::f16; -use rkyv::{Archive, ArchiveUnsized, CheckBytes, Deserialize, Serialize}; -use simd::Floating; -use vector::VectorOwned; -use vector::vect::VectOwned; - -pub trait Vector: VectorOwned { - type Metadata: Copy - + Serialize< - rkyv::ser::serializers::CompositeSerializer< - rkyv::ser::serializers::AlignedSerializer, - rkyv::ser::serializers::FallbackScratch< - rkyv::ser::serializers::HeapScratch<8192>, - rkyv::ser::serializers::AllocScratch, - >, - rkyv::ser::serializers::SharedSerializeMap, - >, - > + for<'a> CheckBytes>; - type Element: Copy - + Serialize< - rkyv::ser::serializers::CompositeSerializer< - rkyv::ser::serializers::AlignedSerializer, - rkyv::ser::serializers::FallbackScratch< - rkyv::ser::serializers::HeapScratch<8192>, - rkyv::ser::serializers::AllocScratch, - >, - rkyv::ser::serializers::SharedSerializeMap, - >, - > + for<'a> CheckBytes> - + Archive; - - fn metadata_from_archived( - archived: &::Archived, - ) -> Self::Metadata; - - fn vector_split(vector: Self::Borrowed<'_>) -> (Self::Metadata, Vec<&[Self::Element]>); - fn vector_merge(metadata: Self::Metadata, slice: &[Self::Element]) -> Self; - fn from_owned(vector: OwnedVector) -> Self; - - type DistanceAccumulator; - fn distance_begin(distance_kind: DistanceKind) -> Self::DistanceAccumulator; - fn distance_next( - accumulator: &mut Self::DistanceAccumulator, - left: &[Self::Element], - right: &[Self::Element], - ); - fn distance_end( - accumulator: Self::DistanceAccumulator, - left: Self::Metadata, - right: Self::Metadata, - ) -> f32; - - fn random_projection(vector: Self::Borrowed<'_>) -> Self; - - fn residual(vector: Self::Borrowed<'_>, center: Self::Borrowed<'_>) -> Self; - - fn rabitq_fscan_preprocess(vector: Self::Borrowed<'_>) -> Lut; - - fn rabitq_code(dims: u32, vector: Self::Borrowed<'_>) -> Code; - - fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec; - - fn build_from_vecf32(x: &[f32]) -> Self; -} - -impl Vector for VectOwned { - type Metadata = (); - - type Element = f32; - - fn metadata_from_archived(_: &::Archived) -> Self::Metadata {} - - fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f32]>) { - let vector = vector.slice(); - ((), match vector.len() { - 0..=960 => vec![vector], - 961..=1280 => vec![&vector[..640], &vector[640..]], - 1281.. => vector.chunks(1920).collect(), - }) - } - - fn vector_merge((): Self::Metadata, slice: &[Self::Element]) -> Self { - VectOwned::new(slice.to_vec()) - } - - fn from_owned(vector: OwnedVector) -> Self { - match vector { - OwnedVector::Vecf32(x) => x, - _ => unreachable!(), - } - } - - type DistanceAccumulator = (DistanceKind, f32); - fn distance_begin(distance_kind: DistanceKind) -> Self::DistanceAccumulator { - (distance_kind, 0.0) - } - fn distance_next( - accumulator: &mut Self::DistanceAccumulator, - left: &[Self::Element], - right: &[Self::Element], - ) { - match accumulator.0 { - DistanceKind::L2 => accumulator.1 += f32::reduce_sum_of_d2(left, right), - DistanceKind::Dot => accumulator.1 += -f32::reduce_sum_of_xy(left, right), - } - } - fn distance_end( - accumulator: Self::DistanceAccumulator, - (): Self::Metadata, - (): Self::Metadata, - ) -> f32 { - accumulator.1 - } - - fn random_projection(vector: Self::Borrowed<'_>) -> Self { - Self::new(crate::projection::project(vector.slice())) - } - - fn residual(vector: Self::Borrowed<'_>, center: Self::Borrowed<'_>) -> Self { - Self::new(Floating::vector_sub(vector.slice(), center.slice())) - } - - fn rabitq_fscan_preprocess(vector: Self::Borrowed<'_>) -> Lut { - rabitq::preprocess(vector.slice()) - } - - fn rabitq_code(dims: u32, vector: Self::Borrowed<'_>) -> Code { - rabitq::code(dims, vector.slice()) - } - - fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { - vector.slice().to_vec() - } - - fn build_from_vecf32(x: &[f32]) -> Self { - Self::new(x.to_vec()) - } -} - -impl Vector for VectOwned { - type Metadata = (); - - type Element = f16; - - fn metadata_from_archived(_: &::Archived) -> Self::Metadata {} - - fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f16]>) { - let vector = vector.slice(); - ((), match vector.len() { - 0..=1920 => vec![vector], - 1921..=2560 => vec![&vector[..1280], &vector[1280..]], - 2561.. => vector.chunks(3840).collect(), - }) - } - - fn vector_merge((): Self::Metadata, slice: &[Self::Element]) -> Self { - VectOwned::new(slice.to_vec()) - } - - fn from_owned(vector: OwnedVector) -> Self { - match vector { - OwnedVector::Vecf16(x) => x, - _ => unreachable!(), - } - } - - type DistanceAccumulator = (DistanceKind, f32); - fn distance_begin(distance_kind: DistanceKind) -> Self::DistanceAccumulator { - (distance_kind, 0.0) - } - fn distance_next( - accumulator: &mut Self::DistanceAccumulator, - left: &[Self::Element], - right: &[Self::Element], - ) { - match accumulator.0 { - DistanceKind::L2 => accumulator.1 += f16::reduce_sum_of_d2(left, right), - DistanceKind::Dot => accumulator.1 += -f16::reduce_sum_of_xy(left, right), - } - } - fn distance_end( - accumulator: Self::DistanceAccumulator, - (): Self::Metadata, - (): Self::Metadata, - ) -> f32 { - accumulator.1 - } - - fn random_projection(vector: Self::Borrowed<'_>) -> Self { - Self::new(f16::vector_from_f32(&crate::projection::project( - &f16::vector_to_f32(vector.slice()), - ))) - } - - fn residual(vector: Self::Borrowed<'_>, center: Self::Borrowed<'_>) -> Self { - Self::new(Floating::vector_sub(vector.slice(), center.slice())) - } - - fn rabitq_fscan_preprocess(vector: Self::Borrowed<'_>) -> Lut { - rabitq::preprocess(&f16::vector_to_f32(vector.slice())) - } - - fn rabitq_code(dims: u32, vector: Self::Borrowed<'_>) -> Code { - rabitq::code(dims, &f16::vector_to_f32(vector.slice())) - } - - fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { - f16::vector_to_f32(vector.slice()) - } - - fn build_from_vecf32(x: &[f32]) -> Self { - Self::new(f16::vector_from_f32(x)) - } -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct MetaTuple { - pub dims: u32, - pub height_of_root: u32, - pub is_residual: bool, - pub vectors_first: u32, - // raw vector - pub mean: (u32, u16), - // for meta tuple, it's pointers to next level - pub first: u32, -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct VectorTuple { - pub slice: Vec, - pub payload: Option, - pub chain: Result<(u32, u16), V::Metadata>, -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct Height1Tuple { - // raw vector - pub mean: (u32, u16), - // for height 1 tuple, it's pointers to next level - pub first: u32, - // RaBitQ algorithm - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub t: Vec, -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct Height0Tuple { - // raw vector - pub mean: (u32, u16), - // for height 0 tuple, it's pointers to heap relation - pub payload: NonZeroU64, - // RaBitQ algorithm - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub t: Vec, -} diff --git a/src/vchordrq/algorithm/vacuum.rs b/src/vchordrq/algorithm/vacuum.rs deleted file mode 100644 index ee97ca6c..00000000 --- a/src/vchordrq/algorithm/vacuum.rs +++ /dev/null @@ -1,112 +0,0 @@ -use crate::vchordrq::algorithm::tuples::*; -use algorithm::{Page, RelationWrite}; -use std::num::NonZeroU64; - -pub fn vacuum( - relation: impl RelationWrite, - delay: impl Fn(), - callback: impl Fn(NonZeroU64) -> bool, -) { - // step 1: vacuum height_0_tuple - { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let mut firsts = vec![meta_tuple.first]; - let make_firsts = |firsts| { - let mut results = Vec::new(); - for first in firsts { - let mut current = first; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - results.push(h1_tuple.first); - } - current = h1_guard.get_opaque().next; - } - } - results - }; - for _ in (1..meta_tuple.height_of_root).rev() { - firsts = make_firsts(firsts); - } - for first in firsts { - let mut current = first; - while current != u32::MAX { - delay(); - let mut h0_guard = relation.write(current, false); - let mut reconstruct_removes = Vec::new(); - for i in 1..=h0_guard.len() { - let h0_tuple = h0_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - if callback(h0_tuple.payload) { - reconstruct_removes.push(i); - } - } - h0_guard.reconstruct(&reconstruct_removes); - current = h0_guard.get_opaque().next; - } - } - } - // step 2: vacuum vector_tuple - { - let mut current = { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - meta_tuple.vectors_first - }; - while current != u32::MAX { - delay(); - let read = relation.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - let Some(vector_tuple) = read.get(i) else { - continue; - }; - let vector_tuple = - unsafe { rkyv::archived_root::>(vector_tuple) }; - if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(payload) { - break 'flag true; - } - } - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current, true); - for i in 1..=write.len() { - let Some(vector_tuple) = write.get(i) else { - continue; - }; - let vector_tuple = - unsafe { rkyv::archived_root::>(vector_tuple) }; - if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(payload) { - write.free(i); - } - } - } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; - } - } - } -} diff --git a/src/vchordrq/algorithm/vectors.rs b/src/vchordrq/algorithm/vectors.rs deleted file mode 100644 index 72e448a1..00000000 --- a/src/vchordrq/algorithm/vectors.rs +++ /dev/null @@ -1,76 +0,0 @@ -use super::tuples::Vector; -use crate::vchordrq::algorithm::tuples::VectorTuple; -use crate::vchordrq::types::DistanceKind; -use algorithm::{Page, RelationRead}; -use distance::Distance; -use std::num::NonZeroU64; - -pub fn vector_dist( - relation: impl RelationRead, - vector: V::Borrowed<'_>, - mean: (u32, u16), - payload: Option, - for_distance: Option, - for_original: bool, -) -> Option<(Option, Option)> { - if for_distance.is_none() && !for_original && payload.is_none() { - return Some((None, None)); - } - let (left_metadata, slices) = V::vector_split(vector); - let mut cursor = Ok(mean); - let mut result = for_distance.map(|x| V::distance_begin(x)); - let mut original = Vec::new(); - for i in 0..slices.len() { - let Ok(mean) = cursor else { - // fails consistency check - return None; - }; - let vector_guard = relation.read(mean.0); - let Some(vector_tuple) = vector_guard.get(mean.1) else { - // fails consistency check - return None; - }; - let vector_tuple = unsafe { rkyv::archived_root::>(vector_tuple) }; - if vector_tuple.payload != payload { - // fails consistency check - return None; - } - if let Some(result) = result.as_mut() { - V::distance_next(result, slices[i], &vector_tuple.slice); - } - if for_original { - original.extend_from_slice(&vector_tuple.slice); - } - cursor = match &vector_tuple.chain { - rkyv::result::ArchivedResult::Ok(x) => Ok(*x), - rkyv::result::ArchivedResult::Err(x) => Err(V::metadata_from_archived(x)), - }; - } - let Err(right_metadata) = cursor else { - panic!("data corruption") - }; - Some(( - result.map(|r| Distance::from_f32(V::distance_end(r, left_metadata, right_metadata))), - for_original.then(|| V::vector_merge(right_metadata, &original)), - )) -} - -pub fn vector_warm(relation: impl RelationRead, mean: (u32, u16)) { - let mut cursor = Ok(mean); - while let Ok(mean) = cursor { - let vector_guard = relation.read(mean.0); - let Some(vector_tuple) = vector_guard.get(mean.1) else { - // fails consistency check - return; - }; - let vector_tuple = unsafe { rkyv::archived_root::>(vector_tuple) }; - if vector_tuple.payload.is_some() { - // fails consistency check - return; - } - cursor = match &vector_tuple.chain { - rkyv::result::ArchivedResult::Ok(x) => Ok(*x), - rkyv::result::ArchivedResult::Err(x) => Err(V::metadata_from_archived(x)), - }; - } -} diff --git a/src/vchordrq/index/utils.rs b/src/vchordrq/index/utils.rs deleted file mode 100644 index 726a5979..00000000 --- a/src/vchordrq/index/utils.rs +++ /dev/null @@ -1,20 +0,0 @@ -use std::num::NonZeroU64; - -pub fn pointer_to_ctid(pointer: NonZeroU64) -> pgrx::pg_sys::ItemPointerData { - let value = pointer.get(); - pgrx::pg_sys::ItemPointerData { - ip_blkid: pgrx::pg_sys::BlockIdData { - bi_hi: ((value >> 32) & 0xffff) as u16, - bi_lo: ((value >> 16) & 0xffff) as u16, - }, - ip_posid: (value & 0xffff) as u16, - } -} - -pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> NonZeroU64 { - let mut value = 0; - value |= (ctid.ip_blkid.bi_hi as u64) << 32; - value |= (ctid.ip_blkid.bi_lo as u64) << 16; - value |= ctid.ip_posid as u64; - NonZeroU64::new(value).expect("invalid pointer") -} diff --git a/src/vchordrq/mod.rs b/src/vchordrq/mod.rs deleted file mode 100644 index c2ae9456..00000000 --- a/src/vchordrq/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -mod algorithm; -mod gucs; -mod index; -mod types; - -pub unsafe fn init() { - unsafe { - index::init(); - gucs::init(); - } -} diff --git a/src/vchordrqfscan/algorithm/build.rs b/src/vchordrqfscan/algorithm/build.rs deleted file mode 100644 index e46d6a7d..00000000 --- a/src/vchordrqfscan/algorithm/build.rs +++ /dev/null @@ -1,445 +0,0 @@ -use crate::vchordrqfscan::algorithm::rabitq; -use crate::vchordrqfscan::algorithm::tuples::*; -use crate::vchordrqfscan::index::am_options::Opfamily; -use crate::vchordrqfscan::types::DistanceKind; -use crate::vchordrqfscan::types::VchordrqfscanBuildOptions; -use crate::vchordrqfscan::types::VchordrqfscanExternalBuildOptions; -use crate::vchordrqfscan::types::VchordrqfscanIndexingOptions; -use crate::vchordrqfscan::types::VchordrqfscanInternalBuildOptions; -use crate::vchordrqfscan::types::VectorOptions; -use algorithm::{Page, PageGuard, RelationWrite}; -use rand::Rng; -use rkyv::ser::serializers::AllocSerializer; -use simd::Floating; -use std::marker::PhantomData; -use std::num::NonZeroU64; -use std::sync::Arc; - -pub trait HeapRelation { - fn traverse(&self, progress: bool, callback: F) - where - F: FnMut((NonZeroU64, Vec)); - fn opfamily(&self) -> Opfamily; -} - -pub trait Reporter { - fn tuples_total(&mut self, tuples_total: u64); -} - -pub fn build( - vector_options: VectorOptions, - vchordrqfscan_options: VchordrqfscanIndexingOptions, - heap_relation: T, - relation: impl RelationWrite, - mut reporter: R, -) { - let dims = vector_options.dims; - let is_residual = - vchordrqfscan_options.residual_quantization && vector_options.d == DistanceKind::L2; - let structures = match vchordrqfscan_options.build { - VchordrqfscanBuildOptions::External(external_build) => Structure::extern_build( - vector_options.clone(), - heap_relation.opfamily(), - external_build.clone(), - ), - VchordrqfscanBuildOptions::Internal(internal_build) => { - let mut tuples_total = 0_u64; - let samples = { - let mut rand = rand::thread_rng(); - let max_number_of_samples = internal_build - .lists - .last() - .unwrap() - .saturating_mul(internal_build.sampling_factor); - let mut samples = Vec::new(); - let mut number_of_samples = 0_u32; - heap_relation.traverse(false, |(_, vector)| { - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); - if number_of_samples < max_number_of_samples { - samples.push(vector); - number_of_samples += 1; - } else { - let index = rand.gen_range(0..max_number_of_samples) as usize; - samples[index] = vector; - } - tuples_total += 1; - }); - samples - }; - reporter.tuples_total(tuples_total); - Structure::internal_build(vector_options.clone(), internal_build.clone(), samples) - } - }; - let mut meta = Tape::create(&relation, false); - assert_eq!(meta.first(), 0); - let mut vectors = Tape::create(&relation, true); - let mut pointer_of_means = Vec::>::new(); - for i in 0..structures.len() { - let mut level = Vec::new(); - for j in 0..structures[i].len() { - let pointer = vectors.push(&VectorTuple { - payload: None, - vector: structures[i].means[j].clone(), - }); - level.push(pointer); - } - pointer_of_means.push(level); - } - let mut pointer_of_firsts = Vec::>::new(); - for i in 0..structures.len() { - let mut level = Vec::new(); - for j in 0..structures[i].len() { - if i == 0 { - let tape = Tape::::create(&relation, false); - level.push(tape.first()); - } else { - let mut tape = Tape::::create(&relation, false); - let mut cache = Vec::new(); - let h2_mean = &structures[i].means[j]; - let h2_children = &structures[i].children[j]; - for child in h2_children.iter().copied() { - let h1_mean = &structures[i - 1].means[child as usize]; - let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(h1_mean, h2_mean)) - } else { - rabitq::code(dims, h1_mean) - }; - cache.push((child, code)); - if cache.len() == 32 { - let group = std::mem::take(&mut cache); - let codes = std::array::from_fn(|k| group[k].1.clone()); - let packed = rabitq::pack_codes(dims, codes); - tape.push(&Height1Tuple { - mask: [true; 32], - mean: std::array::from_fn(|k| { - pointer_of_means[i - 1][group[k].0 as usize] - }), - first: std::array::from_fn(|k| { - pointer_of_firsts[i - 1][group[k].0 as usize] - }), - dis_u_2: packed.dis_u_2, - factor_ppc: packed.factor_ppc, - factor_ip: packed.factor_ip, - factor_err: packed.factor_err, - t: packed.t, - }); - } - } - if !cache.is_empty() { - let group = std::mem::take(&mut cache); - let codes = std::array::from_fn(|k| { - if k < group.len() { - group[k].1.clone() - } else { - rabitq::dummy_code(dims) - } - }); - let packed = rabitq::pack_codes(dims, codes); - tape.push(&Height1Tuple { - mask: std::array::from_fn(|k| k < group.len()), - mean: std::array::from_fn(|k| { - if k < group.len() { - pointer_of_means[i - 1][group[k].0 as usize] - } else { - Default::default() - } - }), - first: std::array::from_fn(|k| { - if k < group.len() { - pointer_of_firsts[i - 1][group[k].0 as usize] - } else { - Default::default() - } - }), - dis_u_2: packed.dis_u_2, - factor_ppc: packed.factor_ppc, - factor_ip: packed.factor_ip, - factor_err: packed.factor_err, - t: packed.t, - }); - } - level.push(tape.first()); - } - } - pointer_of_firsts.push(level); - } - meta.push(&MetaTuple { - dims, - height_of_root: structures.len() as u32, - is_residual, - vectors_first: vectors.first(), - mean: pointer_of_means.last().unwrap()[0], - first: pointer_of_firsts.last().unwrap()[0], - }); -} - -struct Structure { - means: Vec>, - children: Vec>, -} - -impl Structure { - fn len(&self) -> usize { - self.children.len() - } - fn internal_build( - vector_options: VectorOptions, - internal_build: VchordrqfscanInternalBuildOptions, - mut samples: Vec>, - ) -> Vec { - use std::iter::once; - for sample in samples.iter_mut() { - *sample = crate::projection::project(sample); - } - let mut result = Vec::::new(); - for w in internal_build.lists.iter().rev().copied().chain(once(1)) { - let means = crate::utils::parallelism::RayonParallelism::scoped( - internal_build.build_threads as _, - Arc::new(|| { - pgrx::check_for_interrupts!(); - }), - |parallelism| { - crate::utils::k_means::k_means( - parallelism, - w as usize, - vector_options.dims as usize, - if let Some(structure) = result.last() { - &structure.means - } else { - &samples - }, - internal_build.spherical_centroids, - 10, - ) - }, - ) - .expect("failed to create thread pool"); - if let Some(structure) = result.last() { - let mut children = vec![Vec::new(); means.len()]; - for i in 0..structure.len() as u32 { - let target = - crate::utils::k_means::k_means_lookup(&structure.means[i as usize], &means); - children[target].push(i); - } - let (means, children) = std::iter::zip(means, children) - .filter(|(_, x)| !x.is_empty()) - .unzip::<_, _, Vec<_>, Vec<_>>(); - result.push(Structure { means, children }); - } else { - let children = vec![Vec::new(); means.len()]; - result.push(Structure { means, children }); - } - } - result - } - fn extern_build( - vector_options: VectorOptions, - _opfamily: Opfamily, - external_build: VchordrqfscanExternalBuildOptions, - ) -> Vec { - use std::collections::BTreeMap; - let VchordrqfscanExternalBuildOptions { table } = external_build; - let mut parents = BTreeMap::new(); - let mut vectors = BTreeMap::new(); - pgrx::spi::Spi::connect(|client| { - use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; - use pgrx::pg_sys::panic::ErrorReportable; - use vector::VectorBorrowed; - let schema_query = "SELECT n.nspname::TEXT - FROM pg_catalog.pg_extension e - LEFT JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace - WHERE e.extname = 'vector';"; - let pgvector_schema: String = client - .select(schema_query, None, None) - .unwrap_or_report() - .first() - .get_by_name("nspname") - .expect("external build: cannot get schema of pgvector") - .expect("external build: cannot get schema of pgvector"); - let dump_query = - format!("SELECT id, parent, vector::{pgvector_schema}.vector FROM {table};"); - let centroids = client.select(&dump_query, None, None).unwrap_or_report(); - for row in centroids { - let id: Option = row.get_by_name("id").unwrap(); - let parent: Option = row.get_by_name("parent").unwrap(); - let vector: Option = row.get_by_name("vector").unwrap(); - let id = id.expect("external build: id could not be NULL"); - let vector = vector.expect("external build: vector could not be NULL"); - let pop = parents.insert(id, parent); - if pop.is_some() { - pgrx::error!( - "external build: there are at least two lines have same id, id = {id}" - ); - } - if vector_options.dims != vector.as_borrowed().dims() { - pgrx::error!("external build: incorrect dimension, id = {id}"); - } - vectors.insert(id, crate::projection::project(vector.as_borrowed().slice())); - } - }); - if parents.len() >= 2 && parents.values().all(|x| x.is_none()) { - // if there are more than one vertexs and no edges, - // assume there is an implicit root - let n = parents.len(); - let mut result = Vec::new(); - result.push(Structure { - means: vectors.values().cloned().collect::>(), - children: vec![Vec::new(); n], - }); - result.push(Structure { - means: vec![{ - // compute the vector on root, without normalizing it - let mut sum = vec![0.0f32; vector_options.dims as _]; - for vector in vectors.values() { - f32::vector_add_inplace(&mut sum, vector); - } - f32::vector_mul_scalar_inplace(&mut sum, 1.0 / n as f32); - sum - }], - children: vec![(0..n as u32).collect()], - }); - return result; - } - let mut children = parents - .keys() - .map(|x| (*x, Vec::new())) - .collect::>(); - let mut root = None; - for (&id, &parent) in parents.iter() { - if let Some(parent) = parent { - if let Some(parent) = children.get_mut(&parent) { - parent.push(id); - } else { - pgrx::error!( - "external build: parent does not exist, id = {id}, parent = {parent}" - ); - } - } else { - if let Some(root) = root { - pgrx::error!("external build: two root, id = {root}, id = {id}"); - } else { - root = Some(id); - } - } - } - let Some(root) = root else { - pgrx::error!("external build: there are no root"); - }; - let mut heights = BTreeMap::<_, _>::new(); - fn dfs_for_heights( - heights: &mut BTreeMap>, - children: &BTreeMap>, - u: i32, - ) { - if heights.contains_key(&u) { - pgrx::error!("external build: detect a cycle, id = {u}"); - } - heights.insert(u, None); - let mut height = None; - for &v in children[&u].iter() { - dfs_for_heights(heights, children, v); - let new = heights[&v].unwrap() + 1; - if let Some(height) = height { - if height != new { - pgrx::error!("external build: two heights, id = {u}"); - } - } else { - height = Some(new); - } - } - if height.is_none() { - height = Some(1); - } - heights.insert(u, height); - } - dfs_for_heights(&mut heights, &children, root); - let heights = heights - .into_iter() - .map(|(k, v)| (k, v.expect("not a connected graph"))) - .collect::>(); - if !(1..=8).contains(&(heights[&root] - 1)) { - pgrx::error!( - "external build: unexpected tree height, height = {}", - heights[&root] - ); - } - let mut cursors = vec![0_u32; 1 + heights[&root] as usize]; - let mut labels = BTreeMap::new(); - for id in parents.keys().copied() { - let height = heights[&id]; - let cursor = cursors[height as usize]; - labels.insert(id, (height, cursor)); - cursors[height as usize] += 1; - } - fn extract( - height: u32, - labels: &BTreeMap, - vectors: &BTreeMap>, - children: &BTreeMap>, - ) -> (Vec>, Vec>) { - labels - .iter() - .filter(|(_, &(h, _))| h == height) - .map(|(id, _)| { - ( - vectors[id].clone(), - children[id].iter().map(|id| labels[id].1).collect(), - ) - }) - .unzip() - } - let mut result = Vec::new(); - for height in 1..=heights[&root] { - let (means, children) = extract(height, &labels, &vectors, &children); - result.push(Structure { means, children }); - } - result - } -} - -struct Tape<'a: 'b, 'b, T, R: 'b + RelationWrite> { - relation: &'a R, - head: R::WriteGuard<'b>, - first: u32, - tracking_freespace: bool, - _phantom: PhantomData T>, -} - -impl<'a: 'b, 'b, T, R: 'b + RelationWrite> Tape<'a, 'b, T, R> { - fn create(relation: &'a R, tracking_freespace: bool) -> Self { - let mut head = relation.extend(tracking_freespace); - head.get_opaque_mut().skip = head.id(); - let first = head.id(); - Self { - relation, - head, - first, - tracking_freespace, - _phantom: PhantomData, - } - } - fn first(&self) -> u32 { - self.first - } -} - -impl<'a: 'b, 'b, T, R: 'b + RelationWrite> Tape<'a, 'b, T, R> -where - T: rkyv::Serialize>, -{ - fn push(&mut self, x: &T) -> (u32, u16) { - let bytes = rkyv::to_bytes(x).expect("failed to serialize"); - if let Some(i) = self.head.alloc(&bytes) { - (self.head.id(), i) - } else { - let next = self.relation.extend(self.tracking_freespace); - self.head.get_opaque_mut().next = next.id(); - self.head = next; - if let Some(i) = self.head.alloc(&bytes) { - (self.head.id(), i) - } else { - panic!("tuple is too large to fit in a fresh page") - } - } - } -} diff --git a/src/vchordrqfscan/algorithm/insert.rs b/src/vchordrqfscan/algorithm/insert.rs deleted file mode 100644 index cb30577f..00000000 --- a/src/vchordrqfscan/algorithm/insert.rs +++ /dev/null @@ -1,299 +0,0 @@ -use crate::vchordrqfscan::algorithm::rabitq; -use crate::vchordrqfscan::algorithm::rabitq::fscan_process_lowerbound; -use crate::vchordrqfscan::algorithm::tuples::*; -use crate::vchordrqfscan::types::DistanceKind; -use crate::vchordrqfscan::types::distance; -use algorithm::{Page, PageGuard, RelationWrite}; -use always_equal::AlwaysEqual; -use distance::Distance; -use simd::Floating; -use std::cmp::Reverse; -use std::collections::BinaryHeap; -use std::num::NonZeroU64; - -pub fn insert( - relation: impl RelationWrite + Clone, - payload: NonZeroU64, - vector: Vec, - distance_kind: DistanceKind, - in_building: bool, -) { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dims = meta_tuple.dims; - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); - let vector = crate::projection::project(&vector); - let is_residual = meta_tuple.is_residual; - let default_lut = if !is_residual { - Some(rabitq::fscan_preprocess(&vector)) - } else { - None - }; - let h0_vector = { - let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { - vector: vector.clone(), - payload: Some(payload), - }) - .unwrap(); - append( - relation.clone(), - meta_tuple.vectors_first, - &tuple, - true, - true, - true, - ) - }; - let h0_payload = payload; - let mut list = ( - meta_tuple.first, - if is_residual { - let vector_guard = relation.read(meta_tuple.mean.0); - let vector_tuple = vector_guard - .get(meta_tuple.mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - Some(vector_tuple.vector.to_vec()) - } else { - None - }, - ); - let make_list = |list: (u32, Option>)| { - let mut results = Vec::new(); - { - let lut = if is_residual { - &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) - } else { - default_lut.as_ref().unwrap() - }; - let mut current = list.0; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let lowerbounds = fscan_process_lowerbound( - distance_kind, - dims, - lut, - ( - &h1_tuple.dis_u_2, - &h1_tuple.factor_ppc, - &h1_tuple.factor_ip, - &h1_tuple.factor_err, - &h1_tuple.t, - ), - 1.9, - ); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(( - Reverse(lowerbounds[j]), - AlwaysEqual(h1_tuple.mean[j]), - AlwaysEqual(h1_tuple.first[j]), - )); - } - } - } - current = h1_guard.get_opaque().next; - } - } - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get(mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); - cache.push(( - Reverse(dis_u), - AlwaysEqual(first), - AlwaysEqual(if is_residual { - Some(vector_tuple.vector.to_vec()) - } else { - None - }), - )); - } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop().unwrap(); - (first, mean) - } - }; - for _ in (1..meta_tuple.height_of_root).rev() { - list = make_list(list); - } - let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(&vector, list.1.as_ref().unwrap())) - } else { - rabitq::code(dims, &vector) - }; - let dummy = rkyv::to_bytes::<_, 8192>(&Height0Tuple { - mask: [false; 32], - mean: [(0, 0); 32], - payload: [NonZeroU64::MIN; 32], - dis_u_2: [0.0f32; 32], - factor_ppc: [0.0f32; 32], - factor_ip: [0.0f32; 32], - factor_err: [0.0f32; 32], - t: vec![0; (dims.div_ceil(4) * 16) as usize], - }) - .unwrap(); - append_by_update( - relation.clone(), - list.0, - &dummy, - in_building, - in_building, - |bytes| { - let t = rkyv::check_archived_root::(bytes).expect("data corruption"); - t.mask.iter().any(|x| *x) - }, - |bytes| put(bytes, dims, &code, h0_vector, h0_payload), - ); -} - -fn append( - relation: impl RelationWrite, - first: u32, - tuple: &[u8], - tracking_freespace: bool, - skipping_traversal: bool, - updating_skip: bool, -) -> (u32, u16) { - if tracking_freespace { - if let Some(mut write) = relation.search(tuple.len()) { - let i = write.alloc(tuple).unwrap(); - return (write.id(), i); - } - } - assert!(first != u32::MAX); - let mut current = first; - loop { - let read = relation.read(current); - if read.freespace() as usize >= tuple.len() || read.get_opaque().next == u32::MAX { - drop(read); - let mut write = relation.write(current, tracking_freespace); - if let Some(i) = write.alloc(tuple) { - return (current, i); - } - if write.get_opaque().next == u32::MAX { - let mut extend = relation.extend(tracking_freespace); - write.get_opaque_mut().next = extend.id(); - drop(write); - if let Some(i) = extend.alloc(tuple) { - let result = (extend.id(), i); - drop(extend); - if updating_skip { - let mut past = relation.write(first, tracking_freespace); - let skip = &mut past.get_opaque_mut().skip; - assert!(*skip != u32::MAX); - *skip = std::cmp::max(*skip, result.0); - } - return result; - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - } - if skipping_traversal && current == first && write.get_opaque().skip != first { - current = write.get_opaque().skip; - } else { - current = write.get_opaque().next; - } - } else { - if skipping_traversal && current == first && read.get_opaque().skip != first { - current = read.get_opaque().skip; - } else { - current = read.get_opaque().next; - } - } - } -} - -fn append_by_update( - relation: impl RelationWrite, - first: u32, - tuple: &[u8], - skipping_traversal: bool, - updating_skip: bool, - can_update: impl Fn(&[u8]) -> bool, - mut update: impl FnMut(&mut [u8]) -> bool, -) { - assert!(first != u32::MAX); - let mut current = first; - loop { - let read = relation.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - if can_update(read.get(i).expect("data corruption")) { - break 'flag true; - } - } - if read.freespace() as usize >= tuple.len() { - break 'flag true; - } - if read.get_opaque().next == u32::MAX { - break 'flag true; - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current, false); - for i in 1..=write.len() { - if update(write.get_mut(i).expect("data corruption")) { - return; - } - } - if let Some(i) = write.alloc(tuple) { - if update(write.get_mut(i).expect("data corruption")) { - return; - } - panic!("an update fails on a fresh tuple"); - } - if write.get_opaque().next == u32::MAX { - let mut extend = relation.extend(false); - write.get_opaque_mut().next = extend.id(); - drop(write); - if let Some(i) = extend.alloc(tuple) { - if update(extend.get_mut(i).expect("data corruption")) { - let id = extend.id(); - drop(extend); - if updating_skip { - let mut past = relation.write(first, false); - let skip = &mut past.get_opaque_mut().skip; - assert!(*skip != u32::MAX); - *skip = std::cmp::max(*skip, id); - } - return; - } - panic!("an update fails on a fresh tuple"); - } - panic!("a tuple cannot even be fit in a fresh page"); - } - if skipping_traversal && current == first && write.get_opaque().skip != first { - current = write.get_opaque().skip; - } else { - current = write.get_opaque().next; - } - } else { - if skipping_traversal && current == first && read.get_opaque().skip != first { - current = read.get_opaque().skip; - } else { - current = read.get_opaque().next; - } - } - } -} diff --git a/src/vchordrqfscan/algorithm/mod.rs b/src/vchordrqfscan/algorithm/mod.rs deleted file mode 100644 index 448d9194..00000000 --- a/src/vchordrqfscan/algorithm/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -pub mod build; -pub mod insert; -pub mod prewarm; -pub mod rabitq; -pub mod scan; -pub mod tuples; -pub mod vacuum; diff --git a/src/vchordrqfscan/algorithm/prewarm.rs b/src/vchordrqfscan/algorithm/prewarm.rs deleted file mode 100644 index c8500d43..00000000 --- a/src/vchordrqfscan/algorithm/prewarm.rs +++ /dev/null @@ -1,102 +0,0 @@ -use crate::vchordrqfscan::algorithm::tuples::*; -use algorithm::{Page, RelationRead}; -use std::fmt::Write; - -pub fn prewarm(relation: impl RelationRead, height: i32) -> String { - let mut message = String::new(); - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - writeln!(message, "height of root: {}", meta_tuple.height_of_root).unwrap(); - let prewarm_max_height = if height < 0 { 0 } else { height as u32 }; - if prewarm_max_height > meta_tuple.height_of_root { - return message; - } - let mut lists = { - let mut results = Vec::new(); - let counter = 1_usize; - { - let vector_guard = relation.read(meta_tuple.mean.0); - let vector_tuple = vector_guard - .get(meta_tuple.mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let _ = vector_tuple; - results.push(meta_tuple.first); - } - writeln!(message, "number of tuples: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); - results - }; - let mut make_lists = |lists| { - let mut counter = 0_usize; - let mut results = Vec::new(); - for list in lists { - let mut current = list; - while current != u32::MAX { - counter += 1; - pgrx::check_for_interrupts!(); - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(h1_tuple.first[j]); - let mean = h1_tuple.mean[j]; - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get(mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let _ = vector_tuple; - } - } - } - current = h1_guard.get_opaque().next; - } - } - writeln!(message, "number of tuples: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); - results - }; - for _ in (std::cmp::max(1, prewarm_max_height)..meta_tuple.height_of_root).rev() { - lists = make_lists(lists); - } - if prewarm_max_height == 0 { - let mut counter = 0_usize; - let mut results = Vec::new(); - for list in lists { - let mut current = list; - while current != u32::MAX { - counter += 1; - pgrx::check_for_interrupts!(); - let h0_guard = relation.read(current); - for i in 1..=h0_guard.len() { - let h0_tuple = h0_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - for j in 0..32 { - if h0_tuple.mask[j] { - results.push(()); - } - } - } - current = h0_guard.get_opaque().next; - } - } - writeln!(message, "number of tuples: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); - } - message -} diff --git a/src/vchordrqfscan/algorithm/rabitq.rs b/src/vchordrqfscan/algorithm/rabitq.rs deleted file mode 100644 index 707d81c6..00000000 --- a/src/vchordrqfscan/algorithm/rabitq.rs +++ /dev/null @@ -1,22 +0,0 @@ -use crate::vchordrqfscan::types::DistanceKind; -use distance::Distance; - -pub use rabitq::block::Code; -pub use rabitq::block::code; -pub use rabitq::block::dummy_code; -pub use rabitq::block::fscan_preprocess; -pub use rabitq::block::pack_codes; -pub use rabitq::block::{fscan_process_lowerbound_dot, fscan_process_lowerbound_l2}; - -pub fn fscan_process_lowerbound( - distance_kind: DistanceKind, - dims: u32, - lut: &(f32, f32, f32, f32, Vec), - code: (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], &[u8]), - epsilon: f32, -) -> [Distance; 32] { - match distance_kind { - DistanceKind::L2 => fscan_process_lowerbound_l2(dims, lut, code, epsilon), - DistanceKind::Dot => fscan_process_lowerbound_dot(dims, lut, code, epsilon), - } -} diff --git a/src/vchordrqfscan/algorithm/scan.rs b/src/vchordrqfscan/algorithm/scan.rs deleted file mode 100644 index 5546af93..00000000 --- a/src/vchordrqfscan/algorithm/scan.rs +++ /dev/null @@ -1,193 +0,0 @@ -use crate::vchordrqfscan::algorithm::rabitq; -use crate::vchordrqfscan::algorithm::rabitq::fscan_process_lowerbound; -use crate::vchordrqfscan::algorithm::tuples::*; -use crate::vchordrqfscan::types::DistanceKind; -use crate::vchordrqfscan::types::distance; -use algorithm::{Page, RelationWrite}; -use always_equal::AlwaysEqual; -use distance::Distance; -use simd::Floating; -use std::cmp::Reverse; -use std::collections::BinaryHeap; -use std::num::NonZeroU64; - -pub fn scan( - relation: impl RelationWrite + Clone, - vector: Vec, - distance_kind: DistanceKind, - probes: Vec, - epsilon: f32, -) -> impl Iterator { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dims = meta_tuple.dims; - let height_of_root = meta_tuple.height_of_root; - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); - assert_eq!(height_of_root as usize, 1 + probes.len(), "invalid probes"); - let vector = crate::projection::project(&vector); - let is_residual = meta_tuple.is_residual; - let default_lut = if !is_residual { - Some(rabitq::fscan_preprocess(&vector)) - } else { - None - }; - let mut lists: Vec<_> = vec![( - meta_tuple.first, - if is_residual { - let vector_guard = relation.read(meta_tuple.mean.0); - let vector_tuple = vector_guard - .get(meta_tuple.mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - Some(vector_tuple.vector.to_vec()) - } else { - None - }, - )]; - let make_lists = |lists: Vec<(u32, Option>)>, probes| { - let mut results = Vec::new(); - for list in lists { - let lut = if is_residual { - &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) - } else { - default_lut.as_ref().unwrap() - }; - let mut current = list.0; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let lowerbounds = fscan_process_lowerbound( - distance_kind, - dims, - lut, - ( - &h1_tuple.dis_u_2, - &h1_tuple.factor_ppc, - &h1_tuple.factor_ip, - &h1_tuple.factor_err, - &h1_tuple.t, - ), - epsilon, - ); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(( - Reverse(lowerbounds[j]), - AlwaysEqual(h1_tuple.mean[j]), - AlwaysEqual(h1_tuple.first[j]), - )); - } - } - } - current = h1_guard.get_opaque().next; - } - } - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - std::iter::from_fn(|| { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get(mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); - cache.push(( - Reverse(dis_u), - AlwaysEqual(first), - AlwaysEqual(if is_residual { - Some(vector_tuple.vector.to_vec()) - } else { - None - }), - )); - } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; - Some((first, mean)) - }) - .take(probes as usize) - .collect() - }; - for i in (1..meta_tuple.height_of_root).rev() { - lists = make_lists(lists, probes[i as usize - 1]); - } - drop(meta_guard); - { - let mut results = Vec::new(); - for list in lists { - let lut = if is_residual { - &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) - } else { - default_lut.as_ref().unwrap() - }; - let mut current = list.0; - while current != u32::MAX { - let h0_guard = relation.read(current); - for i in 1..=h0_guard.len() { - let h0_tuple = h0_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let lowerbounds = fscan_process_lowerbound( - distance_kind, - dims, - lut, - ( - &h0_tuple.dis_u_2, - &h0_tuple.factor_ppc, - &h0_tuple.factor_ip, - &h0_tuple.factor_err, - &h0_tuple.t, - ), - epsilon, - ); - for j in 0..32 { - if h0_tuple.mask[j] { - results.push(( - Reverse(lowerbounds[j]), - AlwaysEqual(h0_tuple.mean[j]), - AlwaysEqual(h0_tuple.payload[j]), - )); - } - } - } - current = h0_guard.get_opaque().next; - } - } - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _)>::new(); - std::iter::from_fn(move || { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); - let vector_guard = relation.read(mean.0); - let Some(vector_tuple) = vector_guard.get(mean.1) else { - // fails consistency check - continue; - }; - let vector_tuple = rkyv::check_archived_root::(vector_tuple) - .expect("data corruption"); - if vector_tuple.payload != Some(pay_u) { - // fails consistency check - continue; - } - let dis_u = distance(distance_kind, &vector, &vector_tuple.vector); - cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); - } - let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; - Some((dis_u, pay_u)) - }) - } -} diff --git a/src/vchordrqfscan/algorithm/tuples.rs b/src/vchordrqfscan/algorithm/tuples.rs deleted file mode 100644 index 7d1c97af..00000000 --- a/src/vchordrqfscan/algorithm/tuples.rs +++ /dev/null @@ -1,140 +0,0 @@ -use std::num::NonZeroU64; - -use crate::vchordrqfscan::algorithm::rabitq; -use rkyv::{Archive, Deserialize, Serialize}; - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct MetaTuple { - pub dims: u32, - pub height_of_root: u32, - pub is_residual: bool, - pub vectors_first: u32, - // raw vector - pub mean: (u32, u16), - // for meta tuple, it's pointers to next level - pub first: u32, -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct VectorTuple { - pub vector: Vec, - // this field is saved only for vacuum - pub payload: Option, -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct Height1Tuple { - pub mask: [bool; 32], - // raw vector - pub mean: [(u32, u16); 32], - // for height 1 tuple, it's pointers to next level - pub first: [u32; 32], - // RaBitQ algorithm - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct Height0Tuple { - pub mask: [bool; 32], - // raw vector - pub mean: [(u32, u16); 32], - // for height 0 tuple, it's pointers to heap relation - pub payload: [NonZeroU64; 32], - // RaBitQ algorithm - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, -} - -pub fn put( - bytes: &mut [u8], - dims: u32, - code: &rabitq::Code, - vector: (u32, u16), - payload: NonZeroU64, -) -> bool { - // todo: use mutable api - let mut x = rkyv::from_bytes::(bytes).expect("data corruption"); - for j in 0..32 { - if !x.mask[j] { - x.mean[j] = vector; - x.payload[j] = payload; - x.mask[j] = true; - x.dis_u_2[j] = code.dis_u_2; - x.factor_ppc[j] = code.factor_ppc; - x.factor_ip[j] = code.factor_ip; - x.factor_err[j] = code.factor_err; - let width = dims.div_ceil(4) as usize; - let table = [ - (0, 0), - (2, 0), - (4, 0), - (6, 0), - (8, 0), - (10, 0), - (12, 0), - (14, 0), - (1, 0), - (3, 0), - (5, 0), - (7, 0), - (9, 0), - (11, 0), - (13, 0), - (15, 0), - (0, 1), - (2, 1), - (4, 1), - (6, 1), - (8, 1), - (10, 1), - (12, 1), - (14, 1), - (1, 1), - (3, 1), - (5, 1), - (7, 1), - (9, 1), - (11, 1), - (13, 1), - (15, 1), - ]; - let pos = table[j].0; - let mask = match table[j].1 { - 0 => 0xf0, - 1 => 0x0f, - _ => unreachable!(), - }; - let shift = match table[j].1 { - 0 => 0, - 1 => 4, - _ => unreachable!(), - }; - let mut buffer = vec![0u8; width]; - for j in 0..width { - let b0 = code.signs.get(4 * j + 0).copied().unwrap_or_default(); - let b1 = code.signs.get(4 * j + 1).copied().unwrap_or_default(); - let b2 = code.signs.get(4 * j + 2).copied().unwrap_or_default(); - let b3 = code.signs.get(4 * j + 3).copied().unwrap_or_default(); - buffer[j] = b0 | b1 << 1 | b2 << 2 | b3 << 3; - } - for j in 0..width { - x.t[16 * j + pos] &= mask; - x.t[16 * j + pos] |= buffer[j] << shift; - } - bytes.copy_from_slice(&rkyv::to_bytes::<_, 8192>(&x).unwrap()); - return true; - } - } - false -} diff --git a/src/vchordrqfscan/algorithm/vacuum.rs b/src/vchordrqfscan/algorithm/vacuum.rs deleted file mode 100644 index 7773ed27..00000000 --- a/src/vchordrqfscan/algorithm/vacuum.rs +++ /dev/null @@ -1,139 +0,0 @@ -use crate::vchordrqfscan::algorithm::tuples::VectorTuple; -use crate::vchordrqfscan::algorithm::tuples::*; -use algorithm::{Page, RelationWrite}; -use std::num::NonZeroU64; - -pub fn vacuum( - relation: impl RelationWrite, - delay: impl Fn(), - callback: impl Fn(NonZeroU64) -> bool, -) { - // step 1: vacuum height_0_tuple - { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let mut firsts = vec![meta_tuple.first]; - let make_firsts = |firsts| { - let mut results = Vec::new(); - for first in firsts { - let mut current = first; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(h1_tuple.first[j]); - } - } - } - current = h1_guard.get_opaque().next; - } - } - results - }; - for _ in (1..meta_tuple.height_of_root).rev() { - firsts = make_firsts(firsts); - } - for first in firsts { - let mut current = first; - while current != u32::MAX { - delay(); - let mut h0_guard = relation.write(current, false); - for i in 1..=h0_guard.len() { - let h0_tuple = h0_guard - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let flag = 'flag: { - for j in 0..32 { - if h0_tuple.mask[j] && callback(h0_tuple.payload[j]) { - break 'flag true; - } - } - false - }; - if flag { - // todo: use mutable API - let mut temp = h0_guard - .get(i) - .map(rkyv::from_bytes::) - .expect("data corruption") - .expect("data corruption"); - for j in 0..32 { - if temp.mask[j] && callback(temp.payload[j]) { - temp.mask[j] = false; - } - } - let temp = rkyv::to_bytes::<_, 8192>(&temp).expect("failed to serialize"); - h0_guard - .get_mut(i) - .expect("data corruption") - .copy_from_slice(&temp); - } - } - // todo: cross-tuple vacuum so that we can skip a tuple - current = h0_guard.get_opaque().next; - } - } - } - // step 2: vacuum vector_tuple - { - let mut current = { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - meta_tuple.vectors_first - }; - while current != u32::MAX { - delay(); - let read = relation.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - let Some(vector_tuple) = read.get(i) else { - continue; - }; - let vector_tuple = rkyv::check_archived_root::(vector_tuple) - .expect("data corruption"); - if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(payload) { - break 'flag true; - } - } - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current, true); - for i in 1..=write.len() { - let Some(vector_tuple) = write.get(i) else { - continue; - }; - let vector_tuple = rkyv::check_archived_root::(vector_tuple) - .expect("data corruption"); - if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(payload) { - write.free(i); - } - } - } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; - } - } - } -} diff --git a/src/vchordrqfscan/gucs/executing.rs b/src/vchordrqfscan/gucs/executing.rs deleted file mode 100644 index ba204b9f..00000000 --- a/src/vchordrqfscan/gucs/executing.rs +++ /dev/null @@ -1,72 +0,0 @@ -use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; -use std::ffi::CStr; - -static PROBES: GucSetting> = GucSetting::>::new(Some(c"10")); -static EPSILON: GucSetting = GucSetting::::new(1.9); -static MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); - -pub unsafe fn init() { - GucRegistry::define_string_guc( - "vchordrqfscan.probes", - "`probes` argument of vchordrqfscan.", - "`probes` argument of vchordrqfscan.", - &PROBES, - GucContext::Userset, - GucFlags::default(), - ); - GucRegistry::define_float_guc( - "vchordrqfscan.epsilon", - "`epsilon` argument of vchordrqfscan.", - "`epsilon` argument of vchordrqfscan.", - &EPSILON, - 0.0, - 4.0, - GucContext::Userset, - GucFlags::default(), - ); - GucRegistry::define_int_guc( - "vchordrqfscan.max_scan_tuples", - "`max_scan_tuples` argument of vchordrqfscan.", - "`max_scan_tuples` argument of vchordrqfscan.", - &MAX_SCAN_TUPLES, - -1, - u16::MAX as _, - GucContext::Userset, - GucFlags::default(), - ); -} - -pub fn probes() -> Vec { - match PROBES.get() { - None => Vec::new(), - Some(probes) => { - let mut result = Vec::new(); - let mut current = None; - for &c in probes.to_bytes() { - match c { - b' ' => continue, - b',' => result.push(current.take().expect("empty probes")), - b'0'..=b'9' => { - if let Some(x) = current.as_mut() { - *x = *x * 10 + (c - b'0') as u32; - } else { - current = Some((c - b'0') as u32); - } - } - c => pgrx::error!("unknown character in probes: ASCII = {c}"), - } - } - result.push(current.take().expect("empty probes")); - result - } - } -} - -pub fn epsilon() -> f32 { - EPSILON.get() as f32 -} - -pub fn max_scan_tuples() -> Option { - let x = MAX_SCAN_TUPLES.get(); - if x < 0 { None } else { Some(x as u32) } -} diff --git a/src/vchordrqfscan/gucs/mod.rs b/src/vchordrqfscan/gucs/mod.rs deleted file mode 100644 index 48cc060d..00000000 --- a/src/vchordrqfscan/gucs/mod.rs +++ /dev/null @@ -1,14 +0,0 @@ -pub mod executing; -pub mod prewarm; - -pub unsafe fn init() { - unsafe { - executing::init(); - prewarm::init(); - prewarm::prewarm(); - #[cfg(any(feature = "pg13", feature = "pg14"))] - pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrqfscan".as_ptr()); - #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] - pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrqfscan".as_ptr()); - } -} diff --git a/src/vchordrqfscan/gucs/prewarm.rs b/src/vchordrqfscan/gucs/prewarm.rs deleted file mode 100644 index ae9180a5..00000000 --- a/src/vchordrqfscan/gucs/prewarm.rs +++ /dev/null @@ -1,32 +0,0 @@ -use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; -use std::ffi::CStr; - -static PREWARM_DIM: GucSetting> = - GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); - -pub unsafe fn init() { - GucRegistry::define_string_guc( - "vchordrqfscan.prewarm_dim", - "prewarm_dim when the extension is loading.", - "prewarm_dim when the extension is loading.", - &PREWARM_DIM, - GucContext::Userset, - GucFlags::default(), - ); -} - -pub fn prewarm() { - if let Some(prewarm_dim) = PREWARM_DIM.get() { - if let Ok(prewarm_dim) = prewarm_dim.to_str() { - for dim in prewarm_dim.split(',') { - if let Ok(dim) = dim.trim().parse::() { - crate::projection::prewarm(dim as _); - } else { - pgrx::warning!("{dim:?} is not a valid integer"); - } - } - } else { - pgrx::warning!("vchordrqfscan.prewarm_dim is not a valid UTF-8 string"); - } - } -} diff --git a/src/vchordrqfscan/index/am.rs b/src/vchordrqfscan/index/am.rs deleted file mode 100644 index 339b315c..00000000 --- a/src/vchordrqfscan/index/am.rs +++ /dev/null @@ -1,856 +0,0 @@ -use crate::postgres::PostgresRelation; -use crate::vchordrqfscan::algorithm; -use crate::vchordrqfscan::algorithm::build::{HeapRelation, Reporter}; -use crate::vchordrqfscan::index::am_options::{Opfamily, Reloption}; -use crate::vchordrqfscan::index::am_scan::Scanner; -use crate::vchordrqfscan::index::utils::{ctid_to_pointer, pointer_to_ctid}; -use crate::vchordrqfscan::index::{am_options, am_scan}; -use pgrx::datum::Internal; -use pgrx::pg_sys::Datum; -use std::num::NonZeroU64; - -static mut RELOPT_KIND_VCHORDRQFSCAN: pgrx::pg_sys::relopt_kind::Type = 0; - -pub unsafe fn init() { - unsafe { - (&raw mut RELOPT_KIND_VCHORDRQFSCAN).write(pgrx::pg_sys::add_reloption_kind()); - pgrx::pg_sys::add_string_reloption( - (&raw const RELOPT_KIND_VCHORDRQFSCAN).read(), - c"options".as_ptr(), - c"Vector index options, represented as a TOML string.".as_ptr(), - c"".as_ptr(), - None, - pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, - ); - } -} - -#[pgrx::pg_extern(sql = "")] -fn _vchordrqfscan_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { - type T = pgrx::pg_sys::IndexAmRoutine; - unsafe { - let index_am_routine = pgrx::pg_sys::palloc0(size_of::()) as *mut T; - index_am_routine.write(AM_HANDLER); - Internal::from(Some(Datum::from(index_am_routine))) - } -} - -const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = { - let mut am_routine = - unsafe { std::mem::MaybeUninit::::zeroed().assume_init() }; - - am_routine.type_ = pgrx::pg_sys::NodeTag::T_IndexAmRoutine; - - am_routine.amsupport = 1; - am_routine.amcanorderbyop = true; - - #[cfg(feature = "pg17")] - { - am_routine.amcanbuildparallel = true; - } - - // Index access methods that set `amoptionalkey` to `false` - // must index all tuples, even if the first column is `NULL`. - // However, PostgreSQL does not generate a path if there is no - // index clauses, even if there is a `ORDER BY` clause. - // So we have to set it to `true` and set costs of every path - // for vector index scans without `ORDER BY` clauses a large number - // and throw errors if someone really wants such a path. - am_routine.amoptionalkey = true; - - am_routine.amvalidate = Some(amvalidate); - am_routine.amoptions = Some(amoptions); - am_routine.amcostestimate = Some(amcostestimate); - - am_routine.ambuild = Some(ambuild); - am_routine.ambuildempty = Some(ambuildempty); - am_routine.aminsert = Some(aminsert); - am_routine.ambulkdelete = Some(ambulkdelete); - am_routine.amvacuumcleanup = Some(amvacuumcleanup); - - am_routine.ambeginscan = Some(ambeginscan); - am_routine.amrescan = Some(amrescan); - am_routine.amgettuple = Some(amgettuple); - am_routine.amendscan = Some(amendscan); - - am_routine -}; - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { - true -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amoptions(reloptions: Datum, validate: bool) -> *mut pgrx::pg_sys::bytea { - let rdopts = unsafe { - pgrx::pg_sys::build_reloptions( - reloptions, - validate, - (&raw const RELOPT_KIND_VCHORDRQFSCAN).read(), - size_of::(), - Reloption::TAB.as_ptr(), - Reloption::TAB.len() as _, - ) - }; - rdopts as *mut pgrx::pg_sys::bytea -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amcostestimate( - _root: *mut pgrx::pg_sys::PlannerInfo, - path: *mut pgrx::pg_sys::IndexPath, - _loop_count: f64, - index_startup_cost: *mut pgrx::pg_sys::Cost, - index_total_cost: *mut pgrx::pg_sys::Cost, - index_selectivity: *mut pgrx::pg_sys::Selectivity, - index_correlation: *mut f64, - index_pages: *mut f64, -) { - unsafe { - if (*path).indexorderbys.is_null() && (*path).indexclauses.is_null() { - *index_startup_cost = f64::MAX; - *index_total_cost = f64::MAX; - *index_selectivity = 0.0; - *index_correlation = 0.0; - *index_pages = 0.0; - return; - } - *index_startup_cost = 0.0; - *index_total_cost = 0.0; - *index_selectivity = 1.0; - *index_correlation = 1.0; - *index_pages = 0.0; - } -} - -#[derive(Debug, Clone)] -struct PgReporter {} - -impl Reporter for PgReporter { - fn tuples_total(&mut self, tuples_total: u64) { - unsafe { - pgrx::pg_sys::pgstat_progress_update_param( - pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, - tuples_total as _, - ); - } - } -} - -impl PgReporter { - fn tuples_done(&mut self, tuples_done: u64) { - unsafe { - pgrx::pg_sys::pgstat_progress_update_param( - pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, - tuples_done as _, - ); - } - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambuild( - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, -) -> *mut pgrx::pg_sys::IndexBuildResult { - use validator::Validate; - #[derive(Debug, Clone)] - pub struct Heap { - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, - opfamily: Opfamily, - } - impl HeapRelation for Heap { - fn traverse(&self, progress: bool, callback: F) - where - F: FnMut((NonZeroU64, Vec)), - { - pub struct State<'a, F> { - pub this: &'a Heap, - pub callback: F, - } - #[pgrx::pg_guard] - unsafe extern "C" fn call( - _index: pgrx::pg_sys::Relation, - ctid: pgrx::pg_sys::ItemPointer, - values: *mut Datum, - is_null: *mut bool, - _tuple_is_alive: bool, - state: *mut core::ffi::c_void, - ) where - F: FnMut((NonZeroU64, Vec)), - { - use crate::vchordrqfscan::types::OwnedVector; - let state = unsafe { &mut *state.cast::>() }; - let opfamily = state.this.opfamily; - let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; - let pointer = unsafe { ctid_to_pointer(ctid.read()) }; - if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - }; - (state.callback)((pointer, vector.into_vec())); - } - } - let table_am = unsafe { &*(*self.heap).rd_tableam }; - let mut state = State { - this: self, - callback, - }; - unsafe { - table_am.index_build_range_scan.unwrap()( - self.heap, - self.index, - self.index_info, - true, - false, - progress, - 0, - pgrx::pg_sys::InvalidBlockNumber, - Some(call::), - (&mut state) as *mut State as *mut _, - std::ptr::null_mut(), - ); - } - } - - fn opfamily(&self) -> Opfamily { - self.opfamily - } - } - let (vector_options, vchordrqfscan_options) = unsafe { am_options::options(index) }; - if let Err(errors) = Validate::validate(&vector_options) { - pgrx::error!("error while validating options: {}", errors); - } - if vector_options.dims == 0 { - pgrx::error!("error while validating options: dimension cannot be 0"); - } - if vector_options.dims > 1600 { - pgrx::error!("error while validating options: dimension is too large"); - } - if let Err(errors) = Validate::validate(&vchordrqfscan_options) { - pgrx::error!("error while validating options: {}", errors); - } - let opfamily = unsafe { am_options::opfamily(index) }; - let heap_relation = Heap { - heap, - index, - index_info, - opfamily, - }; - let mut reporter = PgReporter {}; - let index_relation = unsafe { PostgresRelation::new(index) }; - algorithm::build::build( - vector_options, - vchordrqfscan_options, - heap_relation.clone(), - index_relation.clone(), - reporter.clone(), - ); - if let Some(leader) = - unsafe { VchordrqfscanLeader::enter(heap, index, (*index_info).ii_Concurrent) } - { - unsafe { - parallel_build( - index, - heap, - index_info, - leader.tablescandesc, - leader.vchordrqfscanshared, - Some(reporter), - ); - leader.wait(); - let nparticipants = leader.nparticipants; - loop { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordrqfscanshared).mutex); - if (*leader.vchordrqfscanshared).nparticipantsdone == nparticipants { - pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqfscanshared).mutex); - break; - } - pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqfscanshared).mutex); - pgrx::pg_sys::ConditionVariableSleep( - &raw mut (*leader.vchordrqfscanshared).workersdonecv, - pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN, - ); - } - pgrx::pg_sys::ConditionVariableCancelSleep(); - } - } else { - let mut indtuples = 0; - reporter.tuples_done(indtuples); - heap_relation.traverse(true, |(payload, vector)| { - algorithm::insert::insert( - index_relation.clone(), - payload, - vector, - opfamily.distance_kind(), - true, - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }); - } - unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } -} - -struct VchordrqfscanShared { - /* Immutable state */ - heaprelid: pgrx::pg_sys::Oid, - indexrelid: pgrx::pg_sys::Oid, - isconcurrent: bool, - - /* Worker progress */ - workersdonecv: pgrx::pg_sys::ConditionVariable, - - /* Mutex for mutable state */ - mutex: pgrx::pg_sys::slock_t, - - /* Mutable state */ - nparticipantsdone: i32, - indtuples: u64, -} - -fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { - matches!( - unsafe { (*snapshot).snapshot_type }, - pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC - | pgrx::pg_sys::SnapshotType::SNAPSHOT_HISTORIC_MVCC - ) -} - -struct VchordrqfscanLeader { - pcxt: *mut pgrx::pg_sys::ParallelContext, - nparticipants: i32, - vchordrqfscanshared: *mut VchordrqfscanShared, - tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, - snapshot: pgrx::pg_sys::Snapshot, -} - -impl VchordrqfscanLeader { - pub unsafe fn enter( - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - isconcurrent: bool, - ) -> Option { - unsafe fn compute_parallel_workers( - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - ) -> i32 { - unsafe { - if pgrx::pg_sys::plan_create_index_workers((*heap).rd_id, (*index).rd_id) == 0 { - return 0; - } - if !(*heap).rd_options.is_null() { - let std_options = (*heap).rd_options.cast::(); - std::cmp::min( - (*std_options).parallel_workers, - pgrx::pg_sys::max_parallel_maintenance_workers, - ) - } else { - pgrx::pg_sys::max_parallel_maintenance_workers - } - } - } - - let request = unsafe { compute_parallel_workers(heap, index) }; - if request <= 0 { - return None; - } - - unsafe { - pgrx::pg_sys::EnterParallelMode(); - } - let pcxt = unsafe { - pgrx::pg_sys::CreateParallelContext( - c"vchord".as_ptr(), - c"vchordrqfscan_parallel_build_main".as_ptr(), - request, - ) - }; - - let snapshot = if isconcurrent { - unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } - } else { - &raw mut pgrx::pg_sys::SnapshotAnyData - }; - - fn estimate_chunk(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { - e.space_for_chunks += x.next_multiple_of(pgrx::pg_sys::ALIGNOF_BUFFER as _); - } - fn estimate_keys(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { - e.number_of_keys += x; - } - let est_tablescandesc = - unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap, snapshot) }; - unsafe { - estimate_chunk(&mut (*pcxt).estimator, size_of::()); - estimate_keys(&mut (*pcxt).estimator, 1); - estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); - estimate_keys(&mut (*pcxt).estimator, 1); - } - - unsafe { - pgrx::pg_sys::InitializeParallelDSM(pcxt); - if (*pcxt).seg.is_null() { - if is_mvcc_snapshot(snapshot) { - pgrx::pg_sys::UnregisterSnapshot(snapshot); - } - pgrx::pg_sys::DestroyParallelContext(pcxt); - pgrx::pg_sys::ExitParallelMode(); - return None; - } - } - - let vchordrqfscanshared = unsafe { - let vchordrqfscanshared = - pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) - .cast::(); - vchordrqfscanshared.write(VchordrqfscanShared { - heaprelid: (*heap).rd_id, - indexrelid: (*index).rd_id, - isconcurrent, - workersdonecv: std::mem::zeroed(), - mutex: std::mem::zeroed(), - nparticipantsdone: 0, - indtuples: 0, - }); - pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqfscanshared).workersdonecv); - pgrx::pg_sys::SpinLockInit(&raw mut (*vchordrqfscanshared).mutex); - vchordrqfscanshared - }; - - let tablescandesc = unsafe { - let tablescandesc = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, est_tablescandesc) - .cast::(); - pgrx::pg_sys::table_parallelscan_initialize(heap, tablescandesc, snapshot); - tablescandesc - }; - - unsafe { - pgrx::pg_sys::shm_toc_insert( - (*pcxt).toc, - 0xA000000000000001, - vchordrqfscanshared.cast(), - ); - pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); - } - - unsafe { - pgrx::pg_sys::LaunchParallelWorkers(pcxt); - } - - let nworkers_launched = unsafe { (*pcxt).nworkers_launched }; - - unsafe { - if nworkers_launched == 0 { - pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); - if is_mvcc_snapshot(snapshot) { - pgrx::pg_sys::UnregisterSnapshot(snapshot); - } - pgrx::pg_sys::DestroyParallelContext(pcxt); - pgrx::pg_sys::ExitParallelMode(); - return None; - } - } - - Some(Self { - pcxt, - nparticipants: nworkers_launched + 1, - vchordrqfscanshared, - tablescandesc, - snapshot, - }) - } - - pub fn wait(&self) { - unsafe { - pgrx::pg_sys::WaitForParallelWorkersToAttach(self.pcxt); - } - } -} - -impl Drop for VchordrqfscanLeader { - fn drop(&mut self) { - if !std::thread::panicking() { - unsafe { - pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); - if is_mvcc_snapshot(self.snapshot) { - pgrx::pg_sys::UnregisterSnapshot(self.snapshot); - } - pgrx::pg_sys::DestroyParallelContext(self.pcxt); - pgrx::pg_sys::ExitParallelMode(); - } - } - } -} - -#[pgrx::pg_guard] -#[unsafe(no_mangle)] -pub unsafe extern "C" fn vchordrqfscan_parallel_build_main( - _seg: *mut pgrx::pg_sys::dsm_segment, - toc: *mut pgrx::pg_sys::shm_toc, -) { - let vchordrqfscanshared = unsafe { - pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() - }; - let tablescandesc = unsafe { - pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) - .cast::() - }; - let heap_lockmode; - let index_lockmode; - if unsafe { !(*vchordrqfscanshared).isconcurrent } { - heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; - } else { - heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; - } - let heap = unsafe { pgrx::pg_sys::table_open((*vchordrqfscanshared).heaprelid, heap_lockmode) }; - let index = - unsafe { pgrx::pg_sys::index_open((*vchordrqfscanshared).indexrelid, index_lockmode) }; - let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; - unsafe { - (*index_info).ii_Concurrent = (*vchordrqfscanshared).isconcurrent; - } - - unsafe { - parallel_build( - index, - heap, - index_info, - tablescandesc, - vchordrqfscanshared, - None, - ); - } - - unsafe { - pgrx::pg_sys::index_close(index, index_lockmode); - pgrx::pg_sys::table_close(heap, heap_lockmode); - } -} - -unsafe fn parallel_build( - index: *mut pgrx::pg_sys::RelationData, - heap: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, - tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, - vchordrqfscanshared: *mut VchordrqfscanShared, - mut reporter: Option, -) { - #[derive(Debug, Clone)] - pub struct Heap { - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, - opfamily: Opfamily, - scan: *mut pgrx::pg_sys::TableScanDescData, - } - impl HeapRelation for Heap { - fn traverse(&self, progress: bool, callback: F) - where - F: FnMut((NonZeroU64, Vec)), - { - pub struct State<'a, F> { - pub this: &'a Heap, - pub callback: F, - } - #[pgrx::pg_guard] - unsafe extern "C" fn call( - _index: pgrx::pg_sys::Relation, - ctid: pgrx::pg_sys::ItemPointer, - values: *mut Datum, - is_null: *mut bool, - _tuple_is_alive: bool, - state: *mut core::ffi::c_void, - ) where - F: FnMut((NonZeroU64, Vec)), - { - use crate::vchordrqfscan::types::OwnedVector; - let state = unsafe { &mut *state.cast::>() }; - let opfamily = state.this.opfamily; - let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; - let pointer = unsafe { ctid_to_pointer(ctid.read()) }; - if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - }; - (state.callback)((pointer, vector.into_vec())); - } - } - let table_am = unsafe { &*(*self.heap).rd_tableam }; - let mut state = State { - this: self, - callback, - }; - unsafe { - table_am.index_build_range_scan.unwrap()( - self.heap, - self.index, - self.index_info, - true, - false, - progress, - 0, - pgrx::pg_sys::InvalidBlockNumber, - Some(call::), - (&mut state) as *mut State as *mut _, - self.scan, - ); - } - } - - fn opfamily(&self) -> Opfamily { - self.opfamily - } - } - - let index_relation = unsafe { PostgresRelation::new(index) }; - let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap, tablescandesc) }; - let opfamily = unsafe { am_options::opfamily(index) }; - let heap_relation = Heap { - heap, - index, - index_info, - opfamily, - scan, - }; - heap_relation.traverse(reporter.is_some(), |(payload, vector)| { - algorithm::insert::insert( - index_relation.clone(), - payload, - vector, - opfamily.distance_kind(), - true, - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqfscanshared).mutex); - (*vchordrqfscanshared).indtuples += 1; - indtuples = (*vchordrqfscanshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqfscanshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }); - - unsafe { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqfscanshared).mutex); - (*vchordrqfscanshared).nparticipantsdone += 1; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqfscanshared).mutex); - pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordrqfscanshared).workersdonecv); - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambuildempty(_index: pgrx::pg_sys::Relation) { - pgrx::error!("Unlogged indexes are not supported."); -} - -#[cfg(feature = "pg13")] -#[pgrx::pg_guard] -pub unsafe extern "C" fn aminsert( - index: pgrx::pg_sys::Relation, - values: *mut Datum, - is_null: *mut bool, - heap_tid: pgrx::pg_sys::ItemPointer, - _heap: pgrx::pg_sys::Relation, - _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, - _index_info: *mut pgrx::pg_sys::IndexInfo, -) -> bool { - use crate::vchordrqfscan::types::OwnedVector; - let opfamily = unsafe { am_options::opfamily(index) }; - let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; - if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - }; - let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - algorithm::insert::insert( - unsafe { PostgresRelation::new(index) }, - pointer, - vector.into_vec(), - opfamily.distance_kind(), - false, - ); - } - false -} - -#[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] -#[pgrx::pg_guard] -pub unsafe extern "C" fn aminsert( - index: pgrx::pg_sys::Relation, - values: *mut Datum, - is_null: *mut bool, - heap_tid: pgrx::pg_sys::ItemPointer, - _heap: pgrx::pg_sys::Relation, - _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, - _index_unchanged: bool, - _index_info: *mut pgrx::pg_sys::IndexInfo, -) -> bool { - use crate::vchordrqfscan::types::OwnedVector; - let opfamily = unsafe { am_options::opfamily(index) }; - let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; - if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - }; - let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - algorithm::insert::insert( - unsafe { PostgresRelation::new(index) }, - pointer, - vector.into_vec(), - opfamily.distance_kind(), - false, - ); - } - false -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambeginscan( - index: pgrx::pg_sys::Relation, - n_keys: std::os::raw::c_int, - n_orderbys: std::os::raw::c_int, -) -> pgrx::pg_sys::IndexScanDesc { - use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; - - let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index, n_keys, n_orderbys) }; - unsafe { - let scanner = am_scan::scan_make(None, None, false); - (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); - } - scan -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amrescan( - scan: pgrx::pg_sys::IndexScanDesc, - keys: pgrx::pg_sys::ScanKey, - _n_keys: std::os::raw::c_int, - orderbys: pgrx::pg_sys::ScanKey, - _n_orderbys: std::os::raw::c_int, -) { - unsafe { - if !keys.is_null() && (*scan).numberOfKeys > 0 { - std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); - } - if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { - std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); - } - let opfamily = am_options::opfamily((*scan).indexRelation); - let (orderbys, spheres) = { - let mut orderbys = Vec::new(); - let mut spheres = Vec::new(); - if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { - pgrx::error!( - "vector search with no WHERE clause and no ORDER BY clause is not supported" - ); - } - for i in 0..(*scan).numberOfOrderBys { - let data = (*scan).orderByData.add(i as usize); - let value = (*data).sk_argument; - let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; - match (*data).sk_strategy { - 1 => orderbys.push(opfamily.datum_to_vector(value, is_null)), - _ => unreachable!(), - } - } - for i in 0..(*scan).numberOfKeys { - let data = (*scan).keyData.add(i as usize); - let value = (*data).sk_argument; - let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; - match (*data).sk_strategy { - 2 => spheres.push(opfamily.datum_to_sphere(value, is_null)), - _ => unreachable!(), - } - } - (orderbys, spheres) - }; - let (vector, threshold, recheck) = am_scan::scan_build(orderbys, spheres, opfamily); - let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace(scanner, am_scan::scan_make(vector, threshold, recheck)); - am_scan::scan_release(scanner); - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amgettuple( - scan: pgrx::pg_sys::IndexScanDesc, - direction: pgrx::pg_sys::ScanDirection::Type, -) -> bool { - if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { - pgrx::error!("vector search without a forward scan direction is not supported"); - } - // https://www.postgresql.org/docs/current/index-locking.html - // If heap entries referenced physical pointers are deleted before - // they are consumed by PostgreSQL, PostgreSQL will received wrong - // physical pointers: no rows or irreverent rows are referenced. - if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC - { - pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); - } - let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; - let relation = unsafe { PostgresRelation::new((*scan).indexRelation) }; - if let Some((pointer, recheck)) = am_scan::scan_next(scanner, relation) { - let ctid = pointer_to_ctid(pointer); - unsafe { - (*scan).xs_heaptid = ctid; - (*scan).xs_recheckorderby = false; - (*scan).xs_recheck = recheck; - } - true - } else { - false - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { - unsafe { - let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace(scanner, am_scan::scan_make(None, None, false)); - am_scan::scan_release(scanner); - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambulkdelete( - info: *mut pgrx::pg_sys::IndexVacuumInfo, - stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, - callback: pgrx::pg_sys::IndexBulkDeleteCallback, - callback_state: *mut std::os::raw::c_void, -) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { - let mut stats = stats; - if stats.is_null() { - stats = unsafe { - pgrx::pg_sys::palloc0(size_of::()).cast() - }; - } - let callback = callback.unwrap(); - let callback = |p: NonZeroU64| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; - algorithm::vacuum::vacuum( - unsafe { PostgresRelation::new((*info).index) }, - || unsafe { - pgrx::pg_sys::vacuum_delay_point(); - }, - callback, - ); - stats -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amvacuumcleanup( - _info: *mut pgrx::pg_sys::IndexVacuumInfo, - _stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, -) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { - std::ptr::null_mut() -} diff --git a/src/vchordrqfscan/index/am_options.rs b/src/vchordrqfscan/index/am_options.rs deleted file mode 100644 index be4154ef..00000000 --- a/src/vchordrqfscan/index/am_options.rs +++ /dev/null @@ -1,218 +0,0 @@ -use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; -use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; -use crate::datatype::typmod::Typmod; -use crate::vchordrqfscan::types::{BorrowedVector, OwnedVector}; -use crate::vchordrqfscan::types::{DistanceKind, VectorKind}; -use crate::vchordrqfscan::types::{VchordrqfscanIndexingOptions, VectorOptions}; -use distance::Distance; -use pgrx::datum::FromDatum; -use pgrx::heap_tuple::PgHeapTuple; -use serde::Deserialize; -use std::ffi::CStr; -use std::num::NonZero; -use vector::VectorBorrowed; - -#[derive(Copy, Clone, Debug, Default)] -#[repr(C)] -pub struct Reloption { - vl_len_: i32, - pub options: i32, -} - -impl Reloption { - pub const TAB: &'static [pgrx::pg_sys::relopt_parse_elt] = &[pgrx::pg_sys::relopt_parse_elt { - optname: c"options".as_ptr(), - opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, - offset: std::mem::offset_of!(Reloption, options) as i32, - }]; - unsafe fn options(&self) -> &CStr { - unsafe { - let ptr = (&raw const *self) - .cast::() - .offset(self.options as _); - CStr::from_ptr(ptr) - } - } -} - -#[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -pub enum PgDistanceKind { - L2, - Dot, - Cos, -} - -impl PgDistanceKind { - pub fn to_distance(self) -> DistanceKind { - match self { - PgDistanceKind::L2 => DistanceKind::L2, - PgDistanceKind::Dot | PgDistanceKind::Cos => DistanceKind::Dot, - } - } -} - -fn convert_name_to_vd(name: &str) -> Option<(VectorKind, PgDistanceKind)> { - match name.strip_suffix("_ops") { - Some("vector_l2") => Some((VectorKind::Vecf32, PgDistanceKind::L2)), - Some("vector_ip") => Some((VectorKind::Vecf32, PgDistanceKind::Dot)), - Some("vector_cosine") => Some((VectorKind::Vecf32, PgDistanceKind::Cos)), - _ => None, - } -} - -unsafe fn convert_reloptions_to_options( - reloptions: *const pgrx::pg_sys::varlena, -) -> VchordrqfscanIndexingOptions { - #[derive(Debug, Clone, Deserialize, Default)] - #[serde(deny_unknown_fields)] - struct Parsed { - #[serde(flatten)] - rabitq: VchordrqfscanIndexingOptions, - } - let reloption = reloptions as *const Reloption; - if reloption.is_null() || unsafe { (*reloption).options == 0 } { - return Default::default(); - } - let s = unsafe { (*reloption).options() }.to_string_lossy(); - match toml::from_str::(&s) { - Ok(p) => p.rabitq, - Err(e) => pgrx::error!("failed to parse options: {}", e), - } -} - -pub unsafe fn options( - index: pgrx::pg_sys::Relation, -) -> (VectorOptions, VchordrqfscanIndexingOptions) { - let att = unsafe { &mut *(*index).rd_att }; - let atts = unsafe { att.attrs.as_slice(att.natts as _) }; - if atts.is_empty() { - pgrx::error!("indexing on no columns is not supported"); - } - if atts.len() != 1 { - pgrx::error!("multicolumn index is not supported"); - } - // get dims - let typmod = Typmod::parse_from_i32(atts[0].type_mod()).unwrap(); - let dims = if let Some(dims) = typmod.dims() { - dims.get() - } else { - pgrx::error!( - "Dimensions type modifier of a vector column is needed for building the index." - ); - }; - // get v, d - let opfamily = unsafe { opfamily(index) }; - let vector = VectorOptions { - dims, - v: opfamily.vector, - d: opfamily.distance_kind(), - }; - // get indexing, segment, optimizing - let rabitq = unsafe { convert_reloptions_to_options((*index).rd_options) }; - (vector, rabitq) -} - -#[derive(Debug, Clone, Copy)] -pub struct Opfamily { - vector: VectorKind, - pg_distance: PgDistanceKind, -} - -impl Opfamily { - pub unsafe fn datum_to_vector( - self, - datum: pgrx::pg_sys::Datum, - is_null: bool, - ) -> Option { - if is_null || datum.is_null() { - return None; - } - let vector = match self.vector { - VectorKind::Vecf32 => { - let vector = unsafe { PgvectorVectorInput::from_datum(datum, false).unwrap() }; - self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed())) - } - }; - Some(vector) - } - pub unsafe fn datum_to_sphere( - self, - datum: pgrx::pg_sys::Datum, - is_null: bool, - ) -> (Option, Option) { - if is_null || datum.is_null() { - return (None, None); - } - let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; - let center = match self.vector { - VectorKind::Vecf32 => tuple - .get_by_index::(NonZero::new(1).unwrap()) - .unwrap() - .map(|vector| self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed()))), - }; - let radius = tuple.get_by_index::(NonZero::new(2).unwrap()).unwrap(); - (center, radius) - } - pub fn preprocess(self, vector: BorrowedVector<'_>) -> OwnedVector { - use BorrowedVector as B; - use OwnedVector as O; - match (vector, self.pg_distance) { - (B::Vecf32(x), PgDistanceKind::L2) => O::Vecf32(x.own()), - (B::Vecf32(x), PgDistanceKind::Dot) => O::Vecf32(x.own()), - (B::Vecf32(x), PgDistanceKind::Cos) => O::Vecf32(x.function_normalize()), - } - } - pub fn process(self, x: Distance) -> f32 { - match self.pg_distance { - PgDistanceKind::Cos => f32::from(x) + 1.0f32, - PgDistanceKind::L2 => f32::from(x).sqrt(), - _ => f32::from(x), - } - } - pub fn distance_kind(self) -> DistanceKind { - self.pg_distance.to_distance() - } -} - -pub unsafe fn opfamily(index: pgrx::pg_sys::Relation) -> Opfamily { - use pgrx::pg_sys::Oid; - - let proc = unsafe { pgrx::pg_sys::index_getprocid(index, 1, 1) }; - - if proc == Oid::INVALID { - pgrx::error!("support function 1 is not found"); - } - - let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); - unsafe { - pgrx::pg_sys::fmgr_info(proc, &mut flinfo); - } - - let fn_addr = flinfo.fn_addr.expect("null function pointer"); - - let mut fcinfo = unsafe { std::mem::zeroed::() }; - fcinfo.flinfo = &mut flinfo; - fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; - fcinfo.context = std::ptr::null_mut(); - fcinfo.resultinfo = std::ptr::null_mut(); - fcinfo.isnull = true; - fcinfo.nargs = 0; - - let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; - - let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; - - let result_string = result_option.expect("null string"); - - let (vector, pg_distance) = convert_name_to_vd(&result_string).unwrap(); - - unsafe { - pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); - } - - Opfamily { - vector, - pg_distance, - } -} diff --git a/src/vchordrqfscan/index/am_scan.rs b/src/vchordrqfscan/index/am_scan.rs deleted file mode 100644 index da049ab4..00000000 --- a/src/vchordrqfscan/index/am_scan.rs +++ /dev/null @@ -1,125 +0,0 @@ -use super::am_options::Opfamily; -use crate::postgres::PostgresRelation; -use crate::vchordrqfscan::algorithm::scan::scan; -use crate::vchordrqfscan::gucs::executing::epsilon; -use crate::vchordrqfscan::gucs::executing::max_scan_tuples; -use crate::vchordrqfscan::gucs::executing::probes; -use crate::vchordrqfscan::types::OwnedVector; -use distance::Distance; -use std::num::NonZeroU64; - -pub enum Scanner { - Initial { - vector: Option<(OwnedVector, Opfamily)>, - threshold: Option, - recheck: bool, - }, - Vbase { - vbase: Box>, - threshold: Option, - recheck: bool, - opfamily: Opfamily, - }, - Empty {}, -} - -pub fn scan_build( - orderbys: Vec>, - spheres: Vec<(Option, Option)>, - opfamily: Opfamily, -) -> (Option<(OwnedVector, Opfamily)>, Option, bool) { - let mut pair = None; - let mut threshold = None; - let mut recheck = false; - for orderby_vector in orderbys { - if pair.is_none() { - pair = orderby_vector; - } else if orderby_vector.is_some() { - pgrx::error!("vector search with multiple vectors is not supported"); - } - } - for (sphere_vector, sphere_threshold) in spheres { - if pair.is_none() { - pair = sphere_vector; - threshold = sphere_threshold; - } else { - recheck = true; - break; - } - } - (pair.map(|x| (x, opfamily)), threshold, recheck) -} - -pub fn scan_make( - vector: Option<(OwnedVector, Opfamily)>, - threshold: Option, - recheck: bool, -) -> Scanner { - Scanner::Initial { - vector, - threshold, - recheck, - } -} - -pub fn scan_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(NonZeroU64, bool)> { - if let Scanner::Initial { - vector, - threshold, - recheck, - } = scanner - { - if let Some((vector, opfamily)) = vector.as_ref() { - let vbase = scan( - relation, - match vector { - OwnedVector::Vecf32(x) => x.slice().to_vec(), - }, - opfamily.distance_kind(), - probes(), - epsilon(), - ); - *scanner = Scanner::Vbase { - vbase: if let Some(max_scan_tuples) = max_scan_tuples() { - Box::new(vbase.take(max_scan_tuples as usize)) - } else { - Box::new(vbase) - }, - threshold: *threshold, - recheck: *recheck, - opfamily: *opfamily, - }; - } else { - *scanner = Scanner::Empty {}; - } - } - match scanner { - Scanner::Initial { .. } => unreachable!(), - Scanner::Vbase { - vbase, - threshold, - recheck, - opfamily, - } => match ( - vbase.next().map(|(d, p)| (opfamily.process(d), p)), - threshold, - ) { - (Some((_, ptr)), None) => Some((ptr, *recheck)), - (Some((distance, ptr)), Some(t)) if distance < *t => Some((ptr, *recheck)), - _ => { - let scanner = std::mem::replace(scanner, Scanner::Empty {}); - scan_release(scanner); - None - } - }, - Scanner::Empty {} => None, - } -} - -pub fn scan_release(scanner: Scanner) { - match scanner { - Scanner::Initial { .. } => {} - Scanner::Vbase { .. } => {} - Scanner::Empty {} => {} - } -} diff --git a/src/vchordrqfscan/index/functions.rs b/src/vchordrqfscan/index/functions.rs deleted file mode 100644 index 27bd9ac2..00000000 --- a/src/vchordrqfscan/index/functions.rs +++ /dev/null @@ -1,26 +0,0 @@ -use crate::postgres::PostgresRelation; -use crate::vchordrqfscan::algorithm::prewarm::prewarm; -use pgrx::pg_sys::Oid; -use pgrx_catalog::{PgAm, PgClass}; - -#[pgrx::pg_extern(sql = "")] -fn _vchordrqfscan_prewarm(indexrelid: Oid, height: i32) -> String { - let pg_am = PgAm::search_amname(c"vchordrqfscan").unwrap(); - let Some(pg_am) = pg_am.get() else { - pgrx::error!("vchord is not installed"); - }; - let pg_class = PgClass::search_reloid(indexrelid).unwrap(); - let Some(pg_class) = pg_class.get() else { - pgrx::error!("there is no such index"); - }; - if pg_class.relam() != pg_am.oid() { - pgrx::error!("{:?} is not a vchordrqfscan index", pg_class.relname()); - } - let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; - let relation = unsafe { PostgresRelation::new(index) }; - let message = prewarm(relation, height); - unsafe { - pgrx::pg_sys::index_close(index, pgrx::pg_sys::ShareLock as _); - } - message -} diff --git a/src/vchordrqfscan/index/mod.rs b/src/vchordrqfscan/index/mod.rs deleted file mode 100644 index 5203e4fb..00000000 --- a/src/vchordrqfscan/index/mod.rs +++ /dev/null @@ -1,12 +0,0 @@ -pub mod am; -pub mod am_options; -pub mod am_scan; -pub mod functions; -pub mod opclass; -pub mod utils; - -pub unsafe fn init() { - unsafe { - am::init(); - } -} diff --git a/src/vchordrqfscan/index/opclass.rs b/src/vchordrqfscan/index/opclass.rs deleted file mode 100644 index d095b9a6..00000000 --- a/src/vchordrqfscan/index/opclass.rs +++ /dev/null @@ -1,14 +0,0 @@ -#[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchordrqfscan_support_vector_l2_ops() -> String { - "vector_l2_ops".to_string() -} - -#[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchordrqfscan_support_vector_ip_ops() -> String { - "vector_ip_ops".to_string() -} - -#[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchordrqfscan_support_vector_cosine_ops() -> String { - "vector_cosine_ops".to_string() -} diff --git a/src/vchordrqfscan/index/utils.rs b/src/vchordrqfscan/index/utils.rs deleted file mode 100644 index 726a5979..00000000 --- a/src/vchordrqfscan/index/utils.rs +++ /dev/null @@ -1,20 +0,0 @@ -use std::num::NonZeroU64; - -pub fn pointer_to_ctid(pointer: NonZeroU64) -> pgrx::pg_sys::ItemPointerData { - let value = pointer.get(); - pgrx::pg_sys::ItemPointerData { - ip_blkid: pgrx::pg_sys::BlockIdData { - bi_hi: ((value >> 32) & 0xffff) as u16, - bi_lo: ((value >> 16) & 0xffff) as u16, - }, - ip_posid: (value & 0xffff) as u16, - } -} - -pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> NonZeroU64 { - let mut value = 0; - value |= (ctid.ip_blkid.bi_hi as u64) << 32; - value |= (ctid.ip_blkid.bi_lo as u64) << 16; - value |= ctid.ip_posid as u64; - NonZeroU64::new(value).expect("invalid pointer") -} diff --git a/src/vchordrqfscan/mod.rs b/src/vchordrqfscan/mod.rs deleted file mode 100644 index c2ae9456..00000000 --- a/src/vchordrqfscan/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -mod algorithm; -mod gucs; -mod index; -mod types; - -pub unsafe fn init() { - unsafe { - index::init(); - gucs::init(); - } -} diff --git a/src/vchordrqfscan/types.rs b/src/vchordrqfscan/types.rs deleted file mode 100644 index 91ca7690..00000000 --- a/src/vchordrqfscan/types.rs +++ /dev/null @@ -1,153 +0,0 @@ -use distance::Distance; -use serde::{Deserialize, Serialize}; -use validator::{Validate, ValidationError, ValidationErrors}; -use vector::vect::{VectBorrowed, VectOwned}; - -#[derive(Debug, Clone, Serialize, Deserialize, Validate)] -#[serde(deny_unknown_fields)] -pub struct VchordrqfscanInternalBuildOptions { - #[serde(default = "VchordrqfscanInternalBuildOptions::default_lists")] - #[validate(length(min = 1, max = 8), custom(function = VchordrqfscanInternalBuildOptions::validate_lists))] - pub lists: Vec, - #[serde(default = "VchordrqfscanInternalBuildOptions::default_spherical_centroids")] - pub spherical_centroids: bool, - #[serde(default = "VchordrqfscanInternalBuildOptions::default_sampling_factor")] - #[validate(range(min = 1, max = 1024))] - pub sampling_factor: u32, - #[serde(default = "VchordrqfscanInternalBuildOptions::default_build_threads")] - #[validate(range(min = 1, max = 255))] - pub build_threads: u16, -} - -impl VchordrqfscanInternalBuildOptions { - fn default_lists() -> Vec { - vec![1000] - } - fn validate_lists(lists: &[u32]) -> Result<(), ValidationError> { - if !lists.is_sorted() { - return Err(ValidationError::new("`lists` should be in ascending order")); - } - if !lists.iter().all(|x| (1..=1 << 24).contains(x)) { - return Err(ValidationError::new("list is too long or too short")); - } - Ok(()) - } - fn default_spherical_centroids() -> bool { - false - } - fn default_sampling_factor() -> u32 { - 256 - } - fn default_build_threads() -> u16 { - 1 - } -} - -impl Default for VchordrqfscanInternalBuildOptions { - fn default() -> Self { - Self { - lists: Self::default_lists(), - spherical_centroids: Self::default_spherical_centroids(), - sampling_factor: Self::default_sampling_factor(), - build_threads: Self::default_build_threads(), - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize, Validate)] -#[serde(deny_unknown_fields)] -pub struct VchordrqfscanExternalBuildOptions { - pub table: String, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -#[serde(rename_all = "snake_case")] -pub enum VchordrqfscanBuildOptions { - Internal(VchordrqfscanInternalBuildOptions), - External(VchordrqfscanExternalBuildOptions), -} - -impl Default for VchordrqfscanBuildOptions { - fn default() -> Self { - Self::Internal(Default::default()) - } -} - -impl Validate for VchordrqfscanBuildOptions { - fn validate(&self) -> Result<(), ValidationErrors> { - use VchordrqfscanBuildOptions::*; - match self { - Internal(internal_build) => internal_build.validate(), - External(external_build) => external_build.validate(), - } - } -} - -#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] -#[serde(deny_unknown_fields)] -pub struct VchordrqfscanIndexingOptions { - #[serde(default = "VchordrqfscanIndexingOptions::default_residual_quantization")] - pub residual_quantization: bool, - pub build: VchordrqfscanBuildOptions, -} - -impl VchordrqfscanIndexingOptions { - fn default_residual_quantization() -> bool { - false - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub enum OwnedVector { - Vecf32(VectOwned), -} - -#[derive(Debug, Clone, Copy)] -pub enum BorrowedVector<'a> { - Vecf32(VectBorrowed<'a, f32>), -} - -#[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub enum DistanceKind { - L2, - Dot, -} - -#[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub enum VectorKind { - Vecf32, -} - -#[derive(Debug, Clone, Serialize, Deserialize, Validate)] -#[serde(deny_unknown_fields)] -#[validate(schema(function = "Self::validate_self"))] -pub struct VectorOptions { - #[validate(range(min = 1, max = 1_048_575))] - #[serde(rename = "dimensions")] - pub dims: u32, - #[serde(rename = "vector")] - pub v: VectorKind, - #[serde(rename = "distance")] - pub d: DistanceKind, -} - -impl VectorOptions { - pub fn validate_self(&self) -> Result<(), ValidationError> { - match (self.v, self.d, self.dims) { - (VectorKind::Vecf32, DistanceKind::L2, 1..65536) => Ok(()), - (VectorKind::Vecf32, DistanceKind::Dot, 1..65536) => Ok(()), - _ => Err(ValidationError::new("not valid vector options")), - } - } -} - -pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance { - use simd::Floating; - match d { - DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)), - DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)), - } -} From fef3e28f6c43021cdc37ceb26b073a9912146e64 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 20 Jan 2025 16:27:30 +0800 Subject: [PATCH 097/324] fix: respect aliasing rule by not reading past of reference (#169) Signed-off-by: usamoi --- src/algorithm/build.rs | 4 +- src/datatype/functions_scalar8.rs | 8 +- src/datatype/memory_halfvec.rs | 210 ++++++++++++++++++ src/datatype/memory_pgvector_halfvec.rs | 205 ----------------- src/datatype/memory_pgvector_vector.rs | 204 ----------------- src/datatype/memory_scalar8.rs | 153 +++++++------ src/datatype/memory_vector.rs | 209 +++++++++++++++++ src/datatype/mod.rs | 8 +- ...vector_halfvec.rs => operators_halfvec.rs} | 14 +- ...pgvector_vector.rs => operators_vector.rs} | 14 +- src/index/am_options.rs | 16 +- 11 files changed, 531 insertions(+), 514 deletions(-) create mode 100644 src/datatype/memory_halfvec.rs delete mode 100644 src/datatype/memory_pgvector_halfvec.rs delete mode 100644 src/datatype/memory_pgvector_vector.rs create mode 100644 src/datatype/memory_vector.rs rename src/datatype/{operators_pgvector_halfvec.rs => operators_halfvec.rs} (83%) rename src/datatype/{operators_pgvector_vector.rs => operators_vector.rs} (83%) diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs index a0810a71..4c893b7f 100644 --- a/src/algorithm/build.rs +++ b/src/algorithm/build.rs @@ -217,7 +217,7 @@ impl Structure { let mut parents = BTreeMap::new(); let mut vectors = BTreeMap::new(); pgrx::spi::Spi::connect(|client| { - use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; + use crate::datatype::memory_vector::VectorOutput; use pgrx::pg_sys::panic::ErrorReportable; use vector::VectorBorrowed; let schema_query = "SELECT n.nspname::TEXT @@ -237,7 +237,7 @@ impl Structure { for row in centroids { let id: Option = row.get_by_name("id").unwrap(); let parent: Option = row.get_by_name("parent").unwrap(); - let vector: Option = row.get_by_name("vector").unwrap(); + let vector: Option = row.get_by_name("vector").unwrap(); let id = id.expect("external build: id could not be NULL"); let vector = vector.expect("external build: vector could not be NULL"); let pop = parents.insert(id, parent); diff --git a/src/datatype/functions_scalar8.rs b/src/datatype/functions_scalar8.rs index adcdb30c..bc73bd81 100644 --- a/src/datatype/functions_scalar8.rs +++ b/src/datatype/functions_scalar8.rs @@ -1,12 +1,12 @@ -use crate::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput; -use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; +use crate::datatype::memory_halfvec::HalfvecInput; use crate::datatype::memory_scalar8::Scalar8Output; +use crate::datatype::memory_vector::VectorInput; use half::f16; use simd::Floating; use vector::scalar8::Scalar8Borrowed; #[pgrx::pg_extern(sql = "")] -fn _vchord_vector_quantize_to_scalar8(vector: PgvectorVectorInput) -> Scalar8Output { +fn _vchord_vector_quantize_to_scalar8(vector: VectorInput) -> Scalar8Output { let vector = vector.as_borrowed(); let sum_of_x2 = f32::reduce_sum_of_x2(vector.slice()); let (k, b, code) = @@ -16,7 +16,7 @@ fn _vchord_vector_quantize_to_scalar8(vector: PgvectorVectorInput) -> Scalar8Out } #[pgrx::pg_extern(sql = "")] -fn _vchord_halfvec_quantize_to_scalar8(vector: PgvectorHalfvecInput) -> Scalar8Output { +fn _vchord_halfvec_quantize_to_scalar8(vector: HalfvecInput) -> Scalar8Output { let vector = vector.as_borrowed(); let sum_of_x2 = f16::reduce_sum_of_x2(vector.slice()); let (k, b, code) = diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs new file mode 100644 index 00000000..b60f6c5f --- /dev/null +++ b/src/datatype/memory_halfvec.rs @@ -0,0 +1,210 @@ +use half::f16; +use pgrx::datum::FromDatum; +use pgrx::datum::IntoDatum; +use pgrx::pg_sys::Datum; +use pgrx::pg_sys::Oid; +use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; +use pgrx::pgrx_sql_entity_graph::metadata::Returns; +use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; +use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; +use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; +use std::marker::PhantomData; +use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; + +#[repr(C, align(8))] +pub struct HalfvecHeader { + varlena: u32, + dims: u16, + unused: u16, + elements: [f16; 0], +} + +impl HalfvecHeader { + fn size_of(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + (size_of::() + size_of::() * len).next_multiple_of(8) + } + pub unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f16> { + unsafe { + let this = this.as_ptr(); + VectBorrowed::new_unchecked(std::slice::from_raw_parts( + (&raw const (*this).elements).cast(), + (&raw const (*this).dims).read() as usize, + )) + } + } +} + +pub struct HalfvecInput<'a>(NonNull, PhantomData<&'a ()>, bool); + +impl HalfvecInput<'_> { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() + }; + HalfvecInput(q, PhantomData, p != q) + } + pub fn as_borrowed(&self) -> VectBorrowed<'_, f16> { + unsafe { HalfvecHeader::as_borrowed(self.0) } + } +} + +impl Drop for HalfvecInput<'_> { + fn drop(&mut self) { + if self.2 { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } + } +} + +pub struct HalfvecOutput(NonNull); + +impl HalfvecOutput { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() + }; + Self(q) + } + #[allow(dead_code)] + pub fn new(vector: VectBorrowed<'_, f16>) -> Self { + unsafe { + let slice = vector.slice(); + let size = HalfvecHeader::size_of(slice.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut HalfvecHeader; + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dims).write(vector.dims() as _); + (&raw mut (*ptr).unused).write(0); + std::ptr::copy_nonoverlapping( + slice.as_ptr(), + (&raw mut (*ptr).elements).cast(), + slice.len(), + ); + Self(NonNull::new(ptr).unwrap()) + } + } + pub fn as_borrowed(&self) -> VectBorrowed<'_, f16> { + unsafe { HalfvecHeader::as_borrowed(self.0) } + } + pub fn into_raw(self) -> *mut HalfvecHeader { + let result = self.0.as_ptr(); + std::mem::forget(self); + result + } +} + +impl Drop for HalfvecOutput { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } +} + +// FromDatum + +impl FromDatum for HalfvecInput<'_> { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +impl FromDatum for HalfvecOutput { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +// IntoDatum + +impl IntoDatum for HalfvecOutput { + fn into_datum(self) -> Option { + Some(Datum::from(self.into_raw())) + } + + fn type_oid() -> Oid { + Oid::INVALID + } + + fn is_compatible_with(_: Oid) -> bool { + true + } +} + +// UnboxDatum + +unsafe impl pgrx::datum::UnboxDatum for HalfvecOutput { + type As<'src> = HalfvecOutput; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +// SqlTranslatable + +unsafe impl SqlTranslatable for HalfvecInput<'_> { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("halfvec"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("halfvec")))) + } +} + +unsafe impl SqlTranslatable for HalfvecOutput { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("halfvec"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("halfvec")))) + } +} + +// ArgAbi + +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for HalfvecInput<'fcx> { + unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { + let index = arg.index(); + unsafe { + arg.unbox_arg_using_from_datum() + .unwrap_or_else(|| panic!("argument {index} must not be null")) + } + } +} + +// BoxRet + +unsafe impl pgrx::callconv::BoxRet for HalfvecOutput { + unsafe fn box_into<'fcx>( + self, + fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, + ) -> pgrx::datum::Datum<'fcx> { + match self.into_datum() { + Some(datum) => unsafe { fcinfo.return_raw_datum(datum) }, + None => fcinfo.return_null(), + } + } +} diff --git a/src/datatype/memory_pgvector_halfvec.rs b/src/datatype/memory_pgvector_halfvec.rs deleted file mode 100644 index 1c065d94..00000000 --- a/src/datatype/memory_pgvector_halfvec.rs +++ /dev/null @@ -1,205 +0,0 @@ -use half::f16; -use pgrx::datum::FromDatum; -use pgrx::datum::IntoDatum; -use pgrx::pg_sys::Datum; -use pgrx::pg_sys::Oid; -use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; -use pgrx::pgrx_sql_entity_graph::metadata::Returns; -use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; -use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; -use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; -use std::ops::Deref; -use std::ptr::NonNull; -use vector::VectorBorrowed; -use vector::vect::VectBorrowed; - -#[repr(C, align(8))] -pub struct PgvectorHalfvecHeader { - varlena: u32, - dims: u16, - unused: u16, - phantom: [f16; 0], -} - -impl PgvectorHalfvecHeader { - fn size_of(len: usize) -> usize { - if len > 65535 { - panic!("vector is too large"); - } - (size_of::() + size_of::() * len).next_multiple_of(8) - } - pub fn as_borrowed(&self) -> VectBorrowed<'_, f16> { - unsafe { - VectBorrowed::new_unchecked(std::slice::from_raw_parts( - self.phantom.as_ptr(), - self.dims as usize, - )) - } - } -} - -pub enum PgvectorHalfvecInput<'a> { - Owned(PgvectorHalfvecOutput), - Borrowed(&'a PgvectorHalfvecHeader), -} - -impl PgvectorHalfvecInput<'_> { - unsafe fn new(p: NonNull) -> Self { - let q = unsafe { - NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() - }; - if p != q { - PgvectorHalfvecInput::Owned(PgvectorHalfvecOutput(q)) - } else { - unsafe { PgvectorHalfvecInput::Borrowed(p.as_ref()) } - } - } -} - -impl Deref for PgvectorHalfvecInput<'_> { - type Target = PgvectorHalfvecHeader; - - fn deref(&self) -> &Self::Target { - match self { - PgvectorHalfvecInput::Owned(x) => x, - PgvectorHalfvecInput::Borrowed(x) => x, - } - } -} - -pub struct PgvectorHalfvecOutput(NonNull); - -impl PgvectorHalfvecOutput { - pub fn new(vector: VectBorrowed<'_, f16>) -> PgvectorHalfvecOutput { - unsafe { - let slice = vector.slice(); - let size = PgvectorHalfvecHeader::size_of(slice.len()); - - let ptr = pgrx::pg_sys::palloc0(size) as *mut PgvectorHalfvecHeader; - (&raw mut (*ptr).varlena).write((size << 2) as u32); - (&raw mut (*ptr).dims).write(vector.dims() as _); - (&raw mut (*ptr).unused).write(0); - std::ptr::copy_nonoverlapping(slice.as_ptr(), (*ptr).phantom.as_mut_ptr(), slice.len()); - PgvectorHalfvecOutput(NonNull::new(ptr).unwrap()) - } - } - pub fn into_raw(self) -> *mut PgvectorHalfvecHeader { - let result = self.0.as_ptr(); - std::mem::forget(self); - result - } -} - -impl Deref for PgvectorHalfvecOutput { - type Target = PgvectorHalfvecHeader; - - fn deref(&self) -> &Self::Target { - unsafe { self.0.as_ref() } - } -} - -impl Drop for PgvectorHalfvecOutput { - fn drop(&mut self) { - unsafe { - pgrx::pg_sys::pfree(self.0.as_ptr() as _); - } - } -} - -impl FromDatum for PgvectorHalfvecInput<'_> { - unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { - if is_null { - None - } else { - let ptr = NonNull::new(datum.cast_mut_ptr::()).unwrap(); - unsafe { Some(PgvectorHalfvecInput::new(ptr)) } - } - } -} - -impl IntoDatum for PgvectorHalfvecOutput { - fn into_datum(self) -> Option { - Some(Datum::from(self.into_raw() as *mut ())) - } - - fn type_oid() -> Oid { - Oid::INVALID - } - - fn is_compatible_with(_: Oid) -> bool { - true - } -} - -impl FromDatum for PgvectorHalfvecOutput { - unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { - if is_null { - None - } else { - let p = NonNull::new(datum.cast_mut_ptr::())?; - let q = - unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast())? }; - if p != q { - Some(PgvectorHalfvecOutput(q)) - } else { - let header = p.as_ptr(); - let vector = unsafe { (*header).as_borrowed() }; - Some(PgvectorHalfvecOutput::new(vector)) - } - } - } -} - -unsafe impl pgrx::datum::UnboxDatum for PgvectorHalfvecOutput { - type As<'src> = PgvectorHalfvecOutput; - #[inline] - unsafe fn unbox<'src>(d: pgrx::datum::Datum<'src>) -> Self::As<'src> - where - Self: 'src, - { - let p = NonNull::new(d.sans_lifetime().cast_mut_ptr::()).unwrap(); - let q = unsafe { - NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() - }; - if p != q { - PgvectorHalfvecOutput(q) - } else { - let header = p.as_ptr(); - let vector = unsafe { (*header).as_borrowed() }; - PgvectorHalfvecOutput::new(vector) - } - } -} - -unsafe impl SqlTranslatable for PgvectorHalfvecInput<'_> { - fn argument_sql() -> Result { - Ok(SqlMapping::As(String::from("halfvec"))) - } - fn return_sql() -> Result { - Ok(Returns::One(SqlMapping::As(String::from("halfvec")))) - } -} - -unsafe impl SqlTranslatable for PgvectorHalfvecOutput { - fn argument_sql() -> Result { - Ok(SqlMapping::As(String::from("halfvec"))) - } - fn return_sql() -> Result { - Ok(Returns::One(SqlMapping::As(String::from("halfvec")))) - } -} - -unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for PgvectorHalfvecInput<'fcx> { - unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { - unsafe { arg.unbox_arg_using_from_datum().unwrap() } - } -} - -unsafe impl pgrx::callconv::BoxRet for PgvectorHalfvecOutput { - unsafe fn box_into<'fcx>( - self, - fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, - ) -> pgrx::datum::Datum<'fcx> { - unsafe { fcinfo.return_raw_datum(Datum::from(self.into_raw() as *mut ())) } - } -} diff --git a/src/datatype/memory_pgvector_vector.rs b/src/datatype/memory_pgvector_vector.rs deleted file mode 100644 index e3ab9f90..00000000 --- a/src/datatype/memory_pgvector_vector.rs +++ /dev/null @@ -1,204 +0,0 @@ -use pgrx::datum::FromDatum; -use pgrx::datum::IntoDatum; -use pgrx::pg_sys::Datum; -use pgrx::pg_sys::Oid; -use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; -use pgrx::pgrx_sql_entity_graph::metadata::Returns; -use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; -use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; -use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; -use std::ops::Deref; -use std::ptr::NonNull; -use vector::VectorBorrowed; -use vector::vect::VectBorrowed; - -#[repr(C, align(8))] -pub struct PgvectorVectorHeader { - varlena: u32, - dims: u16, - unused: u16, - phantom: [f32; 0], -} - -impl PgvectorVectorHeader { - fn size_of(len: usize) -> usize { - if len > 65535 { - panic!("vector is too large"); - } - (size_of::() + size_of::() * len).next_multiple_of(8) - } - pub fn as_borrowed(&self) -> VectBorrowed<'_, f32> { - unsafe { - VectBorrowed::new_unchecked(std::slice::from_raw_parts( - self.phantom.as_ptr(), - self.dims as usize, - )) - } - } -} - -pub enum PgvectorVectorInput<'a> { - Owned(PgvectorVectorOutput), - Borrowed(&'a PgvectorVectorHeader), -} - -impl PgvectorVectorInput<'_> { - unsafe fn new(p: NonNull) -> Self { - let q = unsafe { - NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() - }; - if p != q { - PgvectorVectorInput::Owned(PgvectorVectorOutput(q)) - } else { - unsafe { PgvectorVectorInput::Borrowed(p.as_ref()) } - } - } -} - -impl Deref for PgvectorVectorInput<'_> { - type Target = PgvectorVectorHeader; - - fn deref(&self) -> &Self::Target { - match self { - PgvectorVectorInput::Owned(x) => x, - PgvectorVectorInput::Borrowed(x) => x, - } - } -} - -pub struct PgvectorVectorOutput(NonNull); - -impl PgvectorVectorOutput { - pub fn new(vector: VectBorrowed<'_, f32>) -> PgvectorVectorOutput { - unsafe { - let slice = vector.slice(); - let size = PgvectorVectorHeader::size_of(slice.len()); - - let ptr = pgrx::pg_sys::palloc0(size) as *mut PgvectorVectorHeader; - (&raw mut (*ptr).varlena).write((size << 2) as u32); - (&raw mut (*ptr).dims).write(vector.dims() as _); - (&raw mut (*ptr).unused).write(0); - std::ptr::copy_nonoverlapping(slice.as_ptr(), (*ptr).phantom.as_mut_ptr(), slice.len()); - PgvectorVectorOutput(NonNull::new(ptr).unwrap()) - } - } - pub fn into_raw(self) -> *mut PgvectorVectorHeader { - let result = self.0.as_ptr(); - std::mem::forget(self); - result - } -} - -impl Deref for PgvectorVectorOutput { - type Target = PgvectorVectorHeader; - - fn deref(&self) -> &Self::Target { - unsafe { self.0.as_ref() } - } -} - -impl Drop for PgvectorVectorOutput { - fn drop(&mut self) { - unsafe { - pgrx::pg_sys::pfree(self.0.as_ptr() as _); - } - } -} - -impl FromDatum for PgvectorVectorInput<'_> { - unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { - if is_null { - None - } else { - let ptr = NonNull::new(datum.cast_mut_ptr::()).unwrap(); - unsafe { Some(PgvectorVectorInput::new(ptr)) } - } - } -} - -impl IntoDatum for PgvectorVectorOutput { - fn into_datum(self) -> Option { - Some(Datum::from(self.into_raw() as *mut ())) - } - - fn type_oid() -> Oid { - Oid::INVALID - } - - fn is_compatible_with(_: Oid) -> bool { - true - } -} - -impl FromDatum for PgvectorVectorOutput { - unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { - if is_null { - None - } else { - let p = NonNull::new(datum.cast_mut_ptr::())?; - let q = - unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast())? }; - if p != q { - Some(PgvectorVectorOutput(q)) - } else { - let header = p.as_ptr(); - let vector = unsafe { (*header).as_borrowed() }; - Some(PgvectorVectorOutput::new(vector)) - } - } - } -} - -unsafe impl pgrx::datum::UnboxDatum for PgvectorVectorOutput { - type As<'src> = PgvectorVectorOutput; - #[inline] - unsafe fn unbox<'src>(d: pgrx::datum::Datum<'src>) -> Self::As<'src> - where - Self: 'src, - { - let p = NonNull::new(d.sans_lifetime().cast_mut_ptr::()).unwrap(); - let q = unsafe { - NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() - }; - if p != q { - PgvectorVectorOutput(q) - } else { - let header = p.as_ptr(); - let vector = unsafe { (*header).as_borrowed() }; - PgvectorVectorOutput::new(vector) - } - } -} - -unsafe impl SqlTranslatable for PgvectorVectorInput<'_> { - fn argument_sql() -> Result { - Ok(SqlMapping::As(String::from("vector"))) - } - fn return_sql() -> Result { - Ok(Returns::One(SqlMapping::As(String::from("vector")))) - } -} - -unsafe impl SqlTranslatable for PgvectorVectorOutput { - fn argument_sql() -> Result { - Ok(SqlMapping::As(String::from("vector"))) - } - fn return_sql() -> Result { - Ok(Returns::One(SqlMapping::As(String::from("vector")))) - } -} - -unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for PgvectorVectorInput<'fcx> { - unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { - unsafe { arg.unbox_arg_using_from_datum().unwrap() } - } -} - -unsafe impl pgrx::callconv::BoxRet for PgvectorVectorOutput { - unsafe fn box_into<'fcx>( - self, - fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, - ) -> pgrx::datum::Datum<'fcx> { - unsafe { fcinfo.return_raw_datum(Datum::from(self.into_raw() as *mut ())) } - } -} diff --git a/src/datatype/memory_scalar8.rs b/src/datatype/memory_scalar8.rs index 1641c63d..4f306548 100644 --- a/src/datatype/memory_scalar8.rs +++ b/src/datatype/memory_scalar8.rs @@ -7,7 +7,7 @@ use pgrx::pgrx_sql_entity_graph::metadata::Returns; use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; -use std::ops::Deref; +use std::marker::PhantomData; use std::ptr::NonNull; use vector::VectorBorrowed; use vector::scalar8::Scalar8Borrowed; @@ -21,7 +21,7 @@ pub struct Scalar8Header { k: f32, b: f32, sum_of_code: f32, - phantom: [u8; 0], + elements: [u8; 0], } impl Scalar8Header { @@ -31,44 +31,43 @@ impl Scalar8Header { } (size_of::() + size_of::() * len).next_multiple_of(8) } - pub fn as_borrowed(&self) -> Scalar8Borrowed<'_> { + pub unsafe fn as_borrowed<'a>(this: NonNull) -> Scalar8Borrowed<'a> { unsafe { + let this = this.as_ptr(); Scalar8Borrowed::new_unchecked( - self.sum_of_x2, - self.k, - self.b, - self.sum_of_code, - std::slice::from_raw_parts(self.phantom.as_ptr(), self.dims as usize), + (&raw const (*this).sum_of_x2).read(), + (&raw const (*this).k).read(), + (&raw const (*this).b).read(), + (&raw const (*this).sum_of_code).read(), + std::slice::from_raw_parts( + (&raw const (*this).elements).cast(), + (&raw const (*this).dims).read() as usize, + ), ) } } } -pub enum Scalar8Input<'a> { - Owned(Scalar8Output), - Borrowed(&'a Scalar8Header), -} +pub struct Scalar8Input<'a>(NonNull, PhantomData<&'a ()>, bool); impl Scalar8Input<'_> { - unsafe fn new(p: NonNull) -> Self { + unsafe fn from_ptr(p: NonNull) -> Self { let q = unsafe { - NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() }; - if p != q { - Scalar8Input::Owned(Scalar8Output(q)) - } else { - unsafe { Scalar8Input::Borrowed(p.as_ref()) } - } + Scalar8Input(q, PhantomData, p != q) + } + pub fn as_borrowed(&self) -> Scalar8Borrowed<'_> { + unsafe { Scalar8Header::as_borrowed(self.0) } } } -impl Deref for Scalar8Input<'_> { - type Target = Scalar8Header; - - fn deref(&self) -> &Self::Target { - match self { - Scalar8Input::Owned(x) => x, - Scalar8Input::Borrowed(x) => x, +impl Drop for Scalar8Input<'_> { + fn drop(&mut self) { + if self.2 { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } } } } @@ -76,7 +75,13 @@ impl Deref for Scalar8Input<'_> { pub struct Scalar8Output(NonNull); impl Scalar8Output { - pub fn new(vector: Scalar8Borrowed<'_>) -> Scalar8Output { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() + }; + Self(q) + } + pub fn new(vector: Scalar8Borrowed<'_>) -> Self { unsafe { let code = vector.code(); let size = Scalar8Header::size_of(code.len()); @@ -89,10 +94,17 @@ impl Scalar8Output { (&raw mut (*ptr).k).write(vector.k()); (&raw mut (*ptr).b).write(vector.b()); (&raw mut (*ptr).sum_of_code).write(vector.sum_of_code()); - std::ptr::copy_nonoverlapping(code.as_ptr(), (*ptr).phantom.as_mut_ptr(), code.len()); - Scalar8Output(NonNull::new(ptr).unwrap()) + std::ptr::copy_nonoverlapping( + code.as_ptr(), + (&raw mut (*ptr).elements).cast(), + code.len(), + ); + Self(NonNull::new(ptr).unwrap()) } } + pub fn as_borrowed(&self) -> Scalar8Borrowed<'_> { + unsafe { Scalar8Header::as_borrowed(self.0) } + } pub fn into_raw(self) -> *mut Scalar8Header { let result = self.0.as_ptr(); std::mem::forget(self); @@ -100,36 +112,43 @@ impl Scalar8Output { } } -impl Deref for Scalar8Output { - type Target = Scalar8Header; - - fn deref(&self) -> &Self::Target { - unsafe { self.0.as_ref() } - } -} - impl Drop for Scalar8Output { fn drop(&mut self) { unsafe { - pgrx::pg_sys::pfree(self.0.as_ptr() as _); + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); } } } +// FromDatum + impl FromDatum for Scalar8Input<'_> { unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { if is_null { None } else { - let ptr = NonNull::new(datum.cast_mut_ptr::()).unwrap(); - unsafe { Some(Scalar8Input::new(ptr)) } + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +impl FromDatum for Scalar8Output { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } } } } +// IntoDatum + impl IntoDatum for Scalar8Output { fn into_datum(self) -> Option { - Some(Datum::from(self.into_raw() as *mut ())) + Some(Datum::from(self.into_raw())) } fn type_oid() -> Oid { @@ -141,46 +160,23 @@ impl IntoDatum for Scalar8Output { } } -impl FromDatum for Scalar8Output { - unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { - if is_null { - None - } else { - let p = NonNull::new(datum.cast_mut_ptr::())?; - let q = - unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast())? }; - if p != q { - Some(Scalar8Output(q)) - } else { - let header = p.as_ptr(); - let vector = unsafe { (*header).as_borrowed() }; - Some(Scalar8Output::new(vector)) - } - } - } -} +// UnboxDatum unsafe impl pgrx::datum::UnboxDatum for Scalar8Output { type As<'src> = Scalar8Output; #[inline] - unsafe fn unbox<'src>(d: pgrx::datum::Datum<'src>) -> Self::As<'src> + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> where Self: 'src, { - let p = NonNull::new(d.sans_lifetime().cast_mut_ptr::()).unwrap(); - let q = unsafe { - NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() - }; - if p != q { - Scalar8Output(q) - } else { - let header = p.as_ptr(); - let vector = unsafe { (*header).as_borrowed() }; - Scalar8Output::new(vector) - } + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } } } +// SqlTranslatable + unsafe impl SqlTranslatable for Scalar8Input<'_> { fn argument_sql() -> Result { Ok(SqlMapping::As(String::from("scalar8"))) @@ -199,17 +195,28 @@ unsafe impl SqlTranslatable for Scalar8Output { } } +// ArgAbi + unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Scalar8Input<'fcx> { unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { - unsafe { arg.unbox_arg_using_from_datum().unwrap() } + let index = arg.index(); + unsafe { + arg.unbox_arg_using_from_datum() + .unwrap_or_else(|| panic!("argument {index} must not be null")) + } } } +// BoxRet + unsafe impl pgrx::callconv::BoxRet for Scalar8Output { unsafe fn box_into<'fcx>( self, fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, ) -> pgrx::datum::Datum<'fcx> { - unsafe { fcinfo.return_raw_datum(Datum::from(self.into_raw() as *mut ())) } + match self.into_datum() { + Some(datum) => unsafe { fcinfo.return_raw_datum(datum) }, + None => fcinfo.return_null(), + } } } diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs new file mode 100644 index 00000000..de70ba16 --- /dev/null +++ b/src/datatype/memory_vector.rs @@ -0,0 +1,209 @@ +use pgrx::datum::FromDatum; +use pgrx::datum::IntoDatum; +use pgrx::pg_sys::Datum; +use pgrx::pg_sys::Oid; +use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; +use pgrx::pgrx_sql_entity_graph::metadata::Returns; +use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; +use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; +use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; +use std::marker::PhantomData; +use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; + +#[repr(C, align(8))] +pub struct VectorHeader { + varlena: u32, + dims: u16, + unused: u16, + elements: [f32; 0], +} + +impl VectorHeader { + fn size_of(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + (size_of::() + size_of::() * len).next_multiple_of(8) + } + pub unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f32> { + unsafe { + let this = this.as_ptr(); + VectBorrowed::new_unchecked(std::slice::from_raw_parts( + (&raw const (*this).elements).cast(), + (&raw const (*this).dims).read() as usize, + )) + } + } +} + +pub struct VectorInput<'a>(NonNull, PhantomData<&'a ()>, bool); + +impl VectorInput<'_> { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() + }; + VectorInput(q, PhantomData, p != q) + } + pub fn as_borrowed(&self) -> VectBorrowed<'_, f32> { + unsafe { VectorHeader::as_borrowed(self.0) } + } +} + +impl Drop for VectorInput<'_> { + fn drop(&mut self) { + if self.2 { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } + } +} + +pub struct VectorOutput(NonNull); + +impl VectorOutput { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() + }; + Self(q) + } + #[allow(dead_code)] + pub fn new(vector: VectBorrowed<'_, f32>) -> Self { + unsafe { + let slice = vector.slice(); + let size = VectorHeader::size_of(slice.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut VectorHeader; + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dims).write(vector.dims() as _); + (&raw mut (*ptr).unused).write(0); + std::ptr::copy_nonoverlapping( + slice.as_ptr(), + (&raw mut (*ptr).elements).cast(), + slice.len(), + ); + Self(NonNull::new(ptr).unwrap()) + } + } + pub fn as_borrowed(&self) -> VectBorrowed<'_, f32> { + unsafe { VectorHeader::as_borrowed(self.0) } + } + pub fn into_raw(self) -> *mut VectorHeader { + let result = self.0.as_ptr(); + std::mem::forget(self); + result + } +} + +impl Drop for VectorOutput { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } +} + +// FromDatum + +impl FromDatum for VectorInput<'_> { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +impl FromDatum for VectorOutput { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +// IntoDatum + +impl IntoDatum for VectorOutput { + fn into_datum(self) -> Option { + Some(Datum::from(self.into_raw())) + } + + fn type_oid() -> Oid { + Oid::INVALID + } + + fn is_compatible_with(_: Oid) -> bool { + true + } +} + +// UnboxDatum + +unsafe impl pgrx::datum::UnboxDatum for VectorOutput { + type As<'src> = VectorOutput; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +// SqlTranslatable + +unsafe impl SqlTranslatable for VectorInput<'_> { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("vector"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("vector")))) + } +} + +unsafe impl SqlTranslatable for VectorOutput { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("vector"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("vector")))) + } +} + +// ArgAbi + +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for VectorInput<'fcx> { + unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { + let index = arg.index(); + unsafe { + arg.unbox_arg_using_from_datum() + .unwrap_or_else(|| panic!("argument {index} must not be null")) + } + } +} + +// BoxAbi + +unsafe impl pgrx::callconv::BoxRet for VectorOutput { + unsafe fn box_into<'fcx>( + self, + fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, + ) -> pgrx::datum::Datum<'fcx> { + match self.into_datum() { + Some(datum) => unsafe { fcinfo.return_raw_datum(datum) }, + None => fcinfo.return_null(), + } + } +} diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 98b6650e..cd08970c 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -1,10 +1,10 @@ pub mod binary_scalar8; pub mod functions_scalar8; -pub mod memory_pgvector_halfvec; -pub mod memory_pgvector_vector; +pub mod memory_halfvec; pub mod memory_scalar8; -pub mod operators_pgvector_halfvec; -pub mod operators_pgvector_vector; +pub mod memory_vector; +pub mod operators_halfvec; pub mod operators_scalar8; +pub mod operators_vector; pub mod text_scalar8; pub mod typmod; diff --git a/src/datatype/operators_pgvector_halfvec.rs b/src/datatype/operators_halfvec.rs similarity index 83% rename from src/datatype/operators_pgvector_halfvec.rs rename to src/datatype/operators_halfvec.rs index fb0492ac..907d415d 100644 --- a/src/datatype/operators_pgvector_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -1,14 +1,14 @@ -use crate::datatype::memory_pgvector_halfvec::{PgvectorHalfvecInput, PgvectorHalfvecOutput}; +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; use std::num::NonZero; use vector::VectorBorrowed; use vector::vect::VectBorrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_halfvec_sphere_l2_in( - lhs: PgvectorHalfvecInput<'_>, + lhs: HalfvecInput<'_>, rhs: pgrx::composite_type!("sphere_halfvec"), ) -> bool { - let center: PgvectorHalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + let center: HalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), @@ -29,10 +29,10 @@ fn _vchord_halfvec_sphere_l2_in( #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_halfvec_sphere_ip_in( - lhs: PgvectorHalfvecInput<'_>, + lhs: HalfvecInput<'_>, rhs: pgrx::composite_type!("sphere_halfvec"), ) -> bool { - let center: PgvectorHalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + let center: HalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), @@ -53,10 +53,10 @@ fn _vchord_halfvec_sphere_ip_in( #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_halfvec_sphere_cosine_in( - lhs: PgvectorHalfvecInput<'_>, + lhs: HalfvecInput<'_>, rhs: pgrx::composite_type!("sphere_halfvec"), ) -> bool { - let center: PgvectorHalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + let center: HalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), diff --git a/src/datatype/operators_pgvector_vector.rs b/src/datatype/operators_vector.rs similarity index 83% rename from src/datatype/operators_pgvector_vector.rs rename to src/datatype/operators_vector.rs index e6b4be10..b011c3ae 100644 --- a/src/datatype/operators_pgvector_vector.rs +++ b/src/datatype/operators_vector.rs @@ -1,14 +1,14 @@ -use crate::datatype::memory_pgvector_vector::{PgvectorVectorInput, PgvectorVectorOutput}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use std::num::NonZero; use vector::VectorBorrowed; use vector::vect::VectBorrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_vector_sphere_l2_in( - lhs: PgvectorVectorInput<'_>, + lhs: VectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { - let center: PgvectorVectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + let center: VectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), @@ -29,10 +29,10 @@ fn _vchord_vector_sphere_l2_in( #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_vector_sphere_ip_in( - lhs: PgvectorVectorInput<'_>, + lhs: VectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { - let center: PgvectorVectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + let center: VectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), @@ -53,10 +53,10 @@ fn _vchord_vector_sphere_ip_in( #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchord_vector_sphere_cosine_in( - lhs: PgvectorVectorInput<'_>, + lhs: VectorInput<'_>, rhs: pgrx::composite_type!("sphere_vector"), ) -> bool { - let center: PgvectorVectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + let center: VectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), diff --git a/src/index/am_options.rs b/src/index/am_options.rs index 34c76a9e..06ca8e50 100644 --- a/src/index/am_options.rs +++ b/src/index/am_options.rs @@ -1,7 +1,7 @@ -use crate::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput; -use crate::datatype::memory_pgvector_halfvec::PgvectorHalfvecOutput; -use crate::datatype::memory_pgvector_vector::PgvectorVectorInput; -use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput; +use crate::datatype::memory_halfvec::HalfvecInput; +use crate::datatype::memory_halfvec::HalfvecOutput; +use crate::datatype::memory_vector::VectorInput; +use crate::datatype::memory_vector::VectorOutput; use crate::datatype::typmod::Typmod; use crate::types::{BorrowedVector, OwnedVector}; use crate::types::{DistanceKind, VectorKind}; @@ -133,11 +133,11 @@ impl Opfamily { } let vector = match self.vector { VectorKind::Vecf32 => { - let vector = unsafe { PgvectorVectorInput::from_datum(datum, false).unwrap() }; + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed())) } VectorKind::Vecf16 => { - let vector = unsafe { PgvectorHalfvecInput::from_datum(datum, false).unwrap() }; + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; self.preprocess(BorrowedVector::Vecf16(vector.as_borrowed())) } }; @@ -154,11 +154,11 @@ impl Opfamily { let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; let center = match self.vector { VectorKind::Vecf32 => tuple - .get_by_index::(NonZero::new(1).unwrap()) + .get_by_index::(NonZero::new(1).unwrap()) .unwrap() .map(|vector| self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed()))), VectorKind::Vecf16 => tuple - .get_by_index::(NonZero::new(1).unwrap()) + .get_by_index::(NonZero::new(1).unwrap()) .unwrap() .map(|vector| self.preprocess(BorrowedVector::Vecf16(vector.as_borrowed()))), }; From 94e23616fb96e6bc5b7629c8abd0e229088d6761 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 22 Jan 2025 16:41:10 +0800 Subject: [PATCH 098/324] fix: correct output of prewarm (#173) https://github.com/tensorchord/VectorChord/issues/171#issuecomment-2604327460 Signed-off-by: usamoi --- src/algorithm/prewarm.rs | 57 ++++++++++++++++++++++++++-------------- 1 file changed, 37 insertions(+), 20 deletions(-) diff --git a/src/algorithm/prewarm.rs b/src/algorithm/prewarm.rs index 9373f4a9..ab2731fa 100644 --- a/src/algorithm/prewarm.rs +++ b/src/algorithm/prewarm.rs @@ -21,26 +21,29 @@ pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> } type State = Vec; let mut state: State = { - let mut results = Vec::new(); - let counter = 1_usize; + let mut nodes = Vec::new(); { vectors::vector_access_1::(relation.clone(), root_mean, ()); - results.push(root_first); + nodes.push(root_first); } - writeln!(message, "number of tuples: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); - results + writeln!(message, "------------------------").unwrap(); + writeln!(message, "number of nodes: {}", nodes.len()).unwrap(); + writeln!(message, "number of tuples: {}", 1).unwrap(); + writeln!(message, "number of pages: {}", 1).unwrap(); + nodes }; let mut step = |state: State| { - let mut counter = 0_usize; - let mut results = Vec::new(); + let mut counter_pages = 0_usize; + let mut counter_tuples = 0_usize; + let mut nodes = Vec::new(); for list in state { let mut current = list; while current != u32::MAX { - counter += 1; + counter_pages += 1; pgrx::check_for_interrupts!(); let h1_guard = relation.read(current); for i in 1..=h1_guard.len() { + counter_tuples += 1; let h1_tuple = h1_guard .get(i) .expect("data corruption") @@ -51,7 +54,7 @@ pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> vectors::vector_access_1::(relation.clone(), mean, ()); } for first in h1_tuple.first().iter().copied() { - results.push(first); + nodes.push(first); } } H1TupleReader::_1(_) => (), @@ -60,16 +63,19 @@ pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> current = h1_guard.get_opaque().next; } } - writeln!(message, "number of tuples: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); - results + writeln!(message, "------------------------").unwrap(); + writeln!(message, "number of nodes: {}", nodes.len()).unwrap(); + writeln!(message, "number of tuples: {}", counter_tuples).unwrap(); + writeln!(message, "number of pages: {}", counter_pages).unwrap(); + nodes }; for _ in (std::cmp::max(1, prewarm_max_height)..height_of_root).rev() { state = step(state); } if prewarm_max_height == 0 { - let mut counter = 0_usize; - let mut results = Vec::new(); + let mut counter_pages = 0_usize; + let mut counter_tuples = 0_usize; + let mut counter_nodes = 0_usize; for list in state { let jump_guard = relation.read(list); let jump_tuple = jump_guard @@ -79,21 +85,32 @@ pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> let first = jump_tuple.first(); let mut current = first; while current != u32::MAX { - counter += 1; + counter_pages += 1; pgrx::check_for_interrupts!(); let h0_guard = relation.read(current); for i in 1..=h0_guard.len() { - let _h0_tuple = h0_guard + counter_tuples += 1; + let h0_tuple = h0_guard .get(i) .expect("data corruption") .pipe(read_tuple::); - results.push(()); + match h0_tuple { + H0TupleReader::_0(_h0_tuple) => { + counter_nodes += 1; + } + H0TupleReader::_1(_h0_tuple) => { + counter_nodes += 32; + } + H0TupleReader::_2(_) => (), + } } current = h0_guard.get_opaque().next; } } - writeln!(message, "number of tuples: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); + writeln!(message, "------------------------").unwrap(); + writeln!(message, "number of nodes: {}", counter_nodes).unwrap(); + writeln!(message, "number of tuples: {}", counter_tuples).unwrap(); + writeln!(message, "number of pages: {}", counter_pages).unwrap(); } message } From 44827d2ba0bb8f5d8af55ee3ea3931115597b45c Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 23 Jan 2025 17:07:09 +0800 Subject: [PATCH 099/324] fix: add magic number and version number to meta tuple (#174) Signed-off-by: usamoi --- src/algorithm/tuples.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/algorithm/tuples.rs b/src/algorithm/tuples.rs index 6ecef02e..06d7cedb 100644 --- a/src/algorithm/tuples.rs +++ b/src/algorithm/tuples.rs @@ -5,6 +5,8 @@ use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; +const MAGIC: u64 = u64::from_ne_bytes(*b"vchordrq"); +const VERSION: u64 = 1; pub trait Tuple: 'static { type Reader<'a>: TupleReader<'a, Tuple = Self>; @@ -42,6 +44,8 @@ pub fn write_tuple(source: &mut [u8]) -> T::Writer<'_> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct MetaTupleHeader { + magic: u64, + version: u64, dims: u32, height_of_root: u32, is_residual: Bool, @@ -69,6 +73,8 @@ impl Tuple for MetaTuple { fn serialize(&self) -> Vec { MetaTupleHeader { + magic: MAGIC, + version: VERSION, dims: self.dims, height_of_root: self.height_of_root, is_residual: self.is_residual.into(), @@ -91,6 +97,17 @@ pub struct MetaTupleReader<'a> { impl<'a> TupleReader<'a> for MetaTupleReader<'a> { type Tuple = MetaTuple; fn deserialize_ref(source: &'a [u8]) -> Self { + if source.len() < 16 { + panic!("bad bytes") + } + let magic = u64::from_ne_bytes(std::array::from_fn(|i| source[i + 0])); + if magic != MAGIC { + panic!("bad magic number"); + } + let version = u64::from_ne_bytes(std::array::from_fn(|i| source[i + 8])); + if version != VERSION { + panic!("bad version number"); + } let checker = RefChecker::new(source); let header = checker.prefix(0); Self { header } From af38f7c29a0522dcb899ff6e9add2a56465dd70a Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Mon, 27 Jan 2025 10:12:10 +0800 Subject: [PATCH 100/324] Release 0.2.0 (#177) Close #175 ## Upgrade - `sql/*`: schema for 0.1.0, 0.2.0 and 0.1.0->0.2.0 - `tools/dev.md`: document about how to release new version in the future ## Chore - `scripts/bench.py`: additional argument `-t` for choosing top 10/100 - `scripts/index.py`: support more than 8 cores for server in external index build --------- Signed-off-by: cutecutecat --- scripts/README.md | 2 +- scripts/bench.py | 144 ++++---- scripts/index.py | 6 +- sql/install/vchord--0.1.0.sql | 164 +++++++++ sql/install/vchord--0.2.0.sql | 500 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.1.0--0.2.0.sql | 343 ++++++++++++++++++ tools/dev.md | 87 +++++ tools/dump-codegen.sh | 44 +++ tools/dump.sh | 44 +++ tools/package.sh | 1 + 10 files changed, 1259 insertions(+), 76 deletions(-) create mode 100644 sql/install/vchord--0.1.0.sql create mode 100644 sql/install/vchord--0.2.0.sql create mode 100644 sql/upgrade/vchord--0.1.0--0.2.0.sql create mode 100644 tools/dev.md create mode 100755 tools/dump-codegen.sh create mode 100755 tools/dump.sh diff --git a/scripts/README.md b/scripts/README.md index 9ed71720..56ca24ee 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -55,7 +55,7 @@ docker run --name vchord -e POSTGRES_PASSWORD=123 -p 5432:5432 -d vchord:pg16-la # When using CPU to train k-means clustering conda install conda-forge::pgvector-python numpy pytorch::faiss-cpu conda-forge::psycopg h5py tqdm # or -pip install pgvector-python numpy faiss-cpu psycopg h5py tqdm +pip install pgvector numpy faiss-cpu psycopg h5py tqdm # When using GPU to train k-means clustering conda install conda-forge::pgvector-python numpy pytorch::faiss-gpu conda-forge::psycopg h5py tqdm diff --git a/scripts/bench.py b/scripts/bench.py index 0d7670cc..4e6eca5e 100644 --- a/scripts/bench.py +++ b/scripts/bench.py @@ -9,8 +9,6 @@ import h5py from pgvector.psycopg import register_vector -TOP = [10] - def build_arg_parse(): parser = argparse.ArgumentParser() @@ -26,6 +24,9 @@ def build_arg_parse(): parser.add_argument( "-p", "--password", help="Database password", default="password" ) + parser.add_argument( + "-t", "--top", help="Dimension", type=int, choices=[10, 100], default=10 + ) parser.add_argument( "--nprob", help="argument probes for query", default=100, type=int ) @@ -116,90 +117,86 @@ def calculate_metrics(all_results, k, m): def parallel_bench( - name, test, answer, metric_ops, num_processes, password, nprob, epsilon + name, test, answer, metric_ops, num_processes, password, top, nprob, epsilon ): """Run benchmark in parallel using multiple processes""" m = test.shape[0] - for k in TOP: - # Split data into batches for each process - batch_size = m // num_processes - batches = [] - - for i in range(num_processes): - start_idx = i * batch_size - end_idx = start_idx + batch_size if i < num_processes - 1 else m - - batch = ( - test[start_idx:end_idx], - answer[start_idx:end_idx], - k, - metric_ops, - password, - name, - nprob, - epsilon, - ) - batches.append(batch) - - # Create process pool and execute batches - with mp.Pool(processes=num_processes) as pool: - batch_results = list( - tqdm( - pool.imap(process_batch, batches), - total=len(batches), - desc=f"Processing k={k}", - ) + # Split data into batches for each process + batch_size = m // num_processes + batches = [] + + for i in range(num_processes): + start_idx = i * batch_size + end_idx = start_idx + batch_size if i < num_processes - 1 else m + + batch = ( + test[start_idx:end_idx], + answer[start_idx:end_idx], + top, + metric_ops, + password, + name, + nprob, + epsilon, + ) + batches.append(batch) + + # Create process pool and execute batches + with mp.Pool(processes=num_processes) as pool: + batch_results = list( + tqdm( + pool.imap(process_batch, batches), + total=len(batches), + desc=f"Processing k={top}", ) + ) - # Flatten results from all batches - all_results = [result for batch in batch_results for result in batch] + # Flatten results from all batches + all_results = [result for batch in batch_results for result in batch] - # Calculate metrics - recall, qps, p50, p99 = calculate_metrics(all_results, k, m) + # Calculate metrics + recall, qps, p50, p99 = calculate_metrics(all_results, top, m) - print(f"Top: {k}") - print(f" Recall: {recall:.4f}") - print(f" QPS: {qps*num_processes:.2f}") - print(f" P50 latency: {p50:.2f}ms") - print(f" P99 latency: {p99:.2f}ms") + print(f"Top: {top}") + print(f" Recall: {recall:.4f}") + print(f" QPS: {qps*num_processes:.2f}") + print(f" P50 latency: {p50:.2f}ms") + print(f" P99 latency: {p99:.2f}ms") -def sequential_bench(name, test, answer, metric_ops, conn): +def sequential_bench(name, test, answer, metric_ops, conn, top): """Original sequential benchmark implementation with latency tracking""" m = test.shape[0] - for k in TOP: - results = [] - pbar = tqdm(enumerate(test), total=m) - for i, query in pbar: - start = time.perf_counter() - result = conn.execute( - f"SELECT id FROM {name} ORDER BY embedding {metric_ops} %s LIMIT {k}", - (query,), - ).fetchall() - end = time.perf_counter() - - query_time = end - start - hit = len(set([p[0] for p in result[:k]]) & set(answer[i][:k].tolist())) - results.append((hit, query_time)) - - # Update progress bar with running metrics - curr_results = results[: i + 1] - curr_recall, curr_qps, curr_p50, _ = calculate_metrics( - curr_results, k, i + 1 - ) - pbar.set_description( - f"recall: {curr_recall:.4f} QPS: {curr_qps:.2f} P50: {curr_p50:.2f}ms" - ) + results = [] + pbar = tqdm(enumerate(test), total=m) + for i, query in pbar: + start = time.perf_counter() + result = conn.execute( + f"SELECT id FROM {name} ORDER BY embedding {metric_ops} %s LIMIT {top}", + (query,), + ).fetchall() + end = time.perf_counter() + + query_time = end - start + hit = len(set([p[0] for p in result[:top]]) & set(answer[i][:top].tolist())) + results.append((hit, query_time)) + + # Update progress bar with running metrics + curr_results = results[: i + 1] + curr_recall, curr_qps, curr_p50, _ = calculate_metrics(curr_results, top, i + 1) + pbar.set_description( + f"recall: {curr_recall:.4f} QPS: {curr_qps:.2f} P50: {curr_p50:.2f}ms" + ) - # Calculate final metrics - recall, qps, p50, p99 = calculate_metrics(results, k, m) + # Calculate final metrics + recall, qps, p50, p99 = calculate_metrics(results, top, m) - print(f"Top: {k}") - print(f" Recall: {recall:.4f}") - print(f" QPS: {qps:.2f}") - print(f" P50 latency: {p50:.2f}ms") - print(f" P99 latency: {p99:.2f}ms") + print(f"Top: {top}") + print(f" Recall: {recall:.4f}") + print(f" QPS: {qps:.2f}") + print(f" P50 latency: {p50:.2f}ms") + print(f" P99 latency: {p99:.2f}ms") if __name__ == "__main__": @@ -228,9 +225,10 @@ def sequential_bench(name, test, answer, metric_ops, conn): metric_ops, args.processes, args.password, + args.top, args.nprob, args.epsilon, ) else: conn = create_connection(args.password, args.nprob, args.epsilon) - sequential_bench(args.name, test, answer, metric_ops, conn) + sequential_bench(args.name, test, answer, metric_ops, conn, args.top) diff --git a/scripts/index.py b/scripts/index.py index 77e38399..cb4f111f 100644 --- a/scripts/index.py +++ b/scripts/index.py @@ -35,6 +35,7 @@ def build_arg_parse(): "-p", "--password", help="Database password", default="password" ) parser.add_argument("-d", "--dim", help="Dimension", type=int, required=True) + # Remember to set `max_worker_processes` at server start parser.add_argument( "-w", "--workers", @@ -132,9 +133,10 @@ async def add_centroids(conn, name, centroids): await asyncio.sleep(0) -async def add_embeddings(conn, name, dim, train, chunks): +async def add_embeddings(conn, name, dim, train, chunks, workers): await conn.execute(f"DROP TABLE IF EXISTS {name}") await conn.execute(f"CREATE TABLE {name} (id integer, embedding vector({dim}))") + await conn.execute(f"ALTER TABLE {name} SET (parallel_workers = {workers})") n, dim = train.shape chunk_size = math.ceil(n / chunks) @@ -201,7 +203,7 @@ async def main(dataset): metric_ops, ivf_config = get_ivf_ops_config( args.metric, args.workers, args.lists, args.name if args.centroids else None ) - await add_embeddings(conn, args.name, args.dim, dataset["train"], args.chunks) + await add_embeddings(conn, args.name, args.dim, dataset["train"], args.chunks, args.workers) index_finish = asyncio.Event() # Need a separate connection for monitor process diff --git a/sql/install/vchord--0.1.0.sql b/sql/install/vchord--0.1.0.sql new file mode 100644 index 00000000..8a8fd495 --- /dev/null +++ b/sql/install/vchord--0.1.0.sql @@ -0,0 +1,164 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:17 +-- bootstrap +-- List of shell types + +CREATE TYPE sphere_vector; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:53 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:29 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:5 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am.rs:29 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:6 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:6 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:1 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:18 +-- finalize +-- List of data types + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); +/* */ + diff --git a/sql/install/vchord--0.2.0.sql b/sql/install/vchord--0.2.0.sql new file mode 100644 index 00000000..dca70ae1 --- /dev/null +++ b/sql/install/vchord--0.2.0.sql @@ -0,0 +1,500 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:18 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:54 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:30 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:6 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:54 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:30 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:6 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am.rs:33 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:12 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:6 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:1 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:19 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.1.0--0.2.0.sql b/sql/upgrade/vchord--0.1.0--0.2.0.sql new file mode 100644 index 00000000..4360f27f --- /dev/null +++ b/sql/upgrade/vchord--0.1.0--0.2.0.sql @@ -0,0 +1,343 @@ +-- src/lib.rs:18 +CREATE TYPE scalar8; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:54 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:30 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:6 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:54 +-- src/datatype/operators_pgvector_vector.rs:30 +-- src/datatype/operators_pgvector_vector.rs:6 +-- src/index/am.rs:33 +-- src/index/functions.rs:12 +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:19 +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); diff --git a/tools/dev.md b/tools/dev.md new file mode 100644 index 00000000..8558ad1c --- /dev/null +++ b/tools/dev.md @@ -0,0 +1,87 @@ +## Release a new version + +### Pre-requisite + +```shell +sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config clang postgresql-server-dev-16 +cargo install --locked cargo-pgrx +cargo pgrx init +``` + +### Step 1 - generate install schema +```shell +# Change to version to release +SEMVER=0.2.0 +PREV_VERSION=0.1.0 + +mkdir temp + +cargo build --package vchord --lib --features pg16 --release +./tools/schema.sh --features pg16 --release + +cp ./target/release/schema.sql ./sql/install/vchord--$SEMVER.sql +``` + +### Step 2 - generate upgrade schema +```shell +PREV_VERSION_FILE=sql/install/vchord--$PREV_VERSION.sql + +# New lines redirect to install.sql, revised lines & deleted lines redirect to terminal +diff -u $PREV_VERSION_FILE install.sql | awk ' + /^\+/ && !/^+++/ { + print substr($0, 2) > "temp/upgrade.sql" + next + } + /^-/ && !/^---/ || /^@/ { print } + { next } +' +cp temp/upgrade.sql ./sql/upgrade/vchord--$PREV_VERSION--$SEMVER.sql +``` + +### Step 3 - validate +```shell + +# sudo pg_dropcluster --stop 16 main +sudo pg_createcluster 16 main --start +sudo -u postgres createdb vchord +sudo -u postgres psql -d vchord -c "ALTER USER postgres WITH PASSWORD '123';" + +# Dump upgraded schema and compare it +export PGHOST=localhost +export PGPASSWORD=123 +export PGUSER=postgres +./tools/dump.sh $PREV_VERSION $SEMVER > temp/dump1.sql +./tools/dump.sh $SEMVER > temp/dump2.sql +code --diff temp/dump1.sql temp/dump2.sql +``` + +### Step 4 - further test + +```shell +sudo apt remove -y vchord-pg16 +SEMVER=$SEMVER VERSION="16" ARCH="x86_64" PLATFORM="amd64" ./tools/package.sh +cd temp +wget https://github.com/tensorchord/VectorChord/releases/download/"$PREV_VERSION"/vchord-pg16_"$PREV_VERSION"_amd64.deb +sudo apt install -y ./vchord-pg16_"$PREV_VERSION"_amd64.deb +sudo -u postgres psql -d vchord -c "ALTER SYSTEM SET SHARED_PRELOAD_LIBRARIES='vchord.so';" +sudo systemctl restart postgresql@16-main.service +sudo -u postgres psql -d vchord -c 'CREATE EXTENSION IF NOT EXISTS vector; CREATE EXTENSION IF NOT EXISTS vchord;' +sudo -u postgres psql -d vchord -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord';" + +# Run Test -- Upgrade +sudo apt install -f ../build/vchord-pg16_"$SEMVER"_amd64.deb +sudo systemctl restart postgresql@16-main.service +sudo -u postgres psql -d vchord -c 'ALTER EXTENSION vchord UPDATE;' +sudo -u postgres psql -d vchord -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord';" +sqllogictest -h localhost -u postgres -d vchord -w 123 '../tests/**/*.slt' + +# Run Test -- Install +sudo -u postgres psql -d vchord -c 'DROP EXTENSION vchord CASCADE;' +sudo -u postgres psql -d vchord -c 'CREATE EXTENSION vchord;' +sudo -u postgres psql -d vchord -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord';" +sqllogictest -h localhost -u postgres -d vchord -w 123 '../tests/**/*.slt' + +sudo -u postgres psql -d vchord -c "DROP SCHEMA IF EXISTS public CASCADE;" +sudo apt remove -y vchord-pg16 +cd .. && rm -rf temp +``` \ No newline at end of file diff --git a/tools/dump-codegen.sh b/tools/dump-codegen.sh new file mode 100755 index 00000000..5d3718fd --- /dev/null +++ b/tools/dump-codegen.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -e + +cat << EOF +#include "postgres.h" + +#include "access/amapi.h" +#include "fmgr.h" + +#include + +#define DECLARE(funcname) \ + extern Datum funcname() { exit(1); } \ + extern const Pg_finfo_record *pg_finfo_##funcname(void); \ + const Pg_finfo_record *pg_finfo_##funcname(void) { return &my_finfo; } \ + extern int no_such_variable + +#define DECLARE_AMHANDLER(funcname) \ + extern Datum funcname() { return amhandler(); } \ + extern const Pg_finfo_record *pg_finfo_##funcname(void); \ + const Pg_finfo_record *pg_finfo_##funcname(void) { return &my_finfo; } \ + extern int no_such_variable + +static const Pg_finfo_record my_finfo = {1}; + +PG_MODULE_MAGIC; + +Datum amhandler() { + IndexAmRoutine *amroutine = makeNode(IndexAmRoutine); + amroutine->amsupport = 1; + amroutine->amcanorderbyop = true; + (Datum) amroutine; +} +EOF + +printf "\n" + +while read -r line; do + if [[ $line == *"amhandler"* ]]; then + echo "DECLARE_AMHANDLER($line);" + else + echo "DECLARE($line);" + fi +done <<< $(grep -ohr "'\w\+_wrapper'" $(dirname "$0")/../sql | sort | uniq | sed "s/'//g") \ No newline at end of file diff --git a/tools/dump.sh b/tools/dump.sh new file mode 100755 index 00000000..110884fd --- /dev/null +++ b/tools/dump.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -e + +mkdir -p $(dirname "$0")/../target/tools + +if [ $# -eq 0 ] ; then + echo "Usage: $0 INITIAL_VERSION VERSIONS.." + echo "Dump the extension members. Install INITIAL_VERSION first and upgrade to every version in VERSIONS." + echo "The extension members are installed in database \"vchord\", so use \"createdb vchord\" to setup." + echo "Examples:" + echo " ./tools/dump.sh 0.1.11 0.2.0 > ./dump_upgrade.sql" + echo " ./tools/dump.sh 0.2.0 > ./dump_install.sql" + echo " diff ./dump_upgrade.sql ./dump_install.sql" + exit 0 +fi + +f=() +prev_arg="" +for arg in "$@"; do + if [ "$prev_arg" = "" ]; then + x=$(realpath "$(dirname "$0")/../sql/install/vchord--${arg}.sql") + else + x=$(realpath "$(dirname "$0")/../sql/upgrade/vchord--${prev_arg}--${arg}.sql") + fi + prev_arg=$arg + f+=("$x") +done + +so=$(realpath $(dirname "$0")/../target/tools/vchord.so) +$(dirname "$0")/dump-codegen.sh | gcc -I $(pg_config --includedir-server) -fPIC -shared -o $so -x c - + +sql=$(mktemp) +echo "BEGIN;" >> $sql +echo "CREATE SCHEMA public;" >> $sql +echo "CREATE EXTENSION vector;" >> $sql +cat ${f[@]} \ + | grep -v '^\\' \ + | sed "s|MODULE_PATHNAME|$so|g" \ + >> $sql +echo "END;" >> $sql + +psql -d vchord -f $sql 1>&2 +pg_dump -d vchord +psql -d vchord -c "DROP SCHEMA IF EXISTS public CASCADE;" 1>&2 \ No newline at end of file diff --git a/tools/package.sh b/tools/package.sh index c6955ae2..1b4b052b 100755 --- a/tools/package.sh +++ b/tools/package.sh @@ -12,6 +12,7 @@ rm -rf ./build/dir_deb rm -rf ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb mkdir -p ./build/dir_zip +cp -a ./sql/upgrade/. ./build/dir_zip/ cp ./target/release/schema.sql ./build/dir_zip/vchord--$SEMVER.sql sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/dir_zip/vchord.control cp ./target/release/libvchord.so ./build/dir_zip/vchord.so From beb69c08d9052dc1aeaa7fd478575ffee696fe35 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 27 Jan 2025 12:35:42 +0800 Subject: [PATCH 101/324] chore: install zip in tensorchord-pgrx (#178) Signed-off-by: usamoi --- docker/pgrx.Dockerfile | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile index 60b6225b..93ec4aff 100644 --- a/docker/pgrx.Dockerfile +++ b/docker/pgrx.Dockerfile @@ -20,7 +20,8 @@ RUN set -eux; \ ca-certificates \ build-essential \ postgresql-common gnupg \ - libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config + libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config \ + zip RUN set -eux; \ apt -y install lsb-release wget software-properties-common gnupg; \ From 87594a43a75d51cc15531e99b9690c70de78e90e Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 27 Jan 2025 13:46:57 +0800 Subject: [PATCH 102/324] fix: release package name, version and licenses (#179) Signed-off-by: usamoi --- .github/workflows/pgrx.yml | 2 +- .github/workflows/psql.yml | 28 +++++++++++------- .github/workflows/release.yml | 30 ++++++++++++------- .github/workflows/rust.yml | 12 ++++++-- docker/Dockerfile | 4 +-- docker/binary.Dockerfile | 2 +- docker/pg-cnpg/Dockerfile | 2 +- tools/package.sh | 54 +++++++++++++++++------------------ tools/schema-codegen.sh | 43 ---------------------------- tools/schema.sh | 43 ---------------------------- 10 files changed, 80 insertions(+), 140 deletions(-) delete mode 100755 tools/schema-codegen.sh delete mode 100755 tools/schema.sh diff --git a/.github/workflows/pgrx.yml b/.github/workflows/pgrx.yml index 212de577..9bd8489b 100644 --- a/.github/workflows/pgrx.yml +++ b/.github/workflows/pgrx.yml @@ -25,7 +25,7 @@ permissions: jobs: build: - runs-on: ubicloud-standard-8 + runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v4 - name: Set up QEMU diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml index a6311893..42c108c0 100644 --- a/.github/workflows/psql.yml +++ b/.github/workflows/psql.yml @@ -41,14 +41,22 @@ jobs: strategy: matrix: version: ["14", "15", "16", "17"] - runner: ["ubicloud-standard-4", "ubicloud-standard-4-arm"] + runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] env: PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" SQLLOGICTEST: "0.25.0" - ARCH: ${{ matrix.runner == 'ubicloud-standard-4' && 'x86_64' || 'aarch64' }} - PLATFORM: ${{ matrix.runner == 'ubicloud-standard-4' && 'amd64' || 'arm64' }} + ARCH: ${{ matrix.runner == 'ubuntu-22.04' && 'x86_64' || 'aarch64' }} + PLATFORM: ${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} steps: + - name: Set up Environment + run: | + sudo apt-get remove -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' + + sudo apt-get install -y postgresql-common + sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo apt-get install -y postgresql-client-17 - uses: actions/checkout@v4 - name: Configure sccache uses: actions/github-script@v7 @@ -67,14 +75,14 @@ jobs: sudo chmod -R 777 . - name: Build - env: - SEMVER: "0.0.0" - VERSION: ${{ matrix.version }} run: | - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg${{ matrix.version }} --release - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE ./tools/schema.sh --features pg${{ matrix.version }} --release - ./tools/package.sh - docker build -t vchord:pg${{ matrix.version }} --build-arg PG_VERSION=${{ matrix.version }} -f ./docker/Dockerfile . + docker run --rm -v .:/workspace $CACHE_ENVS \ + -e SEMVER=0.0.0 \ + -e VERSION=${{ matrix.version }} \ + -e ARCH=$ARCH \ + -e PLATFORM=$PLATFORM \ + $PGRX_IMAGE ./tools/package.sh + docker build -t vchord:pg${{ matrix.version }} --build-arg PG_VERSION=${{ matrix.version }} --build-arg SEMVER=0.0.0 -f ./docker/Dockerfile . - name: Setup SQL Logic Test run: | diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 96088248..1662bbd0 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -38,14 +38,22 @@ jobs: strategy: matrix: version: ["14", "15", "16", "17"] - runner: ["ubicloud-standard-4", "ubicloud-standard-4-arm"] + runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] env: PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" SEMVER: ${{ needs.semver.outputs.SEMVER }} - ARCH: ${{ matrix.runner == 'ubicloud-standard-4' && 'x86_64' || 'aarch64' }} - PLATFORM: ${{ matrix.runner == 'ubicloud-standard-4' && 'amd64' || 'arm64' }} + ARCH: ${{ matrix.runner == 'ubuntu-22.04' && 'x86_64' || 'aarch64' }} + PLATFORM: ${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} steps: + - name: Set up Environment + run: | + sudo apt-get remove -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' + + sudo apt-get install -y postgresql-common + sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo apt-get install -y postgresql-client-17 - uses: actions/checkout@v4 - name: Configure sccache uses: actions/github-script@v7 @@ -65,15 +73,17 @@ jobs: - name: Build env: - VERSION: ${{ matrix.version }} GH_TOKEN: ${{ github.token }} run: | - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --features pg$VERSION --release - docker run --rm -v .:/workspace $CACHE_ENVS -e SEMVER=${SEMVER} $PGRX_IMAGE ./tools/schema.sh --features pg$VERSION --release - ./tools/package.sh + docker run --rm -v .:/workspace $CACHE_ENVS \ + -e SEMVER=$SEMVER \ + -e VERSION=${{ matrix.version }} \ + -e ARCH=$ARCH \ + -e PLATFORM=$PLATFORM \ + $PGRX_IMAGE ./tools/package.sh ls ./build - gh release upload --clobber $SEMVER ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb - gh release upload --clobber $SEMVER ./build/vchord-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip + gh release upload --clobber $SEMVER ./build/postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${PLATFORM}.deb + gh release upload --clobber $SEMVER ./build/postgresql-${{ matrix.version }}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip docker: runs-on: ubuntu-latest @@ -92,7 +102,7 @@ jobs: run: | mkdir -p build for arch in amd64 arm64; do - gh release download $SEMVER --pattern "vchord-pg${{ matrix.version }}_${SEMVER}_${arch}.deb" --output ./build/vchord-pg${{ matrix.version }}_${SEMVER}_${arch}.deb + gh release download $SEMVER --pattern "postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${arch}.deb" --output ./build/postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${arch}.deb done - name: Set up QEMU uses: docker/setup-qemu-action@v3 diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index df5547a4..8b1b4e2d 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -31,15 +31,23 @@ jobs: strategy: matrix: include: - - runner: "ubicloud-standard-4" + - runner: "ubuntu-22.04" arch: "x86_64" - - runner: "ubicloud-standard-4-arm" + - runner: "ubuntu-22.04-arm" arch: "aarch64" runs-on: ${{ matrix.runner }} env: PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" steps: + - name: Set up Environment + run: | + sudo apt-get remove -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' + + sudo apt-get install -y postgresql-common + sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo apt-get install -y postgresql-client-17 - uses: actions/checkout@v4 - name: Configure sccache uses: actions/github-script@v7 diff --git a/docker/Dockerfile b/docker/Dockerfile index ac5b1aad..48576eec 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -4,11 +4,11 @@ ARG PGVECTOR=0.8.0 FROM pgvector/pgvector:${PGVECTOR}-pg${PG_VERSION} ARG PG_VERSION -ARG SEMVER=0.0.0 +ARG SEMVER ARG TARGETARCH RUN echo ${PG_VERSION} -COPY ./build/vchord-pg${PG_VERSION}_${SEMVER}_${TARGETARCH}.deb /tmp/vchord.deb +COPY ./build/postgresql-${PG_VERSION}-vchord_${SEMVER}-1_${TARGETARCH}.deb /tmp/vchord.deb RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb CMD ["postgres", "-c" ,"shared_preload_libraries=vchord.so"] diff --git a/docker/binary.Dockerfile b/docker/binary.Dockerfile index 001bcc99..438031d7 100644 --- a/docker/binary.Dockerfile +++ b/docker/binary.Dockerfile @@ -5,4 +5,4 @@ ARG PG_VERSION ARG TARGETARCH WORKDIR /workspace -COPY ./build/vchord-pg${PG_VERSION}_${SEMVER}_${TARGETARCH}.deb /workspace/ +COPY ./build/postgresql-${PG_VERSION}-vchord_${SEMVER}-1_${TARGETARCH}.deb /workspace/ diff --git a/docker/pg-cnpg/Dockerfile b/docker/pg-cnpg/Dockerfile index 7f7c621c..4871e1d1 100644 --- a/docker/pg-cnpg/Dockerfile +++ b/docker/pg-cnpg/Dockerfile @@ -19,7 +19,7 @@ ARG ALTDIR=/var/lib/postgresql/data/tensorchord USER root -COPY --from=binary /workspace/vchord-pg${PG_MAJOR}_${SEMVER}_${TARGETARCH}.deb /tmp/vchord.deb +COPY --from=binary /workspace/postgresql-${PG_MAJOR}-vchord_${SEMVER}-1_${TARGETARCH}.deb /tmp/vchord.deb RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb # PGDATA is set in pg-slim and used by dependents on this image. diff --git a/tools/package.sh b/tools/package.sh index 1b4b052b..e8c92bf1 100755 --- a/tools/package.sh +++ b/tools/package.sh @@ -6,40 +6,40 @@ printf "VERSION = ${VERSION}\n" printf "ARCH = ${ARCH}\n" printf "PLATFORM = ${PLATFORM}\n" -rm -rf ./build/dir_zip -rm -rf ./build/vchord-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip -rm -rf ./build/dir_deb -rm -rf ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb +cargo build --lib --features pg$VERSION --release +cargo pgrx schema --features pg$VERSION --out ./target/schema.sql -mkdir -p ./build/dir_zip -cp -a ./sql/upgrade/. ./build/dir_zip/ -cp ./target/release/schema.sql ./build/dir_zip/vchord--$SEMVER.sql -sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/dir_zip/vchord.control -cp ./target/release/libvchord.so ./build/dir_zip/vchord.so -zip ./build/vchord-pg${VERSION}_${ARCH}-unknown-linux-gnu_${SEMVER}.zip -j ./build/dir_zip/* +rm -rf ./build -mkdir -p ./build/dir_deb -mkdir -p ./build/dir_deb/DEBIAN -mkdir -p ./build/dir_deb/usr/share/postgresql/$VERSION/extension/ -mkdir -p ./build/dir_deb/usr/lib/postgresql/$VERSION/lib/ -for file in $(ls ./build/dir_zip/*.sql | xargs -n 1 basename); do - cp ./build/dir_zip/$file ./build/dir_deb/usr/share/postgresql/$VERSION/extension/$file +mkdir -p ./build/zip +cp -a ./sql/upgrade/. ./build/zip/ +cp ./target/schema.sql ./build/zip/vchord--$SEMVER.sql +sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control +cp ./target/release/libvchord.so ./build/zip/vchord.so +zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* + +mkdir -p ./build/deb +mkdir -p ./build/deb/DEBIAN +mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ +mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ +for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file done -for file in $(ls ./build/dir_zip/*.control | xargs -n 1 basename); do - cp ./build/dir_zip/$file ./build/dir_deb/usr/share/postgresql/$VERSION/extension/$file +for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file done -for file in $(ls ./build/dir_zip/*.so | xargs -n 1 basename); do - cp ./build/dir_zip/$file ./build/dir_deb/usr/lib/postgresql/$VERSION/lib/$file +for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file done -echo "Package: vchord-pg${VERSION} -Version: ${SEMVER} +echo "Package: postgresql-${VERSION}-vchord +Version: ${SEMVER}-1 Section: database Priority: optional Architecture: ${PLATFORM} Maintainer: Tensorchord Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM -Homepage: https://pgvecto.rs/ -License: apache2" \ -> ./build/dir_deb/DEBIAN/control -(cd ./build/dir_deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/dir_deb/DEBIAN/md5sums -dpkg-deb -Zxz --build ./build/dir_deb/ ./build/vchord-pg${VERSION}_${SEMVER}_${PLATFORM}.deb +Homepage: https://vectorchord.ai/ +License: AGPL-3 or Elastic-2" \ +> ./build/deb/DEBIAN/control +(cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums +dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb diff --git a/tools/schema-codegen.sh b/tools/schema-codegen.sh deleted file mode 100755 index 12f8289e..00000000 --- a/tools/schema-codegen.sh +++ /dev/null @@ -1,43 +0,0 @@ -#!/usr/bin/env bash -set -e - -# CONTROL_FILEPATH -# SO_FILEPATH - -printf "fn main() {\n" - -cat << EOF - vchord::__pgrx_marker(); - - let mut entities = Vec::new(); - let control_file_path = std::path::PathBuf::from("$CONTROL_FILEPATH"); - let control_file = ::pgrx::pgrx_sql_entity_graph::ControlFile::try_from(control_file_path).expect(".control file should properly formatted"); - let control_file_entity = ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity::ExtensionRoot(control_file); - - entities.push(control_file_entity); -EOF - -while read -r name_ident; do -cat << EOF - extern "Rust" { - fn $name_ident() -> ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity; - } - let entity = unsafe { $name_ident() }; - entities.push(entity); -EOF -done <<< $(nm -D -g $SO_FILEPATH | grep "T __pgrx_internals_" | awk '{print $3}') - -cat << EOF - let pgrx_sql = ::pgrx::pgrx_sql_entity_graph::PgrxSql::build( - entities.into_iter(), - "vchord".to_string(), - false, - ) - .expect("SQL generation error"); - eprintln!("SQL entities to {}", "/dev/stdout",); - pgrx_sql - .write(&mut std::io::stdout()) - .expect("Could not write SQL to stdout"); -EOF - -printf "}\n" \ No newline at end of file diff --git a/tools/schema.sh b/tools/schema.sh deleted file mode 100755 index c43f7f80..00000000 --- a/tools/schema.sh +++ /dev/null @@ -1,43 +0,0 @@ -#!/usr/bin/env bash -set -e -if [[ " $@ " =~ --target' '([^ ]+) ]]; then - TARGET="${BASH_REMATCH[1]}" - if [[ " $@ " =~ " --release " ]]; then - DIR="./target/$TARGET/release" - else - DIR="./target/$TARGET/debug" - fi -else - if [[ " $@ " =~ " --release " ]]; then - DIR="./target/release" - else - DIR="./target/debug" - fi -fi - -if [ "$TARGET" = "" ]; then - printf "Target: [not specified]\n" 1>&2 - RUNNER=() -elif [ "$TARGET" = $(rustc -vV | awk '/^host/ { print $2 }') ]; then - printf "Target: [host]\n" 1>&2 - RUNNER=() -elif [ "$TARGET" = "aarch64-unknown-linux-gnu" ]; then - printf "Target: $TARGET\n" 1>&2 - QEMU_LD_PREFIX="/usr/aarch64-linux-gnu" - RUNNER=("qemu-aarch64-static") -elif [ "$TARGET" = "riscv64gc-unknown-linux-gnu" ]; then - printf "Target: $TARGET\n" 1>&2 - QEMU_LD_PREFIX="/usr/riscv64-linux-gnu" - RUNNER=("qemu-riscv64-static") -else - printf "Unknown target: $TARGET\n" 1>&2 - exit 1 -fi - -code=$(mktemp) -chmod 700 $code -CONTROL_FILEPATH="./vchord.control" SO_FILEPATH="$DIR/libvchord.so" $(dirname "$0")/schema-codegen.sh >> $code - -PGRX_EMBED=$code cargo rustc --package vchord --bin pgrx_embed_vchord "$@" -- --cfg pgrx_embed - -CARGO_PKG_VERSION=${SEMVER} QEMU_LD_PREFIX=$QEMU_LD_PREFIX "${RUNNER[@]}" "$DIR/pgrx_embed_vchord" | expand -t 4 > $DIR/schema.sql From 49637c5139a24a4a1ec05fb44e8e841d5f7d62da Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 5 Feb 2025 16:30:42 +0800 Subject: [PATCH 103/324] refactor: move algorithm to a crate (#172) Signed-off-by: usamoi --- .taplo.toml | 19 +- Cargo.lock | 52 +- Cargo.toml | 17 +- crates/algorithm/Cargo.toml | 25 + crates/algorithm/src/build.rs | 101 ++ crates/algorithm/src/bulkdelete.rs | 167 +++ .../algorithm/src}/freepages.rs | 16 +- .../algorithm/src}/insert.rs | 89 +- .../mod.rs => crates/algorithm/src/lib.rs | 36 +- crates/algorithm/src/maintain.rs | 147 +++ .../algorithm/src}/operator.rs | 63 +- {src/utils => crates/algorithm/src}/pipe.rs | 0 .../algorithm/src}/prewarm.rs | 27 +- .../scan.rs => crates/algorithm/src/search.rs | 48 +- .../algorithm/src}/tape.rs | 96 +- .../algorithm/src}/tuples.rs | 91 +- {src => crates/algorithm/src}/types.rs | 13 +- crates/algorithm/src/vectors.rs | 92 ++ crates/k_means/Cargo.toml | 15 + .../k_means.rs => crates/k_means/src/lib.rs | 77 +- crates/rabitq/src/block.rs | 57 +- crates/rabitq/src/lib.rs | 2 +- crates/simd/Cargo.toml | 2 +- crates/simd/src/f16.rs | 2 - crates/simd/src/f32.rs | 8 +- crates/simd/src/fast_scan/mod.rs | 245 ++-- crates/simd/src/lib.rs | 15 +- crates/vector/Cargo.toml | 1 - crates/vector/src/bvect.rs | 15 +- crates/vector/src/lib.rs | 2 +- crates/vector/src/scalar8.rs | 7 +- crates/vector/src/svect.rs | 15 +- crates/vector/src/vect.rs | 7 +- rustfmt.toml | 1 + src/algorithm/build.rs | 373 ------ src/algorithm/vacuum.rs | 311 ----- src/algorithm/vectors.rs | 133 -- src/bin/pgrx_embed.rs | 1 + src/datatype/memory_halfvec.rs | 20 +- src/datatype/memory_scalar8.rs | 20 +- src/datatype/memory_vector.rs | 20 +- src/gucs/mod.rs | 14 - src/gucs/prewarm.rs | 32 - src/index/am.rs | 1103 ----------------- src/index/am/am_build.rs | 950 ++++++++++++++ src/index/am/am_scan.rs | 287 +++++ src/index/am/mod.rs | 324 +++++ src/index/am_options.rs | 235 ---- src/index/am_scan.rs | 186 --- src/index/functions.rs | 27 +- src/{gucs/executing.rs => index/gucs.rs} | 39 +- src/index/mod.rs | 14 +- src/index/opclass.rs | 146 +++ src/{ => index}/projection.rs | 24 + src/{postgres.rs => index/storage.rs} | 37 +- src/index/utils.rs | 34 - src/lib.rs | 16 +- src/{upgrade/symbols.rs => upgrade.rs} | 0 src/upgrade/mod.rs | 1 - src/utils/mod.rs | 3 - src/utils/parallelism.rs | 62 - 61 files changed, 2911 insertions(+), 3071 deletions(-) create mode 100644 crates/algorithm/Cargo.toml create mode 100644 crates/algorithm/src/build.rs create mode 100644 crates/algorithm/src/bulkdelete.rs rename {src/algorithm => crates/algorithm/src}/freepages.rs (80%) rename {src/algorithm => crates/algorithm/src}/insert.rs (61%) rename src/algorithm/mod.rs => crates/algorithm/src/lib.rs (70%) create mode 100644 crates/algorithm/src/maintain.rs rename {src/algorithm => crates/algorithm/src}/operator.rs (91%) rename {src/utils => crates/algorithm/src}/pipe.rs (100%) rename {src/algorithm => crates/algorithm/src}/prewarm.rs (84%) rename src/algorithm/scan.rs => crates/algorithm/src/search.rs (83%) rename {src/algorithm => crates/algorithm/src}/tape.rs (82%) rename {src/algorithm => crates/algorithm/src}/tuples.rs (94%) rename {src => crates/algorithm/src}/types.rs (95%) create mode 100644 crates/algorithm/src/vectors.rs create mode 100644 crates/k_means/Cargo.toml rename src/utils/k_means.rs => crates/k_means/src/lib.rs (80%) delete mode 100644 src/algorithm/build.rs delete mode 100644 src/algorithm/vacuum.rs delete mode 100644 src/algorithm/vectors.rs delete mode 100644 src/gucs/mod.rs delete mode 100644 src/gucs/prewarm.rs delete mode 100644 src/index/am.rs create mode 100644 src/index/am/am_build.rs create mode 100644 src/index/am/am_scan.rs create mode 100644 src/index/am/mod.rs delete mode 100644 src/index/am_options.rs delete mode 100644 src/index/am_scan.rs rename src/{gucs/executing.rs => index/gucs.rs} (61%) rename src/{ => index}/projection.rs (51%) rename src/{postgres.rs => index/storage.rs} (91%) delete mode 100644 src/index/utils.rs rename src/{upgrade/symbols.rs => upgrade.rs} (100%) delete mode 100644 src/upgrade/mod.rs delete mode 100644 src/utils/mod.rs delete mode 100644 src/utils/parallelism.rs diff --git a/.taplo.toml b/.taplo.toml index d9a9fdac..41af6e14 100644 --- a/.taplo.toml +++ b/.taplo.toml @@ -2,9 +2,16 @@ indent_string = " " [[rule]] -keys = ["dependencies", "*-denpendencies", "lints", "patch.*", "profile.*"] - -[rule.formatting] -reorder_keys = true -reorder_arrays = true -align_comments = true +include = ["**/Cargo.toml"] +keys = [ + "dependencies", + "dev-dependencies", + "build-dependencies", + "target.*.dependencies", + "lints", + "patch.*", + "profile.*", + "workspace.dependencies", + "lints.dependencies", +] +formatting = { reorder_keys = true, reorder_arrays = true, align_comments = true } diff --git a/Cargo.lock b/Cargo.lock index f5214d8e..62d7b795 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,27 @@ dependencies = [ "memchr", ] +[[package]] +name = "algorithm" +version = "0.0.0" +dependencies = [ + "always_equal", + "distance", + "half 2.4.1", + "k_means", + "paste", + "rabitq", + "rand", + "random_orthogonal_matrix", + "serde", + "simd", + "toml", + "validator", + "vector", + "zerocopy 0.8.14", + "zerocopy-derive 0.8.14", +] + [[package]] name = "always_equal" version = "0.0.0" @@ -587,6 +608,18 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674" +[[package]] +name = "k_means" +version = "0.0.0" +dependencies = [ + "half 2.4.1", + "log", + "rabitq", + "rand", + "rayon", + "simd", +] + [[package]] name = "libc" version = "0.2.169" @@ -1209,8 +1242,8 @@ dependencies = [ "cc", "half 2.4.1", "rand", - "serde", "simd_macros", + "zerocopy 0.8.14", ] [[package]] @@ -1447,9 +1480,9 @@ dependencies = [ [[package]] name = "validator" -version = "0.19.0" +version = "0.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0b4a29d8709210980a09379f27ee31549b73292c87ab9899beee1c0d3be6303" +checksum = "43fb22e1a008ece370ce08a3e9e4447a910e92621bb49b85d6e48a45397e7cfa" dependencies = [ "idna", "once_cell", @@ -1463,9 +1496,9 @@ dependencies = [ [[package]] name = "validator_derive" -version = "0.19.0" +version = "0.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bac855a2ce6f843beb229757e6e570a42e837bcb15e5f449dd48d5747d41bf77" +checksum = "b7df16e474ef958526d1205f6dda359fdfab79d9aa6d54bafcb92dcd07673dca" dependencies = [ "darling", "once_cell", @@ -1479,24 +1512,20 @@ dependencies = [ name = "vchord" version = "0.0.0" dependencies = [ - "always_equal", + "algorithm", "distance", "half 2.4.1", - "log", + "k_means", "paste", "pgrx", "pgrx-catalog", - "rabitq", "rand", "random_orthogonal_matrix", - "rayon", "serde", "simd", "toml", "validator", "vector", - "zerocopy 0.8.14", - "zerocopy-derive 0.8.14", ] [[package]] @@ -1505,7 +1534,6 @@ version = "0.0.0" dependencies = [ "distance", "half 2.4.1", - "serde", "simd", ] diff --git a/Cargo.toml b/Cargo.toml index 3583ae26..55cf4c2a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "vchord" version.workspace = true -edition.workspace = true +edition = "2021" [lib] name = "vchord" @@ -20,25 +20,21 @@ pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] -always_equal = { path = "./crates/always_equal" } +algorithm = { path = "./crates/algorithm" } distance = { path = "./crates/distance" } -rabitq = { path = "./crates/rabitq" } +k_means = { path = "./crates/k_means" } random_orthogonal_matrix = { path = "./crates/random_orthogonal_matrix" } simd = { path = "./crates/simd" } vector = { path = "./crates/vector" } half.workspace = true -log = "0.4.25" paste = "1" pgrx = { version = "=0.12.9", default-features = false, features = ["cshim"] } pgrx-catalog = "0.1.0" rand.workspace = true -rayon = "1.10.0" serde.workspace = true toml = "0.8.19" -validator = { version = "0.19.0", features = ["derive"] } -zerocopy = "0.8.14" -zerocopy-derive = "0.8.14" +validator.workspace = true [patch.crates-io] half = { git = "https://github.com/tensorchord/half-rs.git", rev = "3f9a8843d6722bd1833de2289347640ad8770146" } @@ -56,14 +52,19 @@ edition = "2021" [workspace.dependencies] half = { version = "2.4.1", features = ["serde", "zerocopy"] } +log = "0.4.25" rand = "0.8.5" serde = "1" +validator = { version = "0.20.0", features = ["derive"] } +zerocopy = "0.8.14" +zerocopy-derive = "0.8.14" [workspace.lints] clippy.identity_op = "allow" clippy.int_plus_one = "allow" clippy.needless_range_loop = "allow" clippy.nonminimal_bool = "allow" +rust.unsafe_code = "deny" rust.unsafe_op_in_unsafe_fn = "deny" rust.unused_lifetimes = "warn" rust.unused_qualifications = "warn" diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml new file mode 100644 index 00000000..1180ba96 --- /dev/null +++ b/crates/algorithm/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "algorithm" +version.workspace = true +edition.workspace = true + +[dependencies] +always_equal = { path = "../always_equal" } +distance = { path = "../distance" } +k_means = { path = "../k_means" } +rabitq = { path = "../rabitq" } +random_orthogonal_matrix = { path = "../random_orthogonal_matrix" } +simd = { path = "../simd" } +vector = { path = "../vector" } + +half.workspace = true +paste = "1" +rand.workspace = true +serde.workspace = true +toml = "0.8.19" +validator.workspace = true +zerocopy.workspace = true +zerocopy-derive.workspace = true + +[lints] +workspace = true diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs new file mode 100644 index 00000000..685aaf13 --- /dev/null +++ b/crates/algorithm/src/build.rs @@ -0,0 +1,101 @@ +use crate::RelationWrite; +use crate::operator::{Accessor2, Operator, Vector}; +use crate::tape::*; +use crate::tuples::*; +use crate::types::*; +use vector::VectorOwned; + +pub fn build( + vector_options: VectorOptions, + vchordrq_options: VchordrqIndexingOptions, + index: impl RelationWrite, + structures: Vec>, +) { + let dims = vector_options.dims; + let is_residual = vchordrq_options.residual_quantization && O::SUPPORTS_RESIDUAL; + let mut meta = TapeWriter::<_, _, MetaTuple>::create(|| index.extend(false)); + assert_eq!(meta.first(), 0); + let freepage = TapeWriter::<_, _, FreepageTuple>::create(|| index.extend(false)); + let mut vectors = TapeWriter::<_, _, VectorTuple>::create(|| index.extend(true)); + let mut pointer_of_means = Vec::>::new(); + for i in 0..structures.len() { + let mut level = Vec::new(); + for j in 0..structures[i].len() { + let vector = structures[i].means[j].as_borrowed(); + let (metadata, slices) = O::Vector::vector_split(vector); + let mut chain = Ok(metadata); + for i in (0..slices.len()).rev() { + chain = Err(vectors.push(match chain { + Ok(metadata) => VectorTuple::_0 { + payload: None, + elements: slices[i].to_vec(), + metadata, + }, + Err(pointer) => VectorTuple::_1 { + payload: None, + elements: slices[i].to_vec(), + pointer, + }, + })); + } + level.push(chain.err().unwrap()); + } + pointer_of_means.push(level); + } + let mut pointer_of_firsts = Vec::>::new(); + for i in 0..structures.len() { + let mut level = Vec::new(); + for j in 0..structures[i].len() { + if i == 0 { + let tape = TapeWriter::<_, _, H0Tuple>::create(|| index.extend(false)); + let mut jump = TapeWriter::<_, _, JumpTuple>::create(|| index.extend(false)); + jump.push(JumpTuple { + first: tape.first(), + }); + level.push(jump.first()); + } else { + let mut tape = H1TapeWriter::<_, _>::create(|| index.extend(false)); + let h2_mean = structures[i].means[j].as_borrowed(); + let h2_children = structures[i].children[j].as_slice(); + for child in h2_children.iter().copied() { + let h1_mean = structures[i - 1].means[child as usize].as_borrowed(); + let code = if is_residual { + let mut residual_accessor = O::ResidualAccessor::default(); + residual_accessor.push( + O::Vector::elements_and_metadata(h1_mean).0, + O::Vector::elements_and_metadata(h2_mean).0, + ); + let residual = residual_accessor.finish( + O::Vector::elements_and_metadata(h1_mean).1, + O::Vector::elements_and_metadata(h2_mean).1, + ); + O::Vector::code(residual.as_borrowed()) + } else { + O::Vector::code(h1_mean) + }; + tape.push(H1Branch { + mean: pointer_of_means[i - 1][child as usize], + dis_u_2: code.dis_u_2, + factor_ppc: code.factor_ppc, + factor_ip: code.factor_ip, + factor_err: code.factor_err, + signs: code.signs, + first: pointer_of_firsts[i - 1][child as usize], + }); + } + let tape = tape.into_inner(); + level.push(tape.first()); + } + } + pointer_of_firsts.push(level); + } + meta.push(MetaTuple { + dims, + height_of_root: structures.len() as u32, + is_residual, + vectors_first: vectors.first(), + root_mean: pointer_of_means.last().unwrap()[0], + root_first: pointer_of_firsts.last().unwrap()[0], + freepage_first: freepage.first(), + }); +} diff --git a/crates/algorithm/src/bulkdelete.rs b/crates/algorithm/src/bulkdelete.rs new file mode 100644 index 00000000..524ec908 --- /dev/null +++ b/crates/algorithm/src/bulkdelete.rs @@ -0,0 +1,167 @@ +use crate::operator::Operator; +use crate::pipe::Pipe; +use crate::tuples::*; +use crate::{Page, RelationWrite}; +use std::num::NonZeroU64; + +pub fn bulkdelete( + index: impl RelationWrite, + check: impl Fn(), + callback: impl Fn(NonZeroU64) -> bool, +) { + let meta_guard = index.read(0); + let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let height_of_root = meta_tuple.height_of_root(); + let root_first = meta_tuple.root_first(); + let vectors_first = meta_tuple.vectors_first(); + drop(meta_guard); + { + type State = Vec; + let mut state: State = vec![root_first]; + let step = |state: State| { + let mut results = Vec::new(); + for first in state { + let mut current = first; + while current != u32::MAX { + let h1_guard = index.read(current); + for i in 1..=h1_guard.len() { + let h1_tuple = h1_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h1_tuple { + H1TupleReader::_0(h1_tuple) => { + for first in h1_tuple.first().iter().copied() { + results.push(first); + } + } + H1TupleReader::_1(_) => (), + } + } + current = h1_guard.get_opaque().next; + } + } + results + }; + for _ in (1..height_of_root).rev() { + state = step(state); + } + for first in state { + let jump_guard = index.read(first); + let jump_tuple = jump_guard + .get(1) + .expect("data corruption") + .pipe(read_tuple::); + let first = jump_tuple.first(); + let mut current = first; + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + let h0_tuple = read + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h0_tuple { + H0TupleReader::_0(h0_tuple) => { + let p = h0_tuple.payload(); + if let Some(payload) = p { + if callback(payload) { + break 'flag true; + } + } + } + H0TupleReader::_1(h0_tuple) => { + let p = h0_tuple.payload(); + for j in 0..32 { + if let Some(payload) = p[j] { + if callback(payload) { + break 'flag true; + } + } + } + } + H0TupleReader::_2(_) => (), + } + } + false + }; + if flag { + drop(read); + let mut write = index.write(current, false); + for i in 1..=write.len() { + let h0_tuple = write + .get_mut(i) + .expect("data corruption") + .pipe(write_tuple::); + match h0_tuple { + H0TupleWriter::_0(mut h0_tuple) => { + let p = h0_tuple.payload(); + if let Some(payload) = *p { + if callback(payload) { + *p = None; + } + } + } + H0TupleWriter::_1(mut h0_tuple) => { + let p = h0_tuple.payload(); + for j in 0..32 { + if let Some(payload) = p[j] { + if callback(payload) { + p[j] = None; + } + } + } + } + H0TupleWriter::_2(_) => (), + } + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; + } + } + } + } + { + let first = vectors_first; + let mut current = first; + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + if let Some(vector_bytes) = read.get(i) { + let vector_tuple = vector_bytes.pipe(read_tuple::>); + let p = vector_tuple.payload(); + if let Some(payload) = p { + if callback(payload) { + break 'flag true; + } + } + } + } + false + }; + if flag { + drop(read); + let mut write = index.write(current, true); + for i in 1..=write.len() { + if let Some(vector_bytes) = write.get(i) { + let vector_tuple = vector_bytes.pipe(read_tuple::>); + let p = vector_tuple.payload(); + if let Some(payload) = p { + if callback(payload) { + write.free(i); + } + } + }; + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; + } + } + } +} diff --git a/src/algorithm/freepages.rs b/crates/algorithm/src/freepages.rs similarity index 80% rename from src/algorithm/freepages.rs rename to crates/algorithm/src/freepages.rs index 8984d3ce..6dc505c2 100644 --- a/src/algorithm/freepages.rs +++ b/crates/algorithm/src/freepages.rs @@ -1,9 +1,9 @@ -use crate::algorithm::tuples::*; -use crate::algorithm::*; -use crate::utils::pipe::Pipe; +use crate::pipe::Pipe; +use crate::tuples::*; +use crate::*; use std::cmp::Reverse; -pub fn mark(relation: impl RelationWrite, freepage_first: u32, pages: &[u32]) { +pub fn mark(index: impl RelationWrite, freepage_first: u32, pages: &[u32]) { let mut pages = pages.to_vec(); pages.sort_by_key(|x| Reverse(*x)); pages.dedup(); @@ -18,7 +18,7 @@ pub fn mark(relation: impl RelationWrite, freepage_first: u32, pages: &[u32]) { } local }; - let mut freespace_guard = relation.write(current, false); + let mut freespace_guard = index.write(current, false); if freespace_guard.len() == 0 { freespace_guard.alloc(&serialize(&FreepageTuple {})); } @@ -30,19 +30,19 @@ pub fn mark(relation: impl RelationWrite, freepage_first: u32, pages: &[u32]) { freespace_tuple.mark(local as _); } if freespace_guard.get_opaque().next == u32::MAX { - let extend = relation.extend(false); + let extend = index.extend(false); freespace_guard.get_opaque_mut().next = extend.id(); } (current, offset) = (freespace_guard.get_opaque().next, offset + 32768); } } -pub fn fetch(relation: impl RelationWrite, freepage_first: u32) -> Option { +pub fn fetch(index: impl RelationWrite, freepage_first: u32) -> Option { let first = freepage_first; assert!(first != u32::MAX); let (mut current, mut offset) = (first, 0_u32); loop { - let mut freespace_guard = relation.write(current, false); + let mut freespace_guard = index.write(current, false); if freespace_guard.len() == 0 { return None; } diff --git a/src/algorithm/insert.rs b/crates/algorithm/src/insert.rs similarity index 61% rename from src/algorithm/insert.rs rename to crates/algorithm/src/insert.rs index f2b8cfbc..433e6002 100644 --- a/src/algorithm/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -1,24 +1,18 @@ -use crate::algorithm::operator::*; -use crate::algorithm::tape::read_h1_tape; -use crate::algorithm::tuples::*; -use crate::algorithm::vectors::{self}; -use crate::algorithm::{Page, PageGuard, RelationWrite}; -use crate::utils::pipe::Pipe; +use crate::operator::*; +use crate::pipe::Pipe; +use crate::tape::{access_1, append}; +use crate::tuples::*; +use crate::vectors::{self}; +use crate::{Page, RelationWrite}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZeroU64; -use vector::VectorBorrowed; -use vector::VectorOwned; +use vector::{VectorBorrowed, VectorOwned}; -pub fn insert( - relation: impl RelationWrite + Clone, - payload: NonZeroU64, - vector: O::Vector, -) { - let vector = O::Vector::random_projection(vector.as_borrowed()); - let meta_guard = relation.read(0); +pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vector: O::Vector) { + let meta_guard = index.read(0); let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); @@ -35,19 +29,14 @@ pub fn insert( None }; - let mean = vectors::vector_append::( - relation.clone(), - vectors_first, - vector.as_borrowed(), - payload, - ); + let mean = vectors::append::(index.clone(), vectors_first, vector.as_borrowed(), payload); type State = (u32, Option<::Vector>); let mut state: State = { let mean = root_mean; if is_residual { - let residual_u = vectors::vector_access_1::( - relation.clone(), + let residual_u = vectors::access_1::( + index.clone(), mean, LAccess::new( O::Vector::elements_and_metadata(vector.as_borrowed()), @@ -68,8 +57,8 @@ pub fn insert( } else { default_lut_block.as_ref().unwrap() }; - read_h1_tape( - relation.clone(), + access_1( + index.clone(), first, || { RAccess::new( @@ -88,8 +77,8 @@ pub fn insert( while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); if is_residual { - let (dis_u, residual_u) = vectors::vector_access_1::( - relation.clone(), + let (dis_u, residual_u) = vectors::access_1::( + index.clone(), mean, LAccess::new( O::Vector::elements_and_metadata(vector.as_borrowed()), @@ -105,8 +94,8 @@ pub fn insert( AlwaysEqual(Some(residual_u)), )); } else { - let dis_u = vectors::vector_access_1::( - relation.clone(), + let dis_u = vectors::access_1::( + index.clone(), mean, LAccess::new( O::Vector::elements_and_metadata(vector.as_borrowed()), @@ -140,7 +129,7 @@ pub fn insert( elements: rabitq::pack_to_u64(&code.signs), }); - let jump_guard = relation.read(first); + let jump_guard = index.read(first); let jump_tuple = jump_guard .get(1) .expect("data corruption") @@ -148,43 +137,5 @@ pub fn insert( let first = jump_tuple.first(); - assert!(first != u32::MAX); - let mut current = first; - loop { - let read = relation.read(current); - if read.get_opaque().next == u32::MAX { - drop(read); - let mut write = relation.write(current, false); - if write.get_opaque().next == u32::MAX { - if write.alloc(&bytes).is_some() { - return; - } - let mut extend = relation.extend(false); - write.get_opaque_mut().next = extend.id(); - drop(write); - let fresh = extend.id(); - if extend.alloc(&bytes).is_some() { - drop(extend); - let mut past = relation.write(first, false); - past.get_opaque_mut().skip = std::cmp::max(past.get_opaque_mut().skip, fresh); - drop(past); - return; - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - } else { - if current == first && write.get_opaque().skip != first { - current = write.get_opaque().skip; - } else { - current = write.get_opaque().next; - } - } - } else { - if current == first && read.get_opaque().skip != first { - current = read.get_opaque().skip; - } else { - current = read.get_opaque().next; - } - } - } + append(index.clone(), first, &bytes, false); } diff --git a/src/algorithm/mod.rs b/crates/algorithm/src/lib.rs similarity index 70% rename from src/algorithm/mod.rs rename to crates/algorithm/src/lib.rs index 7a9c3ff2..dccc4594 100644 --- a/src/algorithm/mod.rs +++ b/crates/algorithm/src/lib.rs @@ -1,13 +1,29 @@ -pub mod build; -pub mod freepages; -pub mod insert; +#![allow(clippy::collapsible_else_if)] +#![allow(clippy::type_complexity)] +#![allow(clippy::len_without_is_empty)] +#![feature(vec_pop_if)] + +mod build; +mod bulkdelete; +mod freepages; +mod insert; +mod maintain; +mod pipe; +mod prewarm; +mod search; +mod tape; +mod tuples; +mod vectors; + pub mod operator; -pub mod prewarm; -pub mod scan; -pub mod tape; -pub mod tuples; -pub mod vacuum; -pub mod vectors; +pub mod types; + +pub use build::build; +pub use bulkdelete::bulkdelete; +pub use insert::insert; +pub use maintain::maintain; +pub use prewarm::prewarm; +pub use search::search; use std::ops::{Deref, DerefMut}; @@ -34,7 +50,7 @@ pub trait PageGuard { fn id(&self) -> u32; } -pub trait RelationRead { +pub trait RelationRead: Clone { type Page: Page; type ReadGuard<'a>: PageGuard + Deref where diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs new file mode 100644 index 00000000..fec057b6 --- /dev/null +++ b/crates/algorithm/src/maintain.rs @@ -0,0 +1,147 @@ +use crate::operator::Operator; +use crate::pipe::Pipe; +use crate::tape::*; +use crate::tuples::*; +use crate::{Page, RelationWrite, freepages}; +use simd::fast_scan::unpack; + +pub fn maintain(index: impl RelationWrite, check: impl Fn()) { + let meta_guard = index.read(0); + let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let dims = meta_tuple.dims(); + let height_of_root = meta_tuple.height_of_root(); + let root_first = meta_tuple.root_first(); + let freepage_first = meta_tuple.freepage_first(); + drop(meta_guard); + + let firsts = { + type State = Vec; + let mut state: State = vec![root_first]; + let step = |state: State| { + let mut results = Vec::new(); + for first in state { + let mut current = first; + while current != u32::MAX { + check(); + let h1_guard = index.read(current); + for i in 1..=h1_guard.len() { + let h1_tuple = h1_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h1_tuple { + H1TupleReader::_0(h1_tuple) => { + for first in h1_tuple.first().iter().copied() { + results.push(first); + } + } + H1TupleReader::_1(_) => (), + } + } + current = h1_guard.get_opaque().next; + } + } + results + }; + for _ in (1..height_of_root).rev() { + state = step(state); + } + state + }; + + for first in firsts { + let mut jump_guard = index.write(first, false); + let mut jump_tuple = jump_guard + .get_mut(1) + .expect("data corruption") + .pipe(write_tuple::); + + let mut tape = H0TapeWriter::<_, _>::create(|| { + if let Some(id) = freepages::fetch(index.clone(), freepage_first) { + let mut write = index.write(id, false); + write.clear(); + write + } else { + index.extend(false) + } + }); + + let mut trace = Vec::new(); + + let first = *jump_tuple.first(); + let mut current = first; + let mut computing = None; + while current != u32::MAX { + check(); + trace.push(current); + let h0_guard = index.read(current); + for i in 1..=h0_guard.len() { + let h0_tuple = h0_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h0_tuple { + H0TupleReader::_0(h0_tuple) => { + if let Some(payload) = h0_tuple.payload() { + tape.push(H0Branch { + mean: h0_tuple.mean(), + dis_u_2: h0_tuple.code().0, + factor_ppc: h0_tuple.code().1, + factor_ip: h0_tuple.code().2, + factor_err: h0_tuple.code().3, + signs: h0_tuple + .code() + .4 + .iter() + .flat_map(|x| { + std::array::from_fn::<_, 64, _>(|i| *x & (1 << i) != 0) + }) + .take(dims as _) + .collect::>(), + payload, + }); + } + } + H0TupleReader::_1(h0_tuple) => { + let computing = &mut computing.take().unwrap_or_else(Vec::new); + computing.extend_from_slice(h0_tuple.elements()); + let unpacked = unpack(computing); + for j in 0..32 { + if let Some(payload) = h0_tuple.payload()[j] { + tape.push(H0Branch { + mean: h0_tuple.mean()[j], + dis_u_2: h0_tuple.metadata().0[j], + factor_ppc: h0_tuple.metadata().1[j], + factor_ip: h0_tuple.metadata().2[j], + factor_err: h0_tuple.metadata().3[j], + signs: unpacked[j] + .iter() + .flat_map(|&x| { + [x & 1 != 0, x & 2 != 0, x & 4 != 0, x & 8 != 0] + }) + .collect(), + payload, + }); + } + } + } + H0TupleReader::_2(h0_tuple) => { + let computing = computing.get_or_insert_with(Vec::new); + computing.extend_from_slice(h0_tuple.elements()); + } + } + } + current = h0_guard.get_opaque().next; + drop(h0_guard); + } + + let tape = tape.into_inner(); + let new = tape.first(); + drop(tape); + + *jump_tuple.first() = new; + drop(jump_guard); + + freepages::mark(index.clone(), freepage_first, &trace); + } +} diff --git a/src/algorithm/operator.rs b/crates/algorithm/src/operator.rs similarity index 91% rename from src/algorithm/operator.rs rename to crates/algorithm/src/operator.rs index 9506d7a2..1d70a5c8 100644 --- a/src/algorithm/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -1,4 +1,4 @@ -use crate::types::{DistanceKind, OwnedVector}; +use crate::types::*; use distance::Distance; use half::f16; use simd::Floating; @@ -174,15 +174,15 @@ impl Default for Block { impl Accessor2< - [u64; 2], - [u64; 2], + [u8; 16], + [u8; 16], (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), (f32, f32, f32, f32, f32), > for Block { type Output = [Distance; 32]; - fn push(&mut self, input: &[[u64; 2]], target: &[[u64; 2]]) { + fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { let t = simd::fast_scan::fast_scan(input, target); for i in 0..32 { self.0[i] += t[i]; @@ -212,15 +212,15 @@ impl impl Accessor2< - [u64; 2], - [u64; 2], + [u8; 16], + [u8; 16], (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), (f32, f32, f32, f32, f32), > for Block { type Output = [Distance; 32]; - fn push(&mut self, input: &[[u64; 2]], target: &[[u64; 2]]) { + fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { let t = simd::fast_scan::fast_scan(input, target); for i in 0..32 { self.0[i] += t[i]; @@ -324,7 +324,6 @@ pub struct RAccess<'a, E, M, A> { } impl<'a, E, M, A> RAccess<'a, E, M, A> { - #[allow(dead_code)] pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { Self { elements, @@ -356,22 +355,16 @@ pub trait Vector: VectorOwned { fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); fn from_owned(vector: OwnedVector) -> Self; - fn random_projection(vector: Self::Borrowed<'_>) -> Self; - - fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u64; 2]>); + fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u8; 16]>); fn compute_lut( vector: Self::Borrowed<'_>, ) -> ( - (f32, f32, f32, f32, Vec<[u64; 2]>), + (f32, f32, f32, f32, Vec<[u8; 16]>), (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), ); fn code(vector: Self::Borrowed<'_>) -> rabitq::Code; - - fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec; - - fn build_from_vecf32(x: &[f32]) -> Self; } impl Vector for VectOwned { @@ -399,18 +392,14 @@ impl Vector for VectOwned { } } - fn random_projection(vector: Self::Borrowed<'_>) -> Self { - Self::new(crate::projection::project(vector.slice())) - } - - fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u64; 2]>) { + fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u8; 16]>) { rabitq::block::preprocess(vector.slice()) } fn compute_lut( vector: Self::Borrowed<'_>, ) -> ( - (f32, f32, f32, f32, Vec<[u64; 2]>), + (f32, f32, f32, f32, Vec<[u8; 16]>), (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), ) { rabitq::compute_lut(vector.slice()) @@ -419,14 +408,6 @@ impl Vector for VectOwned { fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { rabitq::code(vector.dims(), vector.slice()) } - - fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { - vector.slice().to_vec() - } - - fn build_from_vecf32(x: &[f32]) -> Self { - Self::new(x.to_vec()) - } } impl Vector for VectOwned { @@ -454,20 +435,14 @@ impl Vector for VectOwned { } } - fn random_projection(vector: Self::Borrowed<'_>) -> Self { - Self::new(f16::vector_from_f32(&crate::projection::project( - &f16::vector_to_f32(vector.slice()), - ))) - } - - fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u64; 2]>) { + fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u8; 16]>) { rabitq::block::preprocess(&f16::vector_to_f32(vector.slice())) } fn compute_lut( vector: Self::Borrowed<'_>, ) -> ( - (f32, f32, f32, f32, Vec<[u64; 2]>), + (f32, f32, f32, f32, Vec<[u8; 16]>), (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), ) { rabitq::compute_lut(&f16::vector_to_f32(vector.slice())) @@ -476,14 +451,6 @@ impl Vector for VectOwned { fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { rabitq::code(vector.dims(), &f16::vector_to_f32(vector.slice())) } - - fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { - f16::vector_to_f32(vector.slice()) - } - - fn build_from_vecf32(x: &[f32]) -> Self { - Self::new(f16::vector_from_f32(x)) - } } pub trait OperatorDistance: 'static + Debug + Copy { @@ -496,8 +463,8 @@ pub trait OperatorDistance: 'static + Debug + Copy { ) -> Distance; type BlockAccessor: for<'a> Accessor2< - [u64; 2], - [u64; 2], + [u8; 16], + [u8; 16], (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), (f32, f32, f32, f32, f32), Output = [Distance; 32], diff --git a/src/utils/pipe.rs b/crates/algorithm/src/pipe.rs similarity index 100% rename from src/utils/pipe.rs rename to crates/algorithm/src/pipe.rs diff --git a/src/algorithm/prewarm.rs b/crates/algorithm/src/prewarm.rs similarity index 84% rename from src/algorithm/prewarm.rs rename to crates/algorithm/src/prewarm.rs index ab2731fa..587f7528 100644 --- a/src/algorithm/prewarm.rs +++ b/crates/algorithm/src/prewarm.rs @@ -1,12 +1,11 @@ -use crate::algorithm::operator::Operator; -use crate::algorithm::tuples::*; -use crate::algorithm::vectors; -use crate::algorithm::{Page, RelationRead}; -use crate::utils::pipe::Pipe; +use crate::operator::Operator; +use crate::pipe::Pipe; +use crate::tuples::*; +use crate::{Page, RelationRead, vectors}; use std::fmt::Write; -pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> String { - let meta_guard = relation.read(0); +pub fn prewarm(index: impl RelationRead, height: i32, check: impl Fn()) -> String { + let meta_guard = index.read(0); let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); let height_of_root = meta_tuple.height_of_root(); let root_mean = meta_tuple.root_mean(); @@ -23,7 +22,7 @@ pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> let mut state: State = { let mut nodes = Vec::new(); { - vectors::vector_access_1::(relation.clone(), root_mean, ()); + vectors::access_1::(index.clone(), root_mean, ()); nodes.push(root_first); } writeln!(message, "------------------------").unwrap(); @@ -40,8 +39,8 @@ pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> let mut current = list; while current != u32::MAX { counter_pages += 1; - pgrx::check_for_interrupts!(); - let h1_guard = relation.read(current); + check(); + let h1_guard = index.read(current); for i in 1..=h1_guard.len() { counter_tuples += 1; let h1_tuple = h1_guard @@ -51,7 +50,7 @@ pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> match h1_tuple { H1TupleReader::_0(h1_tuple) => { for mean in h1_tuple.mean().iter().copied() { - vectors::vector_access_1::(relation.clone(), mean, ()); + vectors::access_1::(index.clone(), mean, ()); } for first in h1_tuple.first().iter().copied() { nodes.push(first); @@ -77,7 +76,7 @@ pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> let mut counter_tuples = 0_usize; let mut counter_nodes = 0_usize; for list in state { - let jump_guard = relation.read(list); + let jump_guard = index.read(list); let jump_tuple = jump_guard .get(1) .expect("data corruption") @@ -86,8 +85,8 @@ pub fn prewarm(relation: impl RelationRead + Clone, height: i32) -> let mut current = first; while current != u32::MAX { counter_pages += 1; - pgrx::check_for_interrupts!(); - let h0_guard = relation.read(current); + check(); + let h0_guard = index.read(current); for i in 1..=h0_guard.len() { counter_tuples += 1; let h0_tuple = h0_guard diff --git a/src/algorithm/scan.rs b/crates/algorithm/src/search.rs similarity index 83% rename from src/algorithm/scan.rs rename to crates/algorithm/src/search.rs index fee6df3b..bff53e38 100644 --- a/src/algorithm/scan.rs +++ b/crates/algorithm/src/search.rs @@ -1,26 +1,22 @@ -use crate::algorithm::operator::*; -use crate::algorithm::tape::read_h0_tape; -use crate::algorithm::tape::read_h1_tape; -use crate::algorithm::tuples::*; -use crate::algorithm::vectors; -use crate::algorithm::{Page, RelationRead}; -use crate::utils::pipe::Pipe; +use crate::operator::*; +use crate::pipe::Pipe; +use crate::tape::{access_0, access_1}; +use crate::tuples::*; +use crate::{Page, RelationRead, vectors}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZeroU64; -use vector::VectorBorrowed; -use vector::VectorOwned; +use vector::{VectorBorrowed, VectorOwned}; -pub fn scan( - relation: impl RelationRead + Clone, +pub fn search( + index: impl RelationRead, vector: O::Vector, probes: Vec, epsilon: f32, ) -> impl Iterator { - let vector = O::Vector::random_projection(vector.as_borrowed()); - let meta_guard = relation.read(0); + let meta_guard = index.read(0); let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); @@ -41,8 +37,8 @@ pub fn scan( let mut state: State = vec![{ let mean = root_mean; if is_residual { - let residual_u = vectors::vector_access_1::( - relation.clone(), + let residual_u = vectors::access_1::( + index.clone(), mean, LAccess::new( O::Vector::elements_and_metadata(vector.as_borrowed()), @@ -62,8 +58,8 @@ pub fn scan( } else { default_lut.as_ref().map(|x| &x.0).unwrap() }; - read_h1_tape( - relation.clone(), + access_1( + index.clone(), first, || { RAccess::new( @@ -82,8 +78,8 @@ pub fn scan( while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); if is_residual { - let (dis_u, residual_u) = vectors::vector_access_1::( - relation.clone(), + let (dis_u, residual_u) = vectors::access_1::( + index.clone(), mean, LAccess::new( O::Vector::elements_and_metadata(vector.as_borrowed()), @@ -99,8 +95,8 @@ pub fn scan( AlwaysEqual(Some(residual_u)), )); } else { - let dis_u = vectors::vector_access_1::( - relation.clone(), + let dis_u = vectors::access_1::( + index.clone(), mean, LAccess::new( O::Vector::elements_and_metadata(vector.as_borrowed()), @@ -127,14 +123,14 @@ pub fn scan( } else { default_lut.as_ref().unwrap() }; - let jump_guard = relation.read(first); + let jump_guard = index.read(first); let jump_tuple = jump_guard .get(1) .expect("data corruption") .pipe(read_tuple::); let first = jump_tuple.first(); - read_h0_tape( - relation.clone(), + access_0( + index.clone(), first, || { RAccess::new( @@ -153,8 +149,8 @@ pub fn scan( std::iter::from_fn(move || { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); - if let Some(dis_u) = vectors::vector_access_0::( - relation.clone(), + if let Some(dis_u) = vectors::access_0::( + index.clone(), mean, pay_u, LAccess::new( diff --git a/src/algorithm/tape.rs b/crates/algorithm/src/tape.rs similarity index 82% rename from src/algorithm/tape.rs rename to crates/algorithm/src/tape.rs index 4ee722a5..edc1c4ee 100644 --- a/src/algorithm/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -1,12 +1,9 @@ -use super::RelationRead; -use super::operator::Accessor1; -use crate::algorithm::Page; -use crate::algorithm::PageGuard; -use crate::algorithm::tuples::*; -use crate::utils::pipe::Pipe; +use crate::operator::Accessor1; +use crate::pipe::Pipe; +use crate::tuples::*; +use crate::{Page, PageGuard, RelationRead, RelationWrite}; use distance::Distance; -use simd::fast_scan::any_pack; -use simd::fast_scan::padding_pack; +use simd::fast_scan::{any_pack, padding_pack}; use std::marker::PhantomData; use std::num::NonZeroU64; use std::ops::DerefMut; @@ -178,7 +175,7 @@ where } } -pub struct H0BranchWriter { +pub struct H0Branch { pub mean: IndexPointer, pub dis_u_2: f32, pub factor_ppc: f32, @@ -188,12 +185,12 @@ pub struct H0BranchWriter { pub payload: NonZeroU64, } -pub struct H0Tape { +pub struct H0TapeWriter { tape: TapeWriter, - branches: Vec, + branches: Vec, } -impl H0Tape +impl H0TapeWriter where G: PageGuard + DerefMut, G::Target: Page, @@ -205,7 +202,7 @@ where branches: Vec::new(), } } - pub fn push(&mut self, branch: H0BranchWriter) { + pub fn push(&mut self, branch: H0Branch) { self.branches.push(branch); if self.branches.len() == 32 { let chunk = std::array::from_fn::<_, 32, _>(|_| self.branches.pop().unwrap()); @@ -253,14 +250,14 @@ where } } -pub fn read_h1_tape( - relation: impl RelationRead, +pub fn access_1( + index: impl RelationRead, first: u32, - compute_block: impl Fn() -> A + Copy, + make_block_accessor: impl Fn() -> A + Copy, mut callback: impl FnMut(Distance, IndexPointer, u32), ) where A: for<'a> Accessor1< - [u64; 2], + [u8; 16], (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), Output = [Distance; 32], >, @@ -269,7 +266,7 @@ pub fn read_h1_tape( let mut current = first; let mut computing = None; while current != u32::MAX { - let h1_guard = relation.read(current); + let h1_guard = index.read(current); for i in 1..=h1_guard.len() { let h1_tuple = h1_guard .get(i) @@ -277,7 +274,7 @@ pub fn read_h1_tape( .pipe(read_tuple::); match h1_tuple { H1TupleReader::_0(h1_tuple) => { - let mut compute = computing.take().unwrap_or_else(compute_block); + let mut compute = computing.take().unwrap_or_else(make_block_accessor); compute.push(h1_tuple.elements()); let lowerbounds = compute.finish(h1_tuple.metadata()); for i in 0..h1_tuple.len() { @@ -289,7 +286,7 @@ pub fn read_h1_tape( } } H1TupleReader::_1(h1_tuple) => { - let computing = computing.get_or_insert_with(compute_block); + let computing = computing.get_or_insert_with(make_block_accessor); computing.push(h1_tuple.elements()); } } @@ -298,15 +295,15 @@ pub fn read_h1_tape( } } -pub fn read_h0_tape( - relation: impl RelationRead, +pub fn access_0( + index: impl RelationRead, first: u32, - compute_block: impl Fn() -> A + Copy, + make_block_accessor: impl Fn() -> A + Copy, compute_binary: impl Fn((f32, f32, f32, f32, &[u64])) -> Distance, mut callback: impl FnMut(Distance, IndexPointer, NonZeroU64), ) where A: for<'a> Accessor1< - [u64; 2], + [u8; 16], (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), Output = [Distance; 32], >, @@ -315,7 +312,7 @@ pub fn read_h0_tape( let mut current = first; let mut computing = None; while current != u32::MAX { - let h0_guard = relation.read(current); + let h0_guard = index.read(current); for i in 1..=h0_guard.len() { let h0_tuple = h0_guard .get(i) @@ -329,7 +326,7 @@ pub fn read_h0_tape( } } H0TupleReader::_1(h0_tuple) => { - let mut compute = computing.take().unwrap_or_else(compute_block); + let mut compute = computing.take().unwrap_or_else(make_block_accessor); compute.push(h0_tuple.elements()); let lowerbounds = compute.finish(h0_tuple.metadata()); for j in 0..32 { @@ -339,7 +336,7 @@ pub fn read_h0_tape( } } H0TupleReader::_2(h0_tuple) => { - let computing = computing.get_or_insert_with(compute_block); + let computing = computing.get_or_insert_with(make_block_accessor); computing.push(h0_tuple.elements()); } } @@ -347,3 +344,48 @@ pub fn read_h0_tape( current = h0_guard.get_opaque().next; } } + +pub fn append( + index: impl RelationWrite, + first: u32, + bytes: &[u8], + tracking_freespace: bool, +) -> IndexPointer { + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = index.read(current); + if read.freespace() as usize >= bytes.len() || read.get_opaque().next == u32::MAX { + drop(read); + let mut write = index.write(current, tracking_freespace); + if write.get_opaque().next == u32::MAX { + if let Some(i) = write.alloc(bytes) { + return pair_to_pointer((current, i)); + } + let mut extend = index.extend(tracking_freespace); + write.get_opaque_mut().next = extend.id(); + drop(write); + let fresh = extend.id(); + if let Some(i) = extend.alloc(bytes) { + drop(extend); + let mut past = index.write(first, tracking_freespace); + past.get_opaque_mut().skip = fresh.max(past.get_opaque().skip); + return pair_to_pointer((fresh, i)); + } else { + panic!("a tuple cannot even be fit in a fresh page"); + } + } + if current == first && write.get_opaque().skip != first { + current = write.get_opaque().skip; + } else { + current = write.get_opaque().next; + } + } else { + if current == first && read.get_opaque().skip != first { + current = read.get_opaque().skip; + } else { + current = read.get_opaque().next; + } + } + } +} diff --git a/src/algorithm/tuples.rs b/crates/algorithm/src/tuples.rs similarity index 94% rename from src/algorithm/tuples.rs rename to crates/algorithm/src/tuples.rs index 06d7cedb..54c1707f 100644 --- a/src/algorithm/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -1,4 +1,5 @@ -use crate::algorithm::operator::Vector; +use crate::operator::Vector; +use std::marker::PhantomData; use std::num::{NonZeroU8, NonZeroU64}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -141,7 +142,7 @@ impl MetaTupleReader<'_> { // freepage tuple #[repr(C, align(8))] -#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FreepageTupleHeader { a: [u32; 1], b: [u32; 32], @@ -233,16 +234,18 @@ impl FreepageTupleWriter<'_> { // vector tuple #[repr(C, align(8))] -#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader0 { payload: Option, metadata_s: usize, elements_s: usize, elements_e: usize, + #[cfg(target_pointer_width = "32")] + _padding_0: [ZeroU8; 4], } #[repr(C, align(8))] -#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader1 { payload: Option, pointer: IndexPointer, @@ -294,6 +297,8 @@ impl Tuple for VectorTuple { metadata_s, elements_s, elements_e, + #[cfg(target_pointer_width = "32")] + _padding_0: Default::default(), } .as_bytes(), ); @@ -435,10 +440,10 @@ pub enum H1Tuple { factor_err: [f32; 32], first: [u32; 32], len: u32, - elements: Vec<[u64; 2]>, + elements: Vec<[u8; 16]>, }, _1 { - elements: Vec<[u64; 2]>, + elements: Vec<[u8; 16]>, }, } @@ -448,7 +453,7 @@ impl H1Tuple { freespace -= size_of::() as isize; freespace -= size_of::() as isize; if freespace >= 0 { - Some(freespace as usize / size_of::<[u64; 2]>()) + Some(freespace as usize / size_of::<[u8; 16]>()) } else { None } @@ -458,7 +463,7 @@ impl H1Tuple { freespace -= size_of::() as isize; freespace -= size_of::() as isize; if freespace >= 0 { - Some(freespace as usize / size_of::<[u64; 2]>()) + Some(freespace as usize / size_of::<[u8; 16]>()) } else { None } @@ -536,13 +541,13 @@ pub enum H1TupleReader<'a> { #[derive(Debug, Clone, Copy, PartialEq)] pub struct H1TupleReader0<'a> { header: &'a H1TupleHeader0, - elements: &'a [[u64; 2]], + elements: &'a [[u8; 16]], } #[derive(Debug, Clone, Copy, PartialEq)] pub struct H1TupleReader1<'a> { header: &'a H1TupleHeader1, - elements: &'a [[u64; 2]], + elements: &'a [[u8; 16]], } impl<'a> TupleReader<'a> for H1TupleReader<'a> { @@ -586,13 +591,13 @@ impl<'a> H1TupleReader0<'a> { pub fn first(self) -> &'a [u32] { &self.header.first[..self.header.len as usize] } - pub fn elements(&self) -> &'a [[u64; 2]] { + pub fn elements(&self) -> &'a [[u8; 16]] { self.elements } } impl<'a> H1TupleReader1<'a> { - pub fn elements(&self) -> &'a [[u64; 2]] { + pub fn elements(&self) -> &'a [[u8; 16]] { self.elements } } @@ -720,10 +725,10 @@ pub enum H0Tuple { factor_ip: [f32; 32], factor_err: [f32; 32], payload: [Option; 32], - elements: Vec<[u64; 2]>, + elements: Vec<[u8; 16]>, }, _2 { - elements: Vec<[u64; 2]>, + elements: Vec<[u8; 16]>, }, } @@ -733,7 +738,7 @@ impl H0Tuple { freespace -= size_of::() as isize; freespace -= size_of::() as isize; if freespace >= 0 { - Some(freespace as usize / size_of::<[u64; 2]>()) + Some(freespace as usize / size_of::<[u8; 16]>()) } else { None } @@ -743,7 +748,7 @@ impl H0Tuple { freespace -= size_of::() as isize; freespace -= size_of::() as isize; if freespace >= 0 { - Some(freespace as usize / size_of::<[u64; 2]>()) + Some(freespace as usize / size_of::<[u8; 16]>()) } else { None } @@ -878,7 +883,7 @@ impl<'a> H0TupleReader0<'a> { #[derive(Debug, Clone, Copy, PartialEq)] pub struct H0TupleReader1<'a> { header: &'a H0TupleHeader1, - elements: &'a [[u64; 2]], + elements: &'a [[u8; 16]], } impl<'a> H0TupleReader1<'a> { @@ -893,7 +898,7 @@ impl<'a> H0TupleReader1<'a> { &self.header.factor_err, ) } - pub fn elements(self) -> &'a [[u64; 2]] { + pub fn elements(self) -> &'a [[u8; 16]] { self.elements } pub fn payload(self) -> &'a [Option; 32] { @@ -904,11 +909,11 @@ impl<'a> H0TupleReader1<'a> { #[derive(Debug, Clone, Copy, PartialEq)] pub struct H0TupleReader2<'a> { header: &'a H0TupleHeader2, - elements: &'a [[u64; 2]], + elements: &'a [[u8; 16]], } impl<'a> H0TupleReader2<'a> { - pub fn elements(self) -> &'a [[u64; 2]] { + pub fn elements(self) -> &'a [[u8; 16]] { self.elements } } @@ -961,7 +966,7 @@ pub struct H0TupleWriter0<'a> { pub struct H0TupleWriter1<'a> { header: &'a mut H0TupleHeader1, #[allow(dead_code)] - elements: &'a mut [[u64; 2]], + elements: &'a mut [[u8; 16]], } #[derive(Debug)] @@ -969,7 +974,7 @@ pub struct H0TupleWriter2<'a> { #[allow(dead_code)] header: &'a mut H0TupleHeader2, #[allow(dead_code)] - elements: &'a mut [[u64; 2]], + elements: &'a mut [[u8; 16]], } impl<'a> TupleWriter<'a> for H0TupleWriter<'a> { @@ -1025,16 +1030,15 @@ pub const fn pair_to_pointer(pair: (u32, u16)) -> IndexPointer { IndexPointer(value) } -#[allow(dead_code)] -const fn soundness_check(a: (u32, u16)) { +#[test] +const fn soundness_check() { + let a = (111, 222); let b = pair_to_pointer(a); let c = pointer_to_pair(b); assert!(a.0 == c.0); assert!(a.1 == c.1); } -const _: () = soundness_check((111, 222)); - #[repr(transparent)] #[derive( Debug, @@ -1123,15 +1127,17 @@ impl<'a> RefChecker<'a> { } pub struct MutChecker<'a> { - flag: Vec, - bytes: &'a mut [u8], + flag: usize, + bytes: *mut [u8], + phantom: PhantomData<&'a mut [u8]>, } impl<'a> MutChecker<'a> { pub fn new(bytes: &'a mut [u8]) -> Self { Self { - flag: vec![0u64; bytes.len().div_ceil(64)], + flag: 0, bytes, + phantom: PhantomData, } } pub fn prefix( @@ -1143,15 +1149,14 @@ impl<'a> MutChecker<'a> { if !(start <= end && end <= self.bytes.len()) { panic!("bad bytes"); } - for i in start..end { - if (self.flag[i / 64] & (1 << (i % 64))) != 0 { - panic!("bad bytes"); - } else { - self.flag[i / 64] |= 1 << (i % 64); - } + if !(self.flag <= start) { + panic!("bad bytes"); + } else { + self.flag = end; } + #[allow(unsafe_code)] let bytes = unsafe { - std::slice::from_raw_parts_mut(self.bytes.as_mut_ptr().add(start), end - start) + std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) }; FromBytes::mut_from_bytes(bytes).expect("bad bytes") } @@ -1165,21 +1170,19 @@ impl<'a> MutChecker<'a> { if !(start <= end && end <= self.bytes.len()) { panic!("bad bytes"); } - for i in start..end { - if (self.flag[i / 64] & (1 << (i % 64))) != 0 { - panic!("bad bytes"); - } else { - self.flag[i / 64] |= 1 << (i % 64); - } + if !(self.flag <= start) { + panic!("bad bytes"); + } else { + self.flag = end; } + #[allow(unsafe_code)] let bytes = unsafe { - std::slice::from_raw_parts_mut(self.bytes.as_mut_ptr().add(start), end - start) + std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) }; FromBytes::mut_from_bytes(bytes).expect("bad bytes") } } -// this test only passes if `MIRIFLAGS="-Zmiri-tree-borrows"` is set #[test] fn aliasing_test() { #[repr(C, align(8))] diff --git a/src/types.rs b/crates/algorithm/src/types.rs similarity index 95% rename from src/types.rs rename to crates/algorithm/src/types.rs index 4ef2171d..a8f3b8ee 100644 --- a/src/types.rs +++ b/crates/algorithm/src/types.rs @@ -98,7 +98,7 @@ impl VchordrqIndexingOptions { } } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone)] pub enum OwnedVector { Vecf32(VectOwned), Vecf16(VectOwned), @@ -148,3 +148,14 @@ impl VectorOptions { } } } + +pub struct Structure { + pub means: Vec, + pub children: Vec>, +} + +impl Structure { + pub fn len(&self) -> usize { + self.children.len() + } +} diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs new file mode 100644 index 00000000..dcc11af0 --- /dev/null +++ b/crates/algorithm/src/vectors.rs @@ -0,0 +1,92 @@ +use crate::operator::*; +use crate::pipe::Pipe; +use crate::tuples::*; +use crate::{Page, PageGuard, RelationRead, RelationWrite, tape}; +use std::num::NonZeroU64; +use vector::VectorOwned; + +pub fn access_1< + O: Operator, + A: Accessor1<::Element, ::Metadata>, +>( + index: impl RelationRead, + mean: IndexPointer, + accessor: A, +) -> A::Output { + let mut cursor = Err(mean); + let mut result = accessor; + while let Err(mean) = cursor.map_err(pointer_to_pair) { + let vector_guard = index.read(mean.0); + let vector_tuple = vector_guard + .get(mean.1) + .expect("data corruption") + .pipe(read_tuple::>); + if vector_tuple.payload().is_some() { + panic!("data corruption"); + } + result.push(vector_tuple.elements()); + cursor = vector_tuple.metadata_or_pointer(); + } + result.finish(cursor.expect("data corruption")) +} + +pub fn access_0< + O: Operator, + A: Accessor1<::Element, ::Metadata>, +>( + index: impl RelationRead, + mean: IndexPointer, + payload: NonZeroU64, + accessor: A, +) -> Option { + let mut cursor = Err(mean); + let mut result = accessor; + while let Err(mean) = cursor.map_err(pointer_to_pair) { + let vector_guard = index.read(mean.0); + let vector_tuple = vector_guard + .get(mean.1)? + .pipe(read_tuple::>); + if vector_tuple.payload().is_none() { + panic!("data corruption"); + } + if vector_tuple.payload() != Some(payload) { + return None; + } + result.push(vector_tuple.elements()); + cursor = vector_tuple.metadata_or_pointer(); + } + Some(result.finish(cursor.ok()?)) +} + +pub fn append( + index: impl RelationWrite, + vectors_first: u32, + vector: ::Borrowed<'_>, + payload: NonZeroU64, +) -> IndexPointer { + fn append(index: impl RelationWrite, first: u32, bytes: &[u8]) -> IndexPointer { + if let Some(mut write) = index.search(bytes.len()) { + let i = write.alloc(bytes).unwrap(); + return pair_to_pointer((write.id(), i)); + } + tape::append(index, first, bytes, true) + } + let (metadata, slices) = O::Vector::vector_split(vector); + let mut chain = Ok(metadata); + for i in (0..slices.len()).rev() { + let bytes = serialize::>(&match chain { + Ok(metadata) => VectorTuple::_0 { + elements: slices[i].to_vec(), + payload: Some(payload), + metadata, + }, + Err(pointer) => VectorTuple::_1 { + elements: slices[i].to_vec(), + payload: Some(payload), + pointer, + }, + }); + chain = Err(append(index.clone(), vectors_first, &bytes)); + } + chain.err().unwrap() +} diff --git a/crates/k_means/Cargo.toml b/crates/k_means/Cargo.toml new file mode 100644 index 00000000..93b82ed0 --- /dev/null +++ b/crates/k_means/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "k_means" +version.workspace = true +edition.workspace = true + +[dependencies] +half.workspace = true +log.workspace = true +rabitq = { path = "../rabitq" } +rand.workspace = true +rayon = "1.10.0" +simd = { path = "../simd" } + +[lints] +workspace = true diff --git a/src/utils/k_means.rs b/crates/k_means/src/lib.rs similarity index 80% rename from src/utils/k_means.rs rename to crates/k_means/src/lib.rs index b1808c90..b4b2e5c6 100644 --- a/src/utils/k_means.rs +++ b/crates/k_means/src/lib.rs @@ -1,9 +1,72 @@ -use super::parallelism::{ParallelIterator, Parallelism}; +#![allow(clippy::type_complexity)] + use half::f16; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use simd::Floating; use simd::fast_scan::{any_pack, padding_pack}; +use std::any::Any; +use std::panic::AssertUnwindSafe; +use std::sync::Arc; + +pub use rayon::iter::ParallelIterator; + +pub trait Parallelism: Send + Sync { + fn check(&self); + + fn rayon_into_par_iter(&self, x: I) -> I::Iter; +} + +struct ParallelismCheckPanic(Box); + +pub struct RayonParallelism { + stop: Arc, +} + +impl RayonParallelism { + pub fn scoped( + num_threads: usize, + stop: Arc, + f: impl FnOnce(&Self) -> R, + ) -> Result { + match std::panic::catch_unwind(AssertUnwindSafe(|| { + rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .panic_handler(|e| { + if e.downcast_ref::().is_some() { + return; + } + log::error!("Asynchronous task panickied."); + }) + .build_scoped( + |thread| thread.run(), + |_| { + let pool = Self { stop: stop.clone() }; + f(&pool) + }, + ) + })) { + Ok(x) => x, + Err(e) => match e.downcast::() { + Ok(payload) => std::panic::resume_unwind((*payload).0), + Err(e) => std::panic::resume_unwind(e), + }, + } + } +} + +impl Parallelism for RayonParallelism { + fn check(&self) { + match std::panic::catch_unwind(AssertUnwindSafe(|| (self.stop)())) { + Ok(()) => (), + Err(payload) => std::panic::panic_any(ParallelismCheckPanic(payload)), + } + } + + fn rayon_into_par_iter(&self, x: I) -> I::Iter { + x.into_par_iter() + } +} pub fn k_means( parallelism: &P, @@ -108,10 +171,10 @@ fn rabitq_index( factor_ppc: [f32; 32], factor_ip: [f32; 32], factor_err: [f32; 32], - elements: Vec<[u64; 2]>, + elements: Vec<[u8; 16]>, } impl Block { - fn code(&self) -> (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], &[[u64; 2]]) { + fn code(&self) -> (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], &[[u8; 16]]) { ( &self.dis_u_2, &self.factor_ppc, @@ -134,16 +197,14 @@ fn rabitq_index( parallelism .rayon_into_par_iter(0..n) .map(|i| { - use distance::Distance; let lut = rabitq::block::preprocess(&samples[i]); - let mut result = (Distance::INFINITY, 0); + let mut result = (f32::INFINITY, 0); for block in 0..c.div_ceil(32) { let lowerbound = rabitq::block::process_lowerbound_l2(&lut, blocks[block].code(), 1.9); for j in block * 32..std::cmp::min(block * 32 + 32, c) { - if lowerbound[j - block * 32] < result.0 { - let dis = - Distance::from_f32(f32::reduce_sum_of_d2(&samples[i], ¢roids[j])); + if lowerbound[j - block * 32].to_f32() < result.0 { + let dis = f32::reduce_sum_of_d2(&samples[i], ¢roids[j]); if dis <= result.0 { result = (dis, j); } diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs index 9f26fce1..52d4a82d 100644 --- a/crates/rabitq/src/block.rs +++ b/crates/rabitq/src/block.rs @@ -1,7 +1,7 @@ use distance::Distance; use simd::Floating; -pub fn preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec<[u64; 2]>) { +pub fn preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec<[u8; 16]>) { let dis_v_2 = f32::reduce_sum_of_x2(vector); let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); let qvector_sum = if vector.len() <= 4369 { @@ -13,13 +13,13 @@ pub fn preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec<[u64; 2]>) { } pub fn process_lowerbound_l2( - lut: &(f32, f32, f32, f32, Vec<[u64; 2]>), + lut: &(f32, f32, f32, f32, Vec<[u8; 16]>), (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( &[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], - &[[u64; 2]], + &[[u8; 16]], ), epsilon: f32, ) -> [Distance; 32] { @@ -36,13 +36,13 @@ pub fn process_lowerbound_l2( } pub fn process_lowerbound_dot( - lut: &(f32, f32, f32, f32, Vec<[u64; 2]>), + lut: &(f32, f32, f32, f32, Vec<[u8; 16]>), (_, factor_ppc, factor_ip, factor_err, t): ( &[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], - &[[u64; 2]], + &[[u8; 16]], ), epsilon: f32, ) -> [Distance; 32] { @@ -56,11 +56,12 @@ pub fn process_lowerbound_dot( }) } -pub fn compress(mut vector: Vec) -> Vec<[u64; 2]> { - let width = vector.len().div_ceil(4); - vector.resize(width * 4, 0); - let mut result = vec![[0u64, 0u64]; width]; - for i in 0..width { +pub fn compress(mut vector: Vec) -> Vec<[u8; 16]> { + let n = vector.len().div_ceil(4); + vector.resize(n * 4, 0); + let mut result = vec![[0u8; 16]; n]; + for i in 0..n { + #[allow(unsafe_code)] unsafe { // this hint is used to skip bound checks std::hint::assert_unchecked(4 * i + 3 < vector.len()); @@ -70,26 +71,22 @@ pub fn compress(mut vector: Vec) -> Vec<[u64; 2]> { let t_2 = vector[4 * i + 2]; let t_3 = vector[4 * i + 3]; result[i] = [ - u64::from_le_bytes([ - 0, - t_0, - t_1, - t_1 + t_0, - t_2, - t_2 + t_0, - t_2 + t_1, - t_2 + t_1 + t_0, - ]), - u64::from_le_bytes([ - t_3, - t_3 + t_0, - t_3 + t_1, - t_3 + t_1 + t_0, - t_3 + t_2, - t_3 + t_2 + t_0, - t_3 + t_2 + t_1, - t_3 + t_2 + t_1 + t_0, - ]), + 0, + t_0, + t_1, + t_1 + t_0, + t_2, + t_2 + t_0, + t_2 + t_1, + t_2 + t_1 + t_0, + t_3, + t_3 + t_0, + t_3 + t_1, + t_3 + t_1 + t_0, + t_3 + t_2, + t_3 + t_2 + t_0, + t_3 + t_2 + t_1, + t_3 + t_2 + t_1 + t_0, ]; } result diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index 1e379a30..a90f4453 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -77,7 +77,7 @@ pub fn code(dims: u32, vector: &[f32]) -> Code { pub fn compute_lut( vector: &[f32], ) -> ( - (f32, f32, f32, f32, Vec<[u64; 2]>), + (f32, f32, f32, f32, Vec<[u8; 16]>), (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), ) { use simd::Floating; diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 0905fb14..adbd3e1d 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -5,8 +5,8 @@ edition.workspace = true [dependencies] half.workspace = true -serde.workspace = true simd_macros = { path = "../simd_macros" } +zerocopy.workspace = true [dev-dependencies] rand.workspace = true diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index ce906187..3d7fbd5d 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -129,8 +129,6 @@ impl Floating for f16 { } mod reduce_or_of_is_zero_x { - // FIXME: add manually-implemented SIMD version - use half::f16; #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index 555da5c7..d0dc2cdb 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -119,8 +119,6 @@ impl Floating for f32 { } mod reduce_or_of_is_zero_x { - // FIXME: add manually-implemented SIMD version - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] pub fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { for &x in this { @@ -1024,8 +1022,7 @@ mod reduce_min_max_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn reduce_min_max_of_x_v3(this: &[f32]) -> (f32, f32) { - use crate::emulate::emulate_mm256_reduce_max_ps; - use crate::emulate::emulate_mm256_reduce_min_ps; + use crate::emulate::{emulate_mm256_reduce_max_ps, emulate_mm256_reduce_min_ps}; unsafe { use std::arch::x86_64::*; let mut n = this.len(); @@ -1081,8 +1078,7 @@ mod reduce_min_max_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] fn reduce_min_max_of_x_v2(this: &[f32]) -> (f32, f32) { - use crate::emulate::emulate_mm_reduce_max_ps; - use crate::emulate::emulate_mm_reduce_min_ps; + use crate::emulate::{emulate_mm_reduce_max_ps, emulate_mm_reduce_min_ps}; unsafe { use std::arch::x86_64::*; let mut n = this.len(); diff --git a/crates/simd/src/fast_scan/mod.rs b/crates/simd/src/fast_scan/mod.rs index 84ad9901..735d3e87 100644 --- a/crates/simd/src/fast_scan/mod.rs +++ b/crates/simd/src/fast_scan/mod.rs @@ -32,7 +32,7 @@ bits 4..7 | code (n-1),vector 16 | code (n-1),vector 24 | ... | code (n-1),vecto */ -pub fn pack(x: [&[u8]; 32]) -> Vec<[u64; 2]> { +pub fn pack(x: [&[u8]; 32]) -> Vec<[u8; 16]> { let n = { let l = x.each_ref().map(|i| i.len()); for i in 1..32 { @@ -43,74 +43,68 @@ pub fn pack(x: [&[u8]; 32]) -> Vec<[u64; 2]> { let mut result = Vec::with_capacity(n); for i in 0..n { result.push([ - u64::from_le_bytes([ - x[0][i] | (x[16][i] << 4), - x[8][i] | (x[24][i] << 4), - x[1][i] | (x[17][i] << 4), - x[9][i] | (x[25][i] << 4), - x[2][i] | (x[18][i] << 4), - x[10][i] | (x[26][i] << 4), - x[3][i] | (x[19][i] << 4), - x[11][i] | (x[27][i] << 4), - ]), - u64::from_le_bytes([ - x[4][i] | (x[20][i] << 4), - x[12][i] | (x[28][i] << 4), - x[5][i] | (x[21][i] << 4), - x[13][i] | (x[29][i] << 4), - x[6][i] | (x[22][i] << 4), - x[14][i] | (x[30][i] << 4), - x[7][i] | (x[23][i] << 4), - x[15][i] | (x[31][i] << 4), - ]), + x[0][i] | (x[16][i] << 4), + x[8][i] | (x[24][i] << 4), + x[1][i] | (x[17][i] << 4), + x[9][i] | (x[25][i] << 4), + x[2][i] | (x[18][i] << 4), + x[10][i] | (x[26][i] << 4), + x[3][i] | (x[19][i] << 4), + x[11][i] | (x[27][i] << 4), + x[4][i] | (x[20][i] << 4), + x[12][i] | (x[28][i] << 4), + x[5][i] | (x[21][i] << 4), + x[13][i] | (x[29][i] << 4), + x[6][i] | (x[22][i] << 4), + x[14][i] | (x[30][i] << 4), + x[7][i] | (x[23][i] << 4), + x[15][i] | (x[31][i] << 4), ]); } result } -pub fn unpack(x: &[[u64; 2]]) -> [Vec; 32] { +pub fn unpack(x: &[[u8; 16]]) -> [Vec; 32] { let n = x.len(); let mut result = std::array::from_fn(|_| Vec::with_capacity(n)); for i in 0..n { - let a = x[i][0].to_le_bytes(); - let b = x[i][1].to_le_bytes(); - result[0].push(a[0] & 0xf); - result[1].push(a[2] & 0xf); - result[2].push(a[4] & 0xf); - result[3].push(a[6] & 0xf); - result[4].push(b[0] & 0xf); - result[5].push(b[2] & 0xf); - result[6].push(b[4] & 0xf); - result[7].push(b[6] & 0xf); - result[8].push(a[1] & 0xf); - result[9].push(a[3] & 0xf); - result[10].push(a[5] & 0xf); - result[11].push(a[7] & 0xf); - result[12].push(b[1] & 0xf); - result[13].push(b[3] & 0xf); - result[14].push(b[5] & 0xf); - result[15].push(b[7] & 0xf); - result[16].push(a[0] >> 4); - result[17].push(a[2] >> 4); - result[18].push(a[4] >> 4); - result[19].push(a[6] >> 4); - result[20].push(b[0] >> 4); - result[21].push(b[2] >> 4); - result[22].push(b[4] >> 4); - result[23].push(b[6] >> 4); - result[24].push(a[1] >> 4); - result[25].push(a[3] >> 4); - result[26].push(a[5] >> 4); - result[27].push(a[7] >> 4); - result[28].push(b[1] >> 4); - result[29].push(b[3] >> 4); - result[30].push(b[5] >> 4); - result[31].push(b[7] >> 4); + result[0].push(x[i][0] & 0xf); + result[1].push(x[i][2] & 0xf); + result[2].push(x[i][4] & 0xf); + result[3].push(x[i][6] & 0xf); + result[4].push(x[i][8] & 0xf); + result[5].push(x[i][10] & 0xf); + result[6].push(x[i][12] & 0xf); + result[7].push(x[i][14] & 0xf); + result[8].push(x[i][1] & 0xf); + result[9].push(x[i][3] & 0xf); + result[10].push(x[i][5] & 0xf); + result[11].push(x[i][7] & 0xf); + result[12].push(x[i][9] & 0xf); + result[13].push(x[i][11] & 0xf); + result[14].push(x[i][13] & 0xf); + result[15].push(x[i][15] & 0xf); + result[16].push(x[i][0] >> 4); + result[17].push(x[i][2] >> 4); + result[18].push(x[i][4] >> 4); + result[19].push(x[i][6] >> 4); + result[20].push(x[i][8] >> 4); + result[21].push(x[i][10] >> 4); + result[22].push(x[i][12] >> 4); + result[23].push(x[i][14] >> 4); + result[24].push(x[i][1] >> 4); + result[25].push(x[i][3] >> 4); + result[26].push(x[i][5] >> 4); + result[27].push(x[i][7] >> 4); + result[28].push(x[i][9] >> 4); + result[29].push(x[i][11] >> 4); + result[30].push(x[i][13] >> 4); + result[31].push(x[i][15] >> 4); } result } -pub fn padding_pack(x: impl IntoIterator>) -> Vec<[u64; 2]> { +pub fn padding_pack(x: impl IntoIterator>) -> Vec<[u8; 16]> { let x = x.into_iter().collect::>(); let x = x.iter().map(|x| x.as_ref()).collect::>(); if x.is_empty() || x.len() > 32 { @@ -132,7 +126,7 @@ mod fast_scan { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - fn fast_scan_v4(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { + fn fast_scan_v4(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -171,11 +165,11 @@ mod fast_scan { let mut i = 0_usize; while i + 4 <= n { - let c = _mm512_loadu_si512(code.as_ptr().add(i).cast()); + let code = _mm512_loadu_si512(code.as_ptr().add(i).cast()); let mask = _mm512_set1_epi8(0xf); - let clo = _mm512_and_si512(c, mask); - let chi = _mm512_and_si512(_mm512_srli_epi16(c, 4), mask); + let clo = _mm512_and_si512(code, mask); + let chi = _mm512_and_si512(_mm512_srli_epi16(code, 4), mask); let lut = _mm512_loadu_si512(lut.as_ptr().add(i).cast()); let res_lo = _mm512_shuffle_epi8(lut, clo); @@ -188,11 +182,11 @@ mod fast_scan { i += 4; } if i + 2 <= n { - let c = _mm256_loadu_si256(code.as_ptr().add(i).cast()); + let code = _mm256_loadu_si256(code.as_ptr().add(i).cast()); let mask = _mm256_set1_epi8(0xf); - let clo = _mm256_and_si256(c, mask); - let chi = _mm256_and_si256(_mm256_srli_epi16(c, 4), mask); + let clo = _mm256_and_si256(code, mask); + let chi = _mm256_and_si256(_mm256_srli_epi16(code, 4), mask); let lut = _mm256_loadu_si256(lut.as_ptr().add(i).cast()); let res_lo = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, clo)); @@ -205,11 +199,11 @@ mod fast_scan { i += 2; } if i < n { - let c = _mm_loadu_si128(code.as_ptr().add(i).cast()); + let code = _mm_loadu_si128(code.as_ptr().add(i).cast()); let mask = _mm_set1_epi8(0xf); - let clo = _mm_and_si128(c, mask); - let chi = _mm_and_si128(_mm_srli_epi16(c, 4), mask); + let clo = _mm_and_si128(code, mask); + let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); let lut = _mm_loadu_si128(lut.as_ptr().add(i).cast()); let res_lo = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, clo)); @@ -251,11 +245,11 @@ mod fast_scan { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { for n in 90..110 { let code = (0..n) - .map(|_| [rand::random(), rand::random()]) - .collect::>(); + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); let lut = (0..n) - .map(|_| [rand::random(), rand::random()]) - .collect::>(); + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); unsafe { assert_eq!(fast_scan_v4(&code, &lut), fast_scan_fallback(&code, &lut)); } @@ -266,7 +260,7 @@ mod fast_scan { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - fn fast_scan_v3(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { + fn fast_scan_v3(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -291,11 +285,11 @@ mod fast_scan { let mut i = 0_usize; while i + 2 <= n { - let c = _mm256_loadu_si256(code.as_ptr().add(i).cast()); + let code = _mm256_loadu_si256(code.as_ptr().add(i).cast()); let mask = _mm256_set1_epi8(0xf); - let clo = _mm256_and_si256(c, mask); - let chi = _mm256_and_si256(_mm256_srli_epi16(c, 4), mask); + let clo = _mm256_and_si256(code, mask); + let chi = _mm256_and_si256(_mm256_srli_epi16(code, 4), mask); let lut = _mm256_loadu_si256(lut.as_ptr().add(i).cast()); let res_lo = _mm256_shuffle_epi8(lut, clo); @@ -308,11 +302,11 @@ mod fast_scan { i += 2; } if i < n { - let c = _mm_loadu_si128(code.as_ptr().add(i).cast()); + let code = _mm_loadu_si128(code.as_ptr().add(i).cast()); let mask = _mm_set1_epi8(0xf); - let clo = _mm_and_si128(c, mask); - let chi = _mm_and_si128(_mm_srli_epi16(c, 4), mask); + let clo = _mm_and_si128(code, mask); + let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); let lut = _mm_loadu_si128(lut.as_ptr().add(i).cast()); let res_lo = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, clo)); @@ -354,11 +348,11 @@ mod fast_scan { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { for n in 90..110 { let code = (0..n) - .map(|_| [rand::random(), rand::random()]) - .collect::>(); + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); let lut = (0..n) - .map(|_| [rand::random(), rand::random()]) - .collect::>(); + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); unsafe { assert_eq!(fast_scan_v3(&code, &lut), fast_scan_fallback(&code, &lut)); } @@ -368,7 +362,7 @@ mod fast_scan { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - fn fast_scan_v2(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { + fn fast_scan_v2(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -383,11 +377,11 @@ mod fast_scan { let mut i = 0_usize; while i < n { - let c = _mm_loadu_si128(code.as_ptr().add(i).cast()); + let code = _mm_loadu_si128(code.as_ptr().add(i).cast()); let mask = _mm_set1_epi8(0xf); - let clo = _mm_and_si128(c, mask); - let chi = _mm_and_si128(_mm_srli_epi16(c, 4), mask); + let clo = _mm_and_si128(code, mask); + let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); let lut = _mm_loadu_si128(lut.as_ptr().add(i).cast()); let res_lo = _mm_shuffle_epi8(lut, clo); @@ -425,11 +419,11 @@ mod fast_scan { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { for n in 90..110 { let code = (0..n) - .map(|_| [rand::random(), rand::random()]) - .collect::>(); + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); let lut = (0..n) - .map(|_| [rand::random(), rand::random()]) - .collect::>(); + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); unsafe { assert_eq!(fast_scan_v2(&code, &lut), fast_scan_fallback(&code, &lut)); } @@ -439,7 +433,7 @@ mod fast_scan { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "v8.3a")] - fn fast_scan_v8_3a(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { + fn fast_scan_v8_3a(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -454,11 +448,10 @@ mod fast_scan { let mut i = 0_usize; while i < n { - let c = vld1q_u8(code.as_ptr().add(i).cast()); + let code = vld1q_u8(code.as_ptr().add(i).cast()); - let mask = vdupq_n_u8(0xf); - let clo = vandq_u8(c, mask); - let chi = vandq_u8(vshrq_n_u8(c, 4), mask); + let clo = vandq_u8(code, vdupq_n_u8(0xf)); + let chi = vshrq_n_u8(code, 4); let lut = vld1q_u8(lut.as_ptr().add(i).cast()); let res_lo = vreinterpretq_u16_u8(vqtbl1q_u8(lut, clo)); @@ -496,11 +489,11 @@ mod fast_scan { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { for n in 90..110 { let code = (0..n) - .map(|_| [rand::random(), rand::random()]) - .collect::>(); + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); let lut = (0..n) - .map(|_| [rand::random(), rand::random()]) - .collect::>(); + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); unsafe { assert_eq!( fast_scan_v8_3a(&code, &lut), @@ -512,25 +505,16 @@ mod fast_scan { } #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a")] - pub fn fast_scan(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { - assert_eq!(code.len(), lut.len()); - let n = code.len(); - - fn unary(op: impl Fn(T) -> U, a: [T; N]) -> [U; N] { - std::array::from_fn(|i| op(a[i])) - } - fn binary(op: impl Fn(T, T) -> T, a: [T; N], b: [T; N]) -> [T; N] { + pub fn fast_scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + fn binary(op: impl Fn(u16, u16) -> u16, a: [u16; 8], b: [u16; 8]) -> [u16; 8] { std::array::from_fn(|i| op(a[i], b[i])) } - fn shuffle(a: [T; N], b: [u8; N]) -> [T; N] { + fn shuffle(a: [u8; 16], b: [u8; 16]) -> [u8; 16] { std::array::from_fn(|i| a[b[i] as usize]) } - fn cast(x: [u8; 16]) -> [u16; 8] { - std::array::from_fn(|i| u16::from_le_bytes([x[i << 1 | 0], x[i << 1 | 1]])) - } - fn setr(x: [[T; 8]; 4]) -> [T; 32] { - std::array::from_fn(|i| x[i >> 3][i & 7]) - } + + assert_eq!(code.len(), lut.len()); + let n = code.len(); let mut a_0 = [0u16; 8]; let mut a_1 = [0u16; 8]; @@ -538,29 +522,28 @@ mod fast_scan { let mut a_3 = [0u16; 8]; for i in 0..n { - let c = unsafe { std::mem::transmute::<[u64; 2], [u8; 16]>(code[i]) }; - - let mask = [0xfu8; 16]; - let clo = binary(std::ops::BitAnd::bitand, c, mask); - let chi = binary(std::ops::BitAnd::bitand, unary(|x| x >> 4, c), mask); - - let lut = unsafe { std::mem::transmute::<[u64; 2], [u8; 16]>(lut[i]) }; - let res_lo = cast(shuffle(lut, clo)); - a_0 = binary(u16::wrapping_add, a_0, res_lo); - a_1 = binary(u16::wrapping_add, a_1, unary(|x| x >> 8, res_lo)); - let res_hi = cast(shuffle(lut, chi)); - a_2 = binary(u16::wrapping_add, a_2, res_hi); - a_3 = binary(u16::wrapping_add, a_3, unary(|x| x >> 8, res_hi)); + let code = code[i]; + + let clo = code.map(|x| x & 0xf); + let chi = code.map(|x| x >> 4); + + let lut = lut[i]; + let res_lo = zerocopy::transmute!(shuffle(lut, clo)); + a_0 = binary(|x, y| x + y, a_0, res_lo); + a_1 = binary(|x, y| x + y, a_1, res_lo.map(|x| x >> 8)); + let res_hi = zerocopy::transmute!(shuffle(lut, chi)); + a_2 = binary(|x, y| x + y, a_2, res_hi); + a_3 = binary(|x, y| x + y, a_3, res_hi.map(|x| x >> 8)); } - a_0 = binary(u16::wrapping_sub, a_0, unary(|x| x.wrapping_shl(8), a_1)); - a_2 = binary(u16::wrapping_sub, a_2, unary(|x| x.wrapping_shl(8), a_3)); + a_0 = binary(|x, y| x - y, a_0, a_1.map(|x| x << 8)); + a_2 = binary(|x, y| x - y, a_2, a_3.map(|x| x << 8)); - setr([a_0, a_1, a_2, a_3]) + zerocopy::transmute!([a_0, a_1, a_2, a_3]) } } #[inline(always)] -pub fn fast_scan(code: &[[u64; 2]], lut: &[[u64; 2]]) -> [u16; 32] { +pub fn fast_scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { fast_scan::fast_scan(code, lut) } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index aec5c523..2f03d53c 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -2,6 +2,7 @@ #![feature(avx512_target_feature)] #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512_f16))] +#![allow(unsafe_code)] mod aligned; mod emulate; @@ -15,16 +16,7 @@ pub mod quantize; pub mod u8; pub trait Floating: - Copy - + Send - + Sync - + std::fmt::Debug - + serde::Serialize - + for<'a> serde::Deserialize<'a> - + Default - + 'static - + PartialEq - + PartialOrd + Copy + Send + Sync + std::fmt::Debug + Default + 'static + PartialEq + PartialOrd { fn zero() -> Self; fn infinity() -> Self; @@ -79,8 +71,7 @@ mod internal { pub use is_riscv64_cpu_detected; } -pub use simd_macros::multiversion; -pub use simd_macros::target_cpu; +pub use simd_macros::{multiversion, target_cpu}; #[cfg(target_arch = "x86_64")] #[allow(unused_imports)] diff --git a/crates/vector/Cargo.toml b/crates/vector/Cargo.toml index b910d361..186daf82 100644 --- a/crates/vector/Cargo.toml +++ b/crates/vector/Cargo.toml @@ -6,7 +6,6 @@ edition.workspace = true [dependencies] distance = { path = "../distance" } half.workspace = true -serde.workspace = true simd = { path = "../simd" } [lints] diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs index fe80164c..7877addb 100644 --- a/crates/vector/src/bvect.rs +++ b/crates/vector/src/bvect.rs @@ -1,12 +1,11 @@ use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; -use serde::{Deserialize, Serialize}; use std::ops::{Bound, RangeBounds}; pub const BVECTOR_WIDTH: u32 = u64::BITS; // When using binary vector, please ensure that the padding bits are always zero. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone)] pub struct BVectOwned { dims: u32, data: Vec, @@ -29,7 +28,10 @@ impl BVectOwned { if dims % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dims % BVECTOR_WIDTH) != 0 { return None; } - unsafe { Some(Self::new_unchecked(dims, data)) } + #[allow(unsafe_code)] + unsafe { + Some(Self::new_unchecked(dims, data)) + } } /// # Safety @@ -37,6 +39,7 @@ impl BVectOwned { /// * `dims` must be in `1..=65535`. /// * `data` must be of the correct length. /// * The padding bits must be zero. + #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked(dims: u32, data: Vec) -> Self { Self { dims, data } @@ -83,7 +86,10 @@ impl<'a> BVectBorrowed<'a> { if dims % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dims % BVECTOR_WIDTH) != 0 { return None; } - unsafe { Some(Self::new_unchecked(dims, data)) } + #[allow(unsafe_code)] + unsafe { + Some(Self::new_unchecked(dims, data)) + } } /// # Safety @@ -91,6 +97,7 @@ impl<'a> BVectBorrowed<'a> { /// * `dims` must be in `1..=65535`. /// * `data` must be of the correct length. /// * The padding bits must be zero. + #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked(dims: u32, data: &'a [u64]) -> Self { Self { dims, data } diff --git a/crates/vector/src/lib.rs b/crates/vector/src/lib.rs index e82e4a6c..64128c7a 100644 --- a/crates/vector/src/lib.rs +++ b/crates/vector/src/lib.rs @@ -3,7 +3,7 @@ pub mod scalar8; pub mod svect; pub mod vect; -pub trait VectorOwned: Clone + serde::Serialize + for<'a> serde::Deserialize<'a> + 'static { +pub trait VectorOwned: Clone + 'static { type Borrowed<'a>: VectorBorrowed; fn as_borrowed(&self) -> Self::Borrowed<'_>; diff --git a/crates/vector/src/scalar8.rs b/crates/vector/src/scalar8.rs index ff9095aa..3792e270 100644 --- a/crates/vector/src/scalar8.rs +++ b/crates/vector/src/scalar8.rs @@ -1,9 +1,8 @@ use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; -use serde::{Deserialize, Serialize}; use std::ops::RangeBounds; -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone)] pub struct Scalar8Owned { sum_of_x2: f32, k: f32, @@ -29,12 +28,14 @@ impl Scalar8Owned { if !(1..=65535).contains(&code.len()) { return None; } + #[allow(unsafe_code)] Some(unsafe { Self::new_unchecked(sum_of_x2, k, b, sum_of_code, code) }) } /// # Safety /// /// * `code.len()` must not be zero. + #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked( sum_of_x2: f32, @@ -105,12 +106,14 @@ impl<'a> Scalar8Borrowed<'a> { if !(1..=65535).contains(&code.len()) { return None; } + #[allow(unsafe_code)] Some(unsafe { Self::new_unchecked(sum_of_x2, k, b, sum_of_code, code) }) } /// # Safety /// /// * `code.len()` must not be zero. + #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked( sum_of_x2: f32, diff --git a/crates/vector/src/svect.rs b/crates/vector/src/svect.rs index 08f678d3..26cbf197 100644 --- a/crates/vector/src/svect.rs +++ b/crates/vector/src/svect.rs @@ -1,10 +1,9 @@ use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; -use serde::{Deserialize, Serialize}; use simd::Floating; use std::ops::{Bound, RangeBounds}; -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone)] pub struct SVectOwned { dims: u32, indexes: Vec, @@ -37,7 +36,10 @@ impl SVectOwned { if S::reduce_or_of_is_zero_x(&values) { return None; } - unsafe { Some(Self::new_unchecked(dims, indexes, values)) } + #[allow(unsafe_code)] + unsafe { + Some(Self::new_unchecked(dims, indexes, values)) + } } /// # Safety @@ -46,6 +48,7 @@ impl SVectOwned { /// * `indexes.len()` must be equal to `values.len()`. /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dims`. /// * A floating number in `values` must not be positive zero or negative zero. + #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked(dims: u32, indexes: Vec, values: Vec) -> Self { Self { @@ -119,7 +122,10 @@ impl<'a, S: Floating> SVectBorrowed<'a, S> { return None; } } - unsafe { Some(Self::new_unchecked(dims, indexes, values)) } + #[allow(unsafe_code)] + unsafe { + Some(Self::new_unchecked(dims, indexes, values)) + } } /// # Safety @@ -129,6 +135,7 @@ impl<'a, S: Floating> SVectBorrowed<'a, S> { /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dims`. /// * A floating number in `values` must not be positive zero or negative zero. #[inline(always)] + #[allow(unsafe_code)] pub unsafe fn new_unchecked(dims: u32, indexes: &'a [u32], values: &'a [S]) -> Self { Self { dims, diff --git a/crates/vector/src/vect.rs b/crates/vector/src/vect.rs index 34b186f7..527fc6ad 100644 --- a/crates/vector/src/vect.rs +++ b/crates/vector/src/vect.rs @@ -1,11 +1,10 @@ use super::{VectorBorrowed, VectorOwned}; use distance::Distance; -use serde::{Deserialize, Serialize}; use simd::Floating; use std::cmp::Ordering; use std::ops::RangeBounds; -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone)] #[repr(transparent)] pub struct VectOwned(Vec); @@ -20,12 +19,14 @@ impl VectOwned { if !(1..=65535).contains(&slice.len()) { return None; } + #[allow(unsafe_code)] Some(unsafe { Self::new_unchecked(slice) }) } /// # Safety /// /// * `slice.len()` must not be zero. + #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked(slice: Vec) -> Self { Self(slice) @@ -76,12 +77,14 @@ impl<'a, S: Floating> VectBorrowed<'a, S> { if !(1..=65535).contains(&slice.len()) { return None; } + #[allow(unsafe_code)] Some(unsafe { Self::new_unchecked(slice) }) } /// # Safety /// /// * `slice.len()` must not be zero. + #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked(slice: &'a [S]) -> Self { Self(slice) diff --git a/rustfmt.toml b/rustfmt.toml index 35011368..c32b643a 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -1 +1,2 @@ style_edition = "2024" +imports_granularity = "Module" diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs deleted file mode 100644 index 4c893b7f..00000000 --- a/src/algorithm/build.rs +++ /dev/null @@ -1,373 +0,0 @@ -use crate::algorithm::RelationWrite; -use crate::algorithm::operator::{Operator, Vector}; -use crate::algorithm::tape::*; -use crate::algorithm::tuples::*; -use crate::index::am_options::Opfamily; -use crate::types::VchordrqBuildOptions; -use crate::types::VchordrqExternalBuildOptions; -use crate::types::VchordrqIndexingOptions; -use crate::types::VchordrqInternalBuildOptions; -use crate::types::VectorOptions; -use rand::Rng; -use simd::Floating; -use std::num::NonZeroU64; -use std::sync::Arc; -use vector::VectorBorrowed; -use vector::VectorOwned; - -pub trait HeapRelation { - fn traverse(&self, progress: bool, callback: F) - where - F: FnMut((NonZeroU64, O::Vector)); - fn opfamily(&self) -> Opfamily; -} - -pub trait Reporter { - fn tuples_total(&mut self, tuples_total: u64); -} - -pub fn build, R: Reporter>( - vector_options: VectorOptions, - vchordrq_options: VchordrqIndexingOptions, - heap_relation: T, - relation: impl RelationWrite, - mut reporter: R, -) { - let dims = vector_options.dims; - let is_residual = vchordrq_options.residual_quantization && O::SUPPORTS_RESIDUAL; - let structures = match vchordrq_options.build { - VchordrqBuildOptions::External(external_build) => Structure::extern_build( - vector_options.clone(), - heap_relation.opfamily(), - external_build.clone(), - ), - VchordrqBuildOptions::Internal(internal_build) => { - let mut tuples_total = 0_u64; - let samples = { - let mut rand = rand::thread_rng(); - let max_number_of_samples = internal_build - .lists - .last() - .unwrap() - .saturating_mul(internal_build.sampling_factor); - let mut samples = Vec::new(); - let mut number_of_samples = 0_u32; - heap_relation.traverse(false, |(_, vector)| { - let vector = vector.as_borrowed(); - assert_eq!(dims, vector.dims(), "invalid vector dimensions"); - if number_of_samples < max_number_of_samples { - samples.push(O::Vector::build_to_vecf32(vector)); - number_of_samples += 1; - } else { - let index = rand.gen_range(0..max_number_of_samples) as usize; - samples[index] = O::Vector::build_to_vecf32(vector); - } - tuples_total += 1; - }); - samples - }; - reporter.tuples_total(tuples_total); - Structure::internal_build(vector_options.clone(), internal_build.clone(), samples) - } - }; - let mut meta = TapeWriter::<_, _, MetaTuple>::create(|| relation.extend(false)); - assert_eq!(meta.first(), 0); - let freepage = TapeWriter::<_, _, FreepageTuple>::create(|| relation.extend(false)); - let mut vectors = TapeWriter::<_, _, VectorTuple>::create(|| relation.extend(true)); - let mut pointer_of_means = Vec::>::new(); - for i in 0..structures.len() { - let mut level = Vec::new(); - for j in 0..structures[i].len() { - let vector = O::Vector::build_from_vecf32(&structures[i].means[j]); - let (metadata, slices) = O::Vector::vector_split(vector.as_borrowed()); - let mut chain = Ok(metadata); - for i in (0..slices.len()).rev() { - chain = Err(vectors.push(match chain { - Ok(metadata) => VectorTuple::_0 { - payload: None, - elements: slices[i].to_vec(), - metadata, - }, - Err(pointer) => VectorTuple::_1 { - payload: None, - elements: slices[i].to_vec(), - pointer, - }, - })); - } - level.push(chain.err().unwrap()); - } - pointer_of_means.push(level); - } - let mut pointer_of_firsts = Vec::>::new(); - for i in 0..structures.len() { - let mut level = Vec::new(); - for j in 0..structures[i].len() { - if i == 0 { - let tape = TapeWriter::<_, _, H0Tuple>::create(|| relation.extend(false)); - let mut jump = TapeWriter::<_, _, JumpTuple>::create(|| relation.extend(false)); - jump.push(JumpTuple { - first: tape.first(), - }); - level.push(jump.first()); - } else { - let mut tape = H1TapeWriter::<_, _>::create(|| relation.extend(false)); - let h2_mean = &structures[i].means[j]; - let h2_children = &structures[i].children[j]; - for child in h2_children.iter().copied() { - let h1_mean = &structures[i - 1].means[child as usize]; - let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(h1_mean, h2_mean)) - } else { - rabitq::code(dims, h1_mean) - }; - tape.push(H1Branch { - mean: pointer_of_means[i - 1][child as usize], - dis_u_2: code.dis_u_2, - factor_ppc: code.factor_ppc, - factor_ip: code.factor_ip, - factor_err: code.factor_err, - signs: code.signs, - first: pointer_of_firsts[i - 1][child as usize], - }); - } - let tape = tape.into_inner(); - level.push(tape.first()); - } - } - pointer_of_firsts.push(level); - } - meta.push(MetaTuple { - dims, - height_of_root: structures.len() as u32, - is_residual, - vectors_first: vectors.first(), - root_mean: pointer_of_means.last().unwrap()[0], - root_first: pointer_of_firsts.last().unwrap()[0], - freepage_first: freepage.first(), - }); -} - -struct Structure { - means: Vec>, - children: Vec>, -} - -impl Structure { - fn len(&self) -> usize { - self.children.len() - } - fn internal_build( - vector_options: VectorOptions, - internal_build: VchordrqInternalBuildOptions, - mut samples: Vec>, - ) -> Vec { - use std::iter::once; - for sample in samples.iter_mut() { - *sample = crate::projection::project(sample); - } - let mut result = Vec::::new(); - for w in internal_build.lists.iter().rev().copied().chain(once(1)) { - let means = crate::utils::parallelism::RayonParallelism::scoped( - internal_build.build_threads as _, - Arc::new(|| { - pgrx::check_for_interrupts!(); - }), - |parallelism| { - crate::utils::k_means::k_means( - parallelism, - w as usize, - vector_options.dims as usize, - if let Some(structure) = result.last() { - &structure.means - } else { - &samples - }, - internal_build.spherical_centroids, - 10, - ) - }, - ) - .expect("failed to create thread pool"); - if let Some(structure) = result.last() { - let mut children = vec![Vec::new(); means.len()]; - for i in 0..structure.len() as u32 { - let target = - crate::utils::k_means::k_means_lookup(&structure.means[i as usize], &means); - children[target].push(i); - } - let (means, children) = std::iter::zip(means, children) - .filter(|(_, x)| !x.is_empty()) - .unzip::<_, _, Vec<_>, Vec<_>>(); - result.push(Structure { means, children }); - } else { - let children = vec![Vec::new(); means.len()]; - result.push(Structure { means, children }); - } - } - result - } - fn extern_build( - vector_options: VectorOptions, - _opfamily: Opfamily, - external_build: VchordrqExternalBuildOptions, - ) -> Vec { - use std::collections::BTreeMap; - let VchordrqExternalBuildOptions { table } = external_build; - let mut parents = BTreeMap::new(); - let mut vectors = BTreeMap::new(); - pgrx::spi::Spi::connect(|client| { - use crate::datatype::memory_vector::VectorOutput; - use pgrx::pg_sys::panic::ErrorReportable; - use vector::VectorBorrowed; - let schema_query = "SELECT n.nspname::TEXT - FROM pg_catalog.pg_extension e - LEFT JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace - WHERE e.extname = 'vector';"; - let pgvector_schema: String = client - .select(schema_query, None, None) - .unwrap_or_report() - .first() - .get_by_name("nspname") - .expect("external build: cannot get schema of pgvector") - .expect("external build: cannot get schema of pgvector"); - let dump_query = - format!("SELECT id, parent, vector::{pgvector_schema}.vector FROM {table};"); - let centroids = client.select(&dump_query, None, None).unwrap_or_report(); - for row in centroids { - let id: Option = row.get_by_name("id").unwrap(); - let parent: Option = row.get_by_name("parent").unwrap(); - let vector: Option = row.get_by_name("vector").unwrap(); - let id = id.expect("external build: id could not be NULL"); - let vector = vector.expect("external build: vector could not be NULL"); - let pop = parents.insert(id, parent); - if pop.is_some() { - pgrx::error!( - "external build: there are at least two lines have same id, id = {id}" - ); - } - if vector_options.dims != vector.as_borrowed().dims() { - pgrx::error!("external build: incorrect dimension, id = {id}"); - } - vectors.insert(id, crate::projection::project(vector.as_borrowed().slice())); - } - }); - if parents.len() >= 2 && parents.values().all(|x| x.is_none()) { - // if there are more than one vertexs and no edges, - // assume there is an implicit root - let n = parents.len(); - let mut result = Vec::new(); - result.push(Structure { - means: vectors.values().cloned().collect::>(), - children: vec![Vec::new(); n], - }); - result.push(Structure { - means: vec![{ - // compute the vector on root, without normalizing it - let mut sum = vec![0.0f32; vector_options.dims as _]; - for vector in vectors.values() { - f32::vector_add_inplace(&mut sum, vector); - } - f32::vector_mul_scalar_inplace(&mut sum, 1.0 / n as f32); - sum - }], - children: vec![(0..n as u32).collect()], - }); - return result; - } - let mut children = parents - .keys() - .map(|x| (*x, Vec::new())) - .collect::>(); - let mut root = None; - for (&id, &parent) in parents.iter() { - if let Some(parent) = parent { - if let Some(parent) = children.get_mut(&parent) { - parent.push(id); - } else { - pgrx::error!( - "external build: parent does not exist, id = {id}, parent = {parent}" - ); - } - } else { - if let Some(root) = root { - pgrx::error!("external build: two root, id = {root}, id = {id}"); - } else { - root = Some(id); - } - } - } - let Some(root) = root else { - pgrx::error!("external build: there are no root"); - }; - let mut heights = BTreeMap::<_, _>::new(); - fn dfs_for_heights( - heights: &mut BTreeMap>, - children: &BTreeMap>, - u: i32, - ) { - if heights.contains_key(&u) { - pgrx::error!("external build: detect a cycle, id = {u}"); - } - heights.insert(u, None); - let mut height = None; - for &v in children[&u].iter() { - dfs_for_heights(heights, children, v); - let new = heights[&v].unwrap() + 1; - if let Some(height) = height { - if height != new { - pgrx::error!("external build: two heights, id = {u}"); - } - } else { - height = Some(new); - } - } - if height.is_none() { - height = Some(1); - } - heights.insert(u, height); - } - dfs_for_heights(&mut heights, &children, root); - let heights = heights - .into_iter() - .map(|(k, v)| (k, v.expect("not a connected graph"))) - .collect::>(); - if !(1..=8).contains(&(heights[&root] - 1)) { - pgrx::error!( - "external build: unexpected tree height, height = {}", - heights[&root] - ); - } - let mut cursors = vec![0_u32; 1 + heights[&root] as usize]; - let mut labels = BTreeMap::new(); - for id in parents.keys().copied() { - let height = heights[&id]; - let cursor = cursors[height as usize]; - labels.insert(id, (height, cursor)); - cursors[height as usize] += 1; - } - fn extract( - height: u32, - labels: &BTreeMap, - vectors: &BTreeMap>, - children: &BTreeMap>, - ) -> (Vec>, Vec>) { - labels - .iter() - .filter(|(_, (h, _))| *h == height) - .map(|(id, _)| { - ( - vectors[id].clone(), - children[id].iter().map(|id| labels[id].1).collect(), - ) - }) - .unzip() - } - let mut result = Vec::new(); - for height in 1..=heights[&root] { - let (means, children) = extract(height, &labels, &vectors, &children); - result.push(Structure { means, children }); - } - result - } -} diff --git a/src/algorithm/vacuum.rs b/src/algorithm/vacuum.rs deleted file mode 100644 index 17366256..00000000 --- a/src/algorithm/vacuum.rs +++ /dev/null @@ -1,311 +0,0 @@ -use crate::algorithm::freepages; -use crate::algorithm::operator::Operator; -use crate::algorithm::tape::*; -use crate::algorithm::tuples::*; -use crate::algorithm::{Page, RelationWrite}; -use crate::utils::pipe::Pipe; -use simd::fast_scan::unpack; -use std::num::NonZeroU64; - -pub fn bulkdelete( - relation: impl RelationWrite, - delay: impl Fn(), - callback: impl Fn(NonZeroU64) -> bool, -) { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); - let height_of_root = meta_tuple.height_of_root(); - let root_first = meta_tuple.root_first(); - let vectors_first = meta_tuple.vectors_first(); - drop(meta_guard); - { - type State = Vec; - let mut state: State = vec![root_first]; - let step = |state: State| { - let mut results = Vec::new(); - for first in state { - let mut current = first; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h1_tuple { - H1TupleReader::_0(h1_tuple) => { - for first in h1_tuple.first().iter().copied() { - results.push(first); - } - } - H1TupleReader::_1(_) => (), - } - } - current = h1_guard.get_opaque().next; - } - } - results - }; - for _ in (1..height_of_root).rev() { - state = step(state); - } - for first in state { - let jump_guard = relation.read(first); - let jump_tuple = jump_guard - .get(1) - .expect("data corruption") - .pipe(read_tuple::); - let first = jump_tuple.first(); - let mut current = first; - while current != u32::MAX { - delay(); - let read = relation.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - let h0_tuple = read - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h0_tuple { - H0TupleReader::_0(h0_tuple) => { - let p = h0_tuple.payload(); - if let Some(payload) = p { - if callback(payload) { - break 'flag true; - } - } - } - H0TupleReader::_1(h0_tuple) => { - let p = h0_tuple.payload(); - for j in 0..32 { - if let Some(payload) = p[j] { - if callback(payload) { - break 'flag true; - } - } - } - } - H0TupleReader::_2(_) => (), - } - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current, false); - for i in 1..=write.len() { - let h0_tuple = write - .get_mut(i) - .expect("data corruption") - .pipe(write_tuple::); - match h0_tuple { - H0TupleWriter::_0(mut h0_tuple) => { - let p = h0_tuple.payload(); - if let Some(payload) = *p { - if callback(payload) { - *p = None; - } - } - } - H0TupleWriter::_1(mut h0_tuple) => { - let p = h0_tuple.payload(); - for j in 0..32 { - if let Some(payload) = p[j] { - if callback(payload) { - p[j] = None; - } - } - } - } - H0TupleWriter::_2(_) => (), - } - } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; - } - } - } - } - { - let first = vectors_first; - let mut current = first; - while current != u32::MAX { - delay(); - let read = relation.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - if let Some(vector_bytes) = read.get(i) { - let vector_tuple = vector_bytes.pipe(read_tuple::>); - let p = vector_tuple.payload(); - if let Some(payload) = p { - if callback(payload) { - break 'flag true; - } - } - } - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current, true); - for i in 1..=write.len() { - if let Some(vector_bytes) = write.get(i) { - let vector_tuple = vector_bytes.pipe(read_tuple::>); - let p = vector_tuple.payload(); - if let Some(payload) = p { - if callback(payload) { - write.free(i); - } - } - }; - } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; - } - } - } -} - -pub fn maintain(relation: impl RelationWrite + Clone, delay: impl Fn()) { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); - let dims = meta_tuple.dims(); - let height_of_root = meta_tuple.height_of_root(); - let root_first = meta_tuple.root_first(); - let freepage_first = meta_tuple.freepage_first(); - drop(meta_guard); - - let firsts = { - type State = Vec; - let mut state: State = vec![root_first]; - let step = |state: State| { - let mut results = Vec::new(); - for first in state { - let mut current = first; - while current != u32::MAX { - delay(); - let h1_guard = relation.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h1_tuple { - H1TupleReader::_0(h1_tuple) => { - for first in h1_tuple.first().iter().copied() { - results.push(first); - } - } - H1TupleReader::_1(_) => (), - } - } - current = h1_guard.get_opaque().next; - } - } - results - }; - for _ in (1..height_of_root).rev() { - state = step(state); - } - state - }; - - for first in firsts { - let mut jump_guard = relation.write(first, false); - let mut jump_tuple = jump_guard - .get_mut(1) - .expect("data corruption") - .pipe(write_tuple::); - - let mut tape = H0Tape::<_, _>::create(|| { - if let Some(id) = freepages::fetch(relation.clone(), freepage_first) { - let mut write = relation.write(id, false); - write.clear(); - write - } else { - relation.extend(false) - } - }); - - let mut trace = Vec::new(); - - let first = *jump_tuple.first(); - let mut current = first; - let mut computing = None; - while current != u32::MAX { - delay(); - trace.push(current); - let h0_guard = relation.read(current); - for i in 1..=h0_guard.len() { - let h0_tuple = h0_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h0_tuple { - H0TupleReader::_0(h0_tuple) => { - if let Some(payload) = h0_tuple.payload() { - tape.push(H0BranchWriter { - mean: h0_tuple.mean(), - dis_u_2: h0_tuple.code().0, - factor_ppc: h0_tuple.code().1, - factor_ip: h0_tuple.code().2, - factor_err: h0_tuple.code().3, - signs: h0_tuple - .code() - .4 - .iter() - .flat_map(|x| { - std::array::from_fn::<_, 64, _>(|i| *x & (1 << i) != 0) - }) - .take(dims as _) - .collect::>(), - payload, - }); - } - } - H0TupleReader::_1(h0_tuple) => { - let computing = &mut computing.take().unwrap_or_else(Vec::new); - computing.extend_from_slice(h0_tuple.elements()); - let unpacked = unpack(computing); - for j in 0..32 { - if let Some(payload) = h0_tuple.payload()[j] { - tape.push(H0BranchWriter { - mean: h0_tuple.mean()[j], - dis_u_2: h0_tuple.metadata().0[j], - factor_ppc: h0_tuple.metadata().1[j], - factor_ip: h0_tuple.metadata().2[j], - factor_err: h0_tuple.metadata().3[j], - signs: unpacked[j] - .iter() - .flat_map(|&x| { - [x & 1 != 0, x & 2 != 0, x & 4 != 0, x & 8 != 0] - }) - .collect(), - payload, - }); - } - } - } - H0TupleReader::_2(h0_tuple) => { - let computing = computing.get_or_insert_with(Vec::new); - computing.extend_from_slice(h0_tuple.elements()); - } - } - } - current = h0_guard.get_opaque().next; - drop(h0_guard); - } - - let tape = tape.into_inner(); - let new = tape.first(); - drop(tape); - - *jump_tuple.first() = new; - drop(jump_guard); - - freepages::mark(relation.clone(), freepage_first, &trace); - } -} diff --git a/src/algorithm/vectors.rs b/src/algorithm/vectors.rs deleted file mode 100644 index d71499bb..00000000 --- a/src/algorithm/vectors.rs +++ /dev/null @@ -1,133 +0,0 @@ -use crate::algorithm::operator::*; -use crate::algorithm::tuples::*; -use crate::algorithm::{Page, PageGuard, RelationRead, RelationWrite}; -use crate::utils::pipe::Pipe; -use std::num::NonZeroU64; -use vector::VectorOwned; - -pub fn vector_access_1< - O: Operator, - A: Accessor1<::Element, ::Metadata>, ->( - relation: impl RelationRead, - mean: IndexPointer, - accessor: A, -) -> A::Output { - let mut cursor = Err(mean); - let mut result = accessor; - while let Err(mean) = cursor.map_err(pointer_to_pair) { - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get(mean.1) - .expect("data corruption") - .pipe(read_tuple::>); - if vector_tuple.payload().is_some() { - panic!("data corruption"); - } - result.push(vector_tuple.elements()); - cursor = vector_tuple.metadata_or_pointer(); - } - result.finish(cursor.expect("data corruption")) -} - -pub fn vector_access_0< - O: Operator, - A: Accessor1<::Element, ::Metadata>, ->( - relation: impl RelationRead, - mean: IndexPointer, - payload: NonZeroU64, - accessor: A, -) -> Option { - let mut cursor = Err(mean); - let mut result = accessor; - while let Err(mean) = cursor.map_err(pointer_to_pair) { - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get(mean.1)? - .pipe(read_tuple::>); - if vector_tuple.payload().is_none() { - panic!("data corruption"); - } - if vector_tuple.payload() != Some(payload) { - return None; - } - result.push(vector_tuple.elements()); - cursor = vector_tuple.metadata_or_pointer(); - } - Some(result.finish(cursor.ok()?)) -} - -pub fn vector_append( - relation: impl RelationWrite + Clone, - vectors_first: u32, - vector: ::Borrowed<'_>, - payload: NonZeroU64, -) -> IndexPointer { - fn append(relation: impl RelationWrite, first: u32, bytes: &[u8]) -> IndexPointer { - if let Some(mut write) = relation.search(bytes.len()) { - let i = write.alloc(bytes).unwrap(); - return pair_to_pointer((write.id(), i)); - } - assert!(first != u32::MAX); - let mut current = first; - loop { - let read = relation.read(current); - if read.freespace() as usize >= bytes.len() || read.get_opaque().next == u32::MAX { - drop(read); - let mut write = relation.write(current, true); - if let Some(i) = write.alloc(bytes) { - return pair_to_pointer((current, i)); - } - if write.get_opaque().next == u32::MAX { - let mut extend = relation.extend(true); - write.get_opaque_mut().next = extend.id(); - drop(write); - if let Some(i) = extend.alloc(bytes) { - let result = (extend.id(), i); - drop(extend); - let mut past = relation.write(first, true); - let skip = &mut past.get_opaque_mut().skip; - assert!(*skip != u32::MAX); - *skip = std::cmp::max(*skip, result.0); - return pair_to_pointer(result); - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - } - if current == first && write.get_opaque().skip != first { - current = write.get_opaque().skip; - } else { - current = write.get_opaque().next; - } - } else { - if current == first && read.get_opaque().skip != first { - current = read.get_opaque().skip; - } else { - current = read.get_opaque().next; - } - } - } - } - let (metadata, slices) = O::Vector::vector_split(vector); - let mut chain = Ok(metadata); - for i in (0..slices.len()).rev() { - chain = Err(append( - relation.clone(), - vectors_first, - &serialize::>(&match chain { - Ok(metadata) => VectorTuple::_0 { - elements: slices[i].to_vec(), - payload: Some(payload), - metadata, - }, - Err(pointer) => VectorTuple::_1 { - elements: slices[i].to_vec(), - payload: Some(payload), - pointer, - }, - }), - )); - } - chain.err().unwrap() -} diff --git a/src/bin/pgrx_embed.rs b/src/bin/pgrx_embed.rs index 5f5c4d85..afd0164c 100644 --- a/src/bin/pgrx_embed.rs +++ b/src/bin/pgrx_embed.rs @@ -1 +1,2 @@ +#![allow(unsafe_code)] ::pgrx::pgrx_embed!(); diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs index b60f6c5f..3e9fe09f 100644 --- a/src/datatype/memory_halfvec.rs +++ b/src/datatype/memory_halfvec.rs @@ -1,20 +1,14 @@ use half::f16; -use pgrx::datum::FromDatum; -use pgrx::datum::IntoDatum; -use pgrx::pg_sys::Datum; -use pgrx::pg_sys::Oid; -use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; -use pgrx::pgrx_sql_entity_graph::metadata::Returns; -use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; -use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; -use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; +use pgrx::datum::{FromDatum, IntoDatum}; +use pgrx::pg_sys::{Datum, Oid}; +use pgrx::pgrx_sql_entity_graph::metadata::*; use std::marker::PhantomData; use std::ptr::NonNull; use vector::VectorBorrowed; use vector::vect::VectBorrowed; #[repr(C, align(8))] -pub struct HalfvecHeader { +struct HalfvecHeader { varlena: u32, dims: u16, unused: u16, @@ -28,10 +22,10 @@ impl HalfvecHeader { } (size_of::() + size_of::() * len).next_multiple_of(8) } - pub unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f16> { + unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f16> { unsafe { let this = this.as_ptr(); - VectBorrowed::new_unchecked(std::slice::from_raw_parts( + VectBorrowed::new(std::slice::from_raw_parts( (&raw const (*this).elements).cast(), (&raw const (*this).dims).read() as usize, )) @@ -93,7 +87,7 @@ impl HalfvecOutput { pub fn as_borrowed(&self) -> VectBorrowed<'_, f16> { unsafe { HalfvecHeader::as_borrowed(self.0) } } - pub fn into_raw(self) -> *mut HalfvecHeader { + fn into_raw(self) -> *mut HalfvecHeader { let result = self.0.as_ptr(); std::mem::forget(self); result diff --git a/src/datatype/memory_scalar8.rs b/src/datatype/memory_scalar8.rs index 4f306548..19e4ff47 100644 --- a/src/datatype/memory_scalar8.rs +++ b/src/datatype/memory_scalar8.rs @@ -1,19 +1,13 @@ -use pgrx::datum::FromDatum; -use pgrx::datum::IntoDatum; -use pgrx::pg_sys::Datum; -use pgrx::pg_sys::Oid; -use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; -use pgrx::pgrx_sql_entity_graph::metadata::Returns; -use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; -use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; -use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; +use pgrx::datum::{FromDatum, IntoDatum}; +use pgrx::pg_sys::{Datum, Oid}; +use pgrx::pgrx_sql_entity_graph::metadata::*; use std::marker::PhantomData; use std::ptr::NonNull; use vector::VectorBorrowed; use vector::scalar8::Scalar8Borrowed; #[repr(C, align(8))] -pub struct Scalar8Header { +struct Scalar8Header { varlena: u32, dims: u16, unused: u16, @@ -31,10 +25,10 @@ impl Scalar8Header { } (size_of::() + size_of::() * len).next_multiple_of(8) } - pub unsafe fn as_borrowed<'a>(this: NonNull) -> Scalar8Borrowed<'a> { + unsafe fn as_borrowed<'a>(this: NonNull) -> Scalar8Borrowed<'a> { unsafe { let this = this.as_ptr(); - Scalar8Borrowed::new_unchecked( + Scalar8Borrowed::new( (&raw const (*this).sum_of_x2).read(), (&raw const (*this).k).read(), (&raw const (*this).b).read(), @@ -105,7 +99,7 @@ impl Scalar8Output { pub fn as_borrowed(&self) -> Scalar8Borrowed<'_> { unsafe { Scalar8Header::as_borrowed(self.0) } } - pub fn into_raw(self) -> *mut Scalar8Header { + fn into_raw(self) -> *mut Scalar8Header { let result = self.0.as_ptr(); std::mem::forget(self); result diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs index de70ba16..4d9f9f21 100644 --- a/src/datatype/memory_vector.rs +++ b/src/datatype/memory_vector.rs @@ -1,19 +1,13 @@ -use pgrx::datum::FromDatum; -use pgrx::datum::IntoDatum; -use pgrx::pg_sys::Datum; -use pgrx::pg_sys::Oid; -use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; -use pgrx::pgrx_sql_entity_graph::metadata::Returns; -use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; -use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; -use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; +use pgrx::datum::{FromDatum, IntoDatum}; +use pgrx::pg_sys::{Datum, Oid}; +use pgrx::pgrx_sql_entity_graph::metadata::*; use std::marker::PhantomData; use std::ptr::NonNull; use vector::VectorBorrowed; use vector::vect::VectBorrowed; #[repr(C, align(8))] -pub struct VectorHeader { +struct VectorHeader { varlena: u32, dims: u16, unused: u16, @@ -27,10 +21,10 @@ impl VectorHeader { } (size_of::() + size_of::() * len).next_multiple_of(8) } - pub unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f32> { + unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f32> { unsafe { let this = this.as_ptr(); - VectBorrowed::new_unchecked(std::slice::from_raw_parts( + VectBorrowed::new(std::slice::from_raw_parts( (&raw const (*this).elements).cast(), (&raw const (*this).dims).read() as usize, )) @@ -92,7 +86,7 @@ impl VectorOutput { pub fn as_borrowed(&self) -> VectBorrowed<'_, f32> { unsafe { VectorHeader::as_borrowed(self.0) } } - pub fn into_raw(self) -> *mut VectorHeader { + fn into_raw(self) -> *mut VectorHeader { let result = self.0.as_ptr(); std::mem::forget(self); result diff --git a/src/gucs/mod.rs b/src/gucs/mod.rs deleted file mode 100644 index 2fb489e1..00000000 --- a/src/gucs/mod.rs +++ /dev/null @@ -1,14 +0,0 @@ -pub mod executing; -pub mod prewarm; - -pub unsafe fn init() { - unsafe { - executing::init(); - prewarm::init(); - prewarm::prewarm(); - #[cfg(any(feature = "pg13", feature = "pg14"))] - pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); - #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] - pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrq".as_ptr()); - } -} diff --git a/src/gucs/prewarm.rs b/src/gucs/prewarm.rs deleted file mode 100644 index bc484367..00000000 --- a/src/gucs/prewarm.rs +++ /dev/null @@ -1,32 +0,0 @@ -use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; -use std::ffi::CStr; - -static PREWARM_DIM: GucSetting> = - GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); - -pub unsafe fn init() { - GucRegistry::define_string_guc( - "vchordrq.prewarm_dim", - "prewarm_dim when the extension is loading.", - "prewarm_dim when the extension is loading.", - &PREWARM_DIM, - GucContext::Userset, - GucFlags::default(), - ); -} - -pub fn prewarm() { - if let Some(prewarm_dim) = PREWARM_DIM.get() { - if let Ok(prewarm_dim) = prewarm_dim.to_str() { - for dim in prewarm_dim.split(',') { - if let Ok(dim) = dim.trim().parse::() { - crate::projection::prewarm(dim as _); - } else { - pgrx::warning!("{dim:?} is not a valid integer"); - } - } - } else { - pgrx::warning!("vchordrq.prewarm_dim is not a valid UTF-8 string"); - } - } -} diff --git a/src/index/am.rs b/src/index/am.rs deleted file mode 100644 index 5db07d24..00000000 --- a/src/index/am.rs +++ /dev/null @@ -1,1103 +0,0 @@ -use crate::algorithm; -use crate::algorithm::build::{HeapRelation, Reporter}; -use crate::algorithm::operator::{Dot, L2, Op}; -use crate::algorithm::operator::{Operator, Vector}; -use crate::index::am_options::{Opfamily, Reloption}; -use crate::index::am_scan::Scanner; -use crate::index::utils::{ctid_to_pointer, pointer_to_ctid}; -use crate::index::{am_options, am_scan}; -use crate::postgres::PostgresRelation; -use crate::types::{DistanceKind, VectorKind}; -use half::f16; -use pgrx::datum::Internal; -use pgrx::pg_sys::Datum; -use std::num::NonZeroU64; -use vector::vect::VectOwned; - -static mut RELOPT_KIND_VCHORDRQ: pgrx::pg_sys::relopt_kind::Type = 0; - -pub unsafe fn init() { - unsafe { - (&raw mut RELOPT_KIND_VCHORDRQ).write(pgrx::pg_sys::add_reloption_kind()); - pgrx::pg_sys::add_string_reloption( - (&raw const RELOPT_KIND_VCHORDRQ).read(), - c"options".as_ptr(), - c"Vector index options, represented as a TOML string.".as_ptr(), - c"".as_ptr(), - None, - pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, - ); - } -} - -#[pgrx::pg_extern(sql = "")] -fn _vchordrq_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { - type T = pgrx::pg_sys::IndexAmRoutine; - unsafe { - let index_am_routine = pgrx::pg_sys::palloc0(size_of::()) as *mut T; - index_am_routine.write(AM_HANDLER); - Internal::from(Some(Datum::from(index_am_routine))) - } -} - -const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = { - let mut am_routine = - unsafe { std::mem::MaybeUninit::::zeroed().assume_init() }; - - am_routine.type_ = pgrx::pg_sys::NodeTag::T_IndexAmRoutine; - - am_routine.amsupport = 1; - am_routine.amcanorderbyop = true; - - #[cfg(feature = "pg17")] - { - am_routine.amcanbuildparallel = true; - } - - // Index access methods that set `amoptionalkey` to `false` - // must index all tuples, even if the first column is `NULL`. - // However, PostgreSQL does not generate a path if there is no - // index clauses, even if there is a `ORDER BY` clause. - // So we have to set it to `true` and set costs of every path - // for vector index scans without `ORDER BY` clauses a large number - // and throw errors if someone really wants such a path. - am_routine.amoptionalkey = true; - - am_routine.amvalidate = Some(amvalidate); - am_routine.amoptions = Some(amoptions); - am_routine.amcostestimate = Some(amcostestimate); - - am_routine.ambuild = Some(ambuild); - am_routine.ambuildempty = Some(ambuildempty); - am_routine.aminsert = Some(aminsert); - am_routine.ambulkdelete = Some(ambulkdelete); - am_routine.amvacuumcleanup = Some(amvacuumcleanup); - - am_routine.ambeginscan = Some(ambeginscan); - am_routine.amrescan = Some(amrescan); - am_routine.amgettuple = Some(amgettuple); - am_routine.amendscan = Some(amendscan); - - am_routine -}; - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { - true -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amoptions(reloptions: Datum, validate: bool) -> *mut pgrx::pg_sys::bytea { - let rdopts = unsafe { - pgrx::pg_sys::build_reloptions( - reloptions, - validate, - (&raw const RELOPT_KIND_VCHORDRQ).read(), - size_of::(), - Reloption::TAB.as_ptr(), - Reloption::TAB.len() as _, - ) - }; - rdopts as *mut pgrx::pg_sys::bytea -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amcostestimate( - _root: *mut pgrx::pg_sys::PlannerInfo, - path: *mut pgrx::pg_sys::IndexPath, - _loop_count: f64, - index_startup_cost: *mut pgrx::pg_sys::Cost, - index_total_cost: *mut pgrx::pg_sys::Cost, - index_selectivity: *mut pgrx::pg_sys::Selectivity, - index_correlation: *mut f64, - index_pages: *mut f64, -) { - unsafe { - if (*path).indexorderbys.is_null() && (*path).indexclauses.is_null() { - *index_startup_cost = f64::MAX; - *index_total_cost = f64::MAX; - *index_selectivity = 0.0; - *index_correlation = 0.0; - *index_pages = 0.0; - return; - } - *index_startup_cost = 0.0; - *index_total_cost = 0.0; - *index_selectivity = 1.0; - *index_correlation = 1.0; - *index_pages = 0.0; - } -} - -#[derive(Debug, Clone)] -struct PgReporter {} - -impl Reporter for PgReporter { - fn tuples_total(&mut self, tuples_total: u64) { - unsafe { - pgrx::pg_sys::pgstat_progress_update_param( - pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, - tuples_total as _, - ); - } - } -} - -impl PgReporter { - fn tuples_done(&mut self, tuples_done: u64) { - unsafe { - pgrx::pg_sys::pgstat_progress_update_param( - pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, - tuples_done as _, - ); - } - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambuild( - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, -) -> *mut pgrx::pg_sys::IndexBuildResult { - use validator::Validate; - #[derive(Debug, Clone)] - pub struct Heap { - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, - opfamily: Opfamily, - } - impl HeapRelation for Heap { - fn traverse(&self, progress: bool, callback: F) - where - F: FnMut((NonZeroU64, O::Vector)), - { - pub struct State<'a, F> { - pub this: &'a Heap, - pub callback: F, - } - #[pgrx::pg_guard] - unsafe extern "C" fn call( - _index: pgrx::pg_sys::Relation, - ctid: pgrx::pg_sys::ItemPointer, - values: *mut Datum, - is_null: *mut bool, - _tuple_is_alive: bool, - state: *mut core::ffi::c_void, - ) where - F: FnMut((NonZeroU64, O::Vector)), - { - let state = unsafe { &mut *state.cast::>() }; - let opfamily = state.this.opfamily; - let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; - let pointer = unsafe { ctid_to_pointer(ctid.read()) }; - if let Some(vector) = vector { - (state.callback)((pointer, O::Vector::from_owned(vector))); - } - } - let table_am = unsafe { &*(*self.heap).rd_tableam }; - let mut state = State { - this: self, - callback, - }; - unsafe { - table_am.index_build_range_scan.unwrap()( - self.heap, - self.index, - self.index_info, - true, - false, - progress, - 0, - pgrx::pg_sys::InvalidBlockNumber, - Some(call::), - (&mut state) as *mut State as *mut _, - std::ptr::null_mut(), - ); - } - } - - fn opfamily(&self) -> Opfamily { - self.opfamily - } - } - let (vector_options, vchordrq_options) = unsafe { am_options::options(index) }; - if let Err(errors) = Validate::validate(&vector_options) { - pgrx::error!("error while validating options: {}", errors); - } - if vector_options.dims == 0 { - pgrx::error!("error while validating options: dimension cannot be 0"); - } - if vector_options.dims > 60000 { - pgrx::error!("error while validating options: dimension is too large"); - } - if let Err(errors) = Validate::validate(&vchordrq_options) { - pgrx::error!("error while validating options: {}", errors); - } - let opfamily = unsafe { am_options::opfamily(index) }; - let heap_relation = Heap { - heap, - index, - index_info, - opfamily, - }; - let mut reporter = PgReporter {}; - let index_relation = unsafe { PostgresRelation::new(index) }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::build::build::, L2>, Heap, _>( - vector_options, - vchordrq_options, - heap_relation.clone(), - index_relation.clone(), - reporter.clone(), - ) - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::build::build::, Dot>, Heap, _>( - vector_options, - vchordrq_options, - heap_relation.clone(), - index_relation.clone(), - reporter.clone(), - ) - } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::build::build::, L2>, Heap, _>( - vector_options, - vchordrq_options, - heap_relation.clone(), - index_relation.clone(), - reporter.clone(), - ) - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::build::build::, Dot>, Heap, _>( - vector_options, - vchordrq_options, - heap_relation.clone(), - index_relation.clone(), - reporter.clone(), - ) - } - } - if let Some(leader) = unsafe { VchordrqLeader::enter(heap, index, (*index_info).ii_Concurrent) } - { - unsafe { - parallel_build( - index, - heap, - index_info, - leader.tablescandesc, - leader.vchordrqshared, - Some(reporter), - ); - leader.wait(); - let nparticipants = leader.nparticipants; - loop { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordrqshared).mutex); - if (*leader.vchordrqshared).nparticipantsdone == nparticipants { - pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); - break; - } - pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); - pgrx::pg_sys::ConditionVariableSleep( - &raw mut (*leader.vchordrqshared).workersdonecv, - pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN, - ); - } - pgrx::pg_sys::ConditionVariableCancelSleep(); - } - } else { - let mut indtuples = 0; - reporter.tuples_done(indtuples); - let relation = unsafe { PostgresRelation::new(index) }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - HeapRelation::, L2>>::traverse( - &heap_relation, - true, - |(pointer, vector)| { - algorithm::insert::insert::, L2>>( - relation.clone(), - pointer, - vector, - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }, - ); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - HeapRelation::, Dot>>::traverse( - &heap_relation, - true, - |(pointer, vector)| { - algorithm::insert::insert::, Dot>>( - relation.clone(), - pointer, - vector, - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }, - ); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - HeapRelation::, L2>>::traverse( - &heap_relation, - true, - |(pointer, vector)| { - algorithm::insert::insert::, L2>>( - relation.clone(), - pointer, - vector, - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }, - ); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - HeapRelation::, Dot>>::traverse( - &heap_relation, - true, - |(pointer, vector)| { - algorithm::insert::insert::, Dot>>( - relation.clone(), - pointer, - vector, - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }, - ); - } - } - } - let relation = unsafe { PostgresRelation::new(index) }; - let delay = || { - pgrx::check_for_interrupts!(); - }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - type O = Op, L2>; - algorithm::vacuum::maintain::(relation, delay); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - type O = Op, Dot>; - algorithm::vacuum::maintain::(relation, delay); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - type O = Op, L2>; - algorithm::vacuum::maintain::(relation, delay); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - type O = Op, Dot>; - algorithm::vacuum::maintain::(relation, delay); - } - } - unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } -} - -struct VchordrqShared { - /* Immutable state */ - heaprelid: pgrx::pg_sys::Oid, - indexrelid: pgrx::pg_sys::Oid, - isconcurrent: bool, - - /* Worker progress */ - workersdonecv: pgrx::pg_sys::ConditionVariable, - - /* Mutex for mutable state */ - mutex: pgrx::pg_sys::slock_t, - - /* Mutable state */ - nparticipantsdone: i32, - indtuples: u64, -} - -fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { - matches!( - unsafe { (*snapshot).snapshot_type }, - pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC - | pgrx::pg_sys::SnapshotType::SNAPSHOT_HISTORIC_MVCC - ) -} - -struct VchordrqLeader { - pcxt: *mut pgrx::pg_sys::ParallelContext, - nparticipants: i32, - vchordrqshared: *mut VchordrqShared, - tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, - snapshot: pgrx::pg_sys::Snapshot, -} - -impl VchordrqLeader { - pub unsafe fn enter( - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - isconcurrent: bool, - ) -> Option { - unsafe fn compute_parallel_workers( - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - ) -> i32 { - unsafe { - if pgrx::pg_sys::plan_create_index_workers((*heap).rd_id, (*index).rd_id) == 0 { - return 0; - } - if !(*heap).rd_options.is_null() { - let std_options = (*heap).rd_options.cast::(); - std::cmp::min( - (*std_options).parallel_workers, - pgrx::pg_sys::max_parallel_maintenance_workers, - ) - } else { - pgrx::pg_sys::max_parallel_maintenance_workers - } - } - } - - let request = unsafe { compute_parallel_workers(heap, index) }; - if request <= 0 { - return None; - } - - unsafe { - pgrx::pg_sys::EnterParallelMode(); - } - let pcxt = unsafe { - pgrx::pg_sys::CreateParallelContext( - c"vchord".as_ptr(), - c"vchordrq_parallel_build_main".as_ptr(), - request, - ) - }; - - let snapshot = if isconcurrent { - unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } - } else { - &raw mut pgrx::pg_sys::SnapshotAnyData - }; - - fn estimate_chunk(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { - e.space_for_chunks += x.next_multiple_of(pgrx::pg_sys::ALIGNOF_BUFFER as _); - } - fn estimate_keys(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { - e.number_of_keys += x; - } - let est_tablescandesc = - unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap, snapshot) }; - unsafe { - estimate_chunk(&mut (*pcxt).estimator, size_of::()); - estimate_keys(&mut (*pcxt).estimator, 1); - estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); - estimate_keys(&mut (*pcxt).estimator, 1); - } - - unsafe { - pgrx::pg_sys::InitializeParallelDSM(pcxt); - if (*pcxt).seg.is_null() { - if is_mvcc_snapshot(snapshot) { - pgrx::pg_sys::UnregisterSnapshot(snapshot); - } - pgrx::pg_sys::DestroyParallelContext(pcxt); - pgrx::pg_sys::ExitParallelMode(); - return None; - } - } - - let vchordrqshared = unsafe { - let vchordrqshared = - pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) - .cast::(); - vchordrqshared.write(VchordrqShared { - heaprelid: (*heap).rd_id, - indexrelid: (*index).rd_id, - isconcurrent, - workersdonecv: std::mem::zeroed(), - mutex: std::mem::zeroed(), - nparticipantsdone: 0, - indtuples: 0, - }); - pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).workersdonecv); - pgrx::pg_sys::SpinLockInit(&raw mut (*vchordrqshared).mutex); - vchordrqshared - }; - - let tablescandesc = unsafe { - let tablescandesc = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, est_tablescandesc) - .cast::(); - pgrx::pg_sys::table_parallelscan_initialize(heap, tablescandesc, snapshot); - tablescandesc - }; - - unsafe { - pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000001, vchordrqshared.cast()); - pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); - } - - unsafe { - pgrx::pg_sys::LaunchParallelWorkers(pcxt); - } - - let nworkers_launched = unsafe { (*pcxt).nworkers_launched }; - - unsafe { - if nworkers_launched == 0 { - pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); - if is_mvcc_snapshot(snapshot) { - pgrx::pg_sys::UnregisterSnapshot(snapshot); - } - pgrx::pg_sys::DestroyParallelContext(pcxt); - pgrx::pg_sys::ExitParallelMode(); - return None; - } - } - - Some(Self { - pcxt, - nparticipants: nworkers_launched + 1, - vchordrqshared, - tablescandesc, - snapshot, - }) - } - - pub fn wait(&self) { - unsafe { - pgrx::pg_sys::WaitForParallelWorkersToAttach(self.pcxt); - } - } -} - -impl Drop for VchordrqLeader { - fn drop(&mut self) { - if !std::thread::panicking() { - unsafe { - pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); - if is_mvcc_snapshot(self.snapshot) { - pgrx::pg_sys::UnregisterSnapshot(self.snapshot); - } - pgrx::pg_sys::DestroyParallelContext(self.pcxt); - pgrx::pg_sys::ExitParallelMode(); - } - } - } -} - -#[pgrx::pg_guard] -#[unsafe(no_mangle)] -pub unsafe extern "C" fn vchordrq_parallel_build_main( - _seg: *mut pgrx::pg_sys::dsm_segment, - toc: *mut pgrx::pg_sys::shm_toc, -) { - let vchordrqshared = unsafe { - pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() - }; - let tablescandesc = unsafe { - pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) - .cast::() - }; - let heap_lockmode; - let index_lockmode; - if unsafe { !(*vchordrqshared).isconcurrent } { - heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; - } else { - heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; - } - let heap = unsafe { pgrx::pg_sys::table_open((*vchordrqshared).heaprelid, heap_lockmode) }; - let index = unsafe { pgrx::pg_sys::index_open((*vchordrqshared).indexrelid, index_lockmode) }; - let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; - unsafe { - (*index_info).ii_Concurrent = (*vchordrqshared).isconcurrent; - } - - unsafe { - parallel_build(index, heap, index_info, tablescandesc, vchordrqshared, None); - } - - unsafe { - pgrx::pg_sys::index_close(index, index_lockmode); - pgrx::pg_sys::table_close(heap, heap_lockmode); - } -} - -unsafe fn parallel_build( - index: *mut pgrx::pg_sys::RelationData, - heap: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, - tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, - vchordrqshared: *mut VchordrqShared, - mut reporter: Option, -) { - #[derive(Debug, Clone)] - pub struct Heap { - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, - opfamily: Opfamily, - scan: *mut pgrx::pg_sys::TableScanDescData, - } - impl HeapRelation for Heap { - fn traverse(&self, progress: bool, callback: F) - where - F: FnMut((NonZeroU64, O::Vector)), - { - pub struct State<'a, F> { - pub this: &'a Heap, - pub callback: F, - } - #[pgrx::pg_guard] - unsafe extern "C" fn call( - _index: pgrx::pg_sys::Relation, - ctid: pgrx::pg_sys::ItemPointer, - values: *mut Datum, - is_null: *mut bool, - _tuple_is_alive: bool, - state: *mut core::ffi::c_void, - ) where - F: FnMut((NonZeroU64, O::Vector)), - { - let state = unsafe { &mut *state.cast::>() }; - let opfamily = state.this.opfamily; - let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; - let pointer = unsafe { ctid_to_pointer(ctid.read()) }; - if let Some(vector) = vector { - (state.callback)((pointer, O::Vector::from_owned(vector))); - } - } - let table_am = unsafe { &*(*self.heap).rd_tableam }; - let mut state = State { - this: self, - callback, - }; - unsafe { - table_am.index_build_range_scan.unwrap()( - self.heap, - self.index, - self.index_info, - true, - false, - progress, - 0, - pgrx::pg_sys::InvalidBlockNumber, - Some(call::), - (&mut state) as *mut State as *mut _, - self.scan, - ); - } - } - - fn opfamily(&self) -> Opfamily { - self.opfamily - } - } - - let index_relation = unsafe { PostgresRelation::new(index) }; - - let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap, tablescandesc) }; - let opfamily = unsafe { am_options::opfamily(index) }; - let heap_relation = Heap { - heap, - index, - index_info, - opfamily, - scan, - }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - HeapRelation::, L2>>::traverse( - &heap_relation, - true, - |(pointer, vector)| { - algorithm::insert::insert::, L2>>( - index_relation.clone(), - pointer, - vector, - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }, - ); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - HeapRelation::, Dot>>::traverse( - &heap_relation, - true, - |(pointer, vector)| { - algorithm::insert::insert::, Dot>>( - index_relation.clone(), - pointer, - vector, - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }, - ); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - HeapRelation::, L2>>::traverse( - &heap_relation, - true, - |(pointer, vector)| { - algorithm::insert::insert::, L2>>( - index_relation.clone(), - pointer, - vector, - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }, - ); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - HeapRelation::, Dot>>::traverse( - &heap_relation, - true, - |(pointer, vector)| { - algorithm::insert::insert::, Dot>>( - index_relation.clone(), - pointer, - vector, - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }, - ); - } - } - unsafe { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).nparticipantsdone += 1; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordrqshared).workersdonecv); - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambuildempty(_index: pgrx::pg_sys::Relation) { - pgrx::error!("Unlogged indexes are not supported."); -} - -#[cfg(feature = "pg13")] -#[pgrx::pg_guard] -pub unsafe extern "C" fn aminsert( - index: pgrx::pg_sys::Relation, - values: *mut Datum, - is_null: *mut bool, - heap_tid: pgrx::pg_sys::ItemPointer, - _heap: pgrx::pg_sys::Relation, - _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, - _index_info: *mut pgrx::pg_sys::IndexInfo, -) -> bool { - let opfamily = unsafe { am_options::opfamily(index) }; - let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; - if let Some(vector) = vector { - let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::insert::insert::, L2>>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - ) - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::insert::insert::, Dot>>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - ) - } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::insert::insert::, L2>>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - ) - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::insert::insert::, Dot>>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - ) - } - } - } - false -} - -#[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] -#[pgrx::pg_guard] -pub unsafe extern "C" fn aminsert( - index: pgrx::pg_sys::Relation, - values: *mut Datum, - is_null: *mut bool, - heap_tid: pgrx::pg_sys::ItemPointer, - _heap: pgrx::pg_sys::Relation, - _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, - _index_unchanged: bool, - _index_info: *mut pgrx::pg_sys::IndexInfo, -) -> bool { - let opfamily = unsafe { am_options::opfamily(index) }; - let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; - if let Some(vector) = vector { - let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::insert::insert::, L2>>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - ) - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::insert::insert::, Dot>>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - ) - } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::insert::insert::, L2>>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - ) - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::insert::insert::, Dot>>( - unsafe { PostgresRelation::new(index) }, - pointer, - VectOwned::::from_owned(vector), - ) - } - } - } - false -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambeginscan( - index: pgrx::pg_sys::Relation, - n_keys: std::os::raw::c_int, - n_orderbys: std::os::raw::c_int, -) -> pgrx::pg_sys::IndexScanDesc { - use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; - - let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index, n_keys, n_orderbys) }; - unsafe { - let scanner = am_scan::scan_make(None, None, false); - (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); - } - scan -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amrescan( - scan: pgrx::pg_sys::IndexScanDesc, - keys: pgrx::pg_sys::ScanKey, - _n_keys: std::os::raw::c_int, - orderbys: pgrx::pg_sys::ScanKey, - _n_orderbys: std::os::raw::c_int, -) { - unsafe { - if !keys.is_null() && (*scan).numberOfKeys > 0 { - std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); - } - if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { - std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); - } - let opfamily = am_options::opfamily((*scan).indexRelation); - let (orderbys, spheres) = { - let mut orderbys = Vec::new(); - let mut spheres = Vec::new(); - if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { - pgrx::error!( - "vector search with no WHERE clause and no ORDER BY clause is not supported" - ); - } - for i in 0..(*scan).numberOfOrderBys { - let data = (*scan).orderByData.add(i as usize); - let value = (*data).sk_argument; - let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; - match (*data).sk_strategy { - 1 => orderbys.push(opfamily.datum_to_vector(value, is_null)), - _ => unreachable!(), - } - } - for i in 0..(*scan).numberOfKeys { - let data = (*scan).keyData.add(i as usize); - let value = (*data).sk_argument; - let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; - match (*data).sk_strategy { - 2 => spheres.push(opfamily.datum_to_sphere(value, is_null)), - _ => unreachable!(), - } - } - (orderbys, spheres) - }; - let (vector, threshold, recheck) = am_scan::scan_build(orderbys, spheres, opfamily); - let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace(scanner, am_scan::scan_make(vector, threshold, recheck)); - am_scan::scan_release(scanner); - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amgettuple( - scan: pgrx::pg_sys::IndexScanDesc, - direction: pgrx::pg_sys::ScanDirection::Type, -) -> bool { - if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { - pgrx::error!("vector search without a forward scan direction is not supported"); - } - // https://www.postgresql.org/docs/current/index-locking.html - // If heap entries referenced physical pointers are deleted before - // they are consumed by PostgreSQL, PostgreSQL will received wrong - // physical pointers: no rows or irreverent rows are referenced. - if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC - { - pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); - } - let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; - let relation = unsafe { PostgresRelation::new((*scan).indexRelation) }; - if let Some((pointer, recheck)) = am_scan::scan_next(scanner, relation) { - let ctid = pointer_to_ctid(pointer); - unsafe { - (*scan).xs_heaptid = ctid; - (*scan).xs_recheckorderby = false; - (*scan).xs_recheck = recheck; - } - true - } else { - false - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { - unsafe { - let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace(scanner, am_scan::scan_make(None, None, false)); - am_scan::scan_release(scanner); - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambulkdelete( - info: *mut pgrx::pg_sys::IndexVacuumInfo, - stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, - callback: pgrx::pg_sys::IndexBulkDeleteCallback, - callback_state: *mut std::os::raw::c_void, -) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { - let mut stats = stats; - if stats.is_null() { - stats = unsafe { - pgrx::pg_sys::palloc0(size_of::()).cast() - }; - } - let opfamily = unsafe { am_options::opfamily((*info).index) }; - let callback = callback.unwrap(); - let callback = |p: NonZeroU64| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; - let index = unsafe { PostgresRelation::new((*info).index) }; - let delay = || unsafe { - pgrx::pg_sys::vacuum_delay_point(); - }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - type O = Op, L2>; - algorithm::vacuum::bulkdelete::(index.clone(), delay, callback); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - type O = Op, Dot>; - algorithm::vacuum::bulkdelete::(index.clone(), delay, callback); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - type O = Op, L2>; - algorithm::vacuum::bulkdelete::(index.clone(), delay, callback); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - type O = Op, Dot>; - algorithm::vacuum::bulkdelete::(index.clone(), delay, callback); - } - } - stats -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amvacuumcleanup( - info: *mut pgrx::pg_sys::IndexVacuumInfo, - _stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, -) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { - let opfamily = unsafe { am_options::opfamily((*info).index) }; - let index = unsafe { PostgresRelation::new((*info).index) }; - let delay = || unsafe { - pgrx::pg_sys::vacuum_delay_point(); - }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - type O = Op, L2>; - algorithm::vacuum::maintain::(index, delay); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - type O = Op, Dot>; - algorithm::vacuum::maintain::(index, delay); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - type O = Op, L2>; - algorithm::vacuum::maintain::(index, delay); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - type O = Op, Dot>; - algorithm::vacuum::maintain::(index, delay); - } - } - std::ptr::null_mut() -} diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs new file mode 100644 index 00000000..e7d65587 --- /dev/null +++ b/src/index/am/am_build.rs @@ -0,0 +1,950 @@ +use crate::datatype::typmod::Typmod; +use crate::index::am::{Reloption, ctid_to_pointer}; +use crate::index::opclass::{Opfamily, opfamily}; +use crate::index::projection::RandomProject; +use crate::index::storage::PostgresRelation; +use algorithm::operator::{Dot, L2, Op, Vector}; +use algorithm::types::*; +use half::f16; +use pgrx::pg_sys::Datum; +use rand::Rng; +use simd::Floating; +use std::num::NonZeroU64; +use std::sync::Arc; +use vector::vect::VectOwned; +use vector::{VectorBorrowed, VectorOwned}; + +#[derive(Debug, Clone)] +struct Heap { + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + opfamily: Opfamily, + scan: *mut pgrx::pg_sys::TableScanDescData, +} + +impl Heap { + fn traverse(&self, progress: bool, callback: F) { + pub struct State<'a, F> { + pub this: &'a Heap, + pub callback: F, + } + #[pgrx::pg_guard] + unsafe extern "C" fn call( + _index_relation: pgrx::pg_sys::Relation, + ctid: pgrx::pg_sys::ItemPointer, + values: *mut Datum, + is_null: *mut bool, + _tuple_is_alive: bool, + state: *mut core::ffi::c_void, + ) where + F: FnMut((NonZeroU64, V)), + { + let state = unsafe { &mut *state.cast::>() }; + let opfamily = state.this.opfamily; + let vector = unsafe { opfamily.input_vector(*values.add(0), *is_null.add(0)) }; + let pointer = unsafe { ctid_to_pointer(ctid.read()) }; + if let Some(vector) = vector { + (state.callback)((pointer, V::from_owned(vector))); + } + } + let table_am = unsafe { &*(*self.heap_relation).rd_tableam }; + let mut state = State { + this: self, + callback, + }; + unsafe { + table_am.index_build_range_scan.unwrap()( + self.heap_relation, + self.index_relation, + self.index_info, + true, + false, + progress, + 0, + pgrx::pg_sys::InvalidBlockNumber, + Some(call::), + (&mut state) as *mut State as *mut _, + self.scan, + ); + } + } +} + +#[derive(Debug, Clone)] +struct PostgresReporter {} + +impl PostgresReporter { + fn tuples_total(&mut self, tuples_total: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, + tuples_total as _, + ); + } + } + fn tuples_done(&mut self, tuples_done: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, + tuples_done as _, + ); + } + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn ambuild( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, +) -> *mut pgrx::pg_sys::IndexBuildResult { + use validator::Validate; + let (vector_options, vchordrq_options) = unsafe { options(index_relation) }; + if let Err(errors) = Validate::validate(&vector_options) { + pgrx::error!("error while validating options: {}", errors); + } + if vector_options.dims == 0 { + pgrx::error!("error while validating options: dimension cannot be 0"); + } + if vector_options.dims > 60000 { + pgrx::error!("error while validating options: dimension is too large"); + } + if let Err(errors) = Validate::validate(&vchordrq_options) { + pgrx::error!("error while validating options: {}", errors); + } + let opfamily = unsafe { opfamily(index_relation) }; + let heap = Heap { + heap_relation, + index_relation, + index_info, + opfamily, + scan: std::ptr::null_mut(), + }; + let index = unsafe { PostgresRelation::new(index_relation) }; + let mut reporter = PostgresReporter {}; + let structures = match vchordrq_options.build.clone() { + VchordrqBuildOptions::External(external_build) => { + make_external_build(vector_options.clone(), opfamily, external_build.clone()) + } + VchordrqBuildOptions::Internal(internal_build) => { + let mut tuples_total = 0_u64; + let samples = { + let mut rand = rand::thread_rng(); + let max_number_of_samples = internal_build + .lists + .last() + .unwrap() + .saturating_mul(internal_build.sampling_factor); + let mut samples = Vec::new(); + let mut number_of_samples = 0_u32; + match opfamily.vector_kind() { + VectorKind::Vecf32 => { + heap.traverse(false, |(_, vector): (_, VectOwned)| { + let vector = vector.as_borrowed(); + assert_eq!( + vector_options.dims, + vector.dims(), + "invalid vector dimensions" + ); + if number_of_samples < max_number_of_samples { + samples.push(VectOwned::::build_to_vecf32(vector)); + number_of_samples += 1; + } else { + let index = rand.gen_range(0..max_number_of_samples) as usize; + samples[index] = VectOwned::::build_to_vecf32(vector); + } + tuples_total += 1; + }); + } + VectorKind::Vecf16 => { + heap.traverse(false, |(_, vector): (_, VectOwned)| { + let vector = vector.as_borrowed(); + assert_eq!( + vector_options.dims, + vector.dims(), + "invalid vector dimensions" + ); + if number_of_samples < max_number_of_samples { + samples.push(VectOwned::::build_to_vecf32(vector)); + number_of_samples += 1; + } else { + let index = rand.gen_range(0..max_number_of_samples) as usize; + samples[index] = VectOwned::::build_to_vecf32(vector); + } + tuples_total += 1; + }); + } + } + samples + }; + reporter.tuples_total(tuples_total); + make_internal_build(vector_options.clone(), internal_build.clone(), samples) + } + }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => algorithm::build::, L2>>( + vector_options, + vchordrq_options, + index.clone(), + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + (VectorKind::Vecf32, DistanceKind::Dot) => algorithm::build::, Dot>>( + vector_options, + vchordrq_options, + index.clone(), + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + (VectorKind::Vecf16, DistanceKind::L2) => algorithm::build::, L2>>( + vector_options, + vchordrq_options, + index.clone(), + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + (VectorKind::Vecf16, DistanceKind::Dot) => algorithm::build::, Dot>>( + vector_options, + vchordrq_options, + index.clone(), + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + } + if let Some(leader) = + unsafe { VchordrqLeader::enter(heap_relation, index_relation, (*index_info).ii_Concurrent) } + { + unsafe { + parallel_build( + index_relation, + heap_relation, + index_info, + leader.tablescandesc, + leader.vchordrqshared, + Some(reporter), + ); + leader.wait(); + let nparticipants = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordrqshared).mutex); + if (*leader.vchordrqshared).nparticipantsdone == nparticipants { + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*leader.vchordrqshared).workersdonecv, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + } + } else { + let mut indtuples = 0; + reporter.tuples_done(indtuples); + let relation = unsafe { PostgresRelation::new(index_relation) }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, L2>>( + relation.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + indtuples += 1; + reporter.tuples_done(indtuples); + }); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, Dot>>( + relation.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + indtuples += 1; + reporter.tuples_done(indtuples); + }); + } + (VectorKind::Vecf16, DistanceKind::L2) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, L2>>( + relation.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + indtuples += 1; + reporter.tuples_done(indtuples); + }); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, Dot>>( + relation.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + indtuples += 1; + reporter.tuples_done(indtuples); + }); + } + } + } + let check = || { + pgrx::check_for_interrupts!(); + }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + algorithm::maintain::, L2>>(index, check); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + algorithm::maintain::, Dot>>(index, check); + } + (VectorKind::Vecf16, DistanceKind::L2) => { + algorithm::maintain::, L2>>(index, check); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + algorithm::maintain::, Dot>>(index, check); + } + } + unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } +} + +struct VchordrqShared { + /* Immutable state */ + heaprelid: pgrx::pg_sys::Oid, + indexrelid: pgrx::pg_sys::Oid, + isconcurrent: bool, + + /* Worker progress */ + workersdonecv: pgrx::pg_sys::ConditionVariable, + + /* Mutex for mutable state */ + mutex: pgrx::pg_sys::slock_t, + + /* Mutable state */ + nparticipantsdone: i32, + indtuples: u64, +} + +fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { + matches!( + unsafe { (*snapshot).snapshot_type }, + pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + | pgrx::pg_sys::SnapshotType::SNAPSHOT_HISTORIC_MVCC + ) +} + +struct VchordrqLeader { + pcxt: *mut pgrx::pg_sys::ParallelContext, + nparticipants: i32, + vchordrqshared: *mut VchordrqShared, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + snapshot: pgrx::pg_sys::Snapshot, +} + +impl VchordrqLeader { + pub unsafe fn enter( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + isconcurrent: bool, + ) -> Option { + unsafe fn compute_parallel_workers( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + ) -> i32 { + unsafe { + if pgrx::pg_sys::plan_create_index_workers( + (*heap_relation).rd_id, + (*index_relation).rd_id, + ) == 0 + { + return 0; + } + if !(*heap_relation).rd_options.is_null() { + let std_options = (*heap_relation) + .rd_options + .cast::(); + std::cmp::min( + (*std_options).parallel_workers, + pgrx::pg_sys::max_parallel_maintenance_workers, + ) + } else { + pgrx::pg_sys::max_parallel_maintenance_workers + } + } + } + + let request = unsafe { compute_parallel_workers(heap_relation, index_relation) }; + if request <= 0 { + return None; + } + + unsafe { + pgrx::pg_sys::EnterParallelMode(); + } + let pcxt = unsafe { + pgrx::pg_sys::CreateParallelContext( + c"vchord".as_ptr(), + c"vchordrq_parallel_build_main".as_ptr(), + request, + ) + }; + + let snapshot = if isconcurrent { + unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } + } else { + &raw mut pgrx::pg_sys::SnapshotAnyData + }; + + fn estimate_chunk(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.space_for_chunks += x.next_multiple_of(pgrx::pg_sys::ALIGNOF_BUFFER as _); + } + fn estimate_keys(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.number_of_keys += x; + } + let est_tablescandesc = + unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap_relation, snapshot) }; + unsafe { + estimate_chunk(&mut (*pcxt).estimator, size_of::()); + estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); + estimate_keys(&mut (*pcxt).estimator, 1); + } + + unsafe { + pgrx::pg_sys::InitializeParallelDSM(pcxt); + if (*pcxt).seg.is_null() { + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + let vchordrqshared = unsafe { + let vchordrqshared = + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) + .cast::(); + vchordrqshared.write(VchordrqShared { + heaprelid: (*heap_relation).rd_id, + indexrelid: (*index_relation).rd_id, + isconcurrent, + workersdonecv: std::mem::zeroed(), + mutex: std::mem::zeroed(), + nparticipantsdone: 0, + indtuples: 0, + }); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).workersdonecv); + pgrx::pg_sys::SpinLockInit(&raw mut (*vchordrqshared).mutex); + vchordrqshared + }; + + let tablescandesc = unsafe { + let tablescandesc = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, est_tablescandesc) + .cast::(); + pgrx::pg_sys::table_parallelscan_initialize(heap_relation, tablescandesc, snapshot); + tablescandesc + }; + + unsafe { + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000001, vchordrqshared.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); + } + + unsafe { + pgrx::pg_sys::LaunchParallelWorkers(pcxt); + } + + let nworkers_launched = unsafe { (*pcxt).nworkers_launched }; + + unsafe { + if nworkers_launched == 0 { + pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + Some(Self { + pcxt, + nparticipants: nworkers_launched + 1, + vchordrqshared, + tablescandesc, + snapshot, + }) + } + + pub fn wait(&self) { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToAttach(self.pcxt); + } + } +} + +impl Drop for VchordrqLeader { + fn drop(&mut self) { + if !std::thread::panicking() { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); + if is_mvcc_snapshot(self.snapshot) { + pgrx::pg_sys::UnregisterSnapshot(self.snapshot); + } + pgrx::pg_sys::DestroyParallelContext(self.pcxt); + pgrx::pg_sys::ExitParallelMode(); + } + } + } +} + +#[pgrx::pg_guard] +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vchordrq_parallel_build_main( + _seg: *mut pgrx::pg_sys::dsm_segment, + toc: *mut pgrx::pg_sys::shm_toc, +) { + let vchordrqshared = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() + }; + let tablescandesc = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) + .cast::() + }; + let heap_lockmode; + let index_lockmode; + if unsafe { !(*vchordrqshared).isconcurrent } { + heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; + index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; + } else { + heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; + index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; + } + let heap = unsafe { pgrx::pg_sys::table_open((*vchordrqshared).heaprelid, heap_lockmode) }; + let index = unsafe { pgrx::pg_sys::index_open((*vchordrqshared).indexrelid, index_lockmode) }; + let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; + unsafe { + (*index_info).ii_Concurrent = (*vchordrqshared).isconcurrent; + } + + unsafe { + parallel_build(index, heap, index_info, tablescandesc, vchordrqshared, None); + } + + unsafe { + pgrx::pg_sys::index_close(index, index_lockmode); + pgrx::pg_sys::table_close(heap, heap_lockmode); + } +} + +unsafe fn parallel_build( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + vchordrqshared: *mut VchordrqShared, + mut reporter: Option, +) { + let index = unsafe { PostgresRelation::new(index_relation) }; + + let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap_relation, tablescandesc) }; + let opfamily = unsafe { opfamily(index_relation) }; + let heap = Heap { + heap_relation, + index_relation, + index_info, + opfamily, + scan, + }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, L2>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, Dot>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); + } + (VectorKind::Vecf16, DistanceKind::L2) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, L2>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, Dot>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); + } + } + unsafe { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).nparticipantsdone += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordrqshared).workersdonecv); + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn ambuildempty(_index_relation: pgrx::pg_sys::Relation) { + pgrx::error!("Unlogged indexes are not supported."); +} + +unsafe fn options( + index_relation: pgrx::pg_sys::Relation, +) -> (VectorOptions, VchordrqIndexingOptions) { + let att = unsafe { &mut *(*index_relation).rd_att }; + let atts = unsafe { att.attrs.as_slice(att.natts as _) }; + if atts.is_empty() { + pgrx::error!("indexing on no columns is not supported"); + } + if atts.len() != 1 { + pgrx::error!("multicolumn index is not supported"); + } + // get dims + let typmod = Typmod::parse_from_i32(atts[0].type_mod()).unwrap(); + let dims = if let Some(dims) = typmod.dims() { + dims.get() + } else { + pgrx::error!( + "Dimensions type modifier of a vector column is needed for building the index." + ); + }; + // get v, d + let opfamily = unsafe { opfamily(index_relation) }; + let vector = VectorOptions { + dims, + v: opfamily.vector_kind(), + d: opfamily.distance_kind(), + }; + // get indexing, segment, optimizing + let rabitq = 'rabitq: { + let reloption = unsafe { (*index_relation).rd_options as *const Reloption }; + if reloption.is_null() || unsafe { (*reloption).options == 0 } { + break 'rabitq Default::default(); + } + let s = unsafe { Reloption::options(reloption) }.to_string_lossy(); + match toml::from_str::(&s) { + Ok(p) => p, + Err(e) => pgrx::error!("failed to parse options: {}", e), + } + }; + (vector, rabitq) +} + +pub fn make_internal_build( + vector_options: VectorOptions, + internal_build: VchordrqInternalBuildOptions, + mut samples: Vec>, +) -> Vec>> { + use std::iter::once; + for sample in samples.iter_mut() { + *sample = crate::index::projection::project(sample); + } + let mut result = Vec::>>::new(); + for w in internal_build.lists.iter().rev().copied().chain(once(1)) { + let means = k_means::RayonParallelism::scoped( + internal_build.build_threads as _, + Arc::new(|| { + pgrx::check_for_interrupts!(); + }), + |parallelism| { + k_means::k_means( + parallelism, + w as usize, + vector_options.dims as usize, + if let Some(structure) = result.last() { + &structure.means + } else { + &samples + }, + internal_build.spherical_centroids, + 10, + ) + }, + ) + .expect("failed to create thread pool"); + if let Some(structure) = result.last() { + let mut children = vec![Vec::new(); means.len()]; + for i in 0..structure.len() as u32 { + let target = k_means::k_means_lookup(&structure.means[i as usize], &means); + children[target].push(i); + } + let (means, children) = std::iter::zip(means, children) + .filter(|(_, x)| !x.is_empty()) + .unzip::<_, _, Vec<_>, Vec<_>>(); + result.push(Structure { means, children }); + } else { + let children = vec![Vec::new(); means.len()]; + result.push(Structure { means, children }); + } + } + result +} + +pub fn make_external_build( + vector_options: VectorOptions, + _opfamily: Opfamily, + external_build: VchordrqExternalBuildOptions, +) -> Vec>> { + use std::collections::BTreeMap; + let VchordrqExternalBuildOptions { table } = external_build; + let mut parents = BTreeMap::new(); + let mut vectors = BTreeMap::new(); + pgrx::spi::Spi::connect(|client| { + use crate::datatype::memory_vector::VectorOutput; + use pgrx::pg_sys::panic::ErrorReportable; + use vector::VectorBorrowed; + let schema_query = "SELECT n.nspname::TEXT + FROM pg_catalog.pg_extension e + LEFT JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vector';"; + let pgvector_schema: String = client + .select(schema_query, None, None) + .unwrap_or_report() + .first() + .get_by_name("nspname") + .expect("external build: cannot get schema of pgvector") + .expect("external build: cannot get schema of pgvector"); + let dump_query = + format!("SELECT id, parent, vector::{pgvector_schema}.vector FROM {table};"); + let centroids = client.select(&dump_query, None, None).unwrap_or_report(); + for row in centroids { + let id: Option = row.get_by_name("id").unwrap(); + let parent: Option = row.get_by_name("parent").unwrap(); + let vector: Option = row.get_by_name("vector").unwrap(); + let id = id.expect("external build: id could not be NULL"); + let vector = vector.expect("external build: vector could not be NULL"); + let pop = parents.insert(id, parent); + if pop.is_some() { + pgrx::error!( + "external build: there are at least two lines have same id, id = {id}" + ); + } + if vector_options.dims != vector.as_borrowed().dims() { + pgrx::error!("external build: incorrect dimension, id = {id}"); + } + vectors.insert( + id, + crate::index::projection::project(vector.as_borrowed().slice()), + ); + } + }); + if parents.len() >= 2 && parents.values().all(|x| x.is_none()) { + // if there are more than one vertexs and no edges, + // assume there is an implicit root + let n = parents.len(); + let mut result = Vec::new(); + result.push(Structure { + means: vectors.values().cloned().collect::>(), + children: vec![Vec::new(); n], + }); + result.push(Structure { + means: vec![{ + // compute the vector on root, without normalizing it + let mut sum = vec![0.0f32; vector_options.dims as _]; + for vector in vectors.values() { + f32::vector_add_inplace(&mut sum, vector); + } + f32::vector_mul_scalar_inplace(&mut sum, 1.0 / n as f32); + sum + }], + children: vec![(0..n as u32).collect()], + }); + return result; + } + let mut children = parents + .keys() + .map(|x| (*x, Vec::new())) + .collect::>(); + let mut root = None; + for (&id, &parent) in parents.iter() { + if let Some(parent) = parent { + if let Some(parent) = children.get_mut(&parent) { + parent.push(id); + } else { + pgrx::error!("external build: parent does not exist, id = {id}, parent = {parent}"); + } + } else { + if let Some(root) = root { + pgrx::error!("external build: two root, id = {root}, id = {id}"); + } else { + root = Some(id); + } + } + } + let Some(root) = root else { + pgrx::error!("external build: there are no root"); + }; + let mut heights = BTreeMap::<_, _>::new(); + fn dfs_for_heights( + heights: &mut BTreeMap>, + children: &BTreeMap>, + u: i32, + ) { + if heights.contains_key(&u) { + pgrx::error!("external build: detect a cycle, id = {u}"); + } + heights.insert(u, None); + let mut height = None; + for &v in children[&u].iter() { + dfs_for_heights(heights, children, v); + let new = heights[&v].unwrap() + 1; + if let Some(height) = height { + if height != new { + pgrx::error!("external build: two heights, id = {u}"); + } + } else { + height = Some(new); + } + } + if height.is_none() { + height = Some(1); + } + heights.insert(u, height); + } + dfs_for_heights(&mut heights, &children, root); + let heights = heights + .into_iter() + .map(|(k, v)| (k, v.expect("not a connected graph"))) + .collect::>(); + if !(1..=8).contains(&(heights[&root] - 1)) { + pgrx::error!( + "external build: unexpected tree height, height = {}", + heights[&root] + ); + } + let mut cursors = vec![0_u32; 1 + heights[&root] as usize]; + let mut labels = BTreeMap::new(); + for id in parents.keys().copied() { + let height = heights[&id]; + let cursor = cursors[height as usize]; + labels.insert(id, (height, cursor)); + cursors[height as usize] += 1; + } + fn extract( + height: u32, + labels: &BTreeMap, + vectors: &BTreeMap>, + children: &BTreeMap>, + ) -> (Vec>, Vec>) { + labels + .iter() + .filter(|(_, (h, _))| *h == height) + .map(|(id, _)| { + ( + vectors[id].clone(), + children[id].iter().map(|id| labels[id].1).collect(), + ) + }) + .unzip() + } + let mut result = Vec::new(); + for height in 1..=heights[&root] { + let (means, children) = extract(height, &labels, &vectors, &children); + result.push(Structure { means, children }); + } + result +} + +pub fn map_structures(x: Vec>, f: impl Fn(T) -> U + Copy) -> Vec> { + x.into_iter() + .map(|Structure { means, children }| Structure { + means: means.into_iter().map(f).collect(), + children, + }) + .collect() +} + +pub trait InternalBuild: VectorOwned { + fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec; + + fn build_from_vecf32(x: &[f32]) -> Self; +} + +impl InternalBuild for VectOwned { + fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { + vector.slice().to_vec() + } + + fn build_from_vecf32(x: &[f32]) -> Self { + Self::new(x.to_vec()) + } +} + +impl InternalBuild for VectOwned { + fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { + f16::vector_to_f32(vector.slice()) + } + + fn build_from_vecf32(x: &[f32]) -> Self { + Self::new(f16::vector_from_f32(x)) + } +} diff --git a/src/index/am/am_scan.rs b/src/index/am/am_scan.rs new file mode 100644 index 00000000..39ed021b --- /dev/null +++ b/src/index/am/am_scan.rs @@ -0,0 +1,287 @@ +use crate::index::am::pointer_to_ctid; +use crate::index::gucs::{epsilon, max_scan_tuples, probes}; +use crate::index::opclass::{Opfamily, Sphere, opfamily}; +use crate::index::projection::RandomProject; +use crate::index::storage::PostgresRelation; +use algorithm::operator::{Dot, L2, Op, Vector}; +use algorithm::types::*; +use half::f16; +use std::num::NonZeroU64; +use vector::VectorOwned; +use vector::vect::VectOwned; + +#[pgrx::pg_guard] +pub unsafe extern "C" fn ambeginscan( + index_relation: pgrx::pg_sys::Relation, + n_keys: std::os::raw::c_int, + n_orderbys: std::os::raw::c_int, +) -> pgrx::pg_sys::IndexScanDesc { + use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; + + let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index_relation, n_keys, n_orderbys) }; + unsafe { + let scanner = Scanner { + opfamily: opfamily(index_relation), + scanning: Scanning::Empty {}, + }; + (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); + } + scan +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amrescan( + scan: pgrx::pg_sys::IndexScanDesc, + keys: pgrx::pg_sys::ScanKey, + _n_keys: std::os::raw::c_int, + orderbys: pgrx::pg_sys::ScanKey, + _n_orderbys: std::os::raw::c_int, +) { + unsafe { + if !keys.is_null() && (*scan).numberOfKeys > 0 { + std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); + } + if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { + std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); + } + let opfamily = opfamily((*scan).indexRelation); + let (orderbys, spheres) = { + let mut orderbys = Vec::new(); + let mut spheres = Vec::new(); + if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { + pgrx::error!( + "vector search with no WHERE clause and no ORDER BY clause is not supported" + ); + } + for i in 0..(*scan).numberOfOrderBys { + let data = (*scan).orderByData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + match (*data).sk_strategy { + 1 => orderbys.push(opfamily.input_vector(value, is_null)), + _ => unreachable!(), + } + } + for i in 0..(*scan).numberOfKeys { + let data = (*scan).keyData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + match (*data).sk_strategy { + 2 => spheres.push(opfamily.input_sphere(value, is_null)), + _ => unreachable!(), + } + } + (orderbys, spheres) + }; + let scanning; + if let Some((vector, threshold, recheck)) = scanner_build(orderbys, spheres) { + scanning = Scanning::Initial { + vector, + threshold, + recheck, + }; + } else { + scanning = Scanning::Empty {}; + }; + let scanner = &mut *(*scan).opaque.cast::(); + scanner.scanning = scanning; + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amgettuple( + scan: pgrx::pg_sys::IndexScanDesc, + direction: pgrx::pg_sys::ScanDirection::Type, +) -> bool { + if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { + pgrx::error!("vector search without a forward scan direction is not supported"); + } + // https://www.postgresql.org/docs/current/index-locking.html + // If heap entries referenced physical pointers are deleted before + // they are consumed by PostgreSQL, PostgreSQL will received wrong + // physical pointers: no rows or irreverent rows are referenced. + if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + { + pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); + } + let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; + let relation = unsafe { PostgresRelation::new((*scan).indexRelation) }; + if let Some((pointer, recheck)) = scanner_next(scanner, relation) { + let ctid = pointer_to_ctid(pointer); + unsafe { + (*scan).xs_heaptid = ctid; + (*scan).xs_recheckorderby = false; + (*scan).xs_recheck = recheck; + } + true + } else { + false + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { + let scanner = unsafe { &mut *(*scan).opaque.cast::() }; + scanner.scanning = Scanning::Empty {}; +} + +struct Scanner { + opfamily: Opfamily, + scanning: Scanning, +} + +enum Scanning { + Initial { + vector: OwnedVector, + threshold: Option, + recheck: bool, + }, + Vbase { + vbase: Box>, + recheck: bool, + }, + Empty {}, +} + +fn scanner_build( + orderbys: Vec>, + spheres: Vec>>, +) -> Option<(OwnedVector, Option, bool)> { + let mut vector = None; + let mut threshold = None; + let mut recheck = false; + for orderby_vector in orderbys.into_iter().flatten() { + if vector.is_none() { + vector = Some(orderby_vector); + } else { + pgrx::error!("vector search with multiple vectors is not supported"); + } + } + for Sphere { center, radius } in spheres.into_iter().flatten() { + if vector.is_none() { + (vector, threshold) = (Some(center), Some(radius)); + } else { + recheck = true; + } + } + Some((vector?, threshold, recheck)) +} + +fn scanner_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(NonZeroU64, bool)> { + if let Scanning::Initial { + vector, + threshold, + recheck, + } = &scanner.scanning + { + let opfamily = scanner.opfamily; + let vector = vector.clone(); + let threshold = *threshold; + let recheck = *recheck; + let max_scan_tuples = max_scan_tuples(); + let probes = probes(); + let epsilon = epsilon(); + scanner.scanning = Scanning::Vbase { + vbase: match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + let vector = RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ); + let vbase = algorithm::search::, L2>>( + relation, vector, probes, epsilon, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + match (max_scan_tuples, threshold) { + (None, None) => { + Box::new(vbase.fuse()) as Box> + } + (None, Some(threshold)) => { + Box::new(vbase.take_while(move |(x, _)| *x < threshold)) + } + (Some(max_scan_tuples), None) => Box::new(vbase.take(max_scan_tuples as _)), + (Some(max_scan_tuples), Some(threshold)) => Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ), + } + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + let vector = RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ); + let vbase = algorithm::search::, Dot>>( + relation, vector, probes, epsilon, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + match (max_scan_tuples, threshold) { + (None, None) => { + Box::new(vbase) as Box> + } + (None, Some(threshold)) => { + Box::new(vbase.take_while(move |(x, _)| *x < threshold)) + } + (Some(max_scan_tuples), None) => Box::new(vbase.take(max_scan_tuples as _)), + (Some(max_scan_tuples), Some(threshold)) => Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ), + } + } + (VectorKind::Vecf16, DistanceKind::L2) => { + let vector = RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ); + let vbase = algorithm::search::, L2>>( + relation, vector, probes, epsilon, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + match (max_scan_tuples, threshold) { + (None, None) => { + Box::new(vbase) as Box> + } + (None, Some(threshold)) => { + Box::new(vbase.take_while(move |(x, _)| *x < threshold)) + } + (Some(max_scan_tuples), None) => Box::new(vbase.take(max_scan_tuples as _)), + (Some(max_scan_tuples), Some(threshold)) => Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ), + } + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + let vector = RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ); + let vbase = algorithm::search::, Dot>>( + relation, vector, probes, epsilon, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + match (max_scan_tuples, threshold) { + (None, None) => { + Box::new(vbase) as Box> + } + (None, Some(threshold)) => { + Box::new(vbase.take_while(move |(x, _)| *x < threshold)) + } + (Some(max_scan_tuples), None) => Box::new(vbase.take(max_scan_tuples as _)), + (Some(max_scan_tuples), Some(threshold)) => Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ), + } + } + }, + recheck, + }; + } + match &mut scanner.scanning { + Scanning::Initial { .. } => unreachable!(), + Scanning::Vbase { vbase, recheck } => vbase.next().map(|(_, x)| (x, *recheck)), + Scanning::Empty {} => None, + } +} diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs new file mode 100644 index 00000000..dcac2314 --- /dev/null +++ b/src/index/am/mod.rs @@ -0,0 +1,324 @@ +pub mod am_build; +pub mod am_scan; + +use crate::index::projection::RandomProject; +use crate::index::storage::PostgresRelation; +use algorithm::operator::{Dot, L2, Op, Vector}; +use algorithm::types::*; +use half::f16; +use pgrx::datum::Internal; +use pgrx::pg_sys::Datum; +use std::ffi::CStr; +use std::num::NonZeroU64; +use std::sync::OnceLock; +use vector::VectorOwned; +use vector::vect::VectOwned; + +#[repr(C)] +struct Reloption { + vl_len_: i32, + pub options: i32, +} + +impl Reloption { + unsafe fn options<'a>(this: *const Self) -> &'a CStr { + unsafe { + let ptr = this + .cast::() + .add((&raw const (*this).options).read() as _); + CStr::from_ptr(ptr.cast()) + } + } +} + +const TABLE: &[pgrx::pg_sys::relopt_parse_elt] = &[pgrx::pg_sys::relopt_parse_elt { + optname: c"options".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, + offset: std::mem::offset_of!(Reloption, options) as i32, +}]; + +static RELOPT_KIND: OnceLock = OnceLock::new(); + +pub fn init() { + RELOPT_KIND.get_or_init(|| { + let kind; + unsafe { + kind = pgrx::pg_sys::add_reloption_kind(); + pgrx::pg_sys::add_string_reloption( + kind, + c"options".as_ptr(), + c"Vector index options, represented as a TOML string.".as_ptr(), + c"".as_ptr(), + None, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + } + kind + }); +} + +#[pgrx::pg_extern(sql = "")] +fn _vchordrq_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { + type T = pgrx::pg_sys::IndexAmRoutine; + unsafe { + let index_am_routine = pgrx::pg_sys::palloc0(size_of::()) as *mut T; + index_am_routine.write(AM_HANDLER); + Internal::from(Some(Datum::from(index_am_routine))) + } +} + +const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { + let mut am_routine = unsafe { std::mem::zeroed::() }; + + am_routine.type_ = pgrx::pg_sys::NodeTag::T_IndexAmRoutine; + + am_routine.amsupport = 1; + am_routine.amcanorderbyop = true; + + #[cfg(feature = "pg17")] + { + am_routine.amcanbuildparallel = true; + } + + // Index access methods that set `amoptionalkey` to `false` + // must index all tuples, even if the first column is `NULL`. + // However, PostgreSQL does not generate a path if there is no + // index clauses, even if there is a `ORDER BY` clause. + // So we have to set it to `true` and set costs of every path + // for vector index scans without `ORDER BY` clauses a large number + // and throw errors if someone really wants such a path. + am_routine.amoptionalkey = true; + + am_routine.amvalidate = Some(amvalidate); + am_routine.amoptions = Some(amoptions); + am_routine.amcostestimate = Some(amcostestimate); + + am_routine.ambuild = Some(am_build::ambuild); + am_routine.ambuildempty = Some(am_build::ambuildempty); + am_routine.aminsert = Some(aminsert); + am_routine.ambulkdelete = Some(ambulkdelete); + am_routine.amvacuumcleanup = Some(amvacuumcleanup); + + am_routine.ambeginscan = Some(am_scan::ambeginscan); + am_routine.amrescan = Some(am_scan::amrescan); + am_routine.amgettuple = Some(am_scan::amgettuple); + am_routine.amendscan = Some(am_scan::amendscan); + + am_routine +}; + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { + true +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amoptions(reloptions: Datum, validate: bool) -> *mut pgrx::pg_sys::bytea { + let relopt_kind = RELOPT_KIND.get().copied().expect("init is not called"); + let rdopts = unsafe { + pgrx::pg_sys::build_reloptions( + reloptions, + validate, + relopt_kind, + size_of::(), + TABLE.as_ptr(), + TABLE.len() as _, + ) + }; + rdopts as *mut pgrx::pg_sys::bytea +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amcostestimate( + _root: *mut pgrx::pg_sys::PlannerInfo, + path: *mut pgrx::pg_sys::IndexPath, + _loop_count: f64, + index_startup_cost: *mut pgrx::pg_sys::Cost, + index_total_cost: *mut pgrx::pg_sys::Cost, + index_selectivity: *mut pgrx::pg_sys::Selectivity, + index_correlation: *mut f64, + index_pages: *mut f64, +) { + unsafe { + if (*path).indexorderbys.is_null() && (*path).indexclauses.is_null() { + *index_startup_cost = f64::MAX; + *index_total_cost = f64::MAX; + *index_selectivity = 0.0; + *index_correlation = 0.0; + *index_pages = 0.0; + return; + } + *index_startup_cost = 0.0; + *index_total_cost = 0.0; + *index_selectivity = 1.0; + *index_correlation = 1.0; + *index_pages = 0.0; + } +} + +#[cfg(feature = "pg13")] +#[pgrx::pg_guard] +pub unsafe extern "C" fn aminsert( + index_relation: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, + _heap_relation: pgrx::pg_sys::Relation, + _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, + _index_info: *mut pgrx::pg_sys::IndexInfo, +) -> bool { + unsafe { aminsertinner(index_relation, values, is_null, heap_tid) } +} + +#[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] +#[pgrx::pg_guard] +pub unsafe extern "C" fn aminsert( + index_relation: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, + _heap_relation: pgrx::pg_sys::Relation, + _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, + _index_unchanged: bool, + _index_info: *mut pgrx::pg_sys::IndexInfo, +) -> bool { + unsafe { aminsertinner(index_relation, values, is_null, heap_tid) } +} + +unsafe fn aminsertinner( + index_relation: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, +) -> bool { + let opfamily = unsafe { crate::index::opclass::opfamily(index_relation) }; + let index = unsafe { PostgresRelation::new(index_relation) }; + let payload = ctid_to_pointer(unsafe { heap_tid.read() }); + let vector = unsafe { opfamily.input_vector(*values.add(0), *is_null.add(0)) }; + let Some(vector) = vector else { return false }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => algorithm::insert::, L2>>( + index, + payload, + RandomProject::project(VectOwned::::from_owned(vector).as_borrowed()), + ), + (VectorKind::Vecf32, DistanceKind::Dot) => algorithm::insert::, Dot>>( + index, + payload, + RandomProject::project(VectOwned::::from_owned(vector).as_borrowed()), + ), + (VectorKind::Vecf16, DistanceKind::L2) => algorithm::insert::, L2>>( + index, + payload, + RandomProject::project(VectOwned::::from_owned(vector).as_borrowed()), + ), + (VectorKind::Vecf16, DistanceKind::Dot) => algorithm::insert::, Dot>>( + index, + payload, + RandomProject::project(VectOwned::::from_owned(vector).as_borrowed()), + ), + } + false +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn ambulkdelete( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, + callback: pgrx::pg_sys::IndexBulkDeleteCallback, + callback_state: *mut std::os::raw::c_void, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let opfamily = unsafe { crate::index::opclass::opfamily((*info).index) }; + let index = unsafe { PostgresRelation::new((*info).index) }; + let check = || unsafe { + pgrx::pg_sys::vacuum_delay_point(); + }; + let callback = callback.expect("null function pointer"); + let callback = |p: NonZeroU64| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + algorithm::bulkdelete::, L2>>(index, check, callback); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + algorithm::bulkdelete::, Dot>>(index, check, callback); + } + (VectorKind::Vecf16, DistanceKind::L2) => { + algorithm::bulkdelete::, L2>>(index, check, callback); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + algorithm::bulkdelete::, Dot>>(index, check, callback); + } + } + stats +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amvacuumcleanup( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let opfamily = unsafe { crate::index::opclass::opfamily((*info).index) }; + let index = unsafe { PostgresRelation::new((*info).index) }; + let check = || unsafe { + pgrx::pg_sys::vacuum_delay_point(); + }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + algorithm::maintain::, L2>>(index, check); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + algorithm::maintain::, Dot>>(index, check); + } + (VectorKind::Vecf16, DistanceKind::L2) => { + algorithm::maintain::, L2>>(index, check); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + algorithm::maintain::, Dot>>(index, check); + } + } + stats +} + +const fn pointer_to_ctid(pointer: NonZeroU64) -> pgrx::pg_sys::ItemPointerData { + let value = pointer.get(); + pgrx::pg_sys::ItemPointerData { + ip_blkid: pgrx::pg_sys::BlockIdData { + bi_hi: ((value >> 32) & 0xffff) as u16, + bi_lo: ((value >> 16) & 0xffff) as u16, + }, + ip_posid: (value & 0xffff) as u16, + } +} + +const fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> NonZeroU64 { + let mut value = 0; + value |= (ctid.ip_blkid.bi_hi as u64) << 32; + value |= (ctid.ip_blkid.bi_lo as u64) << 16; + value |= ctid.ip_posid as u64; + NonZeroU64::new(value).expect("invalid pointer") +} + +#[test] +const fn soundness_check() { + let a = pgrx::pg_sys::ItemPointerData { + ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi: 1, bi_lo: 2 }, + ip_posid: 3, + }; + let b = ctid_to_pointer(a); + let c = pointer_to_ctid(b); + assert!(a.ip_blkid.bi_hi == c.ip_blkid.bi_hi); + assert!(a.ip_blkid.bi_lo == c.ip_blkid.bi_lo); + assert!(a.ip_posid == c.ip_posid); +} diff --git a/src/index/am_options.rs b/src/index/am_options.rs deleted file mode 100644 index 06ca8e50..00000000 --- a/src/index/am_options.rs +++ /dev/null @@ -1,235 +0,0 @@ -use crate::datatype::memory_halfvec::HalfvecInput; -use crate::datatype::memory_halfvec::HalfvecOutput; -use crate::datatype::memory_vector::VectorInput; -use crate::datatype::memory_vector::VectorOutput; -use crate::datatype::typmod::Typmod; -use crate::types::{BorrowedVector, OwnedVector}; -use crate::types::{DistanceKind, VectorKind}; -use crate::types::{VchordrqIndexingOptions, VectorOptions}; -use distance::Distance; -use pgrx::datum::FromDatum; -use pgrx::heap_tuple::PgHeapTuple; -use serde::Deserialize; -use std::ffi::CStr; -use std::num::NonZero; -use vector::VectorBorrowed; - -#[derive(Copy, Clone, Debug, Default)] -#[repr(C)] -pub struct Reloption { - vl_len_: i32, - pub options: i32, -} - -impl Reloption { - pub const TAB: &'static [pgrx::pg_sys::relopt_parse_elt] = &[pgrx::pg_sys::relopt_parse_elt { - optname: c"options".as_ptr(), - opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, - offset: std::mem::offset_of!(Reloption, options) as i32, - }]; - unsafe fn options(&self) -> &CStr { - unsafe { - let ptr = (&raw const *self) - .cast::() - .offset(self.options as _); - CStr::from_ptr(ptr) - } - } -} - -#[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -pub enum PgDistanceKind { - L2, - Dot, - Cos, -} - -impl PgDistanceKind { - pub fn to_distance(self) -> DistanceKind { - match self { - PgDistanceKind::L2 => DistanceKind::L2, - PgDistanceKind::Dot | PgDistanceKind::Cos => DistanceKind::Dot, - } - } -} - -fn convert_name_to_vd(name: &str) -> Option<(VectorKind, PgDistanceKind)> { - match name.strip_suffix("_ops") { - Some("vector_l2") => Some((VectorKind::Vecf32, PgDistanceKind::L2)), - Some("vector_ip") => Some((VectorKind::Vecf32, PgDistanceKind::Dot)), - Some("vector_cosine") => Some((VectorKind::Vecf32, PgDistanceKind::Cos)), - Some("halfvec_l2") => Some((VectorKind::Vecf16, PgDistanceKind::L2)), - Some("halfvec_ip") => Some((VectorKind::Vecf16, PgDistanceKind::Dot)), - Some("halfvec_cosine") => Some((VectorKind::Vecf16, PgDistanceKind::Cos)), - _ => None, - } -} - -unsafe fn convert_reloptions_to_options( - reloptions: *const pgrx::pg_sys::varlena, -) -> VchordrqIndexingOptions { - #[derive(Debug, Clone, Deserialize, Default)] - #[serde(deny_unknown_fields)] - struct Parsed { - #[serde(flatten)] - rabitq: VchordrqIndexingOptions, - } - let reloption = reloptions as *const Reloption; - if reloption.is_null() || unsafe { (*reloption).options == 0 } { - return Default::default(); - } - let s = unsafe { (*reloption).options() }.to_string_lossy(); - match toml::from_str::(&s) { - Ok(p) => p.rabitq, - Err(e) => pgrx::error!("failed to parse options: {}", e), - } -} - -pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, VchordrqIndexingOptions) { - let att = unsafe { &mut *(*index).rd_att }; - let atts = unsafe { att.attrs.as_slice(att.natts as _) }; - if atts.is_empty() { - pgrx::error!("indexing on no columns is not supported"); - } - if atts.len() != 1 { - pgrx::error!("multicolumn index is not supported"); - } - // get dims - let typmod = Typmod::parse_from_i32(atts[0].type_mod()).unwrap(); - let dims = if let Some(dims) = typmod.dims() { - dims.get() - } else { - pgrx::error!( - "Dimensions type modifier of a vector column is needed for building the index." - ); - }; - // get v, d - let opfamily = unsafe { opfamily(index) }; - let vector = VectorOptions { - dims, - v: opfamily.vector, - d: opfamily.distance_kind(), - }; - // get indexing, segment, optimizing - let rabitq = unsafe { convert_reloptions_to_options((*index).rd_options) }; - (vector, rabitq) -} - -#[derive(Debug, Clone, Copy)] -pub struct Opfamily { - vector: VectorKind, - pg_distance: PgDistanceKind, -} - -impl Opfamily { - pub unsafe fn datum_to_vector( - self, - datum: pgrx::pg_sys::Datum, - is_null: bool, - ) -> Option { - if is_null || datum.is_null() { - return None; - } - let vector = match self.vector { - VectorKind::Vecf32 => { - let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; - self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed())) - } - VectorKind::Vecf16 => { - let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; - self.preprocess(BorrowedVector::Vecf16(vector.as_borrowed())) - } - }; - Some(vector) - } - pub unsafe fn datum_to_sphere( - self, - datum: pgrx::pg_sys::Datum, - is_null: bool, - ) -> (Option, Option) { - if is_null || datum.is_null() { - return (None, None); - } - let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; - let center = match self.vector { - VectorKind::Vecf32 => tuple - .get_by_index::(NonZero::new(1).unwrap()) - .unwrap() - .map(|vector| self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed()))), - VectorKind::Vecf16 => tuple - .get_by_index::(NonZero::new(1).unwrap()) - .unwrap() - .map(|vector| self.preprocess(BorrowedVector::Vecf16(vector.as_borrowed()))), - }; - let radius = tuple.get_by_index::(NonZero::new(2).unwrap()).unwrap(); - (center, radius) - } - pub fn preprocess(self, vector: BorrowedVector<'_>) -> OwnedVector { - use BorrowedVector as B; - use OwnedVector as O; - match (vector, self.pg_distance) { - (B::Vecf32(x), PgDistanceKind::L2) => O::Vecf32(x.own()), - (B::Vecf32(x), PgDistanceKind::Dot) => O::Vecf32(x.own()), - (B::Vecf32(x), PgDistanceKind::Cos) => O::Vecf32(x.function_normalize()), - (B::Vecf16(x), PgDistanceKind::L2) => O::Vecf16(x.own()), - (B::Vecf16(x), PgDistanceKind::Dot) => O::Vecf16(x.own()), - (B::Vecf16(x), PgDistanceKind::Cos) => O::Vecf16(x.function_normalize()), - } - } - pub fn process(self, x: Distance) -> f32 { - match self.pg_distance { - PgDistanceKind::Cos => f32::from(x) + 1.0f32, - PgDistanceKind::L2 => f32::from(x).sqrt(), - PgDistanceKind::Dot => x.into(), - } - } - pub fn distance_kind(self) -> DistanceKind { - self.pg_distance.to_distance() - } - pub fn vector_kind(self) -> VectorKind { - self.vector - } -} - -pub unsafe fn opfamily(index: pgrx::pg_sys::Relation) -> Opfamily { - use pgrx::pg_sys::Oid; - - let proc = unsafe { pgrx::pg_sys::index_getprocid(index, 1, 1) }; - - if proc == Oid::INVALID { - pgrx::error!("support function 1 is not found"); - } - - let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); - unsafe { - pgrx::pg_sys::fmgr_info(proc, &mut flinfo); - } - - let fn_addr = flinfo.fn_addr.expect("null function pointer"); - - let mut fcinfo = unsafe { std::mem::zeroed::() }; - fcinfo.flinfo = &mut flinfo; - fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; - fcinfo.context = std::ptr::null_mut(); - fcinfo.resultinfo = std::ptr::null_mut(); - fcinfo.isnull = true; - fcinfo.nargs = 0; - - let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; - - let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; - - let result_string = result_option.expect("null string"); - - let (vector, pg_distance) = convert_name_to_vd(&result_string).unwrap(); - - unsafe { - pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); - } - - Opfamily { - vector, - pg_distance, - } -} diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs deleted file mode 100644 index 83e62f33..00000000 --- a/src/index/am_scan.rs +++ /dev/null @@ -1,186 +0,0 @@ -use super::am_options::Opfamily; -use crate::algorithm::operator::Vector; -use crate::algorithm::operator::{Dot, L2, Op}; -use crate::algorithm::scan::scan; -use crate::gucs::executing::epsilon; -use crate::gucs::executing::max_scan_tuples; -use crate::gucs::executing::probes; -use crate::postgres::PostgresRelation; -use crate::types::DistanceKind; -use crate::types::OwnedVector; -use crate::types::VectorKind; -use distance::Distance; -use half::f16; -use std::num::NonZeroU64; -use vector::vect::VectOwned; - -pub enum Scanner { - Initial { - vector: Option<(OwnedVector, Opfamily)>, - threshold: Option, - recheck: bool, - }, - Vbase { - vbase: Box>, - threshold: Option, - recheck: bool, - opfamily: Opfamily, - }, - Empty {}, -} - -pub fn scan_build( - orderbys: Vec>, - spheres: Vec<(Option, Option)>, - opfamily: Opfamily, -) -> (Option<(OwnedVector, Opfamily)>, Option, bool) { - let mut pair = None; - let mut threshold = None; - let mut recheck = false; - for orderby_vector in orderbys { - if pair.is_none() { - pair = orderby_vector; - } else if orderby_vector.is_some() { - pgrx::error!("vector search with multiple vectors is not supported"); - } - } - for (sphere_vector, sphere_threshold) in spheres { - if pair.is_none() { - pair = sphere_vector; - threshold = sphere_threshold; - } else { - recheck = true; - break; - } - } - (pair.map(|x| (x, opfamily)), threshold, recheck) -} - -pub fn scan_make( - vector: Option<(OwnedVector, Opfamily)>, - threshold: Option, - recheck: bool, -) -> Scanner { - Scanner::Initial { - vector, - threshold, - recheck, - } -} - -pub fn scan_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(NonZeroU64, bool)> { - if let Scanner::Initial { - vector, - threshold, - recheck, - } = scanner - { - if let Some((vector, opfamily)) = vector.as_ref() { - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - let vbase = scan::, L2>>( - relation, - VectOwned::::from_owned(vector.clone()), - probes(), - epsilon(), - ); - *scanner = Scanner::Vbase { - vbase: if let Some(max_scan_tuples) = max_scan_tuples() { - Box::new(vbase.take(max_scan_tuples as usize)) - } else { - Box::new(vbase) - }, - threshold: *threshold, - recheck: *recheck, - opfamily: *opfamily, - }; - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - let vbase = scan::, Dot>>( - relation, - VectOwned::::from_owned(vector.clone()), - probes(), - epsilon(), - ); - *scanner = Scanner::Vbase { - vbase: if let Some(max_scan_tuples) = max_scan_tuples() { - Box::new(vbase.take(max_scan_tuples as usize)) - } else { - Box::new(vbase) - }, - threshold: *threshold, - recheck: *recheck, - opfamily: *opfamily, - }; - } - (VectorKind::Vecf16, DistanceKind::L2) => { - let vbase = scan::, L2>>( - relation, - VectOwned::::from_owned(vector.clone()), - probes(), - epsilon(), - ); - *scanner = Scanner::Vbase { - vbase: if let Some(max_scan_tuples) = max_scan_tuples() { - Box::new(vbase.take(max_scan_tuples as usize)) - } else { - Box::new(vbase) - }, - threshold: *threshold, - recheck: *recheck, - opfamily: *opfamily, - }; - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - let vbase = scan::, Dot>>( - relation, - VectOwned::::from_owned(vector.clone()), - probes(), - epsilon(), - ); - *scanner = Scanner::Vbase { - vbase: if let Some(max_scan_tuples) = max_scan_tuples() { - Box::new(vbase.take(max_scan_tuples as usize)) - } else { - Box::new(vbase) - }, - threshold: *threshold, - recheck: *recheck, - opfamily: *opfamily, - }; - } - } - } else { - *scanner = Scanner::Empty {}; - } - } - match scanner { - Scanner::Initial { .. } => unreachable!(), - Scanner::Vbase { - vbase, - threshold, - recheck, - opfamily, - } => match ( - vbase.next().map(|(d, p)| (opfamily.process(d), p)), - threshold, - ) { - (Some((_, ptr)), None) => Some((ptr, *recheck)), - (Some((distance, ptr)), Some(t)) if distance < *t => Some((ptr, *recheck)), - _ => { - let scanner = std::mem::replace(scanner, Scanner::Empty {}); - scan_release(scanner); - None - } - }, - Scanner::Empty {} => None, - } -} - -pub fn scan_release(scanner: Scanner) { - match scanner { - Scanner::Initial { .. } => {} - Scanner::Vbase { .. } => {} - Scanner::Empty {} => {} - } -} diff --git a/src/index/functions.rs b/src/index/functions.rs index 1f3b4e28..be2f9633 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -1,9 +1,6 @@ -use super::am_options; -use crate::algorithm::operator::{Dot, L2, Op}; -use crate::algorithm::prewarm::prewarm; -use crate::postgres::PostgresRelation; -use crate::types::DistanceKind; -use crate::types::VectorKind; +use crate::index::storage::PostgresRelation; +use algorithm::operator::{Dot, L2, Op}; +use algorithm::types::*; use half::f16; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass}; @@ -24,19 +21,27 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { } let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; let relation = unsafe { PostgresRelation::new(index) }; - let opfamily = unsafe { am_options::opfamily(index) }; + let opfamily = unsafe { crate::index::opclass::opfamily(index) }; let message = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { - prewarm::, L2>>(relation, height) + algorithm::prewarm::, L2>>(relation, height, || { + pgrx::check_for_interrupts!(); + }) } (VectorKind::Vecf32, DistanceKind::Dot) => { - prewarm::, Dot>>(relation, height) + algorithm::prewarm::, Dot>>(relation, height, || { + pgrx::check_for_interrupts!(); + }) } (VectorKind::Vecf16, DistanceKind::L2) => { - prewarm::, L2>>(relation, height) + algorithm::prewarm::, L2>>(relation, height, || { + pgrx::check_for_interrupts!(); + }) } (VectorKind::Vecf16, DistanceKind::Dot) => { - prewarm::, Dot>>(relation, height) + algorithm::prewarm::, Dot>>(relation, height, || { + pgrx::check_for_interrupts!(); + }) } }; unsafe { diff --git a/src/gucs/executing.rs b/src/index/gucs.rs similarity index 61% rename from src/gucs/executing.rs rename to src/index/gucs.rs index af6cce7d..5e4216b6 100644 --- a/src/gucs/executing.rs +++ b/src/index/gucs.rs @@ -4,8 +4,10 @@ use std::ffi::CStr; static PROBES: GucSetting> = GucSetting::>::new(Some(c"10")); static EPSILON: GucSetting = GucSetting::::new(1.9); static MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); +static PREWARM_DIM: GucSetting> = + GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); -pub unsafe fn init() { +pub fn init() { GucRegistry::define_string_guc( "vchordrq.probes", "`probes` argument of vchordrq.", @@ -34,6 +36,20 @@ pub unsafe fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_string_guc( + "vchordrq.prewarm_dim", + "prewarm_dim when the extension is loading.", + "prewarm_dim when the extension is loading.", + &PREWARM_DIM, + GucContext::Userset, + GucFlags::default(), + ); + unsafe { + #[cfg(any(feature = "pg13", feature = "pg14"))] + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] + pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrq".as_ptr()); + } } pub fn probes() -> Vec { @@ -70,3 +86,24 @@ pub fn max_scan_tuples() -> Option { let x = MAX_SCAN_TUPLES.get(); if x < 0 { None } else { Some(x as u32) } } + +pub fn prewarm_dim() -> Vec { + if let Some(prewarm_dim) = PREWARM_DIM.get() { + if let Ok(prewarm_dim) = prewarm_dim.to_str() { + let mut result = Vec::new(); + for dim in prewarm_dim.split(',') { + if let Ok(dim) = dim.trim().parse::() { + result.push(dim); + } else { + pgrx::warning!("{dim:?} is not a valid integer"); + } + } + result + } else { + pgrx::warning!("vchordrq.prewarm_dim is not a valid UTF-8 string"); + Vec::new() + } + } else { + Vec::new() + } +} diff --git a/src/index/mod.rs b/src/index/mod.rs index 5203e4fb..511a4f57 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -1,12 +1,14 @@ pub mod am; -pub mod am_options; -pub mod am_scan; pub mod functions; +pub mod gucs; pub mod opclass; -pub mod utils; +pub mod projection; +pub mod storage; -pub unsafe fn init() { - unsafe { - am::init(); +pub fn init() { + am::init(); + gucs::init(); + for x in gucs::prewarm_dim() { + projection::prewarm(x as _); } } diff --git a/src/index/opclass.rs b/src/index/opclass.rs index a2dc8618..63a6be5d 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -1,3 +1,13 @@ +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; +use algorithm::types::*; +use distance::Distance; +use pgrx::datum::FromDatum; +use pgrx::heap_tuple::PgHeapTuple; +use pgrx::pg_sys::Datum; +use std::num::NonZero; +use vector::VectorBorrowed; + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_vector_l2_ops() -> String { "vector_l2_ops".to_string() @@ -27,3 +37,139 @@ fn _vchordrq_support_halfvec_ip_ops() -> String { fn _vchordrq_support_halfvec_cosine_ops() -> String { "halfvec_cosine_ops".to_string() } + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum PostgresDistanceKind { + L2, + Ip, + Cosine, +} + +pub struct Sphere { + pub center: T, + pub radius: f32, +} + +#[derive(Debug, Clone, Copy)] +pub struct Opfamily { + vector: VectorKind, + postgres_distance: PostgresDistanceKind, +} + +impl Opfamily { + fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { + use {BorrowedVector as B, OwnedVector as O, PostgresDistanceKind as D}; + match (vector, self.postgres_distance) { + (B::Vecf32(x), D::L2) => O::Vecf32(x.own()), + (B::Vecf32(x), D::Ip) => O::Vecf32(x.own()), + (B::Vecf32(x), D::Cosine) => O::Vecf32(x.function_normalize()), + (B::Vecf16(x), D::L2) => O::Vecf16(x.own()), + (B::Vecf16(x), D::Ip) => O::Vecf16(x.own()), + (B::Vecf16(x), D::Cosine) => O::Vecf16(x.function_normalize()), + } + } + pub unsafe fn input_vector(self, datum: Datum, is_null: bool) -> Option { + if is_null || datum.is_null() { + return None; + } + let vector = match self.vector { + VectorKind::Vecf32 => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + VectorKind::Vecf16 => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + }; + Some(vector) + } + pub unsafe fn input_sphere(self, datum: Datum, is_null: bool) -> Option> { + if is_null || datum.is_null() { + return None; + } + let attno_1 = NonZero::new(1_usize).unwrap(); + let attno_2 = NonZero::new(2_usize).unwrap(); + let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; + let center = match self.vector { + VectorKind::Vecf32 => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + VectorKind::Vecf16 => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + }; + let radius = tuple.get_by_index::(attno_2).unwrap()?; + Some(Sphere { center, radius }) + } + pub fn output(self, x: Distance) -> f32 { + match self.postgres_distance { + PostgresDistanceKind::Cosine => x.to_f32() + 1.0f32, + PostgresDistanceKind::L2 => x.to_f32().sqrt(), + PostgresDistanceKind::Ip => x.to_f32(), + } + } + pub const fn distance_kind(self) -> DistanceKind { + match self.postgres_distance { + PostgresDistanceKind::L2 => DistanceKind::L2, + PostgresDistanceKind::Ip | PostgresDistanceKind::Cosine => DistanceKind::Dot, + } + } + pub const fn vector_kind(self) -> VectorKind { + self.vector + } +} + +pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { + use pgrx::pg_sys::Oid; + + let proc = unsafe { pgrx::pg_sys::index_getprocid(index_relation, 1, 1) }; + + if proc == Oid::INVALID { + pgrx::error!("support function 1 is not found"); + } + + let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); + + unsafe { + pgrx::pg_sys::fmgr_info(proc, &mut flinfo); + } + + let fn_addr = flinfo.fn_addr.expect("null function pointer"); + + let mut fcinfo = unsafe { std::mem::zeroed::() }; + fcinfo.flinfo = &mut flinfo; + fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; + fcinfo.context = std::ptr::null_mut(); + fcinfo.resultinfo = std::ptr::null_mut(); + fcinfo.isnull = true; + fcinfo.nargs = 0; + + let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; + + let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; + + let result_string = result_option.expect("null return value"); + + let (vector, postgres_distance) = match result_string.as_str() { + "vector_l2_ops" => (VectorKind::Vecf32, PostgresDistanceKind::L2), + "vector_ip_ops" => (VectorKind::Vecf32, PostgresDistanceKind::Ip), + "vector_cosine_ops" => (VectorKind::Vecf32, PostgresDistanceKind::Cosine), + "halfvec_l2_ops" => (VectorKind::Vecf16, PostgresDistanceKind::L2), + "halfvec_ip_ops" => (VectorKind::Vecf16, PostgresDistanceKind::Ip), + "halfvec_cosine_ops" => (VectorKind::Vecf16, PostgresDistanceKind::Cosine), + _ => pgrx::error!("unknown operator class"), + }; + + unsafe { + pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); + } + + Opfamily { + vector, + postgres_distance, + } +} diff --git a/src/projection.rs b/src/index/projection.rs similarity index 51% rename from src/projection.rs rename to src/index/projection.rs index fbcaeffa..ca07e24f 100644 --- a/src/projection.rs +++ b/src/index/projection.rs @@ -1,5 +1,7 @@ +use half::f16; use random_orthogonal_matrix::random_orthogonal_matrix; use std::sync::OnceLock; +use vector::vect::{VectBorrowed, VectOwned}; fn matrix(n: usize) -> Option<&'static Vec>> { static MATRIXS: [OnceLock>>; 1 + 60000] = [const { OnceLock::new() }; 1 + 60000]; @@ -20,3 +22,25 @@ pub fn project(vector: &[f32]) -> Vec { .map(|i| f32::reduce_sum_of_xy(vector, &matrix[i])) .collect() } + +pub trait RandomProject { + type Output; + fn project(self) -> Self::Output; +} + +impl RandomProject for VectBorrowed<'_, f32> { + type Output = VectOwned; + fn project(self) -> VectOwned { + VectOwned::new(project(self.slice())) + } +} + +impl RandomProject for VectBorrowed<'_, f16> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use simd::Floating; + VectOwned::new(f16::vector_from_f32(&project(&f16::vector_to_f32( + self.slice(), + )))) + } +} diff --git a/src/postgres.rs b/src/index/storage.rs similarity index 91% rename from src/postgres.rs rename to src/index/storage.rs index f68d0fa7..a190ac08 100644 --- a/src/postgres.rs +++ b/src/index/storage.rs @@ -1,4 +1,4 @@ -use crate::algorithm::{Opaque, Page, PageGuard, RelationRead, RelationWrite}; +use algorithm::{Opaque, Page, PageGuard, RelationRead, RelationWrite}; use std::mem::{MaybeUninit, offset_of}; use std::ops::{Deref, DerefMut}; use std::ptr::NonNull; @@ -42,31 +42,6 @@ impl PostgresPage { assert_eq!(offset_of!(Self, opaque), this.header.pd_special as usize); this } - #[allow(dead_code)] - fn clone_into_boxed(&self) -> Box { - let mut result = Box::new_uninit(); - unsafe { - std::ptr::copy(self as *const Self, result.as_mut_ptr(), 1); - result.assume_init() - } - } - #[allow(dead_code)] - fn reconstruct(&mut self, removes: &[u16]) { - let mut removes = removes.to_vec(); - removes.sort(); - removes.dedup(); - let n = removes.len(); - if n > 0 { - assert!(removes[n - 1] <= self.len()); - unsafe { - pgrx::pg_sys::PageIndexMultiDelete( - (self as *mut Self).cast(), - removes.as_ptr().cast_mut(), - removes.len() as _, - ); - } - } - } } impl Page for PostgresPage { @@ -257,16 +232,6 @@ impl PostgresRelation { pub unsafe fn new(raw: pgrx::pg_sys::Relation) -> Self { Self { raw } } - - #[allow(dead_code)] - pub fn len(&self) -> u32 { - unsafe { - pgrx::pg_sys::RelationGetNumberOfBlocksInFork( - self.raw, - pgrx::pg_sys::ForkNumber::MAIN_FORKNUM, - ) - } - } } impl RelationRead for PostgresRelation { diff --git a/src/index/utils.rs b/src/index/utils.rs deleted file mode 100644 index 18234ac0..00000000 --- a/src/index/utils.rs +++ /dev/null @@ -1,34 +0,0 @@ -use std::num::NonZeroU64; - -pub const fn pointer_to_ctid(pointer: NonZeroU64) -> pgrx::pg_sys::ItemPointerData { - let value = pointer.get(); - pgrx::pg_sys::ItemPointerData { - ip_blkid: pgrx::pg_sys::BlockIdData { - bi_hi: ((value >> 32) & 0xffff) as u16, - bi_lo: ((value >> 16) & 0xffff) as u16, - }, - ip_posid: (value & 0xffff) as u16, - } -} - -pub const fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> NonZeroU64 { - let mut value = 0; - value |= (ctid.ip_blkid.bi_hi as u64) << 32; - value |= (ctid.ip_blkid.bi_lo as u64) << 16; - value |= ctid.ip_posid as u64; - NonZeroU64::new(value).expect("invalid pointer") -} - -#[allow(dead_code)] -const fn soundness_check(a: pgrx::pg_sys::ItemPointerData) { - let b = ctid_to_pointer(a); - let c = pointer_to_ctid(b); - assert!(a.ip_blkid.bi_hi == c.ip_blkid.bi_hi); - assert!(a.ip_blkid.bi_lo == c.ip_blkid.bi_lo); - assert!(a.ip_posid == c.ip_posid); -} - -const _: () = soundness_check(pgrx::pg_sys::ItemPointerData { - ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi: 1, bi_lo: 2 }, - ip_posid: 3, -}); diff --git a/src/lib.rs b/src/lib.rs index 187e3cda..20ce1836 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,32 +1,22 @@ -#![feature(vec_pop_if)] #![allow(clippy::collapsible_else_if)] -#![allow(clippy::infallible_destructuring_match)] #![allow(clippy::too_many_arguments)] -#![allow(clippy::type_complexity)] +#![allow(unsafe_code)] -mod algorithm; mod datatype; -mod gucs; mod index; -mod postgres; -mod projection; -mod types; mod upgrade; -mod utils; pgrx::pg_module_magic!(); pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap); pgrx::extension_sql_file!("./sql/finalize.sql", finalize); #[pgrx::pg_guard] -unsafe extern "C" fn _PG_init() { +extern "C" fn _PG_init() { if unsafe { pgrx::pg_sys::IsUnderPostmaster } { pgrx::error!("vchord must be loaded via shared_preload_libraries."); } + index::init(); unsafe { - index::init(); - gucs::init(); - #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] diff --git a/src/upgrade/symbols.rs b/src/upgrade.rs similarity index 100% rename from src/upgrade/symbols.rs rename to src/upgrade.rs diff --git a/src/upgrade/mod.rs b/src/upgrade/mod.rs deleted file mode 100644 index 6eb441db..00000000 --- a/src/upgrade/mod.rs +++ /dev/null @@ -1 +0,0 @@ -mod symbols; diff --git a/src/utils/mod.rs b/src/utils/mod.rs deleted file mode 100644 index 85a84e0e..00000000 --- a/src/utils/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub mod k_means; -pub mod parallelism; -pub mod pipe; diff --git a/src/utils/parallelism.rs b/src/utils/parallelism.rs deleted file mode 100644 index b960b568..00000000 --- a/src/utils/parallelism.rs +++ /dev/null @@ -1,62 +0,0 @@ -use std::any::Any; -use std::panic::AssertUnwindSafe; -use std::sync::Arc; - -pub use rayon::iter::ParallelIterator; - -pub trait Parallelism: Send + Sync { - fn check(&self); - - fn rayon_into_par_iter(&self, x: I) -> I::Iter; -} - -struct ParallelismCheckPanic(Box); - -pub struct RayonParallelism { - stop: Arc, -} - -impl RayonParallelism { - pub fn scoped( - num_threads: usize, - stop: Arc, - f: impl FnOnce(&Self) -> R, - ) -> Result { - match std::panic::catch_unwind(AssertUnwindSafe(|| { - rayon::ThreadPoolBuilder::new() - .num_threads(num_threads) - .panic_handler(|e| { - if e.downcast_ref::().is_some() { - return; - } - log::error!("Asynchronous task panickied."); - }) - .build_scoped( - |thread| thread.run(), - |_| { - let pool = Self { stop: stop.clone() }; - f(&pool) - }, - ) - })) { - Ok(x) => x, - Err(e) => match e.downcast::() { - Ok(payload) => std::panic::resume_unwind((*payload).0), - Err(e) => std::panic::resume_unwind(e), - }, - } - } -} - -impl Parallelism for RayonParallelism { - fn check(&self) { - match std::panic::catch_unwind(AssertUnwindSafe(|| (self.stop)())) { - Ok(()) => (), - Err(payload) => std::panic::panic_any(ParallelismCheckPanic(payload)), - } - } - - fn rayon_into_par_iter(&self, x: I) -> I::Iter { - x.into_par_iter() - } -} From a3861ed1effbfb7ea6ab6a778fc0aab8626f4a56 Mon Sep 17 00:00:00 2001 From: usamoi Date: Sat, 8 Feb 2025 15:15:45 +0800 Subject: [PATCH 104/324] feat: pinning index in memory when building, second try (#181) issue: #117 first try: #150 --------- Signed-off-by: usamoi --- Cargo.lock | 2 + Cargo.toml | 2 + crates/algorithm/src/build.rs | 2 +- crates/algorithm/src/cache.rs | 51 +++ crates/algorithm/src/lib.rs | 4 + crates/algorithm/src/types.rs | 90 +----- src/index/am/am_build.rs | 588 ++++++++++++++++++++++++++++------ src/index/mod.rs | 1 + src/index/storage.rs | 8 + src/index/types.rs | 108 +++++++ tests/logic/pin.slt | 18 ++ 11 files changed, 696 insertions(+), 178 deletions(-) create mode 100644 crates/algorithm/src/cache.rs create mode 100644 src/index/types.rs create mode 100644 tests/logic/pin.slt diff --git a/Cargo.lock b/Cargo.lock index 62d7b795..aa5d7f37 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1526,6 +1526,8 @@ dependencies = [ "toml", "validator", "vector", + "zerocopy 0.8.14", + "zerocopy-derive 0.8.14", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 55cf4c2a..1b4f79dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,8 @@ rand.workspace = true serde.workspace = true toml = "0.8.19" validator.workspace = true +zerocopy.workspace = true +zerocopy-derive.workspace = true [patch.crates-io] half = { git = "https://github.com/tensorchord/half-rs.git", rev = "3f9a8843d6722bd1833de2289347640ad8770146" } diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 685aaf13..3db4e271 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -7,7 +7,7 @@ use vector::VectorOwned; pub fn build( vector_options: VectorOptions, - vchordrq_options: VchordrqIndexingOptions, + vchordrq_options: VchordrqIndexOptions, index: impl RelationWrite, structures: Vec>, ) { diff --git a/crates/algorithm/src/cache.rs b/crates/algorithm/src/cache.rs new file mode 100644 index 00000000..0405708e --- /dev/null +++ b/crates/algorithm/src/cache.rs @@ -0,0 +1,51 @@ +use crate::pipe::Pipe; +use crate::tuples::{H1Tuple, H1TupleReader, MetaTuple, read_tuple}; +use crate::{Page, RelationRead}; + +pub fn cache(index: impl RelationRead) -> Vec { + let mut trace = Vec::::new(); + let mut read = |id| { + let result = index.read(id); + trace.push(id); + result + }; + let meta_guard = read(0); + let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let height_of_root = meta_tuple.height_of_root(); + let root_first = meta_tuple.root_first(); + drop(meta_guard); + type State = Vec; + let mut state: State = vec![root_first]; + let mut step = |state: State| { + let mut results = Vec::new(); + for first in state { + let mut current = first; + while current != u32::MAX { + let h1_guard = read(current); + for i in 1..=h1_guard.len() { + let h1_tuple = h1_guard + .get(i) + .expect("data corruption") + .pipe(read_tuple::); + match h1_tuple { + H1TupleReader::_0(h1_tuple) => { + for first in h1_tuple.first().iter().copied() { + results.push(first); + } + } + H1TupleReader::_1(_) => (), + } + } + current = h1_guard.get_opaque().next; + } + } + results + }; + for _ in (1..height_of_root).rev() { + state = step(state); + } + for first in state { + let _ = read(first); + } + trace +} diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index dccc4594..4ed0da10 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -5,6 +5,7 @@ mod build; mod bulkdelete; +mod cache; mod freepages; mod insert; mod maintain; @@ -20,14 +21,17 @@ pub mod types; pub use build::build; pub use bulkdelete::bulkdelete; +pub use cache::cache; pub use insert::insert; pub use maintain::maintain; pub use prewarm::prewarm; pub use search::search; use std::ops::{Deref, DerefMut}; +use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] pub struct Opaque { pub next: u32, pub skip: u32, diff --git a/crates/algorithm/src/types.rs b/crates/algorithm/src/types.rs index a8f3b8ee..c2545670 100644 --- a/crates/algorithm/src/types.rs +++ b/crates/algorithm/src/types.rs @@ -1,98 +1,16 @@ use half::f16; use serde::{Deserialize, Serialize}; -use validator::{Validate, ValidationError, ValidationErrors}; +use validator::{Validate, ValidationError}; use vector::vect::{VectBorrowed, VectOwned}; -#[derive(Debug, Clone, Serialize, Deserialize, Validate)] -#[serde(deny_unknown_fields)] -pub struct VchordrqInternalBuildOptions { - #[serde(default = "VchordrqInternalBuildOptions::default_lists")] - #[validate(length(min = 1, max = 8), custom(function = VchordrqInternalBuildOptions::validate_lists))] - pub lists: Vec, - #[serde(default = "VchordrqInternalBuildOptions::default_spherical_centroids")] - pub spherical_centroids: bool, - #[serde(default = "VchordrqInternalBuildOptions::default_sampling_factor")] - #[validate(range(min = 1, max = 1024))] - pub sampling_factor: u32, - #[serde(default = "VchordrqInternalBuildOptions::default_build_threads")] - #[validate(range(min = 1, max = 255))] - pub build_threads: u16, -} - -impl VchordrqInternalBuildOptions { - fn default_lists() -> Vec { - vec![1000] - } - fn validate_lists(lists: &[u32]) -> Result<(), ValidationError> { - if !lists.is_sorted() { - return Err(ValidationError::new("`lists` should be in ascending order")); - } - if !lists.iter().all(|x| (1..=1 << 24).contains(x)) { - return Err(ValidationError::new("list is too long or too short")); - } - Ok(()) - } - fn default_spherical_centroids() -> bool { - false - } - fn default_sampling_factor() -> u32 { - 256 - } - fn default_build_threads() -> u16 { - 1 - } -} - -impl Default for VchordrqInternalBuildOptions { - fn default() -> Self { - Self { - lists: Self::default_lists(), - spherical_centroids: Self::default_spherical_centroids(), - sampling_factor: Self::default_sampling_factor(), - build_threads: Self::default_build_threads(), - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize, Validate)] -#[serde(deny_unknown_fields)] -pub struct VchordrqExternalBuildOptions { - pub table: String, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -#[serde(rename_all = "snake_case")] -pub enum VchordrqBuildOptions { - Internal(VchordrqInternalBuildOptions), - External(VchordrqExternalBuildOptions), -} - -impl Default for VchordrqBuildOptions { - fn default() -> Self { - Self::Internal(Default::default()) - } -} - -impl Validate for VchordrqBuildOptions { - fn validate(&self) -> Result<(), ValidationErrors> { - use VchordrqBuildOptions::*; - match self { - Internal(internal_build) => internal_build.validate(), - External(external_build) => external_build.validate(), - } - } -} - #[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] -pub struct VchordrqIndexingOptions { - #[serde(default = "VchordrqIndexingOptions::default_residual_quantization")] +pub struct VchordrqIndexOptions { + #[serde(default = "VchordrqIndexOptions::default_residual_quantization")] pub residual_quantization: bool, - pub build: VchordrqBuildOptions, } -impl VchordrqIndexingOptions { +impl VchordrqIndexOptions { fn default_residual_quantization() -> bool { false } diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index e7d65587..d0661879 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -2,14 +2,17 @@ use crate::datatype::typmod::Typmod; use crate::index::am::{Reloption, ctid_to_pointer}; use crate::index::opclass::{Opfamily, opfamily}; use crate::index::projection::RandomProject; -use crate::index::storage::PostgresRelation; +use crate::index::storage::{PostgresPage, PostgresRelation}; +use crate::index::types::*; use algorithm::operator::{Dot, L2, Op, Vector}; use algorithm::types::*; +use algorithm::{PageGuard, RelationRead, RelationWrite}; use half::f16; use pgrx::pg_sys::Datum; use rand::Rng; use simd::Floating; use std::num::NonZeroU64; +use std::ops::Deref; use std::sync::Arc; use vector::vect::VectOwned; use vector::{VectorBorrowed, VectorOwned}; @@ -123,11 +126,11 @@ pub unsafe extern "C" fn ambuild( }; let index = unsafe { PostgresRelation::new(index_relation) }; let mut reporter = PostgresReporter {}; - let structures = match vchordrq_options.build.clone() { - VchordrqBuildOptions::External(external_build) => { + let structures = match vchordrq_options.build.source.clone() { + VchordrqBuildSourceOptions::External(external_build) => { make_external_build(vector_options.clone(), opfamily, external_build.clone()) } - VchordrqBuildOptions::Internal(internal_build) => { + VchordrqBuildSourceOptions::Internal(internal_build) => { let mut tuples_total = 0_u64; let samples = { let mut rand = rand::thread_rng(); @@ -185,32 +188,55 @@ pub unsafe extern "C" fn ambuild( match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => algorithm::build::, L2>>( vector_options, - vchordrq_options, + vchordrq_options.index, index.clone(), map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), ), (VectorKind::Vecf32, DistanceKind::Dot) => algorithm::build::, Dot>>( vector_options, - vchordrq_options, + vchordrq_options.index, index.clone(), map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), ), (VectorKind::Vecf16, DistanceKind::L2) => algorithm::build::, L2>>( vector_options, - vchordrq_options, + vchordrq_options.index, index.clone(), map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), ), (VectorKind::Vecf16, DistanceKind::Dot) => algorithm::build::, Dot>>( vector_options, - vchordrq_options, + vchordrq_options.index, index.clone(), map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), ), } - if let Some(leader) = - unsafe { VchordrqLeader::enter(heap_relation, index_relation, (*index_info).ii_Concurrent) } - { + let cache = if vchordrq_options.build.pin { + let mut trace = algorithm::cache(index.clone()); + trace.sort(); + trace.dedup(); + if let Some(max) = trace.last().copied() { + let mut mapping = vec![u32::MAX; 1 + max as usize]; + let mut pages = Vec::>::with_capacity(trace.len()); + for id in trace { + mapping[id as usize] = pages.len() as u32; + pages.push(index.read(id).clone_into_boxed()); + } + vchordrq_cached::VchordrqCached::_1 { mapping, pages } + } else { + vchordrq_cached::VchordrqCached::_0 {} + } + } else { + vchordrq_cached::VchordrqCached::_0 {} + }; + if let Some(leader) = unsafe { + VchordrqLeader::enter( + heap_relation, + index_relation, + (*index_info).ii_Concurrent, + cache, + ) + } { unsafe { parallel_build( index_relation, @@ -218,6 +244,7 @@ pub unsafe extern "C" fn ambuild( index_info, leader.tablescandesc, leader.vchordrqshared, + leader.vchordrqcached, Some(reporter), ); leader.wait(); @@ -324,6 +351,179 @@ struct VchordrqShared { indtuples: u64, } +mod vchordrq_cached { + pub const ALIGN: usize = 8; + pub type Tag = u64; + + use crate::index::storage::PostgresPage; + use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; + + #[repr(C, align(8))] + #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct VchordrqCachedHeader0 {} + + #[repr(C, align(8))] + #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct VchordrqCachedHeader1 { + mapping_s: usize, + mapping_e: usize, + pages_s: usize, + pages_e: usize, + } + + pub enum VchordrqCached { + _0 {}, + _1 { + mapping: Vec, + pages: Vec>, + }, + } + + impl VchordrqCached { + pub fn serialize(&self) -> Vec { + let mut buffer = Vec::new(); + match self { + VchordrqCached::_0 {} => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer[size_of::()..][..size_of::()] + .copy_from_slice(VchordrqCachedHeader0 {}.as_bytes()); + } + VchordrqCached::_1 { mapping, pages } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat(0).take(size_of::())); + let mapping_s = buffer.len(); + buffer.extend(mapping.as_bytes()); + let mapping_e = buffer.len(); + while buffer.len() % ALIGN != 0 { + buffer.push(0u8); + } + let pages_s = buffer.len(); + buffer.extend(pages.iter().flat_map(|x| unsafe { + std::mem::transmute::<&PostgresPage, &[u8; 8192]>(x.as_ref()) + })); + let pages_e = buffer.len(); + while buffer.len() % ALIGN != 0 { + buffer.push(0u8); + } + buffer[size_of::()..][..size_of::()] + .copy_from_slice( + VchordrqCachedHeader1 { + mapping_s, + mapping_e, + pages_s, + pages_e, + } + .as_bytes(), + ); + } + } + buffer + } + } + + #[derive(Debug, Clone, Copy)] + pub enum VchordrqCachedReader<'a> { + _0(#[allow(dead_code)] VchordrqCachedReader0<'a>), + _1(VchordrqCachedReader1<'a>), + } + + #[derive(Debug, Clone, Copy)] + pub struct VchordrqCachedReader0<'a> { + #[allow(dead_code)] + header: &'a VchordrqCachedHeader0, + } + + #[derive(Debug, Clone, Copy)] + pub struct VchordrqCachedReader1<'a> { + #[allow(dead_code)] + header: &'a VchordrqCachedHeader1, + mapping: &'a [u32], + pages: &'a [PostgresPage], + } + + impl<'a> VchordrqCachedReader1<'a> { + pub fn get(&self, id: u32) -> Option<&'a PostgresPage> { + let index = *self.mapping.get(id as usize)?; + if index == u32::MAX { + return None; + } + Some(&self.pages[index as usize]) + } + } + + impl<'a> VchordrqCachedReader<'a> { + pub fn deserialize_ref(source: &'a [u8]) -> Self { + let tag = u64::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &VchordrqCachedHeader0 = checker.prefix(size_of::()); + Self::_0(VchordrqCachedReader0 { header }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &VchordrqCachedHeader1 = checker.prefix(size_of::()); + let mapping = checker.bytes(header.mapping_s, header.mapping_e); + let pages = + unsafe { checker.bytes_slice_unchecked(header.pages_s, header.pages_e) }; + Self::_1(VchordrqCachedReader1 { + header, + mapping, + pages, + }) + } + _ => panic!("bad bytes"), + } + } + } + + pub struct RefChecker<'a> { + bytes: &'a [u8], + } + + impl<'a> RefChecker<'a> { + pub fn new(bytes: &'a [u8]) -> Self { + Self { bytes } + } + pub fn prefix( + &self, + s: usize, + ) -> &'a T { + let start = s; + let end = s + size_of::(); + let bytes = &self.bytes[start..end]; + FromBytes::ref_from_bytes(bytes).expect("bad bytes") + } + pub fn bytes( + &self, + s: usize, + e: usize, + ) -> &'a T { + let start = s; + let end = e; + let bytes = &self.bytes[start..end]; + FromBytes::ref_from_bytes(bytes).expect("bad bytes") + } + pub unsafe fn bytes_slice_unchecked(&self, s: usize, e: usize) -> &'a [T] { + let start = s; + let end = e; + let bytes = &self.bytes[start..end]; + if size_of::() == 0 || bytes.len() % size_of::() == 0 { + let ptr = bytes as *const [u8] as *const T; + if ptr.is_aligned() { + unsafe { std::slice::from_raw_parts(ptr, bytes.len() / size_of::()) } + } else { + panic!("bad bytes") + } + } else { + panic!("bad bytes") + } + } + } +} + fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { matches!( unsafe { (*snapshot).snapshot_type }, @@ -335,9 +535,10 @@ fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { struct VchordrqLeader { pcxt: *mut pgrx::pg_sys::ParallelContext, nparticipants: i32, + snapshot: pgrx::pg_sys::Snapshot, vchordrqshared: *mut VchordrqShared, tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, - snapshot: pgrx::pg_sys::Snapshot, + vchordrqcached: *const u8, } impl VchordrqLeader { @@ -345,7 +546,12 @@ impl VchordrqLeader { heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, isconcurrent: bool, + cache: vchordrq_cached::VchordrqCached, ) -> Option { + let _cache = cache.serialize(); + drop(cache); + let cache = _cache; + unsafe fn compute_parallel_workers( heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, @@ -407,6 +613,8 @@ impl VchordrqLeader { estimate_keys(&mut (*pcxt).estimator, 1); estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, 8 + cache.len()); + estimate_keys(&mut (*pcxt).estimator, 1); } unsafe { @@ -446,9 +654,17 @@ impl VchordrqLeader { tablescandesc }; + let vchordrqcached = unsafe { + let x = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, 8 + cache.len()).cast::(); + (x as *mut u64).write_unaligned(cache.len() as _); + std::ptr::copy(cache.as_ptr(), x.add(8), cache.len()); + x + }; + unsafe { pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000001, vchordrqshared.cast()); pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000003, vchordrqcached.cast()); } unsafe { @@ -472,9 +688,10 @@ impl VchordrqLeader { Some(Self { pcxt, nparticipants: nworkers_launched + 1, + snapshot, vchordrqshared, tablescandesc, - snapshot, + vchordrqcached, }) } @@ -513,6 +730,11 @@ pub unsafe extern "C" fn vchordrq_parallel_build_main( pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) .cast::() }; + let vchordrqcached = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000003, false) + .cast::() + .cast_const() + }; let heap_lockmode; let index_lockmode; if unsafe { !(*vchordrqshared).isconcurrent } { @@ -530,7 +752,15 @@ pub unsafe extern "C" fn vchordrq_parallel_build_main( } unsafe { - parallel_build(index, heap, index_info, tablescandesc, vchordrqshared, None); + parallel_build( + index, + heap, + index_info, + tablescandesc, + vchordrqshared, + vchordrqcached, + None, + ); } unsafe { @@ -545,8 +775,15 @@ unsafe fn parallel_build( index_info: *mut pgrx::pg_sys::IndexInfo, tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, vchordrqshared: *mut VchordrqShared, + vchordrqcached: *const u8, mut reporter: Option, ) { + use vchordrq_cached::VchordrqCachedReader; + let cached = VchordrqCachedReader::deserialize_ref(unsafe { + let bytes = (vchordrqcached as *const u64).read_unaligned(); + std::slice::from_raw_parts(vchordrqcached.add(8), bytes as _) + }); + let index = unsafe { PostgresRelation::new(index_relation) }; let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap_relation, tablescandesc) }; @@ -558,90 +795,184 @@ unsafe fn parallel_build( opfamily, scan, }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, L2>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + match cached { + VchordrqCachedReader::_0(_) => match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, L2>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); + }); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, Dot>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } } - } - }); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, Dot>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + }); + } + (VectorKind::Vecf16, DistanceKind::L2) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, L2>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); + }); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, Dot>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } } + }); + } + }, + VchordrqCachedReader::_1(cached) => { + let index = CachingRelation { + cache: cached, + relation: index, + }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, L2>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); } - }); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, L2>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } + (VectorKind::Vecf32, DistanceKind::Dot) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, Dot>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); } - }); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, Dot>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } + (VectorKind::Vecf16, DistanceKind::L2) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, L2>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); } - }); + (VectorKind::Vecf16, DistanceKind::Dot) => { + heap.traverse(true, |(pointer, vector): (_, VectOwned)| { + algorithm::insert::, Dot>>( + index.clone(), + pointer, + RandomProject::project(vector.as_borrowed()), + ); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } + } + }); + } + } } } unsafe { @@ -948,3 +1279,78 @@ impl InternalBuild for VectOwned { Self::new(f16::vector_from_f32(x)) } } + +struct CachingRelation<'a, R> { + cache: vchordrq_cached::VchordrqCachedReader1<'a>, + relation: R, +} + +impl Clone for CachingRelation<'_, R> { + fn clone(&self) -> Self { + Self { + cache: self.cache, + relation: self.relation.clone(), + } + } +} + +enum CachingRelationReadGuard<'a, G: Deref> { + Wrapping(G), + Cached(u32, &'a G::Target), +} + +impl PageGuard for CachingRelationReadGuard<'_, G> { + fn id(&self) -> u32 { + match self { + CachingRelationReadGuard::Wrapping(x) => x.id(), + CachingRelationReadGuard::Cached(id, _) => *id, + } + } +} + +impl Deref for CachingRelationReadGuard<'_, G> { + type Target = G::Target; + + fn deref(&self) -> &Self::Target { + match self { + CachingRelationReadGuard::Wrapping(x) => x, + CachingRelationReadGuard::Cached(_, page) => page, + } + } +} + +impl> RelationRead for CachingRelation<'_, R> { + type Page = R::Page; + + type ReadGuard<'a> + = CachingRelationReadGuard<'a, R::ReadGuard<'a>> + where + Self: 'a; + + fn read(&self, id: u32) -> Self::ReadGuard<'_> { + if let Some(x) = self.cache.get(id) { + CachingRelationReadGuard::Cached(id, x) + } else { + CachingRelationReadGuard::Wrapping(self.relation.read(id)) + } + } +} + +impl> RelationWrite for CachingRelation<'_, R> { + type WriteGuard<'a> + = R::WriteGuard<'a> + where + Self: 'a; + + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.relation.write(id, tracking_freespace) + } + + fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.relation.extend(tracking_freespace) + } + + fn search(&self, freespace: usize) -> Option> { + self.relation.search(freespace) + } +} diff --git a/src/index/mod.rs b/src/index/mod.rs index 511a4f57..e7a309ae 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -4,6 +4,7 @@ pub mod gucs; pub mod opclass; pub mod projection; pub mod storage; +pub mod types; pub fn init() { am::init(); diff --git a/src/index/storage.rs b/src/index/storage.rs index a190ac08..5d0cfc94 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -16,6 +16,7 @@ const fn size_of_contents() -> usize { } #[repr(C, align(8))] +#[derive(Debug)] pub struct PostgresPage { header: pgrx::pg_sys::PageHeaderData, content: [u8; size_of_contents()], @@ -42,6 +43,13 @@ impl PostgresPage { assert_eq!(offset_of!(Self, opaque), this.header.pd_special as usize); this } + pub fn clone_into_boxed(&self) -> Box { + let mut result = Box::new_uninit(); + unsafe { + std::ptr::copy(self as *const Self, result.as_mut_ptr(), 1); + result.assume_init() + } + } } impl Page for PostgresPage { diff --git a/src/index/types.rs b/src/index/types.rs new file mode 100644 index 00000000..6d6bb3ef --- /dev/null +++ b/src/index/types.rs @@ -0,0 +1,108 @@ +use algorithm::types::VchordrqIndexOptions; +use serde::{Deserialize, Serialize}; +use validator::{Validate, ValidationError, ValidationErrors}; + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqInternalBuildOptions { + #[serde(default = "VchordrqInternalBuildOptions::default_lists")] + #[validate(length(min = 1, max = 8), custom(function = VchordrqInternalBuildOptions::validate_lists))] + pub lists: Vec, + #[serde(default = "VchordrqInternalBuildOptions::default_spherical_centroids")] + pub spherical_centroids: bool, + #[serde(default = "VchordrqInternalBuildOptions::default_sampling_factor")] + #[validate(range(min = 1, max = 1024))] + pub sampling_factor: u32, + #[serde(default = "VchordrqInternalBuildOptions::default_build_threads")] + #[validate(range(min = 1, max = 255))] + pub build_threads: u16, +} + +impl VchordrqInternalBuildOptions { + fn default_lists() -> Vec { + vec![1000] + } + fn validate_lists(lists: &[u32]) -> Result<(), ValidationError> { + if !lists.is_sorted() { + return Err(ValidationError::new("`lists` should be in ascending order")); + } + if !lists.iter().all(|x| (1..=1 << 24).contains(x)) { + return Err(ValidationError::new("list is too long or too short")); + } + Ok(()) + } + fn default_spherical_centroids() -> bool { + false + } + fn default_sampling_factor() -> u32 { + 256 + } + fn default_build_threads() -> u16 { + 1 + } +} + +impl Default for VchordrqInternalBuildOptions { + fn default() -> Self { + Self { + lists: Self::default_lists(), + spherical_centroids: Self::default_spherical_centroids(), + sampling_factor: Self::default_sampling_factor(), + build_threads: Self::default_build_threads(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqExternalBuildOptions { + pub table: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +#[serde(rename_all = "snake_case")] +pub enum VchordrqBuildSourceOptions { + Internal(VchordrqInternalBuildOptions), + External(VchordrqExternalBuildOptions), +} + +impl Default for VchordrqBuildSourceOptions { + fn default() -> Self { + Self::Internal(Default::default()) + } +} + +impl Validate for VchordrqBuildSourceOptions { + fn validate(&self) -> Result<(), ValidationErrors> { + use VchordrqBuildSourceOptions::*; + match self { + Internal(internal_build) => internal_build.validate(), + External(external_build) => external_build.validate(), + } + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +#[serde(rename_all = "snake_case")] +pub struct VchordrqBuildOptions { + #[serde(flatten)] + pub source: VchordrqBuildSourceOptions, + #[serde(default = "VchordrqBuildOptions::default_pin")] + pub pin: bool, +} + +impl VchordrqBuildOptions { + pub fn default_pin() -> bool { + false + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqIndexingOptions { + #[serde(flatten)] + pub index: VchordrqIndexOptions, + pub build: VchordrqBuildOptions, +} diff --git a/tests/logic/pin.slt b/tests/logic/pin.slt new file mode 100644 index 00000000..fbe923c3 --- /dev/null +++ b/tests/logic/pin.slt @@ -0,0 +1,18 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_ip_ops) +WITH (options = $$ +residual_quantization = false +build.pin = true +[build.internal] +lists = [32] +spherical_centroids = true +$$); + +statement ok +DROP TABLE t; From 4b4f3639bfacc6e547f9a5940a835317a80d1e96 Mon Sep 17 00:00:00 2001 From: usamoi Date: Sat, 8 Feb 2025 18:12:53 +0800 Subject: [PATCH 105/324] fix: use linked list of vectors to skip realloc (#182) Signed-off-by: usamoi --- .github/workflows/psql.yml | 2 ++ .github/workflows/rust.yml | 2 ++ crates/algorithm/src/insert.rs | 5 +++-- crates/algorithm/src/lib.rs | 1 + crates/algorithm/src/linkedvec.rs | 36 +++++++++++++++++++++++++++++++ crates/algorithm/src/search.rs | 9 ++++---- crates/simd/cshim.c | 5 ++--- 7 files changed, 51 insertions(+), 9 deletions(-) create mode 100644 crates/algorithm/src/linkedvec.rs diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml index 42c108c0..4bd9e90b 100644 --- a/.github/workflows/psql.yml +++ b/.github/workflows/psql.yml @@ -4,6 +4,7 @@ on: pull_request: paths: - '.github/workflows/psql.yml' + - 'crates/**' - 'src/**' - 'Cargo.lock' - 'Cargo.toml' @@ -16,6 +17,7 @@ on: - main paths: - '.github/workflows/psql.yml' + - 'crates/**' - 'src/**' - 'Cargo.lock' - 'Cargo.toml' diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 8b1b4e2d..f026d699 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -4,6 +4,7 @@ on: pull_request: paths: - '.github/workflows/rust.yml' + - 'crates/**' - 'src/**' - 'Cargo.lock' - 'Cargo.toml' @@ -14,6 +15,7 @@ on: - main paths: - '.github/workflows/rust.yml' + - 'crates/**' - 'src/**' - 'Cargo.lock' - 'Cargo.toml' diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 433e6002..06f18231 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -1,3 +1,4 @@ +use crate::linkedvec::LinkedVec; use crate::operator::*; use crate::pipe::Pipe; use crate::tape::{access_1, append}; @@ -49,7 +50,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto } }; let step = |state: State| { - let mut results = Vec::new(); + let mut results = LinkedVec::new(); { let (first, residual) = state; let lut = if let Some(residual) = residual { @@ -71,7 +72,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto }, ); } - let mut heap = BinaryHeap::from(results); + let mut heap = BinaryHeap::from(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 4ed0da10..f3968995 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -8,6 +8,7 @@ mod bulkdelete; mod cache; mod freepages; mod insert; +mod linkedvec; mod maintain; mod pipe; mod prewarm; diff --git a/crates/algorithm/src/linkedvec.rs b/crates/algorithm/src/linkedvec.rs new file mode 100644 index 00000000..179a0489 --- /dev/null +++ b/crates/algorithm/src/linkedvec.rs @@ -0,0 +1,36 @@ +pub struct LinkedVec { + inner: Vec>, + last: Vec, +} + +impl LinkedVec { + pub fn new() -> Self { + Self { + inner: Vec::new(), + last: Vec::with_capacity(4096), + } + } + pub fn push(&mut self, value: T) { + if self.last.len() == self.last.capacity() { + self.reserve(); + } + #[allow(unsafe_code)] + unsafe { + std::hint::assert_unchecked(self.last.len() != self.last.capacity()); + } + self.last.push(value); + } + #[cold] + fn reserve(&mut self) { + let fresh = Vec::with_capacity(self.last.capacity() * 4); + self.inner.push(core::mem::replace(&mut self.last, fresh)); + } + pub fn into_vec(self) -> Vec { + let mut last = self.last; + last.reserve(self.inner.iter().map(Vec::len).sum::()); + for x in self.inner { + last.extend(x); + } + last + } +} diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index bff53e38..fd200ac7 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -1,3 +1,4 @@ +use crate::linkedvec::LinkedVec; use crate::operator::*; use crate::pipe::Pipe; use crate::tape::{access_0, access_1}; @@ -51,7 +52,7 @@ pub fn search( } }]; let step = |state: State, probes| { - let mut results = Vec::new(); + let mut results = LinkedVec::new(); for (first, residual) in state { let lut = if let Some(residual) = residual { &O::Vector::compute_lut_block(residual.as_borrowed()) @@ -72,7 +73,7 @@ pub fn search( }, ); } - let mut heap = BinaryHeap::from(results); + let mut heap = BinaryHeap::from(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); std::iter::from_fn(|| { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { @@ -116,7 +117,7 @@ pub fn search( state = step(state, probes[i as usize - 1]); } - let mut results = Vec::new(); + let mut results = LinkedVec::new(); for (first, residual) in state { let lut = if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { &O::Vector::compute_lut(residual) @@ -144,7 +145,7 @@ pub fn search( }, ); } - let mut heap = BinaryHeap::from(results); + let mut heap = BinaryHeap::from(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _)>::new(); std::iter::from_fn(move || { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { diff --git a/crates/simd/cshim.c b/crates/simd/cshim.c index 1374bbfe..10b5ceea 100644 --- a/crates/simd/cshim.c +++ b/crates/simd/cshim.c @@ -2,11 +2,10 @@ #error "clang version must be >= 16" #endif -#include -#include - #ifdef __aarch64__ +#include +#include #include #include From 527991d4a4bb92c62a7b092dc8e28459a1c40f99 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 10 Feb 2025 17:03:23 +0800 Subject: [PATCH 106/324] feat: use select algorithm to replace heap, if k in top-k is expected to be small (#183) Signed-off-by: usamoi --- Cargo.lock | 13 ++++ crates/algorithm/Cargo.toml | 2 + crates/algorithm/src/insert.rs | 5 +- crates/algorithm/src/lib.rs | 3 +- .../src/{linkedvec.rs => linked_vec.rs} | 0 crates/algorithm/src/search.rs | 2 +- crates/algorithm/src/select_heap.rs | 63 +++++++++++++++++++ 7 files changed, 84 insertions(+), 4 deletions(-) rename crates/algorithm/src/{linkedvec.rs => linked_vec.rs} (100%) create mode 100644 crates/algorithm/src/select_heap.rs diff --git a/Cargo.lock b/Cargo.lock index aa5d7f37..7e992b88 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -18,6 +18,7 @@ dependencies = [ "always_equal", "distance", "half 2.4.1", + "heapify", "k_means", "paste", "rabitq", @@ -26,6 +27,7 @@ dependencies = [ "serde", "simd", "toml", + "turboselect", "validator", "vector", "zerocopy 0.8.14", @@ -390,6 +392,12 @@ version = "0.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" +[[package]] +name = "heapify" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0049b265b7f201ca9ab25475b22b47fe444060126a51abe00f77d986fc5cc52e" + [[package]] name = "heapless" version = "0.8.0" @@ -1410,6 +1418,11 @@ dependencies = [ "winnow", ] +[[package]] +name = "turboselect" +version = "0.1.0" +source = "git+https://github.com/tensorchord/turboselect.git?rev=d8753c4ffe5b47f28670fea21d56cf3658d51b9b#d8753c4ffe5b47f28670fea21d56cf3658d51b9b" + [[package]] name = "typenum" version = "1.17.0" diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml index 1180ba96..006fe9bd 100644 --- a/crates/algorithm/Cargo.toml +++ b/crates/algorithm/Cargo.toml @@ -13,10 +13,12 @@ simd = { path = "../simd" } vector = { path = "../vector" } half.workspace = true +heapify = "0.2.0" paste = "1" rand.workspace = true serde.workspace = true toml = "0.8.19" +turboselect = { git = "https://github.com/tensorchord/turboselect.git", rev = "d8753c4ffe5b47f28670fea21d56cf3658d51b9b" } validator.workspace = true zerocopy.workspace = true zerocopy-derive.workspace = true diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 06f18231..4fb947fd 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -1,6 +1,7 @@ -use crate::linkedvec::LinkedVec; +use crate::linked_vec::LinkedVec; use crate::operator::*; use crate::pipe::Pipe; +use crate::select_heap::SelectHeap; use crate::tape::{access_1, append}; use crate::tuples::*; use crate::vectors::{self}; @@ -72,7 +73,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto }, ); } - let mut heap = BinaryHeap::from(results.into_vec()); + let mut heap = SelectHeap::from_vec(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index f3968995..b395f36d 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -8,11 +8,12 @@ mod bulkdelete; mod cache; mod freepages; mod insert; -mod linkedvec; +mod linked_vec; mod maintain; mod pipe; mod prewarm; mod search; +mod select_heap; mod tape; mod tuples; mod vectors; diff --git a/crates/algorithm/src/linkedvec.rs b/crates/algorithm/src/linked_vec.rs similarity index 100% rename from crates/algorithm/src/linkedvec.rs rename to crates/algorithm/src/linked_vec.rs diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index fd200ac7..c6bdc5cc 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -1,4 +1,4 @@ -use crate::linkedvec::LinkedVec; +use crate::linked_vec::LinkedVec; use crate::operator::*; use crate::pipe::Pipe; use crate::tape::{access_0, access_1}; diff --git a/crates/algorithm/src/select_heap.rs b/crates/algorithm/src/select_heap.rs new file mode 100644 index 00000000..86723881 --- /dev/null +++ b/crates/algorithm/src/select_heap.rs @@ -0,0 +1,63 @@ +pub struct SelectHeap { + threshold: usize, + inner: Vec, +} + +impl SelectHeap { + pub fn from_vec(mut vec: Vec) -> Self { + let n = vec.len(); + if n != 0 { + let threshold = n.saturating_sub(n.div_ceil(384)); + turboselect::select_nth_unstable(&mut vec, threshold); + vec[threshold..].sort_unstable(); + Self { + threshold, + inner: vec, + } + } else { + Self { + threshold: 0, + inner: Vec::new(), + } + } + } + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + pub fn pop(&mut self) -> Option { + if self.inner.len() <= self.threshold { + heapify::pop_heap(&mut self.inner); + } + let t = self.inner.pop(); + if self.inner.len() == self.threshold { + heapify::make_heap(&mut self.inner); + } + t + } + pub fn peek(&self) -> Option<&T> { + if self.inner.len() <= self.threshold { + self.inner.first() + } else { + self.inner.last() + } + } +} + +#[test] +fn test_select_heap() { + for _ in 0..1000 { + let sequence = (0..10000) + .map(|_| rand::random::()) + .collect::>(); + let answer = { + let mut x = sequence.clone(); + x.sort_by_key(|x| std::cmp::Reverse(*x)); + x + }; + let result = { + let mut x = SelectHeap::from_vec(sequence.clone()); + std::iter::from_fn(|| x.pop()).collect::>() + }; + assert_eq!(answer, result); + } +} From b6e1bef545f71edfb568158f30e386ff13ed5706 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 11 Feb 2025 14:12:21 +0800 Subject: [PATCH 107/324] ci: install pg13 in docker image (#186) Signed-off-by: usamoi --- docker/pgrx.Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile index 93ec4aff..d9f8f50c 100644 --- a/docker/pgrx.Dockerfile +++ b/docker/pgrx.Dockerfile @@ -38,7 +38,7 @@ RUN set -ex; \ RUN /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y # install all the PostgresQL RUN set -ex; \ - for v in $(seq 14 17); do \ + for v in $(seq 13 17); do \ apt install -y --no-install-recommends postgresql-$v postgresql-server-dev-$v postgresql-$v-pgvector; \ done; \ rm -rf /var/lib/apt/lists/*; @@ -57,7 +57,7 @@ RUN rustup target add $(uname -m)-unknown-linux-gnu RUN cargo install cargo-pgrx --locked --version=${PGRX_VERSION} RUN set -ex; \ - for v in $(seq 14 17); do \ + for v in $(seq 13 17); do \ cargo pgrx init --pg$v=/usr/lib/postgresql/$v/bin/pg_config; \ done; From 347f1b9791795493f08407da4960b74423218ed8 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 11 Feb 2025 17:23:05 +0800 Subject: [PATCH 108/324] ci: use less docker (#187) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tests in `./crates` do not run before。。。 Signed-off-by: usamoi --- .github/workflows/check.yml | 141 +++++++++++++++++++++++++++++++ .github/workflows/psql.yml | 103 ---------------------- .github/workflows/rust.yml | 83 ------------------ .github/workflows/style.yml | 40 --------- crates/simd/src/f16.rs | 2 + crates/simd/src/fast_scan/mod.rs | 12 +-- .taplo.toml => taplo.toml | 0 7 files changed, 149 insertions(+), 232 deletions(-) create mode 100644 .github/workflows/check.yml delete mode 100644 .github/workflows/psql.yml delete mode 100644 .github/workflows/rust.yml delete mode 100644 .github/workflows/style.yml rename .taplo.toml => taplo.toml (100%) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml new file mode 100644 index 00000000..f6721115 --- /dev/null +++ b/.github/workflows/check.yml @@ -0,0 +1,141 @@ +name: Check + +on: + pull_request: + paths: + - ".cargo" + - ".github/workflows/check.yml" + - "crates/**" + - "src/**" + - "build.rs" + - "Cargo.lock" + - "Cargo.toml" + - "rust-toolchain.toml" + - "rustfmt.toml" + - "taplo.toml" + - "vchord.control" + push: + paths: + - ".cargo" + - ".github/workflows/lint.yml" + - "crates/**" + - "src/**" + - "build.rs" + - "Cargo.lock" + - "Cargo.toml" + - "rust-toolchain.toml" + - "rustfmt.toml" + - "taplo.toml" + - "vchord.control" + merge_group: + workflow_dispatch: + +concurrency: + group: ${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +env: + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 + SCCACHE_GHA_ENABLED: true + RUSTC_WRAPPER: sccache + RUSTFLAGS: "-Dwarnings" + +jobs: + style: + runs-on: "ubuntu-24.04" + + steps: + - name: Set up Environment + run: | + curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo + + - name: Checkout + uses: actions/checkout@v4 + + - name: Typos + uses: crate-ci/typos@master + + - name: Taplo + run: taplo fmt --check + + - name: Ruff + uses: astral-sh/ruff-action@v1 + + - name: Rustfmt + run: cargo fmt --check + + lint: + strategy: + matrix: + runner: ["ubuntu-24.04", "ubuntu-24.04-arm"] + runs-on: ${{ matrix.runner }} + + steps: + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.7 + + - name: Checkout + uses: actions/checkout@v4 + + - name: Clippy + run: cargo clippy --workspace --exclude vchord + + - name: Cargo Test + run: cargo test --workspace --exclude vchord --no-fail-fast + + psql: + strategy: + matrix: + version: ["13", "14", "15", "16", "17"] + runner: ["ubuntu-24.04", "ubuntu-24.04-arm"] + exclude: + - version: "13" + runner: "ubuntu-24.04" + - version: "13" + runner: "ubuntu-24.04-arm" + runs-on: ${{ matrix.runner }} + + steps: + - name: Set up Environment + run: | + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' + + sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 + + sudo apt-get install -y postgresql-common + sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} + + sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector + echo "local all all trust" | sudo tee /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + echo "host all all 127.0.0.1/32 trust" | sudo tee -a /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + echo "host all all ::1/128 trust" | sudo tee -a /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + sudo -iu postgres createuser -s -r $USER + sudo -iu postgres createdb -O $USER $USER + sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.so"' + sudo systemctl stop postgresql + + curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.12.9/cargo-pgrx-v0.12.9-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx + cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.7 + + - name: Checkout + uses: actions/checkout@v4 + + - name: Clippy + run: cargo clippy -p vchord --features pg${{ matrix.version }} -- --no-deps + + - name: Install + run: cargo pgrx install -p vchord --features pg${{ matrix.version }} --release --sudo + + - name: Sqllogictest + run: | + sudo systemctl start postgresql + psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + sqllogictest --db $USER --user $USER './tests/**/*.slt' diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml deleted file mode 100644 index 4bd9e90b..00000000 --- a/.github/workflows/psql.yml +++ /dev/null @@ -1,103 +0,0 @@ -name: PostgresSQL - -on: - pull_request: - paths: - - '.github/workflows/psql.yml' - - 'crates/**' - - 'src/**' - - 'Cargo.lock' - - 'Cargo.toml' - - '*.control' - - 'rust-toolchain.toml' - - 'tests/**' - - 'tools/**' - push: - branches: - - main - paths: - - '.github/workflows/psql.yml' - - 'crates/**' - - 'src/**' - - 'Cargo.lock' - - 'Cargo.toml' - - '*.control' - - 'rust-toolchain.toml' - - 'tests/**' - - 'tools/**' - merge_group: - workflow_dispatch: - -concurrency: - group: ${{ github.ref }}-${{ github.workflow }} - cancel-in-progress: true - -env: - CARGO_TERM_COLOR: always - RUST_BACKTRACE: 1 - RUSTFLAGS: "-Dwarnings" - -jobs: - test: - runs-on: ${{ matrix.runner }} - strategy: - matrix: - version: ["14", "15", "16", "17"] - runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] - env: - PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" - SQLLOGICTEST: "0.25.0" - ARCH: ${{ matrix.runner == 'ubuntu-22.04' && 'x86_64' || 'aarch64' }} - PLATFORM: ${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} - - steps: - - name: Set up Environment - run: | - sudo apt-get remove -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' - sudo apt-get purge -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' - - sudo apt-get install -y postgresql-common - sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y - sudo apt-get install -y postgresql-client-17 - - uses: actions/checkout@v4 - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - const url = process.env.ACTIONS_CACHE_URL || ''; - const token = process.env.ACTIONS_RUNTIME_TOKEN || ''; - core.exportVariable( - 'CACHE_ENVS', - `-e CARGO_INCREMENTAL=0 -e SCCACHE_GHA_ENABLED=true -e RUSTC_WRAPPER=sccache -e ACTIONS_CACHE_URL=${url} -e ACTIONS_RUNTIME_TOKEN=${token}`, - ); - - name: Set up pgrx docker images and permissions - run: | - docker pull $PGRX_IMAGE - echo "Default user: $(id -u):$(id -g)" - sudo chmod -R 777 . - - - name: Build - run: | - docker run --rm -v .:/workspace $CACHE_ENVS \ - -e SEMVER=0.0.0 \ - -e VERSION=${{ matrix.version }} \ - -e ARCH=$ARCH \ - -e PLATFORM=$PLATFORM \ - $PGRX_IMAGE ./tools/package.sh - docker build -t vchord:pg${{ matrix.version }} --build-arg PG_VERSION=${{ matrix.version }} --build-arg SEMVER=0.0.0 -f ./docker/Dockerfile . - - - name: Setup SQL Logic Test - run: | - curl -fsSL -o sqllogictest.tar.gz https://github.com/risinglightdb/sqllogictest-rs/releases/download/v${SQLLOGICTEST}/sqllogictest-bin-v${SQLLOGICTEST}-$ARCH-unknown-linux-musl.tar.gz - tar -xzf sqllogictest.tar.gz - mv sqllogictest /usr/local/bin/ - - - name: SQL Test - env: - PGPASSWORD: postgres - run: | - docker run --rm --name test -d -e POSTGRES_PASSWORD=${PGPASSWORD} -p 5432:5432 vchord:pg${{ matrix.version }} - sleep 5 - psql -h localhost -U postgres -d postgres -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - sqllogictest './tests/**/*.slt' - docker stop test diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml deleted file mode 100644 index f026d699..00000000 --- a/.github/workflows/rust.yml +++ /dev/null @@ -1,83 +0,0 @@ -name: Rust - -on: - pull_request: - paths: - - '.github/workflows/rust.yml' - - 'crates/**' - - 'src/**' - - 'Cargo.lock' - - 'Cargo.toml' - - '*.control' - - 'rust-toolchain.toml' - push: - branches: - - main - paths: - - '.github/workflows/rust.yml' - - 'crates/**' - - 'src/**' - - 'Cargo.lock' - - 'Cargo.toml' - - '*.control' - - 'rust-toolchain.toml' - merge_group: - workflow_dispatch: - -concurrency: - group: ${{ github.ref }}-${{ github.workflow }} - cancel-in-progress: true - -jobs: - test: - strategy: - matrix: - include: - - runner: "ubuntu-22.04" - arch: "x86_64" - - runner: "ubuntu-22.04-arm" - arch: "aarch64" - runs-on: ${{ matrix.runner }} - env: - PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" - - steps: - - name: Set up Environment - run: | - sudo apt-get remove -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' - sudo apt-get purge -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' - - sudo apt-get install -y postgresql-common - sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y - sudo apt-get install -y postgresql-client-17 - - uses: actions/checkout@v4 - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - const url = process.env.ACTIONS_CACHE_URL || ''; - const token = process.env.ACTIONS_RUNTIME_TOKEN || ''; - core.exportVariable( - 'CACHE_ENVS', - `-e CARGO_INCREMENTAL=0 -e SCCACHE_GHA_ENABLED=true -e RUSTC_WRAPPER=sccache -e ACTIONS_CACHE_URL=${url} -e ACTIONS_RUNTIME_TOKEN=${token}`, - ); - - name: Set up docker images and permissions - run: | - docker pull $PGRX_IMAGE - echo "Default user: $(id -u):$(id -g)" - sudo chown -R 1000:1000 . - - - name: Clippy - run: | - for v in {14..17}; do - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo clippy --target ${{ matrix.arch }}-unknown-linux-gnu --features "pg$v" -- -D warnings - done - - name: Build - run: | - for v in {14..17}; do - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo build --lib --target ${{ matrix.arch }}-unknown-linux-gnu --features "pg$v" - done - - name: Test - run: | - # pg agnostic tests - docker run --rm -v .:/workspace $CACHE_ENVS $PGRX_IMAGE cargo test --no-fail-fast --target ${{ matrix.arch }}-unknown-linux-gnu --features pg17 diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml deleted file mode 100644 index ac5f385e..00000000 --- a/.github/workflows/style.yml +++ /dev/null @@ -1,40 +0,0 @@ -name: Style - -on: - pull_request: - branches: - - main - push: - branches: - - main - merge_group: - workflow_dispatch: - -concurrency: - group: ${{ github.ref }}-${{ github.workflow }} - cancel-in-progress: true - -env: - CARGO_TERM_COLOR: always - -jobs: - lint: - runs-on: ubuntu-latest - timeout-minutes: 5 - steps: - - uses: actions/checkout@v4 - - - name: Typos - uses: crate-ci/typos@master - - - name: TOML lint - run: | - curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-x86_64.gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo - taplo fmt --check - - - name: Cargo Lint - run: | - cargo fmt --check - - - name: Python lint - uses: astral-sh/ruff-action@v1 diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 3d7fbd5d..4de1ed88 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -532,6 +532,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] + #[ignore] fn reduce_sum_of_xy_v8_3a_sve_test() { use rand::Rng; const EPSILON: f32 = 2.0; @@ -909,6 +910,7 @@ mod reduce_sum_of_d2 { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] + #[ignore] fn reduce_sum_of_d2_v8_3a_sve_test() { use rand::Rng; const EPSILON: f32 = 6.0; diff --git a/crates/simd/src/fast_scan/mod.rs b/crates/simd/src/fast_scan/mod.rs index 735d3e87..44cf2fc4 100644 --- a/crates/simd/src/fast_scan/mod.rs +++ b/crates/simd/src/fast_scan/mod.rs @@ -529,15 +529,15 @@ mod fast_scan { let lut = lut[i]; let res_lo = zerocopy::transmute!(shuffle(lut, clo)); - a_0 = binary(|x, y| x + y, a_0, res_lo); - a_1 = binary(|x, y| x + y, a_1, res_lo.map(|x| x >> 8)); + a_0 = binary(u16::wrapping_add, a_0, res_lo); + a_1 = binary(u16::wrapping_add, a_1, res_lo.map(|x| x >> 8)); let res_hi = zerocopy::transmute!(shuffle(lut, chi)); - a_2 = binary(|x, y| x + y, a_2, res_hi); - a_3 = binary(|x, y| x + y, a_3, res_hi.map(|x| x >> 8)); + a_2 = binary(u16::wrapping_add, a_2, res_hi); + a_3 = binary(u16::wrapping_add, a_3, res_hi.map(|x| x >> 8)); } - a_0 = binary(|x, y| x - y, a_0, a_1.map(|x| x << 8)); - a_2 = binary(|x, y| x - y, a_2, a_3.map(|x| x << 8)); + a_0 = binary(u16::wrapping_sub, a_0, a_1.map(|x| x << 8)); + a_2 = binary(u16::wrapping_sub, a_2, a_3.map(|x| x << 8)); zerocopy::transmute!([a_0, a_1, a_2, a_3]) } diff --git a/.taplo.toml b/taplo.toml similarity index 100% rename from .taplo.toml rename to taplo.toml From e3893b7bbdb08eaec2a8407fa5c21f5dd80047b8 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 13 Feb 2025 11:00:22 +0800 Subject: [PATCH 109/324] feat: rerank by fetching vectors in heap table (#189) closes #136 Signed-off-by: usamoi --- .github/workflows/check.yml | 14 +- crates/algorithm/src/build.rs | 1 + crates/algorithm/src/insert.rs | 7 +- crates/algorithm/src/lib.rs | 8 + crates/algorithm/src/rerank.rs | 71 +++++ crates/algorithm/src/search.rs | 48 +-- crates/algorithm/src/tuples.rs | 10 +- crates/algorithm/src/types.rs | 5 + crates/k_means/src/lib.rs | 3 + src/index/am/am_build.rs | 10 +- src/index/am/am_scan.rs | 525 ++++++++++++++++++++++++++++---- src/index/gucs.rs | 4 +- src/index/types.rs | 2 +- tests/logic/rerank_in_table.slt | 67 ++++ 14 files changed, 680 insertions(+), 95 deletions(-) create mode 100644 crates/algorithm/src/rerank.rs create mode 100644 tests/logic/rerank_in_table.slt diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index f6721115..63595972 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -99,8 +99,10 @@ jobs: steps: - name: Set up Environment run: | - sudo apt-get remove -y '^postgres.*' '^libpq.*' - sudo apt-get purge -y '^postgres.*' '^libpq.*' + if [ "${{ matrix.runner }}" = "ubuntu-24.04" ]; then + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' + fi sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 @@ -118,7 +120,13 @@ jobs: sudo systemctl stop postgresql curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.12.9/cargo-pgrx-v0.12.9-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx - cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) + if [ "${{ matrix.runner }}" = "ubuntu-24.04" ]; then + cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) + fi + if [ "${{ matrix.runner }}" = "ubuntu-24.04-arm" ]; then + mkdir $HOME/.pgrx + echo "configs.pg${{ matrix.version }} = \"$(which pg_config)\"" > $HOME/.pgrx/config.toml + fi curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 3db4e271..d8c9b62d 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -93,6 +93,7 @@ pub fn build( dims, height_of_root: structures.len() as u32, is_residual, + rerank_in_heap: vchordrq_options.rerank_in_table, vectors_first: vectors.first(), root_mean: pointer_of_means.last().unwrap()[0], root_first: pointer_of_firsts.last().unwrap()[0], diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 4fb947fd..9b71bae4 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -18,6 +18,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); + let rerank_in_heap = meta_tuple.rerank_in_heap(); let height_of_root = meta_tuple.height_of_root(); assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); let root_mean = meta_tuple.root_mean(); @@ -31,7 +32,11 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto None }; - let mean = vectors::append::(index.clone(), vectors_first, vector.as_borrowed(), payload); + let mean = if !rerank_in_heap { + vectors::append::(index.clone(), vectors_first, vector.as_borrowed(), payload) + } else { + IndexPointer::default() + }; type State = (u32, Option<::Vector>); let mut state: State = { diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index b395f36d..f946c48b 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -12,6 +12,7 @@ mod linked_vec; mod maintain; mod pipe; mod prewarm; +mod rerank; mod search; mod select_heap; mod tape; @@ -27,6 +28,7 @@ pub use cache::cache; pub use insert::insert; pub use maintain::maintain; pub use prewarm::prewarm; +pub use rerank::{rerank_heap, rerank_index}; pub use search::search; use std::ops::{Deref, DerefMut}; @@ -72,3 +74,9 @@ pub trait RelationWrite: RelationRead { fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_>; fn search(&self, freespace: usize) -> Option>; } + +#[derive(Debug, Clone, Copy)] +pub enum RerankMethod { + Index, + Heap, +} diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs new file mode 100644 index 00000000..840fdd4c --- /dev/null +++ b/crates/algorithm/src/rerank.rs @@ -0,0 +1,71 @@ +use crate::operator::*; +use crate::tuples::*; +use crate::{RelationRead, vectors}; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::num::NonZeroU64; +use vector::VectorOwned; + +pub fn rerank_index( + index: impl RelationRead, + vector: O::Vector, + results: Vec<( + Reverse, + AlwaysEqual, + AlwaysEqual, + )>, +) -> impl Iterator { + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _)>::new(); + std::iter::from_fn(move || { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); + if let Some(dis_u) = vectors::access_0::( + index.clone(), + mean, + pay_u, + LAccess::new( + O::Vector::elements_and_metadata(vector.as_borrowed()), + O::DistanceAccessor::default(), + ), + ) { + cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); + }; + } + let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; + Some((dis_u, pay_u)) + }) +} + +pub fn rerank_heap( + vector: O::Vector, + results: Vec<( + Reverse, + AlwaysEqual, + AlwaysEqual, + )>, + fetch: F, +) -> impl Iterator +where + F: Fn(NonZeroU64) -> Option, +{ + let mut heap = BinaryHeap::from(results); + let mut cache = BinaryHeap::<(Reverse, _)>::new(); + std::iter::from_fn(move || { + let vector = O::Vector::elements_and_metadata(vector.as_borrowed()); + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(_), AlwaysEqual(pay_u)) = heap.pop().unwrap(); + if let Some(vec_u) = fetch(pay_u) { + let vec_u = O::Vector::elements_and_metadata(vec_u.as_borrowed()); + let mut accessor = O::DistanceAccessor::default(); + accessor.push(vector.0, vec_u.0); + let dis_u = accessor.finish(vector.1, vec_u.1); + cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); + } + } + let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; + Some((dis_u, pay_u)) + }) +} diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index c6bdc5cc..9972e861 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -3,7 +3,7 @@ use crate::operator::*; use crate::pipe::Pipe; use crate::tape::{access_0, access_1}; use crate::tuples::*; -use crate::{Page, RelationRead, vectors}; +use crate::{Page, RelationRead, RerankMethod, vectors}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -16,14 +16,28 @@ pub fn search( vector: O::Vector, probes: Vec, epsilon: f32, -) -> impl Iterator { +) -> ( + RerankMethod, + Vec<( + Reverse, + AlwaysEqual, + AlwaysEqual, + )>, +) { let meta_guard = index.read(0); let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); + let rerank_in_heap = meta_tuple.rerank_in_heap(); let height_of_root = meta_tuple.height_of_root(); assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); - assert_eq!(height_of_root as usize, 1 + probes.len(), "invalid probes"); + if height_of_root as usize != 1 + probes.len() { + panic!( + "need {} probes, but {} probes provided", + height_of_root - 1, + probes.len() + ); + } let root_mean = meta_tuple.root_mean(); let root_first = meta_tuple.root_first(); drop(meta_guard); @@ -145,24 +159,12 @@ pub fn search( }, ); } - let mut heap = BinaryHeap::from(results.into_vec()); - let mut cache = BinaryHeap::<(Reverse, _)>::new(); - std::iter::from_fn(move || { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); - if let Some(dis_u) = vectors::access_0::( - index.clone(), - mean, - pay_u, - LAccess::new( - O::Vector::elements_and_metadata(vector.as_borrowed()), - O::DistanceAccessor::default(), - ), - ) { - cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); - }; - } - let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; - Some((dis_u, pay_u)) - }) + ( + if rerank_in_heap { + RerankMethod::Heap + } else { + RerankMethod::Index + }, + results.into_vec(), + ) } diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 54c1707f..0f3de427 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -7,7 +7,7 @@ use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: u64 = u64::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 1; +const VERSION: u64 = 2; pub trait Tuple: 'static { type Reader<'a>: TupleReader<'a, Tuple = Self>; @@ -50,7 +50,8 @@ struct MetaTupleHeader { dims: u32, height_of_root: u32, is_residual: Bool, - _padding_0: [ZeroU8; 3], + rerank_in_heap: Bool, + _padding_0: [ZeroU8; 2], vectors_first: u32, // raw vector root_mean: IndexPointer, @@ -63,6 +64,7 @@ pub struct MetaTuple { pub dims: u32, pub height_of_root: u32, pub is_residual: bool, + pub rerank_in_heap: bool, pub vectors_first: u32, pub root_mean: IndexPointer, pub root_first: u32, @@ -79,6 +81,7 @@ impl Tuple for MetaTuple { dims: self.dims, height_of_root: self.height_of_root, is_residual: self.is_residual.into(), + rerank_in_heap: self.rerank_in_heap.into(), _padding_0: Default::default(), vectors_first: self.vectors_first, root_mean: self.root_mean, @@ -125,6 +128,9 @@ impl MetaTupleReader<'_> { pub fn is_residual(self) -> bool { self.header.is_residual.into() } + pub fn rerank_in_heap(self) -> bool { + self.header.rerank_in_heap.into() + } pub fn vectors_first(self) -> u32 { self.header.vectors_first } diff --git a/crates/algorithm/src/types.rs b/crates/algorithm/src/types.rs index c2545670..ea9784af 100644 --- a/crates/algorithm/src/types.rs +++ b/crates/algorithm/src/types.rs @@ -8,12 +8,17 @@ use vector::vect::{VectBorrowed, VectOwned}; pub struct VchordrqIndexOptions { #[serde(default = "VchordrqIndexOptions::default_residual_quantization")] pub residual_quantization: bool, + #[serde(default = "VchordrqIndexOptions::default_rerank_in_table")] + pub rerank_in_table: bool, } impl VchordrqIndexOptions { fn default_residual_quantization() -> bool { false } + fn default_rerank_in_table() -> bool { + false + } } #[derive(Debug, Clone)] diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index b4b2e5c6..53e46d72 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -121,6 +121,9 @@ fn quick_centers( ) -> Vec> { let n = samples.len(); assert!(c >= n); + if c == 1 && n == 0 { + return vec![vec![0.0; dims]]; + } let mut rng = rand::thread_rng(); let mut centroids = samples; for _ in n..c { diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index d0661879..a9a10ab0 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -132,13 +132,15 @@ pub unsafe extern "C" fn ambuild( } VchordrqBuildSourceOptions::Internal(internal_build) => { let mut tuples_total = 0_u64; - let samples = { + let samples = 'a: { let mut rand = rand::thread_rng(); - let max_number_of_samples = internal_build + let Some(max_number_of_samples) = internal_build .lists .last() - .unwrap() - .saturating_mul(internal_build.sampling_factor); + .map(|x| x.saturating_mul(internal_build.sampling_factor)) + else { + break 'a Vec::new(); + }; let mut samples = Vec::new(); let mut number_of_samples = 0_u32; match opfamily.vector_kind() { diff --git a/src/index/am/am_scan.rs b/src/index/am/am_scan.rs index 39ed021b..02b13b6f 100644 --- a/src/index/am/am_scan.rs +++ b/src/index/am/am_scan.rs @@ -3,9 +3,12 @@ use crate::index::gucs::{epsilon, max_scan_tuples, probes}; use crate::index::opclass::{Opfamily, Sphere, opfamily}; use crate::index::projection::RandomProject; use crate::index::storage::PostgresRelation; +use algorithm::RerankMethod; use algorithm::operator::{Dot, L2, Op, Vector}; use algorithm::types::*; use half::f16; +use pgrx::pg_sys::Datum; +use std::cell::LazyCell; use std::num::NonZeroU64; use vector::VectorOwned; use vector::vect::VectOwned; @@ -106,7 +109,41 @@ pub unsafe extern "C" fn amgettuple( } let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; let relation = unsafe { PostgresRelation::new((*scan).indexRelation) }; - if let Some((pointer, recheck)) = scanner_next(scanner, relation) { + if let Some((pointer, recheck)) = scanner_next( + scanner, + relation, + LazyCell::new(move || unsafe { + let index_info = pgrx::pg_sys::BuildIndexInfo((*scan).indexRelation); + let heap_relation = (*scan).heapRelation; + let estate = Scopeguard::new(pgrx::pg_sys::CreateExecutorState(), |estate| { + pgrx::pg_sys::FreeExecutorState(estate); + }); + let econtext = pgrx::pg_sys::MakePerTupleExprContext(*estate.get()); + move |opfamily: Opfamily, payload| { + let slot = Scopeguard::new( + pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), + |slot| pgrx::pg_sys::ExecDropSingleTupleTableSlot(slot), + ); + (*econtext).ecxt_scantuple = *slot.get(); + let table_am = (*heap_relation).rd_tableam; + let fetch_row_version = (*table_am).tuple_fetch_row_version.unwrap(); + let mut ctid = pointer_to_ctid(payload); + if !fetch_row_version(heap_relation, &mut ctid, (*scan).xs_snapshot, *slot.get()) { + return None; + } + let mut values = [Datum::from(0); 32]; + let mut is_null = [true; 32]; + pgrx::pg_sys::FormIndexDatum( + index_info, + *slot.get(), + *estate.get(), + values.as_mut_ptr(), + is_null.as_mut_ptr(), + ); + opfamily.input_vector(values[0], is_null[0]) + } + }), + ) { let ctid = pointer_to_ctid(pointer); unsafe { (*scan).xs_heaptid = ctid; @@ -167,7 +204,15 @@ fn scanner_build( Some((vector?, threshold, recheck)) } -fn scanner_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(NonZeroU64, bool)> { +fn scanner_next( + scanner: &mut Scanner, + relation: PostgresRelation, + fetch: LazyCell, +) -> Option<(NonZeroU64, bool)> +where + F: Fn(Opfamily, NonZeroU64) -> Option + 'static, + I: FnOnce() -> F + 'static, +{ if let Scanning::Initial { vector, threshold, @@ -187,92 +232,420 @@ fn scanner_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(No let vector = RandomProject::project( VectOwned::::from_owned(vector.clone()).as_borrowed(), ); - let vbase = algorithm::search::, L2>>( - relation, vector, probes, epsilon, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - match (max_scan_tuples, threshold) { - (None, None) => { + let (method, results) = algorithm::search::, L2>>( + relation.clone(), + vector.clone(), + probes, + epsilon, + ); + match (method, max_scan_tuples, threshold) { + (RerankMethod::Index, None, None) => { + let vbase = algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.fuse()) as Box> + } + (RerankMethod::Index, None, Some(threshold)) => { + let vbase = algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take_while(move |(x, _)| *x < threshold)) + } + (RerankMethod::Index, Some(max_scan_tuples), None) => { + let vbase = algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take(max_scan_tuples as _)) + } + (RerankMethod::Index, Some(max_scan_tuples), Some(threshold)) => { + let vbase = algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ) + } + (RerankMethod::Heap, None, None) => { + let vbase = algorithm::rerank_heap::, L2>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); Box::new(vbase.fuse()) as Box> } - (None, Some(threshold)) => { + (RerankMethod::Heap, None, Some(threshold)) => { + let vbase = algorithm::rerank_heap::, L2>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); Box::new(vbase.take_while(move |(x, _)| *x < threshold)) } - (Some(max_scan_tuples), None) => Box::new(vbase.take(max_scan_tuples as _)), - (Some(max_scan_tuples), Some(threshold)) => Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ), + (RerankMethod::Heap, Some(max_scan_tuples), None) => { + let vbase = algorithm::rerank_heap::, L2>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take(max_scan_tuples as _)) + } + (RerankMethod::Heap, Some(max_scan_tuples), Some(threshold)) => { + let vbase = algorithm::rerank_heap::, L2>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ) + } } } (VectorKind::Vecf32, DistanceKind::Dot) => { let vector = RandomProject::project( VectOwned::::from_owned(vector.clone()).as_borrowed(), ); - let vbase = algorithm::search::, Dot>>( - relation, vector, probes, epsilon, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - match (max_scan_tuples, threshold) { - (None, None) => { - Box::new(vbase) as Box> - } - (None, Some(threshold)) => { + let (method, results) = algorithm::search::, Dot>>( + relation.clone(), + vector.clone(), + probes, + epsilon, + ); + match (method, max_scan_tuples, threshold) { + (RerankMethod::Index, None, None) => { + let vbase = algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.fuse()) as Box> + } + (RerankMethod::Index, None, Some(threshold)) => { + let vbase = algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take_while(move |(x, _)| *x < threshold)) + } + (RerankMethod::Index, Some(max_scan_tuples), None) => { + let vbase = algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take(max_scan_tuples as _)) + } + (RerankMethod::Index, Some(max_scan_tuples), Some(threshold)) => { + let vbase = algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ) + } + (RerankMethod::Heap, None, None) => { + let vbase = algorithm::rerank_heap::, Dot>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.fuse()) as Box> + } + (RerankMethod::Heap, None, Some(threshold)) => { + let vbase = algorithm::rerank_heap::, Dot>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); Box::new(vbase.take_while(move |(x, _)| *x < threshold)) } - (Some(max_scan_tuples), None) => Box::new(vbase.take(max_scan_tuples as _)), - (Some(max_scan_tuples), Some(threshold)) => Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ), + (RerankMethod::Heap, Some(max_scan_tuples), None) => { + let vbase = algorithm::rerank_heap::, Dot>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take(max_scan_tuples as _)) + } + (RerankMethod::Heap, Some(max_scan_tuples), Some(threshold)) => { + let vbase = algorithm::rerank_heap::, Dot>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ) + } } } (VectorKind::Vecf16, DistanceKind::L2) => { let vector = RandomProject::project( VectOwned::::from_owned(vector.clone()).as_borrowed(), ); - let vbase = algorithm::search::, L2>>( - relation, vector, probes, epsilon, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - match (max_scan_tuples, threshold) { - (None, None) => { - Box::new(vbase) as Box> - } - (None, Some(threshold)) => { + let (method, results) = algorithm::search::, L2>>( + relation.clone(), + vector.clone(), + probes, + epsilon, + ); + match (method, max_scan_tuples, threshold) { + (RerankMethod::Index, None, None) => { + let vbase = algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.fuse()) as Box> + } + (RerankMethod::Index, None, Some(threshold)) => { + let vbase = algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); Box::new(vbase.take_while(move |(x, _)| *x < threshold)) } - (Some(max_scan_tuples), None) => Box::new(vbase.take(max_scan_tuples as _)), - (Some(max_scan_tuples), Some(threshold)) => Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ), + (RerankMethod::Index, Some(max_scan_tuples), None) => { + let vbase = algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take(max_scan_tuples as _)) + } + (RerankMethod::Index, Some(max_scan_tuples), Some(threshold)) => { + let vbase = algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ) + } + (RerankMethod::Heap, None, None) => { + let vbase = algorithm::rerank_heap::, L2>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.fuse()) as Box> + } + (RerankMethod::Heap, None, Some(threshold)) => { + let vbase = algorithm::rerank_heap::, L2>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take_while(move |(x, _)| *x < threshold)) + } + (RerankMethod::Heap, Some(max_scan_tuples), None) => { + let vbase = algorithm::rerank_heap::, L2>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take(max_scan_tuples as _)) + } + (RerankMethod::Heap, Some(max_scan_tuples), Some(threshold)) => { + let vbase = algorithm::rerank_heap::, L2>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ) + } } } (VectorKind::Vecf16, DistanceKind::Dot) => { let vector = RandomProject::project( VectOwned::::from_owned(vector.clone()).as_borrowed(), ); - let vbase = algorithm::search::, Dot>>( - relation, vector, probes, epsilon, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - match (max_scan_tuples, threshold) { - (None, None) => { - Box::new(vbase) as Box> - } - (None, Some(threshold)) => { + let (method, results) = algorithm::search::, Dot>>( + relation.clone(), + vector.clone(), + probes, + epsilon, + ); + match (method, max_scan_tuples, threshold) { + (RerankMethod::Index, None, None) => { + let vbase = algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.fuse()) as Box> + } + (RerankMethod::Index, None, Some(threshold)) => { + let vbase = algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); Box::new(vbase.take_while(move |(x, _)| *x < threshold)) } - (Some(max_scan_tuples), None) => Box::new(vbase.take(max_scan_tuples as _)), - (Some(max_scan_tuples), Some(threshold)) => Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ), + (RerankMethod::Index, Some(max_scan_tuples), None) => { + let vbase = algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take(max_scan_tuples as _)) + } + (RerankMethod::Index, Some(max_scan_tuples), Some(threshold)) => { + let vbase = algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ) + } + (RerankMethod::Heap, None, None) => { + let vbase = algorithm::rerank_heap::, Dot>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.fuse()) as Box> + } + (RerankMethod::Heap, None, Some(threshold)) => { + let vbase = algorithm::rerank_heap::, Dot>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take_while(move |(x, _)| *x < threshold)) + } + (RerankMethod::Heap, Some(max_scan_tuples), None) => { + let vbase = algorithm::rerank_heap::, Dot>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new(vbase.take(max_scan_tuples as _)) + } + (RerankMethod::Heap, Some(max_scan_tuples), Some(threshold)) => { + let vbase = algorithm::rerank_heap::, Dot>, _>( + vector, + results, + move |payload| { + Some(RandomProject::project( + VectOwned::::from_owned(fetch(opfamily, payload)?) + .as_borrowed(), + )) + }, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)); + Box::new( + vbase + .take_while(move |(x, _)| *x < threshold) + .take(max_scan_tuples as _), + ) + } } } }, @@ -285,3 +658,35 @@ fn scanner_next(scanner: &mut Scanner, relation: PostgresRelation) -> Option<(No Scanning::Empty {} => None, } } + +struct Scopeguard +where + T: Copy, + F: FnMut(T), +{ + t: T, + f: F, +} + +impl Scopeguard +where + T: Copy, + F: FnMut(T), +{ + fn new(t: T, f: F) -> Self { + Scopeguard { t, f } + } + fn get(&self) -> &T { + &self.t + } +} + +impl Drop for Scopeguard +where + T: Copy, + F: FnMut(T), +{ + fn drop(&mut self) { + (self.f)(self.t); + } +} diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 5e4216b6..26d061ab 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -72,7 +72,9 @@ pub fn probes() -> Vec { c => pgrx::error!("unknown character in probes: ASCII = {c}"), } } - result.push(current.take().expect("empty probes")); + if let Some(current) = current { + result.push(current); + } result } } diff --git a/src/index/types.rs b/src/index/types.rs index 6d6bb3ef..df238462 100644 --- a/src/index/types.rs +++ b/src/index/types.rs @@ -6,7 +6,7 @@ use validator::{Validate, ValidationError, ValidationErrors}; #[serde(deny_unknown_fields)] pub struct VchordrqInternalBuildOptions { #[serde(default = "VchordrqInternalBuildOptions::default_lists")] - #[validate(length(min = 1, max = 8), custom(function = VchordrqInternalBuildOptions::validate_lists))] + #[validate(length(min = 0, max = 8), custom(function = VchordrqInternalBuildOptions::validate_lists))] pub lists: Vec, #[serde(default = "VchordrqInternalBuildOptions::default_spherical_centroids")] pub spherical_centroids: bool, diff --git a/tests/logic/rerank_in_table.slt b/tests/logic/rerank_in_table.slt new file mode 100644 index 00000000..381335c1 --- /dev/null +++ b/tests/logic/rerank_in_table.slt @@ -0,0 +1,67 @@ +statement ok +CREATE TABLE t_singlecolumn (id integer, val vector(3)); + +statement ok +INSERT INTO t_singlecolumn (id, val) SELECT id, ARRAY[id, id, id]::real[] FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_singlecolumn USING vchordrq (val vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = true +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +query I +SELECT id FROM t_singlecolumn ORDER BY val <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +statement ok +DROP TABLE t_singlecolumn; + +statement ok +CREATE TABLE t_multicolumn (id integer); + +statement ok +INSERT INTO t_multicolumn (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_multicolumn USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = true +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +query I +SELECT id FROM t_multicolumn ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +statement ok +DROP TABLE t_multicolumn; From 130d9c38173e29d0ac98cf94eea68f529d8838e5 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 13 Feb 2025 14:23:39 +0800 Subject: [PATCH 110/324] ci: enable CI for pg13 (#185) Signed-off-by: usamoi --- .github/workflows/check.yml | 5 ----- tests/logic/pushdown_plan.slt | 3 +++ tests/logic/rerank_in_table.slt | 20 ++++++++++---------- 3 files changed, 13 insertions(+), 15 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 63595972..77fe0360 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -89,11 +89,6 @@ jobs: matrix: version: ["13", "14", "15", "16", "17"] runner: ["ubuntu-24.04", "ubuntu-24.04-arm"] - exclude: - - version: "13" - runner: "ubuntu-24.04" - - version: "13" - runner: "ubuntu-24.04-arm" runs-on: ${{ matrix.runner }} steps: diff --git a/tests/logic/pushdown_plan.slt b/tests/logic/pushdown_plan.slt index 78bfe440..f6aecc48 100644 --- a/tests/logic/pushdown_plan.slt +++ b/tests/logic/pushdown_plan.slt @@ -12,6 +12,9 @@ FROM generate_series(1, 10000); statement ok CREATE INDEX ind0 ON t USING vchordrq (val0 vector_l2_ops); +statement ok +SET enable_seqscan TO off; + # statement ok # CREATE INDEX ind1 ON t USING vchordrq (val1 halfvec_dot_ops); diff --git a/tests/logic/rerank_in_table.slt b/tests/logic/rerank_in_table.slt index 381335c1..f6a138a1 100644 --- a/tests/logic/rerank_in_table.slt +++ b/tests/logic/rerank_in_table.slt @@ -1,11 +1,11 @@ statement ok -CREATE TABLE t_singlecolumn (id integer, val vector(3)); +CREATE TABLE t_column (id integer, val vector(3)); statement ok -INSERT INTO t_singlecolumn (id, val) SELECT id, ARRAY[id, id, id]::real[] FROM generate_series(1, 10000) s(id); +INSERT INTO t_column (id, val) SELECT id, ARRAY[id, id, id]::real[] FROM generate_series(1, 10000) s(id); statement ok -CREATE INDEX ON t_singlecolumn USING vchordrq (val vector_l2_ops) +CREATE INDEX ON t_column USING vchordrq (val vector_l2_ops) WITH (options = $$ residual_quantization = false rerank_in_table = true @@ -17,7 +17,7 @@ statement ok SET vchordrq.probes = ''; query I -SELECT id FROM t_singlecolumn ORDER BY val <-> '[1.9, 1.9, 1.9]' limit 9; +SELECT id FROM t_column ORDER BY val <-> '[1.9, 1.9, 1.9]' limit 9; ---- 2 1 @@ -30,16 +30,16 @@ SELECT id FROM t_singlecolumn ORDER BY val <-> '[1.9, 1.9, 1.9]' limit 9; 9 statement ok -DROP TABLE t_singlecolumn; +DROP TABLE t_column; statement ok -CREATE TABLE t_multicolumn (id integer); +CREATE TABLE t_expr (id integer); statement ok -INSERT INTO t_multicolumn (id) SELECT id FROM generate_series(1, 10000) s(id); +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); statement ok -CREATE INDEX ON t_multicolumn USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) WITH (options = $$ residual_quantization = false rerank_in_table = true @@ -51,7 +51,7 @@ statement ok SET vchordrq.probes = ''; query I -SELECT id FROM t_multicolumn ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; ---- 2 1 @@ -64,4 +64,4 @@ SELECT id FROM t_multicolumn ORDER BY ARRAY[id::real, id::real, id::real]::vecto 9 statement ok -DROP TABLE t_multicolumn; +DROP TABLE t_expr; From 1f4a648278d59bf8c92ba38cdf667d03fd2a559d Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 14 Feb 2025 14:19:58 +0800 Subject: [PATCH 111/324] chore: update dependencies (#190) Signed-off-by: usamoi --- .github/workflows/check.yml | 22 +-- Cargo.lock | 143 +++++++++-------- Cargo.toml | 12 +- crates/algorithm/Cargo.toml | 3 +- crates/algorithm/src/lib.rs | 1 - crates/algorithm/src/operator.rs | 26 ++-- crates/algorithm/src/tuples.rs | 16 +- crates/k_means/Cargo.toml | 6 +- crates/k_means/src/lib.rs | 173 +++++++-------------- crates/random_orthogonal_matrix/Cargo.toml | 4 +- crates/random_orthogonal_matrix/src/lib.rs | 21 +-- crates/simd/Cargo.toml | 5 +- crates/simd/src/f16.rs | 72 ++++----- crates/simd/src/f32.rs | 152 +++++++++--------- crates/simd/src/lib.rs | 1 - crates/simd/src/u8.rs | 24 +-- crates/simd_macros/Cargo.toml | 6 +- crates/vector/Cargo.toml | 3 +- rust-toolchain.toml | 2 +- src/index/am/am_build.rs | 41 ++--- 20 files changed, 330 insertions(+), 403 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 77fe0360..1445349e 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -7,6 +7,7 @@ on: - ".github/workflows/check.yml" - "crates/**" - "src/**" + - "tests/**" - "build.rs" - "Cargo.lock" - "Cargo.toml" @@ -17,9 +18,10 @@ on: push: paths: - ".cargo" - - ".github/workflows/lint.yml" + - ".github/workflows/check.yml" - "crates/**" - "src/**" + - "tests/**" - "build.rs" - "Cargo.lock" - "Cargo.toml" @@ -68,7 +70,7 @@ jobs: lint: strategy: matrix: - runner: ["ubuntu-24.04", "ubuntu-24.04-arm"] + runner: ["ubuntu-24.04", "ubicloud-standard-8-arm-ubuntu-2404"] runs-on: ${{ matrix.runner }} steps: @@ -88,16 +90,14 @@ jobs: strategy: matrix: version: ["13", "14", "15", "16", "17"] - runner: ["ubuntu-24.04", "ubuntu-24.04-arm"] + runner: ["ubuntu-24.04", "ubicloud-standard-8-arm-ubuntu-2404"] runs-on: ${{ matrix.runner }} steps: - name: Set up Environment run: | - if [ "${{ matrix.runner }}" = "ubuntu-24.04" ]; then - sudo apt-get remove -y '^postgres.*' '^libpq.*' - sudo apt-get purge -y '^postgres.*' '^libpq.*' - fi + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 @@ -115,13 +115,7 @@ jobs: sudo systemctl stop postgresql curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.12.9/cargo-pgrx-v0.12.9-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx - if [ "${{ matrix.runner }}" = "ubuntu-24.04" ]; then - cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) - fi - if [ "${{ matrix.runner }}" = "ubuntu-24.04-arm" ]; then - mkdir $HOME/.pgrx - echo "configs.pg${{ matrix.version }} = \"$(which pg_config)\"" > $HOME/.pgrx/config.toml - fi + cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest diff --git a/Cargo.lock b/Cargo.lock index 7e992b88..a5a23736 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -26,12 +26,11 @@ dependencies = [ "random_orthogonal_matrix", "serde", "simd", - "toml", "turboselect", "validator", "vector", - "zerocopy 0.8.14", - "zerocopy-derive 0.8.14", + "zerocopy 0.8.17", + "zerocopy-derive 0.8.17", ] [[package]] @@ -100,9 +99,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.7.0" +version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1be3f42a67d6d345ecd59f675f3f012d6974981560836e938c22b424b85ce1be" +checksum = "8f68f53c83ab957f72c32642f3868eec03eb974d1fb82e453128456482613d36" [[package]] name = "bitvec" @@ -140,9 +139,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.9" +version = "1.2.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8293772165d9345bdaaa39b45b2109591e63fe5e6fbc23c6ff930a048aa310b" +checksum = "c7777341816418c02e033934a09f20dc0ccaf65a5201ef8a450ae0105a573fda" dependencies = [ "shlex", ] @@ -219,9 +218,9 @@ checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] name = "crunchy" -version = "0.2.2" +version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" +checksum = "43da5946c66ffcc7745f48db692ffbb10a83bfe0afd96235c5c2a4fb23994929" [[package]] name = "darling" @@ -344,13 +343,14 @@ checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" [[package]] name = "getrandom" -version = "0.2.15" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8" dependencies = [ "cfg-if", "libc", "wasi", + "windows-targets", ] [[package]] @@ -373,8 +373,8 @@ dependencies = [ "cfg-if", "crunchy", "serde", - "zerocopy 0.8.14", - "zerocopy-derive 0.8.14", + "zerocopy 0.8.17", + "zerocopy-derive 0.8.17", ] [[package]] @@ -420,7 +420,7 @@ version = "0.5.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589533453244b0995c858700322199b2becb13b627df2851f64a2775d024abcf" dependencies = [ - "windows-sys 0.59.0", + "windows-sys", ] [[package]] @@ -576,9 +576,9 @@ checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" [[package]] name = "indexmap" -version = "2.7.0" +version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62f822373a4fe84d4bb149bf54e584a7f4abec90e072ed49cda0edea5b95471f" +checksum = "8c9c992b02b5b4c94ea26e32fe5bccb7aa7d9f390ab5c1221ff895bc7ea8b652" dependencies = [ "equivalent", "hashbrown", @@ -586,13 +586,13 @@ dependencies = [ [[package]] name = "is-terminal" -version = "0.4.13" +version = "0.4.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "261f68e344040fbd0edea105bef17c66edf46f984ddb1115b775ce31be948f4b" +checksum = "e19b23d53f35ce9f56aebc7d1bb4e6ac1e9c0db7ac85c8d1760c04379edced37" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys", ] [[package]] @@ -621,7 +621,6 @@ name = "k_means" version = "0.0.0" dependencies = [ "half 2.4.1", - "log", "rabitq", "rand", "rayon", @@ -656,12 +655,6 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4ee93343901ab17bd981295f2cf0026d4ad018c7c31ba84549a4ddbb47a45104" -[[package]] -name = "log" -version = "0.4.25" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f" - [[package]] name = "matrixmultiply" version = "0.3.9" @@ -772,9 +765,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.20.2" +version = "1.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" +checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e" [[package]] name = "owo-colors" @@ -1007,20 +1000,20 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "rand" -version = "0.8.5" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" dependencies = [ - "libc", "rand_chacha", "rand_core", + "zerocopy 0.8.17", ] [[package]] name = "rand_chacha" -version = "0.3.1" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", "rand_core", @@ -1028,18 +1021,19 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.6.4" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +checksum = "b08f3c9802962f7e1b25113931d94f43ed9725bebc59db9d0c3e9a23b67e15ff" dependencies = [ "getrandom", + "zerocopy 0.8.17", ] [[package]] name = "rand_distr" -version = "0.4.3" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" +checksum = "ddc3b5afe4c995c44540865b8ca5c52e6a59fa362da96c5d30886930ddc8da1c" dependencies = [ "num-traits", "rand", @@ -1127,9 +1121,9 @@ dependencies = [ [[package]] name = "ryu" -version = "1.0.18" +version = "1.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" +checksum = "6ea1a2d0a644769cc99faa24c3ad26b379b786fe7c36fd3c546254801650e6dd" [[package]] name = "safe_arch" @@ -1205,9 +1199,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.135" +version = "1.0.138" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b0d7ba2887406110130a978386c4e1befb98c674b4fba677954e4db976630d9" +checksum = "d434192e7da787e94a6ea7e9670b26a036d0ca41e0b7efb2676dd32bae872949" dependencies = [ "itoa", "memchr", @@ -1251,7 +1245,7 @@ dependencies = [ "half 2.4.1", "rand", "simd_macros", - "zerocopy 0.8.14", + "zerocopy 0.8.17", ] [[package]] @@ -1308,9 +1302,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.96" +version = "2.0.98" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5d0adab1ae378d7f53bdebc67a39f1f151407ef230f0ce2883572f5d8985c80" +checksum = "36147f1a48ae0ec2b5b3bc5b537d267457555a10dc06f3dbc8cb11ba3006d3b1" dependencies = [ "proc-macro2", "quote", @@ -1386,9 +1380,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.8.19" +version = "0.8.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1ed1f98e3fdc28d6d910e6737ae6ab1a93bf1985935a1193e68f93eeb68d24e" +checksum = "cd87a5cdd6ffab733b2f74bc4fd7ee5fff6634124999ac278c35fc78c6120148" dependencies = [ "serde", "serde_spanned", @@ -1407,9 +1401,9 @@ dependencies = [ [[package]] name = "toml_edit" -version = "0.22.22" +version = "0.22.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ae48d6208a266e853d946088ed816055e556cc6028c5e8e2b84d9fa5dd7c7f5" +checksum = "17b4795ff5edd201c7cd6dca065ae59972ce77d1b80fa0a84d94950ece7d1474" dependencies = [ "indexmap", "serde", @@ -1443,9 +1437,9 @@ checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] name = "unicode-ident" -version = "1.0.14" +version = "1.0.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" +checksum = "a210d160f08b701c8721ba1c726c11662f877ea6b7094007e1ca9a1041945034" [[package]] name = "unicode-segmentation" @@ -1484,9 +1478,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.11.1" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b913a3b5fe84142e269d63cc62b64319ccaf89b748fc31fe025177f767a756c4" +checksum = "ced87ca4be083373936a67f8de945faa23b6b42384bd5b64434850802c6dccd0" dependencies = [ "getrandom", ] @@ -1539,8 +1533,8 @@ dependencies = [ "toml", "validator", "vector", - "zerocopy 0.8.14", - "zerocopy-derive 0.8.14", + "zerocopy 0.8.17", + "zerocopy-derive 0.8.17", ] [[package]] @@ -1564,9 +1558,12 @@ dependencies = [ [[package]] name = "wasi" -version = "0.11.0+wasi-snapshot-preview1" +version = "0.13.3+wasi-0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +checksum = "26816d2e1a4a36a2940b96c5296ce403917633dff8f3440e9b236ed6f6bacad2" +dependencies = [ + "wit-bindgen-rt", +] [[package]] name = "wide" @@ -1600,7 +1597,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys", ] [[package]] @@ -1609,15 +1606,6 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" -[[package]] -name = "windows-sys" -version = "0.52.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" -dependencies = [ - "windows-targets", -] - [[package]] name = "windows-sys" version = "0.59.0" @@ -1693,13 +1681,22 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "winnow" -version = "0.6.24" +version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8d71a593cc5c42ad7876e2c1fda56f314f3754c084128833e64f1345ff8a03a" +checksum = "59690dea168f2198d1a3b0cac23b8063efcd11012f10ae4698f284808c8ef603" dependencies = [ "memchr", ] +[[package]] +name = "wit-bindgen-rt" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c" +dependencies = [ + "bitflags", +] + [[package]] name = "write16" version = "1.0.0" @@ -1766,11 +1763,11 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.14" +version = "0.8.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a367f292d93d4eab890745e75a778da40909cab4d6ff8173693812f79c4a2468" +checksum = "aa91407dacce3a68c56de03abe2760159582b846c6a4acd2f456618087f12713" dependencies = [ - "zerocopy-derive 0.8.14", + "zerocopy-derive 0.8.17", ] [[package]] @@ -1786,9 +1783,9 @@ dependencies = [ [[package]] name = "zerocopy-derive" -version = "0.8.14" +version = "0.8.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3931cb58c62c13adec22e38686b559c86a30565e16ad6e8510a337cedc611e1" +checksum = "06718a168365cad3d5ff0bb133aad346959a2074bd4a85c121255a11304a8626" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index 1b4f79dc..395f98a5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,12 +28,12 @@ simd = { path = "./crates/simd" } vector = { path = "./crates/vector" } half.workspace = true -paste = "1" +paste.workspace = true pgrx = { version = "=0.12.9", default-features = false, features = ["cshim"] } pgrx-catalog = "0.1.0" rand.workspace = true serde.workspace = true -toml = "0.8.19" +toml = "0.8.20" validator.workspace = true zerocopy.workspace = true zerocopy-derive.workspace = true @@ -54,12 +54,12 @@ edition = "2021" [workspace.dependencies] half = { version = "2.4.1", features = ["serde", "zerocopy"] } -log = "0.4.25" -rand = "0.8.5" +paste = "1" +rand = "0.9.0" serde = "1" validator = { version = "0.20.0", features = ["derive"] } -zerocopy = "0.8.14" -zerocopy-derive = "0.8.14" +zerocopy = "0.8.17" +zerocopy-derive = "0.8.17" [workspace.lints] clippy.identity_op = "allow" diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml index 006fe9bd..83adc488 100644 --- a/crates/algorithm/Cargo.toml +++ b/crates/algorithm/Cargo.toml @@ -14,10 +14,9 @@ vector = { path = "../vector" } half.workspace = true heapify = "0.2.0" -paste = "1" +paste.workspace = true rand.workspace = true serde.workspace = true -toml = "0.8.19" turboselect = { git = "https://github.com/tensorchord/turboselect.git", rev = "d8753c4ffe5b47f28670fea21d56cf3658d51b9b" } validator.workspace = true zerocopy.workspace = true diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index f946c48b..372774f3 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -1,7 +1,6 @@ #![allow(clippy::collapsible_else_if)] #![allow(clippy::type_complexity)] #![allow(clippy::len_without_is_empty)] -#![feature(vec_pop_if)] mod build; mod bulkdelete; diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index 1d70a5c8..e8a57bd4 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -374,11 +374,14 @@ impl Vector for VectOwned { fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f32]>) { let vector = vector.slice(); - ((), match vector.len() { - 0..=960 => vec![vector], - 961..=1280 => vec![&vector[..640], &vector[640..]], - 1281.. => vector.chunks(1920).collect(), - }) + ( + (), + match vector.len() { + 0..=960 => vec![vector], + 961..=1280 => vec![&vector[..640], &vector[640..]], + 1281.. => vector.chunks(1920).collect(), + }, + ) } fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { @@ -417,11 +420,14 @@ impl Vector for VectOwned { fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f16]>) { let vector = vector.slice(); - ((), match vector.len() { - 0..=1920 => vec![vector], - 1921..=2560 => vec![&vector[..1280], &vector[1280..]], - 2561.. => vector.chunks(3840).collect(), - }) + ( + (), + match vector.len() { + 0..=1920 => vec![vector], + 1921..=2560 => vec![&vector[..1280], &vector[1280..]], + 2561.. => vector.chunks(3840).collect(), + }, + ) } fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 0f3de427..f60d5527 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -285,7 +285,7 @@ impl Tuple for VectorTuple { elements, } => { buffer.extend((0 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); while buffer.len() % ALIGN != 0 { buffer.push(0); } @@ -315,7 +315,7 @@ impl Tuple for VectorTuple { elements, } => { buffer.extend((1 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); while buffer.len() % ALIGN != 0 { buffer.push(0); } @@ -493,7 +493,7 @@ impl Tuple for H1Tuple { elements, } => { buffer.extend((0 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); while buffer.len() % ALIGN != 0 { buffer.push(0); } @@ -518,7 +518,7 @@ impl Tuple for H1Tuple { } Self::_1 { elements } => { buffer.extend((1 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); while buffer.len() % ALIGN != 0 { buffer.push(0); } @@ -777,7 +777,7 @@ impl Tuple for H0Tuple { elements, } => { buffer.extend((0 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); while buffer.len() % ALIGN != 0 { buffer.push(0); } @@ -808,7 +808,7 @@ impl Tuple for H0Tuple { elements, } => { buffer.extend((1 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); while buffer.len() % ALIGN != 0 { buffer.push(0); } @@ -831,7 +831,7 @@ impl Tuple for H0Tuple { } Self::_2 { elements } => { buffer.extend((2 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); while buffer.len() % ALIGN != 0 { buffer.push(0); } @@ -1200,7 +1200,7 @@ fn aliasing_test() { let serialized = { let elements = (0u32..1111).collect::>(); let mut buffer = Vec::::new(); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); while buffer.len() % ALIGN != 0 { buffer.push(0); } diff --git a/crates/k_means/Cargo.toml b/crates/k_means/Cargo.toml index 93b82ed0..6ba2ccf7 100644 --- a/crates/k_means/Cargo.toml +++ b/crates/k_means/Cargo.toml @@ -4,12 +4,12 @@ version.workspace = true edition.workspace = true [dependencies] -half.workspace = true -log.workspace = true rabitq = { path = "../rabitq" } +simd = { path = "../simd" } + +half.workspace = true rand.workspace = true rayon = "1.10.0" -simd = { path = "../simd" } [lints] workspace = true diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index 53e46d72..ecbea29b 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -3,73 +3,13 @@ use half::f16; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; use simd::Floating; use simd::fast_scan::{any_pack, padding_pack}; -use std::any::Any; -use std::panic::AssertUnwindSafe; -use std::sync::Arc; -pub use rayon::iter::ParallelIterator; - -pub trait Parallelism: Send + Sync { - fn check(&self); - - fn rayon_into_par_iter(&self, x: I) -> I::Iter; -} - -struct ParallelismCheckPanic(Box); - -pub struct RayonParallelism { - stop: Arc, -} - -impl RayonParallelism { - pub fn scoped( - num_threads: usize, - stop: Arc, - f: impl FnOnce(&Self) -> R, - ) -> Result { - match std::panic::catch_unwind(AssertUnwindSafe(|| { - rayon::ThreadPoolBuilder::new() - .num_threads(num_threads) - .panic_handler(|e| { - if e.downcast_ref::().is_some() { - return; - } - log::error!("Asynchronous task panickied."); - }) - .build_scoped( - |thread| thread.run(), - |_| { - let pool = Self { stop: stop.clone() }; - f(&pool) - }, - ) - })) { - Ok(x) => x, - Err(e) => match e.downcast::() { - Ok(payload) => std::panic::resume_unwind((*payload).0), - Err(e) => std::panic::resume_unwind(e), - }, - } - } -} - -impl Parallelism for RayonParallelism { - fn check(&self) { - match std::panic::catch_unwind(AssertUnwindSafe(|| (self.stop)())) { - Ok(()) => (), - Err(payload) => std::panic::panic_any(ParallelismCheckPanic(payload)), - } - } - - fn rayon_into_par_iter(&self, x: I) -> I::Iter { - x.into_par_iter() - } -} - -pub fn k_means( - parallelism: &P, +pub fn k_means( + num_threads: usize, + check: impl Fn(), c: usize, dims: usize, samples: &[Vec], @@ -82,22 +22,30 @@ pub fn k_means( if n <= c { quick_centers(c, dims, samples.to_vec(), is_spherical) } else { - let compute = |parallelism: &P, centroids: &[Vec]| { - if n >= 1000 && c >= 1000 { - rabitq_index(parallelism, dims, n, c, samples, centroids) - } else { - flat_index(parallelism, dims, n, c, samples, centroids) - } - }; - let mut lloyd_k_means = - LloydKMeans::new(parallelism, c, dims, samples, is_spherical, compute); - for _ in 0..iterations { - parallelism.check(); - if lloyd_k_means.iterate() { - break; - } - } - lloyd_k_means.finish() + rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .build_scoped( + |thread| thread.run(), + move |_| { + let compute = |centroids: &[Vec]| { + if n >= 1000 && c >= 1000 { + rabitq_index(dims, n, c, samples, centroids) + } else { + flat_index(dims, n, c, samples, centroids) + } + }; + let mut lloyd_k_means = + LloydKMeans::new(c, dims, samples, is_spherical, compute); + for _ in 0..iterations { + check(); + if lloyd_k_means.iterate() { + break; + } + } + lloyd_k_means.finish() + }, + ) + .expect("failed to build thread pool") } } @@ -124,10 +72,12 @@ fn quick_centers( if c == 1 && n == 0 { return vec![vec![0.0; dims]]; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); let mut centroids = samples; for _ in n..c { - let r = (0..dims).map(|_| rng.gen_range(-1.0f32..1.0f32)).collect(); + let r = (0..dims) + .map(|_| rng.random_range(-1.0f32..1.0f32)) + .collect(); centroids.push(r); } if is_spherical { @@ -140,8 +90,7 @@ fn quick_centers( centroids } -fn rabitq_index( - parallelism: &P, +fn rabitq_index( dims: usize, n: usize, c: usize, @@ -197,8 +146,8 @@ fn rabitq_index( elements: padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))), }); } - parallelism - .rayon_into_par_iter(0..n) + (0..n) + .into_par_iter() .map(|i| { let lut = rabitq::block::preprocess(&samples[i]); let mut result = (f32::INFINITY, 0); @@ -219,16 +168,15 @@ fn rabitq_index( .collect::>() } -fn flat_index( - parallelism: &P, +fn flat_index( _dims: usize, n: usize, c: usize, samples: &[Vec], centroids: &[Vec], ) -> Vec { - parallelism - .rayon_into_par_iter(0..n) + (0..n) + .into_par_iter() .map(|i| { let mut result = (f32::INFINITY, 0); for j in 0..c { @@ -242,8 +190,7 @@ fn flat_index( .collect::>() } -struct LloydKMeans<'a, P, F> { - parallelism: &'a P, +struct LloydKMeans<'a, F> { dims: usize, c: usize, is_spherical: bool, @@ -256,26 +203,19 @@ struct LloydKMeans<'a, P, F> { const DELTA: f32 = f16::EPSILON.to_f32_const(); -impl<'a, P: Parallelism, F: Fn(&P, &[Vec]) -> Vec> LloydKMeans<'a, P, F> { - fn new( - parallelism: &'a P, - c: usize, - dims: usize, - samples: &'a [Vec], - is_spherical: bool, - compute: F, - ) -> Self { +impl<'a, F: Fn(&[Vec]) -> Vec> LloydKMeans<'a, F> { + fn new(c: usize, dims: usize, samples: &'a [Vec], is_spherical: bool, compute: F) -> Self { let n = samples.len(); - let mut rng = StdRng::from_entropy(); + let mut rng = StdRng::from_seed([7; 32]); let mut centroids = Vec::with_capacity(c); for index in rand::seq::index::sample(&mut rng, n, c).into_iter() { centroids.push(samples[index].clone()); } - let assign = parallelism - .rayon_into_par_iter(0..n) + let assign = (0..n) + .into_par_iter() .map(|i| { let mut result = (f32::INFINITY, 0); for j in 0..c { @@ -289,7 +229,6 @@ impl<'a, P: Parallelism, F: Fn(&P, &[Vec]) -> Vec> LloydKMeans<'a, P .collect::>(); Self { - parallelism, dims, c, is_spherical, @@ -308,9 +247,8 @@ impl<'a, P: Parallelism, F: Fn(&P, &[Vec]) -> Vec> LloydKMeans<'a, P let samples = self.samples; let n = samples.len(); - let (sum, mut count) = self - .parallelism - .rayon_into_par_iter(0..n) + let (sum, mut count) = (0..n) + .into_par_iter() .fold( || (vec![vec![f32::zero(); dims]; c], vec![0.0f32; c]), |(mut sum, mut count), i| { @@ -330,9 +268,8 @@ impl<'a, P: Parallelism, F: Fn(&P, &[Vec]) -> Vec> LloydKMeans<'a, P }, ); - let mut centroids = self - .parallelism - .rayon_into_par_iter(0..c) + let mut centroids = (0..c) + .into_par_iter() .map(|i| f32::vector_mul_scalar(&sum[i], 1.0 / count[i])) .collect::>(); @@ -342,7 +279,7 @@ impl<'a, P: Parallelism, F: Fn(&P, &[Vec]) -> Vec> LloydKMeans<'a, P } let mut o = 0; loop { - let alpha = rand.gen_range(0.0..1.0f32); + let alpha = rand.random_range(0.0..1.0f32); let beta = (count[o] - 1.0) / (n - c) as f32; if alpha < beta { break; @@ -357,15 +294,13 @@ impl<'a, P: Parallelism, F: Fn(&P, &[Vec]) -> Vec> LloydKMeans<'a, P } if self.is_spherical { - self.parallelism - .rayon_into_par_iter(&mut centroids) - .for_each(|centroid| { - let l = f32::reduce_sum_of_x2(centroid).sqrt(); - f32::vector_mul_scalar_inplace(centroid, 1.0 / l); - }); + (&mut centroids).into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); } - let assign = (self.compute)(self.parallelism, ¢roids); + let assign = (self.compute)(¢roids); let result = (0..n).all(|i| assign[i] == self.assign[i]); diff --git a/crates/random_orthogonal_matrix/Cargo.toml b/crates/random_orthogonal_matrix/Cargo.toml index 16c3a39a..4e8f733c 100644 --- a/crates/random_orthogonal_matrix/Cargo.toml +++ b/crates/random_orthogonal_matrix/Cargo.toml @@ -8,8 +8,8 @@ edition.workspace = true nalgebra = "=0.33.0" rand.workspace = true -rand_chacha = "0.3.1" -rand_distr = "0.4.3" +rand_chacha = "0.9.0" +rand_distr = "0.5.0" [lints] workspace = true diff --git a/crates/random_orthogonal_matrix/src/lib.rs b/crates/random_orthogonal_matrix/src/lib.rs index e2711504..a9cbb5c8 100644 --- a/crates/random_orthogonal_matrix/src/lib.rs +++ b/crates/random_orthogonal_matrix/src/lib.rs @@ -30,15 +30,18 @@ fn random_full_rank_matrix(n: usize) -> DMatrix { #[test] fn check_random_orthogonal_matrix() { - assert_eq!(random_orthogonal_matrix(2), vec![ - vec![-0.5424608, -0.8400813], - vec![0.8400813, -0.54246056] - ]); - assert_eq!(random_orthogonal_matrix(3), vec![ - vec![-0.5309615, -0.69094884, -0.49058124], - vec![0.8222731, -0.56002235, -0.10120347], - vec![0.20481002, 0.45712686, -0.86549866] - ]); + assert_eq!( + random_orthogonal_matrix(2), + vec![vec![-0.5424608, -0.8400813], vec![0.8400813, -0.54246056]] + ); + assert_eq!( + random_orthogonal_matrix(3), + vec![ + vec![-0.5309615, -0.69094884, -0.49058124], + vec![0.8222731, -0.56002235, -0.10120347], + vec![0.20481002, 0.45712686, -0.86549866] + ] + ); } pub fn random_orthogonal_matrix(n: usize) -> Vec> { diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index adbd3e1d..bed14a33 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -4,15 +4,16 @@ version.workspace = true edition.workspace = true [dependencies] -half.workspace = true simd_macros = { path = "../simd_macros" } + +half.workspace = true zerocopy.workspace = true [dev-dependencies] rand.workspace = true [build-dependencies] -cc = "1.2.6" +cc = "1.2.13" [lints] workspace = true diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 4de1ed88..01eaa645 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -250,14 +250,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v4:avx512fp16)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -310,14 +310,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let specialized = unsafe { reduce_sum_of_xy_v4(&lhs, &rhs) }; let fallback = reduce_sum_of_xy_fallback(&lhs, &rhs); @@ -370,14 +370,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -439,14 +439,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v2:f16c:fma)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -492,14 +492,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v8.3a:fp16)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -540,14 +540,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v8.3a:sve)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -618,14 +618,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v4:avx512fp16)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -680,14 +680,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -746,14 +746,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -817,14 +817,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v2:f16c:fma)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -870,14 +870,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v8.3a:fp16)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -918,14 +918,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v8.3a:sve)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.gen_range(-1.0..=1.0))) + .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index d0dc2cdb..cdcb944c 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -164,11 +164,11 @@ mod reduce_sum_of_x { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -225,11 +225,11 @@ mod reduce_sum_of_x { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -280,11 +280,11 @@ mod reduce_sum_of_x { println!("test {} ... skipped (v2)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -334,11 +334,11 @@ mod reduce_sum_of_x { println!("test {} ... skipped (v8_3a)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -374,11 +374,11 @@ mod reduce_sum_of_x { println!("test {} ... skipped (v8_3a:sve)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -439,11 +439,11 @@ mod reduce_sum_of_abs_x { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -504,11 +504,11 @@ mod reduce_sum_of_abs_x { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -562,11 +562,11 @@ mod reduce_sum_of_abs_x { println!("test {} ... skipped (v2)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -618,11 +618,11 @@ mod reduce_sum_of_abs_x { println!("test {} ... skipped (v8.3a)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -658,11 +658,11 @@ mod reduce_sum_of_abs_x { println!("test {} ... skipped (v8.3a:sve)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -721,11 +721,11 @@ mod reduce_sum_of_x2 { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -782,11 +782,11 @@ mod reduce_sum_of_x2 { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -838,11 +838,11 @@ mod reduce_sum_of_x2 { println!("test {} ... skipped (v2:fma)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -892,11 +892,11 @@ mod reduce_sum_of_x2 { println!("test {} ... skipped (v8.3a)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -932,11 +932,11 @@ mod reduce_sum_of_x2 { println!("test {} ... skipped (v8.3a:sve)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let this = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let this = &this[..z]; @@ -1002,11 +1002,11 @@ mod reduce_min_max_of_x { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 200; let x = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 50..200 { let x = &x[..z]; @@ -1058,11 +1058,11 @@ mod reduce_min_max_of_x { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 200; let x = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 50..200 { let x = &x[..z]; @@ -1114,11 +1114,11 @@ mod reduce_min_max_of_x { println!("test {} ... skipped (v2)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 200; let x = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 50..200 { let x = &x[..z]; @@ -1169,11 +1169,11 @@ mod reduce_min_max_of_x { println!("test {} ... skipped (v8.3a)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 200; let x = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 50..200 { let x = &x[..z]; @@ -1214,11 +1214,11 @@ mod reduce_min_max_of_x { println!("test {} ... skipped (v8.3a:sve)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 200; let x = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 50..200 { let x = &x[..z]; @@ -1282,14 +1282,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1355,14 +1355,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1421,14 +1421,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v2:fma)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1485,14 +1485,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v8.3a)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1530,14 +1530,14 @@ mod reduce_sum_of_xy { println!("test {} ... skipped (v8.3a:sve)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1605,14 +1605,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1681,14 +1681,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1749,14 +1749,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v2:fma)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1815,14 +1815,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v8.3a)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1860,14 +1860,14 @@ mod reduce_sum_of_d2 { println!("test {} ... skipped (v8.3a:sve)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rhs = (0..n) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -1951,7 +1951,7 @@ mod reduce_sum_of_xy_sparse { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); pub fn sample_u32_sorted( rng: &mut (impl Rng + ?Sized), length: u32, @@ -1968,12 +1968,12 @@ mod reduce_sum_of_xy_sparse { let lm = 300; let lidx = sample_u32_sorted(&mut rng, 10000, lm); let lval = (0..lm) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rm = 350; let ridx = sample_u32_sorted(&mut rng, 10000, rm); let rval = (0..rm) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let specialized = unsafe { reduce_sum_of_xy_sparse_v4(&lidx, &lval, &ridx, &rval) }; let fallback = reduce_sum_of_xy_sparse_fallback(&lidx, &lval, &ridx, &rval); @@ -2101,7 +2101,7 @@ mod reduce_sum_of_d2_sparse { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); pub fn sample_u32_sorted( rng: &mut (impl Rng + ?Sized), length: u32, @@ -2118,12 +2118,12 @@ mod reduce_sum_of_d2_sparse { let lm = 300; let lidx = sample_u32_sorted(&mut rng, 10000, lm); let lval = (0..lm) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let rm = 350; let ridx = sample_u32_sorted(&mut rng, 10000, rm); let rval = (0..rm) - .map(|_| rng.gen_range(-1.0..=1.0)) + .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); let specialized = unsafe { reduce_sum_of_d2_sparse_v4(&lidx, &lval, &ridx, &rval) }; let fallback = reduce_sum_of_d2_sparse_fallback(&lidx, &lval, &ridx, &rval); diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 2f03d53c..fd132347 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -1,4 +1,3 @@ -#![feature(target_feature_11)] #![feature(avx512_target_feature)] #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512_f16))] diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index bc845332..4e0855f8 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -51,10 +51,10 @@ mod reduce_sum_of_x_as_u16 { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_as_u16_v4(this) }; @@ -101,10 +101,10 @@ mod reduce_sum_of_x_as_u16 { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_as_u16_v3(this) }; @@ -151,10 +151,10 @@ mod reduce_sum_of_x_as_u16 { println!("test {} ... skipped (v2)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_as_u16_v2(this) }; @@ -200,10 +200,10 @@ mod reduce_sum_of_x_as_u16 { println!("test {} ... skipped (v8.3a)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_as_u16_v8_3a(this) }; @@ -263,10 +263,10 @@ mod reduce_sum_of_x { println!("test {} ... skipped (v4)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_v4(this) }; @@ -313,10 +313,10 @@ mod reduce_sum_of_x { println!("test {} ... skipped (v3)", module_path!()); return; } - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.gen_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_v3(this) }; diff --git a/crates/simd_macros/Cargo.toml b/crates/simd_macros/Cargo.toml index 799db570..c608b3e4 100644 --- a/crates/simd_macros/Cargo.toml +++ b/crates/simd_macros/Cargo.toml @@ -7,9 +7,9 @@ edition.workspace = true proc-macro = true [dependencies] -proc-macro2 = { version = "1.0.79", features = ["proc-macro"] } -quote = "1.0.35" -syn = { version = "2.0.53", default-features = false, features = [ +proc-macro2 = { version = "1.0.93", features = ["proc-macro"] } +quote = "1.0.38" +syn = { version = "2.0.98", default-features = false, features = [ "clone-impls", "full", "parsing", diff --git a/crates/vector/Cargo.toml b/crates/vector/Cargo.toml index 186daf82..692bf642 100644 --- a/crates/vector/Cargo.toml +++ b/crates/vector/Cargo.toml @@ -5,8 +5,9 @@ edition.workspace = true [dependencies] distance = { path = "../distance" } -half.workspace = true simd = { path = "../simd" } +half.workspace = true + [lints] workspace = true diff --git a/rust-toolchain.toml b/rust-toolchain.toml index eb254600..d3e25b13 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2024-12-25" +channel = "nightly-2025-02-14" diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index a9a10ab0..e8dab2d1 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -13,7 +13,6 @@ use rand::Rng; use simd::Floating; use std::num::NonZeroU64; use std::ops::Deref; -use std::sync::Arc; use vector::vect::VectOwned; use vector::{VectorBorrowed, VectorOwned}; @@ -133,7 +132,7 @@ pub unsafe extern "C" fn ambuild( VchordrqBuildSourceOptions::Internal(internal_build) => { let mut tuples_total = 0_u64; let samples = 'a: { - let mut rand = rand::thread_rng(); + let mut rand = rand::rng(); let Some(max_number_of_samples) = internal_build .lists .last() @@ -156,7 +155,7 @@ pub unsafe extern "C" fn ambuild( samples.push(VectOwned::::build_to_vecf32(vector)); number_of_samples += 1; } else { - let index = rand.gen_range(0..max_number_of_samples) as usize; + let index = rand.random_range(0..max_number_of_samples) as usize; samples[index] = VectOwned::::build_to_vecf32(vector); } tuples_total += 1; @@ -174,7 +173,7 @@ pub unsafe extern "C" fn ambuild( samples.push(VectOwned::::build_to_vecf32(vector)); number_of_samples += 1; } else { - let index = rand.gen_range(0..max_number_of_samples) as usize; + let index = rand.random_range(0..max_number_of_samples) as usize; samples[index] = VectOwned::::build_to_vecf32(vector); } tuples_total += 1; @@ -388,13 +387,13 @@ mod vchordrq_cached { match self { VchordrqCached::_0 {} => { buffer.extend((0 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); buffer[size_of::()..][..size_of::()] .copy_from_slice(VchordrqCachedHeader0 {}.as_bytes()); } VchordrqCached::_1 { mapping, pages } => { buffer.extend((1 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat(0).take(size_of::())); + buffer.extend(std::iter::repeat_n(0, size_of::())); let mapping_s = buffer.len(); buffer.extend(mapping.as_bytes()); let mapping_e = buffer.len(); @@ -1043,27 +1042,21 @@ pub fn make_internal_build( } let mut result = Vec::>>::new(); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { - let means = k_means::RayonParallelism::scoped( + let means = k_means::k_means( internal_build.build_threads as _, - Arc::new(|| { + || { pgrx::check_for_interrupts!(); - }), - |parallelism| { - k_means::k_means( - parallelism, - w as usize, - vector_options.dims as usize, - if let Some(structure) = result.last() { - &structure.means - } else { - &samples - }, - internal_build.spherical_centroids, - 10, - ) }, - ) - .expect("failed to create thread pool"); + w as usize, + vector_options.dims as usize, + if let Some(structure) = result.last() { + &structure.means + } else { + &samples + }, + internal_build.spherical_centroids, + 10, + ); if let Some(structure) = result.last() { let mut children = vec![Vec::new(); means.len()]; for i in 0..structure.len() as u32 { From 0eecdc184238c1b004a5158e0ce6420926c47e16 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 19 Feb 2025 14:07:22 +0800 Subject: [PATCH 112/324] refactor: remove meaningless target feature requirements (#192) 1. only check simd-related features when dispatching 2. choose SVE implementation only if SVE bits >= 256 3. choose SVE f16 `sum_of_d2` and `sum_of_xy` implementation only if SVE bits >= 512 --------- Signed-off-by: usamoi --- crates/simd/build.rs | 2 +- crates/simd/cshim.c | 77 ++++---- crates/simd/src/bit.rs | 101 ++++++---- crates/simd/src/emulate.rs | 4 +- crates/simd/src/f16.rs | 322 ++++++++---------------------- crates/simd/src/f32.rs | 328 +++++++++++++++---------------- crates/simd/src/fast_scan/mod.rs | 33 ++-- crates/simd/src/lib.rs | 68 +++++++ crates/simd/src/packed_u4.rs | 2 +- crates/simd/src/quantize.rs | 30 +-- crates/simd/src/u8.rs | 46 ++--- crates/simd_macros/src/lib.rs | 19 +- crates/simd_macros/src/target.rs | 76 +++---- src/index/functions.rs | 4 +- 14 files changed, 499 insertions(+), 613 deletions(-) diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 22ebf935..c7a49d65 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -8,5 +8,5 @@ fn main() { .flag("-ffp-contract=fast") .flag("-freciprocal-math") .flag("-fno-signed-zeros") - .compile("base_cshim"); + .compile("simd_cshim"); } diff --git a/crates/simd/cshim.c b/crates/simd/cshim.c index 10b5ceea..7c0f044a 100644 --- a/crates/simd/cshim.c +++ b/crates/simd/cshim.c @@ -4,14 +4,16 @@ #ifdef __aarch64__ -#include -#include #include #include +#include +#include + +typedef __fp16 f16; +typedef float f32; -__attribute__((target("v8.3a,fp16"))) float -fp16_reduce_sum_of_xy_v8_3a_fp16_unroll(__fp16 *__restrict a, - __fp16 *__restrict b, size_t n) { +__attribute__((target("fp16"))) float +fp16_reduce_sum_of_xy_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { float16x8_t xy_0 = vdupq_n_f16(0.0); float16x8_t xy_1 = vdupq_n_f16(0.0); float16x8_t xy_2 = vdupq_n_f16(0.0); @@ -34,8 +36,8 @@ fp16_reduce_sum_of_xy_v8_3a_fp16_unroll(__fp16 *__restrict a, xy_3 = vfmaq_f16(xy_3, x_3, y_3); } if (n > 0) { - __fp16 A[32] = {}; - __fp16 B[32] = {}; + f16 A[32] = {}; + f16 B[32] = {}; for (size_t i = 0; i < n; i += 1) { A[i] = a[i]; B[i] = b[i]; @@ -54,14 +56,13 @@ fp16_reduce_sum_of_xy_v8_3a_fp16_unroll(__fp16 *__restrict a, xy_3 = vfmaq_f16(xy_3, x_3, y_3); } float16x8_t xy = vaddq_f16(vaddq_f16(xy_0, xy_1), vaddq_f16(xy_2, xy_3)); - return vgetq_lane_f16(xy, 0) + vgetq_lane_f16(xy, 1) + vgetq_lane_f16(xy, 2) + - vgetq_lane_f16(xy, 3) + vgetq_lane_f16(xy, 4) + vgetq_lane_f16(xy, 5) + - vgetq_lane_f16(xy, 6) + vgetq_lane_f16(xy, 7); + float32x4_t lo = vcvt_f32_f16(vget_low_f16(xy)); + float32x4_t hi = vcvt_f32_f16(vget_high_f16(xy)); + return vaddvq_f32(lo) + vaddvq_f32(hi); } -__attribute__((target("v8.3a,sve"))) float -fp16_reduce_sum_of_xy_v8_3a_sve(__fp16 *__restrict a, __fp16 *__restrict b, - size_t n) { +__attribute__((target("sve"))) float +fp16_reduce_sum_of_xy_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svfloat16_t xy = svdup_f16(0.0); for (size_t i = 0; i < n; i += svcnth()) { svbool_t mask = svwhilelt_b16(i, n); @@ -72,9 +73,8 @@ fp16_reduce_sum_of_xy_v8_3a_sve(__fp16 *__restrict a, __fp16 *__restrict b, return svaddv_f16(svptrue_b16(), xy); } -__attribute__((target("v8.3a,fp16"))) float -fp16_reduce_sum_of_d2_v8_3a_fp16_unroll(__fp16 *__restrict a, - __fp16 *__restrict b, size_t n) { +__attribute__((target("fp16"))) float +fp16_reduce_sum_of_d2_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { float16x8_t d2_0 = vdupq_n_f16(0.0); float16x8_t d2_1 = vdupq_n_f16(0.0); float16x8_t d2_2 = vdupq_n_f16(0.0); @@ -101,8 +101,8 @@ fp16_reduce_sum_of_d2_v8_3a_fp16_unroll(__fp16 *__restrict a, d2_3 = vfmaq_f16(d2_3, d_3, d_3); } if (n > 0) { - __fp16 A[32] = {}; - __fp16 B[32] = {}; + f16 A[32] = {}; + f16 B[32] = {}; for (size_t i = 0; i < n; i += 1) { A[i] = a[i]; B[i] = b[i]; @@ -125,14 +125,13 @@ fp16_reduce_sum_of_d2_v8_3a_fp16_unroll(__fp16 *__restrict a, d2_3 = vfmaq_f16(d2_3, d_3, d_3); } float16x8_t d2 = vaddq_f16(vaddq_f16(d2_0, d2_1), vaddq_f16(d2_2, d2_3)); - return vgetq_lane_f16(d2, 0) + vgetq_lane_f16(d2, 1) + vgetq_lane_f16(d2, 2) + - vgetq_lane_f16(d2, 3) + vgetq_lane_f16(d2, 4) + vgetq_lane_f16(d2, 5) + - vgetq_lane_f16(d2, 6) + vgetq_lane_f16(d2, 7); + float32x4_t lo = vcvt_f32_f16(vget_low_f16(d2)); + float32x4_t hi = vcvt_f32_f16(vget_high_f16(d2)); + return vaddvq_f32(lo) + vaddvq_f32(hi); } -__attribute__((target("v8.3a,sve"))) float -fp16_reduce_sum_of_d2_v8_3a_sve(__fp16 *__restrict a, __fp16 *__restrict b, - size_t n) { +__attribute__((target("sve"))) float +fp16_reduce_sum_of_d2_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svfloat16_t d2 = svdup_f16(0.0); for (size_t i = 0; i < n; i += svcnth()) { svbool_t mask = svwhilelt_b16(i, n); @@ -144,8 +143,8 @@ fp16_reduce_sum_of_d2_v8_3a_sve(__fp16 *__restrict a, __fp16 *__restrict b, return svaddv_f16(svptrue_b16(), d2); } -__attribute__((target("v8.3a,sve"))) float -fp32_reduce_sum_of_x_v8_3a_sve(float *__restrict this, size_t n) { +__attribute__((target("sve"))) float +fp32_reduce_sum_of_x_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32(i, n); @@ -155,8 +154,8 @@ fp32_reduce_sum_of_x_v8_3a_sve(float *__restrict this, size_t n) { return svaddv_f32(svptrue_b32(), sum); } -__attribute__((target("v8.3a,sve"))) float -fp32_reduce_sum_of_abs_x_v8_3a_sve(float *__restrict this, size_t n) { +__attribute__((target("sve"))) float +fp32_reduce_sum_of_abs_x_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32(i, n); @@ -166,8 +165,8 @@ fp32_reduce_sum_of_abs_x_v8_3a_sve(float *__restrict this, size_t n) { return svaddv_f32(svptrue_b32(), sum); } -__attribute__((target("v8.3a,sve"))) float -fp32_reduce_sum_of_x2_v8_3a_sve(float *__restrict this, size_t n) { +__attribute__((target("sve"))) float +fp32_reduce_sum_of_x2_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32(i, n); @@ -177,9 +176,9 @@ fp32_reduce_sum_of_x2_v8_3a_sve(float *__restrict this, size_t n) { return svaddv_f32(svptrue_b32(), sum); } -__attribute__((target("v8.3a,sve"))) void -fp32_reduce_min_max_of_x_v8_3a_sve(float *__restrict this, size_t n, - float *out_min, float *out_max) { +__attribute__((target("sve"))) void +fp32_reduce_min_max_of_x_a3_256(float *restrict this, size_t n, float *out_min, + float *out_max) { svfloat32_t min = svdup_f32(1.0 / 0.0); svfloat32_t max = svdup_f32(-1.0 / 0.0); for (size_t i = 0; i < n; i += svcntw()) { @@ -192,9 +191,9 @@ fp32_reduce_min_max_of_x_v8_3a_sve(float *__restrict this, size_t n, *out_max = svmaxv_f32(svptrue_b32(), max); } -__attribute__((target("v8.3a,sve"))) float -fp32_reduce_sum_of_xy_v8_3a_sve(float *__restrict lhs, float *__restrict rhs, - size_t n) { +__attribute__((target("sve"))) float +fp32_reduce_sum_of_xy_a3_256(float *restrict lhs, float *restrict rhs, + size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32(i, n); @@ -205,9 +204,9 @@ fp32_reduce_sum_of_xy_v8_3a_sve(float *__restrict lhs, float *__restrict rhs, return svaddv_f32(svptrue_b32(), sum); } -__attribute__((target("v8.3a,sve"))) float -fp32_reduce_sum_of_d2_v8_3a_sve(float *__restrict lhs, float *__restrict rhs, - size_t n) { +__attribute__((target("sve"))) float +fp32_reduce_sum_of_d2_a3_256(float *restrict lhs, float *restrict rhs, + size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32(i, n); diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index dc417bda..d8321dc4 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -8,9 +8,9 @@ mod sum_of_and { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] + #[crate::target_cpu(enable = "v4.512")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_and_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + fn sum_of_and_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -38,21 +38,24 @@ mod sum_of_and { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_and_v4_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + fn sum_of_and_v4_512_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!( + "test {} ... skipped (v4.512:avx512vpopcntdq)", + module_path!() + ); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_and_v4_avx512vpopcntdq(&lhs, &rhs) }; - let fallback = sum_of_and_fallback(&lhs, &rhs); + let specialized = unsafe { sum_of_and_v4_512_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] pub fn sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -74,9 +77,9 @@ mod sum_of_or { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] + #[crate::target_cpu(enable = "v4.512")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + fn sum_of_or_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -104,21 +107,24 @@ mod sum_of_or { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_or_v4_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + fn sum_of_or_v4_512_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!( + "test {} ... skipped (v4.512:avx512vpopcntdq)", + module_path!() + ); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_or_v4_avx512vpopcntdq(&lhs, &rhs) }; - let fallback = sum_of_or_fallback(&lhs, &rhs); + let specialized = unsafe { sum_of_or_v4_512_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] pub fn sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -140,9 +146,9 @@ mod sum_of_xor { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] + #[crate::target_cpu(enable = "v4.512")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_xor_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + fn sum_of_xor_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -170,21 +176,24 @@ mod sum_of_xor { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_xor_v4_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + fn sum_of_xor_v4_512_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!( + "test {} ... skipped (v4.512:avx512vpopcntdq)", + module_path!() + ); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_xor_v4_avx512vpopcntdq(&lhs, &rhs) }; - let fallback = sum_of_xor_fallback(&lhs, &rhs); + let specialized = unsafe { sum_of_xor_v4_512_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] pub fn sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -206,9 +215,9 @@ mod sum_of_and_or { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] + #[crate::target_cpu(enable = "v4.512")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_and_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + fn sum_of_and_or_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -242,21 +251,24 @@ mod sum_of_and_or { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_xor_v4_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + fn sum_of_xor_v4_512_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!( + "test {} ... skipped (v4.512:avx512vpopcntdq)", + module_path!() + ); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_and_or_v4_avx512vpopcntdq(&lhs, &rhs) }; - let fallback = sum_of_and_or_fallback(&lhs, &rhs); + let specialized = unsafe { sum_of_and_or_v4_512_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] pub fn sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -280,9 +292,9 @@ mod sum_of_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] + #[crate::target_cpu(enable = "v4.512")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_x_v4_avx512vpopcntdq(this: &[u64]) -> u32 { + fn sum_of_x_v4_512_avx512vpopcntdq(this: &[u64]) -> u32 { unsafe { use std::arch::x86_64::*; let mut and = _mm512_setzero_si512(); @@ -305,20 +317,23 @@ mod sum_of_x { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_x_v4_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + fn sum_of_x_v4_512_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!( + "test {} ... skipped (v4.512:avx512vpopcntdq)", + module_path!() + ); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let this = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_x_v4_avx512vpopcntdq(&this) }; - let fallback = sum_of_x_fallback(&this); + let specialized = unsafe { sum_of_x_v4_512_avx512vpopcntdq(&this) }; + let fallback = fallback(&this); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] pub fn sum_of_x(this: &[u64]) -> u32 { let n = this.len(); let mut and = 0; @@ -335,7 +350,7 @@ pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_and { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -358,7 +373,7 @@ pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_or { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -381,7 +396,7 @@ pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_xor { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index 520b6a1d..2715d860 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -3,7 +3,7 @@ // Instructions. arXiv preprint arXiv:2112.06342. #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v4")] +#[crate::target_cpu(enable = "v4.512")] pub fn emulate_mm512_2intersect_epi32( a: std::arch::x86_64::__m512i, b: std::arch::x86_64::__m512i, @@ -85,7 +85,7 @@ pub fn emulate_mm_reduce_add_ps(mut x: std::arch::x86_64::__m128) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v4")] +#[crate::target_cpu(enable = "v4.512")] pub fn emulate_mm512_reduce_add_epi16(x: std::arch::x86_64::__m512i) -> i16 { unsafe { use std::arch::x86_64::*; diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 01eaa645..2fb176be 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -131,7 +131,7 @@ impl Floating for f16 { mod reduce_or_of_is_zero_x { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { for &x in this { if x == f16::ZERO { @@ -147,7 +147,7 @@ mod reduce_sum_of_x { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_sum_of_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; @@ -163,7 +163,7 @@ mod reduce_sum_of_abs_x { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; @@ -179,7 +179,7 @@ mod reduce_sum_of_x2 { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_sum_of_x2(this: &[f16]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; @@ -195,7 +195,7 @@ mod reduce_min_max_of_x { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -213,9 +213,9 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] + #[crate::target_cpu(enable = "v4.512")] #[target_feature(enable = "avx512fp16")] - pub fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + pub fn reduce_sum_of_xy_v4_512_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -243,11 +243,11 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_xy_v4_avx512fp16_test() { + fn reduce_sum_of_xy_v4_512_avx512fp16_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { - println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4_512:avx512fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -262,8 +262,8 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v4_avx512fp16(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_xy_v4_512_avx512fp16(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -274,8 +274,8 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - pub fn reduce_sum_of_xy_v4(lhs: &[f16], rhs: &[f16]) -> f32 { + #[crate::target_cpu(enable = "v4.512")] + pub fn reduce_sum_of_xy_v4_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -306,7 +306,7 @@ mod reduce_sum_of_xy { fn reduce_sum_of_xy_v4_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -319,8 +319,8 @@ mod reduce_sum_of_xy { let rhs = (0..n) .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); - let specialized = unsafe { reduce_sum_of_xy_v4(&lhs, &rhs) }; - let fallback = reduce_sum_of_xy_fallback(&lhs, &rhs); + let specialized = unsafe { reduce_sum_of_xy_v4_512(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -383,76 +383,7 @@ mod reduce_sum_of_xy { let lhs = &lhs[..z]; let rhs = &rhs[..z]; let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); - assert!( - (specialized - fallback).abs() < EPSILON, - "specialized = {specialized}, fallback = {fallback}." - ); - } - } - } - - #[inline] - #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v2")] - #[target_feature(enable = "f16c")] - #[target_feature(enable = "fma")] - pub fn reduce_sum_of_xy_v2_f16c_fma(lhs: &[f16], rhs: &[f16]) -> f32 { - use crate::emulate::emulate_mm_reduce_add_ps; - assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut xy = _mm_setzero_ps(); - while n >= 4 { - let x = _mm_cvtph_ps(_mm_loadu_si128(a.cast())); - let y = _mm_cvtph_ps(_mm_loadu_si128(b.cast())); - a = a.add(4); - b = b.add(4); - n -= 4; - xy = _mm_fmadd_ps(x, y, xy); - } - let mut xy = emulate_mm_reduce_add_ps(xy); - while n > 0 { - let x = a.read().to_f32(); - let y = b.read().to_f32(); - a = a.add(1); - b = b.add(1); - n -= 1; - xy += x * y; - } - xy - } - } - - #[cfg(all(target_arch = "x86_64", test, not(miri)))] - #[test] - fn reduce_sum_of_xy_v2_f16c_fma_test() { - use rand::Rng; - const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v2") - || !crate::is_feature_detected!("f16c") - || !crate::is_feature_detected!("fma") - { - println!("test {} ... skipped (v2:f16c:fma)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) - .collect::>(); - let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) - .collect::>(); - for z in 3984..4016 { - let lhs = &lhs[..z]; - let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v2_f16c_fma(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -463,33 +394,25 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "fp16")] - pub fn reduce_sum_of_xy_v8_3a_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + pub fn reduce_sum_of_xy_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { extern "C" { - fn fp16_reduce_sum_of_xy_v8_3a_fp16_unroll( - a: *const (), - b: *const (), - n: usize, - ) -> f32; + fn fp16_reduce_sum_of_xy_a2_fp16(a: *const (), b: *const (), n: usize) -> f32; } - fp16_reduce_sum_of_xy_v8_3a_fp16_unroll( - lhs.as_ptr().cast(), - rhs.as_ptr().cast(), - lhs.len(), - ) + fp16_reduce_sum_of_xy_a2_fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_xy_v8_3a_fp16_test() { + fn reduce_sum_of_xy_a2_fp16_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("fp16") { - println!("test {} ... skipped (v8.3a:fp16)", module_path!()); + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -504,8 +427,8 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v8_3a_fp16(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_xy_a2_fp16(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -514,30 +437,27 @@ mod reduce_sum_of_xy { } } - // temporarily disables this for uncertain precision - #[cfg_attr(not(test), expect(dead_code))] #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "sve")] - pub fn reduce_sum_of_xy_v8_3a_sve(lhs: &[f16], rhs: &[f16]) -> f32 { + pub fn reduce_sum_of_xy_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { extern "C" { - fn fp16_reduce_sum_of_xy_v8_3a_sve(a: *const (), b: *const (), n: usize) -> f32; + fn fp16_reduce_sum_of_xy_a3_512(a: *const (), b: *const (), n: usize) -> f32; } - fp16_reduce_sum_of_xy_v8_3a_sve(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + fp16_reduce_sum_of_xy_a3_512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - #[ignore] - fn reduce_sum_of_xy_v8_3a_sve_test() { + fn reduce_sum_of_xy_a3_512_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { - println!("test {} ... skipped (v8.3a:sve)", module_path!()); + if !crate::is_cpu_detected!("a3.512") { + println!("test {} ... skipped (a3.512)", module_path!()); return; } let mut rng = rand::rng(); @@ -552,8 +472,8 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v8_3a_sve(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_xy_a3_512(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -562,7 +482,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"v2:f16c:fma", @"v8.3a:fp16")] + #[crate::multiversion(@"v4.512:avx512fp16", @"v4.512", @"v3", @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -579,9 +499,9 @@ mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] + #[crate::target_cpu(enable = "v4.512")] #[target_feature(enable = "avx512fp16")] - pub fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + pub fn reduce_sum_of_d2_v4_512_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -611,11 +531,11 @@ mod reduce_sum_of_d2 { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_d2_v4_avx512fp16_test() { + fn reduce_sum_of_d2_v4_512_avx512fp16_test() { use rand::Rng; const EPSILON: f32 = 6.0; - if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { - println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4_512:avx512fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -630,8 +550,8 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v4_avx512fp16(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_d2_v4_512_avx512fp16(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -642,8 +562,8 @@ mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - pub fn reduce_sum_of_d2_v4(lhs: &[f16], rhs: &[f16]) -> f32 { + #[crate::target_cpu(enable = "v4.512")] + pub fn reduce_sum_of_d2_v4_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -676,7 +596,7 @@ mod reduce_sum_of_d2 { fn reduce_sum_of_d2_v4_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -692,8 +612,8 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v4(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_d2_v4_512(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -759,78 +679,7 @@ mod reduce_sum_of_d2 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; let specialized = unsafe { reduce_sum_of_d2_v3(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); - assert!( - (specialized - fallback).abs() < EPSILON, - "specialized = {specialized}, fallback = {fallback}." - ); - } - } - } - - #[inline] - #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v2")] - #[target_feature(enable = "f16c")] - #[target_feature(enable = "fma")] - pub fn reduce_sum_of_d2_v2_f16c_fma(lhs: &[f16], rhs: &[f16]) -> f32 { - use crate::emulate::emulate_mm_reduce_add_ps; - assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len() as u32; - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut d2 = _mm_setzero_ps(); - while n >= 4 { - let x = _mm_cvtph_ps(_mm_loadu_si128(a.cast())); - let y = _mm_cvtph_ps(_mm_loadu_si128(b.cast())); - a = a.add(4); - b = b.add(4); - n -= 4; - let d = _mm_sub_ps(x, y); - d2 = _mm_fmadd_ps(d, d, d2); - } - let mut d2 = emulate_mm_reduce_add_ps(d2); - while n > 0 { - let x = a.read().to_f32(); - let y = b.read().to_f32(); - a = a.add(1); - b = b.add(1); - n -= 1; - let d = x - y; - d2 += d * d; - } - d2 - } - } - - #[cfg(all(target_arch = "x86_64", test, not(miri)))] - #[test] - fn reduce_sum_of_d2_v2_f16c_fma_test() { - use rand::Rng; - const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v2") - || !crate::is_feature_detected!("f16c") - || !crate::is_feature_detected!("fma") - { - println!("test {} ... skipped (v2:f16c:fma)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) - .collect::>(); - let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) - .collect::>(); - for z in 3984..4016 { - let lhs = &lhs[..z]; - let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v2_f16c_fma(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -841,33 +690,25 @@ mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "fp16")] - pub fn reduce_sum_of_d2_v8_3a_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + pub fn reduce_sum_of_d2_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { extern "C" { - fn fp16_reduce_sum_of_d2_v8_3a_fp16_unroll( - a: *const (), - b: *const (), - n: usize, - ) -> f32; + fn fp16_reduce_sum_of_d2_a2_fp16(a: *const (), b: *const (), n: usize) -> f32; } - fp16_reduce_sum_of_d2_v8_3a_fp16_unroll( - lhs.as_ptr().cast(), - rhs.as_ptr().cast(), - lhs.len(), - ) + fp16_reduce_sum_of_d2_a2_fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_d2_v8_3a_fp16_test() { + fn reduce_sum_of_d2_a2_fp16_test() { use rand::Rng; const EPSILON: f32 = 6.0; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("fp16") { - println!("test {} ... skipped (v8.3a:fp16)", module_path!()); + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -882,8 +723,8 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v8_3a_fp16(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_d2_a2_fp16(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -892,30 +733,27 @@ mod reduce_sum_of_d2 { } } - // temporarily disables this for uncertain precision - #[cfg_attr(not(test), expect(dead_code))] #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "sve")] - pub fn reduce_sum_of_d2_v8_3a_sve(lhs: &[f16], rhs: &[f16]) -> f32 { + pub fn reduce_sum_of_d2_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { extern "C" { - fn fp16_reduce_sum_of_d2_v8_3a_sve(a: *const (), b: *const (), n: usize) -> f32; + fn fp16_reduce_sum_of_d2_a3_512(a: *const (), b: *const (), n: usize) -> f32; } - fp16_reduce_sum_of_d2_v8_3a_sve(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + fp16_reduce_sum_of_d2_a3_512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - #[ignore] - fn reduce_sum_of_d2_v8_3a_sve_test() { + fn reduce_sum_of_d2_a3_512_test() { use rand::Rng; const EPSILON: f32 = 6.0; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { - println!("test {} ... skipped (v8.3a:sve)", module_path!()); + if !crate::is_cpu_detected!("a3.512") { + println!("test {} ... skipped (a3.512)", module_path!()); return; } let mut rng = rand::rng(); @@ -930,8 +768,8 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v8_3a_sve(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_d2_a3_512(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -940,7 +778,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"v2:f16c:fma", @"v8.3a:fp16")] + #[crate::multiversion(@"v4.512:avx512fp16", @"v4.512", @"v3", @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -959,7 +797,7 @@ mod reduce_sum_of_xy_sparse { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -992,7 +830,7 @@ mod reduce_sum_of_d2_sparse { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -1031,7 +869,7 @@ mod reduce_sum_of_d2_sparse { mod vector_add { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_add(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -1051,7 +889,7 @@ mod vector_add { mod vector_add_inplace { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_add_inplace(lhs: &mut [f16], rhs: &[f16]) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -1064,7 +902,7 @@ mod vector_add_inplace { mod vector_sub { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_sub(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -1084,7 +922,7 @@ mod vector_sub { mod vector_mul { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_mul(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -1104,7 +942,7 @@ mod vector_mul { mod vector_mul_scalar { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_mul_scalar(lhs: &[f16], rhs: f32) -> Vec { let rhs = f16::from_f32(rhs); let n = lhs.len(); @@ -1124,7 +962,7 @@ mod vector_mul_scalar { mod vector_mul_scalar_inplace { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_mul_scalar_inplace(lhs: &mut [f16], rhs: f32) { let rhs = f16::from_f32(rhs); let n = lhs.len(); @@ -1137,7 +975,7 @@ mod vector_mul_scalar_inplace { mod vector_abs_inplace { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_abs_inplace(this: &mut [f16]) { let n = this.len(); for i in 0..n { @@ -1149,7 +987,7 @@ mod vector_abs_inplace { mod vector_from_f32 { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_from_f32(this: &[f32]) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); @@ -1168,7 +1006,7 @@ mod vector_from_f32 { mod vector_to_f32 { use half::f16; - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_to_f32(this: &[f16]) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index cdcb944c..54670deb 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -119,7 +119,7 @@ impl Floating for f32 { } mod reduce_or_of_is_zero_x { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { for &x in this { if x == 0.0f32 { @@ -133,8 +133,8 @@ mod reduce_or_of_is_zero_x { mod reduce_sum_of_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_x_v4(this: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_sum_of_x_v4_512(this: &[f32]) -> f32 { unsafe { use std::arch::x86_64::*; let mut n = this.len(); @@ -160,7 +160,7 @@ mod reduce_sum_of_x { fn reduce_sum_of_x_v4_test() { use rand::Rng; const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -172,8 +172,8 @@ mod reduce_sum_of_x { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_v4(this) }; - let fallback = reduce_sum_of_x_fallback(this); + let specialized = unsafe { reduce_sum_of_x_v4_512(this) }; + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -234,7 +234,7 @@ mod reduce_sum_of_x { for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_v3(this) }; - let fallback = reduce_sum_of_x_fallback(this); + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -289,7 +289,7 @@ mod reduce_sum_of_x { for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_v2(this) }; - let fallback = reduce_sum_of_x_fallback(this); + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -300,8 +300,8 @@ mod reduce_sum_of_x { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] - fn reduce_sum_of_x_v8_3a(this: &[f32]) -> f32 { + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x_a2(this: &[f32]) -> f32 { unsafe { use std::arch::aarch64::*; let mut n = this.len(); @@ -327,11 +327,11 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_x_v8_3a_test() { + fn reduce_sum_of_x_a2_test() { use rand::Rng; const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("v8.3a") { - println!("test {} ... skipped (v8_3a)", module_path!()); + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } let mut rng = rand::rng(); @@ -342,8 +342,8 @@ mod reduce_sum_of_x { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_v8_3a(this) }; - let fallback = reduce_sum_of_x_fallback(this); + let specialized = unsafe { reduce_sum_of_x_a2(this) }; + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -354,24 +354,24 @@ mod reduce_sum_of_x { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "sve")] - fn reduce_sum_of_x_v8_3a_sve(this: &[f32]) -> f32 { + fn reduce_sum_of_x_a3_256(this: &[f32]) -> f32 { unsafe { extern "C" { - fn fp32_reduce_sum_of_x_v8_3a_sve(this: *const f32, n: usize) -> f32; + fn fp32_reduce_sum_of_x_a3_256(this: *const f32, n: usize) -> f32; } - fp32_reduce_sum_of_x_v8_3a_sve(this.as_ptr(), this.len()) + fp32_reduce_sum_of_x_a3_256(this.as_ptr(), this.len()) } } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_x_v8_3a_sve_test() { + fn reduce_sum_of_x_a3_256_test() { use rand::Rng; const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { - println!("test {} ... skipped (v8_3a:sve)", module_path!()); + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); return; } let mut rng = rand::rng(); @@ -382,8 +382,8 @@ mod reduce_sum_of_x { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_v8_3a_sve(this) }; - let fallback = reduce_sum_of_x_fallback(this); + let specialized = unsafe { reduce_sum_of_x_a3_256(this) }; + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -392,7 +392,7 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a:sve", @"v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a3.256", @"a2")] pub fn reduce_sum_of_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -406,8 +406,8 @@ mod reduce_sum_of_x { mod reduce_sum_of_abs_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_abs_x_v4(this: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_sum_of_abs_x_v4_512(this: &[f32]) -> f32 { unsafe { use std::arch::x86_64::*; let mut n = this.len(); @@ -435,7 +435,7 @@ mod reduce_sum_of_abs_x { fn reduce_sum_of_abs_x_v4_test() { use rand::Rng; const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -447,8 +447,8 @@ mod reduce_sum_of_abs_x { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_abs_x_v4(this) }; - let fallback = reduce_sum_of_abs_x_fallback(this); + let specialized = unsafe { reduce_sum_of_abs_x_v4_512(this) }; + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -513,7 +513,7 @@ mod reduce_sum_of_abs_x { for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_abs_x_v3(this) }; - let fallback = reduce_sum_of_abs_x_fallback(this); + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -571,7 +571,7 @@ mod reduce_sum_of_abs_x { for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_abs_x_v2(this) }; - let fallback = reduce_sum_of_abs_x_fallback(this); + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -582,8 +582,8 @@ mod reduce_sum_of_abs_x { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] - fn reduce_sum_of_abs_x_v8_3a(this: &[f32]) -> f32 { + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_abs_x_a2(this: &[f32]) -> f32 { unsafe { use std::arch::aarch64::*; let mut n = this.len(); @@ -611,11 +611,11 @@ mod reduce_sum_of_abs_x { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_abs_x_v8_3a_test() { + fn reduce_sum_of_abs_x_a2_test() { use rand::Rng; const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("v8.3a") { - println!("test {} ... skipped (v8.3a)", module_path!()); + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } let mut rng = rand::rng(); @@ -626,8 +626,8 @@ mod reduce_sum_of_abs_x { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_abs_x_v8_3a(this) }; - let fallback = reduce_sum_of_abs_x_fallback(this); + let specialized = unsafe { reduce_sum_of_abs_x_a2(this) }; + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -638,24 +638,24 @@ mod reduce_sum_of_abs_x { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "sve")] - fn reduce_sum_of_abs_x_v8_3a_sve(this: &[f32]) -> f32 { + fn reduce_sum_of_abs_x_a3_256(this: &[f32]) -> f32 { unsafe { extern "C" { - fn fp32_reduce_sum_of_abs_x_v8_3a_sve(this: *const f32, n: usize) -> f32; + fn fp32_reduce_sum_of_abs_x_a3_256(this: *const f32, n: usize) -> f32; } - fp32_reduce_sum_of_abs_x_v8_3a_sve(this.as_ptr(), this.len()) + fp32_reduce_sum_of_abs_x_a3_256(this.as_ptr(), this.len()) } } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_abs_x_v8_3a_sve_test() { + fn reduce_sum_of_abs_x_a3_256_test() { use rand::Rng; const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { - println!("test {} ... skipped (v8.3a:sve)", module_path!()); + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); return; } let mut rng = rand::rng(); @@ -666,8 +666,8 @@ mod reduce_sum_of_abs_x { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_abs_x_v8_3a_sve(this) }; - let fallback = reduce_sum_of_abs_x_fallback(this); + let specialized = unsafe { reduce_sum_of_abs_x_a3_256(this) }; + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -676,7 +676,7 @@ mod reduce_sum_of_abs_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a:sve", @"v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a3.256", @"a2")] pub fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -690,8 +690,8 @@ mod reduce_sum_of_abs_x { mod reduce_sum_of_x2 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_x2_v4(this: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_sum_of_x2_v4_512(this: &[f32]) -> f32 { unsafe { use std::arch::x86_64::*; let mut n = this.len(); @@ -717,7 +717,7 @@ mod reduce_sum_of_x2 { fn reduce_sum_of_x2_v4_test() { use rand::Rng; const EPSILON: f32 = 0.006; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -729,8 +729,8 @@ mod reduce_sum_of_x2 { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x2_v4(this) }; - let fallback = reduce_sum_of_x2_fallback(this); + let specialized = unsafe { reduce_sum_of_x2_v4_512(this) }; + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -791,7 +791,7 @@ mod reduce_sum_of_x2 { for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x2_v3(this) }; - let fallback = reduce_sum_of_x2_fallback(this); + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -847,7 +847,7 @@ mod reduce_sum_of_x2 { for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x2_v2_fma(this) }; - let fallback = reduce_sum_of_x2_fallback(this); + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -858,8 +858,8 @@ mod reduce_sum_of_x2 { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] - fn reduce_sum_of_x2_v8_3a(this: &[f32]) -> f32 { + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x2_a2(this: &[f32]) -> f32 { unsafe { use std::arch::aarch64::*; let mut n = this.len(); @@ -885,11 +885,11 @@ mod reduce_sum_of_x2 { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_x2_v8_3a_test() { + fn reduce_sum_of_x2_a2_test() { use rand::Rng; const EPSILON: f32 = 0.006; - if !crate::is_cpu_detected!("v8.3a") { - println!("test {} ... skipped (v8.3a)", module_path!()); + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } let mut rng = rand::rng(); @@ -900,8 +900,8 @@ mod reduce_sum_of_x2 { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x2_v8_3a(this) }; - let fallback = reduce_sum_of_x2_fallback(this); + let specialized = unsafe { reduce_sum_of_x2_a2(this) }; + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -912,24 +912,24 @@ mod reduce_sum_of_x2 { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "sve")] - fn reduce_sum_of_x2_v8_3a_sve(this: &[f32]) -> f32 { + fn reduce_sum_of_x2_a3_256(this: &[f32]) -> f32 { unsafe { extern "C" { - fn fp32_reduce_sum_of_x2_v8_3a_sve(this: *const f32, n: usize) -> f32; + fn fp32_reduce_sum_of_x2_a3_256(this: *const f32, n: usize) -> f32; } - fp32_reduce_sum_of_x2_v8_3a_sve(this.as_ptr(), this.len()) + fp32_reduce_sum_of_x2_a3_256(this.as_ptr(), this.len()) } } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_x2_v8_3a_sve_test() { + fn reduce_sum_of_x2_a3_256_test() { use rand::Rng; const EPSILON: f32 = 0.006; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { - println!("test {} ... skipped (v8.3a:sve)", module_path!()); + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); return; } let mut rng = rand::rng(); @@ -940,8 +940,8 @@ mod reduce_sum_of_x2 { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x2_v8_3a_sve(this) }; - let fallback = reduce_sum_of_x2_fallback(this); + let specialized = unsafe { reduce_sum_of_x2_a3_256(this) }; + let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -950,7 +950,7 @@ mod reduce_sum_of_x2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"v8.3a:sve", @"v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", @"v2:fma", @"a3.256", @"a2")] pub fn reduce_sum_of_x2(this: &[f32]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; @@ -967,8 +967,8 @@ mod reduce_min_max_of_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_min_max_of_x_v4(this: &[f32]) -> (f32, f32) { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_min_max_of_x_v4_512(this: &[f32]) -> (f32, f32) { unsafe { use std::arch::x86_64::*; let mut n = this.len(); @@ -998,7 +998,7 @@ mod reduce_min_max_of_x { #[test] fn reduce_min_max_of_x_v4_test() { use rand::Rng; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -1010,8 +1010,8 @@ mod reduce_min_max_of_x { .collect::>(); for z in 50..200 { let x = &x[..z]; - let specialized = unsafe { reduce_min_max_of_x_v4(x) }; - let fallback = reduce_min_max_of_x_fallback(x); + let specialized = unsafe { reduce_min_max_of_x_v4_512(x) }; + let fallback = fallback(x); assert_eq!(specialized.0, fallback.0); assert_eq!(specialized.1, fallback.1); } @@ -1067,7 +1067,7 @@ mod reduce_min_max_of_x { for z in 50..200 { let x = &x[..z]; let specialized = unsafe { reduce_min_max_of_x_v3(x) }; - let fallback = reduce_min_max_of_x_fallback(x); + let fallback = fallback(x); assert_eq!(specialized.0, fallback.0,); assert_eq!(specialized.1, fallback.1,); } @@ -1123,7 +1123,7 @@ mod reduce_min_max_of_x { for z in 50..200 { let x = &x[..z]; let specialized = unsafe { reduce_min_max_of_x_v2(x) }; - let fallback = reduce_min_max_of_x_fallback(x); + let fallback = fallback(x); assert_eq!(specialized.0, fallback.0,); assert_eq!(specialized.1, fallback.1,); } @@ -1132,8 +1132,8 @@ mod reduce_min_max_of_x { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] - fn reduce_min_max_of_x_v8_3a(this: &[f32]) -> (f32, f32) { + #[crate::target_cpu(enable = "a2")] + fn reduce_min_max_of_x_a2(this: &[f32]) -> (f32, f32) { unsafe { use std::arch::aarch64::*; let mut n = this.len(); @@ -1163,10 +1163,10 @@ mod reduce_min_max_of_x { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_min_max_of_x_v8_3a_test() { + fn reduce_min_max_of_x_a2_test() { use rand::Rng; - if !crate::is_cpu_detected!("v8.3a") { - println!("test {} ... skipped (v8.3a)", module_path!()); + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } let mut rng = rand::rng(); @@ -1177,8 +1177,8 @@ mod reduce_min_max_of_x { .collect::>(); for z in 50..200 { let x = &x[..z]; - let specialized = unsafe { reduce_min_max_of_x_v8_3a(x) }; - let fallback = reduce_min_max_of_x_fallback(x); + let specialized = unsafe { reduce_min_max_of_x_a2(x) }; + let fallback = fallback(x); assert_eq!(specialized.0, fallback.0,); assert_eq!(specialized.1, fallback.1,); } @@ -1187,31 +1187,31 @@ mod reduce_min_max_of_x { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "sve")] - fn reduce_min_max_of_x_v8_3a_sve(this: &[f32]) -> (f32, f32) { + fn reduce_min_max_of_x_a3_256(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = -f32::INFINITY; unsafe { extern "C" { - fn fp32_reduce_min_max_of_x_v8_3a_sve( + fn fp32_reduce_min_max_of_x_a3_256( this: *const f32, n: usize, out_min: &mut f32, out_max: &mut f32, ); } - fp32_reduce_min_max_of_x_v8_3a_sve(this.as_ptr(), this.len(), &mut min, &mut max); + fp32_reduce_min_max_of_x_a3_256(this.as_ptr(), this.len(), &mut min, &mut max); } (min, max) } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_min_max_of_x_v8_3a_sve_test() { + fn reduce_min_max_of_x_a3_256_test() { use rand::Rng; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { - println!("test {} ... skipped (v8.3a:sve)", module_path!()); + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); return; } let mut rng = rand::rng(); @@ -1222,15 +1222,15 @@ mod reduce_min_max_of_x { .collect::>(); for z in 50..200 { let x = &x[..z]; - let specialized = unsafe { reduce_min_max_of_x_v8_3a_sve(x) }; - let fallback = reduce_min_max_of_x_fallback(x); + let specialized = unsafe { reduce_min_max_of_x_a3_256(x) }; + let fallback = fallback(x); assert_eq!(specialized.0, fallback.0,); assert_eq!(specialized.1, fallback.1,); } } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a:sve", @"v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a3.256", @"a2")] pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -1246,8 +1246,8 @@ mod reduce_min_max_of_x { mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_xy_v4(lhs: &[f32], rhs: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_sum_of_xy_v4_512(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -1278,7 +1278,7 @@ mod reduce_sum_of_xy { fn reduce_sum_of_xy_v4_test() { use rand::Rng; const EPSILON: f32 = 0.004; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -1294,8 +1294,8 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_xy_v4_512(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1368,7 +1368,7 @@ mod reduce_sum_of_xy { let lhs = &lhs[..z]; let rhs = &rhs[..z]; let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1434,7 +1434,7 @@ mod reduce_sum_of_xy { let lhs = &lhs[..z]; let rhs = &rhs[..z]; let specialized = unsafe { reduce_sum_of_xy_v2_fma(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1445,8 +1445,8 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] - fn reduce_sum_of_xy_v8_3a(lhs: &[f32], rhs: &[f32]) -> f32 { + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_xy_a2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::aarch64::*; @@ -1478,11 +1478,11 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_xy_v8_3a_test() { + fn reduce_sum_of_xy_a2_test() { use rand::Rng; const EPSILON: f32 = 0.004; - if !crate::is_cpu_detected!("v8.3a") { - println!("test {} ... skipped (v8.3a)", module_path!()); + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } let mut rng = rand::rng(); @@ -1497,8 +1497,8 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v8_3a(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1509,25 +1509,25 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "sve")] - fn reduce_sum_of_xy_v8_3a_sve(lhs: &[f32], rhs: &[f32]) -> f32 { + fn reduce_sum_of_xy_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { extern "C" { - fn fp32_reduce_sum_of_xy_v8_3a_sve(a: *const f32, b: *const f32, n: usize) -> f32; + fn fp32_reduce_sum_of_xy_a3_256(a: *const f32, b: *const f32, n: usize) -> f32; } - fp32_reduce_sum_of_xy_v8_3a_sve(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) + fp32_reduce_sum_of_xy_a3_256(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) } } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_xy_v8_3a_sve_test() { + fn reduce_sum_of_xy_a3_256_test() { use rand::Rng; const EPSILON: f32 = 0.004; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { - println!("test {} ... skipped (v8.3a:sve)", module_path!()); + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); return; } let mut rng = rand::rng(); @@ -1542,8 +1542,8 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v8_3a_sve(lhs, rhs) }; - let fallback = reduce_sum_of_xy_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_xy_a3_256(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1552,7 +1552,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"v8.3a:sve", @"v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", @"v2:fma", @"a3.256", @"a2")] pub fn reduce_sum_of_xy(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -1567,8 +1567,8 @@ mod reduce_sum_of_xy { mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_d2_v4(lhs: &[f32], rhs: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_sum_of_d2_v4_512(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::x86_64::*; @@ -1601,7 +1601,7 @@ mod reduce_sum_of_d2 { fn reduce_sum_of_d2_v4_test() { use rand::Rng; const EPSILON: f32 = 0.02; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -1617,8 +1617,8 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v4(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_d2_v4_512(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1694,7 +1694,7 @@ mod reduce_sum_of_d2 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; let specialized = unsafe { reduce_sum_of_d2_v3(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1762,7 +1762,7 @@ mod reduce_sum_of_d2 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; let specialized = unsafe { reduce_sum_of_d2_v2_fma(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1773,8 +1773,8 @@ mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] - fn reduce_sum_of_d2_v8_3a(lhs: &[f32], rhs: &[f32]) -> f32 { + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_d2_a2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { use std::arch::aarch64::*; @@ -1808,11 +1808,11 @@ mod reduce_sum_of_d2 { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_d2_v8_3a_test() { + fn reduce_sum_of_d2_a2_test() { use rand::Rng; const EPSILON: f32 = 0.02; - if !crate::is_cpu_detected!("v8.3a") { - println!("test {} ... skipped (v8.3a)", module_path!()); + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } let mut rng = rand::rng(); @@ -1827,8 +1827,8 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v8_3a(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_d2_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1839,25 +1839,25 @@ mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] + #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "sve")] - fn reduce_sum_of_d2_v8_3a_sve(lhs: &[f32], rhs: &[f32]) -> f32 { + fn reduce_sum_of_d2_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { extern "C" { - fn fp32_reduce_sum_of_d2_v8_3a_sve(a: *const f32, b: *const f32, n: usize) -> f32; + fn fp32_reduce_sum_of_d2_a3_256(a: *const f32, b: *const f32, n: usize) -> f32; } - fp32_reduce_sum_of_d2_v8_3a_sve(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) + fp32_reduce_sum_of_d2_a3_256(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) } } #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_d2_v8_3a_sve_test() { + fn reduce_sum_of_d2_a3_256_test() { use rand::Rng; const EPSILON: f32 = 0.02; - if !crate::is_cpu_detected!("v8.3a") || !crate::is_feature_detected!("sve") { - println!("test {} ... skipped (v8.3a:sve)", module_path!()); + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); return; } let mut rng = rand::rng(); @@ -1872,8 +1872,8 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v8_3a_sve(lhs, rhs) }; - let fallback = reduce_sum_of_d2_fallback(lhs, rhs); + let specialized = unsafe { reduce_sum_of_d2_a3_256(lhs, rhs) }; + let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1882,7 +1882,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"v8.3a:sve", @"v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", @"v2:fma", @"a3.256", @"a2")] pub fn reduce_sum_of_d2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -1898,8 +1898,8 @@ mod reduce_sum_of_d2 { mod reduce_sum_of_xy_sparse { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_xy_sparse_v4(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_sum_of_xy_sparse_v4_512(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { use crate::emulate::emulate_mm512_2intersect_epi32; assert_eq!(li.len(), lv.len()); assert_eq!(ri.len(), rv.len()); @@ -1947,7 +1947,7 @@ mod reduce_sum_of_xy_sparse { fn reduce_sum_of_xy_sparse_v4_test() { use rand::Rng; const EPSILON: f32 = 0.000001; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -1975,8 +1975,8 @@ mod reduce_sum_of_xy_sparse { let rval = (0..rm) .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); - let specialized = unsafe { reduce_sum_of_xy_sparse_v4(&lidx, &lval, &ridx, &rval) }; - let fallback = reduce_sum_of_xy_sparse_fallback(&lidx, &lval, &ridx, &rval); + let specialized = unsafe { reduce_sum_of_xy_sparse_v4_512(&lidx, &lval, &ridx, &rval) }; + let fallback = fallback(&lidx, &lval, &ridx, &rval); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -1984,7 +1984,7 @@ mod reduce_sum_of_xy_sparse { } } - #[crate::multiversion(@"v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion(@"v4.512", "v3", "v2", "a2")] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -2014,8 +2014,8 @@ mod reduce_sum_of_xy_sparse { mod reduce_sum_of_d2_sparse { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_d2_sparse_v4(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_sum_of_d2_sparse_v4_512(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { use crate::emulate::emulate_mm512_2intersect_epi32; assert_eq!(li.len(), lv.len()); assert_eq!(ri.len(), rv.len()); @@ -2097,7 +2097,7 @@ mod reduce_sum_of_d2_sparse { fn reduce_sum_of_d2_sparse_v4_test() { use rand::Rng; const EPSILON: f32 = 0.0004; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -2125,8 +2125,8 @@ mod reduce_sum_of_d2_sparse { let rval = (0..rm) .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); - let specialized = unsafe { reduce_sum_of_d2_sparse_v4(&lidx, &lval, &ridx, &rval) }; - let fallback = reduce_sum_of_d2_sparse_fallback(&lidx, &lval, &ridx, &rval); + let specialized = unsafe { reduce_sum_of_d2_sparse_v4_512(&lidx, &lval, &ridx, &rval) }; + let fallback = fallback(&lidx, &lval, &ridx, &rval); assert!( (specialized - fallback).abs() < EPSILON, "specialized = {specialized}, fallback = {fallback}." @@ -2134,7 +2134,7 @@ mod reduce_sum_of_d2_sparse { } } - #[crate::multiversion(@"v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion(@"v4.512", "v3", "v2", "a2")] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -2171,7 +2171,7 @@ mod reduce_sum_of_d2_sparse { } mod vector_add { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2189,7 +2189,7 @@ mod vector_add { } mod vector_add_inplace { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2200,7 +2200,7 @@ mod vector_add_inplace { } mod vector_sub { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2218,7 +2218,7 @@ mod vector_sub { } mod vector_mul { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2236,7 +2236,7 @@ mod vector_mul { } mod vector_mul_scalar { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { let n = lhs.len(); let mut r = Vec::::with_capacity(n); @@ -2253,7 +2253,7 @@ mod vector_mul_scalar { } mod vector_mul_scalar_inplace { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { let n = lhs.len(); for i in 0..n { @@ -2263,7 +2263,7 @@ mod vector_mul_scalar_inplace { } mod vector_abs_inplace { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn vector_abs_inplace(this: &mut [f32]) { let n = this.len(); for i in 0..n { diff --git a/crates/simd/src/fast_scan/mod.rs b/crates/simd/src/fast_scan/mod.rs index 44cf2fc4..a99c227f 100644 --- a/crates/simd/src/fast_scan/mod.rs +++ b/crates/simd/src/fast_scan/mod.rs @@ -125,8 +125,8 @@ pub fn any_pack(mut x: impl Iterator) -> [T; 32] { mod fast_scan { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn fast_scan_v4(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + #[crate::target_cpu(enable = "v4.512")] + fn fast_scan_v4_512(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -135,7 +135,7 @@ mod fast_scan { use std::arch::x86_64::*; #[inline] - #[crate::target_cpu(enable = "v4")] + #[crate::target_cpu(enable = "v4.512")] fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { unsafe { let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); @@ -145,7 +145,7 @@ mod fast_scan { } #[inline] - #[crate::target_cpu(enable = "v4")] + #[crate::target_cpu(enable = "v4.512")] fn combine4x2(x0x1x2x3: __m512i, y0y1y2y3: __m512i) -> __m256i { unsafe { let x0x1 = _mm512_castsi512_si256(x0x1x2x3); @@ -238,7 +238,7 @@ mod fast_scan { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] fn fast_scan_v4_test() { - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -251,7 +251,7 @@ mod fast_scan { .map(|_| std::array::from_fn(|_| rand::random())) .collect::>(); unsafe { - assert_eq!(fast_scan_v4(&code, &lut), fast_scan_fallback(&code, &lut)); + assert_eq!(fast_scan_v4_512(&code, &lut), fallback(&code, &lut)); } } } @@ -354,7 +354,7 @@ mod fast_scan { .map(|_| std::array::from_fn(|_| rand::random())) .collect::>(); unsafe { - assert_eq!(fast_scan_v3(&code, &lut), fast_scan_fallback(&code, &lut)); + assert_eq!(fast_scan_v3(&code, &lut), fallback(&code, &lut)); } } } @@ -425,15 +425,15 @@ mod fast_scan { .map(|_| std::array::from_fn(|_| rand::random())) .collect::>(); unsafe { - assert_eq!(fast_scan_v2(&code, &lut), fast_scan_fallback(&code, &lut)); + assert_eq!(fast_scan_v2(&code, &lut), fallback(&code, &lut)); } } } } #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] - fn fast_scan_v8_3a(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + #[crate::target_cpu(enable = "a2")] + fn fast_scan_a2(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -481,9 +481,9 @@ mod fast_scan { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn fast_scan_v8_3a_test() { - if !crate::is_cpu_detected!("v8.3a") { - println!("test {} ... skipped (v8.3a)", module_path!()); + fn fast_scan_a2_test() { + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { @@ -495,16 +495,13 @@ mod fast_scan { .map(|_| std::array::from_fn(|_| rand::random())) .collect::>(); unsafe { - assert_eq!( - fast_scan_v8_3a(&code, &lut), - fast_scan_fallback(&code, &lut) - ); + assert_eq!(fast_scan_a2(&code, &lut), fallback(&code, &lut)); } } } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a2")] pub fn fast_scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { fn binary(op: impl Fn(u16, u16) -> u16, a: [u16; 8], b: [u16; 8]) -> [u16; 8] { std::array::from_fn(|i| op(a[i], b[i])) diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index fd132347..5fbcd9eb 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -68,6 +68,74 @@ mod internal { #[cfg(target_arch = "riscv64")] #[allow(unused_imports)] pub use is_riscv64_cpu_detected; + + #[cfg(target_arch = "x86_64")] + pub fn is_v4_512_detected() -> bool { + std::arch::is_x86_feature_detected!("avx512bw") + && std::arch::is_x86_feature_detected!("avx512cd") + && std::arch::is_x86_feature_detected!("avx512dq") + && std::arch::is_x86_feature_detected!("avx512vl") + && std::arch::is_x86_feature_detected!("bmi1") + && std::arch::is_x86_feature_detected!("bmi2") + && std::arch::is_x86_feature_detected!("lzcnt") + && std::arch::is_x86_feature_detected!("movbe") + && std::arch::is_x86_feature_detected!("popcnt") + } + + #[cfg(target_arch = "x86_64")] + pub fn is_v3_detected() -> bool { + std::arch::is_x86_feature_detected!("avx2") + && std::arch::is_x86_feature_detected!("f16c") + && std::arch::is_x86_feature_detected!("fma") + && std::arch::is_x86_feature_detected!("bmi1") + && std::arch::is_x86_feature_detected!("bmi2") + && std::arch::is_x86_feature_detected!("lzcnt") + && std::arch::is_x86_feature_detected!("movbe") + && std::arch::is_x86_feature_detected!("popcnt") + } + + #[cfg(target_arch = "x86_64")] + pub fn is_v2_detected() -> bool { + std::arch::is_x86_feature_detected!("sse4.2") + && std::arch::is_x86_feature_detected!("popcnt") + } + + #[cfg(target_arch = "aarch64")] + pub fn is_a3_512_detected() -> bool { + #[target_feature(enable = "sve")] + fn is_512_detected() -> bool { + let vl: u64; + unsafe { + core::arch::asm!( + "rdvl {0}, #2", + out(reg) vl + ); + } + vl >= 4 + } + std::arch::is_aarch64_feature_detected!("sve") && unsafe { is_512_detected() } + } + + #[cfg(target_arch = "aarch64")] + pub fn is_a3_256_detected() -> bool { + #[target_feature(enable = "sve")] + fn is_256_detected() -> bool { + let vl: u64; + unsafe { + core::arch::asm!( + "rdvl {0}, #2", + out(reg) vl + ); + } + vl >= 2 + } + std::arch::is_aarch64_feature_detected!("sve") && unsafe { is_256_detected() } + } + + #[cfg(target_arch = "aarch64")] + pub fn is_a2_detected() -> bool { + std::arch::is_aarch64_feature_detected!("neon") + } } pub use simd_macros::{multiversion, target_cpu}; diff --git a/crates/simd/src/packed_u4.rs b/crates/simd/src/packed_u4.rs index d68333c7..ae32aebb 100644 --- a/crates/simd/src/packed_u4.rs +++ b/crates/simd/src/packed_u4.rs @@ -3,7 +3,7 @@ pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { } mod reduce_sum_of_xy { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index baaf594a..880f64bc 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -1,8 +1,8 @@ mod mul_add_round { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn mul_add_round_v4(this: &[f32], k: f32, b: f32) -> Vec { + #[crate::target_cpu(enable = "v4.512")] + fn mul_add_round_v4_512(this: &[f32], k: f32, b: f32) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); unsafe { @@ -42,7 +42,7 @@ mod mul_add_round { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] fn mul_add_round_v4_test() { - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -53,8 +53,8 @@ mod mul_add_round { let x = &x[..z]; let k = 20.0; let b = 20.0; - let specialized = unsafe { mul_add_round_v4(x, k, b) }; - let fallback = mul_add_round_fallback(x, k, b); + let specialized = unsafe { mul_add_round_v4_512(x, k, b) }; + let fallback = fallback(x, k, b); assert_eq!(specialized, fallback); } } @@ -123,7 +123,7 @@ mod mul_add_round { let k = 20.0; let b = 20.0; let specialized = unsafe { mul_add_round_v3(x, k, b) }; - let fallback = mul_add_round_fallback(x, k, b); + let fallback = fallback(x, k, b); assert_eq!(specialized, fallback); } } @@ -189,7 +189,7 @@ mod mul_add_round { let k = 20.0; let b = 20.0; let specialized = unsafe { mul_add_round_v2_fma(x, k, b) }; - let fallback = mul_add_round_fallback(x, k, b); + let fallback = fallback(x, k, b); assert_eq!(specialized, fallback); } } @@ -197,8 +197,8 @@ mod mul_add_round { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] - fn mul_add_round_v8_3a(this: &[f32], k: f32, b: f32) -> Vec { + #[crate::target_cpu(enable = "a2")] + fn mul_add_round_a2(this: &[f32], k: f32, b: f32) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); unsafe { @@ -244,9 +244,9 @@ mod mul_add_round { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn mul_add_round_v8_3a_test() { - if !crate::is_cpu_detected!("v8.3a") { - println!("test {} ... skipped (v8.3a)", module_path!()); + fn mul_add_round_a2_test() { + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { @@ -256,14 +256,14 @@ mod mul_add_round { let x = &x[..z]; let k = 20.0; let b = 20.0; - let specialized = unsafe { mul_add_round_v8_3a(x, k, b) }; - let fallback = mul_add_round_fallback(x, k, b); + let specialized = unsafe { mul_add_round_a2(x, k, b) }; + let fallback = fallback(x, k, b); assert_eq!(specialized, fallback); } } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", @"v2:fma", @"a2")] pub fn mul_add_round(this: &[f32], k: f32, b: f32) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index 4e0855f8..b7f50e93 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -1,5 +1,5 @@ mod reduce_sum_of_xy { - #[crate::multiversion("v4", "v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion("v4.512", "v3", "v2", "a2")] pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); @@ -19,8 +19,8 @@ pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { mod reduce_sum_of_x_as_u16 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_x_as_u16_v4(this: &[u8]) -> u16 { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_sum_of_x_as_u16_v4_512(this: &[u8]) -> u16 { use crate::emulate::emulate_mm512_reduce_add_epi16; unsafe { use std::arch::x86_64::*; @@ -47,7 +47,7 @@ mod reduce_sum_of_x_as_u16 { #[test] fn reduce_sum_of_x_as_u16_v4_test() { use rand::Rng; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -57,8 +57,8 @@ mod reduce_sum_of_x_as_u16 { let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u16_v4(this) }; - let fallback = reduce_sum_of_x_as_u16_fallback(this); + let specialized = unsafe { reduce_sum_of_x_as_u16_v4_512(this) }; + let fallback = fallback(this); assert_eq!(specialized, fallback); } } @@ -108,7 +108,7 @@ mod reduce_sum_of_x_as_u16 { for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_as_u16_v3(this) }; - let fallback = reduce_sum_of_x_as_u16_fallback(this); + let fallback = fallback(this); assert_eq!(specialized, fallback); } } @@ -158,7 +158,7 @@ mod reduce_sum_of_x_as_u16 { for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_as_u16_v2(this) }; - let fallback = reduce_sum_of_x_as_u16_fallback(this); + let fallback = fallback(this); assert_eq!(specialized, fallback); } } @@ -166,8 +166,8 @@ mod reduce_sum_of_x_as_u16 { #[inline] #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "v8.3a")] - fn reduce_sum_of_x_as_u16_v8_3a(this: &[u8]) -> u16 { + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x_as_u16_a2(this: &[u8]) -> u16 { unsafe { use std::arch::aarch64::*; let us = vdupq_n_u16(255); @@ -194,10 +194,10 @@ mod reduce_sum_of_x_as_u16 { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_x_as_u16_v8_3a_test() { + fn reduce_sum_of_x_as_u16_a2_test() { use rand::Rng; - if !crate::is_cpu_detected!("v8.3a") { - println!("test {} ... skipped (v8.3a)", module_path!()); + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } let mut rng = rand::rng(); @@ -206,14 +206,14 @@ mod reduce_sum_of_x_as_u16 { let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u16_v8_3a(this) }; - let fallback = reduce_sum_of_x_as_u16_fallback(this); + let specialized = unsafe { reduce_sum_of_x_as_u16_a2(this) }; + let fallback = fallback(this); assert_eq!(specialized, fallback); } } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a2")] pub fn reduce_sum_of_x_as_u16(this: &[u8]) -> u16 { let n = this.len(); let mut sum = 0; @@ -232,8 +232,8 @@ pub fn reduce_sum_of_x_as_u16(vector: &[u8]) -> u16 { mod reduce_sum_of_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_x_v4(this: &[u8]) -> u32 { + #[crate::target_cpu(enable = "v4.512")] + fn reduce_sum_of_x_v4_512(this: &[u8]) -> u32 { unsafe { use std::arch::x86_64::*; let us = _mm512_set1_epi32(255); @@ -259,7 +259,7 @@ mod reduce_sum_of_x { #[test] fn reduce_sum_of_x_v4_test() { use rand::Rng; - if !crate::is_cpu_detected!("v4") { + if !crate::is_cpu_detected!("v4.512") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -269,8 +269,8 @@ mod reduce_sum_of_x { let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_v4(this) }; - let fallback = reduce_sum_of_x_fallback(this); + let specialized = unsafe { reduce_sum_of_x_v4_512(this) }; + let fallback = fallback(this); assert_eq!(specialized, fallback); } } @@ -320,13 +320,13 @@ mod reduce_sum_of_x { for z in 3984..4016 { let this = &this[..z]; let specialized = unsafe { reduce_sum_of_x_v3(this) }; - let fallback = reduce_sum_of_x_fallback(this); + let fallback = fallback(this); assert_eq!(specialized, fallback); } } } - #[crate::multiversion(@"v4", @"v3", "v2", "v8.3a:sve", "v8.3a")] + #[crate::multiversion(@"v4.512", @"v3", "v2", "a2")] pub fn reduce_sum_of_x(this: &[u8]) -> u32 { let n = this.len(); let mut sum = 0; diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs index fe455646..68a5b8c6 100644 --- a/crates/simd_macros/src/lib.rs +++ b/crates/simd_macros/src/lib.rs @@ -118,11 +118,9 @@ pub fn multiversion( } }); } - let fallback_name = - syn::Ident::new(&format!("{name}_fallback"), proc_macro2::Span::mixed_site()); quote::quote! { #versions - fn #fallback_name < #generics_params > (#inputs) #output #generics_where { #block } + fn fallback < #generics_params > (#inputs) #output #generics_where { #block } #[inline(always)] #(#attrs)* #vis #sig { static CACHE: core::sync::atomic::AtomicPtr<()> = core::sync::atomic::AtomicPtr::new(core::ptr::null_mut()); @@ -132,7 +130,7 @@ pub fn multiversion( return unsafe { f(#(#arguments,)*) }; } #branches - let _multiversion_internal: unsafe fn(#inputs) #output = #fallback_name; + let _multiversion_internal: unsafe fn(#inputs) #output = fallback; CACHE.store(_multiversion_internal as *mut (), core::sync::atomic::Ordering::Relaxed); unsafe { _multiversion_internal(#(#arguments,)*) } } @@ -184,21 +182,22 @@ pub fn define_is_cpu_detected(input: proc_macro::TokenStream) -> proc_macro::Tok if target_cpu.target_arch != target_arch { continue; } - let target_features = target_cpu.target_features; let target_cpu = target_cpu.target_cpu; + let ident = syn::Ident::new( + &format!("is_{}_detected", target_cpu.replace('.', "_")), + proc_macro2::Span::mixed_site(), + ); arms.extend(quote::quote! { - (#target_cpu) => { - true #(&& $crate::is_feature_detected!(#target_features))* - }; + (#target_cpu) => { $crate::internal::#ident() }; }); } - let name = syn::Ident::new( + let ident = syn::Ident::new( &format!("is_{target_arch}_cpu_detected"), proc_macro2::Span::mixed_site(), ); quote::quote! { #[macro_export] - macro_rules! #name { + macro_rules! #ident { #arms } } diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs index c2584013..e2be8534 100644 --- a/crates/simd_macros/src/target.rs +++ b/crates/simd_macros/src/target.rs @@ -6,78 +6,48 @@ pub struct TargetCpu { pub const TARGET_CPUS: &[TargetCpu] = &[ TargetCpu { - target_cpu: "v4", + target_cpu: "v4.512", target_arch: "x86_64", target_features: &[ - "avx", - "avx2", - "avx512bw", - "avx512cd", - "avx512dq", - "avx512f", - "avx512vl", - "bmi1", - "bmi2", - "cmpxchg16b", - "f16c", - "fma", - "fxsr", - "lzcnt", - "movbe", - "popcnt", - "sse", - "sse2", - "sse3", - "sse4.1", - "sse4.2", - "ssse3", - "xsave", + "avx512bw", "avx512cd", "avx512dq", "avx512vl", // simd + "bmi1", "bmi2", "lzcnt", "movbe", "popcnt", // bit-operations ], }, TargetCpu { target_cpu: "v3", target_arch: "x86_64", target_features: &[ - "avx", - "avx2", - "bmi1", - "bmi2", - "cmpxchg16b", - "f16c", - "fma", - "fxsr", - "lzcnt", - "movbe", - "popcnt", - "sse", - "sse2", - "sse3", - "sse4.1", - "sse4.2", - "ssse3", - "xsave", + "avx2", "f16c", "fma", // simd + "bmi1", "bmi2", "lzcnt", "movbe", "popcnt", // bit-operations ], }, TargetCpu { target_cpu: "v2", target_arch: "x86_64", target_features: &[ - "cmpxchg16b", - "fxsr", - "popcnt", - "sse", - "sse2", - "sse3", - "sse4.1", - "sse4.2", - "ssse3", + "sse4.2", // simd + "popcnt", // bit-operations ], }, TargetCpu { - target_cpu: "v8.3a", + target_cpu: "a3.512", target_arch: "aarch64", target_features: &[ - "crc", "dpb", "fcma", "jsconv", "lse", "neon", "paca", "pacg", "rcpc", "rdm", + "sve", // simd + ], + }, + TargetCpu { + target_cpu: "a3.256", + target_arch: "aarch64", + target_features: &[ + "sve", // simd + ], + }, + TargetCpu { + target_cpu: "a2", + target_arch: "aarch64", + target_features: &[ + "neon", // simd ], }, ]; diff --git a/src/index/functions.rs b/src/index/functions.rs index be2f9633..64acde82 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -19,7 +19,7 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { if pg_class.relam() != pg_am.oid() { pgrx::error!("{:?} is not a vchordrq index", pg_class.relname()); } - let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::ShareLock as _) }; + let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::AccessShareLock as _) }; let relation = unsafe { PostgresRelation::new(index) }; let opfamily = unsafe { crate::index::opclass::opfamily(index) }; let message = match (opfamily.vector_kind(), opfamily.distance_kind()) { @@ -45,7 +45,7 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { } }; unsafe { - pgrx::pg_sys::index_close(index, pgrx::pg_sys::ShareLock as _); + pgrx::pg_sys::index_close(index, pgrx::pg_sys::AccessShareLock as _); } message } From a9a04217845006ab46e8341c7d97c571253978f9 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 19 Feb 2025 15:23:43 +0800 Subject: [PATCH 113/324] fix: test simd operations in emulator (#193) 1. x86_64 runner may not have avx512fp16 feature, so intel SDE is necessary for testing 2. aarch64 runner has short SVE bits, and cannot be used for testing SVE-256 implementation and SVE-512 implementation, so qemu is necessary for testing Signed-off-by: usamoi --- .github/workflows/check.yml | 33 ++++++++++++++++++++++++++++++++- crates/simd/src/lib.rs | 8 ++++---- 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 1445349e..40b173c5 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -74,6 +74,20 @@ jobs: runs-on: ${{ matrix.runner }} steps: + - name: Set up Environment + run: | + sudo apt-get update + + if [ "$(uname -m)" == "x86_64" ]; then + wget https://downloadmirror.intel.com/843185/sde-external-9.48.0-2024-11-25-lin.tar.xz -O /tmp/sde-external.tar.xz + sudo tar -xf /tmp/sde-external.tar.xz -C /opt + sudo mv /opt/sde-external-9.48.0-2024-11-25-lin /opt/sde + fi + + if [ "$(uname -m)" == "aarch64" ]; then + sudo apt-get install -y qemu-user-static + fi + - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.7 @@ -84,7 +98,22 @@ jobs: run: cargo clippy --workspace --exclude vchord - name: Cargo Test - run: cargo test --workspace --exclude vchord --no-fail-fast + run: cargo test --workspace --exclude vchord --exclude simd --no-fail-fast + + - name: Cargo Test (simd) + run: | + cargo test -p simd --release -- --nocapture + + if [ "$(uname -m)" == "x86_64" ]; then + cargo \ + --config 'target.'\''cfg(all())'\''.runner = ["/opt/sde/sde64", "-spr", "--"]' \ + test -p simd --release -- --nocapture + fi + if [ "$(uname -m)" == "aarch64" ]; then + cargo \ + --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max"]' \ + test -p simd --release -- --nocapture + fi psql: strategy: @@ -96,6 +125,8 @@ jobs: steps: - name: Set up Environment run: | + sudo apt-get update + sudo apt-get remove -y '^postgres.*' '^libpq.*' sudo apt-get purge -y '^postgres.*' '^libpq.*' diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 5fbcd9eb..e2ce8304 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -107,11 +107,11 @@ mod internal { let vl: u64; unsafe { core::arch::asm!( - "rdvl {0}, #2", + "rdvl {0}, #8", out(reg) vl ); } - vl >= 4 + vl >= 512 } std::arch::is_aarch64_feature_detected!("sve") && unsafe { is_512_detected() } } @@ -123,11 +123,11 @@ mod internal { let vl: u64; unsafe { core::arch::asm!( - "rdvl {0}, #2", + "rdvl {0}, #8", out(reg) vl ); } - vl >= 2 + vl >= 256 } std::arch::is_aarch64_feature_detected!("sve") && unsafe { is_256_detected() } } From eca11f3e490089b90df1f7acd5176bd744cce709 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 19 Feb 2025 17:26:24 +0800 Subject: [PATCH 114/324] feat: neon impl of u8::reduce_sum_of_x (#194) Signed-off-by: usamoi --- crates/simd/src/emulate.rs | 1 - crates/simd/src/u8.rs | 104 ++++++++++++++++++++++++++++++++++++- 2 files changed, 102 insertions(+), 3 deletions(-) diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index 2715d860..293aa746 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -134,7 +134,6 @@ pub fn emulate_mm256_reduce_add_epi32(mut x: std::arch::x86_64::__m256i) -> i32 } } -#[expect(dead_code)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index b7f50e93..4c8e238e 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -307,7 +307,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u16_v3_test() { + fn reduce_sum_of_x_v3_test() { use rand::Rng; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -326,7 +326,107 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4.512", @"v3", "v2", "a2")] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_x_v2(this: &[u8]) -> u32 { + use crate::emulate::emulate_mm_reduce_add_epi32; + unsafe { + use std::arch::x86_64::*; + let us = _mm_set1_epi32(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_si128(); + while n >= 4 { + let x = _mm_cvtsi32_si128(a.cast::().read_unaligned()); + a = a.add(4); + n -= 4; + sum = _mm_add_epi32(_mm_and_si128(us, _mm_cvtepi8_epi32(x)), sum); + } + let mut sum = emulate_mm_reduce_add_epi32(sum) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + sum += x as u32; + } + sum + } + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v2_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v2(this) }; + let fallback = fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x_a2(this: &[u8]) -> u32 { + unsafe { + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum_0 = vdupq_n_u32(0); + let mut sum_1 = vdupq_n_u32(0); + while n >= 8 { + let x = vmovl_u8(vld1_u8(a.cast())); + a = a.add(8); + n -= 8; + sum_0 = vaddq_u32(vmovl_u16(vget_low_u16(x)), sum_0); + sum_1 = vaddq_u32(vmovl_u16(vget_high_u16(x)), sum_1); + } + let mut sum = vaddvq_u32(vaddq_u32(sum_0, sum_1)); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = a.read(); + a = a.add(1); + n -= 1; + sum += x as u32; + } + sum + } + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + fn reduce_sum_of_x_a2_test() { + use rand::Rng; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_a2(this) }; + let fallback = fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a2")] pub fn reduce_sum_of_x(this: &[u8]) -> u32 { let n = this.len(); let mut sum = 0; From d37e31f1b5fcd9db96c96a01526acd73c84ebea0 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 19 Feb 2025 18:32:19 +0800 Subject: [PATCH 115/324] chore: update 0.2.1 schema (upgrade) script (#195) Signed-off-by: usamoi --- .github/workflows/pgrx.yml | 53 --- .github/workflows/release.yml | 115 +++--- docker/pgrx.Dockerfile | 64 ---- sql/install/vchord--0.2.1.sql | 500 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.2.0--0.2.1.sql | 1 + tools/package.sh | 45 --- 6 files changed, 575 insertions(+), 203 deletions(-) delete mode 100644 .github/workflows/pgrx.yml delete mode 100644 docker/pgrx.Dockerfile create mode 100644 sql/install/vchord--0.2.1.sql create mode 100644 sql/upgrade/vchord--0.2.0--0.2.1.sql delete mode 100755 tools/package.sh diff --git a/.github/workflows/pgrx.yml b/.github/workflows/pgrx.yml deleted file mode 100644 index 9bd8489b..00000000 --- a/.github/workflows/pgrx.yml +++ /dev/null @@ -1,53 +0,0 @@ -name: Build pgrx Image - -on: - workflow_dispatch: - inputs: - version: - description: 'pgrx version' - required: true - type: string - toolchain: - description: 'additional rust toolchain' - required: true - type: string - -concurrency: - group: ${{ github.ref }}-${{ github.workflow }} - cancel-in-progress: true - -env: - IMAGE_NAME: "ghcr.io/tensorchord/vectorchord-pgrx" - -permissions: - contents: write - packages: write - -jobs: - build: - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v4 - - name: Set up QEMU - uses: docker/setup-qemu-action@v3 - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 - - name: Login to ghcr.io - uses: docker/login-action@v3 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} - - name: Build and push - uses: docker/build-push-action@v6 - with: - context: . - file: ./docker/pgrx.Dockerfile - push: true - cache-from: type=gha - cache-to: type=gha,mode=max - platforms: "linux/amd64,linux/arm64" - build-args: | - PGRX_VERSION=${{ github.event.inputs.version }} - RUST_TOOLCHAIN=${{ github.event.inputs.toolchain }} - tags: ${{ env.IMAGE_NAME }}:${{ github.event.inputs.version }}-${{ github.event.inputs.toolchain }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 1662bbd0..c057afad 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -6,7 +6,7 @@ on: workflow_dispatch: inputs: tag: - description: 'tag name (semver without v-prefix)' + description: "tag name (semver without v-prefix)" required: true type: string @@ -16,12 +16,12 @@ concurrency: jobs: semver: - runs-on: ubuntu-latest - outputs: - SEMVER: ${{ steps.semver.outputs.SEMVER }} + runs-on: "ubuntu-latest" + steps: - - uses: actions/github-script@v7 + - name: Semver id: semver + uses: actions/github-script@v7 with: script: | const tag = "${{ github.event.inputs.tag }}" || "${{ github.event.release.tag_name }}"; @@ -32,58 +32,92 @@ jobs: } core.setOutput('SEMVER', tag); + outputs: + SEMVER: ${{ steps.semver.outputs.SEMVER }} + build: - runs-on: ${{ matrix.runner }} needs: ["semver"] strategy: matrix: - version: ["14", "15", "16", "17"] + version: ["13", "14", "15", "16", "17"] runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] + runs-on: ${{ matrix.runner }} + env: - PGRX_IMAGE: "ghcr.io/tensorchord/vectorchord-pgrx:0.12.9-nightly-2024-12-25" - SEMVER: ${{ needs.semver.outputs.SEMVER }} - ARCH: ${{ matrix.runner == 'ubuntu-22.04' && 'x86_64' || 'aarch64' }} - PLATFORM: ${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 + RUSTFLAGS: "-Dwarnings" steps: - name: Set up Environment run: | - sudo apt-get remove -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' - sudo apt-get purge -y '^postgres.*' '^libpq.*' '^clang.*' '^llvm.*' '^libclang.*' '^libllvm.*' '^mono-llvm.*' + sudo apt-get update + + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' + + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 + sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 sudo apt-get install -y postgresql-common sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y - sudo apt-get install -y postgresql-client-17 - - uses: actions/checkout@v4 - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - const url = process.env.ACTIONS_CACHE_URL || ''; - const token = process.env.ACTIONS_RUNTIME_TOKEN || ''; - core.exportVariable( - 'CACHE_ENVS', - `-e CARGO_INCREMENTAL=0 -e SCCACHE_GHA_ENABLED=true -e RUSTC_WRAPPER=sccache -e ACTIONS_CACHE_URL=${url} -e ACTIONS_RUNTIME_TOKEN=${token}`, - ); - - name: Set up pgrx docker images and permissions - run: | - docker pull $PGRX_IMAGE - echo "Default user: $(id -u):$(id -g)" - sudo chmod -R 777 . + sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} + + sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector + + curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.12.9/cargo-pgrx-v0.12.9-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx + cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) + + - name: Checkout + uses: actions/checkout@v4 - name: Build env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + VERSION: ${{ matrix.version }} + ARCH: ${{ matrix.runner == 'ubuntu-22.04' && 'x86_64' || 'aarch64' }} + PLATFORM: ${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} GH_TOKEN: ${{ github.token }} run: | - docker run --rm -v .:/workspace $CACHE_ENVS \ - -e SEMVER=$SEMVER \ - -e VERSION=${{ matrix.version }} \ - -e ARCH=$ARCH \ - -e PLATFORM=$PLATFORM \ - $PGRX_IMAGE ./tools/package.sh + cargo build --lib --features pg$VERSION --release + + mkdir -p ./build/zip + cp -a ./sql/upgrade/. ./build/zip/ + cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql + sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control + cp ./target/release/libvchord.so ./build/zip/vchord.so + zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* + + mkdir -p ./build/deb + mkdir -p ./build/deb/DEBIAN + mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ + mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ + for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file + done + for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file + done + for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file + done + echo "Package: postgresql-${VERSION}-vchord + Version: ${SEMVER}-1 + Section: database + Priority: optional + Architecture: ${PLATFORM} + Maintainer: Tensorchord + Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM + Homepage: https://vectorchord.ai/ + License: AGPL-3 or Elastic-2" \ + > ./build/deb/DEBIAN/control + (cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums + dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb + ls ./build - gh release upload --clobber $SEMVER ./build/postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${PLATFORM}.deb - gh release upload --clobber $SEMVER ./build/postgresql-${{ matrix.version }}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip + + gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb + gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip docker: runs-on: ubuntu-latest @@ -156,7 +190,7 @@ jobs: TARGETARCH=amd64 PGVECTOR=0.8.0 tags: modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }} - + test: name: Run tests runs-on: @@ -190,5 +224,4 @@ jobs: sleep 5 curl https://registry.pgtrunk.io/extensions/all | jq -r ".[] | .name" > /tmp/extensions.txt trunk-install.sh | tee /tmp/output.txt - cat /tmp/output.txt - + cat /tmp/output.txt diff --git a/docker/pgrx.Dockerfile b/docker/pgrx.Dockerfile deleted file mode 100644 index d9f8f50c..00000000 --- a/docker/pgrx.Dockerfile +++ /dev/null @@ -1,64 +0,0 @@ -# CNPG only support Debian 12 (Bookworm) -FROM ubuntu:22.04 - -ARG PGRX_VERSION -ARG RUST_TOOLCHAIN -ARG TARGETARCH - -ENV DEBIAN_FRONTEND=noninteractive \ - LANG=en_US.UTF-8 \ - LC_ALL=en_US.UTF-8 \ - RUSTFLAGS="-Dwarnings" \ - RUST_BACKTRACE=1 \ - CARGO_TERM_COLOR=always \ - SCCACHE_VERSION=0.9.0 - -RUN set -eux; \ - apt update; \ - apt install -y --no-install-recommends \ - curl \ - ca-certificates \ - build-essential \ - postgresql-common gnupg \ - libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config \ - zip - -RUN set -eux; \ - apt -y install lsb-release wget software-properties-common gnupg; \ - curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | bash -s -- 18; \ - update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 - -# set up sccache -RUN set -ex; \ - curl -fsSL -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v${SCCACHE_VERSION}/sccache-v${SCCACHE_VERSION}-$(uname -m)-unknown-linux-musl.tar.gz; \ - tar -xzf sccache.tar.gz --strip-components=1; \ - rm sccache.tar.gz; \ - mv sccache /usr/local/bin/ - -RUN /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y -# install all the PostgresQL -RUN set -ex; \ - for v in $(seq 13 17); do \ - apt install -y --no-install-recommends postgresql-$v postgresql-server-dev-$v postgresql-$v-pgvector; \ - done; \ - rm -rf /var/lib/apt/lists/*; - -# create a non-root user (make it compatible with Ubuntu 24.04) -RUN useradd -u 1000 -U -m ubuntu -RUN chown -R ubuntu:ubuntu /usr/share/postgresql/ /usr/lib/postgresql/ -USER ubuntu -ENV PATH="$PATH:/home/ubuntu/.cargo/bin" -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y - -WORKDIR /workspace -RUN rustup toolchain install ${RUST_TOOLCHAIN} -RUN rustup target add $(uname -m)-unknown-linux-gnu - -RUN cargo install cargo-pgrx --locked --version=${PGRX_VERSION} - -RUN set -ex; \ - for v in $(seq 13 17); do \ - cargo pgrx init --pg$v=/usr/lib/postgresql/$v/bin/pg_config; \ - done; - -CMD [ "/bin/bash" ] diff --git a/sql/install/vchord--0.2.1.sql b/sql/install/vchord--0.2.1.sql new file mode 100644 index 00000000..e5cf547b --- /dev/null +++ b/sql/install/vchord--0.2.1.sql @@ -0,0 +1,500 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:10 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:54 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:30 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:6 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:54 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:30 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:6 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:60 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:9 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:36 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:31 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:11 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.2.0--0.2.1.sql b/sql/upgrade/vchord--0.2.0--0.2.1.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.2.0--0.2.1.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/tools/package.sh b/tools/package.sh deleted file mode 100755 index e8c92bf1..00000000 --- a/tools/package.sh +++ /dev/null @@ -1,45 +0,0 @@ -#!/usr/bin/env bash -set -eu - -printf "SEMVER = ${SEMVER}\n" -printf "VERSION = ${VERSION}\n" -printf "ARCH = ${ARCH}\n" -printf "PLATFORM = ${PLATFORM}\n" - -cargo build --lib --features pg$VERSION --release -cargo pgrx schema --features pg$VERSION --out ./target/schema.sql - -rm -rf ./build - -mkdir -p ./build/zip -cp -a ./sql/upgrade/. ./build/zip/ -cp ./target/schema.sql ./build/zip/vchord--$SEMVER.sql -sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control -cp ./target/release/libvchord.so ./build/zip/vchord.so -zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* - -mkdir -p ./build/deb -mkdir -p ./build/deb/DEBIAN -mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ -mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ -for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file -done -for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file -done -for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file -done -echo "Package: postgresql-${VERSION}-vchord -Version: ${SEMVER}-1 -Section: database -Priority: optional -Architecture: ${PLATFORM} -Maintainer: Tensorchord -Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM -Homepage: https://vectorchord.ai/ -License: AGPL-3 or Elastic-2" \ -> ./build/deb/DEBIAN/control -(cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums -dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb From 2c534eee86c0ba69c9cae4acefdb1394efba53c4 Mon Sep 17 00:00:00 2001 From: xieydd Date: Fri, 21 Feb 2025 09:37:01 +0800 Subject: [PATCH 116/324] Support linux/arm64 for vchord-cnpg docker image (#184) I have re-released vchord-cnpg:0.2.0 docker image in https://github.com/tensorchord/VectorChord/actions/runs/13234388732. --------- Signed-off-by: xieydd --- .github/workflows/release.yml | 101 +++++++++--- .github/workflows/release_pg_slim.yml | 35 +++- docker/pg-cnpg/Dockerfile | 40 ++--- docker/pg-cnpg/Dockerfile.arm | 70 ++++++++ ...{trunk-install.sh => extension-install.sh} | 2 +- docker/pg-slim/Dockerfile | 151 +++++++++--------- 6 files changed, 279 insertions(+), 120 deletions(-) create mode 100644 docker/pg-cnpg/Dockerfile.arm rename docker/pg-cnpg/{trunk-install.sh => extension-install.sh} (99%) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c057afad..c8b85688 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -120,13 +120,15 @@ jobs: gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip docker: - runs-on: ubuntu-latest + runs-on: ${{ matrix.runner }} needs: ["semver", "build"] strategy: matrix: version: ["14", "15", "16", "17"] + runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] env: SEMVER: ${{ needs.semver.outputs.SEMVER }} + PLATFORM: ${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} steps: - name: Checkout uses: actions/checkout@v4 @@ -135,9 +137,7 @@ jobs: GH_TOKEN: ${{ github.token }} run: | mkdir -p build - for arch in amd64 arm64; do - gh release download $SEMVER --pattern "postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${arch}.deb" --output ./build/postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${arch}.deb - done + gh release download $SEMVER --pattern "postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${PLATFORM}.deb" --output ./build/postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${PLATFORM}.deb - name: Set up QEMU uses: docker/setup-qemu-action@v3 - name: Set up Docker Buildx @@ -152,9 +152,10 @@ jobs: with: context: . push: true - platforms: "linux/amd64,linux/arm64" + platforms: ${{ matrix.runner == 'ubuntu-22.04' && 'linux/amd64' || 'linux/arm64' }} file: ./docker/binary.Dockerfile - tags: tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }} + provenance: false + tags: tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} build-args: | PG_VERSION=${{ matrix.version }} SEMVER=${{ env.SEMVER }} @@ -163,9 +164,10 @@ jobs: with: context: . push: true - platforms: "linux/amd64,linux/arm64" + platforms: ${{ matrix.runner == 'ubuntu-22.04' && 'linux/amd64' || 'linux/arm64' }} file: ./docker/Dockerfile - tags: tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} + provenance: false + tags: tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} build-args: | PG_VERSION=${{ matrix.version }} SEMVER=${{ env.SEMVER }} @@ -177,29 +179,85 @@ jobs: password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} - name: Build and push Enterprise image to Docker Registry uses: docker/build-push-action@v6 - if: ${{ matrix.version != '17' }} with: context: ./docker/pg-cnpg push: true - platforms: "linux/amd64" - file: ./docker/pg-cnpg/Dockerfile + platforms: ${{ matrix.runner == 'ubuntu-22.04' && 'linux/amd64' || 'linux/arm64' }} + file: ${{ matrix.runner == 'ubuntu-22.04' && './docker/pg-cnpg/Dockerfile' || './docker/pg-cnpg/Dockerfile.arm' }} + provenance: false build-args: | PG_MAJOR=${{ matrix.version }} SEMVER=${{ env.SEMVER }} - LIB_DIR=/usr/lib/x86_64-linux-gnu - TARGETARCH=amd64 + LIB_DIR=${{ matrix.runner == 'ubuntu-22.04' && '/usr/lib/x86_64-linux-gnu' || '/usr/lib/aarch64-linux-gnu' }} + TARGETARCH=${{ env.PLATFORM }} PGVECTOR=0.8.0 - tags: modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }} + tags: modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} + + create-manifests: + runs-on: ubuntu-latest + needs: ["semver", "build", "docker"] + strategy: + matrix: + version: ["14", "15", "16", "17"] + env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERIO_USERNAME }} + password: ${{ secrets.DOCKERIO_TOKEN }} + - name: Create manifest and push + run: | + docker manifest create \ + tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }} \ + --amend tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ + --amend tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 + docker manifest push tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }} + docker manifest create \ + tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} \ + --amend tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ + --amend tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 + docker manifest push tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} + create-manifests-modelzai: + runs-on: ubuntu-latest + needs: ["semver", "build", "docker"] + strategy: + matrix: + version: ["14", "15", "16", "17"] + env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + - name: Login to modelzai Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} + password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} + - name: Create manifest and push + run: | + docker manifest create \ + modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }} \ + --amend modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ + --amend modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }}-arm64 + docker manifest push modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }} + test: name: Run tests - runs-on: - - ubuntu-latest - needs: ["semver", "build", "docker"] + runs-on: ${{ matrix.runner }} + needs: ["semver", "build", "docker", "create-manifests"] strategy: matrix: - version: [14, 15, 16] - platform: ["amd64"] + version: [14, 15, 16, 17] + runner: ["ubuntu-22.04"] container: image: modelzai/vchord-cnpg:${{ matrix.version }}-v${{ needs.semver.outputs.SEMVER }} options: --user root @@ -223,5 +281,6 @@ jobs: bash /usr/local/bin/docker-entrypoint.sh postgres & sleep 5 curl https://registry.pgtrunk.io/extensions/all | jq -r ".[] | .name" > /tmp/extensions.txt - trunk-install.sh | tee /tmp/output.txt - cat /tmp/output.txt + extension-install.sh | tee /tmp/output.txt + cat /tmp/output.txt + diff --git a/.github/workflows/release_pg_slim.yml b/.github/workflows/release_pg_slim.yml index 779b35db..74ae65f3 100644 --- a/.github/workflows/release_pg_slim.yml +++ b/.github/workflows/release_pg_slim.yml @@ -8,8 +8,8 @@ jobs: strategy: matrix: version: [14, 15, 16, 17] - platform: ["amd64"] - runs-on: ubuntu-latest + runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] + runs-on: ${{ matrix.runner }} env: PG_MAJOR: ${{ matrix.version }} steps: @@ -24,13 +24,38 @@ jobs: with: username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} - - name: Push binary release to Docker Registry + - name: Push Postgres Slim to Docker Registry uses: docker/build-push-action@v4 with: context: ./docker/pg-slim push: true - platforms: "linux/${{ matrix.platform }}" + platforms: ${{ matrix.runner == 'ubuntu-22.04' && 'linux/amd64' || 'linux/arm64' }} file: ./docker/pg-slim/Dockerfile + provenance: false build-args: | PG_MAJOR=${{ matrix.version }} - tags: modelzai/pg-slim:${{ matrix.version }}-${{ matrix.platform }} \ No newline at end of file + tags: modelzai/pg-slim:${{ matrix.version }}-${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} + create-manifests: + needs: ["pg-slim"] + strategy: + matrix: + version: ["14", "15", "16", "17"] + runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] + runs-on: ${{ matrix.runner }} + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + - name: Login to modelzai Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} + password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} + - name: Create manifest and push + run: | + docker manifest create \ + modelzai/pg-slim:${{ matrix.version }} \ + --amend modelzai/pg-slim:${{ matrix.version }}-amd64 \ + --amend modelzai/pg-slim:${{ matrix.version }}-arm64 + docker manifest push modelzai/pg-slim:${{ matrix.version }} \ No newline at end of file diff --git a/docker/pg-cnpg/Dockerfile b/docker/pg-cnpg/Dockerfile index 4871e1d1..7f9404ee 100644 --- a/docker/pg-cnpg/Dockerfile +++ b/docker/pg-cnpg/Dockerfile @@ -4,8 +4,8 @@ ARG TARGETARCH FROM tensorchord/vchord-binary:pg${PG_MAJOR}-v${SEMVER} as binary -FROM rust:1.82-bookworm as builder -ARG TRUNK_VER=0.15.6 +FROM rust:1.83-bookworm as builder +ARG TRUNK_VER=0.15.7 ENV CARGO_REGISTRIES_CRATES_IO_PROTOCOL sparse RUN cargo install --version $TRUNK_VER pg-trunk @@ -31,15 +31,15 @@ COPY requirements.txt . # Install barman-cloud RUN set -xe; \ - apt-get update; \ - apt-get install -y --no-install-recommends \ - python3-pip \ - python3-psycopg2 \ - python3-setuptools \ - ; \ - pip3 install --upgrade pip; \ + apt-get update; \ + apt-get install -y --no-install-recommends \ + python3-pip \ + python3-psycopg2 \ + python3-setuptools \ + ; \ + pip3 install --upgrade pip; \ # TODO: Remove --no-deps once https://github.com/pypa/pip/issues/9644 is solved - pip3 install --no-deps -r requirements.txt; \ + pip3 install --no-deps -r requirements.txt; \ apt-get autoremove -y; \ apt-get clean; \ rm -rf /var/lib/apt/lists/*; @@ -106,13 +106,13 @@ RUN wget https://download.oracle.com/otn_software/linux/instantclient/1920000/in # Install zhparser dependency RUN wget http://www.xunsearch.com/scws/down/scws-1.2.3.tar.bz2 && \ - tar xvf scws-1.2.3.tar.bz2 && \ - cd scws-1.2.3 && \ - ./configure && \ - make install && \ - cd .. && \ - rm -rf scws-1.2.3.tar.bz2 scws-1.2.3 && \ - ln -s /usr/local/lib/libscws.so ${LIB_DIR}/libscws.so + tar xvf scws-1.2.3.tar.bz2 && \ + cd scws-1.2.3 && \ + ./configure && \ + make install && \ + cd .. && \ + rm -rf scws-1.2.3.tar.bz2 scws-1.2.3 && \ + ln -s /usr/local/lib/libscws.so ${LIB_DIR}/libscws.so # Install duckdb libs RUN wget https://github.com/duckdb/duckdb/releases/download/v0.8.1/libduckdb-linux-amd64.zip && \ @@ -176,8 +176,8 @@ RUN git clone https://github.com/EnterpriseDB/pg_failover_slots.git && \ # cache all extensions ENV LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH -# Test trunk -COPY trunk-install.sh /usr/local/bin/ +# Test trunk install extension +COPY extension-install.sh /usr/local/bin/ # Change the uid of postgres to 26 RUN usermod -u 26 postgres @@ -193,4 +193,4 @@ RUN set -eux; \ RUN chown -R postgres:postgres /tmp USER 26 -ENV PATH $PATH:/usr/lib/postgresql/${PG_MAJOR}/bin +ENV PATH $PATH:/usr/lib/postgresql/${PG_MAJOR}/bin \ No newline at end of file diff --git a/docker/pg-cnpg/Dockerfile.arm b/docker/pg-cnpg/Dockerfile.arm new file mode 100644 index 00000000..dc76570e --- /dev/null +++ b/docker/pg-cnpg/Dockerfile.arm @@ -0,0 +1,70 @@ +ARG PG_MAJOR +ARG SEMVER +ARG TARGETARCH + +FROM tensorchord/vchord-binary:pg${PG_MAJOR}-v${SEMVER}-${TARGETARCH} as binary + +# From https://github.com/cloudnative-pg/postgres-containers/blob/main/Debian/17/bookworm/Dockerfile +FROM postgres:${PG_MAJOR}-bookworm +ARG PG_MAJOR +ARG SEMVER +ARG TARGETARCH +ARG PGVECTOR + +COPY requirements.txt / + +# Install additional extensions +RUN set -xe; \ + apt-get update; \ + apt-get install -y --no-install-recommends \ + "postgresql-${PG_MAJOR}-pgaudit" \ + "postgresql-${PG_MAJOR}-pg-failover-slots" \ + ; \ + rm -fr /tmp/* ; \ + rm -rf /var/lib/apt/lists/*; + +# Install barman-cloud +RUN set -xe; \ + apt-get update; \ + apt-get install -y --no-install-recommends \ + python3-pip \ + python3-psycopg2 \ + python3-setuptools \ + ; \ + pip3 install --break-system-packages --upgrade pip; \ + # TODO: Remove --no-deps once https://github.com/pypa/pip/issues/9644 is solved + pip3 install --break-system-packages --no-deps -r requirements.txt; \ + rm -rf /var/lib/apt/lists/*; + +COPY --from=binary /workspace/postgresql-${PG_MAJOR}-vchord_${SEMVER}-1_${TARGETARCH}.deb /tmp/vchord.deb +RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb + +RUN apt-get update && apt-get install -y \ + jq \ + curl \ + wget \ + sudo \ + && rm -rf /var/lib/apt/lists/* + +# Install pig +RUN curl -fsSL https://repo.pigsty.io/pig | bash && \ + pig repo add pigsty pgdg -u + +# Install pgvector +RUN pig ext install -y pgvector=${PGVECTOR} + +# Install pg_stat_statements +RUN pig ext install -y pg_stat_statements + +# Install auto_explain +RUN pig ext install -y auto_explain + +# Install plpython3u +RUN pig ext install -y plpython3u + +# Install pg_later +RUN pig ext install -y pg_later=0.3.0 + +# Change the uid of postgres to 26 +RUN usermod -u 26 postgres +USER 26 diff --git a/docker/pg-cnpg/trunk-install.sh b/docker/pg-cnpg/extension-install.sh similarity index 99% rename from docker/pg-cnpg/trunk-install.sh rename to docker/pg-cnpg/extension-install.sh index d64bd495..e76f5f5c 100755 --- a/docker/pg-cnpg/trunk-install.sh +++ b/docker/pg-cnpg/extension-install.sh @@ -77,4 +77,4 @@ printf "\n\n***FAILED EXTENSIONS***\n" for failed in "${failed_extensions[@]}" do echo $failed -done +done \ No newline at end of file diff --git a/docker/pg-slim/Dockerfile b/docker/pg-slim/Dockerfile index 06049c9e..89b1629f 100644 --- a/docker/pg-slim/Dockerfile +++ b/docker/pg-slim/Dockerfile @@ -9,22 +9,22 @@ ENV DEBIAN_FRONTEND=noninteractive \ # Get latest package updates RUN set -eux; \ - apt-get update; \ - apt-get upgrade -y + apt-get update; \ + apt-get upgrade -y # explicitly set user/group IDs RUN set -eux; \ - groupadd -r postgres --gid=999; \ -# https://salsa.debian.org/postgresql/postgresql-common/blob/997d842ee744687d99a2b2d95c1083a2615c79e8/debian/postgresql-common.postinst#L32-35 - useradd -r -g postgres --uid=999 --home-dir=/var/lib/postgresql --shell=/bin/bash postgres; \ -# also create the postgres user's home directory with appropriate permissions -# see https://github.com/docker-library/postgres/issues/274 - install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql + groupadd -r postgres --gid=999; \ + # https://salsa.debian.org/postgresql/postgresql-common/blob/997d842ee744687d99a2b2d95c1083a2615c79e8/debian/postgresql-common.postinst#L32-35 + useradd -r -g postgres --uid=999 --home-dir=/var/lib/postgresql --shell=/bin/bash postgres; \ + # also create the postgres user's home directory with appropriate permissions + # see https://github.com/docker-library/postgres/issues/274 + install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql RUN set -eux; \ - apt-get update; \ - apt-get install -y --no-install-recommends \ - locales curl ca-certificates gnupg lsb-release lbzip2 git cmake \ + apt-get update; \ + apt-get install -y --no-install-recommends \ + locales curl ca-certificates gnupg lsb-release lbzip2 git cmake \ ; \ rm -rf /var/lib/apt/lists/*; \ localedef -i en_US -c -f UTF-8 -A /usr/share/locale/locale.alias en_US.UTF-8 @@ -39,72 +39,81 @@ RUN git clone https://github.com/postgres/postgres.git && \ git checkout REL_$PG_RELEASE RUN set -eux; \ - apt-get update && apt-get install -y \ - libreadline-dev \ - zlib1g-dev \ - libpq-dev \ - build-essential \ - python3-dev \ - tcl-dev \ - libxslt1-dev \ - libperl-dev \ - libpam0g-dev \ - libreadline-dev \ - libssl-dev \ - xz-utils \ - libnss-wrapper \ - llvm \ - clang \ - icu-devtools \ - pkg-config \ - libgss-dev \ - libkrb5-dev \ - uuid-dev \ - gettext \ - liblz4-dev \ - libsystemd-dev \ - libselinux1-dev \ - libzstd-dev \ - vim \ - flex \ - bison; \ - apt-get autoremove -y; \ - apt-get clean -y; \ - rm -rf /var/lib/apt/lists/* + apt-get update && apt-get install -y \ + libreadline-dev \ + zlib1g-dev \ + libpq-dev \ + build-essential \ + python3-dev \ + tcl-dev \ + libxslt1-dev \ + libperl-dev \ + libpam0g-dev \ + libreadline-dev \ + libssl-dev \ + xz-utils \ + libnss-wrapper \ + llvm \ + clang \ + icu-devtools \ + pkg-config \ + libgss-dev \ + libkrb5-dev \ + uuid-dev \ + gettext \ + liblz4-dev \ + libsystemd-dev \ + libselinux1-dev \ + libzstd-dev \ + vim \ + flex \ + bison; \ + apt-get autoremove -y; \ + apt-get clean -y; \ + rm -rf /var/lib/apt/lists/* WORKDIR postgres ENV CFLAGS "-g -O2 -fstack-protector-strong -Wformat -Werror=format-security -fno-omit-frame-pointer" ENV LDFLAGS "-Wl,-z,relro -Wl,-z,now" RUN ./configure --prefix=/usr/lib/postgresql/${PG_MAJOR} \ - --datarootdir=${ALTDIR} \ - --libdir=${ALTDIR}/${PG_MAJOR}/lib \ - --with-perl \ - --with-python \ - --with-tcl \ - --with-pam \ - --with-libxml \ - --with-libxslt \ - --with-openssl \ - --enable-nls \ - --enable-thread-safety \ - --enable-debug \ - --disable-rpath \ - --with-uuid=e2fs \ - --with-gnu-ld \ - --with-gssapi \ - --with-pgport=5432 \ - --with-system-tzdata=/usr/share/zoneinfo \ - --with-icu \ - --with-llvm \ - --with-lz4 \ - --with-zstd \ - --with-systemd \ - --with-selinux + --datarootdir=${ALTDIR} \ + --libdir=${ALTDIR}/${PG_MAJOR}/lib \ + --with-perl \ + --with-python \ + --with-tcl \ + --with-pam \ + --with-libxml \ + --with-libxslt \ + --with-openssl \ + --enable-nls \ + --enable-thread-safety \ + --enable-debug \ + --disable-rpath \ + --with-uuid=e2fs \ + --with-gnu-ld \ + --with-gssapi \ + --with-pgport=5432 \ + --with-system-tzdata=/usr/share/zoneinfo \ + --with-icu \ + --with-llvm \ + --with-lz4 \ + --with-zstd \ + --with-systemd \ + --with-selinux RUN make -j$(nproc) RUN make install +# Remove libpq-dev provided from Ubuntu repos and set the newly-compiled one to system path +RUN set -eux; \ + apt-get purge -y libpq-dev && \ + apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +RUN echo "${ALTDIR}/${PG_MAJOR}/lib" > /etc/ld.so.conf.d/postgres.conf && \ + ldconfig + WORKDIR / RUN rm -rf /postgres @@ -115,9 +124,8 @@ COPY --from=tianon/gosu /gosu /usr/local/bin/ # make the sample config easier to munge (and "correct by default") RUN set -eux; \ - echo "unix_socket_directories = '/var/run/postgresql'" >> ${ALTDIR}/postgresql.conf.sample; \ - sed -ri "s!^#?(listen_addresses)\s*=\s*\S+.*!\1 = '*'!" ${ALTDIR}/postgresql.conf.sample; \ - grep -F "listen_addresses = '*'" ${ALTDIR}/postgresql.conf.sample + sed -ri "s!^#?(listen_addresses)\s*=\s*\S+.*!\1 = '*'!" ${ALTDIR}/postgresql.conf.sample; \ + grep -F "listen_addresses = '*'" ${ALTDIR}/postgresql.conf.sample RUN install --verbose --directory --owner postgres --group postgres --mode 3777 /var/run/postgresql @@ -129,9 +137,6 @@ COPY docker-entrypoint.sh docker-ensure-initdb.sh /usr/local/bin/ RUN ln -sT docker-ensure-initdb.sh /usr/local/bin/docker-enforce-initdb.sh ENTRYPOINT ["docker-entrypoint.sh"] -# Remove pre-installed pg_config -RUN rm /usr/bin/pg_config - STOPSIGNAL SIGINT EXPOSE 5432 CMD ["postgres"] \ No newline at end of file From 6b004c83f80c69b536d93b807d669f5abdb2bfa5 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 21 Feb 2025 15:15:31 +0800 Subject: [PATCH 117/324] refactor: split quantized data vectors to two tapes (#196) split data vectors to two tapes, one is immutable since `maintain`, and another one is appendable for newly-inserted vectors Signed-off-by: usamoi --- crates/algorithm/src/build.rs | 101 ++++- crates/algorithm/src/bulkdelete.rs | 183 ++++---- crates/algorithm/src/cache.rs | 54 ++- crates/algorithm/src/freepages.rs | 23 +- crates/algorithm/src/insert.rs | 28 +- crates/algorithm/src/lib.rs | 12 +- crates/algorithm/src/maintain.rs | 243 ++++++----- crates/algorithm/src/operator.rs | 28 ++ crates/algorithm/src/pipe.rs | 14 - crates/algorithm/src/prewarm.rs | 147 +++---- crates/algorithm/src/rerank.rs | 2 +- crates/algorithm/src/search.rs | 41 +- crates/algorithm/src/tape.rs | 308 ++++--------- crates/algorithm/src/tuples.rs | 665 +++++++++++++---------------- crates/algorithm/src/vectors.rs | 36 +- 15 files changed, 883 insertions(+), 1002 deletions(-) delete mode 100644 crates/algorithm/src/pipe.rs diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index d8c9b62d..0a8221f3 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -1,8 +1,9 @@ -use crate::RelationWrite; use crate::operator::{Accessor2, Operator, Vector}; -use crate::tape::*; +use crate::tape::TapeWriter; use crate::tuples::*; use crate::types::*; +use crate::{Branch, DerefMut, Page, PageGuard, RelationWrite}; +use simd::fast_scan::{any_pack, padding_pack}; use vector::VectorOwned; pub fn build( @@ -47,10 +48,13 @@ pub fn build( let mut level = Vec::new(); for j in 0..structures[i].len() { if i == 0 { - let tape = TapeWriter::<_, _, H0Tuple>::create(|| index.extend(false)); + let frozen_tape = TapeWriter::<_, _, FrozenTuple>::create(|| index.extend(false)); + let appendable_tape = + TapeWriter::<_, _, AppendableTuple>::create(|| index.extend(false)); let mut jump = TapeWriter::<_, _, JumpTuple>::create(|| index.extend(false)); jump.push(JumpTuple { - first: tape.first(), + frozen_first: frozen_tape.first(), + appendable_first: appendable_tape.first(), }); level.push(jump.first()); } else { @@ -73,17 +77,46 @@ pub fn build( } else { O::Vector::code(h1_mean) }; - tape.push(H1Branch { + tape.push(Branch { mean: pointer_of_means[i - 1][child as usize], dis_u_2: code.dis_u_2, factor_ppc: code.factor_ppc, factor_ip: code.factor_ip, factor_err: code.factor_err, signs: code.signs, - first: pointer_of_firsts[i - 1][child as usize], + extra: pointer_of_firsts[i - 1][child as usize], }); } - let tape = tape.into_inner(); + let (mut tape, branches) = tape.into_inner(); + if !branches.is_empty() { + let mut remain = + padding_pack(branches.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + loop { + let freespace = tape.freespace(); + if H1Tuple::estimate_size_0(remain.len()) <= freespace as usize { + tape.tape_put(H1Tuple::_0 { + mean: any_pack(branches.iter().map(|x| x.mean)), + dis_u_2: any_pack(branches.iter().map(|x| x.dis_u_2)), + factor_ppc: any_pack(branches.iter().map(|x| x.factor_ppc)), + factor_ip: any_pack(branches.iter().map(|x| x.factor_ip)), + factor_err: any_pack(branches.iter().map(|x| x.factor_err)), + first: any_pack(branches.iter().map(|x| x.extra)), + len: branches.len() as _, + elements: remain, + }); + break; + } + if let Some(w) = H1Tuple::fit_1(freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + tape.tape_put(H1Tuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } else { + tape.tape_move(); + } + } + } level.push(tape.first()); } } @@ -100,3 +133,57 @@ pub fn build( freepage_first: freepage.first(), }); } + +pub struct H1TapeWriter { + tape: TapeWriter, + branches: Vec>, +} + +impl H1TapeWriter +where + G: PageGuard + DerefMut, + G::Target: Page, + E: Fn() -> G, +{ + fn create(extend: E) -> Self { + Self { + tape: TapeWriter::create(extend), + branches: Vec::new(), + } + } + fn push(&mut self, branch: Branch) { + self.branches.push(branch); + if self.branches.len() == 32 { + let chunk = std::array::from_fn::<_, 32, _>(|_| self.branches.pop().unwrap()); + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + loop { + let freespace = self.tape.freespace(); + if H1Tuple::estimate_size_0(remain.len()) <= freespace as usize { + self.tape.tape_put(H1Tuple::_0 { + mean: chunk.each_ref().map(|x| x.mean), + dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), + factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), + factor_ip: chunk.each_ref().map(|x| x.factor_ip), + factor_err: chunk.each_ref().map(|x| x.factor_err), + first: chunk.each_ref().map(|x| x.extra), + len: chunk.len() as _, + elements: remain, + }); + break; + } + if let Some(w) = H1Tuple::fit_1(freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(H1Tuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } else { + self.tape.tape_move(); + } + } + } + } + fn into_inner(self) -> (TapeWriter, Vec>) { + (self.tape, self.branches) + } +} diff --git a/crates/algorithm/src/bulkdelete.rs b/crates/algorithm/src/bulkdelete.rs index 524ec908..506d2f05 100644 --- a/crates/algorithm/src/bulkdelete.rs +++ b/crates/algorithm/src/bulkdelete.rs @@ -1,7 +1,6 @@ -use crate::operator::Operator; -use crate::pipe::Pipe; +use crate::operator::{FunctionalAccessor, Operator}; use crate::tuples::*; -use crate::{Page, RelationWrite}; +use crate::{Page, RelationWrite, tape}; use std::num::NonZeroU64; pub fn bulkdelete( @@ -10,7 +9,8 @@ pub fn bulkdelete( callback: impl Fn(NonZeroU64) -> bool, ) { let meta_guard = index.read(0); - let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let height_of_root = meta_tuple.height_of_root(); let root_first = meta_tuple.root_first(); let vectors_first = meta_tuple.vectors_first(); @@ -21,25 +21,19 @@ pub fn bulkdelete( let step = |state: State| { let mut results = Vec::new(); for first in state { - let mut current = first; - while current != u32::MAX { - let h1_guard = index.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h1_tuple { - H1TupleReader::_0(h1_tuple) => { - for first in h1_tuple.first().iter().copied() { - results.push(first); - } - } - H1TupleReader::_1(_) => (), + tape::read_h1_tape( + index.clone(), + first, + || { + fn push(_: &mut (), _: &[T]) {} + fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { + [(); 32] } - } - current = h1_guard.get_opaque().next; - } + FunctionalAccessor::new((), push, finish) + }, + |(), _, first| results.push(first), + |_| check(), + ); } results }; @@ -48,97 +42,94 @@ pub fn bulkdelete( } for first in state { let jump_guard = index.read(first); - let jump_tuple = jump_guard - .get(1) - .expect("data corruption") - .pipe(read_tuple::); - let first = jump_tuple.first(); - let mut current = first; - while current != u32::MAX { - check(); - let read = index.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - let h0_tuple = read - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h0_tuple { - H0TupleReader::_0(h0_tuple) => { - let p = h0_tuple.payload(); - if let Some(payload) = p { - if callback(payload) { + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + { + let mut current = jump_tuple.frozen_first(); + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + let bytes = read.get(i).expect("data corruption"); + let tuple = FrozenTuple::deserialize_ref(bytes); + if let FrozenTupleReader::_0(tuple) = tuple { + for p in tuple.payload().iter() { + if Some(true) == p.map(&callback) { break 'flag true; } } } - H0TupleReader::_1(h0_tuple) => { - let p = h0_tuple.payload(); - for j in 0..32 { - if let Some(payload) = p[j] { - if callback(payload) { - break 'flag true; - } + } + false + }; + if flag { + drop(read); + let mut write = index.write(current, false); + for i in 1..=write.len() { + let bytes = write.get_mut(i).expect("data corruption"); + let tuple = FrozenTuple::deserialize_mut(bytes); + if let FrozenTupleWriter::_0(mut tuple) = tuple { + for p in tuple.payload().iter_mut() { + if Some(true) == p.map(&callback) { + *p = None; } } } - H0TupleReader::_2(_) => (), } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; } - false - }; - if flag { - drop(read); - let mut write = index.write(current, false); - for i in 1..=write.len() { - let h0_tuple = write - .get_mut(i) - .expect("data corruption") - .pipe(write_tuple::); - match h0_tuple { - H0TupleWriter::_0(mut h0_tuple) => { - let p = h0_tuple.payload(); - if let Some(payload) = *p { - if callback(payload) { - *p = None; - } - } + } + } + { + let mut current = jump_tuple.appendable_first(); + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + let bytes = read.get(i).expect("data corruption"); + let tuple = AppendableTuple::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + break 'flag true; } - H0TupleWriter::_1(mut h0_tuple) => { - let p = h0_tuple.payload(); - for j in 0..32 { - if let Some(payload) = p[j] { - if callback(payload) { - p[j] = None; - } - } - } + } + false + }; + if flag { + drop(read); + let mut write = index.write(current, false); + for i in 1..=write.len() { + let bytes = write.get_mut(i).expect("data corruption"); + let mut tuple = AppendableTuple::deserialize_mut(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + *p = None; } - H0TupleWriter::_2(_) => (), } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; } } } } { - let first = vectors_first; - let mut current = first; + let mut current = vectors_first; while current != u32::MAX { check(); let read = index.read(current); let flag = 'flag: { for i in 1..=read.len() { - if let Some(vector_bytes) = read.get(i) { - let vector_tuple = vector_bytes.pipe(read_tuple::>); - let p = vector_tuple.payload(); - if let Some(payload) = p { - if callback(payload) { - break 'flag true; - } + if let Some(bytes) = read.get(i) { + let tuple = VectorTuple::::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + break 'flag true; } } } @@ -148,13 +139,11 @@ pub fn bulkdelete( drop(read); let mut write = index.write(current, true); for i in 1..=write.len() { - if let Some(vector_bytes) = write.get(i) { - let vector_tuple = vector_bytes.pipe(read_tuple::>); - let p = vector_tuple.payload(); - if let Some(payload) = p { - if callback(payload) { - write.free(i); - } + if let Some(bytes) = write.get(i) { + let tuple = VectorTuple::::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + write.free(i); } }; } diff --git a/crates/algorithm/src/cache.rs b/crates/algorithm/src/cache.rs index 0405708e..d3d0e5d6 100644 --- a/crates/algorithm/src/cache.rs +++ b/crates/algorithm/src/cache.rs @@ -1,16 +1,12 @@ -use crate::pipe::Pipe; -use crate::tuples::{H1Tuple, H1TupleReader, MetaTuple, read_tuple}; -use crate::{Page, RelationRead}; +use crate::operator::FunctionalAccessor; +use crate::tuples::{MetaTuple, WithReader}; +use crate::{Page, RelationRead, tape}; pub fn cache(index: impl RelationRead) -> Vec { - let mut trace = Vec::::new(); - let mut read = |id| { - let result = index.read(id); - trace.push(id); - result - }; - let meta_guard = read(0); - let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let mut trace = vec![0_u32]; + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let height_of_root = meta_tuple.height_of_root(); let root_first = meta_tuple.root_first(); drop(meta_guard); @@ -19,25 +15,23 @@ pub fn cache(index: impl RelationRead) -> Vec { let mut step = |state: State| { let mut results = Vec::new(); for first in state { - let mut current = first; - while current != u32::MAX { - let h1_guard = read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h1_tuple { - H1TupleReader::_0(h1_tuple) => { - for first in h1_tuple.first().iter().copied() { - results.push(first); - } - } - H1TupleReader::_1(_) => (), + tape::read_h1_tape( + index.clone(), + first, + || { + fn push(_: &mut (), _: &[T]) {} + fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { + [(); 32] } - } - current = h1_guard.get_opaque().next; - } + FunctionalAccessor::new((), push, finish) + }, + |(), _, first| { + results.push(first); + }, + |id| { + trace.push(id); + }, + ); } results }; @@ -45,7 +39,7 @@ pub fn cache(index: impl RelationRead) -> Vec { state = step(state); } for first in state { - let _ = read(first); + trace.push(first); } trace } diff --git a/crates/algorithm/src/freepages.rs b/crates/algorithm/src/freepages.rs index 6dc505c2..95b63aea 100644 --- a/crates/algorithm/src/freepages.rs +++ b/crates/algorithm/src/freepages.rs @@ -1,4 +1,3 @@ -use crate::pipe::Pipe; use crate::tuples::*; use crate::*; use std::cmp::Reverse; @@ -7,9 +6,7 @@ pub fn mark(index: impl RelationWrite, freepage_first: u32, pages: &[u32]) { let mut pages = pages.to_vec(); pages.sort_by_key(|x| Reverse(*x)); pages.dedup(); - let first = freepage_first; - assert!(first != u32::MAX); - let (mut current, mut offset) = (first, 0_u32); + let (mut current, mut offset) = (freepage_first, 0_u32); while pages.is_empty() { let locals = { let mut local = Vec::new(); @@ -20,12 +17,10 @@ pub fn mark(index: impl RelationWrite, freepage_first: u32, pages: &[u32]) { }; let mut freespace_guard = index.write(current, false); if freespace_guard.len() == 0 { - freespace_guard.alloc(&serialize(&FreepageTuple {})); + freespace_guard.alloc(&FreepageTuple::serialize(&FreepageTuple {})); } - let mut freespace_tuple = freespace_guard - .get_mut(1) - .expect("data corruption") - .pipe(write_tuple::); + let freespace_bytes = freespace_guard.get_mut(1).expect("data corruption"); + let mut freespace_tuple = FreepageTuple::deserialize_mut(freespace_bytes); for local in locals { freespace_tuple.mark(local as _); } @@ -38,18 +33,14 @@ pub fn mark(index: impl RelationWrite, freepage_first: u32, pages: &[u32]) { } pub fn fetch(index: impl RelationWrite, freepage_first: u32) -> Option { - let first = freepage_first; - assert!(first != u32::MAX); - let (mut current, mut offset) = (first, 0_u32); + let (mut current, mut offset) = (freepage_first, 0_u32); loop { let mut freespace_guard = index.write(current, false); if freespace_guard.len() == 0 { return None; } - let mut freespace_tuple = freespace_guard - .get_mut(1) - .expect("data corruption") - .pipe(write_tuple::); + let freespace_bytes = freespace_guard.get_mut(1).expect("data corruption"); + let mut freespace_tuple = FreepageTuple::deserialize_mut(freespace_bytes); if let Some(local) = freespace_tuple.fetch() { return Some(local as u32 + offset); } diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 9b71bae4..d876aede 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -1,11 +1,9 @@ use crate::linked_vec::LinkedVec; use crate::operator::*; -use crate::pipe::Pipe; use crate::select_heap::SelectHeap; -use crate::tape::{access_1, append}; use crate::tuples::*; use crate::vectors::{self}; -use crate::{Page, RelationWrite}; +use crate::{Page, RelationWrite, tape}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -15,7 +13,8 @@ use vector::{VectorBorrowed, VectorOwned}; pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vector: O::Vector) { let meta_guard = index.read(0); - let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); let rerank_in_heap = meta_tuple.rerank_in_heap(); @@ -42,7 +41,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto let mut state: State = { let mean = root_mean; if is_residual { - let residual_u = vectors::access_1::( + let residual_u = vectors::read_for_h1_tuple::( index.clone(), mean, LAccess::new( @@ -64,7 +63,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto } else { default_lut_block.as_ref().unwrap() }; - access_1( + tape::read_h1_tape( index.clone(), first, || { @@ -76,6 +75,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto |lowerbound, mean, first| { results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(first))); }, + |_| (), ); } let mut heap = SelectHeap::from_vec(results.into_vec()); @@ -84,7 +84,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); if is_residual { - let (dis_u, residual_u) = vectors::access_1::( + let (dis_u, residual_u) = vectors::read_for_h1_tuple::( index.clone(), mean, LAccess::new( @@ -101,7 +101,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto AlwaysEqual(Some(residual_u)), )); } else { - let dis_u = vectors::access_1::( + let dis_u = vectors::read_for_h1_tuple::( index.clone(), mean, LAccess::new( @@ -126,7 +126,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto } else { O::Vector::code(vector.as_borrowed()) }; - let bytes = serialize(&H0Tuple::_0 { + let bytes = AppendableTuple::serialize(&AppendableTuple { mean, dis_u_2: code.dis_u_2, factor_ppc: code.factor_ppc, @@ -137,12 +137,8 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto }); let jump_guard = index.read(first); - let jump_tuple = jump_guard - .get(1) - .expect("data corruption") - .pipe(read_tuple::); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let first = jump_tuple.first(); - - append(index.clone(), first, &bytes, false); + tape::append(index.clone(), jump_tuple.appendable_first(), &bytes, false); } diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 372774f3..3bba11de 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -9,7 +9,6 @@ mod freepages; mod insert; mod linked_vec; mod maintain; -mod pipe; mod prewarm; mod rerank; mod search; @@ -30,6 +29,7 @@ pub use prewarm::prewarm; pub use rerank::{rerank_heap, rerank_index}; pub use search::search; +use crate::tuples::IndexPointer; use std::ops::{Deref, DerefMut}; use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -79,3 +79,13 @@ pub enum RerankMethod { Index, Heap, } + +pub(crate) struct Branch { + pub mean: IndexPointer, + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub signs: Vec, + pub extra: T, +} diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index fec057b6..2749a997 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -1,45 +1,39 @@ -use crate::operator::Operator; -use crate::pipe::Pipe; -use crate::tape::*; +use crate::operator::{FunctionalAccessor, Operator}; +use crate::tape::{self, TapeWriter}; use crate::tuples::*; -use crate::{Page, RelationWrite, freepages}; -use simd::fast_scan::unpack; +use crate::{Branch, DerefMut, Page, PageGuard, RelationWrite, freepages}; +use simd::fast_scan::{padding_pack, unpack}; +use std::num::NonZeroU64; pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let meta_guard = index.read(0); - let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let height_of_root = meta_tuple.height_of_root(); let root_first = meta_tuple.root_first(); let freepage_first = meta_tuple.freepage_first(); drop(meta_guard); - let firsts = { + let state = { type State = Vec; let mut state: State = vec![root_first]; let step = |state: State| { let mut results = Vec::new(); for first in state { - let mut current = first; - while current != u32::MAX { - check(); - let h1_guard = index.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h1_tuple { - H1TupleReader::_0(h1_tuple) => { - for first in h1_tuple.first().iter().copied() { - results.push(first); - } - } - H1TupleReader::_1(_) => (), + tape::read_h1_tape( + index.clone(), + first, + || { + fn push(_: &mut (), _: &[T]) {} + fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { + [(); 32] } - } - current = h1_guard.get_opaque().next; - } + FunctionalAccessor::new((), push, finish) + }, + |(), _, first| results.push(first), + |_| check(), + ); } results }; @@ -49,14 +43,12 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { state }; - for first in firsts { + for first in state { let mut jump_guard = index.write(first, false); - let mut jump_tuple = jump_guard - .get_mut(1) - .expect("data corruption") - .pipe(write_tuple::); + let jump_bytes = jump_guard.get_mut(1).expect("data corruption"); + let mut jump_tuple = JumpTuple::deserialize_mut(jump_bytes); - let mut tape = H0TapeWriter::<_, _>::create(|| { + let mut tape = FrozenTapeWriter::<_, _>::create(|| { if let Some(id) = freepages::fetch(index.clone(), freepage_first) { let mut write = index.write(id, false); write.clear(); @@ -68,80 +60,131 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let mut trace = Vec::new(); - let first = *jump_tuple.first(); - let mut current = first; - let mut computing = None; - while current != u32::MAX { + let mut callback = |code: (f32, f32, f32, f32, Vec<_>), mean, payload| { + tape.push(Branch { + mean, + dis_u_2: code.0, + factor_ppc: code.1, + factor_ip: code.2, + factor_err: code.3, + signs: code.4, + extra: payload, + }); + }; + let mut step = |id| { check(); - trace.push(current); - let h0_guard = index.read(current); - for i in 1..=h0_guard.len() { - let h0_tuple = h0_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h0_tuple { - H0TupleReader::_0(h0_tuple) => { - if let Some(payload) = h0_tuple.payload() { - tape.push(H0Branch { - mean: h0_tuple.mean(), - dis_u_2: h0_tuple.code().0, - factor_ppc: h0_tuple.code().1, - factor_ip: h0_tuple.code().2, - factor_err: h0_tuple.code().3, - signs: h0_tuple - .code() - .4 - .iter() - .flat_map(|x| { - std::array::from_fn::<_, 64, _>(|i| *x & (1 << i) != 0) - }) - .take(dims as _) - .collect::>(), - payload, - }); - } - } - H0TupleReader::_1(h0_tuple) => { - let computing = &mut computing.take().unwrap_or_else(Vec::new); - computing.extend_from_slice(h0_tuple.elements()); - let unpacked = unpack(computing); - for j in 0..32 { - if let Some(payload) = h0_tuple.payload()[j] { - tape.push(H0Branch { - mean: h0_tuple.mean()[j], - dis_u_2: h0_tuple.metadata().0[j], - factor_ppc: h0_tuple.metadata().1[j], - factor_ip: h0_tuple.metadata().2[j], - factor_err: h0_tuple.metadata().3[j], - signs: unpacked[j] - .iter() - .flat_map(|&x| { - [x & 1 != 0, x & 2 != 0, x & 4 != 0, x & 8 != 0] - }) - .collect(), - payload, - }); - } - } - } - H0TupleReader::_2(h0_tuple) => { - let computing = computing.get_or_insert_with(Vec::new); - computing.extend_from_slice(h0_tuple.elements()); - } - } - } - current = h0_guard.get_opaque().next; - drop(h0_guard); + trace.push(id); + }; + tape::read_frozen_tape( + index.clone(), + *jump_tuple.frozen_first(), + || { + FunctionalAccessor::new( + Vec::<[u8; 16]>::new(), + Vec::<[u8; 16]>::extend_from_slice, + |elements: Vec<_>, input: (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32])| { + let unpacked = unpack(&elements); + std::array::from_fn(|i| { + let f = |&x| [x & 1 != 0, x & 2 != 0, x & 4 != 0, x & 8 != 0]; + let signs = unpacked[i].iter().flat_map(f).collect::>(); + (input.0[i], input.1[i], input.2[i], input.3[i], signs) + }) + }, + ) + }, + &mut callback, + &mut step, + ); + tape::read_appendable_tape( + index.clone(), + *jump_tuple.appendable_first(), + |code| { + let signs = code + .4 + .iter() + .flat_map(|x| std::array::from_fn::<_, 64, _>(|i| *x & (1 << i) != 0)) + .take(dims as _) + .collect::>(); + (code.0, code.1, code.2, code.3, signs) + }, + &mut callback, + &mut step, + ); + + let (frozen_tape, branches) = tape.into_inner(); + + let mut appendable_tape = TapeWriter::create(|| index.extend(false)); + + for branch in branches { + appendable_tape.push(AppendableTuple { + mean: branch.mean, + dis_u_2: branch.dis_u_2, + factor_ppc: branch.factor_ppc, + factor_ip: branch.factor_ip, + factor_err: branch.factor_err, + payload: Some(branch.extra), + elements: rabitq::pack_to_u64(&branch.signs), + }); } - let tape = tape.into_inner(); - let new = tape.first(); - drop(tape); + *jump_tuple.frozen_first() = { frozen_tape }.first(); + *jump_tuple.appendable_first() = { appendable_tape }.first(); - *jump_tuple.first() = new; drop(jump_guard); freepages::mark(index.clone(), freepage_first, &trace); } } + +struct FrozenTapeWriter { + tape: TapeWriter, + branches: Vec>, +} + +impl FrozenTapeWriter +where + G: PageGuard + DerefMut, + G::Target: Page, + E: Fn() -> G, +{ + fn create(extend: E) -> Self { + Self { + tape: TapeWriter::create(extend), + branches: Vec::new(), + } + } + fn push(&mut self, branch: Branch) { + self.branches.push(branch); + if self.branches.len() == 32 { + let chunk = std::array::from_fn::<_, 32, _>(|_| self.branches.pop().unwrap()); + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + loop { + let freespace = self.tape.freespace(); + if FrozenTuple::estimate_size_0(remain.len()) <= freespace as usize { + self.tape.tape_put(FrozenTuple::_0 { + mean: chunk.each_ref().map(|x| x.mean), + dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), + factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), + factor_ip: chunk.each_ref().map(|x| x.factor_ip), + factor_err: chunk.each_ref().map(|x| x.factor_err), + payload: chunk.each_ref().map(|x| Some(x.extra)), + elements: remain, + }); + break; + } + if let Some(w) = FrozenTuple::fit_1(freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(FrozenTuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } else { + self.tape.tape_move(); + } + } + } + } + fn into_inner(self) -> (TapeWriter, Vec>) { + (self.tape, self.branches) + } +} diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index e8a57bd4..ae80180d 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -287,6 +287,34 @@ where } } +pub struct FunctionalAccessor { + data: T, + p: P, + f: F, +} + +impl FunctionalAccessor { + pub fn new(data: T, p: P, f: F) -> Self { + Self { data, p, f } + } +} + +impl Accessor1 for FunctionalAccessor +where + P: for<'a> FnMut(&'a mut T, &'a [E]), + F: FnOnce(T, M) -> R, +{ + type Output = R; + + fn push(&mut self, input: &[E]) { + (self.p)(&mut self.data, input); + } + + fn finish(self, input: M) -> Self::Output { + (self.f)(self.data, input) + } +} + pub struct LAccess<'a, E, M, A> { elements: &'a [E], metadata: M, diff --git a/crates/algorithm/src/pipe.rs b/crates/algorithm/src/pipe.rs deleted file mode 100644 index 18cfc376..00000000 --- a/crates/algorithm/src/pipe.rs +++ /dev/null @@ -1,14 +0,0 @@ -pub trait Pipe { - fn pipe(self, f: impl FnOnce(Self) -> T) -> T - where - Self: Sized; -} - -impl Pipe for S { - fn pipe(self, f: impl FnOnce(Self) -> T) -> T - where - Self: Sized, - { - f(self) - } -} diff --git a/crates/algorithm/src/prewarm.rs b/crates/algorithm/src/prewarm.rs index 587f7528..9395c060 100644 --- a/crates/algorithm/src/prewarm.rs +++ b/crates/algorithm/src/prewarm.rs @@ -1,12 +1,12 @@ -use crate::operator::Operator; -use crate::pipe::Pipe; +use crate::operator::{FunctionalAccessor, Operator}; use crate::tuples::*; -use crate::{Page, RelationRead, vectors}; +use crate::{Page, RelationRead, tape, vectors}; use std::fmt::Write; pub fn prewarm(index: impl RelationRead, height: i32, check: impl Fn()) -> String { let meta_guard = index.read(0); - let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let height_of_root = meta_tuple.height_of_root(); let root_mean = meta_tuple.root_mean(); let root_first = meta_tuple.root_first(); @@ -20,96 +20,89 @@ pub fn prewarm(index: impl RelationRead, height: i32, check: impl F } type State = Vec; let mut state: State = { - let mut nodes = Vec::new(); + let mut results = Vec::new(); { - vectors::access_1::(index.clone(), root_mean, ()); - nodes.push(root_first); + vectors::read_for_h1_tuple::(index.clone(), root_mean, ()); + results.push(root_first); } writeln!(message, "------------------------").unwrap(); - writeln!(message, "number of nodes: {}", nodes.len()).unwrap(); - writeln!(message, "number of tuples: {}", 1).unwrap(); + writeln!(message, "number of nodes: {}", results.len()).unwrap(); writeln!(message, "number of pages: {}", 1).unwrap(); - nodes + results }; let mut step = |state: State| { - let mut counter_pages = 0_usize; - let mut counter_tuples = 0_usize; - let mut nodes = Vec::new(); - for list in state { - let mut current = list; - while current != u32::MAX { - counter_pages += 1; - check(); - let h1_guard = index.read(current); - for i in 1..=h1_guard.len() { - counter_tuples += 1; - let h1_tuple = h1_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h1_tuple { - H1TupleReader::_0(h1_tuple) => { - for mean in h1_tuple.mean().iter().copied() { - vectors::access_1::(index.clone(), mean, ()); - } - for first in h1_tuple.first().iter().copied() { - nodes.push(first); - } - } - H1TupleReader::_1(_) => (), + let mut counter = 0_usize; + let mut results = Vec::new(); + for first in state { + tape::read_h1_tape( + index.clone(), + first, + || { + fn push(_: &mut (), _: &[T]) {} + fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { + [(); 32] } - } - current = h1_guard.get_opaque().next; - } + FunctionalAccessor::new((), push, finish) + }, + |(), mean, first| { + results.push(first); + vectors::read_for_h1_tuple::(index.clone(), mean, ()); + }, + |_| { + check(); + counter += 1; + }, + ); } writeln!(message, "------------------------").unwrap(); - writeln!(message, "number of nodes: {}", nodes.len()).unwrap(); - writeln!(message, "number of tuples: {}", counter_tuples).unwrap(); - writeln!(message, "number of pages: {}", counter_pages).unwrap(); - nodes + writeln!(message, "number of nodes: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", counter).unwrap(); + results }; for _ in (std::cmp::max(1, prewarm_max_height)..height_of_root).rev() { state = step(state); } if prewarm_max_height == 0 { - let mut counter_pages = 0_usize; - let mut counter_tuples = 0_usize; - let mut counter_nodes = 0_usize; - for list in state { - let jump_guard = index.read(list); - let jump_tuple = jump_guard - .get(1) - .expect("data corruption") - .pipe(read_tuple::); - let first = jump_tuple.first(); - let mut current = first; - while current != u32::MAX { - counter_pages += 1; - check(); - let h0_guard = index.read(current); - for i in 1..=h0_guard.len() { - counter_tuples += 1; - let h0_tuple = h0_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h0_tuple { - H0TupleReader::_0(_h0_tuple) => { - counter_nodes += 1; - } - H0TupleReader::_1(_h0_tuple) => { - counter_nodes += 32; - } - H0TupleReader::_2(_) => (), + let mut counter = 0_usize; + let mut results = Vec::new(); + for first in state { + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + tape::read_frozen_tape( + index.clone(), + jump_tuple.frozen_first(), + || { + fn push(_: &mut (), _: &[T]) {} + fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { + [(); 32] } - } - current = h0_guard.get_opaque().next; - } + FunctionalAccessor::new((), push, finish) + }, + |(), _, _| { + results.push(()); + }, + |_| { + check(); + counter += 1; + }, + ); + tape::read_appendable_tape( + index.clone(), + jump_tuple.frozen_first(), + |_| (), + |(), _, _| { + results.push(()); + }, + |_| { + check(); + counter += 1; + }, + ); } writeln!(message, "------------------------").unwrap(); - writeln!(message, "number of nodes: {}", counter_nodes).unwrap(); - writeln!(message, "number of tuples: {}", counter_tuples).unwrap(); - writeln!(message, "number of pages: {}", counter_pages).unwrap(); + writeln!(message, "number of nodes: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", counter).unwrap(); } message } diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index 840fdd4c..2e409cb8 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -22,7 +22,7 @@ pub fn rerank_index( std::iter::from_fn(move || { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); - if let Some(dis_u) = vectors::access_0::( + if let Some(dis_u) = vectors::read_for_h0_tuple::( index.clone(), mean, pay_u, diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index 9972e861..7ed279af 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -1,9 +1,7 @@ use crate::linked_vec::LinkedVec; use crate::operator::*; -use crate::pipe::Pipe; -use crate::tape::{access_0, access_1}; use crate::tuples::*; -use crate::{Page, RelationRead, RerankMethod, vectors}; +use crate::{Page, RelationRead, RerankMethod, tape, vectors}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -25,7 +23,8 @@ pub fn search( )>, ) { let meta_guard = index.read(0); - let meta_tuple = meta_guard.get(1).unwrap().pipe(read_tuple::); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); let rerank_in_heap = meta_tuple.rerank_in_heap(); @@ -52,7 +51,7 @@ pub fn search( let mut state: State = vec![{ let mean = root_mean; if is_residual { - let residual_u = vectors::access_1::( + let residual_u = vectors::read_for_h1_tuple::( index.clone(), mean, LAccess::new( @@ -73,7 +72,7 @@ pub fn search( } else { default_lut.as_ref().map(|x| &x.0).unwrap() }; - access_1( + tape::read_h1_tape( index.clone(), first, || { @@ -85,6 +84,7 @@ pub fn search( |lowerbound, mean, first| { results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(first))); }, + |_| (), ); } let mut heap = BinaryHeap::from(results.into_vec()); @@ -93,7 +93,7 @@ pub fn search( while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); if is_residual { - let (dis_u, residual_u) = vectors::access_1::( + let (dis_u, residual_u) = vectors::read_for_h1_tuple::( index.clone(), mean, LAccess::new( @@ -110,7 +110,7 @@ pub fn search( AlwaysEqual(Some(residual_u)), )); } else { - let dis_u = vectors::access_1::( + let dis_u = vectors::read_for_h1_tuple::( index.clone(), mean, LAccess::new( @@ -139,24 +139,29 @@ pub fn search( default_lut.as_ref().unwrap() }; let jump_guard = index.read(first); - let jump_tuple = jump_guard - .get(1) - .expect("data corruption") - .pipe(read_tuple::); - let first = jump_tuple.first(); - access_0( + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + let mut callback = |lowerbound, mean, payload| { + results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(payload))); + }; + tape::read_frozen_tape( index.clone(), - first, + jump_tuple.frozen_first(), || { RAccess::new( (&lut.0.4, (lut.0.0, lut.0.1, lut.0.2, lut.0.3, epsilon)), O::Distance::block_accessor(), ) }, + &mut callback, + |_| (), + ); + tape::read_appendable_tape( + index.clone(), + jump_tuple.appendable_first(), |code| O::Distance::compute_lowerbound_binary(&lut.1, code, epsilon), - |lowerbound, mean, payload| { - results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(payload))); - }, + &mut callback, + |_| (), ); } ( diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index edc1c4ee..e90be1f6 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -1,9 +1,6 @@ use crate::operator::Accessor1; -use crate::pipe::Pipe; use crate::tuples::*; use crate::{Page, PageGuard, RelationRead, RelationWrite}; -use distance::Distance; -use simd::fast_scan::{any_pack, padding_pack}; use std::marker::PhantomData; use std::num::NonZeroU64; use std::ops::DerefMut; @@ -35,10 +32,10 @@ where pub fn first(&self) -> u32 { self.first } - fn freespace(&self) -> u16 { + pub fn freespace(&self) -> u16 { self.head.freespace() } - fn tape_move(&mut self) { + pub fn tape_move(&mut self) { if self.head.len() == 0 { panic!("tuple is too large to fit in a fresh page"); } @@ -56,7 +53,7 @@ where T: Tuple, { pub fn push(&mut self, x: T) -> IndexPointer { - let bytes = serialize(&x); + let bytes = T::serialize(&x); if let Some(i) = self.head.alloc(&bytes) { pair_to_pointer((self.head.id(), i)) } else { @@ -70,8 +67,8 @@ where } } } - fn tape_put(&mut self, x: T) -> IndexPointer { - let bytes = serialize(&x); + pub fn tape_put(&mut self, x: T) -> IndexPointer { + let bytes = T::serialize(&x); if let Some(i) = self.head.alloc(&bytes) { pair_to_pointer((self.head.id(), i)) } else { @@ -80,268 +77,111 @@ where } } -pub struct H1Branch { - pub mean: IndexPointer, - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub signs: Vec, - pub first: u32, -} - -pub struct H1TapeWriter { - tape: TapeWriter, - branches: Vec, -} - -impl H1TapeWriter -where - G: PageGuard + DerefMut, - G::Target: Page, - E: Fn() -> G, -{ - pub fn create(extend: E) -> Self { - Self { - tape: TapeWriter::create(extend), - branches: Vec::new(), - } - } - pub fn push(&mut self, branch: H1Branch) { - self.branches.push(branch); - if self.branches.len() == 32 { - let chunk = std::array::from_fn::<_, 32, _>(|_| self.branches.pop().unwrap()); - let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); - loop { - let freespace = self.tape.freespace(); - match (H1Tuple::fit_0(freespace), H1Tuple::fit_1(freespace)) { - (Some(w), _) if w >= remain.len() => { - self.tape.tape_put(H1Tuple::_0 { - mean: chunk.each_ref().map(|x| x.mean), - dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), - factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), - factor_ip: chunk.each_ref().map(|x| x.factor_ip), - factor_err: chunk.each_ref().map(|x| x.factor_err), - first: chunk.each_ref().map(|x| x.first), - len: chunk.len() as _, - elements: remain, - }); - break; - } - (_, Some(w)) => { - let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); - self.tape.tape_put(H1Tuple::_1 { - elements: left.to_vec(), - }); - remain = right.to_vec(); - } - (_, None) => self.tape.tape_move(), - } - } - } - } - pub fn into_inner(mut self) -> TapeWriter { - let chunk = self.branches; - if !chunk.is_empty() { - let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); - loop { - let freespace = self.tape.freespace(); - match (H1Tuple::fit_0(freespace), H1Tuple::fit_1(freespace)) { - (Some(w), _) if w >= remain.len() => { - self.tape.push(H1Tuple::_0 { - mean: any_pack(chunk.iter().map(|x| x.mean)), - dis_u_2: any_pack(chunk.iter().map(|x| x.dis_u_2)), - factor_ppc: any_pack(chunk.iter().map(|x| x.factor_ppc)), - factor_ip: any_pack(chunk.iter().map(|x| x.factor_ip)), - factor_err: any_pack(chunk.iter().map(|x| x.factor_err)), - first: any_pack(chunk.iter().map(|x| x.first)), - len: chunk.len() as _, - elements: remain, - }); - break; - } - (_, Some(w)) => { - let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); - self.tape.tape_put(H1Tuple::_1 { - elements: left.to_vec(), - }); - remain = right.to_vec(); - } - (_, None) => self.tape.tape_move(), - } - } - } - self.tape - } -} - -pub struct H0Branch { - pub mean: IndexPointer, - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub signs: Vec, - pub payload: NonZeroU64, -} - -pub struct H0TapeWriter { - tape: TapeWriter, - branches: Vec, -} - -impl H0TapeWriter -where - G: PageGuard + DerefMut, - G::Target: Page, - E: Fn() -> G, -{ - pub fn create(extend: E) -> Self { - Self { - tape: TapeWriter::create(extend), - branches: Vec::new(), - } - } - pub fn push(&mut self, branch: H0Branch) { - self.branches.push(branch); - if self.branches.len() == 32 { - let chunk = std::array::from_fn::<_, 32, _>(|_| self.branches.pop().unwrap()); - let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); - loop { - let freespace = self.tape.freespace(); - match (H0Tuple::fit_1(freespace), H0Tuple::fit_2(freespace)) { - (Some(w), _) if w >= remain.len() => { - self.tape.push(H0Tuple::_1 { - mean: chunk.each_ref().map(|x| x.mean), - dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), - factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), - factor_ip: chunk.each_ref().map(|x| x.factor_ip), - factor_err: chunk.each_ref().map(|x| x.factor_err), - payload: chunk.each_ref().map(|x| Some(x.payload)), - elements: remain, - }); - break; - } - (_, Some(w)) => { - let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); - self.tape.tape_put(H0Tuple::_2 { - elements: left.to_vec(), - }); - remain = right.to_vec(); - } - (_, None) => self.tape.tape_move(), - } - } - } - } - pub fn into_inner(mut self) -> TapeWriter { - for x in self.branches { - self.tape.push(H0Tuple::_0 { - mean: x.mean, - dis_u_2: x.dis_u_2, - factor_ppc: x.factor_ppc, - factor_ip: x.factor_ip, - factor_err: x.factor_err, - payload: Some(x.payload), - elements: rabitq::pack_to_u64(&x.signs), - }); - } - self.tape - } -} - -pub fn access_1( +pub fn read_h1_tape( index: impl RelationRead, first: u32, - make_block_accessor: impl Fn() -> A + Copy, - mut callback: impl FnMut(Distance, IndexPointer, u32), + accessor: impl Fn() -> A, + mut callback: impl FnMut(T, IndexPointer, u32), + mut step: impl FnMut(u32), ) where A: for<'a> Accessor1< [u8; 16], (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), - Output = [Distance; 32], + Output = [T; 32], >, { assert!(first != u32::MAX); let mut current = first; - let mut computing = None; + let mut x = None; while current != u32::MAX { - let h1_guard = index.read(current); - for i in 1..=h1_guard.len() { - let h1_tuple = h1_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h1_tuple { - H1TupleReader::_0(h1_tuple) => { - let mut compute = computing.take().unwrap_or_else(make_block_accessor); - compute.push(h1_tuple.elements()); - let lowerbounds = compute.finish(h1_tuple.metadata()); - for i in 0..h1_tuple.len() { - callback( - lowerbounds[i as usize], - h1_tuple.mean()[i as usize], - h1_tuple.first()[i as usize], - ); + step(current); + let guard = index.read(current); + for i in 1..=guard.len() { + let bytes = guard.get(i).expect("data corruption"); + let tuple = H1Tuple::deserialize_ref(bytes); + match tuple { + H1TupleReader::_0(tuple) => { + let mut x = x.take().unwrap_or_else(&accessor); + x.push(tuple.elements()); + let values = x.finish(tuple.metadata()); + for (j, value) in values.into_iter().enumerate() { + if j < tuple.len() as usize { + callback(value, tuple.mean()[j], tuple.first()[j]); + } } } - H1TupleReader::_1(h1_tuple) => { - let computing = computing.get_or_insert_with(make_block_accessor); - computing.push(h1_tuple.elements()); + H1TupleReader::_1(tuple) => { + x.get_or_insert_with(&accessor).push(tuple.elements()); } } } - current = h1_guard.get_opaque().next; + current = guard.get_opaque().next; } } -pub fn access_0( +pub fn read_frozen_tape( index: impl RelationRead, first: u32, - make_block_accessor: impl Fn() -> A + Copy, - compute_binary: impl Fn((f32, f32, f32, f32, &[u64])) -> Distance, - mut callback: impl FnMut(Distance, IndexPointer, NonZeroU64), + accessor: impl Fn() -> A, + mut callback: impl FnMut(T, IndexPointer, NonZeroU64), + mut step: impl FnMut(u32), ) where A: for<'a> Accessor1< [u8; 16], (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), - Output = [Distance; 32], + Output = [T; 32], >, { assert!(first != u32::MAX); let mut current = first; - let mut computing = None; + let mut x = None; while current != u32::MAX { - let h0_guard = index.read(current); - for i in 1..=h0_guard.len() { - let h0_tuple = h0_guard - .get(i) - .expect("data corruption") - .pipe(read_tuple::); - match h0_tuple { - H0TupleReader::_0(h0_tuple) => { - let lowerbound = compute_binary(h0_tuple.code()); - if let Some(payload) = h0_tuple.payload() { - callback(lowerbound, h0_tuple.mean(), payload); - } - } - H0TupleReader::_1(h0_tuple) => { - let mut compute = computing.take().unwrap_or_else(make_block_accessor); - compute.push(h0_tuple.elements()); - let lowerbounds = compute.finish(h0_tuple.metadata()); - for j in 0..32 { - if let Some(payload) = h0_tuple.payload()[j] { - callback(lowerbounds[j], h0_tuple.mean()[j], payload); + step(current); + let guard = index.read(current); + for i in 1..=guard.len() { + let bytes = guard.get(i).expect("data corruption"); + let tuple = FrozenTuple::deserialize_ref(bytes); + match tuple { + FrozenTupleReader::_0(tuple) => { + let mut x = x.take().unwrap_or_else(&accessor); + x.push(tuple.elements()); + let values = x.finish(tuple.metadata()); + for (j, value) in values.into_iter().enumerate() { + if let Some(payload) = tuple.payload()[j] { + callback(value, tuple.mean()[j], payload); } } } - H0TupleReader::_2(h0_tuple) => { - let computing = computing.get_or_insert_with(make_block_accessor); - computing.push(h0_tuple.elements()); + FrozenTupleReader::_1(tuple) => { + x.get_or_insert_with(&accessor).push(tuple.elements()); } } } - current = h0_guard.get_opaque().next; + current = guard.get_opaque().next; + } +} + +pub fn read_appendable_tape( + index: impl RelationRead, + first: u32, + mut access: impl FnMut((f32, f32, f32, f32, &[u64])) -> T, + mut callback: impl FnMut(T, IndexPointer, NonZeroU64), + mut step: impl FnMut(u32), +) { + assert!(first != u32::MAX); + let mut current = first; + while current != u32::MAX { + step(current); + let guard = index.read(current); + for i in 1..=guard.len() { + let bytes = guard.get(i).expect("data corruption"); + let tuple = AppendableTuple::deserialize_ref(bytes); + if let Some(payload) = tuple.payload() { + let value = access(tuple.code()); + callback(value, tuple.mean(), payload); + } + } + current = guard.get_opaque().next; } } diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index f60d5527..60a92dd7 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -7,41 +7,22 @@ use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: u64 = u64::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 2; +const VERSION: u64 = 3; pub trait Tuple: 'static { - type Reader<'a>: TupleReader<'a, Tuple = Self>; fn serialize(&self) -> Vec; } -pub trait WithWriter: Tuple { - type Writer<'a>: TupleWriter<'a, Tuple = Self>; -} - -pub trait TupleReader<'a>: Copy { - type Tuple: Tuple; - fn deserialize_ref(source: &'a [u8]) -> Self; -} - -pub trait TupleWriter<'a> { - type Tuple: Tuple; - fn deserialize_mut(source: &'a mut [u8]) -> Self; -} - -pub fn serialize(tuple: &T) -> Vec { - Tuple::serialize(tuple) +pub trait WithReader: Tuple { + type Reader<'a>; + fn deserialize_ref(source: &[u8]) -> Self::Reader<'_>; } -pub fn read_tuple(source: &[u8]) -> T::Reader<'_> { - TupleReader::deserialize_ref(source) -} - -pub fn write_tuple(source: &mut [u8]) -> T::Writer<'_> { - TupleWriter::deserialize_mut(source) +pub trait WithWriter: Tuple { + type Writer<'a>; + fn deserialize_mut(source: &mut [u8]) -> Self::Writer<'_>; } -// meta tuple - #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct MetaTupleHeader { @@ -72,8 +53,6 @@ pub struct MetaTuple { } impl Tuple for MetaTuple { - type Reader<'a> = MetaTupleReader<'a>; - fn serialize(&self) -> Vec { MetaTupleHeader { magic: MAGIC, @@ -93,14 +72,9 @@ impl Tuple for MetaTuple { } } -#[derive(Debug, Clone, Copy)] -pub struct MetaTupleReader<'a> { - header: &'a MetaTupleHeader, -} - -impl<'a> TupleReader<'a> for MetaTupleReader<'a> { - type Tuple = MetaTuple; - fn deserialize_ref(source: &'a [u8]) -> Self { +impl WithReader for MetaTuple { + type Reader<'a> = MetaTupleReader<'a>; + fn deserialize_ref(source: &[u8]) -> MetaTupleReader<'_> { if source.len() < 16 { panic!("bad bytes") } @@ -114,10 +88,15 @@ impl<'a> TupleReader<'a> for MetaTupleReader<'a> { } let checker = RefChecker::new(source); let header = checker.prefix(0); - Self { header } + MetaTupleReader { header } } } +#[derive(Debug, Clone, Copy)] +pub struct MetaTupleReader<'a> { + header: &'a MetaTupleHeader, +} + impl MetaTupleReader<'_> { pub fn dims(self) -> u32 { self.header.dims @@ -145,8 +124,6 @@ impl MetaTupleReader<'_> { } } -// freepage tuple - #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FreepageTupleHeader { @@ -162,8 +139,6 @@ const _: () = assert!(size_of::() == 4232); pub struct FreepageTuple {} impl Tuple for FreepageTuple { - type Reader<'a> = FreepageTupleReader<'a>; - fn serialize(&self) -> Vec { FreepageTupleHeader { a: std::array::from_fn(|_| 0), @@ -178,21 +153,11 @@ impl Tuple for FreepageTuple { impl WithWriter for FreepageTuple { type Writer<'a> = FreepageTupleWriter<'a>; -} - -#[derive(Debug, Clone, Copy)] -pub struct FreepageTupleReader<'a> { - #[allow(dead_code)] - header: &'a FreepageTupleHeader, -} - -impl<'a> TupleReader<'a> for FreepageTupleReader<'a> { - type Tuple = FreepageTuple; - fn deserialize_ref(source: &'a [u8]) -> Self { - let checker = RefChecker::new(source); + fn deserialize_mut(source: &mut [u8]) -> FreepageTupleWriter<'_> { + let mut checker = MutChecker::new(source); let header = checker.prefix(0); - Self { header } + FreepageTupleWriter { header } } } @@ -200,16 +165,6 @@ pub struct FreepageTupleWriter<'a> { header: &'a mut FreepageTupleHeader, } -impl<'a> TupleWriter<'a> for FreepageTupleWriter<'a> { - type Tuple = FreepageTuple; - - fn deserialize_mut(source: &'a mut [u8]) -> Self { - let mut checker = MutChecker::new(source); - let header = checker.prefix(0); - Self { header } - } -} - impl FreepageTupleWriter<'_> { pub fn mark(&mut self, i: usize) { let c_i = i; @@ -237,8 +192,6 @@ impl FreepageTupleWriter<'_> { } } -// vector tuple - #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader0 { @@ -274,8 +227,6 @@ pub enum VectorTuple { } impl Tuple for VectorTuple { - type Reader<'a> = VectorTupleReader<'a, V>; - fn serialize(&self) -> Vec { let mut buffer = Vec::::new(); match self { @@ -286,9 +237,6 @@ impl Tuple for VectorTuple { } => { buffer.extend((0 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } let metadata_s = buffer.len(); buffer.extend(metadata.as_bytes()); while buffer.len() % ALIGN != 0 { @@ -297,6 +245,9 @@ impl Tuple for VectorTuple { let elements_s = buffer.len(); buffer.extend(elements.as_bytes()); let elements_e = buffer.len(); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } buffer[size_of::()..][..size_of::()].copy_from_slice( VectorTupleHeader0 { payload: *payload, @@ -316,12 +267,12 @@ impl Tuple for VectorTuple { } => { buffer.extend((1 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } let elements_s = buffer.len(); buffer.extend(elements.as_bytes()); let elements_e = buffer.len(); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } buffer[size_of::()..][..size_of::()].copy_from_slice( VectorTupleHeader1 { payload: *payload, @@ -337,6 +288,34 @@ impl Tuple for VectorTuple { } } +impl WithReader for VectorTuple { + type Reader<'a> = VectorTupleReader<'a, V>; + + fn deserialize_ref(source: &[u8]) -> VectorTupleReader<'_, V> { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader0 = checker.prefix(size_of::()); + let metadata = checker.prefix(header.metadata_s); + let elements = checker.bytes(header.elements_s, header.elements_e); + VectorTupleReader::_0(VectorTupleReader0 { + header, + elements, + metadata, + }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + VectorTupleReader::_1(VectorTupleReader1 { header, elements }) + } + _ => panic!("bad bytes"), + } + } +} + #[derive(Clone)] pub struct VectorTupleReader0<'a, V: Vector> { header: &'a VectorTupleHeader0, @@ -362,34 +341,6 @@ pub enum VectorTupleReader<'a, V: Vector> { impl Copy for VectorTupleReader<'_, V> {} -impl<'a, V: Vector> TupleReader<'a> for VectorTupleReader<'a, V> { - type Tuple = VectorTuple; - - fn deserialize_ref(source: &'a [u8]) -> Self { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); - match tag { - 0 => { - let checker = RefChecker::new(source); - let header: &VectorTupleHeader0 = checker.prefix(size_of::()); - let metadata = checker.prefix(header.metadata_s); - let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_0(VectorTupleReader0 { - header, - elements, - metadata, - }) - } - 1 => { - let checker = RefChecker::new(source); - let header: &VectorTupleHeader1 = checker.prefix(size_of::()); - let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_1(VectorTupleReader1 { header, elements }) - } - _ => panic!("bad bytes"), - } - } -} - impl<'a, V: Vector> VectorTupleReader<'a, V> { pub fn payload(self) -> Option { match self { @@ -411,8 +362,6 @@ impl<'a, V: Vector> VectorTupleReader<'a, V> { } } -// height1tuple - #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct H1TupleHeader0 { @@ -454,15 +403,12 @@ pub enum H1Tuple { } impl H1Tuple { - pub fn fit_0(freespace: u16) -> Option { - let mut freespace = freespace as isize; - freespace -= size_of::() as isize; - freespace -= size_of::() as isize; - if freespace >= 0 { - Some(freespace as usize / size_of::<[u8; 16]>()) - } else { - None - } + pub fn estimate_size_0(elements: usize) -> usize { + let mut size = 0_usize; + size += size_of::(); + size += size_of::(); + size += elements * size_of::<[u8; 16]>(); + size } pub fn fit_1(freespace: u16) -> Option { let mut freespace = freespace as isize; @@ -477,8 +423,6 @@ impl H1Tuple { } impl Tuple for H1Tuple { - type Reader<'a> = H1TupleReader<'a>; - fn serialize(&self) -> Vec { let mut buffer = Vec::::new(); match self { @@ -494,9 +438,6 @@ impl Tuple for H1Tuple { } => { buffer.extend((0 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } let elements_s = buffer.len(); buffer.extend(elements.as_bytes()); let elements_e = buffer.len(); @@ -519,9 +460,6 @@ impl Tuple for H1Tuple { Self::_1 { elements } => { buffer.extend((1 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } let elements_s = buffer.len(); buffer.extend(elements.as_bytes()); let elements_e = buffer.len(); @@ -538,47 +476,47 @@ impl Tuple for H1Tuple { } } -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum H1TupleReader<'a> { - _0(H1TupleReader0<'a>), - _1(H1TupleReader1<'a>), -} - -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct H1TupleReader0<'a> { - header: &'a H1TupleHeader0, - elements: &'a [[u8; 16]], -} - -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct H1TupleReader1<'a> { - header: &'a H1TupleHeader1, - elements: &'a [[u8; 16]], -} - -impl<'a> TupleReader<'a> for H1TupleReader<'a> { - type Tuple = H1Tuple; +impl WithReader for H1Tuple { + type Reader<'a> = H1TupleReader<'a>; - fn deserialize_ref(source: &'a [u8]) -> Self { + fn deserialize_ref(source: &[u8]) -> H1TupleReader<'_> { let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); match tag { 0 => { let checker = RefChecker::new(source); let header: &H1TupleHeader0 = checker.prefix(size_of::()); let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_0(H1TupleReader0 { header, elements }) + H1TupleReader::_0(H1TupleReader0 { header, elements }) } 1 => { let checker = RefChecker::new(source); let header: &H1TupleHeader1 = checker.prefix(size_of::()); let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_1(H1TupleReader1 { header, elements }) + H1TupleReader::_1(H1TupleReader1 { header, elements }) } _ => panic!("bad bytes"), } } } +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum H1TupleReader<'a> { + _0(H1TupleReader0<'a>), + _1(H1TupleReader1<'a>), +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct H1TupleReader0<'a> { + header: &'a H1TupleHeader0, + elements: &'a [[u8; 16]], +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct H1TupleReader1<'a> { + header: &'a H1TupleHeader1, + elements: &'a [[u8; 16]], +} + impl<'a> H1TupleReader0<'a> { pub fn len(self) -> u32 { self.header.len @@ -611,30 +549,43 @@ impl<'a> H1TupleReader1<'a> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct JumpTupleHeader { - first: u32, - _padding_0: [ZeroU8; 4], + frozen_first: u32, + appendable_first: u32, } #[derive(Debug, Clone, PartialEq)] pub struct JumpTuple { - pub first: u32, + pub frozen_first: u32, + pub appendable_first: u32, } impl Tuple for JumpTuple { - type Reader<'a> = JumpTupleReader<'a>; - fn serialize(&self) -> Vec { JumpTupleHeader { - first: self.first, - _padding_0: Default::default(), + frozen_first: self.frozen_first, + appendable_first: self.appendable_first, } .as_bytes() .to_vec() } } +impl WithReader for JumpTuple { + type Reader<'a> = JumpTupleReader<'a>; + fn deserialize_ref(source: &[u8]) -> JumpTupleReader<'_> { + let checker = RefChecker::new(source); + let header: &JumpTupleHeader = checker.prefix(0); + JumpTupleReader { header } + } +} + impl WithWriter for JumpTuple { type Writer<'a> = JumpTupleWriter<'a>; + fn deserialize_mut(source: &mut [u8]) -> JumpTupleWriter<'_> { + let mut checker = MutChecker::new(source); + let header: &mut JumpTupleHeader = checker.prefix(0); + JumpTupleWriter { header } + } } #[derive(Debug, Clone, Copy)] @@ -642,19 +593,12 @@ pub struct JumpTupleReader<'a> { header: &'a JumpTupleHeader, } -impl<'a> TupleReader<'a> for JumpTupleReader<'a> { - type Tuple = JumpTuple; - - fn deserialize_ref(source: &'a [u8]) -> Self { - let checker = RefChecker::new(source); - let header: &JumpTupleHeader = checker.prefix(0); - Self { header } - } -} - impl JumpTupleReader<'_> { - pub fn first(self) -> u32 { - self.header.first + pub fn frozen_first(self) -> u32 { + self.header.frozen_first + } + pub fn appendable_first(self) -> u32 { + self.header.appendable_first } } @@ -663,38 +607,18 @@ pub struct JumpTupleWriter<'a> { header: &'a mut JumpTupleHeader, } -impl<'a> TupleWriter<'a> for JumpTupleWriter<'a> { - type Tuple = JumpTuple; - - fn deserialize_mut(source: &'a mut [u8]) -> Self { - let mut checker = MutChecker::new(source); - let header: &mut JumpTupleHeader = checker.prefix(0); - Self { header } - } -} - impl JumpTupleWriter<'_> { - pub fn first(&mut self) -> &mut u32 { - &mut self.header.first + pub fn frozen_first(&mut self) -> &mut u32 { + &mut self.header.frozen_first + } + pub fn appendable_first(&mut self) -> &mut u32 { + &mut self.header.appendable_first } } #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] -struct H0TupleHeader0 { - mean: IndexPointer, - dis_u_2: f32, - factor_ppc: f32, - factor_ip: f32, - factor_err: f32, - payload: Option, - elements_s: usize, - elements_e: usize, -} - -#[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] -struct H0TupleHeader1 { +struct FrozenTupleHeader0 { mean: [IndexPointer; 32], dis_u_2: [f32; 32], factor_ppc: [f32; 32], @@ -707,24 +631,15 @@ struct H0TupleHeader1 { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] -struct H0TupleHeader2 { +struct FrozenTupleHeader1 { elements_s: usize, elements_e: usize, } #[allow(clippy::large_enum_variant)] #[derive(Debug, Clone, PartialEq)] -pub enum H0Tuple { +pub enum FrozenTuple { _0 { - mean: IndexPointer, - dis_u_2: f32, - factor_ppc: f32, - factor_ip: f32, - factor_err: f32, - payload: Option, - elements: Vec, - }, - _1 { mean: [IndexPointer; 32], dis_u_2: [f32; 32], factor_ppc: [f32; 32], @@ -733,26 +648,23 @@ pub enum H0Tuple { payload: [Option; 32], elements: Vec<[u8; 16]>, }, - _2 { + _1 { elements: Vec<[u8; 16]>, }, } -impl H0Tuple { - pub fn fit_1(freespace: u16) -> Option { - let mut freespace = freespace as isize; - freespace -= size_of::() as isize; - freespace -= size_of::() as isize; - if freespace >= 0 { - Some(freespace as usize / size_of::<[u8; 16]>()) - } else { - None - } +impl FrozenTuple { + pub fn estimate_size_0(elements: usize) -> usize { + let mut size = 0_usize; + size += size_of::(); + size += size_of::(); + size += elements * size_of::<[u8; 16]>(); + size } - pub fn fit_2(freespace: u16) -> Option { + pub fn fit_1(freespace: u16) -> Option { let mut freespace = freespace as isize; freespace -= size_of::() as isize; - freespace -= size_of::() as isize; + freespace -= size_of::() as isize; if freespace >= 0 { Some(freespace as usize / size_of::<[u8; 16]>()) } else { @@ -761,13 +673,11 @@ impl H0Tuple { } } -impl Tuple for H0Tuple { - type Reader<'a> = H0TupleReader<'a>; - +impl Tuple for FrozenTuple { fn serialize(&self) -> Vec { let mut buffer = Vec::::new(); match self { - H0Tuple::_0 { + FrozenTuple::_0 { mean, dis_u_2, factor_ppc, @@ -777,15 +687,12 @@ impl Tuple for H0Tuple { elements, } => { buffer.extend((0 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat_n(0, size_of::())); - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } + buffer.extend(std::iter::repeat_n(0, size_of::())); let elements_s = buffer.len(); buffer.extend(elements.as_bytes()); let elements_e = buffer.len(); - buffer[size_of::()..][..size_of::()].copy_from_slice( - H0TupleHeader0 { + buffer[size_of::()..][..size_of::()].copy_from_slice( + FrozenTupleHeader0 { mean: *mean, dis_u_2: *dis_u_2, factor_ppc: *factor_ppc, @@ -798,48 +705,14 @@ impl Tuple for H0Tuple { .as_bytes(), ); } - H0Tuple::_1 { - mean, - dis_u_2, - factor_ppc, - factor_ip, - factor_err, - payload, - elements, - } => { + Self::_1 { elements } => { buffer.extend((1 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat_n(0, size_of::())); - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } - let elements_s = buffer.len(); - buffer.extend(elements.as_bytes()); - let elements_e = buffer.len(); - buffer[size_of::()..][..size_of::()].copy_from_slice( - H0TupleHeader1 { - mean: *mean, - dis_u_2: *dis_u_2, - factor_ppc: *factor_ppc, - factor_ip: *factor_ip, - factor_err: *factor_err, - payload: *payload, - elements_s, - elements_e, - } - .as_bytes(), - ); - } - Self::_2 { elements } => { - buffer.extend((2 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat_n(0, size_of::())); - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } + buffer.extend(std::iter::repeat_n(0, size_of::())); let elements_s = buffer.len(); buffer.extend(elements.as_bytes()); let elements_e = buffer.len(); - buffer[size_of::()..][..size_of::()].copy_from_slice( - H0TupleHeader2 { + buffer[size_of::()..][..size_of::()].copy_from_slice( + FrozenTupleHeader1 { elements_s, elements_e, } @@ -851,48 +724,65 @@ impl Tuple for H0Tuple { } } -impl WithWriter for H0Tuple { - type Writer<'a> = H0TupleWriter<'a>; -} +impl WithReader for FrozenTuple { + type Reader<'a> = FrozenTupleReader<'a>; -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum H0TupleReader<'a> { - _0(H0TupleReader0<'a>), - _1(H0TupleReader1<'a>), - _2(H0TupleReader2<'a>), + fn deserialize_ref(source: &[u8]) -> FrozenTupleReader<'_> { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &FrozenTupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + FrozenTupleReader::_0(FrozenTupleReader0 { header, elements }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &FrozenTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + FrozenTupleReader::_1(FrozenTupleReader1 { header, elements }) + } + _ => panic!("bad bytes"), + } + } } -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct H0TupleReader0<'a> { - header: &'a H0TupleHeader0, - elements: &'a [u64], -} +impl WithWriter for FrozenTuple { + type Writer<'a> = FrozenTupleWriter<'a>; -impl<'a> H0TupleReader0<'a> { - pub fn mean(self) -> IndexPointer { - self.header.mean - } - pub fn code(self) -> (f32, f32, f32, f32, &'a [u64]) { - ( - self.header.dis_u_2, - self.header.factor_ppc, - self.header.factor_ip, - self.header.factor_err, - self.elements, - ) - } - pub fn payload(self) -> Option { - self.header.payload + fn deserialize_mut(source: &mut [u8]) -> FrozenTupleWriter<'_> { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let mut checker = MutChecker::new(source); + let header: &mut FrozenTupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + FrozenTupleWriter::_0(FrozenTupleWriter0 { header, elements }) + } + 1 => { + let mut checker = MutChecker::new(source); + let header: &mut FrozenTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + FrozenTupleWriter::_1(FrozenTupleWriter1 { header, elements }) + } + _ => panic!("bad bytes"), + } } } #[derive(Debug, Clone, Copy, PartialEq)] -pub struct H0TupleReader1<'a> { - header: &'a H0TupleHeader1, +pub enum FrozenTupleReader<'a> { + _0(FrozenTupleReader0<'a>), + _1(FrozenTupleReader1<'a>), +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct FrozenTupleReader0<'a> { + header: &'a FrozenTupleHeader0, elements: &'a [[u8; 16]], } -impl<'a> H0TupleReader1<'a> { +impl<'a> FrozenTupleReader0<'a> { pub fn mean(self) -> &'a [IndexPointer; 32] { &self.header.mean } @@ -913,113 +803,149 @@ impl<'a> H0TupleReader1<'a> { } #[derive(Debug, Clone, Copy, PartialEq)] -pub struct H0TupleReader2<'a> { - header: &'a H0TupleHeader2, +pub struct FrozenTupleReader1<'a> { + header: &'a FrozenTupleHeader1, elements: &'a [[u8; 16]], } -impl<'a> H0TupleReader2<'a> { +impl<'a> FrozenTupleReader1<'a> { pub fn elements(self) -> &'a [[u8; 16]] { self.elements } } -impl<'a> TupleReader<'a> for H0TupleReader<'a> { - type Tuple = H0Tuple; - - fn deserialize_ref(source: &'a [u8]) -> Self { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); - match tag { - 0 => { - let checker = RefChecker::new(source); - let header: &H0TupleHeader0 = checker.prefix(size_of::()); - let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_0(H0TupleReader0 { header, elements }) - } - 1 => { - let checker = RefChecker::new(source); - let header: &H0TupleHeader1 = checker.prefix(size_of::()); - let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_1(H0TupleReader1 { header, elements }) - } - 2 => { - let checker = RefChecker::new(source); - let header: &H0TupleHeader2 = checker.prefix(size_of::()); - let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_2(H0TupleReader2 { header, elements }) - } - _ => panic!("bad bytes"), - } - } -} - #[derive(Debug)] -pub enum H0TupleWriter<'a> { - _0(H0TupleWriter0<'a>), - _1(H0TupleWriter1<'a>), +pub enum FrozenTupleWriter<'a> { + _0(FrozenTupleWriter0<'a>), #[allow(dead_code)] - _2(H0TupleWriter2<'a>), + _1(FrozenTupleWriter1<'a>), } #[derive(Debug)] -pub struct H0TupleWriter0<'a> { - header: &'a mut H0TupleHeader0, - #[allow(dead_code)] - elements: &'a mut [u64], -} - -#[derive(Debug)] -pub struct H0TupleWriter1<'a> { - header: &'a mut H0TupleHeader1, +pub struct FrozenTupleWriter0<'a> { + header: &'a mut FrozenTupleHeader0, #[allow(dead_code)] elements: &'a mut [[u8; 16]], } #[derive(Debug)] -pub struct H0TupleWriter2<'a> { +pub struct FrozenTupleWriter1<'a> { #[allow(dead_code)] - header: &'a mut H0TupleHeader2, + header: &'a mut FrozenTupleHeader1, #[allow(dead_code)] elements: &'a mut [[u8; 16]], } -impl<'a> TupleWriter<'a> for H0TupleWriter<'a> { - type Tuple = H0Tuple; +impl FrozenTupleWriter0<'_> { + pub fn payload(&mut self) -> &mut [Option; 32] { + &mut self.header.payload + } +} - fn deserialize_mut(source: &'a mut [u8]) -> Self { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); - match tag { - 0 => { - let mut checker = MutChecker::new(source); - let header: &mut H0TupleHeader0 = checker.prefix(size_of::()); - let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_0(H0TupleWriter0 { header, elements }) - } - 1 => { - let mut checker = MutChecker::new(source); - let header: &mut H0TupleHeader1 = checker.prefix(size_of::()); - let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_1(H0TupleWriter1 { header, elements }) - } - 2 => { - let mut checker = MutChecker::new(source); - let header: &mut H0TupleHeader2 = checker.prefix(size_of::()); - let elements = checker.bytes(header.elements_s, header.elements_e); - Self::_2(H0TupleWriter2 { header, elements }) +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct AppendableTupleHeader { + mean: IndexPointer, + dis_u_2: f32, + factor_ppc: f32, + factor_ip: f32, + factor_err: f32, + payload: Option, + elements_s: usize, + elements_e: usize, +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug, Clone, PartialEq)] +pub struct AppendableTuple { + pub mean: IndexPointer, + pub dis_u_2: f32, + pub factor_ppc: f32, + pub factor_ip: f32, + pub factor_err: f32, + pub payload: Option, + pub elements: Vec, +} + +impl Tuple for AppendableTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + buffer.extend(std::iter::repeat_n(0, size_of::())); + let elements_s = buffer.len(); + buffer.extend(self.elements.as_bytes()); + let elements_e = buffer.len(); + buffer[..size_of::()].copy_from_slice( + AppendableTupleHeader { + mean: self.mean, + dis_u_2: self.dis_u_2, + factor_ppc: self.factor_ppc, + factor_ip: self.factor_ip, + factor_err: self.factor_err, + payload: self.payload, + elements_s, + elements_e, } - _ => panic!("bad bytes"), - } + .as_bytes(), + ); + buffer } } -impl H0TupleWriter0<'_> { - pub fn payload(&mut self) -> &mut Option { - &mut self.header.payload +impl WithReader for AppendableTuple { + type Reader<'a> = AppendableTupleReader<'a>; + + fn deserialize_ref(source: &[u8]) -> AppendableTupleReader<'_> { + let checker = RefChecker::new(source); + let header: &AppendableTupleHeader = checker.prefix(0); + let elements = checker.bytes(header.elements_s, header.elements_e); + AppendableTupleReader { header, elements } } } -impl H0TupleWriter1<'_> { - pub fn payload(&mut self) -> &mut [Option; 32] { +impl WithWriter for AppendableTuple { + type Writer<'a> = AppendableTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> AppendableTupleWriter<'_> { + let mut checker = MutChecker::new(source); + let header: &mut AppendableTupleHeader = checker.prefix(0); + let elements = checker.bytes(header.elements_s, header.elements_e); + AppendableTupleWriter { header, elements } + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct AppendableTupleReader<'a> { + header: &'a AppendableTupleHeader, + elements: &'a [u64], +} + +impl<'a> AppendableTupleReader<'a> { + pub fn mean(self) -> IndexPointer { + self.header.mean + } + pub fn code(self) -> (f32, f32, f32, f32, &'a [u64]) { + ( + self.header.dis_u_2, + self.header.factor_ppc, + self.header.factor_ip, + self.header.factor_err, + self.elements, + ) + } + pub fn payload(self) -> Option { + self.header.payload + } +} + +#[derive(Debug)] +pub struct AppendableTupleWriter<'a> { + header: &'a mut AppendableTupleHeader, + #[allow(dead_code)] + elements: &'a mut [u64], +} + +impl AppendableTupleWriter<'_> { + pub fn payload(&mut self) -> &mut Option { &mut self.header.payload } } @@ -1201,9 +1127,6 @@ fn aliasing_test() { let elements = (0u32..1111).collect::>(); let mut buffer = Vec::::new(); buffer.extend(std::iter::repeat_n(0, size_of::())); - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } let elements_s = buffer.len(); buffer.extend(elements.as_bytes()); let elements_e = buffer.len(); diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index dcc11af0..42b83b0b 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -1,11 +1,10 @@ use crate::operator::*; -use crate::pipe::Pipe; use crate::tuples::*; use crate::{Page, PageGuard, RelationRead, RelationWrite, tape}; use std::num::NonZeroU64; use vector::VectorOwned; -pub fn access_1< +pub fn read_for_h1_tuple< O: Operator, A: Accessor1<::Element, ::Metadata>, >( @@ -16,21 +15,19 @@ pub fn access_1< let mut cursor = Err(mean); let mut result = accessor; while let Err(mean) = cursor.map_err(pointer_to_pair) { - let vector_guard = index.read(mean.0); - let vector_tuple = vector_guard - .get(mean.1) - .expect("data corruption") - .pipe(read_tuple::>); - if vector_tuple.payload().is_some() { + let guard = index.read(mean.0); + let bytes = guard.get(mean.1).expect("data corruption"); + let tuple = VectorTuple::::deserialize_ref(bytes); + if tuple.payload().is_some() { panic!("data corruption"); } - result.push(vector_tuple.elements()); - cursor = vector_tuple.metadata_or_pointer(); + result.push(tuple.elements()); + cursor = tuple.metadata_or_pointer(); } result.finish(cursor.expect("data corruption")) } -pub fn access_0< +pub fn read_for_h0_tuple< O: Operator, A: Accessor1<::Element, ::Metadata>, >( @@ -42,18 +39,17 @@ pub fn access_0< let mut cursor = Err(mean); let mut result = accessor; while let Err(mean) = cursor.map_err(pointer_to_pair) { - let vector_guard = index.read(mean.0); - let vector_tuple = vector_guard - .get(mean.1)? - .pipe(read_tuple::>); - if vector_tuple.payload().is_none() { + let guard = index.read(mean.0); + let bytes = guard.get(mean.1)?; + let tuple = VectorTuple::::deserialize_ref(bytes); + if tuple.payload().is_none() { panic!("data corruption"); } - if vector_tuple.payload() != Some(payload) { + if tuple.payload() != Some(payload) { return None; } - result.push(vector_tuple.elements()); - cursor = vector_tuple.metadata_or_pointer(); + result.push(tuple.elements()); + cursor = tuple.metadata_or_pointer(); } Some(result.finish(cursor.ok()?)) } @@ -74,7 +70,7 @@ pub fn append( let (metadata, slices) = O::Vector::vector_split(vector); let mut chain = Ok(metadata); for i in (0..slices.len()).rev() { - let bytes = serialize::>(&match chain { + let bytes = VectorTuple::::serialize(&match chain { Ok(metadata) => VectorTuple::_0 { elements: slices[i].to_vec(), payload: Some(payload), From 12763078b36d4e558c62c563d927bf2c7043dc8b Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Tue, 4 Mar 2025 09:12:33 +0800 Subject: [PATCH 118/324] chore: add ghcr image (#200) Will build it in next release Signed-off-by: cutecutecat --- .github/workflows/release.yml | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c8b85688..f115b4c3 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -159,6 +159,12 @@ jobs: build-args: | PG_VERSION=${{ matrix.version }} SEMVER=${{ env.SEMVER }} + - name: Login to GHCR + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} - name: Push PostgreSQL release to Docker Registry uses: docker/build-push-action@v6 with: @@ -167,13 +173,15 @@ jobs: platforms: ${{ matrix.runner == 'ubuntu-22.04' && 'linux/amd64' || 'linux/arm64' }} file: ./docker/Dockerfile provenance: false - tags: tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} + tags: | + tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} + ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} build-args: | PG_VERSION=${{ matrix.version }} SEMVER=${{ env.SEMVER }} PGVECTOR=0.8.0 - name: Login to modelzai Docker Hub - uses: docker/login-action@v2 + uses: docker/login-action@v3 with: username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} @@ -211,6 +219,12 @@ jobs: with: username: ${{ secrets.DOCKERIO_USERNAME }} password: ${{ secrets.DOCKERIO_TOKEN }} + - name: Login to GHCR + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} - name: Create manifest and push run: | docker manifest create \ @@ -223,6 +237,12 @@ jobs: --amend tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ --amend tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 docker manifest push tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} + docker manifest create \ + ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} \ + --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ + --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 + docker manifest push ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} + create-manifests-modelzai: runs-on: ubuntu-latest From 869440140b5bcd324340f5e271a0ad4b8d95208f Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Wed, 5 Mar 2025 16:33:40 +0800 Subject: [PATCH 119/324] chore: update readme and scripts with url (#198) ## docs(README) - add dependency of pgvector - fix `vector_cos_ops` to `vector_cosine_ops` - update image version ## chore(bench scripts) - fix incorrect QPS count at parallel requests - when `workers` > 32, QPS will be overestimated - send the same whole data to all requests, instead of split data for parallel requests - for dataset `laion-100`, 1000 split to 128 parts will be have a huge negative impact on query startup - use dburl instead of local password Signed-off-by: cutecutecat --- README.md | 314 ++++++++++++++++++++++++++++++++++++++-------- scripts/README.md | 10 +- scripts/bench.py | 99 +++++++++------ scripts/index.py | 7 +- 4 files changed, 326 insertions(+), 104 deletions(-) diff --git a/README.md b/README.md index 8b2f92c2..c2e9ebd4 100644 --- a/README.md +++ b/README.md @@ -3,36 +3,76 @@

Effortlessly host 100 million 768-dimensional vectors (250GB+) on an AWS i4i.xlarge instance ($250/month), featuring 4 vCPUs and 32GB of RAM with VectorChord.

-

-discord invitation link -Twitter -Docker pulls -

Docker pull for pgvecto.rs: Previous Docker pulls

-

+
-VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search, and serves as the successor to [pgvecto.rs](https://github.com/tensorchord/pgvecto.rs). -With VectorChord, you can store 400,000 vectors for just $1, enabling significant savings: 6x more vectors compared to Pinecone's optimized storage and 26x more than pgvector/pgvecto.rs for the same price[^1]. For further insights, check out our [launch blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql). +[Official Site][official-site-link] · [Blog][blog-link] · [Feedback][github-issues-link] · [Contact Us][email-link] -[^1]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/#/) with 768-dimensional vectors and 95% recall. + + +[![][github-release-shield]][github-release-link] +[![][docker-release-shield]][docker-release-link] +[![][docker-pulls-shield]][docker-pulls-link] +[![][discord-shield]][discord-link] +[![][X-shield]][X-link] +[![][cloud-shield]][cloud-link] +[![][license-1-shield]][license-1-link] +[![][license-2-shield]][license-2-link] +
+ + +> [!NOTE] +> VectorChord serves as the successor to [pgvecto.rs](https://github.com/tensorchord/pgvecto.rs) [![][previous-docker-pulls-shield]][previous-docker-pulls-link] with better stability and performance. If you are interested in this new solution, you may find the [migration guide](https://docs.vectorchord.ai/vectorchord/admin/migration.html) helpful. + +VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search. + +With VectorChord, you can store 400,000 vectors for just $1, enabling significant savings: 6x more vectors compared to Pinecone's optimized storage and 26x more than pgvector/pgvecto.rs for the same price[^1]. + +![][image-compare] ## Features VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: -**⚡ Enhanced Performance**: Delivering optimized operations with up to 5x faster queries, 16x higher insert throughput, and 16x quicker[^3] index building compared to pgvector's HNSW implementation. +**⚡ Enhanced Performance**: Delivering optimized operations with up to 5x faster queries, 16x higher insert throughput, and 16x quicker[^1] index building compared to pgvector's HNSW implementation. -[^3]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/#/) with 768-dimensional vectors. Please checkout our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. +[^1]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/#/) with 768-dimensional vectors and 95% recall. Please checkout our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. **💰 Affordable Vector Search**: Query 100M 768-dimensional vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%, helping you keep infrastructure costs down while maintaining high search quality. **🔌 Seamless Integration**: Fully compatible with pgvector data types and syntax while providing optimal defaults out of the box - no manual parameter tuning needed. Just drop in VectorChord for enhanced performance. -**🔧 External Index Build**: Leverage IVF to build indexes externally (e.g., on GPU) for faster KMeans clustering, combined with RaBitQ[^2] compression to efficiently store vectors while maintaining search quality through autonomous reranking. +**🔧 Accelerated Index Build**: Leverage IVF to build indexes externally (e.g., on GPU) for faster KMeans clustering, combined with RaBitQ[^3] compression to efficiently store vectors while maintaining search quality through autonomous reranking. + +[^3]: Gao, Jianyang, and Cheng Long. "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search." Proceedings of the ACM on Management of Data 2.3 (2024): 1-27. + +**📏 Long Vector Support**: Store and search vectors up to 60,000[^4] dimensions, enabling the use of the best high-dimensional models like text-embedding-3-large with ease. + +[^4]: There is a [limitation](https://github.com/pgvector/pgvector#vector-type) at pgvector of 16,000 dimensions now. If you really have a large dimension(`16,000 [!TIP] +> If you are using the official [Docker image](https://hub.docker.com/r/tensorchord/vchord-postgres), you can skip this step. + +VectorChord depends on [pgvector](https://github.com/pgvector/pgvector), ensure the pgvector extension is available: -**📏 Long Vector Support**: Store and search vectors up to 65,535 dimensions, enabling the use of the best high-dimensional models like text-embedding-3-large with ease. +```SQL +SELECT * FROM pg_available_extensions WHERE name = 'vector'; +``` +If pgvector is not available, install it using the [pgvector installation instructions](https://github.com/pgvector/pgvector#installation). -[^2]: Gao, Jianyang, and Cheng Long. "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search." Proceedings of the ACM on Management of Data 2.3 (2024): 1-27. +And make sure to add `vchord.so` to the `shared_preload_libraries` in `postgresql.conf`. + +```SQL +-- Add vchord and pgvector to shared_preload_libraries -- +-- Note: A restart is required for this setting to take effect. +ALTER SYSTEM SET shared_preload_libraries = 'vchord.so'; +``` ## Quick Start For new users, we recommend using the Docker image to get started quickly. @@ -41,7 +81,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d tensorchord/vchord-postgres:pg17-v0.1.0 + -d tensorchord/vchord-postgres:pg17-v0.2.1 ``` Then you can connect to the database using the `psql` command line tool. The default username is `postgres`, and the default password is `mysecretpassword`. @@ -49,41 +89,141 @@ Then you can connect to the database using the `psql` command line tool. The def ```bash psql -h localhost -p 5432 -U postgres ``` -Run the following SQL to ensure the extension is enabled. + +Now you can play with VectorChord! + +## Documentation + +- [Installation](#installation) + - [Docker](#installation) + - [APT](#apt) + - [More Methods](#more-methods) +- [Usage](#usage) + - [Storing](#storing) + - [Indexing](#indexing) + - [Query](#query) +- [Performance Tuning](#performance-tuning) + - [Index Build Time](#index-build-time) + - [Query Performance](#query-performance) +- [Advanced Features](#advanced-features) + - [Indexing Prewarm](#indexing-prewarm) + - [Indexing Progress](#indexing-progress) + - [External Index Precomputation](#external-index-precomputation) + + - [Range Query](#range-query) + +## Installation + +### [Docker](https://docs.vectorchord.ai/vectorchord/getting-started/installation.html#docker) + +You can easily get the Docker image from: + +```bash +docker pull tensorchord/vchord-postgres:pg17-v0.2.1 +``` + +### [APT](https://docs.vectorchord.ai/vectorchord/getting-started/installation.html#from-debian-package) + +Debian and Ubuntu packages can be found on [release page](https://github.com/tensorchord/VectorChord/releases). + +To install it: +```bash +wget https://github.com/tensorchord/VectorChord/releases/download/0.2.1/postgresql-17-vchord_0.2.1-1_amd64.deb +sudo apt install postgresql-17-vchord_*.deb +``` + +### More Methods + +VectorChord also supports other installation methods, including: +- [From ZIP package](https://docs.vectorchord.ai/vectorchord/getting-started/installation.html#from-zip-package) + +## Usage + +VectorChord depends on pgvector, including the vector representation. +This way, we can keep the maximum compatibility of `pgvector` for both: +- [vector storing](https://github.com/pgvector/pgvector#storing) +- [vector query](https://github.com/pgvector/pgvector#querying). + +Since you can use them directly, your application can be easily migrated without pain! + +Before all, you need to run the following SQL to ensure the extension is enabled. ```SQL CREATE EXTENSION IF NOT EXISTS vchord CASCADE; ``` +It will install both `pgvector` and `VectorChord`, see [requirements](#requirements) for more detail. -And make sure to add `vchord.so` to the `shared_preload_libraries` in `postgresql.conf`. +### Storing -```SQL --- Add vchord and pgvector to shared_preload_libraries -- -ALTER SYSTEM SET shared_preload_libraries = 'vchord.so'; +Similar to pgvector, you can create a table with vector column in VectorChord and insert some rows to it. + +```sql +CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3)); +INSERT INTO items (embedding) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); ``` +### Indexing + +Similar to [ivfflat](https://github.com/pgvector/pgvector#ivfflat), the index type of VectorChord, RaBitQ(vchordrq) also divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It inherits the advantages of `ivfflat`, such as fast build times and less memory usage, but has [much better performance](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql#heading-ivf-vs-hnsw) than hnsw and ivfflat. + +The RaBitQ(vchordrq) index is supported on some pgvector types and metrics: + +| | vector | halfvec | bit(n) | sparsevec | +| ----------------------- | ------ | ------- | ------ | --------- | +| L2 distance / `<->` | ✅ | ✅ | 🆖 | ❌ | +| inner product / `<#>` | ✅ | ✅ | 🆖 | ❌ | +| cosine distance / `<=>` | ✅ | ✅ | 🆖 | ❌ | +| L1 distance / `<+>` | ❌ | ❌ | 🆖 | ❌ | +| Hamming distance/ `<~>` | 🆖 | 🆖 | ❌ | 🆖 | +| Jaccard distance/ `<%>` | 🆖 | 🆖 | ❌ | 🆖 | + +Where: +- ✅ means supported by pgvector and VectorChord +- ❌ means supported by pgvector but not by VectorChord +- 🆖 means not planned by pgvector and VectorChord +- 🔜 means supported by pgvector now and will be supported by VectorChord soon + To create the VectorChord RaBitQ(vchordrq) index, you can use the following SQL. -```SQL --- Set residual_quantization to true and spherical_centroids to false for L2 distance -- -CREATE INDEX ON gist_train USING vchordrq (embedding vector_l2_ops) WITH (options = $$ +L2 distance +```sql +CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops) WITH (options = $$ residual_quantization = true [build.internal] -lists = [4096] +lists = [1000] spherical_centroids = false $$); +``` +> [!NOTE] +> - Set `residual_quantization` to true and `spherical_centroids` to false for L2 distance +> - Use `halfvec_l2_ops` for `halfvec` +> - The recommend `lists` could be rows / 1000 for up to 1M rows and 4 * sqrt(rows) for over 1M rows --- Set residual_quantization to false and spherical_centroids to true for cos/dot distance -- -CREATE INDEX ON laion USING vchordrq (embedding vector_cos_ops) WITH (options = $$ +Inner product +```sql +CREATE INDEX ON items USING vchordrq (embedding vector_ip_ops) WITH (options = $$ residual_quantization = false [build.internal] -lists = [4096] +lists = [1000] spherical_centroids = true $$); ``` -## Documentation +Cosine distance +```sql +CREATE INDEX ON items USING vchordrq (embedding vector_cosine_ops) WITH (options = $$ +residual_quantization = false +[build.internal] +lists = [1000] +spherical_centroids = true +$$); +``` + +> [!NOTE] +> - Set `residual_quantization` to false and `spherical_centroids` to true for inner product/cosine distance +> - Use `vector_cosine_ops`/`vector_ip_ops` for `halfvec` ### Query @@ -96,20 +236,23 @@ Supported distance functions are: - <#> - (negative) inner product - <=> - cosine distance - +-- Set this to the number of CPU cores available for parallel operations. +SET max_parallel_maintenance_workers = 8; +SET max_parallel_workers = 8; -### Query Performance Tuning +-- Adjust the total number of worker processes. +-- Note: A restart is required for this setting to take effect. +ALTER SYSTEM SET max_worker_processes = 8; +``` + +## Query Performance You can fine-tune the search performance by adjusting the `probes` and `epsilon` parameters: ```sql @@ -118,8 +261,8 @@ You can fine-tune the search performance by adjusting the `probes` and `epsilon` SET vchordrq.probes = 100; -- Set epsilon to control the reranking precision. --- Larger value means more rerank for higher recall rate. --- Don't change it unless you only have limited memory. +-- Larger value means more rerank for higher recall rate and latency. +-- If you need a less precise query, set it to 1.0 might be appropriate. -- Recommended range: 1.0–1.9. Default value is 1.9. SET vchordrq.epsilon = 1.9; @@ -146,25 +289,18 @@ SET jit = off; ALTER SYSTEM SET shared_buffers = '8GB'; ``` +## Advanced Features + ### Indexing prewarm -To prewarm the index, you can use the following SQL. It will significantly improve performance when using limited memory. -```SQL --- vchordrq_prewarm(index_name::regclass) to prewarm the index into the shared buffer -SELECT vchordrq_prewarm('gist_train_embedding_idx'::regclass)" -``` +For disk-first indexing, RaBitQ(vchordrq) is loaded from disk for the first query, +and then cached in memory if `shared_buffer` is sufficient, resulting in a significant cold-start slowdown. -### Index Build Time -Index building can parallelized, and with external centroid precomputation, the total time is primarily limited by disk speed. Optimize parallelism using the following settings: +To improve performance for the first query, you can try the following SQL that preloads the index into memory. ```SQL --- Set this to the number of CPU cores available for parallel operations. -SET max_parallel_maintenance_workers = 8; -SET max_parallel_workers = 8; - --- Adjust the total number of worker processes. --- Note: A restart is required for this setting to take effect. -ALTER SYSTEM SET max_worker_processes = 8; +-- vchordrq_prewarm(index_name::regclass) to prewarm the index into the shared buffer +SELECT vchordrq_prewarm('gist_train_embedding_idx'::regclass)" ``` ### Indexing Progress @@ -204,6 +340,50 @@ $$); To simplify the workflow, we provide end-to-end scripts for external index pre-computation, see [scripts](./scripts/README.md#run-external-index-precomputation-toolkit). + + +### Range Query + +To query vectors within a certain distance range, you can use the following syntax. +```SQL +-- Query vectors within a certain distance range +SELECT vec FROM t WHERE vec <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012) +ORDER BY embedding <-> '[0.24, 0.24, 0.24]' LIMIT 5; +``` + +In this expression, `vec <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012)` is equals to `vec <-> '[0.24, 0.24, 0.24]' < 0.012`. However, the latter one will trigger a **exact nearest neighbor search** as the grammar could not be pushed down. + +Supported range functions are: +- `<<->>` - L2 distance +- `<<#>>` - (negative) inner product +- `<<=>>` - cosine distance + +## Development + ### Build the Postgres Docker Image with VectorChord extension Follow the steps in [Dev Guidance](./scripts/README.md#build-docker). @@ -230,3 +410,29 @@ This software is licensed under a dual license model: 2. **Elastic License v2 (ELv2)**: You may also use, modify, and distribute this software under the Elastic License v2, which has specific restrictions. You may choose either license based on your needs. We welcome any commercial collaboration or support, so please email us with any questions or requests regarding the licenses. + + +[image-compare]: https://github.com/user-attachments/assets/2d985f1e-7093-4c3a-8bf3-9f0b92c0e7e7 +[license-1-link]: https://github.com/tensorchord/VectorChord#license +[license-1-shield]: https://img.shields.io/badge/License-AGPLv3-green?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCIgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiBmaWxsPSIjZmZmZmZmIj48cGF0aCBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik0xMi43NSAyLjc1YS43NS43NSAwIDAwLTEuNSAwVjQuNUg5LjI3NmExLjc1IDEuNzUgMCAwMC0uOTg1LjMwM0w2LjU5NiA1Ljk1N0EuMjUuMjUgMCAwMTYuNDU1IDZIMi4zNTNhLjc1Ljc1IDAgMTAwIDEuNUgzLjkzTC41NjMgMTUuMThhLjc2Mi43NjIgMCAwMC4yMS44OGMuMDguMDY0LjE2MS4xMjUuMzA5LjIyMS4xODYuMTIxLjQ1Mi4yNzguNzkyLjQzMy42OC4zMTEgMS42NjIuNjIgMi44NzYuNjJhNi45MTkgNi45MTkgMCAwMDIuODc2LS42MmMuMzQtLjE1NS42MDYtLjMxMi43OTItLjQzMy4xNS0uMDk3LjIzLS4xNTguMzEtLjIyM2EuNzUuNzUgMCAwMC4yMDktLjg3OEw1LjU2OSA3LjVoLjg4NmMuMzUxIDAgLjY5NC0uMTA2Ljk4NC0uMzAzbDEuNjk2LTEuMTU0QS4yNS4yNSAwIDAxOS4yNzUgNmgxLjk3NXYxNC41SDYuNzYzYS43NS43NSAwIDAwMCAxLjVoMTAuNDc0YS43NS43NSAwIDAwMC0xLjVIMTIuNzVWNmgxLjk3NGMuMDUgMCAuMS4wMTUuMTQuMDQzbDEuNjk3IDEuMTU0Yy4yOS4xOTcuNjMzLjMwMy45ODQuMzAzaC44ODZsLTMuMzY4IDcuNjhhLjc1Ljc1IDAgMDAuMjMuODk2Yy4wMTIuMDA5IDAgMCAuMDAyIDBhMy4xNTQgMy4xNTQgMCAwMC4zMS4yMDZjLjE4NS4xMTIuNDUuMjU2Ljc5LjRhNy4zNDMgNy4zNDMgMCAwMDIuODU1LjU2OCA3LjM0MyA3LjM0MyAwIDAwMi44NTYtLjU2OWMuMzM4LS4xNDMuNjA0LS4yODcuNzktLjM5OWEzLjUgMy41IDAgMDAuMzEtLjIwNi43NS43NSAwIDAwLjIzLS44OTZMMjAuMDcgNy41aDEuNTc4YS43NS43NSAwIDAwMC0xLjVoLTQuMTAyYS4yNS4yNSAwIDAxLS4xNC0uMDQzbC0xLjY5Ny0xLjE1NGExLjc1IDEuNzUgMCAwMC0uOTg0LS4zMDNIMTIuNzVWMi43NXpNMi4xOTMgMTUuMTk4YTUuNDE4IDUuNDE4IDAgMDAyLjU1Ny42MzUgNS40MTggNS40MTggMCAwMDIuNTU3LS42MzVMNC43NSA5LjM2OGwtMi41NTcgNS44M3ptMTQuNTEtLjAyNGMuMDgyLjA0LjE3NC4wODMuMjc1LjEyNi41My4yMjMgMS4zMDUuNDUgMi4yNzIuNDVhNS44NDYgNS44NDYgMCAwMDIuNTQ3LS41NzZMMTkuMjUgOS4zNjdsLTIuNTQ3IDUuODA3eiI+PC9wYXRoPjwvc3ZnPgo= +[license-2-link]: https://github.com/tensorchord/VectorChord#license +[license-2-shield]: https://img.shields.io/badge/License-ELv2-green?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCIgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiBmaWxsPSIjZmZmZmZmIj48cGF0aCBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik0xMi43NSAyLjc1YS43NS43NSAwIDAwLTEuNSAwVjQuNUg5LjI3NmExLjc1IDEuNzUgMCAwMC0uOTg1LjMwM0w2LjU5NiA1Ljk1N0EuMjUuMjUgMCAwMTYuNDU1IDZIMi4zNTNhLjc1Ljc1IDAgMTAwIDEuNUgzLjkzTC41NjMgMTUuMThhLjc2Mi43NjIgMCAwMC4yMS44OGMuMDguMDY0LjE2MS4xMjUuMzA5LjIyMS4xODYuMTIxLjQ1Mi4yNzguNzkyLjQzMy42OC4zMTEgMS42NjIuNjIgMi44NzYuNjJhNi45MTkgNi45MTkgMCAwMDIuODc2LS42MmMuMzQtLjE1NS42MDYtLjMxMi43OTItLjQzMy4xNS0uMDk3LjIzLS4xNTguMzEtLjIyM2EuNzUuNzUgMCAwMC4yMDktLjg3OEw1LjU2OSA3LjVoLjg4NmMuMzUxIDAgLjY5NC0uMTA2Ljk4NC0uMzAzbDEuNjk2LTEuMTU0QS4yNS4yNSAwIDAxOS4yNzUgNmgxLjk3NXYxNC41SDYuNzYzYS43NS43NSAwIDAwMCAxLjVoMTAuNDc0YS43NS43NSAwIDAwMC0xLjVIMTIuNzVWNmgxLjk3NGMuMDUgMCAuMS4wMTUuMTQuMDQzbDEuNjk3IDEuMTU0Yy4yOS4xOTcuNjMzLjMwMy45ODQuMzAzaC44ODZsLTMuMzY4IDcuNjhhLjc1Ljc1IDAgMDAuMjMuODk2Yy4wMTIuMDA5IDAgMCAuMDAyIDBhMy4xNTQgMy4xNTQgMCAwMC4zMS4yMDZjLjE4NS4xMTIuNDUuMjU2Ljc5LjRhNy4zNDMgNy4zNDMgMCAwMDIuODU1LjU2OCA3LjM0MyA3LjM0MyAwIDAwMi44NTYtLjU2OWMuMzM4LS4xNDMuNjA0LS4yODcuNzktLjM5OWEzLjUgMy41IDAgMDAuMzEtLjIwNi43NS43NSAwIDAwLjIzLS44OTZMMjAuMDcgNy41aDEuNTc4YS43NS43NSAwIDAwMC0xLjVoLTQuMTAyYS4yNS4yNSAwIDAxLS4xNC0uMDQzbC0xLjY5Ny0xLjE1NGExLjc1IDEuNzUgMCAwMC0uOTg0LS4zMDNIMTIuNzVWMi43NXpNMi4xOTMgMTUuMTk4YTUuNDE4IDUuNDE4IDAgMDAyLjU1Ny42MzUgNS40MTggNS40MTggMCAwMDIuNTU3LS42MzVMNC43NSA5LjM2OGwtMi41NTcgNS44M3ptMTQuNTEtLjAyNGMuMDgyLjA0LjE3NC4wODMuMjc1LjEyNi41My4yMjMgMS4zMDUuNDUgMi4yNzIuNDVhNS44NDYgNS44NDYgMCAwMDIuNTQ3LS41NzZMMTkuMjUgOS4zNjdsLTIuNTQ3IDUuODA3eiI+PC9wYXRoPjwvc3ZnPgo= + +[docker-release-link]: https://hub.docker.com/r/tensorchord/vchord-postgres +[docker-release-shield]: https://img.shields.io/docker/v/tensorchord/vchord-postgres?color=369eff&label=docker&labelColor=black&logo=docker&logoColor=white&style=flat&sort=semver +[github-release-link]: https://github.com/tensorchord/VectorChord/releases +[github-release-shield]: https://img.shields.io/github/v/release/tensorchord/VectorChord?color=369eff&labelColor=black&logo=github&style=flat +[docker-pulls-link]: https://hub.docker.com/r/tensorchord/vchord-postgres +[docker-pulls-shield]: https://img.shields.io/docker/pulls/tensorchord/vchord-postgres?color=45cc11&labelColor=black&style=flat&sort=semver +[previous-docker-pulls-link]: https://hub.docker.com/r/tensorchord/pgvecto-rs +[previous-docker-pulls-shield]: https://img.shields.io/docker/pulls/tensorchord/pgvecto-rs?color=45cc11&labelColor=black&style=flat&sort=semver +[discord-link]: https://discord.gg/KqswhpVgdU +[discord-shield]: https://img.shields.io/discord/974584200327991326?&logoColor=white&color=5865F2&style=flat&logo=discord&cacheSeconds=60 +[X-link]: https://twitter.com/TensorChord +[X-shield]: https://img.shields.io/twitter/follow/tensorchord?style=flat&logo=X&cacheSeconds=60 +[cloud-link]: https://cloud.vectorchord.ai/ +[cloud-shield]: https://img.shields.io/badge/VectorChord_Cloud-Try_For_Free-F2B263.svg?labelColor=DAFDBA&logo=data:image/svg%2bxml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSIxMzMiIGhlaWdodD0iMTMyIiBmaWxsPSJub25lIj48cGF0aCBmaWxsPSIjRTZEQjNEIiBkPSJNNDguNCAzNy41YzAtMSAwLTEuNS0uMi0xLjhhMSAxIDAgMCAwLS44LS40Yy0uMyAwLS43LjMtMS42IDFMMjcuNiA1MC4xYy0xLjIuOC0xLjcgMS4zLTIuMiAxLjhhNSA1IDAgMCAwLS44IDEuNmMtLjIuNy0uMiAxLjQtLjIgMi45djM3LjNjMCAxLjIgMCAxLjguMyAyIC4yLjMuNS41LjguNC40IDAgLjgtLjQgMS42LTEuM2wxOS0xOC42IDEuNS0xLjhjLjMtLjQuNS0xIC42LTEuNC4yLS42LjItMS4yLjItMi40VjM3LjVaTTM1LjIgMTA1LjNjLS44LjgtMS4yIDEuMy0xLjIgMS42IDAgLjQgMCAuNy4zLjkuMy4yLjkuMiAyIC4yaDM3YzEuMyAwIDIgMCAyLjUtLjJhNSA1IDAgMCAwIDEuNS0uNmMuNi0uNCAxLS45IDEuOS0xLjhMOTYuNiA4NmMuNy0uOSAxLjEtMS4zIDEuMS0xLjZhMSAxIDAgMCAwLS4zLS44Yy0uMy0uMy0uOS0uMy0yLS4zaC0zNWMtMS4yIDAtMS44IDAtMi40LjJhNSA1IDAgMCAwLTEuNC42Yy0uNS4zLTEgLjctMS44IDEuNmwtMTkuNiAxOS42Wk05Ni4zIDcwLjFjMSAwIDEuNCAwIDEuNy0uMi40LS4xLjYtLjQuOC0uN2wuMS0xLjdWMzUuM2wtLjEtMS44Yy0uMi0uMy0uNC0uNS0uOC0uNy0uMy0uMi0uOC0uMi0xLjctLjJINjQuMWMtMSAwLTEuNCAwLTEuNy4yLS40LjItLjYuNC0uOC43bC0uMSAxLjh2MzIuMmwuMSAxLjdjLjIuMy40LjYuOC43LjMuMi44LjIgMS43LjJoMzIuMloiLz48cGF0aCBmaWxsPSIjMTAxNTA5IiBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik00My4yIDIxLjVjLTEuMyAwLTIgMC0yLjMtLjMtLjMtLjMtLjUtLjYtLjUtMSAwLS41LjQtMSAxLjEtMi4xTDUzLjEgMi4zYy42LS44LjktMS4yIDEuMi0xLjNoMWMuNC4xLjcuNSAxLjIgMS4zbDExLjYgMTUuOGMuOCAxIDEuMiAxLjYgMS4yIDIgMCAuNS0uMi44LS41IDEtLjQuNC0xIC40LTIuNC40SDU5Yy0uMy4yLS41LjQtLjYuNy0uMi4zLS4yLjYtLjIgMS40VjY4YzAgMS45IDAgMi44LjQgMy41LjMuNi44IDEuMSAxLjQgMS41LjcuMyAxLjcuMyAzLjUuM2g0NGwxLjMtLjFjLjMtLjEuNS0uMy42LS42bC4xLTEuNFY2NWMwLTEuMyAwLTIgLjMtMi4zLjItLjMuNi0uNSAxLS41czEgLjMgMiAxTDEzMCA3NWMuOC42IDEuMy45IDEuNCAxLjJ2MWMtLjEuNC0uNi43LTEuNCAxLjNsLTE3LjMgMTEuN2MtMSAuOC0xLjYgMS4xLTIgMS4xLS40IDAtLjgtLjItMS0uNS0uMy0uNC0uMy0xLS4zLTIuM1Y4MmwtLjEtMS40Yy0uMS0uMi0uMy0uNC0uNi0uNS0uMy0uMi0uNi0uMi0xLjQtLjJINjAuNWMtMS42IDAtMi40IDAtMy4xLjItLjcuMi0xLjMuNC0yIC44bC0yLjMgMi0yOC42IDI4LjNjLS41LjUtLjguOC0uOSAxLjF2LjhjLjEuMy40LjYuOSAxLjFsMy43IDMuN2MuOCAxIDEuMyAxLjMgMS4zIDEuOCAwIC40IDAgLjctLjMgMS0uMy4zLS45LjUtMiAuOGwtMjAgNC43Yy0xLjEuMi0xLjcuMy0yIC4yLS40LS4xLS43LS40LS44LS44LS4xLS4zIDAtMSAuMy0ybDUtMTkuOWMuMy0xLjEuNC0xLjcuOC0yIC4zLS4zLjctLjQgMS0uMy41IDAgLjkuNSAxLjggMS40bDMuNiAzLjcgMSAuOWguOWMuMy0uMS42LS40IDEtLjlsMjguNi0yOC4yYzEuMi0xLjEgMS43LTEuNyAyLjEtMi4zLjQtLjYuNy0xLjMuOC0yIC4yLS43LjItMS41LjItMy4yVjIzLjZsLS4xLTEuNGMtLjEtLjMtLjMtLjUtLjYtLjZsLTEuNC0uMWgtNi4yWiIgY2xpcC1ydWxlPSJldmVub2RkIi8+PC9zdmc+ +[blog-link]: https://blog.vectorchord.ai/ +[official-site-link]: https://vectorchord.ai/ +[github-issues-link]: https://github.com/tensorchord/VectorChord/issues +[email-link]: mailto:support@tensorchord.ai \ No newline at end of file diff --git a/scripts/README.md b/scripts/README.md index 56ca24ee..5d8bae62 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -87,7 +87,7 @@ conda install conda-forge::pgvector-python numpy pytorch::faiss-gpu conda-forge: python script/train.py -i [dataset file(export.hdf5)] -o [centroid filename(centroid.npy)] --lists [lists] -m [metric(l2/cos/dot)] -g --mmap ``` - `lists` is the number of centroids for clustering, and a typical value could range from: + `lists` is the number of centroids for clustering, and a typical value for large datasets(>5M) could range from: $$ 4*\sqrt{len(vectors)} \le lists \le 16*\sqrt{len(vectors)} @@ -96,13 +96,11 @@ conda install conda-forge::pgvector-python numpy pytorch::faiss-gpu conda-forge: 3. To insert vectors and centroids into the database, and then create an index ```shell - python script/index.py -n [table name] -i [dataset file(export.hdf5)] -c [centroid filename(centroid.npy)] -m [metric(l2/cos/dot)] -d [dim] + python script/index.py -n [table name] -i [dataset file(export.hdf5)] -c [centroid filename(centroid.npy)] -m [metric(l2/cos/dot)] -d [dim] --url postgresql://postgres:123@localhost:5432/postgres ``` 4. Let's start our tour to check the benchmark result of VectorChord ```shell - python script/bench.py -n [table name] -i [dataset file(export.hdf5)] -m [metric(l2/cos/dot)] -p [database password] --nprob 100 --epsilon 1.0 - ``` - - Larger `nprobe` and `epsilon` will have a more precise query but at a slower speed. \ No newline at end of file + python script/bench.py -n [table name] -i [dataset file(export.hdf5)] -m [metric(l2/cos/dot)] --nprob 100 --epsilon 1.0 --url postgresql://postgres:123@localhost:5432/postgres + ``` \ No newline at end of file diff --git a/scripts/bench.py b/scripts/bench.py index 4e6eca5e..d6ff03ce 100644 --- a/scripts/bench.py +++ b/scripts/bench.py @@ -22,7 +22,9 @@ def build_arg_parse(): parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) parser.add_argument("-i", "--input", help="input filepath", required=True) parser.add_argument( - "-p", "--password", help="Database password", default="password" + "--url", + help="url, like `postgresql://postgres:123@localhost:5432/postgres`", + required=True, ) parser.add_argument( "-t", "--top", help="Dimension", type=int, choices=[10, 100], default=10 @@ -39,7 +41,7 @@ def build_arg_parse(): return parser -def create_connection(password, nprob, epsilon): +def create_connection(url, nprob, epsilon): keepalive_kwargs = { "keepalives": 1, "keepalives_idle": 30, @@ -47,33 +49,58 @@ def create_connection(password, nprob, epsilon): "keepalives_count": 5, } conn = psycopg.connect( - conninfo=f"postgresql://postgres:{password}@localhost:5432/postgres", + conninfo=url, dbname="postgres", autocommit=True, **keepalive_kwargs, ) - conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") + try: + conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") + except Exception as e: + pass # Tuning conn.execute("SET jit=false") conn.execute("SET effective_io_concurrency=200") conn.execute(f"SET vchordrq.probes={nprob}") conn.execute(f"SET vchordrq.epsilon={epsilon}") - conn.execute(f"SELECT vchordrq_prewarm('{args.name}_embedding_idx'::regclass)") + try: + conn.execute(f"SELECT vchordrq_prewarm('{args.name}_embedding_idx'::regclass)") + except Exception as e: + pass register_vector(conn) return conn +def calculate_coverage(time_intervals): + if not time_intervals: + return 0 + sorted_intervals = sorted(time_intervals, key=lambda x: x[0]) + merged = [] + current_start, current_end = sorted_intervals[0] + for interval in sorted_intervals[1:]: + next_start, next_end = interval + if next_start <= current_end: + current_end = max(current_end, next_end) + else: + merged.append((current_start, current_end)) + current_start, current_end = next_start, next_end + merged.append((current_start, current_end)) + total_length = 0 + for start, end in merged: + total_length += end - start + return total_length + + def process_batch(args): """Process a batch of queries in a single process""" - batch_queries, batch_answers, k, metric_ops, password, name, nprob, epsilon = args + batch_queries, batch_answers, k, metric_ops, url, name, nprob, epsilon = args # Create a new connection for this process - conn = create_connection(password, nprob, epsilon) + conn = create_connection(url, nprob, epsilon) hits = 0 - latencies = [] results = [] for query, ground_truth in zip(batch_queries, batch_answers): @@ -84,29 +111,33 @@ def process_batch(args): ).fetchall() end = time.perf_counter() - query_time = end - start - latencies.append(query_time) - result_ids = set([p[0] for p in result[:k]]) ground_truth_ids = set(ground_truth[:k].tolist()) hit = len(result_ids & ground_truth_ids) hits += hit - results.append((hit, query_time)) + results.append((hit, (start, end))) conn.close() return results -def calculate_metrics(all_results, k, m): +def calculate_metrics(all_results, k, m, num_processes=1): """Calculate recall, QPS, and latency percentiles from results""" hits, latencies = zip(*all_results) + if isinstance(latencies[0], list | tuple): + # parallel_bench + total_time = calculate_coverage(latencies) + latencies = [(end - start) for start, end in latencies] + else: + # sequential_bench + total_time = sum(latencies) + total_hits = sum(hits) - total_time = sum(latencies) - recall = total_hits / (k * m) - qps = m / total_time + recall = total_hits / (k * m * num_processes) + qps = (m * num_processes) / total_time # Calculate latency percentiles (in milliseconds) latencies_ms = np.array(latencies) * 1000 @@ -117,25 +148,19 @@ def calculate_metrics(all_results, k, m): def parallel_bench( - name, test, answer, metric_ops, num_processes, password, top, nprob, epsilon + name, test, answer, metric_ops, num_processes, url, top, nprob, epsilon ): """Run benchmark in parallel using multiple processes""" m = test.shape[0] - - # Split data into batches for each process - batch_size = m // num_processes batches = [] - for i in range(num_processes): - start_idx = i * batch_size - end_idx = start_idx + batch_size if i < num_processes - 1 else m - + for _ in range(num_processes): batch = ( - test[start_idx:end_idx], - answer[start_idx:end_idx], + test, + answer, top, metric_ops, - password, + url, name, nprob, epsilon, @@ -144,23 +169,17 @@ def parallel_bench( # Create process pool and execute batches with mp.Pool(processes=num_processes) as pool: - batch_results = list( - tqdm( - pool.imap(process_batch, batches), - total=len(batches), - desc=f"Processing k={top}", - ) - ) + batch_results = list(pool.map(process_batch, batches)) # Flatten results from all batches all_results = [result for batch in batch_results for result in batch] # Calculate metrics - recall, qps, p50, p99 = calculate_metrics(all_results, top, m) + recall, qps, p50, p99 = calculate_metrics(all_results, top, m, num_processes) print(f"Top: {top}") print(f" Recall: {recall:.4f}") - print(f" QPS: {qps*num_processes:.2f}") + print(f" QPS: {qps:.2f}") print(f" P50 latency: {p50:.2f}ms") print(f" P99 latency: {p99:.2f}ms") @@ -224,11 +243,11 @@ def sequential_bench(name, test, answer, metric_ops, conn, top): answer, metric_ops, args.processes, - args.password, + args.url, args.top, args.nprob, args.epsilon, ) else: - conn = create_connection(args.password, args.nprob, args.epsilon) - sequential_bench(args.name, test, answer, metric_ops, conn, args.top) + conn = create_connection(args.url, args.nprob, args.epsilon) + sequential_bench(args.name, test, answer, metric_ops, conn, args.top) \ No newline at end of file diff --git a/scripts/index.py b/scripts/index.py index cb4f111f..058a25f4 100644 --- a/scripts/index.py +++ b/scripts/index.py @@ -32,7 +32,7 @@ def build_arg_parse(): parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) parser.add_argument("-i", "--input", help="Input filepath", required=True) parser.add_argument( - "-p", "--password", help="Database password", default="password" + "--url", help="url, like `postgresql://postgres:123@localhost:5432/postgres`", required=True ) parser.add_argument("-d", "--dim", help="Dimension", type=int, required=True) # Remember to set `max_worker_processes` at server start @@ -195,8 +195,7 @@ async def monitor_index_build(conn, finish: asyncio.Event): async def main(dataset): dataset = h5py.File(Path(args.input), "r") - url = f"postgresql://postgres:{args.password}@localhost:5432/postgres" - conn = await create_connection(url) + conn = await create_connection(args.url) if args.centroids: centroids = np.load(args.centroids, allow_pickle=False) await add_centroids(conn, args.name, centroids) @@ -207,7 +206,7 @@ async def main(dataset): index_finish = asyncio.Event() # Need a separate connection for monitor process - monitor_conn = await create_connection(url) + monitor_conn = await create_connection(args.url) monitor_task = monitor_index_build( monitor_conn, index_finish, From 860925e7e09d5637e069ae2da32e8f65882f7531 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Thu, 6 Mar 2025 17:19:25 +0800 Subject: [PATCH 120/324] fix: lint by ruff (#201) Signed-off-by: cutecutecat --- .github/workflows/check.yml | 2 ++ scripts/bench.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 40b173c5..22804d5f 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -6,6 +6,7 @@ on: - ".cargo" - ".github/workflows/check.yml" - "crates/**" + - "scripts/**" - "src/**" - "tests/**" - "build.rs" @@ -20,6 +21,7 @@ on: - ".cargo" - ".github/workflows/check.yml" - "crates/**" + - "scripts/**" - "src/**" - "tests/**" - "build.rs" diff --git a/scripts/bench.py b/scripts/bench.py index d6ff03ce..ee35e685 100644 --- a/scripts/bench.py +++ b/scripts/bench.py @@ -57,7 +57,7 @@ def create_connection(url, nprob, epsilon): try: conn.execute("CREATE EXTENSION IF NOT EXISTS vector") conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") - except Exception as e: + except Exception: pass # Tuning conn.execute("SET jit=false") @@ -67,7 +67,7 @@ def create_connection(url, nprob, epsilon): conn.execute(f"SET vchordrq.epsilon={epsilon}") try: conn.execute(f"SELECT vchordrq_prewarm('{args.name}_embedding_idx'::regclass)") - except Exception as e: + except Exception: pass register_vector(conn) return conn From 8d224d78b38da4ab6e821142d907a815ff9f4a17 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Tue, 11 Mar 2025 10:44:52 +0800 Subject: [PATCH 121/324] fix: wrong operators for halfvec (#202) Signed-off-by: cutecutecat --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c2e9ebd4..14edbb2d 100644 --- a/README.md +++ b/README.md @@ -223,7 +223,7 @@ $$); > [!NOTE] > - Set `residual_quantization` to false and `spherical_centroids` to true for inner product/cosine distance -> - Use `vector_cosine_ops`/`vector_ip_ops` for `halfvec` +> - Use `halfvec_cosine_ops`/`halfvec_ip_ops` for `halfvec` ### Query From b961f6bbb9cce2c8108e381e1d77c27b736f3b22 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 17 Mar 2025 09:58:49 +0800 Subject: [PATCH 122/324] doc: add new options (#207) Signed-off-by: usamoi Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- README.md | 52 ++++++++++-- crates/k_means/src/lib.rs | 8 +- scripts/README.md | 104 +++++++++++++++--------- src/index/am/am_build.rs | 2 +- src/index/types.rs | 7 ++ tools/dev.md | 167 +++++++++++++++++++++++++++----------- 6 files changed, 242 insertions(+), 98 deletions(-) diff --git a/README.md b/README.md index 14edbb2d..ab3816e8 100644 --- a/README.md +++ b/README.md @@ -95,7 +95,7 @@ Now you can play with VectorChord! ## Documentation - [Installation](#installation) - - [Docker](#installation) + - [Docker](#docker) - [APT](#apt) - [More Methods](#more-methods) - [Usage](#usage) @@ -105,6 +105,7 @@ Now you can play with VectorChord! - [Performance Tuning](#performance-tuning) - [Index Build Time](#index-build-time) - [Query Performance](#query-performance) + - [Index Internal Build](#index-internal-build) - [Advanced Features](#advanced-features) - [Indexing Prewarm](#indexing-prewarm) - [Indexing Progress](#indexing-progress) @@ -197,7 +198,8 @@ $$); ``` > [!NOTE] -> - Set `residual_quantization` to true and `spherical_centroids` to false for L2 distance +> - `options` are specified using a [TOML: Tom's Obvious Minimal Language](https://toml.io/) string. +> - Set `residual_quantization` to `true` and `spherical_centroids` to `false` for L2 distance embeddings > - Use `halfvec_l2_ops` for `halfvec` > - The recommend `lists` could be rows / 1000 for up to 1M rows and 4 * sqrt(rows) for over 1M rows @@ -222,7 +224,7 @@ $$); ``` > [!NOTE] -> - Set `residual_quantization` to false and `spherical_centroids` to true for inner product/cosine distance +> - Set `residual_quantization` to `false` and `spherical_centroids` to `true` for cosine similarity embeddings > - Use `halfvec_cosine_ops`/`halfvec_ip_ops` for `halfvec` ### Query @@ -252,7 +254,16 @@ SET max_parallel_workers = 8; ALTER SYSTEM SET max_worker_processes = 8; ``` -## Query Performance +When building an index on a table with more than 10 million vectors, you can choose to consume more memory to accelerate the process by setting `build.pin` to `true`: + +```sql +CREATE INDEX ON items USING vchordrq (embedding vector_cosine_ops) WITH (options = $$ +residual_quantization = false +build.pin = true +$$); +``` + +### Query Performance You can fine-tune the search performance by adjusting the `probes` and `epsilon` parameters: ```sql @@ -289,9 +300,35 @@ SET jit = off; ALTER SYSTEM SET shared_buffers = '8GB'; ``` +### Index Internal Build + +There are 4 options for index internal build, except `lists`. + +* `spherical_centroids`: Specifies whether to use spherical K-means algorithm. If the model generates cosine similarity embeddings, set this to `true`; otherwise, set to `false`. Possible values: `true`, `false`. Default: `false`. + +* `sampling_factor`: Specifies the number of vectors sampled by K-means algorithm. The higher this value, the slower the build, the greater the memory consumption, and the better search performance. Possible values: any integer between `1` and `1024`. Default: `256`. + +* `kmeans_iterations`: Specifies the number of iterations for K-means algorithm. The higher this value, the slower the build. Possible values: any integer between `0` and `1024`. Default: `10`. + +* `build_threads`: Specifies the number of threads used by K-means algorithm. The higher this value, the faster the build. Possible values: any integer between `1` and `255`. Default: `1`. + +An example of these options: + +```sql +CREATE INDEX ON t USING vchordrq (embedding vector_cosine_ops) WITH (options = $$ +residual_quantization = false +[build.internal] +lists = [1000] +spherical_centroids = true +sampling_factor = 512 +kmeans_iterations = 25 +build_threads = 16 +$$); +``` + ## Advanced Features -### Indexing prewarm +### Indexing Prewarm For disk-first indexing, RaBitQ(vchordrq) is loaded from disk for the first query, and then cached in memory if `shared_buffer` is sufficient, resulting in a significant cold-start slowdown. @@ -393,11 +430,12 @@ Follow the steps in [Dev Guidance](./scripts/README.md#build-docker). Install pgrx according to [pgrx's instruction](https://github.com/pgcentralfoundation/pgrx?tab=readme-ov-file#getting-started). ```bash cargo install --locked cargo-pgrx -cargo pgrx init --pg17 $(which pg_config) # To init with system postgres, with pg_config in PATH +cargo pgrx init --pg$(pg_config --version | awk '{print $2}' | cut -d. -f1) $(which pg_config) # To init with system postgres, with pg_config in PATH cargo pgrx install --release --sudo # To install the extension into the system postgres with sudo ``` ## Limitations + - KMeans Clustering: The built-in KMeans clustering depends on multi-thread in-memory build and may require substantial memory. We strongly recommend using external centroid precomputation for efficient index construction. @@ -435,4 +473,4 @@ You may choose either license based on your needs. We welcome any commercial col [blog-link]: https://blog.vectorchord.ai/ [official-site-link]: https://vectorchord.ai/ [github-issues-link]: https://github.com/tensorchord/VectorChord/issues -[email-link]: mailto:support@tensorchord.ai \ No newline at end of file +[email-link]: mailto:support@tensorchord.ai diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index ecbea29b..7f61328b 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -26,7 +26,7 @@ pub fn k_means( .num_threads(num_threads) .build_scoped( |thread| thread.run(), - move |_| { + move |pool| { let compute = |centroids: &[Vec]| { if n >= 1000 && c >= 1000 { rabitq_index(dims, n, c, samples, centroids) @@ -35,14 +35,14 @@ pub fn k_means( } }; let mut lloyd_k_means = - LloydKMeans::new(c, dims, samples, is_spherical, compute); + pool.install(|| LloydKMeans::new(c, dims, samples, is_spherical, compute)); for _ in 0..iterations { check(); - if lloyd_k_means.iterate() { + if pool.install(|| lloyd_k_means.iterate()) { break; } } - lloyd_k_means.finish() + pool.install(|| lloyd_k_means.finish()) }, ) .expect("failed to build thread pool") diff --git a/scripts/README.md b/scripts/README.md index 5d8bae62..dbfd8a32 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -1,49 +1,77 @@ ## Build Docker -Users can choose to build the package with the provided docker image or create the development environment by themselves. - -- (option 1) With `vectorchord-pgrx` Image - -```shell -# use the required version of `pgrx` and `rust` -export PGRX_VERSION=$(awk -F'version = "=|"' '/^pgrx\s*=.*version/ {print $2}' Cargo.toml) -export RUST_TOOLCHAIN=$(awk -F'"' '/^\s*channel\s*=/ {print $2}' rust-toolchain.toml) -export PGRX_IMAGE=ghcr.io/tensorchord/vectorchord-pgrx:$PGRX_VERSION-$RUST_TOOLCHAIN - -docker run --rm -v .:/workspace $PGRX_IMAGE cargo build --lib --features pg16 --release -docker run --rm -v .:/workspace $PGRX_IMAGE ./tools/schema.sh --features pg16 --release -``` - -- (option 2) With Local Development Environment - -```shell -sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config clang -cargo install --locked cargo-pgrx -cargo pgrx init -cargo build --package vchord --lib --features pg16 --release -./tools/schema.sh --features pg16 --release -``` - -- build the debian package - -```shell -export SEMVER="0.0.0" -export VERSION="16" -export ARCH="x86_64" -export PLATFORM="amd64" -./tools/package.sh -``` - -- build the docker image - ```shell -docker build -t vchord:pg16-latest --build-arg PG_VERSION=16 -f ./docker/Dockerfile . +SEMVER='0.2.1' +VERSION='17' +ARCH='x86_64' +PLATFORM='amd64' + +git clone https://github.com/tensorchord/VectorChord.git +cd VectorChord +git checkout $SEMVER + +sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config + +sudo apt-get install -y postgresql-common +sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y +sudo apt-get install -y postgresql-server-dev-$VERSION +sudo apt-get install -y postgresql-$VERSION +sudo apt-get install -y postgresql-$VERSION-pgvector + +curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 +sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 + +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + +cargo install cargo-pgrx@$(sed -n 's/.*pgrx = { version = "\(=.*\)",.*/\1/p' Cargo.toml) --locked +cargo pgrx init --pg$VERSION=$(which pg_config) + +cargo build --lib --features pg$VERSION --release + +mkdir -p ./build/zip +cp -a ./sql/upgrade/. ./build/zip/ +cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql +sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control +cp ./target/release/libvchord.so ./build/zip/vchord.so +zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* + +mkdir -p ./build/deb +mkdir -p ./build/deb/DEBIAN +mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ +mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ +for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file +done +for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file +done +for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file +done +echo "Package: postgresql-${VERSION}-vchord +Version: ${SEMVER}-1 +Section: database +Priority: optional +Architecture: ${PLATFORM} +Maintainer: Tensorchord +Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM +Homepage: https://vectorchord.ai/ +License: AGPL-3 or Elastic-2" \ +> ./build/deb/DEBIAN/control +(cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums +dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb + +ls ./build + +docker build -t vchord:pg$VERSION-latest --build-arg PG_VERSION=$VERSION -f ./docker/Dockerfile . ``` ## Run Instance ```shell -docker run --name vchord -e POSTGRES_PASSWORD=123 -p 5432:5432 -d vchord:pg16-latest +VERSION='17' + +docker run --name vchord -e POSTGRES_PASSWORD=123 -p 5432:5432 -d vchord:pg$VERSION-latest ``` ## Run External Index Precomputation Toolkit diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index e8dab2d1..d28058ff 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -1055,7 +1055,7 @@ pub fn make_internal_build( &samples }, internal_build.spherical_centroids, - 10, + internal_build.kmeans_iterations as _, ); if let Some(structure) = result.last() { let mut children = vec![Vec::new(); means.len()]; diff --git a/src/index/types.rs b/src/index/types.rs index df238462..8edec722 100644 --- a/src/index/types.rs +++ b/src/index/types.rs @@ -13,6 +13,9 @@ pub struct VchordrqInternalBuildOptions { #[serde(default = "VchordrqInternalBuildOptions::default_sampling_factor")] #[validate(range(min = 1, max = 1024))] pub sampling_factor: u32, + #[serde(default = "VchordrqInternalBuildOptions::default_kmeans_iterations")] + #[validate(range(min = 0, max = 1024))] + pub kmeans_iterations: u32, #[serde(default = "VchordrqInternalBuildOptions::default_build_threads")] #[validate(range(min = 1, max = 255))] pub build_threads: u16, @@ -37,6 +40,9 @@ impl VchordrqInternalBuildOptions { fn default_sampling_factor() -> u32 { 256 } + fn default_kmeans_iterations() -> u32 { + 10 + } fn default_build_threads() -> u16 { 1 } @@ -48,6 +54,7 @@ impl Default for VchordrqInternalBuildOptions { lists: Self::default_lists(), spherical_centroids: Self::default_spherical_centroids(), sampling_factor: Self::default_sampling_factor(), + kmeans_iterations: Self::default_kmeans_iterations(), build_threads: Self::default_build_threads(), } } diff --git a/tools/dev.md b/tools/dev.md index 8558ad1c..70c0107d 100644 --- a/tools/dev.md +++ b/tools/dev.md @@ -3,85 +3,156 @@ ### Pre-requisite ```shell -sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config clang postgresql-server-dev-16 -cargo install --locked cargo-pgrx -cargo pgrx init +SEMVER='0.2.1' +VERSION='17' + +git clone https://github.com/tensorchord/VectorChord.git +cd VectorChord +git checkout $SEMVER + +sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config + +sudo apt-get install -y postgresql-common +sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y +sudo apt-get install -y postgresql-server-dev-$VERSION +sudo apt-get install -y postgresql-$VERSION +sudo apt-get install -y postgresql-$VERSION-pgvector + +curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 +sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 + +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + +cargo install cargo-pgrx@$(sed -n 's/.*pgrx = { version = "\(=.*\)",.*/\1/p' Cargo.toml) --locked +cargo pgrx init --pg$VERSION=$(which pg_config) ``` ### Step 1 - generate install schema -```shell -# Change to version to release -SEMVER=0.2.0 -PREV_VERSION=0.1.0 - -mkdir temp -cargo build --package vchord --lib --features pg16 --release -./tools/schema.sh --features pg16 --release +```shell +SEMVER='0.2.1' +VERSION='17' -cp ./target/release/schema.sql ./sql/install/vchord--$SEMVER.sql +cargo build --lib --features pg$VERSION --release +cargo pgrx schema --features pg$VERSION | expand -t 4 > ./sql/install/vchord--$SEMVER.sql ``` ### Step 2 - generate upgrade schema + ```shell -PREV_VERSION_FILE=sql/install/vchord--$PREV_VERSION.sql +PREV='0.2.0' +SEMVER='0.2.1' -# New lines redirect to install.sql, revised lines & deleted lines redirect to terminal -diff -u $PREV_VERSION_FILE install.sql | awk ' +diff -u ./sql/install/vchord--$PREV.sql ./sql/install/vchord--$SEMVER.sql | awk ' /^\+/ && !/^+++/ { - print substr($0, 2) > "temp/upgrade.sql" + print substr($0, 2) > "target/upgrade.sql" next } /^-/ && !/^---/ || /^@/ { print } { next } ' -cp temp/upgrade.sql ./sql/upgrade/vchord--$PREV_VERSION--$SEMVER.sql +cp target/upgrade.sql ./sql/upgrade/vchord--$PREV--$SEMVER.sql ``` ### Step 3 - validate + ```shell +PREV='0.2.0' +SEMVER='0.2.1' -# sudo pg_dropcluster --stop 16 main -sudo pg_createcluster 16 main --start sudo -u postgres createdb vchord -sudo -u postgres psql -d vchord -c "ALTER USER postgres WITH PASSWORD '123';" - -# Dump upgraded schema and compare it -export PGHOST=localhost -export PGPASSWORD=123 -export PGUSER=postgres -./tools/dump.sh $PREV_VERSION $SEMVER > temp/dump1.sql -./tools/dump.sh $SEMVER > temp/dump2.sql -code --diff temp/dump1.sql temp/dump2.sql + +sudo -u postgres ./tools/dump.sh $PREV $SEMVER > target/upgrade.sql +sudo -u postgres ./tools/dump.sh $SEMVER > target/install.sql +code --diff target/upgrade.sql target/install.sql + +sudo -u postgres dropdb vchord +``` + +### Step 4 - package and download + +```shell +SEMVER='0.2.1' +VERSION='17' +ARCH='x86_64' +PLATFORM='amd64' + +cargo build --lib --features pg$VERSION --release + +mkdir -p ./build/zip +cp -a ./sql/upgrade/. ./build/zip/ +cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql +sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control +cp ./target/release/libvchord.so ./build/zip/vchord.so +zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* + +mkdir -p ./build/deb +mkdir -p ./build/deb/DEBIAN +mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ +mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ +for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file +done +for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file +done +for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file +done +echo "Package: postgresql-${VERSION}-vchord +Version: ${SEMVER}-1 +Section: database +Priority: optional +Architecture: ${PLATFORM} +Maintainer: Tensorchord +Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM +Homepage: https://vectorchord.ai/ +License: AGPL-3 or Elastic-2" \ +> ./build/deb/DEBIAN/control +(cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums +dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb + +ls ./build + +wget https://github.com/tensorchord/VectorChord/releases/download/${PREV}/postgresql-${VERSION}-vchord_${PREV}-1_${PLATFORM}.deb -O ./build/postgresql-${VERSION}-vchord_${PREV}-1_${PLATFORM}.deb ``` -### Step 4 - further test +### Step 5 - further test ```shell -sudo apt remove -y vchord-pg16 -SEMVER=$SEMVER VERSION="16" ARCH="x86_64" PLATFORM="amd64" ./tools/package.sh -cd temp -wget https://github.com/tensorchord/VectorChord/releases/download/"$PREV_VERSION"/vchord-pg16_"$PREV_VERSION"_amd64.deb -sudo apt install -y ./vchord-pg16_"$PREV_VERSION"_amd64.deb +PREV='0.2.0' +SEMVER='0.2.1' +VERSION='17' +ARCH='x86_64' +PLATFORM='amd64' + +cargo install sqllogictest-bin + +# upgrade test + +sudo apt install -y ./build/postgresql-${VERSION}-vchord_${PREV}-1_${PLATFORM}.deb sudo -u postgres psql -d vchord -c "ALTER SYSTEM SET SHARED_PRELOAD_LIBRARIES='vchord.so';" -sudo systemctl restart postgresql@16-main.service -sudo -u postgres psql -d vchord -c 'CREATE EXTENSION IF NOT EXISTS vector; CREATE EXTENSION IF NOT EXISTS vchord;' +sudo systemctl restart postgresql@$VERSION-main.service +sudo -u postgres psql -d vchord -c 'CREATE EXTENSION vchord CASCADE;' sudo -u postgres psql -d vchord -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord';" +sudo -u postgres sqllogictest -d vchord '../tests/**/*.slt' -# Run Test -- Upgrade -sudo apt install -f ../build/vchord-pg16_"$SEMVER"_amd64.deb -sudo systemctl restart postgresql@16-main.service +sudo apt install -y ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb +sudo systemctl restart postgresql@$VERSION-main.service sudo -u postgres psql -d vchord -c 'ALTER EXTENSION vchord UPDATE;' sudo -u postgres psql -d vchord -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord';" -sqllogictest -h localhost -u postgres -d vchord -w 123 '../tests/**/*.slt' +sudo -u postgres sqllogictest -d vchord '../tests/**/*.slt' + +sudo apt remove -y postgresql-${VERSION}-vchord -# Run Test -- Install -sudo -u postgres psql -d vchord -c 'DROP EXTENSION vchord CASCADE;' -sudo -u postgres psql -d vchord -c 'CREATE EXTENSION vchord;' +# install test + +sudo apt install -y ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb +sudo -u postgres psql -d vchord -c "ALTER SYSTEM SET SHARED_PRELOAD_LIBRARIES='vchord.so';" +sudo systemctl restart postgresql@$VERSION-main.service +sudo -u postgres psql -d vchord -c 'CREATE EXTENSION vchord CASCADE;' sudo -u postgres psql -d vchord -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord';" -sqllogictest -h localhost -u postgres -d vchord -w 123 '../tests/**/*.slt' +sudo -u postgres sqllogictest -d vchord '../tests/**/*.slt' -sudo -u postgres psql -d vchord -c "DROP SCHEMA IF EXISTS public CASCADE;" -sudo apt remove -y vchord-pg16 -cd .. && rm -rf temp -``` \ No newline at end of file +sudo apt remove -y postgresql-${VERSION}-vchord +``` From faa3628bb0d01fcaa04b4d57b03b4f273f8f747c Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 17 Mar 2025 09:59:40 +0800 Subject: [PATCH 123/324] fix: do not leak memory in heap fetch and fix reading tuples in prewarm (#205) Signed-off-by: usamoi --- crates/algorithm/src/prewarm.rs | 2 +- crates/algorithm/src/rerank.rs | 4 +- src/datatype/memory_halfvec.rs | 16 + src/datatype/memory_vector.rs | 16 + src/index/algorithm.rs | 179 +++++++++ src/index/am/am_build.rs | 394 ++++-------------- src/index/am/am_scan.rs | 692 -------------------------------- src/index/am/mod.rs | 325 +++++++++++---- src/index/functions.rs | 29 +- src/index/mod.rs | 2 + src/index/opclass.rs | 145 ++++--- src/index/projection.rs | 24 -- src/index/scanners/default.rs | 227 +++++++++++ src/index/scanners/mod.rs | 38 ++ src/lib.rs | 1 + 15 files changed, 901 insertions(+), 1193 deletions(-) create mode 100644 src/index/algorithm.rs delete mode 100644 src/index/am/am_scan.rs create mode 100644 src/index/scanners/default.rs create mode 100644 src/index/scanners/mod.rs diff --git a/crates/algorithm/src/prewarm.rs b/crates/algorithm/src/prewarm.rs index 9395c060..e611f740 100644 --- a/crates/algorithm/src/prewarm.rs +++ b/crates/algorithm/src/prewarm.rs @@ -89,7 +89,7 @@ pub fn prewarm(index: impl RelationRead, height: i32, check: impl F ); tape::read_appendable_tape( index.clone(), - jump_tuple.frozen_first(), + jump_tuple.appendable_first(), |_| (), |(), _, _| { results.push(()); diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index 2e409cb8..4484c9aa 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -46,10 +46,10 @@ pub fn rerank_heap( AlwaysEqual, AlwaysEqual, )>, - fetch: F, + mut fetch: F, ) -> impl Iterator where - F: Fn(NonZeroU64) -> Option, + F: FnMut(NonZeroU64) -> Option, { let mut heap = BinaryHeap::from(results); let mut cache = BinaryHeap::<(Reverse, _)>::new(); diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs index 3e9fe09f..8fdae0a0 100644 --- a/src/datatype/memory_halfvec.rs +++ b/src/datatype/memory_halfvec.rs @@ -144,6 +144,22 @@ impl IntoDatum for HalfvecOutput { // UnboxDatum +unsafe impl<'a> pgrx::datum::UnboxDatum for HalfvecInput<'a> { + type As<'src> + = HalfvecInput<'src> + where + 'a: 'src; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + unsafe impl pgrx::datum::UnboxDatum for HalfvecOutput { type As<'src> = HalfvecOutput; #[inline] diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs index 4d9f9f21..16628605 100644 --- a/src/datatype/memory_vector.rs +++ b/src/datatype/memory_vector.rs @@ -143,6 +143,22 @@ impl IntoDatum for VectorOutput { // UnboxDatum +unsafe impl<'a> pgrx::datum::UnboxDatum for VectorInput<'a> { + type As<'src> + = VectorInput<'src> + where + 'a: 'src; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + unsafe impl pgrx::datum::UnboxDatum for VectorOutput { type As<'src> = VectorOutput; #[inline] diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs new file mode 100644 index 00000000..bbd5f7aa --- /dev/null +++ b/src/index/algorithm.rs @@ -0,0 +1,179 @@ +use super::opclass::Opfamily; +use crate::index::am::am_build::InternalBuild; +use algorithm::operator::{Dot, L2, Op}; +use algorithm::types::*; +use algorithm::{RelationRead, RelationWrite}; +use half::f16; +use std::num::NonZeroU64; +use vector::VectorOwned; +use vector::vect::{VectBorrowed, VectOwned}; + +pub fn prewarm( + opfamily: Opfamily, + index: impl RelationRead, + height: i32, + check: impl Fn(), +) -> String { + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + algorithm::prewarm::, L2>>(index, height, check) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + algorithm::prewarm::, Dot>>(index, height, check) + } + (VectorKind::Vecf16, DistanceKind::L2) => { + algorithm::prewarm::, L2>>(index, height, check) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + algorithm::prewarm::, Dot>>(index, height, check) + } + } +} + +pub fn bulkdelete( + opfamily: Opfamily, + index: impl RelationWrite, + check: impl Fn(), + callback: impl Fn(NonZeroU64) -> bool, +) { + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + algorithm::bulkdelete::, L2>>(index, check, callback) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + algorithm::bulkdelete::, Dot>>(index, check, callback) + } + (VectorKind::Vecf16, DistanceKind::L2) => { + algorithm::bulkdelete::, L2>>(index, check, callback) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + algorithm::bulkdelete::, Dot>>(index, check, callback) + } + } +} + +pub fn maintain(opfamily: Opfamily, index: impl RelationWrite, check: impl Fn()) { + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + algorithm::maintain::, L2>>(index, check) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + algorithm::maintain::, Dot>>(index, check) + } + (VectorKind::Vecf16, DistanceKind::L2) => { + algorithm::maintain::, L2>>(index, check) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + algorithm::maintain::, Dot>>(index, check) + } + } +} + +pub fn build( + vector_options: VectorOptions, + vchordrq_options: VchordrqIndexOptions, + index: impl RelationWrite, + structures: Vec>>, +) { + match (vector_options.v, vector_options.d) { + (VectorKind::Vecf32, DistanceKind::L2) => algorithm::build::, L2>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + (VectorKind::Vecf32, DistanceKind::Dot) => algorithm::build::, Dot>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + (VectorKind::Vecf16, DistanceKind::L2) => algorithm::build::, L2>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + (VectorKind::Vecf16, DistanceKind::Dot) => algorithm::build::, Dot>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + } +} + +pub fn insert( + opfamily: Opfamily, + index: impl RelationWrite, + payload: NonZeroU64, + vector: OwnedVector, +) { + match (vector, opfamily.distance_kind()) { + (OwnedVector::Vecf32(vector), DistanceKind::L2) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + algorithm::insert::, L2>>( + index, + payload, + RandomProject::project(vector.as_borrowed()), + ) + } + (OwnedVector::Vecf32(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + algorithm::insert::, Dot>>( + index, + payload, + RandomProject::project(vector.as_borrowed()), + ) + } + (OwnedVector::Vecf16(vector), DistanceKind::L2) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + algorithm::insert::, L2>>( + index, + payload, + RandomProject::project(vector.as_borrowed()), + ) + } + (OwnedVector::Vecf16(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + algorithm::insert::, Dot>>( + index, + payload, + RandomProject::project(vector.as_borrowed()), + ) + } + } +} + +fn map_structures(x: Vec>, f: impl Fn(T) -> U + Copy) -> Vec> { + x.into_iter() + .map(|Structure { means, children }| Structure { + means: means.into_iter().map(f).collect(), + children, + }) + .collect() +} + +pub trait RandomProject { + type Output; + fn project(self) -> Self::Output; +} + +impl RandomProject for VectBorrowed<'_, f32> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use crate::index::projection::project; + let input = self.slice(); + VectOwned::new(project(input)) + } +} + +impl RandomProject for VectBorrowed<'_, f16> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use crate::index::projection::project; + use simd::Floating; + let input = f16::vector_to_f32(self.slice()); + VectOwned::new(f16::vector_from_f32(&project(&input))) + } +} diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index d28058ff..18e30f2a 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -1,20 +1,17 @@ use crate::datatype::typmod::Typmod; use crate::index::am::{Reloption, ctid_to_pointer}; use crate::index::opclass::{Opfamily, opfamily}; -use crate::index::projection::RandomProject; use crate::index::storage::{PostgresPage, PostgresRelation}; use crate::index::types::*; -use algorithm::operator::{Dot, L2, Op, Vector}; use algorithm::types::*; use algorithm::{PageGuard, RelationRead, RelationWrite}; use half::f16; -use pgrx::pg_sys::Datum; +use pgrx::pg_sys::{Datum, ItemPointerData}; use rand::Rng; use simd::Floating; -use std::num::NonZeroU64; use std::ops::Deref; +use vector::VectorOwned; use vector::vect::VectOwned; -use vector::{VectorBorrowed, VectorOwned}; #[derive(Debug, Clone)] struct Heap { @@ -26,13 +23,17 @@ struct Heap { } impl Heap { - fn traverse(&self, progress: bool, callback: F) { + fn traverse))>( + &self, + progress: bool, + callback: F, + ) { pub struct State<'a, F> { pub this: &'a Heap, pub callback: F, } #[pgrx::pg_guard] - unsafe extern "C" fn call( + unsafe extern "C" fn call( _index_relation: pgrx::pg_sys::Relation, ctid: pgrx::pg_sys::ItemPointer, values: *mut Datum, @@ -40,14 +41,14 @@ impl Heap { _tuple_is_alive: bool, state: *mut core::ffi::c_void, ) where - F: FnMut((NonZeroU64, V)), + F: FnMut((ItemPointerData, Vec<(OwnedVector, u16)>)), { let state = unsafe { &mut *state.cast::>() }; let opfamily = state.this.opfamily; - let vector = unsafe { opfamily.input_vector(*values.add(0), *is_null.add(0)) }; - let pointer = unsafe { ctid_to_pointer(ctid.read()) }; - if let Some(vector) = vector { - (state.callback)((pointer, V::from_owned(vector))); + let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; + let ctid = unsafe { ctid.read() }; + if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { + (state.callback)((ctid, store)); } } let table_am = unsafe { &*(*self.heap_relation).rd_tableam }; @@ -65,7 +66,7 @@ impl Heap { progress, 0, pgrx::pg_sys::InvalidBlockNumber, - Some(call::), + Some(call::), (&mut state) as *mut State as *mut _, self.scan, ); @@ -142,76 +143,39 @@ pub unsafe extern "C" fn ambuild( }; let mut samples = Vec::new(); let mut number_of_samples = 0_u32; - match opfamily.vector_kind() { - VectorKind::Vecf32 => { - heap.traverse(false, |(_, vector): (_, VectOwned)| { - let vector = vector.as_borrowed(); - assert_eq!( - vector_options.dims, - vector.dims(), - "invalid vector dimensions" - ); - if number_of_samples < max_number_of_samples { - samples.push(VectOwned::::build_to_vecf32(vector)); - number_of_samples += 1; - } else { - let index = rand.random_range(0..max_number_of_samples) as usize; - samples[index] = VectOwned::::build_to_vecf32(vector); - } - tuples_total += 1; - }); - } - VectorKind::Vecf16 => { - heap.traverse(false, |(_, vector): (_, VectOwned)| { - let vector = vector.as_borrowed(); - assert_eq!( - vector_options.dims, - vector.dims(), - "invalid vector dimensions" - ); - if number_of_samples < max_number_of_samples { - samples.push(VectOwned::::build_to_vecf32(vector)); - number_of_samples += 1; - } else { - let index = rand.random_range(0..max_number_of_samples) as usize; - samples[index] = VectOwned::::build_to_vecf32(vector); - } - tuples_total += 1; - }); + heap.traverse(false, |(_, store)| { + for (vector, _) in store { + let x = match vector { + OwnedVector::Vecf32(x) => VectOwned::build_to_vecf32(x.as_borrowed()), + OwnedVector::Vecf16(x) => VectOwned::build_to_vecf32(x.as_borrowed()), + }; + assert_eq!( + vector_options.dims, + x.len() as u32, + "invalid vector dimensions" + ); + if number_of_samples < max_number_of_samples { + samples.push(x); + number_of_samples += 1; + } else { + let index = rand.random_range(0..max_number_of_samples) as usize; + samples[index] = x; + } } - } + tuples_total += 1; + }); samples }; reporter.tuples_total(tuples_total); make_internal_build(vector_options.clone(), internal_build.clone(), samples) } }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => algorithm::build::, L2>>( - vector_options, - vchordrq_options.index, - index.clone(), - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), - ), - (VectorKind::Vecf32, DistanceKind::Dot) => algorithm::build::, Dot>>( - vector_options, - vchordrq_options.index, - index.clone(), - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), - ), - (VectorKind::Vecf16, DistanceKind::L2) => algorithm::build::, L2>>( - vector_options, - vchordrq_options.index, - index.clone(), - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), - ), - (VectorKind::Vecf16, DistanceKind::Dot) => algorithm::build::, Dot>>( - vector_options, - vchordrq_options.index, - index.clone(), - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), - ), - } + crate::index::algorithm::build( + vector_options, + vchordrq_options.index, + index.clone(), + structures, + ); let cache = if vchordrq_options.build.pin { let mut trace = algorithm::cache(index.clone()); trace.sort(); @@ -268,70 +232,19 @@ pub unsafe extern "C" fn ambuild( let mut indtuples = 0; reporter.tuples_done(indtuples); let relation = unsafe { PostgresRelation::new(index_relation) }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, L2>>( - relation.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, Dot>>( - relation.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, L2>>( - relation.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }); + heap.traverse(true, |(ctid, store)| { + for (vector, extra) in store { + let payload = ctid_to_pointer(ctid, extra); + crate::index::algorithm::insert(opfamily, relation.clone(), payload, vector); } - (VectorKind::Vecf16, DistanceKind::Dot) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, Dot>>( - relation.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - indtuples += 1; - reporter.tuples_done(indtuples); - }); - } - } + indtuples += 1; + reporter.tuples_done(indtuples); + }); } let check = || { pgrx::check_for_interrupts!(); }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::maintain::, L2>>(index, check); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::maintain::, Dot>>(index, check); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::maintain::, L2>>(index, check); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::maintain::, Dot>>(index, check); - } - } + crate::index::algorithm::maintain(opfamily, index, check); unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -797,183 +710,49 @@ unsafe fn parallel_build( scan, }; match cached { - VchordrqCachedReader::_0(_) => match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, L2>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, Dot>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } + VchordrqCachedReader::_0(_) => { + heap.traverse(true, |(ctid, store)| { + for (vector, extra) in store { + let payload = ctid_to_pointer(ctid, extra); + crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); + } + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } - }); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, L2>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); } - }); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, Dot>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }); - } - }, + } + }); + } VchordrqCachedReader::_1(cached) => { let index = CachingRelation { cache: cached, relation: index, }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, L2>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }); + heap.traverse(true, |(ctid, store)| { + for (vector, extra) in store { + let payload = ctid_to_pointer(ctid, extra); + crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); } - (VectorKind::Vecf32, DistanceKind::Dot) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, Dot>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, L2>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - heap.traverse(true, |(pointer, vector): (_, VectOwned)| { - algorithm::insert::, Dot>>( - index.clone(), - pointer, - RandomProject::project(vector.as_borrowed()), - ); - unsafe { - let indtuples; - { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).indtuples += 1; - indtuples = (*vchordrqshared).indtuples; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } - } - }); + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + if let Some(reporter) = reporter.as_mut() { + reporter.tuples_done(indtuples); + } } - } + }); } } unsafe { @@ -1240,15 +1019,6 @@ pub fn make_external_build( result } -pub fn map_structures(x: Vec>, f: impl Fn(T) -> U + Copy) -> Vec> { - x.into_iter() - .map(|Structure { means, children }| Structure { - means: means.into_iter().map(f).collect(), - children, - }) - .collect() -} - pub trait InternalBuild: VectorOwned { fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec; diff --git a/src/index/am/am_scan.rs b/src/index/am/am_scan.rs deleted file mode 100644 index 02b13b6f..00000000 --- a/src/index/am/am_scan.rs +++ /dev/null @@ -1,692 +0,0 @@ -use crate::index::am::pointer_to_ctid; -use crate::index::gucs::{epsilon, max_scan_tuples, probes}; -use crate::index::opclass::{Opfamily, Sphere, opfamily}; -use crate::index::projection::RandomProject; -use crate::index::storage::PostgresRelation; -use algorithm::RerankMethod; -use algorithm::operator::{Dot, L2, Op, Vector}; -use algorithm::types::*; -use half::f16; -use pgrx::pg_sys::Datum; -use std::cell::LazyCell; -use std::num::NonZeroU64; -use vector::VectorOwned; -use vector::vect::VectOwned; - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambeginscan( - index_relation: pgrx::pg_sys::Relation, - n_keys: std::os::raw::c_int, - n_orderbys: std::os::raw::c_int, -) -> pgrx::pg_sys::IndexScanDesc { - use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; - - let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index_relation, n_keys, n_orderbys) }; - unsafe { - let scanner = Scanner { - opfamily: opfamily(index_relation), - scanning: Scanning::Empty {}, - }; - (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); - } - scan -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amrescan( - scan: pgrx::pg_sys::IndexScanDesc, - keys: pgrx::pg_sys::ScanKey, - _n_keys: std::os::raw::c_int, - orderbys: pgrx::pg_sys::ScanKey, - _n_orderbys: std::os::raw::c_int, -) { - unsafe { - if !keys.is_null() && (*scan).numberOfKeys > 0 { - std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); - } - if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { - std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); - } - let opfamily = opfamily((*scan).indexRelation); - let (orderbys, spheres) = { - let mut orderbys = Vec::new(); - let mut spheres = Vec::new(); - if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { - pgrx::error!( - "vector search with no WHERE clause and no ORDER BY clause is not supported" - ); - } - for i in 0..(*scan).numberOfOrderBys { - let data = (*scan).orderByData.add(i as usize); - let value = (*data).sk_argument; - let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; - match (*data).sk_strategy { - 1 => orderbys.push(opfamily.input_vector(value, is_null)), - _ => unreachable!(), - } - } - for i in 0..(*scan).numberOfKeys { - let data = (*scan).keyData.add(i as usize); - let value = (*data).sk_argument; - let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; - match (*data).sk_strategy { - 2 => spheres.push(opfamily.input_sphere(value, is_null)), - _ => unreachable!(), - } - } - (orderbys, spheres) - }; - let scanning; - if let Some((vector, threshold, recheck)) = scanner_build(orderbys, spheres) { - scanning = Scanning::Initial { - vector, - threshold, - recheck, - }; - } else { - scanning = Scanning::Empty {}; - }; - let scanner = &mut *(*scan).opaque.cast::(); - scanner.scanning = scanning; - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amgettuple( - scan: pgrx::pg_sys::IndexScanDesc, - direction: pgrx::pg_sys::ScanDirection::Type, -) -> bool { - if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { - pgrx::error!("vector search without a forward scan direction is not supported"); - } - // https://www.postgresql.org/docs/current/index-locking.html - // If heap entries referenced physical pointers are deleted before - // they are consumed by PostgreSQL, PostgreSQL will received wrong - // physical pointers: no rows or irreverent rows are referenced. - if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC - { - pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); - } - let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; - let relation = unsafe { PostgresRelation::new((*scan).indexRelation) }; - if let Some((pointer, recheck)) = scanner_next( - scanner, - relation, - LazyCell::new(move || unsafe { - let index_info = pgrx::pg_sys::BuildIndexInfo((*scan).indexRelation); - let heap_relation = (*scan).heapRelation; - let estate = Scopeguard::new(pgrx::pg_sys::CreateExecutorState(), |estate| { - pgrx::pg_sys::FreeExecutorState(estate); - }); - let econtext = pgrx::pg_sys::MakePerTupleExprContext(*estate.get()); - move |opfamily: Opfamily, payload| { - let slot = Scopeguard::new( - pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), - |slot| pgrx::pg_sys::ExecDropSingleTupleTableSlot(slot), - ); - (*econtext).ecxt_scantuple = *slot.get(); - let table_am = (*heap_relation).rd_tableam; - let fetch_row_version = (*table_am).tuple_fetch_row_version.unwrap(); - let mut ctid = pointer_to_ctid(payload); - if !fetch_row_version(heap_relation, &mut ctid, (*scan).xs_snapshot, *slot.get()) { - return None; - } - let mut values = [Datum::from(0); 32]; - let mut is_null = [true; 32]; - pgrx::pg_sys::FormIndexDatum( - index_info, - *slot.get(), - *estate.get(), - values.as_mut_ptr(), - is_null.as_mut_ptr(), - ); - opfamily.input_vector(values[0], is_null[0]) - } - }), - ) { - let ctid = pointer_to_ctid(pointer); - unsafe { - (*scan).xs_heaptid = ctid; - (*scan).xs_recheckorderby = false; - (*scan).xs_recheck = recheck; - } - true - } else { - false - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { - let scanner = unsafe { &mut *(*scan).opaque.cast::() }; - scanner.scanning = Scanning::Empty {}; -} - -struct Scanner { - opfamily: Opfamily, - scanning: Scanning, -} - -enum Scanning { - Initial { - vector: OwnedVector, - threshold: Option, - recheck: bool, - }, - Vbase { - vbase: Box>, - recheck: bool, - }, - Empty {}, -} - -fn scanner_build( - orderbys: Vec>, - spheres: Vec>>, -) -> Option<(OwnedVector, Option, bool)> { - let mut vector = None; - let mut threshold = None; - let mut recheck = false; - for orderby_vector in orderbys.into_iter().flatten() { - if vector.is_none() { - vector = Some(orderby_vector); - } else { - pgrx::error!("vector search with multiple vectors is not supported"); - } - } - for Sphere { center, radius } in spheres.into_iter().flatten() { - if vector.is_none() { - (vector, threshold) = (Some(center), Some(radius)); - } else { - recheck = true; - } - } - Some((vector?, threshold, recheck)) -} - -fn scanner_next( - scanner: &mut Scanner, - relation: PostgresRelation, - fetch: LazyCell, -) -> Option<(NonZeroU64, bool)> -where - F: Fn(Opfamily, NonZeroU64) -> Option + 'static, - I: FnOnce() -> F + 'static, -{ - if let Scanning::Initial { - vector, - threshold, - recheck, - } = &scanner.scanning - { - let opfamily = scanner.opfamily; - let vector = vector.clone(); - let threshold = *threshold; - let recheck = *recheck; - let max_scan_tuples = max_scan_tuples(); - let probes = probes(); - let epsilon = epsilon(); - scanner.scanning = Scanning::Vbase { - vbase: match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - let vector = RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), - ); - let (method, results) = algorithm::search::, L2>>( - relation.clone(), - vector.clone(), - probes, - epsilon, - ); - match (method, max_scan_tuples, threshold) { - (RerankMethod::Index, None, None) => { - let vbase = algorithm::rerank_index::, L2>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.fuse()) as Box> - } - (RerankMethod::Index, None, Some(threshold)) => { - let vbase = algorithm::rerank_index::, L2>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take_while(move |(x, _)| *x < threshold)) - } - (RerankMethod::Index, Some(max_scan_tuples), None) => { - let vbase = algorithm::rerank_index::, L2>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take(max_scan_tuples as _)) - } - (RerankMethod::Index, Some(max_scan_tuples), Some(threshold)) => { - let vbase = algorithm::rerank_index::, L2>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ) - } - (RerankMethod::Heap, None, None) => { - let vbase = algorithm::rerank_heap::, L2>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.fuse()) as Box> - } - (RerankMethod::Heap, None, Some(threshold)) => { - let vbase = algorithm::rerank_heap::, L2>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take_while(move |(x, _)| *x < threshold)) - } - (RerankMethod::Heap, Some(max_scan_tuples), None) => { - let vbase = algorithm::rerank_heap::, L2>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take(max_scan_tuples as _)) - } - (RerankMethod::Heap, Some(max_scan_tuples), Some(threshold)) => { - let vbase = algorithm::rerank_heap::, L2>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ) - } - } - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - let vector = RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), - ); - let (method, results) = algorithm::search::, Dot>>( - relation.clone(), - vector.clone(), - probes, - epsilon, - ); - match (method, max_scan_tuples, threshold) { - (RerankMethod::Index, None, None) => { - let vbase = algorithm::rerank_index::, Dot>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.fuse()) as Box> - } - (RerankMethod::Index, None, Some(threshold)) => { - let vbase = algorithm::rerank_index::, Dot>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take_while(move |(x, _)| *x < threshold)) - } - (RerankMethod::Index, Some(max_scan_tuples), None) => { - let vbase = algorithm::rerank_index::, Dot>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take(max_scan_tuples as _)) - } - (RerankMethod::Index, Some(max_scan_tuples), Some(threshold)) => { - let vbase = algorithm::rerank_index::, Dot>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ) - } - (RerankMethod::Heap, None, None) => { - let vbase = algorithm::rerank_heap::, Dot>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.fuse()) as Box> - } - (RerankMethod::Heap, None, Some(threshold)) => { - let vbase = algorithm::rerank_heap::, Dot>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take_while(move |(x, _)| *x < threshold)) - } - (RerankMethod::Heap, Some(max_scan_tuples), None) => { - let vbase = algorithm::rerank_heap::, Dot>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take(max_scan_tuples as _)) - } - (RerankMethod::Heap, Some(max_scan_tuples), Some(threshold)) => { - let vbase = algorithm::rerank_heap::, Dot>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ) - } - } - } - (VectorKind::Vecf16, DistanceKind::L2) => { - let vector = RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), - ); - let (method, results) = algorithm::search::, L2>>( - relation.clone(), - vector.clone(), - probes, - epsilon, - ); - match (method, max_scan_tuples, threshold) { - (RerankMethod::Index, None, None) => { - let vbase = algorithm::rerank_index::, L2>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.fuse()) as Box> - } - (RerankMethod::Index, None, Some(threshold)) => { - let vbase = algorithm::rerank_index::, L2>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take_while(move |(x, _)| *x < threshold)) - } - (RerankMethod::Index, Some(max_scan_tuples), None) => { - let vbase = algorithm::rerank_index::, L2>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take(max_scan_tuples as _)) - } - (RerankMethod::Index, Some(max_scan_tuples), Some(threshold)) => { - let vbase = algorithm::rerank_index::, L2>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ) - } - (RerankMethod::Heap, None, None) => { - let vbase = algorithm::rerank_heap::, L2>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.fuse()) as Box> - } - (RerankMethod::Heap, None, Some(threshold)) => { - let vbase = algorithm::rerank_heap::, L2>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take_while(move |(x, _)| *x < threshold)) - } - (RerankMethod::Heap, Some(max_scan_tuples), None) => { - let vbase = algorithm::rerank_heap::, L2>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take(max_scan_tuples as _)) - } - (RerankMethod::Heap, Some(max_scan_tuples), Some(threshold)) => { - let vbase = algorithm::rerank_heap::, L2>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ) - } - } - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - let vector = RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), - ); - let (method, results) = algorithm::search::, Dot>>( - relation.clone(), - vector.clone(), - probes, - epsilon, - ); - match (method, max_scan_tuples, threshold) { - (RerankMethod::Index, None, None) => { - let vbase = algorithm::rerank_index::, Dot>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.fuse()) as Box> - } - (RerankMethod::Index, None, Some(threshold)) => { - let vbase = algorithm::rerank_index::, Dot>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take_while(move |(x, _)| *x < threshold)) - } - (RerankMethod::Index, Some(max_scan_tuples), None) => { - let vbase = algorithm::rerank_index::, Dot>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take(max_scan_tuples as _)) - } - (RerankMethod::Index, Some(max_scan_tuples), Some(threshold)) => { - let vbase = algorithm::rerank_index::, Dot>>( - relation, vector, results, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ) - } - (RerankMethod::Heap, None, None) => { - let vbase = algorithm::rerank_heap::, Dot>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.fuse()) as Box> - } - (RerankMethod::Heap, None, Some(threshold)) => { - let vbase = algorithm::rerank_heap::, Dot>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take_while(move |(x, _)| *x < threshold)) - } - (RerankMethod::Heap, Some(max_scan_tuples), None) => { - let vbase = algorithm::rerank_heap::, Dot>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new(vbase.take(max_scan_tuples as _)) - } - (RerankMethod::Heap, Some(max_scan_tuples), Some(threshold)) => { - let vbase = algorithm::rerank_heap::, Dot>, _>( - vector, - results, - move |payload| { - Some(RandomProject::project( - VectOwned::::from_owned(fetch(opfamily, payload)?) - .as_borrowed(), - )) - }, - ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)); - Box::new( - vbase - .take_while(move |(x, _)| *x < threshold) - .take(max_scan_tuples as _), - ) - } - } - } - }, - recheck, - }; - } - match &mut scanner.scanning { - Scanning::Initial { .. } => unreachable!(), - Scanning::Vbase { vbase, recheck } => vbase.next().map(|(_, x)| (x, *recheck)), - Scanning::Empty {} => None, - } -} - -struct Scopeguard -where - T: Copy, - F: FnMut(T), -{ - t: T, - f: F, -} - -impl Scopeguard -where - T: Copy, - F: FnMut(T), -{ - fn new(t: T, f: F) -> Self { - Scopeguard { t, f } - } - fn get(&self) -> &T { - &self.t - } -} - -impl Drop for Scopeguard -where - T: Copy, - F: FnMut(T), -{ - fn drop(&mut self) { - (self.f)(self.t); - } -} diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index dcac2314..bc38dc36 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -1,18 +1,15 @@ pub mod am_build; -pub mod am_scan; -use crate::index::projection::RandomProject; +use crate::index::gucs::{epsilon, max_scan_tuples, probes}; +use crate::index::opclass::opfamily; +use crate::index::scanners::*; use crate::index::storage::PostgresRelation; -use algorithm::operator::{Dot, L2, Op, Vector}; -use algorithm::types::*; -use half::f16; use pgrx::datum::Internal; -use pgrx::pg_sys::Datum; +use pgrx::pg_sys::{Datum, ItemPointerData}; +use std::cell::LazyCell; use std::ffi::CStr; use std::num::NonZeroU64; use std::sync::OnceLock; -use vector::VectorOwned; -use vector::vect::VectOwned; #[repr(C)] struct Reloption { @@ -99,10 +96,10 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { am_routine.ambulkdelete = Some(ambulkdelete); am_routine.amvacuumcleanup = Some(amvacuumcleanup); - am_routine.ambeginscan = Some(am_scan::ambeginscan); - am_routine.amrescan = Some(am_scan::amrescan); - am_routine.amgettuple = Some(am_scan::amgettuple); - am_routine.amendscan = Some(am_scan::amendscan); + am_routine.ambeginscan = Some(ambeginscan); + am_routine.amrescan = Some(amrescan); + am_routine.amgettuple = Some(amgettuple); + am_routine.amendscan = Some(amendscan); am_routine }; @@ -189,34 +186,17 @@ unsafe fn aminsertinner( index_relation: pgrx::pg_sys::Relation, values: *mut Datum, is_null: *mut bool, - heap_tid: pgrx::pg_sys::ItemPointer, + ctid: pgrx::pg_sys::ItemPointer, ) -> bool { - let opfamily = unsafe { crate::index::opclass::opfamily(index_relation) }; + let opfamily = unsafe { opfamily(index_relation) }; let index = unsafe { PostgresRelation::new(index_relation) }; - let payload = ctid_to_pointer(unsafe { heap_tid.read() }); - let vector = unsafe { opfamily.input_vector(*values.add(0), *is_null.add(0)) }; - let Some(vector) = vector else { return false }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => algorithm::insert::, L2>>( - index, - payload, - RandomProject::project(VectOwned::::from_owned(vector).as_borrowed()), - ), - (VectorKind::Vecf32, DistanceKind::Dot) => algorithm::insert::, Dot>>( - index, - payload, - RandomProject::project(VectOwned::::from_owned(vector).as_borrowed()), - ), - (VectorKind::Vecf16, DistanceKind::L2) => algorithm::insert::, L2>>( - index, - payload, - RandomProject::project(VectOwned::::from_owned(vector).as_borrowed()), - ), - (VectorKind::Vecf16, DistanceKind::Dot) => algorithm::insert::, Dot>>( - index, - payload, - RandomProject::project(VectOwned::::from_owned(vector).as_borrowed()), - ), + let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; + let ctid = unsafe { ctid.read() }; + if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { + for (vector, extra) in store { + let payload = ctid_to_pointer(ctid, extra); + crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); + } } false } @@ -234,27 +214,14 @@ pub unsafe extern "C" fn ambulkdelete( pgrx::pg_sys::palloc0(size_of::()).cast() }; } - let opfamily = unsafe { crate::index::opclass::opfamily((*info).index) }; + let opfamily = unsafe { opfamily((*info).index) }; let index = unsafe { PostgresRelation::new((*info).index) }; let check = || unsafe { pgrx::pg_sys::vacuum_delay_point(); }; let callback = callback.expect("null function pointer"); - let callback = |p: NonZeroU64| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::bulkdelete::, L2>>(index, check, callback); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::bulkdelete::, Dot>>(index, check, callback); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::bulkdelete::, L2>>(index, check, callback); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::bulkdelete::, Dot>>(index, check, callback); - } - } + let callback = |p: NonZeroU64| unsafe { callback(&mut pointer_to_ctid(p).0, callback_state) }; + crate::index::algorithm::bulkdelete(opfamily, index, check, callback); stats } @@ -269,56 +236,244 @@ pub unsafe extern "C" fn amvacuumcleanup( pgrx::pg_sys::palloc0(size_of::()).cast() }; } - let opfamily = unsafe { crate::index::opclass::opfamily((*info).index) }; + let opfamily = unsafe { opfamily((*info).index) }; let index = unsafe { PostgresRelation::new((*info).index) }; let check = || unsafe { pgrx::pg_sys::vacuum_delay_point(); }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::maintain::, L2>>(index, check); + crate::index::algorithm::maintain(opfamily, index, check); + stats +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn ambeginscan( + index_relation: pgrx::pg_sys::Relation, + n_keys: std::os::raw::c_int, + n_orderbys: std::os::raw::c_int, +) -> pgrx::pg_sys::IndexScanDesc { + use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; + + let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index_relation, n_keys, n_orderbys) }; + let scanner: Scanner = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + unsafe { + (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); + } + scan +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amrescan( + scan: pgrx::pg_sys::IndexScanDesc, + keys: pgrx::pg_sys::ScanKey, + _n_keys: std::os::raw::c_int, + orderbys: pgrx::pg_sys::ScanKey, + _n_orderbys: std::os::raw::c_int, +) { + unsafe { + use crate::index::opclass::Opfamily; + if !keys.is_null() && (*scan).numberOfKeys > 0 { + std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); } - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::maintain::, Dot>>(index, check); + if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { + std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::maintain::, L2>>(index, check); + if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { + pgrx::error!( + "vector search with no WHERE clause and no ORDER BY clause is not supported" + ); } - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::maintain::, Dot>>(index, check); + let opfamily = opfamily((*scan).indexRelation); + let relation = PostgresRelation::new((*scan).indexRelation); + let options = SearchOptions { + epsilon: epsilon(), + probes: probes(), + max_scan_tuples: max_scan_tuples(), + }; + let fetcher = LazyCell::new(move || { + HeapFetcher::new( + (*scan).indexRelation, + (*scan).heapRelation, + (*scan).xs_snapshot, + ) + }); + let scanner = &mut *(*scan).opaque.cast::(); + *scanner = match opfamily { + Opfamily::VectorL2 + | Opfamily::VectorIp + | Opfamily::VectorCosine + | Opfamily::HalfvecL2 + | Opfamily::HalfvecIp + | Opfamily::HalfvecCosine => { + let mut builder = DefaultBuilder::new(opfamily); + for i in 0..(*scan).numberOfOrderBys { + let data = (*scan).orderByData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + for i in 0..(*scan).numberOfKeys { + let data = (*scan).keyData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + LazyCell::new(Box::new(move || builder.build(relation, options, fetcher))) + } + }; + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amgettuple( + scan: pgrx::pg_sys::IndexScanDesc, + direction: pgrx::pg_sys::ScanDirection::Type, +) -> bool { + if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { + pgrx::error!("vector search without a forward scan direction is not supported"); + } + // https://www.postgresql.org/docs/current/index-locking.html + // If heap entries referenced physical pointers are deleted before + // they are consumed by PostgreSQL, PostgreSQL will received wrong + // physical pointers: no rows or irreverent rows are referenced. + if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + { + pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); + } + let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; + if let Some((_, ctid, recheck)) = LazyCell::force_mut(scanner).next() { + unsafe { + (*scan).xs_heaptid = ctid; + (*scan).xs_recheck = recheck; + (*scan).xs_recheckorderby = false; + } + true + } else { + false + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { + let scanner = unsafe { &mut *(*scan).opaque.cast::() }; + *scanner = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); +} + +type Iter = Box>; + +type Scanner = LazyCell Iter>>; + +struct HeapFetcher { + index_info: *mut pgrx::pg_sys::IndexInfo, + estate: *mut pgrx::pg_sys::EState, + econtext: *mut pgrx::pg_sys::ExprContext, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + slot: *mut pgrx::pg_sys::TupleTableSlot, + values: [Datum; 32], + is_nulls: [bool; 32], +} + +impl HeapFetcher { + unsafe fn new( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + ) -> Self { + unsafe { + let index_info = pgrx::pg_sys::BuildIndexInfo(index_relation); + let estate = pgrx::pg_sys::CreateExecutorState(); + let econtext = pgrx::pg_sys::MakePerTupleExprContext(estate); + Self { + index_info, + estate, + econtext, + heap_relation, + snapshot, + slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), + values: [Datum::null(); 32], + is_nulls: [true; 32], + } } } - stats } -const fn pointer_to_ctid(pointer: NonZeroU64) -> pgrx::pg_sys::ItemPointerData { +impl Drop for HeapFetcher { + fn drop(&mut self) { + unsafe { + // free resources for last `fetch` call + pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); + // free common resources + pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); + pgrx::pg_sys::FreeExecutorState(self.estate); + } + } +} + +impl SearchFetcher for HeapFetcher { + fn fetch(&mut self, mut ctid: ItemPointerData) -> Option<(&[Datum; 32], &[bool; 32])> { + unsafe { + // free resources for last `fetch` call + pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); + // perform operation + (*self.econtext).ecxt_scantuple = self.slot; + let table_am = (*self.heap_relation).rd_tableam; + let fetch_row_version = (*table_am) + .tuple_fetch_row_version + .expect("unsupported heap access method"); + if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { + return None; + } + pgrx::pg_sys::FormIndexDatum( + self.index_info, + self.slot, + self.estate, + self.values.as_mut_ptr(), + self.is_nulls.as_mut_ptr(), + ); + Some((&self.values, &self.is_nulls)) + } + } +} + +pub const fn pointer_to_ctid(pointer: NonZeroU64) -> (ItemPointerData, u16) { let value = pointer.get(); - pgrx::pg_sys::ItemPointerData { - ip_blkid: pgrx::pg_sys::BlockIdData { - bi_hi: ((value >> 32) & 0xffff) as u16, - bi_lo: ((value >> 16) & 0xffff) as u16, + let bi_hi = ((value >> 48) & 0xffff) as u16; + let bi_lo = ((value >> 32) & 0xffff) as u16; + let ip_posid = ((value >> 16) & 0xffff) as u16; + let extra = value as u16; + ( + ItemPointerData { + ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi, bi_lo }, + ip_posid, }, - ip_posid: (value & 0xffff) as u16, - } + extra, + ) } -const fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> NonZeroU64 { +pub const fn ctid_to_pointer(ctid: ItemPointerData, extra: u16) -> NonZeroU64 { let mut value = 0; - value |= (ctid.ip_blkid.bi_hi as u64) << 32; - value |= (ctid.ip_blkid.bi_lo as u64) << 16; - value |= ctid.ip_posid as u64; + value |= (ctid.ip_blkid.bi_hi as u64) << 48; + value |= (ctid.ip_blkid.bi_lo as u64) << 32; + value |= (ctid.ip_posid as u64) << 16; + value |= extra as u64; NonZeroU64::new(value).expect("invalid pointer") } #[test] const fn soundness_check() { - let a = pgrx::pg_sys::ItemPointerData { - ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi: 1, bi_lo: 2 }, - ip_posid: 3, - }; - let b = ctid_to_pointer(a); - let c = pointer_to_ctid(b); - assert!(a.ip_blkid.bi_hi == c.ip_blkid.bi_hi); - assert!(a.ip_blkid.bi_lo == c.ip_blkid.bi_lo); - assert!(a.ip_posid == c.ip_posid); + let bi_hi = 1; + let bi_lo = 2; + let ip_posid = 3; + let extra = 7; + let r = pointer_to_ctid(ctid_to_pointer( + ItemPointerData { + ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi, bi_lo }, + ip_posid, + }, + extra, + )); + assert!(r.0.ip_blkid.bi_hi == bi_hi); + assert!(r.0.ip_blkid.bi_lo == bi_lo); + assert!(r.0.ip_posid == ip_posid); + assert!(r.1 == extra); } diff --git a/src/index/functions.rs b/src/index/functions.rs index 64acde82..4a472229 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -1,10 +1,6 @@ use crate::index::storage::PostgresRelation; -use algorithm::operator::{Dot, L2, Op}; -use algorithm::types::*; -use half::f16; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass}; -use vector::vect::VectOwned; #[pgrx::pg_extern(sql = "")] fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { @@ -22,28 +18,9 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::AccessShareLock as _) }; let relation = unsafe { PostgresRelation::new(index) }; let opfamily = unsafe { crate::index::opclass::opfamily(index) }; - let message = match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::prewarm::, L2>>(relation, height, || { - pgrx::check_for_interrupts!(); - }) - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::prewarm::, Dot>>(relation, height, || { - pgrx::check_for_interrupts!(); - }) - } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::prewarm::, L2>>(relation, height, || { - pgrx::check_for_interrupts!(); - }) - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::prewarm::, Dot>>(relation, height, || { - pgrx::check_for_interrupts!(); - }) - } - }; + let message = crate::index::algorithm::prewarm(opfamily, relation, height, || { + pgrx::check_for_interrupts!(); + }); unsafe { pgrx::pg_sys::index_close(index, pgrx::pg_sys::AccessShareLock as _); } diff --git a/src/index/mod.rs b/src/index/mod.rs index e7a309ae..925db7d6 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -1,8 +1,10 @@ +pub mod algorithm; pub mod am; pub mod functions; pub mod gucs; pub mod opclass; pub mod projection; +pub mod scanners; pub mod storage; pub mod types; diff --git a/src/index/opclass.rs b/src/index/opclass.rs index 63a6be5d..4b989e85 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -38,66 +38,64 @@ fn _vchordrq_support_halfvec_cosine_ops() -> String { "halfvec_cosine_ops".to_string() } -#[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -enum PostgresDistanceKind { - L2, - Ip, - Cosine, -} - pub struct Sphere { pub center: T, pub radius: f32, } #[derive(Debug, Clone, Copy)] -pub struct Opfamily { - vector: VectorKind, - postgres_distance: PostgresDistanceKind, +pub enum Opfamily { + VectorL2, + VectorIp, + VectorCosine, + HalfvecL2, + HalfvecIp, + HalfvecCosine, } impl Opfamily { fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { - use {BorrowedVector as B, OwnedVector as O, PostgresDistanceKind as D}; - match (vector, self.postgres_distance) { - (B::Vecf32(x), D::L2) => O::Vecf32(x.own()), - (B::Vecf32(x), D::Ip) => O::Vecf32(x.own()), - (B::Vecf32(x), D::Cosine) => O::Vecf32(x.function_normalize()), - (B::Vecf16(x), D::L2) => O::Vecf16(x.own()), - (B::Vecf16(x), D::Ip) => O::Vecf16(x.own()), - (B::Vecf16(x), D::Cosine) => O::Vecf16(x.function_normalize()), + use {BorrowedVector as B, OwnedVector as O}; + match (vector, self) { + (B::Vecf32(x), Self::VectorL2) => O::Vecf32(x.own()), + (B::Vecf32(x), Self::VectorIp) => O::Vecf32(x.own()), + (B::Vecf32(x), Self::VectorCosine) => O::Vecf32(x.function_normalize()), + (B::Vecf32(_), _) => unreachable!(), + (B::Vecf16(x), Self::HalfvecL2) => O::Vecf16(x.own()), + (B::Vecf16(x), Self::HalfvecIp) => O::Vecf16(x.own()), + (B::Vecf16(x), Self::HalfvecCosine) => O::Vecf16(x.function_normalize()), + (B::Vecf16(_), _) => unreachable!(), } } - pub unsafe fn input_vector(self, datum: Datum, is_null: bool) -> Option { - if is_null || datum.is_null() { + pub unsafe fn store(self, datum: Datum) -> Option> { + if datum.is_null() { return None; } - let vector = match self.vector { - VectorKind::Vecf32 => { + let store = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; - self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + vec![(self.input(BorrowedVector::Vecf32(vector.as_borrowed())), 0)] } - VectorKind::Vecf16 => { + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; - self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] } }; - Some(vector) + Some(store) } - pub unsafe fn input_sphere(self, datum: Datum, is_null: bool) -> Option> { - if is_null || datum.is_null() { + pub unsafe fn input_sphere(self, datum: Datum) -> Option> { + if datum.is_null() { return None; } let attno_1 = NonZero::new(1_usize).unwrap(); let attno_2 = NonZero::new(2_usize).unwrap(); let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; - let center = match self.vector { - VectorKind::Vecf32 => { + let center = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Vecf32(vector.as_borrowed())) } - VectorKind::Vecf16 => { + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Vecf16(vector.as_borrowed())) } @@ -105,21 +103,69 @@ impl Opfamily { let radius = tuple.get_by_index::(attno_2).unwrap()?; Some(Sphere { center, radius }) } + pub unsafe fn input_vector(self, datum: Datum) -> Option { + if datum.is_null() { + return None; + } + let vector = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + }; + Some(vector) + } + #[allow(dead_code)] + pub unsafe fn input_vectors(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let vectors = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { + let vectors = + unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Vecf32(vector.as_borrowed()))); + } + result + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { + let vectors = + unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Vecf16(vector.as_borrowed()))); + } + result + } + }; + Some(vectors) + } pub fn output(self, x: Distance) -> f32 { - match self.postgres_distance { - PostgresDistanceKind::Cosine => x.to_f32() + 1.0f32, - PostgresDistanceKind::L2 => x.to_f32().sqrt(), - PostgresDistanceKind::Ip => x.to_f32(), + match self { + Self::VectorCosine | Self::HalfvecCosine => x.to_f32() + 1.0f32, + Self::VectorL2 | Self::HalfvecL2 => x.to_f32().sqrt(), + Self::VectorIp | Self::HalfvecIp => x.to_f32(), } } pub const fn distance_kind(self) -> DistanceKind { - match self.postgres_distance { - PostgresDistanceKind::L2 => DistanceKind::L2, - PostgresDistanceKind::Ip | PostgresDistanceKind::Cosine => DistanceKind::Dot, + match self { + Self::VectorL2 | Self::HalfvecL2 => DistanceKind::L2, + Self::VectorIp | Self::HalfvecIp | Self::VectorCosine | Self::HalfvecCosine => { + DistanceKind::Dot + } } } pub const fn vector_kind(self) -> VectorKind { - self.vector + match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine => VectorKind::Vecf32, + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => VectorKind::Vecf16, + } } } @@ -154,13 +200,13 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { let result_string = result_option.expect("null return value"); - let (vector, postgres_distance) = match result_string.as_str() { - "vector_l2_ops" => (VectorKind::Vecf32, PostgresDistanceKind::L2), - "vector_ip_ops" => (VectorKind::Vecf32, PostgresDistanceKind::Ip), - "vector_cosine_ops" => (VectorKind::Vecf32, PostgresDistanceKind::Cosine), - "halfvec_l2_ops" => (VectorKind::Vecf16, PostgresDistanceKind::L2), - "halfvec_ip_ops" => (VectorKind::Vecf16, PostgresDistanceKind::Ip), - "halfvec_cosine_ops" => (VectorKind::Vecf16, PostgresDistanceKind::Cosine), + let result = match result_string.as_str() { + "vector_l2_ops" => Opfamily::VectorL2, + "vector_ip_ops" => Opfamily::VectorIp, + "vector_cosine_ops" => Opfamily::VectorCosine, + "halfvec_l2_ops" => Opfamily::HalfvecL2, + "halfvec_ip_ops" => Opfamily::HalfvecIp, + "halfvec_cosine_ops" => Opfamily::HalfvecCosine, _ => pgrx::error!("unknown operator class"), }; @@ -168,8 +214,5 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); } - Opfamily { - vector, - postgres_distance, - } + result } diff --git a/src/index/projection.rs b/src/index/projection.rs index ca07e24f..fbcaeffa 100644 --- a/src/index/projection.rs +++ b/src/index/projection.rs @@ -1,7 +1,5 @@ -use half::f16; use random_orthogonal_matrix::random_orthogonal_matrix; use std::sync::OnceLock; -use vector::vect::{VectBorrowed, VectOwned}; fn matrix(n: usize) -> Option<&'static Vec>> { static MATRIXS: [OnceLock>>; 1 + 60000] = [const { OnceLock::new() }; 1 + 60000]; @@ -22,25 +20,3 @@ pub fn project(vector: &[f32]) -> Vec { .map(|i| f32::reduce_sum_of_xy(vector, &matrix[i])) .collect() } - -pub trait RandomProject { - type Output; - fn project(self) -> Self::Output; -} - -impl RandomProject for VectBorrowed<'_, f32> { - type Output = VectOwned; - fn project(self) -> VectOwned { - VectOwned::new(project(self.slice())) - } -} - -impl RandomProject for VectBorrowed<'_, f16> { - type Output = VectOwned; - fn project(self) -> VectOwned { - use simd::Floating; - VectOwned::new(f16::vector_from_f32(&project(&f16::vector_to_f32( - self.slice(), - )))) - } -} diff --git a/src/index/scanners/default.rs b/src/index/scanners/default.rs new file mode 100644 index 00000000..107d201b --- /dev/null +++ b/src/index/scanners/default.rs @@ -0,0 +1,227 @@ +use super::{SearchBuilder, SearchFetcher, SearchOptions}; +use crate::index::algorithm::RandomProject; +use crate::index::am::pointer_to_ctid; +use crate::index::opclass::{Opfamily, Sphere}; +use algorithm::operator::{Dot, L2, Op, Vector}; +use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; +use algorithm::{RelationRead, RerankMethod}; +use half::f16; +use pgrx::pg_sys::ItemPointerData; +use std::num::NonZeroU64; +use vector::VectorOwned; +use vector::vect::VectOwned; + +pub struct DefaultBuilder { + opfamily: Opfamily, + orderbys: Vec>, + spheres: Vec>>, +} + +impl SearchBuilder for DefaultBuilder { + fn new(opfamily: Opfamily) -> Self { + assert!(matches!( + opfamily, + Opfamily::HalfvecCosine + | Opfamily::HalfvecIp + | Opfamily::HalfvecL2 + | Opfamily::VectorCosine + | Opfamily::VectorIp + | Opfamily::VectorL2 + )); + Self { + opfamily, + orderbys: Vec::new(), + spheres: Vec::new(), + } + } + + unsafe fn add(&mut self, strategy: u16, datum: Option) { + match strategy { + 1 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_vector(x)) }; + self.orderbys.push(x); + } + 2 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_sphere(x)) }; + self.spheres.push(x); + } + _ => unreachable!(), + } + } + + fn build<'a>( + self, + relation: impl RelationRead + 'a, + options: SearchOptions, + mut fetcher: impl SearchFetcher + 'a, + ) -> Box + 'a> { + let mut vector = None; + let mut threshold = None; + let mut recheck = false; + for orderby_vector in self.orderbys.into_iter().flatten() { + if vector.is_none() { + vector = Some(orderby_vector); + } else { + pgrx::error!("vector search with multiple vectors is not supported"); + } + } + for Sphere { center, radius } in self.spheres.into_iter().flatten() { + if vector.is_none() { + (vector, threshold) = (Some(center), Some(radius)); + } else { + recheck = true; + } + } + let opfamily = self.opfamily; + let Some(vector) = vector else { + return Box::new(std::iter::empty()) + as Box>; + }; + let iter: Box> = + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2) => { + let vector = RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ); + let (method, results) = algorithm::search::, L2>>( + relation.clone(), + vector.clone(), + options.probes, + options.epsilon, + ); + let fetch = move |payload| { + let (ctid, _) = pointer_to_ctid(payload); + let (datums, is_nulls) = fetcher.fetch(ctid)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = VectOwned::::from_owned(maybe_vector.unwrap()); + Some(RandomProject::project(raw.as_borrowed())) + }; + match method { + RerankMethod::Index => Box::new( + algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)), + ), + RerankMethod::Heap => Box::new( + algorithm::rerank_heap::, L2>, _>( + vector, results, fetch, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)), + ), + } + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + let vector = RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ); + let (method, results) = algorithm::search::, Dot>>( + relation.clone(), + vector.clone(), + options.probes, + options.epsilon, + ); + let fetch = move |payload| { + let (ctid, _) = pointer_to_ctid(payload); + let (datums, is_nulls) = fetcher.fetch(ctid)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = VectOwned::::from_owned(maybe_vector.unwrap()); + Some(RandomProject::project(raw.as_borrowed())) + }; + match method { + RerankMethod::Index => Box::new( + algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)), + ), + RerankMethod::Heap => Box::new( + algorithm::rerank_heap::, Dot>, _>( + vector, results, fetch, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)), + ), + } + } + (VectorKind::Vecf16, DistanceKind::L2) => { + let vector = RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ); + let (method, results) = algorithm::search::, L2>>( + relation.clone(), + vector.clone(), + options.probes, + options.epsilon, + ); + let fetch = move |payload| { + let (ctid, _) = pointer_to_ctid(payload); + let (datums, is_nulls) = fetcher.fetch(ctid)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = VectOwned::::from_owned(maybe_vector.unwrap()); + Some(RandomProject::project(raw.as_borrowed())) + }; + match method { + RerankMethod::Index => Box::new( + algorithm::rerank_index::, L2>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)), + ), + RerankMethod::Heap => Box::new( + algorithm::rerank_heap::, L2>, _>( + vector, results, fetch, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)), + ), + } + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + let vector = RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ); + let (method, results) = algorithm::search::, Dot>>( + relation.clone(), + vector.clone(), + options.probes, + options.epsilon, + ); + let fetch = move |payload| { + let (ctid, _) = pointer_to_ctid(payload); + let (datums, is_nulls) = fetcher.fetch(ctid)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = VectOwned::::from_owned(maybe_vector.unwrap()); + Some(RandomProject::project(raw.as_borrowed())) + }; + match method { + RerankMethod::Index => Box::new( + algorithm::rerank_index::, Dot>>( + relation, vector, results, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)), + ), + RerankMethod::Heap => Box::new( + algorithm::rerank_heap::, Dot>, _>( + vector, results, fetch, + ) + .map(move |(distance, payload)| (opfamily.output(distance), payload)), + ), + } + } + }; + let iter = if let Some(threshold) = threshold { + Box::new(iter.take_while(move |(x, _)| *x < threshold)) + } else { + iter + }; + let iter = if let Some(max_scan_tuples) = options.max_scan_tuples { + Box::new(iter.take(max_scan_tuples as _)) + } else { + iter + }; + Box::new(iter.map(move |(x, y)| (x, pointer_to_ctid(y).0, recheck))) + } +} diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs new file mode 100644 index 00000000..afef6263 --- /dev/null +++ b/src/index/scanners/mod.rs @@ -0,0 +1,38 @@ +mod default; + +use super::opclass::Opfamily; +use algorithm::RelationRead; +use pgrx::pg_sys::{Datum, ItemPointerData}; +use std::cell::LazyCell; + +pub use default::DefaultBuilder; + +#[derive(Debug)] +pub struct SearchOptions { + pub epsilon: f32, + pub probes: Vec, + pub max_scan_tuples: Option, +} + +pub trait SearchBuilder: 'static { + fn new(opfamily: Opfamily) -> Self; + + unsafe fn add(&mut self, strategy: u16, datum: Option); + + fn build<'a>( + self, + relation: impl RelationRead + 'a, + options: SearchOptions, + fetcher: impl SearchFetcher + 'a, + ) -> Box + 'a>; +} + +pub trait SearchFetcher { + fn fetch(&mut self, ctid: ItemPointerData) -> Option<(&[Datum; 32], &[bool; 32])>; +} + +impl T> SearchFetcher for LazyCell { + fn fetch(&mut self, ctid: ItemPointerData) -> Option<(&[Datum; 32], &[bool; 32])> { + LazyCell::force_mut(self).fetch(ctid) + } +} diff --git a/src/lib.rs b/src/lib.rs index 20ce1836..5a98b3d9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,7 @@ #![allow(clippy::collapsible_else_if)] #![allow(clippy::too_many_arguments)] #![allow(unsafe_code)] +#![feature(lazy_get)] mod datatype; mod index; From 318725505a38f976900b446667a0d081ce731344 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 17 Mar 2025 10:01:12 +0800 Subject: [PATCH 124/324] fix: recycle pages in maintainance (#204) Signed-off-by: usamoi --- crates/algorithm/src/freepages.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/algorithm/src/freepages.rs b/crates/algorithm/src/freepages.rs index 95b63aea..a5f95449 100644 --- a/crates/algorithm/src/freepages.rs +++ b/crates/algorithm/src/freepages.rs @@ -7,7 +7,7 @@ pub fn mark(index: impl RelationWrite, freepage_first: u32, pages: &[u32]) { pages.sort_by_key(|x| Reverse(*x)); pages.dedup(); let (mut current, mut offset) = (freepage_first, 0_u32); - while pages.is_empty() { + while !pages.is_empty() { let locals = { let mut local = Vec::new(); while let Some(target) = pages.pop_if(|x| (offset..offset + 32768).contains(x)) { From 178e118f7ebb317da3c9268de11168afab6baedd Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 17 Mar 2025 16:56:26 +0800 Subject: [PATCH 125/324] chore: update v0.2.2 schema script (#210) Signed-off-by: usamoi --- .github/workflows/check.yml | 4 + sql/install/vchord--0.2.2.sql | 500 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.2.1--0.2.2.sql | 1 + 3 files changed, 505 insertions(+) create mode 100644 sql/install/vchord--0.2.2.sql create mode 100644 sql/upgrade/vchord--0.2.1--0.2.2.sql diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 22804d5f..956668dd 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -69,6 +69,10 @@ jobs: - name: Rustfmt run: cargo fmt --check + - name: SQL + run: | + ! grep -P '\t' -r ./sql + lint: strategy: matrix: diff --git a/sql/install/vchord--0.2.2.sql b/sql/install/vchord--0.2.2.sql new file mode 100644 index 00000000..f408248c --- /dev/null +++ b/sql/install/vchord--0.2.2.sql @@ -0,0 +1,500 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:11 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:54 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:30 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:6 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:54 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:30 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:6 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:57 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:5 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:36 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:31 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:12 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.2.1--0.2.2.sql b/sql/upgrade/vchord--0.2.1--0.2.2.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.2.1--0.2.2.sql @@ -0,0 +1 @@ +/* nothing to do */ From 159bd0e2f56e727ac46d622889ff7d46e86e814c Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 17 Mar 2025 19:27:35 +0800 Subject: [PATCH 126/324] fix: correct logic of marking free pages (#211) Signed-off-by: usamoi --- crates/algorithm/src/freepages.rs | 33 +++++++-------- crates/algorithm/src/lib.rs | 1 + crates/algorithm/src/maintain.rs | 8 ++-- crates/algorithm/src/tuples.rs | 69 ++++++++++++++++++++----------- 4 files changed, 66 insertions(+), 45 deletions(-) diff --git a/crates/algorithm/src/freepages.rs b/crates/algorithm/src/freepages.rs index a5f95449..6b54889c 100644 --- a/crates/algorithm/src/freepages.rs +++ b/crates/algorithm/src/freepages.rs @@ -2,27 +2,26 @@ use crate::tuples::*; use crate::*; use std::cmp::Reverse; -pub fn mark(index: impl RelationWrite, freepage_first: u32, pages: &[u32]) { - let mut pages = pages.to_vec(); - pages.sort_by_key(|x| Reverse(*x)); - pages.dedup(); +pub fn mark(index: impl RelationWrite, freepage_first: u32, values: &[u32]) { + let mut values = { + let mut values = values.to_vec(); + values.sort_by_key(|x| Reverse(*x)); + values.dedup(); + values + }; let (mut current, mut offset) = (freepage_first, 0_u32); - while !pages.is_empty() { - let locals = { - let mut local = Vec::new(); - while let Some(target) = pages.pop_if(|x| (offset..offset + 32768).contains(x)) { - local.push(target - offset); - } - local - }; + loop { let mut freespace_guard = index.write(current, false); if freespace_guard.len() == 0 { freespace_guard.alloc(&FreepageTuple::serialize(&FreepageTuple {})); } let freespace_bytes = freespace_guard.get_mut(1).expect("data corruption"); let mut freespace_tuple = FreepageTuple::deserialize_mut(freespace_bytes); - for local in locals { - freespace_tuple.mark(local as _); + while let Some(target) = values.pop_if(|&mut x| x < offset + 32768) { + freespace_tuple.mark((target - offset) as usize); + } + if values.is_empty() { + return; } if freespace_guard.get_opaque().next == u32::MAX { let extend = index.extend(false); @@ -37,12 +36,12 @@ pub fn fetch(index: impl RelationWrite, freepage_first: u32) -> Option { loop { let mut freespace_guard = index.write(current, false); if freespace_guard.len() == 0 { - return None; + freespace_guard.alloc(&FreepageTuple::serialize(&FreepageTuple {})); } let freespace_bytes = freespace_guard.get_mut(1).expect("data corruption"); let mut freespace_tuple = FreepageTuple::deserialize_mut(freespace_bytes); - if let Some(local) = freespace_tuple.fetch() { - return Some(local as u32 + offset); + if let Some(x) = freespace_tuple.fetch() { + return Some(x as u32 + offset); } if freespace_guard.get_opaque().next == u32::MAX { return None; diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 3bba11de..de792ddf 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -1,3 +1,4 @@ +#![feature(generic_arg_infer)] #![allow(clippy::collapsible_else_if)] #![allow(clippy::type_complexity)] #![allow(clippy::len_without_is_empty)] diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index 2749a997..e7640049 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -48,7 +48,7 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let jump_bytes = jump_guard.get_mut(1).expect("data corruption"); let mut jump_tuple = JumpTuple::deserialize_mut(jump_bytes); - let mut tape = FrozenTapeWriter::<_, _>::create(|| { + let expand = || { if let Some(id) = freepages::fetch(index.clone(), freepage_first) { let mut write = index.write(id, false); write.clear(); @@ -56,7 +56,9 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { } else { index.extend(false) } - }); + }; + + let mut tape = FrozenTapeWriter::<_, _>::create(expand); let mut trace = Vec::new(); @@ -113,7 +115,7 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let (frozen_tape, branches) = tape.into_inner(); - let mut appendable_tape = TapeWriter::create(|| index.extend(false)); + let mut appendable_tape = TapeWriter::create(expand); for branch in branches { appendable_tape.push(AppendableTuple { diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 60a92dd7..c85588b4 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -7,7 +7,7 @@ use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: u64 = u64::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 3; +const VERSION: u64 = 4; pub trait Tuple: 'static { fn serialize(&self) -> Vec; @@ -127,9 +127,9 @@ impl MetaTupleReader<'_> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FreepageTupleHeader { - a: [u32; 1], - b: [u32; 32], - c: [u32; 32 * 32], + level_0: [u32; 1 << 10], + level_1: [u32; 1 << 5], + level_2: [u32; 1 << 0], _padding_0: [ZeroU8; 4], } @@ -141,9 +141,9 @@ pub struct FreepageTuple {} impl Tuple for FreepageTuple { fn serialize(&self) -> Vec { FreepageTupleHeader { - a: std::array::from_fn(|_| 0), - b: std::array::from_fn(|_| 0), - c: std::array::from_fn(|_| 0), + level_0: [0; _], + level_1: [0; _], + level_2: [0; _], _padding_0: Default::default(), } .as_bytes() @@ -167,28 +167,37 @@ pub struct FreepageTupleWriter<'a> { impl FreepageTupleWriter<'_> { pub fn mark(&mut self, i: usize) { - let c_i = i; - self.header.c[c_i / 32] |= 1 << (c_i % 32); - let b_i = i / 32; - self.header.b[b_i / 32] |= 1 << (b_i % 32); - let a_i = i / 32 / 32; - self.header.a[a_i / 32] |= 1 << (a_i % 32); + assert!(i < 32768, "out of bound: {i}"); + set(&mut self.header.level_0[i >> 5], (i >> 0) % 32, true); + set(&mut self.header.level_1[i >> 10], (i >> 5) % 32, true); + set(&mut self.header.level_2[i >> 15], (i >> 10) % 32, true); } - pub fn fetch(&mut self) -> Option { - if self.header.a[0].trailing_ones() == 32 { + #[allow(clippy::just_underscores_and_digits)] + fn find(&self) -> Option { + let _3 = 0_usize; + let _2 = self.header.level_2[_3 << 0].trailing_zeros() as usize; + if _2 == 32 { return None; } - let a_i = self.header.a[0].trailing_zeros() as usize; - let b_i = self.header.b[a_i].trailing_zeros() as usize + a_i * 32; - let c_i = self.header.c[b_i].trailing_zeros() as usize + b_i * 32; - self.header.c[c_i / 32] &= !(1 << (c_i % 32)); - if self.header.c[b_i] == 0 { - self.header.b[b_i / 32] &= !(1 << (b_i % 32)); - if self.header.b[a_i] == 0 { - self.header.a[a_i / 32] &= !(1 << (a_i % 32)); - } + let _1 = self.header.level_1[_3 << 5 | _2 << 0].trailing_zeros() as usize; + if _1 == 32 { + panic!("data corruption"); } - Some(c_i) + let _0 = self.header.level_0[_3 << 10 | _2 << 5 | _1 << 0].trailing_zeros() as usize; + if _0 == 32 { + panic!("data corruption"); + } + Some(_3 << 15 | _2 << 10 | _1 << 5 | _0 << 0) + } + pub fn fetch(&mut self) -> Option { + let i = self.find()?; + let x = false; + set(&mut self.header.level_0[i >> 5], (i >> 0) % 32, x); + let x = self.header.level_0[i >> 5] != 0; + set(&mut self.header.level_1[i >> 10], (i >> 5) % 32, x); + let x = self.header.level_1[i >> 10] != 0; + set(&mut self.header.level_2[i >> 15], (i >> 10) % 32, x); + Some(i) } } @@ -1115,6 +1124,16 @@ impl<'a> MutChecker<'a> { } } +#[inline(always)] +fn set(a: &mut u32, i: usize, x: bool) { + assert!(i < 32, "out of bound: {i}"); + if x { + *a |= 1 << i; + } else { + *a &= !(1 << i); + } +} + #[test] fn aliasing_test() { #[repr(C, align(8))] From 3fe800b3e4af02ecdfdbd061046486900c25b3a4 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 18 Mar 2025 09:44:28 +0800 Subject: [PATCH 127/324] fix: remove heapify (#213) fixes #209 Signed-off-by: usamoi --- Cargo.lock | 7 --- crates/algorithm/Cargo.toml | 1 - crates/algorithm/src/select_heap.rs | 66 +++++++++++++++++------------ 3 files changed, 40 insertions(+), 34 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a5a23736..a80cffbe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -18,7 +18,6 @@ dependencies = [ "always_equal", "distance", "half 2.4.1", - "heapify", "k_means", "paste", "rabitq", @@ -392,12 +391,6 @@ version = "0.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" -[[package]] -name = "heapify" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0049b265b7f201ca9ab25475b22b47fe444060126a51abe00f77d986fc5cc52e" - [[package]] name = "heapless" version = "0.8.0" diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml index 83adc488..7ca1c37c 100644 --- a/crates/algorithm/Cargo.toml +++ b/crates/algorithm/Cargo.toml @@ -13,7 +13,6 @@ simd = { path = "../simd" } vector = { path = "../vector" } half.workspace = true -heapify = "0.2.0" paste.workspace = true rand.workspace = true serde.workspace = true diff --git a/crates/algorithm/src/select_heap.rs b/crates/algorithm/src/select_heap.rs index 86723881..a1a5a224 100644 --- a/crates/algorithm/src/select_heap.rs +++ b/crates/algorithm/src/select_heap.rs @@ -1,44 +1,51 @@ -pub struct SelectHeap { - threshold: usize, - inner: Vec, +use std::collections::BinaryHeap; +use std::num::NonZero; + +pub enum SelectHeap { + Sorted { x: Vec, t: NonZero }, + Heap { x: BinaryHeap }, } impl SelectHeap { pub fn from_vec(mut vec: Vec) -> Self { let n = vec.len(); - if n != 0 { - let threshold = n.saturating_sub(n.div_ceil(384)); - turboselect::select_nth_unstable(&mut vec, threshold); - vec[threshold..].sort_unstable(); - Self { - threshold, - inner: vec, - } + if let Some(t) = NonZero::new(n / 384) { + let index = n - t.get(); + turboselect::select_nth_unstable(&mut vec, index); + vec[index..].sort_unstable(); + Self::Sorted { x: vec, t } } else { - Self { - threshold: 0, - inner: Vec::new(), + Self::Heap { + x: BinaryHeap::from(vec), } } } pub fn is_empty(&self) -> bool { - self.inner.is_empty() + match self { + SelectHeap::Sorted { x, .. } => x.is_empty(), + SelectHeap::Heap { x } => x.is_empty(), + } } pub fn pop(&mut self) -> Option { - if self.inner.len() <= self.threshold { - heapify::pop_heap(&mut self.inner); - } - let t = self.inner.pop(); - if self.inner.len() == self.threshold { - heapify::make_heap(&mut self.inner); + match self { + SelectHeap::Sorted { x: inner, t } => { + let x = inner.pop().expect("inconsistent internal structure"); + if let Some(value) = NonZero::new(t.get() - 1) { + *t = value; + } else { + *self = SelectHeap::Heap { + x: std::mem::take(inner).into(), + }; + } + Some(x) + } + SelectHeap::Heap { x } => x.pop(), } - t } pub fn peek(&self) -> Option<&T> { - if self.inner.len() <= self.threshold { - self.inner.first() - } else { - self.inner.last() + match self { + SelectHeap::Sorted { x, .. } => x.last(), + SelectHeap::Heap { x } => x.peek(), } } } @@ -61,3 +68,10 @@ fn test_select_heap() { assert_eq!(answer, result); } } + +#[test] +fn test_issue_209() { + let mut heap = SelectHeap::from_vec(vec![0]); + assert_eq!(heap.pop(), Some(0)); + assert_eq!(heap.pop(), None); +} From 3c749d53d8d834ebb4ac395cbf506159c96e181e Mon Sep 17 00:00:00 2001 From: xieydd Date: Tue, 18 Mar 2025 15:45:16 +0800 Subject: [PATCH 128/324] fix error in cnpg amd64 Dockerfile (#214) Already build v0.2.2 image in action https://github.com/tensorchord/VectorChord/actions/runs/13913843391. --------- Signed-off-by: xieydd --- docker/pg-cnpg/Dockerfile | 13 +- docker/pg-cnpg/requirements.txt | 720 +++++++++++++++++++------------- 2 files changed, 427 insertions(+), 306 deletions(-) diff --git a/docker/pg-cnpg/Dockerfile b/docker/pg-cnpg/Dockerfile index 7f9404ee..8ac07f2a 100644 --- a/docker/pg-cnpg/Dockerfile +++ b/docker/pg-cnpg/Dockerfile @@ -1,12 +1,12 @@ -ARG PG_MAJOR +ARG PG_MAJOR=16 ARG SEMVER -ARG TARGETARCH +ARG TARGETARCH=amd64 -FROM tensorchord/vchord-binary:pg${PG_MAJOR}-v${SEMVER} as binary +FROM tensorchord/vchord-binary:pg${PG_MAJOR}-v${SEMVER}-${TARGETARCH} AS binary -FROM rust:1.83-bookworm as builder +FROM rust:1.83-bookworm AS builder ARG TRUNK_VER=0.15.7 -ENV CARGO_REGISTRIES_CRATES_IO_PROTOCOL sparse +ENV CARGO_REGISTRIES_CRATES_IO_PROTOCOL=sparse RUN cargo install --version $TRUNK_VER pg-trunk FROM modelzai/pg-slim:${PG_MAJOR}-${TARGETARCH} @@ -173,9 +173,6 @@ RUN git clone https://github.com/EnterpriseDB/pg_failover_slots.git && \ make install PG_CONFIG=/usr/lib/postgresql/${PG_MAJOR}/bin/pg_config && \ cd ../ && rm -rf pg_failover_slots -# cache all extensions -ENV LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH - # Test trunk install extension COPY extension-install.sh /usr/local/bin/ diff --git a/docker/pg-cnpg/requirements.txt b/docker/pg-cnpg/requirements.txt index f7c47313..0aac74fd 100644 --- a/docker/pg-cnpg/requirements.txt +++ b/docker/pg-cnpg/requirements.txt @@ -10,32 +10,35 @@ azure-core==1.32.0 \ # via # azure-identity # azure-storage-blob -azure-identity==1.19.0 \ - --hash=sha256:500144dc18197d7019b81501165d4fa92225f03778f17d7ca8a2a180129a9c83 \ - --hash=sha256:e3f6558c181692d7509f09de10cca527c7dce426776454fb97df512a46527e81 -azure-storage-blob==12.23.1 \ - --hash=sha256:1c2238aa841d1545f42714a5017c010366137a44a0605da2d45f770174bfc6b4 \ - --hash=sha256:a587e54d4e39d2a27bd75109db164ffa2058fe194061e5446c5a89bca918272f -barman[azure,cloud,google,snappy]==3.11.1 \ - --hash=sha256:295b9b7e058e064338f66ca0d10e4892e784a2347f06e4a225164995f6114498 \ - --hash=sha256:4f424f3327cb24fb82d6a29dc1cdf02222b950c447c78273273d6eb76d7ce8d7 +azure-identity==1.21.0 \ + --hash=sha256:258ea6325537352440f71b35c3dffe9d240eae4a5126c1b7ce5efd5766bd9fd9 \ + --hash=sha256:ea22ce6e6b0f429bc1b8d9212d5b9f9877bd4c82f1724bfa910760612c07a9a6 +azure-storage-blob==12.25.0 \ + --hash=sha256:42364ca8f9f49dbccd0acc10144ed47bb6770bf78719970b51915f048891abba \ + --hash=sha256:a38e18bf10258fb19028f343db0d3d373280c6427a619c98c06d76485805b755 +barman[azure,cloud,google,lz4,snappy,zstandard]==3.12.1 \ + --hash=sha256:258ef7100717f66032402e0abe03c977089c50fc47143df5708e92aa1d772937 \ + --hash=sha256:9dd7be219b6f74954b80cdc28f9a72f2acb923e7da65edd0f41cdc82fd32e169 # via -r requirements.in -boto3==1.35.54 \ - --hash=sha256:2d5e160b614db55fbee7981001c54476cb827c441cef65b2fcb2c52a62019909 \ - --hash=sha256:7d9c359bbbc858a60b51c86328db813353c8bd1940212cdbd0a7da835291c2e1 -botocore==1.35.54 \ - --hash=sha256:131bb59ce59c8a939b31e8e647242d70cf11d32d4529fa4dca01feea1e891a76 \ - --hash=sha256:9cca1811094b6cdc144c2c063a3ec2db6d7c88194b04d4277cd34fc8e3473aff +boto3==1.35.99 \ + --hash=sha256:83e560faaec38a956dfb3d62e05e1703ee50432b45b788c09e25107c5058bd71 \ + --hash=sha256:e0abd794a7a591d90558e92e29a9f8837d25ece8e3c120e530526fe27eba5fca + # via + # -r requirements.in + # barman +botocore==1.35.99 \ + --hash=sha256:1eab44e969c39c5f3d9a3104a0836c24715579a455f12b3979a31d7cde51b3c3 \ + --hash=sha256:b22d27b6b617fc2d7342090d6129000af2efd20174215948c0d7ae2da0fab445 # via # boto3 # s3transfer -cachetools==5.5.0 \ - --hash=sha256:02134e8439cdc2ffb62023ce1debca2944c3f289d66bb17ead3ab3dede74b292 \ - --hash=sha256:2cc24fb4cbe39633fb7badd9db9ca6295d766d9c2995f245725a46715d050f2a +cachetools==5.5.2 \ + --hash=sha256:1a661caa9175d26759571b2e19580f9d6393969e5dfca11fdb1f947a23e640d4 \ + --hash=sha256:d26a22bcc62eb95c3beabd9f1ee5e820d3d2704fe2967cbe350e20c8ffcd3f0a # via google-auth -certifi==2024.8.30 \ - --hash=sha256:922820b53db7a7257ffbda3f597266d435245903d80737e34f8a45ff3e3230d8 \ - --hash=sha256:bec941d2aa8195e248a60b31ff9f0558284cf01a52591ceda73ea9afffd69fd9 +certifi==2025.1.31 \ + --hash=sha256:3d5da6925056f6f18f119200434a4780a94263f10d1c21d032a6f6b2baa20651 \ + --hash=sha256:ca78db4565a652026a4db2bcdf68f2fb589ea80d0be70e03929ed730746b84fe # via requests cffi==1.17.1 \ --hash=sha256:045d61c734659cc045141be4bae381a41d89b741f795af1dd018bfb532fd0df8 \ @@ -106,260 +109,255 @@ cffi==1.17.1 \ --hash=sha256:f7f5baafcc48261359e14bcd6d9bff6d4b28d9103847c9e136694cb0501aef87 \ --hash=sha256:fc48c783f9c87e60831201f2cce7f3b2e4846bf4d8728eabe54d60700b318a0b # via cryptography -charset-normalizer==3.4.0 \ - --hash=sha256:0099d79bdfcf5c1f0c2c72f91516702ebf8b0b8ddd8905f97a8aecf49712c621 \ - --hash=sha256:0713f3adb9d03d49d365b70b84775d0a0d18e4ab08d12bc46baa6132ba78aaf6 \ - --hash=sha256:07afec21bbbbf8a5cc3651aa96b980afe2526e7f048fdfb7f1014d84acc8b6d8 \ - --hash=sha256:0b309d1747110feb25d7ed6b01afdec269c647d382c857ef4663bbe6ad95a912 \ - --hash=sha256:0d99dd8ff461990f12d6e42c7347fd9ab2532fb70e9621ba520f9e8637161d7c \ - --hash=sha256:0de7b687289d3c1b3e8660d0741874abe7888100efe14bd0f9fd7141bcbda92b \ - --hash=sha256:1110e22af8ca26b90bd6364fe4c763329b0ebf1ee213ba32b68c73de5752323d \ - --hash=sha256:130272c698667a982a5d0e626851ceff662565379baf0ff2cc58067b81d4f11d \ - --hash=sha256:136815f06a3ae311fae551c3df1f998a1ebd01ddd424aa5603a4336997629e95 \ - --hash=sha256:14215b71a762336254351b00ec720a8e85cada43b987da5a042e4ce3e82bd68e \ - --hash=sha256:1db4e7fefefd0f548d73e2e2e041f9df5c59e178b4c72fbac4cc6f535cfb1565 \ - --hash=sha256:1ffd9493de4c922f2a38c2bf62b831dcec90ac673ed1ca182fe11b4d8e9f2a64 \ - --hash=sha256:2006769bd1640bdf4d5641c69a3d63b71b81445473cac5ded39740a226fa88ab \ - --hash=sha256:20587d20f557fe189b7947d8e7ec5afa110ccf72a3128d61a2a387c3313f46be \ - --hash=sha256:223217c3d4f82c3ac5e29032b3f1c2eb0fb591b72161f86d93f5719079dae93e \ - --hash=sha256:27623ba66c183eca01bf9ff833875b459cad267aeeb044477fedac35e19ba907 \ - --hash=sha256:285e96d9d53422efc0d7a17c60e59f37fbf3dfa942073f666db4ac71e8d726d0 \ - --hash=sha256:2de62e8801ddfff069cd5c504ce3bc9672b23266597d4e4f50eda28846c322f2 \ - --hash=sha256:2f6c34da58ea9c1a9515621f4d9ac379871a8f21168ba1b5e09d74250de5ad62 \ - --hash=sha256:309a7de0a0ff3040acaebb35ec45d18db4b28232f21998851cfa709eeff49d62 \ - --hash=sha256:35c404d74c2926d0287fbd63ed5d27eb911eb9e4a3bb2c6d294f3cfd4a9e0c23 \ - --hash=sha256:3710a9751938947e6327ea9f3ea6332a09bf0ba0c09cae9cb1f250bd1f1549bc \ - --hash=sha256:3d59d125ffbd6d552765510e3f31ed75ebac2c7470c7274195b9161a32350284 \ - --hash=sha256:40d3ff7fc90b98c637bda91c89d51264a3dcf210cade3a2c6f838c7268d7a4ca \ - --hash=sha256:425c5f215d0eecee9a56cdb703203dda90423247421bf0d67125add85d0c4455 \ - --hash=sha256:43193c5cda5d612f247172016c4bb71251c784d7a4d9314677186a838ad34858 \ - --hash=sha256:44aeb140295a2f0659e113b31cfe92c9061622cadbc9e2a2f7b8ef6b1e29ef4b \ - --hash=sha256:47334db71978b23ebcf3c0f9f5ee98b8d65992b65c9c4f2d34c2eaf5bcaf0594 \ - --hash=sha256:4796efc4faf6b53a18e3d46343535caed491776a22af773f366534056c4e1fbc \ - --hash=sha256:4a51b48f42d9358460b78725283f04bddaf44a9358197b889657deba38f329db \ - --hash=sha256:4b67fdab07fdd3c10bb21edab3cbfe8cf5696f453afce75d815d9d7223fbe88b \ - --hash=sha256:4ec9dd88a5b71abfc74e9df5ebe7921c35cbb3b641181a531ca65cdb5e8e4dea \ - --hash=sha256:4f9fc98dad6c2eaa32fc3af1417d95b5e3d08aff968df0cd320066def971f9a6 \ - --hash=sha256:54b6a92d009cbe2fb11054ba694bc9e284dad30a26757b1e372a1fdddaf21920 \ - --hash=sha256:55f56e2ebd4e3bc50442fbc0888c9d8c94e4e06a933804e2af3e89e2f9c1c749 \ - --hash=sha256:5726cf76c982532c1863fb64d8c6dd0e4c90b6ece9feb06c9f202417a31f7dd7 \ - --hash=sha256:5d447056e2ca60382d460a604b6302d8db69476fd2015c81e7c35417cfabe4cd \ - --hash=sha256:5ed2e36c3e9b4f21dd9422f6893dec0abf2cca553af509b10cd630f878d3eb99 \ - --hash=sha256:5ff2ed8194587faf56555927b3aa10e6fb69d931e33953943bc4f837dfee2242 \ - --hash=sha256:62f60aebecfc7f4b82e3f639a7d1433a20ec32824db2199a11ad4f5e146ef5ee \ - --hash=sha256:63bc5c4ae26e4bc6be6469943b8253c0fd4e4186c43ad46e713ea61a0ba49129 \ - --hash=sha256:6b40e8d38afe634559e398cc32b1472f376a4099c75fe6299ae607e404c033b2 \ - --hash=sha256:6b493a043635eb376e50eedf7818f2f322eabbaa974e948bd8bdd29eb7ef2a51 \ - --hash=sha256:6dba5d19c4dfab08e58d5b36304b3f92f3bd5d42c1a3fa37b5ba5cdf6dfcbcee \ - --hash=sha256:6fd30dc99682dc2c603c2b315bded2799019cea829f8bf57dc6b61efde6611c8 \ - --hash=sha256:707b82d19e65c9bd28b81dde95249b07bf9f5b90ebe1ef17d9b57473f8a64b7b \ - --hash=sha256:7706f5850360ac01d80c89bcef1640683cc12ed87f42579dab6c5d3ed6888613 \ - --hash=sha256:7782afc9b6b42200f7362858f9e73b1f8316afb276d316336c0ec3bd73312742 \ - --hash=sha256:79983512b108e4a164b9c8d34de3992f76d48cadc9554c9e60b43f308988aabe \ - --hash=sha256:7f683ddc7eedd742e2889d2bfb96d69573fde1d92fcb811979cdb7165bb9c7d3 \ - --hash=sha256:82357d85de703176b5587dbe6ade8ff67f9f69a41c0733cf2425378b49954de5 \ - --hash=sha256:84450ba661fb96e9fd67629b93d2941c871ca86fc38d835d19d4225ff946a631 \ - --hash=sha256:86f4e8cca779080f66ff4f191a685ced73d2f72d50216f7112185dc02b90b9b7 \ - --hash=sha256:8cda06946eac330cbe6598f77bb54e690b4ca93f593dee1568ad22b04f347c15 \ - --hash=sha256:8ce7fd6767a1cc5a92a639b391891bf1c268b03ec7e021c7d6d902285259685c \ - --hash=sha256:8ff4e7cdfdb1ab5698e675ca622e72d58a6fa2a8aa58195de0c0061288e6e3ea \ - --hash=sha256:9289fd5dddcf57bab41d044f1756550f9e7cf0c8e373b8cdf0ce8773dc4bd417 \ - --hash=sha256:92a7e36b000bf022ef3dbb9c46bfe2d52c047d5e3f3343f43204263c5addc250 \ - --hash=sha256:92db3c28b5b2a273346bebb24857fda45601aef6ae1c011c0a997106581e8a88 \ - --hash=sha256:95c3c157765b031331dd4db3c775e58deaee050a3042fcad72cbc4189d7c8dca \ - --hash=sha256:980b4f289d1d90ca5efcf07958d3eb38ed9c0b7676bf2831a54d4f66f9c27dfa \ - --hash=sha256:9ae4ef0b3f6b41bad6366fb0ea4fc1d7ed051528e113a60fa2a65a9abb5b1d99 \ - --hash=sha256:9c98230f5042f4945f957d006edccc2af1e03ed5e37ce7c373f00a5a4daa6149 \ - --hash=sha256:9fa2566ca27d67c86569e8c85297aaf413ffab85a8960500f12ea34ff98e4c41 \ - --hash=sha256:a14969b8691f7998e74663b77b4c36c0337cb1df552da83d5c9004a93afdb574 \ - --hash=sha256:a8aacce6e2e1edcb6ac625fb0f8c3a9570ccc7bfba1f63419b3769ccf6a00ed0 \ - --hash=sha256:a8e538f46104c815be19c975572d74afb53f29650ea2025bbfaef359d2de2f7f \ - --hash=sha256:aa41e526a5d4a9dfcfbab0716c7e8a1b215abd3f3df5a45cf18a12721d31cb5d \ - --hash=sha256:aa693779a8b50cd97570e5a0f343538a8dbd3e496fa5dcb87e29406ad0299654 \ - --hash=sha256:ab22fbd9765e6954bc0bcff24c25ff71dcbfdb185fcdaca49e81bac68fe724d3 \ - --hash=sha256:ab2e5bef076f5a235c3774b4f4028a680432cded7cad37bba0fd90d64b187d19 \ - --hash=sha256:ab973df98fc99ab39080bfb0eb3a925181454d7c3ac8a1e695fddfae696d9e90 \ - --hash=sha256:af73657b7a68211996527dbfeffbb0864e043d270580c5aef06dc4b659a4b578 \ - --hash=sha256:b197e7094f232959f8f20541ead1d9862ac5ebea1d58e9849c1bf979255dfac9 \ - --hash=sha256:b295729485b06c1a0683af02a9e42d2caa9db04a373dc38a6a58cdd1e8abddf1 \ - --hash=sha256:b8831399554b92b72af5932cdbbd4ddc55c55f631bb13ff8fe4e6536a06c5c51 \ - --hash=sha256:b8dcd239c743aa2f9c22ce674a145e0a25cb1566c495928440a181ca1ccf6719 \ - --hash=sha256:bcb4f8ea87d03bc51ad04add8ceaf9b0f085ac045ab4d74e73bbc2dc033f0236 \ - --hash=sha256:bd7af3717683bea4c87acd8c0d3d5b44d56120b26fd3f8a692bdd2d5260c620a \ - --hash=sha256:bf4475b82be41b07cc5e5ff94810e6a01f276e37c2d55571e3fe175e467a1a1c \ - --hash=sha256:c3e446d253bd88f6377260d07c895816ebf33ffffd56c1c792b13bff9c3e1ade \ - --hash=sha256:c57516e58fd17d03ebe67e181a4e4e2ccab1168f8c2976c6a334d4f819fe5944 \ - --hash=sha256:c94057af19bc953643a33581844649a7fdab902624d2eb739738a30e2b3e60fc \ - --hash=sha256:cab5d0b79d987c67f3b9e9c53f54a61360422a5a0bc075f43cab5621d530c3b6 \ - --hash=sha256:ce031db0408e487fd2775d745ce30a7cd2923667cf3b69d48d219f1d8f5ddeb6 \ - --hash=sha256:cee4373f4d3ad28f1ab6290684d8e2ebdb9e7a1b74fdc39e4c211995f77bec27 \ - --hash=sha256:d5b054862739d276e09928de37c79ddeec42a6e1bfc55863be96a36ba22926f6 \ - --hash=sha256:dbe03226baf438ac4fda9e2d0715022fd579cb641c4cf639fa40d53b2fe6f3e2 \ - --hash=sha256:dc15e99b2d8a656f8e666854404f1ba54765871104e50c8e9813af8a7db07f12 \ - --hash=sha256:dcaf7c1524c0542ee2fc82cc8ec337f7a9f7edee2532421ab200d2b920fc97cf \ - --hash=sha256:dd4eda173a9fcccb5f2e2bd2a9f423d180194b1bf17cf59e3269899235b2a114 \ - --hash=sha256:dd9a8bd8900e65504a305bf8ae6fa9fbc66de94178c420791d0293702fce2df7 \ - --hash=sha256:de7376c29d95d6719048c194a9cf1a1b0393fbe8488a22008610b0361d834ecf \ - --hash=sha256:e7fdd52961feb4c96507aa649550ec2a0d527c086d284749b2f582f2d40a2e0d \ - --hash=sha256:e91f541a85298cf35433bf66f3fab2a4a2cff05c127eeca4af174f6d497f0d4b \ - --hash=sha256:e9e3c4c9e1ed40ea53acf11e2a386383c3304212c965773704e4603d589343ed \ - --hash=sha256:ee803480535c44e7f5ad00788526da7d85525cfefaf8acf8ab9a310000be4b03 \ - --hash=sha256:f09cb5a7bbe1ecae6e87901a2eb23e0256bb524a79ccc53eb0b7629fbe7677c4 \ - --hash=sha256:f19c1585933c82098c2a520f8ec1227f20e339e33aca8fa6f956f6691b784e67 \ - --hash=sha256:f1a2f519ae173b5b6a2c9d5fa3116ce16e48b3462c8b96dfdded11055e3d6365 \ - --hash=sha256:f28f891ccd15c514a0981f3b9db9aa23d62fe1a99997512b0491d2ed323d229a \ - --hash=sha256:f3e73a4255342d4eb26ef6df01e3962e73aa29baa3124a8e824c5d3364a65748 \ - --hash=sha256:f606a1881d2663630ea5b8ce2efe2111740df4b687bd78b34a8131baa007f79b \ - --hash=sha256:fe9f97feb71aa9896b81973a7bbada8c49501dc73e58a10fcef6663af95e5079 \ - --hash=sha256:ffc519621dce0c767e96b9c53f09c5d215578e10b02c285809f76509a3931482 +charset-normalizer==3.4.1 \ + --hash=sha256:0167ddc8ab6508fe81860a57dd472b2ef4060e8d378f0cc555707126830f2537 \ + --hash=sha256:01732659ba9b5b873fc117534143e4feefecf3b2078b0a6a2e925271bb6f4cfa \ + --hash=sha256:01ad647cdd609225c5350561d084b42ddf732f4eeefe6e678765636791e78b9a \ + --hash=sha256:04432ad9479fa40ec0f387795ddad4437a2b50417c69fa275e212933519ff294 \ + --hash=sha256:0907f11d019260cdc3f94fbdb23ff9125f6b5d1039b76003b5b0ac9d6a6c9d5b \ + --hash=sha256:0924e81d3d5e70f8126529951dac65c1010cdf117bb75eb02dd12339b57749dd \ + --hash=sha256:09b26ae6b1abf0d27570633b2b078a2a20419c99d66fb2823173d73f188ce601 \ + --hash=sha256:09b5e6733cbd160dcc09589227187e242a30a49ca5cefa5a7edd3f9d19ed53fd \ + --hash=sha256:0af291f4fe114be0280cdd29d533696a77b5b49cfde5467176ecab32353395c4 \ + --hash=sha256:0f55e69f030f7163dffe9fd0752b32f070566451afe180f99dbeeb81f511ad8d \ + --hash=sha256:1a2bc9f351a75ef49d664206d51f8e5ede9da246602dc2d2726837620ea034b2 \ + --hash=sha256:22e14b5d70560b8dd51ec22863f370d1e595ac3d024cb8ad7d308b4cd95f8313 \ + --hash=sha256:234ac59ea147c59ee4da87a0c0f098e9c8d169f4dc2a159ef720f1a61bbe27cd \ + --hash=sha256:2369eea1ee4a7610a860d88f268eb39b95cb588acd7235e02fd5a5601773d4fa \ + --hash=sha256:237bdbe6159cff53b4f24f397d43c6336c6b0b42affbe857970cefbb620911c8 \ + --hash=sha256:28bf57629c75e810b6ae989f03c0828d64d6b26a5e205535585f96093e405ed1 \ + --hash=sha256:2967f74ad52c3b98de4c3b32e1a44e32975e008a9cd2a8cc8966d6a5218c5cb2 \ + --hash=sha256:2a75d49014d118e4198bcee5ee0a6f25856b29b12dbf7cd012791f8a6cc5c496 \ + --hash=sha256:2bdfe3ac2e1bbe5b59a1a63721eb3b95fc9b6817ae4a46debbb4e11f6232428d \ + --hash=sha256:2d074908e1aecee37a7635990b2c6d504cd4766c7bc9fc86d63f9c09af3fa11b \ + --hash=sha256:2fb9bd477fdea8684f78791a6de97a953c51831ee2981f8e4f583ff3b9d9687e \ + --hash=sha256:311f30128d7d333eebd7896965bfcfbd0065f1716ec92bd5638d7748eb6f936a \ + --hash=sha256:329ce159e82018d646c7ac45b01a430369d526569ec08516081727a20e9e4af4 \ + --hash=sha256:345b0426edd4e18138d6528aed636de7a9ed169b4aaf9d61a8c19e39d26838ca \ + --hash=sha256:363e2f92b0f0174b2f8238240a1a30142e3db7b957a5dd5689b0e75fb717cc78 \ + --hash=sha256:3a3bd0dcd373514dcec91c411ddb9632c0d7d92aed7093b8c3bbb6d69ca74408 \ + --hash=sha256:3bed14e9c89dcb10e8f3a29f9ccac4955aebe93c71ae803af79265c9ca5644c5 \ + --hash=sha256:44251f18cd68a75b56585dd00dae26183e102cd5e0f9f1466e6df5da2ed64ea3 \ + --hash=sha256:44ecbf16649486d4aebafeaa7ec4c9fed8b88101f4dd612dcaf65d5e815f837f \ + --hash=sha256:4532bff1b8421fd0a320463030c7520f56a79c9024a4e88f01c537316019005a \ + --hash=sha256:49402233c892a461407c512a19435d1ce275543138294f7ef013f0b63d5d3765 \ + --hash=sha256:4c0907b1928a36d5a998d72d64d8eaa7244989f7aaaf947500d3a800c83a3fd6 \ + --hash=sha256:4d86f7aff21ee58f26dcf5ae81a9addbd914115cdebcbb2217e4f0ed8982e146 \ + --hash=sha256:5777ee0881f9499ed0f71cc82cf873d9a0ca8af166dfa0af8ec4e675b7df48e6 \ + --hash=sha256:5df196eb874dae23dcfb968c83d4f8fdccb333330fe1fc278ac5ceeb101003a9 \ + --hash=sha256:619a609aa74ae43d90ed2e89bdd784765de0a25ca761b93e196d938b8fd1dbbd \ + --hash=sha256:6e27f48bcd0957c6d4cb9d6fa6b61d192d0b13d5ef563e5f2ae35feafc0d179c \ + --hash=sha256:6ff8a4a60c227ad87030d76e99cd1698345d4491638dfa6673027c48b3cd395f \ + --hash=sha256:73d94b58ec7fecbc7366247d3b0b10a21681004153238750bb67bd9012414545 \ + --hash=sha256:7461baadb4dc00fd9e0acbe254e3d7d2112e7f92ced2adc96e54ef6501c5f176 \ + --hash=sha256:75832c08354f595c760a804588b9357d34ec00ba1c940c15e31e96d902093770 \ + --hash=sha256:7709f51f5f7c853f0fb938bcd3bc59cdfdc5203635ffd18bf354f6967ea0f824 \ + --hash=sha256:78baa6d91634dfb69ec52a463534bc0df05dbd546209b79a3880a34487f4b84f \ + --hash=sha256:7974a0b5ecd505609e3b19742b60cee7aa2aa2fb3151bc917e6e2646d7667dcf \ + --hash=sha256:7a4f97a081603d2050bfaffdefa5b02a9ec823f8348a572e39032caa8404a487 \ + --hash=sha256:7b1bef6280950ee6c177b326508f86cad7ad4dff12454483b51d8b7d673a2c5d \ + --hash=sha256:7d053096f67cd1241601111b698f5cad775f97ab25d81567d3f59219b5f1adbd \ + --hash=sha256:804a4d582ba6e5b747c625bf1255e6b1507465494a40a2130978bda7b932c90b \ + --hash=sha256:807f52c1f798eef6cf26beb819eeb8819b1622ddfeef9d0977a8502d4db6d534 \ + --hash=sha256:80ed5e856eb7f30115aaf94e4a08114ccc8813e6ed1b5efa74f9f82e8509858f \ + --hash=sha256:8417cb1f36cc0bc7eaba8ccb0e04d55f0ee52df06df3ad55259b9a323555fc8b \ + --hash=sha256:8436c508b408b82d87dc5f62496973a1805cd46727c34440b0d29d8a2f50a6c9 \ + --hash=sha256:89149166622f4db9b4b6a449256291dc87a99ee53151c74cbd82a53c8c2f6ccd \ + --hash=sha256:8bfa33f4f2672964266e940dd22a195989ba31669bd84629f05fab3ef4e2d125 \ + --hash=sha256:8c60ca7339acd497a55b0ea5d506b2a2612afb2826560416f6894e8b5770d4a9 \ + --hash=sha256:91b36a978b5ae0ee86c394f5a54d6ef44db1de0815eb43de826d41d21e4af3de \ + --hash=sha256:955f8851919303c92343d2f66165294848d57e9bba6cf6e3625485a70a038d11 \ + --hash=sha256:97f68b8d6831127e4787ad15e6757232e14e12060bec17091b85eb1486b91d8d \ + --hash=sha256:9b23ca7ef998bc739bf6ffc077c2116917eabcc901f88da1b9856b210ef63f35 \ + --hash=sha256:9f0b8b1c6d84c8034a44893aba5e767bf9c7a211e313a9605d9c617d7083829f \ + --hash=sha256:aabfa34badd18f1da5ec1bc2715cadc8dca465868a4e73a0173466b688f29dda \ + --hash=sha256:ab36c8eb7e454e34e60eb55ca5d241a5d18b2c6244f6827a30e451c42410b5f7 \ + --hash=sha256:b010a7a4fd316c3c484d482922d13044979e78d1861f0e0650423144c616a46a \ + --hash=sha256:b1ac5992a838106edb89654e0aebfc24f5848ae2547d22c2c3f66454daa11971 \ + --hash=sha256:b7b2d86dd06bfc2ade3312a83a5c364c7ec2e3498f8734282c6c3d4b07b346b8 \ + --hash=sha256:b97e690a2118911e39b4042088092771b4ae3fc3aa86518f84b8cf6888dbdb41 \ + --hash=sha256:bc2722592d8998c870fa4e290c2eec2c1569b87fe58618e67d38b4665dfa680d \ + --hash=sha256:c0429126cf75e16c4f0ad00ee0eae4242dc652290f940152ca8c75c3a4b6ee8f \ + --hash=sha256:c30197aa96e8eed02200a83fba2657b4c3acd0f0aa4bdc9f6c1af8e8962e0757 \ + --hash=sha256:c4c3e6da02df6fa1410a7680bd3f63d4f710232d3139089536310d027950696a \ + --hash=sha256:c75cb2a3e389853835e84a2d8fb2b81a10645b503eca9bcb98df6b5a43eb8886 \ + --hash=sha256:c96836c97b1238e9c9e3fe90844c947d5afbf4f4c92762679acfe19927d81d77 \ + --hash=sha256:d7f50a1f8c450f3925cb367d011448c39239bb3eb4117c36a6d354794de4ce76 \ + --hash=sha256:d973f03c0cb71c5ed99037b870f2be986c3c05e63622c017ea9816881d2dd247 \ + --hash=sha256:d98b1668f06378c6dbefec3b92299716b931cd4e6061f3c875a71ced1780ab85 \ + --hash=sha256:d9c3cdf5390dcd29aa8056d13e8e99526cda0305acc038b96b30352aff5ff2bb \ + --hash=sha256:dad3e487649f498dd991eeb901125411559b22e8d7ab25d3aeb1af367df5efd7 \ + --hash=sha256:dccbe65bd2f7f7ec22c4ff99ed56faa1e9f785482b9bbd7c717e26fd723a1d1e \ + --hash=sha256:dd78cfcda14a1ef52584dbb008f7ac81c1328c0f58184bf9a84c49c605002da6 \ + --hash=sha256:e218488cd232553829be0664c2292d3af2eeeb94b32bea483cf79ac6a694e037 \ + --hash=sha256:e358e64305fe12299a08e08978f51fc21fac060dcfcddd95453eabe5b93ed0e1 \ + --hash=sha256:ea0d8d539afa5eb2728aa1932a988a9a7af94f18582ffae4bc10b3fbdad0626e \ + --hash=sha256:eab677309cdb30d047996b36d34caeda1dc91149e4fdca0b1a039b3f79d9a807 \ + --hash=sha256:eb8178fe3dba6450a3e024e95ac49ed3400e506fd4e9e5c32d30adda88cbd407 \ + --hash=sha256:ecddf25bee22fe4fe3737a399d0d177d72bc22be6913acfab364b40bce1ba83c \ + --hash=sha256:eea6ee1db730b3483adf394ea72f808b6e18cf3cb6454b4d86e04fa8c4327a12 \ + --hash=sha256:f08ff5e948271dc7e18a35641d2f11a4cd8dfd5634f55228b691e62b37125eb3 \ + --hash=sha256:f30bf9fd9be89ecb2360c7d94a711f00c09b976258846efe40db3d05828e8089 \ + --hash=sha256:fa88b843d6e211393a37219e6a1c1df99d35e8fd90446f1118f4216e307e48cd \ + --hash=sha256:fc54db6c8593ef7d4b2a331b58653356cf04f67c960f584edb7c3d8c97e8f39e \ + --hash=sha256:fd4ec41f914fa74ad1b8304bbc634b3de73d2a0889bd32076342a573e0779e00 \ + --hash=sha256:ffc9202a29ab3920fa812879e95a9e78b2465fd10be7fcbd042899695d75e616 # via requests -cramjam==2.9.0 \ - --hash=sha256:00d96f798bc980b29f8e1c3ed7d554050e05d4cde23d1633ffed4cd63110024a \ - --hash=sha256:013cb872205641c6e5269f530ed40aaaa5640d84e0d8f33b89f5a1bf7f655527 \ - --hash=sha256:017b7066f18b7b676068f51b1dbdecc02d76d9af10092252b22dcbd03a78ed33 \ - --hash=sha256:05970fb640f236767003e62c256a085754536169bac863f4a3502ecb59cbf197 \ - --hash=sha256:0711c776750e243ae347d6609c975f0ff4be9ae65b2764d29e4bbdad8e574c3a \ - --hash=sha256:0b1410e68c464666473a89cade17483b94bb4639d9161c440ee54ee1e0eca583 \ - --hash=sha256:0b4f1b5e33915ed591c0c19b8c3bbdd7aa0f6a9bfe2b7246b475d497bda15f18 \ - --hash=sha256:0ed2fef010d1caca9ea63814e9cb5b1d47d907b80302b8cc0b3a1e116ea241e2 \ - --hash=sha256:0ed6362cb6c964f8d0c6e7f790e8961b9242cd3acd87c56169ca14d642653707 \ - --hash=sha256:0fb5ea631dbf998f667766a9e485e757817d66ed559916ba553a0ec2f902d788 \ - --hash=sha256:132db7d3346ea21ba44e7ee23ec73bd6fa9eb1e77133ca6dfe1f7449a69999af \ - --hash=sha256:170a50407f9400073621cc1d5f3200ca3ad9de3000831e3e86f5561ca8048a08 \ - --hash=sha256:1b4ca30c9f27e3b88bc082d4637e7648f93da5cb69a2dbe0c0300bc51353c820 \ - --hash=sha256:1ca28a8f6ab5fca35f163fd7d7a970880ce4fc1a0bead1249ecdaa96ec9ac1f4 \ - --hash=sha256:1d5b5512dc61ea78f32e021e88a5fd5b46a821409479e6657d33614fc9e45677 \ - --hash=sha256:21569f19d5848606b85ac0dde0dc3639319d26fed8522c7103515df875bcb300 \ - --hash=sha256:24afad3ba62774abbb150dc25aab21b047ab999c4143c7a8d96577848baf7af6 \ - --hash=sha256:2addf801c88bead21256ccd87dc97cffead03758c4a4947fad8e454f4abfda0a \ - --hash=sha256:3121e2fbec58907fa70636adaeaf30c27614c867e08a7a5bd2887b33786ff790 \ - --hash=sha256:34578e4c1518b10dad5e0ba40c721e529ef13e7742a528843b40e1f20dd6078c \ - --hash=sha256:35cad507eb02c775e6c5444312f98b28dd8bf122425677ae199484996e838673 \ - --hash=sha256:36426e3f1920f6aa4c644d007bf9cfad06dd9f1a30cd0a921d72b010492d8447 \ - --hash=sha256:37054c73704a3183b60869e7fec1614648752c31d89f44de1ffe1f01ad4d20d5 \ - --hash=sha256:399baf80fea574e3870f233e12e6a12f02c53b054e13d792348b272b0614370a \ - --hash=sha256:441d3875cdffe5df9294b93ef570058837732dd727cd9d18efa0f089f1c2687a \ - --hash=sha256:4714e1ea0c3329368b83fe5ad6e831d5ca11fb794ca7cf491622eb6b2d420d2f \ - --hash=sha256:47d7253b5a10c201cc65aecfb517dfa1c0b5831b2524ac32dd2964fceafc0dc4 \ - --hash=sha256:4a63c4e63319bf7dfc3ab46c06afb76d3d9cc1c94369b609dde480e5cc78e4de \ - --hash=sha256:4adbf4366f8dc29b7c5c731c800cf633be76c9911e928daeb606827d6ae7c599 \ - --hash=sha256:4bd76b654275736fd4f55521981b73751c34dacf70a1dbce96e454a39d43201f \ - --hash=sha256:4d5a39118008bb9f2fba36a0ceea6c41fbd0b55d2647b043ba51a868e5f6de92 \ - --hash=sha256:598eac1713ddbe69c3b30dcc890d69b206ce08903fc3aed58149aae87c61973a \ - --hash=sha256:604c16052cf29d0c796927ed7e107f65429d2036c82c9a8009bd453c94e5e4f0 \ - --hash=sha256:65a097ea765dd4ef2fb868b5b0959d7c93a64c250b2c52f462898c823ae4b950 \ - --hash=sha256:65bded20fd2cef17b22246c336ddd67fac842341ee311042b4a70e65dc745aa7 \ - --hash=sha256:6ebee5f5d7e2b9277895ea4fd94646b72075fe9cfc0e8f4770b65c9e72b1fec1 \ - --hash=sha256:72e9ebc27c557706a3c9964c1d1b4522857760dbd60c105a4f5421f3b66e31a2 \ - --hash=sha256:78e7349f945a83bc48855fb042873092a69b155a088b8c11942eb76418b32705 \ - --hash=sha256:7de19a382bcab93cd4d028d51f6f581920a3b79659a384775188135b7fc64f15 \ - --hash=sha256:7f33a83969fa94ee8e0c1f0aef8eb303ead3e9142338dc543abeb7e1a28734ab \ - --hash=sha256:7f6ef35eba883927af2678b561cc4407e0b3b0d58a251c863bec4b3d8258cc2f \ - --hash=sha256:8982925d179b940efa860513a31b839bb06343501077cca3e67f7a2f7360d355 \ - --hash=sha256:8e33ebe4d709b21bc15e7ddf485ac6b30d7fdc7ed7c3c65130654c007f50c183 \ - --hash=sha256:904be92e3bc25e78343ee52aa0fd5fba3a31d11d474e8af4623a9d00baa84bc2 \ - --hash=sha256:912c94781c8ff318a4d3f3306f8d94d41ae5aa7b9760c4bb0476b01142084845 \ - --hash=sha256:9221297c547d702e1431e96705fce26c6a87df34a681a6b97fe63b536d09c1d8 \ - --hash=sha256:97a6311bd32f301ff1b922bc9de62ace3d9fd845e20efc0f71b4d0239a45b8d2 \ - --hash=sha256:9862ca8ead80857ecfb9b07f02f577733261e981346f31585fe118975eabb738 \ - --hash=sha256:9de33ef3bc006c11fbad1dc8b15341dcc78430df2c5ce1e790dfb729b11ab593 \ - --hash=sha256:9f685fe4e49b2f3e233548e3397b3f9189d71a265718ec631d13eca3d5718ddb \ - --hash=sha256:a0f654c739a6bc4a69a2aaf31463328a208757ed780ff886234532f78e06a864 \ - --hash=sha256:a4156fcefa1dfaa65d35ff82c252d1e32be12820f26d04748be6cd3b461cf85f \ - --hash=sha256:a41b4b10a381be1d42a1a7dd07b8c3faccd3d12c7e98e973a6ec558fd040a607 \ - --hash=sha256:ab17a429a92db90bf40115efb97d10e71b94b0dcacf30cf724552df2794a58fb \ - --hash=sha256:abd8bf9a94e3866215ac181a7dbcfa1ddbedca4f8048494a79934febe88537df \ - --hash=sha256:ad301801afa0eecdacabf353a2802df5e6770f9bfb0a559d6c069813d83cfd42 \ - --hash=sha256:b0078727fe8c28ef1695e5d04aae5c41ac697eb087cba387c6a02b825f9071c0 \ - --hash=sha256:b21e55b5cfdaff96eae1f323ae9a0d36e86852cdf62fe23b60a2481d2fed5571 \ - --hash=sha256:b4a3104022129d7463100dfaf12efd398ebfa4b7e4e50832ccc596754f7c26df \ - --hash=sha256:b4b8d8160685c11ffb4e8e6daaab79cb351a1c54ceec41cc18a0a62c89309fe0 \ - --hash=sha256:b8f8b1117b4e697d39950ecab01700ce0aef66541e4478eb4d7b3ade8703347b \ - --hash=sha256:b99efaf81be8e381de1cde6574e2c89030ed53994e73b0e75b62d6e232f491c5 \ - --hash=sha256:ba7e2d33e1d092dffd0a3ff4bd1b86177594aa3c2901fd478e78e1fb2aee8ed3 \ - --hash=sha256:bafc32f01d4ab64f83fdbc29bc5bd25a920b59c751c12e06e6f4b1e379be7600 \ - --hash=sha256:bd04205b2a87087ffc2257c3ad33f11daabc053956f64ac1ec7bae299cac3f2f \ - --hash=sha256:bd26d71939de5dcf169d479fbc7fcfed21e6675bab33e7f7e9f8405f19711c71 \ - --hash=sha256:c3464d0042a03e8ef38a2b774ef23163cf3c0cdc41b8dfbf7c4aadf93e40b459 \ - --hash=sha256:c48da60a5eb481b412e5e462b81ad307fb2203178a2840a743f0a7c5fc1718c9 \ - --hash=sha256:c4fa6c23e56d48df18f534af921ec936c812743a8972ecdd5e5ff47b464fea00 \ - --hash=sha256:c902e56e60c48f5f15e55257aaa1c2678323df5f18a1b839e8d05cac1107576c \ - --hash=sha256:ca880f555c8db40942acc8a50722c33e229b6be90e598acc1a201f36487b917d \ - --hash=sha256:cb1e86bfea656b51f2e75f2cedb17fc08b552d105b814d19b595294ecbe94d8d \ - --hash=sha256:cd4d4ab9deb5846af0ac6cf1fa139cfa40291ad14d073efa8b8e20c8d1aa90bd \ - --hash=sha256:dbbd6fba677e1cbc9d6bd4ebbe3e8b3667d0295f1731489db2a971c95f0ceca0 \ - --hash=sha256:dc07376aa33b6004ea372ac9b0ba0ed3455aa2fc4e18727414142ecb46b176b8 \ - --hash=sha256:dd70ea5d7b2c5e479e04ac3a00d8bc3deca146d2b5dbfbe3d7b42ed136e19de4 \ - --hash=sha256:ddb9c4db36188a8f08c2303100a83100f26a8572803ae35eadff359bebd3d204 \ - --hash=sha256:df089639983a03070be6eabc60317aa1ffbf2c5409023b57a5fc2e4975163bc4 \ - --hash=sha256:e0b062d261fa3fac00146cf801896c8cfafe1e41332eb047aa0a36558299daa6 \ - --hash=sha256:e248510f8e2dbc71fa99f86238c9023365dbe1a4520eb40e33d73416527349f2 \ - --hash=sha256:e94021c541eb2a199b5a2ffae0ea84fb8b99863dab99a5b154b00bc7a44b5c48 \ - --hash=sha256:e98a18c22a85f321091cc8db6694af1d713a369c2d60ec611c10ccfe24ab103a \ - --hash=sha256:ea9bcaff298f5d35ef67346d474fca388c5cf6d4edab1d06b84868800f88bd36 \ - --hash=sha256:eb16d995e454b0155b166f6e6da7df4ac812d44e0f3b6dc0f344a934609fd5bc \ - --hash=sha256:ed486e57a79ccc7aebaa2ec12517d891fdc5d2fde16915e3db705b8a47570981 \ - --hash=sha256:ed7fd7bc2b86ec3161fe0cc49f5f392e6efa55c91a95397d5047820c38117660 \ - --hash=sha256:ef553d4080368006817c1a935ed619c71987cf10417a32386acc00c5418a2934 \ - --hash=sha256:f103e648aa3ebe9b8e2c1a3a92719288d8f3f41007c319ad298cdce2d0c28641 \ - --hash=sha256:fc49b6575e3cb15da3180c5a3926ec81db33b109e48530708da76614b306904b \ - --hash=sha256:fe9af350dfbdc7ed4c93a8016a8ad7b5492fc116e7197cad7cbce99b434d3fe1 +cramjam==2.9.1 \ + --hash=sha256:008b49b455b396acc5459dfb06fb9d56049c4097ee8e590892a4d3da9a711da3 \ + --hash=sha256:038df668ffb94d64d67b6ecc59cbd206745a425ffc0402897dde12d89fa6a870 \ + --hash=sha256:04828cbfad7384f06a4a7d0d927c3e85ef11dc5a40b9cf5f3e29ac4e23ecd678 \ + --hash=sha256:06068bd191a82ad4fc1ac23d6f8627fb5e37ec4be0431711b9a2dbacaccfeddb \ + --hash=sha256:06e3f97a379386d97debf08638a78b3d3850fdf6124755eb270b54905a169930 \ + --hash=sha256:07ac76b7f992556e7aa910244be11ece578cdf84f4d5d5297461f9a895e18312 \ + --hash=sha256:0bedb84e068b53c944bd08dcb501fd00d67daa8a917922356dd559b484ce7eab \ + --hash=sha256:11118675e9c7952ececabc62f023290ee4f8ecf0bee0d2c7eb8d1c402ee9769d \ + --hash=sha256:1376f6fdbf0b30712413a0b4e51663a4938ae2f6b449f8e4635dbb3694db83cf \ + --hash=sha256:1539fd758f0e57fad7913cebff8baaee871bb561ddf6fa710a427b74da6b6778 \ + --hash=sha256:15955dd75e80f66c1ea271167a5347661d9bdc365f894a57698c383c9b7d465c \ + --hash=sha256:1c33bc095db5733c841a102b8693062be5db8cdac17b9782ebc00577c6a94480 \ + --hash=sha256:217fe22b41f8c3dce03852f828b059abfad11d1344a1df2f43d3eb8634b18d75 \ + --hash=sha256:21ea784e6c3f1843d3523ae0f03651dd06058b39eeb64beb82ee3b100fa83662 \ + --hash=sha256:23b9786d1d17686fb8d600ade2a19374c7188d4b8867efa9af0d8274a220aec7 \ + --hash=sha256:258120cb1e3afc3443f756f9de161ed63eed56a2c31f6093e81c571c0f2dc9f6 \ + --hash=sha256:27571bfa5a5d618604696747d0dc1d2a99b5906c967c8dee53c13a7107edfde6 \ + --hash=sha256:2ceef6e09ee22457997370882aa3c69de01e6dd0aaa2f953e1e87ad11641d042 \ + --hash=sha256:335103317475bf992953c58838152a4761fc3c87354000edbfc4d7e57cf05909 \ + --hash=sha256:336cc591d86cbd225d256813779f46624f857bc9c779db126271eff9ddc524ae \ + --hash=sha256:342fb946f8d3e9e35b837288b03ab23cfbe0bb5a30e582ed805ef79706823a96 \ + --hash=sha256:3b695259e71fde6d5be66b77a4474523ced9ffe9fe8a34cb9b520ec1241a14d3 \ + --hash=sha256:3ba79c7d2cc5adb897b690c05dd9b67c4d401736d207314b99315f7be3cd94fd \ + --hash=sha256:3f815fb0eba625af45139af4f90f5fc2ddda61b171c2cc3ab63d44b40c5c7768 \ + --hash=sha256:4125d8cd86fa08495d310e80926c2f0563f157b76862e7479f9b2cf94823ea0c \ + --hash=sha256:4206ebdd1d1ef0f3f86c8c2f7c426aa4af6094f4f41e274601fd4c4569f37454 \ + --hash=sha256:440b489902bfb7a26d3fec1ca888007615336ff763d2a32a2fc40586548a0dbf \ + --hash=sha256:45c18cc13156e8697a8d3f9e57e49a69b00e14a103196efab0893fae1a5257f8 \ + --hash=sha256:4826d6d81ea490fa7a3ae7a4b9729866a945ffac1f77fe57b71e49d6e1b21efd \ + --hash=sha256:53145fc9f2319c1245d4329e1da8cfacd6e35e27090c07c0b9d453ae2bbdac3e \ + --hash=sha256:56495975401b1821dbe1f29cf222e23556232209a2fdb809fe8156d120ca9c7f \ + --hash=sha256:57ca8f3775324a9de3ee6f05ca172687ba258c0dea79f7e3a6b4112834982f2a \ + --hash=sha256:5925a738b8478f223ab9756fc794e3cabd5917fd7846f66adcf1d5fc2bf9864c \ + --hash=sha256:5a7797a2fff994fc5e323f7a967a35a3e37e3006ed21d64dcded086502f482af \ + --hash=sha256:67040e0fd84404885ec716a806bee6110f9960c3647e0ef1670aab3b7375a70a \ + --hash=sha256:6a2ca4d3c683d28d3217821029eb08d3487d5043d7eb455df11ff3cacfd4c916 \ + --hash=sha256:6b19fc60ead1cae9795a5b359599da3a1c95d38f869bdfb51c441fd76b04e926 \ + --hash=sha256:6b5cef5cf40725fe64592af9ec163e7389855077700678a1d94bec549403a74d \ + --hash=sha256:6b7de6b61b11545570e4d6033713f3599525efc615ee353a822be8f6b0c65b77 \ + --hash=sha256:6d2df8a6511cc08ef1fccd2e0c65e2ebc9f57574ec8376052a76851af5398810 \ + --hash=sha256:6d86b44933aea0151e4a2e1e6935448499849045c38167d288ca4c59d5b8cd4e \ + --hash=sha256:79417957972553502b217a0093532e48893c8b4ca30ccc941cefe9c72379df7c \ + --hash=sha256:7eb032549dec897b942ddcf80c1cdccbcb40629f15fc902731dbe6362da49326 \ + --hash=sha256:8097ee39b61c86848a443c0b25b2df1de6b331fd512b20836a4f5cfde51ab255 \ + --hash=sha256:84d154fbadece82935396eb6bcb502085d944d2fd13b07a94348364344370c2c \ + --hash=sha256:86824c695688fcd06c5ac9bbd3fea9bdfb4cca194b1e706fbf11a629df48d2b4 \ + --hash=sha256:872b00ff83e84bcbdc7e951af291ebe65eed20b09c47e7c4af21c312f90b796f \ + --hash=sha256:8a9f52c27292c21457f43c4ce124939302a9acfb62295e7cda8667310563a5a3 \ + --hash=sha256:8bc9c2c748aaf91863d89c4583f529c1c709485c94f8dfeb3ee48662d88e3258 \ + --hash=sha256:8d1248dfa7f151e893ce819670f00879e4b7650b8d4c01279ce4f12140d68dd2 \ + --hash=sha256:8d47fd41ce260cf4f0ff0e788de961fab9e9c6844a05ce55d06ce31e06107bdc \ + --hash=sha256:8dc5207567459d049696f62a1fdfb220f3fe6aa0d722285d44753e12504dac6c \ + --hash=sha256:8e0c5d98a4e791f0bbd0ffcb7dae879baeb2dcc357348a8dc2be0a8c10403a2a \ + --hash=sha256:8e82464d1e00fbbb12958999b8471ba5e9f3d9711954505a0a7b378762332e6f \ + --hash=sha256:95f3646ddc98af25af25d5692ae65966488a283813336ea9cf41b22e542e7c0d \ + --hash=sha256:9847dd6f288f1c56359f52acb48ff2df848ff3e3bff34d23855bbcf7016427cc \ + --hash=sha256:9da6d970281083bae91b914362de325414aa03c01fc806f6bb2cc006322ec834 \ + --hash=sha256:9e9193cd4bb57e7acd3af24891526299244bfed88168945efdaa09af4e50720f \ + --hash=sha256:a237064a6e2c2256c9a1cf2beb7c971382190c0f1eb2e810e02e971881756132 \ + --hash=sha256:a36adf7d13b7accfa206e1c917f08924eb905b45aa8e62176509afa7b14db71e \ + --hash=sha256:a47de0a68f5f4d9951250ef5af31f2a7228132caa9ed60994234f7eb98090d33 \ + --hash=sha256:ab1e69dc4831bbb79b6d547077aae89074c83e8ad94eba1a3d80e94d2424fd02 \ + --hash=sha256:ab687bef5c493732b9a4ab870542ee43f5eae0025f9c684c7cb399c3a85cb380 \ + --hash=sha256:ac48b978aa0675f62b642750e798c394a64d25ce852e4e541f69bef9a564c2f0 \ + --hash=sha256:af39006faddfc6253beb93ca821d544931cfee7f0177b99ff106dfd8fd6a2cd8 \ + --hash=sha256:b0944a7c3a78f940c06d1b29bdce91a17798d80593dd01ebfeb842761e48a8b5 \ + --hash=sha256:b3291be0d3f73d5774d69013be4ab33978c777363b5312d14f62f77817c2f75a \ + --hash=sha256:b5b1cd7d39242b2b903cf09cd4696b3a6e04dc537ffa9f3ac8668edae76eecb6 \ + --hash=sha256:b7ac273498a2c6772d67707e101b74014c0d9413bb4711c51d8ec311de59b4b1 \ + --hash=sha256:b9db1debe48060e41a5b91af9193c524e473c57f6105462c5524a41f5aabdb88 \ + --hash=sha256:ba560244bc1335b420b74e91e35f9d4e7f307a3be3a4603ce0f0d7e15a0acdf0 \ + --hash=sha256:c60e5996aa02547d12bc2740d44e90e006b0f93100f53206f7abe6732ad56e69 \ + --hash=sha256:ce2b94117f373defc876f88e74e44049a9969223dbca3240415b71752d0422fb \ + --hash=sha256:cf29b4def86ec503e329fe138842a9b79a997e3beb6c7809b05665a0d291edff \ + --hash=sha256:cf4ea758d98b6fad1b4b2d808d0de690d3162ac56c26968aea0af6524e3eb736 \ + --hash=sha256:d14a0efb21e0fec0631bcd66040b06e6a0fe10825f3aacffded38c1c978bdff9 \ + --hash=sha256:d35923fb5411bde30b53c0696dff8e24c8a38b010b89544834c53f4462fd71df \ + --hash=sha256:d51b9b140b1df39a44bff7896d98a10da345b7d5f5ce92368d328c1c2c829167 \ + --hash=sha256:d90a72608c7550cd7eba914668f6277bfb0b24f074d1f1bd9d061fcb6f2adbd6 \ + --hash=sha256:da0cc0efdbfb8ee2361f89f38ded03d11678f37e392afff7a97b09c55dadfc83 \ + --hash=sha256:dda7698b6d7caeae1047adafebc4b43b2a82478234f6c2b45bc3edad854e0600 \ + --hash=sha256:e076fd87089197cb61117c63dbe7712ad5eccb93968860eb3bae09b767bac813 \ + --hash=sha256:e13c9a697881e5e38148958612dc6856967f5ff8cd7bba5ff751f2d6ac020aa4 \ + --hash=sha256:ec769e5b16251704502277a1163dcf2611551452d7590ff4cc422b7b0367fc96 \ + --hash=sha256:f6f18f0242212d3409d26ce3874937b5b979cebd61f08b633a6ea893c32fc7b6 \ + --hash=sha256:f89924858712b8b936f04f3d690e72825a3e5127a140b434c79030c1c5a887ce \ + --hash=sha256:fb01f6e38719818778144d3165a89ea1ad9dc58c6342b7f20aa194c70f34cbd1 \ + --hash=sha256:fbfe35929a61b914de9e5dbacde0cfbba86cbf5122f9285a24c14ed0b645490b \ + --hash=sha256:fd0fa9a0e7f18224b6d2d1d69dbdc3aecec80ef1393c59244159b131604a4395 \ + --hash=sha256:ff362f68bd68ac0eccb445209238d589bba728fb6d7f2e9dc199e0ec3a61d6e0 # via # barman # python-snappy -cryptography==43.0.3 \ - --hash=sha256:0c580952eef9bf68c4747774cde7ec1d85a6e61de97281f2dba83c7d2c806362 \ - --hash=sha256:0f996e7268af62598f2fc1204afa98a3b5712313a55c4c9d434aef49cadc91d4 \ - --hash=sha256:1ec0bcf7e17c0c5669d881b1cd38c4972fade441b27bda1051665faaa89bdcaa \ - --hash=sha256:281c945d0e28c92ca5e5930664c1cefd85efe80e5c0d2bc58dd63383fda29f83 \ - --hash=sha256:2ce6fae5bdad59577b44e4dfed356944fbf1d925269114c28be377692643b4ff \ - --hash=sha256:315b9001266a492a6ff443b61238f956b214dbec9910a081ba5b6646a055a805 \ - --hash=sha256:443c4a81bb10daed9a8f334365fe52542771f25aedaf889fd323a853ce7377d6 \ - --hash=sha256:4a02ded6cd4f0a5562a8887df8b3bd14e822a90f97ac5e544c162899bc467664 \ - --hash=sha256:53a583b6637ab4c4e3591a15bc9db855b8d9dee9a669b550f311480acab6eb08 \ - --hash=sha256:63efa177ff54aec6e1c0aefaa1a241232dcd37413835a9b674b6e3f0ae2bfd3e \ - --hash=sha256:74f57f24754fe349223792466a709f8e0c093205ff0dca557af51072ff47ab18 \ - --hash=sha256:7e1ce50266f4f70bf41a2c6dc4358afadae90e2a1e5342d3c08883df1675374f \ - --hash=sha256:81ef806b1fef6b06dcebad789f988d3b37ccaee225695cf3e07648eee0fc6b73 \ - --hash=sha256:846da004a5804145a5f441b8530b4bf35afbf7da70f82409f151695b127213d5 \ - --hash=sha256:8ac43ae87929a5982f5948ceda07001ee5e83227fd69cf55b109144938d96984 \ - --hash=sha256:9762ea51a8fc2a88b70cf2995e5675b38d93bf36bd67d91721c309df184f49bd \ - --hash=sha256:a2a431ee15799d6db9fe80c82b055bae5a752bef645bba795e8e52687c69efe3 \ - --hash=sha256:bf7a1932ac4176486eab36a19ed4c0492da5d97123f1406cf15e41b05e787d2e \ - --hash=sha256:c2e6fc39c4ab499049df3bdf567f768a723a5e8464816e8f009f121a5a9f4405 \ - --hash=sha256:cbeb489927bd7af4aa98d4b261af9a5bc025bd87f0e3547e11584be9e9427be2 \ - --hash=sha256:d03b5621a135bffecad2c73e9f4deb1a0f977b9a8ffe6f8e002bf6c9d07b918c \ - --hash=sha256:d56e96520b1020449bbace2b78b603442e7e378a9b3bd68de65c782db1507995 \ - --hash=sha256:df6b6c6d742395dd77a23ea3728ab62f98379eff8fb61be2744d4679ab678f73 \ - --hash=sha256:e1be4655c7ef6e1bbe6b5d0403526601323420bcf414598955968c9ef3eb7d16 \ - --hash=sha256:f18c716be16bc1fea8e95def49edf46b82fccaa88587a45f8dc0ff6ab5d8e0a7 \ - --hash=sha256:f46304d6f0c6ab8e52770addfa2fc41e6629495548862279641972b6215451cd \ - --hash=sha256:f7b178f11ed3664fd0e995a47ed2b5ff0a12d893e41dd0494f406d1cf555cab7 +cryptography==44.0.2 \ + --hash=sha256:04abd71114848aa25edb28e225ab5f268096f44cf0127f3d36975bdf1bdf3390 \ + --hash=sha256:0529b1d5a0105dd3731fa65680b45ce49da4d8115ea76e9da77a875396727b41 \ + --hash=sha256:1bc312dfb7a6e5d66082c87c34c8a62176e684b6fe3d90fcfe1568de675e6688 \ + --hash=sha256:268e4e9b177c76d569e8a145a6939eca9a5fec658c932348598818acf31ae9a5 \ + --hash=sha256:29ecec49f3ba3f3849362854b7253a9f59799e3763b0c9d0826259a88efa02f1 \ + --hash=sha256:2bf7bf75f7df9715f810d1b038870309342bff3069c5bd8c6b96128cb158668d \ + --hash=sha256:3b721b8b4d948b218c88cb8c45a01793483821e709afe5f622861fc6182b20a7 \ + --hash=sha256:3c00b6b757b32ce0f62c574b78b939afab9eecaf597c4d624caca4f9e71e7843 \ + --hash=sha256:3dc62975e31617badc19a906481deacdeb80b4bb454394b4098e3f2525a488c5 \ + --hash=sha256:4973da6ca3db4405c54cd0b26d328be54c7747e89e284fcff166132eb7bccc9c \ + --hash=sha256:4e389622b6927d8133f314949a9812972711a111d577a5d1f4bee5e58736b80a \ + --hash=sha256:51e4de3af4ec3899d6d178a8c005226491c27c4ba84101bfb59c901e10ca9f79 \ + --hash=sha256:5f6f90b72d8ccadb9c6e311c775c8305381db88374c65fa1a68250aa8a9cb3a6 \ + --hash=sha256:6210c05941994290f3f7f175a4a57dbbb2afd9273657614c506d5976db061181 \ + --hash=sha256:6f101b1f780f7fc613d040ca4bdf835c6ef3b00e9bd7125a4255ec574c7916e4 \ + --hash=sha256:7bdcd82189759aba3816d1f729ce42ffded1ac304c151d0a8e89b9996ab863d5 \ + --hash=sha256:7ca25849404be2f8e4b3c59483d9d3c51298a22c1c61a0e84415104dacaf5562 \ + --hash=sha256:81276f0ea79a208d961c433a947029e1a15948966658cf6710bbabb60fcc2639 \ + --hash=sha256:8cadc6e3b5a1f144a039ea08a0bdb03a2a92e19c46be3285123d32029f40a922 \ + --hash=sha256:8e0ddd63e6bf1161800592c71ac794d3fb8001f2caebe0966e77c5234fa9efc3 \ + --hash=sha256:909c97ab43a9c0c0b0ada7a1281430e4e5ec0458e6d9244c0e821bbf152f061d \ + --hash=sha256:96e7a5e9d6e71f9f4fca8eebfd603f8e86c5225bb18eb621b2c1e50b290a9471 \ + --hash=sha256:9a1e657c0f4ea2a23304ee3f964db058c9e9e635cc7019c4aa21c330755ef6fd \ + --hash=sha256:9eb9d22b0a5d8fd9925a7764a054dca914000607dff201a24c791ff5c799e1fa \ + --hash=sha256:af4ff3e388f2fa7bff9f7f2b31b87d5651c45731d3e8cfa0944be43dff5cfbdb \ + --hash=sha256:b042d2a275c8cee83a4b7ae30c45a15e6a4baa65a179a0ec2d78ebb90e4f6699 \ + --hash=sha256:bc821e161ae88bfe8088d11bb39caf2916562e0a2dc7b6d56714a48b784ef0bb \ + --hash=sha256:c505d61b6176aaf982c5717ce04e87da5abc9a36a5b39ac03905c4aafe8de7aa \ + --hash=sha256:c63454aa261a0cf0c5b4718349629793e9e634993538db841165b3df74f37ec0 \ + --hash=sha256:c7362add18b416b69d58c910caa217f980c5ef39b23a38a0880dfd87bdf8cd23 \ + --hash=sha256:d03806036b4f89e3b13b6218fefea8d5312e450935b1a2d55f0524e2ed7c59d9 \ + --hash=sha256:d1b3031093a366ac767b3feb8bcddb596671b3aaff82d4050f984da0c248b615 \ + --hash=sha256:d1c3572526997b36f245a96a2b1713bf79ce99b271bbcf084beb6b9b075f29ea \ + --hash=sha256:efcfe97d1b3c79e486554efddeb8f6f53a4cdd4cf6086642784fa31fc384e1d7 \ + --hash=sha256:f514ef4cd14bb6fb484b4a60203e912cfcb64f2ab139e88c2274511514bf7308 # via # azure-identity # azure-storage-blob # msal # pyjwt -google-api-core==2.22.0 \ - --hash=sha256:26f8d76b96477db42b55fd02a33aae4a42ec8b86b98b94969b7333a2c828bf35 \ - --hash=sha256:a6652b6bd51303902494998626653671703c420f6f4c88cfd3f50ed723e9d021 +google-api-core==2.24.2 \ + --hash=sha256:810a63ac95f3c441b7c0e43d344e372887f62ce9071ba972eacf32672e072de9 \ + --hash=sha256:81718493daf06d96d6bc76a91c23874dbf2fac0adbbf542831b805ee6e974696 # via # google-cloud-core # google-cloud-storage -google-auth==2.35.0 \ - --hash=sha256:25df55f327ef021de8be50bad0dfd4a916ad0de96da86cd05661c9297723ad3f \ - --hash=sha256:f4c64ed4e01e8e8b646ef34c018f8bf3338df0c8e37d8b3bba40e7f574a3278a +google-auth==2.38.0 \ + --hash=sha256:8285113607d3b80a3f1543b75962447ba8a09fe85783432a784fdeef6ac094c4 \ + --hash=sha256:e7dae6694313f434a2727bf2906f27ad259bae090d7aa896590d86feec3d9d4a # via # google-api-core # google-cloud-core # google-cloud-storage -google-cloud-core==2.4.1 \ - --hash=sha256:9b7749272a812bde58fff28868d0c5e2f585b82f37e09a1f6ed2d4d10f134073 \ - --hash=sha256:a9e6a4422b9ac5c29f79a0ede9485473338e2ce78d91f2370c01e730eab22e61 +google-cloud-core==2.4.3 \ + --hash=sha256:1fab62d7102844b278fe6dead3af32408b1df3eb06f5c7e8634cbd40edc4da53 \ + --hash=sha256:5130f9f4c14b4fafdff75c79448f9495cfade0d8775facf1b09c3bf67e027f6e # via google-cloud-storage -google-cloud-storage==2.18.2 \ - --hash=sha256:97a4d45c368b7d401ed48c4fdfe86e1e1cb96401c9e199e419d289e2c0370166 \ - --hash=sha256:aaf7acd70cdad9f274d29332673fcab98708d0e1f4dceb5a5356aaef06af4d99 +google-cloud-storage==3.1.0 \ + --hash=sha256:944273179897c7c8a07ee15f2e6466a02da0c7c4b9ecceac2a26017cb2972049 \ + --hash=sha256:eaf36966b68660a9633f03b067e4a10ce09f1377cae3ff9f2c699f69a81c66c6 google-crc32c==1.6.0 \ --hash=sha256:05e2d8c9a2f853ff116db9706b4a27350587f341eda835f46db3c0a8c8ce2f24 \ --hash=sha256:18e311c64008f1f1379158158bb3f0c8d72635b9eb4f9545f8cf990c5668e59d \ @@ -395,9 +393,9 @@ google-resumable-media==2.7.2 \ --hash=sha256:3ce7551e9fe6d99e9a126101d2536612bb73486721951e9562fee0f90c6ababa \ --hash=sha256:5280aed4629f2b60b847b0d42f9857fd4935c11af266744df33d8074cae92fe0 # via google-cloud-storage -googleapis-common-protos==1.65.0 \ - --hash=sha256:2972e6c496f435b92590fd54045060867f3fe9be2c82ab148fc8885035479a63 \ - --hash=sha256:334a29d07cddc3aa01dee4988f9afd9b2916ee2ff49d6b757155dc0d197852c0 +googleapis-common-protos==1.69.1 \ + --hash=sha256:4077f27a6900d5946ee5a369fab9c8ded4c0ef1c6e880458ea2f70c14f7b70d5 \ + --hash=sha256:e20d2d8dda87da6fe7340afbbdf4f0bcb4c8fae7e6cadf55926c31f946b0b9b1 # via google-api-core idna==3.10 \ --hash=sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9 \ @@ -413,36 +411,64 @@ jmespath==1.0.1 \ # via # boto3 # botocore -msal==1.31.0 \ - --hash=sha256:2c4f189cf9cc8f00c80045f66d39b7c0f3ed45873fd3d1f2af9f22db2e12ff4b \ - --hash=sha256:96bc37cff82ebe4b160d5fc0f1196f6ca8b50e274ecd0ec5bf69c438514086e7 +lz4==4.4.3 \ + --hash=sha256:0aea6f283abd6acb1883b70d7a117b913e20c770845559f9421394bc9c522b24 \ + --hash=sha256:174b7ce5456671c73b81bb115defac8a584363d8b38a48ed3ad976e08eea27cd \ + --hash=sha256:1ebf23ffd36b32b980f720a81990fcfdeadacafe7498fbeff7a8e058259d4e58 \ + --hash=sha256:1f25e1b571a8be2c3d60d46679ef2471ae565f7ba9ba8382596695413523b188 \ + --hash=sha256:20e385cb8bd8321593788f11101d8c89a823a56191978e427e3c5141e129f14b \ + --hash=sha256:2ae50a175fb7b900f7aa42575f4fe99c32ca0ff57e5a8c1fd25e1243e67409db \ + --hash=sha256:2b45914f25d916324531d0259072b402c5f99b67c6e9ac8cbc3d49935aeb1d97 \ + --hash=sha256:38df5929ffefa9dda120ba1790a2e94fda81916c5aaa1ee652f4b1e515ebb9ed \ + --hash=sha256:3f21e503c18157512d2e34ae4c301e44a826c7b87e1d8998981367e3c9fe0932 \ + --hash=sha256:43461e439ef71d49bb0ee3a1719494cd952a58d205496698e0cde866f22006bc \ + --hash=sha256:434a1d1547a0547164866f1ccc31bbda235ac5b9087f24a84956756b52371f40 \ + --hash=sha256:447993c4dda0b6b0e1bd862752c855df8745f2910bea5015344f83ff3e99f305 \ + --hash=sha256:5c9e32989df06c57f10aa09ad9b30e8a25baf1aefe850e13b0ea5de600477d6a \ + --hash=sha256:61e08d84e3bf8ca9f43dc6b33f8cd7ba19f49864e2c91eb2160f83b6f9a268fa \ + --hash=sha256:699d26ac579eb42c71d131f9fb7b6e1c495a14e257264206a3c3bfcc146ed9bb \ + --hash=sha256:71ebdaadf546d6d393b9a21796172723724b737e84f68f36caf367d1c87a86a1 \ + --hash=sha256:7f5c05bd4b0909b682608c453acc31f1a9170d55f56d27cd701213e0683fc66a \ + --hash=sha256:848c5b040d2cfe35097b1d65d1095d83a3f86374ce879e189533f61405d8763b \ + --hash=sha256:8fe3caea61427057a9e3697c69b2403510fdccfca4483520d02b98ffae74531e \ + --hash=sha256:91ed5b71f9179bf3dbfe85d92b52d4b53de2e559aa4daa3b7de18e0dd24ad77d \ + --hash=sha256:a46f48740584eab3194fbee91c61f7fa396dbb1c5e7aa76ca08165d4e63fb40f \ + --hash=sha256:ab26b4af13308b8296688b03d74c3b0c8e8ed1f6b2d1454ef97bdb589db409db \ + --hash=sha256:b1b98f0a4137d01b84c680813eef6198e1e00f1f28bc20ce7b5c436459a0d146 \ + --hash=sha256:b1d179bdefd9ddb8d11d7de7825e73fb957511b722a8cb484e417885c210e68c \ + --hash=sha256:c3d2d5df5476b065aae9d1ad551fdc7b17c151b84e8edd9212108946b2337c66 \ + --hash=sha256:c4be1e5d9c8ad61345730c41c9ef21bdbb022cced4df70431110888d3ad5c0fb \ + --hash=sha256:da091dd8c96dbda124d766231f38619afd5c544051fb4424d2566c905957d342 \ + --hash=sha256:de86400c8b60c7707665e63934a82ae6792e7102c17a72e9b361a7f40d3c6049 \ + --hash=sha256:e365850166729fa82be618f476966161d5c47ea081eafc4febfc542bc85bac5d \ + --hash=sha256:e86c7fbe46f6e2e9dfb5377ee690fb8987e8e8363f435886ab91012b88f08a26 \ + --hash=sha256:fe6080299a25fd7cbb1957c921cca6a884acbfcd44cc23de48079389d322e326 +msal==1.32.0 \ + --hash=sha256:5445fe3af1da6be484991a7ab32eaa82461dc2347de105b76af92c610c3335c2 \ + --hash=sha256:9dbac5384a10bbbf4dae5c7ea0d707d14e087b92c5aa4954b3feaa2d1aa0bcb7 # via # azure-identity # msal-extensions -msal-extensions==1.2.0 \ - --hash=sha256:6f41b320bfd2933d631a215c91ca0dd3e67d84bd1a2f50ce917d5874ec646bef \ - --hash=sha256:cf5ba83a2113fa6dc011a254a72f1c223c88d7dfad74cc30617c4679a417704d +msal-extensions==1.3.1 \ + --hash=sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca \ + --hash=sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4 # via azure-identity -portalocker==2.10.1 \ - --hash=sha256:53a5984ebc86a025552264b459b46a2086e269b21823cb572f8f28ee759e45bf \ - --hash=sha256:ef1bf844e878ab08aee7e40184156e1151f228f103aa5c6bd0724cc330960f8f - # via msal-extensions -proto-plus==1.25.0 \ - --hash=sha256:c91fc4a65074ade8e458e95ef8bac34d4008daa7cce4a12d6707066fca648961 \ - --hash=sha256:fbb17f57f7bd05a68b7707e745e26528b0b3c34e378db91eef93912c54982d91 +proto-plus==1.26.1 \ + --hash=sha256:13285478c2dcf2abb829db158e1047e2f1e8d63a077d94263c2b88b043c75a66 \ + --hash=sha256:21a515a4c4c0088a773899e23c7bbade3d18f9c66c73edd4c7ee3816bc96a012 # via google-api-core -protobuf==5.28.3 \ - --hash=sha256:0c4eec6f987338617072592b97943fdbe30d019c56126493111cf24344c1cc24 \ - --hash=sha256:135658402f71bbd49500322c0f736145731b16fc79dc8f367ab544a17eab4535 \ - --hash=sha256:27b246b3723692bf1068d5734ddaf2fccc2cdd6e0c9b47fe099244d80200593b \ - --hash=sha256:3e6101d095dfd119513cde7259aa703d16c6bbdfae2554dfe5cfdbe94e32d548 \ - --hash=sha256:3fa2de6b8b29d12c61911505d893afe7320ce7ccba4df913e2971461fa36d584 \ - --hash=sha256:64badbc49180a5e401f373f9ce7ab1d18b63f7dd4a9cdc43c92b9f0b481cef7b \ - --hash=sha256:70585a70fc2dd4818c51287ceef5bdba6387f88a578c86d47bb34669b5552c36 \ - --hash=sha256:712319fbdddb46f21abb66cd33cb9e491a5763b2febd8f228251add221981135 \ - --hash=sha256:91fba8f445723fcf400fdbe9ca796b19d3b1242cd873907979b9ed71e4afe868 \ - --hash=sha256:a3f6857551e53ce35e60b403b8a27b0295f7d6eb63d10484f12bc6879c715687 \ - --hash=sha256:cee1757663fa32a1ee673434fcf3bf24dd54763c79690201208bafec62f19eed +protobuf==5.29.3 \ + --hash=sha256:0a18ed4a24198528f2333802eb075e59dea9d679ab7a6c5efb017a59004d849f \ + --hash=sha256:0eb32bfa5219fc8d4111803e9a690658aa2e6366384fd0851064b963b6d1f2a7 \ + --hash=sha256:3ea51771449e1035f26069c4c7fd51fba990d07bc55ba80701c78f886bf9c888 \ + --hash=sha256:5da0f41edaf117bde316404bad1a486cb4ededf8e4a54891296f648e8e076620 \ + --hash=sha256:6ce8cc3389a20693bfde6c6562e03474c40851b44975c9b2bf6df7d8c4f864da \ + --hash=sha256:84a57163a0ccef3f96e4b6a20516cedcf5bb3a95a657131c5c3ac62200d23252 \ + --hash=sha256:a4fa6f80816a9a0678429e84973f2f98cbc218cca434abe8db2ad0bffc98503a \ + --hash=sha256:a8434404bbf139aa9e1300dbf989667a83d42ddda9153d8ab76e0d5dcaca484e \ + --hash=sha256:b89c115d877892a512f79a8114564fb435943b59067615894c3b13cd3e1fa107 \ + --hash=sha256:c027e08a08be10b67c06bf2370b99c811c466398c357e615ca88c91c07f0910f \ + --hash=sha256:daaf63f70f25e8689c072cfad4334ca0ac1d1e05a92fc15c54eb9cf23c3efd84 # via # google-api-core # googleapis-common-protos @@ -461,9 +487,9 @@ pycparser==2.22 \ --hash=sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6 \ --hash=sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc # via cffi -pyjwt[crypto]==2.9.0 \ - --hash=sha256:3b02fb0f44517787776cf48f2ae25d8e14f300e6d7545a4315cee571a415e850 \ - --hash=sha256:7e1e5b56cc735432a7369cbfa0efe50fa113ebecdc04ae6922deba8b84582d0c +pyjwt[crypto]==2.10.1 \ + --hash=sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953 \ + --hash=sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb # via # msal # pyjwt @@ -488,13 +514,13 @@ rsa==4.9 \ --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \ --hash=sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21 # via google-auth -s3transfer==0.10.3 \ - --hash=sha256:263ed587a5803c6c708d3ce44dc4dfedaab4c1a32e8329bab818933d79ddcf5d \ - --hash=sha256:4f50ed74ab84d474ce614475e0b8d5047ff080810aac5d01ea25231cfc944b0c +s3transfer==0.10.4 \ + --hash=sha256:244a76a24355363a68164241438de1b72f8781664920260c48465896b712a41e \ + --hash=sha256:29edc09801743c21eb5ecbc617a152df41d3c287f67b615f73e5f750583666a7 # via boto3 -six==1.16.0 \ - --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \ - --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254 +six==1.17.0 \ + --hash=sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274 \ + --hash=sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81 # via # azure-core # python-dateutil @@ -505,9 +531,107 @@ typing-extensions==4.12.2 \ # azure-core # azure-identity # azure-storage-blob -urllib3==2.2.3 \ - --hash=sha256:ca899ca043dcb1bafa3e262d73aa25c465bfb49e0bd9dd5d59f1d0acba2f8fac \ - --hash=sha256:e7d814a81dad81e6caf2ec9fdedb284ecc9c73076b62654547cc64ccdcae26e9 +urllib3==2.3.0 \ + --hash=sha256:1cee9ad369867bfdbbb48b7dd50374c0967a0bb7710050facf0dd6911440e3df \ + --hash=sha256:f8c5449b3cf0861679ce7e0503c7b44b5ec981bec0d1d3795a07f1ba96f0204d # via # botocore # requests +zstandard==0.23.0 \ + --hash=sha256:034b88913ecc1b097f528e42b539453fa82c3557e414b3de9d5632c80439a473 \ + --hash=sha256:0a7f0804bb3799414af278e9ad51be25edf67f78f916e08afdb983e74161b916 \ + --hash=sha256:11e3bf3c924853a2d5835b24f03eeba7fc9b07d8ca499e247e06ff5676461a15 \ + --hash=sha256:12a289832e520c6bd4dcaad68e944b86da3bad0d339ef7989fb7e88f92e96072 \ + --hash=sha256:1516c8c37d3a053b01c1c15b182f3b5f5eef19ced9b930b684a73bad121addf4 \ + --hash=sha256:157e89ceb4054029a289fb504c98c6a9fe8010f1680de0201b3eb5dc20aa6d9e \ + --hash=sha256:1bfe8de1da6d104f15a60d4a8a768288f66aa953bbe00d027398b93fb9680b26 \ + --hash=sha256:1e172f57cd78c20f13a3415cc8dfe24bf388614324d25539146594c16d78fcc8 \ + --hash=sha256:1fd7e0f1cfb70eb2f95a19b472ee7ad6d9a0a992ec0ae53286870c104ca939e5 \ + --hash=sha256:203d236f4c94cd8379d1ea61db2fce20730b4c38d7f1c34506a31b34edc87bdd \ + --hash=sha256:27d3ef2252d2e62476389ca8f9b0cf2bbafb082a3b6bfe9d90cbcbb5529ecf7c \ + --hash=sha256:29a2bc7c1b09b0af938b7a8343174b987ae021705acabcbae560166567f5a8db \ + --hash=sha256:2ef230a8fd217a2015bc91b74f6b3b7d6522ba48be29ad4ea0ca3a3775bf7dd5 \ + --hash=sha256:2ef3775758346d9ac6214123887d25c7061c92afe1f2b354f9388e9e4d48acfc \ + --hash=sha256:2f146f50723defec2975fb7e388ae3a024eb7151542d1599527ec2aa9cacb152 \ + --hash=sha256:2fb4535137de7e244c230e24f9d1ec194f61721c86ebea04e1581d9d06ea1269 \ + --hash=sha256:32ba3b5ccde2d581b1e6aa952c836a6291e8435d788f656fe5976445865ae045 \ + --hash=sha256:34895a41273ad33347b2fc70e1bff4240556de3c46c6ea430a7ed91f9042aa4e \ + --hash=sha256:379b378ae694ba78cef921581ebd420c938936a153ded602c4fea612b7eaa90d \ + --hash=sha256:38302b78a850ff82656beaddeb0bb989a0322a8bbb1bf1ab10c17506681d772a \ + --hash=sha256:3aa014d55c3af933c1315eb4bb06dd0459661cc0b15cd61077afa6489bec63bb \ + --hash=sha256:4051e406288b8cdbb993798b9a45c59a4896b6ecee2f875424ec10276a895740 \ + --hash=sha256:40b33d93c6eddf02d2c19f5773196068d875c41ca25730e8288e9b672897c105 \ + --hash=sha256:43da0f0092281bf501f9c5f6f3b4c975a8a0ea82de49ba3f7100e64d422a1274 \ + --hash=sha256:445e4cb5048b04e90ce96a79b4b63140e3f4ab5f662321975679b5f6360b90e2 \ + --hash=sha256:48ef6a43b1846f6025dde6ed9fee0c24e1149c1c25f7fb0a0585572b2f3adc58 \ + --hash=sha256:50a80baba0285386f97ea36239855f6020ce452456605f262b2d33ac35c7770b \ + --hash=sha256:519fbf169dfac1222a76ba8861ef4ac7f0530c35dd79ba5727014613f91613d4 \ + --hash=sha256:53dd9d5e3d29f95acd5de6802e909ada8d8d8cfa37a3ac64836f3bc4bc5512db \ + --hash=sha256:53ea7cdc96c6eb56e76bb06894bcfb5dfa93b7adcf59d61c6b92674e24e2dd5e \ + --hash=sha256:576856e8594e6649aee06ddbfc738fec6a834f7c85bf7cadd1c53d4a58186ef9 \ + --hash=sha256:59556bf80a7094d0cfb9f5e50bb2db27fefb75d5138bb16fb052b61b0e0eeeb0 \ + --hash=sha256:5d41d5e025f1e0bccae4928981e71b2334c60f580bdc8345f824e7c0a4c2a813 \ + --hash=sha256:61062387ad820c654b6a6b5f0b94484fa19515e0c5116faf29f41a6bc91ded6e \ + --hash=sha256:61f89436cbfede4bc4e91b4397eaa3e2108ebe96d05e93d6ccc95ab5714be512 \ + --hash=sha256:62136da96a973bd2557f06ddd4e8e807f9e13cbb0bfb9cc06cfe6d98ea90dfe0 \ + --hash=sha256:64585e1dba664dc67c7cdabd56c1e5685233fbb1fc1966cfba2a340ec0dfff7b \ + --hash=sha256:65308f4b4890aa12d9b6ad9f2844b7ee42c7f7a4fd3390425b242ffc57498f48 \ + --hash=sha256:66b689c107857eceabf2cf3d3fc699c3c0fe8ccd18df2219d978c0283e4c508a \ + --hash=sha256:6a41c120c3dbc0d81a8e8adc73312d668cd34acd7725f036992b1b72d22c1772 \ + --hash=sha256:6f77fa49079891a4aab203d0b1744acc85577ed16d767b52fc089d83faf8d8ed \ + --hash=sha256:72c68dda124a1a138340fb62fa21b9bf4848437d9ca60bd35db36f2d3345f373 \ + --hash=sha256:752bf8a74412b9892f4e5b58f2f890a039f57037f52c89a740757ebd807f33ea \ + --hash=sha256:76e79bc28a65f467e0409098fa2c4376931fd3207fbeb6b956c7c476d53746dd \ + --hash=sha256:774d45b1fac1461f48698a9d4b5fa19a69d47ece02fa469825b442263f04021f \ + --hash=sha256:77da4c6bfa20dd5ea25cbf12c76f181a8e8cd7ea231c673828d0386b1740b8dc \ + --hash=sha256:77ea385f7dd5b5676d7fd943292ffa18fbf5c72ba98f7d09fc1fb9e819b34c23 \ + --hash=sha256:80080816b4f52a9d886e67f1f96912891074903238fe54f2de8b786f86baded2 \ + --hash=sha256:80a539906390591dd39ebb8d773771dc4db82ace6372c4d41e2d293f8e32b8db \ + --hash=sha256:82d17e94d735c99621bf8ebf9995f870a6b3e6d14543b99e201ae046dfe7de70 \ + --hash=sha256:837bb6764be6919963ef41235fd56a6486b132ea64afe5fafb4cb279ac44f259 \ + --hash=sha256:84433dddea68571a6d6bd4fbf8ff398236031149116a7fff6f777ff95cad3df9 \ + --hash=sha256:8c24f21fa2af4bb9f2c492a86fe0c34e6d2c63812a839590edaf177b7398f700 \ + --hash=sha256:8ed7d27cb56b3e058d3cf684d7200703bcae623e1dcc06ed1e18ecda39fee003 \ + --hash=sha256:9206649ec587e6b02bd124fb7799b86cddec350f6f6c14bc82a2b70183e708ba \ + --hash=sha256:983b6efd649723474f29ed42e1467f90a35a74793437d0bc64a5bf482bedfa0a \ + --hash=sha256:98da17ce9cbf3bfe4617e836d561e433f871129e3a7ac16d6ef4c680f13a839c \ + --hash=sha256:9c236e635582742fee16603042553d276cca506e824fa2e6489db04039521e90 \ + --hash=sha256:9da6bc32faac9a293ddfdcb9108d4b20416219461e4ec64dfea8383cac186690 \ + --hash=sha256:a05e6d6218461eb1b4771d973728f0133b2a4613a6779995df557f70794fd60f \ + --hash=sha256:a0817825b900fcd43ac5d05b8b3079937073d2b1ff9cf89427590718b70dd840 \ + --hash=sha256:a4ae99c57668ca1e78597d8b06d5af837f377f340f4cce993b551b2d7731778d \ + --hash=sha256:a8c86881813a78a6f4508ef9daf9d4995b8ac2d147dcb1a450448941398091c9 \ + --hash=sha256:a8fffdbd9d1408006baaf02f1068d7dd1f016c6bcb7538682622c556e7b68e35 \ + --hash=sha256:a9b07268d0c3ca5c170a385a0ab9fb7fdd9f5fd866be004c4ea39e44edce47dd \ + --hash=sha256:ab19a2d91963ed9e42b4e8d77cd847ae8381576585bad79dbd0a8837a9f6620a \ + --hash=sha256:ac184f87ff521f4840e6ea0b10c0ec90c6b1dcd0bad2f1e4a9a1b4fa177982ea \ + --hash=sha256:b0e166f698c5a3e914947388c162be2583e0c638a4703fc6a543e23a88dea3c1 \ + --hash=sha256:b2170c7e0367dde86a2647ed5b6f57394ea7f53545746104c6b09fc1f4223573 \ + --hash=sha256:b2d8c62d08e7255f68f7a740bae85b3c9b8e5466baa9cbf7f57f1cde0ac6bc09 \ + --hash=sha256:b4567955a6bc1b20e9c31612e615af6b53733491aeaa19a6b3b37f3b65477094 \ + --hash=sha256:b69bb4f51daf461b15e7b3db033160937d3ff88303a7bc808c67bbc1eaf98c78 \ + --hash=sha256:b8c0bd73aeac689beacd4e7667d48c299f61b959475cdbb91e7d3d88d27c56b9 \ + --hash=sha256:be9b5b8659dff1f913039c2feee1aca499cfbc19e98fa12bc85e037c17ec6ca5 \ + --hash=sha256:bf0a05b6059c0528477fba9054d09179beb63744355cab9f38059548fedd46a9 \ + --hash=sha256:c16842b846a8d2a145223f520b7e18b57c8f476924bda92aeee3a88d11cfc391 \ + --hash=sha256:c363b53e257246a954ebc7c488304b5592b9c53fbe74d03bc1c64dda153fb847 \ + --hash=sha256:c7c517d74bea1a6afd39aa612fa025e6b8011982a0897768a2f7c8ab4ebb78a2 \ + --hash=sha256:d20fd853fbb5807c8e84c136c278827b6167ded66c72ec6f9a14b863d809211c \ + --hash=sha256:d2240ddc86b74966c34554c49d00eaafa8200a18d3a5b6ffbf7da63b11d74ee2 \ + --hash=sha256:d477ed829077cd945b01fc3115edd132c47e6540ddcd96ca169facff28173057 \ + --hash=sha256:d50d31bfedd53a928fed6707b15a8dbeef011bb6366297cc435accc888b27c20 \ + --hash=sha256:dc1d33abb8a0d754ea4763bad944fd965d3d95b5baef6b121c0c9013eaf1907d \ + --hash=sha256:dc5d1a49d3f8262be192589a4b72f0d03b72dcf46c51ad5852a4fdc67be7b9e4 \ + --hash=sha256:e2d1a054f8f0a191004675755448d12be47fa9bebbcffa3cdf01db19f2d30a54 \ + --hash=sha256:e7792606d606c8df5277c32ccb58f29b9b8603bf83b48639b7aedf6df4fe8171 \ + --hash=sha256:ed1708dbf4d2e3a1c5c69110ba2b4eb6678262028afd6c6fbcc5a8dac9cda68e \ + --hash=sha256:f2d4380bf5f62daabd7b751ea2339c1a21d1c9463f1feb7fc2bdcea2c29c3160 \ + --hash=sha256:f3513916e8c645d0610815c257cbfd3242adfd5c4cfa78be514e5a3ebb42a41b \ + --hash=sha256:f8346bfa098532bc1fb6c7ef06783e969d87a99dd1d2a5a18a892c1d7a643c58 \ + --hash=sha256:f83fa6cae3fff8e98691248c9320356971b59678a17f20656a9e59cd32cee6d8 \ + --hash=sha256:fa6ce8b52c5987b3e34d5674b0ab529a4602b632ebab0a93b07bfb4dfc8f8a33 \ + --hash=sha256:fb2b1ecfef1e67897d336de3a0e3f52478182d6a47eda86cbd42504c5cbd009a \ + --hash=sha256:fc9ca1c9718cb3b06634c7c8dec57d24e9438b2aa9a0f02b8bb36bf478538880 \ + --hash=sha256:fd30d9c67d13d891f2360b2a120186729c111238ac63b43dbd37a5a40670b8ca \ + --hash=sha256:fd7699e8fd9969f455ef2926221e0233f81a2542921471382e77a9e2f2b57f4b \ + --hash=sha256:fe3b385d996ee0822fd46528d9f0443b880d4d05528fd26a9119a54ec3f91c69 \ No newline at end of file From 5a6756231a63d537497af61fa742842defe773cd Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 18 Mar 2025 20:31:15 +0800 Subject: [PATCH 129/324] ci: fix release docker build (#215) Signed-off-by: xieydd Co-authored-by: xieydd --- docker/pg-cnpg/Dockerfile.arm | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docker/pg-cnpg/Dockerfile.arm b/docker/pg-cnpg/Dockerfile.arm index dc76570e..6529b13d 100644 --- a/docker/pg-cnpg/Dockerfile.arm +++ b/docker/pg-cnpg/Dockerfile.arm @@ -27,9 +27,11 @@ RUN set -xe; \ RUN set -xe; \ apt-get update; \ apt-get install -y --no-install-recommends \ + python3-dev \ python3-pip \ python3-psycopg2 \ python3-setuptools \ + build-essential \ ; \ pip3 install --break-system-packages --upgrade pip; \ # TODO: Remove --no-deps once https://github.com/pypa/pip/issues/9644 is solved From ad8bdacca085539620eb625c75c880a9d48cd860 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 18 Mar 2025 21:01:35 +0800 Subject: [PATCH 130/324] feat: check quals to skip rerank if `rerank_in_table` is enabled (#206) Signed-off-by: usamoi --- src/index/am/mod.rs | 66 +++++++++++++++++------ src/index/hook.rs | 95 +++++++++++++++++++++++++++++++++ src/index/mod.rs | 2 + tests/logic/rerank_in_table.slt | 22 ++++++++ 4 files changed, 169 insertions(+), 16 deletions(-) create mode 100644 src/index/hook.rs diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index bc38dc36..41a21e94 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -9,6 +9,7 @@ use pgrx::pg_sys::{Datum, ItemPointerData}; use std::cell::LazyCell; use std::ffi::CStr; use std::num::NonZeroU64; +use std::ptr::NonNull; use std::sync::OnceLock; #[repr(C)] @@ -254,7 +255,10 @@ pub unsafe extern "C" fn ambeginscan( use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index_relation, n_keys, n_orderbys) }; - let scanner: Scanner = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + let scanner: Scanner = Scanner { + hack: None, + scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), + }; unsafe { (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); } @@ -289,15 +293,19 @@ pub unsafe extern "C" fn amrescan( probes: probes(), max_scan_tuples: max_scan_tuples(), }; - let fetcher = LazyCell::new(move || { - HeapFetcher::new( - (*scan).indexRelation, - (*scan).heapRelation, - (*scan).xs_snapshot, - ) - }); let scanner = &mut *(*scan).opaque.cast::(); - *scanner = match opfamily { + let fetcher = { + let hack = scanner.hack; + LazyCell::new(move || { + HeapFetcher::new( + (*scan).indexRelation, + (*scan).heapRelation, + (*scan).xs_snapshot, + hack, + ) + }) + }; + scanner.scanning = match opfamily { Opfamily::VectorL2 | Opfamily::VectorIp | Opfamily::VectorCosine @@ -340,7 +348,7 @@ pub unsafe extern "C" fn amgettuple( pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); } let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; - if let Some((_, ctid, recheck)) = LazyCell::force_mut(scanner).next() { + if let Some((_, ctid, recheck)) = LazyCell::force_mut(&mut scanner.scanning).next() { unsafe { (*scan).xs_heaptid = ctid; (*scan).xs_recheck = recheck; @@ -355,12 +363,15 @@ pub unsafe extern "C" fn amgettuple( #[pgrx::pg_guard] pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { let scanner = unsafe { &mut *(*scan).opaque.cast::() }; - *scanner = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); } type Iter = Box>; -type Scanner = LazyCell Iter>>; +pub struct Scanner { + pub hack: Option>, + scanning: LazyCell Iter>>, +} struct HeapFetcher { index_info: *mut pgrx::pg_sys::IndexInfo, @@ -371,6 +382,7 @@ struct HeapFetcher { slot: *mut pgrx::pg_sys::TupleTableSlot, values: [Datum; 32], is_nulls: [bool; 32], + hack: Option>, } impl HeapFetcher { @@ -378,6 +390,7 @@ impl HeapFetcher { index_relation: pgrx::pg_sys::Relation, heap_relation: pgrx::pg_sys::Relation, snapshot: pgrx::pg_sys::Snapshot, + hack: Option>, ) -> Self { unsafe { let index_info = pgrx::pg_sys::BuildIndexInfo(index_relation); @@ -392,6 +405,7 @@ impl HeapFetcher { slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), values: [Datum::null(); 32], is_nulls: [true; 32], + hack, } } } @@ -400,7 +414,6 @@ impl HeapFetcher { impl Drop for HeapFetcher { fn drop(&mut self) { unsafe { - // free resources for last `fetch` call pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); // free common resources pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); @@ -412,10 +425,7 @@ impl Drop for HeapFetcher { impl SearchFetcher for HeapFetcher { fn fetch(&mut self, mut ctid: ItemPointerData) -> Option<(&[Datum; 32], &[bool; 32])> { unsafe { - // free resources for last `fetch` call - pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); // perform operation - (*self.econtext).ecxt_scantuple = self.slot; let table_am = (*self.heap_relation).rd_tableam; let fetch_row_version = (*table_am) .tuple_fetch_row_version @@ -423,6 +433,30 @@ impl SearchFetcher for HeapFetcher { if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { return None; } + if let Some(hack) = self.hack { + if let Some(qual) = NonNull::new(hack.as_ref().ss.ps.qual) { + use pgrx::datum::FromDatum; + use pgrx::memcxt::PgMemoryContexts; + assert!(qual.as_ref().flags & pgrx::pg_sys::EEO_FLAG_IS_QUAL as u8 != 0); + let evalfunc = qual.as_ref().evalfunc.expect("no evalfunc for qual"); + if let Some(mut econtext) = NonNull::new(hack.as_ref().ss.ps.ps_ExprContext) { + econtext.as_mut().ecxt_scantuple = self.slot; + pgrx::pg_sys::MemoryContextReset(econtext.as_ref().ecxt_per_tuple_memory); + let result = PgMemoryContexts::For(econtext.as_ref().ecxt_per_tuple_memory) + .switch_to(|_| { + let mut is_null = true; + let datum = + evalfunc(qual.as_ptr(), econtext.as_mut(), &mut is_null); + bool::from_datum(datum, is_null) + }); + if result != Some(true) { + return None; + } + } + } + } + (*self.econtext).ecxt_scantuple = self.slot; + pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); pgrx::pg_sys::FormIndexDatum( self.index_info, self.slot, diff --git a/src/index/hook.rs b/src/index/hook.rs new file mode 100644 index 00000000..ec979cb5 --- /dev/null +++ b/src/index/hook.rs @@ -0,0 +1,95 @@ +use std::sync::atomic::AtomicPtr; + +#[pgrx::pg_guard] +unsafe extern "C" fn rewrite_plan_state( + node: *mut pgrx::pg_sys::PlanState, + context: *mut core::ffi::c_void, +) -> bool { + unsafe fn dirty_check(index_relation: *mut pgrx::pg_sys::RelationData) -> Option { + type FnPtr = unsafe extern "C" fn( + *mut pgrx::pg_sys::RelationData, + i32, + i32, + ) -> *mut pgrx::pg_sys::IndexScanDescData; + unsafe { + let index_relation = index_relation.as_ref()?; + let indam = index_relation.rd_indam.as_ref()?; + let ambeginscan = indam.ambeginscan.as_ref()?; + Some(core::ptr::fn_addr_eq::( + *ambeginscan, + super::am::ambeginscan, + )) + } + } + + unsafe { + if (*node).type_ == pgrx::pg_sys::NodeTag::T_IndexScanState { + let node = node as *mut pgrx::pg_sys::IndexScanState; + let index_relation = (*node).iss_RelationDesc; + if Some(true) == dirty_check(index_relation) && (*node).iss_ScanDesc.is_null() { + use crate::index::am::Scanner; + + (*node).iss_ScanDesc = pgrx::pg_sys::index_beginscan( + (*node).ss.ss_currentRelation, + (*node).iss_RelationDesc, + (*(*node).ss.ps.state).es_snapshot, + (*node).iss_NumScanKeys, + (*node).iss_NumOrderByKeys, + ); + + let scanner = &mut *((*(*node).iss_ScanDesc).opaque as *mut Scanner); + scanner.hack = std::ptr::NonNull::new(node); + + if (*node).iss_NumRuntimeKeys == 0 || (*node).iss_RuntimeKeysReady { + pgrx::pg_sys::index_rescan( + (*node).iss_ScanDesc, + (*node).iss_ScanKeys, + (*node).iss_NumScanKeys, + (*node).iss_OrderByKeys, + (*node).iss_NumOrderByKeys, + ); + } + } + } + pgrx::pg_sys::planstate_tree_walker(node, Some(rewrite_plan_state), context) + } +} + +static PREV_EXECUTOR_START: AtomicPtr<()> = AtomicPtr::new(core::ptr::null_mut()); + +#[pgrx::pg_guard] +unsafe extern "C" fn executor_start( + query_desc: *mut pgrx::pg_sys::QueryDesc, + eflags: ::std::os::raw::c_int, +) { + unsafe { + use core::mem::transmute; + use pgrx::pg_sys::ExecutorStart_hook_type; + use std::sync::atomic::Ordering; + let value = transmute::<*mut (), ExecutorStart_hook_type>( + PREV_EXECUTOR_START.load(Ordering::Relaxed), + ); + if let Some(prev_executor_start) = value { + prev_executor_start(query_desc, eflags); + } else { + pgrx::pg_sys::standard_ExecutorStart(query_desc, eflags); + } + let planstate = (*query_desc).planstate; + let context = core::ptr::null_mut(); + rewrite_plan_state(planstate, context); + } +} + +pub fn init() { + unsafe { + use core::mem::transmute; + use std::sync::atomic::Ordering; + PREV_EXECUTOR_START.store( + transmute::, *mut ()>( + pgrx::pg_sys::ExecutorStart_hook, + ), + Ordering::Relaxed, + ); + pgrx::pg_sys::ExecutorStart_hook = Some(executor_start); + } +} diff --git a/src/index/mod.rs b/src/index/mod.rs index 925db7d6..453de139 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -2,6 +2,7 @@ pub mod algorithm; pub mod am; pub mod functions; pub mod gucs; +pub mod hook; pub mod opclass; pub mod projection; pub mod scanners; @@ -10,6 +11,7 @@ pub mod types; pub fn init() { am::init(); + hook::init(); gucs::init(); for x in gucs::prewarm_dim() { projection::prewarm(x as _); diff --git a/tests/logic/rerank_in_table.slt b/tests/logic/rerank_in_table.slt index f6a138a1..3f049a78 100644 --- a/tests/logic/rerank_in_table.slt +++ b/tests/logic/rerank_in_table.slt @@ -1,3 +1,6 @@ +statement ok +SET enable_seqscan = off; + statement ok CREATE TABLE t_column (id integer, val vector(3)); @@ -63,5 +66,24 @@ SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <- 8 9 +query I +SELECT id FROM t_expr WHERE id <= 5 OR id % 2 = 1 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' LIMIT 9; +---- +2 +1 +3 +4 +5 +7 +9 +11 +13 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + statement ok DROP TABLE t_expr; From 00e1a797d85d0087a17e60f9dc59a103e21726c8 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 18 Mar 2025 21:02:13 +0800 Subject: [PATCH 131/324] feat: maxsim operator and indexing on maxsim operator (#197) Signed-off-by: usamoi --- .github/workflows/check.yml | 65 ++++++++- .github/workflows/release.yml | 13 +- Cargo.lock | 1 + Cargo.toml | 1 + crates/algorithm/src/build.rs | 3 +- crates/algorithm/src/insert.rs | 2 +- crates/algorithm/src/lib.rs | 11 +- crates/algorithm/src/maintain.rs | 3 + crates/algorithm/src/rerank.rs | 18 ++- crates/algorithm/src/search.rs | 218 ++++++++++++++++++++++++++--- crates/algorithm/src/tape.rs | 2 +- crates/algorithm/src/tuples.rs | 16 ++- crates/algorithm/src/vectors.rs | 2 +- crates/distance/src/lib.rs | 1 + crates/simd/src/f16.rs | 6 +- src/datatype/operators_halfvec.rs | 19 +++ src/datatype/operators_vector.rs | 19 +++ src/index/am/mod.rs | 20 ++- src/index/gucs.rs | 32 +++++ src/index/opclass.rs | 76 ++++++++-- src/index/scanners/default.rs | 12 +- src/index/scanners/maxsim.rs | 224 ++++++++++++++++++++++++++++++ src/index/scanners/mod.rs | 4 + src/sql/finalize.sql | 24 ++++ tests/logic/multivector.slt | 55 ++++++++ 25 files changed, 778 insertions(+), 69 deletions(-) create mode 100644 src/index/scanners/maxsim.rs create mode 100644 tests/logic/multivector.slt diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 956668dd..b94a3660 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -47,7 +47,7 @@ env: jobs: style: - runs-on: "ubuntu-24.04" + runs-on: "ubuntu-latest" steps: - name: Set up Environment @@ -76,8 +76,8 @@ jobs: lint: strategy: matrix: - runner: ["ubuntu-24.04", "ubicloud-standard-8-arm-ubuntu-2404"] - runs-on: ${{ matrix.runner }} + arch: ["x86_64", "aarch64"] + runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-24.04' || 'ubuntu-24.04-arm' }} steps: - name: Set up Environment @@ -125,8 +125,8 @@ jobs: strategy: matrix: version: ["13", "14", "15", "16", "17"] - runner: ["ubuntu-24.04", "ubicloud-standard-8-arm-ubuntu-2404"] - runs-on: ${{ matrix.runner }} + arch: ["x86_64", "aarch64"] + runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} steps: - name: Set up Environment @@ -136,6 +136,7 @@ jobs: sudo apt-get remove -y '^postgres.*' '^libpq.*' sudo apt-get purge -y '^postgres.*' '^libpq.*' + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 sudo apt-get install -y postgresql-common @@ -173,3 +174,57 @@ jobs: sudo systemctl start postgresql psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' sqllogictest --db $USER --user $USER './tests/**/*.slt' + + - name: Package + env: + SEMVER: "0.0.0" + VERSION: ${{ matrix.version }} + ARCH: ${{ matrix.arch }} + PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} + run: | + cat $(pg_config --sharedir)/extension/vchord--0.0.0.sql | expand -t 4 > ./sql/install/vchord--$SEMVER.sql + + mkdir -p ./build/zip + cp -a ./sql/upgrade/. ./build/zip/ + cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql + sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control + cp ./target/release/libvchord.so ./build/zip/vchord.so + zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* + + mkdir -p ./build/deb + mkdir -p ./build/deb/DEBIAN + mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ + mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ + for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file + done + for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file + done + for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do + cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file + done + echo "Package: postgresql-${VERSION}-vchord + Version: ${SEMVER}-1 + Section: database + Priority: optional + Architecture: ${PLATFORM} + Maintainer: Tensorchord + Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM + Homepage: https://vectorchord.ai/ + License: AGPL-3 or Elastic-2" \ + > ./build/deb/DEBIAN/control + (cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums + dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb + + ls ./build + + - name: Upload Artifacts + uses: actions/upload-artifact@v4 + with: + name: artifacts-psql-${{ matrix.version }}-${{ matrix.arch }} + path: | + ./build/postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu.zip + ./build/postgresql-${{ matrix.version }}-vchord_0.0.0-1_${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }}.deb + compression-level: 9 + retention-days: 14 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index f115b4c3..75fb4516 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -40,8 +40,8 @@ jobs: strategy: matrix: version: ["13", "14", "15", "16", "17"] - runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] - runs-on: ${{ matrix.runner }} + arch: ["x86_64", "aarch64"] + runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} env: CARGO_TERM_COLOR: always @@ -75,9 +75,8 @@ jobs: env: SEMVER: ${{ needs.semver.outputs.SEMVER }} VERSION: ${{ matrix.version }} - ARCH: ${{ matrix.runner == 'ubuntu-22.04' && 'x86_64' || 'aarch64' }} - PLATFORM: ${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} - GH_TOKEN: ${{ github.token }} + ARCH: ${{ matrix.arch }} + PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} run: | cargo build --lib --features pg$VERSION --release @@ -116,6 +115,10 @@ jobs: ls ./build + - name: Upload Artifacts + env: + GH_TOKEN: ${{ github.token }} + run: | gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip diff --git a/Cargo.lock b/Cargo.lock index a80cffbe..23382925 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1513,6 +1513,7 @@ name = "vchord" version = "0.0.0" dependencies = [ "algorithm", + "always_equal", "distance", "half 2.4.1", "k_means", diff --git a/Cargo.toml b/Cargo.toml index 395f98a5..e9274669 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] [dependencies] algorithm = { path = "./crates/algorithm" } +always_equal = { path = "./crates/always_equal" } distance = { path = "./crates/distance" } k_means = { path = "./crates/k_means" } random_orthogonal_matrix = { path = "./crates/random_orthogonal_matrix" } diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 0a8221f3..735a741e 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -2,7 +2,7 @@ use crate::operator::{Accessor2, Operator, Vector}; use crate::tape::TapeWriter; use crate::tuples::*; use crate::types::*; -use crate::{Branch, DerefMut, Page, PageGuard, RelationWrite}; +use crate::{Branch, DerefMut, IndexPointer, Page, PageGuard, RelationWrite}; use simd::fast_scan::{any_pack, padding_pack}; use vector::VectorOwned; @@ -55,6 +55,7 @@ pub fn build( jump.push(JumpTuple { frozen_first: frozen_tape.first(), appendable_first: appendable_tape.first(), + tuples: 0, }); level.push(jump.first()); } else { diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index d876aede..6ad410a9 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -3,7 +3,7 @@ use crate::operator::*; use crate::select_heap::SelectHeap; use crate::tuples::*; use crate::vectors::{self}; -use crate::{Page, RelationWrite, tape}; +use crate::{IndexPointer, Page, RelationWrite, tape}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index de792ddf..10c861c9 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -27,10 +27,9 @@ pub use cache::cache; pub use insert::insert; pub use maintain::maintain; pub use prewarm::prewarm; -pub use rerank::{rerank_heap, rerank_index}; -pub use search::search; +pub use rerank::{how, rerank_heap, rerank_index}; +pub use search::{search, search_and_estimate}; -use crate::tuples::IndexPointer; use std::ops::{Deref, DerefMut}; use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -90,3 +89,9 @@ pub(crate) struct Branch { pub signs: Vec, pub extra: T, } + +#[repr(transparent)] +#[derive( + Debug, Default, Clone, Copy, PartialEq, Eq, Hash, IntoBytes, FromBytes, Immutable, KnownLayout, +)] +pub struct IndexPointer(pub u64); diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index e7640049..f079e807 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -62,6 +62,7 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let mut trace = Vec::new(); + let mut tuples = 0_u64; let mut callback = |code: (f32, f32, f32, f32, Vec<_>), mean, payload| { tape.push(Branch { mean, @@ -72,6 +73,7 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { signs: code.4, extra: payload, }); + tuples += 1; }; let mut step = |id| { check(); @@ -131,6 +133,7 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { *jump_tuple.frozen_first() = { frozen_tape }.first(); *jump_tuple.appendable_first() = { appendable_tape }.first(); + *jump_tuple.tuples() = tuples; drop(jump_guard); diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index 4484c9aa..e69dd15b 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -1,6 +1,6 @@ use crate::operator::*; -use crate::tuples::*; -use crate::{RelationRead, vectors}; +use crate::tuples::{MetaTuple, WithReader}; +use crate::{IndexPointer, Page, RelationRead, RerankMethod, vectors}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -8,6 +8,18 @@ use std::collections::BinaryHeap; use std::num::NonZeroU64; use vector::VectorOwned; +pub fn how(index: impl RelationRead) -> RerankMethod { + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let rerank_in_heap = meta_tuple.rerank_in_heap(); + if rerank_in_heap { + RerankMethod::Heap + } else { + RerankMethod::Index + } +} + pub fn rerank_index( index: impl RelationRead, vector: O::Vector, @@ -54,9 +66,9 @@ where let mut heap = BinaryHeap::from(results); let mut cache = BinaryHeap::<(Reverse, _)>::new(); std::iter::from_fn(move || { - let vector = O::Vector::elements_and_metadata(vector.as_borrowed()); while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(_), AlwaysEqual(pay_u)) = heap.pop().unwrap(); + let vector = O::Vector::elements_and_metadata(vector.as_borrowed()); if let Some(vec_u) = fetch(pay_u) { let vec_u = O::Vector::elements_and_metadata(vec_u.as_borrowed()); let mut accessor = O::DistanceAccessor::default(); diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index 7ed279af..edb6b7ff 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -1,7 +1,7 @@ use crate::linked_vec::LinkedVec; use crate::operator::*; use crate::tuples::*; -use crate::{Page, RelationRead, RerankMethod, tape, vectors}; +use crate::{IndexPointer, Page, RelationRead, tape, vectors}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -14,20 +14,174 @@ pub fn search( vector: O::Vector, probes: Vec, epsilon: f32, +) -> Vec<( + Reverse, + AlwaysEqual, + AlwaysEqual, +)> { + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dims = meta_tuple.dims(); + let is_residual = meta_tuple.is_residual(); + let height_of_root = meta_tuple.height_of_root(); + assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); + if height_of_root as usize != 1 + probes.len() { + panic!( + "need {} probes, but {} probes provided", + height_of_root - 1, + probes.len() + ); + } + let root_mean = meta_tuple.root_mean(); + let root_first = meta_tuple.root_first(); + drop(meta_guard); + + let default_lut = if !is_residual { + Some(O::Vector::compute_lut(vector.as_borrowed())) + } else { + None + }; + + type State = Vec<(u32, Option<::Vector>)>; + let mut state: State = vec![{ + let mean = root_mean; + if is_residual { + let residual_u = vectors::read_for_h1_tuple::( + index.clone(), + mean, + LAccess::new( + O::Vector::elements_and_metadata(vector.as_borrowed()), + O::ResidualAccessor::default(), + ), + ); + (root_first, Some(residual_u)) + } else { + (root_first, None) + } + }]; + let step = |state: State| { + let mut results = LinkedVec::new(); + for (first, residual) in state { + let lut = if let Some(residual) = residual { + &O::Vector::compute_lut_block(residual.as_borrowed()) + } else { + default_lut.as_ref().map(|x| &x.0).unwrap() + }; + tape::read_h1_tape( + index.clone(), + first, + || { + RAccess::new( + (&lut.4, (lut.0, lut.1, lut.2, lut.3, epsilon)), + O::Distance::block_accessor(), + ) + }, + |lowerbound, mean, first| { + results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(first))); + }, + |_| (), + ); + } + let mut heap = BinaryHeap::from(results.into_vec()); + let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); + let index = index.clone(); + let vector = vector.as_borrowed(); + std::iter::from_fn(move || { + while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { + let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + if is_residual { + let (dis_u, residual_u) = vectors::read_for_h1_tuple::( + index.clone(), + mean, + LAccess::new( + O::Vector::elements_and_metadata(vector), + ( + O::DistanceAccessor::default(), + O::ResidualAccessor::default(), + ), + ), + ); + cache.push(( + Reverse(dis_u), + AlwaysEqual(first), + AlwaysEqual(Some(residual_u)), + )); + } else { + let dis_u = vectors::read_for_h1_tuple::( + index.clone(), + mean, + LAccess::new( + O::Vector::elements_and_metadata(vector), + O::DistanceAccessor::default(), + ), + ); + cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); + } + } + let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; + Some((first, mean)) + }) + }; + for i in (1..height_of_root).rev() { + state = step(state).take(probes[i as usize - 1] as usize).collect(); + } + + let mut results = LinkedVec::new(); + for (first, residual) in state { + let lut = if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { + &O::Vector::compute_lut(residual) + } else { + default_lut.as_ref().unwrap() + }; + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + let mut callback = |lowerbound, mean, payload| { + results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(payload))); + }; + tape::read_frozen_tape( + index.clone(), + jump_tuple.frozen_first(), + || { + RAccess::new( + (&lut.0.4, (lut.0.0, lut.0.1, lut.0.2, lut.0.3, epsilon)), + O::Distance::block_accessor(), + ) + }, + &mut callback, + |_| (), + ); + tape::read_appendable_tape( + index.clone(), + jump_tuple.appendable_first(), + |code| O::Distance::compute_lowerbound_binary(&lut.1, code, epsilon), + &mut callback, + |_| (), + ); + } + results.into_vec() +} + +pub fn search_and_estimate( + index: impl RelationRead, + vector: O::Vector, + probes: Vec, + epsilon: f32, + mut t: u32, ) -> ( - RerankMethod, Vec<( Reverse, AlwaysEqual, AlwaysEqual, )>, + Distance, ) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); - let rerank_in_heap = meta_tuple.rerank_in_heap(); let height_of_root = meta_tuple.height_of_root(); assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); if height_of_root as usize != 1 + probes.len() { @@ -64,7 +218,7 @@ pub fn search( (root_first, None) } }]; - let step = |state: State, probes| { + let step = |state: State| { let mut results = LinkedVec::new(); for (first, residual) in state { let lut = if let Some(residual) = residual { @@ -89,7 +243,9 @@ pub fn search( } let mut heap = BinaryHeap::from(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - std::iter::from_fn(|| { + let index = index.clone(); + let vector = vector.as_borrowed(); + std::iter::from_fn(move || { while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); if is_residual { @@ -97,7 +253,7 @@ pub fn search( index.clone(), mean, LAccess::new( - O::Vector::elements_and_metadata(vector.as_borrowed()), + O::Vector::elements_and_metadata(vector), ( O::DistanceAccessor::default(), O::ResidualAccessor::default(), @@ -114,21 +270,33 @@ pub fn search( index.clone(), mean, LAccess::new( - O::Vector::elements_and_metadata(vector.as_borrowed()), + O::Vector::elements_and_metadata(vector), O::DistanceAccessor::default(), ), ); cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); } } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; - Some((first, mean)) + let (Reverse(distance), AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; + Some((first, mean, distance)) }) - .take(probes as usize) - .collect() }; - for i in (1..height_of_root).rev() { - state = step(state, probes[i as usize - 1]); + for i in (2..height_of_root).rev() { + state = step(state) + .map(|(x, y, _)| (x, y)) + .take(probes[i as usize - 1] as usize) + .collect(); + } + let mut iter: Box>; + if height_of_root > 1 { + let mut it = step(state); + state = std::iter::from_fn(|| it.next()) + .map(|(x, y, _)| (x, y)) + .take(probes[0] as usize) + .collect(); + iter = Box::new(it.map(|(x, _, z)| (x, z))); + } else { + iter = Box::new(std::iter::empty()); } let mut results = LinkedVec::new(); @@ -163,13 +331,23 @@ pub fn search( &mut callback, |_| (), ); + t = t.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as u32); } - ( - if rerank_in_heap { - RerankMethod::Heap + let mut estimation = f32::NAN; + loop { + if t != 0 { + if let Some((first, distance)) = iter.next() { + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + t = t.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as u32); + estimation = distance.to_f32(); + } else { + break; + } } else { - RerankMethod::Index - }, - results.into_vec(), - ) + break; + } + } + (results.into_vec(), Distance::from_f32(estimation)) } diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index e90be1f6..53996734 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -1,6 +1,6 @@ use crate::operator::Accessor1; use crate::tuples::*; -use crate::{Page, PageGuard, RelationRead, RelationWrite}; +use crate::{IndexPointer, Page, PageGuard, RelationRead, RelationWrite}; use std::marker::PhantomData; use std::num::NonZeroU64; use std::ops::DerefMut; diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index c85588b4..ffd64440 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -1,3 +1,4 @@ +use crate::IndexPointer; use crate::operator::Vector; use std::marker::PhantomData; use std::num::{NonZeroU8, NonZeroU64}; @@ -560,12 +561,14 @@ impl<'a> H1TupleReader1<'a> { struct JumpTupleHeader { frozen_first: u32, appendable_first: u32, + tuples: u64, } #[derive(Debug, Clone, PartialEq)] pub struct JumpTuple { pub frozen_first: u32, pub appendable_first: u32, + pub tuples: u64, } impl Tuple for JumpTuple { @@ -573,6 +576,7 @@ impl Tuple for JumpTuple { JumpTupleHeader { frozen_first: self.frozen_first, appendable_first: self.appendable_first, + tuples: self.tuples, } .as_bytes() .to_vec() @@ -609,6 +613,9 @@ impl JumpTupleReader<'_> { pub fn appendable_first(self) -> u32 { self.header.appendable_first } + pub fn tuples(self) -> u64 { + self.header.tuples + } } #[derive(Debug)] @@ -623,6 +630,9 @@ impl JumpTupleWriter<'_> { pub fn appendable_first(&mut self) -> &mut u32 { &mut self.header.appendable_first } + pub fn tuples(&mut self) -> &mut u64 { + &mut self.header.tuples + } } #[repr(C, align(8))] @@ -998,12 +1008,6 @@ const fn soundness_check() { )] pub struct ZeroU8(Option); -#[repr(transparent)] -#[derive( - Debug, Default, Clone, Copy, PartialEq, Eq, Hash, IntoBytes, FromBytes, Immutable, KnownLayout, -)] -pub struct IndexPointer(pub u64); - #[repr(transparent)] #[derive( Debug, diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index 42b83b0b..d1319a62 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -1,6 +1,6 @@ use crate::operator::*; use crate::tuples::*; -use crate::{Page, PageGuard, RelationRead, RelationWrite, tape}; +use crate::{IndexPointer, Page, PageGuard, RelationRead, RelationWrite, tape}; use std::num::NonZeroU64; use vector::VectorOwned; diff --git a/crates/distance/src/lib.rs b/crates/distance/src/lib.rs index 5f21aace..1d8fb05d 100644 --- a/crates/distance/src/lib.rs +++ b/crates/distance/src/lib.rs @@ -6,6 +6,7 @@ impl Distance { pub const ZERO: Self = Distance::from_f32(0.0f32); pub const INFINITY: Self = Distance::from_f32(f32::INFINITY); pub const NEG_INFINITY: Self = Distance::from_f32(f32::NEG_INFINITY); + pub const NAN: Self = Distance::from_f32(f32::NAN); #[inline(always)] pub const fn from_f32(value: f32) -> Self { diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 2fb176be..60e4d045 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -533,7 +533,7 @@ mod reduce_sum_of_d2 { #[test] fn reduce_sum_of_d2_v4_512_avx512fp16_test() { use rand::Rng; - const EPSILON: f32 = 6.0; + const EPSILON: f32 = 6.4; if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512fp16") { println!("test {} ... skipped (v4_512:avx512fp16)", module_path!()); return; @@ -706,7 +706,7 @@ mod reduce_sum_of_d2 { #[test] fn reduce_sum_of_d2_a2_fp16_test() { use rand::Rng; - const EPSILON: f32 = 6.0; + const EPSILON: f32 = 6.4; if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { println!("test {} ... skipped (a2:fp16)", module_path!()); return; @@ -751,7 +751,7 @@ mod reduce_sum_of_d2 { #[test] fn reduce_sum_of_d2_a3_512_test() { use rand::Rng; - const EPSILON: f32 = 6.0; + const EPSILON: f32 = 6.4; if !crate::is_cpu_detected!("a3.512") { println!("test {} ... skipped (a3.512)", module_path!()); return; diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index 907d415d..0be60ccb 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -1,4 +1,5 @@ use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use pgrx::Array; use std::num::NonZero; use vector::VectorBorrowed; use vector::vect::VectBorrowed; @@ -74,3 +75,21 @@ fn _vchord_halfvec_sphere_cosine_in( let d = VectBorrowed::operator_cos(lhs, center).to_f32(); d < radius } + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_halfvec_operator_maxsim_ip( + lhs: Array<'_, HalfvecInput<'_>>, + rhs: Array<'_, HalfvecInput<'_>>, +) -> f32 { + let mut maxsim = 0.0f32; + for rhs in rhs.iter().flatten() { + let mut d = f32::INFINITY; + for lhs in lhs.iter().flatten() { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); + } + maxsim += d; + } + maxsim +} diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index b011c3ae..d8d60b50 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -1,4 +1,5 @@ use crate::datatype::memory_vector::{VectorInput, VectorOutput}; +use pgrx::Array; use std::num::NonZero; use vector::VectorBorrowed; use vector::vect::VectBorrowed; @@ -74,3 +75,21 @@ fn _vchord_vector_sphere_cosine_in( let d = VectBorrowed::operator_cos(lhs, center).to_f32(); d < radius } + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_vector_operator_maxsim_ip( + lhs: Array<'_, VectorInput<'_>>, + rhs: Array<'_, VectorInput<'_>>, +) -> f32 { + let mut maxsim = 0.0f32; + for rhs in rhs.iter().flatten() { + let mut d = f32::INFINITY; + for lhs in lhs.iter().flatten() { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); + } + maxsim += d; + } + maxsim +} diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index 41a21e94..abed8414 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -1,6 +1,6 @@ pub mod am_build; -use crate::index::gucs::{epsilon, max_scan_tuples, probes}; +use crate::index::gucs::{epsilon, max_maxsim_tuples, max_scan_tuples, maxsim_threshold, probes}; use crate::index::opclass::opfamily; use crate::index::scanners::*; use crate::index::storage::PostgresRelation; @@ -292,6 +292,8 @@ pub unsafe extern "C" fn amrescan( epsilon: epsilon(), probes: probes(), max_scan_tuples: max_scan_tuples(), + max_maxsim_tuples: max_maxsim_tuples(), + maxsim_threshold: maxsim_threshold(), }; let scanner = &mut *(*scan).opaque.cast::(); let fetcher = { @@ -327,6 +329,22 @@ pub unsafe extern "C" fn amrescan( } LazyCell::new(Box::new(move || builder.build(relation, options, fetcher))) } + Opfamily::VectorMaxsimIp | Opfamily::HalfvecMaxsimIp => { + let mut builder = MaxsimBuilder::new(opfamily); + for i in 0..(*scan).numberOfOrderBys { + let data = (*scan).orderByData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + for i in 0..(*scan).numberOfKeys { + let data = (*scan).keyData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + LazyCell::new(Box::new(move || builder.build(relation, options, fetcher))) + } }; } } diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 26d061ab..c0849bc4 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -6,6 +6,8 @@ static EPSILON: GucSetting = GucSetting::::new(1.9); static MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); static PREWARM_DIM: GucSetting> = GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); +static MAX_MAXSIM_TUPLES: GucSetting = GucSetting::::new(-1); +static MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); pub fn init() { GucRegistry::define_string_guc( @@ -44,6 +46,26 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_int_guc( + "vchordrq.max_maxsim_tuples", + "`max_maxsim_tuples` argument of vchordrq.", + "`max_maxsim_tuples` argument of vchordrq.", + &MAX_MAXSIM_TUPLES, + -1, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + "vchordrq.maxsim_threshold", + "`maxsim_threshold` argument of vchordrq.", + "`maxsim_threshold` argument of vchordrq.", + &MAXSIM_THRESHOLD, + 0, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); @@ -89,6 +111,16 @@ pub fn max_scan_tuples() -> Option { if x < 0 { None } else { Some(x as u32) } } +pub fn max_maxsim_tuples() -> Option { + let x = MAX_MAXSIM_TUPLES.get(); + if x < 0 { None } else { Some(x as u32) } +} + +pub fn maxsim_threshold() -> u32 { + let x = MAXSIM_THRESHOLD.get(); + x as u32 +} + pub fn prewarm_dim() -> Vec { if let Some(prewarm_dim) = PREWARM_DIM.get() { if let Ok(prewarm_dim) = prewarm_dim.to_str() { diff --git a/src/index/opclass.rs b/src/index/opclass.rs index 4b989e85..9067fbe0 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -38,6 +38,16 @@ fn _vchordrq_support_halfvec_cosine_ops() -> String { "halfvec_cosine_ops".to_string() } +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_vector_maxsim_ip_ops() -> String { + "vector_maxsim_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_halfvec_maxsim_ip_ops() -> String { + "halfvec_maxsim_ip_ops".to_string() +} + pub struct Sphere { pub center: T, pub radius: f32, @@ -51,6 +61,8 @@ pub enum Opfamily { HalfvecL2, HalfvecIp, HalfvecCosine, + VectorMaxsimIp, + HalfvecMaxsimIp, } impl Opfamily { @@ -58,11 +70,11 @@ impl Opfamily { use {BorrowedVector as B, OwnedVector as O}; match (vector, self) { (B::Vecf32(x), Self::VectorL2) => O::Vecf32(x.own()), - (B::Vecf32(x), Self::VectorIp) => O::Vecf32(x.own()), + (B::Vecf32(x), Self::VectorIp | Self::VectorMaxsimIp) => O::Vecf32(x.own()), (B::Vecf32(x), Self::VectorCosine) => O::Vecf32(x.function_normalize()), (B::Vecf32(_), _) => unreachable!(), (B::Vecf16(x), Self::HalfvecL2) => O::Vecf16(x.own()), - (B::Vecf16(x), Self::HalfvecIp) => O::Vecf16(x.own()), + (B::Vecf16(x), Self::HalfvecIp | Self::HalfvecMaxsimIp) => O::Vecf16(x.own()), (B::Vecf16(x), Self::HalfvecCosine) => O::Vecf16(x.function_normalize()), (B::Vecf16(_), _) => unreachable!(), } @@ -80,6 +92,30 @@ impl Opfamily { let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] } + Self::VectorMaxsimIp => { + let vectors = + unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Vecf32(vector.as_borrowed())), + i as u16, + )); + } + result + } + Self::HalfvecMaxsimIp => { + let vectors = + unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Vecf16(vector.as_borrowed())), + i as u16, + )); + } + result + } }; Some(store) } @@ -91,11 +127,11 @@ impl Opfamily { let attno_2 = NonZero::new(2_usize).unwrap(); let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; let center = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsimIp => { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Vecf32(vector.as_borrowed())) } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsimIp => { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Vecf16(vector.as_borrowed())) } @@ -108,24 +144,23 @@ impl Opfamily { return None; } let vector = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsimIp => { let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; self.input(BorrowedVector::Vecf32(vector.as_borrowed())) } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsimIp => { let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; self.input(BorrowedVector::Vecf16(vector.as_borrowed())) } }; Some(vector) } - #[allow(dead_code)] pub unsafe fn input_vectors(self, datum: Datum) -> Option> { if datum.is_null() { return None; } let vectors = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsimIp => { let vectors = unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; let mut result = Vec::with_capacity(vectors.len()); @@ -134,7 +169,7 @@ impl Opfamily { } result } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsimIp => { let vectors = unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; let mut result = Vec::with_capacity(vectors.len()); @@ -150,21 +185,30 @@ impl Opfamily { match self { Self::VectorCosine | Self::HalfvecCosine => x.to_f32() + 1.0f32, Self::VectorL2 | Self::HalfvecL2 => x.to_f32().sqrt(), - Self::VectorIp | Self::HalfvecIp => x.to_f32(), + Self::VectorIp | Self::HalfvecIp | Self::VectorMaxsimIp | Self::HalfvecMaxsimIp => { + x.to_f32() + } } } pub const fn distance_kind(self) -> DistanceKind { match self { Self::VectorL2 | Self::HalfvecL2 => DistanceKind::L2, - Self::VectorIp | Self::HalfvecIp | Self::VectorCosine | Self::HalfvecCosine => { - DistanceKind::Dot - } + Self::VectorIp + | Self::HalfvecIp + | Self::VectorCosine + | Self::HalfvecCosine + | Self::VectorMaxsimIp + | Self::HalfvecMaxsimIp => DistanceKind::Dot, } } pub const fn vector_kind(self) -> VectorKind { match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine => VectorKind::Vecf32, - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => VectorKind::Vecf16, + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsimIp => { + VectorKind::Vecf32 + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsimIp => { + VectorKind::Vecf16 + } } } } @@ -207,6 +251,8 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { "halfvec_l2_ops" => Opfamily::HalfvecL2, "halfvec_ip_ops" => Opfamily::HalfvecIp, "halfvec_cosine_ops" => Opfamily::HalfvecCosine, + "vector_maxsim_ip_ops" => Opfamily::VectorMaxsimIp, + "halfvec_maxsim_ip_ops" => Opfamily::HalfvecMaxsimIp, _ => pgrx::error!("unknown operator class"), }; diff --git a/src/index/scanners/default.rs b/src/index/scanners/default.rs index 107d201b..493987a0 100644 --- a/src/index/scanners/default.rs +++ b/src/index/scanners/default.rs @@ -83,7 +83,7 @@ impl SearchBuilder for DefaultBuilder { let vector = RandomProject::project( VectOwned::::from_owned(vector.clone()).as_borrowed(), ); - let (method, results) = algorithm::search::, L2>>( + let results = algorithm::search::, L2>>( relation.clone(), vector.clone(), options.probes, @@ -97,6 +97,7 @@ impl SearchBuilder for DefaultBuilder { let raw = VectOwned::::from_owned(maybe_vector.unwrap()); Some(RandomProject::project(raw.as_borrowed())) }; + let method = algorithm::how(relation.clone()); match method { RerankMethod::Index => Box::new( algorithm::rerank_index::, L2>>( @@ -116,7 +117,7 @@ impl SearchBuilder for DefaultBuilder { let vector = RandomProject::project( VectOwned::::from_owned(vector.clone()).as_borrowed(), ); - let (method, results) = algorithm::search::, Dot>>( + let results = algorithm::search::, Dot>>( relation.clone(), vector.clone(), options.probes, @@ -130,6 +131,7 @@ impl SearchBuilder for DefaultBuilder { let raw = VectOwned::::from_owned(maybe_vector.unwrap()); Some(RandomProject::project(raw.as_borrowed())) }; + let method = algorithm::how(relation.clone()); match method { RerankMethod::Index => Box::new( algorithm::rerank_index::, Dot>>( @@ -149,7 +151,7 @@ impl SearchBuilder for DefaultBuilder { let vector = RandomProject::project( VectOwned::::from_owned(vector.clone()).as_borrowed(), ); - let (method, results) = algorithm::search::, L2>>( + let results = algorithm::search::, L2>>( relation.clone(), vector.clone(), options.probes, @@ -163,6 +165,7 @@ impl SearchBuilder for DefaultBuilder { let raw = VectOwned::::from_owned(maybe_vector.unwrap()); Some(RandomProject::project(raw.as_borrowed())) }; + let method = algorithm::how(relation.clone()); match method { RerankMethod::Index => Box::new( algorithm::rerank_index::, L2>>( @@ -182,7 +185,7 @@ impl SearchBuilder for DefaultBuilder { let vector = RandomProject::project( VectOwned::::from_owned(vector.clone()).as_borrowed(), ); - let (method, results) = algorithm::search::, Dot>>( + let results = algorithm::search::, Dot>>( relation.clone(), vector.clone(), options.probes, @@ -196,6 +199,7 @@ impl SearchBuilder for DefaultBuilder { let raw = VectOwned::::from_owned(maybe_vector.unwrap()); Some(RandomProject::project(raw.as_borrowed())) }; + let method = algorithm::how(relation.clone()); match method { RerankMethod::Index => Box::new( algorithm::rerank_index::, Dot>>( diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs new file mode 100644 index 00000000..016ff5c3 --- /dev/null +++ b/src/index/scanners/maxsim.rs @@ -0,0 +1,224 @@ +use super::{SearchBuilder, SearchFetcher, SearchOptions}; +use crate::index::algorithm::RandomProject; +use crate::index::am::pointer_to_ctid; +use crate::index::opclass::Opfamily; +use algorithm::operator::{Dot, Op, Vector}; +use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; +use algorithm::{RelationRead, RerankMethod}; +use always_equal::AlwaysEqual; +use distance::Distance; +use half::f16; +use pgrx::pg_sys::{BlockIdData, ItemPointerData}; +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashMap}; +use vector::VectorOwned; +use vector::vect::VectOwned; + +pub struct MaxsimBuilder { + opfamily: Opfamily, + orderbys: Vec>>, +} + +impl SearchBuilder for MaxsimBuilder { + fn new(opfamily: Opfamily) -> Self { + assert!(matches!( + opfamily, + Opfamily::VectorMaxsimIp | Opfamily::HalfvecMaxsimIp + )); + Self { + opfamily, + orderbys: Vec::new(), + } + } + + unsafe fn add(&mut self, strategy: u16, datum: Option) { + match strategy { + 3 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_vectors(x)) }; + self.orderbys.push(x); + } + _ => unreachable!(), + } + } + + fn build<'a>( + self, + relation: impl RelationRead + 'a, + options: SearchOptions, + _: impl SearchFetcher + 'a, + ) -> Box + 'a> { + let mut vectors = None; + for orderby_vectors in self.orderbys.into_iter().flatten() { + if vectors.is_none() { + vectors = Some(orderby_vectors); + } else { + pgrx::error!("maxsim search with multiple vectors is not supported"); + } + } + if let Some(_max_scan_tuples) = options.max_scan_tuples { + pgrx::error!("maxsim search with max_scan_tuples is not supported"); + } + let max_maxsim_tuples = options + .max_maxsim_tuples + .expect("max_maxsim_tuples must be set for maxsim search"); + let maxsim_threshold = options.maxsim_threshold; + let opfamily = self.opfamily; + let Some(vectors) = vectors else { + return Box::new(std::iter::empty()) + as Box>; + }; + let method = algorithm::how(relation.clone()); + if !matches!(method, RerankMethod::Index) { + pgrx::error!("maxsim search with rerank_in_table is not supported"); + } + assert!(matches!(opfamily.distance_kind(), DistanceKind::Dot)); + let (c, estimations) = match opfamily.vector_kind() { + VectorKind::Vecf32 => { + let vectors = vectors + .into_iter() + .map(|vector| { + RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ) + }) + .collect::>(); + let mut estimations = Vec::new(); + let mut c = HashMap::<[u16; 3], States>::new(); + for (query_id, vector) in vectors.iter().enumerate() { + let (results, est) = algorithm::search_and_estimate::, Dot>>( + relation.clone(), + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + ); + if results.is_empty() { + estimations.push(0.0); + continue; + } + let returning = algorithm::rerank_index::, Dot>>( + relation.clone(), + vector.clone(), + results, + ) + .take(max_maxsim_tuples as usize) + .collect::>(); + let mut max = f32::NEG_INFINITY; + for (distance, payload) in returning { + max = max.max(distance.to_f32()); + use std::collections::hash_map::Entry::{Occupied, Vacant}; + let (ctid, _) = pointer_to_ctid(payload); + let key = [ctid.ip_blkid.bi_hi, ctid.ip_blkid.bi_lo, ctid.ip_posid]; + match c.entry(key) { + Vacant(e) => { + let states = e.insert(States::new(vectors.len())); + states.update(query_id, distance); + } + Occupied(mut e) => { + let states = e.get_mut(); + states.update(query_id, distance); + } + } + } + estimations.push(max.max(est.to_f32())); + } + (c, estimations) + } + VectorKind::Vecf16 => { + let vectors = vectors + .into_iter() + .map(|vector| { + RandomProject::project( + VectOwned::::from_owned(vector.clone()).as_borrowed(), + ) + }) + .collect::>(); + let mut estimations = Vec::new(); + let mut c = HashMap::<[u16; 3], States>::new(); + for (query_id, vector) in vectors.iter().enumerate() { + let (results, est) = algorithm::search_and_estimate::, Dot>>( + relation.clone(), + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + ); + if results.is_empty() { + estimations.push(0.0); + continue; + } + let returning = algorithm::rerank_index::, Dot>>( + relation.clone(), + vector.clone(), + results, + ) + .take(max_maxsim_tuples as usize) + .collect::>(); + let mut max = f32::NEG_INFINITY; + for (distance, payload) in returning { + max = max.max(distance.to_f32()); + use std::collections::hash_map::Entry::{Occupied, Vacant}; + let (ctid, _) = pointer_to_ctid(payload); + let key = [ctid.ip_blkid.bi_hi, ctid.ip_blkid.bi_lo, ctid.ip_posid]; + match c.entry(key) { + Vacant(e) => { + let states = e.insert(States::new(vectors.len())); + states.update(query_id, distance); + } + Occupied(mut e) => { + let states = e.get_mut(); + states.update(query_id, distance); + } + } + } + estimations.push(max.max(est.to_f32())); + } + (c, estimations) + } + }; + let c = c + .into_iter() + .map(|(key, states)| { + let mut maxsim = 0.0f32; + for (query_id, distance) in states.into_iter() { + let d = distance.unwrap_or(Distance::from_f32(estimations[query_id])); + maxsim += opfamily.output(d); + } + (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) + }) + .collect::>(); + let mut c = BinaryHeap::from(c); + Box::new(std::iter::from_fn(move || { + let (Reverse(distance), AlwaysEqual(key)) = c.pop()?; + let distance = distance.to_f32(); + let ctid = ItemPointerData { + ip_blkid: BlockIdData { + bi_hi: key[0], + bi_lo: key[1], + }, + ip_posid: key[2], + }; + let recheck = false; + Some((distance, ctid, recheck)) + })) + } +} + +struct States { + inner: Vec>, +} + +impl States { + fn new(len: usize) -> Self { + Self { + inner: vec![None; len], + } + } + fn update(&mut self, query_id: usize, distance: Distance) { + let this = self.inner[query_id].get_or_insert(Distance::INFINITY); + *this = std::cmp::min(*this, distance); + } + fn into_iter(self) -> impl Iterator)> { + self.inner.into_iter().enumerate() + } +} diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs index afef6263..bad06d54 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/scanners/mod.rs @@ -1,4 +1,5 @@ mod default; +mod maxsim; use super::opclass::Opfamily; use algorithm::RelationRead; @@ -6,12 +7,15 @@ use pgrx::pg_sys::{Datum, ItemPointerData}; use std::cell::LazyCell; pub use default::DefaultBuilder; +pub use maxsim::MaxsimBuilder; #[derive(Debug)] pub struct SearchOptions { pub epsilon: f32, pub probes: Vec, pub max_scan_tuples: Option, + pub max_maxsim_tuples: Option, + pub maxsim_threshold: u32, } pub trait SearchBuilder: 'static { diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 7bc36b62..171e8d1c 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -113,6 +113,18 @@ CREATE OPERATOR <<=>> ( COMMUTATOR = <<=>> ); +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim_ip, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim_ip, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + -- List of functions CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector @@ -148,6 +160,8 @@ CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ip_ops USING vchordrq; -- List of operator classes @@ -186,3 +200,13 @@ CREATE OPERATOR CLASS halfvec_cosine_ops OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ip_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ip_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ip_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ip_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ip_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ip_ops(); diff --git a/tests/logic/multivector.slt b/tests/logic/multivector.slt new file mode 100644 index 00000000..5cc772e0 --- /dev/null +++ b/tests/logic/multivector.slt @@ -0,0 +1,55 @@ +statement ok +CREATE TABLE t (id integer, val vector(2)[]); + +statement ok +INSERT INTO t (id, val) +SELECT id, + ARRAY[ + ARRAY[cos(((id + 0) % 10000) / 10000.0 * 6.283185307179586), sin(((id + 0) % 10000) / 10000.0 * 6.283185307179586)]::vector, + ARRAY[cos(((id + 22) % 10000) / 10000.0 * 6.283185307179586), sin(((id + 22) % 10000) / 10000.0 * 6.283185307179586)]::vector, + ARRAY[cos(((id + 777) % 10000) / 10000.0 * 6.283185307179586), sin(((id + 777) % 10000) / 10000.0 * 6.283185307179586)]::vector + ] +FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX t_val_idx ON t USING vchordrq (val vector_maxsim_ip_ops) +WITH (options = $$ +build.internal.lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +statement ok +SET vchordrq.max_maxsim_tuples = 3000; + +statement ok +SET enable_seqscan TO off; + +query I +SELECT id FROM t ORDER BY val @# ARRAY['[0.7197411498053302, 0.6942425205048314]'::vector, '[0.10645067063129976, 0.9943179847122079]'::vector] limit 18; +---- +1387 +1388 +1386 +1389 +1385 +1390 +1384 +1391 +1383 +1392 +1382 +1393 +1381 +1394 +1380 +1395 +1379 +1396 + +statement ok +DROP INDEX t_val_idx; + +statement ok +DROP TABLE t; From c7cc49d036d68b7a5bee03b344438cca3d6372a7 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 19 Mar 2025 19:19:12 +0800 Subject: [PATCH 132/324] chore: update dependencies (#216) Signed-off-by: usamoi --- .cargo/config.toml | 3 + .github/workflows/check.yml | 2 +- .github/workflows/release.yml | 2 +- Cargo.lock | 383 +++--- Cargo.toml | 16 +- crates/algorithm/src/tuples.rs | 1 - crates/random_orthogonal_matrix/Cargo.toml | 2 +- crates/simd/Cargo.toml | 2 +- crates/simd/build.rs | 13 +- crates/simd/cshim.c | 30 +- crates/simd/src/bit.rs | 204 ++- crates/simd/src/emulate.rs | 204 ++- crates/simd/src/f16.rs | 298 +++-- crates/simd/src/f32.rs | 1316 ++++++++++---------- crates/simd/src/fast_scan/mod.rs | 392 +++--- crates/simd/src/quantize.rs | 262 ++-- crates/simd/src/u8.rs | 310 +++-- crates/simd_macros/Cargo.toml | 6 +- rust-toolchain.toml | 2 +- src/index/am/am_build.rs | 5 +- 20 files changed, 1642 insertions(+), 1811 deletions(-) diff --git a/.cargo/config.toml b/.cargo/config.toml index c4db6490..f4f72dfd 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -4,3 +4,6 @@ rustdocflags = ["--document-private-items"] [target.'cfg(target_os="macos")'] # Postgres symbols won't be available until runtime rustflags = ["-Clink-arg=-Wl,-undefined,dynamic_lookup"] + +[env] +CC = "clang" diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index b94a3660..e729fc17 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -152,7 +152,7 @@ jobs: sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.so"' sudo systemctl stop postgresql - curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.12.9/cargo-pgrx-v0.12.9-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx + curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.13.1/cargo-pgrx-v0.13.1-$(uname -m)-unknown-linux-gnu.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 75fb4516..fd36e4ab 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -65,7 +65,7 @@ jobs: sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector - curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.12.9/cargo-pgrx-v0.12.9-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx + curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.13.1/cargo-pgrx-v0.13.1-$(uname -m)-unknown-linux-gnu.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) - name: Checkout diff --git a/Cargo.lock b/Cargo.lock index 23382925..8ed8973a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,7 +17,7 @@ version = "0.0.0" dependencies = [ "always_equal", "distance", - "half 2.4.1", + "half 2.5.0", "k_means", "paste", "rabitq", @@ -28,8 +28,8 @@ dependencies = [ "turboselect", "validator", "vector", - "zerocopy 0.8.17", - "zerocopy-derive 0.8.17", + "zerocopy", + "zerocopy-derive", ] [[package]] @@ -38,19 +38,25 @@ version = "0.0.0" [[package]] name = "annotate-snippets" -version = "0.9.2" +version = "0.11.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ccaf7e9dfbb6ab22c82e473cd1a8a7bd313c19a5b7e40970f3d89ef5a5c9e81e" +checksum = "710e8eae58854cdc1790fcb56cca04d712a17be849eeb81da2a724bf4bae2bc4" dependencies = [ + "anstyle", "unicode-width", - "yansi-term", ] +[[package]] +name = "anstyle" +version = "1.0.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9" + [[package]] name = "anyhow" -version = "1.0.95" +version = "1.0.97" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34ac096ce696dc2fcabef30516bb13c0a68a11d30131d3df6f04711467681b04" +checksum = "dcfed56ad506cb2c684a14971b8861fdc3baaaae314b9e5f9bb532cbe3ba7a4f" [[package]] name = "approx" @@ -63,9 +69,9 @@ dependencies = [ [[package]] name = "atomic-traits" -version = "0.3.0" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b29ec3788e96fb4fdb275ccb9d62811f2fa903d76c5eb4dd6fe7d09a7ed5871f" +checksum = "707f750b93bd1b739cf9ddf85f8fe7c97a4a62c60ccf8b6f232514bd9103bedc" dependencies = [ "cfg-if", "rustc_version", @@ -79,9 +85,9 @@ checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" [[package]] name = "bindgen" -version = "0.70.1" +version = "0.71.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f49d8fed880d473ea71efb9bf597651e77201bdd4893efe54c9e5d65ae04ce6f" +checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" dependencies = [ "annotate-snippets", "bitflags", @@ -98,9 +104,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.8.0" +version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f68f53c83ab957f72c32642f3868eec03eb974d1fb82e453128456482613d36" +checksum = "5c8214115b7bf84099f1309324e63141d4c5d7cc26862f97a0a857dbefe165bd" [[package]] name = "bitvec" @@ -116,9 +122,9 @@ dependencies = [ [[package]] name = "bytemuck" -version = "1.21.0" +version = "1.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef657dfab802224e671f5818e9a4935f9b1957ed18e58292690cc39e7a4092a3" +checksum = "b6b1fc10dbac614ebc03540c9dbd60e83887fda27794998c6528f1782047d540" [[package]] name = "byteorder" @@ -128,9 +134,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "cargo_toml" -version = "0.19.2" +version = "0.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a98356df42a2eb1bd8f1793ae4ee4de48e384dd974ce5eac8eee802edb7492be" +checksum = "5fbd1fe9db3ebf71b89060adaf7b0504c2d6a425cf061313099547e382c2e472" dependencies = [ "serde", "toml", @@ -138,9 +144,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.13" +version = "1.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7777341816418c02e033934a09f20dc0ccaf65a5201ef8a450ae0105a573fda" +checksum = "be714c154be609ec7f5dad223a33bf1482fff90472de28f7362806e6d4832b8c" dependencies = [ "shlex", ] @@ -183,9 +189,9 @@ dependencies = [ [[package]] name = "convert_case" -version = "0.6.0" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec182b0ca2f35d8fc196cf3404988fd8b8c739a4d270ff118a398feb0cbec1ca" +checksum = "bb402b8d4c85569410425650ce3eddc7d698ed96d39a73f941b08fb63082f1e7" dependencies = [ "unicode-segmentation", ] @@ -273,9 +279,9 @@ version = "0.0.0" [[package]] name = "either" -version = "1.13.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" [[package]] name = "enum-map" @@ -299,9 +305,9 @@ dependencies = [ [[package]] name = "equivalent" -version = "1.0.1" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "eyre" @@ -315,9 +321,9 @@ dependencies = [ [[package]] name = "fixedbitset" -version = "0.4.2" +version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" [[package]] name = "fnv" @@ -342,14 +348,14 @@ checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" [[package]] name = "getrandom" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8" +checksum = "73fea8450eea4bac3940448fb7ae50d91f034f941199fcd9d909a5a07aa455f0" dependencies = [ "cfg-if", "libc", + "r-efi", "wasi", - "windows-targets", ] [[package]] @@ -366,14 +372,13 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403" [[package]] name = "half" -version = "2.4.1" -source = "git+https://github.com/tensorchord/half-rs.git?rev=3f9a8843d6722bd1833de2289347640ad8770146#3f9a8843d6722bd1833de2289347640ad8770146" +version = "2.5.0" +source = "git+https://github.com/tensorchord/half-rs.git?rev=9dc2559e95f295aaa8ffef5e0529167731dffd47#9dc2559e95f295aaa8ffef5e0529167731dffd47" dependencies = [ "cfg-if", "crunchy", "serde", - "zerocopy 0.8.17", - "zerocopy-derive 0.8.17", + "zerocopy", ] [[package]] @@ -403,9 +408,9 @@ dependencies = [ [[package]] name = "hermit-abi" -version = "0.4.0" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc" +checksum = "fbd780fe5cc30f81464441920d82ac8740e2e46b29a6fad543ddd075229ce37e" [[package]] name = "home" @@ -569,9 +574,9 @@ checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" [[package]] name = "indexmap" -version = "2.7.1" +version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c9c992b02b5b4c94ea26e32fe5bccb7aa7d9f390ab5c1221ff895bc7ea8b652" +checksum = "3954d50fe15b02142bf25d3b8bdadb634ec3948f103d04ffe3031bc8fe9d7058" dependencies = [ "equivalent", "hashbrown", @@ -579,9 +584,9 @@ dependencies = [ [[package]] name = "is-terminal" -version = "0.4.15" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e19b23d53f35ce9f56aebc7d1bb4e6ac1e9c0db7ac85c8d1760c04379edced37" +checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", @@ -605,15 +610,15 @@ dependencies = [ [[package]] name = "itoa" -version = "1.0.14" +version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674" +checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" [[package]] name = "k_means" version = "0.0.0" dependencies = [ - "half 2.4.1", + "half 2.5.0", "rabitq", "rand", "rayon", @@ -622,9 +627,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.169" +version = "0.2.171" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a" +checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6" [[package]] name = "libloading" @@ -644,9 +649,9 @@ checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" [[package]] name = "litemap" -version = "0.7.4" +version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ee93343901ab17bd981295f2cf0026d4ad018c7c31ba84549a4ddbb47a45104" +checksum = "23fb14cb19457329c82206317a5663005a4d404783dc74f4252769b0d5f42856" [[package]] name = "matrixmultiply" @@ -758,15 +763,15 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.20.3" +version = "1.21.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e" +checksum = "d75b0bedcc4fe52caa0e03d9f1151a323e4aa5e2d78ba3580400cd3c9e2bc4bc" [[package]] name = "owo-colors" -version = "4.1.0" +version = "4.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb37767f6569cd834a413442455e0f066d0d522de8630436e2a1761d9726ba56" +checksum = "1036865bb9422d3300cf723f657c2851d0e9ab12567854b1f4eba3d77decf564" dependencies = [ "supports-color 2.1.0", "supports-color 3.0.2", @@ -794,22 +799,11 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" -[[package]] -name = "pest" -version = "2.7.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b7cafe60d6cf8e62e1b9b2ea516a089c008945bb5a275416789e7db0bc199dc" -dependencies = [ - "memchr", - "thiserror 2.0.11", - "ucd-trie", -] - [[package]] name = "petgraph" -version = "0.6.5" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db" +checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" dependencies = [ "fixedbitset", "indexmap", @@ -817,9 +811,9 @@ dependencies = [ [[package]] name = "pgrx" -version = "0.12.9" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "227bf7e162ce710994306a97bc56bb3fe305f21120ab6692e2151c48416f5c0d" +checksum = "74fb2c1fbb9edc39097200f882ee58550e68f519db52309e94d894a2ad1ddf07" dependencies = [ "atomic-traits", "bitflags", @@ -835,15 +829,15 @@ dependencies = [ "serde", "serde_cbor", "serde_json", - "thiserror 1.0.69", + "thiserror", "uuid", ] [[package]] name = "pgrx-bindgen" -version = "0.12.9" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81cbcd956c2da35baaf0a116e6f6a49a6c2fbc8f6b332f66d6fd060bfd00615f" +checksum = "7bba29f6257193d1795eee0f96e72d368b30e95ab50d93397b5a76cc7784ddef" dependencies = [ "bindgen", "cc", @@ -859,9 +853,9 @@ dependencies = [ [[package]] name = "pgrx-catalog" -version = "0.1.0" +version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38fc7bbd6334277565c322bf9d48c80bb8858cad493ec29e337df39dc259fc6a" +checksum = "494cd323b4bb0b2befdcdecde2db9d56c3a6bb90b4a6bcfd8f40a747beb39fd1" dependencies = [ "paste", "pgrx", @@ -869,9 +863,9 @@ dependencies = [ [[package]] name = "pgrx-macros" -version = "0.12.9" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2f4291450d65e4deb770ce57ea93e22353d97950566222429cd166ebdf6f938" +checksum = "f1f9d2a6a94d8ea1a722ea489048acd203c6ba9b62f70edac587dbe8e1fb4585" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", @@ -881,9 +875,9 @@ dependencies = [ [[package]] name = "pgrx-pg-config" -version = "0.12.9" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "86a64a4c6e4e43e73cf8d3379d9533df98ded45c920e1ba8131c979633d74132" +checksum = "05be49ab66e458f4c164b45daa885e53dbbadc4f37bb47e8b5d85d5b4d9a0be4" dependencies = [ "cargo_toml", "eyre", @@ -892,16 +886,16 @@ dependencies = [ "pathsearch", "serde", "serde_json", - "thiserror 1.0.69", + "thiserror", "toml", "url", ] [[package]] name = "pgrx-pg-sys" -version = "0.12.9" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63a5dc64f2a8226434118aa2c4700450fa42b04f29488ad98268848b21c1a4ec" +checksum = "a360ad31a947674a46e72a39881f1be8afe958157ae44cf8ad6ed70b984c0f40" dependencies = [ "cee-scape", "libc", @@ -914,9 +908,9 @@ dependencies = [ [[package]] name = "pgrx-sql-entity-graph" -version = "0.12.9" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d81cc2e851c7e36b2f47c03e22d64d56c1d0e762fbde0039ba2cd490cfef3615" +checksum = "c8f625bb35d59b2d6f308ed5d1b6da3783c8514d47f353ec23a8f05da0fef6f4" dependencies = [ "convert_case", "eyre", @@ -924,17 +918,17 @@ dependencies = [ "proc-macro2", "quote", "syn", - "thiserror 1.0.69", + "thiserror", "unescape", ] [[package]] name = "ppv-lite86" -version = "0.2.20" +version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" dependencies = [ - "zerocopy 0.7.35", + "zerocopy", ] [[package]] @@ -961,22 +955,28 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.93" +version = "1.0.94" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60946a68e5f9d28b0dc1c21bb8a97ee7d018a8b322fa57838ba31cc878e22d99" +checksum = "a31971752e70b8b2686d7e46ec17fb38dad4051d94024c88df49b667caea9c84" dependencies = [ "unicode-ident", ] [[package]] name = "quote" -version = "1.0.38" +version = "1.0.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e4dccaaaf89514f546c693ddc140f729f958c247918a13380cccc6078391acc" +checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74765f6d916ee2faa39bc8e68e4f3ed8949b48cccdac59983d287a7cb71ce9c5" + [[package]] name = "rabitq" version = "0.0.0" @@ -999,7 +999,7 @@ checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" dependencies = [ "rand_chacha", "rand_core", - "zerocopy 0.8.17", + "zerocopy", ] [[package]] @@ -1014,19 +1014,18 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.9.0" +version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b08f3c9802962f7e1b25113931d94f43ed9725bebc59db9d0c3e9a23b67e15ff" +checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" dependencies = [ "getrandom", - "zerocopy 0.8.17", ] [[package]] name = "rand_distr" -version = "0.5.0" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddc3b5afe4c995c44540865b8ca5c52e6a59fa362da96c5d30886930ddc8da1c" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" dependencies = [ "num-traits", "rand", @@ -1099,24 +1098,24 @@ checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" [[package]] name = "rustc-hash" -version = "1.1.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" [[package]] name = "rustc_version" -version = "0.3.3" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0dfe2087c51c460008730de8b57e6a320782fbfb312e1f4d520e6c6fae155ee" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" dependencies = [ "semver", ] [[package]] name = "ryu" -version = "1.0.19" +version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ea1a2d0a644769cc99faa24c3ad26b379b786fe7c36fd3c546254801650e6dd" +checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" [[package]] name = "safe_arch" @@ -1144,27 +1143,15 @@ checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" [[package]] name = "semver" -version = "0.11.0" +version = "1.0.26" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f301af10236f6df4160f7c3f04eec6dbc70ace82d23326abad5edee88801c6b6" -dependencies = [ - "semver-parser", -] - -[[package]] -name = "semver-parser" -version = "0.10.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9900206b54a3527fdc7b8a938bffd94a568bac4f4aa8113b209df75a09c0dec2" -dependencies = [ - "pest", -] +checksum = "56e6fa9c48d24d85fb3de5ad847117517440f6beceb7798af16b4a87d616b8d0" [[package]] name = "serde" -version = "1.0.217" +version = "1.0.219" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02fc4265df13d6fa1d00ecff087228cc0a2b5f3c0e87e258d8b94a156e984c70" +checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6" dependencies = [ "serde_derive", ] @@ -1181,9 +1168,9 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.217" +version = "1.0.219" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0" +checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00" dependencies = [ "proc-macro2", "quote", @@ -1192,9 +1179,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.138" +version = "1.0.140" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d434192e7da787e94a6ea7e9670b26a036d0ca41e0b7efb2676dd32bae872949" +checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373" dependencies = [ "itoa", "memchr", @@ -1235,10 +1222,10 @@ name = "simd" version = "0.0.0" dependencies = [ "cc", - "half 2.4.1", + "half 2.5.0", "rand", "simd_macros", - "zerocopy 0.8.17", + "zerocopy", ] [[package]] @@ -1252,9 +1239,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.13.2" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd" [[package]] name = "sptr" @@ -1295,9 +1282,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.98" +version = "2.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36147f1a48ae0ec2b5b3bc5b537d267457555a10dc06f3dbc8cb11ba3006d3b1" +checksum = "b09a44accad81e1ba1cd74a32461ba89dee89095ba17b32f5d03683b1b1fc2a0" dependencies = [ "proc-macro2", "quote", @@ -1323,38 +1310,18 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" [[package]] name = "thiserror" -version = "1.0.69" +version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" dependencies = [ - "thiserror-impl 1.0.69", -] - -[[package]] -name = "thiserror" -version = "2.0.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d452f284b73e6d76dd36758a0c8684b1d5be31f92b89d07fd5822175732206fc" -dependencies = [ - "thiserror-impl 2.0.11", + "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.69" +version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "thiserror-impl" -version = "2.0.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26afc1baea8a989337eeb52b6e72a039780ce45c3edfcc9c5b9d112feeb173c2" +checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d" dependencies = [ "proc-macro2", "quote", @@ -1412,15 +1379,9 @@ source = "git+https://github.com/tensorchord/turboselect.git?rev=d8753c4ffe5b47f [[package]] name = "typenum" -version = "1.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" - -[[package]] -name = "ucd-trie" -version = "0.1.7" +version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" +checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f" [[package]] name = "unescape" @@ -1430,9 +1391,9 @@ checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] name = "unicode-ident" -version = "1.0.16" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a210d160f08b701c8721ba1c726c11662f877ea6b7094007e1ca9a1041945034" +checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" [[package]] name = "unicode-segmentation" @@ -1442,9 +1403,9 @@ checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" [[package]] name = "unicode-width" -version = "0.1.14" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" +checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd" [[package]] name = "url" @@ -1471,9 +1432,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.13.1" +version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ced87ca4be083373936a67f8de945faa23b6b42384bd5b64434850802c6dccd0" +checksum = "458f7a779bf54acc9f347480ac654f68407d3aab21269a6e3c9f922acd9e2da9" dependencies = [ "getrandom", ] @@ -1515,7 +1476,7 @@ dependencies = [ "algorithm", "always_equal", "distance", - "half 2.4.1", + "half 2.5.0", "k_means", "paste", "pgrx", @@ -1527,8 +1488,8 @@ dependencies = [ "toml", "validator", "vector", - "zerocopy 0.8.17", - "zerocopy-derive 0.8.17", + "zerocopy", + "zerocopy-derive", ] [[package]] @@ -1536,7 +1497,7 @@ name = "vector" version = "0.0.0" dependencies = [ "distance", - "half 2.4.1", + "half 2.5.0", "simd", ] @@ -1552,9 +1513,9 @@ dependencies = [ [[package]] name = "wasi" -version = "0.13.3+wasi-0.2.2" +version = "0.14.2+wasi-0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26816d2e1a4a36a2940b96c5296ce403917633dff8f3440e9b236ed6f6bacad2" +checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3" dependencies = [ "wit-bindgen-rt", ] @@ -1569,22 +1530,6 @@ dependencies = [ "safe_arch", ] -[[package]] -name = "winapi" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" -dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", -] - -[[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - [[package]] name = "winapi-util" version = "0.1.9" @@ -1594,12 +1539,6 @@ dependencies = [ "windows-sys", ] -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - [[package]] name = "windows-sys" version = "0.59.0" @@ -1675,18 +1614,18 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "winnow" -version = "0.7.2" +version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59690dea168f2198d1a3b0cac23b8063efcd11012f10ae4698f284808c8ef603" +checksum = "0e97b544156e9bebe1a0ffbc03484fc1ffe3100cbce3ffb17eac35f7cdd7ab36" dependencies = [ "memchr", ] [[package]] name = "wit-bindgen-rt" -version = "0.33.0" +version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c" +checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" dependencies = [ "bitflags", ] @@ -1712,15 +1651,6 @@ dependencies = [ "tap", ] -[[package]] -name = "yansi-term" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe5c30ade05e61656247b2e334a031dfd0cc466fadef865bdcdea8d537951bf1" -dependencies = [ - "winapi", -] - [[package]] name = "yoke" version = "0.7.5" @@ -1747,39 +1677,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.7.35" +version = "0.8.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +checksum = "fd97444d05a4328b90e75e503a34bad781f14e28a823ad3557f0750df1ebcbc6" dependencies = [ - "byteorder", - "zerocopy-derive 0.7.35", -] - -[[package]] -name = "zerocopy" -version = "0.8.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa91407dacce3a68c56de03abe2760159582b846c6a4acd2f456618087f12713" -dependencies = [ - "zerocopy-derive 0.8.17", -] - -[[package]] -name = "zerocopy-derive" -version = "0.7.35" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" -dependencies = [ - "proc-macro2", - "quote", - "syn", + "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.17" +version = "0.8.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06718a168365cad3d5ff0bb133aad346959a2074bd4a85c121255a11304a8626" +checksum = "6352c01d0edd5db859a63e2605f4ea3183ddbd15e2c4a9e7d32184df75e4f154" dependencies = [ "proc-macro2", "quote", @@ -1788,18 +1697,18 @@ dependencies = [ [[package]] name = "zerofrom" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cff3ee08c995dee1859d998dea82f7374f2826091dd9cd47def953cae446cd2e" +checksum = "50cc42e0333e05660c3587f3bf9d0478688e15d870fab3346451ce7f8c9fbea5" dependencies = [ "zerofrom-derive", ] [[package]] name = "zerofrom-derive" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808" +checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index e9274669..234cd1f5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "vchord" version.workspace = true -edition = "2021" +edition.workspace = true [lib] name = "vchord" @@ -30,8 +30,8 @@ vector = { path = "./crates/vector" } half.workspace = true paste.workspace = true -pgrx = { version = "=0.12.9", default-features = false, features = ["cshim"] } -pgrx-catalog = "0.1.0" +pgrx = { version = "=0.13.1", default-features = false, features = ["cshim"] } +pgrx-catalog = "0.1.1" rand.workspace = true serde.workspace = true toml = "0.8.20" @@ -40,7 +40,7 @@ zerocopy.workspace = true zerocopy-derive.workspace = true [patch.crates-io] -half = { git = "https://github.com/tensorchord/half-rs.git", rev = "3f9a8843d6722bd1833de2289347640ad8770146" } +half = { git = "https://github.com/tensorchord/half-rs.git", rev = "9dc2559e95f295aaa8ffef5e0529167731dffd47" } [lints] workspace = true @@ -51,16 +51,16 @@ members = ["crates/*"] [workspace.package] version = "0.0.0" -edition = "2021" +edition = "2024" [workspace.dependencies] -half = { version = "2.4.1", features = ["serde", "zerocopy"] } +half = { version = "2.5.0", features = ["serde", "zerocopy"] } paste = "1" rand = "0.9.0" serde = "1" validator = { version = "0.20.0", features = ["derive"] } -zerocopy = "0.8.17" -zerocopy-derive = "0.8.17" +zerocopy = "0.8.23" +zerocopy-derive = "0.8.23" [workspace.lints] clippy.identity_op = "allow" diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index ffd64440..a0fd1b46 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -3,7 +3,6 @@ use crate::operator::Vector; use std::marker::PhantomData; use std::num::{NonZeroU8, NonZeroU64}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; -use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; diff --git a/crates/random_orthogonal_matrix/Cargo.toml b/crates/random_orthogonal_matrix/Cargo.toml index 4e8f733c..348ecb7f 100644 --- a/crates/random_orthogonal_matrix/Cargo.toml +++ b/crates/random_orthogonal_matrix/Cargo.toml @@ -9,7 +9,7 @@ nalgebra = "=0.33.0" rand.workspace = true rand_chacha = "0.9.0" -rand_distr = "0.5.0" +rand_distr = "0.5.1" [lints] workspace = true diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index bed14a33..bb421562 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -13,7 +13,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.13" +cc = "1.2.16" [lints] workspace = true diff --git a/crates/simd/build.rs b/crates/simd/build.rs index c7a49d65..10fad2c6 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -1,12 +1,7 @@ fn main() { println!("cargo::rerun-if-changed=cshim.c"); - cc::Build::new() - .compiler("clang") - .file("cshim.c") - .opt_level(3) - .flag("-fassociative-math") - .flag("-ffp-contract=fast") - .flag("-freciprocal-math") - .flag("-fno-signed-zeros") - .compile("simd_cshim"); + let mut build = cc::Build::new(); + build.file("cshim.c"); + build.opt_level(3); + build.compile("simd_cshim"); } diff --git a/crates/simd/cshim.c b/crates/simd/cshim.c index 7c0f044a..da02e8f4 100644 --- a/crates/simd/cshim.c +++ b/crates/simd/cshim.c @@ -1,5 +1,13 @@ +#if defined(__clang__) #if !(__clang_major__ >= 16) -#error "clang version must be >= 16" +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." #endif #ifdef __aarch64__ @@ -12,7 +20,7 @@ typedef __fp16 f16; typedef float f32; -__attribute__((target("fp16"))) float +__attribute__((target("+fp16"))) float fp16_reduce_sum_of_xy_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { float16x8_t xy_0 = vdupq_n_f16(0.0); float16x8_t xy_1 = vdupq_n_f16(0.0); @@ -61,7 +69,7 @@ fp16_reduce_sum_of_xy_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { return vaddvq_f32(lo) + vaddvq_f32(hi); } -__attribute__((target("sve"))) float +__attribute__((target("+sve"))) float fp16_reduce_sum_of_xy_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svfloat16_t xy = svdup_f16(0.0); for (size_t i = 0; i < n; i += svcnth()) { @@ -73,7 +81,7 @@ fp16_reduce_sum_of_xy_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { return svaddv_f16(svptrue_b16(), xy); } -__attribute__((target("fp16"))) float +__attribute__((target("+fp16"))) float fp16_reduce_sum_of_d2_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { float16x8_t d2_0 = vdupq_n_f16(0.0); float16x8_t d2_1 = vdupq_n_f16(0.0); @@ -130,7 +138,7 @@ fp16_reduce_sum_of_d2_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { return vaddvq_f32(lo) + vaddvq_f32(hi); } -__attribute__((target("sve"))) float +__attribute__((target("+sve"))) float fp16_reduce_sum_of_d2_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svfloat16_t d2 = svdup_f16(0.0); for (size_t i = 0; i < n; i += svcnth()) { @@ -143,7 +151,7 @@ fp16_reduce_sum_of_d2_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { return svaddv_f16(svptrue_b16(), d2); } -__attribute__((target("sve"))) float +__attribute__((target("+sve"))) float fp32_reduce_sum_of_x_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { @@ -154,7 +162,7 @@ fp32_reduce_sum_of_x_a3_256(float *restrict this, size_t n) { return svaddv_f32(svptrue_b32(), sum); } -__attribute__((target("sve"))) float +__attribute__((target("+sve"))) float fp32_reduce_sum_of_abs_x_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { @@ -165,7 +173,7 @@ fp32_reduce_sum_of_abs_x_a3_256(float *restrict this, size_t n) { return svaddv_f32(svptrue_b32(), sum); } -__attribute__((target("sve"))) float +__attribute__((target("+sve"))) float fp32_reduce_sum_of_x2_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { @@ -176,7 +184,7 @@ fp32_reduce_sum_of_x2_a3_256(float *restrict this, size_t n) { return svaddv_f32(svptrue_b32(), sum); } -__attribute__((target("sve"))) void +__attribute__((target("+sve"))) void fp32_reduce_min_max_of_x_a3_256(float *restrict this, size_t n, float *out_min, float *out_max) { svfloat32_t min = svdup_f32(1.0 / 0.0); @@ -191,7 +199,7 @@ fp32_reduce_min_max_of_x_a3_256(float *restrict this, size_t n, float *out_min, *out_max = svmaxv_f32(svptrue_b32(), max); } -__attribute__((target("sve"))) float +__attribute__((target("+sve"))) float fp32_reduce_sum_of_xy_a3_256(float *restrict lhs, float *restrict rhs, size_t n) { svfloat32_t sum = svdup_f32(0.0); @@ -204,7 +212,7 @@ fp32_reduce_sum_of_xy_a3_256(float *restrict lhs, float *restrict rhs, return svaddv_f32(svptrue_b32(), sum); } -__attribute__((target("sve"))) float +__attribute__((target("+sve"))) float fp32_reduce_sum_of_d2_a3_256(float *restrict lhs, float *restrict rhs, size_t n) { svfloat32_t sum = svdup_f32(0.0); diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index d8321dc4..67e09e34 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -12,28 +12,26 @@ mod sum_of_and { #[target_feature(enable = "avx512vpopcntdq")] fn sum_of_and_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut and = _mm512_setzero_si512(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut n = lhs.len(); - while n >= 8 { - let x = _mm512_loadu_si512(a.cast()); - let y = _mm512_loadu_si512(b.cast()); - a = a.add(8); - b = b.add(8); - n -= 8; - and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); - } - if n > 0 { - let mask = _bzhi_u32(0xff, n as u32) as u8; - let x = _mm512_maskz_loadu_epi64(mask, a.cast()); - let y = _mm512_maskz_loadu_epi64(mask, b.cast()); - and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); - } - _mm512_reduce_add_epi64(and) as u32 + use std::arch::x86_64::*; + let mut and = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + } + _mm512_reduce_add_epi64(and) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -81,28 +79,26 @@ mod sum_of_or { #[target_feature(enable = "avx512vpopcntdq")] fn sum_of_or_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut or = _mm512_setzero_si512(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut n = lhs.len(); - while n >= 8 { - let x = _mm512_loadu_si512(a.cast()); - let y = _mm512_loadu_si512(b.cast()); - a = a.add(8); - b = b.add(8); - n -= 8; - or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); - } - if n > 0 { - let mask = _bzhi_u32(0xff, n as u32) as u8; - let x = _mm512_maskz_loadu_epi64(mask, a.cast()); - let y = _mm512_maskz_loadu_epi64(mask, b.cast()); - or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); - } - _mm512_reduce_add_epi64(or) as u32 + use std::arch::x86_64::*; + let mut or = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); } + _mm512_reduce_add_epi64(or) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -150,28 +146,26 @@ mod sum_of_xor { #[target_feature(enable = "avx512vpopcntdq")] fn sum_of_xor_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut xor = _mm512_setzero_si512(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut n = lhs.len(); - while n >= 8 { - let x = _mm512_loadu_si512(a.cast()); - let y = _mm512_loadu_si512(b.cast()); - a = a.add(8); - b = b.add(8); - n -= 8; - xor = _mm512_add_epi64(xor, _mm512_popcnt_epi64(_mm512_xor_si512(x, y))); - } - if n > 0 { - let mask = _bzhi_u32(0xff, n as u32) as u8; - let x = _mm512_maskz_loadu_epi64(mask, a.cast()); - let y = _mm512_maskz_loadu_epi64(mask, b.cast()); - xor = _mm512_add_epi64(xor, _mm512_popcnt_epi64(_mm512_xor_si512(x, y))); - } - _mm512_reduce_add_epi64(xor) as u32 + use std::arch::x86_64::*; + let mut xor = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + xor = _mm512_add_epi64(xor, _mm512_popcnt_epi64(_mm512_xor_si512(x, y))); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + xor = _mm512_add_epi64(xor, _mm512_popcnt_epi64(_mm512_xor_si512(x, y))); } + _mm512_reduce_add_epi64(xor) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -219,34 +213,32 @@ mod sum_of_and_or { #[target_feature(enable = "avx512vpopcntdq")] fn sum_of_and_or_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut and = _mm512_setzero_si512(); - let mut or = _mm512_setzero_si512(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut n = lhs.len(); - while n >= 8 { - let x = _mm512_loadu_si512(a.cast()); - let y = _mm512_loadu_si512(b.cast()); - a = a.add(8); - b = b.add(8); - n -= 8; - and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); - or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); - } - if n > 0 { - let mask = _bzhi_u32(0xff, n as u32) as u8; - let x = _mm512_maskz_loadu_epi64(mask, a.cast()); - let y = _mm512_maskz_loadu_epi64(mask, b.cast()); - and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); - or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); - } - ( - _mm512_reduce_add_epi64(and) as u32, - _mm512_reduce_add_epi64(or) as u32, - ) + use std::arch::x86_64::*; + let mut and = _mm512_setzero_si512(); + let mut or = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + } + ( + _mm512_reduce_add_epi64(and) as u32, + _mm512_reduce_add_epi64(or) as u32, + ) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -295,24 +287,22 @@ mod sum_of_x { #[crate::target_cpu(enable = "v4.512")] #[target_feature(enable = "avx512vpopcntdq")] fn sum_of_x_v4_512_avx512vpopcntdq(this: &[u64]) -> u32 { - unsafe { - use std::arch::x86_64::*; - let mut and = _mm512_setzero_si512(); - let mut a = this.as_ptr(); - let mut n = this.len(); - while n >= 8 { - let x = _mm512_loadu_si512(a.cast()); - a = a.add(8); - n -= 8; - and = _mm512_add_epi64(and, _mm512_popcnt_epi64(x)); - } - if n > 0 { - let mask = _bzhi_u32(0xff, n as u32) as u8; - let x = _mm512_maskz_loadu_epi64(mask, a.cast()); - and = _mm512_add_epi64(and, _mm512_popcnt_epi64(x)); - } - _mm512_reduce_add_epi64(and) as u32 + use std::arch::x86_64::*; + let mut and = _mm512_setzero_si512(); + let mut a = this.as_ptr(); + let mut n = this.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + a = unsafe { a.add(8) }; + n -= 8; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(x)); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(x)); } + _mm512_reduce_add_epi64(and) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index 293aa746..4a3ebe68 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -8,142 +8,126 @@ pub fn emulate_mm512_2intersect_epi32( a: std::arch::x86_64::__m512i, b: std::arch::x86_64::__m512i, ) -> (std::arch::x86_64::__mmask16, std::arch::x86_64::__mmask16) { - unsafe { - use std::arch::x86_64::*; + use std::arch::x86_64::*; - let a1 = _mm512_alignr_epi32(a, a, 4); - let a2 = _mm512_alignr_epi32(a, a, 8); - let a3 = _mm512_alignr_epi32(a, a, 12); - let b1 = _mm512_shuffle_epi32(b, _MM_PERM_ADCB); - let b2 = _mm512_shuffle_epi32(b, _MM_PERM_BADC); - let b3 = _mm512_shuffle_epi32(b, _MM_PERM_CBAD); - let m00 = _mm512_cmpeq_epi32_mask(a, b); - let m01 = _mm512_cmpeq_epi32_mask(a, b1); - let m02 = _mm512_cmpeq_epi32_mask(a, b2); - let m03 = _mm512_cmpeq_epi32_mask(a, b3); - let m10 = _mm512_cmpeq_epi32_mask(a1, b); - let m11 = _mm512_cmpeq_epi32_mask(a1, b1); - let m12 = _mm512_cmpeq_epi32_mask(a1, b2); - let m13 = _mm512_cmpeq_epi32_mask(a1, b3); - let m20 = _mm512_cmpeq_epi32_mask(a2, b); - let m21 = _mm512_cmpeq_epi32_mask(a2, b1); - let m22 = _mm512_cmpeq_epi32_mask(a2, b2); - let m23 = _mm512_cmpeq_epi32_mask(a2, b3); - let m30 = _mm512_cmpeq_epi32_mask(a3, b); - let m31 = _mm512_cmpeq_epi32_mask(a3, b1); - let m32 = _mm512_cmpeq_epi32_mask(a3, b2); - let m33 = _mm512_cmpeq_epi32_mask(a3, b3); + let a1 = _mm512_alignr_epi32(a, a, 4); + let a2 = _mm512_alignr_epi32(a, a, 8); + let a3 = _mm512_alignr_epi32(a, a, 12); + let b1 = _mm512_shuffle_epi32(b, _MM_PERM_ADCB); + let b2 = _mm512_shuffle_epi32(b, _MM_PERM_BADC); + let b3 = _mm512_shuffle_epi32(b, _MM_PERM_CBAD); + let m00 = _mm512_cmpeq_epi32_mask(a, b); + let m01 = _mm512_cmpeq_epi32_mask(a, b1); + let m02 = _mm512_cmpeq_epi32_mask(a, b2); + let m03 = _mm512_cmpeq_epi32_mask(a, b3); + let m10 = _mm512_cmpeq_epi32_mask(a1, b); + let m11 = _mm512_cmpeq_epi32_mask(a1, b1); + let m12 = _mm512_cmpeq_epi32_mask(a1, b2); + let m13 = _mm512_cmpeq_epi32_mask(a1, b3); + let m20 = _mm512_cmpeq_epi32_mask(a2, b); + let m21 = _mm512_cmpeq_epi32_mask(a2, b1); + let m22 = _mm512_cmpeq_epi32_mask(a2, b2); + let m23 = _mm512_cmpeq_epi32_mask(a2, b3); + let m30 = _mm512_cmpeq_epi32_mask(a3, b); + let m31 = _mm512_cmpeq_epi32_mask(a3, b1); + let m32 = _mm512_cmpeq_epi32_mask(a3, b2); + let m33 = _mm512_cmpeq_epi32_mask(a3, b3); - let m0 = m00 | m10 | m20 | m30; - let m1 = m01 | m11 | m21 | m31; - let m2 = m02 | m12 | m22 | m32; - let m3 = m03 | m13 | m23 | m33; + let m0 = m00 | m10 | m20 | m30; + let m1 = m01 | m11 | m21 | m31; + let m2 = m02 | m12 | m22 | m32; + let m3 = m03 | m13 | m23 | m33; - let res_a = m00 - | m01 - | m02 - | m03 - | (m10 | m11 | m12 | m13).rotate_left(4) - | (m20 | m21 | m22 | m23).rotate_left(8) - | (m30 | m31 | m32 | m33).rotate_right(4); + let res_a = m00 + | m01 + | m02 + | m03 + | (m10 | m11 | m12 | m13).rotate_left(4) + | (m20 | m21 | m22 | m23).rotate_left(8) + | (m30 | m31 | m32 | m33).rotate_right(4); - let res_b = m0 - | ((0x7777 & m1) << 1) - | ((m1 >> 3) & 0x1111) - | ((0x3333 & m2) << 2) - | ((m2 >> 2) & 0x3333) - | ((0x1111 & m3) << 3) - | ((m3 >> 1) & 0x7777); - (res_a, res_b) - } + let res_b = m0 + | ((0x7777 & m1) << 1) + | ((m1 >> 3) & 0x1111) + | ((0x3333 & m2) << 2) + | ((m2 >> 2) & 0x3333) + | ((0x1111 & m3) << 3) + | ((m3 >> 1) & 0x7777); + (res_a, res_b) } #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] pub fn emulate_mm256_reduce_add_ps(mut x: std::arch::x86_64::__m256) -> f32 { - unsafe { - use std::arch::x86_64::*; - x = _mm256_add_ps(x, _mm256_permute2f128_ps(x, x, 1)); - x = _mm256_hadd_ps(x, x); - x = _mm256_hadd_ps(x, x); - _mm256_cvtss_f32(x) - } + use std::arch::x86_64::*; + x = _mm256_add_ps(x, _mm256_permute2f128_ps(x, x, 1)); + x = _mm256_hadd_ps(x, x); + x = _mm256_hadd_ps(x, x); + _mm256_cvtss_f32(x) } #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] pub fn emulate_mm_reduce_add_ps(mut x: std::arch::x86_64::__m128) -> f32 { - unsafe { - use std::arch::x86_64::*; - x = _mm_hadd_ps(x, x); - x = _mm_hadd_ps(x, x); - _mm_cvtss_f32(x) - } + use std::arch::x86_64::*; + x = _mm_hadd_ps(x, x); + x = _mm_hadd_ps(x, x); + _mm_cvtss_f32(x) } #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4.512")] pub fn emulate_mm512_reduce_add_epi16(x: std::arch::x86_64::__m512i) -> i16 { - unsafe { - use std::arch::x86_64::*; - _mm256_reduce_add_epi16(_mm512_castsi512_si256(x)) - + _mm256_reduce_add_epi16(_mm512_extracti32x8_epi32(x, 1)) - } + use std::arch::x86_64::*; + _mm256_reduce_add_epi16(_mm512_castsi512_si256(x)) + + _mm256_reduce_add_epi16(_mm512_extracti32x8_epi32(x, 1)) } #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] pub fn emulate_mm256_reduce_add_epi16(mut x: std::arch::x86_64::__m256i) -> i16 { - unsafe { - use std::arch::x86_64::*; - x = _mm256_add_epi16(x, _mm256_permute2f128_si256(x, x, 1)); - x = _mm256_hadd_epi16(x, x); - x = _mm256_hadd_epi16(x, x); - let x = _mm256_cvtsi256_si32(x); - (x as i16) + ((x >> 16) as i16) - } + use std::arch::x86_64::*; + x = _mm256_add_epi16(x, _mm256_permute2f128_si256(x, x, 1)); + x = _mm256_hadd_epi16(x, x); + x = _mm256_hadd_epi16(x, x); + let x = _mm256_cvtsi256_si32(x); + (x as i16) + ((x >> 16) as i16) } #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] pub fn emulate_mm_reduce_add_epi16(mut x: std::arch::x86_64::__m128i) -> i16 { - unsafe { - use std::arch::x86_64::*; - x = _mm_hadd_epi16(x, x); - x = _mm_hadd_epi16(x, x); - let x = _mm_cvtsi128_si32(x); - (x as i16) + ((x >> 16) as i16) - } + use std::arch::x86_64::*; + x = _mm_hadd_epi16(x, x); + x = _mm_hadd_epi16(x, x); + let x = _mm_cvtsi128_si32(x); + (x as i16) + ((x >> 16) as i16) } #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] pub fn emulate_mm256_reduce_add_epi32(mut x: std::arch::x86_64::__m256i) -> i32 { - unsafe { - use std::arch::x86_64::*; - x = _mm256_add_epi32(x, _mm256_permute2f128_si256(x, x, 1)); - x = _mm256_hadd_epi32(x, x); - x = _mm256_hadd_epi32(x, x); - _mm256_cvtsi256_si32(x) - } + use std::arch::x86_64::*; + x = _mm256_add_epi32(x, _mm256_permute2f128_si256(x, x, 1)); + x = _mm256_hadd_epi32(x, x); + x = _mm256_hadd_epi32(x, x); + _mm256_cvtsi256_si32(x) } #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] pub fn emulate_mm_reduce_add_epi32(mut x: std::arch::x86_64::__m128i) -> i32 { - unsafe { - use std::arch::x86_64::*; - x = _mm_hadd_epi32(x, x); - x = _mm_hadd_epi32(x, x); - _mm_cvtsi128_si32(x) - } + use std::arch::x86_64::*; + x = _mm_hadd_epi32(x, x); + x = _mm_hadd_epi32(x, x); + _mm_cvtsi128_si32(x) } #[inline] @@ -151,15 +135,15 @@ pub fn emulate_mm_reduce_add_epi32(mut x: std::arch::x86_64::__m128i) -> i32 { #[crate::target_cpu(enable = "v3")] pub fn emulate_mm256_reduce_min_ps(x: std::arch::x86_64::__m256) -> f32 { use crate::aligned::Aligned16; + use std::arch::x86_64::*; + let lo = _mm256_castps256_ps128(x); + let hi = _mm256_extractf128_ps(x, 1); + let min = _mm_min_ps(lo, hi); + let mut x = Aligned16([0.0f32; 4]); unsafe { - use std::arch::x86_64::*; - let lo = _mm256_castps256_ps128(x); - let hi = _mm256_extractf128_ps(x, 1); - let min = _mm_min_ps(lo, hi); - let mut x = Aligned16([0.0f32; 4]); _mm_store_ps(x.0.as_mut_ptr(), min); - f32::min(f32::min(x.0[0], x.0[1]), f32::min(x.0[2], x.0[3])) } + f32::min(f32::min(x.0[0], x.0[1]), f32::min(x.0[2], x.0[3])) } #[inline] @@ -167,13 +151,13 @@ pub fn emulate_mm256_reduce_min_ps(x: std::arch::x86_64::__m256) -> f32 { #[crate::target_cpu(enable = "v2")] pub fn emulate_mm_reduce_min_ps(x: std::arch::x86_64::__m128) -> f32 { use crate::aligned::Aligned16; + use std::arch::x86_64::*; + let min = x; + let mut x = Aligned16([0.0f32; 4]); unsafe { - use std::arch::x86_64::*; - let min = x; - let mut x = Aligned16([0.0f32; 4]); _mm_store_ps(x.0.as_mut_ptr(), min); - f32::min(f32::min(x.0[0], x.0[1]), f32::min(x.0[2], x.0[3])) } + f32::min(f32::min(x.0[0], x.0[1]), f32::min(x.0[2], x.0[3])) } #[inline] @@ -181,15 +165,15 @@ pub fn emulate_mm_reduce_min_ps(x: std::arch::x86_64::__m128) -> f32 { #[crate::target_cpu(enable = "v3")] pub fn emulate_mm256_reduce_max_ps(x: std::arch::x86_64::__m256) -> f32 { use crate::aligned::Aligned16; + use std::arch::x86_64::*; + let lo = _mm256_castps256_ps128(x); + let hi = _mm256_extractf128_ps(x, 1); + let max = _mm_max_ps(lo, hi); + let mut x = Aligned16([0.0f32; 4]); unsafe { - use std::arch::x86_64::*; - let lo = _mm256_castps256_ps128(x); - let hi = _mm256_extractf128_ps(x, 1); - let max = _mm_max_ps(lo, hi); - let mut x = Aligned16([0.0f32; 4]); _mm_store_ps(x.0.as_mut_ptr(), max); - f32::max(f32::max(x.0[0], x.0[1]), f32::max(x.0[2], x.0[3])) } + f32::max(f32::max(x.0[0], x.0[1]), f32::max(x.0[2], x.0[3])) } #[inline] @@ -197,11 +181,11 @@ pub fn emulate_mm256_reduce_max_ps(x: std::arch::x86_64::__m256) -> f32 { #[crate::target_cpu(enable = "v2")] pub fn emulate_mm_reduce_max_ps(x: std::arch::x86_64::__m128) -> f32 { use crate::aligned::Aligned16; + use std::arch::x86_64::*; + let max = x; + let mut x = Aligned16([0.0f32; 4]); unsafe { - use std::arch::x86_64::*; - let max = x; - let mut x = Aligned16([0.0f32; 4]); _mm_store_ps(x.0.as_mut_ptr(), max); - f32::max(f32::max(x.0[0], x.0[1]), f32::max(x.0[2], x.0[3])) } + f32::max(f32::max(x.0[0], x.0[1]), f32::max(x.0[2], x.0[3])) } diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 60e4d045..d385fd40 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -217,28 +217,26 @@ mod reduce_sum_of_xy { #[target_feature(enable = "avx512fp16")] pub fn reduce_sum_of_xy_v4_512_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut xy = _mm512_setzero_ph(); - while n >= 32 { - let x = _mm512_loadu_ph(a.cast()); - let y = _mm512_loadu_ph(b.cast()); - a = a.add(32); - b = b.add(32); - n -= 32; - xy = _mm512_fmadd_ph(x, y, xy); - } - if n > 0 { - let mask = _bzhi_u32(0xffffffff, n as u32); - let x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())); - let y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())); - xy = _mm512_fmadd_ph(x, y, xy); - } - _mm512_reduce_add_ph(xy) as f32 - } + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm512_setzero_ph(); + while n >= 32 { + let x = unsafe { _mm512_loadu_ph(a.cast()) }; + let y = unsafe { _mm512_loadu_ph(b.cast()) }; + a = unsafe { a.add(32) }; + b = unsafe { b.add(32) }; + n -= 32; + xy = _mm512_fmadd_ph(x, y, xy); + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + let y = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())) }; + xy = _mm512_fmadd_ph(x, y, xy); + } + _mm512_reduce_add_ph(xy) as f32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -277,28 +275,26 @@ mod reduce_sum_of_xy { #[crate::target_cpu(enable = "v4.512")] pub fn reduce_sum_of_xy_v4_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut xy = _mm512_setzero_ps(); - while n >= 16 { - let x = _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())); - let y = _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())); - a = a.add(16); - b = b.add(16); - n -= 16; - xy = _mm512_fmadd_ps(x, y, xy); - } - if n > 0 { - let mask = _bzhi_u32(0xffff, n as u32) as u16; - let x = _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())); - let y = _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())); - xy = _mm512_fmadd_ps(x, y, xy); - } - _mm512_reduce_add_ps(xy) - } + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + let y = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())) }; + a = unsafe { a.add(16) }; + b = unsafe { b.add(16) }; + n -= 16; + xy = _mm512_fmadd_ps(x, y, xy); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + let y = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())) }; + xy = _mm512_fmadd_ps(x, y, xy); + } + _mm512_reduce_add_ps(xy) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -334,31 +330,29 @@ mod reduce_sum_of_xy { pub fn reduce_sum_of_xy_v3(lhs: &[f16], rhs: &[f16]) -> f32 { use crate::emulate::emulate_mm256_reduce_add_ps; assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut xy = _mm256_setzero_ps(); - while n >= 8 { - let x = _mm256_cvtph_ps(_mm_loadu_si128(a.cast())); - let y = _mm256_cvtph_ps(_mm_loadu_si128(b.cast())); - a = a.add(8); - b = b.add(8); - n -= 8; - xy = _mm256_fmadd_ps(x, y, xy); - } - let mut xy = emulate_mm256_reduce_add_ps(xy); - while n > 0 { - let x = a.read().to_f32(); - let y = b.read().to_f32(); - a = a.add(1); - b = b.add(1); - n -= 1; - xy += x * y; - } - xy + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + xy = _mm256_fmadd_ps(x, y, xy); + } + let mut xy = emulate_mm256_reduce_add_ps(xy); + while n > 0 { + let x = unsafe { a.read().to_f32() }; + let y = unsafe { b.read().to_f32() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + xy += x * y; } + xy } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -399,8 +393,12 @@ mod reduce_sum_of_xy { pub fn reduce_sum_of_xy_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { - extern "C" { - fn fp16_reduce_sum_of_xy_a2_fp16(a: *const (), b: *const (), n: usize) -> f32; + unsafe extern "C" { + unsafe fn fp16_reduce_sum_of_xy_a2_fp16( + a: *const (), + b: *const (), + n: usize, + ) -> f32; } fp16_reduce_sum_of_xy_a2_fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } @@ -444,8 +442,9 @@ mod reduce_sum_of_xy { pub fn reduce_sum_of_xy_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { - extern "C" { - fn fp16_reduce_sum_of_xy_a3_512(a: *const (), b: *const (), n: usize) -> f32; + unsafe extern "C" { + unsafe fn fp16_reduce_sum_of_xy_a3_512(a: *const (), b: *const (), n: usize) + -> f32; } fp16_reduce_sum_of_xy_a3_512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } @@ -503,30 +502,28 @@ mod reduce_sum_of_d2 { #[target_feature(enable = "avx512fp16")] pub fn reduce_sum_of_d2_v4_512_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len() as u32; - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut d2 = _mm512_setzero_ph(); - while n >= 32 { - let x = _mm512_loadu_ph(a.cast()); - let y = _mm512_loadu_ph(b.cast()); - a = a.add(32); - b = b.add(32); - n -= 32; - let d = _mm512_sub_ph(x, y); - d2 = _mm512_fmadd_ph(d, d, d2); - } - if n > 0 { - let mask = _bzhi_u32(0xffffffff, n); - let x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())); - let y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())); - let d = _mm512_sub_ph(x, y); - d2 = _mm512_fmadd_ph(d, d, d2); - } - _mm512_reduce_add_ph(d2) as f32 - } + use std::arch::x86_64::*; + let mut n = lhs.len() as u32; + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm512_setzero_ph(); + while n >= 32 { + let x = unsafe { _mm512_loadu_ph(a.cast()) }; + let y = unsafe { _mm512_loadu_ph(b.cast()) }; + a = unsafe { a.add(32) }; + b = unsafe { b.add(32) }; + n -= 32; + let d = _mm512_sub_ph(x, y); + d2 = _mm512_fmadd_ph(d, d, d2); + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n); + let x = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + let y = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())) }; + let d = _mm512_sub_ph(x, y); + d2 = _mm512_fmadd_ph(d, d, d2); + } + _mm512_reduce_add_ph(d2) as f32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -565,30 +562,28 @@ mod reduce_sum_of_d2 { #[crate::target_cpu(enable = "v4.512")] pub fn reduce_sum_of_d2_v4_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len() as u32; - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut d2 = _mm512_setzero_ps(); - while n >= 16 { - let x = _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())); - let y = _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())); - a = a.add(16); - b = b.add(16); - n -= 16; - let d = _mm512_sub_ps(x, y); - d2 = _mm512_fmadd_ps(d, d, d2); - } - if n > 0 { - let mask = _bzhi_u32(0xffff, n) as u16; - let x = _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())); - let y = _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())); - let d = _mm512_sub_ps(x, y); - d2 = _mm512_fmadd_ps(d, d, d2); - } - _mm512_reduce_add_ps(d2) - } + use std::arch::x86_64::*; + let mut n = lhs.len() as u32; + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + let y = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())) }; + a = unsafe { a.add(16) }; + b = unsafe { b.add(16) }; + n -= 16; + let d = _mm512_sub_ps(x, y); + d2 = _mm512_fmadd_ps(d, d, d2); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + let y = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())) }; + let d = _mm512_sub_ps(x, y); + d2 = _mm512_fmadd_ps(d, d, d2); + } + _mm512_reduce_add_ps(d2) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -628,33 +623,31 @@ mod reduce_sum_of_d2 { pub fn reduce_sum_of_d2_v3(lhs: &[f16], rhs: &[f16]) -> f32 { use crate::emulate::emulate_mm256_reduce_add_ps; assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len() as u32; - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut d2 = _mm256_setzero_ps(); - while n >= 8 { - let x = _mm256_cvtph_ps(_mm_loadu_si128(a.cast())); - let y = _mm256_cvtph_ps(_mm_loadu_si128(b.cast())); - a = a.add(8); - b = b.add(8); - n -= 8; - let d = _mm256_sub_ps(x, y); - d2 = _mm256_fmadd_ps(d, d, d2); - } - let mut d2 = emulate_mm256_reduce_add_ps(d2); - while n > 0 { - let x = a.read().to_f32(); - let y = b.read().to_f32(); - a = a.add(1); - b = b.add(1); - n -= 1; - let d = x - y; - d2 += d * d; - } - d2 + use std::arch::x86_64::*; + let mut n = lhs.len() as u32; + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + let d = _mm256_sub_ps(x, y); + d2 = _mm256_fmadd_ps(d, d, d2); + } + let mut d2 = emulate_mm256_reduce_add_ps(d2); + while n > 0 { + let x = unsafe { a.read().to_f32() }; + let y = unsafe { b.read().to_f32() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + let d = x - y; + d2 += d * d; } + d2 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -695,8 +688,12 @@ mod reduce_sum_of_d2 { pub fn reduce_sum_of_d2_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { - extern "C" { - fn fp16_reduce_sum_of_d2_a2_fp16(a: *const (), b: *const (), n: usize) -> f32; + unsafe extern "C" { + unsafe fn fp16_reduce_sum_of_d2_a2_fp16( + a: *const (), + b: *const (), + n: usize, + ) -> f32; } fp16_reduce_sum_of_d2_a2_fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } @@ -740,8 +737,9 @@ mod reduce_sum_of_d2 { pub fn reduce_sum_of_d2_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { - extern "C" { - fn fp16_reduce_sum_of_d2_a3_512(a: *const (), b: *const (), n: usize) -> f32; + unsafe extern "C" { + unsafe fn fp16_reduce_sum_of_d2_a3_512(a: *const (), b: *const (), n: usize) + -> f32; } fp16_reduce_sum_of_d2_a3_512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index 54670deb..f2991df4 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -135,24 +135,22 @@ mod reduce_sum_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4.512")] fn reduce_sum_of_x_v4_512(this: &[f32]) -> f32 { - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm512_setzero_ps(); - while n >= 16 { - let x = _mm512_loadu_ps(a); - a = a.add(16); - n -= 16; - sum = _mm512_add_ps(x, sum); - } - if n > 0 { - let mask = _bzhi_u32(0xffff, n as u32) as u16; - let x = _mm512_maskz_loadu_ps(mask, a); - sum = _mm512_add_ps(x, sum); - } - _mm512_reduce_add_ps(sum) + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + a = unsafe { a.add(16) }; + n -= 16; + sum = _mm512_add_ps(x, sum); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + sum = _mm512_add_ps(x, sum); } + _mm512_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -187,33 +185,31 @@ mod reduce_sum_of_x { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_x_v3(this: &[f32]) -> f32 { use crate::emulate::emulate_mm256_reduce_add_ps; - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm256_setzero_ps(); - while n >= 8 { - let x = _mm256_loadu_ps(a); - a = a.add(8); - n -= 8; - sum = _mm256_add_ps(x, sum); - } - if n >= 4 { - let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); - a = a.add(4); - n -= 4; - sum = _mm256_add_ps(x, sum); - } - let mut sum = emulate_mm256_reduce_add_ps(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - sum += x; - } - sum + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + a = unsafe { a.add(8) }; + n -= 8; + sum = _mm256_add_ps(x, sum); + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + a = unsafe { a.add(4) }; + n -= 4; + sum = _mm256_add_ps(x, sum); + } + let mut sum = emulate_mm256_reduce_add_ps(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + sum += x; } + sum } #[cfg(all(target_arch = "x86_64", test))] @@ -248,27 +244,25 @@ mod reduce_sum_of_x { #[crate::target_cpu(enable = "v2")] fn reduce_sum_of_x_v2(this: &[f32]) -> f32 { use crate::emulate::emulate_mm_reduce_add_ps; - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm_setzero_ps(); - while n >= 4 { - let x = _mm_loadu_ps(a); - a = a.add(4); - n -= 4; - sum = _mm_add_ps(x, sum); - } - let mut sum = emulate_mm_reduce_add_ps(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - sum += x; - } - sum + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + a = unsafe { a.add(4) }; + n -= 4; + sum = _mm_add_ps(x, sum); + } + let mut sum = emulate_mm_reduce_add_ps(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + sum += x; } + sum } #[cfg(all(target_arch = "x86_64", test))] @@ -302,27 +296,25 @@ mod reduce_sum_of_x { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_x_a2(this: &[f32]) -> f32 { - unsafe { - use std::arch::aarch64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = vdupq_n_f32(0.0); - while n >= 4 { - let x = vld1q_f32(a); - a = a.add(4); - n -= 4; - sum = vaddq_f32(x, sum); - } - let mut sum = vaddvq_f32(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - sum += x; - } - sum + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + a = unsafe { a.add(4) }; + n -= 4; + sum = vaddq_f32(x, sum); + } + let mut sum = vaddvq_f32(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + sum += x; } + sum } #[cfg(all(target_arch = "aarch64", test, not(miri)))] @@ -358,8 +350,8 @@ mod reduce_sum_of_x { #[target_feature(enable = "sve")] fn reduce_sum_of_x_a3_256(this: &[f32]) -> f32 { unsafe { - extern "C" { - fn fp32_reduce_sum_of_x_a3_256(this: *const f32, n: usize) -> f32; + unsafe extern "C" { + unsafe fn fp32_reduce_sum_of_x_a3_256(this: *const f32, n: usize) -> f32; } fp32_reduce_sum_of_x_a3_256(this.as_ptr(), this.len()) } @@ -408,26 +400,24 @@ mod reduce_sum_of_abs_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4.512")] fn reduce_sum_of_abs_x_v4_512(this: &[f32]) -> f32 { - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm512_setzero_ps(); - while n >= 16 { - let x = _mm512_loadu_ps(a); - let abs_x = _mm512_abs_ps(x); - a = a.add(16); - n -= 16; - sum = _mm512_add_ps(abs_x, sum); - } - if n > 0 { - let mask = _bzhi_u32(0xffff, n as u32) as u16; - let x = _mm512_maskz_loadu_ps(mask, a); - let abs_x = _mm512_abs_ps(x); - sum = _mm512_add_ps(abs_x, sum); - } - _mm512_reduce_add_ps(sum) - } + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + let abs_x = _mm512_abs_ps(x); + a = unsafe { a.add(16) }; + n -= 16; + sum = _mm512_add_ps(abs_x, sum); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + let abs_x = _mm512_abs_ps(x); + sum = _mm512_add_ps(abs_x, sum); + } + _mm512_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -462,37 +452,35 @@ mod reduce_sum_of_abs_x { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_abs_x_v3(this: &[f32]) -> f32 { use crate::emulate::emulate_mm256_reduce_add_ps; - unsafe { - use std::arch::x86_64::*; - let abs = _mm256_castsi256_ps(_mm256_srli_epi32(_mm256_set1_epi32(-1), 1)); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm256_setzero_ps(); - while n >= 8 { - let x = _mm256_loadu_ps(a); - let abs_x = _mm256_and_ps(abs, x); - a = a.add(8); - n -= 8; - sum = _mm256_add_ps(abs_x, sum); - } - if n >= 4 { - let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); - let abs_x = _mm256_and_ps(abs, x); - a = a.add(4); - n -= 4; - sum = _mm256_add_ps(abs_x, sum); - } - let mut sum = emulate_mm256_reduce_add_ps(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let abs_x = x.abs(); - a = a.add(1); - n -= 1; - sum += abs_x; - } - sum + use std::arch::x86_64::*; + let abs = _mm256_castsi256_ps(_mm256_srli_epi32(_mm256_set1_epi32(-1), 1)); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + let abs_x = _mm256_and_ps(abs, x); + a = unsafe { a.add(8) }; + n -= 8; + sum = _mm256_add_ps(abs_x, sum); + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let abs_x = _mm256_and_ps(abs, x); + a = unsafe { a.add(4) }; + n -= 4; + sum = _mm256_add_ps(abs_x, sum); + } + let mut sum = emulate_mm256_reduce_add_ps(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let abs_x = x.abs(); + a = unsafe { a.add(1) }; + n -= 1; + sum += abs_x; } + sum } #[cfg(all(target_arch = "x86_64", test))] @@ -527,30 +515,28 @@ mod reduce_sum_of_abs_x { #[crate::target_cpu(enable = "v2")] fn reduce_sum_of_abs_x_v2(this: &[f32]) -> f32 { use crate::emulate::emulate_mm_reduce_add_ps; - unsafe { - use std::arch::x86_64::*; - let abs = _mm_castsi128_ps(_mm_srli_epi32(_mm_set1_epi32(-1), 1)); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm_setzero_ps(); - while n >= 4 { - let x = _mm_loadu_ps(a); - let abs_x = _mm_and_ps(abs, x); - a = a.add(4); - n -= 4; - sum = _mm_add_ps(abs_x, sum); - } - let mut sum = emulate_mm_reduce_add_ps(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let abs_x = x.abs(); - a = a.add(1); - n -= 1; - sum += abs_x; - } - sum + use std::arch::x86_64::*; + let abs = _mm_castsi128_ps(_mm_srli_epi32(_mm_set1_epi32(-1), 1)); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + let abs_x = _mm_and_ps(abs, x); + a = unsafe { a.add(4) }; + n -= 4; + sum = _mm_add_ps(abs_x, sum); + } + let mut sum = emulate_mm_reduce_add_ps(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let abs_x = x.abs(); + a = unsafe { a.add(1) }; + n -= 1; + sum += abs_x; } + sum } #[cfg(all(target_arch = "x86_64", test))] @@ -584,29 +570,27 @@ mod reduce_sum_of_abs_x { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_abs_x_a2(this: &[f32]) -> f32 { - unsafe { - use std::arch::aarch64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = vdupq_n_f32(0.0); - while n >= 4 { - let x = vld1q_f32(a); - let abs_x = vabsq_f32(x); - a = a.add(4); - n -= 4; - sum = vaddq_f32(abs_x, sum); - } - let mut sum = vaddvq_f32(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let abs_x = x.abs(); - a = a.add(1); - n -= 1; - sum += abs_x; - } - sum + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + let abs_x = vabsq_f32(x); + a = unsafe { a.add(4) }; + n -= 4; + sum = vaddq_f32(abs_x, sum); + } + let mut sum = vaddvq_f32(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let abs_x = x.abs(); + a = unsafe { a.add(1) }; + n -= 1; + sum += abs_x; } + sum } #[cfg(all(target_arch = "aarch64", test, not(miri)))] @@ -642,8 +626,8 @@ mod reduce_sum_of_abs_x { #[target_feature(enable = "sve")] fn reduce_sum_of_abs_x_a3_256(this: &[f32]) -> f32 { unsafe { - extern "C" { - fn fp32_reduce_sum_of_abs_x_a3_256(this: *const f32, n: usize) -> f32; + unsafe extern "C" { + unsafe fn fp32_reduce_sum_of_abs_x_a3_256(this: *const f32, n: usize) -> f32; } fp32_reduce_sum_of_abs_x_a3_256(this.as_ptr(), this.len()) } @@ -692,24 +676,22 @@ mod reduce_sum_of_x2 { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4.512")] fn reduce_sum_of_x2_v4_512(this: &[f32]) -> f32 { - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut x2 = _mm512_setzero_ps(); - while n >= 16 { - let x = _mm512_loadu_ps(a); - a = a.add(16); - n -= 16; - x2 = _mm512_fmadd_ps(x, x, x2); - } - if n > 0 { - let mask = _bzhi_u32(0xffff, n as u32) as u16; - let x = _mm512_maskz_loadu_ps(mask, a); - x2 = _mm512_fmadd_ps(x, x, x2); - } - _mm512_reduce_add_ps(x2) + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut x2 = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + a = unsafe { a.add(16) }; + n -= 16; + x2 = _mm512_fmadd_ps(x, x, x2); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + x2 = _mm512_fmadd_ps(x, x, x2); } + _mm512_reduce_add_ps(x2) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -744,33 +726,31 @@ mod reduce_sum_of_x2 { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_x2_v3(this: &[f32]) -> f32 { use crate::emulate::emulate_mm256_reduce_add_ps; - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut x2 = _mm256_setzero_ps(); - while n >= 8 { - let x = _mm256_loadu_ps(a); - a = a.add(8); - n -= 8; - x2 = _mm256_fmadd_ps(x, x, x2); - } - if n >= 4 { - let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); - a = a.add(4); - n -= 4; - x2 = _mm256_fmadd_ps(x, x, x2); - } - let mut x2 = emulate_mm256_reduce_add_ps(x2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - x2 += x * x; - } - x2 + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut x2 = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + a = unsafe { a.add(8) }; + n -= 8; + x2 = _mm256_fmadd_ps(x, x, x2); + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + a = unsafe { a.add(4) }; + n -= 4; + x2 = _mm256_fmadd_ps(x, x, x2); + } + let mut x2 = emulate_mm256_reduce_add_ps(x2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + x2 += x * x; } + x2 } #[cfg(all(target_arch = "x86_64", test))] @@ -806,27 +786,25 @@ mod reduce_sum_of_x2 { #[target_feature(enable = "fma")] fn reduce_sum_of_x2_v2_fma(this: &[f32]) -> f32 { use crate::emulate::emulate_mm_reduce_add_ps; - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut x2 = _mm_setzero_ps(); - while n >= 4 { - let x = _mm_loadu_ps(a); - a = a.add(4); - n -= 4; - x2 = _mm_fmadd_ps(x, x, x2); - } - let mut x2 = emulate_mm_reduce_add_ps(x2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - x2 += x * x; - } - x2 + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut x2 = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + a = unsafe { a.add(4) }; + n -= 4; + x2 = _mm_fmadd_ps(x, x, x2); + } + let mut x2 = emulate_mm_reduce_add_ps(x2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + x2 += x * x; } + x2 } #[cfg(all(target_arch = "x86_64", test))] @@ -860,27 +838,25 @@ mod reduce_sum_of_x2 { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_x2_a2(this: &[f32]) -> f32 { - unsafe { - use std::arch::aarch64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut x2 = vdupq_n_f32(0.0); - while n >= 4 { - let x = vld1q_f32(a); - a = a.add(4); - n -= 4; - x2 = vfmaq_f32(x2, x, x); - } - let mut x2 = vaddvq_f32(x2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - x2 += x * x; - } - x2 + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut x2 = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + a = unsafe { a.add(4) }; + n -= 4; + x2 = vfmaq_f32(x2, x, x); + } + let mut x2 = vaddvq_f32(x2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + x2 += x * x; } + x2 } #[cfg(all(target_arch = "aarch64", test, not(miri)))] @@ -916,8 +892,8 @@ mod reduce_sum_of_x2 { #[target_feature(enable = "sve")] fn reduce_sum_of_x2_a3_256(this: &[f32]) -> f32 { unsafe { - extern "C" { - fn fp32_reduce_sum_of_x2_a3_256(this: *const f32, n: usize) -> f32; + unsafe extern "C" { + unsafe fn fp32_reduce_sum_of_x2_a3_256(this: *const f32, n: usize) -> f32; } fp32_reduce_sum_of_x2_a3_256(this.as_ptr(), this.len()) } @@ -969,29 +945,27 @@ mod reduce_min_max_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4.512")] fn reduce_min_max_of_x_v4_512(this: &[f32]) -> (f32, f32) { - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut min = _mm512_set1_ps(f32::INFINITY); - let mut max = _mm512_set1_ps(f32::NEG_INFINITY); - while n >= 16 { - let x = _mm512_loadu_ps(a); - a = a.add(16); - n -= 16; - min = _mm512_min_ps(x, min); - max = _mm512_max_ps(x, max); - } - if n > 0 { - let mask = _bzhi_u32(0xffff, n as u32) as u16; - let x = _mm512_maskz_loadu_ps(mask, a); - min = _mm512_mask_min_ps(min, mask, x, min); - max = _mm512_mask_max_ps(max, mask, x, max); - } - let min = _mm512_reduce_min_ps(min); - let max = _mm512_reduce_max_ps(max); - (min, max) - } + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm512_set1_ps(f32::INFINITY); + let mut max = _mm512_set1_ps(f32::NEG_INFINITY); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + a = unsafe { a.add(16) }; + n -= 16; + min = _mm512_min_ps(x, min); + max = _mm512_max_ps(x, max); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + min = _mm512_mask_min_ps(min, mask, x, min); + max = _mm512_mask_max_ps(max, mask, x, max); + } + let min = _mm512_reduce_min_ps(min); + let max = _mm512_reduce_max_ps(max); + (min, max) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -1023,31 +997,29 @@ mod reduce_min_max_of_x { #[crate::target_cpu(enable = "v3")] fn reduce_min_max_of_x_v3(this: &[f32]) -> (f32, f32) { use crate::emulate::{emulate_mm256_reduce_max_ps, emulate_mm256_reduce_min_ps}; - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut min = _mm256_set1_ps(f32::INFINITY); - let mut max = _mm256_set1_ps(f32::NEG_INFINITY); - while n >= 8 { - let x = _mm256_loadu_ps(a); - a = a.add(8); - n -= 8; - min = _mm256_min_ps(x, min); - max = _mm256_max_ps(x, max); - } - let mut min = emulate_mm256_reduce_min_ps(min); - let mut max = emulate_mm256_reduce_max_ps(max); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - min = x.min(min); - max = x.max(max); - } - (min, max) + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm256_set1_ps(f32::INFINITY); + let mut max = _mm256_set1_ps(f32::NEG_INFINITY); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + a = unsafe { a.add(8) }; + n -= 8; + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + } + let mut min = emulate_mm256_reduce_min_ps(min); + let mut max = emulate_mm256_reduce_max_ps(max); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + min = x.min(min); + max = x.max(max); } + (min, max) } #[cfg(all(target_arch = "x86_64", test))] @@ -1079,31 +1051,29 @@ mod reduce_min_max_of_x { #[crate::target_cpu(enable = "v2")] fn reduce_min_max_of_x_v2(this: &[f32]) -> (f32, f32) { use crate::emulate::{emulate_mm_reduce_max_ps, emulate_mm_reduce_min_ps}; - unsafe { - use std::arch::x86_64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut min = _mm_set1_ps(f32::INFINITY); - let mut max = _mm_set1_ps(f32::NEG_INFINITY); - while n >= 4 { - let x = _mm_loadu_ps(a); - a = a.add(4); - n -= 4; - min = _mm_min_ps(x, min); - max = _mm_max_ps(x, max); - } - let mut min = emulate_mm_reduce_min_ps(min); - let mut max = emulate_mm_reduce_max_ps(max); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - min = x.min(min); - max = x.max(max); - } - (min, max) + use std::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm_set1_ps(f32::INFINITY); + let mut max = _mm_set1_ps(f32::NEG_INFINITY); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + a = unsafe { a.add(4) }; + n -= 4; + min = _mm_min_ps(x, min); + max = _mm_max_ps(x, max); + } + let mut min = emulate_mm_reduce_min_ps(min); + let mut max = emulate_mm_reduce_max_ps(max); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + min = x.min(min); + max = x.max(max); } + (min, max) } #[cfg(all(target_arch = "x86_64", test))] @@ -1134,31 +1104,29 @@ mod reduce_min_max_of_x { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_min_max_of_x_a2(this: &[f32]) -> (f32, f32) { - unsafe { - use std::arch::aarch64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut min = vdupq_n_f32(f32::INFINITY); - let mut max = vdupq_n_f32(f32::NEG_INFINITY); - while n >= 4 { - let x = vld1q_f32(a); - a = a.add(4); - n -= 4; - min = vminq_f32(x, min); - max = vmaxq_f32(x, max); - } - let mut min = vminvq_f32(min); - let mut max = vmaxvq_f32(max); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - min = x.min(min); - max = x.max(max); - } - (min, max) + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = vdupq_n_f32(f32::INFINITY); + let mut max = vdupq_n_f32(f32::NEG_INFINITY); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + a = unsafe { a.add(4) }; + n -= 4; + min = vminq_f32(x, min); + max = vmaxq_f32(x, max); + } + let mut min = vminvq_f32(min); + let mut max = vmaxvq_f32(max); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + min = x.min(min); + max = x.max(max); } + (min, max) } #[cfg(all(target_arch = "aarch64", test, not(miri)))] @@ -1193,8 +1161,8 @@ mod reduce_min_max_of_x { let mut min = f32::INFINITY; let mut max = -f32::INFINITY; unsafe { - extern "C" { - fn fp32_reduce_min_max_of_x_a3_256( + unsafe extern "C" { + unsafe fn fp32_reduce_min_max_of_x_a3_256( this: *const f32, n: usize, out_min: &mut f32, @@ -1249,28 +1217,26 @@ mod reduce_sum_of_xy { #[crate::target_cpu(enable = "v4.512")] fn reduce_sum_of_xy_v4_512(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut xy = _mm512_setzero_ps(); - while n >= 16 { - let x = _mm512_loadu_ps(a); - let y = _mm512_loadu_ps(b); - a = a.add(16); - b = b.add(16); - n -= 16; - xy = _mm512_fmadd_ps(x, y, xy); - } - if n > 0 { - let mask = _bzhi_u32(0xffff, n as u32) as u16; - let x = _mm512_maskz_loadu_ps(mask, a); - let y = _mm512_maskz_loadu_ps(mask, b); - xy = _mm512_fmadd_ps(x, y, xy); - } - _mm512_reduce_add_ps(xy) - } + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + let y = unsafe { _mm512_loadu_ps(b) }; + a = unsafe { a.add(16) }; + b = unsafe { b.add(16) }; + n -= 16; + xy = _mm512_fmadd_ps(x, y, xy); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + let y = unsafe { _mm512_maskz_loadu_ps(mask, b) }; + xy = _mm512_fmadd_ps(x, y, xy); + } + _mm512_reduce_add_ps(xy) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -1310,40 +1276,38 @@ mod reduce_sum_of_xy { fn reduce_sum_of_xy_v3(lhs: &[f32], rhs: &[f32]) -> f32 { use crate::emulate::emulate_mm256_reduce_add_ps; assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut xy = _mm256_setzero_ps(); - while n >= 8 { - let x = _mm256_loadu_ps(a); - let y = _mm256_loadu_ps(b); - a = a.add(8); - b = b.add(8); - n -= 8; - xy = _mm256_fmadd_ps(x, y, xy); - } - if n >= 4 { - let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); - let y = _mm256_zextps128_ps256(_mm_loadu_ps(b)); - a = a.add(4); - b = b.add(4); - n -= 4; - xy = _mm256_fmadd_ps(x, y, xy); - } - let mut xy = emulate_mm256_reduce_add_ps(xy); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let y = b.read(); - a = a.add(1); - b = b.add(1); - n -= 1; - xy += x * y; - } - xy + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + let y = unsafe { _mm256_loadu_ps(b) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + xy = _mm256_fmadd_ps(x, y, xy); + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + xy = _mm256_fmadd_ps(x, y, xy); + } + let mut xy = emulate_mm256_reduce_add_ps(xy); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + xy += x * y; } + xy } #[cfg(all(target_arch = "x86_64", test))] @@ -1384,32 +1348,30 @@ mod reduce_sum_of_xy { fn reduce_sum_of_xy_v2_fma(lhs: &[f32], rhs: &[f32]) -> f32 { use crate::emulate::emulate_mm_reduce_add_ps; assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut xy = _mm_setzero_ps(); - while n >= 4 { - let x = _mm_loadu_ps(a); - let y = _mm_loadu_ps(b); - a = a.add(4); - b = b.add(4); - n -= 4; - xy = _mm_fmadd_ps(x, y, xy); - } - let mut xy = emulate_mm_reduce_add_ps(xy); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let y = b.read(); - a = a.add(1); - b = b.add(1); - n -= 1; - xy += x * y; - } - xy + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + let y = unsafe { _mm_loadu_ps(b) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + xy = _mm_fmadd_ps(x, y, xy); + } + let mut xy = emulate_mm_reduce_add_ps(xy); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + xy += x * y; } + xy } #[cfg(all(target_arch = "x86_64", test))] @@ -1448,32 +1410,30 @@ mod reduce_sum_of_xy { #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_xy_a2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::aarch64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut xy = vdupq_n_f32(0.0); - while n >= 4 { - let x = vld1q_f32(a); - let y = vld1q_f32(b); - a = a.add(4); - b = b.add(4); - n -= 4; - xy = vfmaq_f32(xy, x, y); - } - let mut xy = vaddvq_f32(xy); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let y = b.read(); - a = a.add(1); - b = b.add(1); - n -= 1; - xy += x * y; - } - xy + use std::arch::aarch64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut xy = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + let y = unsafe { vld1q_f32(b) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + xy = vfmaq_f32(xy, x, y); + } + let mut xy = vaddvq_f32(xy); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + xy += x * y; } + xy } #[cfg(all(target_arch = "aarch64", test, not(miri)))] @@ -1514,8 +1474,12 @@ mod reduce_sum_of_xy { fn reduce_sum_of_xy_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { - extern "C" { - fn fp32_reduce_sum_of_xy_a3_256(a: *const f32, b: *const f32, n: usize) -> f32; + unsafe extern "C" { + unsafe fn fp32_reduce_sum_of_xy_a3_256( + a: *const f32, + b: *const f32, + n: usize, + ) -> f32; } fp32_reduce_sum_of_xy_a3_256(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) } @@ -1570,30 +1534,28 @@ mod reduce_sum_of_d2 { #[crate::target_cpu(enable = "v4.512")] fn reduce_sum_of_d2_v4_512(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len() as u32; - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut d2 = _mm512_setzero_ps(); - while n >= 16 { - let x = _mm512_loadu_ps(a); - let y = _mm512_loadu_ps(b); - a = a.add(16); - b = b.add(16); - n -= 16; - let d = _mm512_sub_ps(x, y); - d2 = _mm512_fmadd_ps(d, d, d2); - } - if n > 0 { - let mask = _bzhi_u32(0xffff, n) as u16; - let x = _mm512_maskz_loadu_ps(mask, a); - let y = _mm512_maskz_loadu_ps(mask, b); - let d = _mm512_sub_ps(x, y); - d2 = _mm512_fmadd_ps(d, d, d2); - } - _mm512_reduce_add_ps(d2) - } + use std::arch::x86_64::*; + let mut n = lhs.len() as u32; + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + let y = unsafe { _mm512_loadu_ps(b) }; + a = unsafe { a.add(16) }; + b = unsafe { b.add(16) }; + n -= 16; + let d = _mm512_sub_ps(x, y); + d2 = _mm512_fmadd_ps(d, d, d2); + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + let y = unsafe { _mm512_maskz_loadu_ps(mask, b) }; + let d = _mm512_sub_ps(x, y); + d2 = _mm512_fmadd_ps(d, d, d2); + } + _mm512_reduce_add_ps(d2) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -1633,43 +1595,41 @@ mod reduce_sum_of_d2 { fn reduce_sum_of_d2_v3(lhs: &[f32], rhs: &[f32]) -> f32 { use crate::emulate::emulate_mm256_reduce_add_ps; assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut d2 = _mm256_setzero_ps(); - while n >= 8 { - let x = _mm256_loadu_ps(a); - let y = _mm256_loadu_ps(b); - a = a.add(8); - b = b.add(8); - n -= 8; - let d = _mm256_sub_ps(x, y); - d2 = _mm256_fmadd_ps(d, d, d2); - } - if n >= 4 { - let x = _mm256_zextps128_ps256(_mm_loadu_ps(a)); - let y = _mm256_zextps128_ps256(_mm_loadu_ps(b)); - a = a.add(4); - b = b.add(4); - n -= 4; - let d = _mm256_sub_ps(x, y); - d2 = _mm256_fmadd_ps(d, d, d2); - } - let mut d2 = emulate_mm256_reduce_add_ps(d2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let y = b.read(); - a = a.add(1); - b = b.add(1); - n -= 1; - let d = x - y; - d2 += d * d; - } - d2 + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + let y = unsafe { _mm256_loadu_ps(b) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + let d = _mm256_sub_ps(x, y); + d2 = _mm256_fmadd_ps(d, d, d2); + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + let d = _mm256_sub_ps(x, y); + d2 = _mm256_fmadd_ps(d, d, d2); + } + let mut d2 = emulate_mm256_reduce_add_ps(d2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + let d = x - y; + d2 += d * d; } + d2 } #[cfg(all(target_arch = "x86_64", test))] @@ -1710,34 +1670,32 @@ mod reduce_sum_of_d2 { fn reduce_sum_of_d2_v2_fma(lhs: &[f32], rhs: &[f32]) -> f32 { use crate::emulate::emulate_mm_reduce_add_ps; assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut d2 = _mm_setzero_ps(); - while n >= 4 { - let x = _mm_loadu_ps(a); - let y = _mm_loadu_ps(b); - a = a.add(4); - b = b.add(4); - n -= 4; - let d = _mm_sub_ps(x, y); - d2 = _mm_fmadd_ps(d, d, d2); - } - let mut d2 = emulate_mm_reduce_add_ps(d2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let y = b.read(); - a = a.add(1); - b = b.add(1); - n -= 1; - let d = x - y; - d2 += d * d; - } - d2 + use std::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + let y = unsafe { _mm_loadu_ps(b) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + let d = _mm_sub_ps(x, y); + d2 = _mm_fmadd_ps(d, d, d2); + } + let mut d2 = emulate_mm_reduce_add_ps(d2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + let d = x - y; + d2 += d * d; } + d2 } #[cfg(all(target_arch = "x86_64", test))] @@ -1776,34 +1734,32 @@ mod reduce_sum_of_d2 { #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_d2_a2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); - unsafe { - use std::arch::aarch64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut d2 = vdupq_n_f32(0.0); - while n >= 4 { - let x = vld1q_f32(a); - let y = vld1q_f32(b); - a = a.add(4); - b = b.add(4); - n -= 4; - let d = vsubq_f32(x, y); - d2 = vfmaq_f32(d2, d, d); - } - let mut d2 = vaddvq_f32(d2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let y = b.read(); - a = a.add(1); - b = b.add(1); - n -= 1; - let d = x - y; - d2 += d * d; - } - d2 + use std::arch::aarch64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut d2 = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + let y = unsafe { vld1q_f32(b) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + let d = vsubq_f32(x, y); + d2 = vfmaq_f32(d2, d, d); + } + let mut d2 = vaddvq_f32(d2); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + let d = x - y; + d2 += d * d; } + d2 } #[cfg(all(target_arch = "aarch64", test, not(miri)))] @@ -1844,8 +1800,12 @@ mod reduce_sum_of_d2 { fn reduce_sum_of_d2_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { - extern "C" { - fn fp32_reduce_sum_of_d2_a3_256(a: *const f32, b: *const f32, n: usize) -> f32; + unsafe extern "C" { + unsafe fn fp32_reduce_sum_of_d2_a3_256( + a: *const f32, + b: *const f32, + n: usize, + ) -> f32; } fp32_reduce_sum_of_d2_a3_256(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) } @@ -1907,39 +1867,39 @@ mod reduce_sum_of_xy_sparse { let (mut rp, rn) = (0, ri.len()); let (li, lv) = (li.as_ptr(), lv.as_ptr()); let (ri, rv) = (ri.as_ptr(), rv.as_ptr()); - unsafe { - use std::arch::x86_64::*; - let mut xy = _mm512_setzero_ps(); - while lp + 16 <= ln && rp + 16 <= rn { - let lx = _mm512_loadu_epi32(li.add(lp).cast()); - let rx = _mm512_loadu_epi32(ri.add(rp).cast()); - let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); - let lv = _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))); - let rv = _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))); - xy = _mm512_fmadd_ps(lv, rv, xy); - let lt = li.add(lp + 16 - 1).read(); - let rt = ri.add(rp + 16 - 1).read(); - lp += (lt <= rt) as usize * 16; - rp += (lt >= rt) as usize * 16; - } - while lp < ln && rp < rn { - let lw = 16.min(ln - lp); - let rw = 16.min(rn - rp); - let lm = _bzhi_u32(0xffff, lw as _) as u16; - let rm = _bzhi_u32(0xffff, rw as _) as u16; - let lx = _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), lm, li.add(lp).cast()); - let rx = _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), rm, ri.add(rp).cast()); - let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); - let lv = _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))); - let rv = _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))); - xy = _mm512_fmadd_ps(lv, rv, xy); - let lt = li.add(lp + lw - 1).read(); - let rt = ri.add(rp + rw - 1).read(); - lp += (lt <= rt) as usize * lw; - rp += (lt >= rt) as usize * rw; - } - _mm512_reduce_add_ps(xy) + use std::arch::x86_64::*; + let mut xy = _mm512_setzero_ps(); + while lp + 16 <= ln && rp + 16 <= rn { + let lx = unsafe { _mm512_loadu_epi32(li.add(lp).cast()) }; + let rx = unsafe { _mm512_loadu_epi32(ri.add(rp).cast()) }; + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))) }; + let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))) }; + xy = _mm512_fmadd_ps(lv, rv, xy); + let lt = unsafe { li.add(lp + 16 - 1).read() }; + let rt = unsafe { ri.add(rp + 16 - 1).read() }; + lp += (lt <= rt) as usize * 16; + rp += (lt >= rt) as usize * 16; } + while lp < ln && rp < rn { + let lw = 16.min(ln - lp); + let rw = 16.min(rn - rp); + let lm = _bzhi_u32(0xffff, lw as _) as u16; + let rm = _bzhi_u32(0xffff, rw as _) as u16; + let lx = + unsafe { _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), lm, li.add(lp).cast()) }; + let rx = + unsafe { _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), rm, ri.add(rp).cast()) }; + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))) }; + let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))) }; + xy = _mm512_fmadd_ps(lv, rv, xy); + let lt = unsafe { li.add(lp + lw - 1).read() }; + let rt = unsafe { ri.add(rp + rw - 1).read() }; + lp += (lt <= rt) as usize * lw; + rp += (lt >= rt) as usize * rw; + } + _mm512_reduce_add_ps(xy) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -2023,73 +1983,73 @@ mod reduce_sum_of_d2_sparse { let (mut rp, rn) = (0, ri.len()); let (li, lv) = (li.as_ptr(), lv.as_ptr()); let (ri, rv) = (ri.as_ptr(), rv.as_ptr()); - unsafe { - use std::arch::x86_64::*; - let mut d2 = _mm512_setzero_ps(); - while lp + 16 <= ln && rp + 16 <= rn { - let lx = _mm512_loadu_epi32(li.add(lp).cast()); - let rx = _mm512_loadu_epi32(ri.add(rp).cast()); - let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); - let lv = _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))); - let rv = _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))); - let d = _mm512_sub_ps(lv, rv); + use std::arch::x86_64::*; + let mut d2 = _mm512_setzero_ps(); + while lp + 16 <= ln && rp + 16 <= rn { + let lx = unsafe { _mm512_loadu_epi32(li.add(lp).cast()) }; + let rx = unsafe { _mm512_loadu_epi32(ri.add(rp).cast()) }; + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))) }; + let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))) }; + let d = _mm512_sub_ps(lv, rv); + d2 = _mm512_fmadd_ps(d, d, d2); + d2 = _mm512_sub_ps(d2, _mm512_mul_ps(lv, lv)); + d2 = _mm512_sub_ps(d2, _mm512_mul_ps(rv, rv)); + let lt = unsafe { li.add(lp + 16 - 1).read() }; + let rt = unsafe { ri.add(rp + 16 - 1).read() }; + lp += (lt <= rt) as usize * 16; + rp += (lt >= rt) as usize * 16; + } + while lp < ln && rp < rn { + let lw = 16.min(ln - lp); + let rw = 16.min(rn - rp); + let lm = _bzhi_u32(0xffff, lw as _) as u16; + let rm = _bzhi_u32(0xffff, rw as _) as u16; + let lx = + unsafe { _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), lm, li.add(lp).cast()) }; + let rx = + unsafe { _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), rm, ri.add(rp).cast()) }; + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))) }; + let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))) }; + let d = _mm512_sub_ps(lv, rv); + d2 = _mm512_fmadd_ps(d, d, d2); + d2 = _mm512_sub_ps(d2, _mm512_mul_ps(lv, lv)); + d2 = _mm512_sub_ps(d2, _mm512_mul_ps(rv, rv)); + let lt = unsafe { li.add(lp + lw - 1).read() }; + let rt = unsafe { ri.add(rp + rw - 1).read() }; + lp += (lt <= rt) as usize * lw; + rp += (lt >= rt) as usize * rw; + } + { + let mut lp = 0; + while lp + 16 <= ln { + let d = unsafe { _mm512_loadu_ps(lv.add(lp)) }; d2 = _mm512_fmadd_ps(d, d, d2); - d2 = _mm512_sub_ps(d2, _mm512_mul_ps(lv, lv)); - d2 = _mm512_sub_ps(d2, _mm512_mul_ps(rv, rv)); - let lt = li.add(lp + 16 - 1).read(); - let rt = ri.add(rp + 16 - 1).read(); - lp += (lt <= rt) as usize * 16; - rp += (lt >= rt) as usize * 16; + lp += 16; } - while lp < ln && rp < rn { - let lw = 16.min(ln - lp); - let rw = 16.min(rn - rp); + if lp < ln { + let lw = ln - lp; let lm = _bzhi_u32(0xffff, lw as _) as u16; - let rm = _bzhi_u32(0xffff, rw as _) as u16; - let lx = _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), lm, li.add(lp).cast()); - let rx = _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), rm, ri.add(rp).cast()); - let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); - let lv = _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))); - let rv = _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))); - let d = _mm512_sub_ps(lv, rv); + let d = unsafe { _mm512_maskz_loadu_ps(lm, lv.add(lp)) }; d2 = _mm512_fmadd_ps(d, d, d2); - d2 = _mm512_sub_ps(d2, _mm512_mul_ps(lv, lv)); - d2 = _mm512_sub_ps(d2, _mm512_mul_ps(rv, rv)); - let lt = li.add(lp + lw - 1).read(); - let rt = ri.add(rp + rw - 1).read(); - lp += (lt <= rt) as usize * lw; - rp += (lt >= rt) as usize * rw; } - { - let mut lp = 0; - while lp + 16 <= ln { - let d = _mm512_loadu_ps(lv.add(lp)); - d2 = _mm512_fmadd_ps(d, d, d2); - lp += 16; - } - if lp < ln { - let lw = ln - lp; - let lm = _bzhi_u32(0xffff, lw as _) as u16; - let d = _mm512_maskz_loadu_ps(lm, lv.add(lp)); - d2 = _mm512_fmadd_ps(d, d, d2); - } + } + { + let mut rp = 0; + while rp + 16 <= rn { + let d = unsafe { _mm512_loadu_ps(rv.add(rp)) }; + d2 = _mm512_fmadd_ps(d, d, d2); + rp += 16; } - { - let mut rp = 0; - while rp + 16 <= rn { - let d = _mm512_loadu_ps(rv.add(rp)); - d2 = _mm512_fmadd_ps(d, d, d2); - rp += 16; - } - if rp < rn { - let rw = rn - rp; - let rm = _bzhi_u32(0xffff, rw as _) as u16; - let d = _mm512_maskz_loadu_ps(rm, rv.add(rp)); - d2 = _mm512_fmadd_ps(d, d, d2); - } + if rp < rn { + let rw = rn - rp; + let rm = _bzhi_u32(0xffff, rw as _) as u16; + let d = unsafe { _mm512_maskz_loadu_ps(rm, rv.add(rp)) }; + d2 = _mm512_fmadd_ps(d, d, d2); } - _mm512_reduce_add_ps(d2) } + _mm512_reduce_add_ps(d2) } #[cfg(all(target_arch = "x86_64", test, not(miri)))] diff --git a/crates/simd/src/fast_scan/mod.rs b/crates/simd/src/fast_scan/mod.rs index a99c227f..72b2eb26 100644 --- a/crates/simd/src/fast_scan/mod.rs +++ b/crates/simd/src/fast_scan/mod.rs @@ -131,108 +131,106 @@ mod fast_scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - unsafe { - use std::arch::x86_64::*; - - #[inline] - #[crate::target_cpu(enable = "v4.512")] - fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { - unsafe { - let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); - let x0y1 = _mm256_blend_epi32(x0x1, y0y1, 0xf0); - _mm256_add_epi16(x1y0, x0y1) - } - } + use std::arch::x86_64::*; + + #[inline] + #[crate::target_cpu(enable = "v4.512")] + fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { + let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); + let x0y1 = _mm256_blend_epi32(x0x1, y0y1, 0xf0); + _mm256_add_epi16(x1y0, x0y1) + } - #[inline] - #[crate::target_cpu(enable = "v4.512")] - fn combine4x2(x0x1x2x3: __m512i, y0y1y2y3: __m512i) -> __m256i { - unsafe { - let x0x1 = _mm512_castsi512_si256(x0x1x2x3); - let x2x3 = _mm512_extracti64x4_epi64(x0x1x2x3, 1); - let y0y1 = _mm512_castsi512_si256(y0y1y2y3); - let y2y3 = _mm512_extracti64x4_epi64(y0y1y2y3, 1); - let x01y01 = combine2x2(x0x1, y0y1); - let x23y23 = combine2x2(x2x3, y2y3); - _mm256_add_epi16(x01y01, x23y23) - } - } + #[inline] + #[crate::target_cpu(enable = "v4.512")] + fn combine4x2(x0x1x2x3: __m512i, y0y1y2y3: __m512i) -> __m256i { + let x0x1 = _mm512_castsi512_si256(x0x1x2x3); + let x2x3 = _mm512_extracti64x4_epi64(x0x1x2x3, 1); + let y0y1 = _mm512_castsi512_si256(y0y1y2y3); + let y2y3 = _mm512_extracti64x4_epi64(y0y1y2y3, 1); + let x01y01 = combine2x2(x0x1, y0y1); + let x23y23 = combine2x2(x2x3, y2y3); + _mm256_add_epi16(x01y01, x23y23) + } - let mut accu_0 = _mm512_setzero_si512(); - let mut accu_1 = _mm512_setzero_si512(); - let mut accu_2 = _mm512_setzero_si512(); - let mut accu_3 = _mm512_setzero_si512(); + let mut accu_0 = _mm512_setzero_si512(); + let mut accu_1 = _mm512_setzero_si512(); + let mut accu_2 = _mm512_setzero_si512(); + let mut accu_3 = _mm512_setzero_si512(); - let mut i = 0_usize; - while i + 4 <= n { - let code = _mm512_loadu_si512(code.as_ptr().add(i).cast()); + let mut i = 0_usize; + while i + 4 <= n { + let code = unsafe { _mm512_loadu_si512(code.as_ptr().add(i).cast()) }; - let mask = _mm512_set1_epi8(0xf); - let clo = _mm512_and_si512(code, mask); - let chi = _mm512_and_si512(_mm512_srli_epi16(code, 4), mask); + let mask = _mm512_set1_epi8(0xf); + let clo = _mm512_and_si512(code, mask); + let chi = _mm512_and_si512(_mm512_srli_epi16(code, 4), mask); - let lut = _mm512_loadu_si512(lut.as_ptr().add(i).cast()); - let res_lo = _mm512_shuffle_epi8(lut, clo); - accu_0 = _mm512_add_epi16(accu_0, res_lo); - accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); - let res_hi = _mm512_shuffle_epi8(lut, chi); - accu_2 = _mm512_add_epi16(accu_2, res_hi); - accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + let lut = unsafe { _mm512_loadu_si512(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm512_shuffle_epi8(lut, clo); + accu_0 = _mm512_add_epi16(accu_0, res_lo); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_shuffle_epi8(lut, chi); + accu_2 = _mm512_add_epi16(accu_2, res_hi); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); - i += 4; - } - if i + 2 <= n { - let code = _mm256_loadu_si256(code.as_ptr().add(i).cast()); - - let mask = _mm256_set1_epi8(0xf); - let clo = _mm256_and_si256(code, mask); - let chi = _mm256_and_si256(_mm256_srli_epi16(code, 4), mask); - - let lut = _mm256_loadu_si256(lut.as_ptr().add(i).cast()); - let res_lo = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, clo)); - accu_0 = _mm512_add_epi16(accu_0, res_lo); - accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); - let res_hi = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, chi)); - accu_2 = _mm512_add_epi16(accu_2, res_hi); - accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); - - i += 2; - } - if i < n { - let code = _mm_loadu_si128(code.as_ptr().add(i).cast()); - - let mask = _mm_set1_epi8(0xf); - let clo = _mm_and_si128(code, mask); - let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); - - let lut = _mm_loadu_si128(lut.as_ptr().add(i).cast()); - let res_lo = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, clo)); - accu_0 = _mm512_add_epi16(accu_0, res_lo); - accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); - let res_hi = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, chi)); - accu_2 = _mm512_add_epi16(accu_2, res_hi); - accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); - - i += 1; - } - debug_assert_eq!(i, n); + i += 4; + } + if i + 2 <= n { + let code = unsafe { _mm256_loadu_si256(code.as_ptr().add(i).cast()) }; + + let mask = _mm256_set1_epi8(0xf); + let clo = _mm256_and_si256(code, mask); + let chi = _mm256_and_si256(_mm256_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm256_loadu_si256(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, clo)); + accu_0 = _mm512_add_epi16(accu_0, res_lo); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, chi)); + accu_2 = _mm512_add_epi16(accu_2, res_hi); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + + i += 2; + } + if i < n { + let code = unsafe { _mm_loadu_si128(code.as_ptr().add(i).cast()) }; + + let mask = _mm_set1_epi8(0xf); + let clo = _mm_and_si128(code, mask); + let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm_loadu_si128(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, clo)); + accu_0 = _mm512_add_epi16(accu_0, res_lo); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, chi)); + accu_2 = _mm512_add_epi16(accu_2, res_hi); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, n); - let mut result = [0_u16; 32]; + let mut result = [0_u16; 32]; - accu_0 = _mm512_sub_epi16(accu_0, _mm512_slli_epi16(accu_1, 8)); + accu_0 = _mm512_sub_epi16(accu_0, _mm512_slli_epi16(accu_1, 8)); + unsafe { _mm256_storeu_si256( result.as_mut_ptr().add(0).cast(), combine4x2(accu_0, accu_1), ); + } - accu_2 = _mm512_sub_epi16(accu_2, _mm512_slli_epi16(accu_3, 8)); + accu_2 = _mm512_sub_epi16(accu_2, _mm512_slli_epi16(accu_3, 8)); + unsafe { _mm256_storeu_si256( result.as_mut_ptr().add(16).cast(), combine4x2(accu_2, accu_3), ); - - result } + + result } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -265,77 +263,77 @@ mod fast_scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - unsafe { - use std::arch::x86_64::*; + use std::arch::x86_64::*; - #[inline] - #[crate::target_cpu(enable = "v3")] - fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { - unsafe { - let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); - let x0y1 = _mm256_blend_epi32(x0x1, y0y1, 0xf0); - _mm256_add_epi16(x1y0, x0y1) - } - } + #[inline] + #[crate::target_cpu(enable = "v3")] + fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { + let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); + let x0y1 = _mm256_blend_epi32(x0x1, y0y1, 0xf0); + _mm256_add_epi16(x1y0, x0y1) + } - let mut accu_0 = _mm256_setzero_si256(); - let mut accu_1 = _mm256_setzero_si256(); - let mut accu_2 = _mm256_setzero_si256(); - let mut accu_3 = _mm256_setzero_si256(); + let mut accu_0 = _mm256_setzero_si256(); + let mut accu_1 = _mm256_setzero_si256(); + let mut accu_2 = _mm256_setzero_si256(); + let mut accu_3 = _mm256_setzero_si256(); - let mut i = 0_usize; - while i + 2 <= n { - let code = _mm256_loadu_si256(code.as_ptr().add(i).cast()); + let mut i = 0_usize; + while i + 2 <= n { + let code = unsafe { _mm256_loadu_si256(code.as_ptr().add(i).cast()) }; - let mask = _mm256_set1_epi8(0xf); - let clo = _mm256_and_si256(code, mask); - let chi = _mm256_and_si256(_mm256_srli_epi16(code, 4), mask); + let mask = _mm256_set1_epi8(0xf); + let clo = _mm256_and_si256(code, mask); + let chi = _mm256_and_si256(_mm256_srli_epi16(code, 4), mask); - let lut = _mm256_loadu_si256(lut.as_ptr().add(i).cast()); - let res_lo = _mm256_shuffle_epi8(lut, clo); - accu_0 = _mm256_add_epi16(accu_0, res_lo); - accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); - let res_hi = _mm256_shuffle_epi8(lut, chi); - accu_2 = _mm256_add_epi16(accu_2, res_hi); - accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); + let lut = unsafe { _mm256_loadu_si256(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm256_shuffle_epi8(lut, clo); + accu_0 = _mm256_add_epi16(accu_0, res_lo); + accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); + let res_hi = _mm256_shuffle_epi8(lut, chi); + accu_2 = _mm256_add_epi16(accu_2, res_hi); + accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); - i += 2; - } - if i < n { - let code = _mm_loadu_si128(code.as_ptr().add(i).cast()); - - let mask = _mm_set1_epi8(0xf); - let clo = _mm_and_si128(code, mask); - let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); - - let lut = _mm_loadu_si128(lut.as_ptr().add(i).cast()); - let res_lo = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, clo)); - accu_0 = _mm256_add_epi16(accu_0, res_lo); - accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); - let res_hi = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, chi)); - accu_2 = _mm256_add_epi16(accu_2, res_hi); - accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); - - i += 1; - } - debug_assert_eq!(i, n); + i += 2; + } + if i < n { + let code = unsafe { _mm_loadu_si128(code.as_ptr().add(i).cast()) }; + + let mask = _mm_set1_epi8(0xf); + let clo = _mm_and_si128(code, mask); + let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm_loadu_si128(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, clo)); + accu_0 = _mm256_add_epi16(accu_0, res_lo); + accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); + let res_hi = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, chi)); + accu_2 = _mm256_add_epi16(accu_2, res_hi); + accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, n); - let mut result = [0_u16; 32]; + let mut result = [0_u16; 32]; - accu_0 = _mm256_sub_epi16(accu_0, _mm256_slli_epi16(accu_1, 8)); + accu_0 = _mm256_sub_epi16(accu_0, _mm256_slli_epi16(accu_1, 8)); + unsafe { _mm256_storeu_si256( result.as_mut_ptr().add(0).cast(), combine2x2(accu_0, accu_1), ); + } - accu_2 = _mm256_sub_epi16(accu_2, _mm256_slli_epi16(accu_3, 8)); + accu_2 = _mm256_sub_epi16(accu_2, _mm256_slli_epi16(accu_3, 8)); + unsafe { _mm256_storeu_si256( result.as_mut_ptr().add(16).cast(), combine2x2(accu_2, accu_3), ); - - result } + + result } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -367,46 +365,48 @@ mod fast_scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - unsafe { - use std::arch::x86_64::*; - - let mut accu_0 = _mm_setzero_si128(); - let mut accu_1 = _mm_setzero_si128(); - let mut accu_2 = _mm_setzero_si128(); - let mut accu_3 = _mm_setzero_si128(); - - let mut i = 0_usize; - while i < n { - let code = _mm_loadu_si128(code.as_ptr().add(i).cast()); - - let mask = _mm_set1_epi8(0xf); - let clo = _mm_and_si128(code, mask); - let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); - - let lut = _mm_loadu_si128(lut.as_ptr().add(i).cast()); - let res_lo = _mm_shuffle_epi8(lut, clo); - accu_0 = _mm_add_epi16(accu_0, res_lo); - accu_1 = _mm_add_epi16(accu_1, _mm_srli_epi16(res_lo, 8)); - let res_hi = _mm_shuffle_epi8(lut, chi); - accu_2 = _mm_add_epi16(accu_2, res_hi); - accu_3 = _mm_add_epi16(accu_3, _mm_srli_epi16(res_hi, 8)); - - i += 1; - } - debug_assert_eq!(i, n); + use std::arch::x86_64::*; + + let mut accu_0 = _mm_setzero_si128(); + let mut accu_1 = _mm_setzero_si128(); + let mut accu_2 = _mm_setzero_si128(); + let mut accu_3 = _mm_setzero_si128(); + + let mut i = 0_usize; + while i < n { + let code = unsafe { _mm_loadu_si128(code.as_ptr().add(i).cast()) }; + + let mask = _mm_set1_epi8(0xf); + let clo = _mm_and_si128(code, mask); + let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm_loadu_si128(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm_shuffle_epi8(lut, clo); + accu_0 = _mm_add_epi16(accu_0, res_lo); + accu_1 = _mm_add_epi16(accu_1, _mm_srli_epi16(res_lo, 8)); + let res_hi = _mm_shuffle_epi8(lut, chi); + accu_2 = _mm_add_epi16(accu_2, res_hi); + accu_3 = _mm_add_epi16(accu_3, _mm_srli_epi16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, n); - let mut result = [0_u16; 32]; + let mut result = [0_u16; 32]; - accu_0 = _mm_sub_epi16(accu_0, _mm_slli_epi16(accu_1, 8)); + accu_0 = _mm_sub_epi16(accu_0, _mm_slli_epi16(accu_1, 8)); + unsafe { _mm_storeu_si128(result.as_mut_ptr().add(0).cast(), accu_0); _mm_storeu_si128(result.as_mut_ptr().add(8).cast(), accu_1); + } - accu_2 = _mm_sub_epi16(accu_2, _mm_slli_epi16(accu_3, 8)); + accu_2 = _mm_sub_epi16(accu_2, _mm_slli_epi16(accu_3, 8)); + unsafe { _mm_storeu_si128(result.as_mut_ptr().add(16).cast(), accu_2); _mm_storeu_si128(result.as_mut_ptr().add(24).cast(), accu_3); - - result } + + result } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -438,45 +438,47 @@ mod fast_scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - unsafe { - use std::arch::aarch64::*; + use std::arch::aarch64::*; - let mut accu_0 = vdupq_n_u16(0); - let mut accu_1 = vdupq_n_u16(0); - let mut accu_2 = vdupq_n_u16(0); - let mut accu_3 = vdupq_n_u16(0); + let mut accu_0 = vdupq_n_u16(0); + let mut accu_1 = vdupq_n_u16(0); + let mut accu_2 = vdupq_n_u16(0); + let mut accu_3 = vdupq_n_u16(0); - let mut i = 0_usize; - while i < n { - let code = vld1q_u8(code.as_ptr().add(i).cast()); + let mut i = 0_usize; + while i < n { + let code = unsafe { vld1q_u8(code.as_ptr().add(i).cast()) }; - let clo = vandq_u8(code, vdupq_n_u8(0xf)); - let chi = vshrq_n_u8(code, 4); + let clo = vandq_u8(code, vdupq_n_u8(0xf)); + let chi = vshrq_n_u8(code, 4); - let lut = vld1q_u8(lut.as_ptr().add(i).cast()); - let res_lo = vreinterpretq_u16_u8(vqtbl1q_u8(lut, clo)); - accu_0 = vaddq_u16(accu_0, res_lo); - accu_1 = vaddq_u16(accu_1, vshrq_n_u16(res_lo, 8)); - let res_hi = vreinterpretq_u16_u8(vqtbl1q_u8(lut, chi)); - accu_2 = vaddq_u16(accu_2, res_hi); - accu_3 = vaddq_u16(accu_3, vshrq_n_u16(res_hi, 8)); + let lut = unsafe { vld1q_u8(lut.as_ptr().add(i).cast()) }; + let res_lo = vreinterpretq_u16_u8(vqtbl1q_u8(lut, clo)); + accu_0 = vaddq_u16(accu_0, res_lo); + accu_1 = vaddq_u16(accu_1, vshrq_n_u16(res_lo, 8)); + let res_hi = vreinterpretq_u16_u8(vqtbl1q_u8(lut, chi)); + accu_2 = vaddq_u16(accu_2, res_hi); + accu_3 = vaddq_u16(accu_3, vshrq_n_u16(res_hi, 8)); - i += 1; - } - debug_assert_eq!(i, n); + i += 1; + } + debug_assert_eq!(i, n); - let mut result = [0_u16; 32]; + let mut result = [0_u16; 32]; - accu_0 = vsubq_u16(accu_0, vshlq_n_u16(accu_1, 8)); + accu_0 = vsubq_u16(accu_0, vshlq_n_u16(accu_1, 8)); + unsafe { vst1q_u16(result.as_mut_ptr().add(0).cast(), accu_0); vst1q_u16(result.as_mut_ptr().add(8).cast(), accu_1); + } - accu_2 = vsubq_u16(accu_2, vshlq_n_u16(accu_3, 8)); + accu_2 = vsubq_u16(accu_2, vshlq_n_u16(accu_3, 8)); + unsafe { vst1q_u16(result.as_mut_ptr().add(16).cast(), accu_2); vst1q_u16(result.as_mut_ptr().add(24).cast(), accu_3); - - result } + + result } #[cfg(all(target_arch = "aarch64", test, not(miri)))] diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index 880f64bc..3b6bb770 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -3,38 +3,37 @@ mod mul_add_round { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4.512")] fn mul_add_round_v4_512(this: &[f32], k: f32, b: f32) -> Vec { - let n = this.len(); - let mut r = Vec::::with_capacity(n); - unsafe { - use std::arch::x86_64::*; - let lk = _mm512_set1_ps(k); - let lb = _mm512_set1_ps(b); - let mut n = n; - let mut a = this.as_ptr(); - let mut r = r.as_mut_ptr(); - while n >= 16 { - let x = _mm512_loadu_ps(a); - let v = - _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); - let v = _mm512_cvtps_epi32(v); - let vfl = _mm512_cvtepi32_epi8(v); - _mm_storeu_si128(r.cast(), vfl); - n -= 16; - a = a.add(16); - r = r.add(16); + let mut r = Vec::::with_capacity(this.len()); + use std::arch::x86_64::*; + let lk = _mm512_set1_ps(k); + let lb = _mm512_set1_ps(b); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut p = r.as_mut_ptr(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + let v = _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + let v = _mm512_cvtps_epi32(v); + let vfl = _mm512_cvtepi32_epi8(v); + unsafe { + _mm_storeu_si128(p.cast(), vfl); } - if n > 0 { - let mask = _bzhi_u32(0xffff, n as u32) as u16; - let x = _mm512_maskz_loadu_ps(mask, a); - let v = - _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); - let v = _mm512_cvtps_epi32(v); - let vfl = _mm512_cvtepi32_epi8(v); - _mm_mask_storeu_epi8(r.cast(), mask, vfl); + n -= 16; + a = unsafe { a.add(16) }; + p = unsafe { p.add(16) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + let v = _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + let v = _mm512_cvtps_epi32(v); + let vfl = _mm512_cvtepi32_epi8(v); + unsafe { + _mm_mask_storeu_epi8(p.cast(), mask, vfl); } } unsafe { - r.set_len(n); + r.set_len(this.len()); } r } @@ -64,46 +63,47 @@ mod mul_add_round { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn mul_add_round_v3(this: &[f32], k: f32, b: f32) -> Vec { - let n = this.len(); - let mut r = Vec::::with_capacity(n); - unsafe { - use std::arch::x86_64::*; - let cons = _mm256_setr_epi8( - 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 - -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 - 0, 4, 8, 12, -1, -1, -1, -1, // 16..24 - -1, -1, -1, -1, -1, -1, -1, -1, // 24..32 - ); - let lk = _mm256_set1_ps(k); - let lb = _mm256_set1_ps(b); - let mut n = n; - let mut a = this.as_ptr(); - let mut r = r.as_mut_ptr(); - while n >= 8 { - let x = _mm256_loadu_ps(a); - let v = _mm256_fmadd_ps(x, lk, lb); - let v = _mm256_cvtps_epi32(_mm256_round_ps(v, 0x00)); - let vs = _mm256_shuffle_epi8(v, cons); - let vlo = _mm256_extract_epi32::<0>(vs) as u32; - let vhi = _mm256_extract_epi32::<4>(vs) as u32; - let vfl = vlo as u64 | ((vhi as u64) << 32); - r.cast::().write_unaligned(vfl); - n -= 8; - a = a.add(8); - r = r.add(8); + let mut r = Vec::::with_capacity(this.len()); + use std::arch::x86_64::*; + let cons = _mm256_setr_epi8( + 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 + -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 + 0, 4, 8, 12, -1, -1, -1, -1, // 16..24 + -1, -1, -1, -1, -1, -1, -1, -1, // 24..32 + ); + let lk = _mm256_set1_ps(k); + let lb = _mm256_set1_ps(b); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut p = r.as_mut_ptr(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + let v = _mm256_fmadd_ps(x, lk, lb); + let v = _mm256_cvtps_epi32(_mm256_round_ps(v, 0x00)); + let vs = _mm256_shuffle_epi8(v, cons); + let vlo = _mm256_extract_epi32::<0>(vs) as u32; + let vhi = _mm256_extract_epi32::<4>(vs) as u32; + let vfl = vlo as u64 | ((vhi as u64) << 32); + unsafe { + p.cast::().write_unaligned(vfl); } - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let v = x.mul_add(k, b).round_ties_even() as u8; - r.write(v); - n -= 1; - a = a.add(1); - r = r.add(1); + n -= 8; + a = unsafe { a.add(8) }; + p = unsafe { p.add(8) }; + } + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let v = x.mul_add(k, b).round_ties_even() as u8; + unsafe { + p.write(v); } + n -= 1; + a = unsafe { a.add(1) }; + p = unsafe { p.add(1) }; } unsafe { - r.set_len(n); + r.set_len(this.len()); } r } @@ -134,42 +134,43 @@ mod mul_add_round { #[crate::target_cpu(enable = "v2")] #[target_feature(enable = "fma")] fn mul_add_round_v2_fma(this: &[f32], k: f32, b: f32) -> Vec { - let n = this.len(); - let mut r = Vec::::with_capacity(n); - unsafe { - use std::arch::x86_64::*; - let cons = _mm_setr_epi8( - 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 - -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 - ); - let lk = _mm_set1_ps(k); - let lb = _mm_set1_ps(b); - let mut n = n; - let mut a = this.as_ptr(); - let mut r = r.as_mut_ptr(); - while n >= 4 { - let x = _mm_loadu_ps(a); - let v = _mm_fmadd_ps(x, lk, lb); - let v = _mm_cvtps_epi32(_mm_round_ps(v, 0x00)); - let vs = _mm_shuffle_epi8(v, cons); - let vfl = _mm_extract_epi32::<0>(vs) as u32; - r.cast::().write_unaligned(vfl); - n -= 4; - a = a.add(4); - r = r.add(4); + let mut r = Vec::::with_capacity(this.len()); + use std::arch::x86_64::*; + let cons = _mm_setr_epi8( + 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 + -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 + ); + let lk = _mm_set1_ps(k); + let lb = _mm_set1_ps(b); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut p = r.as_mut_ptr(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + let v = _mm_fmadd_ps(x, lk, lb); + let v = _mm_cvtps_epi32(_mm_round_ps(v, 0x00)); + let vs = _mm_shuffle_epi8(v, cons); + let vfl = _mm_extract_epi32::<0>(vs) as u32; + unsafe { + p.cast::().write_unaligned(vfl); } - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let v = x.mul_add(k, b).round_ties_even() as u8; - r.write(v); - n -= 1; - a = a.add(1); - r = r.add(1); + n -= 4; + a = unsafe { a.add(4) }; + p = unsafe { p.add(4) }; + } + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let v = x.mul_add(k, b).round_ties_even() as u8; + unsafe { + p.write(v); } + n -= 1; + a = unsafe { a.add(1) }; + p = unsafe { p.add(1) }; } unsafe { - r.set_len(n); + r.set_len(this.len()); } r } @@ -199,45 +200,44 @@ mod mul_add_round { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn mul_add_round_a2(this: &[f32], k: f32, b: f32) -> Vec { - let n = this.len(); - let mut r = Vec::::with_capacity(n); - unsafe { - use std::arch::aarch64::*; - let cons = vld1q_u8( - [ - 0, 4, 8, 12, 0xff, 0xff, 0xff, 0xff, // 0..8 - 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // 8..15 - ] - .as_ptr(), - ); - let lk = vdupq_n_f32(k); - let lb = vdupq_n_f32(b); - let mut n = n; - let mut a = this.as_ptr(); - let mut r = r.as_mut_ptr(); - while n >= 4 { - let x = vld1q_f32(a); - let v = vfmaq_f32(lb, x, lk); - let v = vcvtnq_u32_f32(v); - let vs = vqtbl1q_u8(vreinterpretq_u8_u32(v), cons); - let vfl = vgetq_lane_u32::<0>(vreinterpretq_u32_u8(vs)); - r.cast::().write_unaligned(vfl); - n -= 4; - a = a.add(4); - r = r.add(4); + let mut r = Vec::::with_capacity(this.len()); + use std::arch::aarch64::*; + const CONS: [u8; 16] = [ + 0, 4, 8, 12, 0xff, 0xff, 0xff, 0xff, // 0..8 + 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // 8..15 + ]; + let cons = unsafe { vld1q_u8(CONS.as_ptr()) }; + let lk = vdupq_n_f32(k); + let lb = vdupq_n_f32(b); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut p = r.as_mut_ptr(); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + let v = vfmaq_f32(lb, x, lk); + let v = vcvtnq_u32_f32(v); + let vs = vqtbl1q_u8(vreinterpretq_u8_u32(v), cons); + let vfl = vgetq_lane_u32::<0>(vreinterpretq_u32_u8(vs)); + unsafe { + p.cast::().write_unaligned(vfl); } - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - let v = x.mul_add(k, b).round_ties_even() as u8; - r.write(v); - n -= 1; - a = a.add(1); - r = r.add(1); + n -= 4; + a = unsafe { a.add(4) }; + p = unsafe { p.add(4) }; + } + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let v = x.mul_add(k, b).round_ties_even() as u8; + unsafe { + p.write(v); } + n -= 1; + a = unsafe { a.add(1) }; + p = unsafe { p.add(1) }; } unsafe { - r.set_len(n); + r.set_len(this.len()); } r } diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index 4c8e238e..043b3eb2 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -22,25 +22,23 @@ mod reduce_sum_of_x_as_u16 { #[crate::target_cpu(enable = "v4.512")] fn reduce_sum_of_x_as_u16_v4_512(this: &[u8]) -> u16 { use crate::emulate::emulate_mm512_reduce_add_epi16; - unsafe { - use std::arch::x86_64::*; - let us = _mm512_set1_epi16(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm512_setzero_si512(); - while n >= 32 { - let x = _mm256_loadu_si256(a.cast()); - a = a.add(32); - n -= 32; - sum = _mm512_add_epi16(_mm512_and_si512(us, _mm512_cvtepi8_epi16(x)), sum); - } - if n > 0 { - let mask = _bzhi_u32(0xffffffff, n as u32); - let x = _mm256_maskz_loadu_epi8(mask, a.cast()); - sum = _mm512_add_epi16(_mm512_and_si512(us, _mm512_cvtepi8_epi16(x)), sum); - } - emulate_mm512_reduce_add_epi16(sum) as u16 + use std::arch::x86_64::*; + let us = _mm512_set1_epi16(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_si512(); + while n >= 32 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + a = unsafe { a.add(32) }; + n -= 32; + sum = _mm512_add_epi16(_mm512_and_si512(us, _mm512_cvtepi8_epi16(x)), sum); } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x = unsafe { _mm256_maskz_loadu_epi8(mask, a.cast()) }; + sum = _mm512_add_epi16(_mm512_and_si512(us, _mm512_cvtepi8_epi16(x)), sum); + } + emulate_mm512_reduce_add_epi16(sum) as u16 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -69,28 +67,26 @@ mod reduce_sum_of_x_as_u16 { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_x_as_u16_v3(this: &[u8]) -> u16 { use crate::emulate::emulate_mm256_reduce_add_epi16; - unsafe { - use std::arch::x86_64::*; - let us = _mm256_set1_epi16(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm256_setzero_si256(); - while n >= 16 { - let x = _mm_loadu_si128(a.cast()); - a = a.add(16); - n -= 16; - sum = _mm256_add_epi16(_mm256_and_si256(us, _mm256_cvtepi8_epi16(x)), sum); - } - let mut sum = emulate_mm256_reduce_add_epi16(sum) as u16; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - sum += x as u16; - } - sum + use std::arch::x86_64::*; + let us = _mm256_set1_epi16(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_si256(); + while n >= 16 { + let x = unsafe { _mm_loadu_si128(a.cast()) }; + a = unsafe { a.add(16) }; + n -= 16; + sum = _mm256_add_epi16(_mm256_and_si256(us, _mm256_cvtepi8_epi16(x)), sum); + } + let mut sum = emulate_mm256_reduce_add_epi16(sum) as u16; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + sum += x as u16; } + sum } #[cfg(all(target_arch = "x86_64", test))] @@ -119,28 +115,26 @@ mod reduce_sum_of_x_as_u16 { #[crate::target_cpu(enable = "v2")] fn reduce_sum_of_x_as_u16_v2(this: &[u8]) -> u16 { use crate::emulate::emulate_mm_reduce_add_epi16; - unsafe { - use std::arch::x86_64::*; - let us = _mm_set1_epi16(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm_setzero_si128(); - while n >= 8 { - let x = _mm_loadu_si64(a.cast()); - a = a.add(8); - n -= 8; - sum = _mm_add_epi16(_mm_and_si128(us, _mm_cvtepi8_epi16(x)), sum); - } - let mut sum = emulate_mm_reduce_add_epi16(sum) as u16; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - sum += x as u16; - } - sum + use std::arch::x86_64::*; + let us = _mm_set1_epi16(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_si128(); + while n >= 8 { + let x = unsafe { _mm_loadu_si64(a.cast()) }; + a = unsafe { a.add(8) }; + n -= 8; + sum = _mm_add_epi16(_mm_and_si128(us, _mm_cvtepi8_epi16(x)), sum); + } + let mut sum = emulate_mm_reduce_add_epi16(sum) as u16; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + sum += x as u16; } + sum } #[cfg(all(target_arch = "x86_64", test))] @@ -168,28 +162,26 @@ mod reduce_sum_of_x_as_u16 { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_x_as_u16_a2(this: &[u8]) -> u16 { - unsafe { - use std::arch::aarch64::*; - let us = vdupq_n_u16(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = vdupq_n_u16(0); - while n >= 8 { - let x = vld1_u8(a); - a = a.add(8); - n -= 8; - sum = vaddq_u16(vandq_u16(us, vmovl_u8(x)), sum); - } - let mut sum = vaddvq_u16(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - sum += x as u16; - } - sum + use std::arch::aarch64::*; + let us = vdupq_n_u16(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_u16(0); + while n >= 8 { + let x = unsafe { vld1_u8(a) }; + a = unsafe { a.add(8) }; + n -= 8; + sum = vaddq_u16(vandq_u16(us, vmovl_u8(x)), sum); + } + let mut sum = vaddvq_u16(sum); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + sum += x as u16; } + sum } #[cfg(all(target_arch = "aarch64", test, not(miri)))] @@ -234,25 +226,23 @@ mod reduce_sum_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4.512")] fn reduce_sum_of_x_v4_512(this: &[u8]) -> u32 { - unsafe { - use std::arch::x86_64::*; - let us = _mm512_set1_epi32(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm512_setzero_si512(); - while n >= 16 { - let x = _mm_loadu_epi8(a.cast()); - a = a.add(16); - n -= 16; - sum = _mm512_add_epi32(_mm512_and_si512(us, _mm512_cvtepi8_epi32(x)), sum); - } - if n > 0 { - let mask = _bzhi_u32(0xffff, n as u32) as u16; - let x = _mm_maskz_loadu_epi8(mask, a.cast()); - sum = _mm512_add_epi32(_mm512_and_si512(us, _mm512_cvtepi8_epi32(x)), sum); - } - _mm512_reduce_add_epi32(sum) as u32 + use std::arch::x86_64::*; + let us = _mm512_set1_epi32(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_si512(); + while n >= 16 { + let x = unsafe { _mm_loadu_epi8(a.cast()) }; + a = unsafe { a.add(16) }; + n -= 16; + sum = _mm512_add_epi32(_mm512_and_si512(us, _mm512_cvtepi8_epi32(x)), sum); } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm_maskz_loadu_epi8(mask, a.cast()) }; + sum = _mm512_add_epi32(_mm512_and_si512(us, _mm512_cvtepi8_epi32(x)), sum); + } + _mm512_reduce_add_epi32(sum) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -281,28 +271,26 @@ mod reduce_sum_of_x { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_x_v3(this: &[u8]) -> u32 { use crate::emulate::emulate_mm256_reduce_add_epi32; - unsafe { - use std::arch::x86_64::*; - let us = _mm256_set1_epi32(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm256_setzero_si256(); - while n >= 8 { - let x = _mm_loadl_epi64(a.cast()); - a = a.add(8); - n -= 8; - sum = _mm256_add_epi32(_mm256_and_si256(us, _mm256_cvtepi8_epi32(x)), sum); - } - let mut sum = emulate_mm256_reduce_add_epi32(sum) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - sum += x as u32; - } - sum + use std::arch::x86_64::*; + let us = _mm256_set1_epi32(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_si256(); + while n >= 8 { + let x = unsafe { _mm_loadl_epi64(a.cast()) }; + a = unsafe { a.add(8) }; + n -= 8; + sum = _mm256_add_epi32(_mm256_and_si256(us, _mm256_cvtepi8_epi32(x)), sum); + } + let mut sum = emulate_mm256_reduce_add_epi32(sum) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + sum += x as u32; } + sum } #[cfg(all(target_arch = "x86_64", test))] @@ -331,28 +319,26 @@ mod reduce_sum_of_x { #[crate::target_cpu(enable = "v2")] fn reduce_sum_of_x_v2(this: &[u8]) -> u32 { use crate::emulate::emulate_mm_reduce_add_epi32; - unsafe { - use std::arch::x86_64::*; - let us = _mm_set1_epi32(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm_setzero_si128(); - while n >= 4 { - let x = _mm_cvtsi32_si128(a.cast::().read_unaligned()); - a = a.add(4); - n -= 4; - sum = _mm_add_epi32(_mm_and_si128(us, _mm_cvtepi8_epi32(x)), sum); - } - let mut sum = emulate_mm_reduce_add_epi32(sum) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - sum += x as u32; - } - sum + use std::arch::x86_64::*; + let us = _mm_set1_epi32(255); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_si128(); + while n >= 4 { + let x = unsafe { _mm_cvtsi32_si128(a.cast::().read_unaligned()) }; + a = unsafe { a.add(4) }; + n -= 4; + sum = _mm_add_epi32(_mm_and_si128(us, _mm_cvtepi8_epi32(x)), sum); + } + let mut sum = emulate_mm_reduce_add_epi32(sum) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + sum += x as u32; } + sum } #[cfg(all(target_arch = "x86_64", test))] @@ -380,29 +366,27 @@ mod reduce_sum_of_x { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_x_a2(this: &[u8]) -> u32 { - unsafe { - use std::arch::aarch64::*; - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum_0 = vdupq_n_u32(0); - let mut sum_1 = vdupq_n_u32(0); - while n >= 8 { - let x = vmovl_u8(vld1_u8(a.cast())); - a = a.add(8); - n -= 8; - sum_0 = vaddq_u32(vmovl_u16(vget_low_u16(x)), sum_0); - sum_1 = vaddq_u32(vmovl_u16(vget_high_u16(x)), sum_1); - } - let mut sum = vaddvq_u32(vaddq_u32(sum_0, sum_1)); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = a.read(); - a = a.add(1); - n -= 1; - sum += x as u32; - } - sum + use std::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum_0 = vdupq_n_u32(0); + let mut sum_1 = vdupq_n_u32(0); + while n >= 8 { + let x = unsafe { vmovl_u8(vld1_u8(a.cast())) }; + a = unsafe { a.add(8) }; + n -= 8; + sum_0 = vaddq_u32(vmovl_u16(vget_low_u16(x)), sum_0); + sum_1 = vaddq_u32(vmovl_u16(vget_high_u16(x)), sum_1); + } + let mut sum = vaddvq_u32(vaddq_u32(sum_0, sum_1)); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + sum += x as u32; } + sum } #[cfg(all(target_arch = "aarch64", test))] diff --git a/crates/simd_macros/Cargo.toml b/crates/simd_macros/Cargo.toml index c608b3e4..b0d41b15 100644 --- a/crates/simd_macros/Cargo.toml +++ b/crates/simd_macros/Cargo.toml @@ -7,9 +7,9 @@ edition.workspace = true proc-macro = true [dependencies] -proc-macro2 = { version = "1.0.93", features = ["proc-macro"] } -quote = "1.0.38" -syn = { version = "2.0.98", default-features = false, features = [ +proc-macro2 = { version = "1", features = ["proc-macro"] } +quote = "1" +syn = { version = "2", default-features = false, features = [ "clone-impls", "full", "parsing", diff --git a/rust-toolchain.toml b/rust-toolchain.toml index d3e25b13..81febe96 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2025-02-14" +channel = "nightly-2025-03-19" diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 18e30f2a..e13097dd 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -271,7 +271,6 @@ mod vchordrq_cached { use crate::index::storage::PostgresPage; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; - use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] @@ -872,7 +871,7 @@ pub fn make_external_build( LEFT JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace WHERE e.extname = 'vector';"; let pgvector_schema: String = client - .select(schema_query, None, None) + .select(schema_query, None, &[]) .unwrap_or_report() .first() .get_by_name("nspname") @@ -880,7 +879,7 @@ pub fn make_external_build( .expect("external build: cannot get schema of pgvector"); let dump_query = format!("SELECT id, parent, vector::{pgvector_schema}.vector FROM {table};"); - let centroids = client.select(&dump_query, None, None).unwrap_or_report(); + let centroids = client.select(&dump_query, None, &[]).unwrap_or_report(); for row in centroids { let id: Option = row.get_by_name("id").unwrap(); let parent: Option = row.get_by_name("parent").unwrap(); From d66b81289ab92749b12f9849a9ac009a540b360f Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 20 Mar 2025 17:59:09 +0800 Subject: [PATCH 133/324] feat: use mimalloc (#217) Signed-off-by: usamoi --- Cargo.lock | 24 ++++- Cargo.toml | 11 +- crates/algorithm/Cargo.toml | 1 - crates/algorithm/src/lib.rs | 11 +- crates/algorithm/src/maintain.rs | 2 +- crates/algorithm/src/operator.rs | 53 +++------- crates/algorithm/src/search.rs | 21 ++-- crates/algorithm/src/tape.rs | 4 +- crates/algorithm/src/tuples.rs | 29 +++--- crates/algorithm/src/types.rs | 3 + crates/k_means/Cargo.toml | 1 - crates/k_means/src/lib.rs | 8 +- crates/rabitq/src/binary.rs | 15 +-- crates/rabitq/src/block.rs | 35 +++---- crates/rabitq/src/lib.rs | 11 +- crates/simd/src/bit.rs | 91 +++++++---------- crates/simd/src/emulate.rs | 4 +- crates/simd/src/f16.rs | 166 +++++++++++++++++-------------- crates/simd/src/f32.rs | 96 +++++++++--------- crates/simd/src/fast_scan/mod.rs | 43 ++++---- crates/simd/src/lib.rs | 8 +- crates/simd/src/packed_u4.rs | 2 +- crates/simd/src/quantize.rs | 10 +- crates/simd/src/u8.rs | 22 ++-- crates/simd_macros/src/target.rs | 9 +- crates/vector/src/svect.rs | 6 +- src/index/am/am_build.rs | 4 +- src/index/am/mod.rs | 3 + src/lib.rs | 7 +- src/upgrade.rs | 2 +- 30 files changed, 350 insertions(+), 352 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8ed8973a..00ad8dee 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -29,7 +29,6 @@ dependencies = [ "validator", "vector", "zerocopy", - "zerocopy-derive", ] [[package]] @@ -377,7 +376,6 @@ source = "git+https://github.com/tensorchord/half-rs.git?rev=9dc2559e95f295aaa8f dependencies = [ "cfg-if", "crunchy", - "serde", "zerocopy", ] @@ -618,7 +616,6 @@ checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" name = "k_means" version = "0.0.0" dependencies = [ - "half 2.5.0", "rabitq", "rand", "rayon", @@ -647,6 +644,16 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" +[[package]] +name = "libmimalloc-sys" +version = "0.1.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07d0e07885d6a754b9c7993f2625187ad694ee985d60f23355ff0e7077261502" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "litemap" version = "0.7.5" @@ -669,6 +676,15 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +[[package]] +name = "mimalloc" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99585191385958383e13f6b822e6b6d8d9cf928e7d286ceb092da92b43c87bc1" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -1478,6 +1494,7 @@ dependencies = [ "distance", "half 2.5.0", "k_means", + "mimalloc", "paste", "pgrx", "pgrx-catalog", @@ -1489,7 +1506,6 @@ dependencies = [ "validator", "vector", "zerocopy", - "zerocopy-derive", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 234cd1f5..d85cb394 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,7 +37,9 @@ serde.workspace = true toml = "0.8.20" validator.workspace = true zerocopy.workspace = true -zerocopy-derive.workspace = true + +[target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos", target_os = "windows")))'.dependencies] +mimalloc = { version = "*", features = ["local_dynamic_tls"] } [patch.crates-io] half = { git = "https://github.com/tensorchord/half-rs.git", rev = "9dc2559e95f295aaa8ffef5e0529167731dffd47" } @@ -54,13 +56,12 @@ version = "0.0.0" edition = "2024" [workspace.dependencies] -half = { version = "2.5.0", features = ["serde", "zerocopy"] } +half = { version = "2.5.0", features = ["zerocopy"] } paste = "1" rand = "0.9.0" -serde = "1" +serde = { version = "1", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } -zerocopy = "0.8.23" -zerocopy-derive = "0.8.23" +zerocopy = { version = "0.8.23", features = ["derive"] } [workspace.lints] clippy.identity_op = "allow" diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml index 7ca1c37c..13fae996 100644 --- a/crates/algorithm/Cargo.toml +++ b/crates/algorithm/Cargo.toml @@ -19,7 +19,6 @@ serde.workspace = true turboselect = { git = "https://github.com/tensorchord/turboselect.git", rev = "d8753c4ffe5b47f28670fea21d56cf3658d51b9b" } validator.workspace = true zerocopy.workspace = true -zerocopy-derive.workspace = true [lints] workspace = true diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 10c861c9..0baebc8c 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -1,8 +1,3 @@ -#![feature(generic_arg_infer)] -#![allow(clippy::collapsible_else_if)] -#![allow(clippy::type_complexity)] -#![allow(clippy::len_without_is_empty)] - mod build; mod bulkdelete; mod cache; @@ -31,7 +26,7 @@ pub use rerank::{how, rerank_heap, rerank_index}; pub use search::{search, search_and_estimate}; use std::ops::{Deref, DerefMut}; -use zerocopy_derive::{FromBytes, Immutable, IntoBytes, KnownLayout}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] @@ -40,11 +35,13 @@ pub struct Opaque { pub skip: u32, } -#[allow(clippy::len_without_is_empty)] pub trait Page: Sized { fn get_opaque(&self) -> &Opaque; fn get_opaque_mut(&mut self) -> &mut Opaque; fn len(&self) -> u16; + fn is_empty(&self) -> bool { + self.len() == 0 + } fn get(&self, i: u16) -> Option<&[u8]>; fn get_mut(&mut self, i: u16) -> Option<&mut [u8]>; fn alloc(&mut self, data: &[u8]) -> Option; diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index f079e807..db3b6c64 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -63,7 +63,7 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let mut trace = Vec::new(); let mut tuples = 0_u64; - let mut callback = |code: (f32, f32, f32, f32, Vec<_>), mean, payload| { + let mut callback = |code: (_, _, _, _, _), mean, payload| { tape.push(Branch { mean, dis_u_2: code.0, diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index ae80180d..fcd6ff12 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -1,6 +1,8 @@ use crate::types::*; use distance::Distance; use half::f16; +use rabitq::binary::{BinaryCode, BinaryLut}; +use rabitq::block::BlockLut; use simd::Floating; use std::fmt::Debug; use std::marker::PhantomData; @@ -183,7 +185,7 @@ impl type Output = [Distance; 32]; fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { - let t = simd::fast_scan::fast_scan(input, target); + let t = simd::fast_scan::scan(input, target); for i in 0..32 { self.0[i] += t[i]; } @@ -221,7 +223,7 @@ impl type Output = [Distance; 32]; fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { - let t = simd::fast_scan::fast_scan(input, target); + let t = simd::fast_scan::scan(input, target); for i in 0..32 { self.0[i] += t[i]; } @@ -383,14 +385,9 @@ pub trait Vector: VectorOwned { fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); fn from_owned(vector: OwnedVector) -> Self; - fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u8; 16]>); + fn compute_lut_block(vector: Self::Borrowed<'_>) -> BlockLut; - fn compute_lut( - vector: Self::Borrowed<'_>, - ) -> ( - (f32, f32, f32, f32, Vec<[u8; 16]>), - (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), - ); + fn compute_lut(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut); fn code(vector: Self::Borrowed<'_>) -> rabitq::Code; } @@ -423,16 +420,11 @@ impl Vector for VectOwned { } } - fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u8; 16]>) { + fn compute_lut_block(vector: Self::Borrowed<'_>) -> BlockLut { rabitq::block::preprocess(vector.slice()) } - fn compute_lut( - vector: Self::Borrowed<'_>, - ) -> ( - (f32, f32, f32, f32, Vec<[u8; 16]>), - (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), - ) { + fn compute_lut(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { rabitq::compute_lut(vector.slice()) } @@ -469,16 +461,11 @@ impl Vector for VectOwned { } } - fn compute_lut_block(vector: Self::Borrowed<'_>) -> (f32, f32, f32, f32, Vec<[u8; 16]>) { + fn compute_lut_block(vector: Self::Borrowed<'_>) -> BlockLut { rabitq::block::preprocess(&f16::vector_to_f32(vector.slice())) } - fn compute_lut( - vector: Self::Borrowed<'_>, - ) -> ( - (f32, f32, f32, f32, Vec<[u8; 16]>), - (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), - ) { + fn compute_lut(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { rabitq::compute_lut(&f16::vector_to_f32(vector.slice())) } @@ -490,11 +477,7 @@ impl Vector for VectOwned { pub trait OperatorDistance: 'static + Debug + Copy { const KIND: DistanceKind; - fn compute_lowerbound_binary( - lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), - code: (f32, f32, f32, f32, &[u64]), - epsilon: f32, - ) -> Distance; + fn compute_lowerbound_binary(lut: &BinaryLut, code: BinaryCode<'_>, epsilon: f32) -> Distance; type BlockAccessor: for<'a> Accessor2< [u8; 16], @@ -515,11 +498,7 @@ pub struct L2; impl OperatorDistance for L2 { const KIND: DistanceKind = DistanceKind::L2; - fn compute_lowerbound_binary( - lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), - code: (f32, f32, f32, f32, &[u64]), - epsilon: f32, - ) -> Distance { + fn compute_lowerbound_binary(lut: &BinaryLut, code: BinaryCode<'_>, epsilon: f32) -> Distance { rabitq::binary::process_lowerbound_l2(lut, code, epsilon) } @@ -532,11 +511,7 @@ pub struct Dot; impl OperatorDistance for Dot { const KIND: DistanceKind = DistanceKind::Dot; - fn compute_lowerbound_binary( - lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), - code: (f32, f32, f32, f32, &[u64]), - epsilon: f32, - ) -> Distance { + fn compute_lowerbound_binary(lut: &BinaryLut, code: BinaryCode<'_>, epsilon: f32) -> Distance { rabitq::binary::process_lowerbound_dot(lut, code, epsilon) } @@ -570,7 +545,7 @@ pub trait Operator: 'static + Debug + Copy { } #[derive(Debug)] -pub struct Op(PhantomData<(fn(V) -> V, fn(D) -> D)>); +pub struct Op(PhantomData V>, PhantomData D>); impl Clone for Op { fn clone(&self) -> Self { diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index edb6b7ff..76084023 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -9,16 +9,18 @@ use std::collections::BinaryHeap; use std::num::NonZeroU64; use vector::{VectorBorrowed, VectorOwned}; +type Results = Vec<( + Reverse, + AlwaysEqual, + AlwaysEqual, +)>; + pub fn search( index: impl RelationRead, vector: O::Vector, probes: Vec, epsilon: f32, -) -> Vec<( - Reverse, - AlwaysEqual, - AlwaysEqual, -)> { +) -> Results { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -169,14 +171,7 @@ pub fn search_and_estimate( probes: Vec, epsilon: f32, mut t: u32, -) -> ( - Vec<( - Reverse, - AlwaysEqual, - AlwaysEqual, - )>, - Distance, -) { +) -> (Results, Distance) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index 53996734..98b34ffd 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -1,6 +1,7 @@ use crate::operator::Accessor1; use crate::tuples::*; use crate::{IndexPointer, Page, PageGuard, RelationRead, RelationWrite}; +use rabitq::binary::BinaryCode; use std::marker::PhantomData; use std::num::NonZeroU64; use std::ops::DerefMut; @@ -164,7 +165,7 @@ pub fn read_frozen_tape( pub fn read_appendable_tape( index: impl RelationRead, first: u32, - mut access: impl FnMut((f32, f32, f32, f32, &[u64])) -> T, + mut access: impl for<'a> FnMut(BinaryCode<'a>) -> T, mut callback: impl FnMut(T, IndexPointer, NonZeroU64), mut step: impl FnMut(u32), ) { @@ -185,6 +186,7 @@ pub fn read_appendable_tape( } } +#[allow(clippy::collapsible_else_if)] pub fn append( index: impl RelationWrite, first: u32, diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index a0fd1b46..31b72234 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -1,8 +1,9 @@ use crate::IndexPointer; use crate::operator::Vector; +use rabitq::binary::BinaryCode; use std::marker::PhantomData; use std::num::{NonZeroU8, NonZeroU64}; -use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; +use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; @@ -141,9 +142,9 @@ pub struct FreepageTuple {} impl Tuple for FreepageTuple { fn serialize(&self) -> Vec { FreepageTupleHeader { - level_0: [0; _], - level_1: [0; _], - level_2: [0; _], + level_0: FromZeros::new_zeroed(), + level_1: FromZeros::new_zeroed(), + level_2: FromZeros::new_zeroed(), _padding_0: Default::default(), } .as_bytes() @@ -172,22 +173,21 @@ impl FreepageTupleWriter<'_> { set(&mut self.header.level_1[i >> 10], (i >> 5) % 32, true); set(&mut self.header.level_2[i >> 15], (i >> 10) % 32, true); } - #[allow(clippy::just_underscores_and_digits)] fn find(&self) -> Option { - let _3 = 0_usize; - let _2 = self.header.level_2[_3 << 0].trailing_zeros() as usize; - if _2 == 32 { + let i_3 = 0_usize; + let i_2 = self.header.level_2[i_3 << 0].trailing_zeros() as usize; + if i_2 == 32 { return None; } - let _1 = self.header.level_1[_3 << 5 | _2 << 0].trailing_zeros() as usize; - if _1 == 32 { + let i_1 = self.header.level_1[i_3 << 5 | i_2 << 0].trailing_zeros() as usize; + if i_1 == 32 { panic!("data corruption"); } - let _0 = self.header.level_0[_3 << 10 | _2 << 5 | _1 << 0].trailing_zeros() as usize; - if _0 == 32 { + let i_0 = self.header.level_0[i_3 << 10 | i_2 << 5 | i_1 << 0].trailing_zeros() as usize; + if i_0 == 32 { panic!("data corruption"); } - Some(_3 << 15 | _2 << 10 | _1 << 5 | _0 << 0) + Some(i_3 << 15 | i_2 << 10 | i_1 << 5 | i_0 << 0) } pub fn fetch(&mut self) -> Option { let i = self.find()?; @@ -873,7 +873,6 @@ struct AppendableTupleHeader { elements_e: usize, } -#[allow(clippy::large_enum_variant)] #[derive(Debug, Clone, PartialEq)] pub struct AppendableTuple { pub mean: IndexPointer, @@ -941,7 +940,7 @@ impl<'a> AppendableTupleReader<'a> { pub fn mean(self) -> IndexPointer { self.header.mean } - pub fn code(self) -> (f32, f32, f32, f32, &'a [u64]) { + pub fn code(self) -> BinaryCode<'a> { ( self.header.dis_u_2, self.header.factor_ppc, diff --git a/crates/algorithm/src/types.rs b/crates/algorithm/src/types.rs index ea9784af..30299fab 100644 --- a/crates/algorithm/src/types.rs +++ b/crates/algorithm/src/types.rs @@ -81,4 +81,7 @@ impl Structure { pub fn len(&self) -> usize { self.children.len() } + pub fn is_empty(&self) -> bool { + self.children.is_empty() + } } diff --git a/crates/k_means/Cargo.toml b/crates/k_means/Cargo.toml index 6ba2ccf7..00e5efe1 100644 --- a/crates/k_means/Cargo.toml +++ b/crates/k_means/Cargo.toml @@ -7,7 +7,6 @@ edition.workspace = true rabitq = { path = "../rabitq" } simd = { path = "../simd" } -half.workspace = true rand.workspace = true rayon = "1.10.0" diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index 7f61328b..cdc475fd 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -1,6 +1,4 @@ -#![allow(clippy::type_complexity)] - -use half::f16; +use rabitq::block::BlockCode; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use rayon::iter::{IntoParallelIterator, ParallelIterator}; @@ -126,7 +124,7 @@ fn rabitq_index( elements: Vec<[u8; 16]>, } impl Block { - fn code(&self) -> (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32], &[[u8; 16]]) { + fn code(&self) -> BlockCode<'_> { ( &self.dis_u_2, &self.factor_ppc, @@ -201,7 +199,7 @@ struct LloydKMeans<'a, F> { compute: F, } -const DELTA: f32 = f16::EPSILON.to_f32_const(); +const DELTA: f32 = 9.7656e-4_f32; impl<'a, F: Fn(&[Vec]) -> Vec> LloydKMeans<'a, F> { fn new(c: usize, dims: usize, samples: &'a [Vec], is_spherical: bool, compute: F) -> Self { diff --git a/crates/rabitq/src/binary.rs b/crates/rabitq/src/binary.rs index 04bcb6f8..267ecaee 100644 --- a/crates/rabitq/src/binary.rs +++ b/crates/rabitq/src/binary.rs @@ -1,9 +1,10 @@ use distance::Distance; use simd::Floating; -pub fn preprocess( - vector: &[f32], -) -> (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)) { +pub type BinaryLut = (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)); +pub type BinaryCode<'a> = (f32, f32, f32, f32, &'a [u64]); + +pub fn preprocess(vector: &[f32]) -> BinaryLut { let dis_v_2 = f32::reduce_sum_of_x2(vector); let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); let qvector_sum = if vector.len() <= 4369 { @@ -15,8 +16,8 @@ pub fn preprocess( } pub fn process_lowerbound_l2( - lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), - (dis_u_2, factor_ppc, factor_ip, factor_err, t): (f32, f32, f32, f32, &[u64]), + lut: &BinaryLut, + (dis_u_2, factor_ppc, factor_ip, factor_err, t): BinaryCode<'_>, epsilon: f32, ) -> Distance { let &(dis_v_2, b, k, qvector_sum, ref s) = lut; @@ -28,8 +29,8 @@ pub fn process_lowerbound_l2( } pub fn process_lowerbound_dot( - lut: &(f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), - (_, factor_ppc, factor_ip, factor_err, t): (f32, f32, f32, f32, &[u64]), + lut: &BinaryLut, + (_, factor_ppc, factor_ip, factor_err, t): BinaryCode<'_>, epsilon: f32, ) -> Distance { let &(dis_v_2, b, k, qvector_sum, ref s) = lut; diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs index 52d4a82d..174c8fe6 100644 --- a/crates/rabitq/src/block.rs +++ b/crates/rabitq/src/block.rs @@ -1,7 +1,16 @@ use distance::Distance; use simd::Floating; -pub fn preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec<[u8; 16]>) { +pub type BlockLut = (f32, f32, f32, f32, Vec<[u8; 16]>); +pub type BlockCode<'a> = ( + &'a [f32; 32], + &'a [f32; 32], + &'a [f32; 32], + &'a [f32; 32], + &'a [[u8; 16]], +); + +pub fn preprocess(vector: &[f32]) -> BlockLut { let dis_v_2 = f32::reduce_sum_of_x2(vector); let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); let qvector_sum = if vector.len() <= 4369 { @@ -13,18 +22,12 @@ pub fn preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec<[u8; 16]>) { } pub fn process_lowerbound_l2( - lut: &(f32, f32, f32, f32, Vec<[u8; 16]>), - (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[[u8; 16]], - ), + lut: &BlockLut, + (dis_u_2, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, epsilon: f32, ) -> [Distance; 32] { let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let r = simd::fast_scan::fast_scan(t, s); + let r = simd::fast_scan::scan(t, s); std::array::from_fn(|i| { let rough = dis_u_2[i] + dis_v_2 @@ -36,18 +39,12 @@ pub fn process_lowerbound_l2( } pub fn process_lowerbound_dot( - lut: &(f32, f32, f32, f32, Vec<[u8; 16]>), - (_, factor_ppc, factor_ip, factor_err, t): ( - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[[u8; 16]], - ), + lut: &BlockLut, + (_, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, epsilon: f32, ) -> [Distance; 32] { let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let r = simd::fast_scan::fast_scan(t, s); + let r = simd::fast_scan::scan(t, s); std::array::from_fn(|i| { let rough = 0.5 * b * factor_ppc[i] + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index a90f4453..0215d252 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -1,8 +1,8 @@ -#![allow(clippy::type_complexity)] - pub mod binary; pub mod block; +use binary::BinaryLut; +use block::BlockLut; use simd::Floating; #[derive(Debug, Clone)] @@ -74,12 +74,7 @@ pub fn code(dims: u32, vector: &[f32]) -> Code { } } -pub fn compute_lut( - vector: &[f32], -) -> ( - (f32, f32, f32, f32, Vec<[u8; 16]>), - (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)), -) { +pub fn compute_lut(vector: &[f32]) -> (BlockLut, BinaryLut) { use simd::Floating; let dis_v_2 = f32::reduce_sum_of_x2(vector); let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index 67e09e34..1e587bb7 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -8,9 +8,9 @@ mod sum_of_and { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] + #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_and_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + fn sum_of_and_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut and = _mm512_setzero_si512(); @@ -36,24 +36,21 @@ mod sum_of_and { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_and_v4_512_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!( - "test {} ... skipped (v4.512:avx512vpopcntdq)", - module_path!() - ); + fn sum_of_and_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_and_v4_512_avx512vpopcntdq(&lhs, &rhs) }; + let specialized = unsafe { sum_of_and_v4_avx512vpopcntdq(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] pub fn sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -75,9 +72,9 @@ mod sum_of_or { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] + #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_or_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + fn sum_of_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut or = _mm512_setzero_si512(); @@ -103,24 +100,21 @@ mod sum_of_or { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_or_v4_512_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!( - "test {} ... skipped (v4.512:avx512vpopcntdq)", - module_path!() - ); + fn sum_of_or_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_or_v4_512_avx512vpopcntdq(&lhs, &rhs) }; + let specialized = unsafe { sum_of_or_v4_avx512vpopcntdq(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] pub fn sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -142,9 +136,9 @@ mod sum_of_xor { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] + #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_xor_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + fn sum_of_xor_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut xor = _mm512_setzero_si512(); @@ -170,24 +164,21 @@ mod sum_of_xor { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_xor_v4_512_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!( - "test {} ... skipped (v4.512:avx512vpopcntdq)", - module_path!() - ); + fn sum_of_xor_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_xor_v4_512_avx512vpopcntdq(&lhs, &rhs) }; + let specialized = unsafe { sum_of_xor_v4_avx512vpopcntdq(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] pub fn sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -209,9 +200,9 @@ mod sum_of_and_or { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] + #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_and_or_v4_512_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + fn sum_of_and_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut and = _mm512_setzero_si512(); @@ -243,24 +234,21 @@ mod sum_of_and_or { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_xor_v4_512_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!( - "test {} ... skipped (v4.512:avx512vpopcntdq)", - module_path!() - ); + fn sum_of_xor_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_and_or_v4_512_avx512vpopcntdq(&lhs, &rhs) }; + let specialized = unsafe { sum_of_and_or_v4_avx512vpopcntdq(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] pub fn sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -284,9 +272,9 @@ mod sum_of_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] + #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_x_v4_512_avx512vpopcntdq(this: &[u64]) -> u32 { + fn sum_of_x_v4_avx512vpopcntdq(this: &[u64]) -> u32 { use std::arch::x86_64::*; let mut and = _mm512_setzero_si512(); let mut a = this.as_ptr(); @@ -307,23 +295,20 @@ mod sum_of_x { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_x_v4_512_avx512vpopcntdq_test() { - if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512vpopcntdq") { - println!( - "test {} ... skipped (v4.512:avx512vpopcntdq)", - module_path!() - ); + fn sum_of_x_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let this = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_x_v4_512_avx512vpopcntdq(&this) }; + let specialized = unsafe { sum_of_x_v4_avx512vpopcntdq(&this) }; let fallback = fallback(&this); assert_eq!(specialized, fallback); } } - #[crate::multiversion(@"v4.512:avx512vpopcntdq", "v4.512", "v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] pub fn sum_of_x(this: &[u64]) -> u32 { let n = this.len(); let mut and = 0; @@ -340,7 +325,7 @@ pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_and { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -363,7 +348,7 @@ pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_or { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -386,7 +371,7 @@ pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_xor { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index 4a3ebe68..7a75042a 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -3,7 +3,7 @@ // Instructions. arXiv preprint arXiv:2112.06342. #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v4.512")] +#[crate::target_cpu(enable = "v4")] pub fn emulate_mm512_2intersect_epi32( a: std::arch::x86_64::__m512i, b: std::arch::x86_64::__m512i, @@ -79,7 +79,7 @@ pub fn emulate_mm_reduce_add_ps(mut x: std::arch::x86_64::__m128) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v4.512")] +#[crate::target_cpu(enable = "v4")] pub fn emulate_mm512_reduce_add_epi16(x: std::arch::x86_64::__m512i) -> i16 { use std::arch::x86_64::*; _mm256_reduce_add_epi16(_mm512_castsi512_si256(x)) diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index d385fd40..7b5eed71 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -1,10 +1,22 @@ use crate::{Floating, f32}; use half::f16; +use zerocopy::FromZeros; + +trait AsF32 { + #[allow(clippy::wrong_self_convention)] + fn as_f32(self) -> f32; +} + +impl AsF32 for f16 { + fn as_f32(self) -> f32 { + self.into() + } +} impl Floating for f16 { #[inline(always)] fn zero() -> Self { - f16::ZERO + FromZeros::new_zeroed() } #[inline(always)] @@ -129,12 +141,12 @@ impl Floating for f16 { } mod reduce_or_of_is_zero_x { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { for &x in this { - if x == f16::ZERO { + if x == FromZeros::new_zeroed() { return true; } } @@ -145,14 +157,14 @@ mod reduce_or_of_is_zero_x { mod reduce_sum_of_x { // FIXME: add manually-implemented SIMD version - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_sum_of_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; for i in 0..n { - x += this[i].to_f32(); + x += this[i].as_f32(); } x } @@ -161,14 +173,14 @@ mod reduce_sum_of_x { mod reduce_sum_of_abs_x { // FIXME: add manually-implemented SIMD version - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; for i in 0..n { - x += this[i].to_f32().abs(); + x += (this[i].as_f32()).abs(); } x } @@ -177,14 +189,14 @@ mod reduce_sum_of_abs_x { mod reduce_sum_of_x2 { // FIXME: add manually-implemented SIMD version - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_sum_of_x2(this: &[f16]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; for i in 0..n { - x2 += this[i].to_f32() * this[i].to_f32(); + x2 += this[i].as_f32() * this[i].as_f32(); } x2 } @@ -193,29 +205,29 @@ mod reduce_sum_of_x2 { mod reduce_min_max_of_x { // FIXME: add manually-implemented SIMD version - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; let n = this.len(); for i in 0..n { - min = min.min(this[i].to_f32()); - max = max.max(this[i].to_f32()); + min = min.min(this[i].as_f32()); + max = max.max(this[i].as_f32()); } (min, max) } } mod reduce_sum_of_xy { - use half::f16; + use super::*; #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] + #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512fp16")] - pub fn reduce_sum_of_xy_v4_512_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + pub fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut n = lhs.len(); @@ -241,11 +253,11 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_xy_v4_512_avx512fp16_test() { + fn reduce_sum_of_xy_v4_avx512fp16_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512fp16") { - println!("test {} ... skipped (v4_512:avx512fp16)", module_path!()); + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -260,7 +272,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v4_512_avx512fp16(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v4_avx512fp16(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -272,8 +284,8 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - pub fn reduce_sum_of_xy_v4_512(lhs: &[f16], rhs: &[f16]) -> f32 { + #[crate::target_cpu(enable = "v4")] + pub fn reduce_sum_of_xy_v4(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut n = lhs.len(); @@ -302,7 +314,7 @@ mod reduce_sum_of_xy { fn reduce_sum_of_xy_v4_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -315,7 +327,7 @@ mod reduce_sum_of_xy { let rhs = (0..n) .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); - let specialized = unsafe { reduce_sum_of_xy_v4_512(&lhs, &rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v4(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -345,8 +357,8 @@ mod reduce_sum_of_xy { } let mut xy = emulate_mm256_reduce_add_ps(xy); while n > 0 { - let x = unsafe { a.read().to_f32() }; - let y = unsafe { b.read().to_f32() }; + let x = unsafe { a.read().as_f32() }; + let y = unsafe { b.read().as_f32() }; a = unsafe { a.add(1) }; b = unsafe { b.add(1) }; n -= 1; @@ -481,26 +493,26 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4.512:avx512fp16", @"v4.512", @"v3", @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); let mut xy = 0.0f32; for i in 0..n { - xy += lhs[i].to_f32() * rhs[i].to_f32(); + xy += lhs[i].as_f32() * rhs[i].as_f32(); } xy } } mod reduce_sum_of_d2 { - use half::f16; + use super::*; #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] + #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512fp16")] - pub fn reduce_sum_of_d2_v4_512_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + pub fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut n = lhs.len() as u32; @@ -528,11 +540,11 @@ mod reduce_sum_of_d2 { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_d2_v4_512_avx512fp16_test() { + fn reduce_sum_of_d2_v4_avx512fp16_test() { use rand::Rng; const EPSILON: f32 = 6.4; - if !crate::is_cpu_detected!("v4.512") || !crate::is_feature_detected!("avx512fp16") { - println!("test {} ... skipped (v4_512:avx512fp16)", module_path!()); + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -547,7 +559,7 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v4_512_avx512fp16(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_d2_v4_avx512fp16(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -559,8 +571,8 @@ mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - pub fn reduce_sum_of_d2_v4_512(lhs: &[f16], rhs: &[f16]) -> f32 { + #[crate::target_cpu(enable = "v4")] + pub fn reduce_sum_of_d2_v4(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut n = lhs.len() as u32; @@ -591,7 +603,7 @@ mod reduce_sum_of_d2 { fn reduce_sum_of_d2_v4_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -607,7 +619,7 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v4_512(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_d2_v4(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -639,8 +651,8 @@ mod reduce_sum_of_d2 { } let mut d2 = emulate_mm256_reduce_add_ps(d2); while n > 0 { - let x = unsafe { a.read().to_f32() }; - let y = unsafe { b.read().to_f32() }; + let x = unsafe { a.read().as_f32() }; + let y = unsafe { b.read().as_f32() }; a = unsafe { a.add(1) }; b = unsafe { b.add(1) }; n -= 1; @@ -776,13 +788,13 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4.512:avx512fp16", @"v4.512", @"v3", @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); let mut d2 = 0.0; for i in 0..n { - let d = lhs[i].to_f32() - rhs[i].to_f32(); + let d = lhs[i].as_f32() - rhs[i].as_f32(); d2 += d * d; } d2 @@ -793,9 +805,9 @@ mod reduce_sum_of_xy_sparse { // There is no manually-implemented SIMD version. // Add it if `svecf16` is supported. - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -806,7 +818,7 @@ mod reduce_sum_of_xy_sparse { while lp < ln && rp < rn { match Ord::cmp(&lidx[lp], &ridx[rp]) { Ordering::Equal => { - xy += lval[lp].to_f32() * rval[rp].to_f32(); + xy += lval[lp].as_f32() * rval[rp].as_f32(); lp += 1; rp += 1; } @@ -826,9 +838,9 @@ mod reduce_sum_of_d2_sparse { // There is no manually-implemented SIMD version. // Add it if `svecf16` is supported. - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -839,35 +851,35 @@ mod reduce_sum_of_d2_sparse { while lp < ln && rp < rn { match Ord::cmp(&lidx[lp], &ridx[rp]) { Ordering::Equal => { - let d = lval[lp].to_f32() - rval[rp].to_f32(); + let d = lval[lp].as_f32() - rval[rp].as_f32(); d2 += d * d; lp += 1; rp += 1; } Ordering::Less => { - d2 += lval[lp].to_f32() * lval[lp].to_f32(); + d2 += lval[lp].as_f32() * lval[lp].as_f32(); lp += 1; } Ordering::Greater => { - d2 += rval[rp].to_f32() * rval[rp].to_f32(); + d2 += rval[rp].as_f32() * rval[rp].as_f32(); rp += 1; } } } for i in lp..ln { - d2 += lval[i].to_f32() * lval[i].to_f32(); + d2 += lval[i].as_f32() * lval[i].as_f32(); } for i in rp..rn { - d2 += rval[i].to_f32() * rval[i].to_f32(); + d2 += rval[i].as_f32() * rval[i].as_f32(); } d2 } } mod vector_add { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_add(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -885,9 +897,9 @@ mod vector_add { } mod vector_add_inplace { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_add_inplace(lhs: &mut [f16], rhs: &[f16]) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -898,9 +910,9 @@ mod vector_add_inplace { } mod vector_sub { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_sub(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -918,9 +930,9 @@ mod vector_sub { } mod vector_mul { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_mul(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -938,9 +950,9 @@ mod vector_mul { } mod vector_mul_scalar { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_mul_scalar(lhs: &[f16], rhs: f32) -> Vec { let rhs = f16::from_f32(rhs); let n = lhs.len(); @@ -958,9 +970,9 @@ mod vector_mul_scalar { } mod vector_mul_scalar_inplace { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_mul_scalar_inplace(lhs: &mut [f16], rhs: f32) { let rhs = f16::from_f32(rhs); let n = lhs.len(); @@ -971,21 +983,21 @@ mod vector_mul_scalar_inplace { } mod vector_abs_inplace { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_abs_inplace(this: &mut [f16]) { let n = this.len(); for i in 0..n { - this[i] = f16::from_f32(this[i].to_f32().abs()); + this[i] = f16::from_f32(this[i].as_f32().abs()); } } } mod vector_from_f32 { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_from_f32(this: &[f32]) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); @@ -1002,15 +1014,15 @@ mod vector_from_f32 { } mod vector_to_f32 { - use half::f16; + use super::*; - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_to_f32(this: &[f16]) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); for i in 0..n { unsafe { - r.as_mut_ptr().add(i).write(this[i].to_f32()); + r.as_mut_ptr().add(i).write(this[i].as_f32()); } } unsafe { diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index f2991df4..6bd0bfa1 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -119,7 +119,7 @@ impl Floating for f32 { } mod reduce_or_of_is_zero_x { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { for &x in this { if x == 0.0f32 { @@ -133,8 +133,8 @@ mod reduce_or_of_is_zero_x { mod reduce_sum_of_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_sum_of_x_v4_512(this: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[f32]) -> f32 { use std::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); @@ -158,7 +158,7 @@ mod reduce_sum_of_x { fn reduce_sum_of_x_v4_test() { use rand::Rng; const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -170,7 +170,7 @@ mod reduce_sum_of_x { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_v4_512(this) }; + let specialized = unsafe { reduce_sum_of_x_v4(this) }; let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, @@ -384,7 +384,7 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a3.256", @"a2")] pub fn reduce_sum_of_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -398,8 +398,8 @@ mod reduce_sum_of_x { mod reduce_sum_of_abs_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_sum_of_abs_x_v4_512(this: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_abs_x_v4(this: &[f32]) -> f32 { use std::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); @@ -425,7 +425,7 @@ mod reduce_sum_of_abs_x { fn reduce_sum_of_abs_x_v4_test() { use rand::Rng; const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -437,7 +437,7 @@ mod reduce_sum_of_abs_x { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_abs_x_v4_512(this) }; + let specialized = unsafe { reduce_sum_of_abs_x_v4(this) }; let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, @@ -660,7 +660,7 @@ mod reduce_sum_of_abs_x { } } - #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a3.256", @"a2")] pub fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -674,8 +674,8 @@ mod reduce_sum_of_abs_x { mod reduce_sum_of_x2 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_sum_of_x2_v4_512(this: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x2_v4(this: &[f32]) -> f32 { use std::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); @@ -699,7 +699,7 @@ mod reduce_sum_of_x2 { fn reduce_sum_of_x2_v4_test() { use rand::Rng; const EPSILON: f32 = 0.006; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -711,7 +711,7 @@ mod reduce_sum_of_x2 { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x2_v4_512(this) }; + let specialized = unsafe { reduce_sum_of_x2_v4(this) }; let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, @@ -926,7 +926,7 @@ mod reduce_sum_of_x2 { } } - #[crate::multiversion(@"v4.512", @"v3", @"v2:fma", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a3.256", @"a2")] pub fn reduce_sum_of_x2(this: &[f32]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; @@ -943,8 +943,8 @@ mod reduce_min_max_of_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_min_max_of_x_v4_512(this: &[f32]) -> (f32, f32) { + #[crate::target_cpu(enable = "v4")] + fn reduce_min_max_of_x_v4(this: &[f32]) -> (f32, f32) { use std::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); @@ -972,7 +972,7 @@ mod reduce_min_max_of_x { #[test] fn reduce_min_max_of_x_v4_test() { use rand::Rng; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -984,7 +984,7 @@ mod reduce_min_max_of_x { .collect::>(); for z in 50..200 { let x = &x[..z]; - let specialized = unsafe { reduce_min_max_of_x_v4_512(x) }; + let specialized = unsafe { reduce_min_max_of_x_v4(x) }; let fallback = fallback(x); assert_eq!(specialized.0, fallback.0); assert_eq!(specialized.1, fallback.1); @@ -1198,7 +1198,7 @@ mod reduce_min_max_of_x { } } - #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a3.256", @"a2")] pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -1214,8 +1214,8 @@ mod reduce_min_max_of_x { mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_sum_of_xy_v4_512(lhs: &[f32], rhs: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_v4(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut n = lhs.len(); @@ -1244,7 +1244,7 @@ mod reduce_sum_of_xy { fn reduce_sum_of_xy_v4_test() { use rand::Rng; const EPSILON: f32 = 0.004; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -1260,7 +1260,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v4_512(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -1516,7 +1516,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4.512", @"v3", @"v2:fma", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a3.256", @"a2")] pub fn reduce_sum_of_xy(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -1531,8 +1531,8 @@ mod reduce_sum_of_xy { mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_sum_of_d2_v4_512(lhs: &[f32], rhs: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_d2_v4(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut n = lhs.len() as u32; @@ -1563,7 +1563,7 @@ mod reduce_sum_of_d2 { fn reduce_sum_of_d2_v4_test() { use rand::Rng; const EPSILON: f32 = 0.02; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -1579,7 +1579,7 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v4_512(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_d2_v4(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -1842,7 +1842,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4.512", @"v3", @"v2:fma", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a3.256", @"a2")] pub fn reduce_sum_of_d2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -1858,8 +1858,8 @@ mod reduce_sum_of_d2 { mod reduce_sum_of_xy_sparse { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_sum_of_xy_sparse_v4_512(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_sparse_v4(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { use crate::emulate::emulate_mm512_2intersect_epi32; assert_eq!(li.len(), lv.len()); assert_eq!(ri.len(), rv.len()); @@ -1907,7 +1907,7 @@ mod reduce_sum_of_xy_sparse { fn reduce_sum_of_xy_sparse_v4_test() { use rand::Rng; const EPSILON: f32 = 0.000001; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -1935,7 +1935,7 @@ mod reduce_sum_of_xy_sparse { let rval = (0..rm) .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); - let specialized = unsafe { reduce_sum_of_xy_sparse_v4_512(&lidx, &lval, &ridx, &rval) }; + let specialized = unsafe { reduce_sum_of_xy_sparse_v4(&lidx, &lval, &ridx, &rval) }; let fallback = fallback(&lidx, &lval, &ridx, &rval); assert!( (specialized - fallback).abs() < EPSILON, @@ -1944,7 +1944,7 @@ mod reduce_sum_of_xy_sparse { } } - #[crate::multiversion(@"v4.512", "v3", "v2", "a2")] + #[crate::multiversion(@"v4", "v3", "v2", "a2")] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -1974,8 +1974,8 @@ mod reduce_sum_of_xy_sparse { mod reduce_sum_of_d2_sparse { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_sum_of_d2_sparse_v4_512(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_d2_sparse_v4(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { use crate::emulate::emulate_mm512_2intersect_epi32; assert_eq!(li.len(), lv.len()); assert_eq!(ri.len(), rv.len()); @@ -2057,7 +2057,7 @@ mod reduce_sum_of_d2_sparse { fn reduce_sum_of_d2_sparse_v4_test() { use rand::Rng; const EPSILON: f32 = 0.0004; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -2085,7 +2085,7 @@ mod reduce_sum_of_d2_sparse { let rval = (0..rm) .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); - let specialized = unsafe { reduce_sum_of_d2_sparse_v4_512(&lidx, &lval, &ridx, &rval) }; + let specialized = unsafe { reduce_sum_of_d2_sparse_v4(&lidx, &lval, &ridx, &rval) }; let fallback = fallback(&lidx, &lval, &ridx, &rval); assert!( (specialized - fallback).abs() < EPSILON, @@ -2094,7 +2094,7 @@ mod reduce_sum_of_d2_sparse { } } - #[crate::multiversion(@"v4.512", "v3", "v2", "a2")] + #[crate::multiversion(@"v4", "v3", "v2", "a2")] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -2131,7 +2131,7 @@ mod reduce_sum_of_d2_sparse { } mod vector_add { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2149,7 +2149,7 @@ mod vector_add { } mod vector_add_inplace { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2160,7 +2160,7 @@ mod vector_add_inplace { } mod vector_sub { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2178,7 +2178,7 @@ mod vector_sub { } mod vector_mul { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2196,7 +2196,7 @@ mod vector_mul { } mod vector_mul_scalar { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { let n = lhs.len(); let mut r = Vec::::with_capacity(n); @@ -2213,7 +2213,7 @@ mod vector_mul_scalar { } mod vector_mul_scalar_inplace { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { let n = lhs.len(); for i in 0..n { @@ -2223,7 +2223,7 @@ mod vector_mul_scalar_inplace { } mod vector_abs_inplace { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn vector_abs_inplace(this: &mut [f32]) { let n = this.len(); for i in 0..n { diff --git a/crates/simd/src/fast_scan/mod.rs b/crates/simd/src/fast_scan/mod.rs index 72b2eb26..23fc54f3 100644 --- a/crates/simd/src/fast_scan/mod.rs +++ b/crates/simd/src/fast_scan/mod.rs @@ -121,12 +121,11 @@ pub fn any_pack(mut x: impl Iterator) -> [T; 32] { std::array::from_fn(|_| x.next()).map(|x| x.unwrap_or_default()) } -#[allow(clippy::module_inception)] -mod fast_scan { +mod scan { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn fast_scan_v4_512(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + #[crate::target_cpu(enable = "v4")] + fn scan_v4(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -134,7 +133,7 @@ mod fast_scan { use std::arch::x86_64::*; #[inline] - #[crate::target_cpu(enable = "v4.512")] + #[crate::target_cpu(enable = "v4")] fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); let x0y1 = _mm256_blend_epi32(x0x1, y0y1, 0xf0); @@ -142,7 +141,7 @@ mod fast_scan { } #[inline] - #[crate::target_cpu(enable = "v4.512")] + #[crate::target_cpu(enable = "v4")] fn combine4x2(x0x1x2x3: __m512i, y0y1y2y3: __m512i) -> __m256i { let x0x1 = _mm512_castsi512_si256(x0x1x2x3); let x2x3 = _mm512_extracti64x4_epi64(x0x1x2x3, 1); @@ -235,8 +234,8 @@ mod fast_scan { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn fast_scan_v4_test() { - if !crate::is_cpu_detected!("v4.512") { + fn scan_v4_test() { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -249,7 +248,7 @@ mod fast_scan { .map(|_| std::array::from_fn(|_| rand::random())) .collect::>(); unsafe { - assert_eq!(fast_scan_v4_512(&code, &lut), fallback(&code, &lut)); + assert_eq!(scan_v4(&code, &lut), fallback(&code, &lut)); } } } @@ -258,7 +257,7 @@ mod fast_scan { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - fn fast_scan_v3(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + fn scan_v3(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -338,7 +337,7 @@ mod fast_scan { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn fast_scan_v3_test() { + fn scan_v3_test() { if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; @@ -352,7 +351,7 @@ mod fast_scan { .map(|_| std::array::from_fn(|_| rand::random())) .collect::>(); unsafe { - assert_eq!(fast_scan_v3(&code, &lut), fallback(&code, &lut)); + assert_eq!(scan_v3(&code, &lut), fallback(&code, &lut)); } } } @@ -360,7 +359,7 @@ mod fast_scan { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - fn fast_scan_v2(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + fn scan_v2(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -411,7 +410,7 @@ mod fast_scan { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn fast_scan_v2_test() { + fn scan_v2_test() { if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); return; @@ -425,7 +424,7 @@ mod fast_scan { .map(|_| std::array::from_fn(|_| rand::random())) .collect::>(); unsafe { - assert_eq!(fast_scan_v2(&code, &lut), fallback(&code, &lut)); + assert_eq!(scan_v2(&code, &lut), fallback(&code, &lut)); } } } @@ -433,7 +432,7 @@ mod fast_scan { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] - fn fast_scan_a2(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + fn scan_a2(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { // bounds checking is not enforced by compiler, so check it manually assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -483,7 +482,7 @@ mod fast_scan { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn fast_scan_a2_test() { + fn scan_a2_test() { if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); return; @@ -497,14 +496,14 @@ mod fast_scan { .map(|_| std::array::from_fn(|_| rand::random())) .collect::>(); unsafe { - assert_eq!(fast_scan_a2(&code, &lut), fallback(&code, &lut)); + assert_eq!(scan_a2(&code, &lut), fallback(&code, &lut)); } } } } - #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a2")] - pub fn fast_scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] + pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { fn binary(op: impl Fn(u16, u16) -> u16, a: [u16; 8], b: [u16; 8]) -> [u16; 8] { std::array::from_fn(|i| op(a[i], b[i])) } @@ -543,6 +542,6 @@ mod fast_scan { } #[inline(always)] -pub fn fast_scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { - fast_scan::fast_scan(code, lut) +pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + scan::scan(code, lut) } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index e2ce8304..4748c21c 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -70,7 +70,7 @@ mod internal { pub use is_riscv64_cpu_detected; #[cfg(target_arch = "x86_64")] - pub fn is_v4_512_detected() -> bool { + pub fn is_v4_detected() -> bool { std::arch::is_x86_feature_detected!("avx512bw") && std::arch::is_x86_feature_detected!("avx512cd") && std::arch::is_x86_feature_detected!("avx512dq") @@ -132,6 +132,12 @@ mod internal { std::arch::is_aarch64_feature_detected!("sve") && unsafe { is_256_detected() } } + #[expect(dead_code)] + #[cfg(target_arch = "aarch64")] + pub fn is_a3_128_detected() -> bool { + std::arch::is_aarch64_feature_detected!("sve") + } + #[cfg(target_arch = "aarch64")] pub fn is_a2_detected() -> bool { std::arch::is_aarch64_feature_detected!("neon") diff --git a/crates/simd/src/packed_u4.rs b/crates/simd/src/packed_u4.rs index ae32aebb..9ec7a331 100644 --- a/crates/simd/src/packed_u4.rs +++ b/crates/simd/src/packed_u4.rs @@ -3,7 +3,7 @@ pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { } mod reduce_sum_of_xy { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index 3b6bb770..bd6eefb7 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -1,8 +1,8 @@ mod mul_add_round { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn mul_add_round_v4_512(this: &[f32], k: f32, b: f32) -> Vec { + #[crate::target_cpu(enable = "v4")] + fn mul_add_round_v4(this: &[f32], k: f32, b: f32) -> Vec { let mut r = Vec::::with_capacity(this.len()); use std::arch::x86_64::*; let lk = _mm512_set1_ps(k); @@ -41,7 +41,7 @@ mod mul_add_round { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] fn mul_add_round_v4_test() { - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -52,7 +52,7 @@ mod mul_add_round { let x = &x[..z]; let k = 20.0; let b = 20.0; - let specialized = unsafe { mul_add_round_v4_512(x, k, b) }; + let specialized = unsafe { mul_add_round_v4(x, k, b) }; let fallback = fallback(x, k, b); assert_eq!(specialized, fallback); } @@ -263,7 +263,7 @@ mod mul_add_round { } } - #[crate::multiversion(@"v4.512", @"v3", @"v2:fma", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2")] pub fn mul_add_round(this: &[f32], k: f32, b: f32) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index 043b3eb2..c2a6b8e5 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -1,5 +1,5 @@ mod reduce_sum_of_xy { - #[crate::multiversion("v4.512", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2")] pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); @@ -19,8 +19,8 @@ pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { mod reduce_sum_of_x_as_u16 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_sum_of_x_as_u16_v4_512(this: &[u8]) -> u16 { + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_as_u16_v4(this: &[u8]) -> u16 { use crate::emulate::emulate_mm512_reduce_add_epi16; use std::arch::x86_64::*; let us = _mm512_set1_epi16(255); @@ -45,7 +45,7 @@ mod reduce_sum_of_x_as_u16 { #[test] fn reduce_sum_of_x_as_u16_v4_test() { use rand::Rng; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -55,7 +55,7 @@ mod reduce_sum_of_x_as_u16 { let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u16_v4_512(this) }; + let specialized = unsafe { reduce_sum_of_x_as_u16_v4(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } @@ -205,7 +205,7 @@ mod reduce_sum_of_x_as_u16 { } } - #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] pub fn reduce_sum_of_x_as_u16(this: &[u8]) -> u16 { let n = this.len(); let mut sum = 0; @@ -224,8 +224,8 @@ pub fn reduce_sum_of_x_as_u16(vector: &[u8]) -> u16 { mod reduce_sum_of_x { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4.512")] - fn reduce_sum_of_x_v4_512(this: &[u8]) -> u32 { + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[u8]) -> u32 { use std::arch::x86_64::*; let us = _mm512_set1_epi32(255); let mut n = this.len(); @@ -249,7 +249,7 @@ mod reduce_sum_of_x { #[test] fn reduce_sum_of_x_v4_test() { use rand::Rng; - if !crate::is_cpu_detected!("v4.512") { + if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } @@ -259,7 +259,7 @@ mod reduce_sum_of_x { let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_v4_512(this) }; + let specialized = unsafe { reduce_sum_of_x_v4(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } @@ -410,7 +410,7 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4.512", @"v3", @"v2", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] pub fn reduce_sum_of_x(this: &[u8]) -> u32 { let n = this.len(); let mut sum = 0; diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs index e2be8534..0c136a89 100644 --- a/crates/simd_macros/src/target.rs +++ b/crates/simd_macros/src/target.rs @@ -6,7 +6,7 @@ pub struct TargetCpu { pub const TARGET_CPUS: &[TargetCpu] = &[ TargetCpu { - target_cpu: "v4.512", + target_cpu: "v4", target_arch: "x86_64", target_features: &[ "avx512bw", "avx512cd", "avx512dq", "avx512vl", // simd @@ -43,6 +43,13 @@ pub const TARGET_CPUS: &[TargetCpu] = &[ "sve", // simd ], }, + TargetCpu { + target_cpu: "a3.128", + target_arch: "aarch64", + target_features: &[ + "sve", // simd + ], + }, TargetCpu { target_cpu: "a2", target_arch: "aarch64", diff --git a/crates/vector/src/svect.rs b/crates/vector/src/svect.rs index 26cbf197..d83d5f2b 100644 --- a/crates/vector/src/svect.rs +++ b/crates/vector/src/svect.rs @@ -155,10 +155,14 @@ impl<'a, S: Floating> SVectBorrowed<'a, S> { } #[inline(always)] - #[allow(clippy::len_without_is_empty)] pub fn len(&self) -> u32 { self.indexes.len() as u32 } + + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.indexes.is_empty() + } } impl VectorBorrowed for SVectBorrowed<'_, S> { diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index e13097dd..a4180f3f 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -338,7 +338,8 @@ mod vchordrq_cached { #[derive(Debug, Clone, Copy)] pub enum VchordrqCachedReader<'a> { - _0(#[allow(dead_code)] VchordrqCachedReader0<'a>), + #[allow(dead_code)] + _0(VchordrqCachedReader0<'a>), _1(VchordrqCachedReader1<'a>), } @@ -853,6 +854,7 @@ pub fn make_internal_build( result } +#[allow(clippy::collapsible_else_if)] pub fn make_external_build( vector_options: VectorOptions, _opfamily: Opfamily, diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index abed8414..70fae6f2 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -126,6 +126,7 @@ pub unsafe extern "C" fn amoptions(reloptions: Datum, validate: bool) -> *mut pg rdopts as *mut pgrx::pg_sys::bytea } +#[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C" fn amcostestimate( _root: *mut pgrx::pg_sys::PlannerInfo, @@ -155,6 +156,7 @@ pub unsafe extern "C" fn amcostestimate( } #[cfg(feature = "pg13")] +#[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C" fn aminsert( index_relation: pgrx::pg_sys::Relation, @@ -169,6 +171,7 @@ pub unsafe extern "C" fn aminsert( } #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] +#[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C" fn aminsert( index_relation: pgrx::pg_sys::Relation, diff --git a/src/lib.rs b/src/lib.rs index 5a98b3d9..116e90d3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,3 @@ -#![allow(clippy::collapsible_else_if)] -#![allow(clippy::too_many_arguments)] #![allow(unsafe_code)] #![feature(lazy_get)] @@ -27,3 +25,8 @@ extern "C" fn _PG_init() { #[cfg(not(target_endian = "little"))] compile_error!("Target architecture is not supported."); + +#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] +#[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))] +#[global_allocator] +static GLOBAL_ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; diff --git a/src/upgrade.rs b/src/upgrade.rs index 5998c1cb..b474744b 100644 --- a/src/upgrade.rs +++ b/src/upgrade.rs @@ -4,7 +4,7 @@ // * https://www.postgresql.org/message-id/CACX+KaPOzzRHEt4w_=iqKbTpMKjyrUGVng1C749yP3r6dprtcg@mail.gmail.com // * https://github.com/tensorchord/pgvecto.rs/issues/397 -#[allow(unused)] +#[allow(unused_macros)] macro_rules! symbol { ($t:ident) => { paste::paste! { From 8e9730f8154c92aae29050d482a18ad4f572dabd Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 24 Mar 2025 18:18:47 +0800 Subject: [PATCH 134/324] feat: emit INFO when performing kmeans (#219) 1. add build phase names 2. emit kmeans log 3. emit errors if both `Dot` and `residual_quantiztion` are used 4. update `tuples_done` by estimated value before insertions 5. update `tuples_done` after insertions are done 6. emit errors if `vchordrq_prewarm` is given a partition index 7. rename `maxsim_ip` to `maxsim` Signed-off-by: usamoi --- crates/algorithm/src/types.rs | 12 +-- crates/k_means/src/lib.rs | 24 +++-- src/datatype/operators_halfvec.rs | 2 +- src/datatype/operators_vector.rs | 2 +- src/index/am/am_build.rs | 145 ++++++++++++++++++++++++------ src/index/am/mod.rs | 3 +- src/index/functions.rs | 46 +++++++--- src/index/opclass.rs | 46 +++++----- src/index/scanners/maxsim.rs | 2 +- src/sql/finalize.sql | 20 ++--- tests/logic/multivector.slt | 2 +- tests/logic/partition.slt | 3 + 12 files changed, 219 insertions(+), 88 deletions(-) diff --git a/crates/algorithm/src/types.rs b/crates/algorithm/src/types.rs index 30299fab..a59de467 100644 --- a/crates/algorithm/src/types.rs +++ b/crates/algorithm/src/types.rs @@ -51,7 +51,7 @@ pub enum VectorKind { #[serde(deny_unknown_fields)] #[validate(schema(function = "Self::validate_self"))] pub struct VectorOptions { - #[validate(range(min = 1, max = 1_048_575))] + #[validate(range(min = 1))] #[serde(rename = "dimensions")] pub dims: u32, #[serde(rename = "vector")] @@ -63,11 +63,11 @@ pub struct VectorOptions { impl VectorOptions { pub fn validate_self(&self) -> Result<(), ValidationError> { match (self.v, self.d, self.dims) { - (VectorKind::Vecf32, DistanceKind::L2, 1..65536) => Ok(()), - (VectorKind::Vecf32, DistanceKind::Dot, 1..65536) => Ok(()), - (VectorKind::Vecf16, DistanceKind::L2, 1..65536) => Ok(()), - (VectorKind::Vecf16, DistanceKind::Dot, 1..65536) => Ok(()), - _ => Err(ValidationError::new("not valid vector options")), + (VectorKind::Vecf32, DistanceKind::L2, 1..=60000) => Ok(()), + (VectorKind::Vecf32, DistanceKind::Dot, 1..=60000) => Ok(()), + (VectorKind::Vecf16, DistanceKind::L2, 1..=60000) => Ok(()), + (VectorKind::Vecf16, DistanceKind::Dot, 1..=60000) => Ok(()), + _ => Err(ValidationError::new("invalid vector options")), } } } diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index cdc475fd..c6c0df9d 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -1,13 +1,27 @@ use rabitq::block::BlockCode; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; -use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use rayon::iter::{IntoParallelIterator, IntoParallelRefMutIterator, ParallelIterator}; use simd::Floating; use simd::fast_scan::{any_pack, padding_pack}; +pub fn preprocess(num_threads: usize, x: &mut [T], f: impl Fn(&mut T) + Sync) { + rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .build_scoped( + |thread| thread.run(), + move |pool| { + pool.install(|| { + x.par_iter_mut().for_each(&f); + }); + }, + ) + .expect("failed to build thread pool") +} + pub fn k_means( num_threads: usize, - check: impl Fn(), + check: impl Fn(usize), c: usize, dims: usize, samples: &[Vec], @@ -26,7 +40,7 @@ pub fn k_means( |thread| thread.run(), move |pool| { let compute = |centroids: &[Vec]| { - if n >= 1000 && c >= 1000 { + if n >= 1024 && c >= 1024 { rabitq_index(dims, n, c, samples, centroids) } else { flat_index(dims, n, c, samples, centroids) @@ -34,8 +48,8 @@ pub fn k_means( }; let mut lloyd_k_means = pool.install(|| LloydKMeans::new(c, dims, samples, is_spherical, compute)); - for _ in 0..iterations { - check(); + for i in 0..iterations { + check(i); if pool.install(|| lloyd_k_means.iterate()) { break; } diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index 0be60ccb..d96e0945 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -77,7 +77,7 @@ fn _vchord_halfvec_sphere_cosine_in( } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_halfvec_operator_maxsim_ip( +fn _vchord_halfvec_operator_maxsim( lhs: Array<'_, HalfvecInput<'_>>, rhs: Array<'_, HalfvecInput<'_>>, ) -> f32 { diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index d8d60b50..8f05a1c0 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -77,7 +77,7 @@ fn _vchord_vector_sphere_cosine_in( } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_vector_operator_maxsim_ip( +fn _vchord_vector_operator_maxsim( lhs: Array<'_, VectorInput<'_>>, rhs: Array<'_, VectorInput<'_>>, ) -> f32 { diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index a4180f3f..40b54ae3 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -9,10 +9,67 @@ use half::f16; use pgrx::pg_sys::{Datum, ItemPointerData}; use rand::Rng; use simd::Floating; +use std::ffi::CStr; +use std::num::NonZero; use std::ops::Deref; use vector::VectorOwned; use vector::vect::VectOwned; +#[derive(Debug, Clone, Copy)] +#[repr(i64)] +pub enum BuildPhase { + Initializing = 1, + InternalBuild = 2, + ExternalBuild = 3, + Build = 4, + Inserting = 5, + Compacting = 6, +} + +impl TryFrom> for BuildPhase { + type Error = (); + + fn try_from(value: NonZero) -> Result { + const INITIALIZING: NonZero = NonZero::new(BuildPhase::Initializing as _).unwrap(); + const INTERNAL_BUILD: NonZero = NonZero::new(BuildPhase::InternalBuild as _).unwrap(); + const EXTERNAL_BUILD: NonZero = NonZero::new(BuildPhase::ExternalBuild as _).unwrap(); + const BUILD: NonZero = NonZero::new(BuildPhase::Build as _).unwrap(); + const INSERTING: NonZero = NonZero::new(BuildPhase::Inserting as _).unwrap(); + const COMPACTING: NonZero = NonZero::new(BuildPhase::Compacting as _).unwrap(); + match value { + INITIALIZING => Ok(BuildPhase::Initializing), + INTERNAL_BUILD => Ok(BuildPhase::InternalBuild), + EXTERNAL_BUILD => Ok(BuildPhase::ExternalBuild), + BUILD => Ok(BuildPhase::Build), + INSERTING => Ok(BuildPhase::Inserting), + COMPACTING => Ok(BuildPhase::Compacting), + _ => Err(()), + } + } +} + +impl BuildPhase { + fn name(self) -> &'static CStr { + match self { + Self::Initializing => c"initializing", + Self::InternalBuild => c"initializing index, by internal build", + Self::ExternalBuild => c"initializing index, by external build", + Self::Build => c"initializing index", + Self::Inserting => c"inserting tuples from table to index", + Self::Compacting => c"compacting tuples in index", + } + } +} + +#[pgrx::pg_guard] +pub extern "C" fn ambuildphasename(x: i64) -> *mut core::ffi::c_char { + if let Some(x) = NonZero::new(x).and_then(|x| BuildPhase::try_from(x).ok()) { + x.name().as_ptr().cast_mut() + } else { + std::ptr::null_mut() + } +} + #[derive(Debug, Clone)] struct Heap { heap_relation: pgrx::pg_sys::Relation, @@ -78,6 +135,14 @@ impl Heap { struct PostgresReporter {} impl PostgresReporter { + fn phase(&mut self, phase: BuildPhase) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_SUBPHASE as _, + phase as _, + ); + } + } fn tuples_total(&mut self, tuples_total: u64) { unsafe { pgrx::pg_sys::pgstat_progress_update_param( @@ -107,15 +172,14 @@ pub unsafe extern "C" fn ambuild( if let Err(errors) = Validate::validate(&vector_options) { pgrx::error!("error while validating options: {}", errors); } - if vector_options.dims == 0 { - pgrx::error!("error while validating options: dimension cannot be 0"); - } - if vector_options.dims > 60000 { - pgrx::error!("error while validating options: dimension is too large"); - } if let Err(errors) = Validate::validate(&vchordrq_options) { pgrx::error!("error while validating options: {}", errors); } + if vector_options.d != DistanceKind::L2 && vchordrq_options.index.residual_quantization { + pgrx::error!( + "error while validating options: residual_quantization can be enabled only if distance type is L2" + ); + } let opfamily = unsafe { opfamily(index_relation) }; let heap = Heap { heap_relation, @@ -128,9 +192,13 @@ pub unsafe extern "C" fn ambuild( let mut reporter = PostgresReporter {}; let structures = match vchordrq_options.build.source.clone() { VchordrqBuildSourceOptions::External(external_build) => { + reporter.phase(BuildPhase::ExternalBuild); + let reltuples = unsafe { (*(*index_relation).rd_rel).reltuples }; + reporter.tuples_total(reltuples as u64); make_external_build(vector_options.clone(), opfamily, external_build.clone()) } VchordrqBuildSourceOptions::Internal(internal_build) => { + reporter.phase(BuildPhase::InternalBuild); let mut tuples_total = 0_u64; let samples = 'a: { let mut rand = rand::rng(); @@ -170,12 +238,14 @@ pub unsafe extern "C" fn ambuild( make_internal_build(vector_options.clone(), internal_build.clone(), samples) } }; + reporter.phase(BuildPhase::Build); crate::index::algorithm::build( vector_options, vchordrq_options.index, index.clone(), structures, ); + reporter.phase(BuildPhase::Inserting); let cache = if vchordrq_options.build.pin { let mut trace = algorithm::cache(index.clone()); trace.sort(); @@ -210,7 +280,9 @@ pub unsafe extern "C" fn ambuild( leader.tablescandesc, leader.vchordrqshared, leader.vchordrqcached, - Some(reporter), + |indtuples| { + reporter.tuples_done(indtuples); + }, ); leader.wait(); let nparticipants = leader.nparticipants; @@ -226,6 +298,8 @@ pub unsafe extern "C" fn ambuild( pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN, ); } + reporter.tuples_done((*leader.vchordrqshared).indtuples); + reporter.tuples_total((*leader.vchordrqshared).indtuples); pgrx::pg_sys::ConditionVariableCancelSleep(); } } else { @@ -240,10 +314,12 @@ pub unsafe extern "C" fn ambuild( indtuples += 1; reporter.tuples_done(indtuples); }); + reporter.tuples_total(indtuples); } let check = || { pgrx::check_for_interrupts!(); }; + reporter.phase(BuildPhase::Compacting); crate::index::algorithm::maintain(opfamily, index, check); unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -673,7 +749,7 @@ pub unsafe extern "C" fn vchordrq_parallel_build_main( tablescandesc, vchordrqshared, vchordrqcached, - None, + |_| (), ); } @@ -690,7 +766,7 @@ unsafe fn parallel_build( tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, vchordrqshared: *mut VchordrqShared, vchordrqcached: *const u8, - mut reporter: Option, + mut callback: impl FnMut(u64), ) { use vchordrq_cached::VchordrqCachedReader; let cached = VchordrqCachedReader::deserialize_ref(unsafe { @@ -724,9 +800,7 @@ unsafe fn parallel_build( indtuples = (*vchordrqshared).indtuples; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } + callback(indtuples); } }); } @@ -748,9 +822,7 @@ unsafe fn parallel_build( indtuples = (*vchordrqshared).indtuples; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } - if let Some(reporter) = reporter.as_mut() { - reporter.tuples_done(indtuples); - } + callback(indtuples); } }); } @@ -816,26 +888,43 @@ pub fn make_internal_build( mut samples: Vec>, ) -> Vec>> { use std::iter::once; - for sample in samples.iter_mut() { - *sample = crate::index::projection::project(sample); - } + k_means::preprocess(internal_build.build_threads as _, &mut samples, |sample| { + *sample = crate::index::projection::project(sample) + }); let mut result = Vec::>>::new(); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { + let input = if let Some(structure) = result.last() { + &structure.means + } else { + &samples + }; + let num_threads = internal_build.build_threads as _; + let num_points = input.len(); + let num_dims = vector_options.dims as usize; + let num_lists = w as usize; + let num_iterations = internal_build.kmeans_iterations as _; + if num_lists > 1 { + pgrx::info!( + "clustering: starting, using {num_threads} threads, clustering {num_points} vectors of {num_dims} dimension into {num_lists} clusters, in {num_iterations} iterations" + ); + } let means = k_means::k_means( - internal_build.build_threads as _, - || { + num_threads, + |i| { pgrx::check_for_interrupts!(); + if num_lists > 1 { + pgrx::info!("clustering: iteration {}", i + 1); + } }, - w as usize, - vector_options.dims as usize, - if let Some(structure) = result.last() { - &structure.means - } else { - &samples - }, + num_lists, + num_dims, + input, internal_build.spherical_centroids, - internal_build.kmeans_iterations as _, + num_iterations, ); + if num_lists > 1 { + pgrx::info!("clustering: finished"); + } if let Some(structure) = result.last() { let mut children = vec![Vec::new(); means.len()]; for i in 0..structure.len() as u32 { diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index 70fae6f2..a408b892 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -91,6 +91,7 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { am_routine.amoptions = Some(amoptions); am_routine.amcostestimate = Some(amcostestimate); + am_routine.ambuildphasename = Some(am_build::ambuildphasename); am_routine.ambuild = Some(am_build::ambuild); am_routine.ambuildempty = Some(am_build::ambuildempty); am_routine.aminsert = Some(aminsert); @@ -332,7 +333,7 @@ pub unsafe extern "C" fn amrescan( } LazyCell::new(Box::new(move || builder.build(relation, options, fetcher))) } - Opfamily::VectorMaxsimIp | Opfamily::HalfvecMaxsimIp => { + Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim => { let mut builder = MaxsimBuilder::new(opfamily); for i in 0..(*scan).numberOfOrderBys { let data = (*scan).orderByData.add(i as usize); diff --git a/src/index/functions.rs b/src/index/functions.rs index 4a472229..e97e261d 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -1,6 +1,6 @@ use crate::index::storage::PostgresRelation; use pgrx::pg_sys::Oid; -use pgrx_catalog::{PgAm, PgClass}; +use pgrx_catalog::{PgAm, PgClass, PgClassRelkind}; #[pgrx::pg_extern(sql = "")] fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { @@ -10,19 +10,43 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { }; let pg_class = PgClass::search_reloid(indexrelid).unwrap(); let Some(pg_class) = pg_class.get() else { - pgrx::error!("there is no such index"); + pgrx::error!("the relation does not exist"); }; + if pg_class.relkind() != PgClassRelkind::Index { + pgrx::error!("the relation {:?} is not an index", pg_class.relname()); + } if pg_class.relam() != pg_am.oid() { - pgrx::error!("{:?} is not a vchordrq index", pg_class.relname()); + pgrx::error!("the index {:?} is not a vchordrq index", pg_class.relname()); } - let index = unsafe { pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::AccessShareLock as _) }; - let relation = unsafe { PostgresRelation::new(index) }; - let opfamily = unsafe { crate::index::opclass::opfamily(index) }; - let message = crate::index::algorithm::prewarm(opfamily, relation, height, || { + let index = Index::open(indexrelid); + let relation = unsafe { PostgresRelation::new(index.raw()) }; + let opfamily = unsafe { crate::index::opclass::opfamily(index.raw()) }; + crate::index::algorithm::prewarm(opfamily, relation, height, || { pgrx::check_for_interrupts!(); - }); - unsafe { - pgrx::pg_sys::index_close(index, pgrx::pg_sys::AccessShareLock as _); + }) +} + +struct Index { + raw: *mut pgrx::pg_sys::RelationData, +} + +impl Index { + fn open(indexrelid: Oid) -> Self { + Self { + raw: unsafe { + pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::AccessShareLock as _) + }, + } + } + fn raw(&self) -> *mut pgrx::pg_sys::RelationData { + self.raw + } +} + +impl Drop for Index { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::index_close(self.raw, pgrx::pg_sys::AccessShareLock as _); + } } - message } diff --git a/src/index/opclass.rs b/src/index/opclass.rs index 9067fbe0..aebf2342 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -39,13 +39,13 @@ fn _vchordrq_support_halfvec_cosine_ops() -> String { } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchordrq_support_vector_maxsim_ip_ops() -> String { - "vector_maxsim_ip_ops".to_string() +fn _vchordrq_support_vector_maxsim_ops() -> String { + "vector_maxsim_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchordrq_support_halfvec_maxsim_ip_ops() -> String { - "halfvec_maxsim_ip_ops".to_string() +fn _vchordrq_support_halfvec_maxsim_ops() -> String { + "halfvec_maxsim_ops".to_string() } pub struct Sphere { @@ -61,8 +61,8 @@ pub enum Opfamily { HalfvecL2, HalfvecIp, HalfvecCosine, - VectorMaxsimIp, - HalfvecMaxsimIp, + VectorMaxsim, + HalfvecMaxsim, } impl Opfamily { @@ -70,11 +70,11 @@ impl Opfamily { use {BorrowedVector as B, OwnedVector as O}; match (vector, self) { (B::Vecf32(x), Self::VectorL2) => O::Vecf32(x.own()), - (B::Vecf32(x), Self::VectorIp | Self::VectorMaxsimIp) => O::Vecf32(x.own()), + (B::Vecf32(x), Self::VectorIp | Self::VectorMaxsim) => O::Vecf32(x.own()), (B::Vecf32(x), Self::VectorCosine) => O::Vecf32(x.function_normalize()), (B::Vecf32(_), _) => unreachable!(), (B::Vecf16(x), Self::HalfvecL2) => O::Vecf16(x.own()), - (B::Vecf16(x), Self::HalfvecIp | Self::HalfvecMaxsimIp) => O::Vecf16(x.own()), + (B::Vecf16(x), Self::HalfvecIp | Self::HalfvecMaxsim) => O::Vecf16(x.own()), (B::Vecf16(x), Self::HalfvecCosine) => O::Vecf16(x.function_normalize()), (B::Vecf16(_), _) => unreachable!(), } @@ -92,7 +92,7 @@ impl Opfamily { let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] } - Self::VectorMaxsimIp => { + Self::VectorMaxsim => { let vectors = unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; let mut result = Vec::with_capacity(vectors.len()); @@ -104,7 +104,7 @@ impl Opfamily { } result } - Self::HalfvecMaxsimIp => { + Self::HalfvecMaxsim => { let vectors = unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; let mut result = Vec::with_capacity(vectors.len()); @@ -127,11 +127,11 @@ impl Opfamily { let attno_2 = NonZero::new(2_usize).unwrap(); let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; let center = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsimIp => { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Vecf32(vector.as_borrowed())) } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsimIp => { + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Vecf16(vector.as_borrowed())) } @@ -144,11 +144,11 @@ impl Opfamily { return None; } let vector = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsimIp => { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; self.input(BorrowedVector::Vecf32(vector.as_borrowed())) } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsimIp => { + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; self.input(BorrowedVector::Vecf16(vector.as_borrowed())) } @@ -160,7 +160,7 @@ impl Opfamily { return None; } let vectors = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsimIp => { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { let vectors = unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; let mut result = Vec::with_capacity(vectors.len()); @@ -169,7 +169,7 @@ impl Opfamily { } result } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsimIp => { + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { let vectors = unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; let mut result = Vec::with_capacity(vectors.len()); @@ -185,7 +185,7 @@ impl Opfamily { match self { Self::VectorCosine | Self::HalfvecCosine => x.to_f32() + 1.0f32, Self::VectorL2 | Self::HalfvecL2 => x.to_f32().sqrt(), - Self::VectorIp | Self::HalfvecIp | Self::VectorMaxsimIp | Self::HalfvecMaxsimIp => { + Self::VectorIp | Self::HalfvecIp | Self::VectorMaxsim | Self::HalfvecMaxsim => { x.to_f32() } } @@ -197,16 +197,16 @@ impl Opfamily { | Self::HalfvecIp | Self::VectorCosine | Self::HalfvecCosine - | Self::VectorMaxsimIp - | Self::HalfvecMaxsimIp => DistanceKind::Dot, + | Self::VectorMaxsim + | Self::HalfvecMaxsim => DistanceKind::Dot, } } pub const fn vector_kind(self) -> VectorKind { match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsimIp => { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { VectorKind::Vecf32 } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsimIp => { + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { VectorKind::Vecf16 } } @@ -251,8 +251,8 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { "halfvec_l2_ops" => Opfamily::HalfvecL2, "halfvec_ip_ops" => Opfamily::HalfvecIp, "halfvec_cosine_ops" => Opfamily::HalfvecCosine, - "vector_maxsim_ip_ops" => Opfamily::VectorMaxsimIp, - "halfvec_maxsim_ip_ops" => Opfamily::HalfvecMaxsimIp, + "vector_maxsim_ops" => Opfamily::VectorMaxsim, + "halfvec_maxsim_ops" => Opfamily::HalfvecMaxsim, _ => pgrx::error!("unknown operator class"), }; diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs index 016ff5c3..9660d441 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/scanners/maxsim.rs @@ -23,7 +23,7 @@ impl SearchBuilder for MaxsimBuilder { fn new(opfamily: Opfamily) -> Self { assert!(matches!( opfamily, - Opfamily::VectorMaxsimIp | Opfamily::HalfvecMaxsimIp + Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim )); Self { opfamily, diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 171e8d1c..c97f6c46 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -114,13 +114,13 @@ CREATE OPERATOR <<=>> ( ); CREATE OPERATOR @# ( - PROCEDURE = _vchord_vector_operator_maxsim_ip, + PROCEDURE = _vchord_vector_operator_maxsim, LEFTARG = vector[], RIGHTARG = vector[] ); CREATE OPERATOR @# ( - PROCEDURE = _vchord_halfvec_operator_maxsim_ip, + PROCEDURE = _vchord_halfvec_operator_maxsim, LEFTARG = halfvec[], RIGHTARG = halfvec[] ); @@ -160,8 +160,8 @@ CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; -CREATE OPERATOR FAMILY vector_maxsim_ip_ops USING vchordrq; -CREATE OPERATOR FAMILY halfvec_maxsim_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; -- List of operator classes @@ -201,12 +201,12 @@ CREATE OPERATOR CLASS halfvec_cosine_ops OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); -CREATE OPERATOR CLASS vector_maxsim_ip_ops - FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ip_ops AS +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, - FUNCTION 1 _vchordrq_support_vector_maxsim_ip_ops(); + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); -CREATE OPERATOR CLASS halfvec_maxsim_ip_ops - FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ip_ops AS +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, - FUNCTION 1 _vchordrq_support_halfvec_maxsim_ip_ops(); + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); diff --git a/tests/logic/multivector.slt b/tests/logic/multivector.slt index 5cc772e0..1de81a50 100644 --- a/tests/logic/multivector.slt +++ b/tests/logic/multivector.slt @@ -12,7 +12,7 @@ SELECT id, FROM generate_series(1, 10000) s(id); statement ok -CREATE INDEX t_val_idx ON t USING vchordrq (val vector_maxsim_ip_ops) +CREATE INDEX t_val_idx ON t USING vchordrq (val vector_maxsim_ops) WITH (options = $$ build.internal.lists = [] $$); diff --git a/tests/logic/partition.slt b/tests/logic/partition.slt index 68687722..506995ef 100644 --- a/tests/logic/partition.slt +++ b/tests/logic/partition.slt @@ -26,6 +26,9 @@ SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <-> '[0.5,0.5,0.5]' limit ---- 10 +statement error the relation "items_val_idx" is not an index +select vchordrq_prewarm('items_val_idx'); + statement ok CREATE INDEX ON id_123 USING vchordrq (val vector_cosine_ops); From 63090fd08b616e7821032d13a86899094a1cea08 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 25 Mar 2025 16:57:36 +0800 Subject: [PATCH 135/324] feat: encode kmeans progress to phase name (#220) * encode kmeans progress to build subphase names * add a new GUC `vchordrq.prererank_filtering` (default to `off`) to control prererank filtering --------- Signed-off-by: usamoi --- Cargo.lock | 7 ++ Cargo.toml | 1 + README.md | 9 +- crates/k_means/src/lib.rs | 2 +- src/index/am/am_build.rs | 194 ++++++++++++++++++++++++-------- src/index/am/mod.rs | 29 +++-- src/index/gucs.rs | 13 +++ tests/logic/rerank_in_table.slt | 12 ++ 8 files changed, 199 insertions(+), 68 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 00ad8dee..4819ae4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1163,6 +1163,12 @@ version = "1.0.26" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56e6fa9c48d24d85fb3de5ad847117517440f6beceb7798af16b4a87d616b8d0" +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.219" @@ -1500,6 +1506,7 @@ dependencies = [ "pgrx-catalog", "rand", "random_orthogonal_matrix", + "seq-macro", "serde", "simd", "toml", diff --git a/Cargo.toml b/Cargo.toml index d85cb394..c29498d4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ paste.workspace = true pgrx = { version = "=0.13.1", default-features = false, features = ["cshim"] } pgrx-catalog = "0.1.1" rand.workspace = true +seq-macro = "0.3.6" serde.workspace = true toml = "0.8.20" validator.workspace = true diff --git a/README.md b/README.md index ab3816e8..1e70f511 100644 --- a/README.md +++ b/README.md @@ -110,8 +110,7 @@ Now you can play with VectorChord! - [Indexing Prewarm](#indexing-prewarm) - [Indexing Progress](#indexing-progress) - [External Index Precomputation](#external-index-precomputation) - + - [Capacity-optimized Index](#capacity-optimized-index) - [Range Query](#range-query) ## Installation @@ -377,8 +376,6 @@ $$); To simplify the workflow, we provide end-to-end scripts for external index pre-computation, see [scripts](./scripts/README.md#run-external-index-precomputation-toolkit). - - ### Range Query To query vectors within a certain distance range, you can use the following syntax. diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index c6c0df9d..7ce5ed3e 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -21,7 +21,7 @@ pub fn preprocess(num_threads: usize, x: &mut [T], f: impl Fn(&mut T) + pub fn k_means( num_threads: usize, - check: impl Fn(usize), + mut check: impl FnMut(usize), c: usize, dims: usize, samples: &[Vec], diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 40b54ae3..c5e83c70 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -10,61 +10,131 @@ use pgrx::pg_sys::{Datum, ItemPointerData}; use rand::Rng; use simd::Floating; use std::ffi::CStr; -use std::num::NonZero; use std::ops::Deref; use vector::VectorOwned; use vector::vect::VectOwned; #[derive(Debug, Clone, Copy)] -#[repr(i64)] -pub enum BuildPhase { - Initializing = 1, - InternalBuild = 2, - ExternalBuild = 3, - Build = 4, - Inserting = 5, - Compacting = 6, +#[repr(u16)] +pub enum BuildPhaseCode { + Initializing = 0, + InternalBuild = 1, + ExternalBuild = 2, + Build = 3, + Inserting = 4, + Compacting = 5, } -impl TryFrom> for BuildPhase { - type Error = (); - - fn try_from(value: NonZero) -> Result { - const INITIALIZING: NonZero = NonZero::new(BuildPhase::Initializing as _).unwrap(); - const INTERNAL_BUILD: NonZero = NonZero::new(BuildPhase::InternalBuild as _).unwrap(); - const EXTERNAL_BUILD: NonZero = NonZero::new(BuildPhase::ExternalBuild as _).unwrap(); - const BUILD: NonZero = NonZero::new(BuildPhase::Build as _).unwrap(); - const INSERTING: NonZero = NonZero::new(BuildPhase::Inserting as _).unwrap(); - const COMPACTING: NonZero = NonZero::new(BuildPhase::Compacting as _).unwrap(); - match value { - INITIALIZING => Ok(BuildPhase::Initializing), - INTERNAL_BUILD => Ok(BuildPhase::InternalBuild), - EXTERNAL_BUILD => Ok(BuildPhase::ExternalBuild), - BUILD => Ok(BuildPhase::Build), - INSERTING => Ok(BuildPhase::Inserting), - COMPACTING => Ok(BuildPhase::Compacting), - _ => Err(()), - } - } -} +pub struct BuildPhase(BuildPhaseCode, u16); impl BuildPhase { - fn name(self) -> &'static CStr { + pub const fn new(code: BuildPhaseCode, k: u16) -> Option { + match (code, k) { + (BuildPhaseCode::Initializing, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::InternalBuild, 0..102) => Some(BuildPhase(code, k)), + (BuildPhaseCode::ExternalBuild, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::Build, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::Inserting, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::Compacting, 0) => Some(BuildPhase(code, k)), + _ => None, + } + } + pub const fn name(self) -> &'static CStr { match self { - Self::Initializing => c"initializing", - Self::InternalBuild => c"initializing index, by internal build", - Self::ExternalBuild => c"initializing index, by external build", - Self::Build => c"initializing index", - Self::Inserting => c"inserting tuples from table to index", - Self::Compacting => c"compacting tuples in index", + BuildPhase(BuildPhaseCode::Initializing, k) => { + static RAW: [&CStr; 1] = [c"initializing"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::InternalBuild, k) => { + static RAWS: [&[&CStr]; 2] = [ + &[c"initializing index, by internal build"], + seq_macro::seq!( + N in 0..=100 { + &[#( + const { + let bytes = concat!("initializing index, by internal build (", N, " %)\0").as_bytes(); + if let Ok(cstr) = CStr::from_bytes_with_nul(bytes) { + cstr + } else { + panic!("") + } + }, + )*] + } + ), + ]; + static RAW: [&CStr; 102] = { + let mut result = [c""; 102]; + let mut offset = 0_usize; + let mut i = 0_usize; + while i < RAWS.len() { + let mut j = 0_usize; + while j < RAWS[i].len() { + { + result[offset] = RAWS[i][j]; + offset += 1; + } + j += 1; + } + i += 1; + } + assert!(offset == result.len()); + result + }; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::ExternalBuild, k) => { + static RAW: [&CStr; 1] = [c"initializing index, by external build"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::Build, k) => { + static RAW: [&CStr; 1] = [c"initializing index"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::Inserting, k) => { + static RAW: [&CStr; 1] = [c"inserting tuples from table to index"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::Compacting, k) => { + static RAW: [&CStr; 1] = [c"compacting tuples in index"]; + RAW[k as usize] + } } } + pub const fn from_code(code: BuildPhaseCode) -> Self { + Self(code, 0) + } + pub const fn from_value(value: u32) -> Option { + const INITIALIZING: u16 = BuildPhaseCode::Initializing as _; + const INTERNAL_BUILD: u16 = BuildPhaseCode::InternalBuild as _; + const EXTERNAL_BUILD: u16 = BuildPhaseCode::ExternalBuild as _; + const BUILD: u16 = BuildPhaseCode::Build as _; + const INSERTING: u16 = BuildPhaseCode::Inserting as _; + const COMPACTING: u16 = BuildPhaseCode::Compacting as _; + let k = value as u16; + match (value >> 16) as u16 { + INITIALIZING => Self::new(BuildPhaseCode::Initializing, k), + INTERNAL_BUILD => Self::new(BuildPhaseCode::InternalBuild, k), + EXTERNAL_BUILD => Self::new(BuildPhaseCode::ExternalBuild, k), + BUILD => Self::new(BuildPhaseCode::Build, k), + INSERTING => Self::new(BuildPhaseCode::Inserting, k), + COMPACTING => Self::new(BuildPhaseCode::Compacting, k), + _ => None, + } + } + pub const fn into_value(self) -> u32 { + (self.0 as u32) << 16 | (self.1 as u32) + } } #[pgrx::pg_guard] pub extern "C" fn ambuildphasename(x: i64) -> *mut core::ffi::c_char { - if let Some(x) = NonZero::new(x).and_then(|x| BuildPhase::try_from(x).ok()) { - x.name().as_ptr().cast_mut() + if let Ok(x) = u32::try_from(x.wrapping_sub(1)) { + if let Some(x) = BuildPhase::from_value(x) { + x.name().as_ptr().cast_mut() + } else { + std::ptr::null_mut() + } } else { std::ptr::null_mut() } @@ -139,7 +209,7 @@ impl PostgresReporter { unsafe { pgrx::pg_sys::pgstat_progress_update_param( pgrx::pg_sys::PROGRESS_CREATEIDX_SUBPHASE as _, - phase as _, + (phase.into_value() as i64) + 1, ); } } @@ -192,13 +262,13 @@ pub unsafe extern "C" fn ambuild( let mut reporter = PostgresReporter {}; let structures = match vchordrq_options.build.source.clone() { VchordrqBuildSourceOptions::External(external_build) => { - reporter.phase(BuildPhase::ExternalBuild); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::ExternalBuild)); let reltuples = unsafe { (*(*index_relation).rd_rel).reltuples }; reporter.tuples_total(reltuples as u64); make_external_build(vector_options.clone(), opfamily, external_build.clone()) } VchordrqBuildSourceOptions::Internal(internal_build) => { - reporter.phase(BuildPhase::InternalBuild); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::InternalBuild)); let mut tuples_total = 0_u64; let samples = 'a: { let mut rand = rand::rng(); @@ -235,17 +305,22 @@ pub unsafe extern "C" fn ambuild( samples }; reporter.tuples_total(tuples_total); - make_internal_build(vector_options.clone(), internal_build.clone(), samples) + make_internal_build( + vector_options.clone(), + internal_build.clone(), + samples, + &mut reporter, + ) } }; - reporter.phase(BuildPhase::Build); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Build)); crate::index::algorithm::build( vector_options, vchordrq_options.index, index.clone(), structures, ); - reporter.phase(BuildPhase::Inserting); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); let cache = if vchordrq_options.build.pin { let mut trace = algorithm::cache(index.clone()); trace.sort(); @@ -319,7 +394,7 @@ pub unsafe extern "C" fn ambuild( let check = || { pgrx::check_for_interrupts!(); }; - reporter.phase(BuildPhase::Compacting); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); crate::index::algorithm::maintain(opfamily, index, check); unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -882,10 +957,11 @@ unsafe fn options( (vector, rabitq) } -pub fn make_internal_build( +fn make_internal_build( vector_options: VectorOptions, internal_build: VchordrqInternalBuildOptions, mut samples: Vec>, + reporter: &mut PostgresReporter, ) -> Vec>> { use std::iter::once; k_means::preprocess(internal_build.build_threads as _, &mut samples, |sample| { @@ -903,6 +979,13 @@ pub fn make_internal_build( let num_dims = vector_options.dims as usize; let num_lists = w as usize; let num_iterations = internal_build.kmeans_iterations as _; + if result.is_empty() { + let percentage = 0; + let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); + let phase = + BuildPhase::new(BuildPhaseCode::InternalBuild, 1 + percentage).unwrap_or(default); + reporter.phase(phase); + } if num_lists > 1 { pgrx::info!( "clustering: starting, using {num_threads} threads, clustering {num_points} vectors of {num_dims} dimension into {num_lists} clusters, in {num_iterations} iterations" @@ -912,6 +995,14 @@ pub fn make_internal_build( num_threads, |i| { pgrx::check_for_interrupts!(); + if result.is_empty() { + let percentage = + ((i as f64 / num_iterations as f64) * 100.0).clamp(0.0, 100.0) as u16; + let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); + let phase = BuildPhase::new(BuildPhaseCode::InternalBuild, 1 + percentage) + .unwrap_or(default); + reporter.phase(phase); + } if num_lists > 1 { pgrx::info!("clustering: iteration {}", i + 1); } @@ -922,6 +1013,13 @@ pub fn make_internal_build( internal_build.spherical_centroids, num_iterations, ); + if result.is_empty() { + let percentage = 100; + let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); + let phase = + BuildPhase::new(BuildPhaseCode::InternalBuild, 1 + percentage).unwrap_or(default); + reporter.phase(phase); + } if num_lists > 1 { pgrx::info!("clustering: finished"); } @@ -944,7 +1042,7 @@ pub fn make_internal_build( } #[allow(clippy::collapsible_else_if)] -pub fn make_external_build( +fn make_external_build( vector_options: VectorOptions, _opfamily: Opfamily, external_build: VchordrqExternalBuildOptions, diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index a408b892..d8b9ad7b 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -12,6 +12,8 @@ use std::num::NonZeroU64; use std::ptr::NonNull; use std::sync::OnceLock; +use super::gucs::prererank_filtering; + #[repr(C)] struct Reloption { vl_len_: i32, @@ -307,7 +309,11 @@ pub unsafe extern "C" fn amrescan( (*scan).indexRelation, (*scan).heapRelation, (*scan).xs_snapshot, - hack, + if let Some(hack) = hack { + hack.as_ptr() + } else { + std::ptr::null_mut() + }, ) }) }; @@ -404,7 +410,7 @@ struct HeapFetcher { slot: *mut pgrx::pg_sys::TupleTableSlot, values: [Datum; 32], is_nulls: [bool; 32], - hack: Option>, + hack: *mut pgrx::pg_sys::IndexScanState, } impl HeapFetcher { @@ -412,7 +418,7 @@ impl HeapFetcher { index_relation: pgrx::pg_sys::Relation, heap_relation: pgrx::pg_sys::Relation, snapshot: pgrx::pg_sys::Snapshot, - hack: Option>, + hack: *mut pgrx::pg_sys::IndexScanState, ) -> Self { unsafe { let index_info = pgrx::pg_sys::BuildIndexInfo(index_relation); @@ -447,7 +453,6 @@ impl Drop for HeapFetcher { impl SearchFetcher for HeapFetcher { fn fetch(&mut self, mut ctid: ItemPointerData) -> Option<(&[Datum; 32], &[bool; 32])> { unsafe { - // perform operation let table_am = (*self.heap_relation).rd_tableam; let fetch_row_version = (*table_am) .tuple_fetch_row_version @@ -455,20 +460,20 @@ impl SearchFetcher for HeapFetcher { if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { return None; } - if let Some(hack) = self.hack { - if let Some(qual) = NonNull::new(hack.as_ref().ss.ps.qual) { + if !self.hack.is_null() && prererank_filtering() { + if let Some(qual) = NonNull::new((*self.hack).ss.ps.qual) { use pgrx::datum::FromDatum; use pgrx::memcxt::PgMemoryContexts; assert!(qual.as_ref().flags & pgrx::pg_sys::EEO_FLAG_IS_QUAL as u8 != 0); let evalfunc = qual.as_ref().evalfunc.expect("no evalfunc for qual"); - if let Some(mut econtext) = NonNull::new(hack.as_ref().ss.ps.ps_ExprContext) { - econtext.as_mut().ecxt_scantuple = self.slot; - pgrx::pg_sys::MemoryContextReset(econtext.as_ref().ecxt_per_tuple_memory); - let result = PgMemoryContexts::For(econtext.as_ref().ecxt_per_tuple_memory) + if !(*self.hack).ss.ps.ps_ExprContext.is_null() { + let econtext = (*self.hack).ss.ps.ps_ExprContext; + (*econtext).ecxt_scantuple = self.slot; + pgrx::pg_sys::MemoryContextReset((*econtext).ecxt_per_tuple_memory); + let result = PgMemoryContexts::For((*econtext).ecxt_per_tuple_memory) .switch_to(|_| { let mut is_null = true; - let datum = - evalfunc(qual.as_ptr(), econtext.as_mut(), &mut is_null); + let datum = evalfunc(qual.as_ptr(), econtext, &mut is_null); bool::from_datum(datum, is_null) }); if result != Some(true) { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index c0849bc4..f156d3ea 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -8,6 +8,7 @@ static PREWARM_DIM: GucSetting> = GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); static MAX_MAXSIM_TUPLES: GucSetting = GucSetting::::new(-1); static MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); +static PRERERANK_FILTERING: GucSetting = GucSetting::::new(false); pub fn init() { GucRegistry::define_string_guc( @@ -66,6 +67,14 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_bool_guc( + "vchordrq.prererank_filtering", + "`prererank_filtering` argument of vchordrq.", + "`prererank_filtering` argument of vchordrq.", + &PRERERANK_FILTERING, + GucContext::Userset, + GucFlags::default(), + ); unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); @@ -141,3 +150,7 @@ pub fn prewarm_dim() -> Vec { Vec::new() } } + +pub fn prererank_filtering() -> bool { + PRERERANK_FILTERING.get() +} diff --git a/tests/logic/rerank_in_table.slt b/tests/logic/rerank_in_table.slt index 3f049a78..8a2b8c44 100644 --- a/tests/logic/rerank_in_table.slt +++ b/tests/logic/rerank_in_table.slt @@ -79,6 +79,18 @@ SELECT id FROM t_expr WHERE id <= 5 OR id % 2 = 1 ORDER BY ARRAY[id::real, id::r 11 13 +statement ok +SET vchordrq.prererank_filtering to off; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +SET vchordrq.prererank_filtering to on; + query I SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); ---- From a47f5b75075e54285f05e6a75aa2a0ca6d24735b Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 2 Apr 2025 10:27:55 +0800 Subject: [PATCH 136/324] readme: move docs to docs repo (#222) pull request in docs repo: https://github.com/tensorchord/pgvecto.rs-docs/pull/106 [Rendered README.md](https://github.com/usamoi/VectorChord/blob/readme/README.md) [Preview docs](https://pgvecto-rs-docs-83ki0aiml-tensorchord.vercel.app/vectorchord/getting-started/overview.html) Signed-off-by: usamoi --- .github/workflows/check.yml | 4 +- README.md | 367 ++---------------------------------- src/index/am/am_build.rs | 2 +- 3 files changed, 22 insertions(+), 351 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index e729fc17..284024f7 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -95,7 +95,7 @@ jobs: fi - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.7 + uses: mozilla-actions/sccache-action@v0.0.9 - name: Checkout uses: actions/checkout@v4 @@ -158,7 +158,7 @@ jobs: curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.7 + uses: mozilla-actions/sccache-action@v0.0.9 - name: Checkout uses: actions/checkout@v4 diff --git a/README.md b/README.md index 1e70f511..d8a0c280 100644 --- a/README.md +++ b/README.md @@ -42,46 +42,28 @@ VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: **🔌 Seamless Integration**: Fully compatible with pgvector data types and syntax while providing optimal defaults out of the box - no manual parameter tuning needed. Just drop in VectorChord for enhanced performance. -**🔧 Accelerated Index Build**: Leverage IVF to build indexes externally (e.g., on GPU) for faster KMeans clustering, combined with RaBitQ[^3] compression to efficiently store vectors while maintaining search quality through autonomous reranking. +**🔧 Accelerated Index Build**: Leverage IVF to build indexes externally (e.g., on GPU) for faster KMeans clustering, combined with RaBitQ[^2] compression to efficiently store vectors while maintaining search quality through autonomous reranking. -[^3]: Gao, Jianyang, and Cheng Long. "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search." Proceedings of the ACM on Management of Data 2.3 (2024): 1-27. +[^2]: Gao, Jianyang, and Cheng Long. "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search." Proceedings of the ACM on Management of Data 2.3 (2024): 1-27. -**📏 Long Vector Support**: Store and search vectors up to 60,000[^4] dimensions, enabling the use of the best high-dimensional models like text-embedding-3-large with ease. +**📏 Long Vector Support**: Store and search vectors up to 60,000[^3] dimensions, enabling the use of the best high-dimensional models like text-embedding-3-large with ease. -[^4]: There is a [limitation](https://github.com/pgvector/pgvector#vector-type) at pgvector of 16,000 dimensions now. If you really have a large dimension(`16,000 [!TIP] -> If you are using the official [Docker image](https://hub.docker.com/r/tensorchord/vchord-postgres), you can skip this step. - -VectorChord depends on [pgvector](https://github.com/pgvector/pgvector), ensure the pgvector extension is available: - -```SQL -SELECT * FROM pg_available_extensions WHERE name = 'vector'; -``` -If pgvector is not available, install it using the [pgvector installation instructions](https://github.com/pgvector/pgvector#installation). - -And make sure to add `vchord.so` to the `shared_preload_libraries` in `postgresql.conf`. +## Quick Start -```SQL --- Add vchord and pgvector to shared_preload_libraries -- --- Note: A restart is required for this setting to take effect. -ALTER SYSTEM SET shared_preload_libraries = 'vchord.so'; -``` +For new users, we recommend using the Docker image to get started quickly. If you do not prefer Docker, please read [installation guide](https://docs.vectorchord.ai/vectorchord/getting-started/installation.html) for other installation methods. -## Quick Start -For new users, we recommend using the Docker image to get started quickly. ```bash docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d tensorchord/vchord-postgres:pg17-v0.2.1 + -d tensorchord/vchord-postgres:pg17-v0.2.2 ``` Then you can connect to the database using the `psql` command line tool. The default username is `postgres`, and the default password is `mysecretpassword`. @@ -92,347 +74,37 @@ psql -h localhost -p 5432 -U postgres Now you can play with VectorChord! -## Documentation - -- [Installation](#installation) - - [Docker](#docker) - - [APT](#apt) - - [More Methods](#more-methods) -- [Usage](#usage) - - [Storing](#storing) - - [Indexing](#indexing) - - [Query](#query) -- [Performance Tuning](#performance-tuning) - - [Index Build Time](#index-build-time) - - [Query Performance](#query-performance) - - [Index Internal Build](#index-internal-build) -- [Advanced Features](#advanced-features) - - [Indexing Prewarm](#indexing-prewarm) - - [Indexing Progress](#indexing-progress) - - [External Index Precomputation](#external-index-precomputation) - - [Capacity-optimized Index](#capacity-optimized-index) - - [Range Query](#range-query) - -## Installation - -### [Docker](https://docs.vectorchord.ai/vectorchord/getting-started/installation.html#docker) - -You can easily get the Docker image from: - -```bash -docker pull tensorchord/vchord-postgres:pg17-v0.2.1 -``` - -### [APT](https://docs.vectorchord.ai/vectorchord/getting-started/installation.html#from-debian-package) - -Debian and Ubuntu packages can be found on [release page](https://github.com/tensorchord/VectorChord/releases). - -To install it: -```bash -wget https://github.com/tensorchord/VectorChord/releases/download/0.2.1/postgresql-17-vchord_0.2.1-1_amd64.deb -sudo apt install postgresql-17-vchord_*.deb -``` - -### More Methods - -VectorChord also supports other installation methods, including: -- [From ZIP package](https://docs.vectorchord.ai/vectorchord/getting-started/installation.html#from-zip-package) - -## Usage - -VectorChord depends on pgvector, including the vector representation. -This way, we can keep the maximum compatibility of `pgvector` for both: -- [vector storing](https://github.com/pgvector/pgvector#storing) -- [vector query](https://github.com/pgvector/pgvector#querying). - -Since you can use them directly, your application can be easily migrated without pain! - -Before all, you need to run the following SQL to ensure the extension is enabled. +VectorChord depends on pgvector, including the vector representation. Since you can use them directly, your application can be easily migrated without pain! -```SQL -CREATE EXTENSION IF NOT EXISTS vchord CASCADE; -``` -It will install both `pgvector` and `VectorChord`, see [requirements](#requirements) for more detail. - -### Storing - -Similar to pgvector, you can create a table with vector column in VectorChord and insert some rows to it. +Similar to pgvector, you can create a table with vector column and insert some rows to it. ```sql CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3)); INSERT INTO items (embedding) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); ``` -### Indexing - -Similar to [ivfflat](https://github.com/pgvector/pgvector#ivfflat), the index type of VectorChord, RaBitQ(vchordrq) also divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It inherits the advantages of `ivfflat`, such as fast build times and less memory usage, but has [much better performance](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql#heading-ivf-vs-hnsw) than hnsw and ivfflat. - -The RaBitQ(vchordrq) index is supported on some pgvector types and metrics: - -| | vector | halfvec | bit(n) | sparsevec | -| ----------------------- | ------ | ------- | ------ | --------- | -| L2 distance / `<->` | ✅ | ✅ | 🆖 | ❌ | -| inner product / `<#>` | ✅ | ✅ | 🆖 | ❌ | -| cosine distance / `<=>` | ✅ | ✅ | 🆖 | ❌ | -| L1 distance / `<+>` | ❌ | ❌ | 🆖 | ❌ | -| Hamming distance/ `<~>` | 🆖 | 🆖 | ❌ | 🆖 | -| Jaccard distance/ `<%>` | 🆖 | 🆖 | ❌ | 🆖 | - -Where: -- ✅ means supported by pgvector and VectorChord -- ❌ means supported by pgvector but not by VectorChord -- 🆖 means not planned by pgvector and VectorChord -- 🔜 means supported by pgvector now and will be supported by VectorChord soon - -To create the VectorChord RaBitQ(vchordrq) index, you can use the following SQL. - -L2 distance -```sql -CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops) WITH (options = $$ -residual_quantization = true -[build.internal] -lists = [1000] -spherical_centroids = false -$$); -``` - -> [!NOTE] -> - `options` are specified using a [TOML: Tom's Obvious Minimal Language](https://toml.io/) string. -> - Set `residual_quantization` to `true` and `spherical_centroids` to `false` for L2 distance embeddings -> - Use `halfvec_l2_ops` for `halfvec` -> - The recommend `lists` could be rows / 1000 for up to 1M rows and 4 * sqrt(rows) for over 1M rows - -Inner product -```sql -CREATE INDEX ON items USING vchordrq (embedding vector_ip_ops) WITH (options = $$ -residual_quantization = false -[build.internal] -lists = [1000] -spherical_centroids = true -$$); -``` - -Cosine distance -```sql -CREATE INDEX ON items USING vchordrq (embedding vector_cosine_ops) WITH (options = $$ -residual_quantization = false -[build.internal] -lists = [1000] -spherical_centroids = true -$$); -``` - -> [!NOTE] -> - Set `residual_quantization` to `false` and `spherical_centroids` to `true` for cosine similarity embeddings -> - Use `halfvec_cosine_ops`/`halfvec_ip_ops` for `halfvec` - -### Query - -The query statement is exactly the same as pgvector. VectorChord supports any filter operation and WHERE/JOIN clauses like pgvecto.rs with VBASE. -```SQL -SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5; -``` -Supported distance functions are: -- <-> - L2 distance -- <#> - (negative) inner product -- <=> - cosine distance - -## Performance Tuning - -### Index Build Time - -Index building can parallelized, and with external centroid precomputation, the total time is primarily limited by disk speed. Optimize parallelism using the following settings: - -```SQL --- Set this to the number of CPU cores available for parallel operations. -SET max_parallel_maintenance_workers = 8; -SET max_parallel_workers = 8; - --- Adjust the total number of worker processes. --- Note: A restart is required for this setting to take effect. -ALTER SYSTEM SET max_worker_processes = 8; -``` - -When building an index on a table with more than 10 million vectors, you can choose to consume more memory to accelerate the process by setting `build.pin` to `true`: - -```sql -CREATE INDEX ON items USING vchordrq (embedding vector_cosine_ops) WITH (options = $$ -residual_quantization = false -build.pin = true -$$); -``` - -### Query Performance -You can fine-tune the search performance by adjusting the `probes` and `epsilon` parameters: - -```sql --- Set probes to control the number of lists scanned. --- Recommended range: 3%–10% of the total `lists` value. -SET vchordrq.probes = 100; - --- Set epsilon to control the reranking precision. --- Larger value means more rerank for higher recall rate and latency. --- If you need a less precise query, set it to 1.0 might be appropriate. --- Recommended range: 1.0–1.9. Default value is 1.9. -SET vchordrq.epsilon = 1.9; - --- vchordrq relies on a projection matrix to optimize performance. --- Add your vector dimensions to the `prewarm_dim` list to reduce latency. --- If this is not configured, the first query will have higher latency as the matrix is generated on demand. --- Default value: '64,128,256,384,512,768,1024,1536' --- Note: This setting requires a database restart to take effect. -ALTER SYSTEM SET vchordrq.prewarm_dim = '64,128,256,384,512,768,1024,1536'; -``` - -And for postgres's setting -```SQL --- If using SSDs, set `effective_io_concurrency` to 200 for faster disk I/O. -SET effective_io_concurrency = 200; - --- Disable JIT (Just-In-Time Compilation) as it offers minimal benefit (1–2%) --- and adds overhead for single-query workloads. -SET jit = off; - --- Allocate at least 25% of total memory to `shared_buffers`. --- For disk-heavy workloads, you can increase this to up to 90% of total memory. You may also want to disable swap with network storage to avoid io hang. --- Note: A restart is required for this setting to take effect. -ALTER SYSTEM SET shared_buffers = '8GB'; -``` - -### Index Internal Build - -There are 4 options for index internal build, except `lists`. - -* `spherical_centroids`: Specifies whether to use spherical K-means algorithm. If the model generates cosine similarity embeddings, set this to `true`; otherwise, set to `false`. Possible values: `true`, `false`. Default: `false`. - -* `sampling_factor`: Specifies the number of vectors sampled by K-means algorithm. The higher this value, the slower the build, the greater the memory consumption, and the better search performance. Possible values: any integer between `1` and `1024`. Default: `256`. - -* `kmeans_iterations`: Specifies the number of iterations for K-means algorithm. The higher this value, the slower the build. Possible values: any integer between `0` and `1024`. Default: `10`. - -* `build_threads`: Specifies the number of threads used by K-means algorithm. The higher this value, the faster the build. Possible values: any integer between `1` and `255`. Default: `1`. - -An example of these options: - -```sql -CREATE INDEX ON t USING vchordrq (embedding vector_cosine_ops) WITH (options = $$ -residual_quantization = false -[build.internal] -lists = [1000] -spherical_centroids = true -sampling_factor = 512 -kmeans_iterations = 25 -build_threads = 16 -$$); -``` - -## Advanced Features - -### Indexing Prewarm - -For disk-first indexing, RaBitQ(vchordrq) is loaded from disk for the first query, -and then cached in memory if `shared_buffer` is sufficient, resulting in a significant cold-start slowdown. - -To improve performance for the first query, you can try the following SQL that preloads the index into memory. +With VectorChord, you can create `vchordrq` indexes. ```SQL --- vchordrq_prewarm(index_name::regclass) to prewarm the index into the shared buffer -SELECT vchordrq_prewarm('gist_train_embedding_idx'::regclass)" -``` - -### Indexing Progress -You can check the indexing progress by querying the `pg_stat_progress_create_index` view. -```SQL -SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index; -``` - -### External Index Precomputation - -Unlike pure SQL, an external index precomputation will first do clustering outside and insert centroids to a PostgreSQL table. Although it might be more complicated, external build is definitely much faster on larger dataset (>5M). - -To get started, you need to do a clustering of vectors using `faiss`, `scikit-learn` or any other clustering library. - -The centroids should be preset in a table of any name with 3 columns: -- id(integer): id of each centroid, should be unique -- parent(integer, nullable): parent id of each centroid, should be NULL for normal clustering -- vector(vector): representation of each centroid, `pgvector` vector type - -And example could be like this: - -```sql --- Create table of centroids -CREATE TABLE public.centroids (id integer NOT NULL UNIQUE, parent integer, vector vector(768)); --- Insert centroids into it -INSERT INTO public.centroids (id, parent, vector) VALUES (1, NULL, '{0.1, 0.2, 0.3, ..., 0.768}'); -INSERT INTO public.centroids (id, parent, vector) VALUES (2, NULL, '{0.4, 0.5, 0.6, ..., 0.768}'); -INSERT INTO public.centroids (id, parent, vector) VALUES (3, NULL, '{0.7, 0.8, 0.9, ..., 0.768}'); --- ... - --- Create index using the centroid table -CREATE INDEX ON gist_train USING vchordrq (embedding vector_l2_ops) WITH (options = $$ -[build.external] -table = 'public.centroids' -$$); -``` - -To simplify the workflow, we provide end-to-end scripts for external index pre-computation, see [scripts](./scripts/README.md#run-external-index-precomputation-toolkit). - -### Capacity-optimized Index - -The default behavior of Vectorchord is `performance-optimized`, -which uses more disk space but has a better latency: -- About `80G` for `5M` 768 dim vectors -- About `800G` for `100M` 768 dim vectors - -Although it is acceptable for such large data, it could be switched to `capacity-optimized` index and save about **50%** of your disk space. - -For `capacity-optimized` index, just enable the `rerank_in_table` option when creating the index: -```sql CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops) WITH (options = $$ residual_quantization = true -rerank_in_table = true [build.internal] lists = [] $$); ``` -> [!CAUTION] -> Compared to the `performance-optimized` index, the `capacity-optimized` index will have a **30-50%** increase in latency and QPS loss at query. - -### Range Query +And then perform a vector search using `SELECT ... ORDER BY ... LIMIT ...`. -To query vectors within a certain distance range, you can use the following syntax. ```SQL --- Query vectors within a certain distance range -SELECT vec FROM t WHERE vec <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012) -ORDER BY embedding <-> '[0.24, 0.24, 0.24]' LIMIT 5; -``` - -In this expression, `vec <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012)` is equals to `vec <-> '[0.24, 0.24, 0.24]' < 0.012`. However, the latter one will trigger a **exact nearest neighbor search** as the grammar could not be pushed down. - -Supported range functions are: -- `<<->>` - L2 distance -- `<<#>>` - (negative) inner product -- `<<=>>` - cosine distance - -## Development - -### Build the Postgres Docker Image with VectorChord extension - -Follow the steps in [Dev Guidance](./scripts/README.md#build-docker). - -### Installing From Source - -Install pgrx according to [pgrx's instruction](https://github.com/pgcentralfoundation/pgrx?tab=readme-ov-file#getting-started). -```bash -cargo install --locked cargo-pgrx -cargo pgrx init --pg$(pg_config --version | awk '{print $2}' | cut -d. -f1) $(which pg_config) # To init with system postgres, with pg_config in PATH -cargo pgrx install --release --sudo # To install the extension into the system postgres with sudo +SET vchordrq.probes TO ''; +SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5; ``` -## Limitations - -- KMeans Clustering: The built-in KMeans clustering depends on multi-thread in-memory build and may require substantial memory. We strongly recommend using external centroid precomputation for efficient index construction. +For more usage, please read: +* [Indexing](https://docs.vectorchord.ai/vectorchord/usage/indexing.html) +* [Performance Tuning](https://docs.vectorchord.ai/vectorchord/usage/performance-tuning.html) +* [Advanced Features](https://docs.vectorchord.ai/vectorchord/usage/advanced-features.html) ## License @@ -444,7 +116,6 @@ This software is licensed under a dual license model: You may choose either license based on your needs. We welcome any commercial collaboration or support, so please email us with any questions or requests regarding the licenses. - [image-compare]: https://github.com/user-attachments/assets/2d985f1e-7093-4c3a-8bf3-9f0b92c0e7e7 [license-1-link]: https://github.com/tensorchord/VectorChord#license [license-1-shield]: https://img.shields.io/badge/License-AGPLv3-green?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCIgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiBmaWxsPSIjZmZmZmZmIj48cGF0aCBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik0xMi43NSAyLjc1YS43NS43NSAwIDAwLTEuNSAwVjQuNUg5LjI3NmExLjc1IDEuNzUgMCAwMC0uOTg1LjMwM0w2LjU5NiA1Ljk1N0EuMjUuMjUgMCAwMTYuNDU1IDZIMi4zNTNhLjc1Ljc1IDAgMTAwIDEuNUgzLjkzTC41NjMgMTUuMThhLjc2Mi43NjIgMCAwMC4yMS44OGMuMDguMDY0LjE2MS4xMjUuMzA5LjIyMS4xODYuMTIxLjQ1Mi4yNzguNzkyLjQzMy42OC4zMTEgMS42NjIuNjIgMi44NzYuNjJhNi45MTkgNi45MTkgMCAwMDIuODc2LS42MmMuMzQtLjE1NS42MDYtLjMxMi43OTItLjQzMy4xNS0uMDk3LjIzLS4xNTguMzEtLjIyM2EuNzUuNzUgMCAwMC4yMDktLjg3OEw1LjU2OSA3LjVoLjg4NmMuMzUxIDAgLjY5NC0uMTA2Ljk4NC0uMzAzbDEuNjk2LTEuMTU0QS4yNS4yNSAwIDAxOS4yNzUgNmgxLjk3NXYxNC41SDYuNzYzYS43NS43NSAwIDAwMCAxLjVoMTAuNDc0YS43NS43NSAwIDAwMC0xLjVIMTIuNzVWNmgxLjk3NGMuMDUgMCAuMS4wMTUuMTQuMDQzbDEuNjk3IDEuMTU0Yy4yOS4xOTcuNjMzLjMwMy45ODQuMzAzaC44ODZsLTMuMzY4IDcuNjhhLjc1Ljc1IDAgMDAuMjMuODk2Yy4wMTIuMDA5IDAgMCAuMDAyIDBhMy4xNTQgMy4xNTQgMCAwMC4zMS4yMDZjLjE4NS4xMTIuNDUuMjU2Ljc5LjRhNy4zNDMgNy4zNDMgMCAwMDIuODU1LjU2OCA3LjM0MyA3LjM0MyAwIDAwMi44NTYtLjU2OWMuMzM4LS4xNDMuNjA0LS4yODcuNzktLjM5OWEzLjUgMy41IDAgMDAuMzEtLjIwNi43NS43NSAwIDAwLjIzLS44OTZMMjAuMDcgNy41aDEuNTc4YS43NS43NSAwIDAwMC0xLjVoLTQuMTAyYS4yNS4yNSAwIDAxLS4xNC0uMDQzbC0xLjY5Ny0xLjE1NGExLjc1IDEuNzUgMCAwMC0uOTg0LS4zMDNIMTIuNzVWMi43NXpNMi4xOTMgMTUuMTk4YTUuNDE4IDUuNDE4IDAgMDAyLjU1Ny42MzUgNS40MTggNS40MTggMCAwMDIuNTU3LS42MzVMNC43NSA5LjM2OGwtMi41NTcgNS44M3ptMTQuNTEtLjAyNGMuMDgyLjA0LjE3NC4wODMuMjc1LjEyNi41My4yMjMgMS4zMDUuNDUgMi4yNzIuNDVhNS44NDYgNS44NDYgMCAwMDIuNTQ3LS41NzZMMTkuMjUgOS4zNjdsLTIuNTQ3IDUuODA3eiI+PC9wYXRoPjwvc3ZnPgo= diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index c5e83c70..0aafeb60 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -1091,7 +1091,7 @@ fn make_external_build( } }); if parents.len() >= 2 && parents.values().all(|x| x.is_none()) { - // if there are more than one vertexs and no edges, + // if there are more than one vertex and no edges, // assume there is an implicit root let n = parents.len(); let mut result = Vec::new(); From d114e6a87027d510bc28e1cebe03a4a27d00b049 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 2 Apr 2025 16:05:00 +0800 Subject: [PATCH 137/324] chore: change default lists and probes to empty (#223) This changes default `lists` from `[1000]` to `[]`, and default `probes` to `10` to ` `. The default value of `1000/10` is too small for big data and may cause issues for small data. `lists = []` implies "clusters are not used, and all data are in the same set". It's good for small data. Signed-off-by: usamoi --- src/index/gucs.rs | 2 +- src/index/types.rs | 2 +- tests/logic/index.slt | 3 +++ tests/logic/partition.slt | 26 +++++++++++++------------- 4 files changed, 18 insertions(+), 15 deletions(-) diff --git a/src/index/gucs.rs b/src/index/gucs.rs index f156d3ea..97d71a6a 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -1,7 +1,7 @@ use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; use std::ffi::CStr; -static PROBES: GucSetting> = GucSetting::>::new(Some(c"10")); +static PROBES: GucSetting> = GucSetting::>::new(Some(c"")); static EPSILON: GucSetting = GucSetting::::new(1.9); static MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); static PREWARM_DIM: GucSetting> = diff --git a/src/index/types.rs b/src/index/types.rs index 8edec722..e626ec03 100644 --- a/src/index/types.rs +++ b/src/index/types.rs @@ -23,7 +23,7 @@ pub struct VchordrqInternalBuildOptions { impl VchordrqInternalBuildOptions { fn default_lists() -> Vec { - vec![1000] + Vec::new() } fn validate_lists(lists: &[u32]) -> Result<(), ValidationError> { if !lists.is_sorted() { diff --git a/tests/logic/index.slt b/tests/logic/index.slt index c3f87954..f2688752 100644 --- a/tests/logic/index.slt +++ b/tests/logic/index.slt @@ -38,6 +38,9 @@ lists = [32] spherical_centroids = true $$); +statement ok +SET vchordrq.probes = '16'; + query I SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; ---- diff --git a/tests/logic/partition.slt b/tests/logic/partition.slt index 506995ef..1e681c5e 100644 --- a/tests/logic/partition.slt +++ b/tests/logic/partition.slt @@ -1,51 +1,51 @@ # partition table statement ok -CREATE TABLE items (val vector(3), category_id int) PARTITION BY LIST(category_id); +CREATE TABLE t (val vector(3), category_id int) PARTITION BY LIST(category_id); statement ok -CREATE TABLE id_123 PARTITION OF items FOR VALUES IN (1, 2, 3); +CREATE TABLE id_123 PARTITION OF t FOR VALUES IN (1, 2, 3); statement ok -CREATE TABLE id_456 PARTITION OF items FOR VALUES IN (4, 5, 6); +CREATE TABLE id_456 PARTITION OF t FOR VALUES IN (4, 5, 6); statement ok -CREATE TABLE id_789 PARTITION OF items FOR VALUES IN (7, 8, 9); +CREATE TABLE id_789 PARTITION OF t FOR VALUES IN (7, 8, 9); statement ok -INSERT INTO items (val, category_id) +INSERT INTO t (val, category_id) SELECT ARRAY[random(), random(), random()]::real[], (random() * 6 + 1)::int FROM generate_series(1, 1000); statement ok -CREATE INDEX ON items USING vchordrq (val public.vector_l2_ops); +CREATE INDEX ON t USING vchordrq (val public.vector_l2_ops); query I -SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; ---- 10 -statement error the relation "items_val_idx" is not an index -select vchordrq_prewarm('items_val_idx'); +statement error the relation "t_val_idx" is not an index +select vchordrq_prewarm('t_val_idx'); statement ok CREATE INDEX ON id_123 USING vchordrq (val vector_cosine_ops); query I -SELECT COUNT(1) FROM (SELECT 1 FROM items ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; ---- 10 # partial index statement ok -CREATE INDEX ON items USING vchordrq (val public.vector_ip_ops) WHERE (category_id = 1); +CREATE INDEX ON t USING vchordrq (val public.vector_ip_ops) WHERE (category_id = 1); query I SELECT COUNT(1) FROM -(SELECT 1 FROM items WHERE (category_id = 1) ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +(SELECT 1 FROM t WHERE (category_id = 1) ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; ---- 10 statement ok -DROP TABLE id_789, id_456, id_123, items; +DROP TABLE id_789, id_456, id_123, t; From b0514e6045ae3a7d5ea88f3f0a8e20060bf6902e Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 8 Apr 2025 15:14:14 +0800 Subject: [PATCH 138/324] feat: allows skipping rerank (#225) Enable it by `SET vchordrq.allows_skipping_rerank TO on`. Skip rerank by enabling it and `SET vchordrq.epsilon = 0.0`. Signed-off-by: usamoi --- .cargo/config.toml | 7 --- build.rs | 5 +- crates/algorithm/src/build.rs | 14 +++-- crates/algorithm/src/freepages.rs | 8 ++- crates/algorithm/src/insert.rs | 22 ++++++-- crates/algorithm/src/lib.rs | 10 +++- crates/algorithm/src/maintain.rs | 6 +- crates/algorithm/src/operator.rs | 14 ++--- crates/algorithm/src/prewarm.rs | 37 ++++++------ crates/algorithm/src/rerank.rs | 48 +++++++++++----- crates/algorithm/src/search.rs | 87 +++++++++++++++++++---------- crates/algorithm/src/select_heap.rs | 50 +++++++++-------- crates/algorithm/src/tape.rs | 8 +-- crates/algorithm/src/tuples.rs | 34 +++++------ crates/algorithm/src/vectors.rs | 6 +- crates/rabitq/src/binary.rs | 11 ++-- crates/rabitq/src/block.rs | 8 +-- crates/rabitq/src/lib.rs | 4 +- src/index/algorithm.rs | 6 +- src/index/am/mod.rs | 13 +++-- src/index/gucs.rs | 13 +++++ src/index/scanners/maxsim.rs | 40 ++++++++----- src/index/scanners/mod.rs | 1 + 23 files changed, 280 insertions(+), 172 deletions(-) diff --git a/.cargo/config.toml b/.cargo/config.toml index f4f72dfd..2196f5fc 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,9 +1,2 @@ -[build] -rustdocflags = ["--document-private-items"] - -[target.'cfg(target_os="macos")'] -# Postgres symbols won't be available until runtime -rustflags = ["-Clink-arg=-Wl,-undefined,dynamic_lookup"] - [env] CC = "clang" diff --git a/build.rs b/build.rs index 1fadc37c..82386588 100644 --- a/build.rs +++ b/build.rs @@ -1,4 +1,5 @@ fn main() { - println!(r#"cargo::rustc-check-cfg=cfg(pgrx_embed)"#); - println!(r#"cargo::rustc-check-cfg=cfg(feature, values("pg12"))"#); + if std::env::var("CARGO_CFG_TARGET_OS").ok().as_deref() == Some("macos") { + println!("cargo::rustc-link-arg-cdylib=-Wl,-undefined,dynamic_lookup"); + } } diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 735a741e..67fa3723 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -39,7 +39,7 @@ pub fn build( }, })); } - level.push(chain.err().unwrap()); + level.push(chain.expect_err("internal error: 0-dimensional vector")); } pointer_of_means.push(level); } @@ -129,8 +129,12 @@ pub fn build( is_residual, rerank_in_heap: vchordrq_options.rerank_in_table, vectors_first: vectors.first(), - root_mean: pointer_of_means.last().unwrap()[0], - root_first: pointer_of_firsts.last().unwrap()[0], + root_mean: pointer_of_means + .last() + .expect("internal error: empty structure")[0], + root_first: pointer_of_firsts + .last() + .expect("internal error: empty structure")[0], freepage_first: freepage.first(), }); } @@ -154,8 +158,7 @@ where } fn push(&mut self, branch: Branch) { self.branches.push(branch); - if self.branches.len() == 32 { - let chunk = std::array::from_fn::<_, 32, _>(|_| self.branches.pop().unwrap()); + if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); loop { let freespace = self.tape.freespace(); @@ -182,6 +185,7 @@ where self.tape.tape_move(); } } + self.branches.clear(); } } fn into_inner(self) -> (TapeWriter, Vec>) { diff --git a/crates/algorithm/src/freepages.rs b/crates/algorithm/src/freepages.rs index 6b54889c..976ae903 100644 --- a/crates/algorithm/src/freepages.rs +++ b/crates/algorithm/src/freepages.rs @@ -13,7 +13,9 @@ pub fn mark(index: impl RelationWrite, freepage_first: u32, values: &[u32]) { loop { let mut freespace_guard = index.write(current, false); if freespace_guard.len() == 0 { - freespace_guard.alloc(&FreepageTuple::serialize(&FreepageTuple {})); + freespace_guard + .alloc(&FreepageTuple::serialize(&FreepageTuple {})) + .expect("implementation: a clear page cannot accommodate a single tuple"); } let freespace_bytes = freespace_guard.get_mut(1).expect("data corruption"); let mut freespace_tuple = FreepageTuple::deserialize_mut(freespace_bytes); @@ -36,7 +38,9 @@ pub fn fetch(index: impl RelationWrite, freepage_first: u32) -> Option { loop { let mut freespace_guard = index.write(current, false); if freespace_guard.len() == 0 { - freespace_guard.alloc(&FreepageTuple::serialize(&FreepageTuple {})); + freespace_guard + .alloc(&FreepageTuple::serialize(&FreepageTuple {})) + .expect("implementation: a clear page cannot accommodate a single tuple"); } let freespace_bytes = freespace_guard.get_mut(1).expect("data corruption"); let mut freespace_tuple = FreepageTuple::deserialize_mut(freespace_bytes); diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 6ad410a9..63bb80a2 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -58,17 +58,19 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto let mut results = LinkedVec::new(); { let (first, residual) = state; - let lut = if let Some(residual) = residual { + let lut_block = if let Some(residual) = residual { &O::Vector::compute_lut_block(residual.as_borrowed()) + } else if let Some(lut_block) = default_lut_block.as_ref() { + lut_block } else { - default_lut_block.as_ref().unwrap() + unreachable!() }; tape::read_h1_tape( index.clone(), first, || { RAccess::new( - (&lut.4, (lut.0, lut.1, lut.2, lut.3, 1.9f32)), + (&lut_block.1, (lut_block.0, 1.9f32)), O::Distance::block_accessor(), ) }, @@ -81,8 +83,9 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto let mut heap = SelectHeap::from_vec(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + while let Some((_, AlwaysEqual(mean), AlwaysEqual(first))) = + pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) + { if is_residual { let (dis_u, residual_u) = vectors::read_for_h1_tuple::( index.clone(), @@ -112,7 +115,9 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); } } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop().unwrap(); + let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache + .pop() + .expect("invariant is violated: tree is not height-balanced"); (first, mean) } }; @@ -142,3 +147,8 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto tape::append(index.clone(), jump_tuple.appendable_first(), &bytes, false); } + +fn pop_if(heap: &mut SelectHeap, mut predicate: impl FnMut(&T) -> bool) -> Option { + let peek = heap.peek()?; + if predicate(peek) { heap.pop() } else { None } +} diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 0baebc8c..9c9eebbe 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -22,7 +22,7 @@ pub use cache::cache; pub use insert::insert; pub use maintain::maintain; pub use prewarm::prewarm; -pub use rerank::{how, rerank_heap, rerank_index}; +pub use rerank::{how, rerank_heap, rerank_index, skip}; pub use search::{search, search_and_estimate}; use std::ops::{Deref, DerefMut}; @@ -36,16 +36,24 @@ pub struct Opaque { } pub trait Page: Sized { + #[must_use] fn get_opaque(&self) -> &Opaque; + #[must_use] fn get_opaque_mut(&mut self) -> &mut Opaque; + #[must_use] fn len(&self) -> u16; + #[must_use] fn is_empty(&self) -> bool { self.len() == 0 } + #[must_use] fn get(&self, i: u16) -> Option<&[u8]>; + #[must_use] fn get_mut(&mut self, i: u16) -> Option<&mut [u8]>; + #[must_use] fn alloc(&mut self, data: &[u8]) -> Option; fn free(&mut self, i: u16); + #[must_use] fn freespace(&self) -> u16; fn clear(&mut self); } diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index db3b6c64..e3f2ddf1 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -155,13 +155,12 @@ where fn create(extend: E) -> Self { Self { tape: TapeWriter::create(extend), - branches: Vec::new(), + branches: Vec::with_capacity(32), } } fn push(&mut self, branch: Branch) { self.branches.push(branch); - if self.branches.len() == 32 { - let chunk = std::array::from_fn::<_, 32, _>(|_| self.branches.pop().unwrap()); + if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); loop { let freespace = self.tape.freespace(); @@ -187,6 +186,7 @@ where self.tape.tape_move(); } } + self.branches.clear(); } } fn into_inner(self) -> (TapeWriter, Vec>) { diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index fcd6ff12..3ce850cc 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -179,7 +179,7 @@ impl [u8; 16], [u8; 16], (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), - (f32, f32, f32, f32, f32), + ((f32, f32, f32, f32), f32), > for Block { type Output = [Distance; 32]; @@ -199,7 +199,7 @@ impl &[f32; 32], &[f32; 32], ), - (dis_v_2, b, k, qvector_sum, epsilon): (f32, f32, f32, f32, f32), + ((dis_v_2, b, k, qvector_sum), epsilon): ((f32, f32, f32, f32), f32), ) -> Self::Output { std::array::from_fn(|i| { let rough = dis_u_2[i] @@ -217,7 +217,7 @@ impl [u8; 16], [u8; 16], (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), - (f32, f32, f32, f32, f32), + ((f32, f32, f32, f32), f32), > for Block { type Output = [Distance; 32]; @@ -232,7 +232,7 @@ impl fn finish( self, (_, factor_ppc, factor_ip, factor_err): (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), - (dis_v_2, b, k, qvector_sum, epsilon): (f32, f32, f32, f32, f32), + ((dis_v_2, b, k, qvector_sum), epsilon): ((f32, f32, f32, f32), f32), ) -> Self::Output { std::array::from_fn(|i| { let rough = 0.5 * b * factor_ppc[i] @@ -416,7 +416,7 @@ impl Vector for VectOwned { fn from_owned(vector: OwnedVector) -> Self { match vector { OwnedVector::Vecf32(x) => x, - _ => unreachable!(), + _ => panic!("internal error: should be a vector"), } } @@ -457,7 +457,7 @@ impl Vector for VectOwned { fn from_owned(vector: OwnedVector) -> Self { match vector { OwnedVector::Vecf16(x) => x, - _ => unreachable!(), + _ => panic!("internal error: should be a halfvec"), } } @@ -483,7 +483,7 @@ pub trait OperatorDistance: 'static + Debug + Copy { [u8; 16], [u8; 16], (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), - (f32, f32, f32, f32, f32), + ((f32, f32, f32, f32), f32), Output = [Distance; 32], > + Default; diff --git a/crates/algorithm/src/prewarm.rs b/crates/algorithm/src/prewarm.rs index e611f740..77042024 100644 --- a/crates/algorithm/src/prewarm.rs +++ b/crates/algorithm/src/prewarm.rs @@ -1,9 +1,14 @@ use crate::operator::{FunctionalAccessor, Operator}; use crate::tuples::*; use crate::{Page, RelationRead, tape, vectors}; +use std::error::Error; use std::fmt::Write; -pub fn prewarm(index: impl RelationRead, height: i32, check: impl Fn()) -> String { +pub fn prewarm( + index: impl RelationRead, + height: i32, + check: impl Fn(), +) -> Result> { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -13,10 +18,10 @@ pub fn prewarm(index: impl RelationRead, height: i32, check: impl F drop(meta_guard); let mut message = String::new(); - writeln!(message, "height of root: {}", height_of_root).unwrap(); + writeln!(message, "height of root: {}", height_of_root)?; let prewarm_max_height = if height < 0 { 0 } else { height as u32 }; if prewarm_max_height > height_of_root { - return message; + return Ok(message); } type State = Vec; let mut state: State = { @@ -25,12 +30,12 @@ pub fn prewarm(index: impl RelationRead, height: i32, check: impl F vectors::read_for_h1_tuple::(index.clone(), root_mean, ()); results.push(root_first); } - writeln!(message, "------------------------").unwrap(); - writeln!(message, "number of nodes: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", 1).unwrap(); + writeln!(message, "------------------------")?; + writeln!(message, "number of nodes: {}", results.len())?; + writeln!(message, "number of pages: {}", 1)?; results }; - let mut step = |state: State| { + let mut step = |state: State| -> Result<_, Box> { let mut counter = 0_usize; let mut results = Vec::new(); for first in state { @@ -54,13 +59,13 @@ pub fn prewarm(index: impl RelationRead, height: i32, check: impl F }, ); } - writeln!(message, "------------------------").unwrap(); - writeln!(message, "number of nodes: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); - results + writeln!(message, "------------------------")?; + writeln!(message, "number of nodes: {}", results.len())?; + writeln!(message, "number of pages: {}", counter)?; + Ok(results) }; for _ in (std::cmp::max(1, prewarm_max_height)..height_of_root).rev() { - state = step(state); + state = step(state)?; } if prewarm_max_height == 0 { let mut counter = 0_usize; @@ -100,9 +105,9 @@ pub fn prewarm(index: impl RelationRead, height: i32, check: impl F }, ); } - writeln!(message, "------------------------").unwrap(); - writeln!(message, "number of nodes: {}", results.len()).unwrap(); - writeln!(message, "number of pages: {}", counter).unwrap(); + writeln!(message, "------------------------")?; + writeln!(message, "number of nodes: {}", results.len())?; + writeln!(message, "number of pages: {}", counter)?; } - message + Ok(message) } diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index e69dd15b..45979c9d 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -8,6 +8,12 @@ use std::collections::BinaryHeap; use std::num::NonZeroU64; use vector::VectorOwned; +type Results = Vec<( + Reverse, + AlwaysEqual, + AlwaysEqual, +)>; + pub fn how(index: impl RelationRead) -> RerankMethod { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); @@ -23,17 +29,14 @@ pub fn how(index: impl RelationRead) -> RerankMethod { pub fn rerank_index( index: impl RelationRead, vector: O::Vector, - results: Vec<( - Reverse, - AlwaysEqual, - AlwaysEqual, - )>, + results: Results, ) -> impl Iterator { let mut heap = BinaryHeap::from(results); let mut cache = BinaryHeap::<(Reverse, _)>::new(); std::iter::from_fn(move || { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); + while let Some((_, AlwaysEqual(pay_u), AlwaysEqual(mean))) = + pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) + { if let Some(dis_u) = vectors::read_for_h0_tuple::( index.clone(), mean, @@ -53,11 +56,7 @@ pub fn rerank_index( pub fn rerank_heap( vector: O::Vector, - results: Vec<( - Reverse, - AlwaysEqual, - AlwaysEqual, - )>, + results: Results, mut fetch: F, ) -> impl Iterator where @@ -66,8 +65,9 @@ where let mut heap = BinaryHeap::from(results); let mut cache = BinaryHeap::<(Reverse, _)>::new(); std::iter::from_fn(move || { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(_), AlwaysEqual(pay_u)) = heap.pop().unwrap(); + while let Some((_, AlwaysEqual(pay_u), AlwaysEqual(_))) = + pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) + { let vector = O::Vector::elements_and_metadata(vector.as_borrowed()); if let Some(vec_u) = fetch(pay_u) { let vec_u = O::Vector::elements_and_metadata(vec_u.as_borrowed()); @@ -81,3 +81,23 @@ where Some((dis_u, pay_u)) }) } + +pub fn skip(results: Results) -> impl Iterator { + let results = BinaryHeap::from(results); + results + .into_iter() + .map(|(Reverse(x), AlwaysEqual(y), _)| (x, y)) +} + +fn pop_if( + heap: &mut BinaryHeap, + mut predicate: impl FnMut(&mut T) -> bool, +) -> Option { + use std::collections::binary_heap::PeekMut; + let mut peek = heap.peek_mut()?; + if predicate(&mut peek) { + Some(PeekMut::pop(peek)) + } else { + None + } +} diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index 76084023..343eb214 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -11,8 +11,8 @@ use vector::{VectorBorrowed, VectorOwned}; type Results = Vec<( Reverse, - AlwaysEqual, AlwaysEqual, + AlwaysEqual, )>; pub fn search( @@ -30,7 +30,7 @@ pub fn search( assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); if height_of_root as usize != 1 + probes.len() { panic!( - "need {} probes, but {} probes provided", + "usage: need {} probes, but {} probes provided", height_of_root - 1, probes.len() ); @@ -65,22 +65,24 @@ pub fn search( let step = |state: State| { let mut results = LinkedVec::new(); for (first, residual) in state { - let lut = if let Some(residual) = residual { + let lut_block = if let Some(residual) = residual { &O::Vector::compute_lut_block(residual.as_borrowed()) + } else if let Some((lut_block, _)) = default_lut.as_ref() { + lut_block } else { - default_lut.as_ref().map(|x| &x.0).unwrap() + unreachable!() }; tape::read_h1_tape( index.clone(), first, || { RAccess::new( - (&lut.4, (lut.0, lut.1, lut.2, lut.3, epsilon)), + (&lut_block.1, (lut_block.0, epsilon)), O::Distance::block_accessor(), ) }, |lowerbound, mean, first| { - results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(first))); + results.push((Reverse(lowerbound), AlwaysEqual(first), AlwaysEqual(mean))); }, |_| (), ); @@ -90,8 +92,9 @@ pub fn search( let index = index.clone(); let vector = vector.as_borrowed(); std::iter::from_fn(move || { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + while let Some((_, AlwaysEqual(first), AlwaysEqual(mean))) = + pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) + { if is_residual { let (dis_u, residual_u) = vectors::read_for_h1_tuple::( index.clone(), @@ -131,23 +134,26 @@ pub fn search( let mut results = LinkedVec::new(); for (first, residual) in state { - let lut = if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { - &O::Vector::compute_lut(residual) - } else { - default_lut.as_ref().unwrap() - }; + let (lut_block, lut_binary) = + if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { + &O::Vector::compute_lut(residual) + } else if let Some(lut) = default_lut.as_ref() { + lut + } else { + unreachable!() + }; let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); let mut callback = |lowerbound, mean, payload| { - results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(payload))); + results.push((Reverse(lowerbound), AlwaysEqual(payload), AlwaysEqual(mean))); }; tape::read_frozen_tape( index.clone(), jump_tuple.frozen_first(), || { RAccess::new( - (&lut.0.4, (lut.0.0, lut.0.1, lut.0.2, lut.0.3, epsilon)), + (&lut_block.1, (lut_block.0, epsilon)), O::Distance::block_accessor(), ) }, @@ -157,7 +163,7 @@ pub fn search( tape::read_appendable_tape( index.clone(), jump_tuple.appendable_first(), - |code| O::Distance::compute_lowerbound_binary(&lut.1, code, epsilon), + |code| O::Distance::compute_lowerbound_binary(lut_binary, code, epsilon), &mut callback, |_| (), ); @@ -181,7 +187,7 @@ pub fn search_and_estimate( assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); if height_of_root as usize != 1 + probes.len() { panic!( - "need {} probes, but {} probes provided", + "usage: need {} probes, but {} probes provided", height_of_root - 1, probes.len() ); @@ -216,22 +222,24 @@ pub fn search_and_estimate( let step = |state: State| { let mut results = LinkedVec::new(); for (first, residual) in state { - let lut = if let Some(residual) = residual { + let lut_block = if let Some(residual) = residual { &O::Vector::compute_lut_block(residual.as_borrowed()) + } else if let Some((lut_block, _)) = default_lut.as_ref() { + lut_block } else { - default_lut.as_ref().map(|x| &x.0).unwrap() + unreachable!() }; tape::read_h1_tape( index.clone(), first, || { RAccess::new( - (&lut.4, (lut.0, lut.1, lut.2, lut.3, epsilon)), + (&lut_block.1, (lut_block.0, epsilon)), O::Distance::block_accessor(), ) }, |lowerbound, mean, first| { - results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(first))); + results.push((Reverse(lowerbound), AlwaysEqual(first), AlwaysEqual(mean))); }, |_| (), ); @@ -241,8 +249,9 @@ pub fn search_and_estimate( let index = index.clone(); let vector = vector.as_borrowed(); std::iter::from_fn(move || { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); + while let Some((_, AlwaysEqual(first), AlwaysEqual(mean))) = + pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) + { if is_residual { let (dis_u, residual_u) = vectors::read_for_h1_tuple::( index.clone(), @@ -296,23 +305,26 @@ pub fn search_and_estimate( let mut results = LinkedVec::new(); for (first, residual) in state { - let lut = if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { - &O::Vector::compute_lut(residual) - } else { - default_lut.as_ref().unwrap() - }; + let (lut_block, lut_binary) = + if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { + &O::Vector::compute_lut(residual) + } else if let Some(lut) = default_lut.as_ref() { + lut + } else { + unreachable!() + }; let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); let mut callback = |lowerbound, mean, payload| { - results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(payload))); + results.push((Reverse(lowerbound), AlwaysEqual(payload), AlwaysEqual(mean))); }; tape::read_frozen_tape( index.clone(), jump_tuple.frozen_first(), || { RAccess::new( - (&lut.0.4, (lut.0.0, lut.0.1, lut.0.2, lut.0.3, epsilon)), + (&lut_block.1, (lut_block.0, epsilon)), O::Distance::block_accessor(), ) }, @@ -322,7 +334,7 @@ pub fn search_and_estimate( tape::read_appendable_tape( index.clone(), jump_tuple.appendable_first(), - |code| O::Distance::compute_lowerbound_binary(&lut.1, code, epsilon), + |code| O::Distance::compute_lowerbound_binary(lut_binary, code, epsilon), &mut callback, |_| (), ); @@ -346,3 +358,16 @@ pub fn search_and_estimate( } (results.into_vec(), Distance::from_f32(estimation)) } + +fn pop_if( + heap: &mut BinaryHeap, + mut predicate: impl FnMut(&mut T) -> bool, +) -> Option { + use std::collections::binary_heap::PeekMut; + let mut peek = heap.peek_mut()?; + if predicate(&mut peek) { + Some(PeekMut::pop(peek)) + } else { + None + } +} diff --git a/crates/algorithm/src/select_heap.rs b/crates/algorithm/src/select_heap.rs index a1a5a224..803cb495 100644 --- a/crates/algorithm/src/select_heap.rs +++ b/crates/algorithm/src/select_heap.rs @@ -1,51 +1,53 @@ use std::collections::BinaryHeap; use std::num::NonZero; +pub struct SortHeap { + inner: Vec, + t: NonZero, +} + +impl SortHeap { + pub fn peek(&self) -> Option<&T> { + self.inner.last() + } +} + pub enum SelectHeap { - Sorted { x: Vec, t: NonZero }, - Heap { x: BinaryHeap }, + Sorted(SortHeap), + Binary(BinaryHeap), } impl SelectHeap { - pub fn from_vec(mut vec: Vec) -> Self { + pub fn from_vec(vec: Vec) -> Self { let n = vec.len(); if let Some(t) = NonZero::new(n / 384) { + let mut inner = vec; let index = n - t.get(); - turboselect::select_nth_unstable(&mut vec, index); - vec[index..].sort_unstable(); - Self::Sorted { x: vec, t } + turboselect::select_nth_unstable(&mut inner, index); + inner[index..].sort_unstable(); + Self::Sorted(SortHeap { inner, t }) } else { - Self::Heap { - x: BinaryHeap::from(vec), - } - } - } - pub fn is_empty(&self) -> bool { - match self { - SelectHeap::Sorted { x, .. } => x.is_empty(), - SelectHeap::Heap { x } => x.is_empty(), + Self::Binary(BinaryHeap::from(vec)) } } pub fn pop(&mut self) -> Option { match self { - SelectHeap::Sorted { x: inner, t } => { - let x = inner.pop().expect("inconsistent internal structure"); + SelectHeap::Sorted(SortHeap { inner, t }) => { + let Some(k) = inner.pop() else { unreachable!() }; if let Some(value) = NonZero::new(t.get() - 1) { *t = value; } else { - *self = SelectHeap::Heap { - x: std::mem::take(inner).into(), - }; + *self = SelectHeap::Binary(std::mem::take(inner).into()); } - Some(x) + Some(k) } - SelectHeap::Heap { x } => x.pop(), + SelectHeap::Binary(x) => x.pop(), } } pub fn peek(&self) -> Option<&T> { match self { - SelectHeap::Sorted { x, .. } => x.last(), - SelectHeap::Heap { x } => x.peek(), + SelectHeap::Sorted(x) => x.peek(), + SelectHeap::Binary(x) => x.peek(), } } } diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index 98b34ffd..e9e09a0a 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -38,7 +38,7 @@ where } pub fn tape_move(&mut self) { if self.head.len() == 0 { - panic!("tuple is too large to fit in a fresh page"); + panic!("implementation: a clear page cannot accommodate a single tuple"); } let next = (self.extend)(); self.head.get_opaque_mut().next = next.id(); @@ -64,7 +64,7 @@ where if let Some(i) = self.head.alloc(&bytes) { pair_to_pointer((self.head.id(), i)) } else { - panic!("tuple is too large to fit in a fresh page") + panic!("implementation: a free page cannot accommodate a single tuple") } } } @@ -73,7 +73,7 @@ where if let Some(i) = self.head.alloc(&bytes) { pair_to_pointer((self.head.id(), i)) } else { - panic!("tuple is too large to fit in this page") + panic!("implementation: a free page cannot accommodate a single tuple") } } } @@ -214,7 +214,7 @@ pub fn append( past.get_opaque_mut().skip = fresh.max(past.get_opaque().skip); return pair_to_pointer((fresh, i)); } else { - panic!("a tuple cannot even be fit in a fresh page"); + panic!("implementation: a clear page cannot accommodate a single tuple"); } } if current == first && write.get_opaque().skip != first { diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 31b72234..ea00abe3 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -77,15 +77,15 @@ impl WithReader for MetaTuple { type Reader<'a> = MetaTupleReader<'a>; fn deserialize_ref(source: &[u8]) -> MetaTupleReader<'_> { if source.len() < 16 { - panic!("bad bytes") + panic!("deserialization: bad bytes") } let magic = u64::from_ne_bytes(std::array::from_fn(|i| source[i + 0])); if magic != MAGIC { - panic!("bad magic number"); + panic!("deserialization: bad magic number"); } let version = u64::from_ne_bytes(std::array::from_fn(|i| source[i + 8])); if version != VERSION { - panic!("bad version number"); + panic!("deserialization: bad version number"); } let checker = RefChecker::new(source); let header = checker.prefix(0); @@ -181,11 +181,11 @@ impl FreepageTupleWriter<'_> { } let i_1 = self.header.level_1[i_3 << 5 | i_2 << 0].trailing_zeros() as usize; if i_1 == 32 { - panic!("data corruption"); + panic!("deserialization: bad bytes"); } let i_0 = self.header.level_0[i_3 << 10 | i_2 << 5 | i_1 << 0].trailing_zeros() as usize; if i_0 == 32 { - panic!("data corruption"); + panic!("deserialization: bad bytes"); } Some(i_3 << 15 | i_2 << 10 | i_1 << 5 | i_0 << 0) } @@ -320,7 +320,7 @@ impl WithReader for VectorTuple { let elements = checker.bytes(header.elements_s, header.elements_e); VectorTupleReader::_1(VectorTupleReader1 { header, elements }) } - _ => panic!("bad bytes"), + _ => panic!("deserialization: bad bytes"), } } } @@ -503,7 +503,7 @@ impl WithReader for H1Tuple { let elements = checker.bytes(header.elements_s, header.elements_e); H1TupleReader::_1(H1TupleReader1 { header, elements }) } - _ => panic!("bad bytes"), + _ => panic!("deserialization: bad bytes"), } } } @@ -760,7 +760,7 @@ impl WithReader for FrozenTuple { let elements = checker.bytes(header.elements_s, header.elements_e); FrozenTupleReader::_1(FrozenTupleReader1 { header, elements }) } - _ => panic!("bad bytes"), + _ => panic!("deserialization: bad bytes"), } } } @@ -783,7 +783,7 @@ impl WithWriter for FrozenTuple { let elements = checker.bytes(header.elements_s, header.elements_e); FrozenTupleWriter::_1(FrozenTupleWriter1 { header, elements }) } - _ => panic!("bad bytes"), + _ => panic!("deserialization: bad bytes"), } } } @@ -1055,7 +1055,7 @@ impl<'a> RefChecker<'a> { let start = s; let end = s + size_of::(); let bytes = &self.bytes[start..end]; - FromBytes::ref_from_bytes(bytes).expect("bad bytes") + FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") } pub fn bytes( &self, @@ -1065,7 +1065,7 @@ impl<'a> RefChecker<'a> { let start = s; let end = e; let bytes = &self.bytes[start..end]; - FromBytes::ref_from_bytes(bytes).expect("bad bytes") + FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") } } @@ -1090,10 +1090,10 @@ impl<'a> MutChecker<'a> { let start = s; let end = s + size_of::(); if !(start <= end && end <= self.bytes.len()) { - panic!("bad bytes"); + panic!("deserialization: bad bytes"); } if !(self.flag <= start) { - panic!("bad bytes"); + panic!("deserialization: bad bytes"); } else { self.flag = end; } @@ -1101,7 +1101,7 @@ impl<'a> MutChecker<'a> { let bytes = unsafe { std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) }; - FromBytes::mut_from_bytes(bytes).expect("bad bytes") + FromBytes::mut_from_bytes(bytes).expect("deserialization: bad bytes") } pub fn bytes( &mut self, @@ -1111,10 +1111,10 @@ impl<'a> MutChecker<'a> { let start = s; let end = e; if !(start <= end && end <= self.bytes.len()) { - panic!("bad bytes"); + panic!("deserialization: bad bytes"); } if !(self.flag <= start) { - panic!("bad bytes"); + panic!("deserialization: bad bytes"); } else { self.flag = end; } @@ -1122,7 +1122,7 @@ impl<'a> MutChecker<'a> { let bytes = unsafe { std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) }; - FromBytes::mut_from_bytes(bytes).expect("bad bytes") + FromBytes::mut_from_bytes(bytes).expect("deserialization: bad bytes") } } diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index d1319a62..99b78918 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -62,7 +62,9 @@ pub fn append( ) -> IndexPointer { fn append(index: impl RelationWrite, first: u32, bytes: &[u8]) -> IndexPointer { if let Some(mut write) = index.search(bytes.len()) { - let i = write.alloc(bytes).unwrap(); + let i = write + .alloc(bytes) + .expect("implementation: a free page cannot accommodate a single tuple"); return pair_to_pointer((write.id(), i)); } tape::append(index, first, bytes, true) @@ -84,5 +86,5 @@ pub fn append( }); chain = Err(append(index.clone(), vectors_first, &bytes)); } - chain.err().unwrap() + chain.expect_err("internal error: 0-dimensional vector") } diff --git a/crates/rabitq/src/binary.rs b/crates/rabitq/src/binary.rs index 267ecaee..fa1d01aa 100644 --- a/crates/rabitq/src/binary.rs +++ b/crates/rabitq/src/binary.rs @@ -1,7 +1,10 @@ use distance::Distance; use simd::Floating; -pub type BinaryLut = (f32, f32, f32, f32, (Vec, Vec, Vec, Vec)); +pub type BinaryLut = ( + (f32, f32, f32, f32), + (Vec, Vec, Vec, Vec), +); pub type BinaryCode<'a> = (f32, f32, f32, f32, &'a [u64]); pub fn preprocess(vector: &[f32]) -> BinaryLut { @@ -12,7 +15,7 @@ pub fn preprocess(vector: &[f32]) -> BinaryLut { } else { simd::u8::reduce_sum_of_x(&qvector) as f32 }; - (dis_v_2, b, k, qvector_sum, binarize(&qvector)) + ((dis_v_2, b, k, qvector_sum), binarize(&qvector)) } pub fn process_lowerbound_l2( @@ -20,7 +23,7 @@ pub fn process_lowerbound_l2( (dis_u_2, factor_ppc, factor_ip, factor_err, t): BinaryCode<'_>, epsilon: f32, ) -> Distance { - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let &((dis_v_2, b, k, qvector_sum), ref s) = lut; let value = asymmetric_binary_dot_product(t, s) as u16; let rough = dis_u_2 + dis_v_2 + b * factor_ppc + ((2.0 * value as f32) - qvector_sum) * factor_ip * k; @@ -33,7 +36,7 @@ pub fn process_lowerbound_dot( (_, factor_ppc, factor_ip, factor_err, t): BinaryCode<'_>, epsilon: f32, ) -> Distance { - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let &((dis_v_2, b, k, qvector_sum), ref s) = lut; let value = asymmetric_binary_dot_product(t, s) as u16; let rough = 0.5 * b * factor_ppc + 0.5 * ((2.0 * value as f32) - qvector_sum) * factor_ip * k; let err = 0.5 * factor_err * dis_v_2.sqrt(); diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs index 174c8fe6..0c5aaaf1 100644 --- a/crates/rabitq/src/block.rs +++ b/crates/rabitq/src/block.rs @@ -1,7 +1,7 @@ use distance::Distance; use simd::Floating; -pub type BlockLut = (f32, f32, f32, f32, Vec<[u8; 16]>); +pub type BlockLut = ((f32, f32, f32, f32), Vec<[u8; 16]>); pub type BlockCode<'a> = ( &'a [f32; 32], &'a [f32; 32], @@ -18,7 +18,7 @@ pub fn preprocess(vector: &[f32]) -> BlockLut { } else { simd::u8::reduce_sum_of_x(&qvector) as f32 }; - (dis_v_2, b, k, qvector_sum, compress(qvector)) + ((dis_v_2, b, k, qvector_sum), compress(qvector)) } pub fn process_lowerbound_l2( @@ -26,7 +26,7 @@ pub fn process_lowerbound_l2( (dis_u_2, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, epsilon: f32, ) -> [Distance; 32] { - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let &((dis_v_2, b, k, qvector_sum), ref s) = lut; let r = simd::fast_scan::scan(t, s); std::array::from_fn(|i| { let rough = dis_u_2[i] @@ -43,7 +43,7 @@ pub fn process_lowerbound_dot( (_, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, epsilon: f32, ) -> [Distance; 32] { - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; + let &((dis_v_2, b, k, qvector_sum), ref s) = lut; let r = simd::fast_scan::scan(t, s); std::array::from_fn(|i| { let rough = diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index 0215d252..a59c6c2a 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -86,7 +86,7 @@ pub fn compute_lut(vector: &[f32]) -> (BlockLut, BinaryLut) { let binary = binary::binarize(&qvector); let block = block::compress(qvector); ( - (dis_v_2, b, k, qvector_sum, block), - (dis_v_2, b, k, qvector_sum, binary), + ((dis_v_2, b, k, qvector_sum), block), + ((dis_v_2, b, k, qvector_sum), binary), ) } diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index bbd5f7aa..4041be2c 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -14,7 +14,7 @@ pub fn prewarm( height: i32, check: impl Fn(), ) -> String { - match (opfamily.vector_kind(), opfamily.distance_kind()) { + let message = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { algorithm::prewarm::, L2>>(index, height, check) } @@ -27,6 +27,10 @@ pub fn prewarm( (VectorKind::Vecf16, DistanceKind::Dot) => { algorithm::prewarm::, Dot>>(index, height, check) } + }; + match message { + Ok(message) => message, + Err(e) => panic!("{e}"), } } diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index d8b9ad7b..949ba0eb 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -1,6 +1,6 @@ pub mod am_build; -use crate::index::gucs::{epsilon, max_maxsim_tuples, max_scan_tuples, maxsim_threshold, probes}; +use crate::index::gucs; use crate::index::opclass::opfamily; use crate::index::scanners::*; use crate::index::storage::PostgresRelation; @@ -295,11 +295,12 @@ pub unsafe extern "C" fn amrescan( let opfamily = opfamily((*scan).indexRelation); let relation = PostgresRelation::new((*scan).indexRelation); let options = SearchOptions { - epsilon: epsilon(), - probes: probes(), - max_scan_tuples: max_scan_tuples(), - max_maxsim_tuples: max_maxsim_tuples(), - maxsim_threshold: maxsim_threshold(), + epsilon: gucs::epsilon(), + allows_skipping_rerank: gucs::allows_skipping_rerank(), + probes: gucs::probes(), + max_scan_tuples: gucs::max_scan_tuples(), + max_maxsim_tuples: gucs::max_maxsim_tuples(), + maxsim_threshold: gucs::maxsim_threshold(), }; let scanner = &mut *(*scan).opaque.cast::(); let fetcher = { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 97d71a6a..b565225a 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -9,6 +9,7 @@ static PREWARM_DIM: GucSetting> = static MAX_MAXSIM_TUPLES: GucSetting = GucSetting::::new(-1); static MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); static PRERERANK_FILTERING: GucSetting = GucSetting::::new(false); +static ALLOWS_SKIPPING_RERANK: GucSetting = GucSetting::::new(false); pub fn init() { GucRegistry::define_string_guc( @@ -75,6 +76,14 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_bool_guc( + "vchordrq.allows_skipping_rerank", + "`allows_skipping_rerank` argument of vchordrq.", + "`allows_skipping_rerank` argument of vchordrq.", + &ALLOWS_SKIPPING_RERANK, + GucContext::Userset, + GucFlags::default(), + ); unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); @@ -154,3 +163,7 @@ pub fn prewarm_dim() -> Vec { pub fn prererank_filtering() -> bool { PRERERANK_FILTERING.get() } + +pub fn allows_skipping_rerank() -> bool { + ALLOWS_SKIPPING_RERANK.get() +} diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs index 9660d441..3785db9d 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/scanners/maxsim.rs @@ -96,13 +96,19 @@ impl SearchBuilder for MaxsimBuilder { estimations.push(0.0); continue; } - let returning = algorithm::rerank_index::, Dot>>( - relation.clone(), - vector.clone(), - results, - ) - .take(max_maxsim_tuples as usize) - .collect::>(); + let returning = if options.epsilon == 0.0 && options.allows_skipping_rerank { + algorithm::skip(results) + .take(max_maxsim_tuples as usize) + .collect::>() + } else { + algorithm::rerank_index::, Dot>>( + relation.clone(), + vector.clone(), + results, + ) + .take(max_maxsim_tuples as usize) + .collect::>() + }; let mut max = f32::NEG_INFINITY; for (distance, payload) in returning { max = max.max(distance.to_f32()); @@ -147,13 +153,19 @@ impl SearchBuilder for MaxsimBuilder { estimations.push(0.0); continue; } - let returning = algorithm::rerank_index::, Dot>>( - relation.clone(), - vector.clone(), - results, - ) - .take(max_maxsim_tuples as usize) - .collect::>(); + let returning = if options.epsilon == 0.0 && options.allows_skipping_rerank { + algorithm::skip(results) + .take(max_maxsim_tuples as usize) + .collect::>() + } else { + algorithm::rerank_index::, Dot>>( + relation.clone(), + vector.clone(), + results, + ) + .take(max_maxsim_tuples as usize) + .collect::>() + }; let mut max = f32::NEG_INFINITY; for (distance, payload) in returning { max = max.max(distance.to_f32()); diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs index bad06d54..e572bf0d 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/scanners/mod.rs @@ -12,6 +12,7 @@ pub use maxsim::MaxsimBuilder; #[derive(Debug)] pub struct SearchOptions { pub epsilon: f32, + pub allows_skipping_rerank: bool, pub probes: Vec, pub max_scan_tuples: Option, pub max_maxsim_tuples: Option, From cade3c6693cb441cfbdd666ab7bb8e3a2d10bdc6 Mon Sep 17 00:00:00 2001 From: xieydd Date: Tue, 8 Apr 2025 16:34:01 +0800 Subject: [PATCH 139/324] delete cnpg image (#224) vchord-cnpg and pg-slim move to https://github.com/tensorchord/VectorChord-images. Signed-off-by: xieydd --- .github/workflows/release.yml | 84 +--- .github/workflows/release_pg_slim.yml | 61 --- docker/pg-cnpg/Dockerfile | 193 -------- docker/pg-cnpg/Dockerfile.arm | 72 --- docker/pg-cnpg/extension-install.sh | 80 ---- docker/pg-cnpg/requirements.txt | 637 ------------------------- docker/pg-slim/Dockerfile | 142 ------ docker/pg-slim/docker-ensure-initdb.sh | 71 --- docker/pg-slim/docker-entrypoint.sh | 356 -------------- 9 files changed, 1 insertion(+), 1695 deletions(-) delete mode 100644 .github/workflows/release_pg_slim.yml delete mode 100644 docker/pg-cnpg/Dockerfile delete mode 100644 docker/pg-cnpg/Dockerfile.arm delete mode 100755 docker/pg-cnpg/extension-install.sh delete mode 100644 docker/pg-cnpg/requirements.txt delete mode 100644 docker/pg-slim/Dockerfile delete mode 100755 docker/pg-slim/docker-ensure-initdb.sh delete mode 100755 docker/pg-slim/docker-entrypoint.sh diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index fd36e4ab..8c8a06a6 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -183,26 +183,6 @@ jobs: PG_VERSION=${{ matrix.version }} SEMVER=${{ env.SEMVER }} PGVECTOR=0.8.0 - - name: Login to modelzai Docker Hub - uses: docker/login-action@v3 - with: - username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} - password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} - - name: Build and push Enterprise image to Docker Registry - uses: docker/build-push-action@v6 - with: - context: ./docker/pg-cnpg - push: true - platforms: ${{ matrix.runner == 'ubuntu-22.04' && 'linux/amd64' || 'linux/arm64' }} - file: ${{ matrix.runner == 'ubuntu-22.04' && './docker/pg-cnpg/Dockerfile' || './docker/pg-cnpg/Dockerfile.arm' }} - provenance: false - build-args: | - PG_MAJOR=${{ matrix.version }} - SEMVER=${{ env.SEMVER }} - LIB_DIR=${{ matrix.runner == 'ubuntu-22.04' && '/usr/lib/x86_64-linux-gnu' || '/usr/lib/aarch64-linux-gnu' }} - TARGETARCH=${{ env.PLATFORM }} - PGVECTOR=0.8.0 - tags: modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} create-manifests: runs-on: ubuntu-latest @@ -244,66 +224,4 @@ jobs: ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} \ --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 - docker manifest push ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} - - - create-manifests-modelzai: - runs-on: ubuntu-latest - needs: ["semver", "build", "docker"] - strategy: - matrix: - version: ["14", "15", "16", "17"] - env: - SEMVER: ${{ needs.semver.outputs.SEMVER }} - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Set up QEMU - uses: docker/setup-qemu-action@v3 - - name: Login to modelzai Docker Hub - uses: docker/login-action@v2 - with: - username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} - password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} - - name: Create manifest and push - run: | - docker manifest create \ - modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }} \ - --amend modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ - --amend modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }}-arm64 - docker manifest push modelzai/vchord-cnpg:${{ matrix.version }}-v${{ env.SEMVER }} - - test: - name: Run tests - runs-on: ${{ matrix.runner }} - needs: ["semver", "build", "docker", "create-manifests"] - strategy: - matrix: - version: [14, 15, 16, 17] - runner: ["ubuntu-22.04"] - container: - image: modelzai/vchord-cnpg:${{ matrix.version }}-v${{ needs.semver.outputs.SEMVER }} - options: --user root - credentials: - username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} - password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} - env: - PGHOST: "localhost" - PGPORT: "5432" - PGDATABASE: "postgres" - PGUSER: "postgres" - PGPASSWORD: "postgres" - POSTGRES_PASSWORD: "password" - PGDATA: "/var/lib/postgresql/data2" - - steps: - - name: Install all extensions in registry - # Entrypoint is overwritten by GitHub Action. We need to execute it manually in order to start Postgres. - # More information here https://github.com/actions/runner/issues/1964 - run: | - bash /usr/local/bin/docker-entrypoint.sh postgres & - sleep 5 - curl https://registry.pgtrunk.io/extensions/all | jq -r ".[] | .name" > /tmp/extensions.txt - extension-install.sh | tee /tmp/output.txt - cat /tmp/output.txt - + docker manifest push ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} \ No newline at end of file diff --git a/.github/workflows/release_pg_slim.yml b/.github/workflows/release_pg_slim.yml deleted file mode 100644 index 74ae65f3..00000000 --- a/.github/workflows/release_pg_slim.yml +++ /dev/null @@ -1,61 +0,0 @@ -name: Release for Postgres slim - -on: - workflow_dispatch: - -jobs: - pg-slim: - strategy: - matrix: - version: [14, 15, 16, 17] - runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] - runs-on: ${{ matrix.runner }} - env: - PG_MAJOR: ${{ matrix.version }} - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Set up QEMU - uses: docker/setup-qemu-action@v3 - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 - - name: Login to Docker Hub - uses: docker/login-action@v2 - with: - username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} - password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} - - name: Push Postgres Slim to Docker Registry - uses: docker/build-push-action@v4 - with: - context: ./docker/pg-slim - push: true - platforms: ${{ matrix.runner == 'ubuntu-22.04' && 'linux/amd64' || 'linux/arm64' }} - file: ./docker/pg-slim/Dockerfile - provenance: false - build-args: | - PG_MAJOR=${{ matrix.version }} - tags: modelzai/pg-slim:${{ matrix.version }}-${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} - create-manifests: - needs: ["pg-slim"] - strategy: - matrix: - version: ["14", "15", "16", "17"] - runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] - runs-on: ${{ matrix.runner }} - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Set up QEMU - uses: docker/setup-qemu-action@v3 - - name: Login to modelzai Docker Hub - uses: docker/login-action@v2 - with: - username: ${{ secrets.DOCKERIO_MODELZ_USERNAME }} - password: ${{ secrets.DOCKERIO_MODELZ_TOKEN }} - - name: Create manifest and push - run: | - docker manifest create \ - modelzai/pg-slim:${{ matrix.version }} \ - --amend modelzai/pg-slim:${{ matrix.version }}-amd64 \ - --amend modelzai/pg-slim:${{ matrix.version }}-arm64 - docker manifest push modelzai/pg-slim:${{ matrix.version }} \ No newline at end of file diff --git a/docker/pg-cnpg/Dockerfile b/docker/pg-cnpg/Dockerfile deleted file mode 100644 index 8ac07f2a..00000000 --- a/docker/pg-cnpg/Dockerfile +++ /dev/null @@ -1,193 +0,0 @@ -ARG PG_MAJOR=16 -ARG SEMVER -ARG TARGETARCH=amd64 - -FROM tensorchord/vchord-binary:pg${PG_MAJOR}-v${SEMVER}-${TARGETARCH} AS binary - -FROM rust:1.83-bookworm AS builder -ARG TRUNK_VER=0.15.7 -ENV CARGO_REGISTRIES_CRATES_IO_PROTOCOL=sparse -RUN cargo install --version $TRUNK_VER pg-trunk - -FROM modelzai/pg-slim:${PG_MAJOR}-${TARGETARCH} -ARG PG_MAJOR -ARG SEMVER -ARG TARGETARCH -ARG LIB_DIR -ARG PGVECTOR -ARG ALTDIR=/var/lib/postgresql/data/tensorchord - -USER root - -COPY --from=binary /workspace/postgresql-${PG_MAJOR}-vchord_${SEMVER}-1_${TARGETARCH}.deb /tmp/vchord.deb -RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb - -# PGDATA is set in pg-slim and used by dependents on this image. -RUN if [ -z "${PGDATA}" ]; then echo "PGDATA is not set"; exit 1; fi - -# Install trunk -COPY --from=builder /usr/local/cargo/bin/trunk /usr/bin/trunk -COPY requirements.txt . - -# Install barman-cloud -RUN set -xe; \ - apt-get update; \ - apt-get install -y --no-install-recommends \ - python3-pip \ - python3-psycopg2 \ - python3-setuptools \ - ; \ - pip3 install --upgrade pip; \ - # TODO: Remove --no-deps once https://github.com/pypa/pip/issues/9644 is solved - pip3 install --no-deps -r requirements.txt; \ - apt-get autoremove -y; \ - apt-get clean; \ - rm -rf /var/lib/apt/lists/*; - -RUN chown -R postgres:postgres ${ALTDIR}/${PG_MAJOR} && \ - chmod -R 0700 ${ALTDIR}/${PG_MAJOR} -RUN chown postgres /usr/share/postgresql/${PG_MAJOR}/extension - -RUN apt-get update && apt-get install -y \ - jq \ - curl \ - wget \ - && rm -rf /var/lib/apt/lists/* - -# Install extension dependencies -RUN apt-get update && apt-get install -y \ - libmysqlclient-dev \ - libtcl8.6 \ - libgeos-dev \ - libproj-dev \ - libjson-c-dev \ - libprotobuf-c-dev \ - libxml2-dev \ - libboost-serialization1.74-dev \ - libhiredis-dev \ - libsybdb5 \ - libpython3.10-dev \ - r-base-core \ - openssl \ - liblz4-1 \ - libpcre2-8-0 \ - libuuid1 \ - libgroonga0 \ - libopenblas0-pthread \ - libcurl4 \ - libjson-c5 \ - libsodium23 \ - libgcc-s1 \ - libselinux1 \ - librdkafka1 \ - libgdal30 \ - libcrypt1 \ - liburiparser1 \ - libfreetype6 \ - libzstd1 \ - zlib1g \ - libperl5.34 \ - libgomp1 \ - libssl3 \ - libsfcgal1 \ - openjdk-11-jdk \ - libaio1 \ - libbson-dev \ - libgsl-dev \ - && rm -rf /var/lib/apt/lists/* -RUN ln -s /usr/lib/jvm/java-11-openjdk-amd64/lib/server/libjvm.so ${LIB_DIR}/libjvm.so -RUN wget https://download.oracle.com/otn_software/linux/instantclient/1920000/instantclient-basiclite-linux.x64-19.20.0.0.0dbru.zip && \ - unzip instantclient-basiclite-linux.x64-19.20.0.0.0dbru.zip && \ - cp instantclient_19_20/libclntsh.so.19.1 ${LIB_DIR}/ && \ - cp instantclient_19_20/libnnz19.so ${LIB_DIR}/ && \ - cp instantclient_19_20/libclntshcore.so.19.1 ${LIB_DIR}/ && \ - rm -rf instantclient_19_20 && \ - rm instantclient-basiclite-linux.x64-19.20.0.0.0dbru.zip - -# Install zhparser dependency -RUN wget http://www.xunsearch.com/scws/down/scws-1.2.3.tar.bz2 && \ - tar xvf scws-1.2.3.tar.bz2 && \ - cd scws-1.2.3 && \ - ./configure && \ - make install && \ - cd .. && \ - rm -rf scws-1.2.3.tar.bz2 scws-1.2.3 && \ - ln -s /usr/local/lib/libscws.so ${LIB_DIR}/libscws.so - -# Install duckdb libs -RUN wget https://github.com/duckdb/duckdb/releases/download/v0.8.1/libduckdb-linux-amd64.zip && \ - unzip libduckdb-linux-amd64.zip && \ - cp libduckdb.so ${LIB_DIR}/ && \ - rm -rf libduckdb-linux-amd64.zip libduckdb.so - -# Install pg_stat_statements -RUN trunk install pg_stat_statements - -# Install auto_explain -RUN trunk install auto_explain - -# Install plpython3u -RUN trunk install plpython3u - -# Install pgvector -RUN trunk install pgvector --version ${PGVECTOR} - -# Install pgmq -RUN trunk install pgmq --version 1.4.5 - -# Install pg_later -RUN trunk install pg_later --version 0.3.0 - -# Clone and build AWS SDK for C++ -RUN git clone https://github.com/aws/aws-sdk-cpp.git && \ - cd aws-sdk-cpp && \ - git checkout 1.9.263 && \ - git submodule update --init --recursive && \ - mkdir build && cd build && \ - cmake -DBUILD_ONLY="s3;core;config;sts;cognito-identity;transfer;identity-management" -DAUTORUN_UNIT_TESTS=OFF -DCMAKE_CXX_FLAGS=-Wno-error=deprecated-declarations .. && \ - make -j$(nproc) && \ - make install && \ - cd ../../../ && rm -rf aws-sdk-cpp - -# Clone and build Apache Arrow -RUN git clone https://github.com/apache/arrow.git && \ - cd arrow && \ - git checkout apache-arrow-7.0.1 && \ - cd cpp && \ - mkdir build && cd build && \ - cmake -DARROW_PARQUET=ON -DARROW_S3=ON -DARROW_WITH_SNAPPY=ON .. && \ - make -j$(nproc) && \ - make install && \ - cd ../../../ && rm -rf arrow - -# Clone and build pgaudit -RUN git clone https://github.com/pgaudit/pgaudit.git && \ - cd pgaudit && \ - git checkout REL_${PG_MAJOR}_STABLE && \ - make install USE_PGXS=1 PG_CONFIG=/usr/lib/postgresql/${PG_MAJOR}/bin/pg_config && \ - cd ../ && rm -rf pgaudit - -# Clone and build pg_failover_slots -RUN git clone https://github.com/EnterpriseDB/pg_failover_slots.git && \ - cd pg_failover_slots && \ - make install PG_CONFIG=/usr/lib/postgresql/${PG_MAJOR}/bin/pg_config && \ - cd ../ && rm -rf pg_failover_slots - -# Test trunk install extension -COPY extension-install.sh /usr/local/bin/ - -# Change the uid of postgres to 26 -RUN usermod -u 26 postgres -RUN chown -R postgres:postgres ${ALTDIR} -RUN cp /usr/share/postgresql/${PG_MAJOR}/extension/* ${ALTDIR}/extension/ -RUN cp /usr/lib/postgresql/${PG_MAJOR}/lib/* ${ALTDIR}/${PG_MAJOR}/lib/ - -RUN set -eux; \ - mkdir /tmp/pg_pkglibdir; \ - mkdir /tmp/pg_sharedir; \ - cp -r $(pg_config --pkglibdir)/* /tmp/pg_pkglibdir; \ - cp -r $(pg_config --sharedir)/* /tmp/pg_sharedir - -RUN chown -R postgres:postgres /tmp -USER 26 -ENV PATH $PATH:/usr/lib/postgresql/${PG_MAJOR}/bin \ No newline at end of file diff --git a/docker/pg-cnpg/Dockerfile.arm b/docker/pg-cnpg/Dockerfile.arm deleted file mode 100644 index 6529b13d..00000000 --- a/docker/pg-cnpg/Dockerfile.arm +++ /dev/null @@ -1,72 +0,0 @@ -ARG PG_MAJOR -ARG SEMVER -ARG TARGETARCH - -FROM tensorchord/vchord-binary:pg${PG_MAJOR}-v${SEMVER}-${TARGETARCH} as binary - -# From https://github.com/cloudnative-pg/postgres-containers/blob/main/Debian/17/bookworm/Dockerfile -FROM postgres:${PG_MAJOR}-bookworm -ARG PG_MAJOR -ARG SEMVER -ARG TARGETARCH -ARG PGVECTOR - -COPY requirements.txt / - -# Install additional extensions -RUN set -xe; \ - apt-get update; \ - apt-get install -y --no-install-recommends \ - "postgresql-${PG_MAJOR}-pgaudit" \ - "postgresql-${PG_MAJOR}-pg-failover-slots" \ - ; \ - rm -fr /tmp/* ; \ - rm -rf /var/lib/apt/lists/*; - -# Install barman-cloud -RUN set -xe; \ - apt-get update; \ - apt-get install -y --no-install-recommends \ - python3-dev \ - python3-pip \ - python3-psycopg2 \ - python3-setuptools \ - build-essential \ - ; \ - pip3 install --break-system-packages --upgrade pip; \ - # TODO: Remove --no-deps once https://github.com/pypa/pip/issues/9644 is solved - pip3 install --break-system-packages --no-deps -r requirements.txt; \ - rm -rf /var/lib/apt/lists/*; - -COPY --from=binary /workspace/postgresql-${PG_MAJOR}-vchord_${SEMVER}-1_${TARGETARCH}.deb /tmp/vchord.deb -RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb - -RUN apt-get update && apt-get install -y \ - jq \ - curl \ - wget \ - sudo \ - && rm -rf /var/lib/apt/lists/* - -# Install pig -RUN curl -fsSL https://repo.pigsty.io/pig | bash && \ - pig repo add pigsty pgdg -u - -# Install pgvector -RUN pig ext install -y pgvector=${PGVECTOR} - -# Install pg_stat_statements -RUN pig ext install -y pg_stat_statements - -# Install auto_explain -RUN pig ext install -y auto_explain - -# Install plpython3u -RUN pig ext install -y plpython3u - -# Install pg_later -RUN pig ext install -y pg_later=0.3.0 - -# Change the uid of postgres to 26 -RUN usermod -u 26 postgres -USER 26 diff --git a/docker/pg-cnpg/extension-install.sh b/docker/pg-cnpg/extension-install.sh deleted file mode 100755 index e76f5f5c..00000000 --- a/docker/pg-cnpg/extension-install.sh +++ /dev/null @@ -1,80 +0,0 @@ -#!/bin/bash -trunk_install_failed_extensions=() -need_load_shared_preload_libraries_extensions=() -version_not_found_extensions=() -file='/tmp/extensions.txt' -extension_count=$(<$file wc -l) -lines=$(cat $file) -for line in $lines -do - output=$(trunk install $line 2>&1) - - if [ $? -ne 0 ]; then - if [[ $output == *"Failed to find an archive for"* || $output == *"Failed to fetch Trunk archive from"* ]]; then - version_not_found_extensions+=("$line") - else - echo "trunk install command failed" - trunk_install_failed_extensions+=("$line") - fi - fi - echo $output - printf "\n\n" -done -IFS=$'\n' extensions=(`psql postgres://postgres:postgres@localhost:5432 -tA postgres -c 'select name from pg_available_extensions;'`) -for ext in "${extensions[@]}" -do - # drop schema columnar if ext name is columnar - if [ "$ext" == "columnar" ]; then - psql postgres://postgres:postgres@localhost:5432 -c "drop extension if exists citus_columnar cascade;" - fi - # drop type semver if ext name is semver - if [ "$ext" == "semver" ]; then - psql postgres://postgres:postgres@localhost:5432 -c "drop extension if exists pg_text_semver cascade;" - fi - # if extension name is meta_triggers, create extension meta first and create extension meta_triggers - if [ "$ext" == "meta_triggers" ]; then - psql postgres://postgres:postgres@localhost:5432 -c "create extension if not exists hstore cascade;" - psql postgres://postgres:postgres@localhost:5432 -c "create extension if not exists meta cascade;" - psql postgres://postgres:postgres@localhost:5432 -c "create extension if not exists meta_triggers cascade;" - fi - output=$(psql postgres://postgres:postgres@localhost:5432 -c "create extension if not exists \"$ext\" cascade;" 2>&1) - if [ $? -ne 0 ]; then - if [[ $output == *"shared_preload_libraries"* ]]; then - need_load_shared_preload_libraries_extensions+=("$ext") - elif [[ $output == *"already exists"* ]]; then - echo "extension \"$ext\" already exists" - else - echo "CREATE EXTENSION command failed" - failed_extensions+=("$ext") - fi - fi - echo $output - printf "\n\n" -done -available_extensions_count=${#extensions[@]} -failure_count=${#failed_extensions[@]} -need_load_shared_preload_libraries_count=${#need_load_shared_preload_libraries_extensions[@]} -not_found_count=${#version_not_found_extensions[@]} -success=$(($available_extensions_count-$failure_count)) -success_percent=$(awk "BEGIN { pc=100*${success}/${extension_count}; i=int(pc); print (pc-i<0.5)?i:i+1 }") -failure_percent=$(awk "BEGIN { pc=100*${failure_count}/${extension_count}; i=int(pc); print (pc-i<0.5)?i:i+1 }") - -printf "\n\n***TRUNK INSTALL EXTENSIONS THAT VERSION NOT FOUND RATE***\n" -echo "$not_found_count / $extension_count" -printf "***CREATE EXTENSIONS SUCCESS RATE***\n" -echo "$success / $extension_count ($success_percent%)" -printf "\n\n***CREATE EXTENSION FAILURE RATE***\n" -echo "$failure_count / $extension_count ($failure_percent%)" -printf "\n\n***CREATE EXTENSIONS THAT NEED TO BE LOADED IN shared_preload_libraries***\n" -echo "$need_load_shared_preload_libraries_count / $extension_count" -printf "\n\n***NEED TO LOAD shared_preload_libraries EXTENSIONS***\n" -for need in "${need_load_shared_preload_libraries_extensions[@]}" -do - echo $need -done - -printf "\n\n***FAILED EXTENSIONS***\n" -for failed in "${failed_extensions[@]}" -do - echo $failed -done \ No newline at end of file diff --git a/docker/pg-cnpg/requirements.txt b/docker/pg-cnpg/requirements.txt deleted file mode 100644 index 0aac74fd..00000000 --- a/docker/pg-cnpg/requirements.txt +++ /dev/null @@ -1,637 +0,0 @@ -# -# This file is autogenerated by pip-compile with Python 3.11 -# by the following command: -# -# pip-compile --generate-hashes -# -azure-core==1.32.0 \ - --hash=sha256:22b3c35d6b2dae14990f6c1be2912bf23ffe50b220e708a28ab1bb92b1c730e5 \ - --hash=sha256:eac191a0efb23bfa83fddf321b27b122b4ec847befa3091fa736a5c32c50d7b4 - # via - # azure-identity - # azure-storage-blob -azure-identity==1.21.0 \ - --hash=sha256:258ea6325537352440f71b35c3dffe9d240eae4a5126c1b7ce5efd5766bd9fd9 \ - --hash=sha256:ea22ce6e6b0f429bc1b8d9212d5b9f9877bd4c82f1724bfa910760612c07a9a6 -azure-storage-blob==12.25.0 \ - --hash=sha256:42364ca8f9f49dbccd0acc10144ed47bb6770bf78719970b51915f048891abba \ - --hash=sha256:a38e18bf10258fb19028f343db0d3d373280c6427a619c98c06d76485805b755 -barman[azure,cloud,google,lz4,snappy,zstandard]==3.12.1 \ - --hash=sha256:258ef7100717f66032402e0abe03c977089c50fc47143df5708e92aa1d772937 \ - --hash=sha256:9dd7be219b6f74954b80cdc28f9a72f2acb923e7da65edd0f41cdc82fd32e169 - # via -r requirements.in -boto3==1.35.99 \ - --hash=sha256:83e560faaec38a956dfb3d62e05e1703ee50432b45b788c09e25107c5058bd71 \ - --hash=sha256:e0abd794a7a591d90558e92e29a9f8837d25ece8e3c120e530526fe27eba5fca - # via - # -r requirements.in - # barman -botocore==1.35.99 \ - --hash=sha256:1eab44e969c39c5f3d9a3104a0836c24715579a455f12b3979a31d7cde51b3c3 \ - --hash=sha256:b22d27b6b617fc2d7342090d6129000af2efd20174215948c0d7ae2da0fab445 - # via - # boto3 - # s3transfer -cachetools==5.5.2 \ - --hash=sha256:1a661caa9175d26759571b2e19580f9d6393969e5dfca11fdb1f947a23e640d4 \ - --hash=sha256:d26a22bcc62eb95c3beabd9f1ee5e820d3d2704fe2967cbe350e20c8ffcd3f0a - # via google-auth -certifi==2025.1.31 \ - --hash=sha256:3d5da6925056f6f18f119200434a4780a94263f10d1c21d032a6f6b2baa20651 \ - --hash=sha256:ca78db4565a652026a4db2bcdf68f2fb589ea80d0be70e03929ed730746b84fe - # via requests -cffi==1.17.1 \ - --hash=sha256:045d61c734659cc045141be4bae381a41d89b741f795af1dd018bfb532fd0df8 \ - --hash=sha256:0984a4925a435b1da406122d4d7968dd861c1385afe3b45ba82b750f229811e2 \ - --hash=sha256:0e2b1fac190ae3ebfe37b979cc1ce69c81f4e4fe5746bb401dca63a9062cdaf1 \ - --hash=sha256:0f048dcf80db46f0098ccac01132761580d28e28bc0f78ae0d58048063317e15 \ - --hash=sha256:1257bdabf294dceb59f5e70c64a3e2f462c30c7ad68092d01bbbfb1c16b1ba36 \ - --hash=sha256:1c39c6016c32bc48dd54561950ebd6836e1670f2ae46128f67cf49e789c52824 \ - --hash=sha256:1d599671f396c4723d016dbddb72fe8e0397082b0a77a4fab8028923bec050e8 \ - --hash=sha256:28b16024becceed8c6dfbc75629e27788d8a3f9030691a1dbf9821a128b22c36 \ - --hash=sha256:2bb1a08b8008b281856e5971307cc386a8e9c5b625ac297e853d36da6efe9c17 \ - --hash=sha256:30c5e0cb5ae493c04c8b42916e52ca38079f1b235c2f8ae5f4527b963c401caf \ - --hash=sha256:31000ec67d4221a71bd3f67df918b1f88f676f1c3b535a7eb473255fdc0b83fc \ - --hash=sha256:386c8bf53c502fff58903061338ce4f4950cbdcb23e2902d86c0f722b786bbe3 \ - --hash=sha256:3edc8d958eb099c634dace3c7e16560ae474aa3803a5df240542b305d14e14ed \ - --hash=sha256:45398b671ac6d70e67da8e4224a065cec6a93541bb7aebe1b198a61b58c7b702 \ - --hash=sha256:46bf43160c1a35f7ec506d254e5c890f3c03648a4dbac12d624e4490a7046cd1 \ - --hash=sha256:4ceb10419a9adf4460ea14cfd6bc43d08701f0835e979bf821052f1805850fe8 \ - --hash=sha256:51392eae71afec0d0c8fb1a53b204dbb3bcabcb3c9b807eedf3e1e6ccf2de903 \ - --hash=sha256:5da5719280082ac6bd9aa7becb3938dc9f9cbd57fac7d2871717b1feb0902ab6 \ - --hash=sha256:610faea79c43e44c71e1ec53a554553fa22321b65fae24889706c0a84d4ad86d \ - --hash=sha256:636062ea65bd0195bc012fea9321aca499c0504409f413dc88af450b57ffd03b \ - --hash=sha256:6883e737d7d9e4899a8a695e00ec36bd4e5e4f18fabe0aca0efe0a4b44cdb13e \ - --hash=sha256:6b8b4a92e1c65048ff98cfe1f735ef8f1ceb72e3d5f0c25fdb12087a23da22be \ - --hash=sha256:6f17be4345073b0a7b8ea599688f692ac3ef23ce28e5df79c04de519dbc4912c \ - --hash=sha256:706510fe141c86a69c8ddc029c7910003a17353970cff3b904ff0686a5927683 \ - --hash=sha256:72e72408cad3d5419375fc87d289076ee319835bdfa2caad331e377589aebba9 \ - --hash=sha256:733e99bc2df47476e3848417c5a4540522f234dfd4ef3ab7fafdf555b082ec0c \ - --hash=sha256:7596d6620d3fa590f677e9ee430df2958d2d6d6de2feeae5b20e82c00b76fbf8 \ - --hash=sha256:78122be759c3f8a014ce010908ae03364d00a1f81ab5c7f4a7a5120607ea56e1 \ - --hash=sha256:805b4371bf7197c329fcb3ead37e710d1bca9da5d583f5073b799d5c5bd1eee4 \ - --hash=sha256:85a950a4ac9c359340d5963966e3e0a94a676bd6245a4b55bc43949eee26a655 \ - --hash=sha256:8f2cdc858323644ab277e9bb925ad72ae0e67f69e804f4898c070998d50b1a67 \ - --hash=sha256:9755e4345d1ec879e3849e62222a18c7174d65a6a92d5b346b1863912168b595 \ - --hash=sha256:98e3969bcff97cae1b2def8ba499ea3d6f31ddfdb7635374834cf89a1a08ecf0 \ - --hash=sha256:a08d7e755f8ed21095a310a693525137cfe756ce62d066e53f502a83dc550f65 \ - --hash=sha256:a1ed2dd2972641495a3ec98445e09766f077aee98a1c896dcb4ad0d303628e41 \ - --hash=sha256:a24ed04c8ffd54b0729c07cee15a81d964e6fee0e3d4d342a27b020d22959dc6 \ - --hash=sha256:a45e3c6913c5b87b3ff120dcdc03f6131fa0065027d0ed7ee6190736a74cd401 \ - --hash=sha256:a9b15d491f3ad5d692e11f6b71f7857e7835eb677955c00cc0aefcd0669adaf6 \ - --hash=sha256:ad9413ccdeda48c5afdae7e4fa2192157e991ff761e7ab8fdd8926f40b160cc3 \ - --hash=sha256:b2ab587605f4ba0bf81dc0cb08a41bd1c0a5906bd59243d56bad7668a6fc6c16 \ - --hash=sha256:b62ce867176a75d03a665bad002af8e6d54644fad99a3c70905c543130e39d93 \ - --hash=sha256:c03e868a0b3bc35839ba98e74211ed2b05d2119be4e8a0f224fba9384f1fe02e \ - --hash=sha256:c59d6e989d07460165cc5ad3c61f9fd8f1b4796eacbd81cee78957842b834af4 \ - --hash=sha256:c7eac2ef9b63c79431bc4b25f1cd649d7f061a28808cbc6c47b534bd789ef964 \ - --hash=sha256:c9c3d058ebabb74db66e431095118094d06abf53284d9c81f27300d0e0d8bc7c \ - --hash=sha256:ca74b8dbe6e8e8263c0ffd60277de77dcee6c837a3d0881d8c1ead7268c9e576 \ - --hash=sha256:caaf0640ef5f5517f49bc275eca1406b0ffa6aa184892812030f04c2abf589a0 \ - --hash=sha256:cdf5ce3acdfd1661132f2a9c19cac174758dc2352bfe37d98aa7512c6b7178b3 \ - --hash=sha256:d016c76bdd850f3c626af19b0542c9677ba156e4ee4fccfdd7848803533ef662 \ - --hash=sha256:d01b12eeeb4427d3110de311e1774046ad344f5b1a7403101878976ecd7a10f3 \ - --hash=sha256:d63afe322132c194cf832bfec0dc69a99fb9bb6bbd550f161a49e9e855cc78ff \ - --hash=sha256:da95af8214998d77a98cc14e3a3bd00aa191526343078b530ceb0bd710fb48a5 \ - --hash=sha256:dd398dbc6773384a17fe0d3e7eeb8d1a21c2200473ee6806bb5e6a8e62bb73dd \ - --hash=sha256:de2ea4b5833625383e464549fec1bc395c1bdeeb5f25c4a3a82b5a8c756ec22f \ - --hash=sha256:de55b766c7aa2e2a3092c51e0483d700341182f08e67c63630d5b6f200bb28e5 \ - --hash=sha256:df8b1c11f177bc2313ec4b2d46baec87a5f3e71fc8b45dab2ee7cae86d9aba14 \ - --hash=sha256:e03eab0a8677fa80d646b5ddece1cbeaf556c313dcfac435ba11f107ba117b5d \ - --hash=sha256:e221cf152cff04059d011ee126477f0d9588303eb57e88923578ace7baad17f9 \ - --hash=sha256:e31ae45bc2e29f6b2abd0de1cc3b9d5205aa847cafaecb8af1476a609a2f6eb7 \ - --hash=sha256:edae79245293e15384b51f88b00613ba9f7198016a5948b5dddf4917d4d26382 \ - --hash=sha256:f1e22e8c4419538cb197e4dd60acc919d7696e5ef98ee4da4e01d3f8cfa4cc5a \ - --hash=sha256:f3a2b4222ce6b60e2e8b337bb9596923045681d71e5a082783484d845390938e \ - --hash=sha256:f6a16c31041f09ead72d69f583767292f750d24913dadacf5756b966aacb3f1a \ - --hash=sha256:f75c7ab1f9e4aca5414ed4d8e5c0e303a34f4421f8a0d47a4d019ceff0ab6af4 \ - --hash=sha256:f79fc4fc25f1c8698ff97788206bb3c2598949bfe0fef03d299eb1b5356ada99 \ - --hash=sha256:f7f5baafcc48261359e14bcd6d9bff6d4b28d9103847c9e136694cb0501aef87 \ - --hash=sha256:fc48c783f9c87e60831201f2cce7f3b2e4846bf4d8728eabe54d60700b318a0b - # via cryptography -charset-normalizer==3.4.1 \ - --hash=sha256:0167ddc8ab6508fe81860a57dd472b2ef4060e8d378f0cc555707126830f2537 \ - --hash=sha256:01732659ba9b5b873fc117534143e4feefecf3b2078b0a6a2e925271bb6f4cfa \ - --hash=sha256:01ad647cdd609225c5350561d084b42ddf732f4eeefe6e678765636791e78b9a \ - --hash=sha256:04432ad9479fa40ec0f387795ddad4437a2b50417c69fa275e212933519ff294 \ - --hash=sha256:0907f11d019260cdc3f94fbdb23ff9125f6b5d1039b76003b5b0ac9d6a6c9d5b \ - --hash=sha256:0924e81d3d5e70f8126529951dac65c1010cdf117bb75eb02dd12339b57749dd \ - --hash=sha256:09b26ae6b1abf0d27570633b2b078a2a20419c99d66fb2823173d73f188ce601 \ - --hash=sha256:09b5e6733cbd160dcc09589227187e242a30a49ca5cefa5a7edd3f9d19ed53fd \ - --hash=sha256:0af291f4fe114be0280cdd29d533696a77b5b49cfde5467176ecab32353395c4 \ - --hash=sha256:0f55e69f030f7163dffe9fd0752b32f070566451afe180f99dbeeb81f511ad8d \ - --hash=sha256:1a2bc9f351a75ef49d664206d51f8e5ede9da246602dc2d2726837620ea034b2 \ - --hash=sha256:22e14b5d70560b8dd51ec22863f370d1e595ac3d024cb8ad7d308b4cd95f8313 \ - --hash=sha256:234ac59ea147c59ee4da87a0c0f098e9c8d169f4dc2a159ef720f1a61bbe27cd \ - --hash=sha256:2369eea1ee4a7610a860d88f268eb39b95cb588acd7235e02fd5a5601773d4fa \ - --hash=sha256:237bdbe6159cff53b4f24f397d43c6336c6b0b42affbe857970cefbb620911c8 \ - --hash=sha256:28bf57629c75e810b6ae989f03c0828d64d6b26a5e205535585f96093e405ed1 \ - --hash=sha256:2967f74ad52c3b98de4c3b32e1a44e32975e008a9cd2a8cc8966d6a5218c5cb2 \ - --hash=sha256:2a75d49014d118e4198bcee5ee0a6f25856b29b12dbf7cd012791f8a6cc5c496 \ - --hash=sha256:2bdfe3ac2e1bbe5b59a1a63721eb3b95fc9b6817ae4a46debbb4e11f6232428d \ - --hash=sha256:2d074908e1aecee37a7635990b2c6d504cd4766c7bc9fc86d63f9c09af3fa11b \ - --hash=sha256:2fb9bd477fdea8684f78791a6de97a953c51831ee2981f8e4f583ff3b9d9687e \ - --hash=sha256:311f30128d7d333eebd7896965bfcfbd0065f1716ec92bd5638d7748eb6f936a \ - --hash=sha256:329ce159e82018d646c7ac45b01a430369d526569ec08516081727a20e9e4af4 \ - --hash=sha256:345b0426edd4e18138d6528aed636de7a9ed169b4aaf9d61a8c19e39d26838ca \ - --hash=sha256:363e2f92b0f0174b2f8238240a1a30142e3db7b957a5dd5689b0e75fb717cc78 \ - --hash=sha256:3a3bd0dcd373514dcec91c411ddb9632c0d7d92aed7093b8c3bbb6d69ca74408 \ - --hash=sha256:3bed14e9c89dcb10e8f3a29f9ccac4955aebe93c71ae803af79265c9ca5644c5 \ - --hash=sha256:44251f18cd68a75b56585dd00dae26183e102cd5e0f9f1466e6df5da2ed64ea3 \ - --hash=sha256:44ecbf16649486d4aebafeaa7ec4c9fed8b88101f4dd612dcaf65d5e815f837f \ - --hash=sha256:4532bff1b8421fd0a320463030c7520f56a79c9024a4e88f01c537316019005a \ - --hash=sha256:49402233c892a461407c512a19435d1ce275543138294f7ef013f0b63d5d3765 \ - --hash=sha256:4c0907b1928a36d5a998d72d64d8eaa7244989f7aaaf947500d3a800c83a3fd6 \ - --hash=sha256:4d86f7aff21ee58f26dcf5ae81a9addbd914115cdebcbb2217e4f0ed8982e146 \ - --hash=sha256:5777ee0881f9499ed0f71cc82cf873d9a0ca8af166dfa0af8ec4e675b7df48e6 \ - --hash=sha256:5df196eb874dae23dcfb968c83d4f8fdccb333330fe1fc278ac5ceeb101003a9 \ - --hash=sha256:619a609aa74ae43d90ed2e89bdd784765de0a25ca761b93e196d938b8fd1dbbd \ - --hash=sha256:6e27f48bcd0957c6d4cb9d6fa6b61d192d0b13d5ef563e5f2ae35feafc0d179c \ - --hash=sha256:6ff8a4a60c227ad87030d76e99cd1698345d4491638dfa6673027c48b3cd395f \ - --hash=sha256:73d94b58ec7fecbc7366247d3b0b10a21681004153238750bb67bd9012414545 \ - --hash=sha256:7461baadb4dc00fd9e0acbe254e3d7d2112e7f92ced2adc96e54ef6501c5f176 \ - --hash=sha256:75832c08354f595c760a804588b9357d34ec00ba1c940c15e31e96d902093770 \ - --hash=sha256:7709f51f5f7c853f0fb938bcd3bc59cdfdc5203635ffd18bf354f6967ea0f824 \ - --hash=sha256:78baa6d91634dfb69ec52a463534bc0df05dbd546209b79a3880a34487f4b84f \ - --hash=sha256:7974a0b5ecd505609e3b19742b60cee7aa2aa2fb3151bc917e6e2646d7667dcf \ - --hash=sha256:7a4f97a081603d2050bfaffdefa5b02a9ec823f8348a572e39032caa8404a487 \ - --hash=sha256:7b1bef6280950ee6c177b326508f86cad7ad4dff12454483b51d8b7d673a2c5d \ - --hash=sha256:7d053096f67cd1241601111b698f5cad775f97ab25d81567d3f59219b5f1adbd \ - --hash=sha256:804a4d582ba6e5b747c625bf1255e6b1507465494a40a2130978bda7b932c90b \ - --hash=sha256:807f52c1f798eef6cf26beb819eeb8819b1622ddfeef9d0977a8502d4db6d534 \ - --hash=sha256:80ed5e856eb7f30115aaf94e4a08114ccc8813e6ed1b5efa74f9f82e8509858f \ - --hash=sha256:8417cb1f36cc0bc7eaba8ccb0e04d55f0ee52df06df3ad55259b9a323555fc8b \ - --hash=sha256:8436c508b408b82d87dc5f62496973a1805cd46727c34440b0d29d8a2f50a6c9 \ - --hash=sha256:89149166622f4db9b4b6a449256291dc87a99ee53151c74cbd82a53c8c2f6ccd \ - --hash=sha256:8bfa33f4f2672964266e940dd22a195989ba31669bd84629f05fab3ef4e2d125 \ - --hash=sha256:8c60ca7339acd497a55b0ea5d506b2a2612afb2826560416f6894e8b5770d4a9 \ - --hash=sha256:91b36a978b5ae0ee86c394f5a54d6ef44db1de0815eb43de826d41d21e4af3de \ - --hash=sha256:955f8851919303c92343d2f66165294848d57e9bba6cf6e3625485a70a038d11 \ - --hash=sha256:97f68b8d6831127e4787ad15e6757232e14e12060bec17091b85eb1486b91d8d \ - --hash=sha256:9b23ca7ef998bc739bf6ffc077c2116917eabcc901f88da1b9856b210ef63f35 \ - --hash=sha256:9f0b8b1c6d84c8034a44893aba5e767bf9c7a211e313a9605d9c617d7083829f \ - --hash=sha256:aabfa34badd18f1da5ec1bc2715cadc8dca465868a4e73a0173466b688f29dda \ - --hash=sha256:ab36c8eb7e454e34e60eb55ca5d241a5d18b2c6244f6827a30e451c42410b5f7 \ - --hash=sha256:b010a7a4fd316c3c484d482922d13044979e78d1861f0e0650423144c616a46a \ - --hash=sha256:b1ac5992a838106edb89654e0aebfc24f5848ae2547d22c2c3f66454daa11971 \ - --hash=sha256:b7b2d86dd06bfc2ade3312a83a5c364c7ec2e3498f8734282c6c3d4b07b346b8 \ - --hash=sha256:b97e690a2118911e39b4042088092771b4ae3fc3aa86518f84b8cf6888dbdb41 \ - --hash=sha256:bc2722592d8998c870fa4e290c2eec2c1569b87fe58618e67d38b4665dfa680d \ - --hash=sha256:c0429126cf75e16c4f0ad00ee0eae4242dc652290f940152ca8c75c3a4b6ee8f \ - --hash=sha256:c30197aa96e8eed02200a83fba2657b4c3acd0f0aa4bdc9f6c1af8e8962e0757 \ - --hash=sha256:c4c3e6da02df6fa1410a7680bd3f63d4f710232d3139089536310d027950696a \ - --hash=sha256:c75cb2a3e389853835e84a2d8fb2b81a10645b503eca9bcb98df6b5a43eb8886 \ - --hash=sha256:c96836c97b1238e9c9e3fe90844c947d5afbf4f4c92762679acfe19927d81d77 \ - --hash=sha256:d7f50a1f8c450f3925cb367d011448c39239bb3eb4117c36a6d354794de4ce76 \ - --hash=sha256:d973f03c0cb71c5ed99037b870f2be986c3c05e63622c017ea9816881d2dd247 \ - --hash=sha256:d98b1668f06378c6dbefec3b92299716b931cd4e6061f3c875a71ced1780ab85 \ - --hash=sha256:d9c3cdf5390dcd29aa8056d13e8e99526cda0305acc038b96b30352aff5ff2bb \ - --hash=sha256:dad3e487649f498dd991eeb901125411559b22e8d7ab25d3aeb1af367df5efd7 \ - --hash=sha256:dccbe65bd2f7f7ec22c4ff99ed56faa1e9f785482b9bbd7c717e26fd723a1d1e \ - --hash=sha256:dd78cfcda14a1ef52584dbb008f7ac81c1328c0f58184bf9a84c49c605002da6 \ - --hash=sha256:e218488cd232553829be0664c2292d3af2eeeb94b32bea483cf79ac6a694e037 \ - --hash=sha256:e358e64305fe12299a08e08978f51fc21fac060dcfcddd95453eabe5b93ed0e1 \ - --hash=sha256:ea0d8d539afa5eb2728aa1932a988a9a7af94f18582ffae4bc10b3fbdad0626e \ - --hash=sha256:eab677309cdb30d047996b36d34caeda1dc91149e4fdca0b1a039b3f79d9a807 \ - --hash=sha256:eb8178fe3dba6450a3e024e95ac49ed3400e506fd4e9e5c32d30adda88cbd407 \ - --hash=sha256:ecddf25bee22fe4fe3737a399d0d177d72bc22be6913acfab364b40bce1ba83c \ - --hash=sha256:eea6ee1db730b3483adf394ea72f808b6e18cf3cb6454b4d86e04fa8c4327a12 \ - --hash=sha256:f08ff5e948271dc7e18a35641d2f11a4cd8dfd5634f55228b691e62b37125eb3 \ - --hash=sha256:f30bf9fd9be89ecb2360c7d94a711f00c09b976258846efe40db3d05828e8089 \ - --hash=sha256:fa88b843d6e211393a37219e6a1c1df99d35e8fd90446f1118f4216e307e48cd \ - --hash=sha256:fc54db6c8593ef7d4b2a331b58653356cf04f67c960f584edb7c3d8c97e8f39e \ - --hash=sha256:fd4ec41f914fa74ad1b8304bbc634b3de73d2a0889bd32076342a573e0779e00 \ - --hash=sha256:ffc9202a29ab3920fa812879e95a9e78b2465fd10be7fcbd042899695d75e616 - # via requests -cramjam==2.9.1 \ - --hash=sha256:008b49b455b396acc5459dfb06fb9d56049c4097ee8e590892a4d3da9a711da3 \ - --hash=sha256:038df668ffb94d64d67b6ecc59cbd206745a425ffc0402897dde12d89fa6a870 \ - --hash=sha256:04828cbfad7384f06a4a7d0d927c3e85ef11dc5a40b9cf5f3e29ac4e23ecd678 \ - --hash=sha256:06068bd191a82ad4fc1ac23d6f8627fb5e37ec4be0431711b9a2dbacaccfeddb \ - --hash=sha256:06e3f97a379386d97debf08638a78b3d3850fdf6124755eb270b54905a169930 \ - --hash=sha256:07ac76b7f992556e7aa910244be11ece578cdf84f4d5d5297461f9a895e18312 \ - --hash=sha256:0bedb84e068b53c944bd08dcb501fd00d67daa8a917922356dd559b484ce7eab \ - --hash=sha256:11118675e9c7952ececabc62f023290ee4f8ecf0bee0d2c7eb8d1c402ee9769d \ - --hash=sha256:1376f6fdbf0b30712413a0b4e51663a4938ae2f6b449f8e4635dbb3694db83cf \ - --hash=sha256:1539fd758f0e57fad7913cebff8baaee871bb561ddf6fa710a427b74da6b6778 \ - --hash=sha256:15955dd75e80f66c1ea271167a5347661d9bdc365f894a57698c383c9b7d465c \ - --hash=sha256:1c33bc095db5733c841a102b8693062be5db8cdac17b9782ebc00577c6a94480 \ - --hash=sha256:217fe22b41f8c3dce03852f828b059abfad11d1344a1df2f43d3eb8634b18d75 \ - --hash=sha256:21ea784e6c3f1843d3523ae0f03651dd06058b39eeb64beb82ee3b100fa83662 \ - --hash=sha256:23b9786d1d17686fb8d600ade2a19374c7188d4b8867efa9af0d8274a220aec7 \ - --hash=sha256:258120cb1e3afc3443f756f9de161ed63eed56a2c31f6093e81c571c0f2dc9f6 \ - --hash=sha256:27571bfa5a5d618604696747d0dc1d2a99b5906c967c8dee53c13a7107edfde6 \ - --hash=sha256:2ceef6e09ee22457997370882aa3c69de01e6dd0aaa2f953e1e87ad11641d042 \ - --hash=sha256:335103317475bf992953c58838152a4761fc3c87354000edbfc4d7e57cf05909 \ - --hash=sha256:336cc591d86cbd225d256813779f46624f857bc9c779db126271eff9ddc524ae \ - --hash=sha256:342fb946f8d3e9e35b837288b03ab23cfbe0bb5a30e582ed805ef79706823a96 \ - --hash=sha256:3b695259e71fde6d5be66b77a4474523ced9ffe9fe8a34cb9b520ec1241a14d3 \ - --hash=sha256:3ba79c7d2cc5adb897b690c05dd9b67c4d401736d207314b99315f7be3cd94fd \ - --hash=sha256:3f815fb0eba625af45139af4f90f5fc2ddda61b171c2cc3ab63d44b40c5c7768 \ - --hash=sha256:4125d8cd86fa08495d310e80926c2f0563f157b76862e7479f9b2cf94823ea0c \ - --hash=sha256:4206ebdd1d1ef0f3f86c8c2f7c426aa4af6094f4f41e274601fd4c4569f37454 \ - --hash=sha256:440b489902bfb7a26d3fec1ca888007615336ff763d2a32a2fc40586548a0dbf \ - --hash=sha256:45c18cc13156e8697a8d3f9e57e49a69b00e14a103196efab0893fae1a5257f8 \ - --hash=sha256:4826d6d81ea490fa7a3ae7a4b9729866a945ffac1f77fe57b71e49d6e1b21efd \ - --hash=sha256:53145fc9f2319c1245d4329e1da8cfacd6e35e27090c07c0b9d453ae2bbdac3e \ - --hash=sha256:56495975401b1821dbe1f29cf222e23556232209a2fdb809fe8156d120ca9c7f \ - --hash=sha256:57ca8f3775324a9de3ee6f05ca172687ba258c0dea79f7e3a6b4112834982f2a \ - --hash=sha256:5925a738b8478f223ab9756fc794e3cabd5917fd7846f66adcf1d5fc2bf9864c \ - --hash=sha256:5a7797a2fff994fc5e323f7a967a35a3e37e3006ed21d64dcded086502f482af \ - --hash=sha256:67040e0fd84404885ec716a806bee6110f9960c3647e0ef1670aab3b7375a70a \ - --hash=sha256:6a2ca4d3c683d28d3217821029eb08d3487d5043d7eb455df11ff3cacfd4c916 \ - --hash=sha256:6b19fc60ead1cae9795a5b359599da3a1c95d38f869bdfb51c441fd76b04e926 \ - --hash=sha256:6b5cef5cf40725fe64592af9ec163e7389855077700678a1d94bec549403a74d \ - --hash=sha256:6b7de6b61b11545570e4d6033713f3599525efc615ee353a822be8f6b0c65b77 \ - --hash=sha256:6d2df8a6511cc08ef1fccd2e0c65e2ebc9f57574ec8376052a76851af5398810 \ - --hash=sha256:6d86b44933aea0151e4a2e1e6935448499849045c38167d288ca4c59d5b8cd4e \ - --hash=sha256:79417957972553502b217a0093532e48893c8b4ca30ccc941cefe9c72379df7c \ - --hash=sha256:7eb032549dec897b942ddcf80c1cdccbcb40629f15fc902731dbe6362da49326 \ - --hash=sha256:8097ee39b61c86848a443c0b25b2df1de6b331fd512b20836a4f5cfde51ab255 \ - --hash=sha256:84d154fbadece82935396eb6bcb502085d944d2fd13b07a94348364344370c2c \ - --hash=sha256:86824c695688fcd06c5ac9bbd3fea9bdfb4cca194b1e706fbf11a629df48d2b4 \ - --hash=sha256:872b00ff83e84bcbdc7e951af291ebe65eed20b09c47e7c4af21c312f90b796f \ - --hash=sha256:8a9f52c27292c21457f43c4ce124939302a9acfb62295e7cda8667310563a5a3 \ - --hash=sha256:8bc9c2c748aaf91863d89c4583f529c1c709485c94f8dfeb3ee48662d88e3258 \ - --hash=sha256:8d1248dfa7f151e893ce819670f00879e4b7650b8d4c01279ce4f12140d68dd2 \ - --hash=sha256:8d47fd41ce260cf4f0ff0e788de961fab9e9c6844a05ce55d06ce31e06107bdc \ - --hash=sha256:8dc5207567459d049696f62a1fdfb220f3fe6aa0d722285d44753e12504dac6c \ - --hash=sha256:8e0c5d98a4e791f0bbd0ffcb7dae879baeb2dcc357348a8dc2be0a8c10403a2a \ - --hash=sha256:8e82464d1e00fbbb12958999b8471ba5e9f3d9711954505a0a7b378762332e6f \ - --hash=sha256:95f3646ddc98af25af25d5692ae65966488a283813336ea9cf41b22e542e7c0d \ - --hash=sha256:9847dd6f288f1c56359f52acb48ff2df848ff3e3bff34d23855bbcf7016427cc \ - --hash=sha256:9da6d970281083bae91b914362de325414aa03c01fc806f6bb2cc006322ec834 \ - --hash=sha256:9e9193cd4bb57e7acd3af24891526299244bfed88168945efdaa09af4e50720f \ - --hash=sha256:a237064a6e2c2256c9a1cf2beb7c971382190c0f1eb2e810e02e971881756132 \ - --hash=sha256:a36adf7d13b7accfa206e1c917f08924eb905b45aa8e62176509afa7b14db71e \ - --hash=sha256:a47de0a68f5f4d9951250ef5af31f2a7228132caa9ed60994234f7eb98090d33 \ - --hash=sha256:ab1e69dc4831bbb79b6d547077aae89074c83e8ad94eba1a3d80e94d2424fd02 \ - --hash=sha256:ab687bef5c493732b9a4ab870542ee43f5eae0025f9c684c7cb399c3a85cb380 \ - --hash=sha256:ac48b978aa0675f62b642750e798c394a64d25ce852e4e541f69bef9a564c2f0 \ - --hash=sha256:af39006faddfc6253beb93ca821d544931cfee7f0177b99ff106dfd8fd6a2cd8 \ - --hash=sha256:b0944a7c3a78f940c06d1b29bdce91a17798d80593dd01ebfeb842761e48a8b5 \ - --hash=sha256:b3291be0d3f73d5774d69013be4ab33978c777363b5312d14f62f77817c2f75a \ - --hash=sha256:b5b1cd7d39242b2b903cf09cd4696b3a6e04dc537ffa9f3ac8668edae76eecb6 \ - --hash=sha256:b7ac273498a2c6772d67707e101b74014c0d9413bb4711c51d8ec311de59b4b1 \ - --hash=sha256:b9db1debe48060e41a5b91af9193c524e473c57f6105462c5524a41f5aabdb88 \ - --hash=sha256:ba560244bc1335b420b74e91e35f9d4e7f307a3be3a4603ce0f0d7e15a0acdf0 \ - --hash=sha256:c60e5996aa02547d12bc2740d44e90e006b0f93100f53206f7abe6732ad56e69 \ - --hash=sha256:ce2b94117f373defc876f88e74e44049a9969223dbca3240415b71752d0422fb \ - --hash=sha256:cf29b4def86ec503e329fe138842a9b79a997e3beb6c7809b05665a0d291edff \ - --hash=sha256:cf4ea758d98b6fad1b4b2d808d0de690d3162ac56c26968aea0af6524e3eb736 \ - --hash=sha256:d14a0efb21e0fec0631bcd66040b06e6a0fe10825f3aacffded38c1c978bdff9 \ - --hash=sha256:d35923fb5411bde30b53c0696dff8e24c8a38b010b89544834c53f4462fd71df \ - --hash=sha256:d51b9b140b1df39a44bff7896d98a10da345b7d5f5ce92368d328c1c2c829167 \ - --hash=sha256:d90a72608c7550cd7eba914668f6277bfb0b24f074d1f1bd9d061fcb6f2adbd6 \ - --hash=sha256:da0cc0efdbfb8ee2361f89f38ded03d11678f37e392afff7a97b09c55dadfc83 \ - --hash=sha256:dda7698b6d7caeae1047adafebc4b43b2a82478234f6c2b45bc3edad854e0600 \ - --hash=sha256:e076fd87089197cb61117c63dbe7712ad5eccb93968860eb3bae09b767bac813 \ - --hash=sha256:e13c9a697881e5e38148958612dc6856967f5ff8cd7bba5ff751f2d6ac020aa4 \ - --hash=sha256:ec769e5b16251704502277a1163dcf2611551452d7590ff4cc422b7b0367fc96 \ - --hash=sha256:f6f18f0242212d3409d26ce3874937b5b979cebd61f08b633a6ea893c32fc7b6 \ - --hash=sha256:f89924858712b8b936f04f3d690e72825a3e5127a140b434c79030c1c5a887ce \ - --hash=sha256:fb01f6e38719818778144d3165a89ea1ad9dc58c6342b7f20aa194c70f34cbd1 \ - --hash=sha256:fbfe35929a61b914de9e5dbacde0cfbba86cbf5122f9285a24c14ed0b645490b \ - --hash=sha256:fd0fa9a0e7f18224b6d2d1d69dbdc3aecec80ef1393c59244159b131604a4395 \ - --hash=sha256:ff362f68bd68ac0eccb445209238d589bba728fb6d7f2e9dc199e0ec3a61d6e0 - # via - # barman - # python-snappy -cryptography==44.0.2 \ - --hash=sha256:04abd71114848aa25edb28e225ab5f268096f44cf0127f3d36975bdf1bdf3390 \ - --hash=sha256:0529b1d5a0105dd3731fa65680b45ce49da4d8115ea76e9da77a875396727b41 \ - --hash=sha256:1bc312dfb7a6e5d66082c87c34c8a62176e684b6fe3d90fcfe1568de675e6688 \ - --hash=sha256:268e4e9b177c76d569e8a145a6939eca9a5fec658c932348598818acf31ae9a5 \ - --hash=sha256:29ecec49f3ba3f3849362854b7253a9f59799e3763b0c9d0826259a88efa02f1 \ - --hash=sha256:2bf7bf75f7df9715f810d1b038870309342bff3069c5bd8c6b96128cb158668d \ - --hash=sha256:3b721b8b4d948b218c88cb8c45a01793483821e709afe5f622861fc6182b20a7 \ - --hash=sha256:3c00b6b757b32ce0f62c574b78b939afab9eecaf597c4d624caca4f9e71e7843 \ - --hash=sha256:3dc62975e31617badc19a906481deacdeb80b4bb454394b4098e3f2525a488c5 \ - --hash=sha256:4973da6ca3db4405c54cd0b26d328be54c7747e89e284fcff166132eb7bccc9c \ - --hash=sha256:4e389622b6927d8133f314949a9812972711a111d577a5d1f4bee5e58736b80a \ - --hash=sha256:51e4de3af4ec3899d6d178a8c005226491c27c4ba84101bfb59c901e10ca9f79 \ - --hash=sha256:5f6f90b72d8ccadb9c6e311c775c8305381db88374c65fa1a68250aa8a9cb3a6 \ - --hash=sha256:6210c05941994290f3f7f175a4a57dbbb2afd9273657614c506d5976db061181 \ - --hash=sha256:6f101b1f780f7fc613d040ca4bdf835c6ef3b00e9bd7125a4255ec574c7916e4 \ - --hash=sha256:7bdcd82189759aba3816d1f729ce42ffded1ac304c151d0a8e89b9996ab863d5 \ - --hash=sha256:7ca25849404be2f8e4b3c59483d9d3c51298a22c1c61a0e84415104dacaf5562 \ - --hash=sha256:81276f0ea79a208d961c433a947029e1a15948966658cf6710bbabb60fcc2639 \ - --hash=sha256:8cadc6e3b5a1f144a039ea08a0bdb03a2a92e19c46be3285123d32029f40a922 \ - --hash=sha256:8e0ddd63e6bf1161800592c71ac794d3fb8001f2caebe0966e77c5234fa9efc3 \ - --hash=sha256:909c97ab43a9c0c0b0ada7a1281430e4e5ec0458e6d9244c0e821bbf152f061d \ - --hash=sha256:96e7a5e9d6e71f9f4fca8eebfd603f8e86c5225bb18eb621b2c1e50b290a9471 \ - --hash=sha256:9a1e657c0f4ea2a23304ee3f964db058c9e9e635cc7019c4aa21c330755ef6fd \ - --hash=sha256:9eb9d22b0a5d8fd9925a7764a054dca914000607dff201a24c791ff5c799e1fa \ - --hash=sha256:af4ff3e388f2fa7bff9f7f2b31b87d5651c45731d3e8cfa0944be43dff5cfbdb \ - --hash=sha256:b042d2a275c8cee83a4b7ae30c45a15e6a4baa65a179a0ec2d78ebb90e4f6699 \ - --hash=sha256:bc821e161ae88bfe8088d11bb39caf2916562e0a2dc7b6d56714a48b784ef0bb \ - --hash=sha256:c505d61b6176aaf982c5717ce04e87da5abc9a36a5b39ac03905c4aafe8de7aa \ - --hash=sha256:c63454aa261a0cf0c5b4718349629793e9e634993538db841165b3df74f37ec0 \ - --hash=sha256:c7362add18b416b69d58c910caa217f980c5ef39b23a38a0880dfd87bdf8cd23 \ - --hash=sha256:d03806036b4f89e3b13b6218fefea8d5312e450935b1a2d55f0524e2ed7c59d9 \ - --hash=sha256:d1b3031093a366ac767b3feb8bcddb596671b3aaff82d4050f984da0c248b615 \ - --hash=sha256:d1c3572526997b36f245a96a2b1713bf79ce99b271bbcf084beb6b9b075f29ea \ - --hash=sha256:efcfe97d1b3c79e486554efddeb8f6f53a4cdd4cf6086642784fa31fc384e1d7 \ - --hash=sha256:f514ef4cd14bb6fb484b4a60203e912cfcb64f2ab139e88c2274511514bf7308 - # via - # azure-identity - # azure-storage-blob - # msal - # pyjwt -google-api-core==2.24.2 \ - --hash=sha256:810a63ac95f3c441b7c0e43d344e372887f62ce9071ba972eacf32672e072de9 \ - --hash=sha256:81718493daf06d96d6bc76a91c23874dbf2fac0adbbf542831b805ee6e974696 - # via - # google-cloud-core - # google-cloud-storage -google-auth==2.38.0 \ - --hash=sha256:8285113607d3b80a3f1543b75962447ba8a09fe85783432a784fdeef6ac094c4 \ - --hash=sha256:e7dae6694313f434a2727bf2906f27ad259bae090d7aa896590d86feec3d9d4a - # via - # google-api-core - # google-cloud-core - # google-cloud-storage -google-cloud-core==2.4.3 \ - --hash=sha256:1fab62d7102844b278fe6dead3af32408b1df3eb06f5c7e8634cbd40edc4da53 \ - --hash=sha256:5130f9f4c14b4fafdff75c79448f9495cfade0d8775facf1b09c3bf67e027f6e - # via google-cloud-storage -google-cloud-storage==3.1.0 \ - --hash=sha256:944273179897c7c8a07ee15f2e6466a02da0c7c4b9ecceac2a26017cb2972049 \ - --hash=sha256:eaf36966b68660a9633f03b067e4a10ce09f1377cae3ff9f2c699f69a81c66c6 -google-crc32c==1.6.0 \ - --hash=sha256:05e2d8c9a2f853ff116db9706b4a27350587f341eda835f46db3c0a8c8ce2f24 \ - --hash=sha256:18e311c64008f1f1379158158bb3f0c8d72635b9eb4f9545f8cf990c5668e59d \ - --hash=sha256:236c87a46cdf06384f614e9092b82c05f81bd34b80248021f729396a78e55d7e \ - --hash=sha256:35834855408429cecf495cac67ccbab802de269e948e27478b1e47dfb6465e57 \ - --hash=sha256:386122eeaaa76951a8196310432c5b0ef3b53590ef4c317ec7588ec554fec5d2 \ - --hash=sha256:40b05ab32a5067525670880eb5d169529089a26fe35dce8891127aeddc1950e8 \ - --hash=sha256:48abd62ca76a2cbe034542ed1b6aee851b6f28aaca4e6551b5599b6f3ef175cc \ - --hash=sha256:50cf2a96da226dcbff8671233ecf37bf6e95de98b2a2ebadbfdf455e6d05df42 \ - --hash=sha256:51c4f54dd8c6dfeb58d1df5e4f7f97df8abf17a36626a217f169893d1d7f3e9f \ - --hash=sha256:5bcc90b34df28a4b38653c36bb5ada35671ad105c99cfe915fb5bed7ad6924aa \ - --hash=sha256:62f6d4a29fea082ac4a3c9be5e415218255cf11684ac6ef5488eea0c9132689b \ - --hash=sha256:6eceb6ad197656a1ff49ebfbbfa870678c75be4344feb35ac1edf694309413dc \ - --hash=sha256:7aec8e88a3583515f9e0957fe4f5f6d8d4997e36d0f61624e70469771584c760 \ - --hash=sha256:91ca8145b060679ec9176e6de4f89b07363d6805bd4760631ef254905503598d \ - --hash=sha256:a184243544811e4a50d345838a883733461e67578959ac59964e43cca2c791e7 \ - --hash=sha256:a9e4b426c3702f3cd23b933436487eb34e01e00327fac20c9aebb68ccf34117d \ - --hash=sha256:bb0966e1c50d0ef5bc743312cc730b533491d60585a9a08f897274e57c3f70e0 \ - --hash=sha256:bb8b3c75bd157010459b15222c3fd30577042a7060e29d42dabce449c087f2b3 \ - --hash=sha256:bd5e7d2445d1a958c266bfa5d04c39932dc54093fa391736dbfdb0f1929c1fb3 \ - --hash=sha256:c87d98c7c4a69066fd31701c4e10d178a648c2cac3452e62c6b24dc51f9fcc00 \ - --hash=sha256:d2952396dc604544ea7476b33fe87faedc24d666fb0c2d5ac971a2b9576ab871 \ - --hash=sha256:d8797406499f28b5ef791f339594b0b5fdedf54e203b5066675c406ba69d705c \ - --hash=sha256:d9e9913f7bd69e093b81da4535ce27af842e7bf371cde42d1ae9e9bd382dc0e9 \ - --hash=sha256:e2806553238cd076f0a55bddab37a532b53580e699ed8e5606d0de1f856b5205 \ - --hash=sha256:ebab974b1687509e5c973b5c4b8b146683e101e102e17a86bd196ecaa4d099fc \ - --hash=sha256:ed767bf4ba90104c1216b68111613f0d5926fb3780660ea1198fc469af410e9d \ - --hash=sha256:f7a1fc29803712f80879b0806cb83ab24ce62fc8daf0569f2204a0cfd7f68ed4 - # via - # google-cloud-storage - # google-resumable-media -google-resumable-media==2.7.2 \ - --hash=sha256:3ce7551e9fe6d99e9a126101d2536612bb73486721951e9562fee0f90c6ababa \ - --hash=sha256:5280aed4629f2b60b847b0d42f9857fd4935c11af266744df33d8074cae92fe0 - # via google-cloud-storage -googleapis-common-protos==1.69.1 \ - --hash=sha256:4077f27a6900d5946ee5a369fab9c8ded4c0ef1c6e880458ea2f70c14f7b70d5 \ - --hash=sha256:e20d2d8dda87da6fe7340afbbdf4f0bcb4c8fae7e6cadf55926c31f946b0b9b1 - # via google-api-core -idna==3.10 \ - --hash=sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9 \ - --hash=sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3 - # via requests -isodate==0.7.2 \ - --hash=sha256:28009937d8031054830160fce6d409ed342816b543597cece116d966c6d99e15 \ - --hash=sha256:4cd1aa0f43ca76f4a6c6c0292a85f40b35ec2e43e315b59f06e6d32171a953e6 - # via azure-storage-blob -jmespath==1.0.1 \ - --hash=sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980 \ - --hash=sha256:90261b206d6defd58fdd5e85f478bf633a2901798906be2ad389150c5c60edbe - # via - # boto3 - # botocore -lz4==4.4.3 \ - --hash=sha256:0aea6f283abd6acb1883b70d7a117b913e20c770845559f9421394bc9c522b24 \ - --hash=sha256:174b7ce5456671c73b81bb115defac8a584363d8b38a48ed3ad976e08eea27cd \ - --hash=sha256:1ebf23ffd36b32b980f720a81990fcfdeadacafe7498fbeff7a8e058259d4e58 \ - --hash=sha256:1f25e1b571a8be2c3d60d46679ef2471ae565f7ba9ba8382596695413523b188 \ - --hash=sha256:20e385cb8bd8321593788f11101d8c89a823a56191978e427e3c5141e129f14b \ - --hash=sha256:2ae50a175fb7b900f7aa42575f4fe99c32ca0ff57e5a8c1fd25e1243e67409db \ - --hash=sha256:2b45914f25d916324531d0259072b402c5f99b67c6e9ac8cbc3d49935aeb1d97 \ - --hash=sha256:38df5929ffefa9dda120ba1790a2e94fda81916c5aaa1ee652f4b1e515ebb9ed \ - --hash=sha256:3f21e503c18157512d2e34ae4c301e44a826c7b87e1d8998981367e3c9fe0932 \ - --hash=sha256:43461e439ef71d49bb0ee3a1719494cd952a58d205496698e0cde866f22006bc \ - --hash=sha256:434a1d1547a0547164866f1ccc31bbda235ac5b9087f24a84956756b52371f40 \ - --hash=sha256:447993c4dda0b6b0e1bd862752c855df8745f2910bea5015344f83ff3e99f305 \ - --hash=sha256:5c9e32989df06c57f10aa09ad9b30e8a25baf1aefe850e13b0ea5de600477d6a \ - --hash=sha256:61e08d84e3bf8ca9f43dc6b33f8cd7ba19f49864e2c91eb2160f83b6f9a268fa \ - --hash=sha256:699d26ac579eb42c71d131f9fb7b6e1c495a14e257264206a3c3bfcc146ed9bb \ - --hash=sha256:71ebdaadf546d6d393b9a21796172723724b737e84f68f36caf367d1c87a86a1 \ - --hash=sha256:7f5c05bd4b0909b682608c453acc31f1a9170d55f56d27cd701213e0683fc66a \ - --hash=sha256:848c5b040d2cfe35097b1d65d1095d83a3f86374ce879e189533f61405d8763b \ - --hash=sha256:8fe3caea61427057a9e3697c69b2403510fdccfca4483520d02b98ffae74531e \ - --hash=sha256:91ed5b71f9179bf3dbfe85d92b52d4b53de2e559aa4daa3b7de18e0dd24ad77d \ - --hash=sha256:a46f48740584eab3194fbee91c61f7fa396dbb1c5e7aa76ca08165d4e63fb40f \ - --hash=sha256:ab26b4af13308b8296688b03d74c3b0c8e8ed1f6b2d1454ef97bdb589db409db \ - --hash=sha256:b1b98f0a4137d01b84c680813eef6198e1e00f1f28bc20ce7b5c436459a0d146 \ - --hash=sha256:b1d179bdefd9ddb8d11d7de7825e73fb957511b722a8cb484e417885c210e68c \ - --hash=sha256:c3d2d5df5476b065aae9d1ad551fdc7b17c151b84e8edd9212108946b2337c66 \ - --hash=sha256:c4be1e5d9c8ad61345730c41c9ef21bdbb022cced4df70431110888d3ad5c0fb \ - --hash=sha256:da091dd8c96dbda124d766231f38619afd5c544051fb4424d2566c905957d342 \ - --hash=sha256:de86400c8b60c7707665e63934a82ae6792e7102c17a72e9b361a7f40d3c6049 \ - --hash=sha256:e365850166729fa82be618f476966161d5c47ea081eafc4febfc542bc85bac5d \ - --hash=sha256:e86c7fbe46f6e2e9dfb5377ee690fb8987e8e8363f435886ab91012b88f08a26 \ - --hash=sha256:fe6080299a25fd7cbb1957c921cca6a884acbfcd44cc23de48079389d322e326 -msal==1.32.0 \ - --hash=sha256:5445fe3af1da6be484991a7ab32eaa82461dc2347de105b76af92c610c3335c2 \ - --hash=sha256:9dbac5384a10bbbf4dae5c7ea0d707d14e087b92c5aa4954b3feaa2d1aa0bcb7 - # via - # azure-identity - # msal-extensions -msal-extensions==1.3.1 \ - --hash=sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca \ - --hash=sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4 - # via azure-identity -proto-plus==1.26.1 \ - --hash=sha256:13285478c2dcf2abb829db158e1047e2f1e8d63a077d94263c2b88b043c75a66 \ - --hash=sha256:21a515a4c4c0088a773899e23c7bbade3d18f9c66c73edd4c7ee3816bc96a012 - # via google-api-core -protobuf==5.29.3 \ - --hash=sha256:0a18ed4a24198528f2333802eb075e59dea9d679ab7a6c5efb017a59004d849f \ - --hash=sha256:0eb32bfa5219fc8d4111803e9a690658aa2e6366384fd0851064b963b6d1f2a7 \ - --hash=sha256:3ea51771449e1035f26069c4c7fd51fba990d07bc55ba80701c78f886bf9c888 \ - --hash=sha256:5da0f41edaf117bde316404bad1a486cb4ededf8e4a54891296f648e8e076620 \ - --hash=sha256:6ce8cc3389a20693bfde6c6562e03474c40851b44975c9b2bf6df7d8c4f864da \ - --hash=sha256:84a57163a0ccef3f96e4b6a20516cedcf5bb3a95a657131c5c3ac62200d23252 \ - --hash=sha256:a4fa6f80816a9a0678429e84973f2f98cbc218cca434abe8db2ad0bffc98503a \ - --hash=sha256:a8434404bbf139aa9e1300dbf989667a83d42ddda9153d8ab76e0d5dcaca484e \ - --hash=sha256:b89c115d877892a512f79a8114564fb435943b59067615894c3b13cd3e1fa107 \ - --hash=sha256:c027e08a08be10b67c06bf2370b99c811c466398c357e615ca88c91c07f0910f \ - --hash=sha256:daaf63f70f25e8689c072cfad4334ca0ac1d1e05a92fc15c54eb9cf23c3efd84 - # via - # google-api-core - # googleapis-common-protos - # proto-plus -pyasn1==0.6.1 \ - --hash=sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629 \ - --hash=sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034 - # via - # pyasn1-modules - # rsa -pyasn1-modules==0.4.1 \ - --hash=sha256:49bfa96b45a292b711e986f222502c1c9a5e1f4e568fc30e2574a6c7d07838fd \ - --hash=sha256:c28e2dbf9c06ad61c71a075c7e0f9fd0f1b0bb2d2ad4377f240d33ac2ab60a7c - # via google-auth -pycparser==2.22 \ - --hash=sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6 \ - --hash=sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc - # via cffi -pyjwt[crypto]==2.10.1 \ - --hash=sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953 \ - --hash=sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb - # via - # msal - # pyjwt -python-dateutil==2.9.0.post0 \ - --hash=sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3 \ - --hash=sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427 - # via - # barman - # botocore -python-snappy==0.7.3 \ - --hash=sha256:074c0636cfcd97e7251330f428064050ac81a52c62ed884fc2ddebbb60ed7f50 \ - --hash=sha256:40216c1badfb2d38ac781ecb162a1d0ec40f8ee9747e610bcfefdfa79486cee3 -requests==2.32.3 \ - --hash=sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760 \ - --hash=sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6 - # via - # azure-core - # google-api-core - # google-cloud-storage - # msal -rsa==4.9 \ - --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \ - --hash=sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21 - # via google-auth -s3transfer==0.10.4 \ - --hash=sha256:244a76a24355363a68164241438de1b72f8781664920260c48465896b712a41e \ - --hash=sha256:29edc09801743c21eb5ecbc617a152df41d3c287f67b615f73e5f750583666a7 - # via boto3 -six==1.17.0 \ - --hash=sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274 \ - --hash=sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81 - # via - # azure-core - # python-dateutil -typing-extensions==4.12.2 \ - --hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \ - --hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8 - # via - # azure-core - # azure-identity - # azure-storage-blob -urllib3==2.3.0 \ - --hash=sha256:1cee9ad369867bfdbbb48b7dd50374c0967a0bb7710050facf0dd6911440e3df \ - --hash=sha256:f8c5449b3cf0861679ce7e0503c7b44b5ec981bec0d1d3795a07f1ba96f0204d - # via - # botocore - # requests -zstandard==0.23.0 \ - --hash=sha256:034b88913ecc1b097f528e42b539453fa82c3557e414b3de9d5632c80439a473 \ - --hash=sha256:0a7f0804bb3799414af278e9ad51be25edf67f78f916e08afdb983e74161b916 \ - --hash=sha256:11e3bf3c924853a2d5835b24f03eeba7fc9b07d8ca499e247e06ff5676461a15 \ - --hash=sha256:12a289832e520c6bd4dcaad68e944b86da3bad0d339ef7989fb7e88f92e96072 \ - --hash=sha256:1516c8c37d3a053b01c1c15b182f3b5f5eef19ced9b930b684a73bad121addf4 \ - --hash=sha256:157e89ceb4054029a289fb504c98c6a9fe8010f1680de0201b3eb5dc20aa6d9e \ - --hash=sha256:1bfe8de1da6d104f15a60d4a8a768288f66aa953bbe00d027398b93fb9680b26 \ - --hash=sha256:1e172f57cd78c20f13a3415cc8dfe24bf388614324d25539146594c16d78fcc8 \ - --hash=sha256:1fd7e0f1cfb70eb2f95a19b472ee7ad6d9a0a992ec0ae53286870c104ca939e5 \ - --hash=sha256:203d236f4c94cd8379d1ea61db2fce20730b4c38d7f1c34506a31b34edc87bdd \ - --hash=sha256:27d3ef2252d2e62476389ca8f9b0cf2bbafb082a3b6bfe9d90cbcbb5529ecf7c \ - --hash=sha256:29a2bc7c1b09b0af938b7a8343174b987ae021705acabcbae560166567f5a8db \ - --hash=sha256:2ef230a8fd217a2015bc91b74f6b3b7d6522ba48be29ad4ea0ca3a3775bf7dd5 \ - --hash=sha256:2ef3775758346d9ac6214123887d25c7061c92afe1f2b354f9388e9e4d48acfc \ - --hash=sha256:2f146f50723defec2975fb7e388ae3a024eb7151542d1599527ec2aa9cacb152 \ - --hash=sha256:2fb4535137de7e244c230e24f9d1ec194f61721c86ebea04e1581d9d06ea1269 \ - --hash=sha256:32ba3b5ccde2d581b1e6aa952c836a6291e8435d788f656fe5976445865ae045 \ - --hash=sha256:34895a41273ad33347b2fc70e1bff4240556de3c46c6ea430a7ed91f9042aa4e \ - --hash=sha256:379b378ae694ba78cef921581ebd420c938936a153ded602c4fea612b7eaa90d \ - --hash=sha256:38302b78a850ff82656beaddeb0bb989a0322a8bbb1bf1ab10c17506681d772a \ - --hash=sha256:3aa014d55c3af933c1315eb4bb06dd0459661cc0b15cd61077afa6489bec63bb \ - --hash=sha256:4051e406288b8cdbb993798b9a45c59a4896b6ecee2f875424ec10276a895740 \ - --hash=sha256:40b33d93c6eddf02d2c19f5773196068d875c41ca25730e8288e9b672897c105 \ - --hash=sha256:43da0f0092281bf501f9c5f6f3b4c975a8a0ea82de49ba3f7100e64d422a1274 \ - --hash=sha256:445e4cb5048b04e90ce96a79b4b63140e3f4ab5f662321975679b5f6360b90e2 \ - --hash=sha256:48ef6a43b1846f6025dde6ed9fee0c24e1149c1c25f7fb0a0585572b2f3adc58 \ - --hash=sha256:50a80baba0285386f97ea36239855f6020ce452456605f262b2d33ac35c7770b \ - --hash=sha256:519fbf169dfac1222a76ba8861ef4ac7f0530c35dd79ba5727014613f91613d4 \ - --hash=sha256:53dd9d5e3d29f95acd5de6802e909ada8d8d8cfa37a3ac64836f3bc4bc5512db \ - --hash=sha256:53ea7cdc96c6eb56e76bb06894bcfb5dfa93b7adcf59d61c6b92674e24e2dd5e \ - --hash=sha256:576856e8594e6649aee06ddbfc738fec6a834f7c85bf7cadd1c53d4a58186ef9 \ - --hash=sha256:59556bf80a7094d0cfb9f5e50bb2db27fefb75d5138bb16fb052b61b0e0eeeb0 \ - --hash=sha256:5d41d5e025f1e0bccae4928981e71b2334c60f580bdc8345f824e7c0a4c2a813 \ - --hash=sha256:61062387ad820c654b6a6b5f0b94484fa19515e0c5116faf29f41a6bc91ded6e \ - --hash=sha256:61f89436cbfede4bc4e91b4397eaa3e2108ebe96d05e93d6ccc95ab5714be512 \ - --hash=sha256:62136da96a973bd2557f06ddd4e8e807f9e13cbb0bfb9cc06cfe6d98ea90dfe0 \ - --hash=sha256:64585e1dba664dc67c7cdabd56c1e5685233fbb1fc1966cfba2a340ec0dfff7b \ - --hash=sha256:65308f4b4890aa12d9b6ad9f2844b7ee42c7f7a4fd3390425b242ffc57498f48 \ - --hash=sha256:66b689c107857eceabf2cf3d3fc699c3c0fe8ccd18df2219d978c0283e4c508a \ - --hash=sha256:6a41c120c3dbc0d81a8e8adc73312d668cd34acd7725f036992b1b72d22c1772 \ - --hash=sha256:6f77fa49079891a4aab203d0b1744acc85577ed16d767b52fc089d83faf8d8ed \ - --hash=sha256:72c68dda124a1a138340fb62fa21b9bf4848437d9ca60bd35db36f2d3345f373 \ - --hash=sha256:752bf8a74412b9892f4e5b58f2f890a039f57037f52c89a740757ebd807f33ea \ - --hash=sha256:76e79bc28a65f467e0409098fa2c4376931fd3207fbeb6b956c7c476d53746dd \ - --hash=sha256:774d45b1fac1461f48698a9d4b5fa19a69d47ece02fa469825b442263f04021f \ - --hash=sha256:77da4c6bfa20dd5ea25cbf12c76f181a8e8cd7ea231c673828d0386b1740b8dc \ - --hash=sha256:77ea385f7dd5b5676d7fd943292ffa18fbf5c72ba98f7d09fc1fb9e819b34c23 \ - --hash=sha256:80080816b4f52a9d886e67f1f96912891074903238fe54f2de8b786f86baded2 \ - --hash=sha256:80a539906390591dd39ebb8d773771dc4db82ace6372c4d41e2d293f8e32b8db \ - --hash=sha256:82d17e94d735c99621bf8ebf9995f870a6b3e6d14543b99e201ae046dfe7de70 \ - --hash=sha256:837bb6764be6919963ef41235fd56a6486b132ea64afe5fafb4cb279ac44f259 \ - --hash=sha256:84433dddea68571a6d6bd4fbf8ff398236031149116a7fff6f777ff95cad3df9 \ - --hash=sha256:8c24f21fa2af4bb9f2c492a86fe0c34e6d2c63812a839590edaf177b7398f700 \ - --hash=sha256:8ed7d27cb56b3e058d3cf684d7200703bcae623e1dcc06ed1e18ecda39fee003 \ - --hash=sha256:9206649ec587e6b02bd124fb7799b86cddec350f6f6c14bc82a2b70183e708ba \ - --hash=sha256:983b6efd649723474f29ed42e1467f90a35a74793437d0bc64a5bf482bedfa0a \ - --hash=sha256:98da17ce9cbf3bfe4617e836d561e433f871129e3a7ac16d6ef4c680f13a839c \ - --hash=sha256:9c236e635582742fee16603042553d276cca506e824fa2e6489db04039521e90 \ - --hash=sha256:9da6bc32faac9a293ddfdcb9108d4b20416219461e4ec64dfea8383cac186690 \ - --hash=sha256:a05e6d6218461eb1b4771d973728f0133b2a4613a6779995df557f70794fd60f \ - --hash=sha256:a0817825b900fcd43ac5d05b8b3079937073d2b1ff9cf89427590718b70dd840 \ - --hash=sha256:a4ae99c57668ca1e78597d8b06d5af837f377f340f4cce993b551b2d7731778d \ - --hash=sha256:a8c86881813a78a6f4508ef9daf9d4995b8ac2d147dcb1a450448941398091c9 \ - --hash=sha256:a8fffdbd9d1408006baaf02f1068d7dd1f016c6bcb7538682622c556e7b68e35 \ - --hash=sha256:a9b07268d0c3ca5c170a385a0ab9fb7fdd9f5fd866be004c4ea39e44edce47dd \ - --hash=sha256:ab19a2d91963ed9e42b4e8d77cd847ae8381576585bad79dbd0a8837a9f6620a \ - --hash=sha256:ac184f87ff521f4840e6ea0b10c0ec90c6b1dcd0bad2f1e4a9a1b4fa177982ea \ - --hash=sha256:b0e166f698c5a3e914947388c162be2583e0c638a4703fc6a543e23a88dea3c1 \ - --hash=sha256:b2170c7e0367dde86a2647ed5b6f57394ea7f53545746104c6b09fc1f4223573 \ - --hash=sha256:b2d8c62d08e7255f68f7a740bae85b3c9b8e5466baa9cbf7f57f1cde0ac6bc09 \ - --hash=sha256:b4567955a6bc1b20e9c31612e615af6b53733491aeaa19a6b3b37f3b65477094 \ - --hash=sha256:b69bb4f51daf461b15e7b3db033160937d3ff88303a7bc808c67bbc1eaf98c78 \ - --hash=sha256:b8c0bd73aeac689beacd4e7667d48c299f61b959475cdbb91e7d3d88d27c56b9 \ - --hash=sha256:be9b5b8659dff1f913039c2feee1aca499cfbc19e98fa12bc85e037c17ec6ca5 \ - --hash=sha256:bf0a05b6059c0528477fba9054d09179beb63744355cab9f38059548fedd46a9 \ - --hash=sha256:c16842b846a8d2a145223f520b7e18b57c8f476924bda92aeee3a88d11cfc391 \ - --hash=sha256:c363b53e257246a954ebc7c488304b5592b9c53fbe74d03bc1c64dda153fb847 \ - --hash=sha256:c7c517d74bea1a6afd39aa612fa025e6b8011982a0897768a2f7c8ab4ebb78a2 \ - --hash=sha256:d20fd853fbb5807c8e84c136c278827b6167ded66c72ec6f9a14b863d809211c \ - --hash=sha256:d2240ddc86b74966c34554c49d00eaafa8200a18d3a5b6ffbf7da63b11d74ee2 \ - --hash=sha256:d477ed829077cd945b01fc3115edd132c47e6540ddcd96ca169facff28173057 \ - --hash=sha256:d50d31bfedd53a928fed6707b15a8dbeef011bb6366297cc435accc888b27c20 \ - --hash=sha256:dc1d33abb8a0d754ea4763bad944fd965d3d95b5baef6b121c0c9013eaf1907d \ - --hash=sha256:dc5d1a49d3f8262be192589a4b72f0d03b72dcf46c51ad5852a4fdc67be7b9e4 \ - --hash=sha256:e2d1a054f8f0a191004675755448d12be47fa9bebbcffa3cdf01db19f2d30a54 \ - --hash=sha256:e7792606d606c8df5277c32ccb58f29b9b8603bf83b48639b7aedf6df4fe8171 \ - --hash=sha256:ed1708dbf4d2e3a1c5c69110ba2b4eb6678262028afd6c6fbcc5a8dac9cda68e \ - --hash=sha256:f2d4380bf5f62daabd7b751ea2339c1a21d1c9463f1feb7fc2bdcea2c29c3160 \ - --hash=sha256:f3513916e8c645d0610815c257cbfd3242adfd5c4cfa78be514e5a3ebb42a41b \ - --hash=sha256:f8346bfa098532bc1fb6c7ef06783e969d87a99dd1d2a5a18a892c1d7a643c58 \ - --hash=sha256:f83fa6cae3fff8e98691248c9320356971b59678a17f20656a9e59cd32cee6d8 \ - --hash=sha256:fa6ce8b52c5987b3e34d5674b0ab529a4602b632ebab0a93b07bfb4dfc8f8a33 \ - --hash=sha256:fb2b1ecfef1e67897d336de3a0e3f52478182d6a47eda86cbd42504c5cbd009a \ - --hash=sha256:fc9ca1c9718cb3b06634c7c8dec57d24e9438b2aa9a0f02b8bb36bf478538880 \ - --hash=sha256:fd30d9c67d13d891f2360b2a120186729c111238ac63b43dbd37a5a40670b8ca \ - --hash=sha256:fd7699e8fd9969f455ef2926221e0233f81a2542921471382e77a9e2f2b57f4b \ - --hash=sha256:fe3b385d996ee0822fd46528d9f0443b880d4d05528fd26a9119a54ec3f91c69 \ No newline at end of file diff --git a/docker/pg-slim/Dockerfile b/docker/pg-slim/Dockerfile deleted file mode 100644 index 89b1629f..00000000 --- a/docker/pg-slim/Dockerfile +++ /dev/null @@ -1,142 +0,0 @@ -FROM ubuntu:22.04 - -# 14, 15, 16, 17 -ARG PG_MAJOR -ARG ALTDIR=/var/lib/postgresql/data/tensorchord -ENV DEBIAN_FRONTEND=noninteractive \ - TZ=Etc/UTC \ - PGDATA=/var/lib/postgresql/data - -# Get latest package updates -RUN set -eux; \ - apt-get update; \ - apt-get upgrade -y - -# explicitly set user/group IDs -RUN set -eux; \ - groupadd -r postgres --gid=999; \ - # https://salsa.debian.org/postgresql/postgresql-common/blob/997d842ee744687d99a2b2d95c1083a2615c79e8/debian/postgresql-common.postinst#L32-35 - useradd -r -g postgres --uid=999 --home-dir=/var/lib/postgresql --shell=/bin/bash postgres; \ - # also create the postgres user's home directory with appropriate permissions - # see https://github.com/docker-library/postgres/issues/274 - install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql - -RUN set -eux; \ - apt-get update; \ - apt-get install -y --no-install-recommends \ - locales curl ca-certificates gnupg lsb-release lbzip2 git cmake \ - ; \ - rm -rf /var/lib/apt/lists/*; \ - localedef -i en_US -c -f UTF-8 -A /usr/share/locale/locale.alias en_US.UTF-8 -ENV LANG en_US.utf8 - -RUN mkdir /docker-entrypoint-initdb.d - -RUN git clone https://github.com/postgres/postgres.git && \ - cd postgres && \ - PG_RELEASE=$(git tag | grep '^REL_'${PG_MAJOR}'_' | grep -vE 'RC|BETA' | sort -V | tail -1 | sed 's/REL_//') && \ - echo $PG_RELEASE && \ - git checkout REL_$PG_RELEASE - -RUN set -eux; \ - apt-get update && apt-get install -y \ - libreadline-dev \ - zlib1g-dev \ - libpq-dev \ - build-essential \ - python3-dev \ - tcl-dev \ - libxslt1-dev \ - libperl-dev \ - libpam0g-dev \ - libreadline-dev \ - libssl-dev \ - xz-utils \ - libnss-wrapper \ - llvm \ - clang \ - icu-devtools \ - pkg-config \ - libgss-dev \ - libkrb5-dev \ - uuid-dev \ - gettext \ - liblz4-dev \ - libsystemd-dev \ - libselinux1-dev \ - libzstd-dev \ - vim \ - flex \ - bison; \ - apt-get autoremove -y; \ - apt-get clean -y; \ - rm -rf /var/lib/apt/lists/* - -WORKDIR postgres - -ENV CFLAGS "-g -O2 -fstack-protector-strong -Wformat -Werror=format-security -fno-omit-frame-pointer" -ENV LDFLAGS "-Wl,-z,relro -Wl,-z,now" -RUN ./configure --prefix=/usr/lib/postgresql/${PG_MAJOR} \ - --datarootdir=${ALTDIR} \ - --libdir=${ALTDIR}/${PG_MAJOR}/lib \ - --with-perl \ - --with-python \ - --with-tcl \ - --with-pam \ - --with-libxml \ - --with-libxslt \ - --with-openssl \ - --enable-nls \ - --enable-thread-safety \ - --enable-debug \ - --disable-rpath \ - --with-uuid=e2fs \ - --with-gnu-ld \ - --with-gssapi \ - --with-pgport=5432 \ - --with-system-tzdata=/usr/share/zoneinfo \ - --with-icu \ - --with-llvm \ - --with-lz4 \ - --with-zstd \ - --with-systemd \ - --with-selinux - -RUN make -j$(nproc) -RUN make install - -# Remove libpq-dev provided from Ubuntu repos and set the newly-compiled one to system path -RUN set -eux; \ - apt-get purge -y libpq-dev && \ - apt-get autoremove -y && \ - rm -rf /var/lib/apt/lists/* - -RUN echo "${ALTDIR}/${PG_MAJOR}/lib" > /etc/ld.so.conf.d/postgres.conf && \ - ldconfig - -WORKDIR / -RUN rm -rf /postgres - -RUN mkdir -p /var/run/postgresql && chmod 775 /var/run/postgresql -RUN mkdir -p /usr/share/postgresql/${PG_MAJOR}/extension && chmod 775 /usr/share/postgresql/${PG_MAJOR}/extension - -COPY --from=tianon/gosu /gosu /usr/local/bin/ - -# make the sample config easier to munge (and "correct by default") -RUN set -eux; \ - sed -ri "s!^#?(listen_addresses)\s*=\s*\S+.*!\1 = '*'!" ${ALTDIR}/postgresql.conf.sample; \ - grep -F "listen_addresses = '*'" ${ALTDIR}/postgresql.conf.sample - -RUN install --verbose --directory --owner postgres --group postgres --mode 3777 /var/run/postgresql - -# this 1777 will be replaced by 0700 at runtime (allows semi-arbitrary "--user" values) -RUN install --verbose --directory --owner postgres --group postgres --mode 1777 ${PGDATA} - -ENV PATH $PATH:/usr/lib/postgresql/$PG_MAJOR/bin:/usr/local/bin -COPY docker-entrypoint.sh docker-ensure-initdb.sh /usr/local/bin/ -RUN ln -sT docker-ensure-initdb.sh /usr/local/bin/docker-enforce-initdb.sh -ENTRYPOINT ["docker-entrypoint.sh"] - -STOPSIGNAL SIGINT -EXPOSE 5432 -CMD ["postgres"] \ No newline at end of file diff --git a/docker/pg-slim/docker-ensure-initdb.sh b/docker/pg-slim/docker-ensure-initdb.sh deleted file mode 100755 index ae1f6b6b..00000000 --- a/docker/pg-slim/docker-ensure-initdb.sh +++ /dev/null @@ -1,71 +0,0 @@ -#!/usr/bin/env bash -set -Eeuo pipefail - -# -# This script is intended for three main use cases: -# -# 1. (most importantly) as an example of how to use "docker-entrypoint.sh" to extend/reuse the initialization behavior -# -# 2. ("docker-ensure-initdb.sh") as a Kubernetes "init container" to ensure the provided database directory is initialized; see also "startup probes" for an alternative solution -# (no-op if database is already initialized) -# -# 3. ("docker-enforce-initdb.sh") as part of CI to ensure the database is fully initialized before use -# (error if database is already initialized) -# - -source /usr/local/bin/docker-entrypoint.sh - -# arguments to this script are assumed to be arguments to the "postgres" server (same as "docker-entrypoint.sh"), and most "docker-entrypoint.sh" functions assume "postgres" is the first argument (see "_main" over there) -if [ "$#" -eq 0 ] || [ "$1" != 'postgres' ]; then - set -- postgres "$@" -fi - -# see also "_main" in "docker-entrypoint.sh" - -docker_setup_env -# setup data directories and permissions (when run as root) -docker_create_db_directories -if [ "$(id -u)" = '0' ]; then - # then restart script as postgres user - exec gosu postgres "$BASH_SOURCE" "$@" -fi - -# only run initialization on an empty data directory -if [ -z "$DATABASE_ALREADY_EXISTS" ]; then - docker_verify_minimum_env - - # check dir permissions to reduce likelihood of half-initialized database - ls /docker-entrypoint-initdb.d/ > /dev/null - - docker_init_database_dir - pg_setup_hba_conf "$@" - - # PGPASSWORD is required for psql when authentication is required for 'local' connections via pg_hba.conf and is otherwise harmless - # e.g. when '--auth=md5' or '--auth-local=md5' is used in POSTGRES_INITDB_ARGS - export PGPASSWORD="${PGPASSWORD:-$POSTGRES_PASSWORD}" - docker_temp_server_start "$@" - - docker_setup_db - docker_process_init_files /docker-entrypoint-initdb.d/* - - docker_temp_server_stop - unset PGPASSWORD -else - self="$(basename "$0")" - case "$self" in - docker-ensure-initdb.sh) - echo >&2 "$self: note: database already initialized in '$PGDATA'!" - exit 0 - ;; - - docker-enforce-initdb.sh) - echo >&2 "$self: error: (unexpected) database found in '$PGDATA'!" - exit 1 - ;; - - *) - echo >&2 "$self: error: unknown file name: $self" - exit 99 - ;; - esac -fi diff --git a/docker/pg-slim/docker-entrypoint.sh b/docker/pg-slim/docker-entrypoint.sh deleted file mode 100755 index 6f59993e..00000000 --- a/docker/pg-slim/docker-entrypoint.sh +++ /dev/null @@ -1,356 +0,0 @@ -#!/usr/bin/env bash -set -Eeo pipefail -# TODO swap to -Eeuo pipefail above (after handling all potentially-unset variables) - -# usage: file_env VAR [DEFAULT] -# ie: file_env 'XYZ_DB_PASSWORD' 'example' -# (will allow for "$XYZ_DB_PASSWORD_FILE" to fill in the value of -# "$XYZ_DB_PASSWORD" from a file, especially for Docker's secrets feature) -file_env() { - local var="$1" - local fileVar="${var}_FILE" - local def="${2:-}" - if [ "${!var:-}" ] && [ "${!fileVar:-}" ]; then - printf >&2 'error: both %s and %s are set (but are exclusive)\n' "$var" "$fileVar" - exit 1 - fi - local val="$def" - if [ "${!var:-}" ]; then - val="${!var}" - elif [ "${!fileVar:-}" ]; then - val="$(< "${!fileVar}")" - fi - export "$var"="$val" - unset "$fileVar" -} - -# check to see if this file is being run or sourced from another script -_is_sourced() { - # https://unix.stackexchange.com/a/215279 - [ "${#FUNCNAME[@]}" -ge 2 ] \ - && [ "${FUNCNAME[0]}" = '_is_sourced' ] \ - && [ "${FUNCNAME[1]}" = 'source' ] -} - -# used to create initial postgres directories and if run as root, ensure ownership to the "postgres" user -docker_create_db_directories() { - local user; user="$(id -u)" - - mkdir -p "$PGDATA" - # ignore failure since there are cases where we can't chmod (and PostgreSQL might fail later anyhow - it's picky about permissions of this directory) - chmod 00700 "$PGDATA" || : - - # ignore failure since it will be fine when using the image provided directory; see also https://github.com/docker-library/postgres/pull/289 - mkdir -p /var/run/postgresql || : - chmod 03775 /var/run/postgresql || : - - # Create the transaction log directory before initdb is run so the directory is owned by the correct user - if [ -n "${POSTGRES_INITDB_WALDIR:-}" ]; then - mkdir -p "$POSTGRES_INITDB_WALDIR" - if [ "$user" = '0' ]; then - find "$POSTGRES_INITDB_WALDIR" \! -user postgres -exec chown postgres '{}' + - fi - chmod 700 "$POSTGRES_INITDB_WALDIR" - fi - - # allow the container to be started with `--user` - if [ "$user" = '0' ]; then - find "$PGDATA" \! -user postgres -exec chown postgres '{}' + - find /var/run/postgresql \! -user postgres -exec chown postgres '{}' + - fi -} - -# initialize empty PGDATA directory with new database via 'initdb' -# arguments to `initdb` can be passed via POSTGRES_INITDB_ARGS or as arguments to this function -# `initdb` automatically creates the "postgres", "template0", and "template1" dbnames -# this is also where the database user is created, specified by `POSTGRES_USER` env -docker_init_database_dir() { - # "initdb" is particular about the current user existing in "/etc/passwd", so we use "nss_wrapper" to fake that if necessary - # see https://github.com/docker-library/postgres/pull/253, https://github.com/docker-library/postgres/issues/359, https://cwrap.org/nss_wrapper.html - local uid; uid="$(id -u)" - if ! getent passwd "$uid" &> /dev/null; then - # see if we can find a suitable "libnss_wrapper.so" (https://salsa.debian.org/sssd-team/nss-wrapper/-/commit/b9925a653a54e24d09d9b498a2d913729f7abb15) - local wrapper - for wrapper in {/usr,}/lib{/*,}/libnss_wrapper.so; do - if [ -s "$wrapper" ]; then - NSS_WRAPPER_PASSWD="$(mktemp)" - NSS_WRAPPER_GROUP="$(mktemp)" - export LD_PRELOAD="$wrapper" NSS_WRAPPER_PASSWD NSS_WRAPPER_GROUP - local gid; gid="$(id -g)" - printf 'postgres:x:%s:%s:PostgreSQL:%s:/bin/false\n' "$uid" "$gid" "$PGDATA" > "$NSS_WRAPPER_PASSWD" - printf 'postgres:x:%s:\n' "$gid" > "$NSS_WRAPPER_GROUP" - break - fi - done - fi - - if [ -n "${POSTGRES_INITDB_WALDIR:-}" ]; then - set -- --waldir "$POSTGRES_INITDB_WALDIR" "$@" - fi - - # --pwfile refuses to handle a properly-empty file (hence the "\n"): https://github.com/docker-library/postgres/issues/1025 - eval 'initdb --username="$POSTGRES_USER" --pwfile=<(printf "%s\n" "$POSTGRES_PASSWORD") '"$POSTGRES_INITDB_ARGS"' "$@"' - - # unset/cleanup "nss_wrapper" bits - if [[ "${LD_PRELOAD:-}" == */libnss_wrapper.so ]]; then - rm -f "$NSS_WRAPPER_PASSWD" "$NSS_WRAPPER_GROUP" - unset LD_PRELOAD NSS_WRAPPER_PASSWD NSS_WRAPPER_GROUP - fi -} - -# print large warning if POSTGRES_PASSWORD is long -# error if both POSTGRES_PASSWORD is empty and POSTGRES_HOST_AUTH_METHOD is not 'trust' -# print large warning if POSTGRES_HOST_AUTH_METHOD is set to 'trust' -# assumes database is not set up, ie: [ -z "$DATABASE_ALREADY_EXISTS" ] -docker_verify_minimum_env() { - case "${PG_MAJOR:-}" in - 12 | 13) # https://github.com/postgres/postgres/commit/67a472d71c98c3d2fa322a1b4013080b20720b98 - # check password first so we can output the warning before postgres - # messes it up - if [ "${#POSTGRES_PASSWORD}" -ge 100 ]; then - cat >&2 <<-'EOWARN' - - WARNING: The supplied POSTGRES_PASSWORD is 100+ characters. - - This will not work if used via PGPASSWORD with "psql". - - https://www.postgresql.org/message-id/flat/E1Rqxp2-0004Qt-PL%40wrigleys.postgresql.org (BUG #6412) - https://github.com/docker-library/postgres/issues/507 - - EOWARN - fi - ;; - esac - if [ -z "$POSTGRES_PASSWORD" ] && [ 'trust' != "$POSTGRES_HOST_AUTH_METHOD" ]; then - # The - option suppresses leading tabs but *not* spaces. :) - cat >&2 <<-'EOE' - Error: Database is uninitialized and superuser password is not specified. - You must specify POSTGRES_PASSWORD to a non-empty value for the - superuser. For example, "-e POSTGRES_PASSWORD=password" on "docker run". - - You may also use "POSTGRES_HOST_AUTH_METHOD=trust" to allow all - connections without a password. This is *not* recommended. - - See PostgreSQL documentation about "trust": - https://www.postgresql.org/docs/current/auth-trust.html - EOE - exit 1 - fi - if [ 'trust' = "$POSTGRES_HOST_AUTH_METHOD" ]; then - cat >&2 <<-'EOWARN' - ******************************************************************************** - WARNING: POSTGRES_HOST_AUTH_METHOD has been set to "trust". This will allow - anyone with access to the Postgres port to access your database without - a password, even if POSTGRES_PASSWORD is set. See PostgreSQL - documentation about "trust": - https://www.postgresql.org/docs/current/auth-trust.html - In Docker's default configuration, this is effectively any other - container on the same system. - - It is not recommended to use POSTGRES_HOST_AUTH_METHOD=trust. Replace - it with "-e POSTGRES_PASSWORD=password" instead to set a password in - "docker run". - ******************************************************************************** - EOWARN - fi -} - -# usage: docker_process_init_files [file [file [...]]] -# ie: docker_process_init_files /always-initdb.d/* -# process initializer files, based on file extensions and permissions -docker_process_init_files() { - # psql here for backwards compatibility "${psql[@]}" - psql=( docker_process_sql ) - - printf '\n' - local f - for f; do - case "$f" in - *.sh) - # https://github.com/docker-library/postgres/issues/450#issuecomment-393167936 - # https://github.com/docker-library/postgres/pull/452 - if [ -x "$f" ]; then - printf '%s: running %s\n' "$0" "$f" - "$f" - else - printf '%s: sourcing %s\n' "$0" "$f" - . "$f" - fi - ;; - *.sql) printf '%s: running %s\n' "$0" "$f"; docker_process_sql -f "$f"; printf '\n' ;; - *.sql.gz) printf '%s: running %s\n' "$0" "$f"; gunzip -c "$f" | docker_process_sql; printf '\n' ;; - *.sql.xz) printf '%s: running %s\n' "$0" "$f"; xzcat "$f" | docker_process_sql; printf '\n' ;; - *.sql.zst) printf '%s: running %s\n' "$0" "$f"; zstd -dc "$f" | docker_process_sql; printf '\n' ;; - *) printf '%s: ignoring %s\n' "$0" "$f" ;; - esac - printf '\n' - done -} - -# Execute sql script, passed via stdin (or -f flag of pqsl) -# usage: docker_process_sql [psql-cli-args] -# ie: docker_process_sql --dbname=mydb <<<'INSERT ...' -# ie: docker_process_sql -f my-file.sql -# ie: docker_process_sql > "$PGDATA/pg_hba.conf" -} - -# start socket-only postgresql server for setting up or running scripts -# all arguments will be passed along as arguments to `postgres` (via pg_ctl) -docker_temp_server_start() { - if [ "$1" = 'postgres' ]; then - shift - fi - - # internal start of server in order to allow setup using psql client - # does not listen on external TCP/IP and waits until start finishes - set -- "$@" -c listen_addresses='' -p "${PGPORT:-5432}" - - PGUSER="${PGUSER:-$POSTGRES_USER}" \ - pg_ctl -D "$PGDATA" \ - -o "$(printf '%q ' "$@")" \ - -w start -} - -# stop postgresql server after done setting up user and running scripts -docker_temp_server_stop() { - PGUSER="${PGUSER:-postgres}" \ - pg_ctl -D "$PGDATA" -m fast -w stop -} - -# check arguments for an option that would cause postgres to stop -# return true if there is one -_pg_want_help() { - local arg - for arg; do - case "$arg" in - # postgres --help | grep 'then exit' - # leaving out -C on purpose since it always fails and is unhelpful: - # postgres: could not access the server configuration file "/var/lib/postgresql/data/postgresql.conf": No such file or directory - -'?'|--help|--describe-config|-V|--version) - return 0 - ;; - esac - done - return 1 -} - -_main() { - # if first arg looks like a flag, assume we want to run postgres server - if [ "${1:0:1}" = '-' ]; then - set -- postgres "$@" - fi - - if [ "$1" = 'postgres' ] && ! _pg_want_help "$@"; then - docker_setup_env - # setup data directories and permissions (when run as root) - docker_create_db_directories - if [ "$(id -u)" = '0' ]; then - # then restart script as postgres user - exec gosu postgres "$BASH_SOURCE" "$@" - fi - - # only run initialization on an empty data directory - if [ -z "$DATABASE_ALREADY_EXISTS" ]; then - docker_verify_minimum_env - - # check dir permissions to reduce likelihood of half-initialized database - ls /docker-entrypoint-initdb.d/ > /dev/null - - docker_init_database_dir - pg_setup_hba_conf "$@" - - # PGPASSWORD is required for psql when authentication is required for 'local' connections via pg_hba.conf and is otherwise harmless - # e.g. when '--auth=md5' or '--auth-local=md5' is used in POSTGRES_INITDB_ARGS - export PGPASSWORD="${PGPASSWORD:-$POSTGRES_PASSWORD}" - docker_temp_server_start "$@" - - docker_setup_db - docker_process_init_files /docker-entrypoint-initdb.d/* - - docker_temp_server_stop - unset PGPASSWORD - - cat <<-'EOM' - - PostgreSQL init process complete; ready for start up. - - EOM - else - cat <<-'EOM' - - PostgreSQL Database directory appears to contain a database; Skipping initialization - - EOM - fi - fi - - exec "$@" -} - -if ! _is_sourced; then - _main "$@" -fi From 13f34e54077c0d94636511f4986447237fa24b29 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 9 Apr 2025 16:16:10 +0800 Subject: [PATCH 140/324] refactor: remove allows_skipping_rerank (#226) 1. rename `max_maxsim_tuples` to `maxsim_refine`; leave it a default value `0` 2. remove `allows_skipping_rerank`; use estimated distance for vectors whose distance is not calculated Signed-off-by: usamoi --- crates/algorithm/src/build.rs | 55 ++-- crates/algorithm/src/bulkdelete.rs | 4 +- crates/algorithm/src/insert.rs | 38 ++- crates/algorithm/src/lib.rs | 4 +- crates/algorithm/src/maintain.rs | 107 +++++-- crates/algorithm/src/operator.rs | 465 +++++++++++++---------------- crates/algorithm/src/rerank.rs | 125 ++++---- crates/algorithm/src/search.rs | 134 ++++----- crates/algorithm/src/tape.rs | 40 +-- crates/algorithm/src/tuples.rs | 30 +- crates/algorithm/src/vectors.rs | 8 +- crates/k_means/src/lib.rs | 6 +- crates/rabitq/src/binary.rs | 22 +- crates/rabitq/src/block.rs | 17 +- crates/rabitq/src/lib.rs | 2 +- src/datatype/typmod.rs | 8 +- src/index/algorithm.rs | 6 +- src/index/am/am_build.rs | 11 +- src/index/am/mod.rs | 80 +++-- src/index/gucs.rs | 42 +-- src/index/scanners/default.rs | 103 +++++-- src/index/scanners/maxsim.rs | 248 +++++++-------- src/index/scanners/mod.rs | 13 +- tests/logic/multivector.slt | 2 +- 24 files changed, 805 insertions(+), 765 deletions(-) diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 67fa3723..8bd619e2 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -2,7 +2,7 @@ use crate::operator::{Accessor2, Operator, Vector}; use crate::tape::TapeWriter; use crate::tuples::*; use crate::types::*; -use crate::{Branch, DerefMut, IndexPointer, Page, PageGuard, RelationWrite}; +use crate::{Branch, IndexPointer, RelationWrite}; use simd::fast_scan::{any_pack, padding_pack}; use vector::VectorOwned; @@ -12,18 +12,21 @@ pub fn build( index: impl RelationWrite, structures: Vec>, ) { + if vchordrq_options.residual_quantization && !O::SUPPORTS_RESIDUAL { + panic!("residual_quantization can be enabled only if distance type is L2"); + } let dims = vector_options.dims; - let is_residual = vchordrq_options.residual_quantization && O::SUPPORTS_RESIDUAL; - let mut meta = TapeWriter::<_, _, MetaTuple>::create(|| index.extend(false)); + let is_residual = vchordrq_options.residual_quantization; + let mut meta = TapeWriter::<_, MetaTuple>::create(&index, false); assert_eq!(meta.first(), 0); - let freepage = TapeWriter::<_, _, FreepageTuple>::create(|| index.extend(false)); - let mut vectors = TapeWriter::<_, _, VectorTuple>::create(|| index.extend(true)); + let freepage = TapeWriter::<_, FreepageTuple>::create(&index, false); + let mut vectors = TapeWriter::<_, VectorTuple>::create(&index, true); let mut pointer_of_means = Vec::>::new(); for i in 0..structures.len() { let mut level = Vec::new(); for j in 0..structures[i].len() { let vector = structures[i].means[j].as_borrowed(); - let (metadata, slices) = O::Vector::vector_split(vector); + let (slices, metadata) = O::Vector::split(vector); let mut chain = Ok(metadata); for i in (0..slices.len()).rev() { chain = Err(vectors.push(match chain { @@ -48,10 +51,9 @@ pub fn build( let mut level = Vec::new(); for j in 0..structures[i].len() { if i == 0 { - let frozen_tape = TapeWriter::<_, _, FrozenTuple>::create(|| index.extend(false)); - let appendable_tape = - TapeWriter::<_, _, AppendableTuple>::create(|| index.extend(false)); - let mut jump = TapeWriter::<_, _, JumpTuple>::create(|| index.extend(false)); + let frozen_tape = TapeWriter::<_, FrozenTuple>::create(&index, false); + let appendable_tape = TapeWriter::<_, AppendableTuple>::create(&index, false); + let mut jump = TapeWriter::<_, JumpTuple>::create(&index, false); jump.push(JumpTuple { frozen_first: frozen_tape.first(), appendable_first: appendable_tape.first(), @@ -59,21 +61,17 @@ pub fn build( }); level.push(jump.first()); } else { - let mut tape = H1TapeWriter::<_, _>::create(|| index.extend(false)); + let mut tape = H1TapeWriter::create(&index, false); let h2_mean = structures[i].means[j].as_borrowed(); let h2_children = structures[i].children[j].as_slice(); for child in h2_children.iter().copied() { let h1_mean = structures[i - 1].means[child as usize].as_borrowed(); let code = if is_residual { let mut residual_accessor = O::ResidualAccessor::default(); - residual_accessor.push( - O::Vector::elements_and_metadata(h1_mean).0, - O::Vector::elements_and_metadata(h2_mean).0, - ); - let residual = residual_accessor.finish( - O::Vector::elements_and_metadata(h1_mean).1, - O::Vector::elements_and_metadata(h2_mean).1, - ); + residual_accessor + .push(O::Vector::unpack(h1_mean).0, O::Vector::unpack(h2_mean).0); + let residual = residual_accessor + .finish(O::Vector::unpack(h1_mean).1, O::Vector::unpack(h2_mean).1); O::Vector::code(residual.as_borrowed()) } else { O::Vector::code(h1_mean) @@ -139,20 +137,21 @@ pub fn build( }); } -pub struct H1TapeWriter { - tape: TapeWriter, +pub struct H1TapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + tape: TapeWriter<'a, R, H1Tuple>, branches: Vec>, } -impl H1TapeWriter +impl<'a, R> H1TapeWriter<'a, R> where - G: PageGuard + DerefMut, - G::Target: Page, - E: Fn() -> G, + R: RelationWrite + 'a, { - fn create(extend: E) -> Self { + fn create(index: &'a R, tracking_freespace: bool) -> Self { Self { - tape: TapeWriter::create(extend), + tape: TapeWriter::create(index, tracking_freespace), branches: Vec::new(), } } @@ -188,7 +187,7 @@ where self.branches.clear(); } } - fn into_inner(self) -> (TapeWriter, Vec>) { + fn into_inner(self) -> (TapeWriter<'a, R, H1Tuple>, Vec>) { (self.tape, self.branches) } } diff --git a/crates/algorithm/src/bulkdelete.rs b/crates/algorithm/src/bulkdelete.rs index 506d2f05..c34882af 100644 --- a/crates/algorithm/src/bulkdelete.rs +++ b/crates/algorithm/src/bulkdelete.rs @@ -1,12 +1,12 @@ use crate::operator::{FunctionalAccessor, Operator}; use crate::tuples::*; use crate::{Page, RelationWrite, tape}; -use std::num::NonZeroU64; +use std::num::NonZero; pub fn bulkdelete( index: impl RelationWrite, check: impl Fn(), - callback: impl Fn(NonZeroU64) -> bool, + callback: impl Fn(NonZero) -> bool, ) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 63bb80a2..7b662da3 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -8,10 +8,10 @@ use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; -use std::num::NonZeroU64; +use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vector: O::Vector) { +pub fn insert(index: impl RelationWrite, payload: NonZero, vector: O::Vector) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -25,8 +25,8 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto let vectors_first = meta_tuple.vectors_first(); drop(meta_guard); - let default_lut_block = if !is_residual { - Some(O::Vector::compute_lut_block(vector.as_borrowed())) + let default_block_lut = if !is_residual { + Some(O::Vector::block_preprocess(vector.as_borrowed())) } else { None }; @@ -45,7 +45,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto index.clone(), mean, LAccess::new( - O::Vector::elements_and_metadata(vector.as_borrowed()), + O::Vector::unpack(vector.as_borrowed()), O::ResidualAccessor::default(), ), ); @@ -58,23 +58,19 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto let mut results = LinkedVec::new(); { let (first, residual) = state; - let lut_block = if let Some(residual) = residual { - &O::Vector::compute_lut_block(residual.as_borrowed()) - } else if let Some(lut_block) = default_lut_block.as_ref() { - lut_block + let block_lut = if let Some(residual) = residual { + &O::Vector::block_preprocess(residual.as_borrowed()) + } else if let Some(block_lut) = default_block_lut.as_ref() { + block_lut } else { unreachable!() }; tape::read_h1_tape( index.clone(), first, - || { - RAccess::new( - (&lut_block.1, (lut_block.0, 1.9f32)), - O::Distance::block_accessor(), - ) - }, - |lowerbound, mean, first| { + || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), + |(rough, err), mean, first| { + let lowerbound = Distance::from_f32(rough - err * 1.9); results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(first))); }, |_| (), @@ -83,15 +79,17 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto let mut heap = SelectHeap::from_vec(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); { - while let Some((_, AlwaysEqual(mean), AlwaysEqual(first))) = - pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) + while let Some((Reverse(_), AlwaysEqual(mean), AlwaysEqual(first))) = + pop_if(&mut heap, |(d, ..)| { + Some(*d) > cache.peek().map(|(d, ..)| *d) + }) { if is_residual { let (dis_u, residual_u) = vectors::read_for_h1_tuple::( index.clone(), mean, LAccess::new( - O::Vector::elements_and_metadata(vector.as_borrowed()), + O::Vector::unpack(vector.as_borrowed()), ( O::DistanceAccessor::default(), O::ResidualAccessor::default(), @@ -108,7 +106,7 @@ pub fn insert(index: impl RelationWrite, payload: NonZeroU64, vecto index.clone(), mean, LAccess::new( - O::Vector::elements_and_metadata(vector.as_borrowed()), + O::Vector::unpack(vector.as_borrowed()), O::DistanceAccessor::default(), ), ); diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 9c9eebbe..5cb12c37 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -22,8 +22,8 @@ pub use cache::cache; pub use insert::insert; pub use maintain::maintain; pub use prewarm::prewarm; -pub use rerank::{how, rerank_heap, rerank_index, skip}; -pub use search::{search, search_and_estimate}; +pub use rerank::{how, rerank_heap, rerank_index}; +pub use search::{default_search, maxsim_search}; use std::ops::{Deref, DerefMut}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index e3f2ddf1..4f9a7219 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -1,9 +1,9 @@ use crate::operator::{FunctionalAccessor, Operator}; use crate::tape::{self, TapeWriter}; use crate::tuples::*; -use crate::{Branch, DerefMut, Page, PageGuard, RelationWrite, freepages}; +use crate::{Branch, Page, RelationRead, RelationWrite, freepages}; use simd::fast_scan::{padding_pack, unpack}; -use std::num::NonZeroU64; +use std::num::NonZero; pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let meta_guard = index.read(0); @@ -48,17 +48,9 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let jump_bytes = jump_guard.get_mut(1).expect("data corruption"); let mut jump_tuple = JumpTuple::deserialize_mut(jump_bytes); - let expand = || { - if let Some(id) = freepages::fetch(index.clone(), freepage_first) { - let mut write = index.write(id, false); - write.clear(); - write - } else { - index.extend(false) - } - }; + let hooked_index = RelationHooked(index.clone(), hooked_extend(freepage_first)); - let mut tape = FrozenTapeWriter::<_, _>::create(expand); + let mut tape = FrozenTapeWriter::create(&hooked_index, false); let mut trace = Vec::new(); @@ -117,7 +109,7 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let (frozen_tape, branches) = tape.into_inner(); - let mut appendable_tape = TapeWriter::create(expand); + let mut appendable_tape = TapeWriter::create(&hooked_index, false); for branch in branches { appendable_tape.push(AppendableTuple { @@ -141,24 +133,25 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { } } -struct FrozenTapeWriter { - tape: TapeWriter, - branches: Vec>, +struct FrozenTapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + tape: TapeWriter<'a, R, FrozenTuple>, + branches: Vec>>, } -impl FrozenTapeWriter +impl<'a, R> FrozenTapeWriter<'a, R> where - G: PageGuard + DerefMut, - G::Target: Page, - E: Fn() -> G, + R: RelationWrite + 'a, { - fn create(extend: E) -> Self { + fn create(index: &'a R, tracking_freespace: bool) -> Self { Self { - tape: TapeWriter::create(extend), - branches: Vec::with_capacity(32), + tape: TapeWriter::create(index, tracking_freespace), + branches: Vec::new(), } } - fn push(&mut self, branch: Branch) { + fn push(&mut self, branch: Branch>) { self.branches.push(branch); if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); @@ -189,7 +182,71 @@ where self.branches.clear(); } } - fn into_inner(self) -> (TapeWriter, Vec>) { + fn into_inner(self) -> (TapeWriter<'a, R, FrozenTuple>, Vec>>) { (self.tape, self.branches) } } + +#[derive(Clone)] +struct RelationHooked(R, E); + +impl RelationRead for RelationHooked +where + R: RelationRead, + E: Clone, +{ + type Page = R::Page; + + type ReadGuard<'a> + = R::ReadGuard<'a> + where + Self: 'a; + + fn read(&self, id: u32) -> Self::ReadGuard<'_> { + self.0.read(id) + } +} + +impl RelationWrite for RelationHooked +where + R: RelationWrite, + E: Clone + for<'a> Fn(&'a R, bool) -> R::WriteGuard<'a>, +{ + type WriteGuard<'a> + = R::WriteGuard<'a> + where + Self: 'a; + + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.0.write(id, tracking_freespace) + } + + fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_> { + (self.1)(&self.0, tracking_freespace) + } + + fn search(&self, freespace: usize) -> Option> { + self.0.search(freespace) + } +} + +fn hooked_extend( + freepage_first: u32, +) -> impl Clone + for<'a> Fn(&'a R, bool) -> R::WriteGuard<'a> +where + R: RelationWrite, +{ + move |index, tracking_freespace| { + if !tracking_freespace { + if let Some(id) = freepages::fetch(index.clone(), freepage_first) { + let mut write = index.write(id, false); + write.clear(); + write + } else { + index.extend(false) + } + } else { + index.extend(true) + } + } +} diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index 3ce850cc..7eda147f 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -1,4 +1,3 @@ -use crate::types::*; use distance::Distance; use half::f16; use rabitq::binary::{BinaryCode, BinaryLut}; @@ -51,124 +50,144 @@ impl, B: Accessor2(f32, PhantomData O>); - -impl Default for Sum { - fn default() -> Self { - Self(0.0, PhantomData) - } +pub trait Accessor1 { + type Output; + fn push(&mut self, input: &[E]); + fn finish(self, input: M) -> Self::Output; } -impl Accessor2 for Sum, L2>> { - type Output = Distance; +impl Accessor1 for () { + type Output = (); - fn push(&mut self, target: &[f32], input: &[f32]) { - self.0 += f32::reduce_sum_of_d2(target, input) - } + fn push(&mut self, _: &[E]) {} - fn finish(self, (): (), (): ()) -> Self::Output { - Distance::from_f32(self.0) - } + fn finish(self, _: M) -> Self::Output {} } -impl Accessor2 for Sum, Dot>> { - type Output = Distance; +impl Accessor1 for (A,) +where + A: Accessor1, +{ + type Output = (A::Output,); - fn push(&mut self, target: &[f32], input: &[f32]) { - self.0 += f32::reduce_sum_of_xy(target, input) + fn push(&mut self, input: &[E]) { + self.0.push(input); } - fn finish(self, (): (), (): ()) -> Self::Output { - Distance::from_f32(-self.0) + fn finish(self, input: M) -> Self::Output { + (self.0.finish(input),) } } -impl Accessor2 for Sum, L2>> { - type Output = Distance; +impl Accessor1 for (A, B) +where + A: Accessor1, + B: Accessor1, +{ + type Output = (A::Output, B::Output); - fn push(&mut self, target: &[f16], input: &[f16]) { - self.0 += f16::reduce_sum_of_d2(target, input) + fn push(&mut self, input: &[E]) { + self.0.push(input); + self.1.push(input); } - fn finish(self, (): (), (): ()) -> Self::Output { - Distance::from_f32(self.0) + fn finish(self, input: M) -> Self::Output { + (self.0.finish(input), self.1.finish(input)) } } -impl Accessor2 for Sum, Dot>> { - type Output = Distance; - - fn push(&mut self, target: &[f16], input: &[f16]) { - self.0 += f16::reduce_sum_of_xy(target, input) - } +pub struct FunctionalAccessor { + data: T, + p: P, + f: F, +} - fn finish(self, (): (), (): ()) -> Self::Output { - Distance::from_f32(-self.0) +impl FunctionalAccessor { + pub fn new(data: T, p: P, f: F) -> Self { + Self { data, p, f } } } -#[derive(Debug, Clone)] -pub struct Diff(Vec<::Element>); +impl Accessor1 for FunctionalAccessor +where + P: for<'a> FnMut(&'a mut T, &'a [E]), + F: FnOnce(T, M) -> R, +{ + type Output = R; -impl Default for Diff { - fn default() -> Self { - Self(Vec::new()) + fn push(&mut self, input: &[E]) { + (self.p)(&mut self.data, input); } -} -impl Accessor2 for Diff, L2>> { - type Output = VectOwned; - - fn push(&mut self, target: &[f32], input: &[f32]) { - self.0.extend(f32::vector_sub(target, input)); + fn finish(self, input: M) -> Self::Output { + (self.f)(self.data, input) } +} - fn finish(self, (): (), (): ()) -> Self::Output { - VectOwned::new(self.0) +pub struct LAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> LAccess<'a, E, M, A> { + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } } } -impl Accessor2 for Diff, Dot>> { - type Output = VectOwned; +impl> Accessor1 for LAccess<'_, E0, M0, A> { + type Output = A::Output; - fn push(&mut self, target: &[f32], input: &[f32]) { - self.0.extend(f32::vector_sub(target, input)); + fn push(&mut self, rhs: &[E1]) { + let (lhs, elements) = self.elements.split_at(rhs.len()); + self.accessor.push(lhs, rhs); + self.elements = elements; } - fn finish(self, (): (), (): ()) -> Self::Output { - VectOwned::new(self.0) + fn finish(self, rhs: M1) -> Self::Output { + self.accessor.finish(self.metadata, rhs) } } -impl Accessor2 for Diff, L2>> { - type Output = VectOwned; - - fn push(&mut self, target: &[f16], input: &[f16]) { - self.0.extend(f16::vector_sub(target, input)); - } +pub struct RAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} - fn finish(self, (): (), (): ()) -> Self::Output { - VectOwned::new(self.0) +impl<'a, E, M, A> RAccess<'a, E, M, A> { + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } } } -impl Accessor2 for Diff, Dot>> { - type Output = VectOwned; +impl> Accessor1 for RAccess<'_, E1, M1, A> { + type Output = A::Output; - fn push(&mut self, target: &[f16], input: &[f16]) { - self.0.extend(f16::vector_sub(target, input)); + fn push(&mut self, lhs: &[E0]) { + let (rhs, elements) = self.elements.split_at(lhs.len()); + self.accessor.push(lhs, rhs); + self.elements = elements; } - fn finish(self, (): (), (): ()) -> Self::Output { - VectOwned::new(self.0) + fn finish(self, lhs: M0) -> Self::Output { + self.accessor.finish(lhs, self.metadata) } } #[derive(Debug)] -pub struct Block([u16; 32], PhantomData D>); +pub struct BlockAccessor([u16; 32], PhantomData D>); -impl Default for Block { +impl Default for BlockAccessor { fn default() -> Self { Self([0u16; 32], PhantomData) } @@ -179,10 +198,10 @@ impl [u8; 16], [u8; 16], (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), - ((f32, f32, f32, f32), f32), - > for Block + (f32, f32, f32, f32), + > for BlockAccessor { - type Output = [Distance; 32]; + type Output = [(f32, f32); 32]; fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { let t = simd::fast_scan::scan(input, target); @@ -199,7 +218,7 @@ impl &[f32; 32], &[f32; 32], ), - ((dis_v_2, b, k, qvector_sum), epsilon): ((f32, f32, f32, f32), f32), + (dis_v_2, b, k, qvector_sum): (f32, f32, f32, f32), ) -> Self::Output { std::array::from_fn(|i| { let rough = dis_u_2[i] @@ -207,7 +226,7 @@ impl + b * factor_ppc[i] + ((2.0 * self.0[i] as f32) - qvector_sum) * factor_ip[i] * k; let err = factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) + (rough, err) }) } } @@ -217,10 +236,10 @@ impl [u8; 16], [u8; 16], (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), - ((f32, f32, f32, f32), f32), - > for Block + (f32, f32, f32, f32), + > for BlockAccessor { - type Output = [Distance; 32]; + type Output = [(f32, f32); 32]; fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { let t = simd::fast_scan::scan(input, target); @@ -232,162 +251,119 @@ impl fn finish( self, (_, factor_ppc, factor_ip, factor_err): (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), - ((dis_v_2, b, k, qvector_sum), epsilon): ((f32, f32, f32, f32), f32), + (dis_v_2, b, k, qvector_sum): (f32, f32, f32, f32), ) -> Self::Output { std::array::from_fn(|i| { let rough = 0.5 * b * factor_ppc[i] + 0.5 * ((2.0 * self.0[i] as f32) - qvector_sum) * factor_ip[i] * k; let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) + (rough, err) }) } } -pub trait Accessor1 { - type Output; - fn push(&mut self, input: &[E]); - fn finish(self, input: M) -> Self::Output; -} - -impl Accessor1 for () { - type Output = (); - - fn push(&mut self, _: &[E]) {} +#[derive(Debug)] +pub struct DistanceAccessor(f32, PhantomData V>, PhantomData D>); - fn finish(self, _: M) -> Self::Output {} +impl Default for DistanceAccessor { + fn default() -> Self { + Self(0.0, PhantomData, PhantomData) + } } -impl Accessor1 for (A,) -where - A: Accessor1, -{ - type Output = (A::Output,); +impl Accessor2 for DistanceAccessor, L2> { + type Output = Distance; - fn push(&mut self, input: &[E]) { - self.0.push(input); + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0 += f32::reduce_sum_of_d2(target, input) } - fn finish(self, input: M) -> Self::Output { - (self.0.finish(input),) + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(self.0) } } -impl Accessor1 for (A, B) -where - A: Accessor1, - B: Accessor1, -{ - type Output = (A::Output, B::Output); +impl Accessor2 for DistanceAccessor, Dot> { + type Output = Distance; - fn push(&mut self, input: &[E]) { - self.0.push(input); - self.1.push(input); + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0 += f32::reduce_sum_of_xy(target, input) } - fn finish(self, input: M) -> Self::Output { - (self.0.finish(input), self.1.finish(input)) + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(-self.0) } } -pub struct FunctionalAccessor { - data: T, - p: P, - f: F, -} +impl Accessor2 for DistanceAccessor, L2> { + type Output = Distance; -impl FunctionalAccessor { - pub fn new(data: T, p: P, f: F) -> Self { - Self { data, p, f } + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0 += f16::reduce_sum_of_d2(target, input) + } + + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(self.0) } } -impl Accessor1 for FunctionalAccessor -where - P: for<'a> FnMut(&'a mut T, &'a [E]), - F: FnOnce(T, M) -> R, -{ - type Output = R; +impl Accessor2 for DistanceAccessor, Dot> { + type Output = Distance; - fn push(&mut self, input: &[E]) { - (self.p)(&mut self.data, input); + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0 += f16::reduce_sum_of_xy(target, input) } - fn finish(self, input: M) -> Self::Output { - (self.f)(self.data, input) + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(-self.0) } } -pub struct LAccess<'a, E, M, A> { - elements: &'a [E], - metadata: M, - accessor: A, -} +#[derive(Debug, Clone)] +pub struct ResidualAccessor(Vec); -impl<'a, E, M, A> LAccess<'a, E, M, A> { - pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { - Self { - elements, - metadata, - accessor, - } +impl Default for ResidualAccessor { + fn default() -> Self { + Self(Vec::new()) } } -impl> Accessor1 for LAccess<'_, E0, M0, A> { - type Output = A::Output; - - fn push(&mut self, rhs: &[E1]) { - let (lhs, elements) = self.elements.split_at(rhs.len()); - self.accessor.push(lhs, rhs); - self.elements = elements; - } +impl Accessor2 for ResidualAccessor> { + type Output = VectOwned; - fn finish(self, rhs: M1) -> Self::Output { - self.accessor.finish(self.metadata, rhs) + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0.extend(f32::vector_sub(target, input)); } -} - -pub struct RAccess<'a, E, M, A> { - elements: &'a [E], - metadata: M, - accessor: A, -} -impl<'a, E, M, A> RAccess<'a, E, M, A> { - pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { - Self { - elements, - metadata, - accessor, - } + fn finish(self, (): (), (): ()) -> Self::Output { + VectOwned::new(self.0) } } -impl> Accessor1 for RAccess<'_, E1, M1, A> { - type Output = A::Output; +impl Accessor2 for ResidualAccessor> { + type Output = VectOwned; - fn push(&mut self, lhs: &[E0]) { - let (rhs, elements) = self.elements.split_at(lhs.len()); - self.accessor.push(lhs, rhs); - self.elements = elements; + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0.extend(f16::vector_sub(target, input)); } - fn finish(self, lhs: M0) -> Self::Output { - self.accessor.finish(lhs, self.metadata) + fn finish(self, (): (), (): ()) -> Self::Output { + VectOwned::new(self.0) } } pub trait Vector: VectorOwned { type Element: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; + type Metadata: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; - fn vector_split(vector: Self::Borrowed<'_>) -> (Self::Metadata, Vec<&[Self::Element]>); - fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); - fn from_owned(vector: OwnedVector) -> Self; + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[Self::Element]>, Self::Metadata); - fn compute_lut_block(vector: Self::Borrowed<'_>) -> BlockLut; + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); - fn compute_lut(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut); + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut; + + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut); fn code(vector: Self::Borrowed<'_>) -> rabitq::Code; } @@ -397,35 +373,28 @@ impl Vector for VectOwned { type Element = f32; - fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f32]>) { + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[f32]>, ()) { let vector = vector.slice(); ( - (), match vector.len() { 0..=960 => vec![vector], 961..=1280 => vec![&vector[..640], &vector[640..]], 1281.. => vector.chunks(1920).collect(), }, + (), ) } - fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { (vector.slice(), ()) } - fn from_owned(vector: OwnedVector) -> Self { - match vector { - OwnedVector::Vecf32(x) => x, - _ => panic!("internal error: should be a vector"), - } - } - - fn compute_lut_block(vector: Self::Borrowed<'_>) -> BlockLut { + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { rabitq::block::preprocess(vector.slice()) } - fn compute_lut(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { - rabitq::compute_lut(vector.slice()) + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { + rabitq::preprocess(vector.slice()) } fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { @@ -438,35 +407,28 @@ impl Vector for VectOwned { type Element = f16; - fn vector_split(vector: Self::Borrowed<'_>) -> ((), Vec<&[f16]>) { + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[f16]>, ()) { let vector = vector.slice(); ( - (), match vector.len() { 0..=1920 => vec![vector], 1921..=2560 => vec![&vector[..1280], &vector[1280..]], 2561.. => vector.chunks(3840).collect(), }, + (), ) } - fn elements_and_metadata(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { (vector.slice(), ()) } - fn from_owned(vector: OwnedVector) -> Self { - match vector { - OwnedVector::Vecf16(x) => x, - _ => panic!("internal error: should be a halfvec"), - } - } - - fn compute_lut_block(vector: Self::Borrowed<'_>) -> BlockLut { + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { rabitq::block::preprocess(&f16::vector_to_f32(vector.slice())) } - fn compute_lut(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { - rabitq::compute_lut(&f16::vector_to_f32(vector.slice())) + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { + rabitq::preprocess(&f16::vector_to_f32(vector.slice())) } fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { @@ -474,54 +436,23 @@ impl Vector for VectOwned { } } -pub trait OperatorDistance: 'static + Debug + Copy { - const KIND: DistanceKind; - - fn compute_lowerbound_binary(lut: &BinaryLut, code: BinaryCode<'_>, epsilon: f32) -> Distance; - - type BlockAccessor: for<'a> Accessor2< - [u8; 16], - [u8; 16], - (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), - ((f32, f32, f32, f32), f32), - Output = [Distance; 32], - > + Default; - - fn block_accessor() -> Self::BlockAccessor { - Default::default() - } -} - #[derive(Debug, Clone, Copy)] pub struct L2; -impl OperatorDistance for L2 { - const KIND: DistanceKind = DistanceKind::L2; - - fn compute_lowerbound_binary(lut: &BinaryLut, code: BinaryCode<'_>, epsilon: f32) -> Distance { - rabitq::binary::process_lowerbound_l2(lut, code, epsilon) - } - - type BlockAccessor = Block; -} - #[derive(Debug, Clone, Copy)] pub struct Dot; -impl OperatorDistance for Dot { - const KIND: DistanceKind = DistanceKind::Dot; - - fn compute_lowerbound_binary(lut: &BinaryLut, code: BinaryCode<'_>, epsilon: f32) -> Distance { - rabitq::binary::process_lowerbound_dot(lut, code, epsilon) - } - - type BlockAccessor = Block; -} - pub trait Operator: 'static + Debug + Copy { type Vector: Vector; - type Distance: OperatorDistance; + type BlockAccessor: Default + + for<'a> Accessor2< + [u8; 16], + [u8; 16], + (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), + (f32, f32, f32, f32), + Output = [(f32, f32); 32], + >; type DistanceAccessor: Default + Accessor2< @@ -532,8 +463,6 @@ pub trait Operator: 'static + Debug + Copy { Output = Distance, >; - const SUPPORTS_RESIDUAL: bool; - type ResidualAccessor: Default + Accessor2< ::Element, @@ -542,6 +471,10 @@ pub trait Operator: 'static + Debug + Copy { ::Metadata, Output = Self::Vector, >; + + const SUPPORTS_RESIDUAL: bool; + + fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32); } #[derive(Debug)] @@ -558,47 +491,63 @@ impl Copy for Op {} impl Operator for Op, L2> { type Vector = VectOwned; - type Distance = L2; + type BlockAccessor = BlockAccessor; + + type DistanceAccessor = DistanceAccessor, L2>; - type DistanceAccessor = Sum, L2>>; + type ResidualAccessor = ResidualAccessor>; const SUPPORTS_RESIDUAL: bool = true; - type ResidualAccessor = Diff, L2>>; + fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32) { + rabitq::binary::process_l2(lut, code) + } } impl Operator for Op, Dot> { type Vector = VectOwned; - type Distance = Dot; + type BlockAccessor = BlockAccessor; + + type DistanceAccessor = DistanceAccessor, Dot>; - type DistanceAccessor = Sum, Dot>>; + type ResidualAccessor = ResidualAccessor>; const SUPPORTS_RESIDUAL: bool = false; - type ResidualAccessor = Diff, Dot>>; + fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32) { + rabitq::binary::process_dot(lut, code) + } } impl Operator for Op, L2> { type Vector = VectOwned; - type Distance = L2; + type BlockAccessor = BlockAccessor; + + type DistanceAccessor = DistanceAccessor, L2>; - type DistanceAccessor = Sum, L2>>; + type ResidualAccessor = ResidualAccessor>; const SUPPORTS_RESIDUAL: bool = true; - type ResidualAccessor = Diff, L2>>; + fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32) { + rabitq::binary::process_l2(lut, code) + } } impl Operator for Op, Dot> { type Vector = VectOwned; - type Distance = Dot; + type BlockAccessor = BlockAccessor; - type DistanceAccessor = Sum, Dot>>; + type DistanceAccessor = DistanceAccessor, Dot>; + + type ResidualAccessor = ResidualAccessor>; const SUPPORTS_RESIDUAL: bool = false; - type ResidualAccessor = Diff, Dot>>; + fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32) { + rabitq::binary::process_dot(lut, code) + } } diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index 45979c9d..004b19ea 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -5,14 +5,17 @@ use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; -use std::num::NonZeroU64; +use std::num::NonZero; use vector::VectorOwned; -type Results = Vec<( +type Result = ( Reverse, - AlwaysEqual, + AlwaysEqual, + AlwaysEqual>, AlwaysEqual, -)>; +); + +type Rerank = (Reverse, AlwaysEqual>); pub fn how(index: impl RelationRead) -> RerankMethod { let meta_guard = index.read(0); @@ -26,67 +29,81 @@ pub fn how(index: impl RelationRead) -> RerankMethod { } } -pub fn rerank_index( +pub struct Reranker { + heap: BinaryHeap>, + cache: BinaryHeap<(Reverse, AlwaysEqual>)>, + f: F, +} + +impl) -> Option> Iterator for Reranker { + type Item = (Distance, NonZero); + + fn next(&mut self) -> Option { + while let Some((Reverse(_), AlwaysEqual(_), AlwaysEqual(pay_u), AlwaysEqual(mean))) = + pop_if(&mut self.heap, |(d, ..)| { + Some(*d) > self.cache.peek().map(|(d, ..)| *d) + }) + { + if let Some(dis_u) = (self.f)(mean, pay_u) { + self.cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); + }; + } + let (Reverse(dis_u), AlwaysEqual(pay_u)) = self.cache.pop()?; + Some((dis_u, pay_u)) + } +} + +impl Reranker { + pub fn finish( + self, + ) -> ( + impl Iterator>, + impl Iterator, + ) { + (self.heap.into_iter(), self.cache.into_iter()) + } +} + +pub fn rerank_index( index: impl RelationRead, vector: O::Vector, - results: Results, -) -> impl Iterator { - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _)>::new(); - std::iter::from_fn(move || { - while let Some((_, AlwaysEqual(pay_u), AlwaysEqual(mean))) = - pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) - { - if let Some(dis_u) = vectors::read_for_h0_tuple::( + results: Vec>, +) -> Reranker) -> Option> { + Reranker { + heap: BinaryHeap::from(results), + cache: BinaryHeap::<(Reverse, _)>::new(), + f: move |mean, pay_u| { + vectors::read_for_h0_tuple::( index.clone(), mean, pay_u, LAccess::new( - O::Vector::elements_and_metadata(vector.as_borrowed()), + O::Vector::unpack(vector.as_borrowed()), O::DistanceAccessor::default(), ), - ) { - cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); - }; - } - let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; - Some((dis_u, pay_u)) - }) + ) + }, + } } -pub fn rerank_heap( +pub fn rerank_heap( vector: O::Vector, - results: Results, - mut fetch: F, -) -> impl Iterator -where - F: FnMut(NonZeroU64) -> Option, -{ - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _)>::new(); - std::iter::from_fn(move || { - while let Some((_, AlwaysEqual(pay_u), AlwaysEqual(_))) = - pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) - { - let vector = O::Vector::elements_and_metadata(vector.as_borrowed()); - if let Some(vec_u) = fetch(pay_u) { - let vec_u = O::Vector::elements_and_metadata(vec_u.as_borrowed()); - let mut accessor = O::DistanceAccessor::default(); - accessor.push(vector.0, vec_u.0); - let dis_u = accessor.finish(vector.1, vec_u.1); - cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); - } - } - let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; - Some((dis_u, pay_u)) - }) -} - -pub fn skip(results: Results) -> impl Iterator { - let results = BinaryHeap::from(results); - results - .into_iter() - .map(|(Reverse(x), AlwaysEqual(y), _)| (x, y)) + results: Vec>, + mut fetch: impl FnMut(NonZero) -> Option, +) -> Reranker) -> Option> { + Reranker { + heap: BinaryHeap::from(results), + cache: BinaryHeap::<(Reverse, _)>::new(), + f: move |_: IndexPointer, pay_u| { + let vector = O::Vector::unpack(vector.as_borrowed()); + let vec_u = fetch(pay_u)?; + let vec_u = O::Vector::unpack(vec_u.as_borrowed()); + let mut accessor = O::DistanceAccessor::default(); + accessor.push(vector.0, vec_u.0); + let dis_u = accessor.finish(vector.1, vec_u.1); + Some(dis_u) + }, + } } fn pop_if( diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index 343eb214..922f726e 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -6,21 +6,22 @@ use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; -use std::num::NonZeroU64; +use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -type Results = Vec<( +type Result = ( Reverse, - AlwaysEqual, + AlwaysEqual, + AlwaysEqual>, AlwaysEqual, -)>; +); -pub fn search( +pub fn default_search( index: impl RelationRead, vector: O::Vector, probes: Vec, epsilon: f32, -) -> Results { +) -> Vec> { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -40,7 +41,7 @@ pub fn search( drop(meta_guard); let default_lut = if !is_residual { - Some(O::Vector::compute_lut(vector.as_borrowed())) + Some(O::Vector::preprocess(vector.as_borrowed())) } else { None }; @@ -53,7 +54,7 @@ pub fn search( index.clone(), mean, LAccess::new( - O::Vector::elements_and_metadata(vector.as_borrowed()), + O::Vector::unpack(vector.as_borrowed()), O::ResidualAccessor::default(), ), ); @@ -65,23 +66,19 @@ pub fn search( let step = |state: State| { let mut results = LinkedVec::new(); for (first, residual) in state { - let lut_block = if let Some(residual) = residual { - &O::Vector::compute_lut_block(residual.as_borrowed()) - } else if let Some((lut_block, _)) = default_lut.as_ref() { - lut_block + let block_lut = if let Some(residual) = residual { + &O::Vector::block_preprocess(residual.as_borrowed()) + } else if let Some((block_lut, _)) = default_lut.as_ref() { + block_lut } else { unreachable!() }; tape::read_h1_tape( index.clone(), first, - || { - RAccess::new( - (&lut_block.1, (lut_block.0, epsilon)), - O::Distance::block_accessor(), - ) - }, - |lowerbound, mean, first| { + || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), + |(rough, err), mean, first| { + let lowerbound = Distance::from_f32(rough - err * epsilon); results.push((Reverse(lowerbound), AlwaysEqual(first), AlwaysEqual(mean))); }, |_| (), @@ -92,15 +89,17 @@ pub fn search( let index = index.clone(); let vector = vector.as_borrowed(); std::iter::from_fn(move || { - while let Some((_, AlwaysEqual(first), AlwaysEqual(mean))) = - pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) + while let Some((Reverse(_), AlwaysEqual(first), AlwaysEqual(mean))) = + pop_if(&mut heap, |(d, ..)| { + Some(*d) > cache.peek().map(|(d, ..)| *d) + }) { if is_residual { let (dis_u, residual_u) = vectors::read_for_h1_tuple::( index.clone(), mean, LAccess::new( - O::Vector::elements_and_metadata(vector), + O::Vector::unpack(vector), ( O::DistanceAccessor::default(), O::ResidualAccessor::default(), @@ -116,10 +115,7 @@ pub fn search( let dis_u = vectors::read_for_h1_tuple::( index.clone(), mean, - LAccess::new( - O::Vector::elements_and_metadata(vector), - O::DistanceAccessor::default(), - ), + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); } @@ -134,9 +130,9 @@ pub fn search( let mut results = LinkedVec::new(); for (first, residual) in state { - let (lut_block, lut_binary) = + let (block_lut, binary_lut) = if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { - &O::Vector::compute_lut(residual) + &O::Vector::preprocess(residual) } else if let Some(lut) = default_lut.as_ref() { lut } else { @@ -145,25 +141,26 @@ pub fn search( let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let mut callback = |lowerbound, mean, payload| { - results.push((Reverse(lowerbound), AlwaysEqual(payload), AlwaysEqual(mean))); + let mut callback = |(rough, err), mean, payload| { + let lowerbound = Distance::from_f32(rough - err * 1.9); + results.push(( + Reverse(lowerbound), + AlwaysEqual(()), + AlwaysEqual(payload), + AlwaysEqual(mean), + )); }; tape::read_frozen_tape( index.clone(), jump_tuple.frozen_first(), - || { - RAccess::new( - (&lut_block.1, (lut_block.0, epsilon)), - O::Distance::block_accessor(), - ) - }, + || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), &mut callback, |_| (), ); tape::read_appendable_tape( index.clone(), jump_tuple.appendable_first(), - |code| O::Distance::compute_lowerbound_binary(lut_binary, code, epsilon), + |code| O::binary_process(binary_lut, code), &mut callback, |_| (), ); @@ -171,13 +168,13 @@ pub fn search( results.into_vec() } -pub fn search_and_estimate( +pub fn maxsim_search( index: impl RelationRead, vector: O::Vector, probes: Vec, epsilon: f32, mut t: u32, -) -> (Results, Distance) { +) -> (Vec>, Distance) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -197,7 +194,7 @@ pub fn search_and_estimate( drop(meta_guard); let default_lut = if !is_residual { - Some(O::Vector::compute_lut(vector.as_borrowed())) + Some(O::Vector::preprocess(vector.as_borrowed())) } else { None }; @@ -210,7 +207,7 @@ pub fn search_and_estimate( index.clone(), mean, LAccess::new( - O::Vector::elements_and_metadata(vector.as_borrowed()), + O::Vector::unpack(vector.as_borrowed()), O::ResidualAccessor::default(), ), ); @@ -222,23 +219,19 @@ pub fn search_and_estimate( let step = |state: State| { let mut results = LinkedVec::new(); for (first, residual) in state { - let lut_block = if let Some(residual) = residual { - &O::Vector::compute_lut_block(residual.as_borrowed()) - } else if let Some((lut_block, _)) = default_lut.as_ref() { - lut_block + let block_lut = if let Some(residual) = residual { + &O::Vector::block_preprocess(residual.as_borrowed()) + } else if let Some((block_lut, _)) = default_lut.as_ref() { + block_lut } else { unreachable!() }; tape::read_h1_tape( index.clone(), first, - || { - RAccess::new( - (&lut_block.1, (lut_block.0, epsilon)), - O::Distance::block_accessor(), - ) - }, - |lowerbound, mean, first| { + || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), + |(rough, err), mean, first| { + let lowerbound = Distance::from_f32(rough - err * epsilon); results.push((Reverse(lowerbound), AlwaysEqual(first), AlwaysEqual(mean))); }, |_| (), @@ -249,15 +242,17 @@ pub fn search_and_estimate( let index = index.clone(); let vector = vector.as_borrowed(); std::iter::from_fn(move || { - while let Some((_, AlwaysEqual(first), AlwaysEqual(mean))) = - pop_if(&mut heap, |x| Some(x.0) > cache.peek().map(|x| x.0)) + while let Some((Reverse(_), AlwaysEqual(first), AlwaysEqual(mean))) = + pop_if(&mut heap, |(d, ..)| { + Some(*d) > cache.peek().map(|(d, ..)| *d) + }) { if is_residual { let (dis_u, residual_u) = vectors::read_for_h1_tuple::( index.clone(), mean, LAccess::new( - O::Vector::elements_and_metadata(vector), + O::Vector::unpack(vector), ( O::DistanceAccessor::default(), O::ResidualAccessor::default(), @@ -273,10 +268,7 @@ pub fn search_and_estimate( let dis_u = vectors::read_for_h1_tuple::( index.clone(), mean, - LAccess::new( - O::Vector::elements_and_metadata(vector), - O::DistanceAccessor::default(), - ), + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); } @@ -305,9 +297,9 @@ pub fn search_and_estimate( let mut results = LinkedVec::new(); for (first, residual) in state { - let (lut_block, lut_binary) = + let (block_lut, binary_lut) = if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { - &O::Vector::compute_lut(residual) + &O::Vector::preprocess(residual) } else if let Some(lut) = default_lut.as_ref() { lut } else { @@ -316,25 +308,27 @@ pub fn search_and_estimate( let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let mut callback = |lowerbound, mean, payload| { - results.push((Reverse(lowerbound), AlwaysEqual(payload), AlwaysEqual(mean))); + let mut callback = |(rough, err), mean, payload| { + let lowerbound = Distance::from_f32(rough - err * epsilon); + let rough = Distance::from_f32(rough); + results.push(( + Reverse(lowerbound), + AlwaysEqual(rough), + AlwaysEqual(payload), + AlwaysEqual(mean), + )); }; tape::read_frozen_tape( index.clone(), jump_tuple.frozen_first(), - || { - RAccess::new( - (&lut_block.1, (lut_block.0, epsilon)), - O::Distance::block_accessor(), - ) - }, + || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), &mut callback, |_| (), ); tape::read_appendable_tape( index.clone(), jump_tuple.appendable_first(), - |code| O::Distance::compute_lowerbound_binary(lut_binary, code, epsilon), + |code| O::binary_process(binary_lut, code), &mut callback, |_| (), ); diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index e9e09a0a..093774a5 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -3,30 +3,32 @@ use crate::tuples::*; use crate::{IndexPointer, Page, PageGuard, RelationRead, RelationWrite}; use rabitq::binary::BinaryCode; use std::marker::PhantomData; -use std::num::NonZeroU64; -use std::ops::DerefMut; +use std::num::NonZero; -pub struct TapeWriter { - head: G, +pub struct TapeWriter<'a, R, T> +where + R: RelationWrite + 'a, +{ + head: R::WriteGuard<'a>, first: u32, - extend: E, + index: &'a R, + tracking_freespace: bool, _phantom: PhantomData T>, } -impl TapeWriter +impl<'a, R, T> TapeWriter<'a, R, T> where - G: PageGuard + DerefMut, - G::Target: Page, - E: Fn() -> G, + R: RelationWrite + 'a, { - pub fn create(extend: E) -> Self { - let mut head = extend(); + pub fn create(index: &'a R, tracking_freespace: bool) -> Self { + let mut head = index.extend(tracking_freespace); head.get_opaque_mut().skip = head.id(); let first = head.id(); Self { head, first, - extend, + index, + tracking_freespace, _phantom: PhantomData, } } @@ -40,17 +42,15 @@ where if self.head.len() == 0 { panic!("implementation: a clear page cannot accommodate a single tuple"); } - let next = (self.extend)(); + let next = self.index.extend(self.tracking_freespace); self.head.get_opaque_mut().next = next.id(); self.head = next; } } -impl TapeWriter +impl<'a, R, T> TapeWriter<'a, R, T> where - G: PageGuard + DerefMut, - G::Target: Page, - E: Fn() -> G, + R: RelationWrite + 'a, T: Tuple, { pub fn push(&mut self, x: T) -> IndexPointer { @@ -58,7 +58,7 @@ where if let Some(i) = self.head.alloc(&bytes) { pair_to_pointer((self.head.id(), i)) } else { - let next = (self.extend)(); + let next = self.index.extend(self.tracking_freespace); self.head.get_opaque_mut().next = next.id(); self.head = next; if let Some(i) = self.head.alloc(&bytes) { @@ -124,7 +124,7 @@ pub fn read_frozen_tape( index: impl RelationRead, first: u32, accessor: impl Fn() -> A, - mut callback: impl FnMut(T, IndexPointer, NonZeroU64), + mut callback: impl FnMut(T, IndexPointer, NonZero), mut step: impl FnMut(u32), ) where A: for<'a> Accessor1< @@ -166,7 +166,7 @@ pub fn read_appendable_tape( index: impl RelationRead, first: u32, mut access: impl for<'a> FnMut(BinaryCode<'a>) -> T, - mut callback: impl FnMut(T, IndexPointer, NonZeroU64), + mut callback: impl FnMut(T, IndexPointer, NonZero), mut step: impl FnMut(u32), ) { assert!(first != u32::MAX); diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index ea00abe3..849cb474 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -2,7 +2,7 @@ use crate::IndexPointer; use crate::operator::Vector; use rabitq::binary::BinaryCode; use std::marker::PhantomData; -use std::num::{NonZeroU8, NonZeroU64}; +use std::num::NonZero; use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; @@ -204,7 +204,7 @@ impl FreepageTupleWriter<'_> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader0 { - payload: Option, + payload: Option>, metadata_s: usize, elements_s: usize, elements_e: usize, @@ -215,7 +215,7 @@ struct VectorTupleHeader0 { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader1 { - payload: Option, + payload: Option>, pointer: IndexPointer, elements_s: usize, elements_e: usize, @@ -224,12 +224,12 @@ struct VectorTupleHeader1 { #[derive(Debug, Clone, PartialEq)] pub enum VectorTuple { _0 { - payload: Option, + payload: Option>, metadata: V::Metadata, elements: Vec, }, _1 { - payload: Option, + payload: Option>, pointer: IndexPointer, elements: Vec, }, @@ -351,7 +351,7 @@ pub enum VectorTupleReader<'a, V: Vector> { impl Copy for VectorTupleReader<'_, V> {} impl<'a, V: Vector> VectorTupleReader<'a, V> { - pub fn payload(self) -> Option { + pub fn payload(self) -> Option> { match self { VectorTupleReader::_0(this) => this.header.payload, VectorTupleReader::_1(this) => this.header.payload, @@ -642,7 +642,7 @@ struct FrozenTupleHeader0 { factor_ppc: [f32; 32], factor_ip: [f32; 32], factor_err: [f32; 32], - payload: [Option; 32], + payload: [Option>; 32], elements_s: usize, elements_e: usize, } @@ -663,7 +663,7 @@ pub enum FrozenTuple { factor_ppc: [f32; 32], factor_ip: [f32; 32], factor_err: [f32; 32], - payload: [Option; 32], + payload: [Option>; 32], elements: Vec<[u8; 16]>, }, _1 { @@ -815,7 +815,7 @@ impl<'a> FrozenTupleReader0<'a> { pub fn elements(self) -> &'a [[u8; 16]] { self.elements } - pub fn payload(self) -> &'a [Option; 32] { + pub fn payload(self) -> &'a [Option>; 32] { &self.header.payload } } @@ -855,7 +855,7 @@ pub struct FrozenTupleWriter1<'a> { } impl FrozenTupleWriter0<'_> { - pub fn payload(&mut self) -> &mut [Option; 32] { + pub fn payload(&mut self) -> &mut [Option>; 32] { &mut self.header.payload } } @@ -868,7 +868,7 @@ struct AppendableTupleHeader { factor_ppc: f32, factor_ip: f32, factor_err: f32, - payload: Option, + payload: Option>, elements_s: usize, elements_e: usize, } @@ -880,7 +880,7 @@ pub struct AppendableTuple { pub factor_ppc: f32, pub factor_ip: f32, pub factor_err: f32, - pub payload: Option, + pub payload: Option>, pub elements: Vec, } @@ -949,7 +949,7 @@ impl<'a> AppendableTupleReader<'a> { self.elements, ) } - pub fn payload(self) -> Option { + pub fn payload(self) -> Option> { self.header.payload } } @@ -962,7 +962,7 @@ pub struct AppendableTupleWriter<'a> { } impl AppendableTupleWriter<'_> { - pub fn payload(&mut self) -> &mut Option { + pub fn payload(&mut self) -> &mut Option> { &mut self.header.payload } } @@ -1004,7 +1004,7 @@ const fn soundness_check() { Immutable, KnownLayout, )] -pub struct ZeroU8(Option); +pub struct ZeroU8(Option>); #[repr(transparent)] #[derive( diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index 99b78918..e0053947 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -1,7 +1,7 @@ use crate::operator::*; use crate::tuples::*; use crate::{IndexPointer, Page, PageGuard, RelationRead, RelationWrite, tape}; -use std::num::NonZeroU64; +use std::num::NonZero; use vector::VectorOwned; pub fn read_for_h1_tuple< @@ -33,7 +33,7 @@ pub fn read_for_h0_tuple< >( index: impl RelationRead, mean: IndexPointer, - payload: NonZeroU64, + payload: NonZero, accessor: A, ) -> Option { let mut cursor = Err(mean); @@ -58,7 +58,7 @@ pub fn append( index: impl RelationWrite, vectors_first: u32, vector: ::Borrowed<'_>, - payload: NonZeroU64, + payload: NonZero, ) -> IndexPointer { fn append(index: impl RelationWrite, first: u32, bytes: &[u8]) -> IndexPointer { if let Some(mut write) = index.search(bytes.len()) { @@ -69,7 +69,7 @@ pub fn append( } tape::append(index, first, bytes, true) } - let (metadata, slices) = O::Vector::vector_split(vector); + let (slices, metadata) = O::Vector::split(vector); let mut chain = Ok(metadata); for i in (0..slices.len()).rev() { let bytes = VectorTuple::::serialize(&match chain { diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index 7ce5ed3e..4c04b884 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -164,10 +164,10 @@ fn rabitq_index( let lut = rabitq::block::preprocess(&samples[i]); let mut result = (f32::INFINITY, 0); for block in 0..c.div_ceil(32) { - let lowerbound = - rabitq::block::process_lowerbound_l2(&lut, blocks[block].code(), 1.9); + let returns = rabitq::block::process_l2(&lut, blocks[block].code()); + let lowerbound = returns.map(|(rough, err)| rough - err * 1.9); for j in block * 32..std::cmp::min(block * 32 + 32, c) { - if lowerbound[j - block * 32].to_f32() < result.0 { + if lowerbound[j - block * 32] < result.0 { let dis = f32::reduce_sum_of_d2(&samples[i], ¢roids[j]); if dis <= result.0 { result = (dis, j); diff --git a/crates/rabitq/src/binary.rs b/crates/rabitq/src/binary.rs index fa1d01aa..1f8bc10f 100644 --- a/crates/rabitq/src/binary.rs +++ b/crates/rabitq/src/binary.rs @@ -1,4 +1,3 @@ -use distance::Distance; use simd::Floating; pub type BinaryLut = ( @@ -18,32 +17,30 @@ pub fn preprocess(vector: &[f32]) -> BinaryLut { ((dis_v_2, b, k, qvector_sum), binarize(&qvector)) } -pub fn process_lowerbound_l2( +pub fn process_l2( lut: &BinaryLut, (dis_u_2, factor_ppc, factor_ip, factor_err, t): BinaryCode<'_>, - epsilon: f32, -) -> Distance { +) -> (f32, f32) { let &((dis_v_2, b, k, qvector_sum), ref s) = lut; let value = asymmetric_binary_dot_product(t, s) as u16; let rough = dis_u_2 + dis_v_2 + b * factor_ppc + ((2.0 * value as f32) - qvector_sum) * factor_ip * k; let err = factor_err * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) + (rough, err) } -pub fn process_lowerbound_dot( +pub fn process_dot( lut: &BinaryLut, (_, factor_ppc, factor_ip, factor_err, t): BinaryCode<'_>, - epsilon: f32, -) -> Distance { +) -> (f32, f32) { let &((dis_v_2, b, k, qvector_sum), ref s) = lut; let value = asymmetric_binary_dot_product(t, s) as u16; let rough = 0.5 * b * factor_ppc + 0.5 * ((2.0 * value as f32) - qvector_sum) * factor_ip * k; let err = 0.5 * factor_err * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) + (rough, err) } -pub fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { +pub(crate) fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { let n = vector.len(); let mut t0 = vec![0u64; n.div_ceil(64)]; let mut t1 = vec![0u64; n.div_ceil(64)]; @@ -58,7 +55,10 @@ pub fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { (t0, t1, t2, t3) } -fn asymmetric_binary_dot_product(x: &[u64], y: &(Vec, Vec, Vec, Vec)) -> u32 { +pub(crate) fn asymmetric_binary_dot_product( + x: &[u64], + y: &(Vec, Vec, Vec, Vec), +) -> u32 { let t0 = simd::bit::sum_of_and(x, &y.0); let t1 = simd::bit::sum_of_and(x, &y.1); let t2 = simd::bit::sum_of_and(x, &y.2); diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs index 0c5aaaf1..83496f92 100644 --- a/crates/rabitq/src/block.rs +++ b/crates/rabitq/src/block.rs @@ -1,4 +1,3 @@ -use distance::Distance; use simd::Floating; pub type BlockLut = ((f32, f32, f32, f32), Vec<[u8; 16]>); @@ -21,11 +20,10 @@ pub fn preprocess(vector: &[f32]) -> BlockLut { ((dis_v_2, b, k, qvector_sum), compress(qvector)) } -pub fn process_lowerbound_l2( +pub fn process_l2( lut: &BlockLut, (dis_u_2, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, - epsilon: f32, -) -> [Distance; 32] { +) -> [(f32, f32); 32] { let &((dis_v_2, b, k, qvector_sum), ref s) = lut; let r = simd::fast_scan::scan(t, s); std::array::from_fn(|i| { @@ -34,26 +32,25 @@ pub fn process_lowerbound_l2( + b * factor_ppc[i] + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; let err = factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) + (rough, err) }) } -pub fn process_lowerbound_dot( +pub fn process_dot( lut: &BlockLut, (_, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, - epsilon: f32, -) -> [Distance; 32] { +) -> [(f32, f32); 32] { let &((dis_v_2, b, k, qvector_sum), ref s) = lut; let r = simd::fast_scan::scan(t, s); std::array::from_fn(|i| { let rough = 0.5 * b * factor_ppc[i] + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - epsilon * err) + (rough, err) }) } -pub fn compress(mut vector: Vec) -> Vec<[u8; 16]> { +pub(crate) fn compress(mut vector: Vec) -> Vec<[u8; 16]> { let n = vector.len().div_ceil(4); vector.resize(n * 4, 0); let mut result = vec![[0u8; 16]; n]; diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index a59c6c2a..7d4b81bd 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -74,7 +74,7 @@ pub fn code(dims: u32, vector: &[f32]) -> Code { } } -pub fn compute_lut(vector: &[f32]) -> (BlockLut, BinaryLut) { +pub fn preprocess(vector: &[f32]) -> (BlockLut, BinaryLut) { use simd::Floating; let dis_v_2 = f32::reduce_sum_of_x2(vector); let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index 08ac3754..8e335467 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -1,11 +1,11 @@ use serde::{Deserialize, Serialize}; use std::ffi::{CStr, CString}; -use std::num::{NonZero, NonZeroU32}; +use std::num::NonZero; #[derive(Debug, Clone, Copy, Serialize, Deserialize)] pub enum Typmod { Any, - Dims(NonZeroU32), + Dims(NonZero), } impl Typmod { @@ -14,7 +14,7 @@ impl Typmod { if x == -1 { Some(Any) } else if x >= 1 { - Some(Dims(NonZeroU32::new(x as u32).unwrap())) + Some(Dims(NonZero::new(x as u32).unwrap())) } else { None } @@ -33,7 +33,7 @@ impl Typmod { Dims(x) => x.get() as i32, } } - pub fn dims(self) -> Option { + pub fn dims(self) -> Option> { use Typmod::*; match self { Any => None, diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index 4041be2c..2135f635 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -4,7 +4,7 @@ use algorithm::operator::{Dot, L2, Op}; use algorithm::types::*; use algorithm::{RelationRead, RelationWrite}; use half::f16; -use std::num::NonZeroU64; +use std::num::NonZero; use vector::VectorOwned; use vector::vect::{VectBorrowed, VectOwned}; @@ -38,7 +38,7 @@ pub fn bulkdelete( opfamily: Opfamily, index: impl RelationWrite, check: impl Fn(), - callback: impl Fn(NonZeroU64) -> bool, + callback: impl Fn(NonZero) -> bool, ) { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { @@ -110,7 +110,7 @@ pub fn build( pub fn insert( opfamily: Opfamily, index: impl RelationWrite, - payload: NonZeroU64, + payload: NonZero, vector: OwnedVector, ) { match (vector, opfamily.distance_kind()) { diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 0aafeb60..dd17a3e9 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -1,5 +1,5 @@ use crate::datatype::typmod::Typmod; -use crate::index::am::{Reloption, ctid_to_pointer}; +use crate::index::am::{Reloption, ctid_to_key, kv_to_pointer}; use crate::index::opclass::{Opfamily, opfamily}; use crate::index::storage::{PostgresPage, PostgresRelation}; use crate::index::types::*; @@ -383,7 +383,8 @@ pub unsafe extern "C" fn ambuild( let relation = unsafe { PostgresRelation::new(index_relation) }; heap.traverse(true, |(ctid, store)| { for (vector, extra) in store { - let payload = ctid_to_pointer(ctid, extra); + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); crate::index::algorithm::insert(opfamily, relation.clone(), payload, vector); } indtuples += 1; @@ -864,7 +865,8 @@ unsafe fn parallel_build( VchordrqCachedReader::_0(_) => { heap.traverse(true, |(ctid, store)| { for (vector, extra) in store { - let payload = ctid_to_pointer(ctid, extra); + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); } unsafe { @@ -886,7 +888,8 @@ unsafe fn parallel_build( }; heap.traverse(true, |(ctid, store)| { for (vector, extra) in store { - let payload = ctid_to_pointer(ctid, extra); + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); } unsafe { diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index 949ba0eb..126c26b4 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -5,10 +5,10 @@ use crate::index::opclass::opfamily; use crate::index::scanners::*; use crate::index::storage::PostgresRelation; use pgrx::datum::Internal; -use pgrx::pg_sys::{Datum, ItemPointerData}; +use pgrx::pg_sys::{BlockIdData, Datum, ItemPointerData}; use std::cell::LazyCell; use std::ffi::CStr; -use std::num::NonZeroU64; +use std::num::NonZero; use std::ptr::NonNull; use std::sync::OnceLock; @@ -201,7 +201,8 @@ unsafe fn aminsertinner( let ctid = unsafe { ctid.read() }; if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { for (vector, extra) in store { - let payload = ctid_to_pointer(ctid, extra); + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); } } @@ -227,7 +228,11 @@ pub unsafe extern "C" fn ambulkdelete( pgrx::pg_sys::vacuum_delay_point(); }; let callback = callback.expect("null function pointer"); - let callback = |p: NonZeroU64| unsafe { callback(&mut pointer_to_ctid(p).0, callback_state) }; + let callback = |pointer: NonZero| { + let (key, _) = pointer_to_kv(pointer); + let mut ctid = key_to_ctid(key); + unsafe { callback(&mut ctid, callback_state) } + }; crate::index::algorithm::bulkdelete(opfamily, index, check, callback); stats } @@ -296,10 +301,9 @@ pub unsafe extern "C" fn amrescan( let relation = PostgresRelation::new((*scan).indexRelation); let options = SearchOptions { epsilon: gucs::epsilon(), - allows_skipping_rerank: gucs::allows_skipping_rerank(), probes: gucs::probes(), max_scan_tuples: gucs::max_scan_tuples(), - max_maxsim_tuples: gucs::max_maxsim_tuples(), + maxsim_refine: gucs::maxsim_refine(), maxsim_threshold: gucs::maxsim_threshold(), }; let scanner = &mut *(*scan).opaque.cast::(); @@ -377,9 +381,9 @@ pub unsafe extern "C" fn amgettuple( pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); } let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; - if let Some((_, ctid, recheck)) = LazyCell::force_mut(&mut scanner.scanning).next() { + if let Some((_, key, recheck)) = LazyCell::force_mut(&mut scanner.scanning).next() { unsafe { - (*scan).xs_heaptid = ctid; + (*scan).xs_heaptid = key_to_ctid(key); (*scan).xs_recheck = recheck; (*scan).xs_recheckorderby = false; } @@ -395,7 +399,7 @@ pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); } -type Iter = Box>; +type Iter = Box>; pub struct Scanner { pub hack: Option>, @@ -452,8 +456,9 @@ impl Drop for HeapFetcher { } impl SearchFetcher for HeapFetcher { - fn fetch(&mut self, mut ctid: ItemPointerData) -> Option<(&[Datum; 32], &[bool; 32])> { + fn fetch(&mut self, key: [u16; 3]) -> Option<(&[Datum; 32], &[bool; 32])> { unsafe { + let mut ctid = key_to_ctid(key); let table_am = (*self.heap_relation).rd_tableam; let fetch_row_version = (*table_am) .tuple_fetch_row_version @@ -497,45 +502,32 @@ impl SearchFetcher for HeapFetcher { } } -pub const fn pointer_to_ctid(pointer: NonZeroU64) -> (ItemPointerData, u16) { +pub const fn ctid_to_key( + ItemPointerData { + ip_blkid: BlockIdData { bi_hi, bi_lo }, + ip_posid, + }: ItemPointerData, +) -> [u16; 3] { + [bi_hi, bi_lo, ip_posid] +} + +pub const fn key_to_ctid([bi_hi, bi_lo, ip_posid]: [u16; 3]) -> ItemPointerData { + ItemPointerData { + ip_blkid: BlockIdData { bi_hi, bi_lo }, + ip_posid, + } +} + +pub const fn pointer_to_kv(pointer: NonZero) -> ([u16; 3], u16) { let value = pointer.get(); let bi_hi = ((value >> 48) & 0xffff) as u16; let bi_lo = ((value >> 32) & 0xffff) as u16; let ip_posid = ((value >> 16) & 0xffff) as u16; let extra = value as u16; - ( - ItemPointerData { - ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi, bi_lo }, - ip_posid, - }, - extra, - ) -} - -pub const fn ctid_to_pointer(ctid: ItemPointerData, extra: u16) -> NonZeroU64 { - let mut value = 0; - value |= (ctid.ip_blkid.bi_hi as u64) << 48; - value |= (ctid.ip_blkid.bi_lo as u64) << 32; - value |= (ctid.ip_posid as u64) << 16; - value |= extra as u64; - NonZeroU64::new(value).expect("invalid pointer") + ([bi_hi, bi_lo, ip_posid], extra) } -#[test] -const fn soundness_check() { - let bi_hi = 1; - let bi_lo = 2; - let ip_posid = 3; - let extra = 7; - let r = pointer_to_ctid(ctid_to_pointer( - ItemPointerData { - ip_blkid: pgrx::pg_sys::BlockIdData { bi_hi, bi_lo }, - ip_posid, - }, - extra, - )); - assert!(r.0.ip_blkid.bi_hi == bi_hi); - assert!(r.0.ip_blkid.bi_lo == bi_lo); - assert!(r.0.ip_posid == ip_posid); - assert!(r.1 == extra); +pub const fn kv_to_pointer((key, value): ([u16; 3], u16)) -> NonZero { + let x = (key[0] as u64) << 48 | (key[1] as u64) << 32 | (key[2] as u64) << 16 | value as u64; + NonZero::new(x).expect("invalid key") } diff --git a/src/index/gucs.rs b/src/index/gucs.rs index b565225a..4ba75c02 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -1,15 +1,17 @@ use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; use std::ffi::CStr; +static PREWARM_DIM: GucSetting> = + GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); + static PROBES: GucSetting> = GucSetting::>::new(Some(c"")); static EPSILON: GucSetting = GucSetting::::new(1.9); static MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); -static PREWARM_DIM: GucSetting> = - GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); -static MAX_MAXSIM_TUPLES: GucSetting = GucSetting::::new(-1); + +static MAXSIM_REFINE: GucSetting = GucSetting::::new(0); static MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); + static PRERERANK_FILTERING: GucSetting = GucSetting::::new(false); -static ALLOWS_SKIPPING_RERANK: GucSetting = GucSetting::::new(false); pub fn init() { GucRegistry::define_string_guc( @@ -36,7 +38,7 @@ pub fn init() { "`max_scan_tuples` argument of vchordrq.", &MAX_SCAN_TUPLES, -1, - u16::MAX as _, + i32::MAX, GucContext::Userset, GucFlags::default(), ); @@ -49,11 +51,11 @@ pub fn init() { GucFlags::default(), ); GucRegistry::define_int_guc( - "vchordrq.max_maxsim_tuples", - "`max_maxsim_tuples` argument of vchordrq.", - "`max_maxsim_tuples` argument of vchordrq.", - &MAX_MAXSIM_TUPLES, - -1, + "vchordrq.maxsim_refine", + "`maxsim_refine` argument of vchordrq.", + "`maxsim_refine` argument of vchordrq.", + &MAXSIM_REFINE, + 0, i32::MAX, GucContext::Userset, GucFlags::default(), @@ -76,14 +78,6 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); - GucRegistry::define_bool_guc( - "vchordrq.allows_skipping_rerank", - "`allows_skipping_rerank` argument of vchordrq.", - "`allows_skipping_rerank` argument of vchordrq.", - &ALLOWS_SKIPPING_RERANK, - GucContext::Userset, - GucFlags::default(), - ); unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); @@ -129,14 +123,12 @@ pub fn max_scan_tuples() -> Option { if x < 0 { None } else { Some(x as u32) } } -pub fn max_maxsim_tuples() -> Option { - let x = MAX_MAXSIM_TUPLES.get(); - if x < 0 { None } else { Some(x as u32) } +pub fn maxsim_refine() -> u32 { + MAXSIM_REFINE.get() as u32 } pub fn maxsim_threshold() -> u32 { - let x = MAXSIM_THRESHOLD.get(); - x as u32 + MAXSIM_THRESHOLD.get() as u32 } pub fn prewarm_dim() -> Vec { @@ -163,7 +155,3 @@ pub fn prewarm_dim() -> Vec { pub fn prererank_filtering() -> bool { PRERERANK_FILTERING.get() } - -pub fn allows_skipping_rerank() -> bool { - ALLOWS_SKIPPING_RERANK.get() -} diff --git a/src/index/scanners/default.rs b/src/index/scanners/default.rs index 493987a0..615d961d 100644 --- a/src/index/scanners/default.rs +++ b/src/index/scanners/default.rs @@ -1,13 +1,12 @@ use super::{SearchBuilder, SearchFetcher, SearchOptions}; use crate::index::algorithm::RandomProject; -use crate::index::am::pointer_to_ctid; +use crate::index::am::pointer_to_kv; use crate::index::opclass::{Opfamily, Sphere}; -use algorithm::operator::{Dot, L2, Op, Vector}; +use algorithm::operator::{Dot, L2, Op}; use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; use algorithm::{RelationRead, RerankMethod}; use half::f16; -use pgrx::pg_sys::ItemPointerData; -use std::num::NonZeroU64; +use std::num::NonZero; use vector::VectorOwned; use vector::vect::VectOwned; @@ -54,7 +53,7 @@ impl SearchBuilder for DefaultBuilder { relation: impl RelationRead + 'a, options: SearchOptions, mut fetcher: impl SearchFetcher + 'a, - ) -> Box + 'a> { + ) -> Box + 'a> { let mut vector = None; let mut threshold = None; let mut recheck = false; @@ -74,33 +73,41 @@ impl SearchBuilder for DefaultBuilder { } let opfamily = self.opfamily; let Some(vector) = vector else { - return Box::new(std::iter::empty()) - as Box>; + return Box::new(std::iter::empty()) as Box>; }; - let iter: Box> = + let iter: Box)>> = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { let vector = RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), + if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() + } + .as_borrowed(), ); - let results = algorithm::search::, L2>>( + let results = algorithm::default_search::, L2>>( relation.clone(), vector.clone(), options.probes, options.epsilon, ); let fetch = move |payload| { - let (ctid, _) = pointer_to_ctid(payload); - let (datums, is_nulls) = fetcher.fetch(ctid)?; + let (key, _) = pointer_to_kv(payload); + let (datums, is_nulls) = fetcher.fetch(key)?; let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; - let raw = VectOwned::::from_owned(maybe_vector.unwrap()); + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; Some(RandomProject::project(raw.as_borrowed())) }; let method = algorithm::how(relation.clone()); match method { RerankMethod::Index => Box::new( - algorithm::rerank_index::, L2>>( + algorithm::rerank_index::, L2>, _>( relation, vector, results, ) .map(move |(distance, payload)| (opfamily.output(distance), payload)), @@ -115,26 +122,35 @@ impl SearchBuilder for DefaultBuilder { } (VectorKind::Vecf32, DistanceKind::Dot) => { let vector = RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), + if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() + } + .as_borrowed(), ); - let results = algorithm::search::, Dot>>( + let results = algorithm::default_search::, Dot>>( relation.clone(), vector.clone(), options.probes, options.epsilon, ); let fetch = move |payload| { - let (ctid, _) = pointer_to_ctid(payload); - let (datums, is_nulls) = fetcher.fetch(ctid)?; + let (key, _) = pointer_to_kv(payload); + let (datums, is_nulls) = fetcher.fetch(key)?; let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; - let raw = VectOwned::::from_owned(maybe_vector.unwrap()); + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; Some(RandomProject::project(raw.as_borrowed())) }; let method = algorithm::how(relation.clone()); match method { RerankMethod::Index => Box::new( - algorithm::rerank_index::, Dot>>( + algorithm::rerank_index::, Dot>, _>( relation, vector, results, ) .map(move |(distance, payload)| (opfamily.output(distance), payload)), @@ -149,26 +165,35 @@ impl SearchBuilder for DefaultBuilder { } (VectorKind::Vecf16, DistanceKind::L2) => { let vector = RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), + if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() + } + .as_borrowed(), ); - let results = algorithm::search::, L2>>( + let results = algorithm::default_search::, L2>>( relation.clone(), vector.clone(), options.probes, options.epsilon, ); let fetch = move |payload| { - let (ctid, _) = pointer_to_ctid(payload); - let (datums, is_nulls) = fetcher.fetch(ctid)?; + let (key, _) = pointer_to_kv(payload); + let (datums, is_nulls) = fetcher.fetch(key)?; let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; - let raw = VectOwned::::from_owned(maybe_vector.unwrap()); + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; Some(RandomProject::project(raw.as_borrowed())) }; let method = algorithm::how(relation.clone()); match method { RerankMethod::Index => Box::new( - algorithm::rerank_index::, L2>>( + algorithm::rerank_index::, L2>, _>( relation, vector, results, ) .map(move |(distance, payload)| (opfamily.output(distance), payload)), @@ -183,26 +208,35 @@ impl SearchBuilder for DefaultBuilder { } (VectorKind::Vecf16, DistanceKind::Dot) => { let vector = RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), + if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() + } + .as_borrowed(), ); - let results = algorithm::search::, Dot>>( + let results = algorithm::default_search::, Dot>>( relation.clone(), vector.clone(), options.probes, options.epsilon, ); let fetch = move |payload| { - let (ctid, _) = pointer_to_ctid(payload); - let (datums, is_nulls) = fetcher.fetch(ctid)?; + let (key, _) = pointer_to_kv(payload); + let (datums, is_nulls) = fetcher.fetch(key)?; let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; - let raw = VectOwned::::from_owned(maybe_vector.unwrap()); + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; Some(RandomProject::project(raw.as_borrowed())) }; let method = algorithm::how(relation.clone()); match method { RerankMethod::Index => Box::new( - algorithm::rerank_index::, Dot>>( + algorithm::rerank_index::, Dot>, _>( relation, vector, results, ) .map(move |(distance, payload)| (opfamily.output(distance), payload)), @@ -226,6 +260,9 @@ impl SearchBuilder for DefaultBuilder { } else { iter }; - Box::new(iter.map(move |(x, y)| (x, pointer_to_ctid(y).0, recheck))) + Box::new(iter.map(move |(distance, pointer)| { + let (key, _) = pointer_to_kv(pointer); + (distance, key, recheck) + })) } } diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs index 3785db9d..d1117955 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/scanners/maxsim.rs @@ -1,16 +1,15 @@ use super::{SearchBuilder, SearchFetcher, SearchOptions}; use crate::index::algorithm::RandomProject; -use crate::index::am::pointer_to_ctid; +use crate::index::am::pointer_to_kv; use crate::index::opclass::Opfamily; -use algorithm::operator::{Dot, Op, Vector}; +use algorithm::operator::Dot; use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; use algorithm::{RelationRead, RerankMethod}; use always_equal::AlwaysEqual; use distance::Distance; use half::f16; -use pgrx::pg_sys::{BlockIdData, ItemPointerData}; use std::cmp::Reverse; -use std::collections::{BinaryHeap, HashMap}; +use std::collections::BinaryHeap; use vector::VectorOwned; use vector::vect::VectOwned; @@ -46,7 +45,7 @@ impl SearchBuilder for MaxsimBuilder { relation: impl RelationRead + 'a, options: SearchOptions, _: impl SearchFetcher + 'a, - ) -> Box + 'a> { + ) -> Box + 'a> { let mut vectors = None; for orderby_vectors in self.orderbys.into_iter().flatten() { if vectors.is_none() { @@ -58,179 +57,190 @@ impl SearchBuilder for MaxsimBuilder { if let Some(_max_scan_tuples) = options.max_scan_tuples { pgrx::error!("maxsim search with max_scan_tuples is not supported"); } - let max_maxsim_tuples = options - .max_maxsim_tuples - .expect("max_maxsim_tuples must be set for maxsim search"); + let maxsim_refine = options.maxsim_refine; let maxsim_threshold = options.maxsim_threshold; let opfamily = self.opfamily; let Some(vectors) = vectors else { - return Box::new(std::iter::empty()) - as Box>; + return Box::new(std::iter::empty()) as Box>; }; let method = algorithm::how(relation.clone()); if !matches!(method, RerankMethod::Index) { pgrx::error!("maxsim search with rerank_in_table is not supported"); } assert!(matches!(opfamily.distance_kind(), DistanceKind::Dot)); - let (c, estimations) = match opfamily.vector_kind() { + let n = vectors.len(); + let accu_map = |(Reverse(dis), AlwaysEqual(pay))| (dis, pay); + let rough_map = |(_, AlwaysEqual(dis), AlwaysEqual(pay), _)| (dis, pay); + let iter: Box> = match opfamily.vector_kind() { VectorKind::Vecf32 => { + type Op = algorithm::operator::Op, Dot>; let vectors = vectors .into_iter() .map(|vector| { RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), + if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() + } + .as_borrowed(), ) }) .collect::>(); - let mut estimations = Vec::new(); - let mut c = HashMap::<[u16; 3], States>::new(); - for (query_id, vector) in vectors.iter().enumerate() { - let (results, est) = algorithm::search_and_estimate::, Dot>>( + Box::new(vectors.into_iter().map(|vector| { + let (results, estimation) = algorithm::maxsim_search::( relation.clone(), vector.clone(), options.probes.clone(), options.epsilon, maxsim_threshold, ); - if results.is_empty() { - estimations.push(0.0); - continue; - } - let returning = if options.epsilon == 0.0 && options.allows_skipping_rerank { - algorithm::skip(results) - .take(max_maxsim_tuples as usize) - .collect::>() - } else { - algorithm::rerank_index::, Dot>>( + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let mut reranker = algorithm::rerank_index::( relation.clone(), vector.clone(), results, - ) - .take(max_maxsim_tuples as usize) - .collect::>() - }; - let mut max = f32::NEG_INFINITY; - for (distance, payload) in returning { - max = max.max(distance.to_f32()); - use std::collections::hash_map::Entry::{Occupied, Vacant}; - let (ctid, _) = pointer_to_ctid(payload); - let key = [ctid.ip_blkid.bi_hi, ctid.ip_blkid.bi_lo, ctid.ip_posid]; - match c.entry(key) { - Vacant(e) => { - let states = e.insert(States::new(vectors.len())); - states.update(query_id, distance); - } - Occupied(mut e) => { - let states = e.get_mut(); - states.update(query_id, distance); - } - } + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.map(rough_map)); + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - estimations.push(max.max(est.to_f32())); - } - (c, estimations) + (accu_set, rough_set, estimation) + })) } VectorKind::Vecf16 => { + type Op = algorithm::operator::Op, Dot>; let vectors = vectors .into_iter() .map(|vector| { RandomProject::project( - VectOwned::::from_owned(vector.clone()).as_borrowed(), + if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() + } + .as_borrowed(), ) }) .collect::>(); - let mut estimations = Vec::new(); - let mut c = HashMap::<[u16; 3], States>::new(); - for (query_id, vector) in vectors.iter().enumerate() { - let (results, est) = algorithm::search_and_estimate::, Dot>>( + Box::new(vectors.into_iter().map(|vector| { + let (results, estimation) = algorithm::maxsim_search::( relation.clone(), vector.clone(), options.probes.clone(), options.epsilon, maxsim_threshold, ); - if results.is_empty() { - estimations.push(0.0); - continue; - } - let returning = if options.epsilon == 0.0 && options.allows_skipping_rerank { - algorithm::skip(results) - .take(max_maxsim_tuples as usize) - .collect::>() - } else { - algorithm::rerank_index::, Dot>>( + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let mut reranker = algorithm::rerank_index::( relation.clone(), vector.clone(), results, - ) - .take(max_maxsim_tuples as usize) - .collect::>() - }; - let mut max = f32::NEG_INFINITY; - for (distance, payload) in returning { - max = max.max(distance.to_f32()); - use std::collections::hash_map::Entry::{Occupied, Vacant}; - let (ctid, _) = pointer_to_ctid(payload); - let key = [ctid.ip_blkid.bi_hi, ctid.ip_blkid.bi_lo, ctid.ip_posid]; - match c.entry(key) { - Vacant(e) => { - let states = e.insert(States::new(vectors.len())); - states.update(query_id, distance); - } - Occupied(mut e) => { - let states = e.get_mut(); - states.update(query_id, distance); - } - } + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.map(rough_map)); + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - estimations.push(max.max(est.to_f32())); - } - (c, estimations) + (accu_set, rough_set, estimation) + })) } }; - let c = c - .into_iter() - .map(|(key, states)| { + let mut updates = Vec::new(); + let mut estimations = Vec::new(); + for (query_id, (accu_set, rough_set, estimation)) in iter.enumerate() { + updates.reserve(accu_set.len() + rough_set.len()); + let is_empty = accu_set.is_empty() && rough_set.is_empty(); + let mut estimation_by_accu = f32::NEG_INFINITY; + for (distance, payload) in accu_set { + estimation_by_accu = estimation_by_accu.max(distance.to_f32()); + let (key, _) = pointer_to_kv(payload); + updates.push((key, AlwaysEqual((query_id, distance)))); + } + for (distance, payload) in rough_set { + let (key, _) = pointer_to_kv(payload); + updates.push((key, AlwaysEqual((query_id, distance)))); + } + estimations.push(if !is_empty { + estimation_by_accu.max(estimation.to_f32()) + } else { + 0.0 + }); + } + updates.sort_unstable(); + let iter = updates + .chunk_by(|(x, _), (y, _)| x == y) + .map(|chunk| { + let key = chunk[0].0; + let mut value = vec![None; n]; + for &(_, AlwaysEqual((query_id, distance))) in chunk { + let this = value[query_id].get_or_insert(Distance::INFINITY); + *this = std::cmp::min(*this, distance); + } let mut maxsim = 0.0f32; - for (query_id, distance) in states.into_iter() { + for (query_id, distance) in value.into_iter().enumerate() { let d = distance.unwrap_or(Distance::from_f32(estimations[query_id])); - maxsim += opfamily.output(d); + maxsim += Distance::to_f32(d); } (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) }) - .collect::>(); - let mut c = BinaryHeap::from(c); - Box::new(std::iter::from_fn(move || { - let (Reverse(distance), AlwaysEqual(key)) = c.pop()?; - let distance = distance.to_f32(); - let ctid = ItemPointerData { - ip_blkid: BlockIdData { - bi_hi: key[0], - bi_lo: key[1], - }, - ip_posid: key[2], - }; - let recheck = false; - Some((distance, ctid, recheck)) - })) + .collect::>() + .into_iter_sorted_polyfill() + .map(|(Reverse(distance), AlwaysEqual(key))| { + let distance = distance.to_f32(); + let recheck = false; + (distance, key, recheck) + }); + let iter: Box> = Box::new(iter); + let iter = if let Some(max_scan_tuples) = options.max_scan_tuples { + Box::new(iter.take(max_scan_tuples as _)) + } else { + iter + }; + #[allow(clippy::let_and_return)] + iter } } -struct States { - inner: Vec>, +pub trait IntoIterSortedPolyfill { + fn into_iter_sorted_polyfill(self) -> IntoIterSorted; } -impl States { - fn new(len: usize) -> Self { - Self { - inner: vec![None; len], - } +impl IntoIterSortedPolyfill for BinaryHeap { + fn into_iter_sorted_polyfill(self) -> IntoIterSorted { + IntoIterSorted { inner: self } } - fn update(&mut self, query_id: usize, distance: Distance) { - let this = self.inner[query_id].get_or_insert(Distance::INFINITY); - *this = std::cmp::min(*this, distance); +} + +#[derive(Clone, Debug)] +pub struct IntoIterSorted { + inner: BinaryHeap, +} + +impl Iterator for IntoIterSorted { + type Item = T; + + #[inline] + fn next(&mut self) -> Option { + self.inner.pop() } - fn into_iter(self) -> impl Iterator)> { - self.inner.into_iter().enumerate() + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len(); + (exact, Some(exact)) } } + +impl ExactSizeIterator for IntoIterSorted {} + +impl std::iter::FusedIterator for IntoIterSorted {} diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs index e572bf0d..f7b46d8a 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/scanners/mod.rs @@ -3,7 +3,7 @@ mod maxsim; use super::opclass::Opfamily; use algorithm::RelationRead; -use pgrx::pg_sys::{Datum, ItemPointerData}; +use pgrx::pg_sys::Datum; use std::cell::LazyCell; pub use default::DefaultBuilder; @@ -12,10 +12,9 @@ pub use maxsim::MaxsimBuilder; #[derive(Debug)] pub struct SearchOptions { pub epsilon: f32, - pub allows_skipping_rerank: bool, pub probes: Vec, pub max_scan_tuples: Option, - pub max_maxsim_tuples: Option, + pub maxsim_refine: u32, pub maxsim_threshold: u32, } @@ -29,15 +28,15 @@ pub trait SearchBuilder: 'static { relation: impl RelationRead + 'a, options: SearchOptions, fetcher: impl SearchFetcher + 'a, - ) -> Box + 'a>; + ) -> Box + 'a>; } pub trait SearchFetcher { - fn fetch(&mut self, ctid: ItemPointerData) -> Option<(&[Datum; 32], &[bool; 32])>; + fn fetch(&mut self, ctid: [u16; 3]) -> Option<(&[Datum; 32], &[bool; 32])>; } impl T> SearchFetcher for LazyCell { - fn fetch(&mut self, ctid: ItemPointerData) -> Option<(&[Datum; 32], &[bool; 32])> { - LazyCell::force_mut(self).fetch(ctid) + fn fetch(&mut self, key: [u16; 3]) -> Option<(&[Datum; 32], &[bool; 32])> { + LazyCell::force_mut(self).fetch(key) } } diff --git a/tests/logic/multivector.slt b/tests/logic/multivector.slt index 1de81a50..5271e46c 100644 --- a/tests/logic/multivector.slt +++ b/tests/logic/multivector.slt @@ -21,7 +21,7 @@ statement ok SET vchordrq.probes = ''; statement ok -SET vchordrq.max_maxsim_tuples = 3000; +SET vchordrq.maxsim_refine = 3000; statement ok SET enable_seqscan TO off; From f04846a372a78ae98b6ea53ab58682a92b5e1027 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 9 Apr 2025 18:48:09 +0800 Subject: [PATCH 141/324] fix: apply epsilon for default search (#228) Signed-off-by: usamoi --- crates/algorithm/src/search.rs | 51 ++++++++++++---------------------- src/index/scanners/maxsim.rs | 30 ++++++++++---------- 2 files changed, 32 insertions(+), 49 deletions(-) diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index 922f726e..416bd453 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -125,7 +125,7 @@ pub fn default_search( }) }; for i in (1..height_of_root).rev() { - state = step(state).take(probes[i as usize - 1] as usize).collect(); + state = step(state).take(probes[i as usize - 1] as _).collect(); } let mut results = LinkedVec::new(); @@ -142,7 +142,7 @@ pub fn default_search( let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); let mut callback = |(rough, err), mean, payload| { - let lowerbound = Distance::from_f32(rough - err * 1.9); + let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( Reverse(lowerbound), AlwaysEqual(()), @@ -173,7 +173,7 @@ pub fn maxsim_search( vector: O::Vector, probes: Vec, epsilon: f32, - mut t: u32, + mut threshold: u32, ) -> (Vec>, Distance) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); @@ -277,22 +277,10 @@ pub fn maxsim_search( Some((first, mean, distance)) }) }; - for i in (2..height_of_root).rev() { - state = step(state) - .map(|(x, y, _)| (x, y)) - .take(probes[i as usize - 1] as usize) - .collect(); - } - let mut iter: Box>; - if height_of_root > 1 { - let mut it = step(state); - state = std::iter::from_fn(|| it.next()) - .map(|(x, y, _)| (x, y)) - .take(probes[0] as usize) - .collect(); - iter = Box::new(it.map(|(x, _, z)| (x, z))); - } else { - iter = Box::new(std::iter::empty()); + let mut it = None; + for i in (1..height_of_root).rev() { + let it = it.insert(step(state)).map(|(first, mean, _)| (first, mean)); + state = it.take(probes[i as usize - 1] as _).collect(); } let mut results = LinkedVec::new(); @@ -332,25 +320,20 @@ pub fn maxsim_search( &mut callback, |_| (), ); - t = t.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as u32); + threshold = threshold.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as _); } - let mut estimation = f32::NAN; - loop { - if t != 0 { - if let Some((first, distance)) = iter.next() { - let jump_guard = index.read(first); - let jump_bytes = jump_guard.get(1).expect("data corruption"); - let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - t = t.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as u32); - estimation = distance.to_f32(); - } else { - break; - } - } else { + let mut estimation_by_threshold = Distance::NEG_INFINITY; + for (first, _, distance) in it.into_iter().flatten() { + if threshold == 0 { break; } + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + threshold = threshold.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as _); + estimation_by_threshold = distance; } - (results.into_vec(), Distance::from_f32(estimation)) + (results.into_vec(), estimation_by_threshold) } fn pop_if( diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs index d1117955..8cfc8d16 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/scanners/maxsim.rs @@ -88,7 +88,7 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new(vectors.into_iter().map(|vector| { - let (results, estimation) = algorithm::maxsim_search::( + let (results, estimation_by_threshold) = algorithm::maxsim_search::( relation.clone(), vector.clone(), options.probes.clone(), @@ -110,7 +110,7 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } - (accu_set, rough_set, estimation) + (accu_set, rough_set, estimation_by_threshold) })) } VectorKind::Vecf16 => { @@ -129,7 +129,7 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new(vectors.into_iter().map(|vector| { - let (results, estimation) = algorithm::maxsim_search::( + let (results, estimation_by_threshold) = algorithm::maxsim_search::( relation.clone(), vector.clone(), options.probes.clone(), @@ -151,44 +151,44 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } - (accu_set, rough_set, estimation) + (accu_set, rough_set, estimation_by_threshold) })) } }; let mut updates = Vec::new(); let mut estimations = Vec::new(); - for (query_id, (accu_set, rough_set, estimation)) in iter.enumerate() { + for (query_id, (accu_set, rough_set, estimation_by_threshold)) in iter.enumerate() { updates.reserve(accu_set.len() + rough_set.len()); let is_empty = accu_set.is_empty() && rough_set.is_empty(); - let mut estimation_by_accu = f32::NEG_INFINITY; + let mut estimation_by_scope = Distance::NEG_INFINITY; for (distance, payload) in accu_set { - estimation_by_accu = estimation_by_accu.max(distance.to_f32()); + estimation_by_scope = std::cmp::max(estimation_by_scope, distance); let (key, _) = pointer_to_kv(payload); - updates.push((key, AlwaysEqual((query_id, distance)))); + updates.push((key, query_id, distance)); } for (distance, payload) in rough_set { let (key, _) = pointer_to_kv(payload); - updates.push((key, AlwaysEqual((query_id, distance)))); + updates.push((key, query_id, distance)); } estimations.push(if !is_empty { - estimation_by_accu.max(estimation.to_f32()) + std::cmp::max(estimation_by_scope, estimation_by_threshold) } else { - 0.0 + Distance::ZERO }); } - updates.sort_unstable(); + updates.sort_unstable_by_key(|&(key, ..)| key); let iter = updates - .chunk_by(|(x, _), (y, _)| x == y) + .chunk_by(|(kl, ..), (kr, ..)| kl == kr) .map(|chunk| { let key = chunk[0].0; let mut value = vec![None; n]; - for &(_, AlwaysEqual((query_id, distance))) in chunk { + for &(_, query_id, distance) in chunk { let this = value[query_id].get_or_insert(Distance::INFINITY); *this = std::cmp::min(*this, distance); } let mut maxsim = 0.0f32; for (query_id, distance) in value.into_iter().enumerate() { - let d = distance.unwrap_or(Distance::from_f32(estimations[query_id])); + let d = distance.unwrap_or(estimations[query_id]); maxsim += Distance::to_f32(d); } (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) From 86a968ee9fc101c57cc0293a92364d228b150847 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 11 Apr 2025 12:28:14 +0800 Subject: [PATCH 142/324] chore: remove patch of crate half (#231) Signed-off-by: usamoi --- Cargo.lock | 13 +++++++------ Cargo.toml | 5 +---- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4819ae4d..6fe91902 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,7 +17,7 @@ version = "0.0.0" dependencies = [ "always_equal", "distance", - "half 2.5.0", + "half 2.6.0", "k_means", "paste", "rabitq", @@ -371,8 +371,9 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403" [[package]] name = "half" -version = "2.5.0" -source = "git+https://github.com/tensorchord/half-rs.git?rev=9dc2559e95f295aaa8ffef5e0529167731dffd47#9dc2559e95f295aaa8ffef5e0529167731dffd47" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "459196ed295495a68f7d7fe1d84f6c4b7ff0e21fe3017b2f283c6fac3ad803c9" dependencies = [ "cfg-if", "crunchy", @@ -1244,7 +1245,7 @@ name = "simd" version = "0.0.0" dependencies = [ "cc", - "half 2.5.0", + "half 2.6.0", "rand", "simd_macros", "zerocopy", @@ -1498,7 +1499,7 @@ dependencies = [ "algorithm", "always_equal", "distance", - "half 2.5.0", + "half 2.6.0", "k_means", "mimalloc", "paste", @@ -1520,7 +1521,7 @@ name = "vector" version = "0.0.0" dependencies = [ "distance", - "half 2.5.0", + "half 2.6.0", "simd", ] diff --git a/Cargo.toml b/Cargo.toml index c29498d4..7941c213 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,9 +42,6 @@ zerocopy.workspace = true [target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos", target_os = "windows")))'.dependencies] mimalloc = { version = "*", features = ["local_dynamic_tls"] } -[patch.crates-io] -half = { git = "https://github.com/tensorchord/half-rs.git", rev = "9dc2559e95f295aaa8ffef5e0529167731dffd47" } - [lints] workspace = true @@ -57,7 +54,7 @@ version = "0.0.0" edition = "2024" [workspace.dependencies] -half = { version = "2.5.0", features = ["zerocopy"] } +half = { version = "2.6.0", features = ["zerocopy"] } paste = "1" rand = "0.9.0" serde = { version = "1", features = ["derive"] } From 823056000f9f4bc69c0e4ee8500aa50d09681603 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Fri, 11 Apr 2025 14:57:56 +0800 Subject: [PATCH 143/324] fix: release env vars (#232) Signed-off-by: cutecutecat --- .github/workflows/release.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 8c8a06a6..486ae9e4 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -118,6 +118,10 @@ jobs: - name: Upload Artifacts env: GH_TOKEN: ${{ github.token }} + SEMVER: ${{ needs.semver.outputs.SEMVER }} + VERSION: ${{ matrix.version }} + ARCH: ${{ matrix.arch }} + PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} run: | gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip From ad923d2818fcb9f9c0328d427a24006e5487189e Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 11 Apr 2025 15:52:26 +0800 Subject: [PATCH 144/324] chore: update 0.3.0 schema script (#230) Signed-off-by: usamoi --- .github/workflows/check.yml | 2 +- .github/workflows/release.yml | 12 +- crates/algorithm/src/tuples.rs | 2 +- scripts/README.md | 2 +- sql/install/vchord--0.3.0.sql | 566 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.2.2--0.3.0.sql | 51 +++ tools/dev.md | 4 +- vchord.control | 2 +- 8 files changed, 633 insertions(+), 8 deletions(-) create mode 100644 sql/install/vchord--0.3.0.sql create mode 100644 sql/upgrade/vchord--0.2.2--0.3.0.sql diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 284024f7..5a1f386e 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -187,7 +187,7 @@ jobs: mkdir -p ./build/zip cp -a ./sql/upgrade/. ./build/zip/ cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql - sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control + cp ./vchord.control ./build/zip/vchord.control cp ./target/release/libvchord.so ./build/zip/vchord.so zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 486ae9e4..36e1359f 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -71,6 +71,12 @@ jobs: - name: Checkout uses: actions/checkout@v4 + - name: Check + env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + run: | + grep -q "default_version = '${SEMVER}'" vchord.control || exit 1 + - name: Build env: SEMVER: ${{ needs.semver.outputs.SEMVER }} @@ -83,7 +89,7 @@ jobs: mkdir -p ./build/zip cp -a ./sql/upgrade/. ./build/zip/ cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql - sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control + cp ./vchord.control ./build/zip/vchord.control cp ./target/release/libvchord.so ./build/zip/vchord.so zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* @@ -187,7 +193,7 @@ jobs: PG_VERSION=${{ matrix.version }} SEMVER=${{ env.SEMVER }} PGVECTOR=0.8.0 - + create-manifests: runs-on: ubuntu-latest needs: ["semver", "build", "docker"] @@ -228,4 +234,4 @@ jobs: ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} \ --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 - docker manifest push ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} \ No newline at end of file + docker manifest push ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 849cb474..6b7d0ebf 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -8,7 +8,7 @@ use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: u64 = u64::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 4; +const VERSION: u64 = 5; pub trait Tuple: 'static { fn serialize(&self) -> Vec; diff --git a/scripts/README.md b/scripts/README.md index dbfd8a32..07ca423b 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -31,7 +31,7 @@ cargo build --lib --features pg$VERSION --release mkdir -p ./build/zip cp -a ./sql/upgrade/. ./build/zip/ cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql -sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control +cp ./vchord.control ./build/zip/vchord.control cp ./target/release/libvchord.so ./build/zip/vchord.so zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* diff --git a/sql/install/vchord--0.3.0.sql b/sql/install/vchord--0.3.0.sql new file mode 100644 index 00000000..7d1088fd --- /dev/null +++ b/sql/install/vchord--0.3.0.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:9 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:79 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:55 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:31 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:7 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:79 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:55 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:31 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:7 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:60 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:5 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:36 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:31 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:46 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:41 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:10 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.2.2--0.3.0.sql b/sql/upgrade/vchord--0.2.2--0.3.0.sql new file mode 100644 index 00000000..bbbc98bf --- /dev/null +++ b/sql/upgrade/vchord--0.2.2--0.3.0.sql @@ -0,0 +1,51 @@ +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], + "rhs" halfvec[] +) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; + +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], + "rhs" vector[] +) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; + +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; + +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; + +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); diff --git a/tools/dev.md b/tools/dev.md index 70c0107d..7ce13007 100644 --- a/tools/dev.md +++ b/tools/dev.md @@ -71,6 +71,8 @@ sudo -u postgres dropdb vchord ### Step 4 - package and download +Edit `./vchord.control` and bump the version. + ```shell SEMVER='0.2.1' VERSION='17' @@ -82,7 +84,7 @@ cargo build --lib --features pg$VERSION --release mkdir -p ./build/zip cp -a ./sql/upgrade/. ./build/zip/ cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql -sed -e "s/@CARGO_VERSION@/$SEMVER/g" < ./vchord.control > ./build/zip/vchord.control +cp ./vchord.control ./build/zip/vchord.control cp ./target/release/libvchord.so ./build/zip/vchord.so zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* diff --git a/vchord.control b/vchord.control index 314ed5df..5ce41ec2 100644 --- a/vchord.control +++ b/vchord.control @@ -1,5 +1,5 @@ comment = 'vchord: Vector database plugin for Postgres, written in Rust, specifically designed for LLM' -default_version = '@CARGO_VERSION@' +default_version = '0.0.0' module_pathname = '$libdir/vchord' relocatable = true superuser = true From 7f3f3a6465818b6fb0b798c061092e8de20994fd Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 14 Apr 2025 12:53:53 +0800 Subject: [PATCH 145/324] chore: update pgrx to 0.14.1 (#233) Signed-off-by: usamoi --- .github/workflows/check.yml | 2 +- .github/workflows/release.yml | 2 +- Cargo.lock | 124 ++++++++++++++++++++-------------- Cargo.toml | 6 +- crates/simd/Cargo.toml | 2 +- src/index/am/am_build.rs | 10 +-- src/index/am/mod.rs | 25 ++++--- src/index/hook.rs | 8 +-- src/lib.rs | 2 +- src/upgrade.rs | 2 +- 10 files changed, 103 insertions(+), 80 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 5a1f386e..2927d4a6 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -152,7 +152,7 @@ jobs: sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.so"' sudo systemctl stop postgresql - curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.13.1/cargo-pgrx-v0.13.1-$(uname -m)-unknown-linux-gnu.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx + curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.14.1/cargo-pgrx-v0.14.1-$(uname -m)-unknown-linux-gnu.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 36e1359f..f2002726 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -65,7 +65,7 @@ jobs: sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector - curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.13.1/cargo-pgrx-v0.13.1-$(uname -m)-unknown-linux-gnu.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx + curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.14.1/cargo-pgrx-v0.14.1-$(uname -m)-unknown-linux-gnu.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) - name: Checkout diff --git a/Cargo.lock b/Cargo.lock index 6fe91902..7df3ef07 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -31,6 +31,12 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "always_equal" version = "0.0.0" @@ -53,9 +59,9 @@ checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9" [[package]] name = "anyhow" -version = "1.0.97" +version = "1.0.98" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dcfed56ad506cb2c684a14971b8861fdc3baaaae314b9e5f9bb532cbe3ba7a4f" +checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" [[package]] name = "approx" @@ -133,9 +139,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "cargo_toml" -version = "0.21.0" +version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fbd1fe9db3ebf71b89060adaf7b0504c2d6a425cf061313099547e382c2e472" +checksum = "02260d489095346e5cafd04dea8e8cb54d1d74fcd759022a9b72986ebe9a1257" dependencies = [ "serde", "toml", @@ -143,9 +149,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.16" +version = "1.2.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be714c154be609ec7f5dad223a33bf1482fff90472de28f7362806e6d4832b8c" +checksum = "8e3a13707ac958681c13b39b458c073d0d9bc8a22cb1b2f4c8e55eb72c13f362" dependencies = [ "shlex", ] @@ -188,9 +194,9 @@ dependencies = [ [[package]] name = "convert_case" -version = "0.7.1" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb402b8d4c85569410425650ce3eddc7d698ed96d39a73f941b08fb63082f1e7" +checksum = "baaaa0ecca5b51987b9423ccdc971514dd8b0bb7b4060b983d3664dad3f1f89f" dependencies = [ "unicode-segmentation", ] @@ -228,9 +234,9 @@ checksum = "43da5946c66ffcc7745f48db692ffbb10a83bfe0afd96235c5c2a4fb23994929" [[package]] name = "darling" -version = "0.20.10" +version = "0.20.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f63b86c8a8826a49b8c21f08a2d07338eec8d900540f8630dc76284be802989" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" dependencies = [ "darling_core", "darling_macro", @@ -238,9 +244,9 @@ dependencies = [ [[package]] name = "darling_core" -version = "0.20.10" +version = "0.20.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95133861a8032aaea082871032f5815eb9e98cef03fa916ab4500513994df9e5" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" dependencies = [ "fnv", "ident_case", @@ -252,9 +258,9 @@ dependencies = [ [[package]] name = "darling_macro" -version = "0.20.10" +version = "0.20.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" dependencies = [ "darling_core", "quote", @@ -330,6 +336,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + [[package]] name = "form_urlencoded" version = "1.2.1" @@ -394,6 +406,11 @@ name = "hashbrown" version = "0.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] [[package]] name = "heapless" @@ -461,9 +478,9 @@ dependencies = [ [[package]] name = "icu_locid_transform_data" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdc8ff3388f852bede6b579ad4e978ab004f139284d7b28715f773507b946f6e" +checksum = "7515e6d781098bf9f7205ab3fc7e9709d34554ae0b21ddbcb5febfa4bc7df11d" [[package]] name = "icu_normalizer" @@ -485,9 +502,9 @@ dependencies = [ [[package]] name = "icu_normalizer_data" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8cafbf7aa791e9b22bec55a167906f9e1215fd475cd22adfcf660e03e989516" +checksum = "c5e8338228bdc8ab83303f16b797e177953730f601a96c25d10cb3ab0daa0cb7" [[package]] name = "icu_properties" @@ -506,9 +523,9 @@ dependencies = [ [[package]] name = "icu_properties_data" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67a8effbc3dd3e4ba1afa8ad918d5684b8868b3b26500753effea8d2eed19569" +checksum = "85fb8799753b75aee8d2a21d7c14d9f38921b54b3dbda10f5a3c7a7b82dba5e2" [[package]] name = "icu_provider" @@ -573,9 +590,9 @@ checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" [[package]] name = "indexmap" -version = "2.8.0" +version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3954d50fe15b02142bf25d3b8bdadb634ec3948f103d04ffe3031bc8fe9d7058" +checksum = "cea70ddb795996207ad57735b50c5982d8844f38ba9ee5f1aedcfb708a2aa11e" dependencies = [ "equivalent", "hashbrown", @@ -647,9 +664,9 @@ checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" [[package]] name = "libmimalloc-sys" -version = "0.1.40" +version = "0.1.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07d0e07885d6a754b9c7993f2625187ad694ee985d60f23355ff0e7077261502" +checksum = "ec9d6fac27761dabcd4ee73571cdb06b7022dc99089acbe5435691edffaac0f4" dependencies = [ "cc", "libc", @@ -679,9 +696,9 @@ checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" [[package]] name = "mimalloc" -version = "0.1.44" +version = "0.1.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99585191385958383e13f6b822e6b6d8d9cf928e7d286ceb092da92b43c87bc1" +checksum = "995942f432bbb4822a7e9c3faa87a695185b0d09273ba85f097b54f4e458f2af" dependencies = [ "libmimalloc-sys", ] @@ -780,9 +797,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.21.1" +version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d75b0bedcc4fe52caa0e03d9f1151a323e4aa5e2d78ba3580400cd3c9e2bc4bc" +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" [[package]] name = "owo-colors" @@ -818,19 +835,21 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "petgraph" -version = "0.7.1" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" +checksum = "7a98c6720655620a521dcc722d0ad66cd8afd5d86e34a89ef691c50b7b24de06" dependencies = [ "fixedbitset", + "hashbrown", "indexmap", + "serde", ] [[package]] name = "pgrx" -version = "0.13.1" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74fb2c1fbb9edc39097200f882ee58550e68f519db52309e94d894a2ad1ddf07" +checksum = "ab5a7f1a5bac3a76a4fa056655eb80e394c5216717bf4673c99d9d10583e960d" dependencies = [ "atomic-traits", "bitflags", @@ -852,9 +871,9 @@ dependencies = [ [[package]] name = "pgrx-bindgen" -version = "0.13.1" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7bba29f6257193d1795eee0f96e72d368b30e95ab50d93397b5a76cc7784ddef" +checksum = "09e1a6df74e60ac82d6ad52f8a2a85e5d7004f9dfc7c52c94d461d9193aafa0b" dependencies = [ "bindgen", "cc", @@ -863,6 +882,7 @@ dependencies = [ "pgrx-pg-config", "proc-macro2", "quote", + "regex", "shlex", "syn", "walkdir", @@ -870,9 +890,9 @@ dependencies = [ [[package]] name = "pgrx-catalog" -version = "0.1.1" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "494cd323b4bb0b2befdcdecde2db9d56c3a6bb90b4a6bcfd8f40a747beb39fd1" +checksum = "54c52c31fd64fe1e0cde3826b578bcc343e778e874fc2e520ad324659ff0b120" dependencies = [ "paste", "pgrx", @@ -880,9 +900,9 @@ dependencies = [ [[package]] name = "pgrx-macros" -version = "0.13.1" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1f9d2a6a94d8ea1a722ea489048acd203c6ba9b62f70edac587dbe8e1fb4585" +checksum = "868ae3cb1762dcca0e236d4b5368315d4f377acbfff31d8fad9f47394d2e4bcd" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", @@ -892,9 +912,9 @@ dependencies = [ [[package]] name = "pgrx-pg-config" -version = "0.13.1" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05be49ab66e458f4c164b45daa885e53dbbadc4f37bb47e8b5d85d5b4d9a0be4" +checksum = "421be70840655a502381d4ac84b263da83c6714018e27bfdcd7c48f32ff60ef6" dependencies = [ "cargo_toml", "eyre", @@ -910,9 +930,9 @@ dependencies = [ [[package]] name = "pgrx-pg-sys" -version = "0.13.1" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a360ad31a947674a46e72a39881f1be8afe958157ae44cf8ad6ed70b984c0f40" +checksum = "c844a49ee575cf4bd3340c483e49160a99b1bd8c5fd2e87fc7a9231566cb07f2" dependencies = [ "cee-scape", "libc", @@ -925,9 +945,9 @@ dependencies = [ [[package]] name = "pgrx-sql-entity-graph" -version = "0.13.1" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8f625bb35d59b2d6f308ed5d1b6da3783c8514d47f353ec23a8f05da0fef6f4" +checksum = "391595ab6d13f65728bdd691fdd58e69bb31f1f092886ed0a11c838de1d01ae9" dependencies = [ "convert_case", "eyre", @@ -1262,9 +1282,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.14.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd" +checksum = "8917285742e9f3e1683f0a9c4e6b57960b7314d0b08d30d1ecd426713ee2eee9" [[package]] name = "sptr" @@ -1638,9 +1658,9 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "winnow" -version = "0.7.4" +version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e97b544156e9bebe1a0ffbc03484fc1ffe3100cbce3ffb17eac35f7cdd7ab36" +checksum = "63d3fcd9bba44b03821e7d699eeee959f3126dcc4aa8e4ae18ec617c2a5cea10" dependencies = [ "memchr", ] @@ -1701,18 +1721,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.23" +version = "0.8.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd97444d05a4328b90e75e503a34bad781f14e28a823ad3557f0750df1ebcbc6" +checksum = "2586fea28e186957ef732a5f8b3be2da217d65c5969d4b1e17f973ebbe876879" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.23" +version = "0.8.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6352c01d0edd5db859a63e2605f4ea3183ddbd15e2c4a9e7d32184df75e4f154" +checksum = "a996a8f63c5c4448cd959ac1bab0aaa3306ccfd060472f85943ee0750f0169be" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index 7941c213..9d2d5d0a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,8 +30,8 @@ vector = { path = "./crates/vector" } half.workspace = true paste.workspace = true -pgrx = { version = "=0.13.1", default-features = false, features = ["cshim"] } -pgrx-catalog = "0.1.1" +pgrx = { version = "=0.14.1", default-features = false, features = ["cshim"] } +pgrx-catalog = "0.2.0" rand.workspace = true seq-macro = "0.3.6" serde.workspace = true @@ -59,7 +59,7 @@ paste = "1" rand = "0.9.0" serde = { version = "1", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } -zerocopy = { version = "0.8.23", features = ["derive"] } +zerocopy = { version = "0.8.24", features = ["derive"] } [workspace.lints] clippy.identity_op = "allow" diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index bb421562..01a3b681 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -13,7 +13,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.16" +cc = "1.2.19" [lints] workspace = true diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index dd17a3e9..4f4d9401 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -128,7 +128,7 @@ impl BuildPhase { } #[pgrx::pg_guard] -pub extern "C" fn ambuildphasename(x: i64) -> *mut core::ffi::c_char { +pub extern "C-unwind" fn ambuildphasename(x: i64) -> *mut core::ffi::c_char { if let Ok(x) = u32::try_from(x.wrapping_sub(1)) { if let Some(x) = BuildPhase::from_value(x) { x.name().as_ptr().cast_mut() @@ -160,7 +160,7 @@ impl Heap { pub callback: F, } #[pgrx::pg_guard] - unsafe extern "C" fn call( + unsafe extern "C-unwind" fn call( _index_relation: pgrx::pg_sys::Relation, ctid: pgrx::pg_sys::ItemPointer, values: *mut Datum, @@ -232,7 +232,7 @@ impl PostgresReporter { } #[pgrx::pg_guard] -pub unsafe extern "C" fn ambuild( +pub unsafe extern "C-unwind" fn ambuild( heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, index_info: *mut pgrx::pg_sys::IndexInfo, @@ -785,7 +785,7 @@ impl Drop for VchordrqLeader { #[pgrx::pg_guard] #[unsafe(no_mangle)] -pub unsafe extern "C" fn vchordrq_parallel_build_main( +pub unsafe extern "C-unwind" fn vchordrq_parallel_build_main( _seg: *mut pgrx::pg_sys::dsm_segment, toc: *mut pgrx::pg_sys::shm_toc, ) { @@ -914,7 +914,7 @@ unsafe fn parallel_build( } #[pgrx::pg_guard] -pub unsafe extern "C" fn ambuildempty(_index_relation: pgrx::pg_sys::Relation) { +pub unsafe extern "C-unwind" fn ambuildempty(_index_relation: pgrx::pg_sys::Relation) { pgrx::error!("Unlogged indexes are not supported."); } diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index 126c26b4..74f07612 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -109,12 +109,15 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { }; #[pgrx::pg_guard] -pub unsafe extern "C" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { +pub unsafe extern "C-unwind" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { true } #[pgrx::pg_guard] -pub unsafe extern "C" fn amoptions(reloptions: Datum, validate: bool) -> *mut pgrx::pg_sys::bytea { +pub unsafe extern "C-unwind" fn amoptions( + reloptions: Datum, + validate: bool, +) -> *mut pgrx::pg_sys::bytea { let relopt_kind = RELOPT_KIND.get().copied().expect("init is not called"); let rdopts = unsafe { pgrx::pg_sys::build_reloptions( @@ -131,7 +134,7 @@ pub unsafe extern "C" fn amoptions(reloptions: Datum, validate: bool) -> *mut pg #[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] -pub unsafe extern "C" fn amcostestimate( +pub unsafe extern "C-unwind" fn amcostestimate( _root: *mut pgrx::pg_sys::PlannerInfo, path: *mut pgrx::pg_sys::IndexPath, _loop_count: f64, @@ -161,7 +164,7 @@ pub unsafe extern "C" fn amcostestimate( #[cfg(feature = "pg13")] #[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] -pub unsafe extern "C" fn aminsert( +pub unsafe extern "C-unwind" fn aminsert( index_relation: pgrx::pg_sys::Relation, values: *mut Datum, is_null: *mut bool, @@ -176,7 +179,7 @@ pub unsafe extern "C" fn aminsert( #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] #[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] -pub unsafe extern "C" fn aminsert( +pub unsafe extern "C-unwind" fn aminsert( index_relation: pgrx::pg_sys::Relation, values: *mut Datum, is_null: *mut bool, @@ -210,7 +213,7 @@ unsafe fn aminsertinner( } #[pgrx::pg_guard] -pub unsafe extern "C" fn ambulkdelete( +pub unsafe extern "C-unwind" fn ambulkdelete( info: *mut pgrx::pg_sys::IndexVacuumInfo, stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, callback: pgrx::pg_sys::IndexBulkDeleteCallback, @@ -238,7 +241,7 @@ pub unsafe extern "C" fn ambulkdelete( } #[pgrx::pg_guard] -pub unsafe extern "C" fn amvacuumcleanup( +pub unsafe extern "C-unwind" fn amvacuumcleanup( info: *mut pgrx::pg_sys::IndexVacuumInfo, stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, ) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { @@ -258,7 +261,7 @@ pub unsafe extern "C" fn amvacuumcleanup( } #[pgrx::pg_guard] -pub unsafe extern "C" fn ambeginscan( +pub unsafe extern "C-unwind" fn ambeginscan( index_relation: pgrx::pg_sys::Relation, n_keys: std::os::raw::c_int, n_orderbys: std::os::raw::c_int, @@ -277,7 +280,7 @@ pub unsafe extern "C" fn ambeginscan( } #[pgrx::pg_guard] -pub unsafe extern "C" fn amrescan( +pub unsafe extern "C-unwind" fn amrescan( scan: pgrx::pg_sys::IndexScanDesc, keys: pgrx::pg_sys::ScanKey, _n_keys: std::os::raw::c_int, @@ -365,7 +368,7 @@ pub unsafe extern "C" fn amrescan( } #[pgrx::pg_guard] -pub unsafe extern "C" fn amgettuple( +pub unsafe extern "C-unwind" fn amgettuple( scan: pgrx::pg_sys::IndexScanDesc, direction: pgrx::pg_sys::ScanDirection::Type, ) -> bool { @@ -394,7 +397,7 @@ pub unsafe extern "C" fn amgettuple( } #[pgrx::pg_guard] -pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { +pub unsafe extern "C-unwind" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { let scanner = unsafe { &mut *(*scan).opaque.cast::() }; scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); } diff --git a/src/index/hook.rs b/src/index/hook.rs index ec979cb5..748e939c 100644 --- a/src/index/hook.rs +++ b/src/index/hook.rs @@ -1,12 +1,12 @@ use std::sync::atomic::AtomicPtr; #[pgrx::pg_guard] -unsafe extern "C" fn rewrite_plan_state( +unsafe extern "C-unwind" fn rewrite_plan_state( node: *mut pgrx::pg_sys::PlanState, context: *mut core::ffi::c_void, ) -> bool { unsafe fn dirty_check(index_relation: *mut pgrx::pg_sys::RelationData) -> Option { - type FnPtr = unsafe extern "C" fn( + type FnPtr = unsafe extern "C-unwind" fn( *mut pgrx::pg_sys::RelationData, i32, i32, @@ -58,7 +58,7 @@ unsafe extern "C" fn rewrite_plan_state( static PREV_EXECUTOR_START: AtomicPtr<()> = AtomicPtr::new(core::ptr::null_mut()); #[pgrx::pg_guard] -unsafe extern "C" fn executor_start( +unsafe extern "C-unwind" fn executor_start( query_desc: *mut pgrx::pg_sys::QueryDesc, eflags: ::std::os::raw::c_int, ) { @@ -85,7 +85,7 @@ pub fn init() { use core::mem::transmute; use std::sync::atomic::Ordering; PREV_EXECUTOR_START.store( - transmute::, *mut ()>( + transmute::, *mut ()>( pgrx::pg_sys::ExecutorStart_hook, ), Ordering::Relaxed, diff --git a/src/lib.rs b/src/lib.rs index 116e90d3..e51d6145 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,7 +10,7 @@ pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap); pgrx::extension_sql_file!("./sql/finalize.sql", finalize); #[pgrx::pg_guard] -extern "C" fn _PG_init() { +extern "C-unwind" fn _PG_init() { if unsafe { pgrx::pg_sys::IsUnderPostmaster } { pgrx::error!("vchord must be loaded via shared_preload_libraries."); } diff --git a/src/upgrade.rs b/src/upgrade.rs index b474744b..9879439f 100644 --- a/src/upgrade.rs +++ b/src/upgrade.rs @@ -11,7 +11,7 @@ macro_rules! symbol { #[unsafe(no_mangle)] #[doc(hidden)] #[pgrx::pg_guard] - extern "C" fn [<$t _wrapper>](_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> pgrx::pg_sys::Datum { + extern "C-unwind" fn [<$t _wrapper>](_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> pgrx::pg_sys::Datum { pgrx::error!( "the symbol {} is removed in the extension; please run extension update scripts", stringify!($t), From 16b827ef6f111c55bfd3ac93653faba92ca1e0b0 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 14 Apr 2025 17:47:38 +0800 Subject: [PATCH 146/324] feat: cost estimate (#234) Signed-off-by: usamoi --- crates/algorithm/src/build.rs | 1 + crates/algorithm/src/cost.rs | 26 +++++++ crates/algorithm/src/lib.rs | 2 + crates/algorithm/src/tuples.rs | 9 ++- crates/algorithm/src/types.rs | 9 +++ src/index/am/am_build.rs | 5 +- src/index/am/mod.rs | 127 ++++++++++++++++++++++++++++++--- src/index/functions.rs | 18 ++--- 8 files changed, 173 insertions(+), 24 deletions(-) create mode 100644 crates/algorithm/src/cost.rs diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 8bd619e2..41eab332 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -134,6 +134,7 @@ pub fn build( .last() .expect("internal error: empty structure")[0], freepage_first: freepage.first(), + cells: std::array::from_fn(|i| structures.get(i).map(|s| s.len() as _).unwrap_or_default()), }); } diff --git a/crates/algorithm/src/cost.rs b/crates/algorithm/src/cost.rs new file mode 100644 index 00000000..06c17f09 --- /dev/null +++ b/crates/algorithm/src/cost.rs @@ -0,0 +1,26 @@ +use crate::tuples::{MetaTuple, WithReader}; +use crate::{Page, RelationRead}; +use std::num::NonZero; + +pub struct Cost { + pub dims: u32, + pub is_residual: bool, + pub cells: Vec>, +} + +#[must_use] +pub fn cost(index: impl RelationRead) -> Cost { + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dims = meta_tuple.dims(); + let is_residual = meta_tuple.is_residual(); + let cells = meta_tuple.cells(); + drop(meta_guard); + + Cost { + dims, + is_residual, + cells: cells.into_iter().map_while(NonZero::new).collect(), + } +} diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 5cb12c37..4ba4f456 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -1,6 +1,7 @@ mod build; mod bulkdelete; mod cache; +mod cost; mod freepages; mod insert; mod linked_vec; @@ -19,6 +20,7 @@ pub mod types; pub use build::build; pub use bulkdelete::bulkdelete; pub use cache::cache; +pub use cost::cost; pub use insert::insert; pub use maintain::maintain; pub use prewarm::prewarm; diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 6b7d0ebf..204f14c5 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -8,7 +8,7 @@ use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: u64 = u64::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 5; +const VERSION: u64 = 6; pub trait Tuple: 'static { fn serialize(&self) -> Vec; @@ -40,6 +40,8 @@ struct MetaTupleHeader { // for meta tuple, it's pointers to next level root_first: u32, freepage_first: u32, + // statistics + cells: [u32; 8], } pub struct MetaTuple { @@ -51,6 +53,7 @@ pub struct MetaTuple { pub root_mean: IndexPointer, pub root_first: u32, pub freepage_first: u32, + pub cells: [u32; 8], } impl Tuple for MetaTuple { @@ -67,6 +70,7 @@ impl Tuple for MetaTuple { root_mean: self.root_mean, root_first: self.root_first, freepage_first: self.freepage_first, + cells: self.cells, } .as_bytes() .to_vec() @@ -123,6 +127,9 @@ impl MetaTupleReader<'_> { pub fn freepage_first(self) -> u32 { self.header.freepage_first } + pub fn cells(self) -> [u32; 8] { + self.header.cells + } } #[repr(C, align(8))] diff --git a/crates/algorithm/src/types.rs b/crates/algorithm/src/types.rs index a59de467..f889477b 100644 --- a/crates/algorithm/src/types.rs +++ b/crates/algorithm/src/types.rs @@ -47,6 +47,15 @@ pub enum VectorKind { Vecf16, } +impl VectorKind { + pub fn element_size(self) -> u32 { + match self { + VectorKind::Vecf32 => size_of::() as _, + VectorKind::Vecf16 => size_of::() as _, + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] #[validate(schema(function = "Self::validate_self"))] diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 4f4d9401..913dcead 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -251,6 +251,7 @@ pub unsafe extern "C-unwind" fn ambuild( ); } let opfamily = unsafe { opfamily(index_relation) }; + let index = unsafe { PostgresRelation::new(index_relation) }; let heap = Heap { heap_relation, index_relation, @@ -258,7 +259,6 @@ pub unsafe extern "C-unwind" fn ambuild( opfamily, scan: std::ptr::null_mut(), }; - let index = unsafe { PostgresRelation::new(index_relation) }; let mut reporter = PostgresReporter {}; let structures = match vchordrq_options.build.source.clone() { VchordrqBuildSourceOptions::External(external_build) => { @@ -380,12 +380,11 @@ pub unsafe extern "C-unwind" fn ambuild( } else { let mut indtuples = 0; reporter.tuples_done(indtuples); - let relation = unsafe { PostgresRelation::new(index_relation) }; heap.traverse(true, |(ctid, store)| { for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, relation.clone(), payload, vector); + crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); } indtuples += 1; reporter.tuples_done(indtuples); diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index 74f07612..ebfca1d5 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -1,7 +1,7 @@ pub mod am_build; use crate::index::gucs; -use crate::index::opclass::opfamily; +use crate::index::opclass::{Opfamily, opfamily}; use crate::index::scanners::*; use crate::index::storage::PostgresRelation; use pgrx::datum::Internal; @@ -135,7 +135,7 @@ pub unsafe extern "C-unwind" fn amoptions( #[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn amcostestimate( - _root: *mut pgrx::pg_sys::PlannerInfo, + root: *mut pgrx::pg_sys::PlannerInfo, path: *mut pgrx::pg_sys::IndexPath, _loop_count: f64, index_startup_cost: *mut pgrx::pg_sys::Cost, @@ -145,19 +145,99 @@ pub unsafe extern "C-unwind" fn amcostestimate( index_pages: *mut f64, ) { unsafe { + use pgrx::pg_sys::disable_cost; + let index_opt_info = (*path).indexinfo; + // do not use index, if there are no orderbys or clauses if (*path).indexorderbys.is_null() && (*path).indexclauses.is_null() { - *index_startup_cost = f64::MAX; - *index_total_cost = f64::MAX; + *index_startup_cost = disable_cost; + *index_total_cost = disable_cost; *index_selectivity = 0.0; *index_correlation = 0.0; - *index_pages = 0.0; + *index_pages = 1.0; + return; + } + let selectivity = { + use pgrx::pg_sys::{ + JoinType, add_predicate_to_index_quals, clauselist_selectivity, + get_quals_from_indexclauses, + }; + let index_quals = get_quals_from_indexclauses((*path).indexclauses); + let selectivity_quals = add_predicate_to_index_quals(index_opt_info, index_quals); + clauselist_selectivity( + root, + selectivity_quals, + (*(*index_opt_info).rel).relid as _, + JoinType::JOIN_INNER, + std::ptr::null_mut(), + ) + }; + // index exists + if !(*index_opt_info).hypothetical { + let relation = Index::open((*index_opt_info).indexoid, pgrx::pg_sys::NoLock as _); + let opfamily = opfamily(relation.raw()); + if !matches!( + opfamily, + Opfamily::HalfvecCosine + | Opfamily::HalfvecIp + | Opfamily::HalfvecL2 + | Opfamily::VectorCosine + | Opfamily::VectorIp + | Opfamily::VectorL2 + ) { + *index_startup_cost = 0.0; + *index_total_cost = 0.0; + *index_selectivity = 1.0; + *index_correlation = 0.0; + *index_pages = 1.0; + return; + } + let index = PostgresRelation::new(relation.raw()); + let probes = gucs::probes(); + let cost = algorithm::cost(index); + if cost.cells.len() != 1 + probes.len() { + panic!( + "need {} probes, but {} probes provided", + cost.cells.len() - 1, + probes.len() + ); + } + let node_count = { + let tuples = (*index_opt_info).tuples as u32; + let mut count = 0.0; + let r = cost.cells.iter().copied().rev().map(NonZero::get); + let numerator = std::iter::once(1).chain(probes.clone()); + let denumerator = r.clone(); + let scale = r.skip(1).chain(std::iter::once(tuples)); + for (scale, (numerator, denumerator)) in scale.zip(numerator.zip(denumerator)) { + count += (scale as f64) * ((numerator as f64) / (denumerator as f64)); + } + count + }; + let page_count = { + let mut pages = 0_f64; + pages += 1.0; + pages += node_count * cost.dims as f64 / 60000.0; + pages += probes.iter().sum::() as f64 * { + let x = opfamily.vector_kind().element_size() * cost.dims; + x.div_ceil(3840 * x.div_ceil(5120).min(2)) as f64 + }; + pages += cost.cells[0].get() as f64; + pages + }; + let next_count = + f64::max(1.0, (*root).limit_tuples) * f64::min(1000.0, 1.0 / selectivity); + *index_startup_cost = 0.001 * node_count; + *index_total_cost = 0.001 * node_count + next_count; + *index_selectivity = selectivity; + *index_correlation = 0.0; + *index_pages = page_count; return; } *index_startup_cost = 0.0; *index_total_cost = 0.0; - *index_selectivity = 1.0; - *index_correlation = 1.0; - *index_pages = 0.0; + *index_selectivity = selectivity; + *index_correlation = 0.0; + *index_pages = 1.0; } } @@ -301,7 +381,7 @@ pub unsafe extern "C-unwind" fn amrescan( ); } let opfamily = opfamily((*scan).indexRelation); - let relation = PostgresRelation::new((*scan).indexRelation); + let index = PostgresRelation::new((*scan).indexRelation); let options = SearchOptions { epsilon: gucs::epsilon(), probes: gucs::probes(), @@ -345,7 +425,7 @@ pub unsafe extern "C-unwind" fn amrescan( let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; builder.add((*data).sk_strategy, (!is_null).then_some(value)); } - LazyCell::new(Box::new(move || builder.build(relation, options, fetcher))) + LazyCell::new(Box::new(move || builder.build(index, options, fetcher))) } Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim => { let mut builder = MaxsimBuilder::new(opfamily); @@ -361,7 +441,7 @@ pub unsafe extern "C-unwind" fn amrescan( let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; builder.add((*data).sk_strategy, (!is_null).then_some(value)); } - LazyCell::new(Box::new(move || builder.build(relation, options, fetcher))) + LazyCell::new(Box::new(move || builder.build(index, options, fetcher))) } }; } @@ -505,6 +585,31 @@ impl SearchFetcher for HeapFetcher { } } +struct Index { + raw: *mut pgrx::pg_sys::RelationData, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl Index { + fn open(indexrelid: pgrx::pg_sys::Oid, lockmode: pgrx::pg_sys::LOCKMASK) -> Self { + Self { + raw: unsafe { pgrx::pg_sys::index_open(indexrelid, lockmode) }, + lockmode, + } + } + fn raw(&self) -> *mut pgrx::pg_sys::RelationData { + self.raw + } +} + +impl Drop for Index { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::index_close(self.raw, self.lockmode); + } + } +} + pub const fn ctid_to_key( ItemPointerData { ip_blkid: BlockIdData { bi_hi, bi_lo }, diff --git a/src/index/functions.rs b/src/index/functions.rs index e97e261d..8a0bef49 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -18,24 +18,24 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { if pg_class.relam() != pg_am.oid() { pgrx::error!("the index {:?} is not a vchordrq index", pg_class.relname()); } - let index = Index::open(indexrelid); - let relation = unsafe { PostgresRelation::new(index.raw()) }; - let opfamily = unsafe { crate::index::opclass::opfamily(index.raw()) }; - crate::index::algorithm::prewarm(opfamily, relation, height, || { + let relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); + let opfamily = unsafe { crate::index::opclass::opfamily(relation.raw()) }; + let index = unsafe { PostgresRelation::new(relation.raw()) }; + crate::index::algorithm::prewarm(opfamily, index, height, || { pgrx::check_for_interrupts!(); }) } struct Index { raw: *mut pgrx::pg_sys::RelationData, + lockmode: pgrx::pg_sys::LOCKMODE, } impl Index { - fn open(indexrelid: Oid) -> Self { + fn open(indexrelid: Oid, lockmode: pgrx::pg_sys::LOCKMASK) -> Self { Self { - raw: unsafe { - pgrx::pg_sys::index_open(indexrelid, pgrx::pg_sys::AccessShareLock as _) - }, + raw: unsafe { pgrx::pg_sys::index_open(indexrelid, lockmode) }, + lockmode, } } fn raw(&self) -> *mut pgrx::pg_sys::RelationData { @@ -46,7 +46,7 @@ impl Index { impl Drop for Index { fn drop(&mut self) { unsafe { - pgrx::pg_sys::index_close(self.raw, pgrx::pg_sys::AccessShareLock as _); + pgrx::pg_sys::index_close(self.raw, self.lockmode); } } } From 149e4d70f02842791554865f6c431d6c7f181b01 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 15 Apr 2025 09:38:15 +0800 Subject: [PATCH 147/324] fix: do not error if vector has unexpected length (#236) Consider the following scenario: 1. Process 1 performs a search and finds vector A. It locates `x` (1920 dimensions). 2. Process 2 deletes vector A from the index. `x` (1920 dimensions) and `y` (1 dimension) are deleted. 3. Process 3 inserts vector B into the index. `y` (1920 dimensions) and `z` (1 dimension) are inserted. 4. Process 1 locates `y` (1920 dimensions) and detects a length mismatch, resulting in an error. Since vector a can be deleted and is therefore guaranteed not to be included in the query results, any error related to a during reranking should be silently ignored. Signed-off-by: usamoi --- crates/algorithm/src/operator.rs | 98 +++++++++++++++++++++++++++++++- crates/algorithm/src/rerank.rs | 2 +- crates/algorithm/src/vectors.rs | 6 +- 3 files changed, 100 insertions(+), 6 deletions(-) diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index 7eda147f..f43ef971 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -150,6 +150,7 @@ impl> Accessor1 for LAccess } fn finish(self, rhs: M1) -> Self::Output { + assert!(self.elements.is_empty(), "goal is shorter than expected"); self.accessor.finish(self.metadata, rhs) } } @@ -180,10 +181,101 @@ impl> Accessor1 for RAccess } fn finish(self, lhs: M0) -> Self::Output { + assert!(self.elements.is_empty(), "goal is shorter than expected"); self.accessor.finish(lhs, self.metadata) } } +pub trait TryAccessor1: Sized { + type Output; + #[must_use] + fn push(&mut self, input: &[E]) -> Option<()>; + #[must_use] + fn finish(self, input: M) -> Option; +} + +impl TryAccessor1 for () { + type Output = (); + + fn push(&mut self, _: &[E]) -> Option<()> { + Some(()) + } + + fn finish(self, _: M) -> Option { + Some(()) + } +} + +impl TryAccessor1 for (A,) +where + A: TryAccessor1, +{ + type Output = (A::Output,); + + fn push(&mut self, input: &[E]) -> Option<()> { + self.0.push(input)?; + Some(()) + } + + fn finish(self, input: M) -> Option { + Some((self.0.finish(input)?,)) + } +} + +impl TryAccessor1 for (A, B) +where + A: TryAccessor1, + B: TryAccessor1, +{ + type Output = (A::Output, B::Output); + + fn push(&mut self, input: &[E]) -> Option<()> { + self.0.push(input)?; + self.1.push(input)?; + Some(()) + } + + fn finish(self, input: M) -> Option { + Some((self.0.finish(input)?, self.1.finish(input)?)) + } +} + +pub struct LTryAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> LTryAccess<'a, E, M, A> { + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } + } +} + +impl> TryAccessor1 + for LTryAccess<'_, E0, M0, A> +{ + type Output = A::Output; + + fn push(&mut self, rhs: &[E1]) -> Option<()> { + let (lhs, elements) = self.elements.split_at_checked(rhs.len())?; + self.accessor.push(lhs, rhs); + self.elements = elements; + Some(()) + } + + fn finish(self, rhs: M1) -> Option { + if !self.elements.is_empty() { + return None; + } + Some(self.accessor.finish(self.metadata, rhs)) + } +} + #[derive(Debug)] pub struct BlockAccessor([u16; 32], PhantomData D>); @@ -377,7 +469,8 @@ impl Vector for VectOwned { let vector = vector.slice(); ( match vector.len() { - 0..=960 => vec![vector], + 0 => unreachable!(), + 1..=960 => vec![vector], 961..=1280 => vec![&vector[..640], &vector[640..]], 1281.. => vector.chunks(1920).collect(), }, @@ -411,7 +504,8 @@ impl Vector for VectOwned { let vector = vector.slice(); ( match vector.len() { - 0..=1920 => vec![vector], + 0 => unreachable!(), + 1..=1920 => vec![vector], 1921..=2560 => vec![&vector[..1280], &vector[1280..]], 2561.. => vector.chunks(3840).collect(), }, diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index 004b19ea..d9e03e70 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -77,7 +77,7 @@ pub fn rerank_index( index.clone(), mean, pay_u, - LAccess::new( + LTryAccess::new( O::Vector::unpack(vector.as_borrowed()), O::DistanceAccessor::default(), ), diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index e0053947..a269d60e 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -29,7 +29,7 @@ pub fn read_for_h1_tuple< pub fn read_for_h0_tuple< O: Operator, - A: Accessor1<::Element, ::Metadata>, + A: TryAccessor1<::Element, ::Metadata>, >( index: impl RelationRead, mean: IndexPointer, @@ -48,10 +48,10 @@ pub fn read_for_h0_tuple< if tuple.payload() != Some(payload) { return None; } - result.push(tuple.elements()); + result.push(tuple.elements())?; cursor = tuple.metadata_or_pointer(); } - Some(result.finish(cursor.ok()?)) + result.finish(cursor.ok()?) } pub fn append( From 11c807105dcb899295c893aa495fce0b2d0dd3c6 Mon Sep 17 00:00:00 2001 From: xieydd Date: Tue, 15 Apr 2025 11:36:09 +0800 Subject: [PATCH 148/324] add vchord-suite image (#229) Signed-off-by: xieydd --- README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/README.md b/README.md index d8a0c280..7184dfc2 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,7 @@ docker run \ -p 5432:5432 \ -d tensorchord/vchord-postgres:pg17-v0.2.2 ``` +> [!NOTE] In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). Then you can connect to the database using the `psql` command line tool. The default username is `postgres`, and the default password is `mysecretpassword`. @@ -76,6 +77,10 @@ Now you can play with VectorChord! VectorChord depends on pgvector, including the vector representation. Since you can use them directly, your application can be easily migrated without pain! +```sql +CREATE EXTENSION IF NOT EXISTS vchord CASCADE; +``` + Similar to pgvector, you can create a table with vector column and insert some rows to it. ```sql From 81e0cb7e67a9b55bf18d96e55f7c5f2dd37c1bfe Mon Sep 17 00:00:00 2001 From: Keming Date: Tue, 15 Apr 2025 16:23:43 +0800 Subject: [PATCH 149/324] docs: fix readme callout (#237) cc @xieydd --------- Signed-off-by: Keming --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 7184dfc2..7204079a 100644 --- a/README.md +++ b/README.md @@ -63,9 +63,10 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d tensorchord/vchord-postgres:pg17-v0.2.2 + -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.3.0 ``` -> [!NOTE] In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). +> [!NOTE] +> In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). Then you can connect to the database using the `psql` command line tool. The default username is `postgres`, and the default password is `mysecretpassword`. From f354d02815d1146774181373e6f4190310be7092 Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Mon, 21 Apr 2025 11:34:39 +0800 Subject: [PATCH 150/324] docs: Add a new feature `Production Proven` (#238) --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 7204079a..9aefc1ac 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,8 @@ VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: [^4]: Please check our [blog post](https://blog.vectorchord.ai/vector-search-at-10000-qps-in-postgresql-with-vectorchord) for more details, the PostgreSQL scalability is powered by [CloudNative-PG](https://github.com/cloudnative-pg/cloudnative-pg). +**🏭 Production Proven**: Deployed in mission-critical environments, VectorChord ​​reliably handles 3B+ vectors​​ in production with consistent performance. Please check out [3B vectors in PostgresQL to Protect the Earth](https://blog.vectorchord.ai/3-billion-vectors-in-postgresql-to-protect-the-earth). + ## Quick Start For new users, we recommend using the Docker image to get started quickly. If you do not prefer Docker, please read [installation guide](https://docs.vectorchord.ai/vectorchord/getting-started/installation.html) for other installation methods. From f0ce49527b256e70bf40946bf52299eff70b5a6c Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 25 Apr 2025 17:04:16 +0800 Subject: [PATCH 151/324] feat: prefetch in reranking (#240) adds a GUC `io_rerank`, with the following possible values: * `read_buffer`: indicates a preference for `ReadBuffer`. * `prefetch_buffer`: indicates a preference for both `PrefetchBuffer` and `ReadBuffer`; it's good for disk vector search; this is default on PostgreSQL 13, 14, 15, 16. * `read_stream`: indicates a preference for `read_stream`; it's good for disk vector search; this option is only available in PostgreSQL 17; this is default on PostgreSQL 17. notes: * prefetching distance of `prefetch_buffer` depends on the compile-time constant `WINDOW_SIZE`. * prefetching distance of `read_stream` depends on the GUC parameter `effective_io_concurrency`. * prefetching for the bit vector has not been implemented. * prefetching for the centroid has been implemented, but disabled. Signed-off-by: usamoi --- Cargo.lock | 46 +- Cargo.toml | 6 +- crates/algorithm/src/build.rs | 89 ++-- crates/algorithm/src/bulkdelete.rs | 15 +- crates/algorithm/src/cache.rs | 11 +- .../algorithm/src/closure_lifetime_binder.rs | 44 ++ crates/algorithm/src/cost.rs | 7 +- .../src/{select_heap.rs => fast_heap.rs} | 43 +- crates/algorithm/src/insert.rs | 83 ++-- crates/algorithm/src/lib.rs | 102 ++++- crates/algorithm/src/maintain.rs | 95 ++-- crates/algorithm/src/operator.rs | 20 + crates/algorithm/src/prefetcher.rs | 164 +++++++ crates/algorithm/src/prewarm.rs | 40 +- crates/algorithm/src/rerank.rs | 135 +++--- crates/algorithm/src/search.rs | 165 ++++--- crates/algorithm/src/tape.rs | 48 +- crates/algorithm/src/tuples.rs | 425 +++++++++++------- crates/algorithm/src/vectors.rs | 63 ++- crates/simd/src/lib.rs | 2 +- src/index/algorithm.rs | 226 +++++++++- src/index/am/am_build.rs | 6 +- src/index/am/mod.rs | 33 +- src/index/gucs.rs | 35 ++ src/index/lazy_cell.rs | 312 +++++++++++++ src/index/mod.rs | 1 + src/index/scanners/default.rs | 252 ++++++++--- src/index/scanners/maxsim.rs | 123 +++-- src/index/scanners/mod.rs | 19 +- src/index/storage.rs | 244 +++++++++- src/lib.rs | 8 +- 31 files changed, 2210 insertions(+), 652 deletions(-) create mode 100644 crates/algorithm/src/closure_lifetime_binder.rs rename crates/algorithm/src/{select_heap.rs => fast_heap.rs} (60%) create mode 100644 crates/algorithm/src/prefetcher.rs create mode 100644 src/index/lazy_cell.rs diff --git a/Cargo.lock b/Cargo.lock index 7df3ef07..75da4911 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -630,6 +630,26 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" +[[package]] +name = "jemalloc-sys" +version = "0.5.4+5.3.0-patched" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac6c1946e1cea1788cbfde01c993b52a10e2da07f4bac608228d1bed20bfebf2" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "jemallocator" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0de374a9f8e63150e6f5e8a60cc14c668226d7a347d8aee1a45766e3c4dd3bc" +dependencies = [ + "jemalloc-sys", + "libc", +] + [[package]] name = "k_means" version = "0.0.0" @@ -662,16 +682,6 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" -[[package]] -name = "libmimalloc-sys" -version = "0.1.42" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec9d6fac27761dabcd4ee73571cdb06b7022dc99089acbe5435691edffaac0f4" -dependencies = [ - "cc", - "libc", -] - [[package]] name = "litemap" version = "0.7.5" @@ -694,15 +704,6 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" -[[package]] -name = "mimalloc" -version = "0.1.46" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "995942f432bbb4822a7e9c3faa87a695185b0d09273ba85f097b54f4e458f2af" -dependencies = [ - "libmimalloc-sys", -] - [[package]] name = "minimal-lexical" version = "0.2.1" @@ -1030,13 +1031,12 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "rand" -version = "0.9.0" +version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" +checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97" dependencies = [ "rand_chacha", "rand_core", - "zerocopy", ] [[package]] @@ -1520,8 +1520,8 @@ dependencies = [ "always_equal", "distance", "half 2.6.0", + "jemallocator", "k_means", - "mimalloc", "paste", "pgrx", "pgrx-catalog", diff --git a/Cargo.toml b/Cargo.toml index 9d2d5d0a..7a2a7bad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,8 +39,8 @@ toml = "0.8.20" validator.workspace = true zerocopy.workspace = true -[target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos", target_os = "windows")))'.dependencies] -mimalloc = { version = "*", features = ["local_dynamic_tls"] } +[target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), target_os = "linux"))'.dependencies] +jemallocator = { version = "0.5.4", features = ["disable_initial_exec_tls"] } [lints] workspace = true @@ -56,7 +56,7 @@ edition = "2024" [workspace.dependencies] half = { version = "2.6.0", features = ["zerocopy"] } paste = "1" -rand = "0.9.0" +rand = "0.9.1" serde = { version = "1", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } zerocopy = { version = "0.8.24", features = ["derive"] } diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 41eab332..85772cf5 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -2,7 +2,7 @@ use crate::operator::{Accessor2, Operator, Vector}; use crate::tape::TapeWriter; use crate::tuples::*; use crate::types::*; -use crate::{Branch, IndexPointer, RelationWrite}; +use crate::{Branch, RelationWrite}; use simd::fast_scan::{any_pack, padding_pack}; use vector::VectorOwned; @@ -21,28 +21,35 @@ pub fn build( assert_eq!(meta.first(), 0); let freepage = TapeWriter::<_, FreepageTuple>::create(&index, false); let mut vectors = TapeWriter::<_, VectorTuple>::create(&index, true); - let mut pointer_of_means = Vec::>::new(); + let mut pointer_of_means = Vec::, u16)>>::new(); for i in 0..structures.len() { let mut level = Vec::new(); for j in 0..structures[i].len() { let vector = structures[i].means[j].as_borrowed(); let (slices, metadata) = O::Vector::split(vector); let mut chain = Ok(metadata); + let mut prefetch = Vec::new(); for i in (0..slices.len()).rev() { - chain = Err(vectors.push(match chain { + let (id, head) = vectors.push(match chain { Ok(metadata) => VectorTuple::_0 { payload: None, elements: slices[i].to_vec(), metadata, }, - Err(pointer) => VectorTuple::_1 { + Err(head) => VectorTuple::_1 { payload: None, elements: slices[i].to_vec(), - pointer, + head, }, - })); + }); + chain = Err(head); + prefetch.push(id); } - level.push(chain.expect_err("internal error: 0-dimensional vector")); + prefetch.reverse(); + level.push(( + prefetch, + chain.expect_err("internal error: 0-dimensional vector"), + )); } pointer_of_means.push(level); } @@ -61,7 +68,7 @@ pub fn build( }); level.push(jump.first()); } else { - let mut tape = H1TapeWriter::create(&index, false); + let mut tape = H1TapeWriter::create(&index, O::Vector::count(dims as _), false); let h2_mean = structures[i].means[j].as_borrowed(); let h2_children = structures[i].children[j].as_slice(); for child in h2_children.iter().copied() { @@ -77,35 +84,39 @@ pub fn build( O::Vector::code(h1_mean) }; tape.push(Branch { - mean: pointer_of_means[i - 1][child as usize], + head: pointer_of_means[i - 1][child as usize].1, dis_u_2: code.dis_u_2, factor_ppc: code.factor_ppc, factor_ip: code.factor_ip, factor_err: code.factor_err, signs: code.signs, + prefetch: pointer_of_means[i - 1][child as usize].0.clone(), extra: pointer_of_firsts[i - 1][child as usize], }); } - let (mut tape, branches) = tape.into_inner(); - if !branches.is_empty() { + let (mut tape, chunk) = tape.into_inner(); + if !chunk.is_empty() { let mut remain = - padding_pack(branches.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); loop { let freespace = tape.freespace(); - if H1Tuple::estimate_size_0(remain.len()) <= freespace as usize { + if H1Tuple::estimate_size_0(O::Vector::count(dims as _), remain.len()) + <= freespace as usize + { tape.tape_put(H1Tuple::_0 { - mean: any_pack(branches.iter().map(|x| x.mean)), - dis_u_2: any_pack(branches.iter().map(|x| x.dis_u_2)), - factor_ppc: any_pack(branches.iter().map(|x| x.factor_ppc)), - factor_ip: any_pack(branches.iter().map(|x| x.factor_ip)), - factor_err: any_pack(branches.iter().map(|x| x.factor_err)), - first: any_pack(branches.iter().map(|x| x.extra)), - len: branches.len() as _, + head: any_pack(chunk.iter().map(|x| x.head)), + dis_u_2: any_pack(chunk.iter().map(|x| x.dis_u_2)), + factor_ppc: any_pack(chunk.iter().map(|x| x.factor_ppc)), + factor_ip: any_pack(chunk.iter().map(|x| x.factor_ip)), + factor_err: any_pack(chunk.iter().map(|x| x.factor_err)), + first: any_pack(chunk.iter().map(|x| x.extra)), + prefetch: fix(chunk.iter().map(|x| x.prefetch.as_slice())), + len: chunk.len() as _, elements: remain, }); break; } - if let Some(w) = H1Tuple::fit_1(freespace) { + if let Some(w) = H1Tuple::fit_1(O::Vector::count(dims as _), freespace) { let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); tape.tape_put(H1Tuple::_1 { elements: left.to_vec(), @@ -127,14 +138,20 @@ pub fn build( is_residual, rerank_in_heap: vchordrq_options.rerank_in_table, vectors_first: vectors.first(), - root_mean: pointer_of_means + root_prefetch: pointer_of_means .last() - .expect("internal error: empty structure")[0], + .expect("internal error: empty structure")[0] + .0 + .clone(), + root_head: pointer_of_means + .last() + .expect("internal error: empty structure")[0] + .1, root_first: pointer_of_firsts .last() .expect("internal error: empty structure")[0], freepage_first: freepage.first(), - cells: std::array::from_fn(|i| structures.get(i).map(|s| s.len() as _).unwrap_or_default()), + cells: structures.iter().map(|s| s.len() as _).collect(), }); } @@ -144,16 +161,18 @@ where { tape: TapeWriter<'a, R, H1Tuple>, branches: Vec>, + prefetch: usize, } impl<'a, R> H1TapeWriter<'a, R> where R: RelationWrite + 'a, { - fn create(index: &'a R, tracking_freespace: bool) -> Self { + fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { Self { tape: TapeWriter::create(index, tracking_freespace), branches: Vec::new(), + prefetch, } } fn push(&mut self, branch: Branch) { @@ -162,20 +181,21 @@ where let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); loop { let freespace = self.tape.freespace(); - if H1Tuple::estimate_size_0(remain.len()) <= freespace as usize { + if H1Tuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { self.tape.tape_put(H1Tuple::_0 { - mean: chunk.each_ref().map(|x| x.mean), + head: chunk.each_ref().map(|x| x.head), dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), factor_ip: chunk.each_ref().map(|x| x.factor_ip), factor_err: chunk.each_ref().map(|x| x.factor_err), first: chunk.each_ref().map(|x| x.extra), + prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), len: chunk.len() as _, elements: remain, }); break; } - if let Some(w) = H1Tuple::fit_1(freespace) { + if let Some(w) = H1Tuple::fit_1(self.prefetch, freespace) { let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); self.tape.tape_put(H1Tuple::_1 { elements: left.to_vec(), @@ -192,3 +212,16 @@ where (self.tape, self.branches) } } + +fn fix<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { + use std::array::from_fn; + let mut iter = into_iter.into_iter(); + let mut array: [_; 32] = from_fn(|_| iter.next().map(<[u32]>::to_vec).unwrap_or_default()); + if iter.next().is_some() { + panic!("too many slices"); + } + let step = array.iter().map(Vec::len).max().unwrap_or_default(); + array.iter_mut().for_each(|x| x.resize(step, u32::MAX)); + let flat = array.into_iter().flatten().collect::>(); + (0..step).map(|i| from_fn(|j| flat[i * 32 + j])).collect() +} diff --git a/crates/algorithm/src/bulkdelete.rs b/crates/algorithm/src/bulkdelete.rs index c34882af..61aa0e41 100644 --- a/crates/algorithm/src/bulkdelete.rs +++ b/crates/algorithm/src/bulkdelete.rs @@ -1,10 +1,11 @@ +use crate::closure_lifetime_binder::{id_0, id_1}; use crate::operator::{FunctionalAccessor, Operator}; use crate::tuples::*; -use crate::{Page, RelationWrite, tape}; +use crate::{Page, RelationRead, RelationWrite, tape}; use std::num::NonZero; pub fn bulkdelete( - index: impl RelationWrite, + index: impl RelationRead + RelationWrite, check: impl Fn(), callback: impl Fn(NonZero) -> bool, ) { @@ -24,14 +25,8 @@ pub fn bulkdelete( tape::read_h1_tape( index.clone(), first, - || { - fn push(_: &mut (), _: &[T]) {} - fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { - [(); 32] - } - FunctionalAccessor::new((), push, finish) - }, - |(), _, first| results.push(first), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + |(), _, first, _| results.push(first), |_| check(), ); } diff --git a/crates/algorithm/src/cache.rs b/crates/algorithm/src/cache.rs index d3d0e5d6..839948c8 100644 --- a/crates/algorithm/src/cache.rs +++ b/crates/algorithm/src/cache.rs @@ -1,3 +1,4 @@ +use crate::closure_lifetime_binder::{id_0, id_1}; use crate::operator::FunctionalAccessor; use crate::tuples::{MetaTuple, WithReader}; use crate::{Page, RelationRead, tape}; @@ -18,14 +19,8 @@ pub fn cache(index: impl RelationRead) -> Vec { tape::read_h1_tape( index.clone(), first, - || { - fn push(_: &mut (), _: &[T]) {} - fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { - [(); 32] - } - FunctionalAccessor::new((), push, finish) - }, - |(), _, first| { + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + |(), _, first, _| { results.push(first); }, |id| { diff --git a/crates/algorithm/src/closure_lifetime_binder.rs b/crates/algorithm/src/closure_lifetime_binder.rs new file mode 100644 index 00000000..6db1ca95 --- /dev/null +++ b/crates/algorithm/src/closure_lifetime_binder.rs @@ -0,0 +1,44 @@ +// Use stable language features as an alternative to `closure_lifetime_binder`. +// See https://github.com/rust-lang/rust/issues/97362. + +#[inline(always)] +pub fn id_0(f: F) -> F +where + F: for<'a> FnMut(&'a mut A, &'a B) -> R, +{ + f +} + +#[inline(always)] +pub fn id_1(f: F) -> F +where + F: for<'a> FnMut(A, (&'a B, &'a B, &'a B, &'a B)) -> R, +{ + f +} + +#[inline(always)] +pub fn id_2(f: F) -> F +where + F: for<'a> FnMut(A, B, C, &'a D) -> R, +{ + f +} + +#[inline(always)] +pub fn id_3(f: F) -> F +where + T: crate::RelationWrite, + F: for<'a> Fn(&'a T, A) -> T::WriteGuard<'a>, +{ + f +} + +#[inline(always)] +pub fn id_4(f: F) -> F +where + T: crate::RelationRead, + F: FnMut(A, Vec>, B) -> R, +{ + f +} diff --git a/crates/algorithm/src/cost.rs b/crates/algorithm/src/cost.rs index 06c17f09..9f6c2f53 100644 --- a/crates/algorithm/src/cost.rs +++ b/crates/algorithm/src/cost.rs @@ -1,11 +1,10 @@ use crate::tuples::{MetaTuple, WithReader}; use crate::{Page, RelationRead}; -use std::num::NonZero; pub struct Cost { pub dims: u32, pub is_residual: bool, - pub cells: Vec>, + pub cells: Vec, } #[must_use] @@ -15,12 +14,12 @@ pub fn cost(index: impl RelationRead) -> Cost { let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); - let cells = meta_tuple.cells(); + let cells = meta_tuple.cells().to_vec(); drop(meta_guard); Cost { dims, is_residual, - cells: cells.into_iter().map_while(NonZero::new).collect(), + cells, } } diff --git a/crates/algorithm/src/select_heap.rs b/crates/algorithm/src/fast_heap.rs similarity index 60% rename from crates/algorithm/src/select_heap.rs rename to crates/algorithm/src/fast_heap.rs index 803cb495..a4e88f1e 100644 --- a/crates/algorithm/src/select_heap.rs +++ b/crates/algorithm/src/fast_heap.rs @@ -1,3 +1,4 @@ +use crate::Heap; use std::collections::BinaryHeap; use std::num::NonZero; @@ -12,12 +13,12 @@ impl SortHeap { } } -pub enum SelectHeap { +pub enum FastHeap { Sorted(SortHeap), Binary(BinaryHeap), } -impl SelectHeap { +impl FastHeap { pub fn from_vec(vec: Vec) -> Self { let n = vec.len(); if let Some(t) = NonZero::new(n / 384) { @@ -32,26 +33,50 @@ impl SelectHeap { } pub fn pop(&mut self) -> Option { match self { - SelectHeap::Sorted(SortHeap { inner, t }) => { + FastHeap::Sorted(SortHeap { inner, t }) => { let Some(k) = inner.pop() else { unreachable!() }; if let Some(value) = NonZero::new(t.get() - 1) { *t = value; } else { - *self = SelectHeap::Binary(std::mem::take(inner).into()); + *self = FastHeap::Binary(std::mem::take(inner).into()); } Some(k) } - SelectHeap::Binary(x) => x.pop(), + FastHeap::Binary(x) => x.pop(), } } pub fn peek(&self) -> Option<&T> { match self { - SelectHeap::Sorted(x) => x.peek(), - SelectHeap::Binary(x) => x.peek(), + FastHeap::Sorted(x) => x.peek(), + FastHeap::Binary(x) => x.peek(), } } } +impl IntoIterator for FastHeap { + type Item = T; + + type IntoIter = std::vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + match self { + FastHeap::Sorted(sort_heap) => sort_heap.inner.into_iter(), + FastHeap::Binary(binary_heap) => binary_heap.into_vec().into_iter(), + } + } +} + +impl Heap for FastHeap { + fn make(this: Vec) -> Self { + Self::from_vec(this) + } + + fn pop_if(&mut self, predicate: impl FnOnce(&Self::Item) -> bool) -> Option { + let first = self.peek()?; + if predicate(first) { self.pop() } else { None } + } +} + #[test] fn test_select_heap() { for _ in 0..1000 { @@ -64,7 +89,7 @@ fn test_select_heap() { x }; let result = { - let mut x = SelectHeap::from_vec(sequence.clone()); + let mut x = FastHeap::from_vec(sequence.clone()); std::iter::from_fn(|| x.pop()).collect::>() }; assert_eq!(answer, result); @@ -73,7 +98,7 @@ fn test_select_heap() { #[test] fn test_issue_209() { - let mut heap = SelectHeap::from_vec(vec![0]); + let mut heap = FastHeap::from_vec(vec![0]); assert_eq!(heap.pop(), Some(0)); assert_eq!(heap.pop(), None); } diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 7b662da3..38b066e6 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -1,9 +1,8 @@ use crate::linked_vec::LinkedVec; use crate::operator::*; -use crate::select_heap::SelectHeap; use crate::tuples::*; use crate::vectors::{self}; -use crate::{IndexPointer, Page, RelationWrite, tape}; +use crate::{Bump, Page, Prefetcher, RelationRead, RelationWrite, tape}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -11,7 +10,23 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -pub fn insert(index: impl RelationWrite, payload: NonZero, vector: O::Vector) { +type Item<'b> = ( + Reverse, + AlwaysEqual<&'b mut (u32, u16, &'b mut [u32])>, +); + +pub fn insert< + 'b, + R: RelationRead + RelationWrite, + O: Operator, + P: Prefetcher>, +>( + index: R, + payload: NonZero, + vector: O::Vector, + bump: &'b impl Bump, + mut prefetch: impl FnMut(Vec>) -> P, +) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -20,7 +35,8 @@ pub fn insert(index: impl RelationWrite, payload: NonZero, vec let rerank_in_heap = meta_tuple.rerank_in_heap(); let height_of_root = meta_tuple.height_of_root(); assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); - let root_mean = meta_tuple.root_mean(); + let root_prefetch = meta_tuple.root_prefetch().to_vec(); + let root_head = meta_tuple.root_head(); let root_first = meta_tuple.root_first(); let vectors_first = meta_tuple.vectors_first(); drop(meta_guard); @@ -31,31 +47,31 @@ pub fn insert(index: impl RelationWrite, payload: NonZero, vec None }; - let mean = if !rerank_in_heap { + let (list, head) = if !rerank_in_heap { vectors::append::(index.clone(), vectors_first, vector.as_borrowed(), payload) } else { - IndexPointer::default() + (Vec::new(), 0) }; type State = (u32, Option<::Vector>); let mut state: State = { - let mean = root_mean; if is_residual { - let residual_u = vectors::read_for_h1_tuple::( - index.clone(), - mean, + let list = root_prefetch.into_iter().map(|id| index.read(id)); + let residual = vectors::read_for_h1_tuple::( + root_head, + list, LAccess::new( O::Vector::unpack(vector.as_borrowed()), O::ResidualAccessor::default(), ), ); - (root_first, Some(residual_u)) + (root_first, Some(residual)) } else { (root_first, None) } }; - let step = |state: State| { - let mut results = LinkedVec::new(); + let mut step = |state: State| { + let mut results = LinkedVec::>::new(); { let (first, residual) = state; let block_lut = if let Some(residual) = residual { @@ -69,25 +85,26 @@ pub fn insert(index: impl RelationWrite, payload: NonZero, vec index.clone(), first, || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), - |(rough, err), mean, first| { + |(rough, err), head, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * 1.9); - results.push((Reverse(lowerbound), AlwaysEqual(mean), AlwaysEqual(first))); + results.push(( + Reverse(lowerbound), + AlwaysEqual(bump.alloc((first, head, bump.alloc_slice(prefetch)))), + )); }, |_| (), ); } - let mut heap = SelectHeap::from_vec(results.into_vec()); + let mut heap = (prefetch)(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); { - while let Some((Reverse(_), AlwaysEqual(mean), AlwaysEqual(first))) = - pop_if(&mut heap, |(d, ..)| { - Some(*d) > cache.peek().map(|(d, ..)| *d) - }) + while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, ..))), list)) = + heap.pop_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { if is_residual { - let (dis_u, residual_u) = vectors::read_for_h1_tuple::( - index.clone(), - mean, + let (distance, residual) = vectors::read_for_h1_tuple::( + head, + list.into_iter(), LAccess::new( O::Vector::unpack(vector.as_borrowed()), ( @@ -97,20 +114,20 @@ pub fn insert(index: impl RelationWrite, payload: NonZero, vec ), ); cache.push(( - Reverse(dis_u), + Reverse(distance), AlwaysEqual(first), - AlwaysEqual(Some(residual_u)), + AlwaysEqual(Some(residual)), )); } else { - let dis_u = vectors::read_for_h1_tuple::( - index.clone(), - mean, + let distance = vectors::read_for_h1_tuple::( + head, + list.into_iter(), LAccess::new( O::Vector::unpack(vector.as_borrowed()), O::DistanceAccessor::default(), ), ); - cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); + cache.push((Reverse(distance), AlwaysEqual(first), AlwaysEqual(None))); } } let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache @@ -130,12 +147,13 @@ pub fn insert(index: impl RelationWrite, payload: NonZero, vec O::Vector::code(vector.as_borrowed()) }; let bytes = AppendableTuple::serialize(&AppendableTuple { - mean, + head, dis_u_2: code.dis_u_2, factor_ppc: code.factor_ppc, factor_ip: code.factor_ip, factor_err: code.factor_err, payload: Some(payload), + prefetch: list, elements: rabitq::pack_to_u64(&code.signs), }); @@ -145,8 +163,3 @@ pub fn insert(index: impl RelationWrite, payload: NonZero, vec tape::append(index.clone(), jump_tuple.appendable_first(), &bytes, false); } - -fn pop_if(heap: &mut SelectHeap, mut predicate: impl FnMut(&T) -> bool) -> Option { - let peek = heap.peek()?; - if predicate(peek) { heap.pop() } else { None } -} diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 4ba4f456..a3c25586 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -1,15 +1,20 @@ +#![feature(let_chains)] +#![allow(clippy::type_complexity)] + mod build; mod bulkdelete; mod cache; +mod closure_lifetime_binder; mod cost; +mod fast_heap; mod freepages; mod insert; mod linked_vec; mod maintain; +mod prefetcher; mod prewarm; mod rerank; mod search; -mod select_heap; mod tape; mod tuples; mod vectors; @@ -17,16 +22,20 @@ mod vectors; pub mod operator; pub mod types; +use always_equal::AlwaysEqual; pub use build::build; pub use bulkdelete::bulkdelete; pub use cache::cache; pub use cost::cost; +pub use fast_heap::FastHeap; pub use insert::insert; pub use maintain::maintain; +pub use prefetcher::{PlainPrefetcher, Prefetcher, SimplePrefetcher, StreamPrefetcher}; pub use prewarm::prewarm; pub use rerank::{how, rerank_heap, rerank_index}; pub use search::{default_search, maxsim_search}; +use std::collections::BinaryHeap; use std::ops::{Deref, DerefMut}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -64,15 +73,38 @@ pub trait PageGuard { fn id(&self) -> u32; } -pub trait RelationRead: Clone { +pub trait ReadStream { + type Relation: RelationRead; + type Inner: Iterator; + fn next_if( + &mut self, + predicate: impl FnOnce(&T) -> bool, + ) -> Option<(T, Vec<::ReadGuard<'_>>)>; + fn into_inner(self) -> Self::Inner; +} + +pub trait WriteStream { + type Relation: RelationWrite; + type Inner: Iterator; + fn next_if( + &mut self, + predicate: impl FnOnce(&T) -> bool, + ) -> Option<(T, Vec<::WriteGuard<'_>>)>; + fn into_inner(self) -> Self::Inner; +} + +pub trait Relation: Clone { type Page: Page; +} + +pub trait RelationRead: Relation { type ReadGuard<'a>: PageGuard + Deref where Self: 'a; fn read(&self, id: u32) -> Self::ReadGuard<'_>; } -pub trait RelationWrite: RelationRead { +pub trait RelationWrite: Relation { type WriteGuard<'a>: PageGuard + DerefMut where Self: 'a; @@ -81,6 +113,30 @@ pub trait RelationWrite: RelationRead { fn search(&self, freespace: usize) -> Option>; } +pub trait RelationPrefetch: Relation { + fn prefetch(&self, id: u32); +} + +pub trait RelationReadStream: RelationRead { + type ReadStream<'s, I: Iterator>: ReadStream + where + I::Item: Fetch, + Self: 's; + fn read_stream(&self, iter: I) -> Self::ReadStream<'_, I> + where + I::Item: Fetch; +} + +pub trait RelationWriteStream: RelationWrite { + type WriteStream<'s, I: Iterator>: WriteStream + where + I::Item: Fetch, + Self: 's; + fn write_stream(&self, iter: I) -> Self::WriteStream<'_, I> + where + I::Item: Fetch; +} + #[derive(Debug, Clone, Copy)] pub enum RerankMethod { Index, @@ -88,17 +144,45 @@ pub enum RerankMethod { } pub(crate) struct Branch { - pub mean: IndexPointer, + pub head: u16, pub dis_u_2: f32, pub factor_ppc: f32, pub factor_ip: f32, pub factor_err: f32, pub signs: Vec, + pub prefetch: Vec, pub extra: T, } -#[repr(transparent)] -#[derive( - Debug, Default, Clone, Copy, PartialEq, Eq, Hash, IntoBytes, FromBytes, Immutable, KnownLayout, -)] -pub struct IndexPointer(pub u64); +pub trait Fetch { + fn fetch(&self) -> &[u32]; +} + +impl<'b, T, A, B> Fetch for (T, AlwaysEqual<&'b mut (A, B, &'b mut [u32])>) { + fn fetch(&self) -> &[u32] { + let (_, AlwaysEqual((.., list))) = self; + list + } +} + +pub trait Bump: 'static { + #[allow(clippy::mut_from_ref)] + fn alloc(&self, value: T) -> &mut T; + #[allow(clippy::mut_from_ref)] + fn alloc_slice(&self, slice: &[T]) -> &mut [T]; +} + +pub trait Heap: IntoIterator { + fn make(this: Vec) -> Self; + fn pop_if(&mut self, predicate: impl FnOnce(&Self::Item) -> bool) -> Option; +} + +impl Heap for BinaryHeap { + fn make(this: Vec) -> Self { + Self::from(this) + } + fn pop_if(&mut self, predicate: impl FnOnce(&T) -> bool) -> Option { + let peek = self.peek()?; + if predicate(peek) { self.pop() } else { None } + } +} diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index 4f9a7219..1627253e 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -1,11 +1,12 @@ -use crate::operator::{FunctionalAccessor, Operator}; +use crate::closure_lifetime_binder::{id_0, id_1, id_2, id_3}; +use crate::operator::{FunctionalAccessor, Operator, Vector}; use crate::tape::{self, TapeWriter}; use crate::tuples::*; -use crate::{Branch, Page, RelationRead, RelationWrite, freepages}; +use crate::{Branch, Page, Relation, RelationRead, RelationWrite, freepages}; use simd::fast_scan::{padding_pack, unpack}; use std::num::NonZero; -pub fn maintain(index: impl RelationWrite, check: impl Fn()) { +pub fn maintain(index: R, check: impl Fn()) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -24,14 +25,8 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { tape::read_h1_tape( index.clone(), first, - || { - fn push(_: &mut (), _: &[T]) {} - fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { - [(); 32] - } - FunctionalAccessor::new((), push, finish) - }, - |(), _, first| results.push(first), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + |(), _, first, _| results.push(first), |_| check(), ); } @@ -48,25 +43,41 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { let jump_bytes = jump_guard.get_mut(1).expect("data corruption"); let mut jump_tuple = JumpTuple::deserialize_mut(jump_bytes); - let hooked_index = RelationHooked(index.clone(), hooked_extend(freepage_first)); + let hooked_index = RelationHooked( + index.clone(), + id_3(move |index: &R, tracking_freespace: bool| { + if !tracking_freespace { + if let Some(id) = freepages::fetch(index.clone(), freepage_first) { + let mut write = index.write(id, false); + write.clear(); + write + } else { + index.extend(false) + } + } else { + index.extend(true) + } + }), + ); - let mut tape = FrozenTapeWriter::create(&hooked_index, false); + let mut tape = FrozenTapeWriter::create(&hooked_index, O::Vector::count(dims as _), false); let mut trace = Vec::new(); let mut tuples = 0_u64; - let mut callback = |code: (_, _, _, _, _), mean, payload| { + let mut callback = id_2(|code: (_, _, _, _, _), head, payload, prefetch: &[_]| { tape.push(Branch { - mean, + head, dis_u_2: code.0, factor_ppc: code.1, factor_ip: code.2, factor_err: code.3, signs: code.4, + prefetch: prefetch.to_vec(), extra: payload, }); tuples += 1; - }; + }); let mut step = |id| { check(); trace.push(id); @@ -113,12 +124,13 @@ pub fn maintain(index: impl RelationWrite, check: impl Fn()) { for branch in branches { appendable_tape.push(AppendableTuple { - mean: branch.mean, + head: branch.head, dis_u_2: branch.dis_u_2, factor_ppc: branch.factor_ppc, factor_ip: branch.factor_ip, factor_err: branch.factor_err, payload: Some(branch.extra), + prefetch: branch.prefetch, elements: rabitq::pack_to_u64(&branch.signs), }); } @@ -139,16 +151,18 @@ where { tape: TapeWriter<'a, R, FrozenTuple>, branches: Vec>>, + prefetch: usize, } impl<'a, R> FrozenTapeWriter<'a, R> where R: RelationWrite + 'a, { - fn create(index: &'a R, tracking_freespace: bool) -> Self { + fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { Self { tape: TapeWriter::create(index, tracking_freespace), branches: Vec::new(), + prefetch, } } fn push(&mut self, branch: Branch>) { @@ -157,19 +171,20 @@ where let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); loop { let freespace = self.tape.freespace(); - if FrozenTuple::estimate_size_0(remain.len()) <= freespace as usize { + if FrozenTuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { self.tape.tape_put(FrozenTuple::_0 { - mean: chunk.each_ref().map(|x| x.mean), + head: chunk.each_ref().map(|x| x.head), dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), factor_ip: chunk.each_ref().map(|x| x.factor_ip), factor_err: chunk.each_ref().map(|x| x.factor_err), payload: chunk.each_ref().map(|x| Some(x.extra)), + prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), elements: remain, }); break; } - if let Some(w) = FrozenTuple::fit_1(freespace) { + if let Some(w) = FrozenTuple::fit_1(self.prefetch, freespace) { let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); self.tape.tape_put(FrozenTuple::_1 { elements: left.to_vec(), @@ -190,13 +205,19 @@ where #[derive(Clone)] struct RelationHooked(R, E); -impl RelationRead for RelationHooked +impl Relation for RelationHooked where - R: RelationRead, + R: Relation, E: Clone, { type Page = R::Page; +} +impl RelationRead for RelationHooked +where + R: RelationRead, + E: Clone, +{ type ReadGuard<'a> = R::ReadGuard<'a> where @@ -230,23 +251,15 @@ where } } -fn hooked_extend( - freepage_first: u32, -) -> impl Clone + for<'a> Fn(&'a R, bool) -> R::WriteGuard<'a> -where - R: RelationWrite, -{ - move |index, tracking_freespace| { - if !tracking_freespace { - if let Some(id) = freepages::fetch(index.clone(), freepage_first) { - let mut write = index.write(id, false); - write.clear(); - write - } else { - index.extend(false) - } - } else { - index.extend(true) - } +fn fix<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { + use std::array::from_fn; + let mut iter = into_iter.into_iter(); + let mut array: [_; 32] = from_fn(|_| iter.next().map(<[u32]>::to_vec).unwrap_or_default()); + if iter.next().is_some() { + panic!("too many slices"); } + let step = array.iter().map(Vec::len).max().unwrap_or_default(); + array.iter_mut().for_each(|x| x.resize(step, u32::MAX)); + let flat = array.into_iter().flatten().collect::>(); + (0..step).map(|i| from_fn(|j| flat[i * 32 + j])).collect() } diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index f43ef971..9cf3142d 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -451,6 +451,8 @@ pub trait Vector: VectorOwned { fn split(vector: Self::Borrowed<'_>) -> (Vec<&[Self::Element]>, Self::Metadata); + fn count(n: usize) -> usize; + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut; @@ -478,6 +480,15 @@ impl Vector for VectOwned { ) } + fn count(n: usize) -> usize { + match n { + 0 => unreachable!(), + 1..=960 => 1, + 961..=1280 => 2, + 1281.. => n.div_ceil(1920), + } + } + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { (vector.slice(), ()) } @@ -513,6 +524,15 @@ impl Vector for VectOwned { ) } + fn count(n: usize) -> usize { + match n { + 0 => unreachable!(), + 1..=1920 => 1, + 1921..=2560 => 2, + 2561.. => n.div_ceil(3840), + } + } + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { (vector.slice(), ()) } diff --git a/crates/algorithm/src/prefetcher.rs b/crates/algorithm/src/prefetcher.rs new file mode 100644 index 00000000..5a517ffe --- /dev/null +++ b/crates/algorithm/src/prefetcher.rs @@ -0,0 +1,164 @@ +use crate::{Fetch, Heap, ReadStream, RelationPrefetch, RelationRead, RelationReadStream}; +use std::collections::{BinaryHeap, VecDeque, binary_heap, vec_deque}; +use std::iter::Chain; + +pub const WINDOW_SIZE: usize = 32; +const _: () = assert!(WINDOW_SIZE > 0); + +pub trait Prefetcher: IntoIterator +where + Self::Item: Fetch, +{ + type R: RelationRead; + fn pop_if( + &mut self, + predicate: impl FnOnce(&Self::Item) -> bool, + ) -> Option<(Self::Item, Vec<::ReadGuard<'_>>)>; +} + +pub struct PlainPrefetcher { + relation: R, + heap: H, +} + +impl PlainPrefetcher { + pub fn new(relation: R, vec: Vec) -> Self { + Self { + relation, + heap: Heap::make(vec), + } + } +} + +impl IntoIterator for PlainPrefetcher { + type Item = H::Item; + + type IntoIter = H::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.heap.into_iter() + } +} + +impl Prefetcher for PlainPrefetcher +where + H::Item: Fetch + Ord, +{ + type R = R; + fn pop_if<'s>( + &'s mut self, + predicate: impl FnOnce(&Self::Item) -> bool, + ) -> Option<(H::Item, Vec>)> { + let e = self.heap.pop_if(predicate)?; + let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); + Some((e, list)) + } +} + +pub struct SimplePrefetcher { + relation: R, + window: VecDeque, + heap: BinaryHeap, +} + +impl SimplePrefetcher { + pub fn new(relation: R, vec: Vec) -> Self { + Self { + relation, + window: VecDeque::new(), + heap: BinaryHeap::from(vec), + } + } +} + +impl IntoIterator for SimplePrefetcher { + type Item = T; + + type IntoIter = Chain, binary_heap::IntoIter>; + + fn into_iter(self) -> Self::IntoIter { + self.window.into_iter().chain(self.heap) + } +} + +impl Prefetcher for SimplePrefetcher { + type R = R; + fn pop_if<'s>( + &'s mut self, + predicate: impl FnOnce(&Self::Item) -> bool, + ) -> Option<(T, Vec>)> { + while self.window.len() < WINDOW_SIZE + && let Some(e) = self.heap.pop() + { + for id in e.fetch().iter().copied() { + self.relation.prefetch(id); + } + self.window.push_back(e); + } + let e = vec_deque_pop_front_if(&mut self.window, predicate)?; + let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); + Some((e, list)) + } +} + +pub struct StreamPrefetcherHeap(BinaryHeap); + +impl Iterator for StreamPrefetcherHeap { + type Item = T; + fn next(&mut self) -> Option { + self.0.pop() + } +} + +pub struct StreamPrefetcher<'r, R, T> +where + R: RelationReadStream + 'r, + T: Fetch + Ord, +{ + stream: R::ReadStream<'r, StreamPrefetcherHeap>, +} + +impl<'r, R: RelationReadStream, T: Fetch + Ord> StreamPrefetcher<'r, R, T> { + pub fn new(relation: &'r R, vec: Vec) -> Self { + let stream = relation.read_stream(StreamPrefetcherHeap(BinaryHeap::from(vec))); + Self { stream } + } +} + +impl<'r, R: RelationReadStream, T: Fetch + Ord> IntoIterator for StreamPrefetcher<'r, R, T> { + type Item = T; + + type IntoIter = <::ReadStream< + 'r, + StreamPrefetcherHeap, + > as ReadStream>::Inner; + + fn into_iter(self) -> Self::IntoIter { + self.stream.into_inner() + } +} + +impl<'r, R: RelationReadStream, T: Fetch + Ord> Prefetcher for StreamPrefetcher<'r, R, T> { + type R = R; + fn pop_if<'s>( + &'s mut self, + predicate: impl FnOnce(&Self::Item) -> bool, + ) -> Option<(T, Vec>)> { + self.stream.next_if(predicate) + } +} + +// Emulate unstable library feature `vec_deque_pop_if`. +// See https://github.com/rust-lang/rust/issues/135889. + +fn vec_deque_pop_front_if( + this: &mut VecDeque, + predicate: impl FnOnce(&T) -> bool, +) -> Option { + let first = this.front()?; + if predicate(first) { + this.pop_front() + } else { + None + } +} diff --git a/crates/algorithm/src/prewarm.rs b/crates/algorithm/src/prewarm.rs index 77042024..75466ad3 100644 --- a/crates/algorithm/src/prewarm.rs +++ b/crates/algorithm/src/prewarm.rs @@ -1,11 +1,12 @@ +use crate::closure_lifetime_binder::{id_0, id_1, id_2}; use crate::operator::{FunctionalAccessor, Operator}; use crate::tuples::*; use crate::{Page, RelationRead, tape, vectors}; use std::error::Error; use std::fmt::Write; -pub fn prewarm( - index: impl RelationRead, +pub fn prewarm( + index: R, height: i32, check: impl Fn(), ) -> Result> { @@ -13,7 +14,8 @@ pub fn prewarm( let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let height_of_root = meta_tuple.height_of_root(); - let root_mean = meta_tuple.root_mean(); + let root_prefetch = meta_tuple.root_prefetch().to_vec(); + let root_head = meta_tuple.root_head(); let root_first = meta_tuple.root_first(); drop(meta_guard); @@ -27,7 +29,8 @@ pub fn prewarm( let mut state: State = { let mut results = Vec::new(); { - vectors::read_for_h1_tuple::(index.clone(), root_mean, ()); + let list = root_prefetch.into_iter().map(|id| index.read(id)); + vectors::read_for_h1_tuple::(root_head, list, ()); results.push(root_first); } writeln!(message, "------------------------")?; @@ -42,16 +45,11 @@ pub fn prewarm( tape::read_h1_tape( index.clone(), first, - || { - fn push(_: &mut (), _: &[T]) {} - fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { - [(); 32] - } - FunctionalAccessor::new((), push, finish) - }, - |(), mean, first| { + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + |(), head, first, prefetch| { + let list = prefetch.iter().map(|&id| index.read(id)); + vectors::read_for_h1_tuple::(head, list, ()); results.push(first); - vectors::read_for_h1_tuple::(index.clone(), mean, ()); }, |_| { check(); @@ -77,16 +75,10 @@ pub fn prewarm( tape::read_frozen_tape( index.clone(), jump_tuple.frozen_first(), - || { - fn push(_: &mut (), _: &[T]) {} - fn finish(_: (), _: (&T, &T, &T, &T)) -> [(); 32] { - [(); 32] - } - FunctionalAccessor::new((), push, finish) - }, - |(), _, _| { + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + id_2(|_, _, _, _| { results.push(()); - }, + }), |_| { check(); counter += 1; @@ -96,9 +88,9 @@ pub fn prewarm( index.clone(), jump_tuple.appendable_first(), |_| (), - |(), _, _| { + id_2(|_, _, _, _| { results.push(()); - }, + }), |_| { check(); counter += 1; diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index d9e03e70..2f0c6e8f 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -1,21 +1,19 @@ +use crate::closure_lifetime_binder::id_4; use crate::operator::*; +use crate::prefetcher::Prefetcher; use crate::tuples::{MetaTuple, WithReader}; -use crate::{IndexPointer, Page, RelationRead, RerankMethod, vectors}; +use crate::{Page, RelationRead, RerankMethod, vectors}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; +use std::marker::PhantomData; use std::num::NonZero; use vector::VectorOwned; -type Result = ( - Reverse, - AlwaysEqual, - AlwaysEqual>, - AlwaysEqual, -); +type Extra<'b> = &'b mut (NonZero, u16, &'b mut [u32]); -type Rerank = (Reverse, AlwaysEqual>); +type Result = (Reverse, AlwaysEqual>); pub fn how(index: impl RelationRead) -> RerankMethod { let meta_guard = index.read(0); @@ -29,92 +27,97 @@ pub fn how(index: impl RelationRead) -> RerankMethod { } } -pub struct Reranker { - heap: BinaryHeap>, - cache: BinaryHeap<(Reverse, AlwaysEqual>)>, +pub struct Reranker { + prefetcher: P, + cache: BinaryHeap, f: F, + _phantom: PhantomData T>, } -impl) -> Option> Iterator for Reranker { +impl<'b, T, F, P> Iterator for Reranker +where + F: FnMut(NonZero, Vec<::ReadGuard<'_>>, u16) -> Option, + P: Prefetcher, AlwaysEqual), AlwaysEqual>)>, +{ type Item = (Distance, NonZero); fn next(&mut self) -> Option { - while let Some((Reverse(_), AlwaysEqual(_), AlwaysEqual(pay_u), AlwaysEqual(mean))) = - pop_if(&mut self.heap, |(d, ..)| { - Some(*d) > self.cache.peek().map(|(d, ..)| *d) - }) + while let Some(((_, AlwaysEqual(&mut (payload, head, ..))), list)) = self + .prefetcher + .pop_if(|((d, _), ..)| Some(*d) > self.cache.peek().map(|(d, ..)| *d)) { - if let Some(dis_u) = (self.f)(mean, pay_u) { - self.cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); + if let Some(distance) = (self.f)(payload, list, head) { + self.cache.push((Reverse(distance), AlwaysEqual(payload))); }; } - let (Reverse(dis_u), AlwaysEqual(pay_u)) = self.cache.pop()?; - Some((dis_u, pay_u)) + let (Reverse(distance), AlwaysEqual(payload)) = self.cache.pop()?; + Some((distance, payload)) } } -impl Reranker { - pub fn finish( - self, - ) -> ( - impl Iterator>, - impl Iterator, - ) { - (self.heap.into_iter(), self.cache.into_iter()) +impl Reranker { + pub fn finish(self) -> (P, impl Iterator) { + (self.prefetcher, self.cache.into_iter()) } } -pub fn rerank_index( - index: impl RelationRead, +pub fn rerank_index< + 'b, + O: Operator, + T, + P: Prefetcher, AlwaysEqual), AlwaysEqual>)>, +>( vector: O::Vector, - results: Vec>, -) -> Reranker) -> Option> { + prefetcher: P, +) -> Reranker< + T, + impl FnMut(NonZero, Vec<::ReadGuard<'_>>, u16) -> Option, + P, +> { Reranker { - heap: BinaryHeap::from(results), - cache: BinaryHeap::<(Reverse, _)>::new(), - f: move |mean, pay_u| { - vectors::read_for_h0_tuple::( - index.clone(), - mean, - pay_u, + prefetcher, + cache: BinaryHeap::new(), + f: id_4::<_, P::R, _, _, _>(move |payload, list, head| { + vectors::read_for_h0_tuple::( + head, + list.into_iter(), + payload, LTryAccess::new( O::Vector::unpack(vector.as_borrowed()), O::DistanceAccessor::default(), ), ) - }, + }), + _phantom: PhantomData, } } -pub fn rerank_heap( +pub fn rerank_heap< + 'b, + O: Operator, + T, + P: Prefetcher, AlwaysEqual), AlwaysEqual>)>, +>( vector: O::Vector, - results: Vec>, - mut fetch: impl FnMut(NonZero) -> Option, -) -> Reranker) -> Option> { + prefetcher: P, + mut fetch: impl FnMut(NonZero) -> Option + 'b, +) -> Reranker< + T, + impl FnMut(NonZero, Vec<::ReadGuard<'_>>, u16) -> Option, + P, +> { Reranker { - heap: BinaryHeap::from(results), - cache: BinaryHeap::<(Reverse, _)>::new(), - f: move |_: IndexPointer, pay_u| { + prefetcher, + cache: BinaryHeap::new(), + f: id_4::<_, P::R, _, _, _>(move |payload, _, _| { + let unpack = O::Vector::unpack(vector.as_borrowed()); + let vector = fetch(payload)?; let vector = O::Vector::unpack(vector.as_borrowed()); - let vec_u = fetch(pay_u)?; - let vec_u = O::Vector::unpack(vec_u.as_borrowed()); let mut accessor = O::DistanceAccessor::default(); - accessor.push(vector.0, vec_u.0); - let dis_u = accessor.finish(vector.1, vec_u.1); - Some(dis_u) - }, - } -} - -fn pop_if( - heap: &mut BinaryHeap, - mut predicate: impl FnMut(&mut T) -> bool, -) -> Option { - use std::collections::binary_heap::PeekMut; - let mut peek = heap.peek_mut()?; - if predicate(&mut peek) { - Some(PeekMut::pop(peek)) - } else { - None + accessor.push(unpack.0, vector.0); + let distance = accessor.finish(unpack.1, vector.1); + Some(distance) + }), + _phantom: PhantomData, } } diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index 416bd453..9ba922c0 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -1,7 +1,9 @@ +use crate::closure_lifetime_binder::id_2; use crate::linked_vec::LinkedVec; use crate::operator::*; +use crate::prefetcher::Prefetcher; use crate::tuples::*; -use crate::{IndexPointer, Page, RelationRead, tape, vectors}; +use crate::{Bump, Page, RelationRead, tape, vectors}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -9,19 +11,21 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -type Result = ( +type Item<'b> = ( Reverse, - AlwaysEqual, - AlwaysEqual>, - AlwaysEqual, + AlwaysEqual<&'b mut (u32, u16, &'b mut [u32])>, ); -pub fn default_search( - index: impl RelationRead, +type Extra<'b> = &'b mut (NonZero, u16, &'b mut [u32]); + +pub fn default_search<'b, R: RelationRead, O: Operator, P: Prefetcher>>( + index: R, vector: O::Vector, probes: Vec, epsilon: f32, -) -> Vec> { + bump: &'b impl Bump, + mut prefetch: impl FnMut(Vec>) -> P, +) -> Vec<((Reverse, AlwaysEqual<()>), AlwaysEqual>)> { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -36,7 +40,8 @@ pub fn default_search( probes.len() ); } - let root_mean = meta_tuple.root_mean(); + let root_prefetch = meta_tuple.root_prefetch().to_vec(); + let root_head = meta_tuple.root_head(); let root_first = meta_tuple.root_first(); drop(meta_guard); @@ -48,22 +53,22 @@ pub fn default_search( type State = Vec<(u32, Option<::Vector>)>; let mut state: State = vec![{ - let mean = root_mean; if is_residual { - let residual_u = vectors::read_for_h1_tuple::( - index.clone(), - mean, + let list = root_prefetch.into_iter().map(|id| index.read(id)); + let residual = vectors::read_for_h1_tuple::( + root_head, + list, LAccess::new( O::Vector::unpack(vector.as_borrowed()), O::ResidualAccessor::default(), ), ); - (root_first, Some(residual_u)) + (root_first, Some(residual)) } else { (root_first, None) } }]; - let step = |state: State| { + let mut step = |state: State| { let mut results = LinkedVec::new(); for (first, residual) in state { let block_lut = if let Some(residual) = residual { @@ -77,27 +82,27 @@ pub fn default_search( index.clone(), first, || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), - |(rough, err), mean, first| { + |(rough, err), head, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); - results.push((Reverse(lowerbound), AlwaysEqual(first), AlwaysEqual(mean))); + results.push(( + Reverse(lowerbound), + AlwaysEqual(bump.alloc((first, head, bump.alloc_slice(prefetch)))), + )); }, |_| (), ); } - let mut heap = BinaryHeap::from(results.into_vec()); + let mut heap = (prefetch)(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - let index = index.clone(); let vector = vector.as_borrowed(); std::iter::from_fn(move || { - while let Some((Reverse(_), AlwaysEqual(first), AlwaysEqual(mean))) = - pop_if(&mut heap, |(d, ..)| { - Some(*d) > cache.peek().map(|(d, ..)| *d) - }) + while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, ..))), list)) = + heap.pop_if(|(d, ..)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { if is_residual { - let (dis_u, residual_u) = vectors::read_for_h1_tuple::( - index.clone(), - mean, + let (distance, residual) = vectors::read_for_h1_tuple::( + head, + list.into_iter(), LAccess::new( O::Vector::unpack(vector), ( @@ -107,17 +112,17 @@ pub fn default_search( ), ); cache.push(( - Reverse(dis_u), + Reverse(distance), AlwaysEqual(first), - AlwaysEqual(Some(residual_u)), + AlwaysEqual(Some(residual)), )); } else { - let dis_u = vectors::read_for_h1_tuple::( - index.clone(), - mean, + let distance = vectors::read_for_h1_tuple::( + head, + list.into_iter(), LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); - cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); + cache.push((Reverse(distance), AlwaysEqual(first), AlwaysEqual(None))); } } let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; @@ -141,15 +146,13 @@ pub fn default_search( let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let mut callback = |(rough, err), mean, payload| { + let mut callback = id_2(|(rough, err), mean, payload, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( - Reverse(lowerbound), - AlwaysEqual(()), - AlwaysEqual(payload), - AlwaysEqual(mean), + (Reverse(lowerbound), AlwaysEqual(())), + AlwaysEqual(bump.alloc((payload, mean, bump.alloc_slice(prefetch)))), )); - }; + }); tape::read_frozen_tape( index.clone(), jump_tuple.frozen_first(), @@ -168,13 +171,21 @@ pub fn default_search( results.into_vec() } -pub fn maxsim_search( - index: impl RelationRead, +pub fn maxsim_search<'b, R: RelationRead, O: Operator, P: Prefetcher>>( + index: R, vector: O::Vector, probes: Vec, epsilon: f32, mut threshold: u32, -) -> (Vec>, Distance) { + bump: &'b impl Bump, + mut prefetch: impl FnMut(Vec>) -> P, +) -> ( + Vec<( + (Reverse, AlwaysEqual), + AlwaysEqual>, + )>, + Distance, +) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -189,7 +200,8 @@ pub fn maxsim_search( probes.len() ); } - let root_mean = meta_tuple.root_mean(); + let root_prefetch = meta_tuple.root_prefetch().to_vec(); + let root_head = meta_tuple.root_head(); let root_first = meta_tuple.root_first(); drop(meta_guard); @@ -201,22 +213,22 @@ pub fn maxsim_search( type State = Vec<(u32, Option<::Vector>)>; let mut state: State = vec![{ - let mean = root_mean; if is_residual { - let residual_u = vectors::read_for_h1_tuple::( - index.clone(), - mean, + let list = root_prefetch.into_iter().map(|id| index.read(id)); + let residual = vectors::read_for_h1_tuple::( + root_head, + list, LAccess::new( O::Vector::unpack(vector.as_borrowed()), O::ResidualAccessor::default(), ), ); - (root_first, Some(residual_u)) + (root_first, Some(residual)) } else { (root_first, None) } }]; - let step = |state: State| { + let mut step = |state: State| { let mut results = LinkedVec::new(); for (first, residual) in state { let block_lut = if let Some(residual) = residual { @@ -230,27 +242,27 @@ pub fn maxsim_search( index.clone(), first, || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), - |(rough, err), mean, first| { + |(rough, err), head, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); - results.push((Reverse(lowerbound), AlwaysEqual(first), AlwaysEqual(mean))); + results.push(( + Reverse(lowerbound), + AlwaysEqual(bump.alloc((first, head, bump.alloc_slice(prefetch)))), + )); }, |_| (), ); } - let mut heap = BinaryHeap::from(results.into_vec()); + let mut heap = (prefetch)(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - let index = index.clone(); let vector = vector.as_borrowed(); std::iter::from_fn(move || { - while let Some((Reverse(_), AlwaysEqual(first), AlwaysEqual(mean))) = - pop_if(&mut heap, |(d, ..)| { - Some(*d) > cache.peek().map(|(d, ..)| *d) - }) + while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, ..))), list)) = + heap.pop_if(|(d, ..)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { if is_residual { - let (dis_u, residual_u) = vectors::read_for_h1_tuple::( - index.clone(), - mean, + let (distance, residual) = vectors::read_for_h1_tuple::( + head, + list.into_iter(), LAccess::new( O::Vector::unpack(vector), ( @@ -260,17 +272,17 @@ pub fn maxsim_search( ), ); cache.push(( - Reverse(dis_u), + Reverse(distance), AlwaysEqual(first), - AlwaysEqual(Some(residual_u)), + AlwaysEqual(Some(residual)), )); } else { - let dis_u = vectors::read_for_h1_tuple::( - index.clone(), - mean, + let distance = vectors::read_for_h1_tuple::( + head, + list.into_iter(), LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); - cache.push((Reverse(dis_u), AlwaysEqual(first), AlwaysEqual(None))); + cache.push((Reverse(distance), AlwaysEqual(first), AlwaysEqual(None))); } } let (Reverse(distance), AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; @@ -296,16 +308,14 @@ pub fn maxsim_search( let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let mut callback = |(rough, err), mean, payload| { + let mut callback = id_2(|(rough, err), mean, payload, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); let rough = Distance::from_f32(rough); results.push(( - Reverse(lowerbound), - AlwaysEqual(rough), - AlwaysEqual(payload), - AlwaysEqual(mean), + (Reverse(lowerbound), AlwaysEqual(rough)), + AlwaysEqual(bump.alloc((payload, mean, bump.alloc_slice(prefetch)))), )); - }; + }); tape::read_frozen_tape( index.clone(), jump_tuple.frozen_first(), @@ -335,16 +345,3 @@ pub fn maxsim_search( } (results.into_vec(), estimation_by_threshold) } - -fn pop_if( - heap: &mut BinaryHeap, - mut predicate: impl FnMut(&mut T) -> bool, -) -> Option { - use std::collections::binary_heap::PeekMut; - let mut peek = heap.peek_mut()?; - if predicate(&mut peek) { - Some(PeekMut::pop(peek)) - } else { - None - } -} diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index 093774a5..b941afc9 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -1,6 +1,6 @@ use crate::operator::Accessor1; use crate::tuples::*; -use crate::{IndexPointer, Page, PageGuard, RelationRead, RelationWrite}; +use crate::{Page, PageGuard, RelationRead, RelationWrite}; use rabitq::binary::BinaryCode; use std::marker::PhantomData; use std::num::NonZero; @@ -53,25 +53,25 @@ where R: RelationWrite + 'a, T: Tuple, { - pub fn push(&mut self, x: T) -> IndexPointer { + pub fn push(&mut self, x: T) -> (u32, u16) { let bytes = T::serialize(&x); if let Some(i) = self.head.alloc(&bytes) { - pair_to_pointer((self.head.id(), i)) + (self.head.id(), i) } else { let next = self.index.extend(self.tracking_freespace); self.head.get_opaque_mut().next = next.id(); self.head = next; if let Some(i) = self.head.alloc(&bytes) { - pair_to_pointer((self.head.id(), i)) + (self.head.id(), i) } else { panic!("implementation: a free page cannot accommodate a single tuple") } } } - pub fn tape_put(&mut self, x: T) -> IndexPointer { + pub fn tape_put(&mut self, x: T) -> (u32, u16) { let bytes = T::serialize(&x); if let Some(i) = self.head.alloc(&bytes) { - pair_to_pointer((self.head.id(), i)) + (self.head.id(), i) } else { panic!("implementation: a free page cannot accommodate a single tuple") } @@ -82,7 +82,7 @@ pub fn read_h1_tape( index: impl RelationRead, first: u32, accessor: impl Fn() -> A, - mut callback: impl FnMut(T, IndexPointer, u32), + mut callback: impl for<'a> FnMut(T, u16, u32, &'a [u32]), mut step: impl FnMut(u32), ) where A: for<'a> Accessor1< @@ -105,9 +105,10 @@ pub fn read_h1_tape( let mut x = x.take().unwrap_or_else(&accessor); x.push(tuple.elements()); let values = x.finish(tuple.metadata()); + let prefetch: [_; 32] = fix_0(tuple.prefetch()); for (j, value) in values.into_iter().enumerate() { if j < tuple.len() as usize { - callback(value, tuple.mean()[j], tuple.first()[j]); + callback(value, tuple.head()[j], tuple.first()[j], fix_1(prefetch[j])); } } } @@ -124,7 +125,7 @@ pub fn read_frozen_tape( index: impl RelationRead, first: u32, accessor: impl Fn() -> A, - mut callback: impl FnMut(T, IndexPointer, NonZero), + mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), mut step: impl FnMut(u32), ) where A: for<'a> Accessor1< @@ -147,9 +148,10 @@ pub fn read_frozen_tape( let mut x = x.take().unwrap_or_else(&accessor); x.push(tuple.elements()); let values = x.finish(tuple.metadata()); + let prefetch: [_; 32] = fix_0(tuple.prefetch()); for (j, value) in values.into_iter().enumerate() { if let Some(payload) = tuple.payload()[j] { - callback(value, tuple.mean()[j], payload); + callback(value, tuple.mean()[j], payload, fix_1(prefetch[j])); } } } @@ -166,7 +168,7 @@ pub fn read_appendable_tape( index: impl RelationRead, first: u32, mut access: impl for<'a> FnMut(BinaryCode<'a>) -> T, - mut callback: impl FnMut(T, IndexPointer, NonZero), + mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), mut step: impl FnMut(u32), ) { assert!(first != u32::MAX); @@ -179,7 +181,7 @@ pub fn read_appendable_tape( let tuple = AppendableTuple::deserialize_ref(bytes); if let Some(payload) = tuple.payload() { let value = access(tuple.code()); - callback(value, tuple.mean(), payload); + callback(value, tuple.head(), payload, tuple.prefetch()); } } current = guard.get_opaque().next; @@ -188,11 +190,11 @@ pub fn read_appendable_tape( #[allow(clippy::collapsible_else_if)] pub fn append( - index: impl RelationWrite, + index: impl RelationRead + RelationWrite, first: u32, bytes: &[u8], tracking_freespace: bool, -) -> IndexPointer { +) -> (u32, u16) { assert!(first != u32::MAX); let mut current = first; loop { @@ -202,7 +204,7 @@ pub fn append( let mut write = index.write(current, tracking_freespace); if write.get_opaque().next == u32::MAX { if let Some(i) = write.alloc(bytes) { - return pair_to_pointer((current, i)); + return (current, i); } let mut extend = index.extend(tracking_freespace); write.get_opaque_mut().next = extend.id(); @@ -212,7 +214,7 @@ pub fn append( drop(extend); let mut past = index.write(first, tracking_freespace); past.get_opaque_mut().skip = fresh.max(past.get_opaque().skip); - return pair_to_pointer((fresh, i)); + return (fresh, i); } else { panic!("implementation: a clear page cannot accommodate a single tuple"); } @@ -231,3 +233,17 @@ pub fn append( } } } + +fn fix_0(x: &[[T; 32]]) -> [&[T]; 32] { + let step = x.len(); + let flat = x.as_flattened(); + std::array::from_fn(|i| &flat[i * step..][..step]) +} + +fn fix_1(x: &[u32]) -> &[u32] { + if let Some(i) = x.iter().position(|&x| x == u32::MAX) { + &x[..i] + } else { + x + } +} diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 204f14c5..0be43a7f 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -1,4 +1,3 @@ -use crate::IndexPointer; use crate::operator::Vector; use rabitq::binary::BinaryCode; use std::marker::PhantomData; @@ -7,8 +6,8 @@ use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; -const MAGIC: u64 = u64::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 6; +const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); +const VERSION: u64 = 7; pub trait Tuple: 'static { fn serialize(&self) -> Vec; @@ -27,7 +26,6 @@ pub trait WithWriter: Tuple { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct MetaTupleHeader { - magic: u64, version: u64, dims: u32, height_of_root: u32, @@ -36,12 +34,17 @@ struct MetaTupleHeader { _padding_0: [ZeroU8; 2], vectors_first: u32, // raw vector - root_mean: IndexPointer, + root_prefetch_s: u16, + root_prefetch_e: u16, + root_head: u16, + _padding_1: [ZeroU8; 2], // for meta tuple, it's pointers to next level root_first: u32, freepage_first: u32, // statistics - cells: [u32; 8], + cells_s: u16, + cells_e: u16, + _padding_2: [ZeroU8; 4], } pub struct MetaTuple { @@ -50,59 +53,103 @@ pub struct MetaTuple { pub is_residual: bool, pub rerank_in_heap: bool, pub vectors_first: u32, - pub root_mean: IndexPointer, + pub root_prefetch: Vec, + pub root_head: u16, pub root_first: u32, pub freepage_first: u32, - pub cells: [u32; 8], + pub cells: Vec, } impl Tuple for MetaTuple { + #[allow(clippy::match_single_binding)] fn serialize(&self) -> Vec { - MetaTupleHeader { - magic: MAGIC, - version: VERSION, - dims: self.dims, - height_of_root: self.height_of_root, - is_residual: self.is_residual.into(), - rerank_in_heap: self.rerank_in_heap.into(), - _padding_0: Default::default(), - vectors_first: self.vectors_first, - root_mean: self.root_mean, - root_first: self.root_first, - freepage_first: self.freepage_first, - cells: self.cells, + let mut buffer = Vec::::new(); + match self { + MetaTuple { + dims, + height_of_root, + is_residual, + rerank_in_heap, + vectors_first, + root_prefetch, + root_head, + root_first, + freepage_first, + cells, + } => { + buffer.extend((MAGIC as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + let root_prefetch_s = buffer.len() as u16; + buffer.extend(root_prefetch.as_bytes()); + let root_prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let cells_s = buffer.len() as u16; + buffer.extend(cells.as_bytes()); + let cells_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + buffer[size_of::()..][..size_of::()].copy_from_slice( + MetaTupleHeader { + version: VERSION, + dims: *dims, + height_of_root: *height_of_root, + is_residual: (*is_residual).into(), + rerank_in_heap: (*rerank_in_heap).into(), + _padding_0: Default::default(), + vectors_first: *vectors_first, + root_prefetch_s, + root_prefetch_e, + root_head: *root_head, + _padding_1: Default::default(), + root_first: *root_first, + freepage_first: *freepage_first, + cells_s, + cells_e, + _padding_2: Default::default(), + } + .as_bytes(), + ); + } } - .as_bytes() - .to_vec() + buffer } } impl WithReader for MetaTuple { type Reader<'a> = MetaTupleReader<'a>; fn deserialize_ref(source: &[u8]) -> MetaTupleReader<'_> { - if source.len() < 16 { - panic!("deserialization: bad bytes") - } - let magic = u64::from_ne_bytes(std::array::from_fn(|i| source[i + 0])); - if magic != MAGIC { - panic!("deserialization: bad magic number"); - } - let version = u64::from_ne_bytes(std::array::from_fn(|i| source[i + 8])); - if version != VERSION { - panic!("deserialization: bad version number"); + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + MAGIC => { + let checker = RefChecker::new(source); + if VERSION != *checker.prefix::(size_of::()) { + panic!("deserialization: bad version number"); + } + let header: &MetaTupleHeader = checker.prefix(size_of::()); + let root_prefetch = checker.bytes(header.root_prefetch_s, header.root_prefetch_e); + let cells = checker.bytes(header.cells_s, header.cells_e); + MetaTupleReader { + header, + root_prefetch, + cells, + } + } + _ => panic!("deserialization: bad magic number"), } - let checker = RefChecker::new(source); - let header = checker.prefix(0); - MetaTupleReader { header } } } #[derive(Debug, Clone, Copy)] pub struct MetaTupleReader<'a> { header: &'a MetaTupleHeader, + root_prefetch: &'a [u32], + cells: &'a [u32], } -impl MetaTupleReader<'_> { +impl<'a> MetaTupleReader<'a> { pub fn dims(self) -> u32 { self.header.dims } @@ -118,8 +165,11 @@ impl MetaTupleReader<'_> { pub fn vectors_first(self) -> u32 { self.header.vectors_first } - pub fn root_mean(self) -> IndexPointer { - self.header.root_mean + pub fn root_prefetch(self) -> &'a [u32] { + self.root_prefetch + } + pub fn root_head(self) -> u16 { + self.header.root_head } pub fn root_first(self) -> u32 { self.header.root_first @@ -127,8 +177,8 @@ impl MetaTupleReader<'_> { pub fn freepage_first(self) -> u32 { self.header.freepage_first } - pub fn cells(self) -> [u32; 8] { - self.header.cells + pub fn cells(self) -> &'a [u32] { + self.cells } } @@ -164,7 +214,7 @@ impl WithWriter for FreepageTuple { fn deserialize_mut(source: &mut [u8]) -> FreepageTupleWriter<'_> { let mut checker = MutChecker::new(source); - let header = checker.prefix(0); + let header = checker.prefix(0_u16); FreepageTupleWriter { header } } } @@ -212,20 +262,20 @@ impl FreepageTupleWriter<'_> { #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader0 { payload: Option>, - metadata_s: usize, - elements_s: usize, - elements_e: usize, - #[cfg(target_pointer_width = "32")] - _padding_0: [ZeroU8; 4], + metadata_s: u16, + elements_s: u16, + elements_e: u16, + _padding_0: [ZeroU8; 2], } #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader1 { payload: Option>, - pointer: IndexPointer, - elements_s: usize, - elements_e: usize, + head: u16, + elements_s: u16, + elements_e: u16, + _padding_0: [ZeroU8; 2], } #[derive(Debug, Clone, PartialEq)] @@ -237,7 +287,7 @@ pub enum VectorTuple { }, _1 { payload: Option>, - pointer: IndexPointer, + head: u16, elements: Vec, }, } @@ -253,14 +303,14 @@ impl Tuple for VectorTuple { } => { buffer.extend((0 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - let metadata_s = buffer.len(); + let metadata_s = buffer.len() as u16; buffer.extend(metadata.as_bytes()); while buffer.len() % ALIGN != 0 { buffer.push(0); } - let elements_s = buffer.len(); + let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); - let elements_e = buffer.len(); + let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } @@ -270,7 +320,6 @@ impl Tuple for VectorTuple { metadata_s, elements_s, elements_e, - #[cfg(target_pointer_width = "32")] _padding_0: Default::default(), } .as_bytes(), @@ -278,23 +327,24 @@ impl Tuple for VectorTuple { } VectorTuple::_1 { payload, - pointer, + head, elements, } => { buffer.extend((1 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - let elements_s = buffer.len(); + let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); - let elements_e = buffer.len(); + let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } buffer[size_of::()..][..size_of::()].copy_from_slice( VectorTupleHeader1 { payload: *payload, - pointer: *pointer, + head: *head, elements_s, elements_e, + _padding_0: Default::default(), } .as_bytes(), ); @@ -370,10 +420,10 @@ impl<'a, V: Vector> VectorTupleReader<'a, V> { VectorTupleReader::_1(this) => this.elements, } } - pub fn metadata_or_pointer(self) -> Result { + pub fn metadata_or_head(self) -> Result { match self { VectorTupleReader::_0(this) => Ok(*this.metadata), - VectorTupleReader::_1(this) => Err(this.header.pointer), + VectorTupleReader::_1(this) => Err(this.header.head), } } } @@ -381,35 +431,39 @@ impl<'a, V: Vector> VectorTupleReader<'a, V> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct H1TupleHeader0 { - mean: [IndexPointer; 32], + head: [u16; 32], dis_u_2: [f32; 32], factor_ppc: [f32; 32], factor_ip: [f32; 32], factor_err: [f32; 32], first: [u32; 32], + prefetch_s: u16, + prefetch_e: u16, len: u32, + elements_s: u16, + elements_e: u16, _padding_0: [ZeroU8; 4], - elements_s: usize, - elements_e: usize, } #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct H1TupleHeader1 { - elements_s: usize, - elements_e: usize, + elements_s: u16, + elements_e: u16, + _padding_0: [ZeroU8; 4], } #[allow(clippy::large_enum_variant)] #[derive(Debug, Clone, PartialEq)] pub enum H1Tuple { _0 { - mean: [IndexPointer; 32], + head: [u16; 32], dis_u_2: [f32; 32], factor_ppc: [f32; 32], factor_ip: [f32; 32], factor_err: [f32; 32], first: [u32; 32], + prefetch: Vec<[u32; 32]>, len: u32, elements: Vec<[u8; 16]>, }, @@ -419,17 +473,19 @@ pub enum H1Tuple { } impl H1Tuple { - pub fn estimate_size_0(elements: usize) -> usize { + pub fn estimate_size_0(prefetch: usize, elements: usize) -> usize { let mut size = 0_usize; size += size_of::(); size += size_of::(); - size += elements * size_of::<[u8; 16]>(); + size += (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN); + size += (elements * size_of::<[u8; 16]>()).next_multiple_of(ALIGN); size } - pub fn fit_1(freespace: u16) -> Option { + pub fn fit_1(prefetch: usize, freespace: u16) -> Option { let mut freespace = freespace as isize; freespace -= size_of::() as isize; freespace -= size_of::() as isize; + freespace -= (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN) as isize; if freespace >= 0 { Some(freespace as usize / size_of::<[u8; 16]>()) } else { @@ -443,32 +499,44 @@ impl Tuple for H1Tuple { let mut buffer = Vec::::new(); match self { Self::_0 { - mean, + head, dis_u_2, factor_ppc, factor_ip, factor_err, first, + prefetch, len, elements, } => { buffer.extend((0 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - let elements_s = buffer.len(); + let prefetch_s = buffer.len() as u16; + buffer.extend(prefetch.as_bytes()); + let prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); - let elements_e = buffer.len(); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } buffer[size_of::()..][..size_of::()].copy_from_slice( H1TupleHeader0 { - mean: *mean, + head: *head, dis_u_2: *dis_u_2, factor_ppc: *factor_ppc, factor_ip: *factor_ip, factor_err: *factor_err, first: *first, len: *len, - _padding_0: Default::default(), + prefetch_s, + prefetch_e, elements_s, elements_e, + _padding_0: Default::default(), } .as_bytes(), ); @@ -476,13 +544,17 @@ impl Tuple for H1Tuple { Self::_1 { elements } => { buffer.extend((1 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - let elements_s = buffer.len(); + let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); - let elements_e = buffer.len(); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } buffer[size_of::()..][..size_of::()].copy_from_slice( H1TupleHeader1 { elements_s, elements_e, + _padding_0: Default::default(), } .as_bytes(), ); @@ -501,8 +573,13 @@ impl WithReader for H1Tuple { 0 => { let checker = RefChecker::new(source); let header: &H1TupleHeader0 = checker.prefix(size_of::()); + let prefetch = checker.bytes(header.prefetch_s, header.prefetch_e); let elements = checker.bytes(header.elements_s, header.elements_e); - H1TupleReader::_0(H1TupleReader0 { header, elements }) + H1TupleReader::_0(H1TupleReader0 { + header, + prefetch, + elements, + }) } 1 => { let checker = RefChecker::new(source); @@ -524,6 +601,7 @@ pub enum H1TupleReader<'a> { #[derive(Debug, Clone, Copy, PartialEq)] pub struct H1TupleReader0<'a> { header: &'a H1TupleHeader0, + prefetch: &'a [[u32; 32]], elements: &'a [[u8; 16]], } @@ -537,8 +615,8 @@ impl<'a> H1TupleReader0<'a> { pub fn len(self) -> u32 { self.header.len } - pub fn mean(self) -> &'a [IndexPointer] { - &self.header.mean[..self.header.len as usize] + pub fn head(self) -> &'a [u16] { + &self.header.head[..self.header.len as usize] } pub fn metadata(self) -> (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]) { ( @@ -551,6 +629,9 @@ impl<'a> H1TupleReader0<'a> { pub fn first(self) -> &'a [u32] { &self.header.first[..self.header.len as usize] } + pub fn prefetch(self) -> &'a [[u32; 32]] { + self.prefetch + } pub fn elements(&self) -> &'a [[u8; 16]] { self.elements } @@ -593,7 +674,7 @@ impl WithReader for JumpTuple { type Reader<'a> = JumpTupleReader<'a>; fn deserialize_ref(source: &[u8]) -> JumpTupleReader<'_> { let checker = RefChecker::new(source); - let header: &JumpTupleHeader = checker.prefix(0); + let header: &JumpTupleHeader = checker.prefix(0_u16); JumpTupleReader { header } } } @@ -602,7 +683,7 @@ impl WithWriter for JumpTuple { type Writer<'a> = JumpTupleWriter<'a>; fn deserialize_mut(source: &mut [u8]) -> JumpTupleWriter<'_> { let mut checker = MutChecker::new(source); - let header: &mut JumpTupleHeader = checker.prefix(0); + let header: &mut JumpTupleHeader = checker.prefix(0_u16); JumpTupleWriter { header } } } @@ -644,33 +725,37 @@ impl JumpTupleWriter<'_> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FrozenTupleHeader0 { - mean: [IndexPointer; 32], + head: [u16; 32], dis_u_2: [f32; 32], factor_ppc: [f32; 32], factor_ip: [f32; 32], factor_err: [f32; 32], payload: [Option>; 32], - elements_s: usize, - elements_e: usize, + prefetch_s: u16, + prefetch_e: u16, + elements_s: u16, + elements_e: u16, } #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FrozenTupleHeader1 { - elements_s: usize, - elements_e: usize, + elements_s: u16, + elements_e: u16, + _padding_0: [ZeroU8; 4], } #[allow(clippy::large_enum_variant)] #[derive(Debug, Clone, PartialEq)] pub enum FrozenTuple { _0 { - mean: [IndexPointer; 32], + head: [u16; 32], dis_u_2: [f32; 32], factor_ppc: [f32; 32], factor_ip: [f32; 32], factor_err: [f32; 32], payload: [Option>; 32], + prefetch: Vec<[u32; 32]>, elements: Vec<[u8; 16]>, }, _1 { @@ -679,17 +764,19 @@ pub enum FrozenTuple { } impl FrozenTuple { - pub fn estimate_size_0(elements: usize) -> usize { + pub fn estimate_size_0(prefetch: usize, elements: usize) -> usize { let mut size = 0_usize; size += size_of::(); size += size_of::(); - size += elements * size_of::<[u8; 16]>(); + size += (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN); + size += (elements * size_of::<[u8; 16]>()).next_multiple_of(ALIGN); size } - pub fn fit_1(freespace: u16) -> Option { + pub fn fit_1(prefetch: usize, freespace: u16) -> Option { let mut freespace = freespace as isize; freespace -= size_of::() as isize; freespace -= size_of::() as isize; + freespace -= (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN) as isize; if freespace >= 0 { Some(freespace as usize / size_of::<[u8; 16]>()) } else { @@ -703,22 +790,32 @@ impl Tuple for FrozenTuple { let mut buffer = Vec::::new(); match self { FrozenTuple::_0 { - mean, + head, dis_u_2, factor_ppc, factor_ip, factor_err, payload, + prefetch, elements, } => { buffer.extend((0 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - let elements_s = buffer.len(); + let prefetch_s = buffer.len() as u16; + buffer.extend(prefetch.as_bytes()); + let prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); - let elements_e = buffer.len(); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } buffer[size_of::()..][..size_of::()].copy_from_slice( FrozenTupleHeader0 { - mean: *mean, + head: *head, dis_u_2: *dis_u_2, factor_ppc: *factor_ppc, factor_ip: *factor_ip, @@ -726,6 +823,8 @@ impl Tuple for FrozenTuple { payload: *payload, elements_s, elements_e, + prefetch_s, + prefetch_e, } .as_bytes(), ); @@ -733,13 +832,17 @@ impl Tuple for FrozenTuple { Self::_1 { elements } => { buffer.extend((1 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - let elements_s = buffer.len(); + let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); - let elements_e = buffer.len(); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } buffer[size_of::()..][..size_of::()].copy_from_slice( FrozenTupleHeader1 { elements_s, elements_e, + _padding_0: Default::default(), } .as_bytes(), ); @@ -758,8 +861,13 @@ impl WithReader for FrozenTuple { 0 => { let checker = RefChecker::new(source); let header: &FrozenTupleHeader0 = checker.prefix(size_of::()); + let prefetch = checker.bytes(header.prefetch_s, header.prefetch_e); let elements = checker.bytes(header.elements_s, header.elements_e); - FrozenTupleReader::_0(FrozenTupleReader0 { header, elements }) + FrozenTupleReader::_0(FrozenTupleReader0 { + header, + prefetch, + elements, + }) } 1 => { let checker = RefChecker::new(source); @@ -804,12 +912,13 @@ pub enum FrozenTupleReader<'a> { #[derive(Debug, Clone, Copy, PartialEq)] pub struct FrozenTupleReader0<'a> { header: &'a FrozenTupleHeader0, + prefetch: &'a [[u32; 32]], elements: &'a [[u8; 16]], } impl<'a> FrozenTupleReader0<'a> { - pub fn mean(self) -> &'a [IndexPointer; 32] { - &self.header.mean + pub fn mean(self) -> &'a [u16; 32] { + &self.header.head } pub fn metadata(self) -> (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]) { ( @@ -825,6 +934,9 @@ impl<'a> FrozenTupleReader0<'a> { pub fn payload(self) -> &'a [Option>; 32] { &self.header.payload } + pub fn prefetch(self) -> &'a [[u32; 32]] { + self.prefetch + } } #[derive(Debug, Clone, Copy, PartialEq)] @@ -870,24 +982,28 @@ impl FrozenTupleWriter0<'_> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct AppendableTupleHeader { - mean: IndexPointer, + head: u16, + _padding_0: [ZeroU8; 6], dis_u_2: f32, factor_ppc: f32, factor_ip: f32, factor_err: f32, payload: Option>, - elements_s: usize, - elements_e: usize, + prefetch_s: u16, + prefetch_e: u16, + elements_s: u16, + elements_e: u16, } #[derive(Debug, Clone, PartialEq)] pub struct AppendableTuple { - pub mean: IndexPointer, + pub head: u16, pub dis_u_2: f32, pub factor_ppc: f32, pub factor_ip: f32, pub factor_err: f32, pub payload: Option>, + pub prefetch: Vec, pub elements: Vec, } @@ -895,17 +1011,29 @@ impl Tuple for AppendableTuple { fn serialize(&self) -> Vec { let mut buffer = Vec::::new(); buffer.extend(std::iter::repeat_n(0, size_of::())); - let elements_s = buffer.len(); + let prefetch_s = buffer.len() as u16; + buffer.extend(self.prefetch.as_bytes()); + let prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + let elements_s = buffer.len() as u16; buffer.extend(self.elements.as_bytes()); - let elements_e = buffer.len(); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } buffer[..size_of::()].copy_from_slice( AppendableTupleHeader { - mean: self.mean, + head: self.head, + _padding_0: Default::default(), dis_u_2: self.dis_u_2, factor_ppc: self.factor_ppc, factor_ip: self.factor_ip, factor_err: self.factor_err, payload: self.payload, + prefetch_s, + prefetch_e, elements_s, elements_e, } @@ -920,9 +1048,14 @@ impl WithReader for AppendableTuple { fn deserialize_ref(source: &[u8]) -> AppendableTupleReader<'_> { let checker = RefChecker::new(source); - let header: &AppendableTupleHeader = checker.prefix(0); + let header: &AppendableTupleHeader = checker.prefix(0_u16); + let prefetch = checker.bytes(header.prefetch_s, header.prefetch_e); let elements = checker.bytes(header.elements_s, header.elements_e); - AppendableTupleReader { header, elements } + AppendableTupleReader { + header, + prefetch, + elements, + } } } @@ -931,7 +1064,7 @@ impl WithWriter for AppendableTuple { fn deserialize_mut(source: &mut [u8]) -> AppendableTupleWriter<'_> { let mut checker = MutChecker::new(source); - let header: &mut AppendableTupleHeader = checker.prefix(0); + let header: &mut AppendableTupleHeader = checker.prefix(0_u16); let elements = checker.bytes(header.elements_s, header.elements_e); AppendableTupleWriter { header, elements } } @@ -940,12 +1073,13 @@ impl WithWriter for AppendableTuple { #[derive(Debug, Clone, Copy, PartialEq)] pub struct AppendableTupleReader<'a> { header: &'a AppendableTupleHeader, + prefetch: &'a [u32], elements: &'a [u64], } impl<'a> AppendableTupleReader<'a> { - pub fn mean(self) -> IndexPointer { - self.header.mean + pub fn head(self) -> u16 { + self.header.head } pub fn code(self) -> BinaryCode<'a> { ( @@ -959,6 +1093,9 @@ impl<'a> AppendableTupleReader<'a> { pub fn payload(self) -> Option> { self.header.payload } + pub fn prefetch(self) -> &'a [u32] { + self.prefetch + } } #[derive(Debug)] @@ -974,27 +1111,6 @@ impl AppendableTupleWriter<'_> { } } -pub const fn pointer_to_pair(pointer: IndexPointer) -> (u32, u16) { - let value = pointer.0; - (((value >> 16) & 0xffffffff) as u32, (value & 0xffff) as u16) -} - -pub const fn pair_to_pointer(pair: (u32, u16)) -> IndexPointer { - let mut value = 0; - value |= (pair.0 as u64) << 16; - value |= pair.1 as u64; - IndexPointer(value) -} - -#[test] -const fn soundness_check() { - let a = (111, 222); - let b = pair_to_pointer(a); - let c = pointer_to_pair(b); - assert!(a.0 == c.0); - assert!(a.1 == c.1); -} - #[repr(transparent)] #[derive( Debug, @@ -1057,20 +1173,20 @@ impl<'a> RefChecker<'a> { } pub fn prefix( &self, - s: usize, + s: impl Into + Copy, ) -> &'a T { - let start = s; - let end = s + size_of::(); + let start = Into::::into(s); + let end = Into::::into(s) + size_of::(); let bytes = &self.bytes[start..end]; FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") } pub fn bytes( &self, - s: usize, - e: usize, + s: impl Into + Copy, + e: impl Into + Copy, ) -> &'a T { - let start = s; - let end = e; + let start = Into::::into(s); + let end = Into::::into(e); let bytes = &self.bytes[start..end]; FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") } @@ -1092,10 +1208,10 @@ impl<'a> MutChecker<'a> { } pub fn prefix( &mut self, - s: usize, + s: impl Into + Copy, ) -> &'a mut T { - let start = s; - let end = s + size_of::(); + let start = Into::::into(s); + let end = Into::::into(s) + size_of::(); if !(start <= end && end <= self.bytes.len()) { panic!("deserialization: bad bytes"); } @@ -1112,11 +1228,11 @@ impl<'a> MutChecker<'a> { } pub fn bytes( &mut self, - s: usize, - e: usize, + s: impl Into + Copy, + e: impl Into + Copy, ) -> &'a mut T { - let start = s; - let end = e; + let start = Into::::into(s); + let end = Into::::into(e); if !(start <= end && end <= self.bytes.len()) { panic!("deserialization: bad bytes"); } @@ -1148,20 +1264,25 @@ fn aliasing_test() { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct ExampleHeader { - elements_s: usize, - elements_e: usize, + elements_s: u16, + elements_e: u16, + _padding_0: [ZeroU8; 4], } let serialized = { let elements = (0u32..1111).collect::>(); let mut buffer = Vec::::new(); buffer.extend(std::iter::repeat_n(0, size_of::())); - let elements_s = buffer.len(); + let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); - let elements_e = buffer.len(); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } buffer[..size_of::()].copy_from_slice( ExampleHeader { elements_s, elements_e, + _padding_0: Default::default(), } .as_bytes(), ); @@ -1171,7 +1292,7 @@ fn aliasing_test() { source.as_mut_bytes()[..serialized.len()].copy_from_slice(&serialized); let deserialized = { let mut checker = MutChecker::new(source.as_mut_bytes()); - let header: &mut ExampleHeader = checker.prefix(0); + let header: &mut ExampleHeader = checker.prefix(0_u16); let elements: &mut [u32] = checker.bytes(header.elements_s, header.elements_e); (header, elements) }; diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index a269d60e..92db860b 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -1,46 +1,53 @@ use crate::operator::*; use crate::tuples::*; -use crate::{IndexPointer, Page, PageGuard, RelationRead, RelationWrite, tape}; +use crate::{Page, PageGuard, RelationRead, RelationWrite, tape}; use std::num::NonZero; use vector::VectorOwned; pub fn read_for_h1_tuple< + 'a, + R: RelationRead + 'a, O: Operator, A: Accessor1<::Element, ::Metadata>, >( - index: impl RelationRead, - mean: IndexPointer, + head: u16, + mut list: impl Iterator>, accessor: A, ) -> A::Output { - let mut cursor = Err(mean); + let mut cursor = Err(head); let mut result = accessor; - while let Err(mean) = cursor.map_err(pointer_to_pair) { - let guard = index.read(mean.0); - let bytes = guard.get(mean.1).expect("data corruption"); + while let Err(head) = cursor { + let guard = list.next().expect("data corruption"); + let bytes = guard.get(head).expect("data corruption"); let tuple = VectorTuple::::deserialize_ref(bytes); if tuple.payload().is_some() { panic!("data corruption"); } result.push(tuple.elements()); - cursor = tuple.metadata_or_pointer(); + cursor = tuple.metadata_or_head(); + } + if list.next().is_some() { + panic!("data corruption"); } result.finish(cursor.expect("data corruption")) } pub fn read_for_h0_tuple< + 'a, + R: RelationRead + 'a, O: Operator, A: TryAccessor1<::Element, ::Metadata>, >( - index: impl RelationRead, - mean: IndexPointer, + head: u16, + mut list: impl Iterator>, payload: NonZero, accessor: A, ) -> Option { - let mut cursor = Err(mean); + let mut cursor = Err(head); let mut result = accessor; - while let Err(mean) = cursor.map_err(pointer_to_pair) { - let guard = index.read(mean.0); - let bytes = guard.get(mean.1)?; + while let Err(head) = cursor { + let guard = list.next()?; + let bytes = guard.get(head)?; let tuple = VectorTuple::::deserialize_ref(bytes); if tuple.payload().is_none() { panic!("data corruption"); @@ -49,28 +56,32 @@ pub fn read_for_h0_tuple< return None; } result.push(tuple.elements())?; - cursor = tuple.metadata_or_pointer(); + cursor = tuple.metadata_or_head(); + } + if list.next().is_some() { + return None; } result.finish(cursor.ok()?) } pub fn append( - index: impl RelationWrite, + index: impl RelationRead + RelationWrite, vectors_first: u32, vector: ::Borrowed<'_>, payload: NonZero, -) -> IndexPointer { - fn append(index: impl RelationWrite, first: u32, bytes: &[u8]) -> IndexPointer { +) -> (Vec, u16) { + fn append(index: impl RelationRead + RelationWrite, first: u32, bytes: &[u8]) -> (u32, u16) { if let Some(mut write) = index.search(bytes.len()) { let i = write .alloc(bytes) .expect("implementation: a free page cannot accommodate a single tuple"); - return pair_to_pointer((write.id(), i)); + return (write.id(), i); } tape::append(index, first, bytes, true) } let (slices, metadata) = O::Vector::split(vector); let mut chain = Ok(metadata); + let mut prefetch = Vec::new(); for i in (0..slices.len()).rev() { let bytes = VectorTuple::::serialize(&match chain { Ok(metadata) => VectorTuple::_0 { @@ -78,13 +89,19 @@ pub fn append( payload: Some(payload), metadata, }, - Err(pointer) => VectorTuple::_1 { + Err(head) => VectorTuple::_1 { elements: slices[i].to_vec(), payload: Some(payload), - pointer, + head, }, }); - chain = Err(append(index.clone(), vectors_first, &bytes)); + let (id, head) = append(index.clone(), vectors_first, &bytes); + chain = Err(head); + prefetch.push(id); } - chain.expect_err("internal error: 0-dimensional vector") + prefetch.reverse(); + ( + prefetch, + chain.expect_err("internal error: 0-dimensional vector"), + ) } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 4748c21c..7a448ea3 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -1,4 +1,4 @@ -#![feature(avx512_target_feature)] +#![cfg_attr(target_arch = "x86_64", feature(avx512_target_feature))] #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512_f16))] #![allow(unsafe_code)] diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index 2135f635..8c744852 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -2,12 +2,192 @@ use super::opclass::Opfamily; use crate::index::am::am_build::InternalBuild; use algorithm::operator::{Dot, L2, Op}; use algorithm::types::*; -use algorithm::{RelationRead, RelationWrite}; +use algorithm::{Bump, RelationRead, RelationWrite}; use half::f16; +use std::cell::UnsafeCell; +use std::mem::MaybeUninit; use std::num::NonZero; use vector::VectorOwned; use vector::vect::{VectBorrowed, VectOwned}; +#[repr(C, align(4096))] +struct Chunk([u8; 2 * 1024 * 1024]); + +struct Allocator { + used: Vec<*mut MaybeUninit>, + free: Vec<*mut MaybeUninit>, + this: *mut MaybeUninit, + size: usize, +} + +impl Allocator { + pub fn new() -> Self { + Self { + used: Vec::new(), + free: Vec::new(), + this: Box::into_raw(Box::::new_uninit()).cast(), + size: size_of::(), + } + } + pub fn malloc(&mut self) -> *mut T { + const { + assert!(align_of::() <= align_of::()); + assert!(size_of::() <= size_of::()); + } + if size_of::() <= self.size { + self.size = (self.size - size_of::()) / align_of::() * align_of::(); + unsafe { self.this.add(self.size).cast::() } + } else { + #[cold] + fn cold(sel: &mut Allocator) -> *mut T { + std::panic::abort_unwind(|| { + let raw = std::mem::replace(&mut sel.this, std::ptr::null_mut()); + sel.used.push(raw.cast()); + sel.this = sel + .free + .pop() + .unwrap_or_else(|| Box::into_raw(Box::::new_uninit())) + .cast(); + sel.size = size_of::(); + }); + sel.size = (sel.size - size_of::()) / align_of::() * align_of::(); + unsafe { sel.this.add(sel.size).cast::() } + } + cold(self) + } + } + pub fn malloc_n(&mut self, n: usize) -> *mut T { + const { + assert!(align_of::() <= align_of::()); + } + let limit = const { + if size_of::() > 0 { + size_of::() / size_of::() + } else { + usize::MAX + } + }; + if n <= limit && n * size_of::() <= self.size { + self.size = (self.size - n * size_of::()) / align_of::() * align_of::(); + unsafe { self.this.add(self.size).cast::() } + } else { + #[cold] + fn cold(sel: &mut Allocator, n: usize) -> *mut T { + std::panic::abort_unwind(|| { + let raw = std::mem::replace(&mut sel.this, std::ptr::null_mut()); + sel.used.push(raw.cast()); + sel.this = sel + .free + .pop() + .unwrap_or_else(|| Box::into_raw(Box::::new_uninit())) + .cast(); + sel.size = size_of::(); + }); + sel.size = (sel.size - n * size_of::()) / align_of::() * align_of::(); + unsafe { sel.this.add(sel.size).cast::() } + } + if n > limit { + panic!("failed to allocate memory"); + } + cold(self, n) + } + } + pub fn reset(&mut self) { + std::panic::abort_unwind(|| { + self.free.extend(std::mem::take(&mut self.used)); + self.size = size_of::(); + }); + } +} + +impl Drop for Allocator { + fn drop(&mut self) { + for raw in self.used.iter().copied() { + unsafe { + let _ = Box::>::from_raw(raw); + } + } + for raw in self.free.iter().copied() { + unsafe { + let _ = Box::>::from_raw(raw); + } + } + unsafe { + let _ = Box::>::from_raw(self.this.cast()); + } + } +} + +#[test] +fn test_allocator() { + let mut allocator = Allocator::new(); + for _ in 0..1024 * 8 { + allocator.malloc::<()>(); + allocator.malloc::(); + allocator.malloc::(); + allocator.malloc::(); + allocator.malloc::<[u8; 32]>(); + allocator.malloc_n::(2); + allocator.malloc::<[u8; 32]>(); + allocator.malloc_n::(2); + allocator.malloc_n::(2); + allocator.malloc_n::(2); + allocator.malloc_n::(2); + allocator.malloc_n::(2); + } + let number_of_chunks_0 = 1 + allocator.used.len() + allocator.free.len(); + allocator.reset(); + for _ in 0..1024 * 8 { + allocator.malloc::<()>(); + allocator.malloc::(); + allocator.malloc::(); + allocator.malloc::(); + allocator.malloc::<[u8; 32]>(); + allocator.malloc_n::(2); + allocator.malloc::<[u8; 32]>(); + allocator.malloc_n::(2); + allocator.malloc_n::(2); + allocator.malloc_n::(2); + allocator.malloc_n::(2); + allocator.malloc_n::(2); + } + let number_of_chunks_1 = 1 + allocator.used.len() + allocator.free.len(); + assert_eq!(number_of_chunks_0, number_of_chunks_1); +} + +pub struct BumpAlloc { + inner: UnsafeCell, +} + +impl BumpAlloc { + pub fn new() -> Self { + Self { + inner: UnsafeCell::new(Allocator::new()), + } + } + pub fn reset(&mut self) { + self.inner.get_mut().reset(); + } +} + +impl Bump for BumpAlloc { + fn alloc(&self, value: T) -> &mut T { + unsafe { + let ptr = (*self.inner.get()).malloc::(); + ptr.write(value); + &mut *ptr + } + } + + fn alloc_slice(&self, slice: &[T]) -> &mut [T] { + unsafe { + let ptr = (*self.inner.get()).malloc_n::(slice.len()); + std::ptr::copy_nonoverlapping(slice.as_ptr(), ptr, slice.len()); + std::slice::from_raw_parts_mut(ptr, slice.len()) + } + } +} + pub fn prewarm( opfamily: Opfamily, index: impl RelationRead, @@ -16,16 +196,16 @@ pub fn prewarm( ) -> String { let message = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::prewarm::, L2>>(index, height, check) + algorithm::prewarm::<_, Op, L2>>(index, height, check) } (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::prewarm::, Dot>>(index, height, check) + algorithm::prewarm::<_, Op, Dot>>(index, height, check) } (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::prewarm::, L2>>(index, height, check) + algorithm::prewarm::<_, Op, L2>>(index, height, check) } (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::prewarm::, Dot>>(index, height, check) + algorithm::prewarm::<_, Op, Dot>>(index, height, check) } }; match message { @@ -36,7 +216,7 @@ pub fn prewarm( pub fn bulkdelete( opfamily: Opfamily, - index: impl RelationWrite, + index: impl RelationRead + RelationWrite, check: impl Fn(), callback: impl Fn(NonZero) -> bool, ) { @@ -56,19 +236,19 @@ pub fn bulkdelete( } } -pub fn maintain(opfamily: Opfamily, index: impl RelationWrite, check: impl Fn()) { +pub fn maintain(opfamily: Opfamily, index: impl RelationRead + RelationWrite, check: impl Fn()) { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::maintain::, L2>>(index, check) + algorithm::maintain::, L2>, _>(index, check) } (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::maintain::, Dot>>(index, check) + algorithm::maintain::, Dot>, _>(index, check) } (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::maintain::, L2>>(index, check) + algorithm::maintain::, L2>, _>(index, check) } (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::maintain::, Dot>>(index, check) + algorithm::maintain::, Dot>, _>(index, check) } } } @@ -109,41 +289,55 @@ pub fn build( pub fn insert( opfamily: Opfamily, - index: impl RelationWrite, + index: impl RelationRead + RelationWrite, payload: NonZero, vector: OwnedVector, ) { + use algorithm::{FastHeap, PlainPrefetcher}; + let bump = BumpAlloc::new(); + let prefetch = { + let index = index.clone(); + move |results| PlainPrefetcher::<_, FastHeap<_>>::new(index.clone(), results) + }; match (vector, opfamily.distance_kind()) { (OwnedVector::Vecf32(vector), DistanceKind::L2) => { assert!(opfamily.vector_kind() == VectorKind::Vecf32); - algorithm::insert::, L2>>( + algorithm::insert::<_, Op, L2>, _>( index, payload, RandomProject::project(vector.as_borrowed()), + &bump, + prefetch, ) } (OwnedVector::Vecf32(vector), DistanceKind::Dot) => { assert!(opfamily.vector_kind() == VectorKind::Vecf32); - algorithm::insert::, Dot>>( + algorithm::insert::<_, Op, Dot>, _>( index, payload, RandomProject::project(vector.as_borrowed()), + &bump, + prefetch, ) } (OwnedVector::Vecf16(vector), DistanceKind::L2) => { assert!(opfamily.vector_kind() == VectorKind::Vecf16); - algorithm::insert::, L2>>( + algorithm::insert::<_, Op, L2>, _>( index, payload, RandomProject::project(vector.as_borrowed()), + &bump, + prefetch, ) } (OwnedVector::Vecf16(vector), DistanceKind::Dot) => { assert!(opfamily.vector_kind() == VectorKind::Vecf16); - algorithm::insert::, Dot>>( + algorithm::insert::<_, Op, Dot>, _>( index, payload, RandomProject::project(vector.as_borrowed()), + &bump, + prefetch, ) } } diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 913dcead..1478d740 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -4,7 +4,7 @@ use crate::index::opclass::{Opfamily, opfamily}; use crate::index::storage::{PostgresPage, PostgresRelation}; use crate::index::types::*; use algorithm::types::*; -use algorithm::{PageGuard, RelationRead, RelationWrite}; +use algorithm::{PageGuard, Relation, RelationRead, RelationWrite}; use half::f16; use pgrx::pg_sys::{Datum, ItemPointerData}; use rand::Rng; @@ -1274,9 +1274,11 @@ impl Deref for CachingRelationReadGuard<'_, G> { } } -impl> RelationRead for CachingRelation<'_, R> { +impl Relation for CachingRelation<'_, R> { type Page = R::Page; +} +impl> RelationRead for CachingRelation<'_, R> { type ReadGuard<'a> = CachingRelationReadGuard<'a, R::ReadGuard<'a>> where diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index ebfca1d5..b1922cd1 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -1,19 +1,20 @@ pub mod am_build; +use super::algorithm::BumpAlloc; +use super::gucs::prererank_filtering; use crate::index::gucs; +use crate::index::lazy_cell::LazyCell; use crate::index::opclass::{Opfamily, opfamily}; use crate::index::scanners::*; use crate::index::storage::PostgresRelation; +use algorithm::Bump; use pgrx::datum::Internal; use pgrx::pg_sys::{BlockIdData, Datum, ItemPointerData}; -use std::cell::LazyCell; use std::ffi::CStr; use std::num::NonZero; use std::ptr::NonNull; use std::sync::OnceLock; -use super::gucs::prererank_filtering; - #[repr(C)] struct Reloption { vl_len_: i32, @@ -204,7 +205,7 @@ pub unsafe extern "C-unwind" fn amcostestimate( let node_count = { let tuples = (*index_opt_info).tuples as u32; let mut count = 0.0; - let r = cost.cells.iter().copied().rev().map(NonZero::get); + let r = cost.cells.iter().copied().rev(); let numerator = std::iter::once(1).chain(probes.clone()); let denumerator = r.clone(); let scale = r.skip(1).chain(std::iter::once(tuples)); @@ -221,7 +222,7 @@ pub unsafe extern "C-unwind" fn amcostestimate( let x = opfamily.vector_kind().element_size() * cost.dims; x.div_ceil(3840 * x.div_ceil(5120).min(2)) as f64 }; - pages += cost.cells[0].get() as f64; + pages += cost.cells[0] as f64; pages }; let next_count = @@ -352,6 +353,7 @@ pub unsafe extern "C-unwind" fn ambeginscan( let scanner: Scanner = Scanner { hack: None, scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), + bump: Box::new(BumpAlloc::new()), }; unsafe { (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); @@ -380,6 +382,9 @@ pub unsafe extern "C-unwind" fn amrescan( "vector search with no WHERE clause and no ORDER BY clause is not supported" ); } + let scanner = &mut *(*scan).opaque.cast::(); + scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + scanner.bump.reset(); let opfamily = opfamily((*scan).indexRelation); let index = PostgresRelation::new((*scan).indexRelation); let options = SearchOptions { @@ -388,8 +393,8 @@ pub unsafe extern "C-unwind" fn amrescan( max_scan_tuples: gucs::max_scan_tuples(), maxsim_refine: gucs::maxsim_refine(), maxsim_threshold: gucs::maxsim_threshold(), + io_rerank: gucs::io_rerank(), }; - let scanner = &mut *(*scan).opaque.cast::(); let fetcher = { let hack = scanner.hack; LazyCell::new(move || { @@ -405,6 +410,8 @@ pub unsafe extern "C-unwind" fn amrescan( ) }) }; + // PAY ATTENTATION: `scanning` references `bump`, so `scanning` must be dropped before `bump`. + let bump = scanner.bump.as_ref(); scanner.scanning = match opfamily { Opfamily::VectorL2 | Opfamily::VectorIp @@ -425,7 +432,11 @@ pub unsafe extern "C-unwind" fn amrescan( let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; builder.add((*data).sk_strategy, (!is_null).then_some(value)); } - LazyCell::new(Box::new(move || builder.build(index, options, fetcher))) + LazyCell::new(Box::new(move || { + // only do this since `PostgresRelation` has no destructor + let index = bump.alloc(index.clone()); + builder.build(index, options, fetcher, bump) + })) } Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim => { let mut builder = MaxsimBuilder::new(opfamily); @@ -441,7 +452,11 @@ pub unsafe extern "C-unwind" fn amrescan( let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; builder.add((*data).sk_strategy, (!is_null).then_some(value)); } - LazyCell::new(Box::new(move || builder.build(index, options, fetcher))) + LazyCell::new(Box::new(move || { + // only do this since `PostgresRelation` has no destructor + let index = bump.alloc(index.clone()); + builder.build(index, options, fetcher, bump) + })) } }; } @@ -480,6 +495,7 @@ pub unsafe extern "C-unwind" fn amgettuple( pub unsafe extern "C-unwind" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { let scanner = unsafe { &mut *(*scan).opaque.cast::() }; scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + scanner.bump.reset(); } type Iter = Box>; @@ -487,6 +503,7 @@ type Iter = Box>; pub struct Scanner { pub hack: Option>, scanning: LazyCell Iter>>, + bump: Box, } struct HeapFetcher { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 4ba75c02..18fd1e67 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -1,6 +1,17 @@ +use super::scanners::SearchIo; +use pgrx::PostgresGucEnum; use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; use std::ffi::CStr; +#[allow(non_camel_case_types)] +#[derive(Debug, Clone, Copy, PostgresGucEnum)] +pub enum Io { + read_buffer, + prefetch_buffer, + #[cfg(feature = "pg17")] + read_stream, +} + static PREWARM_DIM: GucSetting> = GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); @@ -13,6 +24,13 @@ static MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); static PRERERANK_FILTERING: GucSetting = GucSetting::::new(false); +static IO_RERANK: GucSetting = GucSetting::::new( + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + Io::prefetch_buffer, + #[cfg(feature = "pg17")] + Io::read_stream, +); + pub fn init() { GucRegistry::define_string_guc( "vchordrq.probes", @@ -78,6 +96,14 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_enum_guc( + "vchordrq.io_rerank", + "`io_rerank` argument of vchordrq.", + "`io_rerank` argument of vchordrq.", + &IO_RERANK, + GucContext::Userset, + GucFlags::default(), + ); unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); @@ -155,3 +181,12 @@ pub fn prewarm_dim() -> Vec { pub fn prererank_filtering() -> bool { PRERERANK_FILTERING.get() } + +pub fn io_rerank() -> SearchIo { + match IO_RERANK.get() { + Io::read_buffer => SearchIo::ReadBuffer, + Io::prefetch_buffer => SearchIo::PrefetchBuffer, + #[cfg(feature = "pg17")] + Io::read_stream => SearchIo::ReadStream, + } +} diff --git a/src/index/lazy_cell.rs b/src/index/lazy_cell.rs new file mode 100644 index 00000000..db616a57 --- /dev/null +++ b/src/index/lazy_cell.rs @@ -0,0 +1,312 @@ +// Emulate unstable library feature `lazy_get`. +// See https://github.com/rust-lang/rust/issues/129333. + +#![allow(dead_code)] + +use core::cell::UnsafeCell; +use core::hint::unreachable_unchecked; +use core::ops::Deref; + +enum State { + Uninit(F), + Init(T), + Poisoned, +} + +/// A value which is initialized on the first access. +/// +/// For a thread-safe version of this struct, see [`std::sync::LazyLock`]. +/// +/// [`std::sync::LazyLock`]: ../../std/sync/struct.LazyLock.html +/// +/// # Examples +/// +/// ``` +/// use std::cell::LazyCell; +/// +/// let lazy: LazyCell = LazyCell::new(|| { +/// println!("initializing"); +/// 92 +/// }); +/// println!("ready"); +/// println!("{}", *lazy); +/// println!("{}", *lazy); +/// +/// // Prints: +/// // ready +/// // initializing +/// // 92 +/// // 92 +/// ``` +pub struct LazyCell T> { + state: UnsafeCell>, +} + +impl T> LazyCell { + /// Creates a new lazy value with the given initializing function. + /// + /// # Examples + /// + /// ``` + /// use std::cell::LazyCell; + /// + /// let hello = "Hello, World!".to_string(); + /// + /// let lazy = LazyCell::new(|| hello.to_uppercase()); + /// + /// assert_eq!(&*lazy, "HELLO, WORLD!"); + /// ``` + #[inline] + pub const fn new(f: F) -> LazyCell { + LazyCell { + state: UnsafeCell::new(State::Uninit(f)), + } + } + + /// Consumes this `LazyCell` returning the stored value. + /// + /// Returns `Ok(value)` if `Lazy` is initialized and `Err(f)` otherwise. + /// + /// # Examples + /// + /// ``` + /// #![feature(lazy_cell_into_inner)] + /// + /// use std::cell::LazyCell; + /// + /// let hello = "Hello, World!".to_string(); + /// + /// let lazy = LazyCell::new(|| hello.to_uppercase()); + /// + /// assert_eq!(&*lazy, "HELLO, WORLD!"); + /// assert_eq!(LazyCell::into_inner(lazy).ok(), Some("HELLO, WORLD!".to_string())); + /// ``` + pub fn into_inner(this: Self) -> Result { + match this.state.into_inner() { + State::Init(data) => Ok(data), + State::Uninit(f) => Err(f), + State::Poisoned => panic_poisoned(), + } + } + + /// Forces the evaluation of this lazy value and returns a reference to + /// the result. + /// + /// This is equivalent to the `Deref` impl, but is explicit. + /// + /// # Examples + /// + /// ``` + /// use std::cell::LazyCell; + /// + /// let lazy = LazyCell::new(|| 92); + /// + /// assert_eq!(LazyCell::force(&lazy), &92); + /// assert_eq!(&*lazy, &92); + /// ``` + #[inline] + pub fn force(this: &LazyCell) -> &T { + // SAFETY: + // This invalidates any mutable references to the data. The resulting + // reference lives either until the end of the borrow of `this` (in the + // initialized case) or is invalidated in `really_init` (in the + // uninitialized case; `really_init` will create and return a fresh reference). + let state = unsafe { &*this.state.get() }; + match state { + State::Init(data) => data, + // SAFETY: The state is uninitialized. + State::Uninit(_) => unsafe { LazyCell::really_init(this) }, + State::Poisoned => panic_poisoned(), + } + } + + /// Forces the evaluation of this lazy value and returns a mutable reference to + /// the result. + /// + /// # Examples + /// + /// ``` + /// #![feature(lazy_get)] + /// use std::cell::LazyCell; + /// + /// let mut lazy = LazyCell::new(|| 92); + /// + /// let p = LazyCell::force_mut(&mut lazy); + /// assert_eq!(*p, 92); + /// *p = 44; + /// assert_eq!(*lazy, 44); + /// ``` + #[inline] + pub fn force_mut(this: &mut LazyCell) -> &mut T { + #[cold] + /// # Safety + /// May only be called when the state is `Uninit`. + unsafe fn really_init_mut T>(state: &mut State) -> &mut T { + // INVARIANT: Always valid, but the value may not be dropped. + struct PoisonOnPanic(*mut State); + impl Drop for PoisonOnPanic { + #[inline] + fn drop(&mut self) { + // SAFETY: Invariant states it is valid, and we don't drop the old value. + unsafe { + self.0.write(State::Poisoned); + } + } + } + + let State::Uninit(f) = state else { + // `unreachable!()` here won't optimize out because the function is cold. + // SAFETY: Precondition. + unsafe { unreachable_unchecked() }; + }; + // SAFETY: We never drop the state after we read `f`, and we write a valid value back + // in any case, panic or success. `f` can't access the `LazyCell` because it is mutably + // borrowed. + let f = unsafe { core::ptr::read(f) }; + // INVARIANT: Initiated from mutable reference, don't drop because we read it. + let guard = PoisonOnPanic(state); + let data = f(); + // SAFETY: `PoisonOnPanic` invariant, and we don't drop the old value. + unsafe { + core::ptr::write(guard.0, State::Init(data)); + } + core::mem::forget(guard); + let State::Init(data) = state else { + unreachable!() + }; + data + } + + let state = this.state.get_mut(); + match state { + State::Init(data) => data, + // SAFETY: `state` is `Uninit`. + State::Uninit(_) => unsafe { really_init_mut(state) }, + State::Poisoned => panic_poisoned(), + } + } + + /// # Safety + /// May only be called when the state is `Uninit`. + #[cold] + unsafe fn really_init(this: &LazyCell) -> &T { + // SAFETY: + // This function is only called when the state is uninitialized, + // so no references to `state` can exist except for the reference + // in `force`, which is invalidated here and not accessed again. + let state = unsafe { &mut *this.state.get() }; + // Temporarily mark the state as poisoned. This prevents reentrant + // accesses and correctly poisons the cell if the closure panicked. + let State::Uninit(f) = core::mem::replace(state, State::Poisoned) else { + unreachable!() + }; + + let data = f(); + + // SAFETY: + // If the closure accessed the cell through something like a reentrant + // mutex, but caught the panic resulting from the state being poisoned, + // the mutable borrow for `state` will be invalidated, so we need to + // go through the `UnsafeCell` pointer here. The state can only be + // poisoned at this point, so using `write` to skip the destructor + // of `State` should help the optimizer. + unsafe { this.state.get().write(State::Init(data)) }; + + // SAFETY: + // The previous references were invalidated by the `write` call above, + // so do a new shared borrow of the state instead. + let state = unsafe { &*this.state.get() }; + let State::Init(data) = state else { + unreachable!() + }; + data + } +} + +impl LazyCell { + /// Returns a mutable reference to the value if initialized, or `None` if not. + /// + /// # Examples + /// + /// ``` + /// #![feature(lazy_get)] + /// + /// use std::cell::LazyCell; + /// + /// let mut lazy = LazyCell::new(|| 92); + /// + /// assert_eq!(LazyCell::get_mut(&mut lazy), None); + /// let _ = LazyCell::force(&lazy); + /// *LazyCell::get_mut(&mut lazy).unwrap() = 44; + /// assert_eq!(*lazy, 44); + /// ``` + #[inline] + pub fn get_mut(this: &mut LazyCell) -> Option<&mut T> { + let state = this.state.get_mut(); + match state { + State::Init(data) => Some(data), + _ => None, + } + } + + /// Returns a reference to the value if initialized, or `None` if not. + /// + /// # Examples + /// + /// ``` + /// #![feature(lazy_get)] + /// + /// use std::cell::LazyCell; + /// + /// let lazy = LazyCell::new(|| 92); + /// + /// assert_eq!(LazyCell::get(&lazy), None); + /// let _ = LazyCell::force(&lazy); + /// assert_eq!(LazyCell::get(&lazy), Some(&92)); + /// ``` + #[inline] + pub fn get(this: &LazyCell) -> Option<&T> { + // SAFETY: + // This is sound for the same reason as in `force`: once the state is + // initialized, it will not be mutably accessed again, so this reference + // will stay valid for the duration of the borrow to `self`. + let state = unsafe { &*this.state.get() }; + match state { + State::Init(data) => Some(data), + _ => None, + } + } +} + +impl T> Deref for LazyCell { + type Target = T; + #[inline] + fn deref(&self) -> &T { + LazyCell::force(self) + } +} + +impl Default for LazyCell { + /// Creates a new lazy value using `Default` as the initializing function. + #[inline] + fn default() -> LazyCell { + LazyCell::new(T::default) + } +} + +impl core::fmt::Debug for LazyCell { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut d = f.debug_tuple("LazyCell"); + match LazyCell::get(self) { + Some(data) => d.field(data), + None => d.field(&format_args!("")), + }; + d.finish() + } +} + +#[cold] +#[inline(never)] +const fn panic_poisoned() -> ! { + panic!("LazyCell instance has previously been poisoned") +} diff --git a/src/index/mod.rs b/src/index/mod.rs index 453de139..433a4c41 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -3,6 +3,7 @@ pub mod am; pub mod functions; pub mod gucs; pub mod hook; +pub mod lazy_cell; pub mod opclass; pub mod projection; pub mod scanners; diff --git a/src/index/scanners/default.rs b/src/index/scanners/default.rs index 615d961d..fa48cc4a 100644 --- a/src/index/scanners/default.rs +++ b/src/index/scanners/default.rs @@ -1,11 +1,12 @@ -use super::{SearchBuilder, SearchFetcher, SearchOptions}; +use super::{SearchBuilder, SearchFetcher, SearchIo, SearchOptions}; use crate::index::algorithm::RandomProject; use crate::index::am::pointer_to_kv; use crate::index::opclass::{Opfamily, Sphere}; use algorithm::operator::{Dot, L2, Op}; use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; -use algorithm::{RelationRead, RerankMethod}; +use algorithm::*; use half::f16; +use std::collections::BinaryHeap; use std::num::NonZero; use vector::VectorOwned; use vector::vect::VectOwned; @@ -50,9 +51,10 @@ impl SearchBuilder for DefaultBuilder { fn build<'a>( self, - relation: impl RelationRead + 'a, + relation: &'a (impl RelationPrefetch + RelationReadStream), options: SearchOptions, mut fetcher: impl SearchFetcher + 'a, + bump: &'a impl Bump, ) -> Box + 'a> { let mut vector = None; let mut threshold = None; @@ -86,11 +88,18 @@ impl SearchBuilder for DefaultBuilder { } .as_borrowed(), ); - let results = algorithm::default_search::, L2>>( + let results = default_search::<_, Op, L2>, _>( relation.clone(), vector.clone(), options.probes, options.epsilon, + bump, + { + let index = relation.clone(); + move |results| { + PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) + } + }, ); let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); @@ -104,20 +113,48 @@ impl SearchBuilder for DefaultBuilder { }; Some(RandomProject::project(raw.as_borrowed())) }; - let method = algorithm::how(relation.clone()); - match method { - RerankMethod::Index => Box::new( - algorithm::rerank_index::, L2>, _>( - relation, vector, results, + let method = how(relation.clone()); + match (method, options.io_rerank) { + (RerankMethod::Index, SearchIo::ReadBuffer) => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + Box::new( + rerank_index::, L2>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)), - ), - RerankMethod::Heap => Box::new( - algorithm::rerank_heap::, L2>, _>( - vector, results, fetch, + } + (RerankMethod::Index, SearchIo::PrefetchBuffer) => { + let prefetcher = SimplePrefetcher::new(relation.clone(), results); + Box::new( + rerank_index::, L2>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), + ) + } + (RerankMethod::Index, SearchIo::ReadStream) => { + let prefetcher = StreamPrefetcher::new(relation, results); + Box::new( + rerank_index::, L2>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)), - ), + } + (RerankMethod::Heap, _) => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + Box::new( + rerank_heap::, L2>, _, _>( + vector, prefetcher, fetch, + ) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), + ) + } } } (VectorKind::Vecf32, DistanceKind::Dot) => { @@ -129,11 +166,18 @@ impl SearchBuilder for DefaultBuilder { } .as_borrowed(), ); - let results = algorithm::default_search::, Dot>>( + let results = default_search::<_, Op, Dot>, _>( relation.clone(), vector.clone(), options.probes, options.epsilon, + bump, + { + let index = relation.clone(); + move |results| { + PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) + } + }, ); let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); @@ -147,20 +191,48 @@ impl SearchBuilder for DefaultBuilder { }; Some(RandomProject::project(raw.as_borrowed())) }; - let method = algorithm::how(relation.clone()); - match method { - RerankMethod::Index => Box::new( - algorithm::rerank_index::, Dot>, _>( - relation, vector, results, + let method = how(relation.clone()); + match (method, options.io_rerank) { + (RerankMethod::Index, SearchIo::ReadBuffer) => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + Box::new( + rerank_index::, Dot>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)), - ), - RerankMethod::Heap => Box::new( - algorithm::rerank_heap::, Dot>, _>( - vector, results, fetch, + } + (RerankMethod::Index, SearchIo::PrefetchBuffer) => { + let prefetcher = SimplePrefetcher::new(relation.clone(), results); + Box::new( + rerank_index::, Dot>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), + ) + } + (RerankMethod::Index, SearchIo::ReadStream) => { + let prefetcher = StreamPrefetcher::new(relation, results); + Box::new( + rerank_index::, Dot>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)), - ), + } + (RerankMethod::Heap, _) => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + Box::new( + rerank_heap::, Dot>, _, _>( + vector, prefetcher, fetch, + ) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), + ) + } } } (VectorKind::Vecf16, DistanceKind::L2) => { @@ -172,11 +244,18 @@ impl SearchBuilder for DefaultBuilder { } .as_borrowed(), ); - let results = algorithm::default_search::, L2>>( + let results = default_search::<_, Op, L2>, _>( relation.clone(), vector.clone(), options.probes, options.epsilon, + bump, + { + let index = relation.clone(); + move |results| { + PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) + } + }, ); let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); @@ -190,20 +269,48 @@ impl SearchBuilder for DefaultBuilder { }; Some(RandomProject::project(raw.as_borrowed())) }; - let method = algorithm::how(relation.clone()); - match method { - RerankMethod::Index => Box::new( - algorithm::rerank_index::, L2>, _>( - relation, vector, results, + let method = how(relation.clone()); + match (method, options.io_rerank) { + (RerankMethod::Index, SearchIo::ReadBuffer) => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + Box::new( + rerank_index::, L2>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)), - ), - RerankMethod::Heap => Box::new( - algorithm::rerank_heap::, L2>, _>( - vector, results, fetch, + } + (RerankMethod::Index, SearchIo::PrefetchBuffer) => { + let prefetcher = SimplePrefetcher::new(relation.clone(), results); + Box::new( + rerank_index::, L2>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), + ) + } + (RerankMethod::Index, SearchIo::ReadStream) => { + let prefetcher = StreamPrefetcher::new(relation, results); + Box::new( + rerank_index::, L2>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)), - ), + } + (RerankMethod::Heap, _) => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + Box::new( + rerank_heap::, L2>, _, _>( + vector, prefetcher, fetch, + ) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), + ) + } } } (VectorKind::Vecf16, DistanceKind::Dot) => { @@ -215,11 +322,18 @@ impl SearchBuilder for DefaultBuilder { } .as_borrowed(), ); - let results = algorithm::default_search::, Dot>>( + let results = default_search::<_, Op, Dot>, _>( relation.clone(), vector.clone(), options.probes, options.epsilon, + bump, + { + let index = relation.clone(); + move |results| { + PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) + } + }, ); let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); @@ -233,20 +347,48 @@ impl SearchBuilder for DefaultBuilder { }; Some(RandomProject::project(raw.as_borrowed())) }; - let method = algorithm::how(relation.clone()); - match method { - RerankMethod::Index => Box::new( - algorithm::rerank_index::, Dot>, _>( - relation, vector, results, + let method = how(relation.clone()); + match (method, options.io_rerank) { + (RerankMethod::Index, SearchIo::ReadBuffer) => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + Box::new( + rerank_index::, Dot>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)), - ), - RerankMethod::Heap => Box::new( - algorithm::rerank_heap::, Dot>, _>( - vector, results, fetch, + } + (RerankMethod::Index, SearchIo::PrefetchBuffer) => { + let prefetcher = SimplePrefetcher::new(relation.clone(), results); + Box::new( + rerank_index::, Dot>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), + ) + } + (RerankMethod::Index, SearchIo::ReadStream) => { + let prefetcher = StreamPrefetcher::new(relation, results); + Box::new( + rerank_index::, Dot>, _, _>(vector, prefetcher) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), ) - .map(move |(distance, payload)| (opfamily.output(distance), payload)), - ), + } + (RerankMethod::Heap, _) => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + Box::new( + rerank_heap::, Dot>, _, _>( + vector, prefetcher, fetch, + ) + .map(move |(distance, payload)| { + (opfamily.output(distance), payload) + }), + ) + } } } }; diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs index 8cfc8d16..78f41f69 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/scanners/maxsim.rs @@ -2,14 +2,16 @@ use super::{SearchBuilder, SearchFetcher, SearchOptions}; use crate::index::algorithm::RandomProject; use crate::index::am::pointer_to_kv; use crate::index::opclass::Opfamily; +use crate::index::scanners::SearchIo; use algorithm::operator::Dot; use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; -use algorithm::{RelationRead, RerankMethod}; +use algorithm::*; use always_equal::AlwaysEqual; use distance::Distance; use half::f16; use std::cmp::Reverse; use std::collections::BinaryHeap; +use std::num::NonZero; use vector::VectorOwned; use vector::vect::VectOwned; @@ -42,9 +44,10 @@ impl SearchBuilder for MaxsimBuilder { fn build<'a>( self, - relation: impl RelationRead + 'a, + relation: &'a (impl RelationPrefetch + RelationReadStream), options: SearchOptions, _: impl SearchFetcher + 'a, + bump: &'a impl Bump, ) -> Box + 'a> { let mut vectors = None; for orderby_vectors in self.orderbys.into_iter().flatten() { @@ -63,17 +66,20 @@ impl SearchBuilder for MaxsimBuilder { let Some(vectors) = vectors else { return Box::new(std::iter::empty()) as Box>; }; - let method = algorithm::how(relation.clone()); + let method = how(relation.clone()); if !matches!(method, RerankMethod::Index) { pgrx::error!("maxsim search with rerank_in_table is not supported"); } assert!(matches!(opfamily.distance_kind(), DistanceKind::Dot)); let n = vectors.len(); - let accu_map = |(Reverse(dis), AlwaysEqual(pay))| (dis, pay); - let rough_map = |(_, AlwaysEqual(dis), AlwaysEqual(pay), _)| (dis, pay); + let accu_map = |(Reverse(distance), AlwaysEqual(payload))| (distance, payload); + let rough_map = |((_, AlwaysEqual(rough)), AlwaysEqual(&mut (payload, ..))): ( + _, + AlwaysEqual<&mut (NonZero, _, _)>, + )| (rough, payload); let iter: Box> = match opfamily.vector_kind() { VectorKind::Vecf32 => { - type Op = algorithm::operator::Op, Dot>; + type Op = operator::Op, Dot>; let vectors = vectors .into_iter() .map(|vector| { @@ -88,24 +94,54 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new(vectors.into_iter().map(|vector| { - let (results, estimation_by_threshold) = algorithm::maxsim_search::( + let (results, estimation_by_threshold) = maxsim_search::<_, Op, _>( relation.clone(), vector.clone(), options.probes.clone(), options.epsilon, maxsim_threshold, + bump, + { + let index = relation.clone(); + move |results| { + PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) + } + }, ); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { - let mut reranker = algorithm::rerank_index::( - relation.clone(), - vector.clone(), - results, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.map(rough_map)); + match options.io_rerank { + SearchIo::ReadBuffer => { + let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new( + relation.clone(), + results, + ); + let mut reranker = + rerank_index::(vector.clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + SearchIo::PrefetchBuffer => { + let prefetcher = SimplePrefetcher::new(relation.clone(), results); + let mut reranker = + rerank_index::(vector.clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + SearchIo::ReadStream => { + let prefetcher = StreamPrefetcher::new(relation, results); + let mut reranker = + rerank_index::(vector.clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + } } else { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); @@ -114,7 +150,7 @@ impl SearchBuilder for MaxsimBuilder { })) } VectorKind::Vecf16 => { - type Op = algorithm::operator::Op, Dot>; + type Op = operator::Op, Dot>; let vectors = vectors .into_iter() .map(|vector| { @@ -129,24 +165,54 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new(vectors.into_iter().map(|vector| { - let (results, estimation_by_threshold) = algorithm::maxsim_search::( + let (results, estimation_by_threshold) = maxsim_search::<_, Op, _>( relation.clone(), vector.clone(), options.probes.clone(), options.epsilon, maxsim_threshold, + bump, + { + let index = relation.clone(); + move |results| { + PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) + } + }, ); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { - let mut reranker = algorithm::rerank_index::( - relation.clone(), - vector.clone(), - results, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.map(rough_map)); + match options.io_rerank { + SearchIo::ReadBuffer => { + let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new( + relation.clone(), + results, + ); + let mut reranker = + rerank_index::(vector.clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + SearchIo::PrefetchBuffer => { + let prefetcher = SimplePrefetcher::new(relation.clone(), results); + let mut reranker = + rerank_index::(vector.clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + SearchIo::ReadStream => { + let prefetcher = StreamPrefetcher::new(relation, results); + let mut reranker = + rerank_index::(vector.clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + } } else { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); @@ -211,6 +277,9 @@ impl SearchBuilder for MaxsimBuilder { } } +// Emulate unstable library feature `binary_heap_into_iter_sorted`. +// See https://github.com/rust-lang/rust/issues/59278. + pub trait IntoIterSortedPolyfill { fn into_iter_sorted_polyfill(self) -> IntoIterSorted; } diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs index f7b46d8a..c5fea0d9 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/scanners/mod.rs @@ -2,13 +2,24 @@ mod default; mod maxsim; use super::opclass::Opfamily; -use algorithm::RelationRead; +use crate::index::lazy_cell::LazyCell; +use algorithm::{Bump, RelationPrefetch, RelationReadStream}; use pgrx::pg_sys::Datum; -use std::cell::LazyCell; pub use default::DefaultBuilder; pub use maxsim::MaxsimBuilder; +#[derive(Debug, Clone, Copy)] +pub enum SearchIo { + ReadBuffer, + PrefetchBuffer, + #[cfg_attr( + any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"), + expect(dead_code) + )] + ReadStream, +} + #[derive(Debug)] pub struct SearchOptions { pub epsilon: f32, @@ -16,6 +27,7 @@ pub struct SearchOptions { pub max_scan_tuples: Option, pub maxsim_refine: u32, pub maxsim_threshold: u32, + pub io_rerank: SearchIo, } pub trait SearchBuilder: 'static { @@ -25,9 +37,10 @@ pub trait SearchBuilder: 'static { fn build<'a>( self, - relation: impl RelationRead + 'a, + relation: &'a (impl RelationPrefetch + RelationReadStream), options: SearchOptions, fetcher: impl SearchFetcher + 'a, + bump: &'a impl Bump, ) -> Box + 'a>; } diff --git a/src/index/storage.rs b/src/index/storage.rs index 5d0cfc94..8008aa2c 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -1,4 +1,6 @@ -use algorithm::{Opaque, Page, PageGuard, RelationRead, RelationWrite}; +use algorithm::*; +use std::collections::VecDeque; +use std::iter::{Chain, Flatten}; use std::mem::{MaybeUninit, offset_of}; use std::ops::{Deref, DerefMut}; use std::ptr::NonNull; @@ -157,6 +159,25 @@ impl Page for PostgresPage { const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); +pub struct PostgresBufferGuard { + buf: i32, + id: u32, +} + +impl PageGuard for PostgresBufferGuard { + fn id(&self) -> u32 { + self.id + } +} + +impl Drop for PostgresBufferGuard { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::ReleaseBuffer(self.buf); + } + } +} + pub struct PostgresBufferReadGuard { buf: i32, page: NonNull, @@ -242,9 +263,11 @@ impl PostgresRelation { } } -impl RelationRead for PostgresRelation { +impl Relation for PostgresRelation { type Page = PostgresPage; +} +impl RelationRead for PostgresRelation { type ReadGuard<'a> = PostgresBufferReadGuard; fn read(&self, id: u32) -> Self::ReadGuard<'_> { @@ -355,3 +378,220 @@ impl RelationWrite for PostgresRelation { } } } + +impl RelationPrefetch for PostgresRelation { + fn prefetch(&self, id: u32) { + assert!(id != u32::MAX, "no such page"); + unsafe { + use pgrx::pg_sys::PrefetchBuffer; + PrefetchBuffer(self.raw, 0, id); + } + } +} + +pub struct Cache { + window: VecDeque, + tail: VecDeque, + iter: Option, +} + +impl Default for Cache { + fn default() -> Self { + Self { + window: Default::default(), + tail: Default::default(), + iter: Default::default(), + } + } +} + +impl Cache +where + I::Item: Fetch, +{ + #[allow(dead_code)] + pub fn pop_id(&mut self) -> Option { + while self.tail.is_empty() + && let Some(iter) = self.iter.as_mut() + && let Some(e) = iter.next() + { + for id in e.fetch().iter().copied() { + self.tail.push_back(id); + } + self.window.push_back(e); + } + self.tail.pop_front() + } + #[allow(dead_code)] + pub fn pop_item_if(&mut self, predicate: impl FnOnce(&I::Item) -> bool) -> Option { + while self.window.is_empty() + && let Some(iter) = self.iter.as_mut() + && let Some(e) = iter.next() + { + for id in e.fetch().iter().copied() { + self.tail.push_back(id); + } + self.window.push_back(e); + } + vec_deque_pop_front_if(&mut self.window, predicate) + } +} + +pub struct PostgresReadStream { + #[cfg(feature = "pg17")] + raw: *mut pgrx::pg_sys::ReadStream, + // Because of `Box`'s special alias rules, `Box` cannot be used here. + cache: NonNull>, +} + +impl ReadStream for PostgresReadStream +where + I::Item: Fetch, +{ + type Relation = PostgresRelation; + + type Inner = + Chain, Flatten>>; + + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + fn next_if( + &mut self, + _predicate: impl FnOnce(&I::Item) -> bool, + ) -> Option<(I::Item, Vec)> { + panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); + } + + #[cfg(feature = "pg17")] + fn next_if( + &mut self, + predicate: impl FnOnce(&I::Item) -> bool, + ) -> Option<(I::Item, Vec)> { + if let Some(e) = unsafe { self.cache.as_mut().pop_item_if(predicate) } { + let list = e + .fetch() + .iter() + .map(|_| unsafe { + use pgrx::pg_sys::{ + BUFFER_LOCK_SHARE, BufferGetPage, LockBuffer, read_stream_next_buffer, + }; + let buf = read_stream_next_buffer(self.raw, core::ptr::null_mut()); + LockBuffer(buf, BUFFER_LOCK_SHARE as _); + let page = NonNull::new(BufferGetPage(buf).cast()).expect("failed to get page"); + PostgresBufferReadGuard { + buf, + page, + id: pgrx::pg_sys::BufferGetBlockNumber(buf), + } + }) + .collect::>(); + Some((e, list)) + } else { + None + } + } + + fn into_inner(mut self) -> Self::Inner { + let cache = unsafe { std::mem::take(self.cache.as_mut()) }; + cache + .window + .into_iter() + .chain(cache.iter.into_iter().flatten()) + } +} + +impl Drop for PostgresReadStream { + fn drop(&mut self) { + unsafe { + let _ = std::mem::take(self.cache.as_mut()); + #[cfg(feature = "pg17")] + if !std::thread::panicking() && pgrx::pg_sys::IsTransactionState() { + pgrx::pg_sys::read_stream_end(self.raw); + } + let _ = box_from_non_null(self.cache); + } + } +} + +impl RelationReadStream for PostgresRelation { + type ReadStream<'s, I: Iterator> + = PostgresReadStream + where + I::Item: Fetch; + + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + fn read_stream(&self, _iter: I) -> Self::ReadStream<'_, I> + where + I::Item: Fetch, + { + panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); + } + + #[cfg(feature = "pg17")] + fn read_stream(&self, iter: I) -> Self::ReadStream<'_, I> + where + I::Item: Fetch, + { + #[pgrx::pg_guard] + unsafe extern "C-unwind" fn callback( + _stream: *mut pgrx::pg_sys::ReadStream, + callback_private_data: *mut core::ffi::c_void, + _per_buffer_data: *mut core::ffi::c_void, + ) -> pgrx::pg_sys::BlockNumber + where + I::Item: Fetch, + { + unsafe { + use pgrx::pg_sys::InvalidBlockNumber; + let inner = callback_private_data.cast::>(); + (*inner).pop_id().unwrap_or(InvalidBlockNumber) + } + } + let cache = box_into_non_null(Box::new(Cache { + window: VecDeque::new(), + tail: VecDeque::new(), + iter: Some(iter), + })); + let raw = unsafe { + use pgrx::pg_sys::{ForkNumber, READ_STREAM_DEFAULT, read_stream_begin_relation}; + read_stream_begin_relation( + READ_STREAM_DEFAULT as i32, + core::ptr::null_mut(), + self.raw, + ForkNumber::MAIN_FORKNUM, + Some(callback::), + cache.as_ptr().cast(), + 0, + ) + }; + PostgresReadStream { raw, cache } + } +} + +// Emulate unstable library feature `vec_deque_pop_if`. +// See https://github.com/rust-lang/rust/issues/135889. + +fn vec_deque_pop_front_if( + this: &mut VecDeque, + predicate: impl FnOnce(&T) -> bool, +) -> Option { + let first = this.front()?; + if predicate(first) { + this.pop_front() + } else { + None + } +} + +// Emulate unstable library feature `box_vec_non_null`. +// See https://github.com/rust-lang/rust/issues/130364. + +#[allow(dead_code)] +#[must_use] +fn box_into_non_null(b: Box) -> NonNull { + unsafe { NonNull::new_unchecked(Box::into_raw(b)) } +} + +#[must_use] +unsafe fn box_from_non_null(ptr: NonNull) -> Box { + unsafe { Box::from_raw(ptr.as_ptr()) } +} diff --git a/src/lib.rs b/src/lib.rs index e51d6145..49b14676 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ #![allow(unsafe_code)] -#![feature(lazy_get)] +#![feature(let_chains)] +#![feature(abort_unwind)] mod datatype; mod index; @@ -26,7 +27,8 @@ extern "C-unwind" fn _PG_init() { #[cfg(not(target_endian = "little"))] compile_error!("Target architecture is not supported."); +#[cfg(not(miri))] #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] -#[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))] +#[cfg(target_os = "linux")] #[global_allocator] -static GLOBAL_ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; +static GLOBAL_ALLOCATOR: jemallocator::Jemalloc = jemallocator::Jemalloc; From 266baf34b53ceb96a1226bd0aaa9b832cfdd6a7b Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 25 Apr 2025 19:01:20 +0800 Subject: [PATCH 152/324] chore: update compiler to nightly-2025-04-25 (#241) AVX512 (probably excluding AVX512FP16) is being stabilized. With the next compiler version bump, we might finally be able to get rid of nightly toolchain. Signed-off-by: usamoi --- crates/algorithm/src/lib.rs | 1 - crates/algorithm/src/prewarm.rs | 6 +- crates/simd/build.rs | 27 ++++++-- crates/simd/{cshim.c => cshim/aarch64.c} | 4 -- crates/simd/cshim/x86_64.c | 56 +++++++++++++++ crates/simd/src/f16.rs | 88 ++++++++---------------- crates/simd/src/f32.rs | 4 +- crates/simd/src/lib.rs | 15 +++- crates/simd_macros/src/target.rs | 9 +++ rust-toolchain.toml | 2 +- src/datatype/text_scalar8.rs | 4 +- src/datatype/typmod.rs | 2 +- src/index/algorithm.rs | 13 +++- src/lib.rs | 2 - 14 files changed, 148 insertions(+), 85 deletions(-) rename crates/simd/{cshim.c => cshim/aarch64.c} (99%) create mode 100644 crates/simd/cshim/x86_64.c diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index a3c25586..a167b634 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -1,4 +1,3 @@ -#![feature(let_chains)] #![allow(clippy::type_complexity)] mod build; diff --git a/crates/algorithm/src/prewarm.rs b/crates/algorithm/src/prewarm.rs index 75466ad3..fa4d0138 100644 --- a/crates/algorithm/src/prewarm.rs +++ b/crates/algorithm/src/prewarm.rs @@ -20,7 +20,7 @@ pub fn prewarm( drop(meta_guard); let mut message = String::new(); - writeln!(message, "height of root: {}", height_of_root)?; + writeln!(message, "height of root: {height_of_root}")?; let prewarm_max_height = if height < 0 { 0 } else { height as u32 }; if prewarm_max_height > height_of_root { return Ok(message); @@ -59,7 +59,7 @@ pub fn prewarm( } writeln!(message, "------------------------")?; writeln!(message, "number of nodes: {}", results.len())?; - writeln!(message, "number of pages: {}", counter)?; + writeln!(message, "number of pages: {counter}")?; Ok(results) }; for _ in (std::cmp::max(1, prewarm_max_height)..height_of_root).rev() { @@ -99,7 +99,7 @@ pub fn prewarm( } writeln!(message, "------------------------")?; writeln!(message, "number of nodes: {}", results.len())?; - writeln!(message, "number of pages: {}", counter)?; + writeln!(message, "number of pages: {counter}")?; } Ok(message) } diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 10fad2c6..f9a84d92 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -1,7 +1,22 @@ -fn main() { - println!("cargo::rerun-if-changed=cshim.c"); - let mut build = cc::Build::new(); - build.file("cshim.c"); - build.opt_level(3); - build.compile("simd_cshim"); +use std::error::Error; + +fn main() -> Result<(), Box> { + println!("cargo::rerun-if-changed=cshim"); + let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH")?; + match target_arch.as_str() { + "aarch64" => { + let mut build = cc::Build::new(); + build.file("./cshim/aarch64.c"); + build.opt_level(3); + build.compile("simd_cshim"); + } + "x86_64" => { + let mut build = cc::Build::new(); + build.file("./cshim/x86_64.c"); + build.opt_level(3); + build.compile("simd_cshim"); + } + _ => (), + } + Ok(()) } diff --git a/crates/simd/cshim.c b/crates/simd/cshim/aarch64.c similarity index 99% rename from crates/simd/cshim.c rename to crates/simd/cshim/aarch64.c index da02e8f4..30b5350c 100644 --- a/crates/simd/cshim.c +++ b/crates/simd/cshim/aarch64.c @@ -10,8 +10,6 @@ #error "This file requires Clang or GCC." #endif -#ifdef __aarch64__ - #include #include #include @@ -225,5 +223,3 @@ fp32_reduce_sum_of_d2_a3_256(float *restrict lhs, float *restrict rhs, } return svaddv_f32(svptrue_b32(), sum); } - -#endif diff --git a/crates/simd/cshim/x86_64.c b/crates/simd/cshim/x86_64.c new file mode 100644 index 00000000..5e0f2524 --- /dev/null +++ b/crates/simd/cshim/x86_64.c @@ -0,0 +1,56 @@ +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#else +#error "This file requires Clang." +#endif + +#include +#include +#include + +typedef _Float16 f16; +typedef float f32; + +__attribute__((target("arch=x86-64-v4,avx512fp16"))) float +fp16_reduce_sum_of_xy_v4fp16(f16 *restrict a, f16 *restrict b, size_t n) { + __m512h xy = _mm512_setzero_ph(); + while (n >= 32) { + __m512h x = _mm512_loadu_ph(a); + __m512h y = _mm512_loadu_ph(b); + a += 32; + b += 32; + n -= 32; + xy = _mm512_fmadd_ph(x, y, xy); + } + if (n > 0) { + unsigned int mask = _bzhi_u32(0xffffffff, n); + __m512h x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a)); + __m512h y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b)); + xy = _mm512_fmadd_ph(x, y, xy); + } + return _mm512_reduce_add_ph(xy); +} + +__attribute__((target("arch=x86-64-v4,avx512fp16"))) float +fp16_reduce_sum_of_d2_v4fp16(f16 *restrict a, f16 *restrict b, size_t n) { + __m512h d2 = _mm512_setzero_ph(); + while (n >= 32) { + __m512h x = _mm512_loadu_ph(a); + __m512h y = _mm512_loadu_ph(b); + a += 32; + b += 32; + n -= 32; + __m512h d = _mm512_sub_ph(x, y); + d2 = _mm512_fmadd_ph(d, d, d2); + } + if (n > 0) { + unsigned int mask = _bzhi_u32(0xffffffff, n); + __m512h x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a)); + __m512h y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b)); + __m512h d = _mm512_sub_ph(x, y); + d2 = _mm512_fmadd_ph(d, d, d2); + } + return _mm512_reduce_add_ph(d2); +} diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 7b5eed71..53962f8a 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -225,39 +225,25 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - #[target_feature(enable = "avx512fp16")] - pub fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + #[crate::target_cpu(enable = "v4fp16")] + pub fn reduce_sum_of_xy_v4fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut xy = _mm512_setzero_ph(); - while n >= 32 { - let x = unsafe { _mm512_loadu_ph(a.cast()) }; - let y = unsafe { _mm512_loadu_ph(b.cast()) }; - a = unsafe { a.add(32) }; - b = unsafe { b.add(32) }; - n -= 32; - xy = _mm512_fmadd_ph(x, y, xy); - } - if n > 0 { - let mask = _bzhi_u32(0xffffffff, n as u32); - let x = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; - let y = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())) }; - xy = _mm512_fmadd_ph(x, y, xy); + unsafe { + unsafe extern "C" { + unsafe fn fp16_reduce_sum_of_xy_v4fp16(a: *const (), b: *const (), n: usize) + -> f32; + } + fp16_reduce_sum_of_xy_v4fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } - _mm512_reduce_add_ph(xy) as f32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_xy_v4_avx512fp16_test() { + fn reduce_sum_of_xy_v4fp16_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { - println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + if !crate::is_cpu_detected!("v4fp16") { + println!("test {} ... skipped (v4fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -272,7 +258,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v4_avx512fp16(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v4fp16(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -493,7 +479,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -510,41 +496,25 @@ mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - #[target_feature(enable = "avx512fp16")] - pub fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + #[crate::target_cpu(enable = "v4fp16")] + pub fn reduce_sum_of_d2_v4fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; - let mut n = lhs.len() as u32; - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let mut d2 = _mm512_setzero_ph(); - while n >= 32 { - let x = unsafe { _mm512_loadu_ph(a.cast()) }; - let y = unsafe { _mm512_loadu_ph(b.cast()) }; - a = unsafe { a.add(32) }; - b = unsafe { b.add(32) }; - n -= 32; - let d = _mm512_sub_ph(x, y); - d2 = _mm512_fmadd_ph(d, d, d2); - } - if n > 0 { - let mask = _bzhi_u32(0xffffffff, n); - let x = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; - let y = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())) }; - let d = _mm512_sub_ph(x, y); - d2 = _mm512_fmadd_ph(d, d, d2); + unsafe { + unsafe extern "C" { + unsafe fn fp16_reduce_sum_of_d2_v4fp16(a: *const (), b: *const (), n: usize) + -> f32; + } + fp16_reduce_sum_of_d2_v4fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } - _mm512_reduce_add_ph(d2) as f32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_d2_v4_avx512fp16_test() { + fn reduce_sum_of_d2_v4fp16_test() { use rand::Rng; const EPSILON: f32 = 6.4; - if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { - println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + if !crate::is_cpu_detected!("v4fp16") { + println!("test {} ... skipped (v4fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -559,7 +529,7 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v4_avx512fp16(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_d2_v4fp16(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -575,7 +545,7 @@ mod reduce_sum_of_d2 { pub fn reduce_sum_of_d2_v4(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; - let mut n = lhs.len() as u32; + let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut d2 = _mm512_setzero_ps(); @@ -589,7 +559,7 @@ mod reduce_sum_of_d2 { d2 = _mm512_fmadd_ps(d, d, d2); } if n > 0 { - let mask = _bzhi_u32(0xffff, n) as u16; + let mask = _bzhi_u32(0xffff, n as u32) as u16; let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; let y = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())) }; let d = _mm512_sub_ps(x, y); @@ -636,7 +606,7 @@ mod reduce_sum_of_d2 { use crate::emulate::emulate_mm256_reduce_add_ps; assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; - let mut n = lhs.len() as u32; + let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut d2 = _mm256_setzero_ps(); @@ -788,7 +758,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index 6bd0bfa1..e04fa74b 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -1535,7 +1535,7 @@ mod reduce_sum_of_d2 { fn reduce_sum_of_d2_v4(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; - let mut n = lhs.len() as u32; + let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut d2 = _mm512_setzero_ps(); @@ -1549,7 +1549,7 @@ mod reduce_sum_of_d2 { d2 = _mm512_fmadd_ps(d, d, d2); } if n > 0 { - let mask = _bzhi_u32(0xffff, n) as u16; + let mask = _bzhi_u32(0xffff, n as u32) as u16; let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; let y = unsafe { _mm512_maskz_loadu_ps(mask, b) }; let d = _mm512_sub_ps(x, y); diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 7a448ea3..bdb58b03 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -1,6 +1,5 @@ #![cfg_attr(target_arch = "x86_64", feature(avx512_target_feature))] #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] -#![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512_f16))] #![allow(unsafe_code)] mod aligned; @@ -69,6 +68,20 @@ mod internal { #[allow(unused_imports)] pub use is_riscv64_cpu_detected; + #[cfg(target_arch = "x86_64")] + pub fn is_v4fp16_detected() -> bool { + std::arch::is_x86_feature_detected!("avx512fp16") + && std::arch::is_x86_feature_detected!("avx512bw") + && std::arch::is_x86_feature_detected!("avx512cd") + && std::arch::is_x86_feature_detected!("avx512dq") + && std::arch::is_x86_feature_detected!("avx512vl") + && std::arch::is_x86_feature_detected!("bmi1") + && std::arch::is_x86_feature_detected!("bmi2") + && std::arch::is_x86_feature_detected!("lzcnt") + && std::arch::is_x86_feature_detected!("movbe") + && std::arch::is_x86_feature_detected!("popcnt") + } + #[cfg(target_arch = "x86_64")] pub fn is_v4_detected() -> bool { std::arch::is_x86_feature_detected!("avx512bw") diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs index 0c136a89..715ac870 100644 --- a/crates/simd_macros/src/target.rs +++ b/crates/simd_macros/src/target.rs @@ -5,6 +5,15 @@ pub struct TargetCpu { } pub const TARGET_CPUS: &[TargetCpu] = &[ + // This is a temporary hack, for using AVX512FP16 without unstable features. + TargetCpu { + target_cpu: "v4fp16", + target_arch: "x86_64", + target_features: &[ + "avx512bw", "avx512cd", "avx512dq", "avx512vl", // simd + "bmi1", "bmi2", "lzcnt", "movbe", "popcnt", // bit-operations + ], + }, TargetCpu { target_cpu: "v4", target_arch: "x86_64", diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 81febe96..ed9009cd 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2025-03-19" +channel = "nightly-2025-04-25" diff --git a/src/datatype/text_scalar8.rs b/src/datatype/text_scalar8.rs index 8669a015..d589ffdc 100644 --- a/src/datatype/text_scalar8.rs +++ b/src/datatype/text_scalar8.rs @@ -128,10 +128,10 @@ fn _vchord_scalar8_out(vector: Scalar8Input<'_>) -> CString { buffer.push(')'); buffer.push('['); if let Some(&x) = vector.code().first() { - buffer.push_str(format!("{}", x).as_str()); + buffer.push_str(format!("{x}").as_str()); } for &x in vector.code().iter().skip(1) { - buffer.push_str(format!(", {}", x).as_str()); + buffer.push_str(format!(", {x}").as_str()); } buffer.push(']'); CString::new(buffer).unwrap() diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index 8e335467..4ce686c5 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -64,7 +64,7 @@ fn _vchord_typmod_in_65535(list: pgrx::datum::Array<&CStr>) -> i32 { fn _vchord_typmod_out(typmod: i32) -> CString { let typmod = Typmod::parse_from_i32(typmod).unwrap(); match typmod.into_option_string() { - Some(s) => CString::new(format!("({})", s)).unwrap(), + Some(s) => CString::new(format!("({s})")).unwrap(), None => CString::new("()").unwrap(), } } diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index 8c744852..c4a8f0c6 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -40,7 +40,7 @@ impl Allocator { } else { #[cold] fn cold(sel: &mut Allocator) -> *mut T { - std::panic::abort_unwind(|| { + abort_unwind(|| { let raw = std::mem::replace(&mut sel.this, std::ptr::null_mut()); sel.used.push(raw.cast()); sel.this = sel @@ -73,7 +73,7 @@ impl Allocator { } else { #[cold] fn cold(sel: &mut Allocator, n: usize) -> *mut T { - std::panic::abort_unwind(|| { + abort_unwind(|| { let raw = std::mem::replace(&mut sel.this, std::ptr::null_mut()); sel.used.push(raw.cast()); sel.this = sel @@ -93,7 +93,7 @@ impl Allocator { } } pub fn reset(&mut self) { - std::panic::abort_unwind(|| { + abort_unwind(|| { self.free.extend(std::mem::take(&mut self.used)); self.size = size_of::(); }); @@ -375,3 +375,10 @@ impl RandomProject for VectBorrowed<'_, f16> { VectOwned::new(f16::vector_from_f32(&project(&input))) } } + +// Emulate unstable library feature `abort_unwind`. +// See https://github.com/rust-lang/rust/issues/130338. + +extern "C" fn abort_unwind R, R>(f: F) -> R { + f() +} diff --git a/src/lib.rs b/src/lib.rs index 49b14676..9badf969 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,4 @@ #![allow(unsafe_code)] -#![feature(let_chains)] -#![feature(abort_unwind)] mod datatype; mod index; From 175b57d7cf6e5f49813be7cf8763deed76511e13 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 6 May 2025 10:25:41 +0800 Subject: [PATCH 153/324] feat: add vectorized version of `u8::reduce_sum_of_xy` (#242) Signed-off-by: usamoi --- crates/simd/src/f16.rs | 6 +- .../src/{fast_scan/mod.rs => fast_scan.rs} | 0 crates/simd/src/u8.rs | 301 +++++++++++++++++- 3 files changed, 304 insertions(+), 3 deletions(-) rename crates/simd/src/{fast_scan/mod.rs => fast_scan.rs} (100%) diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 53962f8a..e17d73dc 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -342,7 +342,8 @@ mod reduce_sum_of_xy { xy = _mm256_fmadd_ps(x, y, xy); } let mut xy = emulate_mm256_reduce_add_ps(xy); - while n > 0 { + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { let x = unsafe { a.read().as_f32() }; let y = unsafe { b.read().as_f32() }; a = unsafe { a.add(1) }; @@ -620,7 +621,8 @@ mod reduce_sum_of_d2 { d2 = _mm256_fmadd_ps(d, d, d2); } let mut d2 = emulate_mm256_reduce_add_ps(d2); - while n > 0 { + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { let x = unsafe { a.read().as_f32() }; let y = unsafe { b.read().as_f32() }; a = unsafe { a.add(1) }; diff --git a/crates/simd/src/fast_scan/mod.rs b/crates/simd/src/fast_scan.rs similarity index 100% rename from crates/simd/src/fast_scan/mod.rs rename to crates/simd/src/fast_scan.rs diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index c2a6b8e5..7f5026cf 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -1,5 +1,304 @@ +#![allow(clippy::just_underscores_and_digits)] mod reduce_sum_of_xy { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm512_set1_epi16(0x00ff_u16 as i16); + let hi = _mm512_set1_epi16(0xff00_u16 as i16); + let mut _0 = _mm512_setzero_si512(); + let mut _1 = _mm512_setzero_si512(); + let mut _2 = _mm512_setzero_si512(); + let mut _3 = _mm512_setzero_si512(); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + let y = unsafe { _mm512_loadu_epi8(b.cast()) }; + let x_l = _mm512_and_si512(x, lo); + let x_h = _mm512_and_si512(_mm512_srli_epi16(_mm512_and_si512(x, hi), 8), lo); + let y_l = _mm512_and_si512(y, lo); + let y_h = _mm512_and_si512(_mm512_srli_epi16(_mm512_and_si512(y, hi), 8), lo); + let l = _mm512_mullo_epi16(x_l, y_l); + let h = _mm512_mullo_epi16(x_h, y_h); + a = unsafe { a.add(64) }; + b = unsafe { b.add(64) }; + n -= 64; + _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(l, 0))); + _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(l, 1))); + _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(h, 0))); + _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(h, 1))); + } + if n > 0 { + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; + let x_l = _mm512_and_si512(x, lo); + let x_h = _mm512_and_si512(_mm512_srli_epi16(_mm512_and_si512(x, hi), 8), lo); + let y_l = _mm512_and_si512(y, lo); + let y_h = _mm512_and_si512(_mm512_srli_epi16(_mm512_and_si512(y, hi), 8), lo); + let l = _mm512_mullo_epi16(x_l, y_l); + let h = _mm512_mullo_epi16(x_h, y_h); + _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(l, 0))); + _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(l, 1))); + _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(h, 0))); + _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(h, 1))); + } + let _5 = _mm512_add_epi32(_0, _1); + let _6 = _mm512_add_epi32(_2, _3); + _mm512_reduce_add_epi32(_mm512_add_epi32(_5, _6)) as u32 + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_v4_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::emulate_mm256_reduce_add_epi32; + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm256_set1_epi16(0x00ff_u16 as i16); + let hi = _mm256_set1_epi16(0xff00_u16 as i16); + let mut _0 = _mm256_setzero_si256(); + let mut _1 = _mm256_setzero_si256(); + let mut _2 = _mm256_setzero_si256(); + let mut _3 = _mm256_setzero_si256(); + while n >= 32 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let x_l = _mm256_and_si256(x, lo); + let x_h = _mm256_and_si256(_mm256_srli_epi16(_mm256_and_si256(x, hi), 8), lo); + let y_l = _mm256_and_si256(y, lo); + let y_h = _mm256_and_si256(_mm256_srli_epi16(_mm256_and_si256(y, hi), 8), lo); + let l = _mm256_mullo_epi16(x_l, y_l); + let h = _mm256_mullo_epi16(x_h, y_h); + a = unsafe { a.add(32) }; + b = unsafe { b.add(32) }; + n -= 32; + _0 = _mm256_add_epi32(_0, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(l, 0))); + _1 = _mm256_add_epi32(_1, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(l, 1))); + _2 = _mm256_add_epi32(_2, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(h, 0))); + _3 = _mm256_add_epi32(_3, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(h, 1))); + } + let mut sum = emulate_mm256_reduce_add_epi32(_mm256_add_epi32( + _mm256_add_epi32(_0, _1), + _mm256_add_epi32(_2, _3), + )) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + sum += x as u32 * y as u32; + } + sum + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v3_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::emulate_mm_reduce_add_epi32; + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm_set1_epi16(0x00ff_u16 as i16); + let hi = _mm_set1_epi16(0xff00_u16 as i16); + let mut _0 = _mm_setzero_si128(); + let mut _1 = _mm_setzero_si128(); + let mut _2 = _mm_setzero_si128(); + let mut _3 = _mm_setzero_si128(); + while n >= 16 { + let x = unsafe { _mm_loadu_si128(a.cast()) }; + let y = unsafe { _mm_loadu_si128(b.cast()) }; + let x_l = _mm_and_si128(x, lo); + let x_h = _mm_and_si128(_mm_srli_epi16(_mm_and_si128(x, hi), 8), lo); + let y_l = _mm_and_si128(y, lo); + let y_h = _mm_and_si128(_mm_srli_epi16(_mm_and_si128(y, hi), 8), lo); + let l = _mm_mullo_epi16(x_l, y_l); + let h = _mm_mullo_epi16(x_h, y_h); + a = unsafe { a.add(16) }; + b = unsafe { b.add(16) }; + n -= 16; + _0 = _mm_add_epi32(_0, _mm_cvtepu16_epi32(l)); + _1 = _mm_add_epi32(_1, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(l, l))); + _2 = _mm_add_epi32(_2, _mm_cvtepu16_epi32(h)); + _3 = _mm_add_epi32(_3, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(h, h))); + } + let mut sum = emulate_mm_reduce_add_epi32(_mm_add_epi32( + _mm_add_epi32(_0, _1), + _mm_add_epi32(_2, _3), + )) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + sum += x as u32 * y as u32; + } + sum + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v2_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_xy_a2(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::aarch64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = vdupq_n_u16(0x00ff_u16); + let mut _0 = vdupq_n_u32(0); + let mut _1 = vdupq_n_u32(0); + let mut _2 = vdupq_n_u32(0); + let mut _3 = vdupq_n_u32(0); + while n >= 16 { + let x = vreinterpretq_u16_u8(unsafe { vld1q_u8(a.cast()) }); + let y = vreinterpretq_u16_u8(unsafe { vld1q_u8(b.cast()) }); + let x_l = vandq_u16(x, lo); + let x_h = vshrq_n_u16(x, 8); + let y_l = vandq_u16(y, lo); + let y_h = vshrq_n_u16(y, 8); + let l = vmulq_u16(x_l, y_l); + let h = vmulq_u16(x_h, y_h); + a = unsafe { a.add(16) }; + b = unsafe { b.add(16) }; + n -= 16; + _0 = vaddq_u32(_0, vmovl_u16(vget_low_u16(l))); + _1 = vaddq_u32(_1, vmovl_u16(vget_high_u16(l))); + _2 = vaddq_u32(_2, vmovl_u16(vget_low_u16(h))); + _3 = vaddq_u32(_3, vmovl_u16(vget_high_u16(h))); + } + let mut sum = vaddvq_u32(vaddq_u32(vaddq_u32(_0, _1), vaddq_u32(_2, _3))); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + sum += x as u32 * y as u32; + } + sum + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_xy_a2_test() { + use rand::Rng; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); From b7ad2a173ec22427658d9fded527721af2ca5753 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 9 May 2025 16:17:38 +0800 Subject: [PATCH 154/324] feat: prefetch in searching (#244) Adds a GUC `io_search`, with the same possible values of `io_rerank`. notes: * prefetching for the bottom-level bit vectors is implemented. * prefetching for the other levels of bit vectors is not implemented yet. Signed-off-by: usamoi --- Cargo.lock | 21 ++ Cargo.toml | 1 + crates/algorithm/Cargo.toml | 2 + crates/algorithm/src/build.rs | 126 +--------- crates/algorithm/src/bulkdelete.rs | 20 +- crates/algorithm/src/cache.rs | 13 +- crates/algorithm/src/cost.rs | 2 +- crates/algorithm/src/fast_heap.rs | 34 +-- crates/algorithm/src/freepages.rs | 4 +- crates/algorithm/src/insert.rs | 28 +-- crates/algorithm/src/lib.rs | 112 ++++++--- crates/algorithm/src/maintain.rs | 183 ++++++-------- crates/algorithm/src/operator.rs | 4 +- crates/algorithm/src/prefetcher.rs | 177 ++++++++----- crates/algorithm/src/prewarm.rs | 38 +-- crates/algorithm/src/rerank.rs | 16 +- crates/algorithm/src/search.rs | 74 +++--- crates/algorithm/src/tape.rs | 162 +++++++++--- crates/algorithm/src/tape_writer.rs | 205 +++++++++++++++ crates/algorithm/src/tuples.rs | 155 +++++++++++- crates/algorithm/src/vectors.rs | 6 +- crates/always_equal/Cargo.toml | 1 + crates/distance/Cargo.toml | 1 + crates/k_means/Cargo.toml | 1 + crates/k_means/src/lib.rs | 2 +- crates/rabitq/Cargo.toml | 1 + crates/rabitq/src/binary.rs | 10 +- crates/rabitq/src/block.rs | 6 +- crates/rabitq/src/lib.rs | 3 +- crates/rabitq/src/packing.rs | 88 +++++++ crates/random_orthogonal_matrix/Cargo.toml | 1 + crates/simd/Cargo.toml | 1 + crates/simd/src/bit.rs | 70 +++--- crates/simd/src/fast_scan.rs | 91 +------ crates/simd/src/lib.rs | 1 - crates/simd/src/packed_u4.rs | 18 -- crates/simd/src/u8.rs | 64 ++--- crates/simd_macros/Cargo.toml | 1 + crates/vector/Cargo.toml | 1 + crates/vector/src/bvect.rs | 8 +- crates/vector/src/scalar8.rs | 29 ++- src/datatype/functions_scalar8.rs | 4 +- src/index/algorithm.rs | 222 ++++++++++++++--- src/index/am/am_build.rs | 17 +- src/index/am/mod.rs | 9 +- src/index/functions.rs | 4 +- src/index/gucs.rs | 42 +++- src/index/scanners/default.rs | 274 ++++++++++++++------- src/index/scanners/maxsim.rs | 162 ++++++++---- src/index/scanners/mod.rs | 21 +- src/index/storage.rs | 86 +++++-- 51 files changed, 1694 insertions(+), 928 deletions(-) create mode 100644 crates/algorithm/src/tape_writer.rs create mode 100644 crates/rabitq/src/packing.rs delete mode 100644 crates/simd/src/packed_u4.rs diff --git a/Cargo.lock b/Cargo.lock index 75da4911..863bd9fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -20,6 +20,7 @@ dependencies = [ "half 2.6.0", "k_means", "paste", + "pin-project", "rabitq", "rand", "random_orthogonal_matrix", @@ -960,6 +961,26 @@ dependencies = [ "unescape", ] +[[package]] +name = "pin-project" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677f1add503faace112b9f1373e43e9e054bfdd22ff1a63c1bc485eaec6a6a8a" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e918e4ff8c4549eb882f14b3a4bc8c8bc93de829416eacf579f1207a8fbf861" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "ppv-lite86" version = "0.2.21" diff --git a/Cargo.toml b/Cargo.toml index 7a2a7bad..b9ac4dd8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,6 +2,7 @@ name = "vchord" version.workspace = true edition.workspace = true +publish = false [lib] name = "vchord" diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml index 13fae996..73c7b6c5 100644 --- a/crates/algorithm/Cargo.toml +++ b/crates/algorithm/Cargo.toml @@ -2,6 +2,7 @@ name = "algorithm" version.workspace = true edition.workspace = true +publish = false [dependencies] always_equal = { path = "../always_equal" } @@ -14,6 +15,7 @@ vector = { path = "../vector" } half.workspace = true paste.workspace = true +pin-project = "1" rand.workspace = true serde.workspace = true turboselect = { git = "https://github.com/tensorchord/turboselect.git", rev = "d8753c4ffe5b47f28670fea21d56cf3658d51b9b" } diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 85772cf5..954e4cd7 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -1,15 +1,15 @@ use crate::operator::{Accessor2, Operator, Vector}; use crate::tape::TapeWriter; +use crate::tape_writer::H1TapeWriter; use crate::tuples::*; use crate::types::*; use crate::{Branch, RelationWrite}; -use simd::fast_scan::{any_pack, padding_pack}; use vector::VectorOwned; -pub fn build( +pub fn build( vector_options: VectorOptions, vchordrq_options: VchordrqIndexOptions, - index: impl RelationWrite, + index: &R, structures: Vec>, ) { if vchordrq_options.residual_quantization && !O::SUPPORTS_RESIDUAL { @@ -17,10 +17,10 @@ pub fn build( } let dims = vector_options.dims; let is_residual = vchordrq_options.residual_quantization; - let mut meta = TapeWriter::<_, MetaTuple>::create(&index, false); + let mut meta = TapeWriter::<_, MetaTuple>::create(index, false); assert_eq!(meta.first(), 0); - let freepage = TapeWriter::<_, FreepageTuple>::create(&index, false); - let mut vectors = TapeWriter::<_, VectorTuple>::create(&index, true); + let freepage = TapeWriter::<_, FreepageTuple>::create(index, false); + let mut vectors = TapeWriter::<_, VectorTuple>::create(index, true); let mut pointer_of_means = Vec::, u16)>>::new(); for i in 0..structures.len() { let mut level = Vec::new(); @@ -58,17 +58,17 @@ pub fn build( let mut level = Vec::new(); for j in 0..structures[i].len() { if i == 0 { - let frozen_tape = TapeWriter::<_, FrozenTuple>::create(&index, false); - let appendable_tape = TapeWriter::<_, AppendableTuple>::create(&index, false); - let mut jump = TapeWriter::<_, JumpTuple>::create(&index, false); + let directory_tape = TapeWriter::<_, DirectoryTuple>::create(index, false); + let appendable_tape = TapeWriter::<_, AppendableTuple>::create(index, false); + let mut jump = TapeWriter::<_, JumpTuple>::create(index, false); jump.push(JumpTuple { - frozen_first: frozen_tape.first(), + directory_first: directory_tape.first(), appendable_first: appendable_tape.first(), tuples: 0, }); level.push(jump.first()); } else { - let mut tape = H1TapeWriter::create(&index, O::Vector::count(dims as _), false); + let mut tape = H1TapeWriter::create(index, O::Vector::count(dims as _), false); let h2_mean = structures[i].means[j].as_borrowed(); let h2_children = structures[i].children[j].as_slice(); for child in h2_children.iter().copied() { @@ -95,38 +95,7 @@ pub fn build( }); } let (mut tape, chunk) = tape.into_inner(); - if !chunk.is_empty() { - let mut remain = - padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); - loop { - let freespace = tape.freespace(); - if H1Tuple::estimate_size_0(O::Vector::count(dims as _), remain.len()) - <= freespace as usize - { - tape.tape_put(H1Tuple::_0 { - head: any_pack(chunk.iter().map(|x| x.head)), - dis_u_2: any_pack(chunk.iter().map(|x| x.dis_u_2)), - factor_ppc: any_pack(chunk.iter().map(|x| x.factor_ppc)), - factor_ip: any_pack(chunk.iter().map(|x| x.factor_ip)), - factor_err: any_pack(chunk.iter().map(|x| x.factor_err)), - first: any_pack(chunk.iter().map(|x| x.extra)), - prefetch: fix(chunk.iter().map(|x| x.prefetch.as_slice())), - len: chunk.len() as _, - elements: remain, - }); - break; - } - if let Some(w) = H1Tuple::fit_1(O::Vector::count(dims as _), freespace) { - let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); - tape.tape_put(H1Tuple::_1 { - elements: left.to_vec(), - }); - remain = right.to_vec(); - } else { - tape.tape_move(); - } - } - } + H1TapeWriter::flush(&mut tape, O::Vector::count(dims as _), chunk); level.push(tape.first()); } } @@ -154,74 +123,3 @@ pub fn build( cells: structures.iter().map(|s| s.len() as _).collect(), }); } - -pub struct H1TapeWriter<'a, R> -where - R: RelationWrite + 'a, -{ - tape: TapeWriter<'a, R, H1Tuple>, - branches: Vec>, - prefetch: usize, -} - -impl<'a, R> H1TapeWriter<'a, R> -where - R: RelationWrite + 'a, -{ - fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { - Self { - tape: TapeWriter::create(index, tracking_freespace), - branches: Vec::new(), - prefetch, - } - } - fn push(&mut self, branch: Branch) { - self.branches.push(branch); - if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { - let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); - loop { - let freespace = self.tape.freespace(); - if H1Tuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { - self.tape.tape_put(H1Tuple::_0 { - head: chunk.each_ref().map(|x| x.head), - dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), - factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), - factor_ip: chunk.each_ref().map(|x| x.factor_ip), - factor_err: chunk.each_ref().map(|x| x.factor_err), - first: chunk.each_ref().map(|x| x.extra), - prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), - len: chunk.len() as _, - elements: remain, - }); - break; - } - if let Some(w) = H1Tuple::fit_1(self.prefetch, freespace) { - let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); - self.tape.tape_put(H1Tuple::_1 { - elements: left.to_vec(), - }); - remain = right.to_vec(); - } else { - self.tape.tape_move(); - } - } - self.branches.clear(); - } - } - fn into_inner(self) -> (TapeWriter<'a, R, H1Tuple>, Vec>) { - (self.tape, self.branches) - } -} - -fn fix<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { - use std::array::from_fn; - let mut iter = into_iter.into_iter(); - let mut array: [_; 32] = from_fn(|_| iter.next().map(<[u32]>::to_vec).unwrap_or_default()); - if iter.next().is_some() { - panic!("too many slices"); - } - let step = array.iter().map(Vec::len).max().unwrap_or_default(); - array.iter_mut().for_each(|x| x.resize(step, u32::MAX)); - let flat = array.into_iter().flatten().collect::>(); - (0..step).map(|i| from_fn(|j| flat[i * 32 + j])).collect() -} diff --git a/crates/algorithm/src/bulkdelete.rs b/crates/algorithm/src/bulkdelete.rs index 61aa0e41..4fc1a87d 100644 --- a/crates/algorithm/src/bulkdelete.rs +++ b/crates/algorithm/src/bulkdelete.rs @@ -1,11 +1,12 @@ use crate::closure_lifetime_binder::{id_0, id_1}; use crate::operator::{FunctionalAccessor, Operator}; +use crate::tape::by_next; use crate::tuples::*; use crate::{Page, RelationRead, RelationWrite, tape}; use std::num::NonZero; -pub fn bulkdelete( - index: impl RelationRead + RelationWrite, +pub fn bulkdelete( + index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, ) { @@ -22,12 +23,10 @@ pub fn bulkdelete( let step = |state: State| { let mut results = Vec::new(); for first in state { - tape::read_h1_tape( - index.clone(), - first, + tape::read_h1_tape::( + by_next(index, first).inspect(|_| check()), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), |(), _, first, _| results.push(first), - |_| check(), ); } results @@ -39,8 +38,11 @@ pub fn bulkdelete( let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + let mut directory = tape::read_directory_tape::( + by_next(index, jump_tuple.directory_first()).inspect(|_| check()), + ); { - let mut current = jump_tuple.frozen_first(); + let mut current = directory.next().unwrap_or(u32::MAX); while current != u32::MAX { check(); let read = index.read(current); @@ -72,10 +74,8 @@ pub fn bulkdelete( } } } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; } + current = directory.next().unwrap_or(u32::MAX); } } { diff --git a/crates/algorithm/src/cache.rs b/crates/algorithm/src/cache.rs index 839948c8..00e938b2 100644 --- a/crates/algorithm/src/cache.rs +++ b/crates/algorithm/src/cache.rs @@ -1,9 +1,10 @@ use crate::closure_lifetime_binder::{id_0, id_1}; use crate::operator::FunctionalAccessor; +use crate::tape::by_next; use crate::tuples::{MetaTuple, WithReader}; -use crate::{Page, RelationRead, tape}; +use crate::{Page, PageGuard, RelationRead, tape}; -pub fn cache(index: impl RelationRead) -> Vec { +pub fn cache(index: &R) -> Vec { let mut trace = vec![0_u32]; let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); @@ -16,16 +17,12 @@ pub fn cache(index: impl RelationRead) -> Vec { let mut step = |state: State| { let mut results = Vec::new(); for first in state { - tape::read_h1_tape( - index.clone(), - first, + tape::read_h1_tape::( + by_next(index, first).inspect(|guard| trace.push(guard.id())), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), |(), _, first, _| { results.push(first); }, - |id| { - trace.push(id); - }, ); } results diff --git a/crates/algorithm/src/cost.rs b/crates/algorithm/src/cost.rs index 9f6c2f53..32274ae6 100644 --- a/crates/algorithm/src/cost.rs +++ b/crates/algorithm/src/cost.rs @@ -8,7 +8,7 @@ pub struct Cost { } #[must_use] -pub fn cost(index: impl RelationRead) -> Cost { +pub fn cost(index: &R) -> Cost { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); diff --git a/crates/algorithm/src/fast_heap.rs b/crates/algorithm/src/fast_heap.rs index a4e88f1e..3bd85ac0 100644 --- a/crates/algorithm/src/fast_heap.rs +++ b/crates/algorithm/src/fast_heap.rs @@ -1,4 +1,4 @@ -use crate::Heap; +use crate::Sequence; use std::collections::BinaryHeap; use std::num::NonZero; @@ -53,32 +53,32 @@ impl FastHeap { } } -impl IntoIterator for FastHeap { - type Item = T; - - type IntoIter = std::vec::IntoIter; - - fn into_iter(self) -> Self::IntoIter { - match self { - FastHeap::Sorted(sort_heap) => sort_heap.inner.into_iter(), - FastHeap::Binary(binary_heap) => binary_heap.into_vec().into_iter(), - } +impl From> for FastHeap { + fn from(value: Vec) -> Self { + Self::from_vec(value) } } -impl Heap for FastHeap { - fn make(this: Vec) -> Self { - Self::from_vec(this) +impl Sequence for FastHeap { + type Item = T; + type Inner = std::vec::IntoIter; + fn next(&mut self) -> Option { + self.pop() } - - fn pop_if(&mut self, predicate: impl FnOnce(&Self::Item) -> bool) -> Option { + fn next_if(&mut self, predicate: impl FnOnce(&T) -> bool) -> Option { let first = self.peek()?; if predicate(first) { self.pop() } else { None } } + fn into_inner(self) -> Self::Inner { + match self { + FastHeap::Sorted(sort_heap) => sort_heap.inner.into_iter(), + FastHeap::Binary(binary_heap) => binary_heap.into_vec().into_iter(), + } + } } #[test] -fn test_select_heap() { +fn test_fast_heap() { for _ in 0..1000 { let sequence = (0..10000) .map(|_| rand::random::()) diff --git a/crates/algorithm/src/freepages.rs b/crates/algorithm/src/freepages.rs index 976ae903..55574fdf 100644 --- a/crates/algorithm/src/freepages.rs +++ b/crates/algorithm/src/freepages.rs @@ -2,7 +2,7 @@ use crate::tuples::*; use crate::*; use std::cmp::Reverse; -pub fn mark(index: impl RelationWrite, freepage_first: u32, values: &[u32]) { +pub fn mark(index: &impl RelationWrite, freepage_first: u32, values: &[u32]) { let mut values = { let mut values = values.to_vec(); values.sort_by_key(|x| Reverse(*x)); @@ -33,7 +33,7 @@ pub fn mark(index: impl RelationWrite, freepage_first: u32, values: &[u32]) { } } -pub fn fetch(index: impl RelationWrite, freepage_first: u32) -> Option { +pub fn fetch(index: &impl RelationWrite, freepage_first: u32) -> Option { let (mut current, mut offset) = (freepage_first, 0_u32); loop { let mut freespace_guard = index.write(current, false); diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 38b066e6..52fdb3ff 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -1,8 +1,9 @@ use crate::linked_vec::LinkedVec; use crate::operator::*; +use crate::tape::by_next; use crate::tuples::*; use crate::vectors::{self}; -use crate::{Bump, Page, Prefetcher, RelationRead, RelationWrite, tape}; +use crate::{Bump, Page, Prefetcher, PrefetcherHeapFamily, RelationRead, RelationWrite, tape}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -15,17 +16,12 @@ type Item<'b> = ( AlwaysEqual<&'b mut (u32, u16, &'b mut [u32])>, ); -pub fn insert< - 'b, - R: RelationRead + RelationWrite, - O: Operator, - P: Prefetcher>, ->( - index: R, +pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( + index: &'r R, payload: NonZero, vector: O::Vector, bump: &'b impl Bump, - mut prefetch: impl FnMut(Vec>) -> P, + mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, ) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); @@ -48,7 +44,7 @@ pub fn insert< }; let (list, head) = if !rerank_in_heap { - vectors::append::(index.clone(), vectors_first, vector.as_borrowed(), payload) + vectors::append::(index, vectors_first, vector.as_borrowed(), payload) } else { (Vec::new(), 0) }; @@ -81,9 +77,8 @@ pub fn insert< } else { unreachable!() }; - tape::read_h1_tape( - index.clone(), - first, + tape::read_h1_tape::( + by_next(index, first), || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), |(rough, err), head, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * 1.9); @@ -92,14 +87,13 @@ pub fn insert< AlwaysEqual(bump.alloc((first, head, bump.alloc_slice(prefetch)))), )); }, - |_| (), ); } - let mut heap = (prefetch)(results.into_vec()); + let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); { while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, ..))), list)) = - heap.pop_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) + heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { if is_residual { let (distance, residual) = vectors::read_for_h1_tuple::( @@ -161,5 +155,5 @@ pub fn insert< let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - tape::append(index.clone(), jump_tuple.appendable_first(), &bytes, false); + tape::append(index, jump_tuple.appendable_first(), &bytes, false); } diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index a167b634..185241d9 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -15,6 +15,7 @@ mod prewarm; mod rerank; mod search; mod tape; +mod tape_writer; mod tuples; mod vectors; @@ -29,12 +30,13 @@ pub use cost::cost; pub use fast_heap::FastHeap; pub use insert::insert; pub use maintain::maintain; -pub use prefetcher::{PlainPrefetcher, Prefetcher, SimplePrefetcher, StreamPrefetcher}; +pub use prefetcher::*; pub use prewarm::prewarm; pub use rerank::{how, rerank_heap, rerank_index}; pub use search::{default_search, maxsim_search}; use std::collections::BinaryHeap; +use std::iter::Peekable; use std::ops::{Deref, DerefMut}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -72,27 +74,41 @@ pub trait PageGuard { fn id(&self) -> u32; } -pub trait ReadStream { +pub trait ReadStream<'s> { type Relation: RelationRead; - type Inner: Iterator; + type Item; + type Inner: Iterator; + fn next( + &mut self, + ) -> Option<( + Self::Item, + Vec<::ReadGuard<'s>>, + )>; fn next_if( &mut self, - predicate: impl FnOnce(&T) -> bool, - ) -> Option<(T, Vec<::ReadGuard<'_>>)>; + predicate: impl FnOnce(&Self::Item) -> bool, + ) -> Option<( + Self::Item, + Vec<::ReadGuard<'s>>, + )>; fn into_inner(self) -> Self::Inner; } -pub trait WriteStream { +pub trait WriteStream<'s> { type Relation: RelationWrite; - type Inner: Iterator; + type Item; + type Inner: Iterator; fn next_if( &mut self, - predicate: impl FnOnce(&T) -> bool, - ) -> Option<(T, Vec<::WriteGuard<'_>>)>; + predicate: impl FnOnce(&Self::Item) -> bool, + ) -> Option<( + Self::Item, + Vec<::WriteGuard<'s>>, + )>; fn into_inner(self) -> Self::Inner; } -pub trait Relation: Clone { +pub trait Relation { type Page: Page; } @@ -116,22 +132,34 @@ pub trait RelationPrefetch: Relation { fn prefetch(&self, id: u32); } +#[derive(Debug, Default, Clone)] +pub struct Hints { + pub full: bool, +} + +impl Hints { + #[allow(clippy::needless_update)] + pub fn full(self, full: bool) -> Self { + Self { full, ..self } + } +} + pub trait RelationReadStream: RelationRead { - type ReadStream<'s, I: Iterator>: ReadStream + type ReadStream<'s, I: Iterator>: ReadStream<'s, Item = I::Item, Relation = Self> where I::Item: Fetch, Self: 's; - fn read_stream(&self, iter: I) -> Self::ReadStream<'_, I> + fn read_stream(&self, iter: I, hints: Hints) -> Self::ReadStream<'_, I> where I::Item: Fetch; } pub trait RelationWriteStream: RelationWrite { - type WriteStream<'s, I: Iterator>: WriteStream + type WriteStream<'s, I: Iterator>: WriteStream<'s, Item = I::Item, Relation = Self> where I::Item: Fetch, Self: 's; - fn write_stream(&self, iter: I) -> Self::WriteStream<'_, I> + fn write_stream(&self, iter: I, hints: Hints) -> Self::WriteStream<'_, I> where I::Item: Fetch; } @@ -153,10 +181,23 @@ pub(crate) struct Branch { pub extra: T, } +pub trait Bump: 'static { + #[allow(clippy::mut_from_ref)] + fn alloc(&self, value: T) -> &mut T; + #[allow(clippy::mut_from_ref)] + fn alloc_slice(&self, slice: &[T]) -> &mut [T]; +} + pub trait Fetch { fn fetch(&self) -> &[u32]; } +impl Fetch for u32 { + fn fetch(&self) -> &[u32] { + std::slice::from_ref(self) + } +} + impl<'b, T, A, B> Fetch for (T, AlwaysEqual<&'b mut (A, B, &'b mut [u32])>) { fn fetch(&self) -> &[u32] { let (_, AlwaysEqual((.., list))) = self; @@ -164,24 +205,39 @@ impl<'b, T, A, B> Fetch for (T, AlwaysEqual<&'b mut (A, B, &'b mut [u32])>) { } } -pub trait Bump: 'static { - #[allow(clippy::mut_from_ref)] - fn alloc(&self, value: T) -> &mut T; - #[allow(clippy::mut_from_ref)] - fn alloc_slice(&self, slice: &[T]) -> &mut [T]; -} - -pub trait Heap: IntoIterator { - fn make(this: Vec) -> Self; - fn pop_if(&mut self, predicate: impl FnOnce(&Self::Item) -> bool) -> Option; +pub trait Sequence { + type Item; + type Inner: Iterator; + fn next(&mut self) -> Option; + fn next_if(&mut self, predicate: impl FnOnce(&Self::Item) -> bool) -> Option; + fn into_inner(self) -> Self::Inner; } -impl Heap for BinaryHeap { - fn make(this: Vec) -> Self { - Self::from(this) +impl Sequence for BinaryHeap { + type Item = T; + type Inner = std::vec::IntoIter; + fn next(&mut self) -> Option { + self.pop() } - fn pop_if(&mut self, predicate: impl FnOnce(&T) -> bool) -> Option { + fn next_if(&mut self, predicate: impl FnOnce(&T) -> bool) -> Option { let peek = self.peek()?; if predicate(peek) { self.pop() } else { None } } + fn into_inner(self) -> Self::Inner { + self.into_vec().into_iter() + } +} + +impl Sequence for Peekable { + type Item = I::Item; + type Inner = Peekable; + fn next(&mut self) -> Option { + Iterator::next(self) + } + fn next_if(&mut self, predicate: impl FnOnce(&I::Item) -> bool) -> Option { + Peekable::next_if(self, predicate) + } + fn into_inner(self) -> Self::Inner { + self + } } diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index 1627253e..5e4cb073 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -1,12 +1,16 @@ use crate::closure_lifetime_binder::{id_0, id_1, id_2, id_3}; use crate::operator::{FunctionalAccessor, Operator, Vector}; -use crate::tape::{self, TapeWriter}; +use crate::tape::{self, TapeWriter, by_directory, by_next}; +use crate::tape_writer::{DirectoryTapeWriter, FrozenTapeWriter}; use crate::tuples::*; -use crate::{Branch, Page, Relation, RelationRead, RelationWrite, freepages}; -use simd::fast_scan::{padding_pack, unpack}; -use std::num::NonZero; - -pub fn maintain(index: R, check: impl Fn()) { +use crate::*; +use rabitq::packing::unpack; + +pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( + index: &'r R, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, + check: impl Fn(), +) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -22,12 +26,10 @@ pub fn maintain(index: R, check: i let step = |state: State| { let mut results = Vec::new(); for first in state { - tape::read_h1_tape( - index.clone(), - first, + tape::read_h1_tape::( + by_next(index, first).inspect(|_| check()), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), |(), _, first, _| results.push(first), - |_| check(), ); } results @@ -43,11 +45,10 @@ pub fn maintain(index: R, check: i let jump_bytes = jump_guard.get_mut(1).expect("data corruption"); let mut jump_tuple = JumpTuple::deserialize_mut(jump_bytes); - let hooked_index = RelationHooked( - index.clone(), + let frozen_tape_hooked_index = RelationHooked(index, { id_3(move |index: &R, tracking_freespace: bool| { if !tracking_freespace { - if let Some(id) = freepages::fetch(index.clone(), freepage_first) { + if let Some(id) = freepages::fetch(index, freepage_first) { let mut write = index.write(id, false); write.clear(); write @@ -57,12 +58,18 @@ pub fn maintain(index: R, check: i } else { index.extend(true) } - }), - ); + }) + }); - let mut tape = FrozenTapeWriter::create(&hooked_index, O::Vector::count(dims as _), false); + let mut tape = FrozenTapeWriter::create( + &frozen_tape_hooked_index, + O::Vector::count(dims as _), + false, + ); - let mut trace = Vec::new(); + let mut trace_directory = Vec::new(); + let mut trace_forzen = Vec::new(); + let mut trace_appendable = Vec::new(); let mut tuples = 0_u64; let mut callback = id_2(|code: (_, _, _, _, _), head, payload, prefetch: &[_]| { @@ -78,13 +85,15 @@ pub fn maintain(index: R, check: i }); tuples += 1; }); - let mut step = |id| { - check(); - trace.push(id); - }; - tape::read_frozen_tape( - index.clone(), - *jump_tuple.frozen_first(), + let directory = tape::read_directory_tape::( + by_next(index, *jump_tuple.directory_first()) + .inspect(|_| check()) + .inspect(|guard| trace_directory.push(guard.id())), + ); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory) + .inspect(|_| check()) + .inspect(|guard| trace_forzen.push(guard.id())), || { FunctionalAccessor::new( Vec::<[u8; 16]>::new(), @@ -100,11 +109,11 @@ pub fn maintain(index: R, check: i ) }, &mut callback, - &mut step, ); - tape::read_appendable_tape( - index.clone(), - *jump_tuple.appendable_first(), + tape::read_appendable_tape::( + by_next(index, *jump_tuple.appendable_first()) + .inspect(|_| check()) + .inspect(|guard| trace_appendable.push(guard.id())), |code| { let signs = code .4 @@ -115,11 +124,26 @@ pub fn maintain(index: R, check: i (code.0, code.1, code.2, code.3, signs) }, &mut callback, - &mut step, ); let (frozen_tape, branches) = tape.into_inner(); + let hooked_index = RelationHooked( + index, + id_3(move |index: &R, tracking_freespace: bool| { + if !tracking_freespace { + if let Some(id) = freepages::fetch(index, freepage_first) { + let mut write = index.write(id, false); + write.clear(); + write + } else { + index.extend(false) + } + } else { + index.extend(true) + } + }), + ); let mut appendable_tape = TapeWriter::create(&hooked_index, false); for branch in branches { @@ -135,77 +159,39 @@ pub fn maintain(index: R, check: i }); } - *jump_tuple.frozen_first() = { frozen_tape }.first(); + let frozen_first = { frozen_tape }.first(); + + let directory = by_next(index, frozen_first) + .inspect(|_| check()) + .map(|guard| guard.id()) + .collect::>(); + + let mut directory_tape = DirectoryTapeWriter::create(index, false); + directory_tape.push(directory.as_slice()); + let directory_tape = directory_tape.into_inner(); + + *jump_tuple.directory_first() = { directory_tape }.first(); *jump_tuple.appendable_first() = { appendable_tape }.first(); *jump_tuple.tuples() = tuples; drop(jump_guard); - freepages::mark(index.clone(), freepage_first, &trace); - } -} - -struct FrozenTapeWriter<'a, R> -where - R: RelationWrite + 'a, -{ - tape: TapeWriter<'a, R, FrozenTuple>, - branches: Vec>>, - prefetch: usize, -} + let trace = { + let mut v = Vec::new(); + v.extend_from_slice(&trace_directory); + v.extend_from_slice(&trace_forzen); + v.extend_from_slice(&trace_appendable); + v + }; -impl<'a, R> FrozenTapeWriter<'a, R> -where - R: RelationWrite + 'a, -{ - fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { - Self { - tape: TapeWriter::create(index, tracking_freespace), - branches: Vec::new(), - prefetch, - } - } - fn push(&mut self, branch: Branch>) { - self.branches.push(branch); - if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { - let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); - loop { - let freespace = self.tape.freespace(); - if FrozenTuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { - self.tape.tape_put(FrozenTuple::_0 { - head: chunk.each_ref().map(|x| x.head), - dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), - factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), - factor_ip: chunk.each_ref().map(|x| x.factor_ip), - factor_err: chunk.each_ref().map(|x| x.factor_err), - payload: chunk.each_ref().map(|x| Some(x.extra)), - prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), - elements: remain, - }); - break; - } - if let Some(w) = FrozenTuple::fit_1(self.prefetch, freespace) { - let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); - self.tape.tape_put(FrozenTuple::_1 { - elements: left.to_vec(), - }); - remain = right.to_vec(); - } else { - self.tape.tape_move(); - } - } - self.branches.clear(); - } - } - fn into_inner(self) -> (TapeWriter<'a, R, FrozenTuple>, Vec>>) { - (self.tape, self.branches) + freepages::mark(index, freepage_first, trace.as_slice()); } } #[derive(Clone)] -struct RelationHooked(R, E); +struct RelationHooked<'r, R, E>(&'r R, E); -impl Relation for RelationHooked +impl<'r, R, E> Relation for RelationHooked<'r, R, E> where R: Relation, E: Clone, @@ -213,7 +199,7 @@ where type Page = R::Page; } -impl RelationRead for RelationHooked +impl<'r, R, E> RelationRead for RelationHooked<'r, R, E> where R: RelationRead, E: Clone, @@ -228,7 +214,7 @@ where } } -impl RelationWrite for RelationHooked +impl<'r, R, E> RelationWrite for RelationHooked<'r, R, E> where R: RelationWrite, E: Clone + for<'a> Fn(&'a R, bool) -> R::WriteGuard<'a>, @@ -243,23 +229,10 @@ where } fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_> { - (self.1)(&self.0, tracking_freespace) + (self.1)(self.0, tracking_freespace) } fn search(&self, freespace: usize) -> Option> { self.0.search(freespace) } } - -fn fix<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { - use std::array::from_fn; - let mut iter = into_iter.into_iter(); - let mut array: [_; 32] = from_fn(|_| iter.next().map(<[u32]>::to_vec).unwrap_or_default()); - if iter.next().is_some() { - panic!("too many slices"); - } - let step = array.iter().map(Vec::len).max().unwrap_or_default(); - array.iter_mut().for_each(|x| x.resize(step, u32::MAX)); - let flat = array.into_iter().flatten().collect::>(); - (0..step).map(|i| from_fn(|j| flat[i * 32 + j])).collect() -} diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index 9cf3142d..757df604 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -296,7 +296,7 @@ impl type Output = [(f32, f32); 32]; fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { - let t = simd::fast_scan::scan(input, target); + let t = simd::fast_scan::fast_scan(input, target); for i in 0..32 { self.0[i] += t[i]; } @@ -334,7 +334,7 @@ impl type Output = [(f32, f32); 32]; fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { - let t = simd::fast_scan::scan(input, target); + let t = simd::fast_scan::fast_scan(input, target); for i in 0..32 { self.0[i] += t[i]; } diff --git a/crates/algorithm/src/prefetcher.rs b/crates/algorithm/src/prefetcher.rs index 5a517ffe..7e88b423 100644 --- a/crates/algorithm/src/prefetcher.rs +++ b/crates/algorithm/src/prefetcher.rs @@ -1,94 +1,117 @@ -use crate::{Fetch, Heap, ReadStream, RelationPrefetch, RelationRead, RelationReadStream}; -use std::collections::{BinaryHeap, VecDeque, binary_heap, vec_deque}; +use crate::*; +use std::collections::{VecDeque, vec_deque}; use std::iter::Chain; pub const WINDOW_SIZE: usize = 32; const _: () = assert!(WINDOW_SIZE > 0); -pub trait Prefetcher: IntoIterator +pub trait Prefetcher<'r>: IntoIterator where Self::Item: Fetch, { - type R: RelationRead; - fn pop_if( + type R: RelationRead + 'r; + + fn next(&mut self) -> Option<(Self::Item, Vec<::ReadGuard<'r>>)>; + fn next_if( &mut self, predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<(Self::Item, Vec<::ReadGuard<'_>>)>; + ) -> Option<(Self::Item, Vec<::ReadGuard<'r>>)>; } -pub struct PlainPrefetcher { - relation: R, - heap: H, +pub struct PlainPrefetcher<'r, R, S: Sequence> { + relation: &'r R, + sequence: S, } -impl PlainPrefetcher { - pub fn new(relation: R, vec: Vec) -> Self { - Self { - relation, - heap: Heap::make(vec), - } +impl<'r, R, S: Sequence> PlainPrefetcher<'r, R, S> { + pub fn new(relation: &'r R, sequence: S) -> Self { + Self { relation, sequence } } } -impl IntoIterator for PlainPrefetcher { - type Item = H::Item; +impl<'r, R, S: Sequence> IntoIterator for PlainPrefetcher<'r, R, S> { + type Item = S::Item; - type IntoIter = H::IntoIter; + type IntoIter = S::Inner; fn into_iter(self) -> Self::IntoIter { - self.heap.into_iter() + self.sequence.into_inner() } } -impl Prefetcher for PlainPrefetcher +impl<'r, R: RelationRead, S: Sequence> Prefetcher<'r> for PlainPrefetcher<'r, R, S> where - H::Item: Fetch + Ord, + S::Item: Fetch, { type R = R; - fn pop_if<'s>( - &'s mut self, + + fn next(&mut self) -> Option<(Self::Item, Vec<::ReadGuard<'r>>)> { + let e = self.sequence.next()?; + let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); + Some((e, list)) + } + fn next_if( + &mut self, predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<(H::Item, Vec>)> { - let e = self.heap.pop_if(predicate)?; + ) -> Option<(S::Item, Vec>)> { + let e = self.sequence.next_if(predicate)?; let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); Some((e, list)) } } -pub struct SimplePrefetcher { - relation: R, - window: VecDeque, - heap: BinaryHeap, +pub struct SimplePrefetcher<'r, R, S: Sequence> { + relation: &'r R, + window: VecDeque, + sequence: S, } -impl SimplePrefetcher { - pub fn new(relation: R, vec: Vec) -> Self { +impl<'r, R, S: Sequence> SimplePrefetcher<'r, R, S> { + pub fn new(relation: &'r R, sequence: S) -> Self { Self { relation, window: VecDeque::new(), - heap: BinaryHeap::from(vec), + sequence, } } } -impl IntoIterator for SimplePrefetcher { - type Item = T; +impl<'r, R, S: Sequence> IntoIterator for SimplePrefetcher<'r, R, S> { + type Item = S::Item; - type IntoIter = Chain, binary_heap::IntoIter>; + type IntoIter = Chain, S::Inner>; fn into_iter(self) -> Self::IntoIter { - self.window.into_iter().chain(self.heap) + self.window.into_iter().chain(self.sequence.into_inner()) } } -impl Prefetcher for SimplePrefetcher { +impl<'r, R: RelationRead + RelationPrefetch, S: Sequence> Prefetcher<'r> + for SimplePrefetcher<'r, R, S> +where + S::Item: Fetch, +{ type R = R; - fn pop_if<'s>( - &'s mut self, - predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<(T, Vec>)> { + + fn next(&mut self) -> Option<(Self::Item, Vec<::ReadGuard<'r>>)> { + while self.window.len() < WINDOW_SIZE + && let Some(e) = self.sequence.next() + { + for id in e.fetch().iter().copied() { + self.relation.prefetch(id); + } + self.window.push_back(e); + } + let e = self.window.pop_front()?; + let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); + Some((e, list)) + } + fn next_if( + &mut self, + predicate: impl FnOnce(&S::Item) -> bool, + ) -> Option<(S::Item, Vec>)> { while self.window.len() < WINDOW_SIZE - && let Some(e) = self.heap.pop() + && let Some(e) = self.sequence.next() { for id in e.fetch().iter().copied() { self.relation.prefetch(id); @@ -101,53 +124,81 @@ impl Prefetcher for SimplePr } } -pub struct StreamPrefetcherHeap(BinaryHeap); +pub struct StreamPrefetcherSequence(S); -impl Iterator for StreamPrefetcherHeap { - type Item = T; +impl Iterator for StreamPrefetcherSequence { + type Item = S::Item; fn next(&mut self) -> Option { - self.0.pop() + self.0.next() } } -pub struct StreamPrefetcher<'r, R, T> +pub struct StreamPrefetcher<'r, R: RelationReadStream + 'r, S: Sequence> where - R: RelationReadStream + 'r, - T: Fetch + Ord, + S::Item: Fetch, { - stream: R::ReadStream<'r, StreamPrefetcherHeap>, + stream: R::ReadStream<'r, StreamPrefetcherSequence>, } -impl<'r, R: RelationReadStream, T: Fetch + Ord> StreamPrefetcher<'r, R, T> { - pub fn new(relation: &'r R, vec: Vec) -> Self { - let stream = relation.read_stream(StreamPrefetcherHeap(BinaryHeap::from(vec))); +impl<'r, R: RelationReadStream, S: Sequence> StreamPrefetcher<'r, R, S> +where + S::Item: Fetch, +{ + pub fn new(relation: &'r R, sequence: S, hints: Hints) -> Self { + let stream = relation.read_stream(StreamPrefetcherSequence(sequence), hints); Self { stream } } } -impl<'r, R: RelationReadStream, T: Fetch + Ord> IntoIterator for StreamPrefetcher<'r, R, T> { - type Item = T; +impl<'r, R: RelationReadStream, S: Sequence> IntoIterator for StreamPrefetcher<'r, R, S> +where + S::Item: Fetch, +{ + type Item = S::Item; - type IntoIter = <::ReadStream< - 'r, - StreamPrefetcherHeap, - > as ReadStream>::Inner; + type IntoIter = > as ReadStream<'r>>::Inner; fn into_iter(self) -> Self::IntoIter { self.stream.into_inner() } } -impl<'r, R: RelationReadStream, T: Fetch + Ord> Prefetcher for StreamPrefetcher<'r, R, T> { +impl<'r, R: RelationReadStream, S: Sequence> Prefetcher<'r> for StreamPrefetcher<'r, R, S> +where + S::Item: Fetch, +{ type R = R; - fn pop_if<'s>( - &'s mut self, + fn next(&mut self) -> Option<(S::Item, Vec>)> { + self.stream.next() + } + fn next_if( + &mut self, predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<(T, Vec>)> { + ) -> Option<(S::Item, Vec>)> { self.stream.next_if(predicate) } } +pub trait PrefetcherSequenceFamily<'r, R> { + type P: Prefetcher<'r, R = R, Item = S::Item> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch; +} + +pub trait PrefetcherHeapFamily<'r, R> { + type P: Prefetcher<'r, R = R, Item = T> + where + T: Ord + Fetch + 'r; + + fn prefetch(&mut self, seq: Vec) -> Self::P + where + T: Ord + Fetch + 'r; +} + // Emulate unstable library feature `vec_deque_pop_if`. // See https://github.com/rust-lang/rust/issues/135889. diff --git a/crates/algorithm/src/prewarm.rs b/crates/algorithm/src/prewarm.rs index fa4d0138..31d3e46c 100644 --- a/crates/algorithm/src/prewarm.rs +++ b/crates/algorithm/src/prewarm.rs @@ -1,14 +1,15 @@ use crate::closure_lifetime_binder::{id_0, id_1, id_2}; use crate::operator::{FunctionalAccessor, Operator}; +use crate::tape::{by_directory, by_next}; use crate::tuples::*; -use crate::{Page, RelationRead, tape, vectors}; +use crate::{Page, PrefetcherSequenceFamily, RelationRead, tape, vectors}; use std::error::Error; use std::fmt::Write; -pub fn prewarm( - index: R, +pub fn prewarm<'r, R: RelationRead, O: Operator>( + index: &'r R, height: i32, - check: impl Fn(), + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, ) -> Result> { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); @@ -42,19 +43,14 @@ pub fn prewarm( let mut counter = 0_usize; let mut results = Vec::new(); for first in state { - tape::read_h1_tape( - index.clone(), - first, + tape::read_h1_tape::( + by_next(index, first).inspect(|_| counter += 1), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), |(), head, first, prefetch| { let list = prefetch.iter().map(|&id| index.read(id)); vectors::read_for_h1_tuple::(head, list, ()); results.push(first); }, - |_| { - check(); - counter += 1; - }, ); } writeln!(message, "------------------------")?; @@ -72,29 +68,21 @@ pub fn prewarm( let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - tape::read_frozen_tape( - index.clone(), - jump_tuple.frozen_first(), + let directory = + tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory).inspect(|_| counter += 1), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), id_2(|_, _, _, _| { results.push(()); }), - |_| { - check(); - counter += 1; - }, ); - tape::read_appendable_tape( - index.clone(), - jump_tuple.appendable_first(), + tape::read_appendable_tape::( + by_next(index, jump_tuple.appendable_first()).inspect(|_| counter += 1), |_| (), id_2(|_, _, _, _| { results.push(()); }), - |_| { - check(); - counter += 1; - }, ); } writeln!(message, "------------------------")?; diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index 2f0c6e8f..66a1b841 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -15,7 +15,7 @@ type Extra<'b> = &'b mut (NonZero, u16, &'b mut [u32]); type Result = (Reverse, AlwaysEqual>); -pub fn how(index: impl RelationRead) -> RerankMethod { +pub fn how(index: &impl RelationRead) -> RerankMethod { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -34,17 +34,17 @@ pub struct Reranker { _phantom: PhantomData T>, } -impl<'b, T, F, P> Iterator for Reranker +impl<'r, 'b, T, F, P> Iterator for Reranker where - F: FnMut(NonZero, Vec<::ReadGuard<'_>>, u16) -> Option, - P: Prefetcher, AlwaysEqual), AlwaysEqual>)>, + F: FnMut(NonZero, Vec<::ReadGuard<'r>>, u16) -> Option, + P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, { type Item = (Distance, NonZero); fn next(&mut self) -> Option { while let Some(((_, AlwaysEqual(&mut (payload, head, ..))), list)) = self .prefetcher - .pop_if(|((d, _), ..)| Some(*d) > self.cache.peek().map(|(d, ..)| *d)) + .next_if(|((d, _), ..)| Some(*d) > self.cache.peek().map(|(d, ..)| *d)) { if let Some(distance) = (self.f)(payload, list, head) { self.cache.push((Reverse(distance), AlwaysEqual(payload))); @@ -62,10 +62,11 @@ impl Reranker { } pub fn rerank_index< + 'r, 'b, O: Operator, T, - P: Prefetcher, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, >( vector: O::Vector, prefetcher: P, @@ -93,10 +94,11 @@ pub fn rerank_index< } pub fn rerank_heap< + 'r, 'b, O: Operator, T, - P: Prefetcher, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, >( vector: O::Vector, prefetcher: P, diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index 9ba922c0..7e4bbe4b 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -1,9 +1,10 @@ use crate::closure_lifetime_binder::id_2; use crate::linked_vec::LinkedVec; use crate::operator::*; -use crate::prefetcher::Prefetcher; +use crate::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use crate::tape::{by_directory, by_next}; use crate::tuples::*; -use crate::{Bump, Page, RelationRead, tape, vectors}; +use crate::{Bump, Page, PrefetcherHeapFamily, RelationRead, tape, vectors}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -11,20 +12,17 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -type Item<'b> = ( - Reverse, - AlwaysEqual<&'b mut (u32, u16, &'b mut [u32])>, -); - type Extra<'b> = &'b mut (NonZero, u16, &'b mut [u32]); -pub fn default_search<'b, R: RelationRead, O: Operator, P: Prefetcher>>( - index: R, +#[allow(clippy::too_many_arguments)] +pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( + index: &'r R, vector: O::Vector, probes: Vec, epsilon: f32, bump: &'b impl Bump, - mut prefetch: impl FnMut(Vec>) -> P, + mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, ) -> Vec<((Reverse, AlwaysEqual<()>), AlwaysEqual>)> { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); @@ -78,9 +76,8 @@ pub fn default_search<'b, R: RelationRead, O: Operator, P: Prefetcher( + by_next(index, first), || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), |(rough, err), head, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); @@ -89,15 +86,14 @@ pub fn default_search<'b, R: RelationRead, O: Operator, P: Prefetcher, _, _)>::new(); let vector = vector.as_borrowed(); std::iter::from_fn(move || { while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, ..))), list)) = - heap.pop_if(|(d, ..)| Some(*d) > cache.peek().map(|(d, ..)| *d)) + heap.next_if(|(d, ..)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { if is_residual { let (distance, residual) = vectors::read_for_h1_tuple::( @@ -153,32 +149,32 @@ pub fn default_search<'b, R: RelationRead, O: Operator, P: Prefetcher(by_next(index, jump_tuple.directory_first())); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory), || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), &mut callback, - |_| (), ); - tape::read_appendable_tape( - index.clone(), - jump_tuple.appendable_first(), + tape::read_appendable_tape::( + by_next(index, jump_tuple.appendable_first()), |code| O::binary_process(binary_lut, code), &mut callback, - |_| (), ); } results.into_vec() } -pub fn maxsim_search<'b, R: RelationRead, O: Operator, P: Prefetcher>>( - index: R, +#[allow(clippy::too_many_arguments)] +pub fn maxsim_search<'r, 'b: 'r, R: RelationRead, O: Operator>( + index: &'r R, vector: O::Vector, probes: Vec, epsilon: f32, mut threshold: u32, bump: &'b impl Bump, - mut prefetch: impl FnMut(Vec>) -> P, + mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, ) -> ( Vec<( (Reverse, AlwaysEqual), @@ -238,9 +234,8 @@ pub fn maxsim_search<'b, R: RelationRead, O: Operator, P: Prefetcher( + by_next(index, first), || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), |(rough, err), head, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); @@ -249,15 +244,14 @@ pub fn maxsim_search<'b, R: RelationRead, O: Operator, P: Prefetcher, _, _)>::new(); let vector = vector.as_borrowed(); std::iter::from_fn(move || { while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, ..))), list)) = - heap.pop_if(|(d, ..)| Some(*d) > cache.peek().map(|(d, ..)| *d)) + heap.next_if(|(d, ..)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { if is_residual { let (distance, residual) = vectors::read_for_h1_tuple::( @@ -316,19 +310,17 @@ pub fn maxsim_search<'b, R: RelationRead, O: Operator, P: Prefetcher(by_next(index, jump_tuple.directory_first())); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory), || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), &mut callback, - |_| (), ); - tape::read_appendable_tape( - index.clone(), - jump_tuple.appendable_first(), + tape::read_appendable_tape::( + by_next(index, jump_tuple.appendable_first()), |code| O::binary_process(binary_lut, code), &mut callback, - |_| (), ); threshold = threshold.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as _); } diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index b941afc9..26f5922b 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -1,6 +1,6 @@ use crate::operator::Accessor1; use crate::tuples::*; -use crate::{Page, PageGuard, RelationRead, RelationWrite}; +use crate::{Page, PageGuard, PrefetcherSequenceFamily, RelationRead, RelationWrite}; use rabitq::binary::BinaryCode; use std::marker::PhantomData; use std::num::NonZero; @@ -78,25 +78,133 @@ where } } -pub fn read_h1_tape( - index: impl RelationRead, - first: u32, +pub fn read_directory_tape<'r, R>( + iter: impl Iterator>, +) -> impl Iterator +where + R: RelationRead + 'r, +{ + use std::pin::Pin; + use std::ptr::NonNull; + + #[pin_project::pin_project] + struct State<'r, R: RelationRead + 'r, I> { + slice: NonNull<[u32]>, + #[pin] + now: Option<(R::ReadGuard<'r>, u16)>, + iter: I, + } + + impl<'r, R: RelationRead + 'r, I: Iterator>> State<'r, R, I> { + fn init(self: Pin<&mut Self>) { + let mut this = self.project(); + let now = this.iter.next().map(|guard| (guard, 0)); + this.now.set(now); + } + + fn next(mut self: Pin<&mut Self>) -> Option { + loop { + let mut this = self.as_mut().project(); + // Safety: If the slice is not empty, the function will return immediately, + // so the guard will not be moved or dropped and the slice remains valid. If + // the slice is empty, a pointer is trivially never dangling, so it's safe + // to use. + #[allow(unsafe_code)] + if let Some((first, more)) = unsafe { this.slice.as_ref() }.split_first() { + *this.slice = more.into(); + return Some(*first); + } + // Safety: `guard` is never moved in this block + #[allow(unsafe_code)] + if let Some((guard, i)) = unsafe { this.now.as_mut().get_unchecked_mut() } { + if *i < guard.len() { + *i += 1; + let bytes = guard.get(*i).expect("data corruption"); + let tuple = DirectoryTuple::deserialize_ref(bytes); + *this.slice = match tuple { + DirectoryTupleReader::_0(tuple) => tuple.elements(), + DirectoryTupleReader::_1(tuple) => tuple.elements(), + } + .into(); + continue; + } + } else { + return None; + } + let now = this.iter.next().map(|guard| (guard, 0)); + this.now.set(now); + } + } + } + + let mut state = Box::pin(State::<'r, R, _> { + slice: NonNull::from(&mut []), + now: None, + iter, + }); + + impl<'r, R: RelationRead + 'r, I: Iterator>> Iterator + for Pin>> + { + type Item = u32; + + fn next(&mut self) -> Option { + self.as_mut().next() + } + } + + state.as_mut().init(); + + state +} + +pub fn by_directory<'r, R>( + p: &mut impl PrefetcherSequenceFamily<'r, R>, + iter: impl Iterator, +) -> impl Iterator> +where + R: RelationRead + 'r, +{ + let mut t = p.prefetch(iter.peekable()); + std::iter::from_fn(move || { + use crate::prefetcher::Prefetcher; + let (_, mut x) = t.next()?; + let ret = x.pop().expect("should be at least one element"); + assert!(x.pop().is_none(), "should be at most one element"); + Some(ret) + }) +} + +pub fn by_next<'r, R>(index: &'r R, first: u32) -> impl Iterator> +where + R: RelationRead + 'r, +{ + let mut current = first; + std::iter::from_fn(move || { + if current != u32::MAX { + let guard = index.read(current); + current = guard.get_opaque().next; + Some(guard) + } else { + None + } + }) +} + +pub fn read_h1_tape<'r, R, A, T>( + iter: impl Iterator>, accessor: impl Fn() -> A, mut callback: impl for<'a> FnMut(T, u16, u32, &'a [u32]), - mut step: impl FnMut(u32), ) where + R: RelationRead + 'r, A: for<'a> Accessor1< [u8; 16], (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), Output = [T; 32], >, { - assert!(first != u32::MAX); - let mut current = first; let mut x = None; - while current != u32::MAX { - step(current); - let guard = index.read(current); + for guard in iter { for i in 1..=guard.len() { let bytes = guard.get(i).expect("data corruption"); let tuple = H1Tuple::deserialize_ref(bytes); @@ -117,29 +225,23 @@ pub fn read_h1_tape( } } } - current = guard.get_opaque().next; } } -pub fn read_frozen_tape( - index: impl RelationRead, - first: u32, +pub fn read_frozen_tape<'r, R, A, T>( + iter: impl Iterator>, accessor: impl Fn() -> A, mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), - mut step: impl FnMut(u32), ) where + R: RelationRead + 'r, A: for<'a> Accessor1< [u8; 16], (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), Output = [T; 32], >, { - assert!(first != u32::MAX); - let mut current = first; let mut x = None; - while current != u32::MAX { - step(current); - let guard = index.read(current); + for guard in iter { for i in 1..=guard.len() { let bytes = guard.get(i).expect("data corruption"); let tuple = FrozenTuple::deserialize_ref(bytes); @@ -160,22 +262,17 @@ pub fn read_frozen_tape( } } } - current = guard.get_opaque().next; } } -pub fn read_appendable_tape( - index: impl RelationRead, - first: u32, +pub fn read_appendable_tape<'r, R, T>( + iter: impl Iterator>, mut access: impl for<'a> FnMut(BinaryCode<'a>) -> T, mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), - mut step: impl FnMut(u32), -) { - assert!(first != u32::MAX); - let mut current = first; - while current != u32::MAX { - step(current); - let guard = index.read(current); +) where + R: RelationRead + 'r, +{ + for guard in iter { for i in 1..=guard.len() { let bytes = guard.get(i).expect("data corruption"); let tuple = AppendableTuple::deserialize_ref(bytes); @@ -184,13 +281,12 @@ pub fn read_appendable_tape( callback(value, tuple.head(), payload, tuple.prefetch()); } } - current = guard.get_opaque().next; } } #[allow(clippy::collapsible_else_if)] pub fn append( - index: impl RelationRead + RelationWrite, + index: &(impl RelationRead + RelationWrite), first: u32, bytes: &[u8], tracking_freespace: bool, diff --git a/crates/algorithm/src/tape_writer.rs b/crates/algorithm/src/tape_writer.rs new file mode 100644 index 00000000..b10d3493 --- /dev/null +++ b/crates/algorithm/src/tape_writer.rs @@ -0,0 +1,205 @@ +use crate::tape::TapeWriter; +use crate::tuples::*; +use crate::{Branch, RelationWrite}; +use rabitq::packing::{any_pack, padding_pack}; +use std::num::NonZero; + +pub struct DirectoryTapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + tape: TapeWriter<'a, R, DirectoryTuple>, +} + +impl<'a, R> DirectoryTapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + pub fn create(index: &'a R, tracking_freespace: bool) -> Self { + Self { + tape: TapeWriter::create(index, tracking_freespace), + } + } + pub fn push(&mut self, branch: &[u32]) { + let mut remain = branch.to_vec(); + loop { + let freespace = self.tape.freespace(); + if DirectoryTuple::estimate_size_0(remain.len()) <= freespace as usize { + self.tape.tape_put(DirectoryTuple::_0 { elements: remain }); + break; + } + if let Some(w) = DirectoryTuple::fit_1(freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(DirectoryTuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } else { + self.tape.tape_move(); + } + } + } + pub fn into_inner(self) -> TapeWriter<'a, R, DirectoryTuple> { + self.tape + } +} + +pub struct H1TapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + tape: TapeWriter<'a, R, H1Tuple>, + branches: Vec>, + prefetch: usize, +} + +impl<'a, R> H1TapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + pub fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { + Self { + tape: TapeWriter::create(index, tracking_freespace), + branches: Vec::new(), + prefetch, + } + } + pub fn push(&mut self, branch: Branch) { + self.branches.push(branch); + if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + loop { + let freespace = self.tape.freespace(); + if H1Tuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { + self.tape.tape_put(H1Tuple::_0 { + head: chunk.each_ref().map(|x| x.head), + dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), + factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), + factor_ip: chunk.each_ref().map(|x| x.factor_ip), + factor_err: chunk.each_ref().map(|x| x.factor_err), + first: chunk.each_ref().map(|x| x.extra), + prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), + len: chunk.len() as _, + elements: remain, + }); + break; + } + if let Some(w) = H1Tuple::fit_1(self.prefetch, freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(H1Tuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } else { + self.tape.tape_move(); + } + } + self.branches.clear(); + } + } + pub fn into_inner(self) -> (TapeWriter<'a, R, H1Tuple>, Vec>) { + (self.tape, self.branches) + } + pub fn flush(tape: &mut TapeWriter<'_, R, H1Tuple>, prefetch: usize, chunk: Vec>) { + if chunk.is_empty() { + return; + } + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + loop { + let freespace = tape.freespace(); + if H1Tuple::estimate_size_0(prefetch, remain.len()) <= freespace as usize { + tape.tape_put(H1Tuple::_0 { + head: any_pack(chunk.iter().map(|x| x.head)), + dis_u_2: any_pack(chunk.iter().map(|x| x.dis_u_2)), + factor_ppc: any_pack(chunk.iter().map(|x| x.factor_ppc)), + factor_ip: any_pack(chunk.iter().map(|x| x.factor_ip)), + factor_err: any_pack(chunk.iter().map(|x| x.factor_err)), + first: any_pack(chunk.iter().map(|x| x.extra)), + prefetch: fix(chunk.iter().map(|x| x.prefetch.as_slice())), + len: chunk.len() as _, + elements: remain, + }); + break; + } + if let Some(w) = H1Tuple::fit_1(prefetch, freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + tape.tape_put(H1Tuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } else { + tape.tape_move(); + } + } + } +} + +pub struct FrozenTapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + tape: TapeWriter<'a, R, FrozenTuple>, + branches: Vec>>, + prefetch: usize, +} + +impl<'a, R> FrozenTapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + pub fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { + Self { + tape: TapeWriter::create(index, tracking_freespace), + branches: Vec::new(), + prefetch, + } + } + pub fn push(&mut self, branch: Branch>) { + self.branches.push(branch); + if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + loop { + let freespace = self.tape.freespace(); + if FrozenTuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { + self.tape.tape_put(FrozenTuple::_0 { + head: chunk.each_ref().map(|x| x.head), + dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), + factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), + factor_ip: chunk.each_ref().map(|x| x.factor_ip), + factor_err: chunk.each_ref().map(|x| x.factor_err), + payload: chunk.each_ref().map(|x| Some(x.extra)), + prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), + elements: remain, + }); + break; + } + if let Some(w) = FrozenTuple::fit_1(self.prefetch, freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(FrozenTuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } else { + self.tape.tape_move(); + } + } + self.branches.clear(); + } + } + pub fn into_inner(self) -> (TapeWriter<'a, R, FrozenTuple>, Vec>>) { + (self.tape, self.branches) + } +} + +fn fix<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { + use std::array::from_fn; + let mut iter = into_iter.into_iter(); + let mut array: [_; 32] = from_fn(|_| iter.next().map(<[u32]>::to_vec).unwrap_or_default()); + if iter.next().is_some() { + panic!("too many slices"); + } + let step = array.iter().map(Vec::len).max().unwrap_or_default(); + array.iter_mut().for_each(|x| x.resize(step, u32::MAX)); + let flat = array.into_iter().flatten().collect::>(); + (0..step).map(|i| from_fn(|j| flat[i * 32 + j])).collect() +} diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 0be43a7f..f2d30114 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -428,6 +428,147 @@ impl<'a, V: Vector> VectorTupleReader<'a, V> { } } +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct DirectoryTupleHeader0 { + elements_s: u16, + elements_e: u16, + _padding_0: [ZeroU8; 4], +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct DirectoryTupleHeader1 { + elements_s: u16, + elements_e: u16, + _padding_0: [ZeroU8; 4], +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug, Clone, PartialEq)] +pub enum DirectoryTuple { + _0 { elements: Vec }, + _1 { elements: Vec }, +} + +impl DirectoryTuple { + pub fn estimate_size_0(elements: usize) -> usize { + let mut size = 0_usize; + size += size_of::(); + size += size_of::(); + size += (elements * size_of::()).next_multiple_of(ALIGN); + size + } + pub fn fit_1(freespace: u16) -> Option { + let mut freespace = freespace as isize; + freespace -= size_of::() as isize; + freespace -= size_of::() as isize; + if freespace >= 0 { + Some(freespace as usize / size_of::()) + } else { + None + } + } +} + +impl Tuple for DirectoryTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + Self::_0 { elements } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + buffer[size_of::()..][..size_of::()].copy_from_slice( + DirectoryTupleHeader0 { + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + Self::_1 { elements } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + buffer[size_of::()..][..size_of::()].copy_from_slice( + DirectoryTupleHeader1 { + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for DirectoryTuple { + type Reader<'a> = DirectoryTupleReader<'a>; + + fn deserialize_ref(source: &[u8]) -> DirectoryTupleReader<'_> { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &DirectoryTupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + DirectoryTupleReader::_0(DirectoryTupleReader0 { header, elements }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &DirectoryTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + DirectoryTupleReader::_1(DirectoryTupleReader1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum DirectoryTupleReader<'a> { + _0(DirectoryTupleReader0<'a>), + _1(DirectoryTupleReader1<'a>), +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct DirectoryTupleReader0<'a> { + header: &'a DirectoryTupleHeader0, + elements: &'a [u32], +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct DirectoryTupleReader1<'a> { + header: &'a DirectoryTupleHeader1, + elements: &'a [u32], +} + +impl<'a> DirectoryTupleReader0<'a> { + pub fn elements(&self) -> &'a [u32] { + self.elements + } +} + +impl<'a> DirectoryTupleReader1<'a> { + pub fn elements(&self) -> &'a [u32] { + self.elements + } +} + #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct H1TupleHeader0 { @@ -646,14 +787,14 @@ impl<'a> H1TupleReader1<'a> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct JumpTupleHeader { - frozen_first: u32, + directory_first: u32, appendable_first: u32, tuples: u64, } #[derive(Debug, Clone, PartialEq)] pub struct JumpTuple { - pub frozen_first: u32, + pub directory_first: u32, pub appendable_first: u32, pub tuples: u64, } @@ -661,7 +802,7 @@ pub struct JumpTuple { impl Tuple for JumpTuple { fn serialize(&self) -> Vec { JumpTupleHeader { - frozen_first: self.frozen_first, + directory_first: self.directory_first, appendable_first: self.appendable_first, tuples: self.tuples, } @@ -694,8 +835,8 @@ pub struct JumpTupleReader<'a> { } impl JumpTupleReader<'_> { - pub fn frozen_first(self) -> u32 { - self.header.frozen_first + pub fn directory_first(self) -> u32 { + self.header.directory_first } pub fn appendable_first(self) -> u32 { self.header.appendable_first @@ -711,8 +852,8 @@ pub struct JumpTupleWriter<'a> { } impl JumpTupleWriter<'_> { - pub fn frozen_first(&mut self) -> &mut u32 { - &mut self.header.frozen_first + pub fn directory_first(&mut self) -> &mut u32 { + &mut self.header.directory_first } pub fn appendable_first(&mut self) -> &mut u32 { &mut self.header.appendable_first diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index 92db860b..8d4c7882 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -65,12 +65,12 @@ pub fn read_for_h0_tuple< } pub fn append( - index: impl RelationRead + RelationWrite, + index: &(impl RelationRead + RelationWrite), vectors_first: u32, vector: ::Borrowed<'_>, payload: NonZero, ) -> (Vec, u16) { - fn append(index: impl RelationRead + RelationWrite, first: u32, bytes: &[u8]) -> (u32, u16) { + fn append(index: &(impl RelationRead + RelationWrite), first: u32, bytes: &[u8]) -> (u32, u16) { if let Some(mut write) = index.search(bytes.len()) { let i = write .alloc(bytes) @@ -95,7 +95,7 @@ pub fn append( head, }, }); - let (id, head) = append(index.clone(), vectors_first, &bytes); + let (id, head) = append(index, vectors_first, &bytes); chain = Err(head); prefetch.push(id); } diff --git a/crates/always_equal/Cargo.toml b/crates/always_equal/Cargo.toml index 923fa8e1..887633ab 100644 --- a/crates/always_equal/Cargo.toml +++ b/crates/always_equal/Cargo.toml @@ -2,6 +2,7 @@ name = "always_equal" version.workspace = true edition.workspace = true +publish = false [lints] workspace = true diff --git a/crates/distance/Cargo.toml b/crates/distance/Cargo.toml index 6c578193..c5a21ed3 100644 --- a/crates/distance/Cargo.toml +++ b/crates/distance/Cargo.toml @@ -2,6 +2,7 @@ name = "distance" version.workspace = true edition.workspace = true +publish = false [lints] workspace = true diff --git a/crates/k_means/Cargo.toml b/crates/k_means/Cargo.toml index 00e5efe1..fd3b43a7 100644 --- a/crates/k_means/Cargo.toml +++ b/crates/k_means/Cargo.toml @@ -2,6 +2,7 @@ name = "k_means" version.workspace = true edition.workspace = true +publish = false [dependencies] rabitq = { path = "../rabitq" } diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index 4c04b884..052cbc6a 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -1,9 +1,9 @@ use rabitq::block::BlockCode; +use rabitq::packing::{any_pack, padding_pack}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use rayon::iter::{IntoParallelIterator, IntoParallelRefMutIterator, ParallelIterator}; use simd::Floating; -use simd::fast_scan::{any_pack, padding_pack}; pub fn preprocess(num_threads: usize, x: &mut [T], f: impl Fn(&mut T) + Sync) { rayon::ThreadPoolBuilder::new() diff --git a/crates/rabitq/Cargo.toml b/crates/rabitq/Cargo.toml index 5ca300c4..f1572994 100644 --- a/crates/rabitq/Cargo.toml +++ b/crates/rabitq/Cargo.toml @@ -2,6 +2,7 @@ name = "rabitq" version.workspace = true edition.workspace = true +publish = false [dependencies] distance = { path = "../distance" } diff --git a/crates/rabitq/src/binary.rs b/crates/rabitq/src/binary.rs index 1f8bc10f..402eeda7 100644 --- a/crates/rabitq/src/binary.rs +++ b/crates/rabitq/src/binary.rs @@ -12,7 +12,7 @@ pub fn preprocess(vector: &[f32]) -> BinaryLut { let qvector_sum = if vector.len() <= 4369 { simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 } else { - simd::u8::reduce_sum_of_x(&qvector) as f32 + simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 }; ((dis_v_2, b, k, qvector_sum), binarize(&qvector)) } @@ -59,9 +59,9 @@ pub(crate) fn asymmetric_binary_dot_product( x: &[u64], y: &(Vec, Vec, Vec, Vec), ) -> u32 { - let t0 = simd::bit::sum_of_and(x, &y.0); - let t1 = simd::bit::sum_of_and(x, &y.1); - let t2 = simd::bit::sum_of_and(x, &y.2); - let t3 = simd::bit::sum_of_and(x, &y.3); + let t0 = simd::bit::reduce_sum_of_and(x, &y.0); + let t1 = simd::bit::reduce_sum_of_and(x, &y.1); + let t2 = simd::bit::reduce_sum_of_and(x, &y.2); + let t3 = simd::bit::reduce_sum_of_and(x, &y.3); (t0 << 0) + (t1 << 1) + (t2 << 2) + (t3 << 3) } diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs index 83496f92..ce3071e3 100644 --- a/crates/rabitq/src/block.rs +++ b/crates/rabitq/src/block.rs @@ -15,7 +15,7 @@ pub fn preprocess(vector: &[f32]) -> BlockLut { let qvector_sum = if vector.len() <= 4369 { simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 } else { - simd::u8::reduce_sum_of_x(&qvector) as f32 + simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 }; ((dis_v_2, b, k, qvector_sum), compress(qvector)) } @@ -25,7 +25,7 @@ pub fn process_l2( (dis_u_2, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, ) -> [(f32, f32); 32] { let &((dis_v_2, b, k, qvector_sum), ref s) = lut; - let r = simd::fast_scan::scan(t, s); + let r = simd::fast_scan::fast_scan(t, s); std::array::from_fn(|i| { let rough = dis_u_2[i] + dis_v_2 @@ -41,7 +41,7 @@ pub fn process_dot( (_, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, ) -> [(f32, f32); 32] { let &((dis_v_2, b, k, qvector_sum), ref s) = lut; - let r = simd::fast_scan::scan(t, s); + let r = simd::fast_scan::fast_scan(t, s); std::array::from_fn(|i| { let rough = 0.5 * b * factor_ppc[i] + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index 7d4b81bd..a4537372 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -1,5 +1,6 @@ pub mod binary; pub mod block; +pub mod packing; use binary::BinaryLut; use block::BlockLut; @@ -81,7 +82,7 @@ pub fn preprocess(vector: &[f32]) -> (BlockLut, BinaryLut) { let qvector_sum = if vector.len() <= 4369 { simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 } else { - simd::u8::reduce_sum_of_x(&qvector) as f32 + simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 }; let binary = binary::binarize(&qvector); let block = block::compress(qvector); diff --git a/crates/rabitq/src/packing.rs b/crates/rabitq/src/packing.rs new file mode 100644 index 00000000..c9463e37 --- /dev/null +++ b/crates/rabitq/src/packing.rs @@ -0,0 +1,88 @@ +pub fn pack(x: [&[u8]; 32]) -> Vec<[u8; 16]> { + let n = { + let l = x.each_ref().map(|i| i.len()); + for i in 1..32 { + assert!(l[0] == l[i]); + } + l[0] + }; + let mut result = Vec::with_capacity(n); + for i in 0..n { + result.push([ + x[0][i] | (x[16][i] << 4), + x[8][i] | (x[24][i] << 4), + x[1][i] | (x[17][i] << 4), + x[9][i] | (x[25][i] << 4), + x[2][i] | (x[18][i] << 4), + x[10][i] | (x[26][i] << 4), + x[3][i] | (x[19][i] << 4), + x[11][i] | (x[27][i] << 4), + x[4][i] | (x[20][i] << 4), + x[12][i] | (x[28][i] << 4), + x[5][i] | (x[21][i] << 4), + x[13][i] | (x[29][i] << 4), + x[6][i] | (x[22][i] << 4), + x[14][i] | (x[30][i] << 4), + x[7][i] | (x[23][i] << 4), + x[15][i] | (x[31][i] << 4), + ]); + } + result +} + +pub fn unpack(x: &[[u8; 16]]) -> [Vec; 32] { + let n = x.len(); + let mut result = std::array::from_fn(|_| Vec::with_capacity(n)); + for i in 0..n { + result[0].push(x[i][0] & 0xf); + result[1].push(x[i][2] & 0xf); + result[2].push(x[i][4] & 0xf); + result[3].push(x[i][6] & 0xf); + result[4].push(x[i][8] & 0xf); + result[5].push(x[i][10] & 0xf); + result[6].push(x[i][12] & 0xf); + result[7].push(x[i][14] & 0xf); + result[8].push(x[i][1] & 0xf); + result[9].push(x[i][3] & 0xf); + result[10].push(x[i][5] & 0xf); + result[11].push(x[i][7] & 0xf); + result[12].push(x[i][9] & 0xf); + result[13].push(x[i][11] & 0xf); + result[14].push(x[i][13] & 0xf); + result[15].push(x[i][15] & 0xf); + result[16].push(x[i][0] >> 4); + result[17].push(x[i][2] >> 4); + result[18].push(x[i][4] >> 4); + result[19].push(x[i][6] >> 4); + result[20].push(x[i][8] >> 4); + result[21].push(x[i][10] >> 4); + result[22].push(x[i][12] >> 4); + result[23].push(x[i][14] >> 4); + result[24].push(x[i][1] >> 4); + result[25].push(x[i][3] >> 4); + result[26].push(x[i][5] >> 4); + result[27].push(x[i][7] >> 4); + result[28].push(x[i][9] >> 4); + result[29].push(x[i][11] >> 4); + result[30].push(x[i][13] >> 4); + result[31].push(x[i][15] >> 4); + } + result +} + +pub fn padding_pack(x: impl IntoIterator>) -> Vec<[u8; 16]> { + let x = x.into_iter().collect::>(); + let x = x.iter().map(|x| x.as_ref()).collect::>(); + if x.is_empty() || x.len() > 32 { + panic!("too few or too many slices"); + } + let n = x[0].len(); + let t = vec![0; n]; + pack(std::array::from_fn(|i| { + if i < x.len() { x[i] } else { t.as_slice() } + })) +} + +pub fn any_pack(mut x: impl Iterator) -> [T; 32] { + std::array::from_fn(|_| x.next()).map(|x| x.unwrap_or_default()) +} diff --git a/crates/random_orthogonal_matrix/Cargo.toml b/crates/random_orthogonal_matrix/Cargo.toml index 348ecb7f..ee869dd1 100644 --- a/crates/random_orthogonal_matrix/Cargo.toml +++ b/crates/random_orthogonal_matrix/Cargo.toml @@ -2,6 +2,7 @@ name = "random_orthogonal_matrix" version.workspace = true edition.workspace = true +publish = false [dependencies] # lock algebra version forever so that the QR decomposition never changes for same input diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 01a3b681..74738125 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -2,6 +2,7 @@ name = "simd" version.workspace = true edition.workspace = true +publish = false [dependencies] simd_macros = { path = "../simd_macros" } diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index 1e587bb7..2ace1b68 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -1,16 +1,16 @@ #[inline(always)] -pub fn sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { - sum_of_and::sum_of_and(lhs, rhs) +pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { + reduce_sum_of_and::reduce_sum_of_and(lhs, rhs) } -mod sum_of_and { +mod reduce_sum_of_and { // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_and_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + fn reduce_sum_of_and_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut and = _mm512_setzero_si512(); @@ -36,7 +36,7 @@ mod sum_of_and { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_and_v4_avx512vpopcntdq_test() { + fn reduce_sum_of_and_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; @@ -44,14 +44,14 @@ mod sum_of_and { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_and_v4_avx512vpopcntdq(&lhs, &rhs) }; + let specialized = unsafe { reduce_sum_of_and_v4_avx512vpopcntdq(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] - pub fn sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { + pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); let mut and = 0; @@ -63,18 +63,18 @@ mod sum_of_and { } #[inline(always)] -pub fn sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { - sum_of_or::sum_of_or(lhs, rhs) +pub fn reduce_sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { + reduce_sum_of_or::reduce_sum_of_or(lhs, rhs) } -mod sum_of_or { +mod reduce_sum_of_or { // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + fn reduce_sum_of_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut or = _mm512_setzero_si512(); @@ -100,7 +100,7 @@ mod sum_of_or { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_or_v4_avx512vpopcntdq_test() { + fn reduce_sum_of_or_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; @@ -108,14 +108,14 @@ mod sum_of_or { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_or_v4_avx512vpopcntdq(&lhs, &rhs) }; + let specialized = unsafe { reduce_sum_of_or_v4_avx512vpopcntdq(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] - pub fn sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { + pub fn reduce_sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); let mut or = 0; @@ -127,18 +127,18 @@ mod sum_of_or { } #[inline(always)] -pub fn sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { - sum_of_xor::sum_of_xor(lhs, rhs) +pub fn reduce_sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { + reduce_sum_of_xor::reduce_sum_of_xor(lhs, rhs) } -mod sum_of_xor { +mod reduce_sum_of_xor { // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_xor_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + fn reduce_sum_of_xor_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut xor = _mm512_setzero_si512(); @@ -164,7 +164,7 @@ mod sum_of_xor { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_xor_v4_avx512vpopcntdq_test() { + fn reduce_sum_of_xor_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; @@ -172,14 +172,14 @@ mod sum_of_xor { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_xor_v4_avx512vpopcntdq(&lhs, &rhs) }; + let specialized = unsafe { reduce_sum_of_xor_v4_avx512vpopcntdq(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] - pub fn sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { + pub fn reduce_sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); let mut xor = 0; @@ -191,18 +191,18 @@ mod sum_of_xor { } #[inline(always)] -pub fn sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { - sum_of_and_or::sum_of_and_or(lhs, rhs) +pub fn reduce_sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + reduce_sum_of_and_or::reduce_sum_of_and_or(lhs, rhs) } -mod sum_of_and_or { +mod reduce_sum_of_and_or { // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_and_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + fn reduce_sum_of_and_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert!(lhs.len() == rhs.len()); use std::arch::x86_64::*; let mut and = _mm512_setzero_si512(); @@ -234,7 +234,7 @@ mod sum_of_and_or { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_xor_v4_avx512vpopcntdq_test() { + fn reduce_sum_of_xor_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; @@ -242,14 +242,14 @@ mod sum_of_and_or { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let lhs = (0..126).map(|_| rand::random::()).collect::>(); let rhs = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_and_or_v4_avx512vpopcntdq(&lhs, &rhs) }; + let specialized = unsafe { reduce_sum_of_and_or_v4_avx512vpopcntdq(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); assert_eq!(specialized, fallback); } } #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] - pub fn sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + pub fn reduce_sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); let mut and = 0; @@ -263,18 +263,18 @@ mod sum_of_and_or { } #[inline(always)] -pub fn sum_of_x(this: &[u64]) -> u32 { - sum_of_x::sum_of_x(this) +pub fn reduce_sum_of_x(this: &[u64]) -> u32 { + reduce_sum_of_x::reduce_sum_of_x(this) } -mod sum_of_x { +mod reduce_sum_of_x { // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] - fn sum_of_x_v4_avx512vpopcntdq(this: &[u64]) -> u32 { + fn reduce_sum_of_x_v4_avx512vpopcntdq(this: &[u64]) -> u32 { use std::arch::x86_64::*; let mut and = _mm512_setzero_si512(); let mut a = this.as_ptr(); @@ -295,21 +295,21 @@ mod sum_of_x { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn sum_of_x_v4_avx512vpopcntdq_test() { + fn reduce_sum_of_x_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let this = (0..126).map(|_| rand::random::()).collect::>(); - let specialized = unsafe { sum_of_x_v4_avx512vpopcntdq(&this) }; + let specialized = unsafe { reduce_sum_of_x_v4_avx512vpopcntdq(&this) }; let fallback = fallback(&this); assert_eq!(specialized, fallback); } } #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] - pub fn sum_of_x(this: &[u64]) -> u32 { + pub fn reduce_sum_of_x(this: &[u64]) -> u32 { let n = this.len(); let mut and = 0; for i in 0..n { diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index 23fc54f3..8d3f9db6 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -32,95 +32,6 @@ bits 4..7 | code (n-1),vector 16 | code (n-1),vector 24 | ... | code (n-1),vecto */ -pub fn pack(x: [&[u8]; 32]) -> Vec<[u8; 16]> { - let n = { - let l = x.each_ref().map(|i| i.len()); - for i in 1..32 { - assert!(l[0] == l[i]); - } - l[0] - }; - let mut result = Vec::with_capacity(n); - for i in 0..n { - result.push([ - x[0][i] | (x[16][i] << 4), - x[8][i] | (x[24][i] << 4), - x[1][i] | (x[17][i] << 4), - x[9][i] | (x[25][i] << 4), - x[2][i] | (x[18][i] << 4), - x[10][i] | (x[26][i] << 4), - x[3][i] | (x[19][i] << 4), - x[11][i] | (x[27][i] << 4), - x[4][i] | (x[20][i] << 4), - x[12][i] | (x[28][i] << 4), - x[5][i] | (x[21][i] << 4), - x[13][i] | (x[29][i] << 4), - x[6][i] | (x[22][i] << 4), - x[14][i] | (x[30][i] << 4), - x[7][i] | (x[23][i] << 4), - x[15][i] | (x[31][i] << 4), - ]); - } - result -} - -pub fn unpack(x: &[[u8; 16]]) -> [Vec; 32] { - let n = x.len(); - let mut result = std::array::from_fn(|_| Vec::with_capacity(n)); - for i in 0..n { - result[0].push(x[i][0] & 0xf); - result[1].push(x[i][2] & 0xf); - result[2].push(x[i][4] & 0xf); - result[3].push(x[i][6] & 0xf); - result[4].push(x[i][8] & 0xf); - result[5].push(x[i][10] & 0xf); - result[6].push(x[i][12] & 0xf); - result[7].push(x[i][14] & 0xf); - result[8].push(x[i][1] & 0xf); - result[9].push(x[i][3] & 0xf); - result[10].push(x[i][5] & 0xf); - result[11].push(x[i][7] & 0xf); - result[12].push(x[i][9] & 0xf); - result[13].push(x[i][11] & 0xf); - result[14].push(x[i][13] & 0xf); - result[15].push(x[i][15] & 0xf); - result[16].push(x[i][0] >> 4); - result[17].push(x[i][2] >> 4); - result[18].push(x[i][4] >> 4); - result[19].push(x[i][6] >> 4); - result[20].push(x[i][8] >> 4); - result[21].push(x[i][10] >> 4); - result[22].push(x[i][12] >> 4); - result[23].push(x[i][14] >> 4); - result[24].push(x[i][1] >> 4); - result[25].push(x[i][3] >> 4); - result[26].push(x[i][5] >> 4); - result[27].push(x[i][7] >> 4); - result[28].push(x[i][9] >> 4); - result[29].push(x[i][11] >> 4); - result[30].push(x[i][13] >> 4); - result[31].push(x[i][15] >> 4); - } - result -} - -pub fn padding_pack(x: impl IntoIterator>) -> Vec<[u8; 16]> { - let x = x.into_iter().collect::>(); - let x = x.iter().map(|x| x.as_ref()).collect::>(); - if x.is_empty() || x.len() > 32 { - panic!("too few or too many slices"); - } - let n = x[0].len(); - let t = vec![0; n]; - pack(std::array::from_fn(|i| { - if i < x.len() { x[i] } else { t.as_slice() } - })) -} - -pub fn any_pack(mut x: impl Iterator) -> [T; 32] { - std::array::from_fn(|_| x.next()).map(|x| x.unwrap_or_default()) -} - mod scan { #[inline] #[cfg(target_arch = "x86_64")] @@ -542,6 +453,6 @@ mod scan { } #[inline(always)] -pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { +pub fn fast_scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { scan::scan(code, lut) } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index bdb58b03..11291f28 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -9,7 +9,6 @@ mod f32; pub mod bit; pub mod fast_scan; -pub mod packed_u4; pub mod quantize; pub mod u8; diff --git a/crates/simd/src/packed_u4.rs b/crates/simd/src/packed_u4.rs deleted file mode 100644 index 9ec7a331..00000000 --- a/crates/simd/src/packed_u4.rs +++ /dev/null @@ -1,18 +0,0 @@ -pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { - reduce_sum_of_xy::reduce_sum_of_xy(s, t) -} - -mod reduce_sum_of_xy { - #[crate::multiversion("v4", "v3", "v2", "a2")] - pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { - assert_eq!(s.len(), t.len()); - let n = s.len(); - let mut result = 0; - for i in 0..n { - let (s, t) = (s[i], t[i]); - result += ((s & 15) as u32) * ((t & 15) as u32); - result += ((s >> 4) as u32) * ((t >> 4) as u32); - } - result - } -} diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index 7f5026cf..53fc3510 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -1,9 +1,9 @@ #![allow(clippy::just_underscores_and_digits)] -mod reduce_sum_of_xy { +mod reduce_sum_of_x_as_u32_y_as_u32 { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_x_as_u32_y_as_u32_v4(lhs: &[u8], rhs: &[u8]) -> u32 { use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); @@ -54,7 +54,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_xy_v4_test() { + fn reduce_sum_of_x_as_u32_y_as_u32_v4_test() { use rand::Rng; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -68,7 +68,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v4(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -81,7 +81,7 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_x_as_u32_y_as_u32_v3(lhs: &[u8], rhs: &[u8]) -> u32 { use crate::emulate::emulate_mm256_reduce_add_epi32; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); @@ -129,7 +129,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_xy_v3_test() { + fn reduce_sum_of_x_as_u32_y_as_u32_v3_test() { use rand::Rng; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -143,7 +143,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v3(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -156,7 +156,7 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_x_as_u32_y_as_u32_v2(lhs: &[u8], rhs: &[u8]) -> u32 { use crate::emulate::emulate_mm_reduce_add_epi32; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); @@ -204,7 +204,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_xy_v2_test() { + fn reduce_sum_of_x_as_u32_y_as_u32_v2_test() { use rand::Rng; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); @@ -218,7 +218,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v2(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v2(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -231,7 +231,7 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] - fn reduce_sum_of_xy_a2(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_x_as_u32_y_as_u32_a2(lhs: &[u8], rhs: &[u8]) -> u32 { use core::arch::aarch64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); @@ -274,7 +274,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] - fn reduce_sum_of_xy_a2_test() { + fn reduce_sum_of_x_as_u32_y_as_u32_a2_test() { use rand::Rng; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -288,7 +288,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_a2(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -299,7 +299,7 @@ mod reduce_sum_of_xy { } #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] - pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); let mut result = 0; @@ -311,8 +311,8 @@ mod reduce_sum_of_xy { } #[inline(always)] -pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { - reduce_sum_of_xy::reduce_sum_of_xy(s, t) +pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { + reduce_sum_of_x_as_u32_y_as_u32::reduce_sum_of_x_as_u32_y_as_u32(s, t) } mod reduce_sum_of_x_as_u16 { @@ -520,11 +520,11 @@ pub fn reduce_sum_of_x_as_u16(vector: &[u8]) -> u16 { reduce_sum_of_x_as_u16::reduce_sum_of_x_as_u16(vector) } -mod reduce_sum_of_x { +mod reduce_sum_of_x_as_u32 { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_x_v4(this: &[u8]) -> u32 { + fn reduce_sum_of_x_as_u32_v4(this: &[u8]) -> u32 { use std::arch::x86_64::*; let us = _mm512_set1_epi32(255); let mut n = this.len(); @@ -546,7 +546,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_x_v4_test() { + fn reduce_sum_of_x_as_u32_v4_test() { use rand::Rng; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -558,7 +558,7 @@ mod reduce_sum_of_x { let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_v4(this) }; + let specialized = unsafe { reduce_sum_of_x_as_u32_v4(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } @@ -568,7 +568,7 @@ mod reduce_sum_of_x { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - fn reduce_sum_of_x_v3(this: &[u8]) -> u32 { + fn reduce_sum_of_x_as_u32_v3(this: &[u8]) -> u32 { use crate::emulate::emulate_mm256_reduce_add_epi32; use std::arch::x86_64::*; let us = _mm256_set1_epi32(255); @@ -594,7 +594,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_v3_test() { + fn reduce_sum_of_x_as_u32_v3_test() { use rand::Rng; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -606,7 +606,7 @@ mod reduce_sum_of_x { let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_v3(this) }; + let specialized = unsafe { reduce_sum_of_x_as_u32_v3(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } @@ -616,7 +616,7 @@ mod reduce_sum_of_x { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - fn reduce_sum_of_x_v2(this: &[u8]) -> u32 { + fn reduce_sum_of_x_as_u32_v2(this: &[u8]) -> u32 { use crate::emulate::emulate_mm_reduce_add_epi32; use std::arch::x86_64::*; let us = _mm_set1_epi32(255); @@ -642,7 +642,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_v2_test() { + fn reduce_sum_of_x_as_u32_v2_test() { use rand::Rng; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); @@ -654,7 +654,7 @@ mod reduce_sum_of_x { let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_v2(this) }; + let specialized = unsafe { reduce_sum_of_x_as_u32_v2(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } @@ -664,7 +664,7 @@ mod reduce_sum_of_x { #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] - fn reduce_sum_of_x_a2(this: &[u8]) -> u32 { + fn reduce_sum_of_x_as_u32_a2(this: &[u8]) -> u32 { use std::arch::aarch64::*; let mut n = this.len(); let mut a = this.as_ptr(); @@ -690,7 +690,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "aarch64", test))] #[test] - fn reduce_sum_of_x_a2_test() { + fn reduce_sum_of_x_as_u32_a2_test() { use rand::Rng; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -702,7 +702,7 @@ mod reduce_sum_of_x { let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_a2(this) }; + let specialized = unsafe { reduce_sum_of_x_as_u32_a2(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } @@ -710,7 +710,7 @@ mod reduce_sum_of_x { } #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] - pub fn reduce_sum_of_x(this: &[u8]) -> u32 { + pub fn reduce_sum_of_x_as_u32(this: &[u8]) -> u32 { let n = this.len(); let mut sum = 0; for i in 0..n { @@ -721,6 +721,6 @@ mod reduce_sum_of_x { } #[inline(always)] -pub fn reduce_sum_of_x(vector: &[u8]) -> u32 { - reduce_sum_of_x::reduce_sum_of_x(vector) +pub fn reduce_sum_of_x_as_u32(vector: &[u8]) -> u32 { + reduce_sum_of_x_as_u32::reduce_sum_of_x_as_u32(vector) } diff --git a/crates/simd_macros/Cargo.toml b/crates/simd_macros/Cargo.toml index b0d41b15..f206bff0 100644 --- a/crates/simd_macros/Cargo.toml +++ b/crates/simd_macros/Cargo.toml @@ -2,6 +2,7 @@ name = "simd_macros" version.workspace = true edition.workspace = true +publish = false [lib] proc-macro = true diff --git a/crates/vector/Cargo.toml b/crates/vector/Cargo.toml index 692bf642..d5dba499 100644 --- a/crates/vector/Cargo.toml +++ b/crates/vector/Cargo.toml @@ -2,6 +2,7 @@ name = "vector" version.workspace = true edition.workspace = true +publish = false [dependencies] distance = { path = "../distance" } diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs index 7877addb..c625e6e3 100644 --- a/crates/vector/src/bvect.rs +++ b/crates/vector/src/bvect.rs @@ -148,12 +148,12 @@ impl VectorBorrowed for BVectBorrowed<'_> { #[inline(always)] fn norm(&self) -> f32 { - (simd::bit::sum_of_x(self.data) as f32).sqrt() + (simd::bit::reduce_sum_of_x(self.data) as f32).sqrt() } #[inline(always)] fn operator_dot(self, rhs: Self) -> Distance { - Distance::from(-(simd::bit::sum_of_and(self.data, rhs.data) as f32)) + Distance::from(-(simd::bit::reduce_sum_of_and(self.data, rhs.data) as f32)) } #[inline(always)] @@ -168,12 +168,12 @@ impl VectorBorrowed for BVectBorrowed<'_> { #[inline(always)] fn operator_hamming(self, rhs: Self) -> Distance { - Distance::from(simd::bit::sum_of_xor(self.data, rhs.data) as f32) + Distance::from(simd::bit::reduce_sum_of_xor(self.data, rhs.data) as f32) } #[inline(always)] fn operator_jaccard(self, rhs: Self) -> Distance { - let (and, or) = simd::bit::sum_of_and_or(self.data, rhs.data); + let (and, or) = simd::bit::reduce_sum_of_and_or(self.data, rhs.data); Distance::from(1.0 - (and as f32 / or as f32)) } diff --git a/crates/vector/src/scalar8.rs b/crates/vector/src/scalar8.rs index 3792e270..21cccbbc 100644 --- a/crates/vector/src/scalar8.rs +++ b/crates/vector/src/scalar8.rs @@ -184,20 +184,22 @@ impl VectorBorrowed for Scalar8Borrowed<'_> { #[inline(always)] fn operator_dot(self, rhs: Self) -> Distance { assert_eq!(self.code.len(), rhs.code.len()); - let xy = self.k * rhs.k * simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 - + self.b * rhs.b * self.code.len() as f32 - + self.k * rhs.b * self.sum_of_code - + self.b * rhs.k * rhs.sum_of_code; + let xy = + self.k * rhs.k * simd::u8::reduce_sum_of_x_as_u32_y_as_u32(self.code, rhs.code) as f32 + + self.b * rhs.b * self.code.len() as f32 + + self.k * rhs.b * self.sum_of_code + + self.b * rhs.k * rhs.sum_of_code; Distance::from(-xy) } #[inline(always)] fn operator_l2(self, rhs: Self) -> Distance { assert_eq!(self.code.len(), rhs.code.len()); - let xy = self.k * rhs.k * simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 - + self.b * rhs.b * self.code.len() as f32 - + self.k * rhs.b * self.sum_of_code - + self.b * rhs.k * rhs.sum_of_code; + let xy = + self.k * rhs.k * simd::u8::reduce_sum_of_x_as_u32_y_as_u32(self.code, rhs.code) as f32 + + self.b * rhs.b * self.code.len() as f32 + + self.k * rhs.b * self.sum_of_code + + self.b * rhs.k * rhs.sum_of_code; let x2 = self.sum_of_x2; let y2 = rhs.sum_of_x2; Distance::from(x2 + y2 - 2.0 * xy) @@ -206,10 +208,11 @@ impl VectorBorrowed for Scalar8Borrowed<'_> { #[inline(always)] fn operator_cos(self, rhs: Self) -> Distance { assert_eq!(self.code.len(), rhs.code.len()); - let xy = self.k * rhs.k * simd::u8::reduce_sum_of_xy(self.code, rhs.code) as f32 - + self.b * rhs.b * self.code.len() as f32 - + self.k * rhs.b * self.sum_of_code - + self.b * rhs.k * rhs.sum_of_code; + let xy = + self.k * rhs.k * simd::u8::reduce_sum_of_x_as_u32_y_as_u32(self.code, rhs.code) as f32 + + self.b * rhs.b * self.code.len() as f32 + + self.k * rhs.b * self.sum_of_code + + self.b * rhs.k * rhs.sum_of_code; let x2 = self.sum_of_x2; let y2 = rhs.sum_of_x2; Distance::from(1.0 - xy / (x2 * y2).sqrt()) @@ -281,7 +284,7 @@ impl VectorBorrowed for Scalar8Borrowed<'_> { }, self.k, self.b, - simd::u8::reduce_sum_of_x(code) as f32, + simd::u8::reduce_sum_of_x_as_u32(code) as f32, code.to_owned(), ) } diff --git a/src/datatype/functions_scalar8.rs b/src/datatype/functions_scalar8.rs index bc73bd81..7cc7b42c 100644 --- a/src/datatype/functions_scalar8.rs +++ b/src/datatype/functions_scalar8.rs @@ -11,7 +11,7 @@ fn _vchord_vector_quantize_to_scalar8(vector: VectorInput) -> Scalar8Output { let sum_of_x2 = f32::reduce_sum_of_x2(vector.slice()); let (k, b, code) = simd::quantize::quantize(f32::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); - let sum_of_code = simd::u8::reduce_sum_of_x(&code) as f32; + let sum_of_code = simd::u8::reduce_sum_of_x_as_u32(&code) as f32; Scalar8Output::new(Scalar8Borrowed::new(sum_of_x2, k, b, sum_of_code, &code)) } @@ -21,6 +21,6 @@ fn _vchord_halfvec_quantize_to_scalar8(vector: HalfvecInput) -> Scalar8Output { let sum_of_x2 = f16::reduce_sum_of_x2(vector.slice()); let (k, b, code) = simd::quantize::quantize(f16::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); - let sum_of_code = simd::u8::reduce_sum_of_x(&code) as f32; + let sum_of_code = simd::u8::reduce_sum_of_x_as_u32(&code) as f32; Scalar8Output::new(Scalar8Borrowed::new(sum_of_x2, k, b, sum_of_code, &code)) } diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index c4a8f0c6..fd363256 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -2,9 +2,10 @@ use super::opclass::Opfamily; use crate::index::am::am_build::InternalBuild; use algorithm::operator::{Dot, L2, Op}; use algorithm::types::*; -use algorithm::{Bump, RelationRead, RelationWrite}; +use algorithm::*; use half::f16; use std::cell::UnsafeCell; +use std::collections::BinaryHeap; use std::mem::MaybeUninit; use std::num::NonZero; use vector::VectorOwned; @@ -188,24 +189,28 @@ impl Bump for BumpAlloc { } } -pub fn prewarm( - opfamily: Opfamily, - index: impl RelationRead, - height: i32, - check: impl Fn(), -) -> String { +pub fn prewarm(opfamily: Opfamily, index: &impl RelationRead, height: i32) -> String { + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; let message = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::prewarm::<_, Op, L2>>(index, height, check) + algorithm::prewarm::<_, Op, L2>>(index, height, make_h0_plain_prefetcher) } (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::prewarm::<_, Op, Dot>>(index, height, check) + algorithm::prewarm::<_, Op, Dot>>( + index, + height, + make_h0_plain_prefetcher, + ) } (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::prewarm::<_, Op, L2>>(index, height, check) + algorithm::prewarm::<_, Op, L2>>(index, height, make_h0_plain_prefetcher) } (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::prewarm::<_, Op, Dot>>(index, height, check) + algorithm::prewarm::<_, Op, Dot>>( + index, + height, + make_h0_plain_prefetcher, + ) } }; match message { @@ -216,39 +221,48 @@ pub fn prewarm( pub fn bulkdelete( opfamily: Opfamily, - index: impl RelationRead + RelationWrite, + index: &(impl RelationRead + RelationWrite), check: impl Fn(), callback: impl Fn(NonZero) -> bool, ) { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::bulkdelete::, L2>>(index, check, callback) + algorithm::bulkdelete::<_, Op, L2>>(index, check, callback) } (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::bulkdelete::, Dot>>(index, check, callback) + algorithm::bulkdelete::<_, Op, Dot>>(index, check, callback) } (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::bulkdelete::, L2>>(index, check, callback) + algorithm::bulkdelete::<_, Op, L2>>(index, check, callback) } (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::bulkdelete::, Dot>>(index, check, callback) + algorithm::bulkdelete::<_, Op, Dot>>(index, check, callback) } } } -pub fn maintain(opfamily: Opfamily, index: impl RelationRead + RelationWrite, check: impl Fn()) { +pub fn maintain(opfamily: Opfamily, index: &(impl RelationRead + RelationWrite), check: impl Fn()) { + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::maintain::, L2>, _>(index, check) + algorithm::maintain::<_, Op, L2>>(index, make_h0_plain_prefetcher, check) } (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::maintain::, Dot>, _>(index, check) + algorithm::maintain::<_, Op, Dot>>( + index, + make_h0_plain_prefetcher, + check, + ) } (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::maintain::, L2>, _>(index, check) + algorithm::maintain::<_, Op, L2>>(index, make_h0_plain_prefetcher, check) } (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::maintain::, Dot>, _>(index, check) + algorithm::maintain::<_, Op, Dot>>( + index, + make_h0_plain_prefetcher, + check, + ) } } } @@ -256,29 +270,29 @@ pub fn maintain(opfamily: Opfamily, index: impl RelationRead + RelationWrite, ch pub fn build( vector_options: VectorOptions, vchordrq_options: VchordrqIndexOptions, - index: impl RelationWrite, + index: &impl RelationWrite, structures: Vec>>, ) { match (vector_options.v, vector_options.d) { - (VectorKind::Vecf32, DistanceKind::L2) => algorithm::build::, L2>>( + (VectorKind::Vecf32, DistanceKind::L2) => algorithm::build::<_, Op, L2>>( vector_options, vchordrq_options, index, map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), ), - (VectorKind::Vecf32, DistanceKind::Dot) => algorithm::build::, Dot>>( + (VectorKind::Vecf32, DistanceKind::Dot) => algorithm::build::<_, Op, Dot>>( vector_options, vchordrq_options, index, map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), ), - (VectorKind::Vecf16, DistanceKind::L2) => algorithm::build::, L2>>( + (VectorKind::Vecf16, DistanceKind::L2) => algorithm::build::<_, Op, L2>>( vector_options, vchordrq_options, index, map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), ), - (VectorKind::Vecf16, DistanceKind::Dot) => algorithm::build::, Dot>>( + (VectorKind::Vecf16, DistanceKind::Dot) => algorithm::build::<_, Op, Dot>>( vector_options, vchordrq_options, index, @@ -289,55 +303,51 @@ pub fn build( pub fn insert( opfamily: Opfamily, - index: impl RelationRead + RelationWrite, + index: &(impl RelationRead + RelationWrite), payload: NonZero, vector: OwnedVector, ) { - use algorithm::{FastHeap, PlainPrefetcher}; let bump = BumpAlloc::new(); - let prefetch = { - let index = index.clone(); - move |results| PlainPrefetcher::<_, FastHeap<_>>::new(index.clone(), results) - }; + let make_h1_plain_prefetcher = MakeH1PlainPrefetcherForInsertion { index }; match (vector, opfamily.distance_kind()) { (OwnedVector::Vecf32(vector), DistanceKind::L2) => { assert!(opfamily.vector_kind() == VectorKind::Vecf32); - algorithm::insert::<_, Op, L2>, _>( + algorithm::insert::<_, Op, L2>>( index, payload, RandomProject::project(vector.as_borrowed()), &bump, - prefetch, + make_h1_plain_prefetcher, ) } (OwnedVector::Vecf32(vector), DistanceKind::Dot) => { assert!(opfamily.vector_kind() == VectorKind::Vecf32); - algorithm::insert::<_, Op, Dot>, _>( + algorithm::insert::<_, Op, Dot>>( index, payload, RandomProject::project(vector.as_borrowed()), &bump, - prefetch, + make_h1_plain_prefetcher, ) } (OwnedVector::Vecf16(vector), DistanceKind::L2) => { assert!(opfamily.vector_kind() == VectorKind::Vecf16); - algorithm::insert::<_, Op, L2>, _>( + algorithm::insert::<_, Op, L2>>( index, payload, RandomProject::project(vector.as_borrowed()), &bump, - prefetch, + make_h1_plain_prefetcher, ) } (OwnedVector::Vecf16(vector), DistanceKind::Dot) => { assert!(opfamily.vector_kind() == VectorKind::Vecf16); - algorithm::insert::<_, Op, Dot>, _>( + algorithm::insert::<_, Op, Dot>>( index, payload, RandomProject::project(vector.as_borrowed()), &bump, - prefetch, + make_h1_plain_prefetcher, ) } } @@ -379,6 +389,138 @@ impl RandomProject for VectBorrowed<'_, f16> { // Emulate unstable library feature `abort_unwind`. // See https://github.com/rust-lang/rust/issues/130338. +#[inline(never)] extern "C" fn abort_unwind R, R>(f: F) -> R { f() } + +#[derive(Debug)] +pub struct MakeH1PlainPrefetcherForInsertion<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakeH1PlainPrefetcherForInsertion<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcherForInsertion<'r, R> { + type P + = PlainPrefetcher<'r, R, FastHeap> + where + T: Ord + Fetch + 'r; + + fn prefetch(&mut self, seq: Vec) -> Self::P + where + T: Ord + Fetch + 'r, + { + PlainPrefetcher::new(self.index, FastHeap::from(seq)) + } +} + +#[derive(Debug)] +pub struct MakeH1PlainPrefetcher<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakeH1PlainPrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcher<'r, R> { + type P + = PlainPrefetcher<'r, R, BinaryHeap> + where + T: Ord + Fetch + 'r; + + fn prefetch(&mut self, seq: Vec) -> Self::P + where + T: Ord + Fetch + 'r, + { + PlainPrefetcher::new(self.index, BinaryHeap::from(seq)) + } +} + +#[derive(Debug)] +pub struct MakeH0PlainPrefetcher<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakeH0PlainPrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakeH0PlainPrefetcher<'r, R> { + type P + = PlainPrefetcher<'r, R, S> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch, + { + PlainPrefetcher::new(self.index, seq) + } +} + +#[derive(Debug)] +pub struct MakeH0SimplePrefetcher<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakeH0SimplePrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> + for MakeH0SimplePrefetcher<'r, R> +{ + type P + = SimplePrefetcher<'r, R, S> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch, + { + SimplePrefetcher::new(self.index, seq) + } +} + +#[derive(Debug)] +pub struct MakeH0StreamPrefetcher<'r, R> { + pub index: &'r R, + pub hints: Hints, +} + +impl<'r, R> Clone for MakeH0StreamPrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { + index: self.index, + hints: self.hints.clone(), + } + } +} + +impl<'r, R: RelationReadStream> PrefetcherSequenceFamily<'r, R> for MakeH0StreamPrefetcher<'r, R> { + type P + = StreamPrefetcher<'r, R, S> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch, + { + StreamPrefetcher::new(self.index, seq, self.hints.clone()) + } +} diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 1478d740..8ec98d66 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -314,15 +314,10 @@ pub unsafe extern "C-unwind" fn ambuild( } }; reporter.phase(BuildPhase::from_code(BuildPhaseCode::Build)); - crate::index::algorithm::build( - vector_options, - vchordrq_options.index, - index.clone(), - structures, - ); + crate::index::algorithm::build(vector_options, vchordrq_options.index, &index, structures); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); let cache = if vchordrq_options.build.pin { - let mut trace = algorithm::cache(index.clone()); + let mut trace = algorithm::cache(&index); trace.sort(); trace.dedup(); if let Some(max) = trace.last().copied() { @@ -384,7 +379,7 @@ pub unsafe extern "C-unwind" fn ambuild( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); + crate::index::algorithm::insert(opfamily, &index, payload, vector); } indtuples += 1; reporter.tuples_done(indtuples); @@ -395,7 +390,7 @@ pub unsafe extern "C-unwind" fn ambuild( pgrx::check_for_interrupts!(); }; reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); - crate::index::algorithm::maintain(opfamily, index, check); + crate::index::algorithm::maintain(opfamily, &index, check); unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -866,7 +861,7 @@ unsafe fn parallel_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); + crate::index::algorithm::insert(opfamily, &index, payload, vector); } unsafe { let indtuples; @@ -889,7 +884,7 @@ unsafe fn parallel_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); + crate::index::algorithm::insert(opfamily, &index, payload, vector); } unsafe { let indtuples; diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index b1922cd1..7ca1754d 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -194,7 +194,7 @@ pub unsafe extern "C-unwind" fn amcostestimate( } let index = PostgresRelation::new(relation.raw()); let probes = gucs::probes(); - let cost = algorithm::cost(index); + let cost = algorithm::cost(&index); if cost.cells.len() != 1 + probes.len() { panic!( "need {} probes, but {} probes provided", @@ -287,7 +287,7 @@ unsafe fn aminsertinner( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, index.clone(), payload, vector); + crate::index::algorithm::insert(opfamily, &index, payload, vector); } } false @@ -317,7 +317,7 @@ pub unsafe extern "C-unwind" fn ambulkdelete( let mut ctid = key_to_ctid(key); unsafe { callback(&mut ctid, callback_state) } }; - crate::index::algorithm::bulkdelete(opfamily, index, check, callback); + crate::index::algorithm::bulkdelete(opfamily, &index, check, callback); stats } @@ -337,7 +337,7 @@ pub unsafe extern "C-unwind" fn amvacuumcleanup( let check = || unsafe { pgrx::pg_sys::vacuum_delay_point(); }; - crate::index::algorithm::maintain(opfamily, index, check); + crate::index::algorithm::maintain(opfamily, &index, check); stats } @@ -393,6 +393,7 @@ pub unsafe extern "C-unwind" fn amrescan( max_scan_tuples: gucs::max_scan_tuples(), maxsim_refine: gucs::maxsim_refine(), maxsim_threshold: gucs::maxsim_threshold(), + io_search: gucs::io_search(), io_rerank: gucs::io_rerank(), }; let fetcher = { diff --git a/src/index/functions.rs b/src/index/functions.rs index 8a0bef49..50026bf5 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -21,9 +21,7 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { let relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); let opfamily = unsafe { crate::index::opclass::opfamily(relation.raw()) }; let index = unsafe { PostgresRelation::new(relation.raw()) }; - crate::index::algorithm::prewarm(opfamily, index, height, || { - pgrx::check_for_interrupts!(); - }) + crate::index::algorithm::prewarm(opfamily, &index, height) } struct Index { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 18fd1e67..4f66dfbd 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -1,11 +1,11 @@ -use super::scanners::SearchIo; +use super::scanners::Io; use pgrx::PostgresGucEnum; use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; use std::ffi::CStr; #[allow(non_camel_case_types)] #[derive(Debug, Clone, Copy, PostgresGucEnum)] -pub enum Io { +pub enum PostgresIo { read_buffer, prefetch_buffer, #[cfg(feature = "pg17")] @@ -24,11 +24,18 @@ static MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); static PRERERANK_FILTERING: GucSetting = GucSetting::::new(false); -static IO_RERANK: GucSetting = GucSetting::::new( +static IO_SEARCH: GucSetting = GucSetting::::new( #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] - Io::prefetch_buffer, + PostgresIo::prefetch_buffer, #[cfg(feature = "pg17")] - Io::read_stream, + PostgresIo::read_stream, +); + +static IO_RERANK: GucSetting = GucSetting::::new( + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + PostgresIo::prefetch_buffer, + #[cfg(feature = "pg17")] + PostgresIo::read_stream, ); pub fn init() { @@ -96,6 +103,14 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_enum_guc( + "vchordrq.io_search", + "`io_search` argument of vchordrq.", + "`io_search` argument of vchordrq.", + &IO_SEARCH, + GucContext::Userset, + GucFlags::default(), + ); GucRegistry::define_enum_guc( "vchordrq.io_rerank", "`io_rerank` argument of vchordrq.", @@ -182,11 +197,20 @@ pub fn prererank_filtering() -> bool { PRERERANK_FILTERING.get() } -pub fn io_rerank() -> SearchIo { +pub fn io_search() -> Io { + match IO_RERANK.get() { + PostgresIo::read_buffer => Io::Plain, + PostgresIo::prefetch_buffer => Io::Simple, + #[cfg(feature = "pg17")] + PostgresIo::read_stream => Io::Stream, + } +} + +pub fn io_rerank() -> Io { match IO_RERANK.get() { - Io::read_buffer => SearchIo::ReadBuffer, - Io::prefetch_buffer => SearchIo::PrefetchBuffer, + PostgresIo::read_buffer => Io::Plain, + PostgresIo::prefetch_buffer => Io::Simple, #[cfg(feature = "pg17")] - Io::read_stream => SearchIo::ReadStream, + PostgresIo::read_stream => Io::Stream, } } diff --git a/src/index/scanners/default.rs b/src/index/scanners/default.rs index fa48cc4a..b994715e 100644 --- a/src/index/scanners/default.rs +++ b/src/index/scanners/default.rs @@ -1,5 +1,5 @@ -use super::{SearchBuilder, SearchFetcher, SearchIo, SearchOptions}; -use crate::index::algorithm::RandomProject; +use super::{Io, SearchBuilder, SearchFetcher, SearchOptions}; +use crate::index::algorithm::*; use crate::index::am::pointer_to_kv; use crate::index::opclass::{Opfamily, Sphere}; use algorithm::operator::{Dot, L2, Op}; @@ -49,13 +49,16 @@ impl SearchBuilder for DefaultBuilder { } } - fn build<'a>( + fn build<'a, R>( self, - relation: &'a (impl RelationPrefetch + RelationReadStream), + index: &'a R, options: SearchOptions, mut fetcher: impl SearchFetcher + 'a, bump: &'a impl Bump, - ) -> Box + 'a> { + ) -> Box + 'a> + where + R: RelationRead + RelationPrefetch + RelationReadStream, + { let mut vector = None; let mut threshold = None; let mut recheck = false; @@ -77,6 +80,13 @@ impl SearchBuilder for DefaultBuilder { let Some(vector) = vector else { return Box::new(std::iter::empty()) as Box>; }; + let make_h1_plain_prefetcher = MakeH1PlainPrefetcher { index }; + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; + let make_h0_simple_prefetcher = MakeH0SimplePrefetcher { index }; + let make_h0_stream_prefetcher = MakeH0StreamPrefetcher { + index, + hints: Hints::default().full(true), + }; let iter: Box)>> = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { @@ -88,19 +98,35 @@ impl SearchBuilder for DefaultBuilder { } .as_borrowed(), ); - let results = default_search::<_, Op, L2>, _>( - relation.clone(), - vector.clone(), - options.probes, - options.epsilon, - bump, - { - let index = relation.clone(); - move |results| { - PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) - } - }, - ); + let results = match options.io_search { + Io::Plain => default_search::<_, Op, L2>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op, L2>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op, L2>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); let (datums, is_nulls) = fetcher.fetch(key)?; @@ -113,11 +139,11 @@ impl SearchBuilder for DefaultBuilder { }; Some(RandomProject::project(raw.as_borrowed())) }; - let method = how(relation.clone()); + let method = how(index); match (method, options.io_rerank) { - (RerankMethod::Index, SearchIo::ReadBuffer) => { + (RerankMethod::Index, Io::Plain) => { let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_index::, L2>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -125,8 +151,11 @@ impl SearchBuilder for DefaultBuilder { }), ) } - (RerankMethod::Index, SearchIo::PrefetchBuffer) => { - let prefetcher = SimplePrefetcher::new(relation.clone(), results); + (RerankMethod::Index, Io::Simple) => { + let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( + index, + results.into(), + ); Box::new( rerank_index::, L2>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -134,8 +163,12 @@ impl SearchBuilder for DefaultBuilder { }), ) } - (RerankMethod::Index, SearchIo::ReadStream) => { - let prefetcher = StreamPrefetcher::new(relation, results); + (RerankMethod::Index, Io::Stream) => { + let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( + index, + results.into(), + Hints::default(), + ); Box::new( rerank_index::, L2>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -145,7 +178,7 @@ impl SearchBuilder for DefaultBuilder { } (RerankMethod::Heap, _) => { let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_heap::, L2>, _, _>( vector, prefetcher, fetch, @@ -166,19 +199,35 @@ impl SearchBuilder for DefaultBuilder { } .as_borrowed(), ); - let results = default_search::<_, Op, Dot>, _>( - relation.clone(), - vector.clone(), - options.probes, - options.epsilon, - bump, - { - let index = relation.clone(); - move |results| { - PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) - } - }, - ); + let results = match options.io_search { + Io::Plain => default_search::<_, Op, Dot>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op, Dot>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op, Dot>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); let (datums, is_nulls) = fetcher.fetch(key)?; @@ -191,11 +240,11 @@ impl SearchBuilder for DefaultBuilder { }; Some(RandomProject::project(raw.as_borrowed())) }; - let method = how(relation.clone()); + let method = how(index); match (method, options.io_rerank) { - (RerankMethod::Index, SearchIo::ReadBuffer) => { + (RerankMethod::Index, Io::Plain) => { let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_index::, Dot>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -203,8 +252,9 @@ impl SearchBuilder for DefaultBuilder { }), ) } - (RerankMethod::Index, SearchIo::PrefetchBuffer) => { - let prefetcher = SimplePrefetcher::new(relation.clone(), results); + (RerankMethod::Index, Io::Simple) => { + let prefetcher = + SimplePrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_index::, Dot>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -212,8 +262,12 @@ impl SearchBuilder for DefaultBuilder { }), ) } - (RerankMethod::Index, SearchIo::ReadStream) => { - let prefetcher = StreamPrefetcher::new(relation, results); + (RerankMethod::Index, Io::Stream) => { + let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( + index, + results.into(), + Hints::default(), + ); Box::new( rerank_index::, Dot>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -223,7 +277,7 @@ impl SearchBuilder for DefaultBuilder { } (RerankMethod::Heap, _) => { let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_heap::, Dot>, _, _>( vector, prefetcher, fetch, @@ -244,19 +298,35 @@ impl SearchBuilder for DefaultBuilder { } .as_borrowed(), ); - let results = default_search::<_, Op, L2>, _>( - relation.clone(), - vector.clone(), - options.probes, - options.epsilon, - bump, - { - let index = relation.clone(); - move |results| { - PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) - } - }, - ); + let results = match options.io_search { + Io::Plain => default_search::<_, Op, L2>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op, L2>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op, L2>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); let (datums, is_nulls) = fetcher.fetch(key)?; @@ -269,11 +339,11 @@ impl SearchBuilder for DefaultBuilder { }; Some(RandomProject::project(raw.as_borrowed())) }; - let method = how(relation.clone()); + let method = how(index); match (method, options.io_rerank) { - (RerankMethod::Index, SearchIo::ReadBuffer) => { + (RerankMethod::Index, Io::Plain) => { let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_index::, L2>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -281,8 +351,9 @@ impl SearchBuilder for DefaultBuilder { }), ) } - (RerankMethod::Index, SearchIo::PrefetchBuffer) => { - let prefetcher = SimplePrefetcher::new(relation.clone(), results); + (RerankMethod::Index, Io::Simple) => { + let prefetcher = + SimplePrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_index::, L2>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -290,8 +361,12 @@ impl SearchBuilder for DefaultBuilder { }), ) } - (RerankMethod::Index, SearchIo::ReadStream) => { - let prefetcher = StreamPrefetcher::new(relation, results); + (RerankMethod::Index, Io::Stream) => { + let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( + index, + results.into(), + Hints::default(), + ); Box::new( rerank_index::, L2>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -301,7 +376,7 @@ impl SearchBuilder for DefaultBuilder { } (RerankMethod::Heap, _) => { let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_heap::, L2>, _, _>( vector, prefetcher, fetch, @@ -322,19 +397,35 @@ impl SearchBuilder for DefaultBuilder { } .as_borrowed(), ); - let results = default_search::<_, Op, Dot>, _>( - relation.clone(), - vector.clone(), - options.probes, - options.epsilon, - bump, - { - let index = relation.clone(); - move |results| { - PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) - } - }, - ); + let results = match options.io_search { + Io::Plain => default_search::<_, Op, Dot>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op, Dot>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op, Dot>>( + index, + vector.clone(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); let (datums, is_nulls) = fetcher.fetch(key)?; @@ -347,11 +438,11 @@ impl SearchBuilder for DefaultBuilder { }; Some(RandomProject::project(raw.as_borrowed())) }; - let method = how(relation.clone()); + let method = how(index); match (method, options.io_rerank) { - (RerankMethod::Index, SearchIo::ReadBuffer) => { + (RerankMethod::Index, Io::Plain) => { let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_index::, Dot>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -359,8 +450,9 @@ impl SearchBuilder for DefaultBuilder { }), ) } - (RerankMethod::Index, SearchIo::PrefetchBuffer) => { - let prefetcher = SimplePrefetcher::new(relation.clone(), results); + (RerankMethod::Index, Io::Simple) => { + let prefetcher = + SimplePrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_index::, Dot>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -368,8 +460,12 @@ impl SearchBuilder for DefaultBuilder { }), ) } - (RerankMethod::Index, SearchIo::ReadStream) => { - let prefetcher = StreamPrefetcher::new(relation, results); + (RerankMethod::Index, Io::Stream) => { + let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( + index, + results.into(), + Hints::default(), + ); Box::new( rerank_index::, Dot>, _, _>(vector, prefetcher) .map(move |(distance, payload)| { @@ -379,7 +475,7 @@ impl SearchBuilder for DefaultBuilder { } (RerankMethod::Heap, _) => { let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(relation.clone(), results); + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( rerank_heap::, Dot>, _, _>( vector, prefetcher, fetch, diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs index 78f41f69..6494195b 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/scanners/maxsim.rs @@ -1,8 +1,8 @@ use super::{SearchBuilder, SearchFetcher, SearchOptions}; -use crate::index::algorithm::RandomProject; +use crate::index::algorithm::{RandomProject, *}; use crate::index::am::pointer_to_kv; use crate::index::opclass::Opfamily; -use crate::index::scanners::SearchIo; +use crate::index::scanners::Io; use algorithm::operator::Dot; use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; use algorithm::*; @@ -42,13 +42,16 @@ impl SearchBuilder for MaxsimBuilder { } } - fn build<'a>( + fn build<'a, R>( self, - relation: &'a (impl RelationPrefetch + RelationReadStream), + index: &'a R, options: SearchOptions, - _: impl SearchFetcher + 'a, + _fetcher: impl SearchFetcher + 'a, bump: &'a impl Bump, - ) -> Box + 'a> { + ) -> Box + 'a> + where + R: RelationRead + RelationPrefetch + RelationReadStream, + { let mut vectors = None; for orderby_vectors in self.orderbys.into_iter().flatten() { if vectors.is_none() { @@ -66,11 +69,18 @@ impl SearchBuilder for MaxsimBuilder { let Some(vectors) = vectors else { return Box::new(std::iter::empty()) as Box>; }; - let method = how(relation.clone()); + let method = how(index); if !matches!(method, RerankMethod::Index) { pgrx::error!("maxsim search with rerank_in_table is not supported"); } assert!(matches!(opfamily.distance_kind(), DistanceKind::Dot)); + let make_h1_plain_prefetcher = MakeH1PlainPrefetcher { index }; + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; + let make_h0_simple_prefetcher = MakeH0SimplePrefetcher { index }; + let make_h0_stream_prefetcher = MakeH0StreamPrefetcher { + index, + hints: Hints::default().full(true), + }; let n = vectors.len(); let accu_map = |(Reverse(distance), AlwaysEqual(payload))| (distance, payload); let rough_map = |((_, AlwaysEqual(rough)), AlwaysEqual(&mut (payload, ..))): ( @@ -94,28 +104,44 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new(vectors.into_iter().map(|vector| { - let (results, estimation_by_threshold) = maxsim_search::<_, Op, _>( - relation.clone(), - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - { - let index = relation.clone(); - move |results| { - PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) - } - }, - ); + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { match options.io_rerank { - SearchIo::ReadBuffer => { - let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new( - relation.clone(), - results, - ); + Io::Plain => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); let mut reranker = rerank_index::(vector.clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); @@ -123,8 +149,11 @@ impl SearchBuilder for MaxsimBuilder { accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } - SearchIo::PrefetchBuffer => { - let prefetcher = SimplePrefetcher::new(relation.clone(), results); + Io::Simple => { + let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( + index, + results.into(), + ); let mut reranker = rerank_index::(vector.clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); @@ -132,8 +161,12 @@ impl SearchBuilder for MaxsimBuilder { accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } - SearchIo::ReadStream => { - let prefetcher = StreamPrefetcher::new(relation, results); + Io::Stream => { + let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( + index, + results.into(), + Hints::default(), + ); let mut reranker = rerank_index::(vector.clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); @@ -165,28 +198,44 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new(vectors.into_iter().map(|vector| { - let (results, estimation_by_threshold) = maxsim_search::<_, Op, _>( - relation.clone(), - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - { - let index = relation.clone(); - move |results| { - PlainPrefetcher::<_, BinaryHeap<_>>::new(index.clone(), results) - } - }, - ); + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { match options.io_rerank { - SearchIo::ReadBuffer => { - let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new( - relation.clone(), - results, - ); + Io::Plain => { + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); let mut reranker = rerank_index::(vector.clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); @@ -194,8 +243,11 @@ impl SearchBuilder for MaxsimBuilder { accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } - SearchIo::PrefetchBuffer => { - let prefetcher = SimplePrefetcher::new(relation.clone(), results); + Io::Simple => { + let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( + index, + results.into(), + ); let mut reranker = rerank_index::(vector.clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); @@ -203,8 +255,12 @@ impl SearchBuilder for MaxsimBuilder { accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } - SearchIo::ReadStream => { - let prefetcher = StreamPrefetcher::new(relation, results); + Io::Stream => { + let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( + index, + results.into(), + Hints::default(), + ); let mut reranker = rerank_index::(vector.clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs index c5fea0d9..b08b13e5 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/scanners/mod.rs @@ -3,21 +3,21 @@ mod maxsim; use super::opclass::Opfamily; use crate::index::lazy_cell::LazyCell; -use algorithm::{Bump, RelationPrefetch, RelationReadStream}; +use algorithm::{Bump, RelationPrefetch, RelationRead, RelationReadStream}; use pgrx::pg_sys::Datum; pub use default::DefaultBuilder; pub use maxsim::MaxsimBuilder; #[derive(Debug, Clone, Copy)] -pub enum SearchIo { - ReadBuffer, - PrefetchBuffer, +pub enum Io { + Plain, + Simple, #[cfg_attr( any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"), expect(dead_code) )] - ReadStream, + Stream, } #[derive(Debug)] @@ -27,7 +27,8 @@ pub struct SearchOptions { pub max_scan_tuples: Option, pub maxsim_refine: u32, pub maxsim_threshold: u32, - pub io_rerank: SearchIo, + pub io_search: Io, + pub io_rerank: Io, } pub trait SearchBuilder: 'static { @@ -35,13 +36,15 @@ pub trait SearchBuilder: 'static { unsafe fn add(&mut self, strategy: u16, datum: Option); - fn build<'a>( + fn build<'a, R>( self, - relation: &'a (impl RelationPrefetch + RelationReadStream), + relation: &'a R, options: SearchOptions, fetcher: impl SearchFetcher + 'a, bump: &'a impl Bump, - ) -> Box + 'a>; + ) -> Box + 'a> + where + R: RelationRead + RelationPrefetch + RelationReadStream; } pub trait SearchFetcher { diff --git a/src/index/storage.rs b/src/index/storage.rs index 8008aa2c..432cd44f 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -435,6 +435,19 @@ where } vec_deque_pop_front_if(&mut self.window, predicate) } + #[allow(dead_code)] + pub fn pop_item(&mut self) -> Option { + while self.window.is_empty() + && let Some(iter) = self.iter.as_mut() + && let Some(e) = iter.next() + { + for id in e.fetch().iter().copied() { + self.tail.push_back(id); + } + self.window.push_back(e); + } + self.window.pop_front() + } } pub struct PostgresReadStream { @@ -444,15 +457,44 @@ pub struct PostgresReadStream { cache: NonNull>, } -impl ReadStream for PostgresReadStream +#[cfg(feature = "pg17")] +impl PostgresReadStream +where + I::Item: Fetch, +{ + fn read(&mut self, fetch: &[u32]) -> impl Iterator { + fetch.iter().map(|_| unsafe { + use pgrx::pg_sys::{ + BUFFER_LOCK_SHARE, BufferGetPage, LockBuffer, read_stream_next_buffer, + }; + let buf = read_stream_next_buffer(self.raw, core::ptr::null_mut()); + LockBuffer(buf, BUFFER_LOCK_SHARE as _); + let page = NonNull::new(BufferGetPage(buf).cast()).expect("failed to get page"); + PostgresBufferReadGuard { + buf, + page, + id: pgrx::pg_sys::BufferGetBlockNumber(buf), + } + }) + } +} + +impl<'r, I: Iterator> ReadStream<'r> for PostgresReadStream where I::Item: Fetch, { type Relation = PostgresRelation; + type Item = I::Item; + type Inner = Chain, Flatten>>; + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + fn next(&mut self) -> Option<(I::Item, Vec)> { + panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); + } + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] fn next_if( &mut self, @@ -461,29 +503,23 @@ where panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); } + #[cfg(feature = "pg17")] + fn next(&mut self) -> Option<(I::Item, Vec)> { + if let Some(e) = unsafe { self.cache.as_mut().pop_item() } { + let list = self.read(e.fetch()).collect(); + Some((e, list)) + } else { + None + } + } + #[cfg(feature = "pg17")] fn next_if( &mut self, predicate: impl FnOnce(&I::Item) -> bool, ) -> Option<(I::Item, Vec)> { if let Some(e) = unsafe { self.cache.as_mut().pop_item_if(predicate) } { - let list = e - .fetch() - .iter() - .map(|_| unsafe { - use pgrx::pg_sys::{ - BUFFER_LOCK_SHARE, BufferGetPage, LockBuffer, read_stream_next_buffer, - }; - let buf = read_stream_next_buffer(self.raw, core::ptr::null_mut()); - LockBuffer(buf, BUFFER_LOCK_SHARE as _); - let page = NonNull::new(BufferGetPage(buf).cast()).expect("failed to get page"); - PostgresBufferReadGuard { - buf, - page, - id: pgrx::pg_sys::BufferGetBlockNumber(buf), - } - }) - .collect::>(); + let list = self.read(e.fetch()).collect(); Some((e, list)) } else { None @@ -519,7 +555,7 @@ impl RelationReadStream for PostgresRelation { I::Item: Fetch; #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] - fn read_stream(&self, _iter: I) -> Self::ReadStream<'_, I> + fn read_stream(&self, _iter: I, _hints: Hints) -> Self::ReadStream<'_, I> where I::Item: Fetch, { @@ -527,7 +563,7 @@ impl RelationReadStream for PostgresRelation { } #[cfg(feature = "pg17")] - fn read_stream(&self, iter: I) -> Self::ReadStream<'_, I> + fn read_stream(&self, iter: I, hints: Hints) -> Self::ReadStream<'_, I> where I::Item: Fetch, { @@ -552,9 +588,15 @@ impl RelationReadStream for PostgresRelation { iter: Some(iter), })); let raw = unsafe { - use pgrx::pg_sys::{ForkNumber, READ_STREAM_DEFAULT, read_stream_begin_relation}; + use pgrx::pg_sys::{ + ForkNumber, READ_STREAM_DEFAULT, READ_STREAM_FULL, read_stream_begin_relation, + }; + let mut flags = READ_STREAM_DEFAULT; + if hints.full { + flags |= READ_STREAM_FULL; + } read_stream_begin_relation( - READ_STREAM_DEFAULT as i32, + flags as i32, core::ptr::null_mut(), self.raw, ForkNumber::MAIN_FORKNUM, From fc2357470e0cda602acca1eef239008ac13fbaac Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 13 May 2025 12:18:38 +0800 Subject: [PATCH 155/324] chore: add license check and license header (#248) Signed-off-by: usamoi --- .github/workflows/check.yml | 68 +++- Cargo.lock | 303 ++++++++++-------- Cargo.toml | 4 +- build.rs | 22 +- crates/algorithm/Cargo.toml | 1 - crates/algorithm/src/build.rs | 14 + crates/algorithm/src/bulkdelete.rs | 14 + crates/algorithm/src/cache.rs | 14 + .../algorithm/src/closure_lifetime_binder.rs | 14 + crates/algorithm/src/cost.rs | 14 + crates/algorithm/src/fast_heap.rs | 16 +- crates/algorithm/src/freepages.rs | 14 + crates/algorithm/src/insert.rs | 14 + crates/algorithm/src/lib.rs | 14 + crates/algorithm/src/linked_vec.rs | 14 + crates/algorithm/src/maintain.rs | 14 + crates/algorithm/src/operator.rs | 14 + crates/algorithm/src/prefetcher.rs | 14 + crates/algorithm/src/prewarm.rs | 47 ++- crates/algorithm/src/rerank.rs | 14 + crates/algorithm/src/search.rs | 14 + crates/algorithm/src/tape.rs | 14 + crates/algorithm/src/tape_writer.rs | 14 + crates/algorithm/src/tuples.rs | 14 + crates/algorithm/src/types.rs | 14 + crates/algorithm/src/vectors.rs | 14 + crates/always_equal/src/lib.rs | 14 + crates/distance/src/lib.rs | 14 + crates/k_means/src/lib.rs | 14 + crates/rabitq/src/binary.rs | 14 + crates/rabitq/src/block.rs | 14 + crates/rabitq/src/lib.rs | 14 + crates/rabitq/src/packing.rs | 14 + crates/random_orthogonal_matrix/src/lib.rs | 14 + crates/simd/Cargo.toml | 2 +- crates/simd/build.rs | 17 +- crates/simd/cshim/aarch64.c | 14 + crates/simd/cshim/x86_64.c | 14 + crates/simd/src/aligned.rs | 14 + crates/simd/src/bit.rs | 14 + crates/simd/src/emulate.rs | 14 + crates/simd/src/f16.rs | 14 + crates/simd/src/f32.rs | 14 + crates/simd/src/fast_scan.rs | 14 + crates/simd/src/lib.rs | 14 + crates/simd/src/quantize.rs | 14 + crates/simd/src/u8.rs | 14 + crates/simd_macros/src/lib.rs | 14 + crates/simd_macros/src/target.rs | 14 + crates/vector/src/bvect.rs | 14 + crates/vector/src/lib.rs | 14 + crates/vector/src/scalar8.rs | 14 + crates/vector/src/svect.rs | 14 + crates/vector/src/vect.rs | 14 + deny.toml | 21 ++ rustfmt.toml | 1 - scripts/bench.py | 14 + scripts/dump.py | 14 + scripts/index.py | 14 + scripts/train.py | 14 + src/bin/pgrx_embed.rs | 14 + src/datatype/binary_scalar8.rs | 14 + src/datatype/functions_scalar8.rs | 14 + src/datatype/memory_halfvec.rs | 14 + src/datatype/memory_scalar8.rs | 14 + src/datatype/memory_vector.rs | 14 + src/datatype/mod.rs | 14 + src/datatype/operators_halfvec.rs | 14 + src/datatype/operators_scalar8.rs | 14 + src/datatype/operators_vector.rs | 14 + src/datatype/text_scalar8.rs | 14 + src/datatype/typmod.rs | 14 + src/index/algorithm.rs | 20 +- src/index/am/am_build.rs | 16 +- src/index/am/mod.rs | 14 + src/index/functions.rs | 14 + src/index/gucs.rs | 14 + src/index/hook.rs | 14 + src/index/lazy_cell.rs | 14 + src/index/mod.rs | 14 + src/index/opclass.rs | 14 + src/index/projection.rs | 14 + src/index/scanners/default.rs | 14 + src/index/scanners/maxsim.rs | 14 + src/index/scanners/mod.rs | 14 + src/index/storage.rs | 14 + src/index/types.rs | 14 + src/lib.rs | 14 + src/upgrade.rs | 14 + 89 files changed, 1436 insertions(+), 166 deletions(-) create mode 100644 deny.toml diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 2927d4a6..741de314 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -5,6 +5,7 @@ on: paths: - ".cargo" - ".github/workflows/check.yml" + - "deny.toml" - "crates/**" - "scripts/**" - "src/**" @@ -20,6 +21,7 @@ on: paths: - ".cargo" - ".github/workflows/check.yml" + - "deny.toml" - "crates/**" - "scripts/**" - "src/**" @@ -41,8 +43,6 @@ concurrency: env: CARGO_TERM_COLOR: always RUST_BACKTRACE: 1 - SCCACHE_GHA_ENABLED: true - RUSTC_WRAPPER: sccache RUSTFLAGS: "-Dwarnings" jobs: @@ -54,6 +54,8 @@ jobs: run: | curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo + curl -fsSL https://github.com/EmbarkStudios/cargo-deny/releases/download/0.18.2/cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl/cargo-deny | install -m 755 /dev/stdin /usr/local/bin/cargo-deny + - name: Checkout uses: actions/checkout@v4 @@ -67,7 +69,59 @@ jobs: uses: astral-sh/ruff-action@v1 - name: Rustfmt - run: cargo fmt --check + run: cargo fmt --check -- --config-path /dev/null --config imports_granularity=Module + + - name: Deny + run: cargo deny check + + - name: License Header + run: | + HEADER=$(cat < Result<(), Box> { + if var("CARGO_CFG_TARGET_OS")? == "macos" { println!("cargo::rustc-link-arg-cdylib=-Wl,-undefined,dynamic_lookup"); } + Ok(()) } diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml index 73c7b6c5..acbda149 100644 --- a/crates/algorithm/Cargo.toml +++ b/crates/algorithm/Cargo.toml @@ -18,7 +18,6 @@ paste.workspace = true pin-project = "1" rand.workspace = true serde.workspace = true -turboselect = { git = "https://github.com/tensorchord/turboselect.git", rev = "d8753c4ffe5b47f28670fea21d56cf3658d51b9b" } validator.workspace = true zerocopy.workspace = true diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 954e4cd7..f4534cf7 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::operator::{Accessor2, Operator, Vector}; use crate::tape::TapeWriter; use crate::tape_writer::H1TapeWriter; diff --git a/crates/algorithm/src/bulkdelete.rs b/crates/algorithm/src/bulkdelete.rs index 4fc1a87d..1f3fcf79 100644 --- a/crates/algorithm/src/bulkdelete.rs +++ b/crates/algorithm/src/bulkdelete.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::closure_lifetime_binder::{id_0, id_1}; use crate::operator::{FunctionalAccessor, Operator}; use crate::tape::by_next; diff --git a/crates/algorithm/src/cache.rs b/crates/algorithm/src/cache.rs index 00e938b2..ca183662 100644 --- a/crates/algorithm/src/cache.rs +++ b/crates/algorithm/src/cache.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::closure_lifetime_binder::{id_0, id_1}; use crate::operator::FunctionalAccessor; use crate::tape::by_next; diff --git a/crates/algorithm/src/closure_lifetime_binder.rs b/crates/algorithm/src/closure_lifetime_binder.rs index 6db1ca95..f24f4414 100644 --- a/crates/algorithm/src/closure_lifetime_binder.rs +++ b/crates/algorithm/src/closure_lifetime_binder.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + // Use stable language features as an alternative to `closure_lifetime_binder`. // See https://github.com/rust-lang/rust/issues/97362. diff --git a/crates/algorithm/src/cost.rs b/crates/algorithm/src/cost.rs index 32274ae6..f7ae8ab7 100644 --- a/crates/algorithm/src/cost.rs +++ b/crates/algorithm/src/cost.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::tuples::{MetaTuple, WithReader}; use crate::{Page, RelationRead}; diff --git a/crates/algorithm/src/fast_heap.rs b/crates/algorithm/src/fast_heap.rs index 3bd85ac0..18ef21b2 100644 --- a/crates/algorithm/src/fast_heap.rs +++ b/crates/algorithm/src/fast_heap.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::Sequence; use std::collections::BinaryHeap; use std::num::NonZero; @@ -24,7 +38,7 @@ impl FastHeap { if let Some(t) = NonZero::new(n / 384) { let mut inner = vec; let index = n - t.get(); - turboselect::select_nth_unstable(&mut inner, index); + inner.select_nth_unstable(index); inner[index..].sort_unstable(); Self::Sorted(SortHeap { inner, t }) } else { diff --git a/crates/algorithm/src/freepages.rs b/crates/algorithm/src/freepages.rs index 55574fdf..fb933be1 100644 --- a/crates/algorithm/src/freepages.rs +++ b/crates/algorithm/src/freepages.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::tuples::*; use crate::*; use std::cmp::Reverse; diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 52fdb3ff..4194eb54 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::linked_vec::LinkedVec; use crate::operator::*; use crate::tape::by_next; diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 185241d9..cc37b57b 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #![allow(clippy::type_complexity)] mod build; diff --git a/crates/algorithm/src/linked_vec.rs b/crates/algorithm/src/linked_vec.rs index 179a0489..650a962c 100644 --- a/crates/algorithm/src/linked_vec.rs +++ b/crates/algorithm/src/linked_vec.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + pub struct LinkedVec { inner: Vec>, last: Vec, diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index 5e4cb073..247f5bdd 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::closure_lifetime_binder::{id_0, id_1, id_2, id_3}; use crate::operator::{FunctionalAccessor, Operator, Vector}; use crate::tape::{self, TapeWriter, by_directory, by_next}; diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index 757df604..572442d7 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use distance::Distance; use half::f16; use rabitq::binary::{BinaryCode, BinaryLut}; diff --git a/crates/algorithm/src/prefetcher.rs b/crates/algorithm/src/prefetcher.rs index 7e88b423..ccffdebe 100644 --- a/crates/algorithm/src/prefetcher.rs +++ b/crates/algorithm/src/prefetcher.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::*; use std::collections::{VecDeque, vec_deque}; use std::iter::Chain; diff --git a/crates/algorithm/src/prewarm.rs b/crates/algorithm/src/prewarm.rs index 31d3e46c..ec7ba8a3 100644 --- a/crates/algorithm/src/prewarm.rs +++ b/crates/algorithm/src/prewarm.rs @@ -1,16 +1,29 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::closure_lifetime_binder::{id_0, id_1, id_2}; use crate::operator::{FunctionalAccessor, Operator}; use crate::tape::{by_directory, by_next}; use crate::tuples::*; use crate::{Page, PrefetcherSequenceFamily, RelationRead, tape, vectors}; -use std::error::Error; use std::fmt::Write; pub fn prewarm<'r, R: RelationRead, O: Operator>( index: &'r R, height: i32, mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, -) -> Result> { +) -> String { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -21,10 +34,10 @@ pub fn prewarm<'r, R: RelationRead, O: Operator>( drop(meta_guard); let mut message = String::new(); - writeln!(message, "height of root: {height_of_root}")?; + writeln!(message, "height of root: {height_of_root}").unwrap(); let prewarm_max_height = if height < 0 { 0 } else { height as u32 }; if prewarm_max_height > height_of_root { - return Ok(message); + return message; } type State = Vec; let mut state: State = { @@ -34,12 +47,12 @@ pub fn prewarm<'r, R: RelationRead, O: Operator>( vectors::read_for_h1_tuple::(root_head, list, ()); results.push(root_first); } - writeln!(message, "------------------------")?; - writeln!(message, "number of nodes: {}", results.len())?; - writeln!(message, "number of pages: {}", 1)?; + writeln!(message, "------------------------").unwrap(); + writeln!(message, "number of nodes: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", 1).unwrap(); results }; - let mut step = |state: State| -> Result<_, Box> { + let mut step = |state: State| { let mut counter = 0_usize; let mut results = Vec::new(); for first in state { @@ -53,13 +66,13 @@ pub fn prewarm<'r, R: RelationRead, O: Operator>( }, ); } - writeln!(message, "------------------------")?; - writeln!(message, "number of nodes: {}", results.len())?; - writeln!(message, "number of pages: {counter}")?; - Ok(results) + writeln!(message, "------------------------").unwrap(); + writeln!(message, "number of nodes: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {counter}").unwrap(); + results }; for _ in (std::cmp::max(1, prewarm_max_height)..height_of_root).rev() { - state = step(state)?; + state = step(state); } if prewarm_max_height == 0 { let mut counter = 0_usize; @@ -85,9 +98,9 @@ pub fn prewarm<'r, R: RelationRead, O: Operator>( }), ); } - writeln!(message, "------------------------")?; - writeln!(message, "number of nodes: {}", results.len())?; - writeln!(message, "number of pages: {counter}")?; + writeln!(message, "------------------------").unwrap(); + writeln!(message, "number of nodes: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {counter}").unwrap(); } - Ok(message) + message } diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index 66a1b841..69b86983 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::closure_lifetime_binder::id_4; use crate::operator::*; use crate::prefetcher::Prefetcher; diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index 7e4bbe4b..b5c2a7f5 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::closure_lifetime_binder::id_2; use crate::linked_vec::LinkedVec; use crate::operator::*; diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index 26f5922b..ff01b45b 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::operator::Accessor1; use crate::tuples::*; use crate::{Page, PageGuard, PrefetcherSequenceFamily, RelationRead, RelationWrite}; diff --git a/crates/algorithm/src/tape_writer.rs b/crates/algorithm/src/tape_writer.rs index b10d3493..5a486ca9 100644 --- a/crates/algorithm/src/tape_writer.rs +++ b/crates/algorithm/src/tape_writer.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::tape::TapeWriter; use crate::tuples::*; use crate::{Branch, RelationWrite}; diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index f2d30114..912d7f6c 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::operator::Vector; use rabitq::binary::BinaryCode; use std::marker::PhantomData; diff --git a/crates/algorithm/src/types.rs b/crates/algorithm/src/types.rs index f889477b..5325300c 100644 --- a/crates/algorithm/src/types.rs +++ b/crates/algorithm/src/types.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use half::f16; use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError}; diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index 8d4c7882..0d273edf 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::operator::*; use crate::tuples::*; use crate::{Page, PageGuard, RelationRead, RelationWrite, tape}; diff --git a/crates/always_equal/src/lib.rs b/crates/always_equal/src/lib.rs index 8cb2cc1c..c05d10f9 100644 --- a/crates/always_equal/src/lib.rs +++ b/crates/always_equal/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use std::cmp::Ordering; use std::hash::Hash; diff --git a/crates/distance/src/lib.rs b/crates/distance/src/lib.rs index 1d8fb05d..39867f52 100644 --- a/crates/distance/src/lib.rs +++ b/crates/distance/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] #[repr(transparent)] pub struct Distance(i32); diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index 052cbc6a..a67ffab3 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use rabitq::block::BlockCode; use rabitq::packing::{any_pack, padding_pack}; use rand::rngs::StdRng; diff --git a/crates/rabitq/src/binary.rs b/crates/rabitq/src/binary.rs index 402eeda7..bce60e8b 100644 --- a/crates/rabitq/src/binary.rs +++ b/crates/rabitq/src/binary.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use simd::Floating; pub type BinaryLut = ( diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs index ce3071e3..91956262 100644 --- a/crates/rabitq/src/block.rs +++ b/crates/rabitq/src/block.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use simd::Floating; pub type BlockLut = ((f32, f32, f32, f32), Vec<[u8; 16]>); diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index a4537372..c5ec66b9 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + pub mod binary; pub mod block; pub mod packing; diff --git a/crates/rabitq/src/packing.rs b/crates/rabitq/src/packing.rs index c9463e37..fc85bbcd 100644 --- a/crates/rabitq/src/packing.rs +++ b/crates/rabitq/src/packing.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + pub fn pack(x: [&[u8]; 32]) -> Vec<[u8; 16]> { let n = { let l = x.each_ref().map(|i| i.len()); diff --git a/crates/random_orthogonal_matrix/src/lib.rs b/crates/random_orthogonal_matrix/src/lib.rs index a9cbb5c8..b611f2b8 100644 --- a/crates/random_orthogonal_matrix/src/lib.rs +++ b/crates/random_orthogonal_matrix/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use nalgebra::DMatrix; #[ignore] diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 74738125..04d106d1 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -14,7 +14,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.19" +cc = "1.2.22" [lints] workspace = true diff --git a/crates/simd/build.rs b/crates/simd/build.rs index f9a84d92..12ce1981 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -1,8 +1,23 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use std::env::var; use std::error::Error; fn main() -> Result<(), Box> { println!("cargo::rerun-if-changed=cshim"); - let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH")?; + let target_arch = var("CARGO_CFG_TARGET_ARCH")?; match target_arch.as_str() { "aarch64" => { let mut build = cc::Build::new(); diff --git a/crates/simd/cshim/aarch64.c b/crates/simd/cshim/aarch64.c index 30b5350c..614eaa55 100644 --- a/crates/simd/cshim/aarch64.c +++ b/crates/simd/cshim/aarch64.c @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #if defined(__clang__) #if !(__clang_major__ >= 16) #error "Clang version must be at least 16." diff --git a/crates/simd/cshim/x86_64.c b/crates/simd/cshim/x86_64.c index 5e0f2524..8ba1357b 100644 --- a/crates/simd/cshim/x86_64.c +++ b/crates/simd/cshim/x86_64.c @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #if defined(__clang__) #if !(__clang_major__ >= 16) #error "Clang version must be at least 16." diff --git a/crates/simd/src/aligned.rs b/crates/simd/src/aligned.rs index 18db9d85..e8891bfa 100644 --- a/crates/simd/src/aligned.rs +++ b/crates/simd/src/aligned.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #[allow(dead_code)] #[derive(Debug, Clone, Copy)] #[repr(C, align(16))] diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index 2ace1b68..e547ec4c 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #[inline(always)] pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { reduce_sum_of_and::reduce_sum_of_and(lhs, rhs) diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index 7a75042a..a7680168 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + // VP2INTERSECT emulation. // Díez-Cañas, G. (2021). Faster-Than-Native Alternatives for x86 VP2INTERSECT // Instructions. arXiv preprint arXiv:2112.06342. diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index e17d73dc..b5415cae 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::{Floating, f32}; use half::f16; use zerocopy::FromZeros; diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index e04fa74b..e97041fa 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::Floating; impl Floating for f32 { diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index 8d3f9db6..f1271f4f 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + /* ## code layout for 4-bit quantizer diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 11291f28..627ec3fa 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #![cfg_attr(target_arch = "x86_64", feature(avx512_target_feature))] #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] #![allow(unsafe_code)] diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index bd6eefb7..cb363df0 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + mod mul_add_round { #[inline] #[cfg(target_arch = "x86_64")] diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index 53fc3510..e7490d2f 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #![allow(clippy::just_underscores_and_digits)] mod reduce_sum_of_x_as_u32_y_as_u32 { #[inline] diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs index 68a5b8c6..cd0eba81 100644 --- a/crates/simd_macros/src/lib.rs +++ b/crates/simd_macros/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + mod target; struct MultiversionVersion { diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs index 715ac870..2978ef5e 100644 --- a/crates/simd_macros/src/target.rs +++ b/crates/simd_macros/src/target.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + pub struct TargetCpu { pub target_cpu: &'static str, pub target_arch: &'static str, diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs index c625e6e3..cb34bd2f 100644 --- a/crates/vector/src/bvect.rs +++ b/crates/vector/src/bvect.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; use std::ops::{Bound, RangeBounds}; diff --git a/crates/vector/src/lib.rs b/crates/vector/src/lib.rs index 64128c7a..740e04ec 100644 --- a/crates/vector/src/lib.rs +++ b/crates/vector/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + pub mod bvect; pub mod scalar8; pub mod svect; diff --git a/crates/vector/src/scalar8.rs b/crates/vector/src/scalar8.rs index 21cccbbc..a2c3efeb 100644 --- a/crates/vector/src/scalar8.rs +++ b/crates/vector/src/scalar8.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; use std::ops::RangeBounds; diff --git a/crates/vector/src/svect.rs b/crates/vector/src/svect.rs index d83d5f2b..081ab021 100644 --- a/crates/vector/src/svect.rs +++ b/crates/vector/src/svect.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; use simd::Floating; diff --git a/crates/vector/src/vect.rs b/crates/vector/src/vect.rs index 527fc6ad..4c1a473f 100644 --- a/crates/vector/src/vect.rs +++ b/crates/vector/src/vect.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use super::{VectorBorrowed, VectorOwned}; use distance::Distance; use simd::Floating; diff --git a/deny.toml b/deny.toml new file mode 100644 index 00000000..0fa41117 --- /dev/null +++ b/deny.toml @@ -0,0 +1,21 @@ +[graph] +all-features = true + +[advisories] +ignore = [ + "RUSTSEC-2021-0127", # serde_cbor is unmaintained + "RUSTSEC-2024-0436", # paste - no longer maintained +] + +[licenses] +allow = [ + "Apache-2.0", + "Apache-2.0 WITH LLVM-exception", + "MIT", + # "MPL-2.0", + "BSD-3-Clause", + "ISC", + "Unicode-3.0", +] +exceptions = [{ allow = ["Zlib"], crate = "foldhash" }] +private = { ignore = true } diff --git a/rustfmt.toml b/rustfmt.toml index c32b643a..c1578aaf 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -1,2 +1 @@ -style_edition = "2024" imports_granularity = "Module" diff --git a/scripts/bench.py b/scripts/bench.py index ee35e685..0517b608 100644 --- a/scripts/bench.py +++ b/scripts/bench.py @@ -1,3 +1,17 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025 TensorChord Inc. + import time import argparse from pathlib import Path diff --git a/scripts/dump.py b/scripts/dump.py index fb5b1ccb..abdcdd3e 100644 --- a/scripts/dump.py +++ b/scripts/dump.py @@ -1,3 +1,17 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025 TensorChord Inc. + import argparse import h5py diff --git a/scripts/index.py b/scripts/index.py index 058a25f4..db238f9b 100644 --- a/scripts/index.py +++ b/scripts/index.py @@ -1,3 +1,17 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025 TensorChord Inc. + import asyncio import math from time import perf_counter diff --git a/scripts/train.py b/scripts/train.py index 6d3c44b9..e5f38925 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1,3 +1,17 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025 TensorChord Inc. + import itertools from multiprocessing import Pool, cpu_count from time import perf_counter diff --git a/src/bin/pgrx_embed.rs b/src/bin/pgrx_embed.rs index afd0164c..c7b6102c 100644 --- a/src/bin/pgrx_embed.rs +++ b/src/bin/pgrx_embed.rs @@ -1,2 +1,16 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #![allow(unsafe_code)] ::pgrx::pgrx_embed!(); diff --git a/src/datatype/binary_scalar8.rs b/src/datatype/binary_scalar8.rs index efb8758d..2179d589 100644 --- a/src/datatype/binary_scalar8.rs +++ b/src/datatype/binary_scalar8.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use super::memory_scalar8::{Scalar8Input, Scalar8Output}; use pgrx::datum::Internal; use pgrx::pg_sys::Oid; diff --git a/src/datatype/functions_scalar8.rs b/src/datatype/functions_scalar8.rs index 7cc7b42c..d37d81b9 100644 --- a/src/datatype/functions_scalar8.rs +++ b/src/datatype/functions_scalar8.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::datatype::memory_halfvec::HalfvecInput; use crate::datatype::memory_scalar8::Scalar8Output; use crate::datatype::memory_vector::VectorInput; diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs index 8fdae0a0..95afee4d 100644 --- a/src/datatype/memory_halfvec.rs +++ b/src/datatype/memory_halfvec.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use half::f16; use pgrx::datum::{FromDatum, IntoDatum}; use pgrx::pg_sys::{Datum, Oid}; diff --git a/src/datatype/memory_scalar8.rs b/src/datatype/memory_scalar8.rs index 19e4ff47..6fa927ef 100644 --- a/src/datatype/memory_scalar8.rs +++ b/src/datatype/memory_scalar8.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use pgrx::datum::{FromDatum, IntoDatum}; use pgrx::pg_sys::{Datum, Oid}; use pgrx::pgrx_sql_entity_graph::metadata::*; diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs index 16628605..5b3b1b3f 100644 --- a/src/datatype/memory_vector.rs +++ b/src/datatype/memory_vector.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use pgrx::datum::{FromDatum, IntoDatum}; use pgrx::pg_sys::{Datum, Oid}; use pgrx::pgrx_sql_entity_graph::metadata::*; diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index cd08970c..4ce1c38f 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + pub mod binary_scalar8; pub mod functions_scalar8; pub mod memory_halfvec; diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index d96e0945..9daf54cc 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; use pgrx::Array; use std::num::NonZero; diff --git a/src/datatype/operators_scalar8.rs b/src/datatype/operators_scalar8.rs index 8b75e592..76dbbf3d 100644 --- a/src/datatype/operators_scalar8.rs +++ b/src/datatype/operators_scalar8.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::datatype::memory_scalar8::{Scalar8Input, Scalar8Output}; use std::num::NonZero; use vector::VectorBorrowed; diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index 8f05a1c0..32ab7129 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use pgrx::Array; use std::num::NonZero; diff --git a/src/datatype/text_scalar8.rs b/src/datatype/text_scalar8.rs index d589ffdc..0b519589 100644 --- a/src/datatype/text_scalar8.rs +++ b/src/datatype/text_scalar8.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use super::memory_scalar8::Scalar8Output; use crate::datatype::memory_scalar8::Scalar8Input; use pgrx::pg_sys::Oid; diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index 4ce686c5..14bbcd94 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use serde::{Deserialize, Serialize}; use std::ffi::{CStr, CString}; use std::num::NonZero; diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index fd363256..95959877 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use super::opclass::Opfamily; use crate::index::am::am_build::InternalBuild; use algorithm::operator::{Dot, L2, Op}; @@ -191,7 +205,7 @@ impl Bump for BumpAlloc { pub fn prewarm(opfamily: Opfamily, index: &impl RelationRead, height: i32) -> String { let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; - let message = match (opfamily.vector_kind(), opfamily.distance_kind()) { + match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { algorithm::prewarm::<_, Op, L2>>(index, height, make_h0_plain_prefetcher) } @@ -212,10 +226,6 @@ pub fn prewarm(opfamily: Opfamily, index: &impl RelationRead, height: i32) -> St make_h0_plain_prefetcher, ) } - }; - match message { - Ok(message) => message, - Err(e) => panic!("{e}"), } } diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 8ec98d66..13f5584a 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::datatype::typmod::Typmod; use crate::index::am::{Reloption, ctid_to_key, kv_to_pointer}; use crate::index::opclass::{Opfamily, opfamily}; @@ -56,7 +70,7 @@ impl BuildPhase { if let Ok(cstr) = CStr::from_bytes_with_nul(bytes) { cstr } else { - panic!("") + unreachable!() } }, )*] diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index 7ca1754d..a6607f14 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + pub mod am_build; use super::algorithm::BumpAlloc; diff --git a/src/index/functions.rs b/src/index/functions.rs index 50026bf5..2cc77dba 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::index::storage::PostgresRelation; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass, PgClassRelkind}; diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 4f66dfbd..7a42cee5 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use super::scanners::Io; use pgrx::PostgresGucEnum; use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; diff --git a/src/index/hook.rs b/src/index/hook.rs index 748e939c..5e8665a0 100644 --- a/src/index/hook.rs +++ b/src/index/hook.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use std::sync::atomic::AtomicPtr; #[pgrx::pg_guard] diff --git a/src/index/lazy_cell.rs b/src/index/lazy_cell.rs index db616a57..5fe5a1a2 100644 --- a/src/index/lazy_cell.rs +++ b/src/index/lazy_cell.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + // Emulate unstable library feature `lazy_get`. // See https://github.com/rust-lang/rust/issues/129333. diff --git a/src/index/mod.rs b/src/index/mod.rs index 433a4c41..370e8897 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + pub mod algorithm; pub mod am; pub mod functions; diff --git a/src/index/opclass.rs b/src/index/opclass.rs index aebf2342..1079fd73 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use algorithm::types::*; diff --git a/src/index/projection.rs b/src/index/projection.rs index fbcaeffa..76ede977 100644 --- a/src/index/projection.rs +++ b/src/index/projection.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use random_orthogonal_matrix::random_orthogonal_matrix; use std::sync::OnceLock; diff --git a/src/index/scanners/default.rs b/src/index/scanners/default.rs index b994715e..57bda84d 100644 --- a/src/index/scanners/default.rs +++ b/src/index/scanners/default.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use super::{Io, SearchBuilder, SearchFetcher, SearchOptions}; use crate::index::algorithm::*; use crate::index::am::pointer_to_kv; diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs index 6494195b..5bf8a63d 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/scanners/maxsim.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use super::{SearchBuilder, SearchFetcher, SearchOptions}; use crate::index::algorithm::{RandomProject, *}; use crate::index::am::pointer_to_kv; diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs index b08b13e5..9be42971 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/scanners/mod.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + mod default; mod maxsim; diff --git a/src/index/storage.rs b/src/index/storage.rs index 432cd44f..2bad142f 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use algorithm::*; use std::collections::VecDeque; use std::iter::{Chain, Flatten}; diff --git a/src/index/types.rs b/src/index/types.rs index e626ec03..4e6e286f 100644 --- a/src/index/types.rs +++ b/src/index/types.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use algorithm::types::VchordrqIndexOptions; use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError, ValidationErrors}; diff --git a/src/lib.rs b/src/lib.rs index 9badf969..654b4d1b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #![allow(unsafe_code)] mod datatype; diff --git a/src/upgrade.rs b/src/upgrade.rs index 9879439f..15f81f67 100644 --- a/src/upgrade.rs +++ b/src/upgrade.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + // Referenced symbols must exist in the dynamic library when dropping functions. // So we should never remove symbols used by schema, otherwise there will be errors in upgrade. // Reference: From cf2d140e55d639520b724eadd9cc33a4a057e8a5 Mon Sep 17 00:00:00 2001 From: xieydd Date: Wed, 14 May 2025 10:10:11 +0800 Subject: [PATCH 156/324] add vectorchord docs link (#246) Signed-off-by: xieydd --- README.md | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 9aefc1ac..916af6b3 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@
-[Official Site][official-site-link] · [Blog][blog-link] · [Feedback][github-issues-link] · [Contact Us][email-link] +[Official Site][official-site-link] · [Blog][blog-link] · [Docs][docs-link] · [Feedback][github-issues-link] · [Contact Us][email-link] @@ -101,6 +101,16 @@ lists = [] $$); ``` +> [!NOTE] +> The `lists` option should be configured based on the number of vectors. Below is a table to assist with your selection. + +| vectors Range | List Calculation Formula | Example Result | +| ------------- | ------------------------------ | ---------------- | +| <128k | list = 1 | 1 | +| ≥128k and <2M | list = (2 * vectors) / 1000 | 256-4000 | +| ≥2M and <100M | list ∈ [4√vectors, 8√vectors] | 4000-80000 | +| ≥100M | list ∈ [8√vectors, 16√vectors] | 80000-160000 | + And then perform a vector search using `SELECT ... ORDER BY ... LIMIT ...`. ```SQL @@ -148,3 +158,4 @@ You may choose either license based on your needs. We welcome any commercial col [official-site-link]: https://vectorchord.ai/ [github-issues-link]: https://github.com/tensorchord/VectorChord/issues [email-link]: mailto:support@tensorchord.ai +[docs-link]: https://docs.vectorchord.ai/vectorchord/ From 7721bac08799f1ff9dfca36c4c6736b73874da05 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Wed, 14 May 2025 14:45:18 +0800 Subject: [PATCH 157/324] feat: support prefilter (#247) ## Feat - Enable `rerank_in_table=true` with `read_stream`(pg17) and `prefetch_buffer`(not pg17) ## Optimization - Remove redundant matrix multiply in both heap query/insert, accerate most searches ## Chore - Rename `vchordrq.prererank_filtering` into `vchordrq.prefilter` ## Test - Add specialized tests for pg16(prefetch_buffer) and pg17(read_stream) --------- Signed-off-by: cutecutecat --- .github/workflows/check.yml | 16 +- crates/algorithm/src/fast_heap.rs | 7 +- crates/algorithm/src/insert.rs | 36 +- crates/algorithm/src/lib.rs | 66 ++- src/index/algorithm.rs | 20 +- src/index/am/mod.rs | 60 +- src/index/gucs.rs | 14 +- src/index/scanners/default.rs | 587 ++++++++++++------- src/index/scanners/maxsim.rs | 451 ++++++++------ src/index/scanners/mod.rs | 4 + tests/{logic => general}/distance.slt | 0 tests/{logic => general}/external_build.slt | 0 tests/{logic => general}/index.slt | 0 tests/{logic => general}/issue427.slt | 0 tests/{logic => general}/multivector.slt | 0 tests/{logic => general}/null.fail | 0 tests/{logic => general}/partition.slt | 0 tests/{logic => general}/pin.slt | 0 tests/{logic => general}/pushdown_plan.slt | 0 tests/{logic => general}/pushdown_range.slt | 0 tests/{logic => general}/reindex.slt | 0 tests/general/rerank_in_index.slt | 101 ++++ tests/{logic => general}/rerank_in_table.slt | 4 +- tests/{logic => general}/vector.slt | 0 tests/pg16/filter_rerank_in_index.slt | 101 ++++ tests/pg16/filter_rerank_in_table.slt | 101 ++++ tests/pg17/filter_rerank_in_index.slt | 101 ++++ tests/pg17/filter_rerank_in_table.slt | 101 ++++ 28 files changed, 1336 insertions(+), 434 deletions(-) rename tests/{logic => general}/distance.slt (100%) rename tests/{logic => general}/external_build.slt (100%) rename tests/{logic => general}/index.slt (100%) rename tests/{logic => general}/issue427.slt (100%) rename tests/{logic => general}/multivector.slt (100%) rename tests/{logic => general}/null.fail (100%) rename tests/{logic => general}/partition.slt (100%) rename tests/{logic => general}/pin.slt (100%) rename tests/{logic => general}/pushdown_plan.slt (100%) rename tests/{logic => general}/pushdown_range.slt (100%) rename tests/{logic => general}/reindex.slt (100%) create mode 100644 tests/general/rerank_in_index.slt rename tests/{logic => general}/rerank_in_table.slt (95%) rename tests/{logic => general}/vector.slt (100%) create mode 100644 tests/pg16/filter_rerank_in_index.slt create mode 100644 tests/pg16/filter_rerank_in_table.slt create mode 100644 tests/pg17/filter_rerank_in_index.slt create mode 100644 tests/pg17/filter_rerank_in_table.slt diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 741de314..2978320b 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -235,7 +235,21 @@ jobs: run: | sudo systemctl start postgresql psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - sqllogictest --db $USER --user $USER './tests/**/*.slt' + sqllogictest --db $USER --user $USER './tests/general/*.slt' + + - name: Sqllogictest(PostgreSQL 17 features) + if: matrix.version == '17' + run: | + sudo systemctl start postgresql + psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + sqllogictest --db $USER --user $USER './tests/pg17/*.slt' + + - name: Sqllogictest(PostgreSQL 16 features) + if: matrix.version == '16' + run: | + sudo systemctl start postgresql + psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + sqllogictest --db $USER --user $USER './tests/pg16/*.slt' - name: Package env: diff --git a/crates/algorithm/src/fast_heap.rs b/crates/algorithm/src/fast_heap.rs index 18ef21b2..f70fd176 100644 --- a/crates/algorithm/src/fast_heap.rs +++ b/crates/algorithm/src/fast_heap.rs @@ -76,13 +76,12 @@ impl From> for FastHeap { impl Sequence for FastHeap { type Item = T; type Inner = std::vec::IntoIter; + fn peek(&mut self) -> Option<&T> { + >::peek(self) + } fn next(&mut self) -> Option { self.pop() } - fn next_if(&mut self, predicate: impl FnOnce(&T) -> bool) -> Option { - let first = self.peek()?; - if predicate(first) { self.pop() } else { None } - } fn into_inner(self) -> Self::Inner { match self { FastHeap::Sorted(sort_heap) => sort_heap.inner.into_iter(), diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 4194eb54..b2e08e70 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -30,25 +30,49 @@ type Item<'b> = ( AlwaysEqual<&'b mut (u32, u16, &'b mut [u32])>, ); -pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( +pub fn insert_vector( + index: &R, + payload: NonZero, + vector: &O::Vector, +) -> (Vec, u16) { + // `insert_vector` returns a tuple `(list, head)` which will be used in `insert_index` later: + // - `list`: Represents the list of elements to be inserted into the index. + // - `head`: Represents the head of the list, used as a starting point for insertion. + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dims = meta_tuple.dims(); + let rerank_in_heap = meta_tuple.rerank_in_heap(); + assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); + let vectors_first = meta_tuple.vectors_first(); + drop(meta_guard); + + if !rerank_in_heap { + vectors::append::(index, vectors_first, vector.as_borrowed(), payload) + } else { + (Vec::new(), 0) + } +} + +pub fn insert_index<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( index: &'r R, payload: NonZero, vector: O::Vector, bump: &'b impl Bump, mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, + list: Vec, + head: u16, ) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); - let rerank_in_heap = meta_tuple.rerank_in_heap(); let height_of_root = meta_tuple.height_of_root(); assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); let root_prefetch = meta_tuple.root_prefetch().to_vec(); let root_head = meta_tuple.root_head(); let root_first = meta_tuple.root_first(); - let vectors_first = meta_tuple.vectors_first(); drop(meta_guard); let default_block_lut = if !is_residual { @@ -57,12 +81,6 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( None }; - let (list, head) = if !rerank_in_heap { - vectors::append::(index, vectors_first, vector.as_borrowed(), payload) - } else { - (Vec::new(), 0) - }; - type State = (u32, Option<::Vector>); let mut state: State = { if is_residual { diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index cc37b57b..e06b9a11 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -42,7 +42,7 @@ pub use bulkdelete::bulkdelete; pub use cache::cache; pub use cost::cost; pub use fast_heap::FastHeap; -pub use insert::insert; +pub use insert::{insert_index, insert_vector}; pub use maintain::maintain; pub use prefetcher::*; pub use prewarm::prewarm; @@ -219,24 +219,64 @@ impl<'b, T, A, B> Fetch for (T, AlwaysEqual<&'b mut (A, B, &'b mut [u32])>) { } } +pub struct Filter { + pub iter: S, + pub filter: P, +} + +impl bool> Sequence for Filter { + type Item = S::Item; + type Inner = S::Inner; + + fn peek(&mut self) -> Option<&Self::Item> { + loop { + let item = self.iter.peek()?; + if (self.filter)(item) { + return self.iter.peek(); + } else { + self.iter.next(); + continue; + } + } + } + fn next(&mut self) -> Option { + loop { + let item = self.iter.peek()?; + if (self.filter)(item) { + return self.iter.next(); + } else { + self.iter.next(); + continue; + } + } + } + + fn into_inner(self) -> Self::Inner { + self.iter.into_inner() + } +} + pub trait Sequence { type Item; type Inner: Iterator; + fn peek(&mut self) -> Option<&Self::Item>; fn next(&mut self) -> Option; - fn next_if(&mut self, predicate: impl FnOnce(&Self::Item) -> bool) -> Option; + fn next_if(&mut self, predicate: impl FnOnce(&Self::Item) -> bool) -> Option { + let peek = self.peek()?; + if predicate(peek) { self.next() } else { None } + } fn into_inner(self) -> Self::Inner; } impl Sequence for BinaryHeap { type Item = T; type Inner = std::vec::IntoIter; + fn peek(&mut self) -> Option<&T> { + >::peek(self) + } fn next(&mut self) -> Option { self.pop() } - fn next_if(&mut self, predicate: impl FnOnce(&T) -> bool) -> Option { - let peek = self.peek()?; - if predicate(peek) { self.pop() } else { None } - } fn into_inner(self) -> Self::Inner { self.into_vec().into_iter() } @@ -245,13 +285,21 @@ impl Sequence for BinaryHeap { impl Sequence for Peekable { type Item = I::Item; type Inner = Peekable; + fn peek(&mut self) -> Option<&I::Item> { + Peekable::peek(self) + } fn next(&mut self) -> Option { Iterator::next(self) } - fn next_if(&mut self, predicate: impl FnOnce(&I::Item) -> bool) -> Option { - Peekable::next_if(self, predicate) - } fn into_inner(self) -> Self::Inner { self } } + +pub fn seq_filter(heap: impl Sequence, filter: F) -> impl Sequence +where + F: FnMut(&T) -> bool, + T: Ord, +{ + Filter { iter: heap, filter } +} diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index 95959877..ed207db0 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -322,42 +322,54 @@ pub fn insert( match (vector, opfamily.distance_kind()) { (OwnedVector::Vecf32(vector), DistanceKind::L2) => { assert!(opfamily.vector_kind() == VectorKind::Vecf32); - algorithm::insert::<_, Op, L2>>( + let (list, head) = insert_vector::<_, Op, L2>>(index, payload, &vector); + insert_index::<_, Op, L2>>( index, payload, RandomProject::project(vector.as_borrowed()), &bump, make_h1_plain_prefetcher, + list, + head, ) } (OwnedVector::Vecf32(vector), DistanceKind::Dot) => { assert!(opfamily.vector_kind() == VectorKind::Vecf32); - algorithm::insert::<_, Op, Dot>>( + let (list, head) = insert_vector::<_, Op, Dot>>(index, payload, &vector); + insert_index::<_, Op, Dot>>( index, payload, RandomProject::project(vector.as_borrowed()), &bump, make_h1_plain_prefetcher, + list, + head, ) } (OwnedVector::Vecf16(vector), DistanceKind::L2) => { assert!(opfamily.vector_kind() == VectorKind::Vecf16); - algorithm::insert::<_, Op, L2>>( + let (list, head) = insert_vector::<_, Op, L2>>(index, payload, &vector); + insert_index::<_, Op, L2>>( index, payload, RandomProject::project(vector.as_borrowed()), &bump, make_h1_plain_prefetcher, + list, + head, ) } (OwnedVector::Vecf16(vector), DistanceKind::Dot) => { assert!(opfamily.vector_kind() == VectorKind::Vecf16); - algorithm::insert::<_, Op, Dot>>( + let (list, head) = insert_vector::<_, Op, Dot>>(index, payload, &vector); + insert_index::<_, Op, Dot>>( index, payload, RandomProject::project(vector.as_borrowed()), &bump, make_h1_plain_prefetcher, + list, + head, ) } } diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index a6607f14..2e1dbaf1 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -15,7 +15,6 @@ pub mod am_build; use super::algorithm::BumpAlloc; -use super::gucs::prererank_filtering; use crate::index::gucs; use crate::index::lazy_cell::LazyCell; use crate::index::opclass::{Opfamily, opfamily}; @@ -581,28 +580,6 @@ impl SearchFetcher for HeapFetcher { if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { return None; } - if !self.hack.is_null() && prererank_filtering() { - if let Some(qual) = NonNull::new((*self.hack).ss.ps.qual) { - use pgrx::datum::FromDatum; - use pgrx::memcxt::PgMemoryContexts; - assert!(qual.as_ref().flags & pgrx::pg_sys::EEO_FLAG_IS_QUAL as u8 != 0); - let evalfunc = qual.as_ref().evalfunc.expect("no evalfunc for qual"); - if !(*self.hack).ss.ps.ps_ExprContext.is_null() { - let econtext = (*self.hack).ss.ps.ps_ExprContext; - (*econtext).ecxt_scantuple = self.slot; - pgrx::pg_sys::MemoryContextReset((*econtext).ecxt_per_tuple_memory); - let result = PgMemoryContexts::For((*econtext).ecxt_per_tuple_memory) - .switch_to(|_| { - let mut is_null = true; - let datum = evalfunc(qual.as_ptr(), econtext, &mut is_null); - bool::from_datum(datum, is_null) - }); - if result != Some(true) { - return None; - } - } - } - } (*self.econtext).ecxt_scantuple = self.slot; pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); pgrx::pg_sys::FormIndexDatum( @@ -615,6 +592,43 @@ impl SearchFetcher for HeapFetcher { Some((&self.values, &self.is_nulls)) } } + + fn filter(&mut self, key: [u16; 3]) -> bool { + if self.hack.is_null() { + return true; + } + unsafe { + let mut ctid = key_to_ctid(key); + let table_am = (*self.heap_relation).rd_tableam; + let fetch_row_version = (*table_am) + .tuple_fetch_row_version + .expect("unsupported heap access method"); + if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { + return false; + } + if let Some(qual) = NonNull::new((*self.hack).ss.ps.qual) { + use pgrx::datum::FromDatum; + use pgrx::memcxt::PgMemoryContexts; + assert!(qual.as_ref().flags & pgrx::pg_sys::EEO_FLAG_IS_QUAL as u8 != 0); + let evalfunc = qual.as_ref().evalfunc.expect("no evalfunc for qual"); + if !(*self.hack).ss.ps.ps_ExprContext.is_null() { + let econtext = (*self.hack).ss.ps.ps_ExprContext; + (*econtext).ecxt_scantuple = self.slot; + pgrx::pg_sys::MemoryContextReset((*econtext).ecxt_per_tuple_memory); + let result = PgMemoryContexts::For((*econtext).ecxt_per_tuple_memory) + .switch_to(|_| { + let mut is_null = true; + let datum = evalfunc(qual.as_ptr(), econtext, &mut is_null); + bool::from_datum(datum, is_null) + }); + if result != Some(true) { + return false; + } + } + } + true + } + } } struct Index { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 7a42cee5..d91a04ee 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -36,7 +36,7 @@ static MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); static MAXSIM_REFINE: GucSetting = GucSetting::::new(0); static MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); -static PRERERANK_FILTERING: GucSetting = GucSetting::::new(false); +static PREFILTER: GucSetting = GucSetting::::new(false); static IO_SEARCH: GucSetting = GucSetting::::new( #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] @@ -110,10 +110,10 @@ pub fn init() { GucFlags::default(), ); GucRegistry::define_bool_guc( - "vchordrq.prererank_filtering", - "`prererank_filtering` argument of vchordrq.", - "`prererank_filtering` argument of vchordrq.", - &PRERERANK_FILTERING, + "vchordrq.prefilter", + "`prefilter` argument of vchordrq.", + "`prefilter` argument of vchordrq.", + &PREFILTER, GucContext::Userset, GucFlags::default(), ); @@ -207,8 +207,8 @@ pub fn prewarm_dim() -> Vec { } } -pub fn prererank_filtering() -> bool { - PRERERANK_FILTERING.get() +pub fn prefilter() -> bool { + PREFILTER.get() } pub fn io_search() -> Io { diff --git a/src/index/scanners/default.rs b/src/index/scanners/default.rs index 57bda84d..021d514c 100644 --- a/src/index/scanners/default.rs +++ b/src/index/scanners/default.rs @@ -15,8 +15,9 @@ use super::{Io, SearchBuilder, SearchFetcher, SearchOptions}; use crate::index::algorithm::*; use crate::index::am::pointer_to_kv; +use crate::index::gucs::prefilter; use crate::index::opclass::{Opfamily, Sphere}; -use algorithm::operator::{Dot, L2, Op}; +use algorithm::operator::{Dot, L2}; use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; use algorithm::*; use half::f16; @@ -104,16 +105,15 @@ impl SearchBuilder for DefaultBuilder { let iter: Box)>> = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { - let vector = RandomProject::project( - if let OwnedVector::Vecf32(vector) = vector { - vector - } else { - unreachable!() - } - .as_borrowed(), - ); + type Op = operator::Op, L2>; + let original_vector = if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() + }; + let vector = RandomProject::project(original_vector.as_borrowed()); let results = match options.io_search { - Io::Plain => default_search::<_, Op, L2>>( + Io::Plain => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -122,7 +122,7 @@ impl SearchBuilder for DefaultBuilder { make_h1_plain_prefetcher, make_h0_plain_prefetcher, ), - Io::Simple => default_search::<_, Op, L2>>( + Io::Simple => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -131,7 +131,7 @@ impl SearchBuilder for DefaultBuilder { make_h1_plain_prefetcher, make_h0_simple_prefetcher, ), - Io::Stream => default_search::<_, Op, L2>>( + Io::Stream => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -141,80 +141,123 @@ impl SearchBuilder for DefaultBuilder { make_h0_stream_prefetcher, ), }; - let fetch = move |payload| { - let (key, _) = pointer_to_kv(payload); - let (datums, is_nulls) = fetcher.fetch(key)?; - let datum = (!is_nulls[0]).then_some(datums[0]); - let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; - let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() { - vector - } else { - unreachable!() - }; - Some(RandomProject::project(raw.as_borrowed())) - }; let method = how(index); - match (method, options.io_rerank) { - (RerankMethod::Index, Io::Plain) => { + match (method, options.io_rerank, prefilter()) { + (RerankMethod::Index, Io::Plain, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Plain, false) => { let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_index::, L2>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), - ) + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Simple, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) } - (RerankMethod::Index, Io::Simple) => { + (RerankMethod::Index, Io::Simple, false) => { let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( index, results.into(), ); - Box::new( - rerank_index::, L2>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), - ) + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) } - (RerankMethod::Index, Io::Stream) => { + (RerankMethod::Index, Io::Stream, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = + StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Stream, false) => { let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( index, results.into(), Hints::default(), ); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + if !fetcher.filter(key) { + return None; + } + let (datums, is_nulls) = fetcher.fetch(key)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( - rerank_index::, L2>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), + rerank_heap::(original_vector, prefetcher, fetch).map( + move |(distance, payload)| (opfamily.output(distance), payload), + ), ) } - (RerankMethod::Heap, _) => { + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let (datums, is_nulls) = fetcher.fetch(key)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( - rerank_heap::, L2>, _, _>( - vector, prefetcher, fetch, - ) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), + rerank_heap::(original_vector, prefetcher, fetch).map( + move |(distance, payload)| (opfamily.output(distance), payload), + ), ) } } } (VectorKind::Vecf32, DistanceKind::Dot) => { - let vector = RandomProject::project( - if let OwnedVector::Vecf32(vector) = vector { - vector - } else { - unreachable!() - } - .as_borrowed(), - ); + type Op = operator::Op, Dot>; + let original_vector = if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() + }; + let vector = RandomProject::project(original_vector.as_borrowed()); let results = match options.io_search { - Io::Plain => default_search::<_, Op, Dot>>( + Io::Plain => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -223,7 +266,7 @@ impl SearchBuilder for DefaultBuilder { make_h1_plain_prefetcher, make_h0_plain_prefetcher, ), - Io::Simple => default_search::<_, Op, Dot>>( + Io::Simple => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -232,7 +275,7 @@ impl SearchBuilder for DefaultBuilder { make_h1_plain_prefetcher, make_h0_simple_prefetcher, ), - Io::Stream => default_search::<_, Op, Dot>>( + Io::Stream => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -242,78 +285,123 @@ impl SearchBuilder for DefaultBuilder { make_h0_stream_prefetcher, ), }; - let fetch = move |payload| { - let (key, _) = pointer_to_kv(payload); - let (datums, is_nulls) = fetcher.fetch(key)?; - let datum = (!is_nulls[0]).then_some(datums[0]); - let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; - let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() { - vector - } else { - unreachable!() - }; - Some(RandomProject::project(raw.as_borrowed())) - }; let method = how(index); - match (method, options.io_rerank) { - (RerankMethod::Index, Io::Plain) => { + match (method, options.io_rerank, prefilter()) { + (RerankMethod::Index, Io::Plain, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Plain, false) => { let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_index::, Dot>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), - ) + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Simple, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( + index, + results.into(), + ); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) } - (RerankMethod::Index, Io::Simple) => { + (RerankMethod::Index, Io::Stream, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); let prefetcher = - SimplePrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_index::, Dot>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), - ) + StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) } - (RerankMethod::Index, Io::Stream) => { + (RerankMethod::Index, Io::Stream, false) => { let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( index, results.into(), Hints::default(), ); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + if !fetcher.filter(key) { + return None; + } + let (datums, is_nulls) = fetcher.fetch(key)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( - rerank_index::, Dot>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), + rerank_heap::(original_vector, prefetcher, fetch).map( + move |(distance, payload)| (opfamily.output(distance), payload), + ), ) } - (RerankMethod::Heap, _) => { + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let (datums, is_nulls) = fetcher.fetch(key)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( - rerank_heap::, Dot>, _, _>( - vector, prefetcher, fetch, - ) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), + rerank_heap::(original_vector, prefetcher, fetch).map( + move |(distance, payload)| (opfamily.output(distance), payload), + ), ) } } } (VectorKind::Vecf16, DistanceKind::L2) => { - let vector = RandomProject::project( - if let OwnedVector::Vecf16(vector) = vector { - vector - } else { - unreachable!() - } - .as_borrowed(), - ); + type Op = operator::Op, L2>; + let original_vector = if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() + }; + let vector = RandomProject::project(original_vector.as_borrowed()); let results = match options.io_search { - Io::Plain => default_search::<_, Op, L2>>( + Io::Plain => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -322,7 +410,7 @@ impl SearchBuilder for DefaultBuilder { make_h1_plain_prefetcher, make_h0_plain_prefetcher, ), - Io::Simple => default_search::<_, Op, L2>>( + Io::Simple => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -331,7 +419,7 @@ impl SearchBuilder for DefaultBuilder { make_h1_plain_prefetcher, make_h0_simple_prefetcher, ), - Io::Stream => default_search::<_, Op, L2>>( + Io::Stream => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -341,78 +429,123 @@ impl SearchBuilder for DefaultBuilder { make_h0_stream_prefetcher, ), }; - let fetch = move |payload| { - let (key, _) = pointer_to_kv(payload); - let (datums, is_nulls) = fetcher.fetch(key)?; - let datum = (!is_nulls[0]).then_some(datums[0]); - let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; - let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() { - vector - } else { - unreachable!() - }; - Some(RandomProject::project(raw.as_borrowed())) - }; let method = how(index); - match (method, options.io_rerank) { - (RerankMethod::Index, Io::Plain) => { + match (method, options.io_rerank, prefilter()) { + (RerankMethod::Index, Io::Plain, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Plain, false) => { let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_index::, L2>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), - ) + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Simple, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) } - (RerankMethod::Index, Io::Simple) => { + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( + index, + results.into(), + ); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Stream, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); let prefetcher = - SimplePrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_index::, L2>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), - ) + StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) } - (RerankMethod::Index, Io::Stream) => { + (RerankMethod::Index, Io::Stream, false) => { let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( index, results.into(), Hints::default(), ); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + if !fetcher.filter(key) { + return None; + } + let (datums, is_nulls) = fetcher.fetch(key)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( - rerank_index::, L2>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), + rerank_heap::(original_vector, prefetcher, fetch).map( + move |(distance, payload)| (opfamily.output(distance), payload), + ), ) } - (RerankMethod::Heap, _) => { + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let (datums, is_nulls) = fetcher.fetch(key)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( - rerank_heap::, L2>, _, _>( - vector, prefetcher, fetch, - ) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), + rerank_heap::(original_vector, prefetcher, fetch).map( + move |(distance, payload)| (opfamily.output(distance), payload), + ), ) } } } (VectorKind::Vecf16, DistanceKind::Dot) => { - let vector = RandomProject::project( - if let OwnedVector::Vecf16(vector) = vector { - vector - } else { - unreachable!() - } - .as_borrowed(), - ); + type Op = operator::Op, Dot>; + let original_vector = if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() + }; + let vector = RandomProject::project(original_vector.as_borrowed()); let results = match options.io_search { - Io::Plain => default_search::<_, Op, Dot>>( + Io::Plain => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -421,7 +554,7 @@ impl SearchBuilder for DefaultBuilder { make_h1_plain_prefetcher, make_h0_plain_prefetcher, ), - Io::Simple => default_search::<_, Op, Dot>>( + Io::Simple => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -430,7 +563,7 @@ impl SearchBuilder for DefaultBuilder { make_h1_plain_prefetcher, make_h0_simple_prefetcher, ), - Io::Stream => default_search::<_, Op, Dot>>( + Io::Stream => default_search::<_, Op>( index, vector.clone(), options.probes, @@ -440,63 +573,109 @@ impl SearchBuilder for DefaultBuilder { make_h0_stream_prefetcher, ), }; - let fetch = move |payload| { - let (key, _) = pointer_to_kv(payload); - let (datums, is_nulls) = fetcher.fetch(key)?; - let datum = (!is_nulls[0]).then_some(datums[0]); - let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; - let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() { - vector - } else { - unreachable!() - }; - Some(RandomProject::project(raw.as_borrowed())) - }; let method = how(index); - match (method, options.io_rerank) { - (RerankMethod::Index, Io::Plain) => { + match (method, options.io_rerank, prefilter()) { + (RerankMethod::Index, Io::Plain, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Plain, false) => { let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_index::, Dot>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), - ) + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) } - (RerankMethod::Index, Io::Simple) => { + (RerankMethod::Index, Io::Simple, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( + index, + results.into(), + ); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Index, Io::Stream, true) => { + let seq = seq_filter(BinaryHeap::from(results), move |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); let prefetcher = - SimplePrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_index::, Dot>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), - ) + StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) } - (RerankMethod::Index, Io::Stream) => { + (RerankMethod::Index, Io::Stream, false) => { let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( index, results.into(), Hints::default(), ); + Box::new(rerank_index::(original_vector, prefetcher).map( + move |(distance, payload)| (opfamily.output(distance), payload), + )) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + if !fetcher.filter(key) { + return None; + } + let (datums, is_nulls) = fetcher.fetch(key)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = + PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( - rerank_index::, Dot>, _, _>(vector, prefetcher) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), + rerank_heap::(original_vector, prefetcher, fetch).map( + move |(distance, payload)| (opfamily.output(distance), payload), + ), ) } - (RerankMethod::Heap, _) => { + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let (datums, is_nulls) = fetcher.fetch(key)?; + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; let prefetcher = PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); Box::new( - rerank_heap::, Dot>, _, _>( - vector, prefetcher, fetch, - ) - .map(move |(distance, payload)| { - (opfamily.output(distance), payload) - }), + rerank_heap::(original_vector, prefetcher, fetch).map( + move |(distance, payload)| (opfamily.output(distance), payload), + ), ) } } diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs index 5bf8a63d..9c39d800 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/scanners/maxsim.rs @@ -15,6 +15,7 @@ use super::{SearchBuilder, SearchFetcher, SearchOptions}; use crate::index::algorithm::{RandomProject, *}; use crate::index::am::pointer_to_kv; +use crate::index::gucs::prefilter; use crate::index::opclass::Opfamily; use crate::index::scanners::Io; use algorithm::operator::Dot; @@ -60,7 +61,7 @@ impl SearchBuilder for MaxsimBuilder { self, index: &'a R, options: SearchOptions, - _fetcher: impl SearchFetcher + 'a, + mut fetcher: impl SearchFetcher + 'a, bump: &'a impl Bump, ) -> Box + 'a> where @@ -104,191 +105,299 @@ impl SearchBuilder for MaxsimBuilder { let iter: Box> = match opfamily.vector_kind() { VectorKind::Vecf32 => { type Op = operator::Op, Dot>; - let vectors = vectors + let original_vectors = vectors .into_iter() .map(|vector| { - RandomProject::project( - if let OwnedVector::Vecf32(vector) = vector { - vector - } else { - unreachable!() - } - .as_borrowed(), - ) + if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() + } }) .collect::>(); - Box::new(vectors.into_iter().map(|vector| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - match options.io_rerank { - Io::Plain => { - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - let mut reranker = - rerank_index::(vector.clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - Io::Simple => { - let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( - index, - results.into(), - ); - let mut reranker = - rerank_index::(vector.clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - Io::Stream => { - let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( - index, - results.into(), - Hints::default(), - ); - let mut reranker = - rerank_index::(vector.clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + let vectors = original_vectors + .clone() + .into_iter() + .map(|vector| RandomProject::project(vector.as_borrowed())); + Box::new(vectors.into_iter().zip(original_vectors).map( + |(vector, original_vector)| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + match (options.io_rerank, prefilter()) { + (Io::Plain, true) => { + let seq = seq_filter(BinaryHeap::from(results), |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::<_, _>::new( + index, + BinaryHeap::from(results), + ); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let seq = seq_filter(BinaryHeap::from(results), |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::<'a, R, _>::new( + index, + BinaryHeap::from(results), + ); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let seq = seq_filter(BinaryHeap::from(results), |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = + StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = StreamPrefetcher::<_, _>::new( + index, + BinaryHeap::from(results), + Hints::default(), + ); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) + (accu_set, rough_set, estimation_by_threshold) + }, + )) } VectorKind::Vecf16 => { type Op = operator::Op, Dot>; - let vectors = vectors + let original_vectors = vectors .into_iter() .map(|vector| { - RandomProject::project( - if let OwnedVector::Vecf16(vector) = vector { - vector - } else { - unreachable!() - } - .as_borrowed(), - ) + if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() + } }) .collect::>(); - Box::new(vectors.into_iter().map(|vector| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - match options.io_rerank { - Io::Plain => { - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - let mut reranker = - rerank_index::(vector.clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - Io::Simple => { - let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( - index, - results.into(), - ); - let mut reranker = - rerank_index::(vector.clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - Io::Stream => { - let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( - index, - results.into(), - Hints::default(), - ); - let mut reranker = - rerank_index::(vector.clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + let vectors = original_vectors + .clone() + .into_iter() + .map(|vector| RandomProject::project(vector.as_borrowed())); + Box::new(vectors.into_iter().zip(original_vectors).map( + |(vector, original_vector)| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + vector.clone(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + match (options.io_rerank, prefilter()) { + (Io::Plain, true) => { + let seq = seq_filter(BinaryHeap::from(results), |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::<_, _>::new( + index, + BinaryHeap::from(results), + ); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let seq = seq_filter(BinaryHeap::from(results), |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::<'a, R, _>::new( + index, + BinaryHeap::from(results), + ); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let seq = seq_filter(BinaryHeap::from(results), |key| { + fetcher.filter(pointer_to_kv(key.1.0.0).0) + }); + let prefetcher = + StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = StreamPrefetcher::<_, _>::new( + index, + BinaryHeap::from(results), + Hints::default(), + ); + let mut reranker = rerank_index::( + original_vector.clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) + (accu_set, rough_set, estimation_by_threshold) + }, + )) } }; let mut updates = Vec::new(); diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs index 9be42971..16ddd5d7 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/scanners/mod.rs @@ -63,10 +63,14 @@ pub trait SearchBuilder: 'static { pub trait SearchFetcher { fn fetch(&mut self, ctid: [u16; 3]) -> Option<(&[Datum; 32], &[bool; 32])>; + fn filter(&mut self, key: [u16; 3]) -> bool; } impl T> SearchFetcher for LazyCell { fn fetch(&mut self, key: [u16; 3]) -> Option<(&[Datum; 32], &[bool; 32])> { LazyCell::force_mut(self).fetch(key) } + fn filter(&mut self, key: [u16; 3]) -> bool { + LazyCell::force_mut(self).filter(key) + } } diff --git a/tests/logic/distance.slt b/tests/general/distance.slt similarity index 100% rename from tests/logic/distance.slt rename to tests/general/distance.slt diff --git a/tests/logic/external_build.slt b/tests/general/external_build.slt similarity index 100% rename from tests/logic/external_build.slt rename to tests/general/external_build.slt diff --git a/tests/logic/index.slt b/tests/general/index.slt similarity index 100% rename from tests/logic/index.slt rename to tests/general/index.slt diff --git a/tests/logic/issue427.slt b/tests/general/issue427.slt similarity index 100% rename from tests/logic/issue427.slt rename to tests/general/issue427.slt diff --git a/tests/logic/multivector.slt b/tests/general/multivector.slt similarity index 100% rename from tests/logic/multivector.slt rename to tests/general/multivector.slt diff --git a/tests/logic/null.fail b/tests/general/null.fail similarity index 100% rename from tests/logic/null.fail rename to tests/general/null.fail diff --git a/tests/logic/partition.slt b/tests/general/partition.slt similarity index 100% rename from tests/logic/partition.slt rename to tests/general/partition.slt diff --git a/tests/logic/pin.slt b/tests/general/pin.slt similarity index 100% rename from tests/logic/pin.slt rename to tests/general/pin.slt diff --git a/tests/logic/pushdown_plan.slt b/tests/general/pushdown_plan.slt similarity index 100% rename from tests/logic/pushdown_plan.slt rename to tests/general/pushdown_plan.slt diff --git a/tests/logic/pushdown_range.slt b/tests/general/pushdown_range.slt similarity index 100% rename from tests/logic/pushdown_range.slt rename to tests/general/pushdown_range.slt diff --git a/tests/logic/reindex.slt b/tests/general/reindex.slt similarity index 100% rename from tests/logic/reindex.slt rename to tests/general/reindex.slt diff --git a/tests/general/rerank_in_index.slt b/tests/general/rerank_in_index.slt new file mode 100644 index 00000000..02d1e954 --- /dev/null +++ b/tests/general/rerank_in_index.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_column (id integer, val vector(3)); + +statement ok +INSERT INTO t_column (id, val) SELECT id, ARRAY[id, id, id]::real[] FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_column USING vchordrq (val vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = false +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +query I +SELECT id FROM t_column ORDER BY val <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +statement ok +DROP TABLE t_column; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = false +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT id FROM t_expr WHERE id <= 5 OR id % 2 = 1 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' LIMIT 9; +---- +2 +1 +3 +4 +5 +7 +9 +11 +13 + +statement ok +SET vchordrq.prefilter to off; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +SET vchordrq.prefilter to on; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/logic/rerank_in_table.slt b/tests/general/rerank_in_table.slt similarity index 95% rename from tests/logic/rerank_in_table.slt rename to tests/general/rerank_in_table.slt index 8a2b8c44..e9287976 100644 --- a/tests/logic/rerank_in_table.slt +++ b/tests/general/rerank_in_table.slt @@ -80,7 +80,7 @@ SELECT id FROM t_expr WHERE id <= 5 OR id % 2 = 1 ORDER BY ARRAY[id::real, id::r 13 statement ok -SET vchordrq.prererank_filtering to off; +SET vchordrq.prefilter to off; query I SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); @@ -89,7 +89,7 @@ SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 OR {2,3,1,4,5,6,7,8,9} statement ok -SET vchordrq.prererank_filtering to on; +SET vchordrq.prefilter to on; query I SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); diff --git a/tests/logic/vector.slt b/tests/general/vector.slt similarity index 100% rename from tests/logic/vector.slt rename to tests/general/vector.slt diff --git a/tests/pg16/filter_rerank_in_index.slt b/tests/pg16/filter_rerank_in_index.slt new file mode 100644 index 00000000..69c6e9e5 --- /dev/null +++ b/tests/pg16/filter_rerank_in_index.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = false +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +# non-heap rerank + no prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + simple prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_search to 'prefetch_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + no prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + simple prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_search to 'prefetch_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/pg16/filter_rerank_in_table.slt b/tests/pg16/filter_rerank_in_table.slt new file mode 100644 index 00000000..c556207b --- /dev/null +++ b/tests/pg16/filter_rerank_in_table.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = true +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +# heap rerank + no prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + simple prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_search to 'prefetch_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + no prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + simple prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_search to 'prefetch_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/pg17/filter_rerank_in_index.slt b/tests/pg17/filter_rerank_in_index.slt new file mode 100644 index 00000000..3b0e9b10 --- /dev/null +++ b/tests/pg17/filter_rerank_in_index.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = false +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +# non-heap rerank + no prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + stream prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_search to 'read_stream'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + no prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + stream prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_search to 'read_stream'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/pg17/filter_rerank_in_table.slt b/tests/pg17/filter_rerank_in_table.slt new file mode 100644 index 00000000..19f0bf60 --- /dev/null +++ b/tests/pg17/filter_rerank_in_table.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = true +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +# heap rerank + no prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + stream prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_search to 'read_stream'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + no prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + stream prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_search to 'read_stream'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; From c0c51282608ed6e47c9a70a71f9c48c7937f20c9 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 21 May 2025 10:16:39 +0800 Subject: [PATCH 158/324] feat: residual quantization without too many lookup tables (#252) note: - residual quantization for inner product is implemented Signed-off-by: usamoi --- .github/workflows/check.yml | 4 + crates/algorithm/src/build.rs | 62 +- crates/algorithm/src/bulkdelete.rs | 229 +++---- crates/algorithm/src/cache.rs | 13 +- .../algorithm/src/closure_lifetime_binder.rs | 4 +- crates/algorithm/src/cost.rs | 9 +- crates/algorithm/src/fast_heap.rs | 6 +- crates/algorithm/src/insert.rs | 186 +++--- crates/algorithm/src/lib.rs | 82 +-- crates/algorithm/src/maintain.rs | 109 ++-- crates/algorithm/src/operator.rs | 557 ++++++++++++++---- crates/algorithm/src/prefetcher.rs | 7 +- crates/algorithm/src/prewarm.rs | 37 +- crates/algorithm/src/rerank.rs | 8 +- crates/algorithm/src/search.rs | 280 ++++----- crates/algorithm/src/tape.rs | 33 +- crates/algorithm/src/tape_writer.rs | 49 +- crates/algorithm/src/tuples.rs | 362 +++++++----- crates/algorithm/src/types.rs | 2 +- crates/algorithm/src/vectors.rs | 12 +- crates/k_means/src/lib.rs | 31 +- crates/rabitq/src/binary.rs | 203 +++++-- crates/rabitq/src/block.rs | 205 +++++-- crates/rabitq/src/lib.rs | 107 ++-- crates/rabitq/src/packing.rs | 28 + crates/simd/src/fast_scan.rs | 16 +- src/index/algorithm.rs | 108 ++-- src/index/am/am_build.rs | 38 +- src/index/am/mod.rs | 86 +-- src/index/gucs.rs | 31 +- src/index/scanners/default.rs | 526 ++++++++--------- src/index/scanners/maxsim.rs | 539 +++++++++-------- src/index/scanners/mod.rs | 69 ++- 33 files changed, 2306 insertions(+), 1732 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 2978320b..aaa8971d 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -52,6 +52,8 @@ jobs: steps: - name: Set up Environment run: | + rustup set profile + curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo curl -fsSL https://github.com/EmbarkStudios/cargo-deny/releases/download/0.18.2/cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl/cargo-deny | install -m 755 /dev/stdin /usr/local/bin/cargo-deny @@ -140,6 +142,7 @@ jobs: steps: - name: Set up Environment run: | + rustup set profile sudo apt-get update if [ "$(uname -m)" == "x86_64" ]; then @@ -193,6 +196,7 @@ jobs: steps: - name: Set up Environment run: | + rustup set profile sudo apt-get update sudo apt-get remove -y '^postgres.*' '^libpq.*' diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index f4534cf7..5fe3a531 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -12,13 +12,13 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::operator::{Accessor2, Operator, Vector}; +use crate::operator::{Operator, Vector}; use crate::tape::TapeWriter; use crate::tape_writer::H1TapeWriter; use crate::tuples::*; use crate::types::*; use crate::{Branch, RelationWrite}; -use vector::VectorOwned; +use vector::{VectorBorrowed, VectorOwned}; pub fn build( vector_options: VectorOptions, @@ -26,20 +26,17 @@ pub fn build( index: &R, structures: Vec>, ) { - if vchordrq_options.residual_quantization && !O::SUPPORTS_RESIDUAL { - panic!("residual_quantization can be enabled only if distance type is L2"); - } let dims = vector_options.dims; let is_residual = vchordrq_options.residual_quantization; let mut meta = TapeWriter::<_, MetaTuple>::create(index, false); assert_eq!(meta.first(), 0); let freepage = TapeWriter::<_, FreepageTuple>::create(index, false); let mut vectors = TapeWriter::<_, VectorTuple>::create(index, true); - let mut pointer_of_means = Vec::, u16)>>::new(); + let mut pointer_of_centroids = Vec::, u16)>>::new(); for i in 0..structures.len() { let mut level = Vec::new(); for j in 0..structures[i].len() { - let vector = structures[i].means[j].as_borrowed(); + let vector = structures[i].centroids[j].as_borrowed(); let (slices, metadata) = O::Vector::split(vector); let mut chain = Ok(metadata); let mut prefetch = Vec::new(); @@ -65,7 +62,7 @@ pub fn build( chain.expect_err("internal error: 0-dimensional vector"), )); } - pointer_of_means.push(level); + pointer_of_centroids.push(level); } let mut pointer_of_firsts = Vec::>::new(); for i in 0..structures.len() { @@ -78,34 +75,24 @@ pub fn build( jump.push(JumpTuple { directory_first: directory_tape.first(), appendable_first: appendable_tape.first(), + centroid_prefetch: pointer_of_centroids[i][j].0.clone(), + centroid_head: pointer_of_centroids[i][j].1, tuples: 0, }); level.push(jump.first()); } else { let mut tape = H1TapeWriter::create(index, O::Vector::count(dims as _), false); - let h2_mean = structures[i].means[j].as_borrowed(); - let h2_children = structures[i].children[j].as_slice(); - for child in h2_children.iter().copied() { - let h1_mean = structures[i - 1].means[child as usize].as_borrowed(); - let code = if is_residual { - let mut residual_accessor = O::ResidualAccessor::default(); - residual_accessor - .push(O::Vector::unpack(h1_mean).0, O::Vector::unpack(h2_mean).0); - let residual = residual_accessor - .finish(O::Vector::unpack(h1_mean).1, O::Vector::unpack(h2_mean).1); - O::Vector::code(residual.as_borrowed()) - } else { - O::Vector::code(h1_mean) - }; + let centroid = structures[i].centroids[j].as_borrowed(); + for child in structures[i].children[j].iter().copied() { + let vector = structures[i - 1].centroids[child as usize].as_borrowed(); + let (code, delta) = O::build(vector, is_residual.then_some(centroid.own())); tape.push(Branch { - head: pointer_of_means[i - 1][child as usize].1, - dis_u_2: code.dis_u_2, - factor_ppc: code.factor_ppc, - factor_ip: code.factor_ip, - factor_err: code.factor_err, - signs: code.signs, - prefetch: pointer_of_means[i - 1][child as usize].0.clone(), + code, + delta, + prefetch: pointer_of_centroids[i - 1][child as usize].0.clone(), + head: pointer_of_centroids[i - 1][child as usize].1, extra: pointer_of_firsts[i - 1][child as usize], + norm: norm::(vector), }); } let (mut tape, chunk) = tape.into_inner(); @@ -121,19 +108,30 @@ pub fn build( is_residual, rerank_in_heap: vchordrq_options.rerank_in_table, vectors_first: vectors.first(), - root_prefetch: pointer_of_means + centroid_prefetch: pointer_of_centroids .last() .expect("internal error: empty structure")[0] .0 .clone(), - root_head: pointer_of_means + centroid_head: pointer_of_centroids .last() .expect("internal error: empty structure")[0] .1, - root_first: pointer_of_firsts + centroid_norm: norm::( + structures + .last() + .expect("internal error: empty structure") + .centroids[0] + .as_borrowed(), + ), + first: pointer_of_firsts .last() .expect("internal error: empty structure")[0], freepage_first: freepage.first(), cells: structures.iter().map(|s| s.len() as _).collect(), }); } + +fn norm(vector: V::Borrowed<'_>) -> f32 { + V::squared_norm(vector).sqrt() +} diff --git a/crates/algorithm/src/bulkdelete.rs b/crates/algorithm/src/bulkdelete.rs index 1f3fcf79..f13c72b5 100644 --- a/crates/algorithm/src/bulkdelete.rs +++ b/crates/algorithm/src/bulkdelete.rs @@ -28,138 +28,149 @@ pub fn bulkdelete( let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let height_of_root = meta_tuple.height_of_root(); - let root_first = meta_tuple.root_first(); - let vectors_first = meta_tuple.vectors_first(); + + type State = Vec; + let mut state: State = vec![meta_tuple.first()]; + drop(meta_guard); - { - type State = Vec; - let mut state: State = vec![root_first]; - let step = |state: State| { - let mut results = Vec::new(); - for first in state { - tape::read_h1_tape::( - by_next(index, first).inspect(|_| check()), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), - |(), _, first, _| results.push(first), - ); - } - results - }; - for _ in (1..height_of_root).rev() { - state = step(state); - } + + let step = |state: State| { + let mut results = Vec::new(); for first in state { - let jump_guard = index.read(first); - let jump_bytes = jump_guard.get(1).expect("data corruption"); - let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let mut directory = tape::read_directory_tape::( - by_next(index, jump_tuple.directory_first()).inspect(|_| check()), + tape::read_h1_tape::( + by_next(index, first).inspect(|_| check()), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + |(), _, _, first, _| results.push(first), ); - { - let mut current = directory.next().unwrap_or(u32::MAX); - while current != u32::MAX { - check(); - let read = index.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - let bytes = read.get(i).expect("data corruption"); - let tuple = FrozenTuple::deserialize_ref(bytes); - if let FrozenTupleReader::_0(tuple) = tuple { - for p in tuple.payload().iter() { - if Some(true) == p.map(&callback) { - break 'flag true; - } - } - } - } - false - }; - if flag { - drop(read); - let mut write = index.write(current, false); - for i in 1..=write.len() { - let bytes = write.get_mut(i).expect("data corruption"); - let tuple = FrozenTuple::deserialize_mut(bytes); - if let FrozenTupleWriter::_0(mut tuple) = tuple { - for p in tuple.payload().iter_mut() { - if Some(true) == p.map(&callback) { - *p = None; - } + } + results + }; + for _ in (1..height_of_root).rev() { + state = step(state); + } + for first in state { + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + let mut directory = tape::read_directory_tape::( + by_next(index, jump_tuple.directory_first()).inspect(|_| check()), + ); + { + let mut current = directory.next().unwrap_or(u32::MAX); + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + let bytes = read.get(i).expect("data corruption"); + let tuple = FrozenTuple::deserialize_ref(bytes); + if let FrozenTupleReader::_0(tuple) = tuple { + for p in tuple.payload().iter() { + if Some(true) == p.map(&callback) { + break 'flag true; } } } } - current = directory.next().unwrap_or(u32::MAX); - } - } - { - let mut current = jump_tuple.appendable_first(); - while current != u32::MAX { - check(); - let read = index.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - let bytes = read.get(i).expect("data corruption"); - let tuple = AppendableTuple::deserialize_ref(bytes); - let p = tuple.payload(); - if Some(true) == p.map(&callback) { - break 'flag true; - } - } - false - }; - if flag { - drop(read); - let mut write = index.write(current, false); - for i in 1..=write.len() { - let bytes = write.get_mut(i).expect("data corruption"); - let mut tuple = AppendableTuple::deserialize_mut(bytes); - let p = tuple.payload(); - if Some(true) == p.map(&callback) { - *p = None; + false + }; + if flag { + drop(read); + let mut write = index.write(current, false); + for i in 1..=write.len() { + let bytes = write.get_mut(i).expect("data corruption"); + let tuple = FrozenTuple::deserialize_mut(bytes); + if let FrozenTupleWriter::_0(mut tuple) = tuple { + for p in tuple.payload().iter_mut() { + if Some(true) == p.map(&callback) { + *p = None; + } } } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; } } + current = directory.next().unwrap_or(u32::MAX); } } - } - { - let mut current = vectors_first; - while current != u32::MAX { - check(); - let read = index.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - if let Some(bytes) = read.get(i) { - let tuple = VectorTuple::::deserialize_ref(bytes); + { + let mut current = jump_tuple.appendable_first(); + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + let bytes = read.get(i).expect("data corruption"); + let tuple = AppendableTuple::deserialize_ref(bytes); let p = tuple.payload(); if Some(true) == p.map(&callback) { break 'flag true; } } - } - false - }; - if flag { - drop(read); - let mut write = index.write(current, true); - for i in 1..=write.len() { - if let Some(bytes) = write.get(i) { - let tuple = VectorTuple::::deserialize_ref(bytes); + false + }; + if flag { + drop(read); + let mut write = index.write(current, false); + for i in 1..=write.len() { + let bytes = write.get_mut(i).expect("data corruption"); + let mut tuple = AppendableTuple::deserialize_mut(bytes); let p = tuple.payload(); if Some(true) == p.map(&callback) { - write.free(i); + *p = None; } - }; + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; } } } } + +pub fn bulkdelete_vectors( + index: &R, + check: impl Fn(), + callback: impl Fn(NonZero) -> bool, +) { + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let vectors_first = meta_tuple.vectors_first(); + + drop(meta_guard); + + let mut current = vectors_first; + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + if let Some(bytes) = read.get(i) { + let tuple = VectorTuple::::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + break 'flag true; + } + } + } + false + }; + if flag { + drop(read); + let mut write = index.write(current, true); + for i in 1..=write.len() { + if let Some(bytes) = write.get(i) { + let tuple = VectorTuple::::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + write.free(i); + } + }; + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; + } + } +} diff --git a/crates/algorithm/src/cache.rs b/crates/algorithm/src/cache.rs index ca183662..ca3f0927 100644 --- a/crates/algorithm/src/cache.rs +++ b/crates/algorithm/src/cache.rs @@ -24,28 +24,33 @@ pub fn cache(index: &R) -> Vec { let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let height_of_root = meta_tuple.height_of_root(); - let root_first = meta_tuple.root_first(); - drop(meta_guard); + type State = Vec; - let mut state: State = vec![root_first]; + let mut state: State = vec![meta_tuple.first()]; + + drop(meta_guard); + let mut step = |state: State| { let mut results = Vec::new(); for first in state { tape::read_h1_tape::( by_next(index, first).inspect(|guard| trace.push(guard.id())), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), - |(), _, first, _| { + |(), _, _, first, _| { results.push(first); }, ); } results }; + for _ in (1..height_of_root).rev() { state = step(state); } + for first in state { trace.push(first); } + trace } diff --git a/crates/algorithm/src/closure_lifetime_binder.rs b/crates/algorithm/src/closure_lifetime_binder.rs index f24f4414..592757cf 100644 --- a/crates/algorithm/src/closure_lifetime_binder.rs +++ b/crates/algorithm/src/closure_lifetime_binder.rs @@ -24,9 +24,9 @@ where } #[inline(always)] -pub fn id_1(f: F) -> F +pub fn id_1(f: F) -> F where - F: for<'a> FnMut(A, (&'a B, &'a B, &'a B, &'a B)) -> R, + F: for<'a> FnMut(A, (&'a B, &'a C)) -> R, { f } diff --git a/crates/algorithm/src/cost.rs b/crates/algorithm/src/cost.rs index f7ae8ab7..12ad027e 100644 --- a/crates/algorithm/src/cost.rs +++ b/crates/algorithm/src/cost.rs @@ -17,7 +17,6 @@ use crate::{Page, RelationRead}; pub struct Cost { pub dims: u32, - pub is_residual: bool, pub cells: Vec, } @@ -27,13 +26,9 @@ pub fn cost(index: &R) -> Cost { let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); - let is_residual = meta_tuple.is_residual(); let cells = meta_tuple.cells().to_vec(); + drop(meta_guard); - Cost { - dims, - is_residual, - cells, - } + Cost { dims, cells } } diff --git a/crates/algorithm/src/fast_heap.rs b/crates/algorithm/src/fast_heap.rs index f70fd176..6222ffb0 100644 --- a/crates/algorithm/src/fast_heap.rs +++ b/crates/algorithm/src/fast_heap.rs @@ -76,12 +76,12 @@ impl From> for FastHeap { impl Sequence for FastHeap { type Item = T; type Inner = std::vec::IntoIter; - fn peek(&mut self) -> Option<&T> { - >::peek(self) - } fn next(&mut self) -> Option { self.pop() } + fn peek(&mut self) -> Option<&Self::Item> { + (self as &Self).peek() + } fn into_inner(self) -> Self::Inner { match self { FastHeap::Sorted(sort_heap) => sort_heap.inner.into_iter(), diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index b2e08e70..66915daa 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -25,43 +25,37 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -type Item<'b> = ( - Reverse, - AlwaysEqual<&'b mut (u32, u16, &'b mut [u32])>, -); +type Extra<'b> = &'b mut (u32, u16, f32, &'b mut [u32]); pub fn insert_vector( index: &R, payload: NonZero, - vector: &O::Vector, + vector: ::Borrowed<'_>, ) -> (Vec, u16) { - // `insert_vector` returns a tuple `(list, head)` which will be used in `insert_index` later: - // - `list`: Represents the list of elements to be inserted into the index. - // - `head`: Represents the head of the list, used as a starting point for insertion. let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let rerank_in_heap = meta_tuple.rerank_in_heap(); - assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); + assert_eq!(dims, vector.dims(), "unmatched dimensions"); let vectors_first = meta_tuple.vectors_first(); + drop(meta_guard); if !rerank_in_heap { - vectors::append::(index, vectors_first, vector.as_borrowed(), payload) + vectors::append::(index, vectors_first, vector, payload) } else { (Vec::new(), 0) } } -pub fn insert_index<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( +pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( index: &'r R, payload: NonZero, - vector: O::Vector, + vector: ::Borrowed<'_>, + key: (Vec, u16), bump: &'b impl Bump, mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, - list: Vec, - head: u16, ) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); @@ -69,123 +63,109 @@ pub fn insert_index<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); let height_of_root = meta_tuple.height_of_root(); - assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); - let root_prefetch = meta_tuple.root_prefetch().to_vec(); - let root_head = meta_tuple.root_head(); - let root_first = meta_tuple.root_first(); - drop(meta_guard); + assert_eq!(dims, vector.dims(), "unmatched dimensions"); + let epsilon = 1.9; - let default_block_lut = if !is_residual { - Some(O::Vector::block_preprocess(vector.as_borrowed())) + type State = (Reverse, AlwaysEqual, AlwaysEqual); + let mut state: State = if !is_residual { + let first = meta_tuple.first(); + // it's safe to leave it a fake value + ( + Reverse(Distance::ZERO), + AlwaysEqual(0.0), + AlwaysEqual(first), + ) } else { - None + let prefetch = bump.alloc_slice(meta_tuple.centroid_prefetch()); + let head = meta_tuple.centroid_head(); + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + let distance = vectors::read_for_h1_tuple::( + prefetch.iter().map(|&id| index.read(id)), + head, + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + ); + (Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first)) }; - type State = (u32, Option<::Vector>); - let mut state: State = { - if is_residual { - let list = root_prefetch.into_iter().map(|id| index.read(id)); - let residual = vectors::read_for_h1_tuple::( - root_head, - list, - LAccess::new( - O::Vector::unpack(vector.as_borrowed()), - O::ResidualAccessor::default(), - ), - ); - (root_first, Some(residual)) - } else { - (root_first, None) - } - }; - let mut step = |state: State| { - let mut results = LinkedVec::>::new(); + drop(meta_guard); + let lut = (O::Vector::block_preprocess(vector),); + + let mut step = |state: State| { + let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); { - let (first, residual) = state; - let block_lut = if let Some(residual) = residual { - &O::Vector::block_preprocess(residual.as_borrowed()) - } else if let Some(block_lut) = default_block_lut.as_ref() { - block_lut - } else { - unreachable!() + let (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) = state; + let process = |value, code, delta, lut| { + O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) }; tape::read_h1_tape::( by_next(index, first), - || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), - |(rough, err), head, first, prefetch| { - let lowerbound = Distance::from_f32(rough - err * 1.9); + || O::block_access(&lut.0, process), + |(rough, err), head, norm, first, prefetch| { + let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( Reverse(lowerbound), - AlwaysEqual(bump.alloc((first, head, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc((first, head, norm, bump.alloc_slice(prefetch)))), )); }, ); } let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); - let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); + let mut cache = BinaryHeap::<(_, _, _)>::new(); { - while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, ..))), list)) = + while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, norm, ..))), prefetch)) = heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { - if is_residual { - let (distance, residual) = vectors::read_for_h1_tuple::( - head, - list.into_iter(), - LAccess::new( - O::Vector::unpack(vector.as_borrowed()), - ( - O::DistanceAccessor::default(), - O::ResidualAccessor::default(), - ), - ), - ); - cache.push(( - Reverse(distance), - AlwaysEqual(first), - AlwaysEqual(Some(residual)), - )); - } else { - let distance = vectors::read_for_h1_tuple::( - head, - list.into_iter(), - LAccess::new( - O::Vector::unpack(vector.as_borrowed()), - O::DistanceAccessor::default(), - ), - ); - cache.push((Reverse(distance), AlwaysEqual(first), AlwaysEqual(None))); - } + let distance = vectors::read_for_h1_tuple::( + prefetch.into_iter(), + head, + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + ); + cache.push((Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))); } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache - .pop() - .expect("invariant is violated: tree is not height-balanced"); - (first, mean) + cache.pop() } + .expect("invariant is violated: tree is not height-balanced") }; + for _ in (1..height_of_root).rev() { state = step(state); } - let (first, residual) = state; - let code = if let Some(residual) = residual { - O::Vector::code(residual.as_borrowed()) - } else { - O::Vector::code(vector.as_borrowed()) - }; - let bytes = AppendableTuple::serialize(&AppendableTuple { - head, - dis_u_2: code.dis_u_2, - factor_ppc: code.factor_ppc, - factor_ip: code.factor_ip, - factor_err: code.factor_err, - payload: Some(payload), - prefetch: list, - elements: rabitq::pack_to_u64(&code.signs), - }); + let (_, _, AlwaysEqual(first)) = state; let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - tape::append(index, jump_tuple.appendable_first(), &bytes, false); + let (code, delta) = O::build( + vector, + is_residual.then(|| { + vectors::read_for_h1_tuple::( + jump_tuple + .centroid_prefetch() + .iter() + .map(|&id| index.read(id)), + jump_tuple.centroid_head(), + FunctionalAccessor::new(Vec::new(), Vec::extend_from_slice, O::Vector::pack), + ) + }), + ); + + let (prefetch, head) = key; + let serialized = AppendableTuple::serialize(&AppendableTuple { + metadata: [ + code.0.dis_u_2, + code.0.factor_cnt, + code.0.factor_ip, + code.0.factor_err, + ], + delta, + payload: Some(payload), + prefetch, + head, + elements: rabitq::packing::pack_to_u64(&code.1), + }); + + tape::append(index, jump_tuple.appendable_first(), &serialized, false); } diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index e06b9a11..853a2803 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -12,6 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. +#![feature(select_unpredictable)] #![allow(clippy::type_complexity)] mod build; @@ -38,11 +39,11 @@ pub mod types; use always_equal::AlwaysEqual; pub use build::build; -pub use bulkdelete::bulkdelete; +pub use bulkdelete::{bulkdelete, bulkdelete_vectors}; pub use cache::cache; pub use cost::cost; pub use fast_heap::FastHeap; -pub use insert::{insert_index, insert_vector}; +pub use insert::{insert, insert_vector}; pub use maintain::maintain; pub use prefetcher::*; pub use prewarm::prewarm; @@ -185,13 +186,11 @@ pub enum RerankMethod { } pub(crate) struct Branch { - pub head: u16, - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub signs: Vec, + pub code: rabitq::Code, + pub delta: f32, pub prefetch: Vec, + pub head: u16, + pub norm: f32, pub extra: T, } @@ -219,64 +218,33 @@ impl<'b, T, A, B> Fetch for (T, AlwaysEqual<&'b mut (A, B, &'b mut [u32])>) { } } -pub struct Filter { - pub iter: S, - pub filter: P, -} - -impl bool> Sequence for Filter { - type Item = S::Item; - type Inner = S::Inner; - - fn peek(&mut self) -> Option<&Self::Item> { - loop { - let item = self.iter.peek()?; - if (self.filter)(item) { - return self.iter.peek(); - } else { - self.iter.next(); - continue; - } - } - } - fn next(&mut self) -> Option { - loop { - let item = self.iter.peek()?; - if (self.filter)(item) { - return self.iter.next(); - } else { - self.iter.next(); - continue; - } - } - } - - fn into_inner(self) -> Self::Inner { - self.iter.into_inner() +impl<'b, T, A, B, C> Fetch for (T, AlwaysEqual<&'b mut (A, B, C, &'b mut [u32])>) { + fn fetch(&self) -> &[u32] { + let (_, AlwaysEqual((.., list))) = self; + list } } pub trait Sequence { type Item; type Inner: Iterator; - fn peek(&mut self) -> Option<&Self::Item>; + #[must_use] fn next(&mut self) -> Option; - fn next_if(&mut self, predicate: impl FnOnce(&Self::Item) -> bool) -> Option { - let peek = self.peek()?; - if predicate(peek) { self.next() } else { None } - } + #[must_use] + fn peek(&mut self) -> Option<&Self::Item>; + #[must_use] fn into_inner(self) -> Self::Inner; } impl Sequence for BinaryHeap { type Item = T; type Inner = std::vec::IntoIter; - fn peek(&mut self) -> Option<&T> { - >::peek(self) - } fn next(&mut self) -> Option { self.pop() } + fn peek(&mut self) -> Option<&T> { + (self as &Self).peek() + } fn into_inner(self) -> Self::Inner { self.into_vec().into_iter() } @@ -285,21 +253,13 @@ impl Sequence for BinaryHeap { impl Sequence for Peekable { type Item = I::Item; type Inner = Peekable; - fn peek(&mut self) -> Option<&I::Item> { - Peekable::peek(self) - } fn next(&mut self) -> Option { Iterator::next(self) } + fn peek(&mut self) -> Option<&I::Item> { + self.peek() + } fn into_inner(self) -> Self::Inner { self } } - -pub fn seq_filter(heap: impl Sequence, filter: F) -> impl Sequence -where - F: FnMut(&T) -> bool, - T: Ord, -{ - Filter { iter: heap, filter } -} diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index 247f5bdd..f09879df 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -30,30 +30,29 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let height_of_root = meta_tuple.height_of_root(); - let root_first = meta_tuple.root_first(); let freepage_first = meta_tuple.freepage_first(); + + type State = Vec; + let mut state: State = vec![meta_tuple.first()]; + drop(meta_guard); - let state = { - type State = Vec; - let mut state: State = vec![root_first]; - let step = |state: State| { - let mut results = Vec::new(); - for first in state { - tape::read_h1_tape::( - by_next(index, first).inspect(|_| check()), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), - |(), _, first, _| results.push(first), - ); - } - results - }; - for _ in (1..height_of_root).rev() { - state = step(state); + let step = |state: State| { + let mut results = Vec::new(); + for first in state { + tape::read_h1_tape::( + by_next(index, first).inspect(|_| check()), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + |(), _, _, first, _| results.push(first), + ); } - state + results }; + for _ in (1..height_of_root).rev() { + state = step(state); + } + for first in state { let mut jump_guard = index.write(first, false); let jump_bytes = jump_guard.get_mut(1).expect("data corruption"); @@ -86,15 +85,13 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( let mut trace_appendable = Vec::new(); let mut tuples = 0_u64; - let mut callback = id_2(|code: (_, _, _, _, _), head, payload, prefetch: &[_]| { + let mut callback = id_2(|(code, delta): (_, _), head, payload, prefetch: &[_]| { tape.push(Branch { - head, - dis_u_2: code.0, - factor_ppc: code.1, - factor_ip: code.2, - factor_err: code.3, - signs: code.4, + code, + delta, prefetch: prefetch.to_vec(), + head, + norm: 0.0, extra: payload, }); tuples += 1; @@ -112,14 +109,27 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( FunctionalAccessor::new( Vec::<[u8; 16]>::new(), Vec::<[u8; 16]>::extend_from_slice, - |elements: Vec<_>, input: (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32])| { - let unpacked = unpack(&elements); - std::array::from_fn(|i| { - let f = |&x| [x & 1 != 0, x & 2 != 0, x & 4 != 0, x & 8 != 0]; - let signs = unpacked[i].iter().flat_map(f).collect::>(); - (input.0[i], input.1[i], input.2[i], input.3[i], signs) - }) - }, + id_1( + |elements: Vec<_>, (metadata, delta): (&[[f32; 32]; 4], &[f32; 32])| { + let unpacked = unpack(&elements); + std::array::from_fn(|i| { + let f = |&x| [x & 1 != 0, x & 2 != 0, x & 4 != 0, x & 8 != 0]; + let signs = unpacked[i].iter().flat_map(f).collect::>(); + ( + ( + rabitq::CodeMetadata { + dis_u_2: metadata[0][i], + factor_cnt: metadata[1][i], + factor_ip: metadata[2][i], + factor_err: metadata[3][i], + }, + signs, + ), + delta[i], + ) + }) + }, + ), ) }, &mut callback, @@ -128,14 +138,24 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( by_next(index, *jump_tuple.appendable_first()) .inspect(|_| check()) .inspect(|guard| trace_appendable.push(guard.id())), - |code| { - let signs = code - .4 + |metadata, elements, delta| { + let signs = elements .iter() .flat_map(|x| std::array::from_fn::<_, 64, _>(|i| *x & (1 << i) != 0)) .take(dims as _) .collect::>(); - (code.0, code.1, code.2, code.3, signs) + ( + ( + rabitq::CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }, + signs, + ), + delta, + ) }, &mut callback, ); @@ -162,14 +182,17 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( for branch in branches { appendable_tape.push(AppendableTuple { + metadata: [ + branch.code.0.dis_u_2, + branch.code.0.factor_cnt, + branch.code.0.factor_ip, + branch.code.0.factor_err, + ], + elements: rabitq::packing::pack_to_u64(&branch.code.1), + delta: branch.delta, + prefetch: branch.prefetch, head: branch.head, - dis_u_2: branch.dis_u_2, - factor_ppc: branch.factor_ppc, - factor_ip: branch.factor_ip, - factor_err: branch.factor_err, payload: Some(branch.extra), - prefetch: branch.prefetch, - elements: rabitq::pack_to_u64(&branch.signs), }); } diff --git a/crates/algorithm/src/operator.rs b/crates/algorithm/src/operator.rs index 572442d7..7de8d079 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/algorithm/src/operator.rs @@ -14,12 +14,13 @@ use distance::Distance; use half::f16; -use rabitq::binary::{BinaryCode, BinaryLut}; -use rabitq::block::BlockLut; +use rabitq::CodeMetadata; +use rabitq::binary::{BinaryLut, BinaryLutMetadata}; +use rabitq::block::{BlockLut, BlockLutMetadata, STEP}; use simd::Floating; use std::fmt::Debug; use std::marker::PhantomData; -use vector::vect::VectOwned; +use vector::vect::{VectBorrowed, VectOwned}; use vector::{VectorBorrowed, VectorOwned}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -291,79 +292,40 @@ impl> TryAccessor1 } #[derive(Debug)] -pub struct BlockAccessor([u16; 32], PhantomData D>); +pub struct BlockAccessor([u32; 32], F); -impl Default for BlockAccessor { - fn default() -> Self { - Self([0u16; 32], PhantomData) - } -} - -impl - Accessor2< - [u8; 16], - [u8; 16], - (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), - (f32, f32, f32, f32), - > for BlockAccessor +impl> + Accessor2<[u8; 16], [u8; 16], (&[[f32; 32]; 4], &[f32; 32]), BlockLutMetadata> + for BlockAccessor { - type Output = [(f32, f32); 32]; + type Output = [F::Output; 32]; fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { - let t = simd::fast_scan::fast_scan(input, target); - for i in 0..32 { - self.0[i] += t[i]; - } - } + use std::iter::zip; - fn finish( - self, - (dis_u_2, factor_ppc, factor_ip, factor_err): ( - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[f32; 32], - ), - (dis_v_2, b, k, qvector_sum): (f32, f32, f32, f32), - ) -> Self::Output { - std::array::from_fn(|i| { - let rough = dis_u_2[i] - + dis_v_2 - + b * factor_ppc[i] - + ((2.0 * self.0[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = factor_err[i] * dis_v_2.sqrt(); - (rough, err) - }) - } -} - -impl - Accessor2< - [u8; 16], - [u8; 16], - (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), - (f32, f32, f32, f32), - > for BlockAccessor -{ - type Output = [(f32, f32); 32]; - - fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { - let t = simd::fast_scan::fast_scan(input, target); - for i in 0..32 { - self.0[i] += t[i]; + for (input, target) in zip(input.chunks(STEP), target.chunks(STEP)) { + let delta = simd::fast_scan::scan(input, target); + simd::fast_scan::accu(&mut self.0, &delta); } } fn finish( - self, - (_, factor_ppc, factor_ip, factor_err): (&[f32; 32], &[f32; 32], &[f32; 32], &[f32; 32]), - (dis_v_2, b, k, qvector_sum): (f32, f32, f32, f32), + mut self, + (metadata, delta): (&[[f32; 32]; 4], &[f32; 32]), + lut: BlockLutMetadata, ) -> Self::Output { std::array::from_fn(|i| { - let rough = 0.5 * b * factor_ppc[i] - + 0.5 * ((2.0 * self.0[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); - (rough, err) + (self.1).call( + self.0[i], + CodeMetadata { + dis_u_2: metadata[0][i], + factor_cnt: metadata[1][i], + factor_ip: metadata[2][i], + factor_err: metadata[3][i], + }, + delta[i], + lut, + ) }) } } @@ -426,35 +388,23 @@ impl Accessor2 for DistanceAccessor, Dot> { } #[derive(Debug, Clone)] -pub struct ResidualAccessor(Vec); +pub struct CloneAccessor(Vec); -impl Default for ResidualAccessor { +impl Default for CloneAccessor { fn default() -> Self { Self(Vec::new()) } } -impl Accessor2 for ResidualAccessor> { - type Output = VectOwned; - - fn push(&mut self, target: &[f32], input: &[f32]) { - self.0.extend(f32::vector_sub(target, input)); - } - - fn finish(self, (): (), (): ()) -> Self::Output { - VectOwned::new(self.0) - } -} - -impl Accessor2 for ResidualAccessor> { - type Output = VectOwned; +impl Accessor1 for CloneAccessor { + type Output = V; - fn push(&mut self, target: &[f16], input: &[f16]) { - self.0.extend(f16::vector_sub(target, input)); + fn push(&mut self, input: &[V::Element]) { + self.0.extend(input); } - fn finish(self, (): (), (): ()) -> Self::Output { - VectOwned::new(self.0) + fn finish(self, metadata: V::Metadata) -> Self::Output { + V::pack(self.0, metadata) } } @@ -469,11 +419,15 @@ pub trait Vector: VectorOwned { fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); + fn pack(elements: Vec, metadata: Self::Metadata) -> Self; + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut; fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut); fn code(vector: Self::Borrowed<'_>) -> rabitq::Code; + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32; } impl Vector for VectOwned { @@ -507,6 +461,10 @@ impl Vector for VectOwned { (vector.slice(), ()) } + fn pack(elements: Vec, (): Self::Metadata) -> Self { + VectOwned::new(elements) + } + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { rabitq::block::preprocess(vector.slice()) } @@ -518,6 +476,10 @@ impl Vector for VectOwned { fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { rabitq::code(vector.dims(), vector.slice()) } + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { + f32::reduce_sum_of_x2(vector.slice()) + } } impl Vector for VectOwned { @@ -551,6 +513,10 @@ impl Vector for VectOwned { (vector.slice(), ()) } + fn pack(elements: Vec, (): Self::Metadata) -> Self { + VectOwned::new(elements) + } + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { rabitq::block::preprocess(&f16::vector_to_f32(vector.slice())) } @@ -562,6 +528,10 @@ impl Vector for VectOwned { fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { rabitq::code(vector.dims(), &f16::vector_to_f32(vector.slice())) } + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { + f16::reduce_sum_of_x2(vector.slice()) + } } #[derive(Debug, Clone, Copy)] @@ -573,15 +543,6 @@ pub struct Dot; pub trait Operator: 'static + Debug + Copy { type Vector: Vector; - type BlockAccessor: Default - + for<'a> Accessor2< - [u8; 16], - [u8; 16], - (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), - (f32, f32, f32, f32), - Output = [(f32, f32); 32], - >; - type DistanceAccessor: Default + Accessor2< ::Element, @@ -591,18 +552,40 @@ pub trait Operator: 'static + Debug + Copy { Output = Distance, >; - type ResidualAccessor: Default - + Accessor2< - ::Element, - ::Element, - ::Metadata, - ::Metadata, - Output = Self::Vector, - >; - - const SUPPORTS_RESIDUAL: bool; - - fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32); + fn block_access>( + lut: &BlockLut, + f: F, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]>; + + fn binary_access>( + lut: &BinaryLut, + f: F, + ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output; + + fn block_process( + value: u32, + code: CodeMetadata, + lut: BlockLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32); + + fn binary_process( + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32); + + fn build( + vector: ::Borrowed<'_>, + centroid: Option, + ) -> (rabitq::Code, f32); } #[derive(Debug)] @@ -619,63 +602,377 @@ impl Copy for Op {} impl Operator for Op, L2> { type Vector = VectOwned; - type BlockAccessor = BlockAccessor; - type DistanceAccessor = DistanceAccessor, L2>; - type ResidualAccessor = ResidualAccessor>; + fn block_access>( + lut: &BlockLut, + f: F, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]> + { + RAccess::new((&lut.1, lut.0), BlockAccessor([0u32; 32], f)) + } + + fn binary_access>( + lut: &BinaryLut, + mut f: F, + ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::binary::asymmetric_binary_dot_product(elements, &lut.1); + f.call( + value, + CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }, + delta, + lut.0, + ) + } + } + + fn block_process( + value: u32, + code: CodeMetadata, + lut: BlockLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + _: f32, + ) -> (f32, f32) { + if !is_residual { + rabitq::block::half_process_l2(value, code, lut) + } else { + rabitq::block::half_process_l2_residual(value, code, lut, dis_f, delta) + } + } - const SUPPORTS_RESIDUAL: bool = true; + fn binary_process( + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + _: f32, + ) -> (f32, f32) { + if !is_residual { + rabitq::binary::half_process_l2(value, code, lut) + } else { + rabitq::binary::half_process_l2_residual(value, code, lut, dis_f, delta) + } + } - fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32) { - rabitq::binary::process_l2(lut, code) + fn build(vector: VectBorrowed<'_, f32>, centroid: Option) -> (rabitq::Code, f32) { + if let Some(centroid) = centroid { + let residual = VectOwned::new(f32::vector_sub(vector.slice(), centroid.slice())); + let code = Self::Vector::code(residual.as_borrowed()); + let delta = { + use std::iter::zip; + let dims = vector.dims(); + let t = zip(&code.1, centroid.slice()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) + .sum::() + / (dims as f32).sqrt(); + let sum_of_x_2 = code.0.dis_u_2; + let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); + 2.0 * dis_u * t / x_0 + }; + (code, delta) + } else { + let code = Self::Vector::code(vector); + let delta = 0.0; + (code, delta) + } } } impl Operator for Op, Dot> { type Vector = VectOwned; - type BlockAccessor = BlockAccessor; - type DistanceAccessor = DistanceAccessor, Dot>; - type ResidualAccessor = ResidualAccessor>; + fn block_access>( + lut: &BlockLut, + f: F, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]> + { + RAccess::new((&lut.1, lut.0), BlockAccessor([0u32; 32], f)) + } + + fn binary_access>( + lut: &BinaryLut, + mut f: F, + ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::binary::asymmetric_binary_dot_product(elements, &lut.1); + f.call( + value, + CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }, + delta, + lut.0, + ) + } + } - const SUPPORTS_RESIDUAL: bool = false; + fn block_process( + value: u32, + code: CodeMetadata, + lut: BlockLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32) { + if !is_residual { + rabitq::block::half_process_dot(value, code, lut) + } else { + rabitq::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + } + } + + fn binary_process( + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32) { + if !is_residual { + rabitq::binary::half_process_dot(value, code, lut) + } else { + rabitq::binary::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + } + } - fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32) { - rabitq::binary::process_dot(lut, code) + fn build(vector: VectBorrowed<'_, f32>, centroid: Option) -> (rabitq::Code, f32) { + if let Some(centroid) = centroid { + let residual = VectOwned::new(f32::vector_sub(vector.slice(), centroid.slice())); + let code = Self::Vector::code(residual.as_borrowed()); + let delta = { + use std::iter::zip; + let dims = vector.dims(); + let t = zip(&code.1, centroid.slice()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) + .sum::() + / (dims as f32).sqrt(); + let sum_of_x_2 = code.0.dis_u_2; + let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); + dis_u * t / x_0 - f32::reduce_sum_of_xy(residual.slice(), centroid.slice()) + }; + (code, delta) + } else { + let code = Self::Vector::code(vector); + let delta = 0.0; + (code, delta) + } } } impl Operator for Op, L2> { type Vector = VectOwned; - type BlockAccessor = BlockAccessor; - type DistanceAccessor = DistanceAccessor, L2>; - type ResidualAccessor = ResidualAccessor>; + fn block_access>( + lut: &BlockLut, + f: F, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]> + { + RAccess::new((&lut.1, lut.0), BlockAccessor([0u32; 32], f)) + } + + fn binary_access>( + lut: &BinaryLut, + mut f: F, + ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::binary::asymmetric_binary_dot_product(elements, &lut.1); + f.call( + value, + CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }, + delta, + lut.0, + ) + } + } - const SUPPORTS_RESIDUAL: bool = true; + fn block_process( + value: u32, + code: CodeMetadata, + lut: BlockLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + _: f32, + ) -> (f32, f32) { + if !is_residual { + rabitq::block::half_process_l2(value, code, lut) + } else { + rabitq::block::half_process_l2_residual(value, code, lut, dis_f, delta) + } + } - fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32) { - rabitq::binary::process_l2(lut, code) + fn binary_process( + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + _: f32, + ) -> (f32, f32) { + if !is_residual { + rabitq::binary::half_process_l2(value, code, lut) + } else { + rabitq::binary::half_process_l2_residual(value, code, lut, dis_f, delta) + } + } + + fn build(vector: VectBorrowed<'_, f16>, centroid: Option) -> (rabitq::Code, f32) { + if let Some(centroid) = centroid { + let residual = VectOwned::new(f16::vector_sub(vector.slice(), centroid.slice())); + let code = Self::Vector::code(residual.as_borrowed()); + let delta = { + use std::iter::zip; + let dims = vector.dims(); + let t = zip(&code.1, centroid.slice()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num).to_f32()) + .sum::() + / (dims as f32).sqrt(); + let sum_of_x_2 = code.0.dis_u_2; + let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); + 2.0 * dis_u * t / x_0 + }; + (code, delta) + } else { + let code = Self::Vector::code(vector); + let delta = 0.0; + (code, delta) + } } } impl Operator for Op, Dot> { type Vector = VectOwned; - type BlockAccessor = BlockAccessor; - type DistanceAccessor = DistanceAccessor, Dot>; - type ResidualAccessor = ResidualAccessor>; + fn block_access>( + lut: &BlockLut, + f: F, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]> + { + RAccess::new((&lut.1, lut.0), BlockAccessor([0u32; 32], f)) + } + + fn binary_access>( + lut: &BinaryLut, + mut f: F, + ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::binary::asymmetric_binary_dot_product(elements, &lut.1); + f.call( + value, + CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }, + delta, + lut.0, + ) + } + } + + fn block_process( + value: u32, + code: CodeMetadata, + lut: BlockLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32) { + if !is_residual { + rabitq::block::half_process_dot(value, code, lut) + } else { + rabitq::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + } + } + + fn binary_process( + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + is_residual: bool, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32) { + if !is_residual { + rabitq::binary::half_process_dot(value, code, lut) + } else { + rabitq::binary::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + } + } + + fn build(vector: VectBorrowed<'_, f16>, centroid: Option) -> (rabitq::Code, f32) { + if let Some(centroid) = centroid { + let residual = VectOwned::new(f16::vector_sub(vector.slice(), centroid.slice())); + let code = Self::Vector::code(residual.as_borrowed()); + let delta = { + use std::iter::zip; + let dims = vector.dims(); + let t = zip(&code.1, centroid.slice()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num).to_f32()) + .sum::() + / (dims as f32).sqrt(); + let sum_of_x_2 = code.0.dis_u_2; + let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); + dis_u * t / x_0 - f16::reduce_sum_of_xy(residual.slice(), centroid.slice()) + }; + (code, delta) + } else { + let code = Self::Vector::code(vector); + let delta = 0.0; + (code, delta) + } + } +} - const SUPPORTS_RESIDUAL: bool = false; +pub trait Call { + type Output; + + fn call(&mut self, a: A, b: B, c: C, d: D) -> Self::Output; +} + +impl R, R> Call for F { + type Output = R; - fn binary_process(lut: &BinaryLut, code: BinaryCode<'_>) -> (f32, f32) { - rabitq::binary::process_dot(lut, code) + fn call(&mut self, a: A, b: B, c: C, d: D) -> R { + (self)(a, b, c, d) } } diff --git a/crates/algorithm/src/prefetcher.rs b/crates/algorithm/src/prefetcher.rs index ccffdebe..a5510e3d 100644 --- a/crates/algorithm/src/prefetcher.rs +++ b/crates/algorithm/src/prefetcher.rs @@ -25,7 +25,9 @@ where { type R: RelationRead + 'r; + #[must_use] fn next(&mut self) -> Option<(Self::Item, Vec<::ReadGuard<'r>>)>; + #[must_use] fn next_if( &mut self, predicate: impl FnOnce(&Self::Item) -> bool, @@ -68,7 +70,10 @@ where &mut self, predicate: impl FnOnce(&Self::Item) -> bool, ) -> Option<(S::Item, Vec>)> { - let e = self.sequence.next_if(predicate)?; + if !predicate(self.sequence.peek()?) { + return None; + } + let e = self.sequence.next()?; let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); Some((e, list)) } diff --git a/crates/algorithm/src/prewarm.rs b/crates/algorithm/src/prewarm.rs index ec7ba8a3..36b332e1 100644 --- a/crates/algorithm/src/prewarm.rs +++ b/crates/algorithm/src/prewarm.rs @@ -16,22 +16,19 @@ use crate::closure_lifetime_binder::{id_0, id_1, id_2}; use crate::operator::{FunctionalAccessor, Operator}; use crate::tape::{by_directory, by_next}; use crate::tuples::*; -use crate::{Page, PrefetcherSequenceFamily, RelationRead, tape, vectors}; +use crate::{Bump, Page, PrefetcherSequenceFamily, RelationRead, tape, vectors}; use std::fmt::Write; -pub fn prewarm<'r, R: RelationRead, O: Operator>( +pub fn prewarm<'r, 'b: 'r, R: RelationRead, O: Operator>( index: &'r R, height: i32, + bump: &'b impl Bump, mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, ) -> String { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let height_of_root = meta_tuple.height_of_root(); - let root_prefetch = meta_tuple.root_prefetch().to_vec(); - let root_head = meta_tuple.root_head(); - let root_first = meta_tuple.root_first(); - drop(meta_guard); let mut message = String::new(); writeln!(message, "height of root: {height_of_root}").unwrap(); @@ -39,19 +36,23 @@ pub fn prewarm<'r, R: RelationRead, O: Operator>( if prewarm_max_height > height_of_root { return message; } + type State = Vec; let mut state: State = { let mut results = Vec::new(); - { - let list = root_prefetch.into_iter().map(|id| index.read(id)); - vectors::read_for_h1_tuple::(root_head, list, ()); - results.push(root_first); - } + let prefetch = bump.alloc_slice(meta_tuple.centroid_prefetch()); + let head = meta_tuple.centroid_head(); + let first = meta_tuple.first(); + vectors::read_for_h1_tuple::(prefetch.iter().map(|&id| index.read(id)), head, ()); + results.push(first); writeln!(message, "------------------------").unwrap(); writeln!(message, "number of nodes: {}", results.len()).unwrap(); writeln!(message, "number of pages: {}", 1).unwrap(); results }; + + drop(meta_guard); + let mut step = |state: State| { let mut counter = 0_usize; let mut results = Vec::new(); @@ -59,9 +60,12 @@ pub fn prewarm<'r, R: RelationRead, O: Operator>( tape::read_h1_tape::( by_next(index, first).inspect(|_| counter += 1), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), - |(), head, first, prefetch| { - let list = prefetch.iter().map(|&id| index.read(id)); - vectors::read_for_h1_tuple::(head, list, ()); + |(), head, _, first, prefetch| { + vectors::read_for_h1_tuple::( + prefetch.iter().map(|&id| index.read(id)), + head, + (), + ); results.push(first); }, ); @@ -71,9 +75,11 @@ pub fn prewarm<'r, R: RelationRead, O: Operator>( writeln!(message, "number of pages: {counter}").unwrap(); results }; + for _ in (std::cmp::max(1, prewarm_max_height)..height_of_root).rev() { state = step(state); } + if prewarm_max_height == 0 { let mut counter = 0_usize; let mut results = Vec::new(); @@ -92,7 +98,7 @@ pub fn prewarm<'r, R: RelationRead, O: Operator>( ); tape::read_appendable_tape::( by_next(index, jump_tuple.appendable_first()).inspect(|_| counter += 1), - |_| (), + |_, _, _| (), id_2(|_, _, _, _| { results.push(()); }), @@ -102,5 +108,6 @@ pub fn prewarm<'r, R: RelationRead, O: Operator>( writeln!(message, "number of nodes: {}", results.len()).unwrap(); writeln!(message, "number of pages: {counter}").unwrap(); } + message } diff --git a/crates/algorithm/src/rerank.rs b/crates/algorithm/src/rerank.rs index 69b86983..d413bea8 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/algorithm/src/rerank.rs @@ -56,11 +56,11 @@ where type Item = (Distance, NonZero); fn next(&mut self) -> Option { - while let Some(((_, AlwaysEqual(&mut (payload, head, ..))), list)) = self + while let Some(((_, AlwaysEqual(&mut (payload, head, ..))), prefetch)) = self .prefetcher .next_if(|((d, _), ..)| Some(*d) > self.cache.peek().map(|(d, ..)| *d)) { - if let Some(distance) = (self.f)(payload, list, head) { + if let Some(distance) = (self.f)(payload, prefetch, head) { self.cache.push((Reverse(distance), AlwaysEqual(payload))); }; } @@ -92,10 +92,10 @@ pub fn rerank_index< Reranker { prefetcher, cache: BinaryHeap::new(), - f: id_4::<_, P::R, _, _, _>(move |payload, list, head| { + f: id_4::<_, P::R, _, _, _>(move |payload, prefetch, head| { vectors::read_for_h0_tuple::( + prefetch.into_iter(), head, - list.into_iter(), payload, LTryAccess::new( O::Vector::unpack(vector.as_borrowed()), diff --git a/crates/algorithm/src/search.rs b/crates/algorithm/src/search.rs index b5c2a7f5..5f59855f 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/algorithm/src/search.rs @@ -31,7 +31,7 @@ type Extra<'b> = &'b mut (NonZero, u16, &'b mut [u32]); #[allow(clippy::too_many_arguments)] pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( index: &'r R, - vector: O::Vector, + vector: ::Borrowed<'_>, probes: Vec, epsilon: f32, bump: &'b impl Bump, @@ -44,7 +44,7 @@ pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); let height_of_root = meta_tuple.height_of_root(); - assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); + assert_eq!(dims, vector.dims(), "unmatched dimensions"); if height_of_root as usize != 1 + probes.len() { panic!( "usage: need {} probes, but {} probes provided", @@ -52,127 +52,99 @@ pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( probes.len() ); } - let root_prefetch = meta_tuple.root_prefetch().to_vec(); - let root_head = meta_tuple.root_head(); - let root_first = meta_tuple.root_first(); - drop(meta_guard); - let default_lut = if !is_residual { - Some(O::Vector::preprocess(vector.as_borrowed())) + type State = Vec<(Reverse, AlwaysEqual, AlwaysEqual)>; + let mut state: State = if !is_residual { + let first = meta_tuple.first(); + // it's safe to leave it a fake value + vec![( + Reverse(Distance::ZERO), + AlwaysEqual(0.0), + AlwaysEqual(first), + )] } else { - None + let prefetch = bump.alloc_slice(meta_tuple.centroid_prefetch()); + let head = meta_tuple.centroid_head(); + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + let distance = vectors::read_for_h1_tuple::( + prefetch.iter().map(|&id| index.read(id)), + head, + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + ); + vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] }; - type State = Vec<(u32, Option<::Vector>)>; - let mut state: State = vec![{ - if is_residual { - let list = root_prefetch.into_iter().map(|id| index.read(id)); - let residual = vectors::read_for_h1_tuple::( - root_head, - list, - LAccess::new( - O::Vector::unpack(vector.as_borrowed()), - O::ResidualAccessor::default(), - ), - ); - (root_first, Some(residual)) - } else { - (root_first, None) - } - }]; - let mut step = |state: State| { + drop(meta_guard); + let lut = O::Vector::preprocess(vector); + + let mut step = |state: State| { let mut results = LinkedVec::new(); - for (first, residual) in state { - let block_lut = if let Some(residual) = residual { - &O::Vector::block_preprocess(residual.as_borrowed()) - } else if let Some((block_lut, _)) = default_lut.as_ref() { - block_lut - } else { - unreachable!() + for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { + let process = |value, code, delta, lut| { + O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) }; tape::read_h1_tape::( by_next(index, first), - || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), - |(rough, err), head, first, prefetch| { + || O::block_access(&lut.0, process), + |(rough, err), head, norm, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( Reverse(lowerbound), - AlwaysEqual(bump.alloc((first, head, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc((first, head, norm, bump.alloc_slice(prefetch)))), )); }, ); } let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); - let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - let vector = vector.as_borrowed(); + let mut cache = BinaryHeap::<(_, _, _)>::new(); std::iter::from_fn(move || { - while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, ..))), list)) = - heap.next_if(|(d, ..)| Some(*d) > cache.peek().map(|(d, ..)| *d)) + while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, norm, ..))), prefetch)) = + heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { - if is_residual { - let (distance, residual) = vectors::read_for_h1_tuple::( - head, - list.into_iter(), - LAccess::new( - O::Vector::unpack(vector), - ( - O::DistanceAccessor::default(), - O::ResidualAccessor::default(), - ), - ), - ); - cache.push(( - Reverse(distance), - AlwaysEqual(first), - AlwaysEqual(Some(residual)), - )); - } else { - let distance = vectors::read_for_h1_tuple::( - head, - list.into_iter(), - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), - ); - cache.push((Reverse(distance), AlwaysEqual(first), AlwaysEqual(None))); - } + let distance = vectors::read_for_h1_tuple::( + prefetch.into_iter(), + head, + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + ); + cache.push((Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))); } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; - Some((first, mean)) + cache.pop() }) }; + for i in (1..height_of_root).rev() { state = step(state).take(probes[i as usize - 1] as _).collect(); } let mut results = LinkedVec::new(); - for (first, residual) in state { - let (block_lut, binary_lut) = - if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { - &O::Vector::preprocess(residual) - } else if let Some(lut) = default_lut.as_ref() { - lut - } else { - unreachable!() - }; + for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { + let block_process = |value, code, delta, lut| { + O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) + }; + let binary_process = |value, code, delta, lut| { + O::binary_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) + }; let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let mut callback = id_2(|(rough, err), mean, payload, prefetch| { + let directory = + tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); + let mut callback = id_2(|(rough, err), head, payload, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( (Reverse(lowerbound), AlwaysEqual(())), - AlwaysEqual(bump.alloc((payload, mean, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc((payload, head, bump.alloc_slice(prefetch)))), )); }); - let directory = - tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); tape::read_frozen_tape::( by_directory(&mut prefetch_h0_tuples, directory), - || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), + || O::block_access(&lut.0, block_process), &mut callback, ); tape::read_appendable_tape::( by_next(index, jump_tuple.appendable_first()), - |code| O::binary_process(binary_lut, code), + O::binary_access(&lut.1, binary_process), &mut callback, ); } @@ -182,7 +154,7 @@ pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( #[allow(clippy::too_many_arguments)] pub fn maxsim_search<'r, 'b: 'r, R: RelationRead, O: Operator>( index: &'r R, - vector: O::Vector, + vector: ::Borrowed<'_>, probes: Vec, epsilon: f32, mut threshold: u32, @@ -202,7 +174,7 @@ pub fn maxsim_search<'r, 'b: 'r, R: RelationRead, O: Operator>( let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); let height_of_root = meta_tuple.height_of_root(); - assert_eq!(dims, vector.as_borrowed().dims(), "unmatched dimensions"); + assert_eq!(dims, vector.dims(), "unmatched dimensions"); if height_of_root as usize != 1 + probes.len() { panic!( "usage: need {} probes, but {} probes provided", @@ -210,136 +182,108 @@ pub fn maxsim_search<'r, 'b: 'r, R: RelationRead, O: Operator>( probes.len() ); } - let root_prefetch = meta_tuple.root_prefetch().to_vec(); - let root_head = meta_tuple.root_head(); - let root_first = meta_tuple.root_first(); - drop(meta_guard); - let default_lut = if !is_residual { - Some(O::Vector::preprocess(vector.as_borrowed())) + type State = Vec<(Reverse, AlwaysEqual, AlwaysEqual)>; + let mut state: State = if !is_residual { + let first = meta_tuple.first(); + // it's safe to leave it a fake value + vec![( + Reverse(Distance::ZERO), + AlwaysEqual(0.0), + AlwaysEqual(first), + )] } else { - None + let prefetch = bump.alloc_slice(meta_tuple.centroid_prefetch()); + let head = meta_tuple.centroid_head(); + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + let distance = vectors::read_for_h1_tuple::( + prefetch.iter().map(|&id| index.read(id)), + head, + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + ); + vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] }; - type State = Vec<(u32, Option<::Vector>)>; - let mut state: State = vec![{ - if is_residual { - let list = root_prefetch.into_iter().map(|id| index.read(id)); - let residual = vectors::read_for_h1_tuple::( - root_head, - list, - LAccess::new( - O::Vector::unpack(vector.as_borrowed()), - O::ResidualAccessor::default(), - ), - ); - (root_first, Some(residual)) - } else { - (root_first, None) - } - }]; - let mut step = |state: State| { + drop(meta_guard); + let lut = O::Vector::preprocess(vector); + + let mut step = |state: State| { let mut results = LinkedVec::new(); - for (first, residual) in state { - let block_lut = if let Some(residual) = residual { - &O::Vector::block_preprocess(residual.as_borrowed()) - } else if let Some((block_lut, _)) = default_lut.as_ref() { - block_lut - } else { - unreachable!() + for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { + let process = |value, code, delta, lut| { + O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) }; tape::read_h1_tape::( by_next(index, first), - || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), - |(rough, err), head, first, prefetch| { + || O::block_access(&lut.0, process), + |(rough, err), head, norm, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( Reverse(lowerbound), - AlwaysEqual(bump.alloc((first, head, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc((first, head, norm, bump.alloc_slice(prefetch)))), )); }, ); } let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); - let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - let vector = vector.as_borrowed(); + let mut cache = BinaryHeap::<(_, _, _)>::new(); std::iter::from_fn(move || { - while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, ..))), list)) = - heap.next_if(|(d, ..)| Some(*d) > cache.peek().map(|(d, ..)| *d)) + while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, norm, ..))), prefetch)) = + heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { - if is_residual { - let (distance, residual) = vectors::read_for_h1_tuple::( - head, - list.into_iter(), - LAccess::new( - O::Vector::unpack(vector), - ( - O::DistanceAccessor::default(), - O::ResidualAccessor::default(), - ), - ), - ); - cache.push(( - Reverse(distance), - AlwaysEqual(first), - AlwaysEqual(Some(residual)), - )); - } else { - let distance = vectors::read_for_h1_tuple::( - head, - list.into_iter(), - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), - ); - cache.push((Reverse(distance), AlwaysEqual(first), AlwaysEqual(None))); - } + let distance = vectors::read_for_h1_tuple::( + prefetch.into_iter(), + head, + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + ); + cache.push((Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))); } - let (Reverse(distance), AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; - Some((first, mean, distance)) + cache.pop() }) }; + let mut it = None; for i in (1..height_of_root).rev() { - let it = it.insert(step(state)).map(|(first, mean, _)| (first, mean)); + let it = it.insert(step(state)); state = it.take(probes[i as usize - 1] as _).collect(); } let mut results = LinkedVec::new(); - for (first, residual) in state { - let (block_lut, binary_lut) = - if let Some(residual) = residual.as_ref().map(|x| x.as_borrowed()) { - &O::Vector::preprocess(residual) - } else if let Some(lut) = default_lut.as_ref() { - lut - } else { - unreachable!() - }; + for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { + let block_process = |value, code, delta, lut| { + O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) + }; + let binary_process = |value, code, delta, lut| { + O::binary_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) + }; let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let mut callback = id_2(|(rough, err), mean, payload, prefetch| { + let directory = + tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); + let mut callback = id_2(|(rough, err), head, payload, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); let rough = Distance::from_f32(rough); results.push(( (Reverse(lowerbound), AlwaysEqual(rough)), - AlwaysEqual(bump.alloc((payload, mean, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc((payload, head, bump.alloc_slice(prefetch)))), )); }); - let directory = - tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); tape::read_frozen_tape::( by_directory(&mut prefetch_h0_tuples, directory), - || RAccess::new((&block_lut.1, block_lut.0), O::BlockAccessor::default()), + || O::block_access(&lut.0, block_process), &mut callback, ); tape::read_appendable_tape::( by_next(index, jump_tuple.appendable_first()), - |code| O::binary_process(binary_lut, code), + O::binary_access(&lut.1, binary_process), &mut callback, ); threshold = threshold.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as _); } let mut estimation_by_threshold = Distance::NEG_INFINITY; - for (first, _, distance) in it.into_iter().flatten() { + for (Reverse(distance), AlwaysEqual(_), AlwaysEqual(first)) in it.into_iter().flatten() { if threshold == 0 { break; } diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index ff01b45b..f1c7e496 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -15,7 +15,6 @@ use crate::operator::Accessor1; use crate::tuples::*; use crate::{Page, PageGuard, PrefetcherSequenceFamily, RelationRead, RelationWrite}; -use rabitq::binary::BinaryCode; use std::marker::PhantomData; use std::num::NonZero; @@ -208,14 +207,10 @@ where pub fn read_h1_tape<'r, R, A, T>( iter: impl Iterator>, accessor: impl Fn() -> A, - mut callback: impl for<'a> FnMut(T, u16, u32, &'a [u32]), + mut callback: impl for<'a> FnMut(T, u16, f32, u32, &'a [u32]), ) where R: RelationRead + 'r, - A: for<'a> Accessor1< - [u8; 16], - (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), - Output = [T; 32], - >, + A: for<'a> Accessor1<[u8; 16], (&'a [[f32; 32]; 4], &'a [f32; 32]), Output = [T; 32]>, { let mut x = None; for guard in iter { @@ -226,11 +221,17 @@ pub fn read_h1_tape<'r, R, A, T>( H1TupleReader::_0(tuple) => { let mut x = x.take().unwrap_or_else(&accessor); x.push(tuple.elements()); - let values = x.finish(tuple.metadata()); + let values = x.finish((tuple.metadata(), tuple.delta())); let prefetch: [_; 32] = fix_0(tuple.prefetch()); for (j, value) in values.into_iter().enumerate() { if j < tuple.len() as usize { - callback(value, tuple.head()[j], tuple.first()[j], fix_1(prefetch[j])); + callback( + value, + tuple.head()[j], + tuple.norm()[j], + tuple.first()[j], + fix_1(prefetch[j]), + ); } } } @@ -248,11 +249,7 @@ pub fn read_frozen_tape<'r, R, A, T>( mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), ) where R: RelationRead + 'r, - A: for<'a> Accessor1< - [u8; 16], - (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]), - Output = [T; 32], - >, + A: for<'a> Accessor1<[u8; 16], (&'a [[f32; 32]; 4], &'a [f32; 32]), Output = [T; 32]>, { let mut x = None; for guard in iter { @@ -263,11 +260,11 @@ pub fn read_frozen_tape<'r, R, A, T>( FrozenTupleReader::_0(tuple) => { let mut x = x.take().unwrap_or_else(&accessor); x.push(tuple.elements()); - let values = x.finish(tuple.metadata()); + let values = x.finish((tuple.metadata(), tuple.delta())); let prefetch: [_; 32] = fix_0(tuple.prefetch()); for (j, value) in values.into_iter().enumerate() { if let Some(payload) = tuple.payload()[j] { - callback(value, tuple.mean()[j], payload, fix_1(prefetch[j])); + callback(value, tuple.head()[j], payload, fix_1(prefetch[j])); } } } @@ -281,7 +278,7 @@ pub fn read_frozen_tape<'r, R, A, T>( pub fn read_appendable_tape<'r, R, T>( iter: impl Iterator>, - mut access: impl for<'a> FnMut(BinaryCode<'a>) -> T, + mut access: impl for<'a> FnMut([f32; 4], &'a [u64], f32) -> T, mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), ) where R: RelationRead + 'r, @@ -291,7 +288,7 @@ pub fn read_appendable_tape<'r, R, T>( let bytes = guard.get(i).expect("data corruption"); let tuple = AppendableTuple::deserialize_ref(bytes); if let Some(payload) = tuple.payload() { - let value = access(tuple.code()); + let value = access(tuple.metadata(), tuple.elements(), tuple.delta()); callback(value, tuple.head(), payload, tuple.prefetch()); } } diff --git a/crates/algorithm/src/tape_writer.rs b/crates/algorithm/src/tape_writer.rs index 5a486ca9..119c7ca5 100644 --- a/crates/algorithm/src/tape_writer.rs +++ b/crates/algorithm/src/tape_writer.rs @@ -81,18 +81,23 @@ where pub fn push(&mut self, branch: Branch) { self.branches.push(branch); if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { - let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + let mut remain = + padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.code.1))); loop { let freespace = self.tape.freespace(); if H1Tuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { self.tape.tape_put(H1Tuple::_0 { + metadata: [ + chunk.each_ref().map(|x| x.code.0.dis_u_2), + chunk.each_ref().map(|x| x.code.0.factor_cnt), + chunk.each_ref().map(|x| x.code.0.factor_ip), + chunk.each_ref().map(|x| x.code.0.factor_err), + ], + delta: chunk.each_ref().map(|x| x.delta), + prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), + norm: chunk.each_ref().map(|x| x.norm), head: chunk.each_ref().map(|x| x.head), - dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), - factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), - factor_ip: chunk.each_ref().map(|x| x.factor_ip), - factor_err: chunk.each_ref().map(|x| x.factor_err), first: chunk.each_ref().map(|x| x.extra), - prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), len: chunk.len() as _, elements: remain, }); @@ -118,18 +123,22 @@ where if chunk.is_empty() { return; } - let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.code.1))); loop { let freespace = tape.freespace(); if H1Tuple::estimate_size_0(prefetch, remain.len()) <= freespace as usize { tape.tape_put(H1Tuple::_0 { + metadata: [ + any_pack(chunk.iter().map(|x| x.code.0.dis_u_2)), + any_pack(chunk.iter().map(|x| x.code.0.factor_cnt)), + any_pack(chunk.iter().map(|x| x.code.0.factor_ip)), + any_pack(chunk.iter().map(|x| x.code.0.factor_err)), + ], + delta: any_pack(chunk.iter().map(|x| x.delta)), + prefetch: fix(chunk.iter().map(|x| x.prefetch.as_slice())), head: any_pack(chunk.iter().map(|x| x.head)), - dis_u_2: any_pack(chunk.iter().map(|x| x.dis_u_2)), - factor_ppc: any_pack(chunk.iter().map(|x| x.factor_ppc)), - factor_ip: any_pack(chunk.iter().map(|x| x.factor_ip)), - factor_err: any_pack(chunk.iter().map(|x| x.factor_err)), + norm: any_pack(chunk.iter().map(|x| x.norm)), first: any_pack(chunk.iter().map(|x| x.extra)), - prefetch: fix(chunk.iter().map(|x| x.prefetch.as_slice())), len: chunk.len() as _, elements: remain, }); @@ -171,18 +180,22 @@ where pub fn push(&mut self, branch: Branch>) { self.branches.push(branch); if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { - let mut remain = padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))); + let mut remain = + padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.code.1))); loop { let freespace = self.tape.freespace(); if FrozenTuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { self.tape.tape_put(FrozenTuple::_0 { + metadata: [ + chunk.each_ref().map(|x| x.code.0.dis_u_2), + chunk.each_ref().map(|x| x.code.0.factor_cnt), + chunk.each_ref().map(|x| x.code.0.factor_ip), + chunk.each_ref().map(|x| x.code.0.factor_err), + ], + delta: chunk.each_ref().map(|x| x.delta), + prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), head: chunk.each_ref().map(|x| x.head), - dis_u_2: chunk.each_ref().map(|x| x.dis_u_2), - factor_ppc: chunk.each_ref().map(|x| x.factor_ppc), - factor_ip: chunk.each_ref().map(|x| x.factor_ip), - factor_err: chunk.each_ref().map(|x| x.factor_err), payload: chunk.each_ref().map(|x| Some(x.extra)), - prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), elements: remain, }); break; diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 912d7f6c..3d103fc6 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -13,7 +13,6 @@ // Copyright (c) 2025 TensorChord Inc. use crate::operator::Vector; -use rabitq::binary::BinaryCode; use std::marker::PhantomData; use std::num::NonZero; use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; @@ -21,7 +20,7 @@ use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 7; +const VERSION: u64 = 8; pub trait Tuple: 'static { fn serialize(&self) -> Vec; @@ -45,20 +44,18 @@ struct MetaTupleHeader { height_of_root: u32, is_residual: Bool, rerank_in_heap: Bool, + cells_s: u16, + cells_e: u16, _padding_0: [ZeroU8; 2], vectors_first: u32, - // raw vector - root_prefetch_s: u16, - root_prefetch_e: u16, - root_head: u16, - _padding_1: [ZeroU8; 2], - // for meta tuple, it's pointers to next level - root_first: u32, freepage_first: u32, - // statistics - cells_s: u16, - cells_e: u16, - _padding_2: [ZeroU8; 4], + _padding_1: [ZeroU8; 2], + // tree + centroid_prefetch_s: u16, + centroid_prefetch_e: u16, + centroid_head: u16, + centroid_norm: f32, + first: u32, } pub struct MetaTuple { @@ -66,12 +63,13 @@ pub struct MetaTuple { pub height_of_root: u32, pub is_residual: bool, pub rerank_in_heap: bool, + pub cells: Vec, pub vectors_first: u32, - pub root_prefetch: Vec, - pub root_head: u16, - pub root_first: u32, pub freepage_first: u32, - pub cells: Vec, + pub centroid_prefetch: Vec, + pub centroid_head: u16, + pub centroid_norm: f32, + pub first: u32, } impl Tuple for MetaTuple { @@ -84,27 +82,31 @@ impl Tuple for MetaTuple { height_of_root, is_residual, rerank_in_heap, + cells, vectors_first, - root_prefetch, - root_head, - root_first, freepage_first, - cells, + centroid_prefetch, + centroid_head, + centroid_norm, + first, } => { buffer.extend((MAGIC as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); - let root_prefetch_s = buffer.len() as u16; - buffer.extend(root_prefetch.as_bytes()); - let root_prefetch_e = buffer.len() as u16; - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } + // cells let cells_s = buffer.len() as u16; buffer.extend(cells.as_bytes()); let cells_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // centroid_prefetch + let centroid_prefetch_s = buffer.len() as u16; + buffer.extend(centroid_prefetch.as_bytes()); + let centroid_prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header buffer[size_of::()..][..size_of::()].copy_from_slice( MetaTupleHeader { version: VERSION, @@ -112,17 +114,17 @@ impl Tuple for MetaTuple { height_of_root: *height_of_root, is_residual: (*is_residual).into(), rerank_in_heap: (*rerank_in_heap).into(), - _padding_0: Default::default(), - vectors_first: *vectors_first, - root_prefetch_s, - root_prefetch_e, - root_head: *root_head, - _padding_1: Default::default(), - root_first: *root_first, - freepage_first: *freepage_first, cells_s, cells_e, - _padding_2: Default::default(), + vectors_first: *vectors_first, + freepage_first: *freepage_first, + centroid_prefetch_s, + centroid_prefetch_e, + centroid_head: *centroid_head, + centroid_norm: *centroid_norm, + first: *first, + _padding_0: Default::default(), + _padding_1: Default::default(), } .as_bytes(), ); @@ -143,11 +145,12 @@ impl WithReader for MetaTuple { panic!("deserialization: bad version number"); } let header: &MetaTupleHeader = checker.prefix(size_of::()); - let root_prefetch = checker.bytes(header.root_prefetch_s, header.root_prefetch_e); + let centroid_prefetch = + checker.bytes(header.centroid_prefetch_s, header.centroid_prefetch_e); let cells = checker.bytes(header.cells_s, header.cells_e); MetaTupleReader { header, - root_prefetch, + centroid_prefetch, cells, } } @@ -159,7 +162,7 @@ impl WithReader for MetaTuple { #[derive(Debug, Clone, Copy)] pub struct MetaTupleReader<'a> { header: &'a MetaTupleHeader, - root_prefetch: &'a [u32], + centroid_prefetch: &'a [u32], cells: &'a [u32], } @@ -176,23 +179,26 @@ impl<'a> MetaTupleReader<'a> { pub fn rerank_in_heap(self) -> bool { self.header.rerank_in_heap.into() } + pub fn cells(self) -> &'a [u32] { + self.cells + } pub fn vectors_first(self) -> u32 { self.header.vectors_first } - pub fn root_prefetch(self) -> &'a [u32] { - self.root_prefetch + pub fn freepage_first(self) -> u32 { + self.header.freepage_first } - pub fn root_head(self) -> u16 { - self.header.root_head + pub fn centroid_prefetch(self) -> &'a [u32] { + self.centroid_prefetch } - pub fn root_first(self) -> u32 { - self.header.root_first + pub fn centroid_head(self) -> u16 { + self.header.centroid_head } - pub fn freepage_first(self) -> u32 { - self.header.freepage_first + pub fn centroid_norm(self) -> f32 { + self.header.centroid_norm } - pub fn cells(self) -> &'a [u32] { - self.cells + pub fn first(self) -> u32 { + self.header.first } } @@ -317,17 +323,20 @@ impl Tuple for VectorTuple { } => { buffer.extend((0 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); + // metadata let metadata_s = buffer.len() as u16; buffer.extend(metadata.as_bytes()); while buffer.len() % ALIGN != 0 { buffer.push(0); } + // elements let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // header buffer[size_of::()..][..size_of::()].copy_from_slice( VectorTupleHeader0 { payload: *payload, @@ -346,12 +355,14 @@ impl Tuple for VectorTuple { } => { buffer.extend((1 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // header buffer[size_of::()..][..size_of::()].copy_from_slice( VectorTupleHeader1 { payload: *payload, @@ -475,8 +486,11 @@ impl DirectoryTuple { } pub fn fit_1(freespace: u16) -> Option { let mut freespace = freespace as isize; + freespace &= !(ALIGN - 1) as isize; freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; if freespace >= 0 { Some(freespace as usize / size_of::()) } else { @@ -492,12 +506,14 @@ impl Tuple for DirectoryTuple { Self::_0 { elements } => { buffer.extend((0 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // header buffer[size_of::()..][..size_of::()].copy_from_slice( DirectoryTupleHeader0 { elements_s, @@ -510,12 +526,14 @@ impl Tuple for DirectoryTuple { Self::_1 { elements } => { buffer.extend((1 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // header buffer[size_of::()..][..size_of::()].copy_from_slice( DirectoryTupleHeader1 { elements_s, @@ -586,14 +604,13 @@ impl<'a> DirectoryTupleReader1<'a> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct H1TupleHeader0 { - head: [u16; 32], - dis_u_2: [f32; 32], - factor_ppc: [f32; 32], - factor_ip: [f32; 32], - factor_err: [f32; 32], - first: [u32; 32], + metadata: [[f32; 32]; 4], + delta: [f32; 32], prefetch_s: u16, prefetch_e: u16, + head: [u16; 32], + norm: [f32; 32], + first: [u32; 32], len: u32, elements_s: u16, elements_e: u16, @@ -612,13 +629,12 @@ struct H1TupleHeader1 { #[derive(Debug, Clone, PartialEq)] pub enum H1Tuple { _0 { + metadata: [[f32; 32]; 4], + delta: [f32; 32], + prefetch: Vec<[u32; 32]>, head: [u16; 32], - dis_u_2: [f32; 32], - factor_ppc: [f32; 32], - factor_ip: [f32; 32], - factor_err: [f32; 32], + norm: [f32; 32], first: [u32; 32], - prefetch: Vec<[u32; 32]>, len: u32, elements: Vec<[u8; 16]>, }, @@ -638,9 +654,13 @@ impl H1Tuple { } pub fn fit_1(prefetch: usize, freespace: u16) -> Option { let mut freespace = freespace as isize; + freespace &= !(ALIGN - 1) as isize; freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; freespace -= (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN) as isize; + freespace &= !(ALIGN - 1) as isize; if freespace >= 0 { Some(freespace as usize / size_of::<[u8; 16]>()) } else { @@ -655,10 +675,9 @@ impl Tuple for H1Tuple { match self { Self::_0 { head, - dis_u_2, - factor_ppc, - factor_ip, - factor_err, + norm, + metadata, + delta, first, prefetch, len, @@ -666,25 +685,27 @@ impl Tuple for H1Tuple { } => { buffer.extend((0 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); + // prefetch let prefetch_s = buffer.len() as u16; buffer.extend(prefetch.as_bytes()); let prefetch_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // elements let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // header buffer[size_of::()..][..size_of::()].copy_from_slice( H1TupleHeader0 { head: *head, - dis_u_2: *dis_u_2, - factor_ppc: *factor_ppc, - factor_ip: *factor_ip, - factor_err: *factor_err, + norm: *norm, + metadata: *metadata, + delta: *delta, first: *first, len: *len, prefetch_s, @@ -699,12 +720,14 @@ impl Tuple for H1Tuple { Self::_1 { elements } => { buffer.extend((1 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // header buffer[size_of::()..][..size_of::()].copy_from_slice( H1TupleHeader1 { elements_s, @@ -767,26 +790,27 @@ pub struct H1TupleReader1<'a> { } impl<'a> H1TupleReader0<'a> { - pub fn len(self) -> u32 { - self.header.len - } - pub fn head(self) -> &'a [u16] { - &self.header.head[..self.header.len as usize] + pub fn metadata(self) -> &'a [[f32; 32]; 4] { + &self.header.metadata } - pub fn metadata(self) -> (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]) { - ( - &self.header.dis_u_2, - &self.header.factor_ppc, - &self.header.factor_ip, - &self.header.factor_err, - ) - } - pub fn first(self) -> &'a [u32] { - &self.header.first[..self.header.len as usize] + pub fn delta(self) -> &'a [f32; 32] { + &self.header.delta } pub fn prefetch(self) -> &'a [[u32; 32]] { self.prefetch } + pub fn head(self) -> &'a [u16; 32] { + &self.header.head + } + pub fn norm(self) -> &'a [f32; 32] { + &self.header.norm + } + pub fn first(self) -> &'a [u32; 32] { + &self.header.first + } + pub fn len(self) -> u32 { + self.header.len + } pub fn elements(&self) -> &'a [[u8; 16]] { self.elements } @@ -801,6 +825,10 @@ impl<'a> H1TupleReader1<'a> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct JumpTupleHeader { + centroid_prefetch_s: u16, + centroid_prefetch_e: u16, + centroid_head: u16, + _padding_0: [ZeroU8; 2], directory_first: u32, appendable_first: u32, tuples: u64, @@ -808,6 +836,8 @@ struct JumpTupleHeader { #[derive(Debug, Clone, PartialEq)] pub struct JumpTuple { + pub centroid_prefetch: Vec, + pub centroid_head: u16, pub directory_first: u32, pub appendable_first: u32, pub tuples: u64, @@ -815,13 +845,29 @@ pub struct JumpTuple { impl Tuple for JumpTuple { fn serialize(&self) -> Vec { - JumpTupleHeader { - directory_first: self.directory_first, - appendable_first: self.appendable_first, - tuples: self.tuples, + let mut buffer = Vec::::new(); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // centroid_prefetch + let centroid_prefetch_s = buffer.len() as u16; + buffer.extend(self.centroid_prefetch.as_bytes()); + let centroid_prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); } - .as_bytes() - .to_vec() + // header + buffer[..size_of::()].copy_from_slice( + JumpTupleHeader { + centroid_prefetch_s, + centroid_prefetch_e, + centroid_head: self.centroid_head, + directory_first: self.directory_first, + appendable_first: self.appendable_first, + tuples: self.tuples, + _padding_0: Default::default(), + } + .as_bytes(), + ); + buffer } } @@ -830,7 +876,12 @@ impl WithReader for JumpTuple { fn deserialize_ref(source: &[u8]) -> JumpTupleReader<'_> { let checker = RefChecker::new(source); let header: &JumpTupleHeader = checker.prefix(0_u16); - JumpTupleReader { header } + let centroid_prefetch = + checker.bytes(header.centroid_prefetch_s, header.centroid_prefetch_e); + JumpTupleReader { + header, + centroid_prefetch, + } } } @@ -846,9 +897,16 @@ impl WithWriter for JumpTuple { #[derive(Debug, Clone, Copy)] pub struct JumpTupleReader<'a> { header: &'a JumpTupleHeader, + centroid_prefetch: &'a [u32], } -impl JumpTupleReader<'_> { +impl<'a> JumpTupleReader<'a> { + pub fn centroid_prefetch(self) -> &'a [u32] { + self.centroid_prefetch + } + pub fn centroid_head(self) -> u16 { + self.header.centroid_head + } pub fn directory_first(self) -> u32 { self.header.directory_first } @@ -880,14 +938,13 @@ impl JumpTupleWriter<'_> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FrozenTupleHeader0 { - head: [u16; 32], - dis_u_2: [f32; 32], - factor_ppc: [f32; 32], - factor_ip: [f32; 32], - factor_err: [f32; 32], + metadata: [[f32; 32]; 4], + delta: [f32; 32], + // it's not last field for reducing padding bytes payload: [Option>; 32], prefetch_s: u16, prefetch_e: u16, + head: [u16; 32], elements_s: u16, elements_e: u16, } @@ -904,13 +961,11 @@ struct FrozenTupleHeader1 { #[derive(Debug, Clone, PartialEq)] pub enum FrozenTuple { _0 { - head: [u16; 32], - dis_u_2: [f32; 32], - factor_ppc: [f32; 32], - factor_ip: [f32; 32], - factor_err: [f32; 32], + metadata: [[f32; 32]; 4], + delta: [f32; 32], payload: [Option>; 32], prefetch: Vec<[u32; 32]>, + head: [u16; 32], elements: Vec<[u8; 16]>, }, _1 { @@ -929,9 +984,13 @@ impl FrozenTuple { } pub fn fit_1(prefetch: usize, freespace: u16) -> Option { let mut freespace = freespace as isize; + freespace &= !(ALIGN - 1) as isize; freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; freespace -= (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN) as isize; + freespace &= !(ALIGN - 1) as isize; if freespace >= 0 { Some(freespace as usize / size_of::<[u8; 16]>()) } else { @@ -945,36 +1004,35 @@ impl Tuple for FrozenTuple { let mut buffer = Vec::::new(); match self { FrozenTuple::_0 { - head, - dis_u_2, - factor_ppc, - factor_ip, - factor_err, + metadata, + delta, payload, prefetch, + head, elements, } => { buffer.extend((0 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); + // prefetch let prefetch_s = buffer.len() as u16; buffer.extend(prefetch.as_bytes()); let prefetch_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // elements let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // header buffer[size_of::()..][..size_of::()].copy_from_slice( FrozenTupleHeader0 { head: *head, - dis_u_2: *dis_u_2, - factor_ppc: *factor_ppc, - factor_ip: *factor_ip, - factor_err: *factor_err, + metadata: *metadata, + delta: *delta, payload: *payload, elements_s, elements_e, @@ -987,12 +1045,14 @@ impl Tuple for FrozenTuple { Self::_1 { elements } => { buffer.extend((1 as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements let elements_s = buffer.len() as u16; buffer.extend(elements.as_bytes()); let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // header buffer[size_of::()..][..size_of::()].copy_from_slice( FrozenTupleHeader1 { elements_s, @@ -1072,19 +1132,11 @@ pub struct FrozenTupleReader0<'a> { } impl<'a> FrozenTupleReader0<'a> { - pub fn mean(self) -> &'a [u16; 32] { - &self.header.head - } - pub fn metadata(self) -> (&'a [f32; 32], &'a [f32; 32], &'a [f32; 32], &'a [f32; 32]) { - ( - &self.header.dis_u_2, - &self.header.factor_ppc, - &self.header.factor_ip, - &self.header.factor_err, - ) + pub fn metadata(self) -> &'a [[f32; 32]; 4] { + &self.header.metadata } - pub fn elements(self) -> &'a [[u8; 16]] { - self.elements + pub fn delta(self) -> &'a [f32; 32] { + &self.header.delta } pub fn payload(self) -> &'a [Option>; 32] { &self.header.payload @@ -1092,6 +1144,12 @@ impl<'a> FrozenTupleReader0<'a> { pub fn prefetch(self) -> &'a [[u32; 32]] { self.prefetch } + pub fn head(self) -> &'a [u16; 32] { + &self.header.head + } + pub fn elements(self) -> &'a [[u8; 16]] { + self.elements + } } #[derive(Debug, Clone, Copy, PartialEq)] @@ -1137,60 +1195,58 @@ impl FrozenTupleWriter0<'_> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct AppendableTupleHeader { - head: u16, - _padding_0: [ZeroU8; 6], - dis_u_2: f32, - factor_ppc: f32, - factor_ip: f32, - factor_err: f32, - payload: Option>, + metadata: [f32; 4], + delta: f32, prefetch_s: u16, prefetch_e: u16, + head: u16, + _padding_0: [ZeroU8; 2], elements_s: u16, elements_e: u16, + // it's the last field for reducing padding bytes + payload: Option>, } #[derive(Debug, Clone, PartialEq)] pub struct AppendableTuple { - pub head: u16, - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub payload: Option>, + pub metadata: [f32; 4], + pub delta: f32, pub prefetch: Vec, + pub head: u16, pub elements: Vec, + pub payload: Option>, } impl Tuple for AppendableTuple { fn serialize(&self) -> Vec { let mut buffer = Vec::::new(); buffer.extend(std::iter::repeat_n(0, size_of::())); + // prefetch let prefetch_s = buffer.len() as u16; buffer.extend(self.prefetch.as_bytes()); let prefetch_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // elements let elements_s = buffer.len() as u16; buffer.extend(self.elements.as_bytes()); let elements_e = buffer.len() as u16; while buffer.len() % ALIGN != 0 { buffer.push(0); } + // header buffer[..size_of::()].copy_from_slice( AppendableTupleHeader { - head: self.head, - _padding_0: Default::default(), - dis_u_2: self.dis_u_2, - factor_ppc: self.factor_ppc, - factor_ip: self.factor_ip, - factor_err: self.factor_err, - payload: self.payload, + metadata: self.metadata, + delta: self.delta, prefetch_s, prefetch_e, + head: self.head, elements_s, elements_e, + payload: self.payload, + _padding_0: Default::default(), } .as_bytes(), ); @@ -1233,23 +1289,23 @@ pub struct AppendableTupleReader<'a> { } impl<'a> AppendableTupleReader<'a> { + pub fn metadata(self) -> [f32; 4] { + self.header.metadata + } + pub fn delta(self) -> f32 { + self.header.delta + } + pub fn prefetch(self) -> &'a [u32] { + self.prefetch + } pub fn head(self) -> u16 { self.header.head } - pub fn code(self) -> BinaryCode<'a> { - ( - self.header.dis_u_2, - self.header.factor_ppc, - self.header.factor_ip, - self.header.factor_err, - self.elements, - ) - } pub fn payload(self) -> Option> { self.header.payload } - pub fn prefetch(self) -> &'a [u32] { - self.prefetch + pub fn elements(self) -> &'a [u64] { + self.elements } } diff --git a/crates/algorithm/src/types.rs b/crates/algorithm/src/types.rs index 5325300c..477bf44b 100644 --- a/crates/algorithm/src/types.rs +++ b/crates/algorithm/src/types.rs @@ -96,7 +96,7 @@ impl VectorOptions { } pub struct Structure { - pub means: Vec, + pub centroids: Vec, pub children: Vec>, } diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index 0d273edf..0fcd9534 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -24,14 +24,14 @@ pub fn read_for_h1_tuple< O: Operator, A: Accessor1<::Element, ::Metadata>, >( + mut prefetch: impl Iterator>, head: u16, - mut list: impl Iterator>, accessor: A, ) -> A::Output { let mut cursor = Err(head); let mut result = accessor; while let Err(head) = cursor { - let guard = list.next().expect("data corruption"); + let guard = prefetch.next().expect("data corruption"); let bytes = guard.get(head).expect("data corruption"); let tuple = VectorTuple::::deserialize_ref(bytes); if tuple.payload().is_some() { @@ -40,7 +40,7 @@ pub fn read_for_h1_tuple< result.push(tuple.elements()); cursor = tuple.metadata_or_head(); } - if list.next().is_some() { + if prefetch.next().is_some() { panic!("data corruption"); } result.finish(cursor.expect("data corruption")) @@ -52,15 +52,15 @@ pub fn read_for_h0_tuple< O: Operator, A: TryAccessor1<::Element, ::Metadata>, >( + mut prefetch: impl Iterator>, head: u16, - mut list: impl Iterator>, payload: NonZero, accessor: A, ) -> Option { let mut cursor = Err(head); let mut result = accessor; while let Err(head) = cursor { - let guard = list.next()?; + let guard = prefetch.next()?; let bytes = guard.get(head)?; let tuple = VectorTuple::::deserialize_ref(bytes); if tuple.payload().is_none() { @@ -72,7 +72,7 @@ pub fn read_for_h0_tuple< result.push(tuple.elements())?; cursor = tuple.metadata_or_head(); } - if list.next().is_some() { + if prefetch.next().is_some() { return None; } result.finish(cursor.ok()?) diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index a67ffab3..0398a4b6 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -123,30 +123,17 @@ fn rabitq_index( samples: &[Vec], centroids: &[Vec], ) -> Vec { - struct Branch { - dis_u_2: f32, - factor_ppc: f32, - factor_ip: f32, - factor_err: f32, - signs: Vec, - } let branches = { let mut branches = Vec::new(); for centroid in centroids { let code = rabitq::code(dims as _, centroid); - branches.push(Branch { - dis_u_2: code.dis_u_2, - factor_ppc: code.factor_ppc, - factor_ip: code.factor_ip, - factor_err: code.factor_err, - signs: code.signs, - }); + branches.push(code); } branches }; struct Block { dis_u_2: [f32; 32], - factor_ppc: [f32; 32], + factor_cnt: [f32; 32], factor_ip: [f32; 32], factor_err: [f32; 32], elements: Vec<[u8; 16]>, @@ -155,7 +142,7 @@ fn rabitq_index( fn code(&self) -> BlockCode<'_> { ( &self.dis_u_2, - &self.factor_ppc, + &self.factor_cnt, &self.factor_ip, &self.factor_err, &self.elements, @@ -165,11 +152,11 @@ fn rabitq_index( let mut blocks = Vec::new(); for chunk in branches.chunks(32) { blocks.push(Block { - dis_u_2: any_pack(chunk.iter().map(|x| x.dis_u_2)), - factor_ppc: any_pack(chunk.iter().map(|x| x.factor_ppc)), - factor_ip: any_pack(chunk.iter().map(|x| x.factor_ip)), - factor_err: any_pack(chunk.iter().map(|x| x.factor_err)), - elements: padding_pack(chunk.iter().map(|x| rabitq::pack_to_u4(&x.signs))), + dis_u_2: any_pack(chunk.iter().map(|x| x.0.dis_u_2)), + factor_cnt: any_pack(chunk.iter().map(|x| x.0.factor_cnt)), + factor_ip: any_pack(chunk.iter().map(|x| x.0.factor_ip)), + factor_err: any_pack(chunk.iter().map(|x| x.0.factor_err)), + elements: padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.1))), }); } (0..n) @@ -178,7 +165,7 @@ fn rabitq_index( let lut = rabitq::block::preprocess(&samples[i]); let mut result = (f32::INFINITY, 0); for block in 0..c.div_ceil(32) { - let returns = rabitq::block::process_l2(&lut, blocks[block].code()); + let returns = rabitq::block::full_process_l2(&lut, blocks[block].code()); let lowerbound = returns.map(|(rough, err)| rough - err * 1.9); for j in block * 32..std::cmp::min(block * 32 + 32, c) { if lowerbound[j - block * 32] < result.0 { diff --git a/crates/rabitq/src/binary.rs b/crates/rabitq/src/binary.rs index bce60e8b..dc1996ab 100644 --- a/crates/rabitq/src/binary.rs +++ b/crates/rabitq/src/binary.rs @@ -14,68 +14,189 @@ use simd::Floating; -pub type BinaryLut = ( - (f32, f32, f32, f32), - (Vec, Vec, Vec, Vec), -); -pub type BinaryCode<'a> = (f32, f32, f32, f32, &'a [u64]); +use crate::CodeMetadata; + +const BITS: usize = 6; + +#[derive(Debug, Clone, Copy)] +pub struct BinaryLutMetadata { + dis_v_2: f32, + b: f32, + k: f32, + qvector_sum: f32, +} + +pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); +pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u64]); pub fn preprocess(vector: &[f32]) -> BinaryLut { let dis_v_2 = f32::reduce_sum_of_x2(vector); - let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); - let qvector_sum = if vector.len() <= 4369 { + preprocess_with_distance(vector, dis_v_2) +} + +pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BinaryLut { + let (k, b, qvector) = simd::quantize::quantize(vector, ((1 << BITS) - 1) as f32); + let qvector_sum = if vector.len() <= (65535_usize / ((1 << BITS) - 1)) { simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 } else { simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 }; - ((dis_v_2, b, k, qvector_sum), binarize(&qvector)) + ( + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }, + binarize(&qvector), + ) } -pub fn process_l2( +pub fn full_process_l2( lut: &BinaryLut, - (dis_u_2, factor_ppc, factor_ip, factor_err, t): BinaryCode<'_>, + ((dis_u_2, factor_cnt, factor_ip, factor_err), t): BinaryCode<'_>, ) -> (f32, f32) { - let &((dis_v_2, b, k, qvector_sum), ref s) = lut; - let value = asymmetric_binary_dot_product(t, s) as u16; - let rough = - dis_u_2 + dis_v_2 + b * factor_ppc + ((2.0 * value as f32) - qvector_sum) * factor_ip * k; - let err = factor_err * dis_v_2.sqrt(); + let &( + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }, + ref s, + ) = lut; + let sum = asymmetric_binary_dot_product(t, s); + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; + let err = 2.0 * factor_err * dis_v_2.sqrt(); (rough, err) } -pub fn process_dot( +pub fn full_process_dot( lut: &BinaryLut, - (_, factor_ppc, factor_ip, factor_err, t): BinaryCode<'_>, + ((_, factor_cnt, factor_ip, factor_err), t): BinaryCode<'_>, +) -> (f32, f32) { + let &( + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }, + ref s, + ) = lut; + let sum = asymmetric_binary_dot_product(t, s); + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = -e * factor_ip; + let err = factor_err * dis_v_2.sqrt(); + (rough, err) +} + +pub fn half_process_l2( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, +) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; + let err = 2.0 * factor_err * dis_v_2.sqrt(); + (rough, err) +} + +pub fn half_process_l2_residual( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2: _, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + dis_f: f32, + delta: f32, ) -> (f32, f32) { - let &((dis_v_2, b, k, qvector_sum), ref s) = lut; - let value = asymmetric_binary_dot_product(t, s) as u16; - let rough = 0.5 * b * factor_ppc + 0.5 * ((2.0 * value as f32) - qvector_sum) * factor_ip * k; - let err = 0.5 * factor_err * dis_v_2.sqrt(); + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = dis_u_2 + dis_f - 2.0 * e * factor_ip + delta; + let err = 2.0 * factor_err * dis_f.sqrt(); (rough, err) } -pub(crate) fn binarize(vector: &[u8]) -> (Vec, Vec, Vec, Vec) { +pub fn half_process_dot( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, +) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = -e * factor_ip; + let err = factor_err * dis_v_2.sqrt(); + (rough, err) +} + +pub fn half_process_dot_residual( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + dis_f: f32, + delta: f32, + norm: f32, +) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = -e * factor_ip + dis_f + delta; + let err = factor_err * (dis_v_2 + norm * norm + 2.0 * dis_f).sqrt(); + (rough, err) +} + +pub(crate) fn binarize(vector: &[u8]) -> [Vec; BITS] { let n = vector.len(); - let mut t0 = vec![0u64; n.div_ceil(64)]; - let mut t1 = vec![0u64; n.div_ceil(64)]; - let mut t2 = vec![0u64; n.div_ceil(64)]; - let mut t3 = vec![0u64; n.div_ceil(64)]; - for i in 0..n { - t0[i / 64] |= (((vector[i] >> 0) & 1) as u64) << (i % 64); - t1[i / 64] |= (((vector[i] >> 1) & 1) as u64) << (i % 64); - t2[i / 64] |= (((vector[i] >> 2) & 1) as u64) << (i % 64); - t3[i / 64] |= (((vector[i] >> 3) & 1) as u64) << (i % 64); + let mut t: [_; BITS] = std::array::from_fn(|_| vec![0_u64; n.div_ceil(64)]); + for i in 0..BITS { + for j in 0..n { + let bit = (vector[j] >> i) & 1; + t[i][j / 64] |= (bit as u64) << (j % 64); + } } - (t0, t1, t2, t3) + t } -pub(crate) fn asymmetric_binary_dot_product( - x: &[u64], - y: &(Vec, Vec, Vec, Vec), -) -> u32 { - let t0 = simd::bit::reduce_sum_of_and(x, &y.0); - let t1 = simd::bit::reduce_sum_of_and(x, &y.1); - let t2 = simd::bit::reduce_sum_of_and(x, &y.2); - let t3 = simd::bit::reduce_sum_of_and(x, &y.3); - (t0 << 0) + (t1 << 1) + (t2 << 2) + (t3 << 3) +pub fn asymmetric_binary_dot_product(x: &[u64], y: &[Vec; BITS]) -> u32 { + let mut result = 0_u32; + for i in 0..BITS { + result += simd::bit::reduce_sum_of_and(x, y[i].as_slice()) << i; + } + result } diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs index 91956262..f6d9477e 100644 --- a/crates/rabitq/src/block.rs +++ b/crates/rabitq/src/block.rs @@ -14,7 +14,19 @@ use simd::Floating; -pub type BlockLut = ((f32, f32, f32, f32), Vec<[u8; 16]>); +use crate::CodeMetadata; + +const BITS: usize = 8; +pub const STEP: usize = 65535_usize / ((1_usize << BITS) - 1); + +#[derive(Debug, Clone, Copy)] +pub struct BlockLutMetadata { + dis_v_2: f32, + k: f32, + c: f32, +} + +pub type BlockLut = (BlockLutMetadata, Vec<[u8; 16]>); pub type BlockCode<'a> = ( &'a [f32; 32], &'a [f32; 32], @@ -25,77 +37,158 @@ pub type BlockCode<'a> = ( pub fn preprocess(vector: &[f32]) -> BlockLut { let dis_v_2 = f32::reduce_sum_of_x2(vector); - let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); - let qvector_sum = if vector.len() <= 4369 { - simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 - } else { - simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 - }; - ((dis_v_2, b, k, qvector_sum), compress(qvector)) + preprocess_with_distance(vector, dis_v_2) } -pub fn process_l2( +pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BlockLut { + let cvector = compress(vector); + let (k, b, cqvector) = + simd::quantize::quantize(cvector.as_flattened(), ((1 << BITS) - 1) as f32); + ( + BlockLutMetadata { + dis_v_2, + k, + c: b * vector.len().div_ceil(4) as f32, + }, + cqvector.as_chunks::<16>().0.to_vec(), + ) +} + +pub fn full_process_l2( lut: &BlockLut, - (dis_u_2, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, + (dis_u_2, _, factor_ip, factor_err, t): BlockCode<'_>, ) -> [(f32, f32); 32] { - let &((dis_v_2, b, k, qvector_sum), ref s) = lut; - let r = simd::fast_scan::fast_scan(t, s); + use std::iter::zip; + let &(BlockLutMetadata { dis_v_2, k, c }, ref s) = lut; + let mut sum = [0_u32; 32]; + for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { + let delta = simd::fast_scan::scan(t, s); + simd::fast_scan::accu(&mut sum, &delta); + } std::array::from_fn(|i| { - let rough = dis_u_2[i] - + dis_v_2 - + b * factor_ppc[i] - + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = factor_err[i] * dis_v_2.sqrt(); + let e = k * (sum[i] as f32) + c; + let rough = dis_u_2[i] + dis_v_2 - 2.0 * e * factor_ip[i]; + let err = 2.0 * factor_err[i] * dis_v_2.sqrt(); (rough, err) }) } -pub fn process_dot( +pub fn full_process_dot( lut: &BlockLut, - (_, factor_ppc, factor_ip, factor_err, t): BlockCode<'_>, + (_, _, factor_ip, factor_err, t): BlockCode<'_>, ) -> [(f32, f32); 32] { - let &((dis_v_2, b, k, qvector_sum), ref s) = lut; - let r = simd::fast_scan::fast_scan(t, s); + use std::iter::zip; + let &(BlockLutMetadata { dis_v_2, k, c }, ref s) = lut; + let mut sum = [0_u32; 32]; + for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { + let delta = simd::fast_scan::scan(t, s); + simd::fast_scan::accu(&mut sum, &delta); + } std::array::from_fn(|i| { - let rough = - 0.5 * b * factor_ppc[i] + 0.5 * ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = 0.5 * factor_err[i] * dis_v_2.sqrt(); + let e = k * (sum[i] as f32) + c; + let rough = -e * factor_ip[i]; + let err = factor_err[i] * dis_v_2.sqrt(); (rough, err) }) } -pub(crate) fn compress(mut vector: Vec) -> Vec<[u8; 16]> { - let n = vector.len().div_ceil(4); - vector.resize(n * 4, 0); - let mut result = vec![[0u8; 16]; n]; - for i in 0..n { - #[allow(unsafe_code)] - unsafe { - // this hint is used to skip bound checks - std::hint::assert_unchecked(4 * i + 3 < vector.len()); - } - let t_0 = vector[4 * i + 0]; - let t_1 = vector[4 * i + 1]; - let t_2 = vector[4 * i + 2]; - let t_3 = vector[4 * i + 3]; - result[i] = [ - 0, - t_0, - t_1, - t_1 + t_0, - t_2, - t_2 + t_0, - t_2 + t_1, - t_2 + t_1 + t_0, - t_3, - t_3 + t_0, - t_3 + t_1, - t_3 + t_1 + t_0, - t_3 + t_2, - t_3 + t_2 + t_0, - t_3 + t_2 + t_1, - t_3 + t_2 + t_1 + t_0, - ]; +pub fn half_process_l2( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, +) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; + let err = 2.0 * factor_err * dis_v_2.sqrt(); + (rough, err) +} + +pub fn half_process_l2_residual( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2: _, k, c }: BlockLutMetadata, + dis_f: f32, + delta: f32, +) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = dis_u_2 + dis_f - 2.0 * e * factor_ip + delta; + let err = 2.0 * factor_err * dis_f.sqrt(); + (rough, err) +} + +pub fn half_process_dot( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, +) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = -e * factor_ip; + let err = factor_err * dis_v_2.sqrt(); + (rough, err) +} + +pub fn half_process_dot_residual( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, + dis_f: f32, + delta: f32, + norm: f32, +) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = -e * factor_ip + dis_f + delta; + let err = factor_err * (dis_v_2 + norm * norm + 2.0 * dis_f).sqrt(); + (rough, err) +} + +fn compress(vector: &[f32]) -> Vec<[f32; 16]> { + let f = |[t_0, t_1, t_2, t_3]: [f32; 4]| { + [ + 0.0 - t_3 - t_2 - t_1 - t_0, + 0.0 - t_3 - t_2 - t_1 + t_0, + 0.0 - t_3 - t_2 + t_1 - t_0, + 0.0 - t_3 - t_2 + t_1 + t_0, + 0.0 - t_3 + t_2 - t_1 - t_0, + 0.0 - t_3 + t_2 - t_1 + t_0, + 0.0 - t_3 + t_2 + t_1 - t_0, + 0.0 - t_3 + t_2 + t_1 + t_0, + 0.0 + t_3 - t_2 - t_1 - t_0, + 0.0 + t_3 - t_2 - t_1 + t_0, + 0.0 + t_3 - t_2 + t_1 - t_0, + 0.0 + t_3 - t_2 + t_1 + t_0, + 0.0 + t_3 + t_2 - t_1 - t_0, + 0.0 + t_3 + t_2 - t_1 + t_0, + 0.0 + t_3 + t_2 + t_1 - t_0, + 0.0 + t_3 + t_2 + t_1 + t_0, + ] + }; + + let (arrays, reminder) = vector.as_chunks::<4>(); + let mut result = arrays.iter().copied().map(f).collect::>(); + if !reminder.is_empty() { + let mut array = [0.0; 4]; + array[..reminder.len()].copy_from_slice(reminder); + result.push(f(array)); } result } diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index c5ec66b9..6274274e 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -12,6 +12,9 @@ // // Copyright (c) 2025 TensorChord Inc. +#![feature(slice_as_chunks)] +#![feature(select_unpredictable)] + pub mod binary; pub mod block; pub mod packing; @@ -20,88 +23,54 @@ use binary::BinaryLut; use block::BlockLut; use simd::Floating; -#[derive(Debug, Clone)] -pub struct Code { +#[derive(Debug, Clone, Copy)] +pub struct CodeMetadata { pub dis_u_2: f32, - pub factor_ppc: f32, + pub factor_cnt: f32, pub factor_ip: f32, pub factor_err: f32, - pub signs: Vec, -} - -pub fn pack_to_u4(signs: &[bool]) -> Vec { - fn f(x: [bool; 4]) -> u8 { - x[0] as u8 | (x[1] as u8) << 1 | (x[2] as u8) << 2 | (x[3] as u8) << 3 - } - let mut result = Vec::with_capacity(signs.len().div_ceil(4)); - for i in 0..signs.len().div_ceil(4) { - let x = std::array::from_fn(|j| signs.get(i * 4 + j).copied().unwrap_or_default()); - result.push(f(x)); - } - result } -pub fn pack_to_u64(signs: &[bool]) -> Vec { - fn f(x: [bool; 64]) -> u64 { - let mut result = 0_u64; - for i in 0..64 { - result |= (x[i] as u64) << i; - } - result - } - let mut result = Vec::with_capacity(signs.len().div_ceil(64)); - for i in 0..signs.len().div_ceil(64) { - let x = std::array::from_fn(|j| signs.get(i * 64 + j).copied().unwrap_or_default()); - result.push(f(x)); - } - result -} +pub type Code = (CodeMetadata, Vec); pub fn code(dims: u32, vector: &[f32]) -> Code { let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); let sum_of_x_2 = f32::reduce_sum_of_x2(vector); - let dis_u = sum_of_x_2.sqrt(); - let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); - let x_x0 = dis_u / x0; - let fac_norm = (dims as f32).sqrt(); - let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); - let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); - let factor_ip = -2.0f32 / fac_norm * x_x0; - let cnt_pos = vector - .iter() - .map(|x| x.is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.is_sign_negative() as i32) - .sum::(); - let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; - let mut signs = Vec::new(); - for i in 0..dims { - signs.push(vector[i as usize].is_sign_positive()); - } - Code { - dis_u_2: sum_of_x_2, - factor_ppc, - factor_ip, - factor_err, - signs, - } + ( + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector + .iter() + .map(|x| x.is_sign_positive() as i32) + .sum::(); + let cnt_neg = vector + .iter() + .map(|x| x.is_sign_negative() as i32) + .sum::(); + (cnt_pos - cnt_neg) as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + factor_err: { + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); + dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (dims as f32 - 1.0).sqrt() + }, + }, + { + let mut signs = Vec::new(); + for i in 0..dims { + signs.push(vector[i as usize].is_sign_positive()); + } + signs + }, + ) } pub fn preprocess(vector: &[f32]) -> (BlockLut, BinaryLut) { - use simd::Floating; let dis_v_2 = f32::reduce_sum_of_x2(vector); - let (k, b, qvector) = simd::quantize::quantize(vector, 15.0); - let qvector_sum = if vector.len() <= 4369 { - simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 - } else { - simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 - }; - let binary = binary::binarize(&qvector); - let block = block::compress(qvector); ( - ((dis_v_2, b, k, qvector_sum), block), - ((dis_v_2, b, k, qvector_sum), binary), + block::preprocess_with_distance(vector, dis_v_2), + binary::preprocess_with_distance(vector, dis_v_2), ) } diff --git a/crates/rabitq/src/packing.rs b/crates/rabitq/src/packing.rs index fc85bbcd..f4e249c2 100644 --- a/crates/rabitq/src/packing.rs +++ b/crates/rabitq/src/packing.rs @@ -100,3 +100,31 @@ pub fn padding_pack(x: impl IntoIterator>) -> Vec<[u8; 1 pub fn any_pack(mut x: impl Iterator) -> [T; 32] { std::array::from_fn(|_| x.next()).map(|x| x.unwrap_or_default()) } + +pub fn pack_to_u4(signs: &[bool]) -> Vec { + fn f(x: [bool; 4]) -> u8 { + x[0] as u8 | (x[1] as u8) << 1 | (x[2] as u8) << 2 | (x[3] as u8) << 3 + } + let mut result = Vec::with_capacity(signs.len().div_ceil(4)); + for i in 0..signs.len().div_ceil(4) { + let x = std::array::from_fn(|j| signs.get(i * 4 + j).copied().unwrap_or_default()); + result.push(f(x)); + } + result +} + +pub fn pack_to_u64(signs: &[bool]) -> Vec { + fn f(x: [bool; 64]) -> u64 { + let mut result = 0_u64; + for i in 0..64 { + result |= (x[i] as u64) << i; + } + result + } + let mut result = Vec::with_capacity(signs.len().div_ceil(64)); + for i in 0..signs.len().div_ceil(64) { + let x = std::array::from_fn(|j| signs.get(i * 64 + j).copied().unwrap_or_default()); + result.push(f(x)); + } + result +} diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index f1271f4f..43673381 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -467,6 +467,20 @@ mod scan { } #[inline(always)] -pub fn fast_scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { +pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { scan::scan(code, lut) } + +mod accu { + #[crate::multiversion("v4", "v3", "v2", "a2")] + pub fn accu(sum: &mut [u32; 32], delta: &[u16; 32]) { + for i in 0..32 { + sum[i] += delta[i] as u32; + } + } +} + +#[inline(always)] +pub fn accu(sum: &mut [u32; 32], delta: &[u16; 32]) { + accu::accu(sum, delta); +} diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index ed207db0..f75b9442 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -16,7 +16,11 @@ use super::opclass::Opfamily; use crate::index::am::am_build::InternalBuild; use algorithm::operator::{Dot, L2, Op}; use algorithm::types::*; -use algorithm::*; +use algorithm::{ + Bump, FastHeap, Fetch, Hints, PlainPrefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily, + RelationPrefetch, RelationRead, RelationReadStream, RelationWrite, Sequence, SimplePrefetcher, + StreamPrefetcher, +}; use half::f16; use std::cell::UnsafeCell; use std::collections::BinaryHeap; @@ -204,25 +208,34 @@ impl Bump for BumpAlloc { } pub fn prewarm(opfamily: Opfamily, index: &impl RelationRead, height: i32) -> String { + let bump = BumpAlloc::new(); let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::prewarm::<_, Op, L2>>(index, height, make_h0_plain_prefetcher) - } + (VectorKind::Vecf32, DistanceKind::L2) => algorithm::prewarm::<_, Op, L2>>( + index, + height, + &bump, + make_h0_plain_prefetcher, + ), (VectorKind::Vecf32, DistanceKind::Dot) => { algorithm::prewarm::<_, Op, Dot>>( index, height, + &bump, make_h0_plain_prefetcher, ) } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::prewarm::<_, Op, L2>>(index, height, make_h0_plain_prefetcher) - } + (VectorKind::Vecf16, DistanceKind::L2) => algorithm::prewarm::<_, Op, L2>>( + index, + height, + &bump, + make_h0_plain_prefetcher, + ), (VectorKind::Vecf16, DistanceKind::Dot) => { algorithm::prewarm::<_, Op, Dot>>( index, height, + &bump, make_h0_plain_prefetcher, ) } @@ -237,16 +250,20 @@ pub fn bulkdelete( ) { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::bulkdelete::<_, Op, L2>>(index, check, callback) + algorithm::bulkdelete::<_, Op, L2>>(index, &check, &callback); + algorithm::bulkdelete_vectors::<_, Op, L2>>(index, &check, &callback); } (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::bulkdelete::<_, Op, Dot>>(index, check, callback) + algorithm::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + algorithm::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); } (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::bulkdelete::<_, Op, L2>>(index, check, callback) + algorithm::bulkdelete::<_, Op, L2>>(index, &check, &callback); + algorithm::bulkdelete_vectors::<_, Op, L2>>(index, &check, &callback); } (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::bulkdelete::<_, Op, Dot>>(index, check, callback) + algorithm::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + algorithm::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); } } } @@ -322,54 +339,70 @@ pub fn insert( match (vector, opfamily.distance_kind()) { (OwnedVector::Vecf32(vector), DistanceKind::L2) => { assert!(opfamily.vector_kind() == VectorKind::Vecf32); - let (list, head) = insert_vector::<_, Op, L2>>(index, payload, &vector); - insert_index::<_, Op, L2>>( + let projected = RandomProject::project(vector.as_borrowed()); + let key = algorithm::insert_vector::<_, Op, L2>>( + index, + payload, + vector.as_borrowed(), + ); + algorithm::insert::<_, Op, L2>>( index, payload, - RandomProject::project(vector.as_borrowed()), + projected.as_borrowed(), + key, &bump, make_h1_plain_prefetcher, - list, - head, ) } (OwnedVector::Vecf32(vector), DistanceKind::Dot) => { assert!(opfamily.vector_kind() == VectorKind::Vecf32); - let (list, head) = insert_vector::<_, Op, Dot>>(index, payload, &vector); - insert_index::<_, Op, Dot>>( + let projected = RandomProject::project(vector.as_borrowed()); + let key = algorithm::insert_vector::<_, Op, Dot>>( + index, + payload, + vector.as_borrowed(), + ); + algorithm::insert::<_, Op, Dot>>( index, payload, - RandomProject::project(vector.as_borrowed()), + projected.as_borrowed(), + key, &bump, make_h1_plain_prefetcher, - list, - head, ) } (OwnedVector::Vecf16(vector), DistanceKind::L2) => { assert!(opfamily.vector_kind() == VectorKind::Vecf16); - let (list, head) = insert_vector::<_, Op, L2>>(index, payload, &vector); - insert_index::<_, Op, L2>>( + let projected = RandomProject::project(vector.as_borrowed()); + let key = algorithm::insert_vector::<_, Op, L2>>( index, payload, - RandomProject::project(vector.as_borrowed()), + vector.as_borrowed(), + ); + algorithm::insert::<_, Op, L2>>( + index, + payload, + projected.as_borrowed(), + key, &bump, make_h1_plain_prefetcher, - list, - head, ) } (OwnedVector::Vecf16(vector), DistanceKind::Dot) => { assert!(opfamily.vector_kind() == VectorKind::Vecf16); - let (list, head) = insert_vector::<_, Op, Dot>>(index, payload, &vector); - insert_index::<_, Op, Dot>>( + let projected = RandomProject::project(vector.as_borrowed()); + let key = algorithm::insert_vector::<_, Op, Dot>>( + index, + payload, + vector.as_borrowed(), + ); + algorithm::insert::<_, Op, Dot>>( index, payload, - RandomProject::project(vector.as_borrowed()), + projected.as_borrowed(), + key, &bump, make_h1_plain_prefetcher, - list, - head, ) } } @@ -377,10 +410,15 @@ pub fn insert( fn map_structures(x: Vec>, f: impl Fn(T) -> U + Copy) -> Vec> { x.into_iter() - .map(|Structure { means, children }| Structure { - means: means.into_iter().map(f).collect(), - children, - }) + .map( + |Structure { + centroids, + children, + }| Structure { + centroids: centroids.into_iter().map(f).collect(), + children, + }, + ) .collect() } diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 13f5584a..32ce4102 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -259,11 +259,6 @@ pub unsafe extern "C-unwind" fn ambuild( if let Err(errors) = Validate::validate(&vchordrq_options) { pgrx::error!("error while validating options: {}", errors); } - if vector_options.d != DistanceKind::L2 && vchordrq_options.index.residual_quantization { - pgrx::error!( - "error while validating options: residual_quantization can be enabled only if distance type is L2" - ); - } let opfamily = unsafe { opfamily(index_relation) }; let index = unsafe { PostgresRelation::new(index_relation) }; let heap = Heap { @@ -981,7 +976,7 @@ fn make_internal_build( let mut result = Vec::>>::new(); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { let input = if let Some(structure) = result.last() { - &structure.means + &structure.centroids } else { &samples }; @@ -1002,7 +997,7 @@ fn make_internal_build( "clustering: starting, using {num_threads} threads, clustering {num_points} vectors of {num_dims} dimension into {num_lists} clusters, in {num_iterations} iterations" ); } - let means = k_means::k_means( + let centroids = k_means::k_means( num_threads, |i| { pgrx::check_for_interrupts!(); @@ -1035,18 +1030,24 @@ fn make_internal_build( pgrx::info!("clustering: finished"); } if let Some(structure) = result.last() { - let mut children = vec![Vec::new(); means.len()]; + let mut children = vec![Vec::new(); centroids.len()]; for i in 0..structure.len() as u32 { - let target = k_means::k_means_lookup(&structure.means[i as usize], &means); + let target = k_means::k_means_lookup(&structure.centroids[i as usize], ¢roids); children[target].push(i); } - let (means, children) = std::iter::zip(means, children) + let (centroids, children) = std::iter::zip(centroids, children) .filter(|(_, x)| !x.is_empty()) .unzip::<_, _, Vec<_>, Vec<_>>(); - result.push(Structure { means, children }); + result.push(Structure { + centroids, + children, + }); } else { - let children = vec![Vec::new(); means.len()]; - result.push(Structure { means, children }); + let children = vec![Vec::new(); centroids.len()]; + result.push(Structure { + centroids, + children, + }); } } result @@ -1107,11 +1108,11 @@ fn make_external_build( let n = parents.len(); let mut result = Vec::new(); result.push(Structure { - means: vectors.values().cloned().collect::>(), + centroids: vectors.values().cloned().collect::>(), children: vec![Vec::new(); n], }); result.push(Structure { - means: vec![{ + centroids: vec![{ // compute the vector on root, without normalizing it let mut sum = vec![0.0f32; vector_options.dims as _]; for vector in vectors.values() { @@ -1212,8 +1213,11 @@ fn make_external_build( } let mut result = Vec::new(); for height in 1..=heights[&root] { - let (means, children) = extract(height, &labels, &vectors, &children); - result.push(Structure { means, children }); + let (centroids, children) = extract(height, &labels, &vectors, &children); + result.push(Structure { + centroids, + children, + }); } result } diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index 2e1dbaf1..67d98733 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -408,6 +408,7 @@ pub unsafe extern "C-unwind" fn amrescan( maxsim_threshold: gucs::maxsim_threshold(), io_search: gucs::io_search(), io_rerank: gucs::io_rerank(), + prefilter: gucs::prefilter(), }; let fetcher = { let hack = scanner.hack; @@ -569,8 +570,10 @@ impl Drop for HeapFetcher { } } -impl SearchFetcher for HeapFetcher { - fn fetch(&mut self, key: [u16; 3]) -> Option<(&[Datum; 32], &[bool; 32])> { +impl Fetcher for HeapFetcher { + type Tuple<'a> = HeapTuple<'a>; + + fn fetch(&mut self, key: [u16; 3]) -> Option> { unsafe { let mut ctid = key_to_ctid(key); let table_am = (*self.heap_relation).rd_tableam; @@ -580,49 +583,54 @@ impl SearchFetcher for HeapFetcher { if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { return None; } - (*self.econtext).ecxt_scantuple = self.slot; - pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); + Some(HeapTuple { this: self }) + } + } +} + +pub struct HeapTuple<'a> { + this: &'a mut HeapFetcher, +} + +impl Tuple for HeapTuple<'_> { + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]) { + unsafe { + let this = &mut self.this; + (*this.econtext).ecxt_scantuple = this.slot; + pgrx::pg_sys::MemoryContextReset((*this.econtext).ecxt_per_tuple_memory); pgrx::pg_sys::FormIndexDatum( - self.index_info, - self.slot, - self.estate, - self.values.as_mut_ptr(), - self.is_nulls.as_mut_ptr(), + this.index_info, + this.slot, + this.estate, + this.values.as_mut_ptr(), + this.is_nulls.as_mut_ptr(), ); - Some((&self.values, &self.is_nulls)) + (&this.values, &this.is_nulls) } } - fn filter(&mut self, key: [u16; 3]) -> bool { - if self.hack.is_null() { - return true; - } + fn filter(&mut self) -> bool { unsafe { - let mut ctid = key_to_ctid(key); - let table_am = (*self.heap_relation).rd_tableam; - let fetch_row_version = (*table_am) - .tuple_fetch_row_version - .expect("unsupported heap access method"); - if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { - return false; - } - if let Some(qual) = NonNull::new((*self.hack).ss.ps.qual) { - use pgrx::datum::FromDatum; - use pgrx::memcxt::PgMemoryContexts; - assert!(qual.as_ref().flags & pgrx::pg_sys::EEO_FLAG_IS_QUAL as u8 != 0); - let evalfunc = qual.as_ref().evalfunc.expect("no evalfunc for qual"); - if !(*self.hack).ss.ps.ps_ExprContext.is_null() { - let econtext = (*self.hack).ss.ps.ps_ExprContext; - (*econtext).ecxt_scantuple = self.slot; - pgrx::pg_sys::MemoryContextReset((*econtext).ecxt_per_tuple_memory); - let result = PgMemoryContexts::For((*econtext).ecxt_per_tuple_memory) - .switch_to(|_| { - let mut is_null = true; - let datum = evalfunc(qual.as_ptr(), econtext, &mut is_null); - bool::from_datum(datum, is_null) - }); - if result != Some(true) { - return false; + let this = &mut self.this; + if !this.hack.is_null() { + if let Some(qual) = NonNull::new((*this.hack).ss.ps.qual) { + use pgrx::datum::FromDatum; + use pgrx::memcxt::PgMemoryContexts; + assert!(qual.as_ref().flags & pgrx::pg_sys::EEO_FLAG_IS_QUAL as u8 != 0); + let evalfunc = qual.as_ref().evalfunc.expect("no evalfunc for qual"); + if !(*this.hack).ss.ps.ps_ExprContext.is_null() { + let econtext = (*this.hack).ss.ps.ps_ExprContext; + (*econtext).ecxt_scantuple = this.slot; + pgrx::pg_sys::MemoryContextReset((*econtext).ecxt_per_tuple_memory); + let result = PgMemoryContexts::For((*econtext).ecxt_per_tuple_memory) + .switch_to(|_| { + let mut is_null = true; + let datum = evalfunc(qual.as_ptr(), econtext, &mut is_null); + bool::from_datum(datum, is_null) + }); + if result != Some(true) { + return false; + } } } } diff --git a/src/index/gucs.rs b/src/index/gucs.rs index d91a04ee..3ccfe94b 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -187,23 +187,30 @@ pub fn maxsim_threshold() -> u32 { } pub fn prewarm_dim() -> Vec { - if let Some(prewarm_dim) = PREWARM_DIM.get() { - if let Ok(prewarm_dim) = prewarm_dim.to_str() { + match PREWARM_DIM.get() { + None => Vec::new(), + Some(probes) => { let mut result = Vec::new(); - for dim in prewarm_dim.split(',') { - if let Ok(dim) = dim.trim().parse::() { - result.push(dim); - } else { - pgrx::warning!("{dim:?} is not a valid integer"); + let mut current = None; + for &c in probes.to_bytes() { + match c { + b' ' => continue, + b',' => result.push(current.take().expect("empty prewarm_dim")), + b'0'..=b'9' => { + if let Some(x) = current.as_mut() { + *x = *x * 10 + (c - b'0') as u32; + } else { + current = Some((c - b'0') as u32); + } + } + c => pgrx::error!("unknown character in prewarm_dim: ASCII = {c}"), } } + if let Some(current) = current { + result.push(current); + } result - } else { - pgrx::warning!("vchordrq.prewarm_dim is not a valid UTF-8 string"); - Vec::new() } - } else { - Vec::new() } } diff --git a/src/index/scanners/default.rs b/src/index/scanners/default.rs index 021d514c..411a51bc 100644 --- a/src/index/scanners/default.rs +++ b/src/index/scanners/default.rs @@ -12,14 +12,15 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::{Io, SearchBuilder, SearchFetcher, SearchOptions}; +use super::{Fetcher, Io, SearchBuilder, SearchOptions, filter}; use crate::index::algorithm::*; use crate::index::am::pointer_to_kv; -use crate::index::gucs::prefilter; use crate::index::opclass::{Opfamily, Sphere}; -use algorithm::operator::{Dot, L2}; +use crate::index::scanners::Tuple; +use algorithm::operator::{self, Dot, L2}; use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; use algorithm::*; +use always_equal::AlwaysEqual; use half::f16; use std::collections::BinaryHeap; use std::num::NonZero; @@ -68,7 +69,7 @@ impl SearchBuilder for DefaultBuilder { self, index: &'a R, options: SearchOptions, - mut fetcher: impl SearchFetcher + 'a, + mut fetcher: impl Fetcher + 'a, bump: &'a impl Bump, ) -> Box + 'a> where @@ -102,20 +103,21 @@ impl SearchBuilder for DefaultBuilder { index, hints: Hints::default().full(true), }; + let f = move |(distance, payload)| (opfamily.output(distance), payload); let iter: Box)>> = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { type Op = operator::Op, L2>; - let original_vector = if let OwnedVector::Vecf32(vector) = vector { + let unprojected = if let OwnedVector::Vecf32(vector) = vector { vector } else { unreachable!() }; - let vector = RandomProject::project(original_vector.as_borrowed()); + let projected = RandomProject::project(unprojected.as_borrowed()); let results = match options.io_search { Io::Plain => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -124,7 +126,7 @@ impl SearchBuilder for DefaultBuilder { ), Io::Simple => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -133,7 +135,7 @@ impl SearchBuilder for DefaultBuilder { ), Io::Stream => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -142,68 +144,63 @@ impl SearchBuilder for DefaultBuilder { ), }; let method = how(index); - match (method, options.io_rerank, prefilter()) { + let sequence = BinaryHeap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } (RerankMethod::Index, Io::Plain, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); - let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Index, Io::Plain, false) => { - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); - let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Index, Io::Simple, false) => { - let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( - index, - results.into(), - ); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, Hints::default()); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); + let sequence = filter(sequence, predicate); let prefetcher = - StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) - } - (RerankMethod::Index, Io::Stream, false) => { - let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( - index, - results.into(), - Hints::default(), - ); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + StreamPrefetcher::new(index, sequence, Hints::default()); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Heap, _, true) => { + (RerankMethod::Heap, _, false) => { let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); - if !fetcher.filter(key) { - return None; - } - let (datums, is_nulls) = fetcher.fetch(key)?; + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; @@ -215,18 +212,17 @@ impl SearchBuilder for DefaultBuilder { }; Some(raw) }; - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_heap::(original_vector, prefetcher, fetch).map( - move |(distance, payload)| (opfamily.output(distance), payload), - ), - ) + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) } - (RerankMethod::Heap, _, false) => { + (RerankMethod::Heap, _, true) => { let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); - let (datums, is_nulls) = fetcher.fetch(key)?; + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; @@ -238,28 +234,23 @@ impl SearchBuilder for DefaultBuilder { }; Some(raw) }; - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_heap::(original_vector, prefetcher, fetch).map( - move |(distance, payload)| (opfamily.output(distance), payload), - ), - ) + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) } } } (VectorKind::Vecf32, DistanceKind::Dot) => { type Op = operator::Op, Dot>; - let original_vector = if let OwnedVector::Vecf32(vector) = vector { + let unprojected = if let OwnedVector::Vecf32(vector) = vector { vector } else { unreachable!() }; - let vector = RandomProject::project(original_vector.as_borrowed()); + let projected = RandomProject::project(unprojected.as_borrowed()); let results = match options.io_search { Io::Plain => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -268,7 +259,7 @@ impl SearchBuilder for DefaultBuilder { ), Io::Simple => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -277,7 +268,7 @@ impl SearchBuilder for DefaultBuilder { ), Io::Stream => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -286,68 +277,63 @@ impl SearchBuilder for DefaultBuilder { ), }; let method = how(index); - match (method, options.io_rerank, prefilter()) { + let sequence = BinaryHeap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } (RerankMethod::Index, Io::Plain, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); - let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Index, Io::Plain, false) => { - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); - let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Index, Io::Simple, false) => { - let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( - index, - results.into(), - ); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, Hints::default()); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); + let sequence = filter(sequence, predicate); let prefetcher = - StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) - } - (RerankMethod::Index, Io::Stream, false) => { - let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( - index, - results.into(), - Hints::default(), - ); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + StreamPrefetcher::new(index, sequence, Hints::default()); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Heap, _, true) => { + (RerankMethod::Heap, _, false) => { let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); - if !fetcher.filter(key) { - return None; - } - let (datums, is_nulls) = fetcher.fetch(key)?; + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; @@ -359,18 +345,17 @@ impl SearchBuilder for DefaultBuilder { }; Some(raw) }; - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_heap::(original_vector, prefetcher, fetch).map( - move |(distance, payload)| (opfamily.output(distance), payload), - ), - ) + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) } - (RerankMethod::Heap, _, false) => { + (RerankMethod::Heap, _, true) => { let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); - let (datums, is_nulls) = fetcher.fetch(key)?; + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; @@ -382,28 +367,23 @@ impl SearchBuilder for DefaultBuilder { }; Some(raw) }; - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_heap::(original_vector, prefetcher, fetch).map( - move |(distance, payload)| (opfamily.output(distance), payload), - ), - ) + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) } } } (VectorKind::Vecf16, DistanceKind::L2) => { type Op = operator::Op, L2>; - let original_vector = if let OwnedVector::Vecf16(vector) = vector { + let unprojected = if let OwnedVector::Vecf16(vector) = vector { vector } else { unreachable!() }; - let vector = RandomProject::project(original_vector.as_borrowed()); + let projected = RandomProject::project(unprojected.as_borrowed()); let results = match options.io_search { Io::Plain => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -412,7 +392,7 @@ impl SearchBuilder for DefaultBuilder { ), Io::Simple => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -421,7 +401,7 @@ impl SearchBuilder for DefaultBuilder { ), Io::Stream => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -430,68 +410,63 @@ impl SearchBuilder for DefaultBuilder { ), }; let method = how(index); - match (method, options.io_rerank, prefilter()) { + let sequence = BinaryHeap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } (RerankMethod::Index, Io::Plain, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); - let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Index, Io::Plain, false) => { - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); - let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Index, Io::Simple, false) => { - let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( - index, - results.into(), - ); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, Hints::default()); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); + let sequence = filter(sequence, predicate); let prefetcher = - StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) - } - (RerankMethod::Index, Io::Stream, false) => { - let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( - index, - results.into(), - Hints::default(), - ); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + StreamPrefetcher::new(index, sequence, Hints::default()); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Heap, _, true) => { + (RerankMethod::Heap, _, false) => { let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); - if !fetcher.filter(key) { - return None; - } - let (datums, is_nulls) = fetcher.fetch(key)?; + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; @@ -503,18 +478,17 @@ impl SearchBuilder for DefaultBuilder { }; Some(raw) }; - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_heap::(original_vector, prefetcher, fetch).map( - move |(distance, payload)| (opfamily.output(distance), payload), - ), - ) + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) } - (RerankMethod::Heap, _, false) => { + (RerankMethod::Heap, _, true) => { let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); - let (datums, is_nulls) = fetcher.fetch(key)?; + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; @@ -526,28 +500,23 @@ impl SearchBuilder for DefaultBuilder { }; Some(raw) }; - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_heap::(original_vector, prefetcher, fetch).map( - move |(distance, payload)| (opfamily.output(distance), payload), - ), - ) + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) } } } (VectorKind::Vecf16, DistanceKind::Dot) => { type Op = operator::Op, Dot>; - let original_vector = if let OwnedVector::Vecf16(vector) = vector { + let unprojected = if let OwnedVector::Vecf16(vector) = vector { vector } else { unreachable!() }; - let vector = RandomProject::project(original_vector.as_borrowed()); + let projected = RandomProject::project(unprojected.as_borrowed()); let results = match options.io_search { Io::Plain => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -556,7 +525,7 @@ impl SearchBuilder for DefaultBuilder { ), Io::Simple => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -565,7 +534,7 @@ impl SearchBuilder for DefaultBuilder { ), Io::Stream => default_search::<_, Op>( index, - vector.clone(), + projected.as_borrowed(), options.probes, options.epsilon, bump, @@ -574,68 +543,63 @@ impl SearchBuilder for DefaultBuilder { ), }; let method = how(index); - match (method, options.io_rerank, prefilter()) { + let sequence = BinaryHeap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } (RerankMethod::Index, Io::Plain, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); - let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Index, Io::Plain, false) => { - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); - let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Index, Io::Simple, false) => { - let prefetcher = SimplePrefetcher::<'a, R, BinaryHeap<_>>::new( - index, - results.into(), - ); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, Hints::default()); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { - let seq = seq_filter(BinaryHeap::from(results), move |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) + let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() }); + let sequence = filter(sequence, predicate); let prefetcher = - StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) + StreamPrefetcher::new(index, sequence, Hints::default()); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } - (RerankMethod::Index, Io::Stream, false) => { - let prefetcher = StreamPrefetcher::<_, BinaryHeap<_>>::new( - index, - results.into(), - Hints::default(), - ); - Box::new(rerank_index::(original_vector, prefetcher).map( - move |(distance, payload)| (opfamily.output(distance), payload), - )) - } - (RerankMethod::Heap, _, true) => { + (RerankMethod::Heap, _, false) => { let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); - if !fetcher.filter(key) { - return None; - } - let (datums, is_nulls) = fetcher.fetch(key)?; + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; @@ -647,18 +611,17 @@ impl SearchBuilder for DefaultBuilder { }; Some(raw) }; - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_heap::(original_vector, prefetcher, fetch).map( - move |(distance, payload)| (opfamily.output(distance), payload), - ), - ) + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) } - (RerankMethod::Heap, _, false) => { + (RerankMethod::Heap, _, true) => { let fetch = move |payload| { let (key, _) = pointer_to_kv(payload); - let (datums, is_nulls) = fetcher.fetch(key)?; + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); let datum = (!is_nulls[0]).then_some(datums[0]); let maybe_vector = unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; @@ -670,13 +633,8 @@ impl SearchBuilder for DefaultBuilder { }; Some(raw) }; - let prefetcher = - PlainPrefetcher::<_, BinaryHeap<_>>::new(index, results.into()); - Box::new( - rerank_heap::(original_vector, prefetcher, fetch).map( - move |(distance, payload)| (opfamily.output(distance), payload), - ), - ) + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) } } } @@ -697,3 +655,11 @@ impl SearchBuilder for DefaultBuilder { })) } } + +#[inline(always)] +pub fn id_0(f: F) -> F +where + F: for<'a> FnMut(&(A, AlwaysEqual<&'a mut (B, C, &'a mut D)>)) -> R, +{ + f +} diff --git a/src/index/scanners/maxsim.rs b/src/index/scanners/maxsim.rs index 9c39d800..f70b8570 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/scanners/maxsim.rs @@ -12,12 +12,11 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::{SearchBuilder, SearchFetcher, SearchOptions}; -use crate::index::algorithm::{RandomProject, *}; +use super::{Fetcher, SearchBuilder, SearchOptions}; +use crate::index::algorithm::*; use crate::index::am::pointer_to_kv; -use crate::index::gucs::prefilter; use crate::index::opclass::Opfamily; -use crate::index::scanners::Io; +use crate::index::scanners::{Io, Tuple, filter}; use algorithm::operator::Dot; use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; use algorithm::*; @@ -61,7 +60,7 @@ impl SearchBuilder for MaxsimBuilder { self, index: &'a R, options: SearchOptions, - mut fetcher: impl SearchFetcher + 'a, + mut fetcher: impl Fetcher + 'a, bump: &'a impl Bump, ) -> Box + 'a> where @@ -105,7 +104,7 @@ impl SearchBuilder for MaxsimBuilder { let iter: Box> = match opfamily.vector_kind() { VectorKind::Vecf32 => { type Op = operator::Op, Dot>; - let original_vectors = vectors + let unprojected = vectors .into_iter() .map(|vector| { if let OwnedVector::Vecf32(vector) = vector { @@ -115,145 +114,138 @@ impl SearchBuilder for MaxsimBuilder { } }) .collect::>(); - let vectors = original_vectors - .clone() - .into_iter() - .map(|vector| RandomProject::project(vector.as_borrowed())); - Box::new(vectors.into_iter().zip(original_vectors).map( - |(vector, original_vector)| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - match (options.io_rerank, prefilter()) { - (Io::Plain, true) => { - let seq = seq_filter(BinaryHeap::from(results), |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) - }); - let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::<_, _>::new( - index, - BinaryHeap::from(results), - ); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let seq = seq_filter(BinaryHeap::from(results), |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) - }); - let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::<'a, R, _>::new( - index, - BinaryHeap::from(results), - ); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, true) => { - let seq = seq_filter(BinaryHeap::from(results), |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) - }); - let prefetcher = - StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = StreamPrefetcher::<_, _>::new( - index, - BinaryHeap::from(results), - Hints::default(), - ); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = BinaryHeap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, Hints::default()); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, Hints::default()); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); } - (accu_set, rough_set, estimation_by_threshold) - }, - )) + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); + } + (accu_set, rough_set, estimation_by_threshold) + })) } VectorKind::Vecf16 => { type Op = operator::Op, Dot>; - let original_vectors = vectors + let unprojected = vectors .into_iter() .map(|vector| { if let OwnedVector::Vecf16(vector) = vector { @@ -263,141 +255,134 @@ impl SearchBuilder for MaxsimBuilder { } }) .collect::>(); - let vectors = original_vectors - .clone() - .into_iter() - .map(|vector| RandomProject::project(vector.as_borrowed())); - Box::new(vectors.into_iter().zip(original_vectors).map( - |(vector, original_vector)| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - vector.clone(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - match (options.io_rerank, prefilter()) { - (Io::Plain, true) => { - let seq = seq_filter(BinaryHeap::from(results), |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) - }); - let prefetcher = PlainPrefetcher::<_, _>::new(index, seq); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::<_, _>::new( - index, - BinaryHeap::from(results), - ); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let seq = seq_filter(BinaryHeap::from(results), |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) - }); - let prefetcher = SimplePrefetcher::<'a, R, _>::new(index, seq); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::<'a, R, _>::new( - index, - BinaryHeap::from(results), - ); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, true) => { - let seq = seq_filter(BinaryHeap::from(results), |key| { - fetcher.filter(pointer_to_kv(key.1.0.0).0) - }); - let prefetcher = - StreamPrefetcher::<_, _>::new(index, seq, Hints::default()); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = StreamPrefetcher::<_, _>::new( - index, - BinaryHeap::from(results), - Hints::default(), - ); - let mut reranker = rerank_index::( - original_vector.clone(), - prefetcher, - ); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = BinaryHeap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, Hints::default()); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, Hints::default()); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); } - (accu_set, rough_set, estimation_by_threshold) - }, - )) + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); + } + (accu_set, rough_set, estimation_by_threshold) + })) } }; let mut updates = Vec::new(); @@ -492,3 +477,11 @@ impl Iterator for IntoIterSorted { impl ExactSizeIterator for IntoIterSorted {} impl std::iter::FusedIterator for IntoIterSorted {} + +#[inline(always)] +pub fn id_0(f: F) -> F +where + F: for<'a> FnMut(&(A, AlwaysEqual<&'a mut (B, C, &'a mut D)>)) -> R, +{ + f +} diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs index 16ddd5d7..960ce6c9 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/scanners/mod.rs @@ -17,7 +17,7 @@ mod maxsim; use super::opclass::Opfamily; use crate::index::lazy_cell::LazyCell; -use algorithm::{Bump, RelationPrefetch, RelationRead, RelationReadStream}; +use algorithm::{Bump, RelationPrefetch, RelationRead, RelationReadStream, Sequence}; use pgrx::pg_sys::Datum; pub use default::DefaultBuilder; @@ -43,6 +43,7 @@ pub struct SearchOptions { pub maxsim_threshold: u32, pub io_search: Io, pub io_rerank: Io, + pub prefilter: bool, } pub trait SearchBuilder: 'static { @@ -54,23 +55,73 @@ pub trait SearchBuilder: 'static { self, relation: &'a R, options: SearchOptions, - fetcher: impl SearchFetcher + 'a, + fetcher: impl Fetcher + 'a, bump: &'a impl Bump, ) -> Box + 'a> where R: RelationRead + RelationPrefetch + RelationReadStream; } -pub trait SearchFetcher { - fn fetch(&mut self, ctid: [u16; 3]) -> Option<(&[Datum; 32], &[bool; 32])>; - fn filter(&mut self, key: [u16; 3]) -> bool; +pub trait Fetcher { + type Tuple<'a>: Tuple + where + Self: 'a; + + fn fetch(&mut self, key: [u16; 3]) -> Option>; } -impl T> SearchFetcher for LazyCell { - fn fetch(&mut self, key: [u16; 3]) -> Option<(&[Datum; 32], &[bool; 32])> { +pub trait Tuple { + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); + fn filter(&mut self) -> bool; +} + +impl T> Fetcher for LazyCell { + type Tuple<'a> + = T::Tuple<'a> + where + Self: 'a; + + fn fetch(&mut self, key: [u16; 3]) -> Option> { LazyCell::force_mut(self).fetch(key) } - fn filter(&mut self, key: [u16; 3]) -> bool { - LazyCell::force_mut(self).filter(key) +} + +pub struct Filter { + sequence: S, + predicate: P, +} + +impl Sequence for Filter +where + S: Sequence, + P: FnMut(&S::Item) -> bool, +{ + type Item = S::Item; + + type Inner = S::Inner; + + fn next(&mut self) -> Option { + while !(self.predicate)(self.sequence.peek()?) { + let _ = self.sequence.next(); + } + self.sequence.next() + } + + fn peek(&mut self) -> Option<&Self::Item> { + while !(self.predicate)(self.sequence.peek()?) { + let _ = self.sequence.next(); + } + self.sequence.peek() + } + + fn into_inner(self) -> Self::Inner { + self.sequence.into_inner() + } +} + +pub fn filter(sequence: S, predicate: P) -> Filter { + Filter { + sequence, + predicate, } } From 94804897626660e1682a1aeb850864d4198cc390 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 22 May 2025 12:55:14 +0800 Subject: [PATCH 159/324] feat: remove random matrix (#255) Signed-off-by: usamoi --- Cargo.lock | 185 +-------------------- Cargo.toml | 5 +- crates/algorithm/Cargo.toml | 1 - crates/algorithm/src/tuples.rs | 2 +- crates/rabitq/Cargo.toml | 6 + crates/rabitq/build.rs | 27 +++ crates/rabitq/src/lib.rs | 1 + crates/rabitq/src/rotate.rs | 72 ++++++++ crates/random_orthogonal_matrix/Cargo.toml | 16 -- crates/random_orthogonal_matrix/src/lib.rs | 72 -------- crates/simd/Cargo.toml | 2 + crates/simd/src/fht.rs | 182 ++++++++++++++++++++ crates/simd/src/lib.rs | 4 + crates/simd/src/rotate.rs | 52 ++++++ src/index/algorithm.rs | 8 +- src/index/am/am_build.rs | 7 +- src/index/gucs.rs | 39 ----- src/index/mod.rs | 4 - src/index/projection.rs | 36 ---- 19 files changed, 362 insertions(+), 359 deletions(-) create mode 100644 crates/rabitq/build.rs create mode 100644 crates/rabitq/src/rotate.rs delete mode 100644 crates/random_orthogonal_matrix/Cargo.toml delete mode 100644 crates/random_orthogonal_matrix/src/lib.rs create mode 100644 crates/simd/src/fht.rs create mode 100644 crates/simd/src/rotate.rs delete mode 100644 src/index/projection.rs diff --git a/Cargo.lock b/Cargo.lock index e1d14d57..c53cf491 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -23,7 +23,6 @@ dependencies = [ "pin-project", "rabitq", "rand", - "random_orthogonal_matrix", "serde", "simd", "validator", @@ -63,15 +62,6 @@ version = "1.0.98" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" -[[package]] -name = "approx" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" -dependencies = [ - "num-traits", -] - [[package]] name = "atomic-traits" version = "0.4.0" @@ -82,12 +72,6 @@ dependencies = [ "rustc_version", ] -[[package]] -name = "autocfg" -version = "1.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" - [[package]] name = "bindgen" version = "0.71.1" @@ -125,12 +109,6 @@ dependencies = [ "wyz", ] -[[package]] -name = "bytemuck" -version = "1.23.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9134a6ef01ce4b366b50689c94f82c14bc72bc5d0386829828a2e2752ef7958c" - [[package]] name = "byteorder" version = "1.5.0" @@ -644,28 +622,12 @@ dependencies = [ "windows-targets 0.53.0", ] -[[package]] -name = "libm" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" - [[package]] name = "litemap" version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956" -[[package]] -name = "matrixmultiply" -version = "0.3.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" -dependencies = [ - "autocfg", - "rawpointer", -] - [[package]] name = "memchr" version = "2.7.4" @@ -678,33 +640,6 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" -[[package]] -name = "nalgebra" -version = "0.33.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c4b5f057b303842cf3262c27e465f4c303572e7f6b0648f60e16248ac3397f4" -dependencies = [ - "approx", - "matrixmultiply", - "nalgebra-macros", - "num-complex", - "num-rational", - "num-traits", - "simba", - "typenum", -] - -[[package]] -name = "nalgebra-macros" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "nom" version = "7.1.3" @@ -715,55 +650,6 @@ dependencies = [ "minimal-lexical", ] -[[package]] -name = "num-bigint" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" -dependencies = [ - "num-integer", - "num-traits", -] - -[[package]] -name = "num-complex" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" -dependencies = [ - "num-traits", -] - -[[package]] -name = "num-integer" -version = "0.1.46" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" -dependencies = [ - "num-traits", -] - -[[package]] -name = "num-rational" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" -dependencies = [ - "num-bigint", - "num-integer", - "num-traits", -] - -[[package]] -name = "num-traits" -version = "0.2.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" -dependencies = [ - "autocfg", - "libm", -] - [[package]] name = "once_cell" version = "1.21.3" @@ -1017,7 +903,10 @@ name = "rabitq" version = "0.0.0" dependencies = [ "distance", + "rand", + "rand_chacha", "simd", + "zerocopy", ] [[package]] @@ -1055,32 +944,6 @@ dependencies = [ "getrandom", ] -[[package]] -name = "rand_distr" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" -dependencies = [ - "num-traits", - "rand", -] - -[[package]] -name = "random_orthogonal_matrix" -version = "0.0.0" -dependencies = [ - "nalgebra", - "rand", - "rand_chacha", - "rand_distr", -] - -[[package]] -name = "rawpointer" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" - [[package]] name = "rayon" version = "1.10.0" @@ -1151,15 +1014,6 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" -[[package]] -name = "safe_arch" -version = "0.7.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96b02de82ddbe1b636e6170c21be622223aea188ef2e139be0a5b219ec215323" -dependencies = [ - "bytemuck", -] - [[package]] name = "same-file" version = "1.0.6" @@ -1244,26 +1098,15 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" -[[package]] -name = "simba" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3a386a501cd104797982c15ae17aafe8b9261315b5d07e3ec803f2ea26be0fa" -dependencies = [ - "approx", - "num-complex", - "num-traits", - "paste", - "wide", -] - [[package]] name = "simd" version = "0.0.0" dependencies = [ "cc", "half 2.6.0", + "paste", "rand", + "seq-macro", "simd_macros", "zerocopy", ] @@ -1419,12 +1262,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bfb942dfe1d8e29a7ee7fcbde5bd2b9a25fb89aa70caea2eba3bee836ff41076" -[[package]] -name = "typenum" -version = "1.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f" - [[package]] name = "unescape" version = "0.1.0" @@ -1518,8 +1355,8 @@ dependencies = [ "paste", "pgrx", "pgrx-catalog", + "rabitq", "rand", - "random_orthogonal_matrix", "seq-macro", "serde", "simd", @@ -1557,16 +1394,6 @@ dependencies = [ "wit-bindgen-rt", ] -[[package]] -name = "wide" -version = "0.7.32" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41b5576b9a81633f3e8df296ce0063042a73507636cbe956c61133dd7034ab22" -dependencies = [ - "bytemuck", - "safe_arch", -] - [[package]] name = "winapi-util" version = "0.1.9" diff --git a/Cargo.toml b/Cargo.toml index 444c74c0..c23d84a3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,7 +25,7 @@ algorithm = { path = "./crates/algorithm" } always_equal = { path = "./crates/always_equal" } distance = { path = "./crates/distance" } k_means = { path = "./crates/k_means" } -random_orthogonal_matrix = { path = "./crates/random_orthogonal_matrix" } +rabitq = { path = "./crates/rabitq" } simd = { path = "./crates/simd" } vector = { path = "./crates/vector" } @@ -34,7 +34,7 @@ paste.workspace = true pgrx = { version = "=0.14.1", default-features = false, features = ["cshim"] } pgrx-catalog = "0.2.0" rand.workspace = true -seq-macro = "0.3.6" +seq-macro.workspace = true serde.workspace = true toml = "0.8.22" validator.workspace = true @@ -58,6 +58,7 @@ edition = "2024" half = { version = "2.6.0", features = ["zerocopy"] } paste = "1" rand = "0.9.1" +seq-macro = "0.3.6" serde = { version = "1", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } zerocopy = { version = "0.8.25", features = ["derive"] } diff --git a/crates/algorithm/Cargo.toml b/crates/algorithm/Cargo.toml index acbda149..ffdef2ee 100644 --- a/crates/algorithm/Cargo.toml +++ b/crates/algorithm/Cargo.toml @@ -9,7 +9,6 @@ always_equal = { path = "../always_equal" } distance = { path = "../distance" } k_means = { path = "../k_means" } rabitq = { path = "../rabitq" } -random_orthogonal_matrix = { path = "../random_orthogonal_matrix" } simd = { path = "../simd" } vector = { path = "../vector" } diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 3d103fc6..1df5f77f 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -20,7 +20,7 @@ use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 8; +const VERSION: u64 = 9; pub trait Tuple: 'static { fn serialize(&self) -> Vec; diff --git a/crates/rabitq/Cargo.toml b/crates/rabitq/Cargo.toml index f1572994..6dff2f62 100644 --- a/crates/rabitq/Cargo.toml +++ b/crates/rabitq/Cargo.toml @@ -8,5 +8,11 @@ publish = false distance = { path = "../distance" } simd = { path = "../simd" } +zerocopy.workspace = true + +[build-dependencies] +rand = "0.9.1" +rand_chacha = "0.9.0" + [lints] workspace = true diff --git a/crates/rabitq/build.rs b/crates/rabitq/build.rs new file mode 100644 index 00000000..a56f6a95 --- /dev/null +++ b/crates/rabitq/build.rs @@ -0,0 +1,27 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use std::env::var; +use std::error::Error; +use std::path::PathBuf; + +fn main() -> Result<(), Box> { + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha12Rng; + let mut rng = ChaCha12Rng::from_seed([7; 32]); + let bits = (0..262144).map(|_| rng.random::()).collect::>(); + let out_dir = var("OUT_DIR")?; + std::fs::write(PathBuf::from(out_dir).join("bits"), bits)?; + Ok(()) +} diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index 6274274e..d96f3676 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -18,6 +18,7 @@ pub mod binary; pub mod block; pub mod packing; +pub mod rotate; use binary::BinaryLut; use block::BlockLut; diff --git a/crates/rabitq/src/rotate.rs b/crates/rabitq/src/rotate.rs new file mode 100644 index 00000000..1bbbb051 --- /dev/null +++ b/crates/rabitq/src/rotate.rs @@ -0,0 +1,72 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +const BITS: &[u8; 262144] = include_bytes!(concat!(env!("OUT_DIR"), "/bits")); + +const _: () = assert!(BITS[0] == 246); +const _: () = assert!(BITS[1] == 133); +const _: () = assert!(BITS[2] == 163); +const _: () = assert!(BITS[3] == 106); +const _: () = assert!(BITS[4] == 54); +const _: () = assert!(BITS[5] == 126); +const _: () = assert!(BITS[6] == 9); +const _: () = assert!(BITS[7] == 115); + +static BITS_0: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[0]); +static BITS_1: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[1]); +static BITS_2: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[2]); +static BITS_3: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[3]); + +fn kacs_walk(result: &mut [f32]) { + let n = result.len(); + let m = n / 2; + let (l, t) = result.split_at_mut(m); + let (_, r) = t.split_at_mut(n - 2 * m); + simd::rotate::givens(l, r); +} + +pub fn rotate(vector: &[f32]) -> Vec { + use simd::Floating; + use std::ops::Bound::{Excluded, Included, Unbounded}; + + let mut result = vector.to_vec(); + let n = vector.len(); + let base = n.ilog2(); + let scale = 1.0 / ((1_usize << base) as f32).sqrt(); + + let l = (Unbounded, Excluded(1_usize << base)); + let r = (Included(n - (1_usize << base)), Unbounded); + + simd::rotate::flip(&BITS_0, &mut result); + simd::fht::fht(&mut result[l]); + f32::vector_mul_scalar_inplace(&mut result[l], scale); + kacs_walk(&mut result); + + simd::rotate::flip(&BITS_1, &mut result); + simd::fht::fht(&mut result[r]); + f32::vector_mul_scalar_inplace(&mut result[r], scale); + kacs_walk(&mut result); + + simd::rotate::flip(&BITS_2, &mut result); + simd::fht::fht(&mut result[l]); + f32::vector_mul_scalar_inplace(&mut result[l], scale); + kacs_walk(&mut result); + + simd::rotate::flip(&BITS_3, &mut result); + simd::fht::fht(&mut result[r]); + f32::vector_mul_scalar_inplace(&mut result[r], scale); + kacs_walk(&mut result); + + result +} diff --git a/crates/random_orthogonal_matrix/Cargo.toml b/crates/random_orthogonal_matrix/Cargo.toml deleted file mode 100644 index ee869dd1..00000000 --- a/crates/random_orthogonal_matrix/Cargo.toml +++ /dev/null @@ -1,16 +0,0 @@ -[package] -name = "random_orthogonal_matrix" -version.workspace = true -edition.workspace = true -publish = false - -[dependencies] -# lock algebra version forever so that the QR decomposition never changes for same input -nalgebra = "=0.33.0" - -rand.workspace = true -rand_chacha = "0.9.0" -rand_distr = "0.5.1" - -[lints] -workspace = true diff --git a/crates/random_orthogonal_matrix/src/lib.rs b/crates/random_orthogonal_matrix/src/lib.rs deleted file mode 100644 index b611f2b8..00000000 --- a/crates/random_orthogonal_matrix/src/lib.rs +++ /dev/null @@ -1,72 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -use nalgebra::DMatrix; - -#[ignore] -#[test] -fn check_full_rank_matrix() { - let parallelism = std::thread::available_parallelism().unwrap().get(); - std::thread::scope(|scope| { - let mut threads = vec![]; - for remainder in 0..parallelism { - threads.push(scope.spawn(move || { - for n in (0..=60000).filter(|x| x % parallelism == remainder) { - let matrix = random_full_rank_matrix(n); - assert!(matrix.is_invertible()); - } - })); - } - for thread in threads { - thread.join().unwrap(); - } - }); -} - -fn random_full_rank_matrix(n: usize) -> DMatrix { - use rand::{Rng, SeedableRng}; - use rand_chacha::ChaCha12Rng; - use rand_distr::StandardNormal; - let mut rng = ChaCha12Rng::from_seed([7; 32]); - DMatrix::from_fn(n, n, |_, _| rng.sample(StandardNormal)) -} - -#[test] -fn check_random_orthogonal_matrix() { - assert_eq!( - random_orthogonal_matrix(2), - vec![vec![-0.5424608, -0.8400813], vec![0.8400813, -0.54246056]] - ); - assert_eq!( - random_orthogonal_matrix(3), - vec![ - vec![-0.5309615, -0.69094884, -0.49058124], - vec![0.8222731, -0.56002235, -0.10120347], - vec![0.20481002, 0.45712686, -0.86549866] - ] - ); -} - -pub fn random_orthogonal_matrix(n: usize) -> Vec> { - use nalgebra::QR; - let matrix = random_full_rank_matrix(n); - // QR decomposition is unique if the matrix is full rank - let qr = QR::new(matrix); - let q = qr.q(); - let mut projection = Vec::new(); - for row in q.row_iter() { - projection.push(row.iter().copied().collect::>()); - } - projection -} diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 04d106d1..32bc9779 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -8,6 +8,8 @@ publish = false simd_macros = { path = "../simd_macros" } half.workspace = true +paste.workspace = true +seq-macro.workspace = true zerocopy.workspace = true [dev-dependencies] diff --git a/crates/simd/src/fht.rs b/crates/simd/src/fht.rs new file mode 100644 index 00000000..be277473 --- /dev/null +++ b/crates/simd/src/fht.rs @@ -0,0 +1,182 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +#[inline(always)] +fn basic_1(x: &mut [f32]) { + assert!(x.len() == (1 << (8 + 1))); + for i in 0..1 << (8 - Q) { + basic_2::(&mut x[i << (Q + 1)..][..1 << (Q + 1)]); + } +} + +#[inline(always)] +fn basic_2(x: &mut [f32]) { + assert!(x.len() == (1 << (Q + 1))); + for j in 0..1 << Q { + let (l, r) = (x[j], x[j + (1 << Q)]); + (x[j], x[j + (1 << Q)]) = (l + r, l - r); + } +} + +mod step_1 { + seq_macro::seq!( + Q in 0..16 { + mod dispatch_~Q { + #[crate::multiversion("v4", "v3", "v2", "a2")] + pub fn f(x: &mut [f32]) { + crate::fht::basic_1::(x); + } + } + #[allow(unused_imports)] + pub use dispatch_~Q::f as dispatch_~Q; + } + ); +} + +mod step_2 { + seq_macro::seq!( + Q in 0..16 { + mod dispatch_~Q { + #[crate::multiversion("v4", "v3", "v2", "a2")] + pub fn f(x: &mut [f32]) { + crate::fht::basic_2::(x); + } + } + #[allow(unused_imports)] + pub use dispatch_~Q::f as dispatch_~Q; + } + ); +} + +macro_rules! fht { + ($p:literal, 0) => { + { + #[crate::multiversion("v4", "v3", "v2", "a2")] + fn walk(x: &mut [f32]) { + assert!(x.len() == (1 << $p)); + seq_macro::seq!( + Q in 0..$p { + for i in 0..1 << ($p - Q - 1) { + basic_2::(&mut x[i << (Q + 1)..][..1 << (Q + 1)]); + } + } + ); + } + walk as fn(&mut [f32]) + } + }; + ($p:literal, 1) => { + { + fn walk(x: &mut [f32]) { + assert!(x.len() == (1 << $p)); + seq_macro::seq!( + Q in 0..8 { + for i in (0..(1 << ($p - Q - 1))).step_by(1 << (8 - Q)) { + step_1::dispatch_~Q(&mut x[i << (Q + 1)..][..1 << (8 + 1)]); + } + } + ); + seq_macro::seq!( + Q in 8..$p { + for i in 0..1 << ($p - Q - 1) { + step_2::dispatch_~Q(&mut x[i << (Q + 1)..][..1 << (Q + 1)]); + } + } + ); + } + walk as fn(&mut [f32]) + } + }; +} + +pub fn fht(x: &mut [f32]) { + const FHT: [fn(&mut [f32]); 1 + 16] = [ + |_| (), + fht!(1, 0), + fht!(2, 0), + fht!(3, 0), + fht!(4, 0), + fht!(5, 0), + fht!(6, 0), + fht!(7, 0), + fht!(8, 0), + fht!(9, 1), + fht!(10, 1), + fht!(11, 1), + fht!(12, 1), + fht!(13, 1), + fht!(14, 1), + fht!(15, 1), + fht!(16, 1), + ]; + let n = x.len(); + let Some(i) = n.checked_ilog2() else { + panic!("the dimension of the vector is 0") + }; + if n != (1 << i) { + panic!("the dimension of the vector is not a power of 2"); + } + if i > 16 { + panic!("the dimension of the vector is too large"); + } + FHT[i as usize](x) +} + +#[cfg(test)] +mod tests { + fn native(x: &mut [f32]) { + let n = x.len(); + assert!(n.is_power_of_two()); + let mut h = 1; + while h < n { + for i in (0..n).step_by(h * 2) { + for j in i..i + h { + (x[j], x[j + h]) = (x[j] + x[j + h], x[j] - x[j + h]); + } + } + h *= 2; + } + } + + #[test] + fn fht() { + use rand::Rng; + use std::iter::zip; + const EPSILON: f32 = 1e-6; + let mut rng = rand::rng(); + let mut n = 1_usize; + while n <= 65536 { + let x = (0..n) + .map(|_| rng.random_range(-1.0_f32..=1.0_f32)) + .collect::>(); + let x_expected = { + let mut x = x.clone(); + native(x.as_mut_slice()); + x + }; + let x_got = { + let mut x = x.clone(); + crate::fht::fht(x.as_mut_slice()); + x + }; + let mse = zip(x_expected, x_got) + .map(|(x, y)| (x - y) * (x - y)) + .sum::() + / n as f32; + eprintln!("n = {n}, mse = {mse:.12}"); + assert!(mse <= EPSILON); + n *= 2; + } + } +} diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 627ec3fa..31307998 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -14,6 +14,8 @@ #![cfg_attr(target_arch = "x86_64", feature(avx512_target_feature))] #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] +#![feature(slice_as_chunks)] +#![feature(select_unpredictable)] #![allow(unsafe_code)] mod aligned; @@ -23,7 +25,9 @@ mod f32; pub mod bit; pub mod fast_scan; +pub mod fht; pub mod quantize; +pub mod rotate; pub mod u8; pub trait Floating: diff --git a/crates/simd/src/rotate.rs b/crates/simd/src/rotate.rs new file mode 100644 index 00000000..7a211e51 --- /dev/null +++ b/crates/simd/src/rotate.rs @@ -0,0 +1,52 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub mod givens { + #[crate::multiversion("v4", "v3", "v2", "a2")] + pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let scale = 1.0 / (2.0_f32).sqrt(); + for i in 0..n { + (lhs[i], rhs[i]) = ((lhs[i] + rhs[i]) * scale, (lhs[i] - rhs[i]) * scale); + } + } +} + +pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { + givens::givens(lhs, rhs) +} + +pub mod flip { + #[crate::multiversion("v4", "v3", "v2", "a2")] + pub fn flip(bits: &[u64; 1024], result: &mut [f32]) { + use std::hint::select_unpredictable; + let result: &mut [u32] = unsafe { std::mem::transmute(result) }; + let (slice, remainder) = result.as_chunks_mut::<64>(); + let n = slice.len(); + assert!(n <= 1024); + for i in 0..n { + for j in 0..64 { + slice[i][j] ^= select_unpredictable((bits[i] & (1 << j)) != 0, 0x80000000, 0); + } + } + for j in 0..remainder.len() { + remainder[j] ^= select_unpredictable((bits[n] & (1 << j)) != 0, 0x80000000, 0); + } + } +} + +pub fn flip(bits: &[u64; 1024], result: &mut [f32]) { + flip::flip(bits, result) +} diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index f75b9442..e1b89058 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -430,19 +430,19 @@ pub trait RandomProject { impl RandomProject for VectBorrowed<'_, f32> { type Output = VectOwned; fn project(self) -> VectOwned { - use crate::index::projection::project; + use rabitq::rotate::rotate; let input = self.slice(); - VectOwned::new(project(input)) + VectOwned::new(rotate(input)) } } impl RandomProject for VectBorrowed<'_, f16> { type Output = VectOwned; fn project(self) -> VectOwned { - use crate::index::projection::project; + use rabitq::rotate::rotate; use simd::Floating; let input = f16::vector_to_f32(self.slice()); - VectOwned::new(f16::vector_from_f32(&project(&input))) + VectOwned::new(f16::vector_from_f32(&rotate(&input))) } } diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 32ce4102..532c5e1d 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -971,7 +971,7 @@ fn make_internal_build( ) -> Vec>> { use std::iter::once; k_means::preprocess(internal_build.build_threads as _, &mut samples, |sample| { - *sample = crate::index::projection::project(sample) + *sample = rabitq::rotate::rotate(sample) }); let mut result = Vec::>>::new(); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { @@ -1096,10 +1096,7 @@ fn make_external_build( if vector_options.dims != vector.as_borrowed().dims() { pgrx::error!("external build: incorrect dimension, id = {id}"); } - vectors.insert( - id, - crate::index::projection::project(vector.as_borrowed().slice()), - ); + vectors.insert(id, rabitq::rotate::rotate(vector.as_borrowed().slice())); } }); if parents.len() >= 2 && parents.values().all(|x| x.is_none()) { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 3ccfe94b..cebfa08a 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -26,9 +26,6 @@ pub enum PostgresIo { read_stream, } -static PREWARM_DIM: GucSetting> = - GucSetting::>::new(Some(c"64,128,256,384,512,768,1024,1536")); - static PROBES: GucSetting> = GucSetting::>::new(Some(c"")); static EPSILON: GucSetting = GucSetting::::new(1.9); static MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); @@ -81,14 +78,6 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); - GucRegistry::define_string_guc( - "vchordrq.prewarm_dim", - "prewarm_dim when the extension is loading.", - "prewarm_dim when the extension is loading.", - &PREWARM_DIM, - GucContext::Userset, - GucFlags::default(), - ); GucRegistry::define_int_guc( "vchordrq.maxsim_refine", "`maxsim_refine` argument of vchordrq.", @@ -186,34 +175,6 @@ pub fn maxsim_threshold() -> u32 { MAXSIM_THRESHOLD.get() as u32 } -pub fn prewarm_dim() -> Vec { - match PREWARM_DIM.get() { - None => Vec::new(), - Some(probes) => { - let mut result = Vec::new(); - let mut current = None; - for &c in probes.to_bytes() { - match c { - b' ' => continue, - b',' => result.push(current.take().expect("empty prewarm_dim")), - b'0'..=b'9' => { - if let Some(x) = current.as_mut() { - *x = *x * 10 + (c - b'0') as u32; - } else { - current = Some((c - b'0') as u32); - } - } - c => pgrx::error!("unknown character in prewarm_dim: ASCII = {c}"), - } - } - if let Some(current) = current { - result.push(current); - } - result - } - } -} - pub fn prefilter() -> bool { PREFILTER.get() } diff --git a/src/index/mod.rs b/src/index/mod.rs index 370e8897..32707d4c 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -19,7 +19,6 @@ pub mod gucs; pub mod hook; pub mod lazy_cell; pub mod opclass; -pub mod projection; pub mod scanners; pub mod storage; pub mod types; @@ -28,7 +27,4 @@ pub fn init() { am::init(); hook::init(); gucs::init(); - for x in gucs::prewarm_dim() { - projection::prewarm(x as _); - } } diff --git a/src/index/projection.rs b/src/index/projection.rs deleted file mode 100644 index 76ede977..00000000 --- a/src/index/projection.rs +++ /dev/null @@ -1,36 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -use random_orthogonal_matrix::random_orthogonal_matrix; -use std::sync::OnceLock; - -fn matrix(n: usize) -> Option<&'static Vec>> { - static MATRIXS: [OnceLock>>; 1 + 60000] = [const { OnceLock::new() }; 1 + 60000]; - MATRIXS - .get(n) - .map(|x| x.get_or_init(|| random_orthogonal_matrix(n))) -} - -pub fn prewarm(n: usize) { - let _ = matrix(n); -} - -pub fn project(vector: &[f32]) -> Vec { - use simd::Floating; - let n = vector.len(); - let matrix = matrix(n).expect("dimension too large"); - (0..n) - .map(|i| f32::reduce_sum_of_xy(vector, &matrix[i])) - .collect() -} From 7a10814c2e3d19be66c9347bbee4efee9ca6c611 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 22 May 2025 12:57:45 +0800 Subject: [PATCH 160/324] chore: upload 0.4.0 schema scripts (#254) Signed-off-by: usamoi --- sql/install/vchord--0.4.0.sql | 566 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.3.0--0.4.0.sql | 1 + 2 files changed, 567 insertions(+) create mode 100644 sql/install/vchord--0.4.0.sql create mode 100644 sql/upgrade/vchord--0.3.0--0.4.0.sql diff --git a/sql/install/vchord--0.4.0.sql b/sql/install/vchord--0.4.0.sql new file mode 100644 index 00000000..522e9c8a --- /dev/null +++ b/sql/install/vchord--0.4.0.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:22 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:21 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:133 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:74 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:23 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.3.0--0.4.0.sql b/sql/upgrade/vchord--0.3.0--0.4.0.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.3.0--0.4.0.sql @@ -0,0 +1 @@ +/* nothing to do */ From 245858148395d18e58dcb6d615cd90fd241af579 Mon Sep 17 00:00:00 2001 From: usamoi Date: Sat, 24 May 2025 17:45:33 +0800 Subject: [PATCH 161/324] fix: do not call kacs_walk if dimension is a power of 2 (#258) Signed-off-by: usamoi --- crates/rabitq/src/rotate.rs | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/crates/rabitq/src/rotate.rs b/crates/rabitq/src/rotate.rs index 1bbbb051..130d7dfa 100644 --- a/crates/rabitq/src/rotate.rs +++ b/crates/rabitq/src/rotate.rs @@ -51,22 +51,30 @@ pub fn rotate(vector: &[f32]) -> Vec { simd::rotate::flip(&BITS_0, &mut result); simd::fht::fht(&mut result[l]); f32::vector_mul_scalar_inplace(&mut result[l], scale); - kacs_walk(&mut result); + if n != (1_usize << base) { + kacs_walk(&mut result); + } simd::rotate::flip(&BITS_1, &mut result); simd::fht::fht(&mut result[r]); f32::vector_mul_scalar_inplace(&mut result[r], scale); - kacs_walk(&mut result); + if n != (1_usize << base) { + kacs_walk(&mut result); + } simd::rotate::flip(&BITS_2, &mut result); simd::fht::fht(&mut result[l]); f32::vector_mul_scalar_inplace(&mut result[l], scale); - kacs_walk(&mut result); + if n != (1_usize << base) { + kacs_walk(&mut result); + } simd::rotate::flip(&BITS_3, &mut result); simd::fht::fht(&mut result[r]); f32::vector_mul_scalar_inplace(&mut result[r], scale); - kacs_walk(&mut result); + if n != (1_usize << base) { + kacs_walk(&mut result); + } result } From 1348a75ba2303879b22a5d738c767d48f925c4a2 Mon Sep 17 00:00:00 2001 From: usamoi Date: Sat, 24 May 2025 17:51:07 +0800 Subject: [PATCH 162/324] chore: upload 0.4.1 schema scripts (#259) Signed-off-by: usamoi --- sql/install/vchord--0.4.1.sql | 566 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.4.0--0.4.1.sql | 1 + 2 files changed, 567 insertions(+) create mode 100644 sql/install/vchord--0.4.1.sql create mode 100644 sql/upgrade/vchord--0.4.0--0.4.1.sql diff --git a/sql/install/vchord--0.4.1.sql b/sql/install/vchord--0.4.1.sql new file mode 100644 index 00000000..522e9c8a --- /dev/null +++ b/sql/install/vchord--0.4.1.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:22 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:21 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:133 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:74 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:23 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.4.0--0.4.1.sql b/sql/upgrade/vchord--0.4.0--0.4.1.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.4.0--0.4.1.sql @@ -0,0 +1 @@ +/* nothing to do */ From ea426e76795e5a36459eb7735906ab2765bd6526 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 26 May 2025 10:51:19 +0800 Subject: [PATCH 163/324] ci: add macos (#262) job will only run if pull request or commit message contains `try-job: psql_macos`, or it's triggered manually Signed-off-by: usamoi --- .github/workflows/check.yml | 120 +++++++++++------- crates/simd/cshim/aarch64.c | 16 +-- rust-toolchain.toml | 2 +- .../filter_rerank_in_index.slt | 0 .../filter_rerank_in_table.slt | 0 5 files changed, 83 insertions(+), 55 deletions(-) rename tests/{pg16 => general}/filter_rerank_in_index.slt (100%) rename tests/{pg16 => general}/filter_rerank_in_table.slt (100%) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index aaa8971d..cdb9addd 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -1,39 +1,8 @@ name: Check on: - pull_request: - paths: - - ".cargo" - - ".github/workflows/check.yml" - - "deny.toml" - - "crates/**" - - "scripts/**" - - "src/**" - - "tests/**" - - "build.rs" - - "Cargo.lock" - - "Cargo.toml" - - "rust-toolchain.toml" - - "rustfmt.toml" - - "taplo.toml" - - "vchord.control" push: - paths: - - ".cargo" - - ".github/workflows/check.yml" - - "deny.toml" - - "crates/**" - - "scripts/**" - - "src/**" - - "tests/**" - - "build.rs" - - "Cargo.lock" - - "Cargo.toml" - - "rust-toolchain.toml" - - "rustfmt.toml" - - "taplo.toml" - - "vchord.control" - merge_group: + pull_request: workflow_dispatch: concurrency: @@ -235,25 +204,17 @@ jobs: - name: Install run: cargo pgrx install -p vchord --features pg${{ matrix.version }} --release --sudo - - name: Sqllogictest + - name: Service run: | sudo systemctl start postgresql psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - sqllogictest --db $USER --user $USER './tests/general/*.slt' - - name: Sqllogictest(PostgreSQL 17 features) - if: matrix.version == '17' - run: | - sudo systemctl start postgresql - psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - sqllogictest --db $USER --user $USER './tests/pg17/*.slt' + - name: Sqllogictest + run: sqllogictest --db $USER --user $USER './tests/general/*.slt' - - name: Sqllogictest(PostgreSQL 16 features) - if: matrix.version == '16' - run: | - sudo systemctl start postgresql - psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - sqllogictest --db $USER --user $USER './tests/pg16/*.slt' + - name: Sqllogictest (PostgreSQL 17) + if: matrix.version == '17' + run: sqllogictest --db $USER --user $USER './tests/pg17/*.slt' - name: Package env: @@ -308,3 +269,70 @@ jobs: ./build/postgresql-${{ matrix.version }}-vchord_0.0.0-1_${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }}.deb compression-level: 9 retention-days: 14 + + psql_macos: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'try-job: psql_macos')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'try-job: psql_macos')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + version: ["14", "17"] + arch: ["aarch64"] + + runs-on: "macos-15" + + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: true + + steps: + - name: Set up Environment + run: | + rustup set profile + + HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install postgresql@${{ matrix.version }} + HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install pgvector + brew services start postgresql@${{ matrix.version }} + for i in {1..60}; do [ -S /tmp/.s.PGSQL.5432 ] && echo "PostgreSQL ready" && break || sleep 1; done + [ -S /tmp/.s.PGSQL.5432 ] || echo "PostgreSQL socket not found after 60 seconds" + $(brew --prefix postgresql@${{ matrix.version }})/bin/createdb -O $USER $USER + if [[ ${{ matrix.version }} != 13 && ${{ matrix.version }} != 14 && ${{ matrix.version }} != 15 ]]; then + $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.dylib"' + else + $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.so"' + fi + brew services stop postgresql@${{ matrix.version }} + + curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.14.1/cargo-pgrx-v0.14.1-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./cargo-pgrx | tee /usr/local/bin/cargo-pgrx > /dev/null && chmod 755 /usr/local/bin/cargo-pgrx + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Checkout + uses: actions/checkout@v4 + + - name: Clippy + run: | + env CC=$(brew --prefix llvm@18)/bin/clang PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config cargo clippy -p vchord --features pg${{ matrix.version }} -- --no-deps + + - name: Install + run: | + env CC=$(brew --prefix llvm@18)/bin/clang PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config cargo pgrx install -p vchord --features pg${{ matrix.version }} --release --sudo --pg-config=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config + + - name: Service + run: | + brew services start postgresql@${{ matrix.version }} + for i in {1..60}; do [ -S /tmp/.s.PGSQL.5432 ] && echo "PostgreSQL ready" && break || sleep 1; done + [ -S /tmp/.s.PGSQL.5432 ] || echo "PostgreSQL socket not found after 60 seconds" + $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: sqllogictest --db $USER --user $USER './tests/general/*.slt' + + - name: Sqllogictest (PostgreSQL 17) + if: matrix.version == '17' + run: sqllogictest --db $USER --user $USER './tests/pg17/*.slt' diff --git a/crates/simd/cshim/aarch64.c b/crates/simd/cshim/aarch64.c index 614eaa55..8e2d9336 100644 --- a/crates/simd/cshim/aarch64.c +++ b/crates/simd/cshim/aarch64.c @@ -85,7 +85,7 @@ __attribute__((target("+sve"))) float fp16_reduce_sum_of_xy_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svfloat16_t xy = svdup_f16(0.0); for (size_t i = 0; i < n; i += svcnth()) { - svbool_t mask = svwhilelt_b16(i, n); + svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); svfloat16_t x = svld1_f16(mask, a + i); svfloat16_t y = svld1_f16(mask, b + i); xy = svmla_f16_x(mask, xy, x, y); @@ -154,7 +154,7 @@ __attribute__((target("+sve"))) float fp16_reduce_sum_of_d2_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svfloat16_t d2 = svdup_f16(0.0); for (size_t i = 0; i < n; i += svcnth()) { - svbool_t mask = svwhilelt_b16(i, n); + svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); svfloat16_t x = svld1_f16(mask, a + i); svfloat16_t y = svld1_f16(mask, b + i); svfloat16_t d = svsub_f16_x(mask, x, y); @@ -167,7 +167,7 @@ __attribute__((target("+sve"))) float fp32_reduce_sum_of_x_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32(i, n); + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, this + i); sum = svadd_f32_x(mask, sum, x); } @@ -178,7 +178,7 @@ __attribute__((target("+sve"))) float fp32_reduce_sum_of_abs_x_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32(i, n); + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, this + i); sum = svadd_f32_x(mask, sum, svabs_f32_x(mask, x)); } @@ -189,7 +189,7 @@ __attribute__((target("+sve"))) float fp32_reduce_sum_of_x2_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32(i, n); + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, this + i); sum = svmla_f32_x(mask, sum, x, x); } @@ -202,7 +202,7 @@ fp32_reduce_min_max_of_x_a3_256(float *restrict this, size_t n, float *out_min, svfloat32_t min = svdup_f32(1.0 / 0.0); svfloat32_t max = svdup_f32(-1.0 / 0.0); for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32(i, n); + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, this + i); min = svmin_f32_x(mask, min, x); max = svmax_f32_x(mask, max, x); @@ -216,7 +216,7 @@ fp32_reduce_sum_of_xy_a3_256(float *restrict lhs, float *restrict rhs, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32(i, n); + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, lhs + i); svfloat32_t y = svld1_f32(mask, rhs + i); sum = svmla_f32_x(mask, sum, x, y); @@ -229,7 +229,7 @@ fp32_reduce_sum_of_d2_a3_256(float *restrict lhs, float *restrict rhs, size_t n) { svfloat32_t sum = svdup_f32(0.0); for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32(i, n); + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, lhs + i); svfloat32_t y = svld1_f32(mask, rhs + i); svfloat32_t d = svsub_f32_x(mask, x, y); diff --git a/rust-toolchain.toml b/rust-toolchain.toml index ed9009cd..1b42d5ea 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2025-04-25" +channel = "nightly-2025-04-26" diff --git a/tests/pg16/filter_rerank_in_index.slt b/tests/general/filter_rerank_in_index.slt similarity index 100% rename from tests/pg16/filter_rerank_in_index.slt rename to tests/general/filter_rerank_in_index.slt diff --git a/tests/pg16/filter_rerank_in_table.slt b/tests/general/filter_rerank_in_table.slt similarity index 100% rename from tests/pg16/filter_rerank_in_table.slt rename to tests/general/filter_rerank_in_table.slt From 76b8316fe794676be5b0beba4716eb58f043b78d Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 27 May 2025 13:48:03 +0800 Subject: [PATCH 164/324] chore: add a script for building (#263) environment variable `PGRX_PG_CONFIG_PATH` must be set to path of `pg_config` package the extension: `cargo make package -o ./build/raw` install the extension by building: `cargo make install --sudo` install the extension from package: `cargo make install --sudo -i ./build/raw` with this script, `bundle` and `bundle_loader` is set correctly on MacOS, to avoid potential symbol resolution bugs try-job: psql_macos Signed-off-by: usamoi --- .cargo/config.toml | 3 + .github/workflows/check.yml | 62 +++---- .github/workflows/release.yml | 32 +--- Cargo.lock | 280 +++++++++++++++++++++++++++++ build.rs | 23 ++- crates/make/Cargo.toml | 12 ++ crates/make/src/main.rs | 322 ++++++++++++++++++++++++++++++++++ 7 files changed, 678 insertions(+), 56 deletions(-) create mode 100644 crates/make/Cargo.toml create mode 100644 crates/make/src/main.rs diff --git a/.cargo/config.toml b/.cargo/config.toml index 2196f5fc..63524241 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,2 +1,5 @@ +[alias] +make = "run -p make --" + [env] CC = "clang" diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index cdb9addd..6d87932f 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -177,6 +177,7 @@ jobs: sudo apt-get install -y postgresql-common sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} + echo PGRX_PG_CONFIG_PATH=pg_config >> $GITHUB_ENV sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector echo "local all all trust" | sudo tee /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf @@ -187,9 +188,6 @@ jobs: sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.so"' sudo systemctl stop postgresql - curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.14.1/cargo-pgrx-v0.14.1-$(uname -m)-unknown-linux-gnu.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx - cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - name: Set up Sccache @@ -202,7 +200,9 @@ jobs: run: cargo clippy -p vchord --features pg${{ matrix.version }} -- --no-deps - name: Install - run: cargo pgrx install -p vchord --features pg${{ matrix.version }} --release --sudo + run: | + cargo make package -o ./build/raw + cargo make install --sudo -i ./build/raw - name: Service run: | @@ -223,28 +223,12 @@ jobs: ARCH: ${{ matrix.arch }} PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} run: | - cat $(pg_config --sharedir)/extension/vchord--0.0.0.sql | expand -t 4 > ./sql/install/vchord--$SEMVER.sql - - mkdir -p ./build/zip - cp -a ./sql/upgrade/. ./build/zip/ - cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql - cp ./vchord.control ./build/zip/vchord.control - cp ./target/release/libvchord.so ./build/zip/vchord.so - zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* + (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) mkdir -p ./build/deb mkdir -p ./build/deb/DEBIAN - mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ - mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ - for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file - done - for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file - done - for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file - done + mkdir -p ./build/deb$(pg_config --pkglibdir) && cp -r ./build/raw/pkglibdir/. ./build/deb$(pg_config --pkglibdir) + mkdir -p ./build/deb$(pg_config --sharedir) && cp -r ./build/raw/sharedir/. ./build/deb$(pg_config --sharedir) echo "Package: postgresql-${VERSION}-vchord Version: ${SEMVER}-1 Section: database @@ -255,11 +239,9 @@ jobs: Homepage: https://vectorchord.ai/ License: AGPL-3 or Elastic-2" \ > ./build/deb/DEBIAN/control - (cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums + (cd ./build/deb && find usr -type f -print0 | xargs -0 md5sum) > ./build/deb/DEBIAN/md5sums dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb - ls ./build - - name: Upload Artifacts uses: actions/upload-artifact@v4 with: @@ -292,6 +274,8 @@ jobs: run: | rustup set profile + echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV + HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install postgresql@${{ matrix.version }} HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install pgvector brew services start postgresql@${{ matrix.version }} @@ -304,8 +288,7 @@ jobs: $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.so"' fi brew services stop postgresql@${{ matrix.version }} - - curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.14.1/cargo-pgrx-v0.14.1-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./cargo-pgrx | tee /usr/local/bin/cargo-pgrx > /dev/null && chmod 755 /usr/local/bin/cargo-pgrx + echo PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest @@ -317,11 +300,12 @@ jobs: - name: Clippy run: | - env CC=$(brew --prefix llvm@18)/bin/clang PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config cargo clippy -p vchord --features pg${{ matrix.version }} -- --no-deps + cargo clippy -p vchord --features pg${{ matrix.version }} -- --no-deps - name: Install run: | - env CC=$(brew --prefix llvm@18)/bin/clang PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config cargo pgrx install -p vchord --features pg${{ matrix.version }} --release --sudo --pg-config=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config + cargo make package -o ./build/raw + cargo make install --sudo -i ./build/raw - name: Service run: | @@ -336,3 +320,21 @@ jobs: - name: Sqllogictest (PostgreSQL 17) if: matrix.version == '17' run: sqllogictest --db $USER --user $USER './tests/pg17/*.slt' + + - name: Package + env: + SEMVER: "0.0.0" + VERSION: ${{ matrix.version }} + ARCH: ${{ matrix.arch }} + run: | + cargo make package -o ./build/raw + (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-apple-darwin.zip .) + + - name: Upload Artifacts + uses: actions/upload-artifact@v4 + with: + name: artifacts-psql_macos-${{ matrix.version }}-${{ matrix.arch }} + path: | + ./build/postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin.zip + compression-level: 9 + retention-days: 14 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index f2002726..121bd63a 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -62,11 +62,12 @@ jobs: sudo apt-get install -y postgresql-common sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} + echo PGRX_PG_CONFIG_PATH=pg_config >> $GITHUB_ENV sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector - curl -fsSL https://github.com/tensorchord/pgrx/releases/download/v0.14.1/cargo-pgrx-v0.14.1-$(uname -m)-unknown-linux-gnu.tar.gz | tar -xOzf - ./cargo-pgrx | install -m 755 /dev/stdin /usr/local/bin/cargo-pgrx - cargo pgrx init --pg${{ matrix.version }}=$(which pg_config) + mkdir -p ~/.pgrx + touch ~/.pgrx/config.toml - name: Checkout uses: actions/checkout@v4 @@ -84,28 +85,13 @@ jobs: ARCH: ${{ matrix.arch }} PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} run: | - cargo build --lib --features pg$VERSION --release - - mkdir -p ./build/zip - cp -a ./sql/upgrade/. ./build/zip/ - cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql - cp ./vchord.control ./build/zip/vchord.control - cp ./target/release/libvchord.so ./build/zip/vchord.so - zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* + cargo make package -o ./build/raw + (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) mkdir -p ./build/deb mkdir -p ./build/deb/DEBIAN - mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ - mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ - for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file - done - for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file - done - for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file - done + mkdir -p ./build/deb$(pg_config --pkglibdir) && cp -r ./build/raw/pkglibdir/. ./build/deb$(pg_config --pkglibdir) + mkdir -p ./build/deb$(pg_config --sharedir) && cp -r ./build/raw/sharedir/. ./build/deb$(pg_config --sharedir) echo "Package: postgresql-${VERSION}-vchord Version: ${SEMVER}-1 Section: database @@ -116,11 +102,9 @@ jobs: Homepage: https://vectorchord.ai/ License: AGPL-3 or Elastic-2" \ > ./build/deb/DEBIAN/control - (cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums + (cd ./build/deb && find usr -type f -print0 | xargs -0 md5sum) > ./build/deb/DEBIAN/md5sums dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb - ls ./build - - name: Upload Artifacts env: GH_TOKEN: ${{ github.token }} diff --git a/Cargo.lock b/Cargo.lock index c53cf491..16a63c54 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627" + [[package]] name = "aho-corasick" version = "1.1.3" @@ -50,12 +56,56 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "anstream" +version = "0.6.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8acc5369981196006228e28809f761875c0327210a891e941f4c683b3a99529b" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + [[package]] name = "anstyle" version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9" +[[package]] +name = "anstyle-parse" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b2d16507662817a6a20a9ea92df6652ee4f94f914589377d69f3b21bc5798a9" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79947af37f4177cfead1110013d678905c37501914fba0efea834c3fe9a8d60c" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6680de5231bd6ee4c6191b8a1325daa282b415391ec9d3a37bd34f2060dc73fa" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", +] + [[package]] name = "anyhow" version = "1.0.98" @@ -170,6 +220,52 @@ dependencies = [ "libloading", ] +[[package]] +name = "clap" +version = "4.5.38" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed93b9805f8ba930df42c2590f05453d5ec36cbb85d018868a5b24d31f6ac000" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.5.38" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "379026ff283facf611b0ea629334361c4211d1b12ee01024eec1591133b04120" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.5.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09176aae279615badda0765c0c0b3f6ed53f4709118af73cf4655d85d1530cd7" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6" + +[[package]] +name = "colorchoice" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b63caa9aa9397e2d9480a9b13673856c78d8ac123288526c37d7839f2a86990" + [[package]] name = "convert_case" version = "0.8.0" @@ -179,6 +275,15 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "crc32fast" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3" +dependencies = [ + "cfg-if", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -286,12 +391,28 @@ dependencies = [ "syn", ] +[[package]] +name = "env_home" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe" + [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.3.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cea14ef9355e3beab063703aa9dab15afd25f0667c341310c1e5274bb1d0da18" +dependencies = [ + "libc", + "windows-sys", +] + [[package]] name = "eyre" version = "0.6.12" @@ -302,12 +423,28 @@ dependencies = [ "once_cell", ] +[[package]] +name = "fastrand" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" + [[package]] name = "fixedbitset" version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" +[[package]] +name = "flate2" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ced92e76e966ca2fd84c8f7aa01a4aea65b0eb6648d72f7c8f3e2764a67fece" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + [[package]] name = "fnv" version = "1.0.7" @@ -400,6 +537,12 @@ dependencies = [ "stable_deref_trait", ] +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hermit-abi" version = "0.5.1" @@ -561,6 +704,12 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7655c9839580ee829dfacba1d1278c2b7883e50a277ff7541299489d6bdfdc45" +[[package]] +name = "is_terminal_polyfill" +version = "1.70.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" + [[package]] name = "itertools" version = "0.13.0" @@ -622,12 +771,29 @@ dependencies = [ "windows-targets 0.53.0", ] +[[package]] +name = "linux-raw-sys" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd945864f07fe9f5371a27ad7b52a172b4b499999f1d97574c9fa68373937e12" + [[package]] name = "litemap" version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956" +[[package]] +name = "make" +version = "0.0.0" +dependencies = [ + "anyhow", + "clap", + "object", + "tempfile", + "which", +] + [[package]] name = "memchr" version = "2.7.4" @@ -640,6 +806,15 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3be647b768db090acb35d5ec5db2b0e1f1de11133ca123b9eacf5137868f892a" +dependencies = [ + "adler2", +] + [[package]] name = "nom" version = "7.1.3" @@ -650,12 +825,33 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "flate2", + "hashbrown", + "indexmap", + "memchr", + "ruzstd", + "wasmparser", +] + [[package]] name = "once_cell" version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "once_cell_polyfill" +version = "1.70.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4895175b425cb1f87721b59f0f286c2092bd4af812243672510e1ac53e2e0ad" + [[package]] name = "owo-colors" version = "4.2.0" @@ -1008,6 +1204,28 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c71e83d6afe7ff64890ec6b71d6a69bb8a610ab78ce364b3352876bb4c801266" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "ruzstd" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fad02996bfc73da3e301efe90b1837be9ed8f4a462b6ed410aa35d00381de89f" +dependencies = [ + "twox-hash", +] + [[package]] name = "ryu" version = "1.0.20" @@ -1138,6 +1356,12 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "strsim" version = "0.11.1" @@ -1191,6 +1415,19 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" +[[package]] +name = "tempfile" +version = "3.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8a64e3985349f2441a1a9ef0b853f869006c3855f2cda6862a94d26ebb9d6a1" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys", +] + [[package]] name = "thiserror" version = "2.0.12" @@ -1262,6 +1499,16 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bfb942dfe1d8e29a7ee7fcbde5bd2b9a25fb89aa70caea2eba3bee836ff41076" +[[package]] +name = "twox-hash" +version = "1.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97fee6b57c6a41524a810daee9286c02d7752c4253064d0b05472833a438f675" +dependencies = [ + "cfg-if", + "static_assertions", +] + [[package]] name = "unescape" version = "0.1.0" @@ -1303,6 +1550,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.16.0" @@ -1394,6 +1647,27 @@ dependencies = [ "wit-bindgen-rt", ] +[[package]] +name = "wasmparser" +version = "0.222.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa210fd1788e6b37a1d1930f3389c48e1d6ebd1a013d34fa4b7f9e3e3bf03146" +dependencies = [ + "bitflags", +] + +[[package]] +name = "which" +version = "7.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d643ce3fd3e5b54854602a080f34fb10ab75e0b813ee32d00ca2b44fa74762" +dependencies = [ + "either", + "env_home", + "rustix", + "winsafe", +] + [[package]] name = "winapi-util" version = "0.1.9" @@ -1549,6 +1823,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "winsafe" +version = "0.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" + [[package]] name = "wit-bindgen-rt" version = "0.39.0" diff --git a/build.rs b/build.rs index 6c379dc3..34c8efec 100644 --- a/build.rs +++ b/build.rs @@ -12,12 +12,31 @@ // // Copyright (c) 2025 TensorChord Inc. -use std::env::var; +use std::collections::HashMap; +use std::env::{var, var_os}; use std::error::Error; +use std::process::Command; fn main() -> Result<(), Box> { if var("CARGO_CFG_TARGET_OS")? == "macos" { - println!("cargo::rustc-link-arg-cdylib=-Wl,-undefined,dynamic_lookup"); + if let Some(path) = var_os("PGRX_PG_CONFIG_PATH") { + let map = { + let command_output = Command::new(&path).output()?; + let command_stdout = String::from_utf8(command_output.stdout)?; + let mut map = HashMap::new(); + for line in command_stdout.lines() { + if let Some((key, value)) = line.split_once(" = ") { + map.insert(key.to_string(), value.to_string()); + eprintln!("Config `{key}`: {value}"); + } + } + map + }; + let bindir = &map["BINDIR"]; + println!("cargo::rustc-link-arg-cdylib=-Wl,-bundle,-bundle_loader,{bindir}/postgres",); + } else { + println!("cargo::rustc-link-arg-cdylib=-Wl,-undefined,dynamic_lookup"); + } } Ok(()) } diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml new file mode 100644 index 00000000..2167f140 --- /dev/null +++ b/crates/make/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "make" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +anyhow = "1.0.98" +clap = { version = "4.5.38", features = ["derive"] } +which = "7.0.3" +tempfile = "3.20.0" +object = { version = "0.36.7", features = ["all"] } diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs new file mode 100644 index 00000000..a86aedac --- /dev/null +++ b/crates/make/src/main.rs @@ -0,0 +1,322 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use anyhow::{Context, Result, bail}; +use clap::{Args, Parser, Subcommand}; +use object::Object; +use std::collections::HashMap; +use std::env::consts::{DLL_PREFIX, DLL_SUFFIX, OS}; +use std::env::var_os; +use std::fs::read_dir; +use std::io::Write; +use std::path::PathBuf; +use std::process::{Command, Stdio}; + +#[derive(Parser)] +struct Cli { + #[command(subcommand)] + command: Commands, +} + +#[derive(Subcommand)] +enum Commands { + Package(PackageArgs), + Install(InstallArgs), +} + +#[derive(Args)] +struct PackageArgs { + #[arg(short, long)] + output: String, +} + +#[derive(Args)] +struct InstallArgs { + #[arg(short, long)] + input: Option, + #[arg(long, default_value = "false")] + sudo: bool, +} + +fn main() -> Result<()> { + let cli = Cli::parse(); + if !std::fs::exists("vchord.control")? { + bail!("The script must be run from the VectorChord source directory.") + } + let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { + eprintln!("Environment variable `PGRX_PG_CONFIG_PATH`: {value:#?}"); + PathBuf::from(value) + } else { + if let Ok(path) = which::which("pg_config") { + eprintln!("Found executable `pg_config` in PATH: {path:#?}"); + eprintln!("Hint: set environment variable `PGRX_PG_CONFIG_PATH` to {path:#?}"); + } + bail!("Environment variable `PGRX_PG_CONFIG_PATH` is not set.") + }; + let map = { + let mut command = Command::new(&path); + command.stderr(Stdio::inherit()); + let command_output = command.output()?; + let command_stdout = String::from_utf8(command_output.stdout)?; + let mut map = HashMap::new(); + for line in command_stdout.lines() { + if let Some((key, value)) = line.split_once(" = ") { + map.insert(key.to_string(), value.to_string()); + eprintln!("Config `{key}`: {value}"); + } + } + map + }; + let fork = { + let version = map["VERSION"].clone(); + let fork = if let Some(prefix_stripped) = version.strip_prefix("PostgreSQL ") { + if let Some((stripped, _)) = prefix_stripped.split_once(|c: char| !c.is_ascii_digit()) { + format!("pg{stripped}",) + } else { + format!("pg{prefix_stripped}",) + } + } else { + bail!("PostgreSQL version is invalid.") + }; + eprintln!("Fork: {fork}"); + fork + }; + let version = 'version: { + for line in std::fs::read_to_string("./vchord.control")?.lines() { + if let Some(prefix_stripped) = line.strip_prefix("default_version = '") { + if let Some(stripped) = prefix_stripped.strip_suffix("'") { + eprintln!("VectorChord version: {stripped}"); + break 'version stripped.to_string(); + } + } + } + bail!("VectorChord version is not defined.") + }; + let build = || { + let mut command = Command::new("cargo"); + command + .args(["build", "--release"]) + .args(["-p", "vchord", "--lib"]) + .args(["--features", fork.as_str()]) + .env("PGRX_PG_CONFIG_PATH", &path) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + let debug = format!("{command:?}"); + let status = command.spawn()?.wait()?; + if !status.success() { + bail!("Cargo build failed: {debug}"); + } + Ok(()) + }; + let schema = || { + let object = std::fs::read(format!("./target/release/{DLL_PREFIX}vchord{DLL_SUFFIX}"))?; + let object = object::File::parse(object.as_slice())?; + let exports = object + .exports()? + .into_iter() + .flat_map(|x| str::from_utf8(x.name())); + let exports = if matches!(object.format(), object::BinaryFormat::MachO) { + exports + .flat_map(|x| x.strip_prefix("_")) + .filter(|x| x.starts_with("__pgrx_internals")) + .collect::>() + } else { + exports + .filter(|x| x.starts_with("__pgrx_internals")) + .collect::>() + }; + let pushes = exports + .into_iter() + .map(|x| { + format!( + r#" + entities.push(unsafe {{ + unsafe extern "Rust" {{ + fn {x}() -> ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity; + }} + {x}() + }}); + "# + ) + }) + .collect::>() + .join("\n"); + let code = format!( + r#" + pub fn main() {{ + extern crate vchord as _; + + let mut entities = Vec::new(); + let control_file = std::fs::read_to_string("./vchord.control").unwrap(); + let control_file = ::pgrx::pgrx_sql_entity_graph::ControlFile::try_from(control_file.as_str()).expect(".control file should properly formatted"); + let control_file_entity = ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity::ExtensionRoot(control_file); + + entities.push(control_file_entity); + + {pushes} + + let pgrx_sql = ::pgrx::pgrx_sql_entity_graph::PgrxSql::build( + entities.into_iter(), + "vchord".to_string(), + false, + ) + .expect("SQL generation error"); + + pgrx_sql + .write(&mut std::io::stdout()) + .expect("Could not write SQL to stdout"); + }} + "# + ); + let mut file = tempfile::NamedTempFile::new()?; + file.write_all(code.as_bytes())?; + let pgrx_embed_path = file.into_temp_path(); + let mut command = Command::new("cargo"); + command + .args(["rustc"]) + .args(["-p", "vchord", "--bin", "pgrx_embed_vchord"]) + .args(["--features", fork.as_str()]) + .args(["--", "--cfg", "pgrx_embed"]) + .env("PGRX_EMBED", &pgrx_embed_path) + .env("PGRX_PG_CONFIG_PATH", &path) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + let debug = format!("{command:?}"); + let status = command.spawn()?.wait()?; + if !status.success() { + bail!("Cargo build failed: {debug}"); + } + let mut command = Command::new("./target/debug/pgrx_embed_vchord"); + command.stderr(Stdio::inherit()); + let command_output = command.output()?; + let command_stdout = String::from_utf8(command_output.stdout)?.replace("\t", " "); + Ok(command_stdout) + }; + let uilts_install = |sudo: bool, permission: &str, src: &str, dst: &str| { + let mut command; + if sudo { + command = Command::new("sudo"); + command.arg("install"); + } else { + command = Command::new("install"); + }; + command.args(["-m", permission, src, dst]); + let debug = format!("{command:?}"); + let status = command.spawn()?.wait()?; + if !status.success() { + bail!("Command execution failed: {debug}"); + } + Ok(()) + }; + let uilts_cp = |sudo: bool, args: &[&str]| { + let mut command; + if sudo { + command = Command::new("sudo"); + command.arg("cp"); + } else { + command = Command::new("cp"); + }; + command.args(args); + let debug = format!("{command:?}"); + let status = command.spawn()?.wait()?; + if !status.success() { + bail!("Command execution failed: {debug}"); + } + Ok(()) + }; + let dll_suffix = if OS == "macos" && matches!(fork.as_str(), "pg13" | "pg14" | "pg15") { + ".so" + } else { + DLL_SUFFIX + }; + let install = |sudo, pkglibdir, sharedir_extension| -> Result<()> { + uilts_install( + sudo, + "755", + &format!("./target/release/{DLL_PREFIX}vchord{DLL_SUFFIX}"), + &format!("{pkglibdir}/vchord{dll_suffix}"), + )?; + uilts_install( + sudo, + "644", + "./vchord.control", + &format!("{sharedir_extension}/vchord.control"), + )?; + if version != "0.0.0" { + for maybe_entry in read_dir("./sql/upgrade")? { + let path = maybe_entry?.path(); + let name = path.file_name().context("broken assets")?; + uilts_install( + sudo, + "644", + &format!("{}", path.display()), + &format!("{sharedir_extension}/{}", name.display()), + )?; + } + uilts_install( + sudo, + "644", + &format!("./sql/install/vchord--{version}.sql"), + &format!("{sharedir_extension}/vchord--{version}.sql"), + )?; + } else { + let contents = schema()?; + let mut file = tempfile::NamedTempFile::new()?; + file.write_all(contents.as_bytes())?; + let path = file.into_temp_path(); + uilts_install( + sudo, + "644", + &format!("{}", path.display()), + &format!("{sharedir_extension}/vchord--0.0.0.sql"), + )?; + } + Ok(()) + }; + match cli.command { + Commands::Package(PackageArgs { output }) => { + build()?; + let pkglibdir = format!("{output}/pkglibdir"); + let sharedir = format!("{output}/sharedir"); + let sharedir_extension = format!("{sharedir}/extension"); + if std::fs::exists(&output)? { + std::fs::remove_dir_all(&output)?; + } + std::fs::create_dir_all(&output)?; + std::fs::create_dir_all(&pkglibdir)?; + std::fs::create_dir_all(&sharedir)?; + std::fs::create_dir_all(&sharedir_extension)?; + install(false, pkglibdir, sharedir_extension)?; + } + Commands::Install(InstallArgs { sudo, input }) => { + if let Some(input) = input { + uilts_cp( + sudo, + &["-r", &format!("{input}/pkglibdir/."), &map["PKGLIBDIR"]], + )?; + uilts_cp( + sudo, + &["-r", &format!("{input}/sharedir/."), &map["SHAREDIR"]], + )?; + } else { + build()?; + let pkglibdir = map["PKGLIBDIR"].clone(); + let sharedir = map["SHAREDIR"].clone(); + let sharedir_extension = format!("{sharedir}/extension"); + install(sudo, pkglibdir, sharedir_extension)?; + } + } + } + Ok(()) +} From 14108d490296cda74bfd0d619d0ea19958b8424d Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 28 May 2025 15:53:23 +0800 Subject: [PATCH 165/324] refactor: use a buffer to reuse recently-freed pages in maintaince (#264) Signed-off-by: usamoi --- build.rs | 6 +- crates/algorithm/src/build.rs | 5 +- crates/algorithm/src/freepages.rs | 70 ++++++++---------------- crates/algorithm/src/insert.rs | 10 +++- crates/algorithm/src/maintain.rs | 90 +++++++++++++++++------------- crates/algorithm/src/tape.rs | 17 +++++- crates/algorithm/src/tuples.rs | 91 ++++++++----------------------- crates/algorithm/src/vectors.rs | 2 +- src/index/algorithm.rs | 21 ++++++- src/index/am/am_build.rs | 6 +- src/index/am/mod.rs | 2 +- 11 files changed, 152 insertions(+), 168 deletions(-) diff --git a/build.rs b/build.rs index 34c8efec..fdf80e9d 100644 --- a/build.rs +++ b/build.rs @@ -15,13 +15,15 @@ use std::collections::HashMap; use std::env::{var, var_os}; use std::error::Error; -use std::process::Command; +use std::process::{Command, Stdio}; fn main() -> Result<(), Box> { if var("CARGO_CFG_TARGET_OS")? == "macos" { if let Some(path) = var_os("PGRX_PG_CONFIG_PATH") { let map = { - let command_output = Command::new(&path).output()?; + let mut command = Command::new(&path); + command.stderr(Stdio::inherit()); + let command_output = command.output()?; let command_stdout = String::from_utf8(command_output.stdout)?; let mut map = HashMap::new(); for line in command_stdout.lines() { diff --git a/crates/algorithm/src/build.rs b/crates/algorithm/src/build.rs index 5fe3a531..04fc2cc7 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/algorithm/src/build.rs @@ -30,7 +30,8 @@ pub fn build( let is_residual = vchordrq_options.residual_quantization; let mut meta = TapeWriter::<_, MetaTuple>::create(index, false); assert_eq!(meta.first(), 0); - let freepage = TapeWriter::<_, FreepageTuple>::create(index, false); + let mut freepages = TapeWriter::<_, FreepagesTuple>::create(index, false); + freepages.push(FreepagesTuple {}); let mut vectors = TapeWriter::<_, VectorTuple>::create(index, true); let mut pointer_of_centroids = Vec::, u16)>>::new(); for i in 0..structures.len() { @@ -127,7 +128,7 @@ pub fn build( first: pointer_of_firsts .last() .expect("internal error: empty structure")[0], - freepage_first: freepage.first(), + freepages_first: freepages.first(), cells: structures.iter().map(|s| s.len() as _).collect(), }); } diff --git a/crates/algorithm/src/freepages.rs b/crates/algorithm/src/freepages.rs index fb933be1..1d0c96ef 100644 --- a/crates/algorithm/src/freepages.rs +++ b/crates/algorithm/src/freepages.rs @@ -14,56 +14,30 @@ use crate::tuples::*; use crate::*; -use std::cmp::Reverse; -pub fn mark(index: &impl RelationWrite, freepage_first: u32, values: &[u32]) { - let mut values = { - let mut values = values.to_vec(); - values.sort_by_key(|x| Reverse(*x)); - values.dedup(); - values - }; - let (mut current, mut offset) = (freepage_first, 0_u32); - loop { - let mut freespace_guard = index.write(current, false); - if freespace_guard.len() == 0 { - freespace_guard - .alloc(&FreepageTuple::serialize(&FreepageTuple {})) - .expect("implementation: a clear page cannot accommodate a single tuple"); - } - let freespace_bytes = freespace_guard.get_mut(1).expect("data corruption"); - let mut freespace_tuple = FreepageTuple::deserialize_mut(freespace_bytes); - while let Some(target) = values.pop_if(|&mut x| x < offset + 32768) { - freespace_tuple.mark((target - offset) as usize); - } - if values.is_empty() { - return; - } - if freespace_guard.get_opaque().next == u32::MAX { - let extend = index.extend(false); - freespace_guard.get_opaque_mut().next = extend.id(); - } - (current, offset) = (freespace_guard.get_opaque().next, offset + 32768); +pub fn alloc(index: &R, freepages_first: u32) -> Option> { + let mut freepages_guard = index.write(freepages_first, false); + let freepages_bytes = freepages_guard.get_mut(1).expect("data corruption"); + let mut freepages_tuple = FreepagesTuple::deserialize_mut(freepages_bytes); + let id = *freepages_tuple.first(); + if id != u32::MAX { + let mut guard = index.write(id, false); + *freepages_tuple.first() = guard.get_opaque_mut().next; + drop(freepages_guard); // write log of freespaces_guard + Some(guard) + } else { + None } } -pub fn fetch(index: &impl RelationWrite, freepage_first: u32) -> Option { - let (mut current, mut offset) = (freepage_first, 0_u32); - loop { - let mut freespace_guard = index.write(current, false); - if freespace_guard.len() == 0 { - freespace_guard - .alloc(&FreepageTuple::serialize(&FreepageTuple {})) - .expect("implementation: a clear page cannot accommodate a single tuple"); - } - let freespace_bytes = freespace_guard.get_mut(1).expect("data corruption"); - let mut freespace_tuple = FreepageTuple::deserialize_mut(freespace_bytes); - if let Some(x) = freespace_tuple.fetch() { - return Some(x as u32 + offset); - } - if freespace_guard.get_opaque().next == u32::MAX { - return None; - } - (current, offset) = (freespace_guard.get_opaque().next, offset + 32768); - } +// the page must be inaccessible in the graph +pub fn free(index: &R, freepages_first: u32, id: u32) { + let mut guard = index.write(id, false); + let mut freepages_guard = index.write(freepages_first, false); + let freepages_bytes = freepages_guard.get_mut(1).expect("data corruption"); + let mut freepages_tuple = FreepagesTuple::deserialize_mut(freepages_bytes); + guard.get_opaque_mut().next = *freepages_tuple.first(); + drop(guard); // write log of guard + *freepages_tuple.first() = id; + drop(freepages_guard); // write log of freepages_guard } diff --git a/crates/algorithm/src/insert.rs b/crates/algorithm/src/insert.rs index 66915daa..75e3911d 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/algorithm/src/insert.rs @@ -56,6 +56,7 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( key: (Vec, u16), bump: &'b impl Bump, mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, + skip_freespaces: bool, ) { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); @@ -63,6 +64,7 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); let height_of_root = meta_tuple.height_of_root(); + let freepages_first = meta_tuple.freepages_first(); assert_eq!(dims, vector.dims(), "unmatched dimensions"); let epsilon = 1.9; @@ -167,5 +169,11 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( elements: rabitq::packing::pack_to_u64(&code.1), }); - tape::append(index, jump_tuple.appendable_first(), &serialized, false); + tape::append( + index, + jump_tuple.appendable_first(), + &serialized, + false, + (!skip_freespaces).then_some(freepages_first), + ); } diff --git a/crates/algorithm/src/maintain.rs b/crates/algorithm/src/maintain.rs index f09879df..1b200022 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/algorithm/src/maintain.rs @@ -19,18 +19,25 @@ use crate::tape_writer::{DirectoryTapeWriter, FrozenTapeWriter}; use crate::tuples::*; use crate::*; use rabitq::packing::unpack; +use std::cell::RefCell; + +pub struct Maintain { + pub number_of_formerly_allocated_pages: usize, + pub number_of_freshly_allocated_pages: usize, + pub number_of_freed_pages: usize, +} pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( index: &'r R, mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, check: impl Fn(), -) { +) -> Maintain { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let height_of_root = meta_tuple.height_of_root(); - let freepage_first = meta_tuple.freepage_first(); + let freepages_first = meta_tuple.freepages_first(); type State = Vec; let mut state: State = vec![meta_tuple.first()]; @@ -53,19 +60,40 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( state = step(state); } + struct Buffers { + pages: Vec, + number_of_formerly_allocated_pages: usize, + number_of_freshly_allocated_pages: usize, + } + + let buffers = RefCell::new(Buffers { + pages: Vec::new(), + number_of_formerly_allocated_pages: 0, + number_of_freshly_allocated_pages: 0, + }); + for first in state { let mut jump_guard = index.write(first, false); let jump_bytes = jump_guard.get_mut(1).expect("data corruption"); let mut jump_tuple = JumpTuple::deserialize_mut(jump_bytes); - let frozen_tape_hooked_index = RelationHooked(index, { - id_3(move |index: &R, tracking_freespace: bool| { + let hooked_index = RelationHooked(index, { + id_3(|index: &R, tracking_freespace: bool| { if !tracking_freespace { - if let Some(id) = freepages::fetch(index, freepage_first) { - let mut write = index.write(id, false); - write.clear(); - write + let mut buffers = buffers.borrow_mut(); + if let Some(id) = buffers.pages.pop() { + drop(buffers); + let mut guard = index.write(id, false); + guard.clear(); + guard + } else if let Some(mut guard) = freepages::alloc(index, freepages_first) { + buffers.number_of_formerly_allocated_pages += 1; + drop(buffers); + guard.clear(); + guard } else { + buffers.number_of_freshly_allocated_pages += 1; + drop(buffers); index.extend(false) } } else { @@ -74,11 +102,7 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( }) }); - let mut tape = FrozenTapeWriter::create( - &frozen_tape_hooked_index, - O::Vector::count(dims as _), - false, - ); + let mut tape = FrozenTapeWriter::create(&hooked_index, O::Vector::count(dims as _), false); let mut trace_directory = Vec::new(); let mut trace_forzen = Vec::new(); @@ -162,22 +186,6 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( let (frozen_tape, branches) = tape.into_inner(); - let hooked_index = RelationHooked( - index, - id_3(move |index: &R, tracking_freespace: bool| { - if !tracking_freespace { - if let Some(id) = freepages::fetch(index, freepage_first) { - let mut write = index.write(id, false); - write.clear(); - write - } else { - index.extend(false) - } - } else { - index.extend(true) - } - }), - ); let mut appendable_tape = TapeWriter::create(&hooked_index, false); for branch in branches { @@ -203,7 +211,7 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( .map(|guard| guard.id()) .collect::>(); - let mut directory_tape = DirectoryTapeWriter::create(index, false); + let mut directory_tape = DirectoryTapeWriter::create(&hooked_index, false); directory_tape.push(directory.as_slice()); let directory_tape = directory_tape.into_inner(); @@ -213,15 +221,21 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( drop(jump_guard); - let trace = { - let mut v = Vec::new(); - v.extend_from_slice(&trace_directory); - v.extend_from_slice(&trace_forzen); - v.extend_from_slice(&trace_appendable); - v - }; + let mut buffers = buffers.borrow_mut(); + buffers.pages.extend_from_slice(&trace_directory); + buffers.pages.extend_from_slice(&trace_forzen); + buffers.pages.extend_from_slice(&trace_appendable); + } + + let buffers = RefCell::into_inner(buffers); + for id in buffers.pages.iter().copied() { + freepages::free(index, freepages_first, id); + } - freepages::mark(index, freepage_first, trace.as_slice()); + Maintain { + number_of_formerly_allocated_pages: buffers.number_of_formerly_allocated_pages, + number_of_freshly_allocated_pages: buffers.number_of_freshly_allocated_pages, + number_of_freed_pages: buffers.pages.len(), } } diff --git a/crates/algorithm/src/tape.rs b/crates/algorithm/src/tape.rs index f1c7e496..cf2b5964 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/algorithm/src/tape.rs @@ -14,7 +14,7 @@ use crate::operator::Accessor1; use crate::tuples::*; -use crate::{Page, PageGuard, PrefetcherSequenceFamily, RelationRead, RelationWrite}; +use crate::{Page, PageGuard, PrefetcherSequenceFamily, RelationRead, RelationWrite, freepages}; use std::marker::PhantomData; use std::num::NonZero; @@ -301,7 +301,9 @@ pub fn append( first: u32, bytes: &[u8], tracking_freespace: bool, + freepages_first: Option, ) -> (u32, u16) { + assert!(!tracking_freespace || freepages_first.is_none()); assert!(first != u32::MAX); let mut current = first; loop { @@ -313,7 +315,18 @@ pub fn append( if let Some(i) = write.alloc(bytes) { return (current, i); } - let mut extend = index.extend(tracking_freespace); + let mut extend = { + if let Some(freepages_first) = freepages_first { + if let Some(mut guard) = freepages::alloc(index, freepages_first) { + guard.clear(); + guard + } else { + index.extend(tracking_freespace) + } + } else { + index.extend(tracking_freespace) + } + }; write.get_opaque_mut().next = extend.id(); drop(write); let fresh = extend.id(); diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 1df5f77f..9ebf5571 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -15,12 +15,12 @@ use crate::operator::Vector; use std::marker::PhantomData; use std::num::NonZero; -use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 9; +const VERSION: u64 = 10; pub trait Tuple: 'static { fn serialize(&self) -> Vec; @@ -48,7 +48,7 @@ struct MetaTupleHeader { cells_e: u16, _padding_0: [ZeroU8; 2], vectors_first: u32, - freepage_first: u32, + freepages_first: u32, _padding_1: [ZeroU8; 2], // tree centroid_prefetch_s: u16, @@ -65,7 +65,7 @@ pub struct MetaTuple { pub rerank_in_heap: bool, pub cells: Vec, pub vectors_first: u32, - pub freepage_first: u32, + pub freepages_first: u32, pub centroid_prefetch: Vec, pub centroid_head: u16, pub centroid_norm: f32, @@ -84,7 +84,7 @@ impl Tuple for MetaTuple { rerank_in_heap, cells, vectors_first, - freepage_first, + freepages_first, centroid_prefetch, centroid_head, centroid_norm, @@ -117,7 +117,7 @@ impl Tuple for MetaTuple { cells_s, cells_e, vectors_first: *vectors_first, - freepage_first: *freepage_first, + freepages_first: *freepages_first, centroid_prefetch_s, centroid_prefetch_e, centroid_head: *centroid_head, @@ -185,8 +185,8 @@ impl<'a> MetaTupleReader<'a> { pub fn vectors_first(self) -> u32 { self.header.vectors_first } - pub fn freepage_first(self) -> u32 { - self.header.freepage_first + pub fn freepages_first(self) -> u32 { + self.header.freepages_first } pub fn centroid_prefetch(self) -> &'a [u32] { self.centroid_prefetch @@ -204,24 +204,18 @@ impl<'a> MetaTupleReader<'a> { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] -struct FreepageTupleHeader { - level_0: [u32; 1 << 10], - level_1: [u32; 1 << 5], - level_2: [u32; 1 << 0], +struct FreepagesTupleHeader { + first: u32, _padding_0: [ZeroU8; 4], } -const _: () = assert!(size_of::() == 4232); - #[derive(Debug, Clone, PartialEq)] -pub struct FreepageTuple {} +pub struct FreepagesTuple {} -impl Tuple for FreepageTuple { +impl Tuple for FreepagesTuple { fn serialize(&self) -> Vec { - FreepageTupleHeader { - level_0: FromZeros::new_zeroed(), - level_1: FromZeros::new_zeroed(), - level_2: FromZeros::new_zeroed(), + FreepagesTupleHeader { + first: u32::MAX, _padding_0: Default::default(), } .as_bytes() @@ -229,52 +223,23 @@ impl Tuple for FreepageTuple { } } -impl WithWriter for FreepageTuple { - type Writer<'a> = FreepageTupleWriter<'a>; +impl WithWriter for FreepagesTuple { + type Writer<'a> = FreepagesTupleWriter<'a>; - fn deserialize_mut(source: &mut [u8]) -> FreepageTupleWriter<'_> { + fn deserialize_mut(source: &mut [u8]) -> FreepagesTupleWriter<'_> { let mut checker = MutChecker::new(source); let header = checker.prefix(0_u16); - FreepageTupleWriter { header } + FreepagesTupleWriter { header } } } -pub struct FreepageTupleWriter<'a> { - header: &'a mut FreepageTupleHeader, +pub struct FreepagesTupleWriter<'a> { + header: &'a mut FreepagesTupleHeader, } -impl FreepageTupleWriter<'_> { - pub fn mark(&mut self, i: usize) { - assert!(i < 32768, "out of bound: {i}"); - set(&mut self.header.level_0[i >> 5], (i >> 0) % 32, true); - set(&mut self.header.level_1[i >> 10], (i >> 5) % 32, true); - set(&mut self.header.level_2[i >> 15], (i >> 10) % 32, true); - } - fn find(&self) -> Option { - let i_3 = 0_usize; - let i_2 = self.header.level_2[i_3 << 0].trailing_zeros() as usize; - if i_2 == 32 { - return None; - } - let i_1 = self.header.level_1[i_3 << 5 | i_2 << 0].trailing_zeros() as usize; - if i_1 == 32 { - panic!("deserialization: bad bytes"); - } - let i_0 = self.header.level_0[i_3 << 10 | i_2 << 5 | i_1 << 0].trailing_zeros() as usize; - if i_0 == 32 { - panic!("deserialization: bad bytes"); - } - Some(i_3 << 15 | i_2 << 10 | i_1 << 5 | i_0 << 0) - } - pub fn fetch(&mut self) -> Option { - let i = self.find()?; - let x = false; - set(&mut self.header.level_0[i >> 5], (i >> 0) % 32, x); - let x = self.header.level_0[i >> 5] != 0; - set(&mut self.header.level_1[i >> 10], (i >> 5) % 32, x); - let x = self.header.level_1[i >> 10] != 0; - set(&mut self.header.level_2[i >> 15], (i >> 10) % 32, x); - Some(i) +impl FreepagesTupleWriter<'_> { + pub fn first(&mut self) -> &mut u32 { + &mut self.header.first } } @@ -1460,16 +1425,6 @@ impl<'a> MutChecker<'a> { } } -#[inline(always)] -fn set(a: &mut u32, i: usize, x: bool) { - assert!(i < 32, "out of bound: {i}"); - if x { - *a |= 1 << i; - } else { - *a &= !(1 << i); - } -} - #[test] fn aliasing_test() { #[repr(C, align(8))] diff --git a/crates/algorithm/src/vectors.rs b/crates/algorithm/src/vectors.rs index 0fcd9534..1fa2fb0d 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/algorithm/src/vectors.rs @@ -91,7 +91,7 @@ pub fn append( .expect("implementation: a free page cannot accommodate a single tuple"); return (write.id(), i); } - tape::append(index, first, bytes, true) + tape::append(index, first, bytes, true, None) } let (slices, metadata) = O::Vector::split(vector); let mut chain = Ok(metadata); diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs index e1b89058..a6f98a2c 100644 --- a/src/index/algorithm.rs +++ b/src/index/algorithm.rs @@ -270,7 +270,7 @@ pub fn bulkdelete( pub fn maintain(opfamily: Opfamily, index: &(impl RelationRead + RelationWrite), check: impl Fn()) { let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { + let maintain = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2) => { algorithm::maintain::<_, Op, L2>>(index, make_h0_plain_prefetcher, check) } @@ -291,7 +291,19 @@ pub fn maintain(opfamily: Opfamily, index: &(impl RelationRead + RelationWrite), check, ) } - } + }; + pgrx::info!( + "maintain: number_of_formerly_allocated_pages = {}", + maintain.number_of_formerly_allocated_pages + ); + pgrx::info!( + "maintain: number_of_freshly_allocated_pages = {}", + maintain.number_of_freshly_allocated_pages + ); + pgrx::info!( + "maintain: number_of_freed_pages = {}", + maintain.number_of_freed_pages + ); } pub fn build( @@ -333,6 +345,7 @@ pub fn insert( index: &(impl RelationRead + RelationWrite), payload: NonZero, vector: OwnedVector, + skip_freespaces: bool, ) { let bump = BumpAlloc::new(); let make_h1_plain_prefetcher = MakeH1PlainPrefetcherForInsertion { index }; @@ -352,6 +365,7 @@ pub fn insert( key, &bump, make_h1_plain_prefetcher, + skip_freespaces, ) } (OwnedVector::Vecf32(vector), DistanceKind::Dot) => { @@ -369,6 +383,7 @@ pub fn insert( key, &bump, make_h1_plain_prefetcher, + skip_freespaces, ) } (OwnedVector::Vecf16(vector), DistanceKind::L2) => { @@ -386,6 +401,7 @@ pub fn insert( key, &bump, make_h1_plain_prefetcher, + skip_freespaces, ) } (OwnedVector::Vecf16(vector), DistanceKind::Dot) => { @@ -403,6 +419,7 @@ pub fn insert( key, &bump, make_h1_plain_prefetcher, + skip_freespaces, ) } } diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index 532c5e1d..d44e0c0a 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -388,7 +388,7 @@ pub unsafe extern "C-unwind" fn ambuild( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, &index, payload, vector); + crate::index::algorithm::insert(opfamily, &index, payload, vector, true); } indtuples += 1; reporter.tuples_done(indtuples); @@ -870,7 +870,7 @@ unsafe fn parallel_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, &index, payload, vector); + crate::index::algorithm::insert(opfamily, &index, payload, vector, true); } unsafe { let indtuples; @@ -893,7 +893,7 @@ unsafe fn parallel_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, &index, payload, vector); + crate::index::algorithm::insert(opfamily, &index, payload, vector, true); } unsafe { let indtuples; diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index 67d98733..dc369d98 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -300,7 +300,7 @@ unsafe fn aminsertinner( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, &index, payload, vector); + crate::index::algorithm::insert(opfamily, &index, payload, vector, false); } } false From 59959f3c443ac3c31ddb91ac59ee4a03ef0c1112 Mon Sep 17 00:00:00 2001 From: xieydd Date: Thu, 29 May 2025 19:18:21 +0800 Subject: [PATCH 166/324] ci: Added pgxn publication workflow (#265) fix #261 --------- Signed-off-by: xieydd Co-authored-by: usamoi --- .github/workflows/pgxn.yml | 75 +++++++++++++++++++++++++++++++++++ .github/workflows/release.yml | 1 + META.json.in | 49 +++++++++++++++++++++++ Makefile | 17 ++++++++ 4 files changed, 142 insertions(+) create mode 100644 .github/workflows/pgxn.yml create mode 100644 META.json.in create mode 100644 Makefile diff --git a/.github/workflows/pgxn.yml b/.github/workflows/pgxn.yml new file mode 100644 index 00000000..64404419 --- /dev/null +++ b/.github/workflows/pgxn.yml @@ -0,0 +1,75 @@ +# workflows/pgxn.yml +# +# Publish vchord (PGXN) +# Build and publish the vchord extension to the PostgreSQL Extension Network (PGXN). This +# workflow is gets triggered on a new release creation or manually via the GitHub UI. + +name: Publish vchord to PGXN + +on: + release: + types: [created] + workflow_dispatch: + inputs: + tag: + description: "tag name (semver without v-prefix)" + required: true + type: string + +jobs: + semver: + runs-on: "ubuntu-latest" + + steps: + - name: Semver + id: semver + uses: actions/github-script@v7 + with: + script: | + const tag = "${{ github.event.inputs.tag }}" || "${{ github.event.release.tag_name }}"; + console.log(`Tag: ${tag}`); + const r = /^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(?:-((?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\+([0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?$/; + if (!r.test(tag)) { + core.setFailed(`Action failed with an invalid semver.`); + } + core.setOutput('SEMVER', tag); + + outputs: + SEMVER: ${{ steps.semver.outputs.SEMVER }} + + publish: + name: Publish to PGXN + runs-on: ubuntu-latest + needs: ["semver"] + env: + DISTNAME: vchord + DISTVERSION: ${{ needs.semver.outputs.SEMVER }} + + steps: + - name: Checkout Git Repository + uses: actions/checkout@v4 + + - name: Bundle the Release + run: | + # Create a PGXN-compatible zip file. + sed -e "s/@DISTVERSION@/${DISTVERSION}/g" META.json.in > META.json + # Create the build directory + mkdir -p ./build + # Output the archive to ./build/ and use DISTNAME--DISTVERSION.zip + git archive --format zip --prefix "${DISTNAME}-${DISTVERSION}/" --add-file META.json -o "./build/${DISTNAME}--${DISTVERSION}.zip" HEAD + echo "Release bundle created: ./build/${DISTNAME}--${DISTVERSION}.zip" + + - name: Release on PGXN + env: + PGXN_USERNAME: ${{ secrets.PGXN_USERNAME }} + PGXN_PASSWORD: ${{ secrets.PGXN_PASSWORD }} + run: | + # DISTNAME and DISTVERSION are available from the job's env context + ARCHIVE_PATH="./build/${DISTNAME}--${DISTVERSION}.zip" + echo "Uploading ${ARCHIVE_PATH} to PGXN..." + curl --fail -sS \ + --user "${PGXN_USERNAME}:${PGXN_PASSWORD}" \ + -F "submit=Release It!" \ + -F "archive=@${ARCHIVE_PATH}" \ + -H "X-Requested-With: XMLHttpRequest" \ + https://manager.pgxn.org/upload diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 121bd63a..3dc7fccc 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -219,3 +219,4 @@ jobs: --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 docker manifest push ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} + \ No newline at end of file diff --git a/META.json.in b/META.json.in new file mode 100644 index 00000000..a823cf15 --- /dev/null +++ b/META.json.in @@ -0,0 +1,49 @@ +{ + "name": "vchord", + "abstract": "Scalable, fast, and disk-friendly vector search in Postgres", + "description": "VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search, it supports L2 distance, inner product, cosine distance and maxsim", + "version": "@DISTVERSION@", + "maintainer": [ + "TensorChord support@tensorchord.ai", + "usamoi usamoi@tensorchord.ai" + ], + "license": "agpl_3", + "prereqs": { + "runtime": { + "requires": { + "PostgreSQL": "13.0.0" + } + } + }, + "provides": { + "vchord": { + "file": "src/lib.rs", + "docfile": "README.md", + "version": "@DISTVERSION@", + "abstract": "Scalable, fast, and disk-friendly vector search in Postgres" + } + }, + "resources": { + "homepage": "https://github.com/tensorchord/VectorChord", + "bugtracker": { + "web": "https://github.com/tensorchord/VectorChord/issues" + }, + "repository": { + "url": "https://github.com/tensorchord/VectorChord.git", + "web": "https://github.com/tensorchord/VectorChord", + "type": "git" + } + }, + "generated_by": "TensorChord", + "meta-spec": { + "version": "1.0.0", + "url": "http://pgxn.org/meta/spec.txt" + }, + "tags": [ + "vector search", + "vectors", + "datatype", + "nearest neighbor search", + "approximate nearest neighbors" + ] +} \ No newline at end of file diff --git a/Makefile b/Makefile new file mode 100644 index 00000000..5dc495f5 --- /dev/null +++ b/Makefile @@ -0,0 +1,17 @@ +PG_CONFIG ?= $(shell which pg_config) + +.PHONY: make package install +.DEFAULT_GOAL: package + +all: build-make package + +build-make: + mkdir -p ./build + cargo build -p make + cp ./target/debug/make ./build/make + +package: + PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" ./build/make package -o ./build/raw + +install: + PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" ./build/make install -i ./build/raw From d4ded96a1d34673bfe36130781ed02a5e599e002 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 30 May 2025 07:39:45 +0800 Subject: [PATCH 167/324] chore: upload 0.4.2 schema scripts (#266) try-job: psql_macos --------- Signed-off-by: usamoi --- .github/workflows/check.yml | 9 +- .github/workflows/pgxn.yml | 75 ---- .github/workflows/release.yml | 29 +- Makefile | 24 +- crates/make/src/main.rs | 68 +--- sql/install/vchord--0.4.2.sql | 566 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.4.1--0.4.2.sql | 1 + 7 files changed, 617 insertions(+), 155 deletions(-) delete mode 100644 .github/workflows/pgxn.yml create mode 100644 sql/install/vchord--0.4.2.sql create mode 100644 sql/upgrade/vchord--0.4.1--0.4.2.sql diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 6d87932f..60c47e02 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -201,8 +201,8 @@ jobs: - name: Install run: | - cargo make package -o ./build/raw - cargo make install --sudo -i ./build/raw + cargo make build -o ./build/raw + sudo make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install - name: Service run: | @@ -304,8 +304,8 @@ jobs: - name: Install run: | - cargo make package -o ./build/raw - cargo make install --sudo -i ./build/raw + cargo make build -o ./build/raw + sudo make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install - name: Service run: | @@ -327,7 +327,6 @@ jobs: VERSION: ${{ matrix.version }} ARCH: ${{ matrix.arch }} run: | - cargo make package -o ./build/raw (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-apple-darwin.zip .) - name: Upload Artifacts diff --git a/.github/workflows/pgxn.yml b/.github/workflows/pgxn.yml deleted file mode 100644 index 64404419..00000000 --- a/.github/workflows/pgxn.yml +++ /dev/null @@ -1,75 +0,0 @@ -# workflows/pgxn.yml -# -# Publish vchord (PGXN) -# Build and publish the vchord extension to the PostgreSQL Extension Network (PGXN). This -# workflow is gets triggered on a new release creation or manually via the GitHub UI. - -name: Publish vchord to PGXN - -on: - release: - types: [created] - workflow_dispatch: - inputs: - tag: - description: "tag name (semver without v-prefix)" - required: true - type: string - -jobs: - semver: - runs-on: "ubuntu-latest" - - steps: - - name: Semver - id: semver - uses: actions/github-script@v7 - with: - script: | - const tag = "${{ github.event.inputs.tag }}" || "${{ github.event.release.tag_name }}"; - console.log(`Tag: ${tag}`); - const r = /^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(?:-((?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\+([0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?$/; - if (!r.test(tag)) { - core.setFailed(`Action failed with an invalid semver.`); - } - core.setOutput('SEMVER', tag); - - outputs: - SEMVER: ${{ steps.semver.outputs.SEMVER }} - - publish: - name: Publish to PGXN - runs-on: ubuntu-latest - needs: ["semver"] - env: - DISTNAME: vchord - DISTVERSION: ${{ needs.semver.outputs.SEMVER }} - - steps: - - name: Checkout Git Repository - uses: actions/checkout@v4 - - - name: Bundle the Release - run: | - # Create a PGXN-compatible zip file. - sed -e "s/@DISTVERSION@/${DISTVERSION}/g" META.json.in > META.json - # Create the build directory - mkdir -p ./build - # Output the archive to ./build/ and use DISTNAME--DISTVERSION.zip - git archive --format zip --prefix "${DISTNAME}-${DISTVERSION}/" --add-file META.json -o "./build/${DISTNAME}--${DISTVERSION}.zip" HEAD - echo "Release bundle created: ./build/${DISTNAME}--${DISTVERSION}.zip" - - - name: Release on PGXN - env: - PGXN_USERNAME: ${{ secrets.PGXN_USERNAME }} - PGXN_PASSWORD: ${{ secrets.PGXN_PASSWORD }} - run: | - # DISTNAME and DISTVERSION are available from the job's env context - ARCHIVE_PATH="./build/${DISTNAME}--${DISTVERSION}.zip" - echo "Uploading ${ARCHIVE_PATH} to PGXN..." - curl --fail -sS \ - --user "${PGXN_USERNAME}:${PGXN_PASSWORD}" \ - -F "submit=Release It!" \ - -F "archive=@${ARCHIVE_PATH}" \ - -H "X-Requested-With: XMLHttpRequest" \ - https://manager.pgxn.org/upload diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 3dc7fccc..f54b5aec 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -85,7 +85,8 @@ jobs: ARCH: ${{ matrix.arch }} PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} run: | - cargo make package -o ./build/raw + cargo make build -o ./build/raw + (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) mkdir -p ./build/deb @@ -219,4 +220,28 @@ jobs: --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 docker manifest push ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} - \ No newline at end of file + + pgxn: + runs-on: "ubuntu-latest" + needs: ["semver", "build"] + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Upload + env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + PGXN_PASSWORD: ${{ secrets.PGXN_PASSWORD }} + run: | + mkdir -p ./build + + sed -e "s/@DISTVERSION@/${SEMVER}/g" META.json.in > META.json + git archive --format zip --prefix "vchord-${SEMVER}/" --add-file META.json -o "./build/vchord--${SEMVER}.zip" HEAD + + curl --fail -sS \ + --user "tensorchord:${PGXN_PASSWORD}" \ + -F "submit=Release It!" \ + -F "archive=@./build/vchord--${SEMVER}.zip" \ + -H "X-Requested-With: XMLHttpRequest" \ + https://manager.pgxn.org/upload diff --git a/Makefile b/Makefile index 5dc495f5..3e10bbd6 100644 --- a/Makefile +++ b/Makefile @@ -1,17 +1,17 @@ -PG_CONFIG ?= $(shell which pg_config) +PG_CONFIG ?= pg_config -.PHONY: make package install -.DEFAULT_GOAL: package +.PHONY: all build install uninstall -all: build-make package +all: build -build-make: - mkdir -p ./build - cargo build -p make - cp ./target/debug/make ./build/make - -package: - PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" ./build/make package -o ./build/raw +build: + PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo make build -o ./build/raw install: - PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" ./build/make install -i ./build/raw + cp -r ./build/raw/pkglibdir/. $(shell $(PG_CONFIG) --pkglibdir) + cp -r ./build/raw/sharedir/. $(shell $(PG_CONFIG) --sharedir) + +uninstall: + rm -f $(shell find $(shell $(PG_CONFIG) --pkglibdir) -type f -name 'vchord.*') + rm -f $(shell find $(shell $(PG_CONFIG) --sharedir)/extension -type f -name 'vchord.*') + rm -f $(shell find $(shell $(PG_CONFIG) --sharedir)/extension -type f -name 'vchord--*.sql') diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index a86aedac..2ecdbce2 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -31,24 +31,15 @@ struct Cli { #[derive(Subcommand)] enum Commands { - Package(PackageArgs), - Install(InstallArgs), + Build(BuildArgs), } #[derive(Args)] -struct PackageArgs { +struct BuildArgs { #[arg(short, long)] output: String, } -#[derive(Args)] -struct InstallArgs { - #[arg(short, long)] - input: Option, - #[arg(long, default_value = "false")] - sudo: bool, -} - fn main() -> Result<()> { let cli = Cli::parse(); if !std::fs::exists("vchord.control")? { @@ -203,14 +194,8 @@ fn main() -> Result<()> { let command_stdout = String::from_utf8(command_output.stdout)?.replace("\t", " "); Ok(command_stdout) }; - let uilts_install = |sudo: bool, permission: &str, src: &str, dst: &str| { - let mut command; - if sudo { - command = Command::new("sudo"); - command.arg("install"); - } else { - command = Command::new("install"); - }; + let uilts_install = |permission: &str, src: &str, dst: &str| { + let mut command = Command::new("install"); command.args(["-m", permission, src, dst]); let debug = format!("{command:?}"); let status = command.spawn()?.wait()?; @@ -219,36 +204,18 @@ fn main() -> Result<()> { } Ok(()) }; - let uilts_cp = |sudo: bool, args: &[&str]| { - let mut command; - if sudo { - command = Command::new("sudo"); - command.arg("cp"); - } else { - command = Command::new("cp"); - }; - command.args(args); - let debug = format!("{command:?}"); - let status = command.spawn()?.wait()?; - if !status.success() { - bail!("Command execution failed: {debug}"); - } - Ok(()) - }; let dll_suffix = if OS == "macos" && matches!(fork.as_str(), "pg13" | "pg14" | "pg15") { ".so" } else { DLL_SUFFIX }; - let install = |sudo, pkglibdir, sharedir_extension| -> Result<()> { + let install = |pkglibdir, sharedir_extension| -> Result<()> { uilts_install( - sudo, "755", &format!("./target/release/{DLL_PREFIX}vchord{DLL_SUFFIX}"), &format!("{pkglibdir}/vchord{dll_suffix}"), )?; uilts_install( - sudo, "644", "./vchord.control", &format!("{sharedir_extension}/vchord.control"), @@ -258,14 +225,12 @@ fn main() -> Result<()> { let path = maybe_entry?.path(); let name = path.file_name().context("broken assets")?; uilts_install( - sudo, "644", &format!("{}", path.display()), &format!("{sharedir_extension}/{}", name.display()), )?; } uilts_install( - sudo, "644", &format!("./sql/install/vchord--{version}.sql"), &format!("{sharedir_extension}/vchord--{version}.sql"), @@ -276,7 +241,6 @@ fn main() -> Result<()> { file.write_all(contents.as_bytes())?; let path = file.into_temp_path(); uilts_install( - sudo, "644", &format!("{}", path.display()), &format!("{sharedir_extension}/vchord--0.0.0.sql"), @@ -285,7 +249,7 @@ fn main() -> Result<()> { Ok(()) }; match cli.command { - Commands::Package(PackageArgs { output }) => { + Commands::Build(BuildArgs { output }) => { build()?; let pkglibdir = format!("{output}/pkglibdir"); let sharedir = format!("{output}/sharedir"); @@ -297,25 +261,7 @@ fn main() -> Result<()> { std::fs::create_dir_all(&pkglibdir)?; std::fs::create_dir_all(&sharedir)?; std::fs::create_dir_all(&sharedir_extension)?; - install(false, pkglibdir, sharedir_extension)?; - } - Commands::Install(InstallArgs { sudo, input }) => { - if let Some(input) = input { - uilts_cp( - sudo, - &["-r", &format!("{input}/pkglibdir/."), &map["PKGLIBDIR"]], - )?; - uilts_cp( - sudo, - &["-r", &format!("{input}/sharedir/."), &map["SHAREDIR"]], - )?; - } else { - build()?; - let pkglibdir = map["PKGLIBDIR"].clone(); - let sharedir = map["SHAREDIR"].clone(); - let sharedir_extension = format!("{sharedir}/extension"); - install(sudo, pkglibdir, sharedir_extension)?; - } + install(pkglibdir, sharedir_extension)?; } } Ok(()) diff --git a/sql/install/vchord--0.4.2.sql b/sql/install/vchord--0.4.2.sql new file mode 100644 index 00000000..522e9c8a --- /dev/null +++ b/sql/install/vchord--0.4.2.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:22 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:21 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:133 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:74 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:23 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.4.1--0.4.2.sql b/sql/upgrade/vchord--0.4.1--0.4.2.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.4.1--0.4.2.sql @@ -0,0 +1 @@ +/* nothing to do */ From 7f5834de3b2a5c7e4c9d40d594a9bad511b9fe40 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Tue, 3 Jun 2025 08:31:19 +0800 Subject: [PATCH 168/324] chore: update readme with version and badges (#267) - Update GHCR recommendation to `0.4.2` - Add release/download badge about GHCR - Introduce DeepWiki badge to enable auto refresh of LLM generated [wiki](https://deepwiki.com/tensorchord/VectorChord) - Update lists `lists` parameter divide from `100M` to `50M` based on `Cohere-50M` experiment result Signed-off-by: cutecutecat --- README.md | 42 +++++++++++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 916af6b3..0879e079 100644 --- a/README.md +++ b/README.md @@ -8,14 +8,16 @@ [Official Site][official-site-link] · [Blog][blog-link] · [Docs][docs-link] · [Feedback][github-issues-link] · [Contact Us][email-link] - [![][github-release-shield]][github-release-link] [![][docker-release-shield]][docker-release-link] [![][docker-pulls-shield]][docker-pulls-link] +[![][ghcr-release-shield]][ghcr-release-link] +[![][github-downloads-shield]][github-downloads-link] [![][discord-shield]][discord-link] [![][X-shield]][X-link] [![][cloud-shield]][cloud-link] +[![][deepwiki-shield]][deepwiki-link] [![][license-1-shield]][license-1-link] [![][license-2-shield]][license-2-link]
@@ -65,7 +67,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.3.0 + -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.4.2 ``` > [!NOTE] > In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). @@ -101,16 +103,6 @@ lists = [] $$); ``` -> [!NOTE] -> The `lists` option should be configured based on the number of vectors. Below is a table to assist with your selection. - -| vectors Range | List Calculation Formula | Example Result | -| ------------- | ------------------------------ | ---------------- | -| <128k | list = 1 | 1 | -| ≥128k and <2M | list = (2 * vectors) / 1000 | 256-4000 | -| ≥2M and <100M | list ∈ [4√vectors, 8√vectors] | 4000-80000 | -| ≥100M | list ∈ [8√vectors, 16√vectors] | 80000-160000 | - And then perform a vector search using `SELECT ... ORDER BY ... LIMIT ...`. ```SQL @@ -124,6 +116,19 @@ For more usage, please read: * [Performance Tuning](https://docs.vectorchord.ai/vectorchord/usage/performance-tuning.html) * [Advanced Features](https://docs.vectorchord.ai/vectorchord/usage/advanced-features.html) +> [!NOTE] +> The partition parameter, `lists`, should be configured based on the number of rows. The following table provides guidance for this selection. When searching, remember to set [`vchordrq.probes`](https://docs.vectorchord.ai/vectorchord/usage/search.html#vchordrq-probes) based on the value of lists. + + + +| Number of Rows $N$ | Recommended Number of Partitions $L$ | Example `lists` | +| -------------------------------------- | ------------------------------------ | --------------- | +| $N \in [0, 10^5)$ | N/A | `[]` | +| $N \in [10^5, 2 \times 10^6)$ | $L = \frac{N}{500}$ | `[2000]` | +| $N \in [2 \times 10^6, 5 \times 10^7)$ | $L \in [4 \sqrt{N}, 8 \sqrt{N}]$ | `[10000]` | +| $N \in [5 \times 10^7, \infty)$ | $L \in [8 \sqrt{N}, 16\sqrt{N}]$ | `[80000]` | + + ## License This software is licensed under a dual license model: @@ -141,19 +146,26 @@ You may choose either license based on your needs. We welcome any commercial col [license-2-shield]: https://img.shields.io/badge/License-ELv2-green?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCIgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiBmaWxsPSIjZmZmZmZmIj48cGF0aCBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik0xMi43NSAyLjc1YS43NS43NSAwIDAwLTEuNSAwVjQuNUg5LjI3NmExLjc1IDEuNzUgMCAwMC0uOTg1LjMwM0w2LjU5NiA1Ljk1N0EuMjUuMjUgMCAwMTYuNDU1IDZIMi4zNTNhLjc1Ljc1IDAgMTAwIDEuNUgzLjkzTC41NjMgMTUuMThhLjc2Mi43NjIgMCAwMC4yMS44OGMuMDguMDY0LjE2MS4xMjUuMzA5LjIyMS4xODYuMTIxLjQ1Mi4yNzguNzkyLjQzMy42OC4zMTEgMS42NjIuNjIgMi44NzYuNjJhNi45MTkgNi45MTkgMCAwMDIuODc2LS42MmMuMzQtLjE1NS42MDYtLjMxMi43OTItLjQzMy4xNS0uMDk3LjIzLS4xNTguMzEtLjIyM2EuNzUuNzUgMCAwMC4yMDktLjg3OEw1LjU2OSA3LjVoLjg4NmMuMzUxIDAgLjY5NC0uMTA2Ljk4NC0uMzAzbDEuNjk2LTEuMTU0QS4yNS4yNSAwIDAxOS4yNzUgNmgxLjk3NXYxNC41SDYuNzYzYS43NS43NSAwIDAwMCAxLjVoMTAuNDc0YS43NS43NSAwIDAwMC0xLjVIMTIuNzVWNmgxLjk3NGMuMDUgMCAuMS4wMTUuMTQuMDQzbDEuNjk3IDEuMTU0Yy4yOS4xOTcuNjMzLjMwMy45ODQuMzAzaC44ODZsLTMuMzY4IDcuNjhhLjc1Ljc1IDAgMDAuMjMuODk2Yy4wMTIuMDA5IDAgMCAuMDAyIDBhMy4xNTQgMy4xNTQgMCAwMC4zMS4yMDZjLjE4NS4xMTIuNDUuMjU2Ljc5LjRhNy4zNDMgNy4zNDMgMCAwMDIuODU1LjU2OCA3LjM0MyA3LjM0MyAwIDAwMi44NTYtLjU2OWMuMzM4LS4xNDMuNjA0LS4yODcuNzktLjM5OWEzLjUgMy41IDAgMDAuMzEtLjIwNi43NS43NSAwIDAwLjIzLS44OTZMMjAuMDcgNy41aDEuNTc4YS43NS43NSAwIDAwMC0xLjVoLTQuMTAyYS4yNS4yNSAwIDAxLS4xNC0uMDQzbC0xLjY5Ny0xLjE1NGExLjc1IDEuNzUgMCAwMC0uOTg0LS4zMDNIMTIuNzVWMi43NXpNMi4xOTMgMTUuMTk4YTUuNDE4IDUuNDE4IDAgMDAyLjU1Ny42MzUgNS40MTggNS40MTggMCAwMDIuNTU3LS42MzVMNC43NSA5LjM2OGwtMi41NTcgNS44M3ptMTQuNTEtLjAyNGMuMDgyLjA0LjE3NC4wODMuMjc1LjEyNi41My4yMjMgMS4zMDUuNDUgMi4yNzIuNDVhNS44NDYgNS44NDYgMCAwMDIuNTQ3LS41NzZMMTkuMjUgOS4zNjdsLTIuNTQ3IDUuODA3eiI+PC9wYXRoPjwvc3ZnPgo= [docker-release-link]: https://hub.docker.com/r/tensorchord/vchord-postgres -[docker-release-shield]: https://img.shields.io/docker/v/tensorchord/vchord-postgres?color=369eff&label=docker&labelColor=black&logo=docker&logoColor=white&style=flat&sort=semver +[docker-release-shield]: https://img.shields.io/docker/v/tensorchord/vchord-postgres?color=369eff&label=docker&labelColor=black&logo=docker&logoColor=white&style=flat [github-release-link]: https://github.com/tensorchord/VectorChord/releases [github-release-shield]: https://img.shields.io/github/v/release/tensorchord/VectorChord?color=369eff&labelColor=black&logo=github&style=flat +[ghcr-release-link]: https://github.com/orgs/tensorchord/packages/container/package/vchord-postgres + +[ghcr-release-shield]: https://img.shields.io/docker/v/tensorchord/vchord-postgres?color=369eff&label=GHCR&labelColor=black&logo=github&logoColor=white&style=flat [docker-pulls-link]: https://hub.docker.com/r/tensorchord/vchord-postgres [docker-pulls-shield]: https://img.shields.io/docker/pulls/tensorchord/vchord-postgres?color=45cc11&labelColor=black&style=flat&sort=semver [previous-docker-pulls-link]: https://hub.docker.com/r/tensorchord/pgvecto-rs [previous-docker-pulls-shield]: https://img.shields.io/docker/pulls/tensorchord/pgvecto-rs?color=45cc11&labelColor=black&style=flat&sort=semver +[github-downloads-link]: https://github.com/tensorchord/VectorChord/releases +[github-downloads-shield]: https://img.shields.io/github/downloads/tensorchord/VectorChord/total?color=45cc11&labelColor=black&style=flat&sort=semver [discord-link]: https://discord.gg/KqswhpVgdU [discord-shield]: https://img.shields.io/discord/974584200327991326?&logoColor=white&color=5865F2&style=flat&logo=discord&cacheSeconds=60 -[X-link]: https://twitter.com/TensorChord -[X-shield]: https://img.shields.io/twitter/follow/tensorchord?style=flat&logo=X&cacheSeconds=60 +[X-link]: https://x.com/TensorChord +[X-shield]: https://img.shields.io/badge/follow-%40tensorchord-1DA1F2?logo=x&style=flat&logoColor=white&color=1da1f2 [cloud-link]: https://cloud.vectorchord.ai/ [cloud-shield]: https://img.shields.io/badge/VectorChord_Cloud-Try_For_Free-F2B263.svg?labelColor=DAFDBA&logo=data:image/svg%2bxml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSIxMzMiIGhlaWdodD0iMTMyIiBmaWxsPSJub25lIj48cGF0aCBmaWxsPSIjRTZEQjNEIiBkPSJNNDguNCAzNy41YzAtMSAwLTEuNS0uMi0xLjhhMSAxIDAgMCAwLS44LS40Yy0uMyAwLS43LjMtMS42IDFMMjcuNiA1MC4xYy0xLjIuOC0xLjcgMS4zLTIuMiAxLjhhNSA1IDAgMCAwLS44IDEuNmMtLjIuNy0uMiAxLjQtLjIgMi45djM3LjNjMCAxLjIgMCAxLjguMyAyIC4yLjMuNS41LjguNC40IDAgLjgtLjQgMS42LTEuM2wxOS0xOC42IDEuNS0xLjhjLjMtLjQuNS0xIC42LTEuNC4yLS42LjItMS4yLjItMi40VjM3LjVaTTM1LjIgMTA1LjNjLS44LjgtMS4yIDEuMy0xLjIgMS42IDAgLjQgMCAuNy4zLjkuMy4yLjkuMiAyIC4yaDM3YzEuMyAwIDIgMCAyLjUtLjJhNSA1IDAgMCAwIDEuNS0uNmMuNi0uNCAxLS45IDEuOS0xLjhMOTYuNiA4NmMuNy0uOSAxLjEtMS4zIDEuMS0xLjZhMSAxIDAgMCAwLS4zLS44Yy0uMy0uMy0uOS0uMy0yLS4zaC0zNWMtMS4yIDAtMS44IDAtMi40LjJhNSA1IDAgMCAwLTEuNC42Yy0uNS4zLTEgLjctMS44IDEuNmwtMTkuNiAxOS42Wk05Ni4zIDcwLjFjMSAwIDEuNCAwIDEuNy0uMi40LS4xLjYtLjQuOC0uN2wuMS0xLjdWMzUuM2wtLjEtMS44Yy0uMi0uMy0uNC0uNS0uOC0uNy0uMy0uMi0uOC0uMi0xLjctLjJINjQuMWMtMSAwLTEuNCAwLTEuNy4yLS40LjItLjYuNC0uOC43bC0uMSAxLjh2MzIuMmwuMSAxLjdjLjIuMy40LjYuOC43LjMuMi44LjIgMS43LjJoMzIuMloiLz48cGF0aCBmaWxsPSIjMTAxNTA5IiBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik00My4yIDIxLjVjLTEuMyAwLTIgMC0yLjMtLjMtLjMtLjMtLjUtLjYtLjUtMSAwLS41LjQtMSAxLjEtMi4xTDUzLjEgMi4zYy42LS44LjktMS4yIDEuMi0xLjNoMWMuNC4xLjcuNSAxLjIgMS4zbDExLjYgMTUuOGMuOCAxIDEuMiAxLjYgMS4yIDIgMCAuNS0uMi44LS41IDEtLjQuNC0xIC40LTIuNC40SDU5Yy0uMy4yLS41LjQtLjYuNy0uMi4zLS4yLjYtLjIgMS40VjY4YzAgMS45IDAgMi44LjQgMy41LjMuNi44IDEuMSAxLjQgMS41LjcuMyAxLjcuMyAzLjUuM2g0NGwxLjMtLjFjLjMtLjEuNS0uMy42LS42bC4xLTEuNFY2NWMwLTEuMyAwLTIgLjMtMi4zLjItLjMuNi0uNSAxLS41czEgLjMgMiAxTDEzMCA3NWMuOC42IDEuMy45IDEuNCAxLjJ2MWMtLjEuNC0uNi43LTEuNCAxLjNsLTE3LjMgMTEuN2MtMSAuOC0xLjYgMS4xLTIgMS4xLS40IDAtLjgtLjItMS0uNS0uMy0uNC0uMy0xLS4zLTIuM1Y4MmwtLjEtMS40Yy0uMS0uMi0uMy0uNC0uNi0uNS0uMy0uMi0uNi0uMi0xLjQtLjJINjAuNWMtMS42IDAtMi40IDAtMy4xLjItLjcuMi0xLjMuNC0yIC44bC0yLjMgMi0yOC42IDI4LjNjLS41LjUtLjguOC0uOSAxLjF2LjhjLjEuMy40LjYuOSAxLjFsMy43IDMuN2MuOCAxIDEuMyAxLjMgMS4zIDEuOCAwIC40IDAgLjctLjMgMS0uMy4zLS45LjUtMiAuOGwtMjAgNC43Yy0xLjEuMi0xLjcuMy0yIC4yLS40LS4xLS43LS40LS44LS44LS4xLS4zIDAtMSAuMy0ybDUtMTkuOWMuMy0xLjEuNC0xLjcuOC0yIC4zLS4zLjctLjQgMS0uMy41IDAgLjkuNSAxLjggMS40bDMuNiAzLjcgMSAuOWguOWMuMy0uMS42LS40IDEtLjlsMjguNi0yOC4yYzEuMi0xLjEgMS43LTEuNyAyLjEtMi4zLjQtLjYuNy0xLjMuOC0yIC4yLS43LjItMS41LjItMy4yVjIzLjZsLS4xLTEuNGMtLjEtLjMtLjMtLjUtLjYtLjZsLTEuNC0uMWgtNi4yWiIgY2xpcC1ydWxlPSJldmVub2RkIi8+PC9zdmc+ +[deepwiki-link]: https://deepwiki.com/tensorchord/VectorChord +[deepwiki-shield]: https://deepwiki.com/badge.svg [blog-link]: https://blog.vectorchord.ai/ [official-site-link]: https://vectorchord.ai/ [github-issues-link]: https://github.com/tensorchord/VectorChord/issues From ebbc2f462c0cd86754e22b9a3d3e6363ac7286ca Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 9 Jun 2025 10:08:57 +0800 Subject: [PATCH 169/324] chore: update rust toolchain (#272) AVX512 is stablized in Rust 1.89. So we could switch to stable Rust toolchain on 7 August. job: +psql_macos job: +psql_windows Signed-off-by: usamoi --- .github/workflows/check.yml | 64 +++++- .github/workflows/release.yml | 3 + crates/algorithm/src/lib.rs | 1 - crates/make/src/main.rs | 1 + crates/rabitq/src/lib.rs | 3 - crates/simd/cshim/x86_64.c | 37 +++- crates/simd/src/f16.rs | 48 +++-- crates/simd/src/lib.rs | 18 -- crates/simd_macros/src/target.rs | 9 - rust-toolchain.toml | 2 +- src/index/am/am_build.rs | 2 +- src/index/am/mod.rs | 8 +- src/index/lazy_cell.rs | 326 ------------------------------- src/index/mod.rs | 1 - src/index/scanners/mod.rs | 5 +- src/index/storage.rs | 19 -- 16 files changed, 138 insertions(+), 409 deletions(-) delete mode 100644 src/index/lazy_cell.rs diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 60c47e02..8ba0878b 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -22,6 +22,7 @@ jobs: - name: Set up Environment run: | rustup set profile + rustup default nightly-2025-06-08 curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo @@ -112,6 +113,8 @@ jobs: - name: Set up Environment run: | rustup set profile + rustup default nightly-2025-06-08 + sudo apt-get update if [ "$(uname -m)" == "x86_64" ]; then @@ -152,6 +155,11 @@ jobs: fi psql: + if: | + (github.event_name == 'push' && !contains(github.event.head_commit.message, 'job: -psql')) || + (github.event_name == 'pull_request' && !contains(github.event.pull_request.body, 'job: -psql')) || + github.event_name == 'workflow_dispatch' + strategy: matrix: version: ["13", "14", "15", "16", "17"] @@ -166,6 +174,8 @@ jobs: - name: Set up Environment run: | rustup set profile + rustup default nightly-2025-06-08 + sudo apt-get update sudo apt-get remove -y '^postgres.*' '^libpq.*' @@ -254,8 +264,8 @@ jobs: psql_macos: if: | - (github.event_name == 'push' && contains(github.event.head_commit.message, 'try-job: psql_macos')) || - (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'try-job: psql_macos')) || + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_macos')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +psql_macos')) || github.event_name == 'workflow_dispatch' strategy: @@ -273,6 +283,7 @@ jobs: - name: Set up Environment run: | rustup set profile + rustup default nightly-2025-06-08 echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV @@ -337,3 +348,52 @@ jobs: ./build/postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin.zip compression-level: 9 retention-days: 14 + + psql_windows: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_windows')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +psql_windows')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + version: ["16"] + arch: ["x86_64"] + + runs-on: "windows-2022" + + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: true + + steps: + - name: Set up Environment + run: | + rustup set profile + rustup default nightly-2025-06-08 + + Invoke-WebRequest -Uri "https://get.enterprisedb.com/postgresql/postgresql-16.9-1-windows-x64-binaries.zip" -OutFile "$env:TEMP\postgresql-install.zip" + Expand-Archive -Path "$env:TEMP\postgresql-install.zip" -DestinationPath "D:\postgresql-install" -Force + Add-Content -Path $env:GITHUB_ENV -Value "PGRX_PG_CONFIG_PATH=D:\postgresql-install\pgsql\bin\pg_config.exe" + Add-Content -Path $env:GITHUB_ENV -Value "PGBIN=D:\postgresql-install\pgsql\bin" + Add-Content -Path $env:GITHUB_ENV -Value "PGDATA=D:\postgresql-install\pgsql\data" + Add-Content -Path $env:GITHUB_ENV -Value "PGROOT=D:\postgresql-install\pgsql" + + Invoke-WebRequest -Uri https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-${{ matrix.arch }}-pc-windows-msvc.zip -OutFile "$env:TEMP\sqllogictest-install.zip" + Expand-Archive -Path "$env:TEMP\sqllogictest-install.zip" -DestinationPath "D:\sqllogictest-install" -Force + Add-Content -Path $env:GITHUB_PATH -Value "D:\sqllogictest-install" + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Checkout + uses: actions/checkout@v4 + + - name: Clippy + run: | + cargo clippy -p vchord --features pg${{ matrix.version }} -- --no-deps + + - name: Install + run: | + cargo make build -o ./build/raw + make PG_CONFIG="$env:PGRX_PG_CONFIG_PATH" install diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index f54b5aec..65a37193 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -51,6 +51,9 @@ jobs: steps: - name: Set up Environment run: | + rustup set profile + rustup default nightly-2025-06-08 + sudo apt-get update sudo apt-get remove -y '^postgres.*' '^libpq.*' diff --git a/crates/algorithm/src/lib.rs b/crates/algorithm/src/lib.rs index 853a2803..68602551 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algorithm/src/lib.rs @@ -12,7 +12,6 @@ // // Copyright (c) 2025 TensorChord Inc. -#![feature(select_unpredictable)] #![allow(clippy::type_complexity)] mod build; diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index 2ecdbce2..cf17d576 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -83,6 +83,7 @@ fn main() -> Result<()> { eprintln!("Fork: {fork}"); fork }; + #[allow(clippy::collapsible_if)] let version = 'version: { for line in std::fs::read_to_string("./vchord.control")?.lines() { if let Some(prefix_stripped) = line.strip_prefix("default_version = '") { diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index d96f3676..30f1856a 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -12,9 +12,6 @@ // // Copyright (c) 2025 TensorChord Inc. -#![feature(slice_as_chunks)] -#![feature(select_unpredictable)] - pub mod binary; pub mod block; pub mod packing; diff --git a/crates/simd/cshim/x86_64.c b/crates/simd/cshim/x86_64.c index 8ba1357b..61a93421 100644 --- a/crates/simd/cshim/x86_64.c +++ b/crates/simd/cshim/x86_64.c @@ -16,19 +16,50 @@ #if !(__clang_major__ >= 16) #error "Clang version must be at least 16." #endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 12) +#error "GCC version must be at least 12." +#endif #else -#error "This file requires Clang." +#error "This file requires Clang or GCC." #endif #include #include #include +#if defined(__clang__) && defined(_MSC_VER) && (__clang_major__ <= 19) +// https://github.com/llvm/llvm-project/issues/53520 +// clang-format off +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +// clang-format on +#endif + typedef _Float16 f16; typedef float f32; __attribute__((target("arch=x86-64-v4,avx512fp16"))) float -fp16_reduce_sum_of_xy_v4fp16(f16 *restrict a, f16 *restrict b, size_t n) { +fp16_reduce_sum_of_xy_v4_avx512fp16(f16 *restrict a, f16 *restrict b, size_t n) { __m512h xy = _mm512_setzero_ph(); while (n >= 32) { __m512h x = _mm512_loadu_ph(a); @@ -48,7 +79,7 @@ fp16_reduce_sum_of_xy_v4fp16(f16 *restrict a, f16 *restrict b, size_t n) { } __attribute__((target("arch=x86-64-v4,avx512fp16"))) float -fp16_reduce_sum_of_d2_v4fp16(f16 *restrict a, f16 *restrict b, size_t n) { +fp16_reduce_sum_of_d2_v4_avx512fp16(f16 *restrict a, f16 *restrict b, size_t n) { __m512h d2 = _mm512_setzero_ph(); while (n >= 32) { __m512h x = _mm512_loadu_ph(a); diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index b5415cae..7d7e184d 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -239,25 +239,29 @@ mod reduce_sum_of_xy { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4fp16")] - pub fn reduce_sum_of_xy_v4fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + pub fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { unsafe extern "C" { - unsafe fn fp16_reduce_sum_of_xy_v4fp16(a: *const (), b: *const (), n: usize) - -> f32; + unsafe fn fp16_reduce_sum_of_xy_v4_avx512fp16( + a: *const (), + b: *const (), + n: usize, + ) -> f32; } - fp16_reduce_sum_of_xy_v4fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + fp16_reduce_sum_of_xy_v4_avx512fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } } #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_xy_v4fp16_test() { + fn reduce_sum_of_xy_v4_avx512fp16_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v4fp16") { - println!("test {} ... skipped (v4fp16)", module_path!()); + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -272,7 +276,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_v4fp16(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v4_avx512fp16(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -494,7 +498,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -511,25 +515,29 @@ mod reduce_sum_of_d2 { #[inline] #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4fp16")] - pub fn reduce_sum_of_d2_v4fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + pub fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { unsafe extern "C" { - unsafe fn fp16_reduce_sum_of_d2_v4fp16(a: *const (), b: *const (), n: usize) - -> f32; + unsafe fn fp16_reduce_sum_of_d2_v4_avx512fp16( + a: *const (), + b: *const (), + n: usize, + ) -> f32; } - fp16_reduce_sum_of_d2_v4fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + fp16_reduce_sum_of_d2_v4_avx512fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) } } #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_d2_v4fp16_test() { + fn reduce_sum_of_d2_v4_avx512fp16_test() { use rand::Rng; const EPSILON: f32 = 6.4; - if !crate::is_cpu_detected!("v4fp16") { - println!("test {} ... skipped (v4fp16)", module_path!()); + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -544,7 +552,7 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_v4fp16(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_d2_v4_avx512fp16(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -774,7 +782,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 31307998..8ddaac42 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -12,10 +12,6 @@ // // Copyright (c) 2025 TensorChord Inc. -#![cfg_attr(target_arch = "x86_64", feature(avx512_target_feature))] -#![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] -#![feature(slice_as_chunks)] -#![feature(select_unpredictable)] #![allow(unsafe_code)] mod aligned; @@ -85,20 +81,6 @@ mod internal { #[allow(unused_imports)] pub use is_riscv64_cpu_detected; - #[cfg(target_arch = "x86_64")] - pub fn is_v4fp16_detected() -> bool { - std::arch::is_x86_feature_detected!("avx512fp16") - && std::arch::is_x86_feature_detected!("avx512bw") - && std::arch::is_x86_feature_detected!("avx512cd") - && std::arch::is_x86_feature_detected!("avx512dq") - && std::arch::is_x86_feature_detected!("avx512vl") - && std::arch::is_x86_feature_detected!("bmi1") - && std::arch::is_x86_feature_detected!("bmi2") - && std::arch::is_x86_feature_detected!("lzcnt") - && std::arch::is_x86_feature_detected!("movbe") - && std::arch::is_x86_feature_detected!("popcnt") - } - #[cfg(target_arch = "x86_64")] pub fn is_v4_detected() -> bool { std::arch::is_x86_feature_detected!("avx512bw") diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs index 2978ef5e..720abc6c 100644 --- a/crates/simd_macros/src/target.rs +++ b/crates/simd_macros/src/target.rs @@ -19,15 +19,6 @@ pub struct TargetCpu { } pub const TARGET_CPUS: &[TargetCpu] = &[ - // This is a temporary hack, for using AVX512FP16 without unstable features. - TargetCpu { - target_cpu: "v4fp16", - target_arch: "x86_64", - target_features: &[ - "avx512bw", "avx512cd", "avx512dq", "avx512vl", // simd - "bmi1", "bmi2", "lzcnt", "movbe", "popcnt", // bit-operations - ], - }, TargetCpu { target_cpu: "v4", target_arch: "x86_64", diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 1b42d5ea..04270ea8 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2025-04-26" +channel = "nightly-2025-06-08" diff --git a/src/index/am/am_build.rs b/src/index/am/am_build.rs index d44e0c0a..7ac85117 100644 --- a/src/index/am/am_build.rs +++ b/src/index/am/am_build.rs @@ -374,7 +374,7 @@ pub unsafe extern "C-unwind" fn ambuild( pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); pgrx::pg_sys::ConditionVariableSleep( &raw mut (*leader.vchordrqshared).workersdonecv, - pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, ); } reporter.tuples_done((*leader.vchordrqshared).indtuples); diff --git a/src/index/am/mod.rs b/src/index/am/mod.rs index dc369d98..80c8da9a 100644 --- a/src/index/am/mod.rs +++ b/src/index/am/mod.rs @@ -16,15 +16,16 @@ pub mod am_build; use super::algorithm::BumpAlloc; use crate::index::gucs; -use crate::index::lazy_cell::LazyCell; use crate::index::opclass::{Opfamily, opfamily}; use crate::index::scanners::*; use crate::index::storage::PostgresRelation; use algorithm::Bump; use pgrx::datum::Internal; use pgrx::pg_sys::{BlockIdData, Datum, ItemPointerData}; +use std::cell::LazyCell; use std::ffi::CStr; use std::num::NonZero; +use std::ops::DerefMut; use std::ptr::NonNull; use std::sync::OnceLock; @@ -59,7 +60,7 @@ pub fn init() { unsafe { kind = pgrx::pg_sys::add_reloption_kind(); pgrx::pg_sys::add_string_reloption( - kind, + kind as _, c"options".as_ptr(), c"Vector index options, represented as a TOML string.".as_ptr(), c"".as_ptr(), @@ -494,7 +495,7 @@ pub unsafe extern "C-unwind" fn amgettuple( pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); } let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; - if let Some((_, key, recheck)) = LazyCell::force_mut(&mut scanner.scanning).next() { + if let Some((_, key, recheck)) = scanner.scanning.deref_mut().next() { unsafe { (*scan).xs_heaptid = key_to_ctid(key); (*scan).xs_recheck = recheck; @@ -609,6 +610,7 @@ impl Tuple for HeapTuple<'_> { } } + #[allow(clippy::collapsible_if)] fn filter(&mut self) -> bool { unsafe { let this = &mut self.this; diff --git a/src/index/lazy_cell.rs b/src/index/lazy_cell.rs deleted file mode 100644 index 5fe5a1a2..00000000 --- a/src/index/lazy_cell.rs +++ /dev/null @@ -1,326 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -// Emulate unstable library feature `lazy_get`. -// See https://github.com/rust-lang/rust/issues/129333. - -#![allow(dead_code)] - -use core::cell::UnsafeCell; -use core::hint::unreachable_unchecked; -use core::ops::Deref; - -enum State { - Uninit(F), - Init(T), - Poisoned, -} - -/// A value which is initialized on the first access. -/// -/// For a thread-safe version of this struct, see [`std::sync::LazyLock`]. -/// -/// [`std::sync::LazyLock`]: ../../std/sync/struct.LazyLock.html -/// -/// # Examples -/// -/// ``` -/// use std::cell::LazyCell; -/// -/// let lazy: LazyCell = LazyCell::new(|| { -/// println!("initializing"); -/// 92 -/// }); -/// println!("ready"); -/// println!("{}", *lazy); -/// println!("{}", *lazy); -/// -/// // Prints: -/// // ready -/// // initializing -/// // 92 -/// // 92 -/// ``` -pub struct LazyCell T> { - state: UnsafeCell>, -} - -impl T> LazyCell { - /// Creates a new lazy value with the given initializing function. - /// - /// # Examples - /// - /// ``` - /// use std::cell::LazyCell; - /// - /// let hello = "Hello, World!".to_string(); - /// - /// let lazy = LazyCell::new(|| hello.to_uppercase()); - /// - /// assert_eq!(&*lazy, "HELLO, WORLD!"); - /// ``` - #[inline] - pub const fn new(f: F) -> LazyCell { - LazyCell { - state: UnsafeCell::new(State::Uninit(f)), - } - } - - /// Consumes this `LazyCell` returning the stored value. - /// - /// Returns `Ok(value)` if `Lazy` is initialized and `Err(f)` otherwise. - /// - /// # Examples - /// - /// ``` - /// #![feature(lazy_cell_into_inner)] - /// - /// use std::cell::LazyCell; - /// - /// let hello = "Hello, World!".to_string(); - /// - /// let lazy = LazyCell::new(|| hello.to_uppercase()); - /// - /// assert_eq!(&*lazy, "HELLO, WORLD!"); - /// assert_eq!(LazyCell::into_inner(lazy).ok(), Some("HELLO, WORLD!".to_string())); - /// ``` - pub fn into_inner(this: Self) -> Result { - match this.state.into_inner() { - State::Init(data) => Ok(data), - State::Uninit(f) => Err(f), - State::Poisoned => panic_poisoned(), - } - } - - /// Forces the evaluation of this lazy value and returns a reference to - /// the result. - /// - /// This is equivalent to the `Deref` impl, but is explicit. - /// - /// # Examples - /// - /// ``` - /// use std::cell::LazyCell; - /// - /// let lazy = LazyCell::new(|| 92); - /// - /// assert_eq!(LazyCell::force(&lazy), &92); - /// assert_eq!(&*lazy, &92); - /// ``` - #[inline] - pub fn force(this: &LazyCell) -> &T { - // SAFETY: - // This invalidates any mutable references to the data. The resulting - // reference lives either until the end of the borrow of `this` (in the - // initialized case) or is invalidated in `really_init` (in the - // uninitialized case; `really_init` will create and return a fresh reference). - let state = unsafe { &*this.state.get() }; - match state { - State::Init(data) => data, - // SAFETY: The state is uninitialized. - State::Uninit(_) => unsafe { LazyCell::really_init(this) }, - State::Poisoned => panic_poisoned(), - } - } - - /// Forces the evaluation of this lazy value and returns a mutable reference to - /// the result. - /// - /// # Examples - /// - /// ``` - /// #![feature(lazy_get)] - /// use std::cell::LazyCell; - /// - /// let mut lazy = LazyCell::new(|| 92); - /// - /// let p = LazyCell::force_mut(&mut lazy); - /// assert_eq!(*p, 92); - /// *p = 44; - /// assert_eq!(*lazy, 44); - /// ``` - #[inline] - pub fn force_mut(this: &mut LazyCell) -> &mut T { - #[cold] - /// # Safety - /// May only be called when the state is `Uninit`. - unsafe fn really_init_mut T>(state: &mut State) -> &mut T { - // INVARIANT: Always valid, but the value may not be dropped. - struct PoisonOnPanic(*mut State); - impl Drop for PoisonOnPanic { - #[inline] - fn drop(&mut self) { - // SAFETY: Invariant states it is valid, and we don't drop the old value. - unsafe { - self.0.write(State::Poisoned); - } - } - } - - let State::Uninit(f) = state else { - // `unreachable!()` here won't optimize out because the function is cold. - // SAFETY: Precondition. - unsafe { unreachable_unchecked() }; - }; - // SAFETY: We never drop the state after we read `f`, and we write a valid value back - // in any case, panic or success. `f` can't access the `LazyCell` because it is mutably - // borrowed. - let f = unsafe { core::ptr::read(f) }; - // INVARIANT: Initiated from mutable reference, don't drop because we read it. - let guard = PoisonOnPanic(state); - let data = f(); - // SAFETY: `PoisonOnPanic` invariant, and we don't drop the old value. - unsafe { - core::ptr::write(guard.0, State::Init(data)); - } - core::mem::forget(guard); - let State::Init(data) = state else { - unreachable!() - }; - data - } - - let state = this.state.get_mut(); - match state { - State::Init(data) => data, - // SAFETY: `state` is `Uninit`. - State::Uninit(_) => unsafe { really_init_mut(state) }, - State::Poisoned => panic_poisoned(), - } - } - - /// # Safety - /// May only be called when the state is `Uninit`. - #[cold] - unsafe fn really_init(this: &LazyCell) -> &T { - // SAFETY: - // This function is only called when the state is uninitialized, - // so no references to `state` can exist except for the reference - // in `force`, which is invalidated here and not accessed again. - let state = unsafe { &mut *this.state.get() }; - // Temporarily mark the state as poisoned. This prevents reentrant - // accesses and correctly poisons the cell if the closure panicked. - let State::Uninit(f) = core::mem::replace(state, State::Poisoned) else { - unreachable!() - }; - - let data = f(); - - // SAFETY: - // If the closure accessed the cell through something like a reentrant - // mutex, but caught the panic resulting from the state being poisoned, - // the mutable borrow for `state` will be invalidated, so we need to - // go through the `UnsafeCell` pointer here. The state can only be - // poisoned at this point, so using `write` to skip the destructor - // of `State` should help the optimizer. - unsafe { this.state.get().write(State::Init(data)) }; - - // SAFETY: - // The previous references were invalidated by the `write` call above, - // so do a new shared borrow of the state instead. - let state = unsafe { &*this.state.get() }; - let State::Init(data) = state else { - unreachable!() - }; - data - } -} - -impl LazyCell { - /// Returns a mutable reference to the value if initialized, or `None` if not. - /// - /// # Examples - /// - /// ``` - /// #![feature(lazy_get)] - /// - /// use std::cell::LazyCell; - /// - /// let mut lazy = LazyCell::new(|| 92); - /// - /// assert_eq!(LazyCell::get_mut(&mut lazy), None); - /// let _ = LazyCell::force(&lazy); - /// *LazyCell::get_mut(&mut lazy).unwrap() = 44; - /// assert_eq!(*lazy, 44); - /// ``` - #[inline] - pub fn get_mut(this: &mut LazyCell) -> Option<&mut T> { - let state = this.state.get_mut(); - match state { - State::Init(data) => Some(data), - _ => None, - } - } - - /// Returns a reference to the value if initialized, or `None` if not. - /// - /// # Examples - /// - /// ``` - /// #![feature(lazy_get)] - /// - /// use std::cell::LazyCell; - /// - /// let lazy = LazyCell::new(|| 92); - /// - /// assert_eq!(LazyCell::get(&lazy), None); - /// let _ = LazyCell::force(&lazy); - /// assert_eq!(LazyCell::get(&lazy), Some(&92)); - /// ``` - #[inline] - pub fn get(this: &LazyCell) -> Option<&T> { - // SAFETY: - // This is sound for the same reason as in `force`: once the state is - // initialized, it will not be mutably accessed again, so this reference - // will stay valid for the duration of the borrow to `self`. - let state = unsafe { &*this.state.get() }; - match state { - State::Init(data) => Some(data), - _ => None, - } - } -} - -impl T> Deref for LazyCell { - type Target = T; - #[inline] - fn deref(&self) -> &T { - LazyCell::force(self) - } -} - -impl Default for LazyCell { - /// Creates a new lazy value using `Default` as the initializing function. - #[inline] - fn default() -> LazyCell { - LazyCell::new(T::default) - } -} - -impl core::fmt::Debug for LazyCell { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - let mut d = f.debug_tuple("LazyCell"); - match LazyCell::get(self) { - Some(data) => d.field(data), - None => d.field(&format_args!("")), - }; - d.finish() - } -} - -#[cold] -#[inline(never)] -const fn panic_poisoned() -> ! { - panic!("LazyCell instance has previously been poisoned") -} diff --git a/src/index/mod.rs b/src/index/mod.rs index 32707d4c..8d2f48e2 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -17,7 +17,6 @@ pub mod am; pub mod functions; pub mod gucs; pub mod hook; -pub mod lazy_cell; pub mod opclass; pub mod scanners; pub mod storage; diff --git a/src/index/scanners/mod.rs b/src/index/scanners/mod.rs index 960ce6c9..5a74485a 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/scanners/mod.rs @@ -16,9 +16,10 @@ mod default; mod maxsim; use super::opclass::Opfamily; -use crate::index::lazy_cell::LazyCell; use algorithm::{Bump, RelationPrefetch, RelationRead, RelationReadStream, Sequence}; use pgrx::pg_sys::Datum; +use std::cell::LazyCell; +use std::ops::DerefMut; pub use default::DefaultBuilder; pub use maxsim::MaxsimBuilder; @@ -82,7 +83,7 @@ impl T> Fetcher for LazyCell { Self: 'a; fn fetch(&mut self, key: [u16; 3]) -> Option> { - LazyCell::force_mut(self).fetch(key) + self.deref_mut().fetch(key) } } diff --git a/src/index/storage.rs b/src/index/storage.rs index 2bad142f..ef88e2b0 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -173,25 +173,6 @@ impl Page for PostgresPage { const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); -pub struct PostgresBufferGuard { - buf: i32, - id: u32, -} - -impl PageGuard for PostgresBufferGuard { - fn id(&self) -> u32 { - self.id - } -} - -impl Drop for PostgresBufferGuard { - fn drop(&mut self) { - unsafe { - pgrx::pg_sys::ReleaseBuffer(self.buf); - } - } -} - pub struct PostgresBufferReadGuard { buf: i32, page: NonNull, From b42e6edd4666355f50450c88916670c658512e81 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 20 Jun 2025 12:31:31 +0800 Subject: [PATCH 170/324] chore: upload 0.4.3 schema scripts (#278) job: +psql_macos job: +psql_windows --------- Signed-off-by: usamoi --- Cargo.lock | 40 +- Cargo.toml | 7 +- README.md | 2 +- crates/algorithm/src/tuples.rs | 5 +- crates/rabitq/src/block.rs | 6 +- sql/install/vchord--0.4.3.sql | 566 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.4.2--0.4.3.sql | 1 + src/lib.rs | 8 +- 8 files changed, 616 insertions(+), 19 deletions(-) create mode 100644 sql/install/vchord--0.4.3.sql create mode 100644 sql/upgrade/vchord--0.4.2--0.4.3.sql diff --git a/Cargo.lock b/Cargo.lock index 16a63c54..0ff475a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -771,6 +771,16 @@ dependencies = [ "windows-targets 0.53.0", ] +[[package]] +name = "libmimalloc-sys" +version = "0.1.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf88cd67e9de251c1781dbe2f641a1a3ad66eaae831b8a2c38fbdc5ddae16d4d" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "linux-raw-sys" version = "0.9.4" @@ -800,6 +810,15 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +[[package]] +name = "mimalloc" +version = "0.1.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1791cbe101e95af5764f06f20f6760521f7158f69dbf9d6baf941ee1bf6bc40" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -1303,9 +1322,9 @@ dependencies = [ [[package]] name = "serde_spanned" -version = "0.6.8" +version = "0.6.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87607cb1398ed59d48732e575a4c28a7a8ebf2454b964fe3f224f2afc07909e1" +checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3" dependencies = [ "serde", ] @@ -1460,9 +1479,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.8.22" +version = "0.8.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05ae329d1f08c4d17a59bed7ff5b5a769d062e64a62d34a3261b219e62cd5aae" +checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" dependencies = [ "serde", "serde_spanned", @@ -1472,18 +1491,18 @@ dependencies = [ [[package]] name = "toml_datetime" -version = "0.6.9" +version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3da5db5a963e24bc68be8b17b6fa82814bb22ee8660f192bb182771d498f09a3" +checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c" dependencies = [ "serde", ] [[package]] name = "toml_edit" -version = "0.22.26" +version = "0.22.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "310068873db2c5b3e7659d2cc35d21855dbafa50d1ce336397c666e3cb08137e" +checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" dependencies = [ "indexmap", "serde", @@ -1495,9 +1514,9 @@ dependencies = [ [[package]] name = "toml_write" -version = "0.1.1" +version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfb942dfe1d8e29a7ee7fcbde5bd2b9a25fb89aa70caea2eba3bee836ff41076" +checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" [[package]] name = "twox-hash" @@ -1605,6 +1624,7 @@ dependencies = [ "half 2.6.0", "jemallocator", "k_means", + "mimalloc", "paste", "pgrx", "pgrx-catalog", diff --git a/Cargo.toml b/Cargo.toml index c23d84a3..1072e19f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,13 +36,16 @@ pgrx-catalog = "0.2.0" rand.workspace = true seq-macro.workspace = true serde.workspace = true -toml = "0.8.22" +toml = "0.8.23" validator.workspace = true zerocopy.workspace = true -[target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), target_os = "linux"))'.dependencies] +[target.'cfg(all(target_arch = "x86_64", target_os = "linux"))'.dependencies] jemallocator = { version = "0.5.4", features = ["disable_initial_exec_tls"] } +[target.'cfg(all(target_arch = "aarch64", target_os = "linux"))'.dependencies] +mimalloc = { version = "0.1.47", features = ["local_dynamic_tls"] } + [lints] workspace = true diff --git a/README.md b/README.md index 0879e079..49f13564 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.4.2 + -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.4.3 ``` > [!NOTE] > In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). diff --git a/crates/algorithm/src/tuples.rs b/crates/algorithm/src/tuples.rs index 9ebf5571..b5f450b6 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/algorithm/src/tuples.rs @@ -142,7 +142,10 @@ impl WithReader for MetaTuple { MAGIC => { let checker = RefChecker::new(source); if VERSION != *checker.prefix::(size_of::()) { - panic!("deserialization: bad version number"); + panic!( + "deserialization: bad version number; {}", + "after upgrading VectorChord, please use REINDEX to rebuild the index." + ); } let header: &MetaTupleHeader = checker.prefix(size_of::()); let centroid_prefetch = diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs index f6d9477e..49bc6079 100644 --- a/crates/rabitq/src/block.rs +++ b/crates/rabitq/src/block.rs @@ -183,11 +183,11 @@ fn compress(vector: &[f32]) -> Vec<[f32; 16]> { ] }; - let (arrays, reminder) = vector.as_chunks::<4>(); + let (arrays, remainder) = vector.as_chunks::<4>(); let mut result = arrays.iter().copied().map(f).collect::>(); - if !reminder.is_empty() { + if !remainder.is_empty() { let mut array = [0.0; 4]; - array[..reminder.len()].copy_from_slice(reminder); + array[..remainder.len()].copy_from_slice(remainder); result.push(f(array)); } result diff --git a/sql/install/vchord--0.4.3.sql b/sql/install/vchord--0.4.3.sql new file mode 100644 index 00000000..583a2f61 --- /dev/null +++ b/sql/install/vchord--0.4.3.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:22 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:21 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:133 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:75 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:23 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.4.2--0.4.3.sql b/sql/upgrade/vchord--0.4.2--0.4.3.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.4.2--0.4.3.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/src/lib.rs b/src/lib.rs index 654b4d1b..89b7a2a5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -40,7 +40,11 @@ extern "C-unwind" fn _PG_init() { compile_error!("Target architecture is not supported."); #[cfg(not(miri))] -#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] -#[cfg(target_os = "linux")] +#[cfg(all(target_arch = "x86_64", target_os = "linux"))] #[global_allocator] static GLOBAL_ALLOCATOR: jemallocator::Jemalloc = jemallocator::Jemalloc; + +#[cfg(not(miri))] +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +#[global_allocator] +static GLOBAL_ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; From 73690fe3fced943ec4affc250476f35214393a02 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 1 Jul 2025 11:27:39 +0800 Subject: [PATCH 171/324] feat: graph index (#280) job: +psql_macos job: +psql_windows Signed-off-by: usamoi --- .github/workflows/check.yml | 13 +- .github/workflows/release.yml | 2 +- Cargo.lock | 399 +++++++--- Cargo.toml | 13 +- crates/algo/Cargo.toml | 17 + crates/algo/src/accessor.rs | 335 ++++++++ crates/{algorithm => algo}/src/lib.rs | 231 +++--- crates/{algorithm => algo}/src/prefetcher.rs | 101 ++- crates/algo/src/tuples.rs | 233 ++++++ crates/distance/Cargo.toml | 3 + crates/distance/src/lib.rs | 23 +- crates/k_means/src/lib.rs | 18 +- crates/make/Cargo.toml | 4 +- crates/rabitq/src/b1.rs | 181 +++++ crates/rabitq/src/b2.rs | 71 ++ crates/rabitq/src/b4.rs | 71 ++ crates/rabitq/src/b8.rs | 62 ++ crates/rabitq/src/binary.rs | 202 ----- crates/rabitq/src/block.rs | 194 ----- crates/rabitq/src/extended.rs | 181 +++++ crates/rabitq/src/lib.rs | 64 +- crates/rabitq/src/original.rs | 456 +++++++++++ crates/rabitq/src/packing.rs | 16 - crates/simd/src/f16.rs | 5 + crates/simd/src/f32.rs | 14 + crates/simd/src/lib.rs | 2 + crates/simd/src/packed.rs | 92 +++ crates/simd/src/rotate.rs | 6 +- crates/vchordg/Cargo.toml | 26 + crates/vchordg/src/build.rs | 72 ++ crates/vchordg/src/bulkdelete.rs | 62 ++ crates/vchordg/src/candidates.rs | 72 ++ crates/vchordg/src/insert.rs | 577 ++++++++++++++ crates/vchordg/src/lib.rs | 44 ++ crates/vchordg/src/maintain.rs | 273 +++++++ crates/vchordg/src/operator.rs | 253 ++++++ crates/vchordg/src/prune.rs | 56 ++ crates/vchordg/src/results.rs | 73 ++ crates/vchordg/src/search.rs | 179 +++++ crates/vchordg/src/tuples.rs | 726 ++++++++++++++++++ crates/vchordg/src/types.rs | 133 ++++ crates/vchordg/src/visited.rs | 34 + crates/{algorithm => vchordrq}/Cargo.toml | 3 +- crates/{algorithm => vchordrq}/src/build.rs | 7 +- .../{algorithm => vchordrq}/src/bulkdelete.rs | 14 +- crates/{algorithm => vchordrq}/src/cache.rs | 12 +- .../src/closure_lifetime_binder.rs | 13 +- crates/{algorithm => vchordrq}/src/cost.rs | 2 +- .../{algorithm => vchordrq}/src/fast_heap.rs | 2 +- .../{algorithm => vchordrq}/src/freepages.rs | 13 +- crates/{algorithm => vchordrq}/src/insert.rs | 20 +- crates/vchordrq/src/lib.rs | 69 ++ .../{algorithm => vchordrq}/src/linked_vec.rs | 0 .../{algorithm => vchordrq}/src/maintain.rs | 66 +- .../{algorithm => vchordrq}/src/operator.rs | 429 ++--------- crates/{algorithm => vchordrq}/src/prewarm.rs | 12 +- crates/{algorithm => vchordrq}/src/rerank.rs | 26 +- crates/{algorithm => vchordrq}/src/search.rs | 20 +- crates/{algorithm => vchordrq}/src/tape.rs | 68 +- .../src/tape_writer.rs | 6 +- crates/{algorithm => vchordrq}/src/tuples.rs | 207 +---- crates/{algorithm => vchordrq}/src/types.rs | 17 +- crates/{algorithm => vchordrq}/src/vectors.rs | 18 +- crates/vector/src/bvect.rs | 2 +- crates/vector/src/lib.rs | 2 +- crates/vector/src/scalar8.rs | 2 +- crates/vector/src/svect.rs | 2 +- crates/vector/src/vect.rs | 2 +- rust-toolchain.toml | 2 +- src/datatype/binary_scalar8.rs | 2 +- src/datatype/operators_halfvec.rs | 2 +- src/datatype/operators_scalar8.rs | 4 +- src/datatype/operators_vector.rs | 2 +- src/index/algorithm.rs | 603 --------------- src/index/fetcher.rs | 192 +++++ src/index/functions.rs | 4 +- src/index/gucs.rs | 276 +++++-- src/index/hook.rs | 111 ++- src/index/mod.rs | 12 +- src/index/opclass.rs | 276 +------ src/index/storage.rs | 287 ++++--- src/index/vchordg/algo.rs | 260 +++++++ src/index/vchordg/am/am_build.rs | 704 +++++++++++++++++ src/index/vchordg/am/mod.rs | 438 +++++++++++ src/index/vchordg/mod.rs | 19 + src/index/vchordg/opclass.rs | 168 ++++ src/index/vchordg/scanners/default.rs | 501 ++++++++++++ src/index/vchordg/scanners/mod.rs | 61 ++ src/index/vchordg/types.rs | 24 + src/index/vchordrq/algo.rs | 421 ++++++++++ src/index/{ => vchordrq}/am/am_build.rs | 134 ++-- src/index/{ => vchordrq}/am/mod.rs | 222 ++---- src/index/vchordrq/mod.rs | 19 + src/index/vchordrq/opclass.rs | 234 ++++++ src/index/{ => vchordrq}/scanners/default.rs | 30 +- src/index/{ => vchordrq}/scanners/maxsim.rs | 21 +- src/index/{ => vchordrq}/scanners/mod.rs | 35 +- src/index/{ => vchordrq}/types.rs | 2 +- src/lib.rs | 3 +- src/sql/finalize.sql | 46 ++ vchord.control | 2 +- 101 files changed, 8960 insertions(+), 2780 deletions(-) create mode 100644 crates/algo/Cargo.toml create mode 100644 crates/algo/src/accessor.rs rename crates/{algorithm => algo}/src/lib.rs (50%) rename crates/{algorithm => algo}/src/prefetcher.rs (67%) create mode 100644 crates/algo/src/tuples.rs create mode 100644 crates/rabitq/src/b1.rs create mode 100644 crates/rabitq/src/b2.rs create mode 100644 crates/rabitq/src/b4.rs create mode 100644 crates/rabitq/src/b8.rs delete mode 100644 crates/rabitq/src/binary.rs delete mode 100644 crates/rabitq/src/block.rs create mode 100644 crates/rabitq/src/extended.rs create mode 100644 crates/rabitq/src/original.rs create mode 100644 crates/simd/src/packed.rs create mode 100644 crates/vchordg/Cargo.toml create mode 100644 crates/vchordg/src/build.rs create mode 100644 crates/vchordg/src/bulkdelete.rs create mode 100644 crates/vchordg/src/candidates.rs create mode 100644 crates/vchordg/src/insert.rs create mode 100644 crates/vchordg/src/lib.rs create mode 100644 crates/vchordg/src/maintain.rs create mode 100644 crates/vchordg/src/operator.rs create mode 100644 crates/vchordg/src/prune.rs create mode 100644 crates/vchordg/src/results.rs create mode 100644 crates/vchordg/src/search.rs create mode 100644 crates/vchordg/src/tuples.rs create mode 100644 crates/vchordg/src/types.rs create mode 100644 crates/vchordg/src/visited.rs rename crates/{algorithm => vchordrq}/Cargo.toml (91%) rename crates/{algorithm => vchordrq}/src/build.rs (98%) rename crates/{algorithm => vchordrq}/src/bulkdelete.rs (96%) rename crates/{algorithm => vchordrq}/src/cache.rs (88%) rename crates/{algorithm => vchordrq}/src/closure_lifetime_binder.rs (79%) rename crates/{algorithm => vchordrq}/src/cost.rs (96%) rename crates/{algorithm => vchordrq}/src/fast_heap.rs (99%) rename crates/{algorithm => vchordrq}/src/freepages.rs (90%) rename crates/{algorithm => vchordrq}/src/insert.rs (92%) create mode 100644 crates/vchordrq/src/lib.rs rename crates/{algorithm => vchordrq}/src/linked_vec.rs (100%) rename crates/{algorithm => vchordrq}/src/maintain.rs (86%) rename crates/{algorithm => vchordrq}/src/operator.rs (61%) rename crates/{algorithm => vchordrq}/src/prewarm.rs (94%) rename crates/{algorithm => vchordrq}/src/rerank.rs (85%) rename crates/{algorithm => vchordrq}/src/search.rs (97%) rename crates/{algorithm => vchordrq}/src/tape.rs (86%) rename crates/{algorithm => vchordrq}/src/tape_writer.rs (98%) rename crates/{algorithm => vchordrq}/src/tuples.rs (88%) rename crates/{algorithm => vchordrq}/src/types.rs (86%) rename crates/{algorithm => vchordrq}/src/vectors.rs (89%) delete mode 100644 src/index/algorithm.rs create mode 100644 src/index/fetcher.rs create mode 100644 src/index/vchordg/algo.rs create mode 100644 src/index/vchordg/am/am_build.rs create mode 100644 src/index/vchordg/am/mod.rs create mode 100644 src/index/vchordg/mod.rs create mode 100644 src/index/vchordg/opclass.rs create mode 100644 src/index/vchordg/scanners/default.rs create mode 100644 src/index/vchordg/scanners/mod.rs create mode 100644 src/index/vchordg/types.rs create mode 100644 src/index/vchordrq/algo.rs rename src/index/{ => vchordrq}/am/am_build.rs (93%) rename src/index/{ => vchordrq}/am/mod.rs (75%) create mode 100644 src/index/vchordrq/mod.rs create mode 100644 src/index/vchordrq/opclass.rs rename src/index/{ => vchordrq}/scanners/default.rs (97%) rename src/index/{ => vchordrq}/scanners/maxsim.rs (98%) rename src/index/{ => vchordrq}/scanners/mod.rs (77%) rename src/index/{ => vchordrq}/types.rs (98%) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 8ba0878b..3d1cd725 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -22,7 +22,7 @@ jobs: - name: Set up Environment run: | rustup set profile - rustup default nightly-2025-06-08 + rustup default 1.89-beta curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo @@ -77,21 +77,18 @@ jobs: for p in $RS_FILES; do if ! head -n "$COUNT" "$p" | cmp -s - <(echo "$RS_HEADER"); then echo "license header mismatch in file $p" - exit 1 fi done for p in $C_FILES; do if ! head -n "$COUNT" "$p" | cmp -s - <(echo "$C_HEADER"); then echo "license header mismatch in file $p" - exit 1 fi done for p in $PY_FILES; do if ! head -n "$COUNT" "$p" | cmp -s - <(echo "$PY_HEADER"); then echo "license header mismatch in file $p" - exit 1 fi done @@ -113,7 +110,7 @@ jobs: - name: Set up Environment run: | rustup set profile - rustup default nightly-2025-06-08 + rustup default 1.89-beta sudo apt-get update @@ -174,7 +171,7 @@ jobs: - name: Set up Environment run: | rustup set profile - rustup default nightly-2025-06-08 + rustup default 1.89-beta sudo apt-get update @@ -283,7 +280,7 @@ jobs: - name: Set up Environment run: | rustup set profile - rustup default nightly-2025-06-08 + rustup default 1.89-beta echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV @@ -370,7 +367,7 @@ jobs: - name: Set up Environment run: | rustup set profile - rustup default nightly-2025-06-08 + rustup default 1.89-beta Invoke-WebRequest -Uri "https://get.enterprisedb.com/postgresql/postgresql-16.9-1-windows-x64-binaries.zip" -OutFile "$env:TEMP\postgresql-install.zip" Expand-Archive -Path "$env:TEMP\postgresql-install.zip" -DestinationPath "D:\postgresql-install" -Force diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 65a37193..7d663b9a 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -52,7 +52,7 @@ jobs: - name: Set up Environment run: | rustup set profile - rustup default nightly-2025-06-08 + rustup default 1.89-beta sudo apt-get update diff --git a/Cargo.lock b/Cargo.lock index 0ff475a8..33cf64de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,9 +4,9 @@ version = 4 [[package]] name = "adler2" -version = "2.0.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "aho-corasick" @@ -18,20 +18,14 @@ dependencies = [ ] [[package]] -name = "algorithm" +name = "algo" version = "0.0.0" dependencies = [ "always_equal", + "bumpalo", "distance", "half 2.6.0", - "k_means", - "paste", - "pin-project", - "rabitq", - "rand", - "serde", "simd", - "validator", "vector", "zerocopy", ] @@ -58,9 +52,9 @@ dependencies = [ [[package]] name = "anstream" -version = "0.6.18" +version = "0.6.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8acc5369981196006228e28809f761875c0327210a891e941f4c683b3a99529b" +checksum = "301af1932e46185686725e0fad2f8f2aa7da69dd70bf6ecc44d6b703844a3933" dependencies = [ "anstyle", "anstyle-parse", @@ -73,37 +67,37 @@ dependencies = [ [[package]] name = "anstyle" -version = "1.0.10" +version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9" +checksum = "862ed96ca487e809f1c8e5a8447f6ee2cf102f846893800b20cebdf541fc6bbd" [[package]] name = "anstyle-parse" -version = "0.2.6" +version = "0.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b2d16507662817a6a20a9ea92df6652ee4f94f914589377d69f3b21bc5798a9" +checksum = "4e7644824f0aa2c7b9384579234ef10eb7efb6a0deb83f9630a49594dd9c15c2" dependencies = [ "utf8parse", ] [[package]] name = "anstyle-query" -version = "1.1.2" +version = "1.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79947af37f4177cfead1110013d678905c37501914fba0efea834c3fe9a8d60c" +checksum = "6c8bdeb6047d8983be085bab0ba1472e6dc604e7041dbf6fcd5e71523014fae9" dependencies = [ - "windows-sys", + "windows-sys 0.59.0", ] [[package]] name = "anstyle-wincon" -version = "3.0.8" +version = "3.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6680de5231bd6ee4c6191b8a1325daa282b415391ec9d3a37bd34f2060dc73fa" +checksum = "403f75924867bb1033c59fbf0797484329750cfbe3c4325cd33127941fabc882" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -143,9 +137,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.9.0" +version = "2.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c8214115b7bf84099f1309324e63141d4c5d7cc26862f97a0a857dbefe165bd" +checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967" [[package]] name = "bitvec" @@ -159,6 +153,12 @@ dependencies = [ "wyz", ] +[[package]] +name = "bumpalo" +version = "3.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43" + [[package]] name = "byteorder" version = "1.5.0" @@ -177,9 +177,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.22" +version = "1.2.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32db95edf998450acc7881c932f94cd9b05c87b4b2599e8bab064753da4acfd1" +checksum = "d487aa071b5f64da6f19a3e848e3578944b726ee5a4854b82172f02aa876bfdc" dependencies = [ "shlex", ] @@ -205,9 +205,9 @@ dependencies = [ [[package]] name = "cfg-if" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268" [[package]] name = "clang-sys" @@ -222,9 +222,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.38" +version = "4.5.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed93b9805f8ba930df42c2590f05453d5ec36cbb85d018868a5b24d31f6ac000" +checksum = "40b6887a1d8685cebccf115538db5c0efe625ccac9696ad45c409d96566e910f" dependencies = [ "clap_builder", "clap_derive", @@ -232,9 +232,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.38" +version = "4.5.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "379026ff283facf611b0ea629334361c4211d1b12ee01024eec1591133b04120" +checksum = "e0c66c08ce9f0c698cbce5c0279d0bb6ac936d8674174fe48f736533b964f59e" dependencies = [ "anstream", "anstyle", @@ -244,9 +244,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.5.32" +version = "4.5.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09176aae279615badda0765c0c0b3f6ed53f4709118af73cf4655d85d1530cd7" +checksum = "d2c7947ae4cc3d851207c1adb5b5e260ff0cca11446b1d6d1423788e442257ce" dependencies = [ "heck", "proc-macro2", @@ -256,15 +256,24 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.7.4" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" + +[[package]] +name = "codepage" +version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6" +checksum = "48f68d061bc2828ae826206326e61251aca94c1e4a5305cf52d9138639c918b4" +dependencies = [ + "encoding_rs", +] [[package]] name = "colorchoice" -version = "1.0.3" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b63caa9aa9397e2d9480a9b13673856c78d8ac123288526c37d7839f2a86990" +checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" [[package]] name = "convert_case" @@ -311,9 +320,9 @@ checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] name = "crunchy" -version = "0.2.3" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43da5946c66ffcc7745f48db692ffbb10a83bfe0afd96235c5c2a4fb23994929" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "darling" @@ -364,6 +373,9 @@ dependencies = [ [[package]] name = "distance" version = "0.0.0" +dependencies = [ + "zerocopy", +] [[package]] name = "either" @@ -371,6 +383,15 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "encoding_rs" +version = "0.8.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" +dependencies = [ + "cfg-if", +] + [[package]] name = "enum-map" version = "2.7.3" @@ -405,12 +426,12 @@ checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "errno" -version = "0.3.12" +version = "0.3.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cea14ef9355e3beab063703aa9dab15afd25f0667c341310c1e5274bb1d0da18" +checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.60.2", ] [[package]] @@ -437,9 +458,9 @@ checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" [[package]] name = "flate2" -version = "1.1.1" +version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ced92e76e966ca2fd84c8f7aa01a4aea65b0eb6648d72f7c8f3e2764a67fece" +checksum = "4a3d7db9596fecd151c5f638c0ee5d5bd487b6e0ea232e5dc96d5250f6f94b1d" dependencies = [ "crc32fast", "miniz_oxide", @@ -518,9 +539,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.15.3" +version = "0.15.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "84b26c544d002229e640969970a2e74021aadf6e2f96372b9c58eff97de08eb3" +checksum = "5971ac85611da7067dbfcabef3c70ebb5606018acd9e2a3903a0da507521e0d5" dependencies = [ "allocator-api2", "equivalent", @@ -545,9 +566,9 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hermit-abi" -version = "0.5.1" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f154ce46856750ed433c8649605bf7ed2de3bc35fd9d2a9f30cddd873c80cb08" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" [[package]] name = "home" @@ -555,7 +576,7 @@ version = "0.5.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589533453244b0995c858700322199b2becb13b627df2851f64a2775d024abcf" dependencies = [ - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -607,9 +628,9 @@ checksum = "00210d6893afc98edb752b664b8890f0ef174c8adbb8d0be9710fa66fbbf72d3" [[package]] name = "icu_properties" -version = "2.0.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2549ca8c7241c82f59c80ba2a6f415d931c5b58d24fb8412caa1a1f02c49139a" +checksum = "016c619c1eeb94efb86809b015c58f479963de65bdb6253345c1a1276f22e32b" dependencies = [ "displaydoc", "icu_collections", @@ -623,9 +644,9 @@ dependencies = [ [[package]] name = "icu_properties_data" -version = "2.0.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8197e866e47b68f8f7d95249e172903bec06004b18b2937f1095d40a0c57de04" +checksum = "298459143998310acd25ffe6810ed544932242d3f07083eee1084d83a71bd632" [[package]] name = "icu_provider" @@ -679,9 +700,9 @@ checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" [[package]] name = "indexmap" -version = "2.9.0" +version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cea70ddb795996207ad57735b50c5982d8844f38ba9ee5f1aedcfb708a2aa11e" +checksum = "fe4cd85333e22411419a0bcae1297d25e58c9443848b11dc6a86fefe8c78a661" dependencies = [ "equivalent", "hashbrown", @@ -695,7 +716,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -745,6 +766,16 @@ dependencies = [ "libc", ] +[[package]] +name = "js-sys" +version = "0.3.77" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1cfaf33c695fc6e08064efbc1f72ec937429614f25eef83af942d0e227c3a28f" +dependencies = [ + "once_cell", + "wasm-bindgen", +] + [[package]] name = "k_means" version = "0.0.0" @@ -757,18 +788,18 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.172" +version = "0.2.174" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d750af042f7ef4f724306de029d18836c26c1765a54a6a3f094cbd23a7267ffa" +checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776" [[package]] name = "libloading" -version = "0.8.7" +version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a793df0d7afeac54f95b471d3af7f0d4fb975699f972341a4b76988d49cdf0c" +checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.53.0", + "windows-targets 0.53.2", ] [[package]] @@ -793,6 +824,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956" +[[package]] +name = "log" +version = "0.4.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" + [[package]] name = "make" version = "0.0.0" @@ -806,9 +843,9 @@ dependencies = [ [[package]] name = "memchr" -version = "2.7.4" +version = "2.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" [[package]] name = "mimalloc" @@ -819,6 +856,12 @@ dependencies = [ "libmimalloc-sys", ] +[[package]] +name = "min-max-heap" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2687e6cf9c00f48e9284cf9fd15f2ef341d03cc7743abf9df4c5f07fdee50b18" + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -827,9 +870,9 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" [[package]] name = "miniz_oxide" -version = "0.8.8" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3be647b768db090acb35d5ec5db2b0e1f1de11133ca123b9eacf5137868f892a" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" dependencies = [ "adler2", ] @@ -873,9 +916,9 @@ checksum = "a4895175b425cb1f87721b59f0f286c2092bd4af812243672510e1ac53e2e0ad" [[package]] name = "owo-colors" -version = "4.2.0" +version = "4.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1036865bb9422d3300cf723f657c2851d0e9ab12567854b1f4eba3d77decf564" +checksum = "48dd4f4a2c8405440fd0462561f0e5806bd0f77e86f51c761481bdd4018b545e" dependencies = [ "supports-color 2.1.0", "supports-color 3.0.2", @@ -905,9 +948,9 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "petgraph" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a98c6720655620a521dcc722d0ad66cd8afd5d86e34a89ef691c50b7b24de06" +checksum = "54acf3a685220b533e437e264e4d932cfbdc4cc7ec0cd232ed73c08d03b8a7ca" dependencies = [ "fixedbitset", "hashbrown", @@ -917,9 +960,9 @@ dependencies = [ [[package]] name = "pgrx" -version = "0.14.1" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab5a7f1a5bac3a76a4fa056655eb80e394c5216717bf4673c99d9d10583e960d" +checksum = "bab5bc1d60d3bc3c966d307a3c7313b1ebfb49a0ec183be3f1a057df0bcc9988" dependencies = [ "atomic-traits", "bitflags", @@ -941,9 +984,9 @@ dependencies = [ [[package]] name = "pgrx-bindgen" -version = "0.14.1" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09e1a6df74e60ac82d6ad52f8a2a85e5d7004f9dfc7c52c94d461d9193aafa0b" +checksum = "9804b74c211a9edd550cd974718f8cc407dec50d8e9cafb906e0b042ba434af0" dependencies = [ "bindgen", "cc", @@ -960,9 +1003,9 @@ dependencies = [ [[package]] name = "pgrx-catalog" -version = "0.2.0" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54c52c31fd64fe1e0cde3826b578bcc343e778e874fc2e520ad324659ff0b120" +checksum = "fad3b930e3f5abd77c689fd9f64dba3bbb97993c493700daa4c1242c27fc141c" dependencies = [ "paste", "pgrx", @@ -970,9 +1013,9 @@ dependencies = [ [[package]] name = "pgrx-macros" -version = "0.14.1" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "868ae3cb1762dcca0e236d4b5368315d4f377acbfff31d8fad9f47394d2e4bcd" +checksum = "f230769493bf567f137de23264d604d267dd72b8a77c596528e43cf423c6208e" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", @@ -982,11 +1025,13 @@ dependencies = [ [[package]] name = "pgrx-pg-config" -version = "0.14.1" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "421be70840655a502381d4ac84b263da83c6714018e27bfdcd7c48f32ff60ef6" +checksum = "49b64c071c2a46a19ab4521120a25b02b598f4abf6e9b4b1769a7922edeee3de" dependencies = [ "cargo_toml", + "codepage", + "encoding_rs", "eyre", "home", "owo-colors", @@ -996,13 +1041,14 @@ dependencies = [ "thiserror", "toml", "url", + "winapi", ] [[package]] name = "pgrx-pg-sys" -version = "0.14.1" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c844a49ee575cf4bd3340c483e49160a99b1bd8c5fd2e87fc7a9231566cb07f2" +checksum = "fcbfa98ec7a90252d13a78ac666541173dbb01a2fc1ba20131db6490c0711125" dependencies = [ "cee-scape", "libc", @@ -1015,9 +1061,9 @@ dependencies = [ [[package]] name = "pgrx-sql-entity-graph" -version = "0.14.1" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "391595ab6d13f65728bdd691fdd58e69bb31f1f092886ed0a11c838de1d01ae9" +checksum = "e79bbf5a33cff6cfdc6dda3a976cd931c995eaa2c073a7c59b8f8fe8f6faa073" dependencies = [ "convert_case", "eyre", @@ -1109,9 +1155,9 @@ dependencies = [ [[package]] name = "r-efi" -version = "5.2.0" +version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74765f6d916ee2faa39bc8e68e4f3ed8949b48cccdac59983d287a7cb71ce9c5" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" [[package]] name = "rabitq" @@ -1233,9 +1279,15 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys", + "windows-sys 0.59.0", ] +[[package]] +name = "rustversion" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a0d197bd2c9dc6e53b84da9556a69ba4cdfab8619eb41a8bd1cc2027a0f6b1d" + [[package]] name = "ruzstd" version = "0.7.3" @@ -1359,9 +1411,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.15.0" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8917285742e9f3e1683f0a9c4e6b57960b7314d0b08d30d1ecd426713ee2eee9" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" [[package]] name = "sptr" @@ -1408,9 +1460,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.101" +version = "2.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ce2b7fc941b3a24138a0a7cf8e858bfc6a992e7978a068a5c760deb0ed43caf" +checksum = "17b6f705963418cdb9927482fa304bc562ece2fdd4f616084c50b7023b435a40" dependencies = [ "proc-macro2", "quote", @@ -1444,7 +1496,7 @@ dependencies = [ "getrandom", "once_cell", "rustix", - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -1548,9 +1600,9 @@ checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" [[package]] name = "unicode-width" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd" +checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" [[package]] name = "url" @@ -1577,11 +1629,13 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.16.0" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "458f7a779bf54acc9f347480ac654f68407d3aab21269a6e3c9f922acd9e2da9" +checksum = "3cf4199d1e5d15ddd86a694e4d0dffa9c323ce759fea589f00fef9d81cc1931d" dependencies = [ "getrandom", + "js-sys", + "wasm-bindgen", ] [[package]] @@ -1618,8 +1672,9 @@ dependencies = [ name = "vchord" version = "0.0.0" dependencies = [ - "algorithm", + "algo", "always_equal", + "bumpalo", "distance", "half 2.6.0", "jemallocator", @@ -1635,6 +1690,49 @@ dependencies = [ "simd", "toml", "validator", + "vchordg", + "vchordrq", + "vector", + "zerocopy", +] + +[[package]] +name = "vchordg" +version = "0.0.0" +dependencies = [ + "algo", + "always_equal", + "distance", + "half 2.6.0", + "k_means", + "min-max-heap", + "paste", + "pin-project", + "rabitq", + "rand", + "serde", + "simd", + "validator", + "vector", + "zerocopy", +] + +[[package]] +name = "vchordrq" +version = "0.0.0" +dependencies = [ + "algo", + "always_equal", + "distance", + "half 2.6.0", + "k_means", + "paste", + "pin-project", + "rabitq", + "rand", + "serde", + "simd", + "validator", "vector", "zerocopy", ] @@ -1667,6 +1765,64 @@ dependencies = [ "wit-bindgen-rt", ] +[[package]] +name = "wasm-bindgen" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1edc8929d7499fc4e8f0be2262a241556cfc54a0bea223790e71446f2aab1ef5" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", +] + +[[package]] +name = "wasm-bindgen-backend" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f0a0651a5c2bc21487bde11ee802ccaf4c51935d0d3d42a6101f98161700bc6" +dependencies = [ + "bumpalo", + "log", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fe63fc6d09ed3792bd0897b314f53de8e16568c2b3f7982f468c0bf9bd0b407" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-backend", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a05d73b933a847d6cccdda8f838a22ff101ad9bf93e33684f39c1f5f0eece3d" +dependencies = [ + "unicode-ident", +] + [[package]] name = "wasmparser" version = "0.222.1" @@ -1688,15 +1844,37 @@ dependencies = [ "winsafe", ] +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + [[package]] name = "winapi-util" version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys", + "windows-sys 0.59.0", ] +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-sys" version = "0.59.0" @@ -1706,6 +1884,15 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.2", +] + [[package]] name = "windows-targets" version = "0.52.6" @@ -1724,9 +1911,9 @@ dependencies = [ [[package]] name = "windows-targets" -version = "0.53.0" +version = "0.53.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1e4c7e8ceaaf9cb7d7507c974735728ab453b67ef8f18febdd7c11fe59dca8b" +checksum = "c66f69fcc9ce11da9966ddb31a40968cad001c5bedeb5c2b82ede4253ab48aef" dependencies = [ "windows_aarch64_gnullvm 0.53.0", "windows_aarch64_msvc 0.53.0", @@ -1836,9 +2023,9 @@ checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486" [[package]] name = "winnow" -version = "0.7.10" +version = "0.7.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c06928c8748d81b05c9be96aad92e1b6ff01833332f281e8cfca3be4b35fc9ec" +checksum = "74c7b26e3480b707944fc872477815d29a8e429d2f93a1ce000f5fa84a15cbcd" dependencies = [ "memchr", ] @@ -1899,18 +2086,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.25" +version = "0.8.26" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1702d9583232ddb9174e01bb7c15a2ab8fb1bc6f227aa1233858c351a3ba0cb" +checksum = "1039dd0d3c310cf05de012d8a39ff557cb0d23087fd44cad61df08fc31907a2f" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.25" +version = "0.8.26" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28a6e20d751156648aa063f3800b706ee209a32c0b4d9f24be3d980b01be55ef" +checksum = "9ecf5b4cc5364572d7f4c329661bcc82724222973f2cab6f050a4e5c22f75181" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index 1072e19f..61e7071a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,20 +19,24 @@ pg14 = ["pgrx/pg14", "pgrx-catalog/pg14"] pg15 = ["pgrx/pg15", "pgrx-catalog/pg15"] pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] +pg18 = ["pgrx/pg18", "pgrx-catalog/pg18"] [dependencies] -algorithm = { path = "./crates/algorithm" } +algo = { path = "./crates/algo" } always_equal = { path = "./crates/always_equal" } distance = { path = "./crates/distance" } k_means = { path = "./crates/k_means" } rabitq = { path = "./crates/rabitq" } simd = { path = "./crates/simd" } +vchordg = { path = "./crates/vchordg" } +vchordrq = { path = "./crates/vchordrq" } vector = { path = "./crates/vector" } +bumpalo.workspace = true half.workspace = true paste.workspace = true -pgrx = { version = "=0.14.1", default-features = false, features = ["cshim"] } -pgrx-catalog = "0.2.0" +pgrx = { version = "=0.15.0", default-features = false, features = ["cshim"] } +pgrx-catalog = "0.3.0" rand.workspace = true seq-macro.workspace = true serde.workspace = true @@ -58,13 +62,14 @@ version = "0.0.0" edition = "2024" [workspace.dependencies] +bumpalo = "3.19.0" half = { version = "2.6.0", features = ["zerocopy"] } paste = "1" rand = "0.9.1" seq-macro = "0.3.6" serde = { version = "1", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } -zerocopy = { version = "0.8.25", features = ["derive"] } +zerocopy = { version = "0.8.26", features = ["derive"] } [workspace.lints] clippy.identity_op = "allow" diff --git a/crates/algo/Cargo.toml b/crates/algo/Cargo.toml new file mode 100644 index 00000000..44ce2ed5 --- /dev/null +++ b/crates/algo/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "algo" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +always_equal = { path = "../always_equal" } +bumpalo.workspace = true +distance = { path = "../distance" } +half.workspace = true +simd = { path = "../simd" } +vector = { path = "../vector" } +zerocopy.workspace = true + +[lints] +workspace = true diff --git a/crates/algo/src/accessor.rs b/crates/algo/src/accessor.rs new file mode 100644 index 00000000..bf94e8b9 --- /dev/null +++ b/crates/algo/src/accessor.rs @@ -0,0 +1,335 @@ +use distance::Distance; +use half::f16; +use simd::Floating; +use std::marker::PhantomData; +use vector::vect::VectOwned; + +#[derive(Debug, Clone, Copy)] +pub struct L2S; + +#[derive(Debug, Clone, Copy)] +pub struct Dot; + +pub trait Accessor2 { + type Output; + fn push(&mut self, input: &[E0], target: &[E1]); + fn finish(self, input: M0, target: M1) -> Self::Output; +} + +impl Accessor2 for () { + type Output = (); + + fn push(&mut self, _: &[E0], _: &[E1]) {} + + fn finish(self, _: M0, _: M1) -> Self::Output {} +} + +impl> Accessor2 for (A,) { + type Output = (A::Output,); + + fn push(&mut self, input: &[E0], target: &[E1]) { + self.0.push(input, target); + } + + fn finish(self, input: M0, target: M1) -> Self::Output { + (self.0.finish(input, target),) + } +} + +impl, B: Accessor2> + Accessor2 for (A, B) +{ + type Output = (A::Output, B::Output); + + fn push(&mut self, input: &[E0], target: &[E1]) { + self.0.push(input, target); + self.1.push(input, target); + } + + fn finish(self, input: M0, target: M1) -> Self::Output { + (self.0.finish(input, target), self.1.finish(input, target)) + } +} + +pub trait Accessor1 { + type Output; + fn push(&mut self, input: &[E]); + fn finish(self, input: M) -> Self::Output; +} + +impl Accessor1 for () { + type Output = (); + + fn push(&mut self, _: &[E]) {} + + fn finish(self, _: M) -> Self::Output {} +} + +impl Accessor1 for (A,) +where + A: Accessor1, +{ + type Output = (A::Output,); + + fn push(&mut self, input: &[E]) { + self.0.push(input); + } + + fn finish(self, input: M) -> Self::Output { + (self.0.finish(input),) + } +} + +impl Accessor1 for (A, B) +where + A: Accessor1, + B: Accessor1, +{ + type Output = (A::Output, B::Output); + + fn push(&mut self, input: &[E]) { + self.0.push(input); + self.1.push(input); + } + + fn finish(self, input: M) -> Self::Output { + (self.0.finish(input), self.1.finish(input)) + } +} + +pub struct FunctionalAccessor { + data: T, + p: P, + f: F, +} + +impl FunctionalAccessor { + pub fn new(data: T, p: P, f: F) -> Self { + Self { data, p, f } + } +} + +impl Accessor1 for FunctionalAccessor +where + P: for<'a> FnMut(&'a mut T, &'a [E]), + F: FnOnce(T, M) -> R, +{ + type Output = R; + + fn push(&mut self, input: &[E]) { + (self.p)(&mut self.data, input); + } + + fn finish(self, input: M) -> Self::Output { + (self.f)(self.data, input) + } +} + +pub struct LAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> LAccess<'a, E, M, A> { + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } + } +} + +impl> Accessor1 for LAccess<'_, E0, M0, A> { + type Output = A::Output; + + fn push(&mut self, rhs: &[E1]) { + let (lhs, elements) = self.elements.split_at(rhs.len()); + self.accessor.push(lhs, rhs); + self.elements = elements; + } + + fn finish(self, rhs: M1) -> Self::Output { + assert!(self.elements.is_empty(), "goal is shorter than expected"); + self.accessor.finish(self.metadata, rhs) + } +} + +pub struct RAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> RAccess<'a, E, M, A> { + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } + } +} + +impl> Accessor1 for RAccess<'_, E1, M1, A> { + type Output = A::Output; + + fn push(&mut self, lhs: &[E0]) { + let (rhs, elements) = self.elements.split_at(lhs.len()); + self.accessor.push(lhs, rhs); + self.elements = elements; + } + + fn finish(self, lhs: M0) -> Self::Output { + assert!(self.elements.is_empty(), "goal is shorter than expected"); + self.accessor.finish(lhs, self.metadata) + } +} + +pub trait TryAccessor1: Sized { + type Output; + #[must_use] + fn push(&mut self, input: &[E]) -> Option<()>; + #[must_use] + fn finish(self, input: M) -> Option; +} + +impl TryAccessor1 for () { + type Output = (); + + fn push(&mut self, _: &[E]) -> Option<()> { + Some(()) + } + + fn finish(self, _: M) -> Option { + Some(()) + } +} + +impl TryAccessor1 for (A,) +where + A: TryAccessor1, +{ + type Output = (A::Output,); + + fn push(&mut self, input: &[E]) -> Option<()> { + self.0.push(input)?; + Some(()) + } + + fn finish(self, input: M) -> Option { + Some((self.0.finish(input)?,)) + } +} + +impl TryAccessor1 for (A, B) +where + A: TryAccessor1, + B: TryAccessor1, +{ + type Output = (A::Output, B::Output); + + fn push(&mut self, input: &[E]) -> Option<()> { + self.0.push(input)?; + self.1.push(input)?; + Some(()) + } + + fn finish(self, input: M) -> Option { + Some((self.0.finish(input)?, self.1.finish(input)?)) + } +} + +pub struct LTryAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> LTryAccess<'a, E, M, A> { + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } + } +} + +impl> TryAccessor1 + for LTryAccess<'_, E0, M0, A> +{ + type Output = A::Output; + + fn push(&mut self, rhs: &[E1]) -> Option<()> { + let (lhs, elements) = self.elements.split_at_checked(rhs.len())?; + self.accessor.push(lhs, rhs); + self.elements = elements; + Some(()) + } + + fn finish(self, rhs: M1) -> Option { + if !self.elements.is_empty() { + return None; + } + Some(self.accessor.finish(self.metadata, rhs)) + } +} + +#[derive(Debug)] +pub struct DistanceAccessor(f32, PhantomData V>, PhantomData D>); + +impl Default for DistanceAccessor { + fn default() -> Self { + Self(0.0, PhantomData, PhantomData) + } +} + +impl Accessor2 for DistanceAccessor, L2S> { + type Output = Distance; + + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0 += f32::reduce_sum_of_d2(target, input) + } + + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(self.0) + } +} + +impl Accessor2 for DistanceAccessor, Dot> { + type Output = Distance; + + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0 += f32::reduce_sum_of_xy(target, input) + } + + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(-self.0) + } +} + +impl Accessor2 for DistanceAccessor, L2S> { + type Output = Distance; + + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0 += f16::reduce_sum_of_d2(target, input) + } + + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(self.0) + } +} + +impl Accessor2 for DistanceAccessor, Dot> { + type Output = Distance; + + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0 += f16::reduce_sum_of_xy(target, input) + } + + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(-self.0) + } +} diff --git a/crates/algorithm/src/lib.rs b/crates/algo/src/lib.rs similarity index 50% rename from crates/algorithm/src/lib.rs rename to crates/algo/src/lib.rs index 68602551..8e5de051 100644 --- a/crates/algorithm/src/lib.rs +++ b/crates/algo/src/lib.rs @@ -14,58 +14,23 @@ #![allow(clippy::type_complexity)] -mod build; -mod bulkdelete; -mod cache; -mod closure_lifetime_binder; -mod cost; -mod fast_heap; -mod freepages; -mod insert; -mod linked_vec; -mod maintain; -mod prefetcher; -mod prewarm; -mod rerank; -mod search; -mod tape; -mod tape_writer; -mod tuples; -mod vectors; - -pub mod operator; -pub mod types; +pub mod accessor; +pub mod prefetcher; +pub mod tuples; use always_equal::AlwaysEqual; -pub use build::build; -pub use bulkdelete::{bulkdelete, bulkdelete_vectors}; -pub use cache::cache; -pub use cost::cost; -pub use fast_heap::FastHeap; -pub use insert::{insert, insert_vector}; -pub use maintain::maintain; -pub use prefetcher::*; -pub use prewarm::prewarm; -pub use rerank::{how, rerank_heap, rerank_index}; -pub use search::{default_search, maxsim_search}; - -use std::collections::BinaryHeap; +use std::collections::{BinaryHeap, VecDeque}; use std::iter::Peekable; use std::ops::{Deref, DerefMut}; -use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; +use zerocopy::{FromBytes, IntoBytes}; -#[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] -pub struct Opaque { - pub next: u32, - pub skip: u32, -} +pub trait Page: Sized + 'static { + type Opaque: Opaque; -pub trait Page: Sized { #[must_use] - fn get_opaque(&self) -> &Opaque; + fn get_opaque(&self) -> &Self::Opaque; #[must_use] - fn get_opaque_mut(&mut self) -> &mut Opaque; + fn get_opaque_mut(&mut self) -> &mut Self::Opaque; #[must_use] fn len(&self) -> u16; #[must_use] @@ -81,44 +46,23 @@ pub trait Page: Sized { fn free(&mut self, i: u16); #[must_use] fn freespace(&self) -> u16; - fn clear(&mut self); + fn clear(&mut self, opaque: Self::Opaque); } pub trait PageGuard { fn id(&self) -> u32; } -pub trait ReadStream<'s> { - type Relation: RelationRead; - type Item; - type Inner: Iterator; - fn next( - &mut self, - ) -> Option<( - Self::Item, - Vec<::ReadGuard<'s>>, - )>; - fn next_if( - &mut self, - predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<( - Self::Item, - Vec<::ReadGuard<'s>>, - )>; - fn into_inner(self) -> Self::Inner; -} - -pub trait WriteStream<'s> { - type Relation: RelationWrite; +pub trait ReadStream<'r> { + type Relation: RelationReadTypes; + type Guards: Iterator::ReadGuard<'r>>; type Item; type Inner: Iterator; - fn next_if( + fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; + fn next_if bool>( &mut self, - predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<( - Self::Item, - Vec<::WriteGuard<'s>>, - )>; + predicate: P, + ) -> Option<(Self::Item, Self::Guards)>; fn into_inner(self) -> Self::Inner; } @@ -126,22 +70,35 @@ pub trait Relation { type Page: Page; } -pub trait RelationRead: Relation { - type ReadGuard<'a>: PageGuard + Deref - where - Self: 'a; +pub trait RelationReadTypes: Relation { + type ReadGuard<'r>: PageGuard + Deref; +} + +pub trait RelationRead: RelationReadTypes { fn read(&self, id: u32) -> Self::ReadGuard<'_>; } -pub trait RelationWrite: Relation { - type WriteGuard<'a>: PageGuard + DerefMut - where - Self: 'a; +pub trait RelationWriteTypes: Relation { + type WriteGuard<'r>: PageGuard + DerefMut; +} + +pub trait RelationWrite: RelationWriteTypes { fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_>; - fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_>; + fn extend( + &self, + opaque: ::Opaque, + tracking_freespace: bool, + ) -> Self::WriteGuard<'_>; fn search(&self, freespace: usize) -> Option>; } +pub trait RelationLength: Relation { + fn len(&self) -> u32; + fn is_empty(&self) -> bool { + self.len() == 0 + } +} + pub trait RelationPrefetch: Relation { fn prefetch(&self, id: u32); } @@ -158,22 +115,14 @@ impl Hints { } } -pub trait RelationReadStream: RelationRead { - type ReadStream<'s, I: Iterator>: ReadStream<'s, Item = I::Item, Relation = Self> - where - I::Item: Fetch, - Self: 's; - fn read_stream(&self, iter: I, hints: Hints) -> Self::ReadStream<'_, I> +pub trait RelationReadStreamTypes: RelationReadTypes { + type ReadStream<'r, I: Iterator>: ReadStream<'r, Item = I::Item, Relation = Self> where I::Item: Fetch; } -pub trait RelationWriteStream: RelationWrite { - type WriteStream<'s, I: Iterator>: WriteStream<'s, Item = I::Item, Relation = Self> - where - I::Item: Fetch, - Self: 's; - fn write_stream(&self, iter: I, hints: Hints) -> Self::WriteStream<'_, I> +pub trait RelationReadStream: RelationReadStreamTypes { + fn read_stream<'r, I: Iterator>(&'r self, iter: I, hints: Hints) -> Self::ReadStream<'r, I> where I::Item: Fetch; } @@ -184,43 +133,85 @@ pub enum RerankMethod { Heap, } -pub(crate) struct Branch { - pub code: rabitq::Code, - pub delta: f32, - pub prefetch: Vec, - pub head: u16, - pub norm: f32, - pub extra: T, -} - pub trait Bump: 'static { #[allow(clippy::mut_from_ref)] fn alloc(&self, value: T) -> &mut T; #[allow(clippy::mut_from_ref)] fn alloc_slice(&self, slice: &[T]) -> &mut [T]; + fn reset(&mut self); +} + +impl Bump for bumpalo::Bump { + fn alloc(&self, value: T) -> &mut T { + self.alloc(value) + } + + fn alloc_slice(&self, slice: &[T]) -> &mut [T] { + self.alloc_slice_copy(slice) + } + + fn reset(&mut self) { + self.reset(); + } } pub trait Fetch { - fn fetch(&self) -> &[u32]; + fn fetch(&self) -> Vec; } impl Fetch for u32 { - fn fetch(&self) -> &[u32] { - std::slice::from_ref(self) + fn fetch(&self) -> Vec { + std::slice::from_ref(self).to_vec() } } impl<'b, T, A, B> Fetch for (T, AlwaysEqual<&'b mut (A, B, &'b mut [u32])>) { - fn fetch(&self) -> &[u32] { + fn fetch(&self) -> Vec { let (_, AlwaysEqual((.., list))) = self; - list + list.to_vec() } } impl<'b, T, A, B, C> Fetch for (T, AlwaysEqual<&'b mut (A, B, C, &'b mut [u32])>) { - fn fetch(&self) -> &[u32] { + fn fetch(&self) -> Vec { let (_, AlwaysEqual((.., list))) = self; - list + list.to_vec() + } +} + +impl Fetch for (T, AlwaysEqual<(u32, u16)>) { + fn fetch(&self) -> Vec { + let (_, AlwaysEqual((x, _))) = self; + std::slice::from_ref(x).to_vec() + } +} + +impl Fetch for (u32, u16) { + fn fetch(&self) -> Vec { + std::slice::from_ref(&self.0).to_vec() + } +} + +impl Fetch for (T, AlwaysEqual<((u32, u16), (u32, u16))>) { + fn fetch(&self) -> Vec { + let (_, AlwaysEqual(((x, _), _))) = self; + std::slice::from_ref(x).to_vec() + } +} + +pub trait Fetch1 { + fn fetch_1(&self) -> u32; +} + +impl Fetch for (T, AlwaysEqual<&mut [F]>) { + fn fetch(&self) -> Vec { + self.1.0.iter().map(|x| x.fetch_1()).collect() + } +} + +impl Fetch for (T, AlwaysEqual<(&mut [F], U)>) { + fn fetch(&self) -> Vec { + self.1.0.0.iter().map(|x| x.fetch_1()).collect() } } @@ -262,3 +253,23 @@ impl Sequence for Peekable { self } } + +impl Sequence for VecDeque { + type Item = T; + type Inner = std::collections::vec_deque::IntoIter; + fn next(&mut self) -> Option { + self.pop_front() + } + fn peek(&mut self) -> Option<&T> { + self.front() + } + fn into_inner(self) -> Self::Inner { + self.into_iter() + } +} + +/// # Safety +/// +/// * `Opaque` must aligned to 8 bytes. +#[allow(unsafe_code)] +pub unsafe trait Opaque: Copy + Send + Sync + FromBytes + IntoBytes + 'static {} diff --git a/crates/algorithm/src/prefetcher.rs b/crates/algo/src/prefetcher.rs similarity index 67% rename from crates/algorithm/src/prefetcher.rs rename to crates/algo/src/prefetcher.rs index a5510e3d..d0eb5915 100644 --- a/crates/algorithm/src/prefetcher.rs +++ b/crates/algo/src/prefetcher.rs @@ -12,7 +12,10 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::*; +use crate::{ + Fetch, Hints, ReadStream, RelationPrefetch, RelationRead, RelationReadStream, + RelationReadTypes, Sequence, +}; use std::collections::{VecDeque, vec_deque}; use std::iter::Chain; @@ -24,14 +27,15 @@ where Self::Item: Fetch, { type R: RelationRead + 'r; + type Guards: Iterator::ReadGuard<'r>>; #[must_use] - fn next(&mut self) -> Option<(Self::Item, Vec<::ReadGuard<'r>>)>; + fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; #[must_use] fn next_if( &mut self, predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<(Self::Item, Vec<::ReadGuard<'r>>)>; + ) -> Option<(Self::Item, Self::Guards)>; } pub struct PlainPrefetcher<'r, R, S: Sequence> { @@ -60,22 +64,50 @@ where S::Item: Fetch, { type R = R; + type Guards = PlainPrefetcherGuards<'r, R>; - fn next(&mut self) -> Option<(Self::Item, Vec<::ReadGuard<'r>>)> { + fn next(&mut self) -> Option<(Self::Item, PlainPrefetcherGuards<'r, R>)> { let e = self.sequence.next()?; - let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); - Some((e, list)) + let list = e.fetch().into_iter(); + Some(( + e, + PlainPrefetcherGuards { + relation: self.relation, + list, + }, + )) } + fn next_if( &mut self, predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<(S::Item, Vec>)> { + ) -> Option<(S::Item, PlainPrefetcherGuards<'r, R>)> { if !predicate(self.sequence.peek()?) { return None; } let e = self.sequence.next()?; - let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); - Some((e, list)) + let list = e.fetch().into_iter(); + Some(( + e, + PlainPrefetcherGuards { + relation: self.relation, + list, + }, + )) + } +} + +pub struct PlainPrefetcherGuards<'r, R> { + relation: &'r R, + list: std::vec::IntoIter, +} + +impl<'r, R: RelationRead> Iterator for PlainPrefetcherGuards<'r, R> { + type Item = R::ReadGuard<'r>; + + fn next(&mut self) -> Option { + let id = self.list.next()?; + Some(self.relation.read(id)) } } @@ -111,35 +143,62 @@ where S::Item: Fetch, { type R = R; + type Guards = SimplePrefetcherGuards<'r, R>; - fn next(&mut self) -> Option<(Self::Item, Vec<::ReadGuard<'r>>)> { + fn next(&mut self) -> Option<(Self::Item, SimplePrefetcherGuards<'r, R>)> { while self.window.len() < WINDOW_SIZE && let Some(e) = self.sequence.next() { - for id in e.fetch().iter().copied() { + for id in e.fetch().into_iter() { self.relation.prefetch(id); } self.window.push_back(e); } let e = self.window.pop_front()?; - let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); - Some((e, list)) + let list = e.fetch().into_iter(); + Some(( + e, + SimplePrefetcherGuards { + relation: self.relation, + list, + }, + )) } fn next_if( &mut self, predicate: impl FnOnce(&S::Item) -> bool, - ) -> Option<(S::Item, Vec>)> { + ) -> Option<(S::Item, SimplePrefetcherGuards<'r, R>)> { while self.window.len() < WINDOW_SIZE && let Some(e) = self.sequence.next() { - for id in e.fetch().iter().copied() { + for id in e.fetch().into_iter() { self.relation.prefetch(id); } self.window.push_back(e); } let e = vec_deque_pop_front_if(&mut self.window, predicate)?; - let list = e.fetch().iter().map(|&id| self.relation.read(id)).collect(); - Some((e, list)) + let list = e.fetch().into_iter(); + Some(( + e, + SimplePrefetcherGuards { + relation: self.relation, + list, + }, + )) + } +} + +pub struct SimplePrefetcherGuards<'r, R> { + relation: &'r R, + list: std::vec::IntoIter, +} + +impl<'r, R: RelationRead> Iterator for SimplePrefetcherGuards<'r, R> { + type Item = R::ReadGuard<'r>; + + fn next(&mut self) -> Option { + let id = self.list.next()?; + Some(self.relation.read(id)) } } @@ -182,18 +241,20 @@ where } } -impl<'r, R: RelationReadStream, S: Sequence> Prefetcher<'r> for StreamPrefetcher<'r, R, S> +impl<'r, R: RelationRead + RelationReadStream, S: Sequence> Prefetcher<'r> + for StreamPrefetcher<'r, R, S> where S::Item: Fetch, { type R = R; - fn next(&mut self) -> Option<(S::Item, Vec>)> { + type Guards = > as ReadStream<'r>>::Guards; + fn next(&mut self) -> Option<(S::Item, Self::Guards)> { self.stream.next() } fn next_if( &mut self, predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<(S::Item, Vec>)> { + ) -> Option<(S::Item, Self::Guards)> { self.stream.next_if(predicate) } } diff --git a/crates/algo/src/tuples.rs b/crates/algo/src/tuples.rs new file mode 100644 index 00000000..fd741f70 --- /dev/null +++ b/crates/algo/src/tuples.rs @@ -0,0 +1,233 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use std::marker::PhantomData; +use std::num::NonZero; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +pub const ALIGN: usize = 8; + +pub struct RefChecker<'a> { + bytes: &'a [u8], +} + +impl<'a> RefChecker<'a> { + pub fn new(bytes: &'a [u8]) -> Self { + Self { bytes } + } + #[inline] + pub fn prefix( + &self, + s: impl Into + Copy, + ) -> &'a T { + let start = Into::::into(s); + let end = Into::::into(s) + size_of::(); + let bytes = &self.bytes[start..end]; + FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") + } + #[inline] + pub fn bytes( + &self, + s: impl Into + Copy, + e: impl Into + Copy, + ) -> &'a T { + let start = Into::::into(s); + let end = Into::::into(e); + let bytes = &self.bytes[start..end]; + FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") + } + /// # Safety + /// + /// * `FromBytes` could be implemented for `T`. + #[inline] + #[allow(unsafe_code)] + pub unsafe fn bytes_slice_unchecked( + &self, + s: impl Into + Copy, + e: impl Into + Copy, + ) -> &'a [T] { + let start = Into::::into(s); + let end = Into::::into(e); + let bytes = &self.bytes[start..end]; + if size_of::() == 0 || bytes.len() % size_of::() == 0 { + let ptr = bytes as *const [u8] as *const T; + if ptr.is_aligned() { + unsafe { std::slice::from_raw_parts(ptr, bytes.len() / size_of::()) } + } else { + panic!("deserialization: bad bytes") + } + } else { + panic!("deserialization: bad bytes") + } + } +} + +pub struct MutChecker<'a> { + flag: usize, + bytes: *mut [u8], + phantom: PhantomData<&'a mut [u8]>, +} + +impl<'a> MutChecker<'a> { + pub fn new(bytes: &'a mut [u8]) -> Self { + Self { + flag: 0, + bytes, + phantom: PhantomData, + } + } + #[inline] + pub fn prefix( + &mut self, + s: impl Into + Copy, + ) -> &'a mut T { + let start = Into::::into(s); + let end = Into::::into(s) + size_of::(); + if !(start <= end && end <= self.bytes.len()) { + panic!("deserialization: bad bytes"); + } + if !(self.flag <= start) { + panic!("deserialization: bad bytes"); + } else { + self.flag = end; + } + #[allow(unsafe_code)] + let bytes = unsafe { + std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) + }; + FromBytes::mut_from_bytes(bytes).expect("deserialization: bad bytes") + } + #[inline] + pub fn bytes( + &mut self, + s: impl Into + Copy, + e: impl Into + Copy, + ) -> &'a mut T { + let start = Into::::into(s); + let end = Into::::into(e); + if !(start <= end && end <= self.bytes.len()) { + panic!("deserialization: bad bytes"); + } + if !(self.flag <= start) { + panic!("deserialization: bad bytes"); + } else { + self.flag = end; + } + #[allow(unsafe_code)] + let bytes = unsafe { + std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) + }; + FromBytes::mut_from_bytes(bytes).expect("deserialization: bad bytes") + } +} + +#[test] +fn aliasing_test() { + #[repr(C, align(8))] + #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct ExampleHeader { + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 4], + } + let serialized = { + let elements = (0u32..1111).collect::>(); + let mut buffer = Vec::::new(); + buffer.extend(std::iter::repeat_n(0, size_of::())); + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + buffer[..size_of::()].copy_from_slice( + ExampleHeader { + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + buffer + }; + let mut source = vec![0u64; serialized.len().next_multiple_of(8)]; + source.as_mut_bytes()[..serialized.len()].copy_from_slice(&serialized); + let deserialized = { + let mut checker = MutChecker::new(source.as_mut_bytes()); + let header: &mut ExampleHeader = checker.prefix(0_u16); + let elements: &mut [u32] = checker.bytes(header.elements_s, header.elements_e); + (header, elements) + }; + assert_eq!( + deserialized.1, + (0u32..1111).collect::>().as_slice() + ); +} + +#[repr(transparent)] +#[derive( + Debug, + Default, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct Padding(Option>); + +impl Padding { + pub const ZERO: Self = Self(None); +} + +#[repr(transparent)] +#[derive( + Debug, + Default, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct Bool(u8); + +impl Bool { + pub const FALSE: Self = Self(0); + pub const TRUE: Self = Self(1); +} + +impl From for bool { + fn from(value: Bool) -> Self { + value != Bool::FALSE + } +} + +impl From for Bool { + fn from(value: bool) -> Self { + std::hint::select_unpredictable(value, Self::TRUE, Self::FALSE) + } +} diff --git a/crates/distance/Cargo.toml b/crates/distance/Cargo.toml index c5a21ed3..348805cf 100644 --- a/crates/distance/Cargo.toml +++ b/crates/distance/Cargo.toml @@ -4,5 +4,8 @@ version.workspace = true edition.workspace = true publish = false +[dependencies] +zerocopy.workspace = true + [lints] workspace = true diff --git a/crates/distance/src/lib.rs b/crates/distance/src/lib.rs index 39867f52..557f84e6 100644 --- a/crates/distance/src/lib.rs +++ b/crates/distance/src/lib.rs @@ -12,10 +12,31 @@ // // Copyright (c) 2025 TensorChord Inc. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +#[derive( + Clone, + Copy, + Default, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] #[repr(transparent)] pub struct Distance(i32); +impl std::fmt::Debug for Distance { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.to_f32()) + } +} + impl Distance { pub const ZERO: Self = Distance::from_f32(0.0f32); pub const INFINITY: Self = Distance::from_f32(f32::INFINITY); diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index 0398a4b6..91e1d678 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use rabitq::block::BlockCode; +use rabitq::original::block::BlockCode; use rabitq::packing::{any_pack, padding_pack}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; @@ -55,7 +55,7 @@ pub fn k_means( move |pool| { let compute = |centroids: &[Vec]| { if n >= 1024 && c >= 1024 { - rabitq_index(dims, n, c, samples, centroids) + rabitq_index(n, c, samples, centroids) } else { flat_index(dims, n, c, samples, centroids) } @@ -116,17 +116,11 @@ fn quick_centers( centroids } -fn rabitq_index( - dims: usize, - n: usize, - c: usize, - samples: &[Vec], - centroids: &[Vec], -) -> Vec { +fn rabitq_index(n: usize, c: usize, samples: &[Vec], centroids: &[Vec]) -> Vec { let branches = { let mut branches = Vec::new(); for centroid in centroids { - let code = rabitq::code(dims as _, centroid); + let code = rabitq::original::code(centroid); branches.push(code); } branches @@ -162,10 +156,10 @@ fn rabitq_index( (0..n) .into_par_iter() .map(|i| { - let lut = rabitq::block::preprocess(&samples[i]); + let lut = rabitq::original::block::preprocess(&samples[i]); let mut result = (f32::INFINITY, 0); for block in 0..c.div_ceil(32) { - let returns = rabitq::block::full_process_l2(&lut, blocks[block].code()); + let returns = rabitq::original::block::full_process_l2(blocks[block].code(), &lut); let lowerbound = returns.map(|(rough, err)| rough - err * 1.9); for j in block * 32..std::cmp::min(block * 32 + 32, c) { if lowerbound[j - block * 32] < result.0 { diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index 2167f140..7e5d20bc 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -7,6 +7,6 @@ publish = false [dependencies] anyhow = "1.0.98" clap = { version = "4.5.38", features = ["derive"] } -which = "7.0.3" -tempfile = "3.20.0" object = { version = "0.36.7", features = ["all"] } +tempfile = "3.20.0" +which = "7.0.3" diff --git a/crates/rabitq/src/b1.rs b/crates/rabitq/src/b1.rs new file mode 100644 index 00000000..25fb8fd4 --- /dev/null +++ b/crates/rabitq/src/b1.rs @@ -0,0 +1,181 @@ +use simd::Floating; + +#[derive(Debug, Clone, Copy)] +pub struct CodeMetadata { + pub dis_u_2: f32, + pub factor_cnt: f32, + pub factor_ip: f32, +} + +impl CodeMetadata { + pub fn into_tuple(self) -> (f32, f32, f32) { + (self.dis_u_2, self.factor_cnt, self.factor_ip) + } + pub fn into_array(self) -> [f32; 3] { + [self.dis_u_2, self.factor_cnt, self.factor_ip] + } + pub fn from_tuple((dis_u_2, factor_cnt, factor_ip): (f32, f32, f32)) -> Self { + Self { + dis_u_2, + factor_cnt, + factor_ip, + } + } + pub fn from_array([dis_u_2, factor_cnt, factor_ip]: [f32; 3]) -> Self { + Self { + dis_u_2, + factor_cnt, + factor_ip, + } + } +} + +pub type Code = (CodeMetadata, Vec); + +pub fn code(vector: &[f32]) -> Code { + let n = vector.len(); + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x_2 = f32::reduce_sum_of_x2(vector); + ( + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector + .iter() + .map(|x| x.is_sign_positive() as i32) + .sum::(); + let cnt_neg = vector + .iter() + .map(|x| x.is_sign_negative() as i32) + .sum::(); + (cnt_pos - cnt_neg) as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + }, + { + let mut signs = Vec::new(); + for i in 0..n { + signs.push(vector[i].is_sign_positive()); + } + signs + }, + ) +} + +pub mod binary { + pub fn pack_code(input: &[bool]) -> Vec { + let f = |t: &[bool; 64]| { + let mut result = 0_u64; + for i in 0..64 { + result |= (t[i] as u64) << i; + } + result + }; + let (arrays, remainder) = input.as_chunks::<64>(); + let mut buffer = [false; 64]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() + } + + use super::CodeMetadata; + use simd::Floating; + + const BITS: usize = 6; + + #[derive(Debug, Clone, Copy)] + pub struct BinaryLutMetadata { + dis_v_2: f32, + b: f32, + k: f32, + qvector_sum: f32, + } + + pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u64]); + + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let dis_v_2 = f32::reduce_sum_of_x2(vector); + preprocess_with_distance(vector, dis_v_2) + } + + pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BinaryLut { + let (k, b, qvector) = simd::quantize::quantize(vector, ((1 << BITS) - 1) as f32); + let qvector_sum = if vector.len() <= (65535_usize / ((1 << BITS) - 1)) { + simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 + } else { + simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 + }; + ( + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }, + binarize(&qvector), + ) + } + + pub fn accumulate(x: &[u64], y: &[impl AsRef<[u64]>; BITS]) -> u32 { + let mut result = 0_u32; + for i in 0..BITS { + result += simd::bit::reduce_sum_of_and(x, y[i].as_ref()) << i; + } + result + } + + pub fn half_process_l2( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt, + factor_ip, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + ) -> (f32,) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; + (rough,) + } + + pub fn half_process_dot( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt, + factor_ip, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2: _, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + ) -> (f32,) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = -e * factor_ip; + (rough,) + } + + pub(crate) fn binarize(vector: &[u8]) -> [Vec; BITS] { + let n = vector.len(); + let mut t: [_; BITS] = std::array::from_fn(|_| vec![0_u64; n.div_ceil(64)]); + for i in 0..BITS { + for j in 0..n { + let bit = (vector[j] >> i) & 1; + t[i][j / 64] |= (bit as u64) << (j % 64); + } + } + t + } +} diff --git a/crates/rabitq/src/b2.rs b/crates/rabitq/src/b2.rs new file mode 100644 index 00000000..a133bfbf --- /dev/null +++ b/crates/rabitq/src/b2.rs @@ -0,0 +1,71 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub use crate::extended::{Code, CodeMetadata}; + +pub fn code(vector: &[f32]) -> Code { + crate::extended::code::<2>(vector) +} + +pub mod binary { + pub fn pack_code(input: &[u8]) -> Vec { + let f = |&[t_0, t_1, t_2, t_3]: &[u8; 4]| t_0 | (t_1 << 2) | (t_2 << 4) | (t_3 << 6); + let (arrays, remainder) = input.as_chunks::<4>(); + let mut buffer = [0u8; 4]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() + } + + use crate::extended::CodeMetadata; + + const BITS: usize = 8; + + pub type BinaryLutMetadata = CodeMetadata; + pub type BinaryLut = (BinaryLutMetadata, Vec); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); + + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::code::(vector); + (metadata, crate::b8::binary::pack_code(&elements)) + } + + pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { + simd::packed::u2_u8_reduce_sum_of_xy(x, y) + } + + pub fn half_process_dot( + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_dot::<2, BITS>(n, value, code, lut); + (rough,) + } + + pub fn half_process_l2( + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_l2::<2, BITS>(n, value, code, lut); + (rough,) + } +} diff --git a/crates/rabitq/src/b4.rs b/crates/rabitq/src/b4.rs new file mode 100644 index 00000000..6aa4236e --- /dev/null +++ b/crates/rabitq/src/b4.rs @@ -0,0 +1,71 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub use crate::extended::{Code, CodeMetadata}; + +pub fn code(vector: &[f32]) -> Code { + crate::extended::code::<4>(vector) +} + +pub mod binary { + pub fn pack_code(input: &[u8]) -> Vec { + let f = |&[t_0, t_1]: &[u8; 2]| t_0 | (t_1 << 4); + let (arrays, remainder) = input.as_chunks::<2>(); + let mut buffer = [0u8; 2]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() + } + + use crate::extended::CodeMetadata; + + const BITS: usize = 4; + + pub type BinaryLutMetadata = CodeMetadata; + pub type BinaryLut = (BinaryLutMetadata, Vec); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); + + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::code::(vector); + (metadata, pack_code(&elements)) + } + + pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { + simd::packed::u4_u4_reduce_sum_of_xy(x, y) + } + + pub fn half_process_dot( + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_dot::<4, BITS>(n, value, code, lut); + (rough,) + } + + pub fn half_process_l2( + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_l2::<4, BITS>(n, value, code, lut); + (rough,) + } +} diff --git a/crates/rabitq/src/b8.rs b/crates/rabitq/src/b8.rs new file mode 100644 index 00000000..3b138890 --- /dev/null +++ b/crates/rabitq/src/b8.rs @@ -0,0 +1,62 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub use crate::extended::{Code, CodeMetadata}; + +pub fn code(vector: &[f32]) -> Code { + crate::extended::code::<8>(vector) +} + +pub mod binary { + pub fn pack_code(input: &[u8]) -> Vec { + input.to_vec() + } + + use crate::extended::CodeMetadata; + + const BITS: usize = 8; + + pub type BinaryLutMetadata = CodeMetadata; + pub type BinaryLut = (BinaryLutMetadata, Vec); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); + + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::code::(vector); + (metadata, pack_code(&elements)) + } + + pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { + simd::u8::reduce_sum_of_x_as_u32_y_as_u32(x, y) + } + + pub fn half_process_dot( + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_dot::<8, BITS>(n, value, code, lut); + (rough,) + } + + pub fn half_process_l2( + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_l2::<8, BITS>(n, value, code, lut); + (rough,) + } +} diff --git a/crates/rabitq/src/binary.rs b/crates/rabitq/src/binary.rs deleted file mode 100644 index dc1996ab..00000000 --- a/crates/rabitq/src/binary.rs +++ /dev/null @@ -1,202 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -use simd::Floating; - -use crate::CodeMetadata; - -const BITS: usize = 6; - -#[derive(Debug, Clone, Copy)] -pub struct BinaryLutMetadata { - dis_v_2: f32, - b: f32, - k: f32, - qvector_sum: f32, -} - -pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); -pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u64]); - -pub fn preprocess(vector: &[f32]) -> BinaryLut { - let dis_v_2 = f32::reduce_sum_of_x2(vector); - preprocess_with_distance(vector, dis_v_2) -} - -pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BinaryLut { - let (k, b, qvector) = simd::quantize::quantize(vector, ((1 << BITS) - 1) as f32); - let qvector_sum = if vector.len() <= (65535_usize / ((1 << BITS) - 1)) { - simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 - } else { - simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 - }; - ( - BinaryLutMetadata { - dis_v_2, - b, - k, - qvector_sum, - }, - binarize(&qvector), - ) -} - -pub fn full_process_l2( - lut: &BinaryLut, - ((dis_u_2, factor_cnt, factor_ip, factor_err), t): BinaryCode<'_>, -) -> (f32, f32) { - let &( - BinaryLutMetadata { - dis_v_2, - b, - k, - qvector_sum, - }, - ref s, - ) = lut; - let sum = asymmetric_binary_dot_product(t, s); - let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; - let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; - let err = 2.0 * factor_err * dis_v_2.sqrt(); - (rough, err) -} - -pub fn full_process_dot( - lut: &BinaryLut, - ((_, factor_cnt, factor_ip, factor_err), t): BinaryCode<'_>, -) -> (f32, f32) { - let &( - BinaryLutMetadata { - dis_v_2, - b, - k, - qvector_sum, - }, - ref s, - ) = lut; - let sum = asymmetric_binary_dot_product(t, s); - let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; - let rough = -e * factor_ip; - let err = factor_err * dis_v_2.sqrt(); - (rough, err) -} - -pub fn half_process_l2( - sum: u32, - CodeMetadata { - dis_u_2, - factor_cnt, - factor_ip, - factor_err, - }: CodeMetadata, - BinaryLutMetadata { - dis_v_2, - b, - k, - qvector_sum, - }: BinaryLutMetadata, -) -> (f32, f32) { - let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; - let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; - let err = 2.0 * factor_err * dis_v_2.sqrt(); - (rough, err) -} - -pub fn half_process_l2_residual( - sum: u32, - CodeMetadata { - dis_u_2, - factor_cnt, - factor_ip, - factor_err, - }: CodeMetadata, - BinaryLutMetadata { - dis_v_2: _, - b, - k, - qvector_sum, - }: BinaryLutMetadata, - dis_f: f32, - delta: f32, -) -> (f32, f32) { - let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; - let rough = dis_u_2 + dis_f - 2.0 * e * factor_ip + delta; - let err = 2.0 * factor_err * dis_f.sqrt(); - (rough, err) -} - -pub fn half_process_dot( - sum: u32, - CodeMetadata { - dis_u_2: _, - factor_cnt, - factor_ip, - factor_err, - }: CodeMetadata, - BinaryLutMetadata { - dis_v_2, - b, - k, - qvector_sum, - }: BinaryLutMetadata, -) -> (f32, f32) { - let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; - let rough = -e * factor_ip; - let err = factor_err * dis_v_2.sqrt(); - (rough, err) -} - -pub fn half_process_dot_residual( - sum: u32, - CodeMetadata { - dis_u_2: _, - factor_cnt, - factor_ip, - factor_err, - }: CodeMetadata, - BinaryLutMetadata { - dis_v_2, - b, - k, - qvector_sum, - }: BinaryLutMetadata, - dis_f: f32, - delta: f32, - norm: f32, -) -> (f32, f32) { - let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; - let rough = -e * factor_ip + dis_f + delta; - let err = factor_err * (dis_v_2 + norm * norm + 2.0 * dis_f).sqrt(); - (rough, err) -} - -pub(crate) fn binarize(vector: &[u8]) -> [Vec; BITS] { - let n = vector.len(); - let mut t: [_; BITS] = std::array::from_fn(|_| vec![0_u64; n.div_ceil(64)]); - for i in 0..BITS { - for j in 0..n { - let bit = (vector[j] >> i) & 1; - t[i][j / 64] |= (bit as u64) << (j % 64); - } - } - t -} - -pub fn asymmetric_binary_dot_product(x: &[u64], y: &[Vec; BITS]) -> u32 { - let mut result = 0_u32; - for i in 0..BITS { - result += simd::bit::reduce_sum_of_and(x, y[i].as_slice()) << i; - } - result -} diff --git a/crates/rabitq/src/block.rs b/crates/rabitq/src/block.rs deleted file mode 100644 index 49bc6079..00000000 --- a/crates/rabitq/src/block.rs +++ /dev/null @@ -1,194 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -use simd::Floating; - -use crate::CodeMetadata; - -const BITS: usize = 8; -pub const STEP: usize = 65535_usize / ((1_usize << BITS) - 1); - -#[derive(Debug, Clone, Copy)] -pub struct BlockLutMetadata { - dis_v_2: f32, - k: f32, - c: f32, -} - -pub type BlockLut = (BlockLutMetadata, Vec<[u8; 16]>); -pub type BlockCode<'a> = ( - &'a [f32; 32], - &'a [f32; 32], - &'a [f32; 32], - &'a [f32; 32], - &'a [[u8; 16]], -); - -pub fn preprocess(vector: &[f32]) -> BlockLut { - let dis_v_2 = f32::reduce_sum_of_x2(vector); - preprocess_with_distance(vector, dis_v_2) -} - -pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BlockLut { - let cvector = compress(vector); - let (k, b, cqvector) = - simd::quantize::quantize(cvector.as_flattened(), ((1 << BITS) - 1) as f32); - ( - BlockLutMetadata { - dis_v_2, - k, - c: b * vector.len().div_ceil(4) as f32, - }, - cqvector.as_chunks::<16>().0.to_vec(), - ) -} - -pub fn full_process_l2( - lut: &BlockLut, - (dis_u_2, _, factor_ip, factor_err, t): BlockCode<'_>, -) -> [(f32, f32); 32] { - use std::iter::zip; - let &(BlockLutMetadata { dis_v_2, k, c }, ref s) = lut; - let mut sum = [0_u32; 32]; - for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { - let delta = simd::fast_scan::scan(t, s); - simd::fast_scan::accu(&mut sum, &delta); - } - std::array::from_fn(|i| { - let e = k * (sum[i] as f32) + c; - let rough = dis_u_2[i] + dis_v_2 - 2.0 * e * factor_ip[i]; - let err = 2.0 * factor_err[i] * dis_v_2.sqrt(); - (rough, err) - }) -} - -pub fn full_process_dot( - lut: &BlockLut, - (_, _, factor_ip, factor_err, t): BlockCode<'_>, -) -> [(f32, f32); 32] { - use std::iter::zip; - let &(BlockLutMetadata { dis_v_2, k, c }, ref s) = lut; - let mut sum = [0_u32; 32]; - for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { - let delta = simd::fast_scan::scan(t, s); - simd::fast_scan::accu(&mut sum, &delta); - } - std::array::from_fn(|i| { - let e = k * (sum[i] as f32) + c; - let rough = -e * factor_ip[i]; - let err = factor_err[i] * dis_v_2.sqrt(); - (rough, err) - }) -} - -pub fn half_process_l2( - sum: u32, - CodeMetadata { - dis_u_2, - factor_cnt: _, - factor_ip, - factor_err, - }: CodeMetadata, - BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, -) -> (f32, f32) { - let e = k * (sum as f32) + c; - let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; - let err = 2.0 * factor_err * dis_v_2.sqrt(); - (rough, err) -} - -pub fn half_process_l2_residual( - sum: u32, - CodeMetadata { - dis_u_2, - factor_cnt: _, - factor_ip, - factor_err, - }: CodeMetadata, - BlockLutMetadata { dis_v_2: _, k, c }: BlockLutMetadata, - dis_f: f32, - delta: f32, -) -> (f32, f32) { - let e = k * (sum as f32) + c; - let rough = dis_u_2 + dis_f - 2.0 * e * factor_ip + delta; - let err = 2.0 * factor_err * dis_f.sqrt(); - (rough, err) -} - -pub fn half_process_dot( - sum: u32, - CodeMetadata { - dis_u_2: _, - factor_cnt: _, - factor_ip, - factor_err, - }: CodeMetadata, - BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, -) -> (f32, f32) { - let e = k * (sum as f32) + c; - let rough = -e * factor_ip; - let err = factor_err * dis_v_2.sqrt(); - (rough, err) -} - -pub fn half_process_dot_residual( - sum: u32, - CodeMetadata { - dis_u_2: _, - factor_cnt: _, - factor_ip, - factor_err, - }: CodeMetadata, - BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, - dis_f: f32, - delta: f32, - norm: f32, -) -> (f32, f32) { - let e = k * (sum as f32) + c; - let rough = -e * factor_ip + dis_f + delta; - let err = factor_err * (dis_v_2 + norm * norm + 2.0 * dis_f).sqrt(); - (rough, err) -} - -fn compress(vector: &[f32]) -> Vec<[f32; 16]> { - let f = |[t_0, t_1, t_2, t_3]: [f32; 4]| { - [ - 0.0 - t_3 - t_2 - t_1 - t_0, - 0.0 - t_3 - t_2 - t_1 + t_0, - 0.0 - t_3 - t_2 + t_1 - t_0, - 0.0 - t_3 - t_2 + t_1 + t_0, - 0.0 - t_3 + t_2 - t_1 - t_0, - 0.0 - t_3 + t_2 - t_1 + t_0, - 0.0 - t_3 + t_2 + t_1 - t_0, - 0.0 - t_3 + t_2 + t_1 + t_0, - 0.0 + t_3 - t_2 - t_1 - t_0, - 0.0 + t_3 - t_2 - t_1 + t_0, - 0.0 + t_3 - t_2 + t_1 - t_0, - 0.0 + t_3 - t_2 + t_1 + t_0, - 0.0 + t_3 + t_2 - t_1 - t_0, - 0.0 + t_3 + t_2 - t_1 + t_0, - 0.0 + t_3 + t_2 + t_1 - t_0, - 0.0 + t_3 + t_2 + t_1 + t_0, - ] - }; - - let (arrays, remainder) = vector.as_chunks::<4>(); - let mut result = arrays.iter().copied().map(f).collect::>(); - if !remainder.is_empty() { - let mut array = [0.0; 4]; - array[..remainder.len()].copy_from_slice(remainder); - result.push(f(array)); - } - result -} diff --git a/crates/rabitq/src/extended.rs b/crates/rabitq/src/extended.rs new file mode 100644 index 00000000..2a6c2814 --- /dev/null +++ b/crates/rabitq/src/extended.rs @@ -0,0 +1,181 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use simd::Floating; + +#[derive(Debug, Clone, Copy)] +pub struct CodeMetadata { + pub dis_u_2: f32, + pub norm_of_lattice: f32, + pub sum_of_code: f32, +} + +impl CodeMetadata { + pub fn into_tuple(self) -> (f32, f32, f32) { + (self.dis_u_2, self.norm_of_lattice, self.sum_of_code) + } + pub fn into_array(self) -> [f32; 3] { + [self.dis_u_2, self.norm_of_lattice, self.sum_of_code] + } + pub fn from_tuple((dis_u_2, norm_of_lattice, sum_of_code): (f32, f32, f32)) -> Self { + Self { + dis_u_2, + norm_of_lattice, + sum_of_code, + } + } + pub fn from_array([dis_u_2, norm_of_lattice, sum_of_code]: [f32; 3]) -> Self { + Self { + dis_u_2, + norm_of_lattice, + sum_of_code, + } + } +} + +pub type Code = (CodeMetadata, Vec); + +pub fn code(vector: &[f32]) -> Code { + assert!((1..=8).contains(&BITS)); + + let n = vector.len(); + let dis_u_2 = f32::reduce_sum_of_x2(vector); + let code = { + let min = -(1 << (BITS - 1)); + let max = (1 << (BITS - 1)) - 1; + let normalized_vector = { + let mut vector = vector.to_vec(); + f32::vector_mul_scalar_inplace(&mut vector, 1.0 / dis_u_2.sqrt()); + vector + }; + let scale = { + let mut o = normalized_vector.clone(); + f32::vector_abs_inplace(&mut o); + find_scale::(&o) as f32 + f32::EPSILON + }; + let mut code = Vec::with_capacity(n as _); + for i in 0..n { + let v = scale * normalized_vector[i]; + let c = v.floor().clamp(min as f32, max as f32) as i32; + code.push((c + (1 << (BITS - 1))) as _); + } + code + }; + let norm_of_lattice = { + let base = -0.5 * ((1 << BITS) - 1) as f32; + let mut y = 0.0; + for i in 0..n { + let x = base + code[i] as f32; + y += x * x; + } + y.sqrt() + }; + let sum_of_code = { + let mut y = 0; + for i in 0..n { + let x = code[i] as u32; + y += x; + } + y as f32 + }; + ( + CodeMetadata { + dis_u_2, + norm_of_lattice, + sum_of_code, + }, + code, + ) +} + +pub fn half_process_l2( + n: u32, + value: u32, + lhs: CodeMetadata, + rhs: CodeMetadata, +) -> f32 { + assert!((1..=8).contains(&X)); + assert!((1..=8).contains(&Y)); + + let c_x = ((1 << X) - 1) as f32 * 0.5; + let c_y = ((1 << Y) - 1) as f32 * 0.5; + + let ip = value as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + n as f32 * c_x * c_y; + lhs.dis_u_2 + rhs.dis_u_2 + - 2.0 + * ip + * (lhs.dis_u_2.sqrt() / lhs.norm_of_lattice) + * (rhs.dis_u_2.sqrt() / rhs.norm_of_lattice) +} + +pub fn half_process_dot( + n: u32, + value: u32, + lhs: CodeMetadata, + rhs: CodeMetadata, +) -> f32 { + assert!((1..=8).contains(&X)); + assert!((1..=8).contains(&Y)); + + let c_x = ((1 << X) - 1) as f32 * 0.5; + let c_y = ((1 << Y) - 1) as f32 * 0.5; + + let ip = value as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + n as f32 * c_x * c_y; + -ip * (lhs.dis_u_2.sqrt() / lhs.norm_of_lattice) * (rhs.dis_u_2.sqrt() / rhs.norm_of_lattice) +} + +fn find_scale(o: &[f32]) -> f64 { + assert!((1..=8).contains(&B)); + + let mask = (1_u32 << (B - 1)) - 1; + let dims = o.len(); + + let mut code = Vec::::with_capacity(dims); + let mut numerator = 0.0f64; + let mut sqr_denominator = 0.0f64; + + let (mut y_m, mut x_m); + + for i in 0..dims { + code.push(0); + numerator += 0.5 * o[i] as f64; + sqr_denominator += 0.5 * 0.5; + } + { + let x = 0.0; + let y = numerator / sqr_denominator.sqrt(); + (y_m, x_m) = (y, x); + } + + let mut events = Vec::<(f64, usize)>::new(); + for i in 0..dims { + for c in 1..=mask { + let x = (c as f64) / o[i] as f64; + events.push((x, i)); + } + } + events.sort_unstable_by(|(x, _), (y, _)| f64::total_cmp(x, y)); + for (x, i) in events.into_iter() { + code[i] += 1; + numerator += o[i] as f64; + sqr_denominator += 2.0 * code[i] as f64; + + let y = numerator / sqr_denominator.sqrt(); + if y > y_m { + (y_m, x_m) = (y, x); + } + } + + x_m +} diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index 30f1856a..0a1925f1 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -12,63 +12,11 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod binary; -pub mod block; +pub mod b1; +pub mod b2; +pub mod b4; +pub mod b8; +mod extended; +pub mod original; pub mod packing; pub mod rotate; - -use binary::BinaryLut; -use block::BlockLut; -use simd::Floating; - -#[derive(Debug, Clone, Copy)] -pub struct CodeMetadata { - pub dis_u_2: f32, - pub factor_cnt: f32, - pub factor_ip: f32, - pub factor_err: f32, -} - -pub type Code = (CodeMetadata, Vec); - -pub fn code(dims: u32, vector: &[f32]) -> Code { - let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); - let sum_of_x_2 = f32::reduce_sum_of_x2(vector); - ( - CodeMetadata { - dis_u_2: sum_of_x_2, - factor_cnt: { - let cnt_pos = vector - .iter() - .map(|x| x.is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.is_sign_negative() as i32) - .sum::(); - (cnt_pos - cnt_neg) as f32 - }, - factor_ip: sum_of_x_2 / sum_of_abs_x, - factor_err: { - let dis_u = sum_of_x_2.sqrt(); - let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); - dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (dims as f32 - 1.0).sqrt() - }, - }, - { - let mut signs = Vec::new(); - for i in 0..dims { - signs.push(vector[i as usize].is_sign_positive()); - } - signs - }, - ) -} - -pub fn preprocess(vector: &[f32]) -> (BlockLut, BinaryLut) { - let dis_v_2 = f32::reduce_sum_of_x2(vector); - ( - block::preprocess_with_distance(vector, dis_v_2), - binary::preprocess_with_distance(vector, dis_v_2), - ) -} diff --git a/crates/rabitq/src/original.rs b/crates/rabitq/src/original.rs new file mode 100644 index 00000000..a24c3111 --- /dev/null +++ b/crates/rabitq/src/original.rs @@ -0,0 +1,456 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use binary::BinaryLut; +use block::BlockLut; +use simd::Floating; + +#[derive(Debug, Clone, Copy)] +pub struct CodeMetadata { + pub dis_u_2: f32, + pub factor_cnt: f32, + pub factor_ip: f32, + pub factor_err: f32, +} + +impl CodeMetadata { + pub fn into_tuple(self) -> (f32, f32, f32, f32) { + ( + self.dis_u_2, + self.factor_cnt, + self.factor_ip, + self.factor_err, + ) + } + pub fn into_array(self) -> [f32; 4] { + [ + self.dis_u_2, + self.factor_cnt, + self.factor_ip, + self.factor_err, + ] + } + pub fn from_tuple((dis_u_2, factor_cnt, factor_ip, factor_err): (f32, f32, f32, f32)) -> Self { + Self { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + } + } + pub fn from_array([dis_u_2, factor_cnt, factor_ip, factor_err]: [f32; 4]) -> Self { + Self { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + } + } +} + +pub type Code = (CodeMetadata, Vec); + +pub fn code(vector: &[f32]) -> Code { + let n = vector.len(); + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x_2 = f32::reduce_sum_of_x2(vector); + ( + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector + .iter() + .map(|x| x.is_sign_positive() as i32) + .sum::(); + let cnt_neg = vector + .iter() + .map(|x| x.is_sign_negative() as i32) + .sum::(); + (cnt_pos - cnt_neg) as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + factor_err: { + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (n as f32).sqrt(); + dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (n as f32 - 1.0).sqrt() + }, + }, + { + let mut signs = Vec::new(); + for i in 0..n { + signs.push(vector[i].is_sign_positive()); + } + signs + }, + ) +} + +pub fn preprocess(vector: &[f32]) -> (BlockLut, BinaryLut) { + let dis_v_2 = f32::reduce_sum_of_x2(vector); + ( + block::preprocess_with_distance(vector, dis_v_2), + binary::preprocess_with_distance(vector, dis_v_2), + ) +} + +pub mod binary { + pub fn pack_code(input: &[bool]) -> Vec { + let f = |t: &[bool; 64]| { + let mut result = 0_u64; + for i in 0..64 { + result |= (t[i] as u64) << i; + } + result + }; + let (arrays, remainder) = input.as_chunks::<64>(); + let mut buffer = [false; 64]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() + } + + use super::CodeMetadata; + use simd::Floating; + + const BITS: usize = 6; + + #[derive(Debug, Clone, Copy)] + pub struct BinaryLutMetadata { + dis_v_2: f32, + b: f32, + k: f32, + qvector_sum: f32, + } + + pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u64]); + + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let dis_v_2 = f32::reduce_sum_of_x2(vector); + preprocess_with_distance(vector, dis_v_2) + } + + pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BinaryLut { + let (k, b, qvector) = simd::quantize::quantize(vector, ((1 << BITS) - 1) as f32); + let qvector_sum = if vector.len() <= (65535_usize / ((1 << BITS) - 1)) { + simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 + } else { + simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 + }; + ( + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }, + binarize(&qvector), + ) + } + + pub fn accumulate(x: &[u64], y: &[impl AsRef<[u64]>; BITS]) -> u32 { + let mut result = 0_u32; + for i in 0..BITS { + result += simd::bit::reduce_sum_of_and(x, y[i].as_ref()) << i; + } + result + } + + pub fn half_process_l2( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + ) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; + let err = 2.0 * factor_err * dis_v_2.sqrt(); + (rough, err) + } + + pub fn half_process_l2_residual( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2: _, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + dis_f: f32, + delta: f32, + ) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = dis_u_2 + dis_f - 2.0 * e * factor_ip + delta; + let err = 2.0 * factor_err * dis_f.sqrt(); + (rough, err) + } + + pub fn half_process_dot( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + ) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = -e * factor_ip; + let err = factor_err * dis_v_2.sqrt(); + (rough, err) + } + + pub fn half_process_dot_residual( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = -e * factor_ip + dis_f + delta; + let err = factor_err * (dis_v_2 + norm * norm + 2.0 * dis_f).sqrt(); + (rough, err) + } + + pub(crate) fn binarize(vector: &[u8]) -> [Vec; BITS] { + let n = vector.len(); + let mut t: [_; BITS] = std::array::from_fn(|_| vec![0_u64; n.div_ceil(64)]); + for i in 0..BITS { + for j in 0..n { + let bit = (vector[j] >> i) & 1; + t[i][j / 64] |= (bit as u64) << (j % 64); + } + } + t + } +} + +pub mod block { + use super::CodeMetadata; + use simd::Floating; + + const BITS: usize = 8; + pub const STEP: usize = 65535_usize / ((1_usize << BITS) - 1); + + #[derive(Debug, Clone, Copy)] + pub struct BlockLutMetadata { + dis_v_2: f32, + k: f32, + c: f32, + } + + pub type BlockLut = (BlockLutMetadata, Vec<[u8; 16]>); + pub type BlockCode<'a> = ( + &'a [f32; 32], + &'a [f32; 32], + &'a [f32; 32], + &'a [f32; 32], + &'a [[u8; 16]], + ); + + pub fn preprocess(vector: &[f32]) -> BlockLut { + let dis_v_2 = f32::reduce_sum_of_x2(vector); + preprocess_with_distance(vector, dis_v_2) + } + + pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BlockLut { + let cvector = compress(vector); + let (k, b, cqvector) = + simd::quantize::quantize(cvector.as_flattened(), ((1 << BITS) - 1) as f32); + ( + BlockLutMetadata { + dis_v_2, + k, + c: b * vector.len().div_ceil(4) as f32, + }, + cqvector.as_chunks::<16>().0.to_vec(), + ) + } + + pub fn full_process_l2( + (dis_u_2, _, factor_ip, factor_err, t): BlockCode<'_>, + lut: &BlockLut, + ) -> [(f32, f32); 32] { + use std::iter::zip; + let &(BlockLutMetadata { dis_v_2, k, c }, ref s) = lut; + let mut sum = [0_u32; 32]; + for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { + let delta = simd::fast_scan::scan(t, s); + simd::fast_scan::accu(&mut sum, &delta); + } + std::array::from_fn(|i| { + let e = k * (sum[i] as f32) + c; + let rough = dis_u_2[i] + dis_v_2 - 2.0 * e * factor_ip[i]; + let err = 2.0 * factor_err[i] * dis_v_2.sqrt(); + (rough, err) + }) + } + + pub fn full_process_dot( + (_, _, factor_ip, factor_err, t): BlockCode<'_>, + lut: &BlockLut, + ) -> [(f32, f32); 32] { + use std::iter::zip; + let &(BlockLutMetadata { dis_v_2, k, c }, ref s) = lut; + let mut sum = [0_u32; 32]; + for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { + let delta = simd::fast_scan::scan(t, s); + simd::fast_scan::accu(&mut sum, &delta); + } + std::array::from_fn(|i| { + let e = k * (sum[i] as f32) + c; + let rough = -e * factor_ip[i]; + let err = factor_err[i] * dis_v_2.sqrt(); + (rough, err) + }) + } + + pub fn half_process_l2( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, + ) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; + let err = 2.0 * factor_err * dis_v_2.sqrt(); + (rough, err) + } + + pub fn half_process_l2_residual( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2: _, k, c }: BlockLutMetadata, + dis_f: f32, + delta: f32, + ) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = dis_u_2 + dis_f - 2.0 * e * factor_ip + delta; + let err = 2.0 * factor_err * dis_f.sqrt(); + (rough, err) + } + + pub fn half_process_dot( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, + ) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = -e * factor_ip; + let err = factor_err * dis_v_2.sqrt(); + (rough, err) + } + + pub fn half_process_dot_residual( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = -e * factor_ip + dis_f + delta; + let err = factor_err * (dis_v_2 + norm * norm + 2.0 * dis_f).sqrt(); + (rough, err) + } + + fn compress(vector: &[f32]) -> Vec<[f32; 16]> { + let f = |&[t_0, t_1, t_2, t_3]: &[f32; 4]| { + [ + 0.0 - t_3 - t_2 - t_1 - t_0, + 0.0 - t_3 - t_2 - t_1 + t_0, + 0.0 - t_3 - t_2 + t_1 - t_0, + 0.0 - t_3 - t_2 + t_1 + t_0, + 0.0 - t_3 + t_2 - t_1 - t_0, + 0.0 - t_3 + t_2 - t_1 + t_0, + 0.0 - t_3 + t_2 + t_1 - t_0, + 0.0 - t_3 + t_2 + t_1 + t_0, + 0.0 + t_3 - t_2 - t_1 - t_0, + 0.0 + t_3 - t_2 - t_1 + t_0, + 0.0 + t_3 - t_2 + t_1 - t_0, + 0.0 + t_3 - t_2 + t_1 + t_0, + 0.0 + t_3 + t_2 - t_1 - t_0, + 0.0 + t_3 + t_2 - t_1 + t_0, + 0.0 + t_3 + t_2 + t_1 - t_0, + 0.0 + t_3 + t_2 + t_1 + t_0, + ] + }; + let (arrays, remainder) = vector.as_chunks::<4>(); + let mut buffer = [0.0f32; 4]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() + } +} diff --git a/crates/rabitq/src/packing.rs b/crates/rabitq/src/packing.rs index f4e249c2..008140dd 100644 --- a/crates/rabitq/src/packing.rs +++ b/crates/rabitq/src/packing.rs @@ -112,19 +112,3 @@ pub fn pack_to_u4(signs: &[bool]) -> Vec { } result } - -pub fn pack_to_u64(signs: &[bool]) -> Vec { - fn f(x: [bool; 64]) -> u64 { - let mut result = 0_u64; - for i in 0..64 { - result |= (x[i] as u64) << i; - } - result - } - let mut result = Vec::with_capacity(signs.len().div_ceil(64)); - for i in 0..signs.len().div_ceil(64) { - let x = std::array::from_fn(|j| signs.get(i * 64 + j).copied().unwrap_or_default()); - result.push(f(x)); - } - result -} diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 7d7e184d..e6828907 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -152,6 +152,11 @@ impl Floating for f16 { fn vector_to_f32_borrowed(this: &[Self]) -> impl AsRef<[f32]> { Self::vector_to_f32(this) } + + #[inline(always)] + fn vector_abs_inplace(this: &mut [Self]) { + vector_abs_inplace::vector_abs_inplace(this); + } } mod reduce_or_of_is_zero_x { diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index e97041fa..70ba66bd 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -95,41 +95,55 @@ impl Floating for f32 { reduce_sum_of_d2_sparse::reduce_sum_of_d2_sparse(lidx, lval, ridx, rval) } + #[inline(always)] fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { vector_add::vector_add(lhs, rhs) } + #[inline(always)] fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { vector_add_inplace::vector_add_inplace(lhs, rhs) } + #[inline(always)] fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { vector_sub::vector_sub(lhs, rhs) } + #[inline(always)] fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { vector_mul::vector_mul(lhs, rhs) } + #[inline(always)] fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { vector_mul_scalar::vector_mul_scalar(lhs, rhs) } + #[inline(always)] fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { vector_mul_scalar_inplace::vector_mul_scalar_inplace(lhs, rhs); } + #[inline(always)] fn vector_from_f32(this: &[f32]) -> Vec { this.to_vec() } + #[inline(always)] fn vector_to_f32(this: &[f32]) -> Vec { this.to_vec() } + #[inline(always)] fn vector_to_f32_borrowed(this: &[Self]) -> impl AsRef<[f32]> { this } + + #[inline(always)] + fn vector_abs_inplace(this: &mut [Self]) { + vector_abs_inplace::vector_abs_inplace(this); + } } mod reduce_or_of_is_zero_x { diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 8ddaac42..f96490ce 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -22,6 +22,7 @@ mod f32; pub mod bit; pub mod fast_scan; pub mod fht; +pub mod packed; pub mod quantize; pub mod rotate; pub mod u8; @@ -57,6 +58,7 @@ pub trait Floating: fn vector_mul(lhs: &[Self], rhs: &[Self]) -> Vec; fn vector_mul_scalar(lhs: &[Self], rhs: f32) -> Vec; fn vector_mul_scalar_inplace(lhs: &mut [Self], rhs: f32); + fn vector_abs_inplace(this: &mut [Self]); } mod internal { diff --git a/crates/simd/src/packed.rs b/crates/simd/src/packed.rs new file mode 100644 index 00000000..645fb8df --- /dev/null +++ b/crates/simd/src/packed.rs @@ -0,0 +1,92 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub fn u2_u2_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + u2_u2_reduce_sum_of_xy::u2_u2_reduce_sum_of_xy(s, t) +} + +mod u2_u2_reduce_sum_of_xy { + #[crate::multiversion("v4", "v3", "v2", "a2")] + pub fn u2_u2_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + assert_eq!(s.len(), t.len()); + let n = s.len(); + let mut result_0 = 0_u32; + let mut result_1 = 0_u32; + let mut result_2 = 0_u32; + let mut result_3 = 0_u32; + for i in 0..n { + let (s, t) = (s[i], t[i]); + result_0 += (((s >> 0) & 3) * ((t >> 0) & 3)) as u32; + result_1 += (((s >> 2) & 3) * ((t >> 2) & 3)) as u32; + result_2 += (((s >> 4) & 3) * ((t >> 4) & 3)) as u32; + result_3 += (((s >> 6) & 3) * ((t >> 6) & 3)) as u32; + } + result_0 + result_1 + result_2 + result_3 + } +} + +pub fn u4_u4_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + u4_u4_reduce_sum_of_xy::u4_u4_reduce_sum_of_xy(s, t) +} + +mod u4_u4_reduce_sum_of_xy { + #[crate::multiversion("v4", "v3", "v2", "a2")] + pub fn u4_u4_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + assert_eq!(s.len(), t.len()); + let n = s.len(); + let mut result_0 = 0; + let mut result_1 = 0; + for i in 0..n { + let (s, t) = (s[i], t[i]); + result_0 += (((s >> 0) & 15) * ((t >> 0) & 15)) as u32; + result_1 += (((s >> 4) & 15) * ((t >> 4) & 15)) as u32; + } + result_0 + result_1 + } +} + +pub fn u2_u8_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + u2_u8_reduce_sum_of_xy::u2_u8_reduce_sum_of_xy(s, t) +} + +mod u2_u8_reduce_sum_of_xy { + #[crate::multiversion("v4", "v3", "v2", "a2")] + pub fn u2_u8_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + assert_eq!(s.len(), t.len().div_ceil(4)); + let mut result_0 = 0_u32; + let mut result_1 = 0_u32; + let mut result_2 = 0_u32; + let mut result_3 = 0_u32; + let (arrays, remainder) = t.as_chunks::<4>(); + assert_eq!(s.len(), arrays.len()); + let n = arrays.len(); + for i in 0..n { + let (s, t) = (s[i], arrays[i]); + result_0 += ((s >> 0) & 3) as u32 * t[0] as u32; + result_1 += ((s >> 2) & 3) as u32 * t[1] as u32; + result_2 += ((s >> 4) & 3) as u32 * t[2] as u32; + result_3 += ((s >> 6) & 3) as u32 * t[3] as u32; + } + let mut buffer = [0u8; 4]; + if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + let (s, t) = (s[n], buffer); + result_0 += ((s >> 0) & 3) as u32 * t[0] as u32; + result_1 += ((s >> 2) & 3) as u32 * t[1] as u32; + result_2 += ((s >> 4) & 3) as u32 * t[2] as u32; + result_3 += ((s >> 6) & 3) as u32 * t[3] as u32; + } + result_0 + result_1 + result_2 + result_3 + } +} diff --git a/crates/simd/src/rotate.rs b/crates/simd/src/rotate.rs index 7a211e51..00642bd3 100644 --- a/crates/simd/src/rotate.rs +++ b/crates/simd/src/rotate.rs @@ -33,12 +33,12 @@ pub mod flip { pub fn flip(bits: &[u64; 1024], result: &mut [f32]) { use std::hint::select_unpredictable; let result: &mut [u32] = unsafe { std::mem::transmute(result) }; - let (slice, remainder) = result.as_chunks_mut::<64>(); - let n = slice.len(); + let (arrays, remainder) = result.as_chunks_mut::<64>(); + let n = arrays.len(); assert!(n <= 1024); for i in 0..n { for j in 0..64 { - slice[i][j] ^= select_unpredictable((bits[i] & (1 << j)) != 0, 0x80000000, 0); + arrays[i][j] ^= select_unpredictable((bits[i] & (1 << j)) != 0, 0x80000000, 0); } } for j in 0..remainder.len() { diff --git a/crates/vchordg/Cargo.toml b/crates/vchordg/Cargo.toml new file mode 100644 index 00000000..a31d3068 --- /dev/null +++ b/crates/vchordg/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "vchordg" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +algo = { path = "../algo" } +always_equal = { path = "../always_equal" } +distance = { path = "../distance" } +k_means = { path = "../k_means" } +rabitq = { path = "../rabitq" } +simd = { path = "../simd" } +vector = { path = "../vector" } + +half.workspace = true +min-max-heap = "1.3.0" +paste.workspace = true +pin-project = "1" +rand.workspace = true +serde.workspace = true +validator.workspace = true +zerocopy.workspace = true + +[lints] +workspace = true diff --git a/crates/vchordg/src/build.rs b/crates/vchordg/src/build.rs new file mode 100644 index 00000000..571575e9 --- /dev/null +++ b/crates/vchordg/src/build.rs @@ -0,0 +1,72 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::Opaque; +use crate::operator::Operator; +use crate::tuples::{MetaTuple, OptionPointer, Tuple}; +use crate::types::{VamanaIndexOptions, VectorOptions}; +use algo::{Page, PageGuard, RelationWrite}; + +pub fn build( + vector_options: VectorOptions, + index_options: VamanaIndexOptions, + index: &R, +) where + R::Page: Page, +{ + let mut meta_guard = index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + link: 1, + _padding_0: Default::default(), + }, + false, + ); + assert_eq!(meta_guard.id(), 0); + let vertex_guard = index.extend( + Opaque { + next: u32::MAX, + skip: 1, + link: 2, + _padding_0: Default::default(), + }, + true, + ); + assert_eq!(vertex_guard.id(), 1); + drop(vertex_guard); + let vector_guard = index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + link: u32::MAX, + _padding_0: Default::default(), + }, + false, + ); + assert_eq!(vector_guard.id(), 2); + drop(vector_guard); + let serialized = MetaTuple::serialize(&MetaTuple { + dims: vector_options.dims, + m: index_options.m, + alpha: index_options.alpha, + ef_construction: index_options.ef_construction, + beam_construction: index_options.beam_construction, + start: OptionPointer::NONE, + }); + let i = meta_guard + .alloc(&serialized) + .expect("implementation: a free page cannot accommodate a single tuple"); + assert_eq!(i, 1); +} diff --git a/crates/vchordg/src/bulkdelete.rs b/crates/vchordg/src/bulkdelete.rs new file mode 100644 index 00000000..ab84ce3f --- /dev/null +++ b/crates/vchordg/src/bulkdelete.rs @@ -0,0 +1,62 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::Opaque; +use crate::operator::Operator; +use crate::tuples::{VertexTuple, WithReader, WithWriter}; +use algo::{Page, RelationLength, RelationRead, RelationWrite}; +use std::num::NonZero; + +pub fn bulkdelete( + index: &R, + check: impl Fn(), + callback: impl Fn(NonZero) -> bool, +) where + R::Page: Page, +{ + let meta_guard = index.read(0); + let link = meta_guard.get_opaque().link; + drop(meta_guard); + let mut current = link; + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + let bytes = read.get(i).expect("data corruption"); + let tuple = VertexTuple::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + break 'flag true; + } + } + false + }; + if flag { + drop(read); + let mut write = index.write(current, false); + for i in 1..=write.len() { + let bytes = write.get_mut(i).expect("data corruption"); + let mut tuple = VertexTuple::deserialize_mut(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + *p = None; + } + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; + } + } +} diff --git a/crates/vchordg/src/candidates.rs b/crates/vchordg/src/candidates.rs new file mode 100644 index 00000000..fb77a012 --- /dev/null +++ b/crates/vchordg/src/candidates.rs @@ -0,0 +1,72 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use algo::{Fetch, RelationRead}; +use std::collections::{BinaryHeap, VecDeque}; +use std::marker::PhantomData; + +pub struct Candidates<'a, P, F, R> +where + P: Prefetcher<'a>, + P::Item: Fetch, +{ + beam: usize, + front: Option

, + heap: BinaryHeap, + prefetch: F, + _phantom: PhantomData &'a R>, +} + +impl<'a, P, F, R> Candidates<'a, P, F, R> +where + P: Prefetcher<'a, R = R>, + P::Item: Fetch + Ord, + F: PrefetcherSequenceFamily<'a, R, P> = P>, + R: RelationRead, +{ + pub fn new(beam: usize, prefetch: F) -> Self { + assert_ne!(beam, 0); + Self { + beam, + front: None, + heap: BinaryHeap::new(), + prefetch, + _phantom: PhantomData, + } + } + pub fn pop(&mut self) -> Option<(P::Item, P::Guards)> { + if let Some(front) = self.front.as_mut() + && let Some(item) = front.next() + { + return Some(item); + } + self.front = Some( + self.prefetch.prefetch( + (0..self.beam) + .flat_map(|_| self.heap.pop()) + .collect::>(), + ), + ); + if let Some(front) = self.front.as_mut() + && let Some(item) = front.next() + { + return Some(item); + } + None + } + pub fn push(&mut self, item: P::Item) { + self.heap.push(item); + } +} diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs new file mode 100644 index 00000000..7d3539cd --- /dev/null +++ b/crates/vchordg/src/insert.rs @@ -0,0 +1,577 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::Opaque; +use crate::candidates::Candidates; +use crate::operator::{CloneAccessor, Operator, Vector}; +use crate::results::Results; +use crate::tuples::*; +use crate::visited::Visited; +use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use algo::{Bump, Page, PageGuard, RelationRead, RelationWrite}; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::VecDeque; +use std::iter::{repeat, zip}; +use std::num::{NonZero, Wrapping}; +use vector::{VectorBorrowed, VectorOwned}; + +pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( + index: &'b R, + vector: ::Borrowed<'b>, + payload: NonZero, + bump: &'b impl Bump, + mut prefetch_vertices: impl PrefetcherSequenceFamily<'b, R> + 'b, + prefetch_vectors: impl PrefetcherSequenceFamily<'b, R> + 'b, +) where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dims = meta_tuple.dims(); + assert_eq!(dims, vector.dims(), "unmatched dimensions"); + let start = meta_tuple.start(); + let m = meta_tuple.m(); + let alpha = meta_tuple.alpha(); + let ef = meta_tuple.ef_construction(); + let beam = meta_tuple.beam_construction(); + drop(meta_guard); + let (vector_pointers, vertex_pointer) = insert_tuples::(index, vector, payload, m); + let start = if start.into_inner().is_none() { + let mut meta_guard = index.write(0, false); + let meta_bytes = meta_guard.get_mut(1).expect("data corruption"); + let mut meta_tuple = MetaTuple::deserialize_mut(meta_bytes); + let dst = meta_tuple.start(); + if dst.into_inner().is_none() { + *dst = OptionPointer::some(vertex_pointer); + return; + } else { + *dst + } + } else { + start + }; + let lut = O::Vector::preprocess(vector); + let mut visited = Visited::new(); + let mut candidates = Candidates::new(beam as usize, prefetch_vectors); + let Some(s) = start.into_inner() else { + return; + }; + { + visited.insert(s); + let vertex_guard = index.read(s.0); + let Some(vertex_bytes) = vertex_guard.get(s.1) else { + // the link is broken + return; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_s = bump.alloc_slice(vertex_tuple.pointers()); + let score_s = O::process((vertex_tuple.metadata(), vertex_tuple.elements()), &lut); + candidates.push((Reverse(score_s), AlwaysEqual((pointers_s, s)))); + } + let mut iter = std::iter::from_fn(|| { + while let Some(((_, AlwaysEqual((pointers_u, u))), guards)) = candidates.pop() { + let Ok((dis_u, _payload_u, outs_u)) = + rerank::(guards, pointers_u, vector, &mut visited) + else { + // the link is broken + continue; + }; + let mut iterator = prefetch_vertices.prefetch(outs_u); + while let Some((v, guards)) = iterator.next() { + visited.insert(v); + let vertex_guard = { + let mut guards = guards; + let r = guards.next().expect("internal"); + assert!(guards.next().is_none(), "internal"); + drop(guards); + r + }; + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); + let score_v = O::process((vertex_tuple.metadata(), vertex_tuple.elements()), &lut); + candidates.push((Reverse(score_v), AlwaysEqual((pointers_v, v)))); + } + return Some((Reverse(dis_u), AlwaysEqual((pointers_u, u)))); + } + None + }); + let mut results = Results::new(ef as _); + for element @ (Reverse(dis_c), _) in iter.by_ref() { + results.push(element); + if results + .peek_ef_th() + .map(|dis_e| dis_e < dis_c) + .unwrap_or_default() + { + break; + } + } + let trace = results + .into_inner() + .0 + .into_iter() + .map(|(dis_v, x)| (Reverse(dis_v), x)) + .collect::>(); + let outs = crate::prune::robust_prune( + |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), + (bump.alloc_slice(&vector_pointers), vertex_pointer), + trace.into_iter().flat_map(|item| { + use algo::accessor::Accessor1; + let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; + let m = strict_sub(pointers_u.len(), 1); + let mut accessor = CloneAccessor::::default(); + for i in 0..m { + let vector_guard = index.read(pointers_u[i].into_inner().0); + let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { + // the link is broken + return None; + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_1(segment) = vector_tuple else { + // the link is broken + return None; + }; + if segment.index() as usize != i { + // the link is broken + return None; + } + accessor.push(segment.elements()); + } + let vector_u; + { + let vector_guard = index.read(pointers_u[m].into_inner().0); + let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { + // the link is broken + return None; + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(segment) = vector_tuple else { + // the link is broken + return None; + }; + accessor.push(segment.elements()); + vector_u = accessor.finish(*segment.metadata()); + } + Some(((Reverse(dis_u), AlwaysEqual((pointers_u, u))), vector_u)) + }), + m, + alpha, + |(_, u)| *u, + ); + { + let mut vector_guard = index.write(vector_pointers.last().unwrap().into_inner().0, false); + let Some(vector_bytes) = + vector_guard.get_mut(vector_pointers.last().unwrap().into_inner().1) + else { + // the link is broken + return; + }; + let vector_tuple = VectorTuple::::deserialize_mut(vector_bytes); + let VectorTupleWriter::_0(mut vector_tuple) = vector_tuple else { + // the link is broken + return; + }; + let iterator = outs + .iter() + .map(|&(Reverse(dis_v), AlwaysEqual((_, v)))| OptionNeighbour::some(v, dis_v)) + .chain(repeat(OptionNeighbour::NONE)); + for (hole, fill) in zip(vector_tuple.neighbours().iter_mut(), iterator) { + *hole = fill; + } + } + for (Reverse(dis_t), AlwaysEqual((pointers_u, u))) in outs { + while add_link::( + index, + (pointers_u, u), + (vertex_pointer, dis_t), + m, + alpha, + bump, + ) == Ok(false) + {} + } +} + +fn rerank<'r, R: RelationRead, O: Operator>( + mut guards: impl Iterator>, + pointers_u: &mut [Pointer], + vector: ::Borrowed<'_>, + visited: &mut Visited, +) -> Result<(Distance, Option>, VecDeque<(u32, u16)>), ()> { + use algo::accessor::{Accessor1, LAccess}; + let m = strict_sub(pointers_u.len(), 1); + let mut accessor = LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()); + for i in 0..m { + let vector_guard = guards.next().expect("internal"); + let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_1(segment) = vector_tuple else { + // the link is broken + return Err(()); + }; + if segment.index() as usize != i { + // the link is broken + return Err(()); + } + accessor.push(segment.elements()); + } + let dis_u; + let payload_u; + let neighbours_u; + { + let vector_guard = guards.next().expect("internal"); + let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(segment) = vector_tuple else { + // the link is broken + return Err(()); + }; + accessor.push(segment.elements()); + dis_u = accessor.finish(*segment.metadata()); + payload_u = segment.payload(); + neighbours_u = segment + .neighbours() + .iter() + .flat_map(|neighbour| neighbour.into_inner()) + .map(|(v, _)| v) + .filter(|&v| !visited.contains(v)) + .collect::>(); + } + Ok((dis_u, payload_u, neighbours_u)) +} + +fn insert_tuples<'b, R: RelationRead + RelationWrite, O: Operator>( + index: &'b R, + vector: ::Borrowed<'b>, + payload: NonZero, + m: u32, +) -> (Vec, (u32, u16)) +where + R::Page: Page, +{ + let vector_bytes = { + let (left, right) = O::Vector::split(vector, m as _); + let left = left.into_iter().enumerate().map(|(index, elements)| { + VectorTuple::serialize(&VectorTuple::::_1 { + payload: Some(payload), + elements: elements.to_vec(), + index: index as u32, + }) + }); + let right = VectorTuple::serialize(&VectorTuple::::_0 { + payload: Some(payload), + elements: right.0.to_vec(), + metadata: right.1, + neighbours: vec![OptionNeighbour::NONE; m as usize], + version: Wrapping(rand::random()), + }); + left.chain(std::iter::once(right)).collect::>() + }; + let mut vertex_bytes = { + let code = O::Vector::code(vector); + VertexTuple::serialize(&VertexTuple { + metadata: code.0.into_array(), + payload: Some(payload), + elements: rabitq::original::binary::pack_code(&code.1), + pointers: vec![Pointer::new((u32::MAX, 0)); vector_bytes.len()], // a sentinel value + }) + }; + append_vertex_tuple(index, 1, vertex_bytes.len(), |mut vertex_guard| { + let vector_pointers = vector_bytes + .iter() + .map(|vector_bytes| { + append_vector_tuple( + index, + vertex_guard.get_opaque().link, + vector_bytes.len(), + |mut vector_guard| { + let i = vector_guard.alloc(vector_bytes).expect( + "implementation: a free page cannot accommodate a single tuple", + ); + Pointer::new((vector_guard.id(), i)) + }, + ) + }) + .collect::>(); + { + let mut vertex_tuple = VertexTuple::deserialize_mut(&mut vertex_bytes); + vertex_tuple.pointers().copy_from_slice(&vector_pointers); + } + let vertex_pointer = { + let i = vertex_guard + .alloc(&vertex_bytes) + .expect("implementation: a free page cannot accommodate a single tuple"); + (vertex_guard.id(), i) + }; + (vector_pointers, vertex_pointer) + }) +} + +#[allow(clippy::collapsible_else_if)] +fn append_vertex_tuple<'r, R: RelationRead + RelationWrite, T>( + index: &'r R, + first: u32, + size: usize, + f: impl FnOnce(R::WriteGuard<'r>) -> T, +) -> T +where + R::Page: Page, +{ + if let Some(guard) = index.search(size) { + return f(guard); + } + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = index.read(current); + if read.freespace() as usize >= size || read.get_opaque().next == u32::MAX { + drop(read); + let mut write = index.write(current, true); + if write.freespace() as usize >= size { + return f(write); + } + if write.get_opaque().next == u32::MAX { + let link = index + .extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + link: u32::MAX, + _padding_0: Default::default(), + }, + false, + ) + .id(); + let extend = index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + link, + _padding_0: Default::default(), + }, + true, + ); + { write }.get_opaque_mut().next = extend.id(); + if extend.freespace() as usize >= size { + let fresh = extend.id(); + let result = f(extend); + let mut past = index.write(first, true); + past.get_opaque_mut().skip = fresh.max(past.get_opaque().skip); + return result; + } else { + panic!("implementation: a clear page cannot accommodate a single tuple"); + } + } + if current == first && write.get_opaque().skip != first { + current = write.get_opaque().skip; + } else { + current = write.get_opaque().next; + } + } else { + if current == first && read.get_opaque().skip != first { + current = read.get_opaque().skip; + } else { + current = read.get_opaque().next; + } + } + } +} + +fn append_vector_tuple<'r, R: RelationRead + RelationWrite, T>( + index: &'r R, + first: u32, + size: usize, + f: impl FnOnce(R::WriteGuard<'r>) -> T, +) -> T +where + R::Page: Page, +{ + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = index.read(current); + if read.freespace() as usize >= size || read.get_opaque().next == u32::MAX { + drop(read); + let mut write = index.write(current, false); + if write.freespace() as usize >= size { + return f(write); + } + if write.get_opaque().next == u32::MAX { + let extend = index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + link: u32::MAX, + _padding_0: Default::default(), + }, + false, + ); + { write }.get_opaque_mut().next = extend.id(); + if extend.freespace() as usize >= size { + return f(extend); + } else { + panic!("implementation: a clear page cannot accommodate a single tuple"); + } + } + current = write.get_opaque().next + } else { + current = read.get_opaque().next + } + } +} + +fn add_link( + index: &R, + (pointers_u, u): (&mut [Pointer], (u32, u16)), + new: ((u32, u16), Distance), + m: u32, + alpha: f32, + bump: &impl Bump, +) -> Result { + let vector_guard = index.read(pointers_u.last().unwrap().into_inner().0); + let Some(vector_bytes) = vector_guard.get(pointers_u.last().unwrap().into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + let check = vector_tuple + .neighbours() + .iter() + .flat_map(|neighbour| neighbour.into_inner()) + .map(|(v, dis_v)| (Reverse(dis_v), AlwaysEqual(v))) + .collect::>(); + let trace = check + .iter() + .copied() + .chain(std::iter::once((Reverse(new.1), AlwaysEqual(new.0)))) + .collect::>(); + let version = vector_tuple.version(); + drop(vector_guard); + let outs = crate::prune::robust_prune( + |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), + (bump.alloc_slice(pointers_u), u), + trace + .into_iter() + .flat_map(|item| { + let (Reverse(dis_u), AlwaysEqual(u)) = item; + let vertex_guard = index.read(u.0); + let Some(vertex_bytes) = vertex_guard.get(u.1) else { + // the link is broken + return None; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_u = bump.alloc_slice(vertex_tuple.pointers()); + Some((Reverse(dis_u), AlwaysEqual((pointers_u, u)))) + }) + .flat_map(|item| { + use algo::accessor::Accessor1; + let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; + let m = strict_sub(pointers_u.len(), 1); + let mut accessor = CloneAccessor::::default(); + for i in 0..m { + let vector_guard = index.read(pointers_u[i].into_inner().0); + let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { + // the link is broken + return None; + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_1(segment) = vector_tuple else { + // the link is broken + return None; + }; + if segment.index() as usize != i { + // the link is broken + return None; + } + accessor.push(segment.elements()); + } + let vector_u; + { + let vector_guard = index.read(pointers_u[m].into_inner().0); + let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { + // the link is broken + return None; + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(segment) = vector_tuple else { + // the link is broken + return None; + }; + accessor.push(segment.elements()); + vector_u = accessor.finish(*segment.metadata()); + } + Some(((Reverse(dis_u), AlwaysEqual((pointers_u, u))), vector_u)) + }), + m, + alpha, + |(_, u)| *u, + ); + // fast path + if outs + .iter() + .map(|&(x, AlwaysEqual((_, v)))| (x, AlwaysEqual(v))) + .eq(check) + { + return Ok(true); + } + let mut vector_guard = index.write(pointers_u.last().unwrap().into_inner().0, false); + let Some(vector_bytes) = vector_guard.get_mut(pointers_u.last().unwrap().into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_mut(vector_bytes); + let VectorTupleWriter::_0(mut vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + if *vector_tuple.version() != version { + return Ok(false); + } else { + *vector_tuple.version() += 1; + } + let filling = outs + .iter() + .map(|&(Reverse(dis_v), AlwaysEqual((_, v)))| OptionNeighbour::some(v, dis_v)) + .chain(repeat(OptionNeighbour::NONE)); + for (hole, fill) in zip(vector_tuple.neighbours().iter_mut(), filling) { + *hole = fill; + } + Ok(true) +} + +// Emulate unstable library feature `strict_overflow_ops`. +// See https://github.com/rust-lang/rust/issues/118260. + +#[inline] +pub const fn strict_sub(lhs: usize, rhs: usize) -> usize { + let (a, b) = lhs.overflowing_sub(rhs); + if b { panic!() } else { a } +} diff --git a/crates/vchordg/src/lib.rs b/crates/vchordg/src/lib.rs new file mode 100644 index 00000000..43550656 --- /dev/null +++ b/crates/vchordg/src/lib.rs @@ -0,0 +1,44 @@ +#![allow(clippy::type_complexity)] + +mod build; +mod bulkdelete; +mod candidates; +mod insert; +mod maintain; +mod prune; +mod results; +mod search; +mod tuples; +mod visited; + +pub mod operator; +pub mod types; + +use algo::tuples::Padding; +pub use build::build; +pub use bulkdelete::bulkdelete; +pub use insert::insert; +pub use maintain::maintain; +pub use search::search; + +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +#[repr(C, align(8))] +#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +pub struct Opaque { + pub next: u32, + pub skip: u32, + pub link: u32, + pub _padding_0: [Padding; 4], +} + +#[allow(unsafe_code)] +unsafe impl algo::Opaque for Opaque {} + +pub type Id = (u32, u16); + +impl algo::Fetch1 for tuples::Pointer { + fn fetch_1(&self) -> u32 { + self.into_inner().0 + } +} diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs new file mode 100644 index 00000000..61799850 --- /dev/null +++ b/crates/vchordg/src/maintain.rs @@ -0,0 +1,273 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::Opaque; +use crate::operator::{CloneAccessor, Operator}; +use crate::tuples::{ + MetaTuple, OptionNeighbour, Pointer, VectorTuple, VectorTupleReader, VectorTupleWriter, + VertexTuple, WithReader, WithWriter, +}; +use algo::{Bump, Page, PageGuard, RelationRead, RelationWrite}; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::VecDeque; +use std::iter::{repeat, zip}; +use std::num::Wrapping; +use vector::VectorOwned; + +pub fn maintain( + index: &R, + check: impl Fn(), + bump: &impl Bump, +) where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let m = meta_tuple.m(); + let alpha = meta_tuple.alpha(); + let start = meta_tuple.start().into_inner(); + let link = meta_guard.get_opaque().link; + drop(meta_guard); + // do it's best to remove broken edges + { + let mut current = link; + while current != u32::MAX { + check(); + let vertex_guard = index.read(current); + let mut members = Vec::new(); + for i in 1..=vertex_guard.len() { + if let Some(vertex_bytes) = vertex_guard.get(i) { + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + if vertex_tuple.payload().is_some() { + let pointers = bump.alloc_slice(vertex_tuple.pointers()); + members.push((pointers, (vertex_guard.id(), i))); + } + } + } + let next = { vertex_guard }.get_opaque().next; + for member in members { + while fix_link::(index, (member.0, member.1), m, alpha, bump) == Ok(false) {} + } + current = next; + } + } + // remove vertices and vectors + { + let mut current = link; + while current != u32::MAX { + check(); + let mut vertex_guard = index.write(current, true); + for i in 1..=vertex_guard.len() { + if let Some(bytes) = vertex_guard.get(i) { + let tuple = VertexTuple::deserialize_ref(bytes); + let p = tuple.payload(); + if p.is_none() && Some((current, i)) != start { + vertex_guard.free(i); + } + }; + } + current = vertex_guard.get_opaque().next; + { + let mut current = { vertex_guard }.get_opaque().link; + while current != u32::MAX { + check(); + let mut vector_guard = index.write(current, true); + for i in 1..=vector_guard.len() { + if let Some(bytes) = vector_guard.get(i) { + use crate::tuples::VectorTupleReader; + let tuple = VectorTuple::::deserialize_ref(bytes); + match tuple { + VectorTupleReader::_0(tuple) => { + let p = tuple.payload(); + if p.is_none() && Some((current, i)) != start { + vector_guard.free(i); + } + } + VectorTupleReader::_1(tuple) => { + let p = tuple.payload(); + if p.is_none() && Some((current, i)) != start { + vector_guard.free(i); + } + } + } + }; + } + current = vector_guard.get_opaque().next; + } + } + } + } +} + +fn read_tuple( + index: &R, + pointers_u: &mut [Pointer], +) -> Result<(O::Vector, VecDeque<((u32, u16), Distance)>, Wrapping), ()> { + use algo::accessor::Accessor1; + let m = strict_sub(pointers_u.len(), 1); + let mut accessor = CloneAccessor::::default(); + for i in 0..m { + let vector_guard = index.read(pointers_u[i].into_inner().0); + let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_1(segment) = vector_tuple else { + // the link is broken + return Err(()); + }; + if segment.index() as usize != i { + // the link is broken + return Err(()); + } + accessor.push(segment.elements()); + } + let vector_u; + let neighbours_u; + let version; + { + let vector_guard = index.read(pointers_u[m].into_inner().0); + let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(segment) = vector_tuple else { + // the link is broken + return Err(()); + }; + accessor.push(segment.elements()); + vector_u = accessor.finish(*segment.metadata()); + neighbours_u = segment + .neighbours() + .iter() + .flat_map(|neighbour| neighbour.into_inner()) + .collect::>(); + version = segment.version(); + } + Ok((vector_u, neighbours_u, version)) +} + +fn fix_link( + index: &R, + (pointers_u, u): (&mut [Pointer], (u32, u16)), + m: u32, + alpha: f32, + bump: &impl Bump, +) -> Result { + let Ok((vector_u, neighbours_u, version)) = read_tuple::(index, pointers_u) else { + // the link is broken + return Err(()); + }; + let mut trace = Vec::new(); + let mut extend = Vec::new(); + for &(v, dis_v) in neighbours_u.iter() { + let vertex_guard = index.read(v.0); + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); + let payload_v = vertex_tuple.payload(); + drop(vertex_guard); + let Ok((vector_v, neighbours_v, _)) = read_tuple::(index, pointers_v) else { + // the link is broken + continue; + }; + if payload_v.is_some() { + trace.push((v, dis_v, vector_v)); + } else { + let outs_v = neighbours_v.iter().map(|(v, _)| *v).collect::>(); + extend.extend(outs_v); + } + } + // fast path + if extend.is_empty() { + return Ok(true); + } + for v in extend { + let vertex_guard = index.read(v.0); + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); + let payload_v = vertex_tuple.payload(); + drop(vertex_guard); + let Ok((vector_v, _, _)) = read_tuple::(index, pointers_v) else { + // the link is broken + continue; + }; + if payload_v.is_some() { + let dis_v = O::distance(vector_u.as_borrowed(), vector_v.as_borrowed()); + trace.push((v, dis_v, vector_v)); + } + } + let outs = crate::prune::robust_prune( + |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), + (bump.alloc_slice(pointers_u), u), + trace.into_iter().flat_map(|item| { + let (u, dis_u, vector_u) = item; + let vertex_guard = index.read(u.0); + let Some(vertex_bytes) = vertex_guard.get(u.1) else { + // the link is broken + return None; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointer_u = bump.alloc_slice(vertex_tuple.pointers()); + Some(((Reverse(dis_u), AlwaysEqual((pointer_u, u))), vector_u)) + }), + m, + alpha, + |(_, u)| *u, + ); + let mut vector_guard = index.write(pointers_u.last().unwrap().into_inner().0, false); + let Some(vector_bytes) = vector_guard.get_mut(pointers_u.last().unwrap().into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_mut(vector_bytes); + let VectorTupleWriter::_0(mut vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + if *vector_tuple.version() != version { + return Ok(false); + } else { + *vector_tuple.version() += 1; + } + let iterator = outs + .iter() + .map(|&(Reverse(dis_v), AlwaysEqual((_, v)))| OptionNeighbour::some(v, dis_v)) + .chain(repeat(OptionNeighbour::NONE)); + for (hole, fill) in zip(vector_tuple.neighbours().iter_mut(), iterator) { + *hole = fill; + } + Ok(true) +} + +// Emulate unstable library feature `strict_overflow_ops`. +// See https://github.com/rust-lang/rust/issues/118260. + +#[inline] +pub const fn strict_sub(lhs: usize, rhs: usize) -> usize { + let (a, b) = lhs.overflowing_sub(rhs); + if b { panic!() } else { a } +} diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs new file mode 100644 index 00000000..63776c22 --- /dev/null +++ b/crates/vchordg/src/operator.rs @@ -0,0 +1,253 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S}; +use distance::Distance; +use half::f16; +use simd::Floating; +use std::fmt::Debug; +use std::marker::PhantomData; +use vector::vect::VectOwned; +use vector::{VectorBorrowed, VectorOwned}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +pub trait Vector: VectorOwned { + type Element: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; + type Metadata: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)); + fn pack(elements: Vec, metadata: Self::Metadata) -> Self; + + fn code(vector: Self::Borrowed<'_>) -> rabitq::b1::Code; + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b1::binary::BinaryLut; +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f32; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + (vector.slice(), ()) + } + + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)) { + let slice = vector.slice(); + let tailing = (size_of::() * m) + .next_multiple_of(crate::tuples::ALIGN); + assert!(tailing <= 8000); + if slice.len() <= (8000 - tailing) / size_of::() { + return (vec![], (slice, ())); + } + let (l, r) = slice.split_at(slice.len() - (8000 - tailing) / size_of::()); + ( + l.chunks(8000 / size_of::()).collect::>(), + (r, ()), + ) + } + + fn pack(elements: Vec, (): Self::Metadata) -> Self { + VectOwned::new(elements) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::b1::Code { + rabitq::b1::code(vector.slice()) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b1::binary::BinaryLut { + rabitq::b1::binary::preprocess(vector.slice()) + } +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f16; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + (vector.slice(), ()) + } + + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)) { + let slice = vector.slice(); + let tailing = (size_of::() * m) + .next_multiple_of(crate::tuples::ALIGN); + assert!(tailing <= 8000); + if slice.len() <= (8000 - tailing) / size_of::() { + return (vec![], (slice, ())); + } + let (l, r) = slice.split_at(slice.len() - (8000 - tailing) / size_of::()); + ( + l.chunks(8000 / size_of::()).collect::>(), + (r, ()), + ) + } + + fn pack(elements: Vec, (): Self::Metadata) -> Self { + VectOwned::new(elements) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::b1::Code { + rabitq::b1::code(&f16::vector_to_f32(vector.slice())) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b1::binary::BinaryLut { + rabitq::b1::binary::preprocess(&f16::vector_to_f32(vector.slice())) + } +} + +pub trait Operator: 'static + Debug + Copy { + type Vector: Vector; + + type DistanceAccessor: Default + + Accessor2< + ::Element, + ::Element, + ::Metadata, + ::Metadata, + Output = Distance, + >; + + fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance; + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance; +} + +#[derive(Debug)] +pub struct Op(PhantomData V>, PhantomData D>); + +impl Clone for Op { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for Op {} + +impl Operator for Op, L2S> { + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, L2S>; + + fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance { + use rabitq::b1::CodeMetadata; + let value = rabitq::b1::binary::accumulate(code.1, &lut.1); + let (distance,) = + rabitq::b1::binary::half_process_l2(value, CodeMetadata::from_array(code.0), lut.0); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_l2s(rhs) + } +} + +impl Operator for Op, Dot> { + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, Dot>; + + fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance { + use rabitq::b1::CodeMetadata; + let value = rabitq::b1::binary::accumulate(code.1, &lut.1); + let (distance,) = + rabitq::b1::binary::half_process_dot(value, CodeMetadata::from_array(code.0), lut.0); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_dot(rhs) + } +} + +impl Operator for Op, L2S> { + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, L2S>; + + fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance { + use rabitq::b1::CodeMetadata; + let value = rabitq::b1::binary::accumulate(code.1, &lut.1); + let (distance,) = + rabitq::b1::binary::half_process_l2(value, CodeMetadata::from_array(code.0), lut.0); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_l2s(rhs) + } +} + +impl Operator for Op, Dot> { + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, Dot>; + + fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance { + use rabitq::b1::CodeMetadata; + let value = rabitq::b1::binary::accumulate(code.1, &lut.1); + let (distance,) = + rabitq::b1::binary::half_process_dot(value, CodeMetadata::from_array(code.0), lut.0); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_dot(rhs) + } +} + +#[derive(Debug, Clone)] +pub struct CloneAccessor(Vec); + +impl Default for CloneAccessor { + fn default() -> Self { + Self(Vec::new()) + } +} + +impl Accessor1 for CloneAccessor { + type Output = V; + + fn push(&mut self, input: &[V::Element]) { + self.0.extend(input); + } + + fn finish(self, metadata: V::Metadata) -> Self::Output { + V::pack(self.0, metadata) + } +} diff --git a/crates/vchordg/src/prune.rs b/crates/vchordg/src/prune.rs new file mode 100644 index 00000000..935120b5 --- /dev/null +++ b/crates/vchordg/src/prune.rs @@ -0,0 +1,56 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; + +pub fn robust_prune( + mut d: impl FnMut(&V, &V) -> Distance, + u: T, + outs: impl Iterator, AlwaysEqual), V)>, + m: u32, + alpha: f32, + key: impl Fn(&T) -> K, +) -> Vec<(Reverse, AlwaysEqual)> { + // V ← (V ∪ Nout(p)) \ {p} + let mut trace = outs.collect::>(); + trace.sort_by_key(|((_, AlwaysEqual(v)), _)| key(v)); + trace.dedup_by_key(|((_, AlwaysEqual(v)), _)| key(v)); + trace.retain(|((_, AlwaysEqual(v)), _)| key(v) != key(&u)); + trace.sort_by_key(|&((Reverse(d), _), _)| d); + // Nout(p) ← ∅ + let mut result = Vec::new(); + let mut pruned = Vec::new(); + for ((Reverse(dis_u), AlwaysEqual(u)), vector_u) in trace { + if result.len() == m as usize { + break; + } + let check = result + .iter() + .map(|(_, vector_v)| d(&vector_u, vector_v)) + .all(|dis| dis_u.to_f32() < alpha * dis.to_f32()); + if check { + result.push(((Reverse(dis_u), AlwaysEqual(u)), vector_u)); + } else { + pruned.push(((Reverse(dis_u), AlwaysEqual(u)), vector_u)); + } + } + result + .into_iter() + .chain(pruned) + .map(|(x, _)| x) + .take(m as _) + .collect() +} diff --git a/crates/vchordg/src/results.rs b/crates/vchordg/src/results.rs new file mode 100644 index 00000000..c04b6fa1 --- /dev/null +++ b/crates/vchordg/src/results.rs @@ -0,0 +1,73 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use always_equal::AlwaysEqual; +use distance::Distance; +use min_max_heap::MinMaxHeap; +use std::cmp::Reverse; +use std::collections::BinaryHeap; + +pub struct Results { + ef: usize, + front: MinMaxHeap<(Distance, AlwaysEqual)>, + back: BinaryHeap<(Reverse, AlwaysEqual)>, + // len(front) < ef, then len(back) == 0 + // len(front) == ef, then max(front) < min(back) +} + +impl Results { + pub fn new(ef: usize) -> Self { + assert!(ef > 0, "ef must be positive integer"); + Self { + ef, + front: MinMaxHeap::with_capacity(ef), + back: BinaryHeap::new(), + } + } + pub fn push(&mut self, item: (Reverse, AlwaysEqual)) { + if self.front.len() < self.ef { + self.front.push((item.0.0, item.1)); + } else { + let item = self.front.push_pop_max((item.0.0, item.1)); + self.back.push((Reverse(item.0), item.1)); + } + } + pub fn peek_ef_th(&mut self) -> Option { + if self.front.len() < self.ef { + None + } else { + self.front.peek_max().map(|&(d, _)| d) + } + } + #[allow(clippy::collapsible_else_if)] + pub fn pop_min(&mut self) -> Option<(Distance, AlwaysEqual)> { + if self.front.len() < self.ef { + self.front.pop_min() + } else { + if let Some(item) = self.back.pop() { + Some(self.front.push_pop_min((item.0.0, item.1))) + } else { + self.front.pop_min() + } + } + } + pub fn into_inner( + self, + ) -> ( + Vec<(Distance, AlwaysEqual)>, + Vec<(Reverse, AlwaysEqual)>, + ) { + (self.front.into_vec(), self.back.into_vec()) + } +} diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs new file mode 100644 index 00000000..8d59246a --- /dev/null +++ b/crates/vchordg/src/search.rs @@ -0,0 +1,179 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::Opaque; +use crate::candidates::Candidates; +use crate::operator::{Operator, Vector}; +use crate::results::Results; +use crate::tuples::*; +use crate::visited::Visited; +use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use algo::{Bump, Page, RelationRead}; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::VecDeque; +use std::num::NonZero; +use vector::{VectorBorrowed, VectorOwned}; + +pub fn search<'b, R: RelationRead, O: Operator>( + index: &'b R, + vector: ::Borrowed<'b>, + ef_search: u32, + beam_search: u32, + bump: &'b impl Bump, + mut prefetch_vertices: impl PrefetcherSequenceFamily<'b, R> + 'b, + prefetch_vectors: impl PrefetcherSequenceFamily<'b, R> + 'b, +) -> Box)> + 'b> +where + R::Page: Page, +{ + let _ = bump; + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dims = meta_tuple.dims(); + let start = meta_tuple.start(); + assert_eq!(dims, vector.dims(), "unmatched dimensions"); + let ef = ef_search; + let beam = beam_search; + drop(meta_guard); + let lut = O::Vector::preprocess(vector); + let mut visited = Visited::new(); + let mut candidates = Candidates::new(beam as usize, prefetch_vectors); + let Some(s) = start.into_inner() else { + return Box::new(std::iter::empty()); + }; + { + visited.insert(s); + let vertex_guard = index.read(s.0); + let Some(vertex_bytes) = vertex_guard.get(s.1) else { + // the link is broken + return Box::new(std::iter::empty()); + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_s = bump.alloc_slice(vertex_tuple.pointers()); + let score_s = O::process((vertex_tuple.metadata(), vertex_tuple.elements()), &lut); + candidates.push((Reverse(score_s), AlwaysEqual(pointers_s))); + } + let mut iter = std::iter::from_fn(move || { + while let Some(((_, AlwaysEqual(pointers_u)), guards)) = candidates.pop() { + let Ok((dis_u, payload_u, outs_u)) = + rerank::(guards, pointers_u, vector, &mut visited) + else { + // the link is broken + continue; + }; + let mut iterator = prefetch_vertices.prefetch(outs_u); + while let Some((v, guards)) = iterator.next() { + visited.insert(v); + let vertex_guard = { + let mut guards = guards; + let r = guards.next().expect("internal"); + assert!(guards.next().is_none(), "internal"); + drop(guards); + r + }; + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); + let score_v = O::process((vertex_tuple.metadata(), vertex_tuple.elements()), &lut); + candidates.push((Reverse(score_v), AlwaysEqual(pointers_v))); + } + return Some((Reverse(dis_u), AlwaysEqual(payload_u))); + } + None + }); + let mut results = Results::new(ef as _); + let search = std::iter::from_fn(move || { + for element @ (Reverse(dis_c), _) in iter.by_ref() { + results.push(element); + if results + .peek_ef_th() + .map(|dis_e| dis_e < dis_c) + .unwrap_or_default() + { + break; + } + } + results.pop_min() + }); + Box::new(search.filter_map(|(dis_u, AlwaysEqual(payload_u))| Some((dis_u, payload_u?)))) +} + +fn rerank<'r, R: RelationRead, O: Operator>( + mut guards: impl Iterator>, + pointers_u: &mut [Pointer], + vector: ::Borrowed<'_>, + visited: &mut Visited, +) -> Result<(Distance, Option>, VecDeque<(u32, u16)>), ()> { + use algo::accessor::{Accessor1, LAccess}; + let m = strict_sub(pointers_u.len(), 1); + let mut accessor = LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()); + for i in 0..m { + let vector_guard = guards.next().expect("internal"); + let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_1(segment) = vector_tuple else { + // the link is broken + return Err(()); + }; + if segment.index() as usize != i { + // the link is broken + return Err(()); + } + accessor.push(segment.elements()); + } + let dis_u; + let payload_u; + let neighbours_u; + { + let vector_guard = guards.next().expect("internal"); + let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(segment) = vector_tuple else { + // the link is broken + return Err(()); + }; + accessor.push(segment.elements()); + dis_u = accessor.finish(*segment.metadata()); + payload_u = segment.payload(); + neighbours_u = segment + .neighbours() + .iter() + .flat_map(|neighbour| neighbour.into_inner()) + .map(|(v, _)| v) + .filter(|&v| !visited.contains(v)) + .collect::>(); + } + Ok((dis_u, payload_u, neighbours_u)) +} + +// Emulate unstable library feature `strict_overflow_ops`. +// See https://github.com/rust-lang/rust/issues/118260. + +#[inline] +pub const fn strict_sub(lhs: usize, rhs: usize) -> usize { + let (a, b) = lhs.overflowing_sub(rhs); + if b { panic!() } else { a } +} diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs new file mode 100644 index 00000000..2d82a130 --- /dev/null +++ b/crates/vchordg/src/tuples.rs @@ -0,0 +1,726 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::operator::Vector; +use algo::tuples::{Bool, MutChecker, Padding, RefChecker}; +use distance::Distance; +use std::num::{NonZero, Wrapping}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +pub const ALIGN: usize = 8; +pub type Tag = u64; +const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordg\0"); +const VERSION: u64 = 10; + +pub trait Tuple: 'static { + fn serialize(&self) -> Vec; +} + +pub trait WithReader: Tuple { + type Reader<'a>; + fn deserialize_ref(source: &[u8]) -> Self::Reader<'_>; +} + +pub trait WithWriter: Tuple { + type Writer<'a>; + fn deserialize_mut(source: &mut [u8]) -> Self::Writer<'_>; +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct MetaTupleHeader { + version: u64, + dims: u32, + m: u32, + alpha: f32, + ef_construction: u32, + beam_construction: u32, + _padding_0: [Padding; 4], + start: OptionPointer, +} + +pub struct MetaTuple { + pub dims: u32, + pub m: u32, + pub alpha: f32, + pub ef_construction: u32, + pub beam_construction: u32, + pub start: OptionPointer, +} + +impl Tuple for MetaTuple { + #[allow(clippy::match_single_binding)] + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + MetaTuple { + dims, + m, + alpha, + ef_construction, + beam_construction, + start, + } => { + buffer.extend((MAGIC as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + MetaTupleHeader { + version: VERSION, + dims: *dims, + m: *m, + alpha: *alpha, + ef_construction: *ef_construction, + beam_construction: *beam_construction, + start: *start, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for MetaTuple { + type Reader<'a> = MetaTupleReader<'a>; + fn deserialize_ref(source: &[u8]) -> MetaTupleReader<'_> { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + MAGIC => { + let checker = RefChecker::new(source); + if VERSION != *checker.prefix::(size_of::()) { + panic!("deserialization: bad version number"); + } + let header: &MetaTupleHeader = checker.prefix(size_of::()); + MetaTupleReader { header } + } + _ => panic!("deserialization: bad magic number"), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct MetaTupleReader<'a> { + header: &'a MetaTupleHeader, +} + +impl<'a> MetaTupleReader<'a> { + pub fn dims(self) -> u32 { + self.header.dims + } + pub fn m(self) -> u32 { + self.header.m + } + pub fn alpha(self) -> f32 { + self.header.alpha + } + pub fn ef_construction(self) -> u32 { + self.header.ef_construction + } + pub fn beam_construction(self) -> u32 { + self.header.beam_construction + } + pub fn start(self) -> OptionPointer { + self.header.start + } +} + +impl WithWriter for MetaTuple { + type Writer<'a> = MetaTupleWriter<'a>; + fn deserialize_mut(source: &mut [u8]) -> MetaTupleWriter<'_> { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + MAGIC => { + let checker = RefChecker::new(source); + if VERSION != *checker.prefix::(size_of::()) { + panic!("deserialization: bad version number"); + } + let mut checker = MutChecker::new(source); + let header: &mut MetaTupleHeader = checker.prefix(size_of::()); + MetaTupleWriter { header } + } + _ => panic!("deserialization: bad magic number"), + } + } +} + +#[derive(Debug)] +pub struct MetaTupleWriter<'a> { + header: &'a mut MetaTupleHeader, +} + +impl<'a> MetaTupleWriter<'a> { + pub fn start(&mut self) -> &mut OptionPointer { + &mut self.header.start + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VertexTupleHeader { + metadata: [f32; 3], + elements_s: u16, + elements_e: u16, + payload: Option>, + pointers_s: u16, + pointers_e: u16, + _padding_0: [Padding; 4], +} + +#[derive(Debug, Clone, PartialEq)] +pub struct VertexTuple { + pub metadata: [f32; 3], + pub elements: Vec, + pub payload: Option>, + pub pointers: Vec, +} + +impl Tuple for VertexTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(self.elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // pointers + let pointers_s = buffer.len() as u16; + buffer.extend(self.pointers.as_bytes()); + let pointers_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[..size_of::()].copy_from_slice( + VertexTupleHeader { + metadata: self.metadata, + elements_s, + elements_e, + payload: self.payload, + pointers_s, + pointers_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + buffer + } +} + +impl WithReader for VertexTuple { + type Reader<'a> = VertexTupleReader<'a>; + + fn deserialize_ref(source: &[u8]) -> VertexTupleReader<'_> { + let checker = RefChecker::new(source); + let header: &VertexTupleHeader = checker.prefix(0_u16); + let elements = checker.bytes(header.elements_s, header.elements_e); + let pointers = checker.bytes(header.pointers_s, header.pointers_e); + VertexTupleReader { + header, + elements, + pointers, + } + } +} + +impl WithWriter for VertexTuple { + type Writer<'a> = VertexTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> VertexTupleWriter<'_> { + let mut checker = MutChecker::new(source); + let header: &mut VertexTupleHeader = checker.prefix(0_u16); + let elements = checker.bytes(header.elements_s, header.elements_e); + let pointers = checker.bytes(header.pointers_s, header.pointers_e); + VertexTupleWriter { + header, + elements, + pointers, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct VertexTupleReader<'a> { + header: &'a VertexTupleHeader, + elements: &'a [u64], + pointers: &'a [Pointer], +} + +impl<'a> VertexTupleReader<'a> { + pub fn metadata(self) -> [f32; 3] { + self.header.metadata + } + pub fn payload(self) -> Option> { + self.header.payload + } + pub fn elements(self) -> &'a [u64] { + self.elements + } + pub fn pointers(self) -> &'a [Pointer] { + self.pointers + } +} + +#[derive(Debug)] +pub struct VertexTupleWriter<'a> { + header: &'a mut VertexTupleHeader, + elements: &'a mut [u64], + pointers: &'a mut [Pointer], +} + +impl VertexTupleWriter<'_> { + #[expect(dead_code)] + pub fn metadata(&mut self) -> &mut [f32; 3] { + &mut self.header.metadata + } + pub fn payload(&mut self) -> &mut Option> { + &mut self.header.payload + } + #[expect(dead_code)] + pub fn elements(&mut self) -> &mut [u64] { + self.elements + } + pub fn pointers(&mut self) -> &mut [Pointer] { + self.pointers + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VectorTupleHeader0 { + payload: Option>, + elements_s: u16, + elements_e: u16, + metadata_s: u16, + neighbours_s: u16, + neighbours_e: u16, + _padding_0: [Padding; 2], + version: Wrapping, +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VectorTupleHeader1 { + payload: Option>, + elements_s: u16, + elements_e: u16, + index: u32, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum VectorTuple { + _0 { + payload: Option>, + metadata: V::Metadata, + elements: Vec, + neighbours: Vec, + version: Wrapping, + }, + _1 { + payload: Option>, + elements: Vec, + index: u32, + }, +} + +impl Tuple for VectorTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + VectorTuple::_0 { + payload, + metadata, + elements, + neighbours, + version, + } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // metadata + let metadata_s = buffer.len() as u16; + buffer.extend(metadata.as_bytes()); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // neighbours + let neighbours_s = buffer.len() as u16; + buffer.extend(neighbours.as_bytes()); + let neighbours_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + VectorTupleHeader0 { + metadata_s, + elements_s, + elements_e, + neighbours_s, + neighbours_e, + payload: *payload, + version: *version, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + VectorTuple::_1 { + payload, + elements, + index, + } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + VectorTupleHeader1 { + elements_s, + elements_e, + payload: *payload, + index: *index, + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for VectorTuple { + type Reader<'a> = VectorTupleReader<'a, V>; + + fn deserialize_ref(source: &[u8]) -> VectorTupleReader<'_, V> { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + let metadata = checker.prefix(header.metadata_s); + let neighbours = checker.bytes(header.neighbours_s, header.neighbours_e); + VectorTupleReader::_0(VectorTupleReader0 { + header, + elements, + metadata, + neighbours, + }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + VectorTupleReader::_1(VectorTupleReader1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +impl WithWriter for VectorTuple { + type Writer<'a> = VectorTupleWriter<'a, V>; + + fn deserialize_mut(source: &mut [u8]) -> VectorTupleWriter<'_, V> { + let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let mut checker = MutChecker::new(source); + let header: &mut VectorTupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + let metadata = checker.prefix(header.metadata_s); + let neighbours = checker.bytes(header.neighbours_s, header.neighbours_e); + VectorTupleWriter::_0(VectorTupleWriter0 { + header, + elements, + metadata, + neighbours, + }) + } + 1 => { + let mut checker = MutChecker::new(source); + let header: &mut VectorTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + VectorTupleWriter::_1(VectorTupleWriter1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +#[derive(Debug)] +pub enum VectorTupleReader<'a, V: Vector> { + _0(VectorTupleReader0<'a, V>), + _1(VectorTupleReader1<'a, V>), +} + +impl<'a, V: Vector> Copy for VectorTupleReader<'a, V> {} + +impl<'a, V: Vector> Clone for VectorTupleReader<'a, V> { + fn clone(&self) -> Self { + *self + } +} + +#[derive(Debug, PartialEq)] +pub struct VectorTupleReader0<'a, V: Vector> { + header: &'a VectorTupleHeader0, + elements: &'a [V::Element], + metadata: &'a V::Metadata, + neighbours: &'a [OptionNeighbour], +} + +impl<'a, V: Vector> Copy for VectorTupleReader0<'a, V> {} + +impl<'a, V: Vector> Clone for VectorTupleReader0<'a, V> { + fn clone(&self) -> Self { + *self + } +} + +impl<'a, V: Vector> VectorTupleReader0<'a, V> { + pub fn payload(self) -> Option> { + self.header.payload + } + pub fn elements(self) -> &'a [V::Element] { + self.elements + } + pub fn metadata(self) -> &'a V::Metadata { + self.metadata + } + pub fn neighbours(self) -> &'a [OptionNeighbour] { + self.neighbours + } + pub fn version(self) -> Wrapping { + self.header.version + } +} + +#[derive(Debug, PartialEq)] +pub struct VectorTupleReader1<'a, V: Vector> { + header: &'a VectorTupleHeader1, + elements: &'a [V::Element], +} + +impl<'a, V: Vector> Copy for VectorTupleReader1<'a, V> {} + +impl<'a, V: Vector> Clone for VectorTupleReader1<'a, V> { + fn clone(&self) -> Self { + *self + } +} + +impl<'a, V: Vector> VectorTupleReader1<'a, V> { + pub fn payload(self) -> Option> { + self.header.payload + } + pub fn elements(self) -> &'a [V::Element] { + self.elements + } + pub fn index(self) -> u32 { + self.header.index + } +} + +#[derive(Debug)] +pub enum VectorTupleWriter<'a, V: Vector> { + _0(VectorTupleWriter0<'a, V>), + #[expect(dead_code)] + _1(VectorTupleWriter1<'a, V>), +} + +#[derive(Debug)] +pub struct VectorTupleWriter0<'a, V: Vector> { + header: &'a mut VectorTupleHeader0, + elements: &'a mut [V::Element], + metadata: &'a mut V::Metadata, + neighbours: &'a mut [OptionNeighbour], +} + +impl VectorTupleWriter0<'_, V> { + #[expect(dead_code)] + pub fn payload(&mut self) -> &mut Option> { + &mut self.header.payload + } + #[expect(dead_code)] + pub fn elements(&mut self) -> &mut [V::Element] { + self.elements + } + #[expect(dead_code)] + pub fn metadata(&mut self) -> &mut V::Metadata { + self.metadata + } + pub fn neighbours(&mut self) -> &mut [OptionNeighbour] { + self.neighbours + } + pub fn version(&mut self) -> &mut Wrapping { + &mut self.header.version + } +} + +#[derive(Debug)] +pub struct VectorTupleWriter1<'a, V: Vector> { + header: &'a mut VectorTupleHeader1, + elements: &'a mut [V::Element], +} + +impl VectorTupleWriter1<'_, V> { + #[expect(dead_code)] + pub fn payload(&mut self) -> &mut Option> { + &mut self.header.payload + } + #[expect(dead_code)] + pub fn elements(&mut self) -> &mut [V::Element] { + self.elements + } + #[expect(dead_code)] + pub fn index(&mut self) -> &mut u32 { + &mut self.header.index + } +} + +#[repr(C)] +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct Pointer { + x: u32, + y: u16, + _padding_0: [Padding; 2], +} + +impl Pointer { + pub fn new((x, y): (u32, u16)) -> Self { + Self { + x, + y, + _padding_0: Default::default(), + } + } + pub fn into_inner(self) -> (u32, u16) { + (self.x, self.y) + } +} + +#[repr(C)] +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct OptionPointer { + x: u32, + y: u16, + _padding_0: [Padding; 1], + validity: Bool, +} + +impl OptionPointer { + pub const NONE: Self = Self { + x: 0, + y: 0, + _padding_0: [Padding::ZERO; 1], + validity: Bool::FALSE, + }; + pub fn some((x, y): (u32, u16)) -> Self { + Self { + x, + y, + _padding_0: Default::default(), + validity: Bool::TRUE, + } + } + pub fn into_inner(self) -> Option<(u32, u16)> { + if self.validity.into() { + Some((self.x, self.y)) + } else { + None + } + } +} + +#[repr(C)] +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct OptionNeighbour { + pointer: OptionPointer, + distance: Distance, +} + +impl OptionNeighbour { + pub const NONE: Self = Self { + pointer: OptionPointer::NONE, + distance: Distance::ZERO, + }; + pub fn some(pointer: (u32, u16), distance: Distance) -> Self { + Self { + pointer: OptionPointer::some(pointer), + distance, + } + } + pub fn into_inner(self) -> Option<((u32, u16), Distance)> { + let inner = self.pointer.into_inner()?; + Some((inner, self.distance)) + } +} diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs new file mode 100644 index 00000000..04e01ddd --- /dev/null +++ b/crates/vchordg/src/types.rs @@ -0,0 +1,133 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use half::f16; +use serde::{Deserialize, Serialize}; +use validator::{Validate, ValidationError}; +use vector::vect::{VectBorrowed, VectOwned}; + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VamanaIndexOptions { + #[serde(default = "VamanaIndexOptions::default_m")] + #[validate(range(min = 1, max = 512))] + pub m: u32, + #[serde(default = "VamanaIndexOptions::default_alpha")] + #[validate(range(min = 1.0, max = 2.0))] + pub alpha: f32, + #[serde(default = "VamanaIndexOptions::default_ef_construction")] + #[validate(range(min = 1, max = 65535))] + pub ef_construction: u32, + #[serde(default = "VamanaIndexOptions::default_beam_construction")] + #[validate(range(min = 1, max = 65535))] + pub beam_construction: u32, +} + +impl VamanaIndexOptions { + fn default_m() -> u32 { + 32 + } + fn default_alpha() -> f32 { + 1.0 + } + fn default_ef_construction() -> u32 { + 500 + } + fn default_beam_construction() -> u32 { + 1 + } +} + +impl Default for VamanaIndexOptions { + fn default() -> Self { + Self { + m: Self::default_m(), + alpha: Self::default_alpha(), + ef_construction: Self::default_ef_construction(), + beam_construction: Self::default_beam_construction(), + } + } +} + +#[derive(Debug, Clone)] +pub enum OwnedVector { + Vecf32(VectOwned), + Vecf16(VectOwned), +} + +#[derive(Debug, Clone, Copy)] +pub enum BorrowedVector<'a> { + Vecf32(VectBorrowed<'a, f32>), + Vecf16(VectBorrowed<'a, f16>), +} + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +pub enum DistanceKind { + L2S, + Dot, +} + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +pub enum VectorKind { + Vecf32, + Vecf16, +} + +impl VectorKind { + pub fn element_size(self) -> u32 { + match self { + VectorKind::Vecf32 => size_of::() as _, + VectorKind::Vecf16 => size_of::() as _, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +#[validate(schema(function = "Self::validate_self"))] +pub struct VectorOptions { + #[validate(range(min = 1))] + #[serde(rename = "dimensions")] + pub dims: u32, + #[serde(rename = "vector")] + pub v: VectorKind, + #[serde(rename = "distance")] + pub d: DistanceKind, +} + +impl VectorOptions { + pub fn validate_self(&self) -> Result<(), ValidationError> { + match (self.v, self.d, self.dims) { + (VectorKind::Vecf32, DistanceKind::L2S, 1..=60000) => Ok(()), + (VectorKind::Vecf16, DistanceKind::L2S, 1..=60000) => Ok(()), + _ => Err(ValidationError::new("invalid vector options")), + } + } +} + +pub struct Structure { + pub centroids: Vec, + pub children: Vec>, +} + +impl Structure { + pub fn len(&self) -> usize { + self.children.len() + } + pub fn is_empty(&self) -> bool { + self.children.is_empty() + } +} diff --git a/crates/vchordg/src/visited.rs b/crates/vchordg/src/visited.rs new file mode 100644 index 00000000..8f0cd169 --- /dev/null +++ b/crates/vchordg/src/visited.rs @@ -0,0 +1,34 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::Id; +use std::collections::HashSet; + +pub struct Visited { + inner: HashSet, +} + +impl Visited { + pub fn new() -> Self { + Self { + inner: HashSet::new(), + } + } + pub fn insert(&mut self, x: Id) { + self.inner.insert(x); + } + pub fn contains(&mut self, x: Id) -> bool { + self.inner.contains(&x) + } +} diff --git a/crates/algorithm/Cargo.toml b/crates/vchordrq/Cargo.toml similarity index 91% rename from crates/algorithm/Cargo.toml rename to crates/vchordrq/Cargo.toml index ffdef2ee..d65cae4b 100644 --- a/crates/algorithm/Cargo.toml +++ b/crates/vchordrq/Cargo.toml @@ -1,10 +1,11 @@ [package] -name = "algorithm" +name = "vchordrq" version.workspace = true edition.workspace = true publish = false [dependencies] +algo = { path = "../algo" } always_equal = { path = "../always_equal" } distance = { path = "../distance" } k_means = { path = "../k_means" } diff --git a/crates/algorithm/src/build.rs b/crates/vchordrq/src/build.rs similarity index 98% rename from crates/algorithm/src/build.rs rename to crates/vchordrq/src/build.rs index 04fc2cc7..c53499e5 100644 --- a/crates/algorithm/src/build.rs +++ b/crates/vchordrq/src/build.rs @@ -17,7 +17,8 @@ use crate::tape::TapeWriter; use crate::tape_writer::H1TapeWriter; use crate::tuples::*; use crate::types::*; -use crate::{Branch, RelationWrite}; +use crate::{Branch, Opaque}; +use algo::{Page, RelationWrite}; use vector::{VectorBorrowed, VectorOwned}; pub fn build( @@ -25,7 +26,9 @@ pub fn build( vchordrq_options: VchordrqIndexOptions, index: &R, structures: Vec>, -) { +) where + R::Page: Page, +{ let dims = vector_options.dims; let is_residual = vchordrq_options.residual_quantization; let mut meta = TapeWriter::<_, MetaTuple>::create(index, false); diff --git a/crates/algorithm/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs similarity index 96% rename from crates/algorithm/src/bulkdelete.rs rename to crates/vchordrq/src/bulkdelete.rs index f13c72b5..00a446ff 100644 --- a/crates/algorithm/src/bulkdelete.rs +++ b/crates/vchordrq/src/bulkdelete.rs @@ -13,17 +13,21 @@ // Copyright (c) 2025 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1}; -use crate::operator::{FunctionalAccessor, Operator}; +use crate::operator::Operator; use crate::tape::by_next; use crate::tuples::*; -use crate::{Page, RelationRead, RelationWrite, tape}; +use crate::{Opaque, Page, tape}; +use algo::accessor::FunctionalAccessor; +use algo::{RelationRead, RelationWrite}; use std::num::NonZero; pub fn bulkdelete( index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, -) { +) where + R::Page: Page, +{ let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -132,7 +136,9 @@ pub fn bulkdelete_vectors( index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, -) { +) where + R::Page: Page, +{ let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); diff --git a/crates/algorithm/src/cache.rs b/crates/vchordrq/src/cache.rs similarity index 88% rename from crates/algorithm/src/cache.rs rename to crates/vchordrq/src/cache.rs index ca3f0927..0a21f3cd 100644 --- a/crates/algorithm/src/cache.rs +++ b/crates/vchordrq/src/cache.rs @@ -13,12 +13,16 @@ // Copyright (c) 2025 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1}; -use crate::operator::FunctionalAccessor; use crate::tape::by_next; use crate::tuples::{MetaTuple, WithReader}; -use crate::{Page, PageGuard, RelationRead, tape}; - -pub fn cache(index: &R) -> Vec { +use crate::{Opaque, Page, PageGuard, tape}; +use algo::RelationRead; +use algo::accessor::FunctionalAccessor; + +pub fn cache(index: &R) -> Vec +where + R::Page: Page, +{ let mut trace = vec![0_u32]; let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); diff --git a/crates/algorithm/src/closure_lifetime_binder.rs b/crates/vchordrq/src/closure_lifetime_binder.rs similarity index 79% rename from crates/algorithm/src/closure_lifetime_binder.rs rename to crates/vchordrq/src/closure_lifetime_binder.rs index 592757cf..4eb2d063 100644 --- a/crates/algorithm/src/closure_lifetime_binder.rs +++ b/crates/vchordrq/src/closure_lifetime_binder.rs @@ -40,19 +40,20 @@ where } #[inline(always)] -pub fn id_3(f: F) -> F +pub fn id_3(f: F) -> F where - T: crate::RelationWrite, - F: for<'a> Fn(&'a T, A) -> T::WriteGuard<'a>, + T: algo::RelationWrite, + F: for<'a> Fn(&'a T, A, B) -> T::WriteGuard<'a>, { f } #[inline(always)] -pub fn id_4(f: F) -> F +pub fn id_4<'r, F, P, A, B, R: ?Sized>(f: F) -> F where - T: crate::RelationRead, - F: FnMut(A, Vec>, B) -> R, + P: algo::prefetcher::Prefetcher<'r>, +

::Item: algo::Fetch, + F: FnMut(A, P::Guards, B) -> R, { f } diff --git a/crates/algorithm/src/cost.rs b/crates/vchordrq/src/cost.rs similarity index 96% rename from crates/algorithm/src/cost.rs rename to crates/vchordrq/src/cost.rs index 12ad027e..c4eefecd 100644 --- a/crates/algorithm/src/cost.rs +++ b/crates/vchordrq/src/cost.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::tuples::{MetaTuple, WithReader}; -use crate::{Page, RelationRead}; +use algo::{Page, RelationRead}; pub struct Cost { pub dims: u32, diff --git a/crates/algorithm/src/fast_heap.rs b/crates/vchordrq/src/fast_heap.rs similarity index 99% rename from crates/algorithm/src/fast_heap.rs rename to crates/vchordrq/src/fast_heap.rs index 6222ffb0..d989ddfe 100644 --- a/crates/algorithm/src/fast_heap.rs +++ b/crates/vchordrq/src/fast_heap.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::Sequence; +use algo::Sequence; use std::collections::BinaryHeap; use std::num::NonZero; diff --git a/crates/algorithm/src/freepages.rs b/crates/vchordrq/src/freepages.rs similarity index 90% rename from crates/algorithm/src/freepages.rs rename to crates/vchordrq/src/freepages.rs index 1d0c96ef..b5420355 100644 --- a/crates/algorithm/src/freepages.rs +++ b/crates/vchordrq/src/freepages.rs @@ -13,9 +13,13 @@ // Copyright (c) 2025 TensorChord Inc. use crate::tuples::*; -use crate::*; +use crate::{Opaque, Page}; +use algo::RelationWrite; -pub fn alloc(index: &R, freepages_first: u32) -> Option> { +pub fn alloc(index: &R, freepages_first: u32) -> Option> +where + R::Page: Page, +{ let mut freepages_guard = index.write(freepages_first, false); let freepages_bytes = freepages_guard.get_mut(1).expect("data corruption"); let mut freepages_tuple = FreepagesTuple::deserialize_mut(freepages_bytes); @@ -31,7 +35,10 @@ pub fn alloc(index: &R, freepages_first: u32) -> Option(index: &R, freepages_first: u32, id: u32) { +pub fn free(index: &R, freepages_first: u32, id: u32) +where + R::Page: Page, +{ let mut guard = index.write(id, false); let mut freepages_guard = index.write(freepages_first, false); let freepages_bytes = freepages_guard.get_mut(1).expect("data corruption"); diff --git a/crates/algorithm/src/insert.rs b/crates/vchordrq/src/insert.rs similarity index 92% rename from crates/algorithm/src/insert.rs rename to crates/vchordrq/src/insert.rs index 75e3911d..563fcb12 100644 --- a/crates/algorithm/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -17,7 +17,10 @@ use crate::operator::*; use crate::tape::by_next; use crate::tuples::*; use crate::vectors::{self}; -use crate::{Bump, Page, Prefetcher, PrefetcherHeapFamily, RelationRead, RelationWrite, tape}; +use crate::{Opaque, Page, tape}; +use algo::accessor::{FunctionalAccessor, LAccess}; +use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily}; +use algo::{Bump, RelationRead, RelationWrite}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -31,7 +34,10 @@ pub fn insert_vector( index: &R, payload: NonZero, vector: ::Borrowed<'_>, -) -> (Vec, u16) { +) -> (Vec, u16) +where + R::Page: Page, +{ let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -43,7 +49,7 @@ pub fn insert_vector( drop(meta_guard); if !rerank_in_heap { - vectors::append::(index, vectors_first, vector, payload) + vectors::append::(index, vectors_first, vector, payload) } else { (Vec::new(), 0) } @@ -57,7 +63,9 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( bump: &'b impl Bump, mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, skip_freespaces: bool, -) { +) where + R::Page: Page, +{ let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -119,7 +127,7 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { let distance = vectors::read_for_h1_tuple::( - prefetch.into_iter(), + prefetch, head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); @@ -166,7 +174,7 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( payload: Some(payload), prefetch, head, - elements: rabitq::packing::pack_to_u64(&code.1), + elements: rabitq::original::binary::pack_code(&code.1), }); tape::append( diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs new file mode 100644 index 00000000..8e7dbd46 --- /dev/null +++ b/crates/vchordrq/src/lib.rs @@ -0,0 +1,69 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +#![allow(clippy::type_complexity)] + +mod build; +mod bulkdelete; +mod cache; +mod closure_lifetime_binder; +mod cost; +mod fast_heap; +mod freepages; +mod insert; +mod linked_vec; +mod maintain; +mod prewarm; +mod rerank; +mod search; +mod tape; +mod tape_writer; +mod tuples; +mod vectors; + +pub mod operator; +pub mod types; + +use algo::{Page, PageGuard}; +pub use build::build; +pub use bulkdelete::{bulkdelete, bulkdelete_vectors}; +pub use cache::cache; +pub use cost::cost; +pub use fast_heap::FastHeap; +pub use insert::{insert, insert_vector}; +pub use maintain::maintain; +pub use prewarm::prewarm; +pub use rerank::{how, rerank_heap, rerank_index}; +pub use search::{default_search, maxsim_search}; + +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +#[repr(C, align(8))] +#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +pub struct Opaque { + pub next: u32, + pub skip: u32, +} + +#[allow(unsafe_code)] +unsafe impl algo::Opaque for Opaque {} + +pub(crate) struct Branch { + pub code: rabitq::original::Code, + pub delta: f32, + pub prefetch: Vec, + pub head: u16, + pub norm: f32, + pub extra: T, +} diff --git a/crates/algorithm/src/linked_vec.rs b/crates/vchordrq/src/linked_vec.rs similarity index 100% rename from crates/algorithm/src/linked_vec.rs rename to crates/vchordrq/src/linked_vec.rs diff --git a/crates/algorithm/src/maintain.rs b/crates/vchordrq/src/maintain.rs similarity index 86% rename from crates/algorithm/src/maintain.rs rename to crates/vchordrq/src/maintain.rs index 1b200022..e2b52186 100644 --- a/crates/algorithm/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -13,11 +13,16 @@ // Copyright (c) 2025 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1, id_2, id_3}; -use crate::operator::{FunctionalAccessor, Operator, Vector}; +use crate::operator::{Operator, Vector}; use crate::tape::{self, TapeWriter, by_directory, by_next}; use crate::tape_writer::{DirectoryTapeWriter, FrozenTapeWriter}; use crate::tuples::*; -use crate::*; +use crate::{Branch, Opaque, Page, freepages}; +use algo::accessor::FunctionalAccessor; +use algo::prefetcher::PrefetcherSequenceFamily; +use algo::{ + PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, +}; use rabitq::packing::unpack; use std::cell::RefCell; @@ -31,7 +36,10 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( index: &'r R, mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, check: impl Fn(), -) -> Maintain { +) -> Maintain +where + R::Page: Page, +{ let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -78,26 +86,26 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( let mut jump_tuple = JumpTuple::deserialize_mut(jump_bytes); let hooked_index = RelationHooked(index, { - id_3(|index: &R, tracking_freespace: bool| { + id_3(|index: &R, opaque: Opaque, tracking_freespace: bool| { if !tracking_freespace { let mut buffers = buffers.borrow_mut(); if let Some(id) = buffers.pages.pop() { drop(buffers); let mut guard = index.write(id, false); - guard.clear(); + guard.clear(opaque); guard } else if let Some(mut guard) = freepages::alloc(index, freepages_first) { buffers.number_of_formerly_allocated_pages += 1; drop(buffers); - guard.clear(); + guard.clear(opaque); guard } else { buffers.number_of_freshly_allocated_pages += 1; drop(buffers); - index.extend(false) + index.extend(opaque, false) } } else { - index.extend(true) + index.extend(opaque, true) } }) }); @@ -141,7 +149,7 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( let signs = unpacked[i].iter().flat_map(f).collect::>(); ( ( - rabitq::CodeMetadata { + rabitq::original::CodeMetadata { dis_u_2: metadata[0][i], factor_cnt: metadata[1][i], factor_ip: metadata[2][i], @@ -170,7 +178,7 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( .collect::>(); ( ( - rabitq::CodeMetadata { + rabitq::original::CodeMetadata { dis_u_2: metadata[0], factor_cnt: metadata[1], factor_ip: metadata[2], @@ -196,7 +204,7 @@ pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( branch.code.0.factor_ip, branch.code.0.factor_err, ], - elements: rabitq::packing::pack_to_u64(&branch.code.1), + elements: rabitq::original::binary::pack_code(&branch.code.1), delta: branch.delta, prefetch: branch.prefetch, head: branch.head, @@ -250,37 +258,47 @@ where type Page = R::Page; } -impl<'r, R, E> RelationRead for RelationHooked<'r, R, E> +impl<'r, R, E> RelationReadTypes for RelationHooked<'r, R, E> where R: RelationRead, E: Clone, { - type ReadGuard<'a> - = R::ReadGuard<'a> - where - Self: 'a; + type ReadGuard<'a> = R::ReadGuard<'a>; +} +impl<'r, R, E> RelationRead for RelationHooked<'r, R, E> +where + R: RelationRead, + E: Clone, +{ fn read(&self, id: u32) -> Self::ReadGuard<'_> { self.0.read(id) } } -impl<'r, R, E> RelationWrite for RelationHooked<'r, R, E> +impl<'r, R, E> RelationWriteTypes for RelationHooked<'r, R, E> where R: RelationWrite, - E: Clone + for<'a> Fn(&'a R, bool) -> R::WriteGuard<'a>, + E: Clone + for<'a> Fn(&'a R, ::Opaque, bool) -> R::WriteGuard<'a>, { - type WriteGuard<'a> - = R::WriteGuard<'a> - where - Self: 'a; + type WriteGuard<'a> = R::WriteGuard<'a>; +} +impl<'r, R, E> RelationWrite for RelationHooked<'r, R, E> +where + R: RelationWrite, + E: Clone + for<'a> Fn(&'a R, ::Opaque, bool) -> R::WriteGuard<'a>, +{ fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { self.0.write(id, tracking_freespace) } - fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_> { - (self.1)(self.0, tracking_freespace) + fn extend( + &self, + opaque: ::Opaque, + tracking_freespace: bool, + ) -> Self::WriteGuard<'_> { + (self.1)(self.0, opaque, tracking_freespace) } fn search(&self, freespace: usize) -> Option> { diff --git a/crates/algorithm/src/operator.rs b/crates/vchordrq/src/operator.rs similarity index 61% rename from crates/algorithm/src/operator.rs rename to crates/vchordrq/src/operator.rs index 7de8d079..04b0bd07 100644 --- a/crates/algorithm/src/operator.rs +++ b/crates/vchordrq/src/operator.rs @@ -12,11 +12,12 @@ // // Copyright (c) 2025 TensorChord Inc. +use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S, RAccess}; use distance::Distance; use half::f16; -use rabitq::CodeMetadata; -use rabitq::binary::{BinaryLut, BinaryLutMetadata}; -use rabitq::block::{BlockLut, BlockLutMetadata, STEP}; +use rabitq::original::CodeMetadata; +use rabitq::original::binary::{BinaryLut, BinaryLutMetadata}; +use rabitq::original::block::{BlockLut, BlockLutMetadata, STEP}; use simd::Floating; use std::fmt::Debug; use std::marker::PhantomData; @@ -24,273 +25,6 @@ use vector::vect::{VectBorrowed, VectOwned}; use vector::{VectorBorrowed, VectorOwned}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; -pub trait Accessor2 { - type Output; - fn push(&mut self, input: &[E0], target: &[E1]); - fn finish(self, input: M0, target: M1) -> Self::Output; -} - -impl Accessor2 for () { - type Output = (); - - fn push(&mut self, _: &[E0], _: &[E1]) {} - - fn finish(self, _: M0, _: M1) -> Self::Output {} -} - -impl> Accessor2 for (A,) { - type Output = (A::Output,); - - fn push(&mut self, input: &[E0], target: &[E1]) { - self.0.push(input, target); - } - - fn finish(self, input: M0, target: M1) -> Self::Output { - (self.0.finish(input, target),) - } -} - -impl, B: Accessor2> - Accessor2 for (A, B) -{ - type Output = (A::Output, B::Output); - - fn push(&mut self, input: &[E0], target: &[E1]) { - self.0.push(input, target); - self.1.push(input, target); - } - - fn finish(self, input: M0, target: M1) -> Self::Output { - (self.0.finish(input, target), self.1.finish(input, target)) - } -} - -pub trait Accessor1 { - type Output; - fn push(&mut self, input: &[E]); - fn finish(self, input: M) -> Self::Output; -} - -impl Accessor1 for () { - type Output = (); - - fn push(&mut self, _: &[E]) {} - - fn finish(self, _: M) -> Self::Output {} -} - -impl Accessor1 for (A,) -where - A: Accessor1, -{ - type Output = (A::Output,); - - fn push(&mut self, input: &[E]) { - self.0.push(input); - } - - fn finish(self, input: M) -> Self::Output { - (self.0.finish(input),) - } -} - -impl Accessor1 for (A, B) -where - A: Accessor1, - B: Accessor1, -{ - type Output = (A::Output, B::Output); - - fn push(&mut self, input: &[E]) { - self.0.push(input); - self.1.push(input); - } - - fn finish(self, input: M) -> Self::Output { - (self.0.finish(input), self.1.finish(input)) - } -} - -pub struct FunctionalAccessor { - data: T, - p: P, - f: F, -} - -impl FunctionalAccessor { - pub fn new(data: T, p: P, f: F) -> Self { - Self { data, p, f } - } -} - -impl Accessor1 for FunctionalAccessor -where - P: for<'a> FnMut(&'a mut T, &'a [E]), - F: FnOnce(T, M) -> R, -{ - type Output = R; - - fn push(&mut self, input: &[E]) { - (self.p)(&mut self.data, input); - } - - fn finish(self, input: M) -> Self::Output { - (self.f)(self.data, input) - } -} - -pub struct LAccess<'a, E, M, A> { - elements: &'a [E], - metadata: M, - accessor: A, -} - -impl<'a, E, M, A> LAccess<'a, E, M, A> { - pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { - Self { - elements, - metadata, - accessor, - } - } -} - -impl> Accessor1 for LAccess<'_, E0, M0, A> { - type Output = A::Output; - - fn push(&mut self, rhs: &[E1]) { - let (lhs, elements) = self.elements.split_at(rhs.len()); - self.accessor.push(lhs, rhs); - self.elements = elements; - } - - fn finish(self, rhs: M1) -> Self::Output { - assert!(self.elements.is_empty(), "goal is shorter than expected"); - self.accessor.finish(self.metadata, rhs) - } -} - -pub struct RAccess<'a, E, M, A> { - elements: &'a [E], - metadata: M, - accessor: A, -} - -impl<'a, E, M, A> RAccess<'a, E, M, A> { - pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { - Self { - elements, - metadata, - accessor, - } - } -} - -impl> Accessor1 for RAccess<'_, E1, M1, A> { - type Output = A::Output; - - fn push(&mut self, lhs: &[E0]) { - let (rhs, elements) = self.elements.split_at(lhs.len()); - self.accessor.push(lhs, rhs); - self.elements = elements; - } - - fn finish(self, lhs: M0) -> Self::Output { - assert!(self.elements.is_empty(), "goal is shorter than expected"); - self.accessor.finish(lhs, self.metadata) - } -} - -pub trait TryAccessor1: Sized { - type Output; - #[must_use] - fn push(&mut self, input: &[E]) -> Option<()>; - #[must_use] - fn finish(self, input: M) -> Option; -} - -impl TryAccessor1 for () { - type Output = (); - - fn push(&mut self, _: &[E]) -> Option<()> { - Some(()) - } - - fn finish(self, _: M) -> Option { - Some(()) - } -} - -impl TryAccessor1 for (A,) -where - A: TryAccessor1, -{ - type Output = (A::Output,); - - fn push(&mut self, input: &[E]) -> Option<()> { - self.0.push(input)?; - Some(()) - } - - fn finish(self, input: M) -> Option { - Some((self.0.finish(input)?,)) - } -} - -impl TryAccessor1 for (A, B) -where - A: TryAccessor1, - B: TryAccessor1, -{ - type Output = (A::Output, B::Output); - - fn push(&mut self, input: &[E]) -> Option<()> { - self.0.push(input)?; - self.1.push(input)?; - Some(()) - } - - fn finish(self, input: M) -> Option { - Some((self.0.finish(input)?, self.1.finish(input)?)) - } -} - -pub struct LTryAccess<'a, E, M, A> { - elements: &'a [E], - metadata: M, - accessor: A, -} - -impl<'a, E, M, A> LTryAccess<'a, E, M, A> { - pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { - Self { - elements, - metadata, - accessor, - } - } -} - -impl> TryAccessor1 - for LTryAccess<'_, E0, M0, A> -{ - type Output = A::Output; - - fn push(&mut self, rhs: &[E1]) -> Option<()> { - let (lhs, elements) = self.elements.split_at_checked(rhs.len())?; - self.accessor.push(lhs, rhs); - self.elements = elements; - Some(()) - } - - fn finish(self, rhs: M1) -> Option { - if !self.elements.is_empty() { - return None; - } - Some(self.accessor.finish(self.metadata, rhs)) - } -} - #[derive(Debug)] pub struct BlockAccessor([u32; 32], F); @@ -330,63 +64,6 @@ impl> } } -#[derive(Debug)] -pub struct DistanceAccessor(f32, PhantomData V>, PhantomData D>); - -impl Default for DistanceAccessor { - fn default() -> Self { - Self(0.0, PhantomData, PhantomData) - } -} - -impl Accessor2 for DistanceAccessor, L2> { - type Output = Distance; - - fn push(&mut self, target: &[f32], input: &[f32]) { - self.0 += f32::reduce_sum_of_d2(target, input) - } - - fn finish(self, (): (), (): ()) -> Self::Output { - Distance::from_f32(self.0) - } -} - -impl Accessor2 for DistanceAccessor, Dot> { - type Output = Distance; - - fn push(&mut self, target: &[f32], input: &[f32]) { - self.0 += f32::reduce_sum_of_xy(target, input) - } - - fn finish(self, (): (), (): ()) -> Self::Output { - Distance::from_f32(-self.0) - } -} - -impl Accessor2 for DistanceAccessor, L2> { - type Output = Distance; - - fn push(&mut self, target: &[f16], input: &[f16]) { - self.0 += f16::reduce_sum_of_d2(target, input) - } - - fn finish(self, (): (), (): ()) -> Self::Output { - Distance::from_f32(self.0) - } -} - -impl Accessor2 for DistanceAccessor, Dot> { - type Output = Distance; - - fn push(&mut self, target: &[f16], input: &[f16]) { - self.0 += f16::reduce_sum_of_xy(target, input) - } - - fn finish(self, (): (), (): ()) -> Self::Output { - Distance::from_f32(-self.0) - } -} - #[derive(Debug, Clone)] pub struct CloneAccessor(Vec); @@ -425,7 +102,7 @@ pub trait Vector: VectorOwned { fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut); - fn code(vector: Self::Borrowed<'_>) -> rabitq::Code; + fn code(vector: Self::Borrowed<'_>) -> rabitq::original::Code; fn squared_norm(vector: Self::Borrowed<'_>) -> f32; } @@ -466,15 +143,15 @@ impl Vector for VectOwned { } fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { - rabitq::block::preprocess(vector.slice()) + rabitq::original::block::preprocess(vector.slice()) } fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { - rabitq::preprocess(vector.slice()) + rabitq::original::preprocess(vector.slice()) } - fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { - rabitq::code(vector.dims(), vector.slice()) + fn code(vector: Self::Borrowed<'_>) -> rabitq::original::Code { + rabitq::original::code(vector.slice()) } fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { @@ -518,15 +195,15 @@ impl Vector for VectOwned { } fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { - rabitq::block::preprocess(&f16::vector_to_f32(vector.slice())) + rabitq::original::block::preprocess(&f16::vector_to_f32(vector.slice())) } fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { - rabitq::preprocess(&f16::vector_to_f32(vector.slice())) + rabitq::original::preprocess(&f16::vector_to_f32(vector.slice())) } - fn code(vector: Self::Borrowed<'_>) -> rabitq::Code { - rabitq::code(vector.dims(), &f16::vector_to_f32(vector.slice())) + fn code(vector: Self::Borrowed<'_>) -> rabitq::original::Code { + rabitq::original::code(&f16::vector_to_f32(vector.slice())) } fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { @@ -534,12 +211,6 @@ impl Vector for VectOwned { } } -#[derive(Debug, Clone, Copy)] -pub struct L2; - -#[derive(Debug, Clone, Copy)] -pub struct Dot; - pub trait Operator: 'static + Debug + Copy { type Vector: Vector; @@ -585,7 +256,7 @@ pub trait Operator: 'static + Debug + Copy { fn build( vector: ::Borrowed<'_>, centroid: Option, - ) -> (rabitq::Code, f32); + ) -> (rabitq::original::Code, f32); } #[derive(Debug)] @@ -599,10 +270,10 @@ impl Clone for Op { impl Copy for Op {} -impl Operator for Op, L2> { +impl Operator for Op, L2S> { type Vector = VectOwned; - type DistanceAccessor = DistanceAccessor, L2>; + type DistanceAccessor = DistanceAccessor, L2S>; fn block_access>( lut: &BlockLut, @@ -617,7 +288,7 @@ impl Operator for Op, L2> { mut f: F, ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::binary::asymmetric_binary_dot_product(elements, &lut.1); + let value = rabitq::original::binary::accumulate(elements, &lut.1); f.call( value, CodeMetadata { @@ -642,9 +313,9 @@ impl Operator for Op, L2> { _: f32, ) -> (f32, f32) { if !is_residual { - rabitq::block::half_process_l2(value, code, lut) + rabitq::original::block::half_process_l2(value, code, lut) } else { - rabitq::block::half_process_l2_residual(value, code, lut, dis_f, delta) + rabitq::original::block::half_process_l2_residual(value, code, lut, dis_f, delta) } } @@ -658,13 +329,16 @@ impl Operator for Op, L2> { _: f32, ) -> (f32, f32) { if !is_residual { - rabitq::binary::half_process_l2(value, code, lut) + rabitq::original::binary::half_process_l2(value, code, lut) } else { - rabitq::binary::half_process_l2_residual(value, code, lut, dis_f, delta) + rabitq::original::binary::half_process_l2_residual(value, code, lut, dis_f, delta) } } - fn build(vector: VectBorrowed<'_, f32>, centroid: Option) -> (rabitq::Code, f32) { + fn build( + vector: VectBorrowed<'_, f32>, + centroid: Option, + ) -> (rabitq::original::Code, f32) { if let Some(centroid) = centroid { let residual = VectOwned::new(f32::vector_sub(vector.slice(), centroid.slice())); let code = Self::Vector::code(residual.as_borrowed()); @@ -708,7 +382,7 @@ impl Operator for Op, Dot> { mut f: F, ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::binary::asymmetric_binary_dot_product(elements, &lut.1); + let value = rabitq::original::binary::accumulate(elements, &lut.1); f.call( value, CodeMetadata { @@ -733,9 +407,9 @@ impl Operator for Op, Dot> { norm: f32, ) -> (f32, f32) { if !is_residual { - rabitq::block::half_process_dot(value, code, lut) + rabitq::original::block::half_process_dot(value, code, lut) } else { - rabitq::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + rabitq::original::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) } } @@ -749,13 +423,18 @@ impl Operator for Op, Dot> { norm: f32, ) -> (f32, f32) { if !is_residual { - rabitq::binary::half_process_dot(value, code, lut) + rabitq::original::binary::half_process_dot(value, code, lut) } else { - rabitq::binary::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + rabitq::original::binary::half_process_dot_residual( + value, code, lut, dis_f, delta, norm, + ) } } - fn build(vector: VectBorrowed<'_, f32>, centroid: Option) -> (rabitq::Code, f32) { + fn build( + vector: VectBorrowed<'_, f32>, + centroid: Option, + ) -> (rabitq::original::Code, f32) { if let Some(centroid) = centroid { let residual = VectOwned::new(f32::vector_sub(vector.slice(), centroid.slice())); let code = Self::Vector::code(residual.as_borrowed()); @@ -781,10 +460,10 @@ impl Operator for Op, Dot> { } } -impl Operator for Op, L2> { +impl Operator for Op, L2S> { type Vector = VectOwned; - type DistanceAccessor = DistanceAccessor, L2>; + type DistanceAccessor = DistanceAccessor, L2S>; fn block_access>( lut: &BlockLut, @@ -799,7 +478,7 @@ impl Operator for Op, L2> { mut f: F, ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::binary::asymmetric_binary_dot_product(elements, &lut.1); + let value = rabitq::original::binary::accumulate(elements, &lut.1); f.call( value, CodeMetadata { @@ -824,9 +503,9 @@ impl Operator for Op, L2> { _: f32, ) -> (f32, f32) { if !is_residual { - rabitq::block::half_process_l2(value, code, lut) + rabitq::original::block::half_process_l2(value, code, lut) } else { - rabitq::block::half_process_l2_residual(value, code, lut, dis_f, delta) + rabitq::original::block::half_process_l2_residual(value, code, lut, dis_f, delta) } } @@ -840,13 +519,16 @@ impl Operator for Op, L2> { _: f32, ) -> (f32, f32) { if !is_residual { - rabitq::binary::half_process_l2(value, code, lut) + rabitq::original::binary::half_process_l2(value, code, lut) } else { - rabitq::binary::half_process_l2_residual(value, code, lut, dis_f, delta) + rabitq::original::binary::half_process_l2_residual(value, code, lut, dis_f, delta) } } - fn build(vector: VectBorrowed<'_, f16>, centroid: Option) -> (rabitq::Code, f32) { + fn build( + vector: VectBorrowed<'_, f16>, + centroid: Option, + ) -> (rabitq::original::Code, f32) { if let Some(centroid) = centroid { let residual = VectOwned::new(f16::vector_sub(vector.slice(), centroid.slice())); let code = Self::Vector::code(residual.as_borrowed()); @@ -890,7 +572,7 @@ impl Operator for Op, Dot> { mut f: F, ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::binary::asymmetric_binary_dot_product(elements, &lut.1); + let value = rabitq::original::binary::accumulate(elements, &lut.1); f.call( value, CodeMetadata { @@ -915,9 +597,9 @@ impl Operator for Op, Dot> { norm: f32, ) -> (f32, f32) { if !is_residual { - rabitq::block::half_process_dot(value, code, lut) + rabitq::original::block::half_process_dot(value, code, lut) } else { - rabitq::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + rabitq::original::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) } } @@ -931,13 +613,18 @@ impl Operator for Op, Dot> { norm: f32, ) -> (f32, f32) { if !is_residual { - rabitq::binary::half_process_dot(value, code, lut) + rabitq::original::binary::half_process_dot(value, code, lut) } else { - rabitq::binary::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + rabitq::original::binary::half_process_dot_residual( + value, code, lut, dis_f, delta, norm, + ) } } - fn build(vector: VectBorrowed<'_, f16>, centroid: Option) -> (rabitq::Code, f32) { + fn build( + vector: VectBorrowed<'_, f16>, + centroid: Option, + ) -> (rabitq::original::Code, f32) { if let Some(centroid) = centroid { let residual = VectOwned::new(f16::vector_sub(vector.slice(), centroid.slice())); let code = Self::Vector::code(residual.as_borrowed()); diff --git a/crates/algorithm/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs similarity index 94% rename from crates/algorithm/src/prewarm.rs rename to crates/vchordrq/src/prewarm.rs index 36b332e1..24a8e8d9 100644 --- a/crates/algorithm/src/prewarm.rs +++ b/crates/vchordrq/src/prewarm.rs @@ -13,10 +13,13 @@ // Copyright (c) 2025 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1, id_2}; -use crate::operator::{FunctionalAccessor, Operator}; +use crate::operator::Operator; use crate::tape::{by_directory, by_next}; use crate::tuples::*; -use crate::{Bump, Page, PrefetcherSequenceFamily, RelationRead, tape, vectors}; +use crate::{Opaque, Page, tape, vectors}; +use algo::accessor::FunctionalAccessor; +use algo::prefetcher::PrefetcherSequenceFamily; +use algo::{Bump, RelationRead}; use std::fmt::Write; pub fn prewarm<'r, 'b: 'r, R: RelationRead, O: Operator>( @@ -24,7 +27,10 @@ pub fn prewarm<'r, 'b: 'r, R: RelationRead, O: Operator>( height: i32, bump: &'b impl Bump, mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, -) -> String { +) -> String +where + R::Page: Page, +{ let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); diff --git a/crates/algorithm/src/rerank.rs b/crates/vchordrq/src/rerank.rs similarity index 85% rename from crates/algorithm/src/rerank.rs rename to crates/vchordrq/src/rerank.rs index d413bea8..4247e622 100644 --- a/crates/algorithm/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -14,9 +14,11 @@ use crate::closure_lifetime_binder::id_4; use crate::operator::*; -use crate::prefetcher::Prefetcher; use crate::tuples::{MetaTuple, WithReader}; -use crate::{Page, RelationRead, RerankMethod, vectors}; +use crate::{Page, vectors}; +use algo::accessor::{Accessor2, LTryAccess}; +use algo::prefetcher::Prefetcher; +use algo::{RelationRead, RerankMethod}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -50,7 +52,7 @@ pub struct Reranker { impl<'r, 'b, T, F, P> Iterator for Reranker where - F: FnMut(NonZero, Vec<::ReadGuard<'r>>, u16) -> Option, + F: FnMut(NonZero, P::Guards, u16) -> Option, P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, { type Item = (Distance, NonZero); @@ -84,17 +86,13 @@ pub fn rerank_index< >( vector: O::Vector, prefetcher: P, -) -> Reranker< - T, - impl FnMut(NonZero, Vec<::ReadGuard<'_>>, u16) -> Option, - P, -> { +) -> Reranker, P::Guards, u16) -> Option, P> { Reranker { prefetcher, cache: BinaryHeap::new(), - f: id_4::<_, P::R, _, _, _>(move |payload, prefetch, head| { + f: id_4::<_, P, _, _, _>(move |payload, prefetch, head| { vectors::read_for_h0_tuple::( - prefetch.into_iter(), + prefetch, head, payload, LTryAccess::new( @@ -117,15 +115,11 @@ pub fn rerank_heap< vector: O::Vector, prefetcher: P, mut fetch: impl FnMut(NonZero) -> Option + 'b, -) -> Reranker< - T, - impl FnMut(NonZero, Vec<::ReadGuard<'_>>, u16) -> Option, - P, -> { +) -> Reranker, P::Guards, u16) -> Option, P> { Reranker { prefetcher, cache: BinaryHeap::new(), - f: id_4::<_, P::R, _, _, _>(move |payload, _, _| { + f: id_4::<_, P, _, _, _>(move |payload, _, _| { let unpack = O::Vector::unpack(vector.as_borrowed()); let vector = fetch(payload)?; let vector = O::Vector::unpack(vector.as_borrowed()); diff --git a/crates/algorithm/src/search.rs b/crates/vchordrq/src/search.rs similarity index 97% rename from crates/algorithm/src/search.rs rename to crates/vchordrq/src/search.rs index 5f59855f..0ab91141 100644 --- a/crates/algorithm/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -15,10 +15,12 @@ use crate::closure_lifetime_binder::id_2; use crate::linked_vec::LinkedVec; use crate::operator::*; -use crate::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use crate::tape::{by_directory, by_next}; use crate::tuples::*; -use crate::{Bump, Page, PrefetcherHeapFamily, RelationRead, tape, vectors}; +use crate::{Opaque, Page, tape, vectors}; +use algo::accessor::LAccess; +use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; +use algo::{Bump, RelationRead}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -37,7 +39,10 @@ pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( bump: &'b impl Bump, mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, -) -> Vec<((Reverse, AlwaysEqual<()>), AlwaysEqual>)> { +) -> Vec<((Reverse, AlwaysEqual<()>), AlwaysEqual>)> +where + R::Page: Page, +{ let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -103,7 +108,7 @@ pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { let distance = vectors::read_for_h1_tuple::( - prefetch.into_iter(), + prefetch, head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); @@ -167,7 +172,10 @@ pub fn maxsim_search<'r, 'b: 'r, R: RelationRead, O: Operator>( AlwaysEqual>, )>, Distance, -) { +) +where + R::Page: Page, +{ let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -233,7 +241,7 @@ pub fn maxsim_search<'r, 'b: 'r, R: RelationRead, O: Operator>( heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { let distance = vectors::read_for_h1_tuple::( - prefetch.into_iter(), + prefetch, head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); diff --git a/crates/algorithm/src/tape.rs b/crates/vchordrq/src/tape.rs similarity index 86% rename from crates/algorithm/src/tape.rs rename to crates/vchordrq/src/tape.rs index cf2b5964..2de5e546 100644 --- a/crates/algorithm/src/tape.rs +++ b/crates/vchordrq/src/tape.rs @@ -12,9 +12,11 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::operator::Accessor1; use crate::tuples::*; -use crate::{Page, PageGuard, PrefetcherSequenceFamily, RelationRead, RelationWrite, freepages}; +use crate::{Opaque, Page, PageGuard, freepages}; +use algo::accessor::Accessor1; +use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use algo::{RelationRead, RelationWrite}; use std::marker::PhantomData; use std::num::NonZero; @@ -32,9 +34,16 @@ where impl<'a, R, T> TapeWriter<'a, R, T> where R: RelationWrite + 'a, + R::Page: Page, { pub fn create(index: &'a R, tracking_freespace: bool) -> Self { - let mut head = index.extend(tracking_freespace); + let mut head = index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + tracking_freespace, + ); head.get_opaque_mut().skip = head.id(); let first = head.id(); Self { @@ -55,7 +64,13 @@ where if self.head.len() == 0 { panic!("implementation: a clear page cannot accommodate a single tuple"); } - let next = self.index.extend(self.tracking_freespace); + let next = self.index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + self.tracking_freespace, + ); self.head.get_opaque_mut().next = next.id(); self.head = next; } @@ -64,6 +79,7 @@ where impl<'a, R, T> TapeWriter<'a, R, T> where R: RelationWrite + 'a, + R::Page: Page, T: Tuple, { pub fn push(&mut self, x: T) -> (u32, u16) { @@ -71,7 +87,13 @@ where if let Some(i) = self.head.alloc(&bytes) { (self.head.id(), i) } else { - let next = self.index.extend(self.tracking_freespace); + let next = self.index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + self.tracking_freespace, + ); self.head.get_opaque_mut().next = next.id(); self.head = next; if let Some(i) = self.head.alloc(&bytes) { @@ -180,10 +202,9 @@ where { let mut t = p.prefetch(iter.peekable()); std::iter::from_fn(move || { - use crate::prefetcher::Prefetcher; let (_, mut x) = t.next()?; - let ret = x.pop().expect("should be at least one element"); - assert!(x.pop().is_none(), "should be at most one element"); + let ret = x.next().expect("should be at least one element"); + assert!(x.next().is_none(), "should be at most one element"); Some(ret) }) } @@ -191,6 +212,7 @@ where pub fn by_next<'r, R>(index: &'r R, first: u32) -> impl Iterator> where R: RelationRead + 'r, + R::Page: Page, { let mut current = first; std::iter::from_fn(move || { @@ -296,13 +318,16 @@ pub fn read_appendable_tape<'r, R, T>( } #[allow(clippy::collapsible_else_if)] -pub fn append( - index: &(impl RelationRead + RelationWrite), +pub fn append( + index: &R, first: u32, bytes: &[u8], tracking_freespace: bool, freepages_first: Option, -) -> (u32, u16) { +) -> (u32, u16) +where + R::Page: Page, +{ assert!(!tracking_freespace || freepages_first.is_none()); assert!(first != u32::MAX); let mut current = first; @@ -318,13 +343,28 @@ pub fn append( let mut extend = { if let Some(freepages_first) = freepages_first { if let Some(mut guard) = freepages::alloc(index, freepages_first) { - guard.clear(); + guard.clear(Opaque { + next: u32::MAX, + skip: u32::MAX, + }); guard } else { - index.extend(tracking_freespace) + index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + tracking_freespace, + ) } } else { - index.extend(tracking_freespace) + index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + tracking_freespace, + ) } }; write.get_opaque_mut().next = extend.id(); diff --git a/crates/algorithm/src/tape_writer.rs b/crates/vchordrq/src/tape_writer.rs similarity index 98% rename from crates/algorithm/src/tape_writer.rs rename to crates/vchordrq/src/tape_writer.rs index 119c7ca5..19caab88 100644 --- a/crates/algorithm/src/tape_writer.rs +++ b/crates/vchordrq/src/tape_writer.rs @@ -14,7 +14,8 @@ use crate::tape::TapeWriter; use crate::tuples::*; -use crate::{Branch, RelationWrite}; +use crate::{Branch, Opaque}; +use algo::{Page, RelationWrite}; use rabitq::packing::{any_pack, padding_pack}; use std::num::NonZero; @@ -28,6 +29,7 @@ where impl<'a, R> DirectoryTapeWriter<'a, R> where R: RelationWrite + 'a, + R::Page: Page, { pub fn create(index: &'a R, tracking_freespace: bool) -> Self { Self { @@ -70,6 +72,7 @@ where impl<'a, R> H1TapeWriter<'a, R> where R: RelationWrite + 'a, + R::Page: Page, { pub fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { Self { @@ -169,6 +172,7 @@ where impl<'a, R> FrozenTapeWriter<'a, R> where R: RelationWrite + 'a, + R::Page: Page, { pub fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { Self { diff --git a/crates/algorithm/src/tuples.rs b/crates/vchordrq/src/tuples.rs similarity index 88% rename from crates/algorithm/src/tuples.rs rename to crates/vchordrq/src/tuples.rs index b5f450b6..fe3b6fdb 100644 --- a/crates/algorithm/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::operator::Vector; -use std::marker::PhantomData; +use algo::tuples::{Bool, MutChecker, Padding, RefChecker}; use std::num::NonZero; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -46,10 +46,10 @@ struct MetaTupleHeader { rerank_in_heap: Bool, cells_s: u16, cells_e: u16, - _padding_0: [ZeroU8; 2], + _padding_0: [Padding; 2], vectors_first: u32, freepages_first: u32, - _padding_1: [ZeroU8; 2], + _padding_1: [Padding; 2], // tree centroid_prefetch_s: u16, centroid_prefetch_e: u16, @@ -209,7 +209,7 @@ impl<'a> MetaTupleReader<'a> { #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FreepagesTupleHeader { first: u32, - _padding_0: [ZeroU8; 4], + _padding_0: [Padding; 4], } #[derive(Debug, Clone, PartialEq)] @@ -253,7 +253,7 @@ struct VectorTupleHeader0 { metadata_s: u16, elements_s: u16, elements_e: u16, - _padding_0: [ZeroU8; 2], + _padding_0: [Padding; 2], } #[repr(C, align(8))] @@ -263,7 +263,7 @@ struct VectorTupleHeader1 { head: u16, elements_s: u16, elements_e: u16, - _padding_0: [ZeroU8; 2], + _padding_0: [Padding; 2], } #[derive(Debug, Clone, PartialEq)] @@ -426,7 +426,7 @@ impl<'a, V: Vector> VectorTupleReader<'a, V> { struct DirectoryTupleHeader0 { elements_s: u16, elements_e: u16, - _padding_0: [ZeroU8; 4], + _padding_0: [Padding; 4], } #[repr(C, align(8))] @@ -434,7 +434,7 @@ struct DirectoryTupleHeader0 { struct DirectoryTupleHeader1 { elements_s: u16, elements_e: u16, - _padding_0: [ZeroU8; 4], + _padding_0: [Padding; 4], } #[allow(clippy::large_enum_variant)] @@ -582,7 +582,7 @@ struct H1TupleHeader0 { len: u32, elements_s: u16, elements_e: u16, - _padding_0: [ZeroU8; 4], + _padding_0: [Padding; 4], } #[repr(C, align(8))] @@ -590,7 +590,7 @@ struct H1TupleHeader0 { struct H1TupleHeader1 { elements_s: u16, elements_e: u16, - _padding_0: [ZeroU8; 4], + _padding_0: [Padding; 4], } #[allow(clippy::large_enum_variant)] @@ -796,7 +796,7 @@ struct JumpTupleHeader { centroid_prefetch_s: u16, centroid_prefetch_e: u16, centroid_head: u16, - _padding_0: [ZeroU8; 2], + _padding_0: [Padding; 2], directory_first: u32, appendable_first: u32, tuples: u64, @@ -922,7 +922,7 @@ struct FrozenTupleHeader0 { struct FrozenTupleHeader1 { elements_s: u16, elements_e: u16, - _padding_0: [ZeroU8; 4], + _padding_0: [Padding; 4], } #[allow(clippy::large_enum_variant)] @@ -1168,7 +1168,7 @@ struct AppendableTupleHeader { prefetch_s: u16, prefetch_e: u16, head: u16, - _padding_0: [ZeroU8; 2], + _padding_0: [Padding; 2], elements_s: u16, elements_e: u16, // it's the last field for reducing padding bytes @@ -1289,184 +1289,3 @@ impl AppendableTupleWriter<'_> { &mut self.header.payload } } - -#[repr(transparent)] -#[derive( - Debug, - Default, - Clone, - Copy, - PartialEq, - Eq, - PartialOrd, - Ord, - Hash, - IntoBytes, - FromBytes, - Immutable, - KnownLayout, -)] -pub struct ZeroU8(Option>); - -#[repr(transparent)] -#[derive( - Debug, - Clone, - Copy, - PartialEq, - Eq, - PartialOrd, - Ord, - Hash, - IntoBytes, - FromBytes, - Immutable, - KnownLayout, -)] -pub struct Bool(u8); - -impl Bool { - pub const FALSE: Self = Self(0); - pub const TRUE: Self = Self(1); -} - -impl From for bool { - fn from(value: Bool) -> Self { - value != Bool::FALSE - } -} - -impl From for Bool { - fn from(value: bool) -> Self { - if value { Self::TRUE } else { Self::FALSE } - } -} - -pub struct RefChecker<'a> { - bytes: &'a [u8], -} - -impl<'a> RefChecker<'a> { - pub fn new(bytes: &'a [u8]) -> Self { - Self { bytes } - } - pub fn prefix( - &self, - s: impl Into + Copy, - ) -> &'a T { - let start = Into::::into(s); - let end = Into::::into(s) + size_of::(); - let bytes = &self.bytes[start..end]; - FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") - } - pub fn bytes( - &self, - s: impl Into + Copy, - e: impl Into + Copy, - ) -> &'a T { - let start = Into::::into(s); - let end = Into::::into(e); - let bytes = &self.bytes[start..end]; - FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") - } -} - -pub struct MutChecker<'a> { - flag: usize, - bytes: *mut [u8], - phantom: PhantomData<&'a mut [u8]>, -} - -impl<'a> MutChecker<'a> { - pub fn new(bytes: &'a mut [u8]) -> Self { - Self { - flag: 0, - bytes, - phantom: PhantomData, - } - } - pub fn prefix( - &mut self, - s: impl Into + Copy, - ) -> &'a mut T { - let start = Into::::into(s); - let end = Into::::into(s) + size_of::(); - if !(start <= end && end <= self.bytes.len()) { - panic!("deserialization: bad bytes"); - } - if !(self.flag <= start) { - panic!("deserialization: bad bytes"); - } else { - self.flag = end; - } - #[allow(unsafe_code)] - let bytes = unsafe { - std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) - }; - FromBytes::mut_from_bytes(bytes).expect("deserialization: bad bytes") - } - pub fn bytes( - &mut self, - s: impl Into + Copy, - e: impl Into + Copy, - ) -> &'a mut T { - let start = Into::::into(s); - let end = Into::::into(e); - if !(start <= end && end <= self.bytes.len()) { - panic!("deserialization: bad bytes"); - } - if !(self.flag <= start) { - panic!("deserialization: bad bytes"); - } else { - self.flag = end; - } - #[allow(unsafe_code)] - let bytes = unsafe { - std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) - }; - FromBytes::mut_from_bytes(bytes).expect("deserialization: bad bytes") - } -} - -#[test] -fn aliasing_test() { - #[repr(C, align(8))] - #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] - struct ExampleHeader { - elements_s: u16, - elements_e: u16, - _padding_0: [ZeroU8; 4], - } - let serialized = { - let elements = (0u32..1111).collect::>(); - let mut buffer = Vec::::new(); - buffer.extend(std::iter::repeat_n(0, size_of::())); - let elements_s = buffer.len() as u16; - buffer.extend(elements.as_bytes()); - let elements_e = buffer.len() as u16; - while buffer.len() % ALIGN != 0 { - buffer.push(0); - } - buffer[..size_of::()].copy_from_slice( - ExampleHeader { - elements_s, - elements_e, - _padding_0: Default::default(), - } - .as_bytes(), - ); - buffer - }; - let mut source = vec![0u64; serialized.len().next_multiple_of(8)]; - source.as_mut_bytes()[..serialized.len()].copy_from_slice(&serialized); - let deserialized = { - let mut checker = MutChecker::new(source.as_mut_bytes()); - let header: &mut ExampleHeader = checker.prefix(0_u16); - let elements: &mut [u32] = checker.bytes(header.elements_s, header.elements_e); - (header, elements) - }; - assert_eq!( - deserialized.1, - (0u32..1111).collect::>().as_slice() - ); -} diff --git a/crates/algorithm/src/types.rs b/crates/vchordrq/src/types.rs similarity index 86% rename from crates/algorithm/src/types.rs rename to crates/vchordrq/src/types.rs index 477bf44b..941b76c2 100644 --- a/crates/algorithm/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -17,7 +17,7 @@ use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError}; use vector::vect::{VectBorrowed, VectOwned}; -#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] pub struct VchordrqIndexOptions { #[serde(default = "VchordrqIndexOptions::default_residual_quantization")] @@ -35,6 +35,15 @@ impl VchordrqIndexOptions { } } +impl Default for VchordrqIndexOptions { + fn default() -> Self { + Self { + residual_quantization: Self::default_residual_quantization(), + rerank_in_table: Self::default_rerank_in_table(), + } + } +} + #[derive(Debug, Clone)] pub enum OwnedVector { Vecf32(VectOwned), @@ -50,7 +59,7 @@ pub enum BorrowedVector<'a> { #[repr(u8)] #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] pub enum DistanceKind { - L2, + L2S, Dot, } @@ -86,9 +95,9 @@ pub struct VectorOptions { impl VectorOptions { pub fn validate_self(&self) -> Result<(), ValidationError> { match (self.v, self.d, self.dims) { - (VectorKind::Vecf32, DistanceKind::L2, 1..=60000) => Ok(()), + (VectorKind::Vecf32, DistanceKind::L2S, 1..=60000) => Ok(()), (VectorKind::Vecf32, DistanceKind::Dot, 1..=60000) => Ok(()), - (VectorKind::Vecf16, DistanceKind::L2, 1..=60000) => Ok(()), + (VectorKind::Vecf16, DistanceKind::L2S, 1..=60000) => Ok(()), (VectorKind::Vecf16, DistanceKind::Dot, 1..=60000) => Ok(()), _ => Err(ValidationError::new("invalid vector options")), } diff --git a/crates/algorithm/src/vectors.rs b/crates/vchordrq/src/vectors.rs similarity index 89% rename from crates/algorithm/src/vectors.rs rename to crates/vchordrq/src/vectors.rs index 1fa2fb0d..b50a71f2 100644 --- a/crates/algorithm/src/vectors.rs +++ b/crates/vchordrq/src/vectors.rs @@ -14,7 +14,9 @@ use crate::operator::*; use crate::tuples::*; -use crate::{Page, PageGuard, RelationRead, RelationWrite, tape}; +use crate::{Opaque, Page, PageGuard, tape}; +use algo::accessor::{Accessor1, TryAccessor1}; +use algo::{RelationRead, RelationWrite}; use std::num::NonZero; use vector::VectorOwned; @@ -78,13 +80,19 @@ pub fn read_for_h0_tuple< result.finish(cursor.ok()?) } -pub fn append( - index: &(impl RelationRead + RelationWrite), +pub fn append( + index: &R, vectors_first: u32, vector: ::Borrowed<'_>, payload: NonZero, -) -> (Vec, u16) { - fn append(index: &(impl RelationRead + RelationWrite), first: u32, bytes: &[u8]) -> (u32, u16) { +) -> (Vec, u16) +where + R::Page: Page, +{ + fn append(index: &R, first: u32, bytes: &[u8]) -> (u32, u16) + where + R::Page: Page, + { if let Some(mut write) = index.search(bytes.len()) { let i = write .alloc(bytes) diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs index cb34bd2f..74b76239 100644 --- a/crates/vector/src/bvect.rs +++ b/crates/vector/src/bvect.rs @@ -171,7 +171,7 @@ impl VectorBorrowed for BVectBorrowed<'_> { } #[inline(always)] - fn operator_l2(self, _: Self) -> Distance { + fn operator_l2s(self, _: Self) -> Distance { unimplemented!() } diff --git a/crates/vector/src/lib.rs b/crates/vector/src/lib.rs index 740e04ec..5f7df24b 100644 --- a/crates/vector/src/lib.rs +++ b/crates/vector/src/lib.rs @@ -36,7 +36,7 @@ pub trait VectorBorrowed: Copy { fn operator_dot(self, rhs: Self) -> distance::Distance; - fn operator_l2(self, rhs: Self) -> distance::Distance; + fn operator_l2s(self, rhs: Self) -> distance::Distance; fn operator_cos(self, rhs: Self) -> distance::Distance; diff --git a/crates/vector/src/scalar8.rs b/crates/vector/src/scalar8.rs index a2c3efeb..87186819 100644 --- a/crates/vector/src/scalar8.rs +++ b/crates/vector/src/scalar8.rs @@ -207,7 +207,7 @@ impl VectorBorrowed for Scalar8Borrowed<'_> { } #[inline(always)] - fn operator_l2(self, rhs: Self) -> Distance { + fn operator_l2s(self, rhs: Self) -> Distance { assert_eq!(self.code.len(), rhs.code.len()); let xy = self.k * rhs.k * simd::u8::reduce_sum_of_x_as_u32_y_as_u32(self.code, rhs.code) as f32 diff --git a/crates/vector/src/svect.rs b/crates/vector/src/svect.rs index 081ab021..19d41104 100644 --- a/crates/vector/src/svect.rs +++ b/crates/vector/src/svect.rs @@ -208,7 +208,7 @@ impl VectorBorrowed for SVectBorrowed<'_, S> { } #[inline(always)] - fn operator_l2(self, rhs: Self) -> Distance { + fn operator_l2s(self, rhs: Self) -> Distance { let d2 = S::reduce_sum_of_d2_sparse(self.indexes, self.values, rhs.indexes, rhs.values); Distance::from(d2) } diff --git a/crates/vector/src/vect.rs b/crates/vector/src/vect.rs index 4c1a473f..8854a055 100644 --- a/crates/vector/src/vect.rs +++ b/crates/vector/src/vect.rs @@ -134,7 +134,7 @@ impl VectorBorrowed for VectBorrowed<'_, S> { } #[inline(always)] - fn operator_l2(self, rhs: Self) -> Distance { + fn operator_l2s(self, rhs: Self) -> Distance { Distance::from(S::reduce_sum_of_d2(self.slice(), rhs.slice())) } diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 04270ea8..55827dfe 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2025-06-08" +channel = "1.89-beta" diff --git a/src/datatype/binary_scalar8.rs b/src/datatype/binary_scalar8.rs index 2179d589..57f510f5 100644 --- a/src/datatype/binary_scalar8.rs +++ b/src/datatype/binary_scalar8.rs @@ -34,7 +34,7 @@ fn _vchord_scalar8_send(vector: Scalar8Input<'_>) -> Vec { } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_recv(internal: Internal, oid: Oid, typmod: i32) -> Scalar8Output { +fn _vchord_scalar8_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Scalar8Output { let _ = (oid, typmod); let buf = unsafe { internal.get_mut::().unwrap() }; diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index 9daf54cc..25727123 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -38,7 +38,7 @@ fn _vchord_halfvec_sphere_l2_in( if lhs.dims() != center.dims() { pgrx::error!("dimension is not matched"); } - let d = VectBorrowed::operator_l2(lhs, center).to_f32().sqrt(); + let d = VectBorrowed::operator_l2s(lhs, center).to_f32().sqrt(); d < radius } diff --git a/src/datatype/operators_scalar8.rs b/src/datatype/operators_scalar8.rs index 76dbbf3d..cf8ee5ef 100644 --- a/src/datatype/operators_scalar8.rs +++ b/src/datatype/operators_scalar8.rs @@ -34,7 +34,7 @@ fn _vchord_scalar8_operator_l2(lhs: Scalar8Input<'_>, rhs: Scalar8Input<'_>) -> if lhs.dims() != rhs.dims() { pgrx::error!("dimension is not matched"); } - Scalar8Borrowed::operator_l2(lhs, rhs).to_f32().sqrt() + Scalar8Borrowed::operator_l2s(lhs, rhs).to_f32().sqrt() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] @@ -91,7 +91,7 @@ fn _vchord_scalar8_sphere_l2_in( if lhs.dims() != center.dims() { pgrx::error!("dimension is not matched"); } - let d = Scalar8Borrowed::operator_l2(lhs, center).to_f32().sqrt(); + let d = Scalar8Borrowed::operator_l2s(lhs, center).to_f32().sqrt(); d < radius } diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index 32ab7129..eb96ac19 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -38,7 +38,7 @@ fn _vchord_vector_sphere_l2_in( if lhs.dims() != center.dims() { pgrx::error!("dimension is not matched"); } - let d = VectBorrowed::operator_l2(lhs, center).to_f32().sqrt(); + let d = VectBorrowed::operator_l2s(lhs, center).to_f32().sqrt(); d < radius } diff --git a/src/index/algorithm.rs b/src/index/algorithm.rs deleted file mode 100644 index a6f98a2c..00000000 --- a/src/index/algorithm.rs +++ /dev/null @@ -1,603 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -use super::opclass::Opfamily; -use crate::index::am::am_build::InternalBuild; -use algorithm::operator::{Dot, L2, Op}; -use algorithm::types::*; -use algorithm::{ - Bump, FastHeap, Fetch, Hints, PlainPrefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily, - RelationPrefetch, RelationRead, RelationReadStream, RelationWrite, Sequence, SimplePrefetcher, - StreamPrefetcher, -}; -use half::f16; -use std::cell::UnsafeCell; -use std::collections::BinaryHeap; -use std::mem::MaybeUninit; -use std::num::NonZero; -use vector::VectorOwned; -use vector::vect::{VectBorrowed, VectOwned}; - -#[repr(C, align(4096))] -struct Chunk([u8; 2 * 1024 * 1024]); - -struct Allocator { - used: Vec<*mut MaybeUninit>, - free: Vec<*mut MaybeUninit>, - this: *mut MaybeUninit, - size: usize, -} - -impl Allocator { - pub fn new() -> Self { - Self { - used: Vec::new(), - free: Vec::new(), - this: Box::into_raw(Box::::new_uninit()).cast(), - size: size_of::(), - } - } - pub fn malloc(&mut self) -> *mut T { - const { - assert!(align_of::() <= align_of::()); - assert!(size_of::() <= size_of::()); - } - if size_of::() <= self.size { - self.size = (self.size - size_of::()) / align_of::() * align_of::(); - unsafe { self.this.add(self.size).cast::() } - } else { - #[cold] - fn cold(sel: &mut Allocator) -> *mut T { - abort_unwind(|| { - let raw = std::mem::replace(&mut sel.this, std::ptr::null_mut()); - sel.used.push(raw.cast()); - sel.this = sel - .free - .pop() - .unwrap_or_else(|| Box::into_raw(Box::::new_uninit())) - .cast(); - sel.size = size_of::(); - }); - sel.size = (sel.size - size_of::()) / align_of::() * align_of::(); - unsafe { sel.this.add(sel.size).cast::() } - } - cold(self) - } - } - pub fn malloc_n(&mut self, n: usize) -> *mut T { - const { - assert!(align_of::() <= align_of::()); - } - let limit = const { - if size_of::() > 0 { - size_of::() / size_of::() - } else { - usize::MAX - } - }; - if n <= limit && n * size_of::() <= self.size { - self.size = (self.size - n * size_of::()) / align_of::() * align_of::(); - unsafe { self.this.add(self.size).cast::() } - } else { - #[cold] - fn cold(sel: &mut Allocator, n: usize) -> *mut T { - abort_unwind(|| { - let raw = std::mem::replace(&mut sel.this, std::ptr::null_mut()); - sel.used.push(raw.cast()); - sel.this = sel - .free - .pop() - .unwrap_or_else(|| Box::into_raw(Box::::new_uninit())) - .cast(); - sel.size = size_of::(); - }); - sel.size = (sel.size - n * size_of::()) / align_of::() * align_of::(); - unsafe { sel.this.add(sel.size).cast::() } - } - if n > limit { - panic!("failed to allocate memory"); - } - cold(self, n) - } - } - pub fn reset(&mut self) { - abort_unwind(|| { - self.free.extend(std::mem::take(&mut self.used)); - self.size = size_of::(); - }); - } -} - -impl Drop for Allocator { - fn drop(&mut self) { - for raw in self.used.iter().copied() { - unsafe { - let _ = Box::>::from_raw(raw); - } - } - for raw in self.free.iter().copied() { - unsafe { - let _ = Box::>::from_raw(raw); - } - } - unsafe { - let _ = Box::>::from_raw(self.this.cast()); - } - } -} - -#[test] -fn test_allocator() { - let mut allocator = Allocator::new(); - for _ in 0..1024 * 8 { - allocator.malloc::<()>(); - allocator.malloc::(); - allocator.malloc::(); - allocator.malloc::(); - allocator.malloc::<[u8; 32]>(); - allocator.malloc_n::(2); - allocator.malloc::<[u8; 32]>(); - allocator.malloc_n::(2); - allocator.malloc_n::(2); - allocator.malloc_n::(2); - allocator.malloc_n::(2); - allocator.malloc_n::(2); - } - let number_of_chunks_0 = 1 + allocator.used.len() + allocator.free.len(); - allocator.reset(); - for _ in 0..1024 * 8 { - allocator.malloc::<()>(); - allocator.malloc::(); - allocator.malloc::(); - allocator.malloc::(); - allocator.malloc::<[u8; 32]>(); - allocator.malloc_n::(2); - allocator.malloc::<[u8; 32]>(); - allocator.malloc_n::(2); - allocator.malloc_n::(2); - allocator.malloc_n::(2); - allocator.malloc_n::(2); - allocator.malloc_n::(2); - } - let number_of_chunks_1 = 1 + allocator.used.len() + allocator.free.len(); - assert_eq!(number_of_chunks_0, number_of_chunks_1); -} - -pub struct BumpAlloc { - inner: UnsafeCell, -} - -impl BumpAlloc { - pub fn new() -> Self { - Self { - inner: UnsafeCell::new(Allocator::new()), - } - } - pub fn reset(&mut self) { - self.inner.get_mut().reset(); - } -} - -impl Bump for BumpAlloc { - fn alloc(&self, value: T) -> &mut T { - unsafe { - let ptr = (*self.inner.get()).malloc::(); - ptr.write(value); - &mut *ptr - } - } - - fn alloc_slice(&self, slice: &[T]) -> &mut [T] { - unsafe { - let ptr = (*self.inner.get()).malloc_n::(slice.len()); - std::ptr::copy_nonoverlapping(slice.as_ptr(), ptr, slice.len()); - std::slice::from_raw_parts_mut(ptr, slice.len()) - } - } -} - -pub fn prewarm(opfamily: Opfamily, index: &impl RelationRead, height: i32) -> String { - let bump = BumpAlloc::new(); - let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => algorithm::prewarm::<_, Op, L2>>( - index, - height, - &bump, - make_h0_plain_prefetcher, - ), - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::prewarm::<_, Op, Dot>>( - index, - height, - &bump, - make_h0_plain_prefetcher, - ) - } - (VectorKind::Vecf16, DistanceKind::L2) => algorithm::prewarm::<_, Op, L2>>( - index, - height, - &bump, - make_h0_plain_prefetcher, - ), - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::prewarm::<_, Op, Dot>>( - index, - height, - &bump, - make_h0_plain_prefetcher, - ) - } - } -} - -pub fn bulkdelete( - opfamily: Opfamily, - index: &(impl RelationRead + RelationWrite), - check: impl Fn(), - callback: impl Fn(NonZero) -> bool, -) { - match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::bulkdelete::<_, Op, L2>>(index, &check, &callback); - algorithm::bulkdelete_vectors::<_, Op, L2>>(index, &check, &callback); - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::bulkdelete::<_, Op, Dot>>(index, &check, &callback); - algorithm::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); - } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::bulkdelete::<_, Op, L2>>(index, &check, &callback); - algorithm::bulkdelete_vectors::<_, Op, L2>>(index, &check, &callback); - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::bulkdelete::<_, Op, Dot>>(index, &check, &callback); - algorithm::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); - } - } -} - -pub fn maintain(opfamily: Opfamily, index: &(impl RelationRead + RelationWrite), check: impl Fn()) { - let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; - let maintain = match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - algorithm::maintain::<_, Op, L2>>(index, make_h0_plain_prefetcher, check) - } - (VectorKind::Vecf32, DistanceKind::Dot) => { - algorithm::maintain::<_, Op, Dot>>( - index, - make_h0_plain_prefetcher, - check, - ) - } - (VectorKind::Vecf16, DistanceKind::L2) => { - algorithm::maintain::<_, Op, L2>>(index, make_h0_plain_prefetcher, check) - } - (VectorKind::Vecf16, DistanceKind::Dot) => { - algorithm::maintain::<_, Op, Dot>>( - index, - make_h0_plain_prefetcher, - check, - ) - } - }; - pgrx::info!( - "maintain: number_of_formerly_allocated_pages = {}", - maintain.number_of_formerly_allocated_pages - ); - pgrx::info!( - "maintain: number_of_freshly_allocated_pages = {}", - maintain.number_of_freshly_allocated_pages - ); - pgrx::info!( - "maintain: number_of_freed_pages = {}", - maintain.number_of_freed_pages - ); -} - -pub fn build( - vector_options: VectorOptions, - vchordrq_options: VchordrqIndexOptions, - index: &impl RelationWrite, - structures: Vec>>, -) { - match (vector_options.v, vector_options.d) { - (VectorKind::Vecf32, DistanceKind::L2) => algorithm::build::<_, Op, L2>>( - vector_options, - vchordrq_options, - index, - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), - ), - (VectorKind::Vecf32, DistanceKind::Dot) => algorithm::build::<_, Op, Dot>>( - vector_options, - vchordrq_options, - index, - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), - ), - (VectorKind::Vecf16, DistanceKind::L2) => algorithm::build::<_, Op, L2>>( - vector_options, - vchordrq_options, - index, - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), - ), - (VectorKind::Vecf16, DistanceKind::Dot) => algorithm::build::<_, Op, Dot>>( - vector_options, - vchordrq_options, - index, - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), - ), - } -} - -pub fn insert( - opfamily: Opfamily, - index: &(impl RelationRead + RelationWrite), - payload: NonZero, - vector: OwnedVector, - skip_freespaces: bool, -) { - let bump = BumpAlloc::new(); - let make_h1_plain_prefetcher = MakeH1PlainPrefetcherForInsertion { index }; - match (vector, opfamily.distance_kind()) { - (OwnedVector::Vecf32(vector), DistanceKind::L2) => { - assert!(opfamily.vector_kind() == VectorKind::Vecf32); - let projected = RandomProject::project(vector.as_borrowed()); - let key = algorithm::insert_vector::<_, Op, L2>>( - index, - payload, - vector.as_borrowed(), - ); - algorithm::insert::<_, Op, L2>>( - index, - payload, - projected.as_borrowed(), - key, - &bump, - make_h1_plain_prefetcher, - skip_freespaces, - ) - } - (OwnedVector::Vecf32(vector), DistanceKind::Dot) => { - assert!(opfamily.vector_kind() == VectorKind::Vecf32); - let projected = RandomProject::project(vector.as_borrowed()); - let key = algorithm::insert_vector::<_, Op, Dot>>( - index, - payload, - vector.as_borrowed(), - ); - algorithm::insert::<_, Op, Dot>>( - index, - payload, - projected.as_borrowed(), - key, - &bump, - make_h1_plain_prefetcher, - skip_freespaces, - ) - } - (OwnedVector::Vecf16(vector), DistanceKind::L2) => { - assert!(opfamily.vector_kind() == VectorKind::Vecf16); - let projected = RandomProject::project(vector.as_borrowed()); - let key = algorithm::insert_vector::<_, Op, L2>>( - index, - payload, - vector.as_borrowed(), - ); - algorithm::insert::<_, Op, L2>>( - index, - payload, - projected.as_borrowed(), - key, - &bump, - make_h1_plain_prefetcher, - skip_freespaces, - ) - } - (OwnedVector::Vecf16(vector), DistanceKind::Dot) => { - assert!(opfamily.vector_kind() == VectorKind::Vecf16); - let projected = RandomProject::project(vector.as_borrowed()); - let key = algorithm::insert_vector::<_, Op, Dot>>( - index, - payload, - vector.as_borrowed(), - ); - algorithm::insert::<_, Op, Dot>>( - index, - payload, - projected.as_borrowed(), - key, - &bump, - make_h1_plain_prefetcher, - skip_freespaces, - ) - } - } -} - -fn map_structures(x: Vec>, f: impl Fn(T) -> U + Copy) -> Vec> { - x.into_iter() - .map( - |Structure { - centroids, - children, - }| Structure { - centroids: centroids.into_iter().map(f).collect(), - children, - }, - ) - .collect() -} - -pub trait RandomProject { - type Output; - fn project(self) -> Self::Output; -} - -impl RandomProject for VectBorrowed<'_, f32> { - type Output = VectOwned; - fn project(self) -> VectOwned { - use rabitq::rotate::rotate; - let input = self.slice(); - VectOwned::new(rotate(input)) - } -} - -impl RandomProject for VectBorrowed<'_, f16> { - type Output = VectOwned; - fn project(self) -> VectOwned { - use rabitq::rotate::rotate; - use simd::Floating; - let input = f16::vector_to_f32(self.slice()); - VectOwned::new(f16::vector_from_f32(&rotate(&input))) - } -} - -// Emulate unstable library feature `abort_unwind`. -// See https://github.com/rust-lang/rust/issues/130338. - -#[inline(never)] -extern "C" fn abort_unwind R, R>(f: F) -> R { - f() -} - -#[derive(Debug)] -pub struct MakeH1PlainPrefetcherForInsertion<'r, R> { - pub index: &'r R, -} - -impl<'r, R> Clone for MakeH1PlainPrefetcherForInsertion<'r, R> { - fn clone(&self) -> Self { - Self { index: self.index } - } -} - -impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcherForInsertion<'r, R> { - type P - = PlainPrefetcher<'r, R, FastHeap> - where - T: Ord + Fetch + 'r; - - fn prefetch(&mut self, seq: Vec) -> Self::P - where - T: Ord + Fetch + 'r, - { - PlainPrefetcher::new(self.index, FastHeap::from(seq)) - } -} - -#[derive(Debug)] -pub struct MakeH1PlainPrefetcher<'r, R> { - pub index: &'r R, -} - -impl<'r, R> Clone for MakeH1PlainPrefetcher<'r, R> { - fn clone(&self) -> Self { - Self { index: self.index } - } -} - -impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcher<'r, R> { - type P - = PlainPrefetcher<'r, R, BinaryHeap> - where - T: Ord + Fetch + 'r; - - fn prefetch(&mut self, seq: Vec) -> Self::P - where - T: Ord + Fetch + 'r, - { - PlainPrefetcher::new(self.index, BinaryHeap::from(seq)) - } -} - -#[derive(Debug)] -pub struct MakeH0PlainPrefetcher<'r, R> { - pub index: &'r R, -} - -impl<'r, R> Clone for MakeH0PlainPrefetcher<'r, R> { - fn clone(&self) -> Self { - Self { index: self.index } - } -} - -impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakeH0PlainPrefetcher<'r, R> { - type P - = PlainPrefetcher<'r, R, S> - where - S::Item: Fetch; - - fn prefetch(&mut self, seq: S) -> Self::P - where - S::Item: Fetch, - { - PlainPrefetcher::new(self.index, seq) - } -} - -#[derive(Debug)] -pub struct MakeH0SimplePrefetcher<'r, R> { - pub index: &'r R, -} - -impl<'r, R> Clone for MakeH0SimplePrefetcher<'r, R> { - fn clone(&self) -> Self { - Self { index: self.index } - } -} - -impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> - for MakeH0SimplePrefetcher<'r, R> -{ - type P - = SimplePrefetcher<'r, R, S> - where - S::Item: Fetch; - - fn prefetch(&mut self, seq: S) -> Self::P - where - S::Item: Fetch, - { - SimplePrefetcher::new(self.index, seq) - } -} - -#[derive(Debug)] -pub struct MakeH0StreamPrefetcher<'r, R> { - pub index: &'r R, - pub hints: Hints, -} - -impl<'r, R> Clone for MakeH0StreamPrefetcher<'r, R> { - fn clone(&self) -> Self { - Self { - index: self.index, - hints: self.hints.clone(), - } - } -} - -impl<'r, R: RelationReadStream> PrefetcherSequenceFamily<'r, R> for MakeH0StreamPrefetcher<'r, R> { - type P - = StreamPrefetcher<'r, R, S> - where - S::Item: Fetch; - - fn prefetch(&mut self, seq: S) -> Self::P - where - S::Item: Fetch, - { - StreamPrefetcher::new(self.index, seq, self.hints.clone()) - } -} diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs new file mode 100644 index 00000000..823a9bf4 --- /dev/null +++ b/src/index/fetcher.rs @@ -0,0 +1,192 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use pgrx::pg_sys::{BlockIdData, Datum, ItemPointerData}; +use std::cell::LazyCell; +use std::num::NonZero; +use std::ops::DerefMut; +use std::ptr::NonNull; + +pub trait Fetcher { + type Tuple<'a>: Tuple + where + Self: 'a; + + fn fetch(&mut self, key: [u16; 3]) -> Option>; +} + +pub trait Tuple { + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); + fn filter(&mut self) -> bool; +} + +impl T> Fetcher for LazyCell { + type Tuple<'a> + = T::Tuple<'a> + where + Self: 'a; + + fn fetch(&mut self, key: [u16; 3]) -> Option> { + self.deref_mut().fetch(key) + } +} + +pub struct HeapFetcher { + index_info: *mut pgrx::pg_sys::IndexInfo, + estate: *mut pgrx::pg_sys::EState, + econtext: *mut pgrx::pg_sys::ExprContext, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + slot: *mut pgrx::pg_sys::TupleTableSlot, + values: [Datum; 32], + is_nulls: [bool; 32], + hack: *mut pgrx::pg_sys::IndexScanState, +} + +impl HeapFetcher { + pub unsafe fn new( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + hack: *mut pgrx::pg_sys::IndexScanState, + ) -> Self { + unsafe { + let index_info = pgrx::pg_sys::BuildIndexInfo(index_relation); + let estate = pgrx::pg_sys::CreateExecutorState(); + let econtext = pgrx::pg_sys::MakePerTupleExprContext(estate); + Self { + index_info, + estate, + econtext, + heap_relation, + snapshot, + slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), + values: [Datum::null(); 32], + is_nulls: [true; 32], + hack, + } + } + } +} + +impl Drop for HeapFetcher { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); + // free common resources + pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); + pgrx::pg_sys::FreeExecutorState(self.estate); + } + } +} + +impl Fetcher for HeapFetcher { + type Tuple<'a> = HeapTuple<'a>; + + fn fetch(&mut self, key: [u16; 3]) -> Option> { + unsafe { + let mut ctid = key_to_ctid(key); + let table_am = (*self.heap_relation).rd_tableam; + let fetch_row_version = (*table_am) + .tuple_fetch_row_version + .expect("unsupported heap access method"); + if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { + return None; + } + Some(HeapTuple { this: self }) + } + } +} + +pub struct HeapTuple<'a> { + this: &'a mut HeapFetcher, +} + +impl Tuple for HeapTuple<'_> { + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]) { + unsafe { + let this = &mut self.this; + (*this.econtext).ecxt_scantuple = this.slot; + pgrx::pg_sys::MemoryContextReset((*this.econtext).ecxt_per_tuple_memory); + pgrx::pg_sys::FormIndexDatum( + this.index_info, + this.slot, + this.estate, + this.values.as_mut_ptr(), + this.is_nulls.as_mut_ptr(), + ); + (&this.values, &this.is_nulls) + } + } + + #[allow(clippy::collapsible_if)] + fn filter(&mut self) -> bool { + unsafe { + let this = &mut self.this; + if !this.hack.is_null() { + if let Some(qual) = NonNull::new((*this.hack).ss.ps.qual) { + use pgrx::datum::FromDatum; + use pgrx::memcxt::PgMemoryContexts; + assert!(qual.as_ref().flags & pgrx::pg_sys::EEO_FLAG_IS_QUAL as u8 != 0); + let evalfunc = qual.as_ref().evalfunc.expect("no evalfunc for qual"); + if !(*this.hack).ss.ps.ps_ExprContext.is_null() { + let econtext = (*this.hack).ss.ps.ps_ExprContext; + (*econtext).ecxt_scantuple = this.slot; + pgrx::pg_sys::MemoryContextReset((*econtext).ecxt_per_tuple_memory); + let result = PgMemoryContexts::For((*econtext).ecxt_per_tuple_memory) + .switch_to(|_| { + let mut is_null = true; + let datum = evalfunc(qual.as_ptr(), econtext, &mut is_null); + bool::from_datum(datum, is_null) + }); + if result != Some(true) { + return false; + } + } + } + } + true + } + } +} + +pub const fn ctid_to_key( + ItemPointerData { + ip_blkid: BlockIdData { bi_hi, bi_lo }, + ip_posid, + }: ItemPointerData, +) -> [u16; 3] { + [bi_hi, bi_lo, ip_posid] +} + +pub const fn key_to_ctid([bi_hi, bi_lo, ip_posid]: [u16; 3]) -> ItemPointerData { + ItemPointerData { + ip_blkid: BlockIdData { bi_hi, bi_lo }, + ip_posid, + } +} + +pub const fn pointer_to_kv(pointer: NonZero) -> ([u16; 3], u16) { + let value = pointer.get(); + let bi_hi = ((value >> 48) & 0xffff) as u16; + let bi_lo = ((value >> 32) & 0xffff) as u16; + let ip_posid = ((value >> 16) & 0xffff) as u16; + let extra = value as u16; + ([bi_hi, bi_lo, ip_posid], extra) +} + +pub const fn kv_to_pointer((key, value): ([u16; 3], u16)) -> NonZero { + let x = (key[0] as u64) << 48 | (key[1] as u64) << 32 | (key[2] as u64) << 16 | value as u64; + NonZero::new(x).expect("invalid key") +} diff --git a/src/index/functions.rs b/src/index/functions.rs index 2cc77dba..b2dbe03b 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -33,9 +33,9 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { pgrx::error!("the index {:?} is not a vchordrq index", pg_class.relname()); } let relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); - let opfamily = unsafe { crate::index::opclass::opfamily(relation.raw()) }; + let opfamily = unsafe { crate::index::vchordrq::opclass::opfamily(relation.raw()) }; let index = unsafe { PostgresRelation::new(relation.raw()) }; - crate::index::algorithm::prewarm(opfamily, &index, height) + crate::index::vchordrq::algo::prewarm(opfamily, &index, height) } struct Index { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index cebfa08a..080dbe84 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -12,126 +12,244 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::scanners::Io; use pgrx::PostgresGucEnum; use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; -use std::ffi::CStr; +use std::ffi::CString; -#[allow(non_camel_case_types)] +#[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] #[derive(Debug, Clone, Copy, PostgresGucEnum)] pub enum PostgresIo { - read_buffer, - prefetch_buffer, - #[cfg(feature = "pg17")] - read_stream, + #[name = c"read_buffer"] + ReadBuffer, + #[name = c"prefetch_buffer"] + PrefetchBuffer, } -static PROBES: GucSetting> = GucSetting::>::new(Some(c"")); -static EPSILON: GucSetting = GucSetting::::new(1.9); -static MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); +#[cfg(any(feature = "pg17", feature = "pg18"))] +#[derive(Debug, Clone, Copy, PostgresGucEnum)] +pub enum PostgresIo { + #[name = c"read_buffer"] + ReadBuffer, + #[name = c"prefetch_buffer"] + PrefetchBuffer, + #[name = c"read_stream"] + ReadStream, +} + +static VCHORDG_EF_SEARCH: GucSetting = GucSetting::::new(64); -static MAXSIM_REFINE: GucSetting = GucSetting::::new(0); -static MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); +static VCHORDG_BEAM_SIZE: GucSetting = GucSetting::::new(1); -static PREFILTER: GucSetting = GucSetting::::new(false); +static VCHORDG_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); + +static VCHORDG_IO_SEARCH: GucSetting = GucSetting::::new( + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + PostgresIo::PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream, +); -static IO_SEARCH: GucSetting = GucSetting::::new( +static VCHORDG_IO_RERANK: GucSetting = GucSetting::::new( #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] - PostgresIo::prefetch_buffer, - #[cfg(feature = "pg17")] - PostgresIo::read_stream, + PostgresIo::PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream, ); -static IO_RERANK: GucSetting = GucSetting::::new( +static VCHORDRQ_PROBES: GucSetting> = GucSetting::>::new(Some(c"")); + +static VCHORDRQ_EPSILON: GucSetting = GucSetting::::new(1.9); + +static VCHORDRQ_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); + +static VCHORDRQ_MAXSIM_REFINE: GucSetting = GucSetting::::new(0); + +static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); + +static VCHORDRQ_PREFILTER: GucSetting = GucSetting::::new(false); + +static VCHORDRQ_IO_SEARCH: GucSetting = GucSetting::::new( #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] - PostgresIo::prefetch_buffer, - #[cfg(feature = "pg17")] - PostgresIo::read_stream, + PostgresIo::PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream, +); + +static VCHORDRQ_IO_RERANK: GucSetting = GucSetting::::new( + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + PostgresIo::PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream, ); pub fn init() { GucRegistry::define_string_guc( - "vchordrq.probes", - "`probes` argument of vchordrq.", - "`probes` argument of vchordrq.", - &PROBES, + c"vchordrq.probes", + c"`probes` argument of vchordrq.", + c"`probes` argument of vchordrq.", + &VCHORDRQ_PROBES, GucContext::Userset, GucFlags::default(), ); GucRegistry::define_float_guc( - "vchordrq.epsilon", - "`epsilon` argument of vchordrq.", - "`epsilon` argument of vchordrq.", - &EPSILON, + c"vchordrq.epsilon", + c"`epsilon` argument of vchordrq.", + c"`epsilon` argument of vchordrq.", + &VCHORDRQ_EPSILON, 0.0, 4.0, GucContext::Userset, GucFlags::default(), ); GucRegistry::define_int_guc( - "vchordrq.max_scan_tuples", - "`max_scan_tuples` argument of vchordrq.", - "`max_scan_tuples` argument of vchordrq.", - &MAX_SCAN_TUPLES, + c"vchordrq.max_scan_tuples", + c"`max_scan_tuples` argument of vchordrq.", + c"`max_scan_tuples` argument of vchordrq.", + &VCHORDRQ_MAX_SCAN_TUPLES, -1, i32::MAX, GucContext::Userset, GucFlags::default(), ); GucRegistry::define_int_guc( - "vchordrq.maxsim_refine", - "`maxsim_refine` argument of vchordrq.", - "`maxsim_refine` argument of vchordrq.", - &MAXSIM_REFINE, + c"vchordrq.maxsim_refine", + c"`maxsim_refine` argument of vchordrq.", + c"`maxsim_refine` argument of vchordrq.", + &VCHORDRQ_MAXSIM_REFINE, 0, i32::MAX, GucContext::Userset, GucFlags::default(), ); GucRegistry::define_int_guc( - "vchordrq.maxsim_threshold", - "`maxsim_threshold` argument of vchordrq.", - "`maxsim_threshold` argument of vchordrq.", - &MAXSIM_THRESHOLD, + c"vchordrq.maxsim_threshold", + c"`maxsim_threshold` argument of vchordrq.", + c"`maxsim_threshold` argument of vchordrq.", + &VCHORDRQ_MAXSIM_THRESHOLD, 0, i32::MAX, GucContext::Userset, GucFlags::default(), ); GucRegistry::define_bool_guc( - "vchordrq.prefilter", - "`prefilter` argument of vchordrq.", - "`prefilter` argument of vchordrq.", - &PREFILTER, + c"vchordrq.prefilter", + c"`prefilter` argument of vchordrq.", + c"`prefilter` argument of vchordrq.", + &VCHORDRQ_PREFILTER, GucContext::Userset, GucFlags::default(), ); GucRegistry::define_enum_guc( - "vchordrq.io_search", - "`io_search` argument of vchordrq.", - "`io_search` argument of vchordrq.", - &IO_SEARCH, + c"vchordrq.io_search", + c"`io_search` argument of vchordrq.", + c"`io_search` argument of vchordrq.", + &VCHORDRQ_IO_SEARCH, GucContext::Userset, GucFlags::default(), ); GucRegistry::define_enum_guc( - "vchordrq.io_rerank", - "`io_rerank` argument of vchordrq.", - "`io_rerank` argument of vchordrq.", - &IO_RERANK, + c"vchordrq.io_rerank", + c"`io_rerank` argument of vchordrq.", + c"`io_rerank` argument of vchordrq.", + &VCHORDRQ_IO_RERANK, GucContext::Userset, GucFlags::default(), ); unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); - #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrq".as_ptr()); } + GucRegistry::define_int_guc( + c"vchordg.ef_search", + c"`ef_search` argument of vchordg.", + c"`ef_search` argument of vchordg.", + &VCHORDG_EF_SEARCH, + 1, + 65535, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordg.beam_search", + c"`beam_search` argument of vchordg.", + c"`beam_search` argument of vchordg.", + &VCHORDG_BEAM_SIZE, + 1, + 65535, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordg.max_scan_tuples", + c"`max_scan_tuples` argument of vchordg.", + c"`max_scan_tuples` argument of vchordg.", + &VCHORDG_MAX_SCAN_TUPLES, + -1, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordg.io_search", + c"`io_search` argument of vchordg.", + c"`io_search` argument of vchordg.", + &VCHORDG_IO_SEARCH, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordg.io_rerank", + c"`io_rerank` argument of vchordg.", + c"`io_rerank` argument of vchordg.", + &VCHORDG_IO_RERANK, + GucContext::Userset, + GucFlags::default(), + ); + unsafe { + #[cfg(any(feature = "pg13", feature = "pg14"))] + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordg".as_ptr()); + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] + pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordg".as_ptr()); + } +} + +pub fn vchordg_ef_search() -> u32 { + VCHORDG_EF_SEARCH.get() as u32 +} + +pub fn vchordg_beam_search() -> u32 { + VCHORDG_BEAM_SIZE.get() as u32 +} + +pub fn vchordg_max_scan_tuples() -> Option { + let x = VCHORDG_MAX_SCAN_TUPLES.get(); + if x < 0 { None } else { Some(x as u32) } +} + +pub fn vchordg_io_search() -> crate::index::vchordg::scanners::Io { + use crate::index::vchordg::scanners::Io; + match VCHORDG_IO_SEARCH.get() { + PostgresIo::ReadBuffer => Io::Plain, + PostgresIo::PrefetchBuffer => Io::Simple, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream => Io::Stream, + } +} + +pub fn vchordg_io_rerank() -> crate::index::vchordg::scanners::Io { + use crate::index::vchordg::scanners::Io; + match VCHORDG_IO_RERANK.get() { + PostgresIo::ReadBuffer => Io::Plain, + PostgresIo::PrefetchBuffer => Io::Simple, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream => Io::Stream, + } } -pub fn probes() -> Vec { - match PROBES.get() { +pub fn vchordrq_probes() -> Vec { + match VCHORDRQ_PROBES.get() { None => Vec::new(), Some(probes) => { let mut result = Vec::new(); @@ -158,41 +276,43 @@ pub fn probes() -> Vec { } } -pub fn epsilon() -> f32 { - EPSILON.get() as f32 +pub fn vchordrq_epsilon() -> f32 { + VCHORDRQ_EPSILON.get() as f32 } -pub fn max_scan_tuples() -> Option { - let x = MAX_SCAN_TUPLES.get(); +pub fn vchordrq_max_scan_tuples() -> Option { + let x = VCHORDRQ_MAX_SCAN_TUPLES.get(); if x < 0 { None } else { Some(x as u32) } } -pub fn maxsim_refine() -> u32 { - MAXSIM_REFINE.get() as u32 +pub fn vchordrq_maxsim_refine() -> u32 { + VCHORDRQ_MAXSIM_REFINE.get() as u32 } -pub fn maxsim_threshold() -> u32 { - MAXSIM_THRESHOLD.get() as u32 +pub fn vchordrq_maxsim_threshold() -> u32 { + VCHORDRQ_MAXSIM_THRESHOLD.get() as u32 } -pub fn prefilter() -> bool { - PREFILTER.get() +pub fn vchordrq_prefilter() -> bool { + VCHORDRQ_PREFILTER.get() } -pub fn io_search() -> Io { - match IO_RERANK.get() { - PostgresIo::read_buffer => Io::Plain, - PostgresIo::prefetch_buffer => Io::Simple, - #[cfg(feature = "pg17")] - PostgresIo::read_stream => Io::Stream, +pub fn vchordrq_io_search() -> crate::index::vchordrq::scanners::Io { + use crate::index::vchordrq::scanners::Io; + match VCHORDRQ_IO_SEARCH.get() { + PostgresIo::ReadBuffer => Io::Plain, + PostgresIo::PrefetchBuffer => Io::Simple, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream => Io::Stream, } } -pub fn io_rerank() -> Io { - match IO_RERANK.get() { - PostgresIo::read_buffer => Io::Plain, - PostgresIo::prefetch_buffer => Io::Simple, - #[cfg(feature = "pg17")] - PostgresIo::read_stream => Io::Stream, +pub fn vchordrq_io_rerank() -> crate::index::vchordrq::scanners::Io { + use crate::index::vchordrq::scanners::Io; + match VCHORDRQ_IO_RERANK.get() { + PostgresIo::ReadBuffer => Io::Plain, + PostgresIo::PrefetchBuffer => Io::Simple, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream => Io::Stream, } } diff --git a/src/index/hook.rs b/src/index/hook.rs index 5e8665a0..636c4c5d 100644 --- a/src/index/hook.rs +++ b/src/index/hook.rs @@ -19,7 +19,7 @@ unsafe extern "C-unwind" fn rewrite_plan_state( node: *mut pgrx::pg_sys::PlanState, context: *mut core::ffi::c_void, ) -> bool { - unsafe fn dirty_check(index_relation: *mut pgrx::pg_sys::RelationData) -> Option { + unsafe fn dirty_check_vchordg(index_relation: *mut pgrx::pg_sys::RelationData) -> Option { type FnPtr = unsafe extern "C-unwind" fn( *mut pgrx::pg_sys::RelationData, i32, @@ -31,7 +31,26 @@ unsafe extern "C-unwind" fn rewrite_plan_state( let ambeginscan = indam.ambeginscan.as_ref()?; Some(core::ptr::fn_addr_eq::( *ambeginscan, - super::am::ambeginscan, + crate::index::vchordg::am::ambeginscan, + )) + } + } + + unsafe fn dirty_check_vchordrq( + index_relation: *mut pgrx::pg_sys::RelationData, + ) -> Option { + type FnPtr = unsafe extern "C-unwind" fn( + *mut pgrx::pg_sys::RelationData, + i32, + i32, + ) -> *mut pgrx::pg_sys::IndexScanDescData; + unsafe { + let index_relation = index_relation.as_ref()?; + let indam = index_relation.rd_indam.as_ref()?; + let ambeginscan = indam.ambeginscan.as_ref()?; + Some(core::ptr::fn_addr_eq::( + *ambeginscan, + crate::index::vchordrq::am::ambeginscan, )) } } @@ -40,28 +59,58 @@ unsafe extern "C-unwind" fn rewrite_plan_state( if (*node).type_ == pgrx::pg_sys::NodeTag::T_IndexScanState { let node = node as *mut pgrx::pg_sys::IndexScanState; let index_relation = (*node).iss_RelationDesc; - if Some(true) == dirty_check(index_relation) && (*node).iss_ScanDesc.is_null() { - use crate::index::am::Scanner; + if (*node).iss_ScanDesc.is_null() { + if Some(true) == dirty_check_vchordg(index_relation) { + use crate::index::vchordg::am::Scanner; + + (*node).iss_ScanDesc = pgrx::pg_sys::index_beginscan( + (*node).ss.ss_currentRelation, + (*node).iss_RelationDesc, + (*(*node).ss.ps.state).es_snapshot, + #[cfg(feature = "pg18")] + std::ptr::null_mut(), + (*node).iss_NumScanKeys, + (*node).iss_NumOrderByKeys, + ); - (*node).iss_ScanDesc = pgrx::pg_sys::index_beginscan( - (*node).ss.ss_currentRelation, - (*node).iss_RelationDesc, - (*(*node).ss.ps.state).es_snapshot, - (*node).iss_NumScanKeys, - (*node).iss_NumOrderByKeys, - ); + let scanner = &mut *((*(*node).iss_ScanDesc).opaque as *mut Scanner); + scanner.hack = std::ptr::NonNull::new(node); - let scanner = &mut *((*(*node).iss_ScanDesc).opaque as *mut Scanner); - scanner.hack = std::ptr::NonNull::new(node); + if (*node).iss_NumRuntimeKeys == 0 || (*node).iss_RuntimeKeysReady { + pgrx::pg_sys::index_rescan( + (*node).iss_ScanDesc, + (*node).iss_ScanKeys, + (*node).iss_NumScanKeys, + (*node).iss_OrderByKeys, + (*node).iss_NumOrderByKeys, + ); + } + } + if Some(true) == dirty_check_vchordrq(index_relation) { + use crate::index::vchordrq::am::Scanner; - if (*node).iss_NumRuntimeKeys == 0 || (*node).iss_RuntimeKeysReady { - pgrx::pg_sys::index_rescan( - (*node).iss_ScanDesc, - (*node).iss_ScanKeys, + (*node).iss_ScanDesc = pgrx::pg_sys::index_beginscan( + (*node).ss.ss_currentRelation, + (*node).iss_RelationDesc, + (*(*node).ss.ps.state).es_snapshot, + #[cfg(feature = "pg18")] + std::ptr::null_mut(), (*node).iss_NumScanKeys, - (*node).iss_OrderByKeys, (*node).iss_NumOrderByKeys, ); + + let scanner = &mut *((*(*node).iss_ScanDesc).opaque as *mut Scanner); + scanner.hack = std::ptr::NonNull::new(node); + + if (*node).iss_NumRuntimeKeys == 0 || (*node).iss_RuntimeKeysReady { + pgrx::pg_sys::index_rescan( + (*node).iss_ScanDesc, + (*node).iss_ScanKeys, + (*node).iss_NumScanKeys, + (*node).iss_OrderByKeys, + (*node).iss_NumOrderByKeys, + ); + } } } } @@ -71,11 +120,23 @@ unsafe extern "C-unwind" fn rewrite_plan_state( static PREV_EXECUTOR_START: AtomicPtr<()> = AtomicPtr::new(core::ptr::null_mut()); +#[cfg(any( + feature = "pg13", + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17" +))] +type ExecutorStartReturnType = (); + +#[cfg(feature = "pg18")] +type ExecutorStartReturnType = bool; + #[pgrx::pg_guard] unsafe extern "C-unwind" fn executor_start( query_desc: *mut pgrx::pg_sys::QueryDesc, eflags: ::std::os::raw::c_int, -) { +) -> ExecutorStartReturnType { unsafe { use core::mem::transmute; use pgrx::pg_sys::ExecutorStart_hook_type; @@ -83,14 +144,16 @@ unsafe extern "C-unwind" fn executor_start( let value = transmute::<*mut (), ExecutorStart_hook_type>( PREV_EXECUTOR_START.load(Ordering::Relaxed), ); - if let Some(prev_executor_start) = value { - prev_executor_start(query_desc, eflags); + #[allow(clippy::let_unit_value)] + let result = if let Some(prev_executor_start) = value { + prev_executor_start(query_desc, eflags) } else { - pgrx::pg_sys::standard_ExecutorStart(query_desc, eflags); - } + pgrx::pg_sys::standard_ExecutorStart(query_desc, eflags) + }; let planstate = (*query_desc).planstate; let context = core::ptr::null_mut(); rewrite_plan_state(planstate, context); + result } } @@ -99,7 +162,7 @@ pub fn init() { use core::mem::transmute; use std::sync::atomic::Ordering; PREV_EXECUTOR_START.store( - transmute::, *mut ()>( + transmute:: _>, *mut ()>( pgrx::pg_sys::ExecutorStart_hook, ), Ordering::Relaxed, diff --git a/src/index/mod.rs b/src/index/mod.rs index 8d2f48e2..55b8d86c 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -12,18 +12,18 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod algorithm; -pub mod am; +pub mod fetcher; pub mod functions; pub mod gucs; pub mod hook; pub mod opclass; -pub mod scanners; pub mod storage; -pub mod types; +pub mod vchordg; +pub mod vchordrq; pub fn init() { - am::init(); - hook::init(); gucs::init(); + hook::init(); + vchordrq::am::init(); + vchordg::am::init(); } diff --git a/src/index/opclass.rs b/src/index/opclass.rs index 1079fd73..6ae30711 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -1,278 +1,74 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_vector_l2_ops() -> String { + "vchordg_vector_l2_ops".to_string() +} -use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; -use crate::datatype::memory_vector::{VectorInput, VectorOutput}; -use algorithm::types::*; -use distance::Distance; -use pgrx::datum::FromDatum; -use pgrx::heap_tuple::PgHeapTuple; -use pgrx::pg_sys::Datum; -use std::num::NonZero; -use vector::VectorBorrowed; +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_vector_cosine_ops() -> String { + "vchordg_vector_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_vector_ip_ops() -> String { + "vchordg_vector_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_halfvec_l2_ops() -> String { + "vchordg_halfvec_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_halfvec_cosine_ops() -> String { + "vchordg_halfvec_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_halfvec_ip_ops() -> String { + "vchordg_halfvec_ip_ops".to_string() +} #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_vector_l2_ops() -> String { - "vector_l2_ops".to_string() + "vchordrq_vector_l2_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_vector_ip_ops() -> String { - "vector_ip_ops".to_string() + "vchordrq_vector_ip_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_vector_cosine_ops() -> String { - "vector_cosine_ops".to_string() + "vchordrq_vector_cosine_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_halfvec_l2_ops() -> String { - "halfvec_l2_ops".to_string() + "vchordrq_halfvec_l2_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_halfvec_ip_ops() -> String { - "halfvec_ip_ops".to_string() + "vchordrq_halfvec_ip_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_halfvec_cosine_ops() -> String { - "halfvec_cosine_ops".to_string() + "vchordrq_halfvec_cosine_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_vector_maxsim_ops() -> String { - "vector_maxsim_ops".to_string() + "vchordrq_vector_maxsim_ops".to_string() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_halfvec_maxsim_ops() -> String { - "halfvec_maxsim_ops".to_string() + "vchordrq_halfvec_maxsim_ops".to_string() } pub struct Sphere { pub center: T, pub radius: f32, } - -#[derive(Debug, Clone, Copy)] -pub enum Opfamily { - VectorL2, - VectorIp, - VectorCosine, - HalfvecL2, - HalfvecIp, - HalfvecCosine, - VectorMaxsim, - HalfvecMaxsim, -} - -impl Opfamily { - fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { - use {BorrowedVector as B, OwnedVector as O}; - match (vector, self) { - (B::Vecf32(x), Self::VectorL2) => O::Vecf32(x.own()), - (B::Vecf32(x), Self::VectorIp | Self::VectorMaxsim) => O::Vecf32(x.own()), - (B::Vecf32(x), Self::VectorCosine) => O::Vecf32(x.function_normalize()), - (B::Vecf32(_), _) => unreachable!(), - (B::Vecf16(x), Self::HalfvecL2) => O::Vecf16(x.own()), - (B::Vecf16(x), Self::HalfvecIp | Self::HalfvecMaxsim) => O::Vecf16(x.own()), - (B::Vecf16(x), Self::HalfvecCosine) => O::Vecf16(x.function_normalize()), - (B::Vecf16(_), _) => unreachable!(), - } - } - pub unsafe fn store(self, datum: Datum) -> Option> { - if datum.is_null() { - return None; - } - let store = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { - let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; - vec![(self.input(BorrowedVector::Vecf32(vector.as_borrowed())), 0)] - } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { - let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; - vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] - } - Self::VectorMaxsim => { - let vectors = - unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; - let mut result = Vec::with_capacity(vectors.len()); - for (i, vector) in vectors.iter_deny_null().enumerate() { - result.push(( - self.input(BorrowedVector::Vecf32(vector.as_borrowed())), - i as u16, - )); - } - result - } - Self::HalfvecMaxsim => { - let vectors = - unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; - let mut result = Vec::with_capacity(vectors.len()); - for (i, vector) in vectors.iter_deny_null().enumerate() { - result.push(( - self.input(BorrowedVector::Vecf16(vector.as_borrowed())), - i as u16, - )); - } - result - } - }; - Some(store) - } - pub unsafe fn input_sphere(self, datum: Datum) -> Option> { - if datum.is_null() { - return None; - } - let attno_1 = NonZero::new(1_usize).unwrap(); - let attno_2 = NonZero::new(2_usize).unwrap(); - let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; - let center = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { - let vector = tuple.get_by_index::(attno_1).unwrap()?; - self.input(BorrowedVector::Vecf32(vector.as_borrowed())) - } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { - let vector = tuple.get_by_index::(attno_1).unwrap()?; - self.input(BorrowedVector::Vecf16(vector.as_borrowed())) - } - }; - let radius = tuple.get_by_index::(attno_2).unwrap()?; - Some(Sphere { center, radius }) - } - pub unsafe fn input_vector(self, datum: Datum) -> Option { - if datum.is_null() { - return None; - } - let vector = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { - let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; - self.input(BorrowedVector::Vecf32(vector.as_borrowed())) - } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { - let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; - self.input(BorrowedVector::Vecf16(vector.as_borrowed())) - } - }; - Some(vector) - } - pub unsafe fn input_vectors(self, datum: Datum) -> Option> { - if datum.is_null() { - return None; - } - let vectors = match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { - let vectors = - unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; - let mut result = Vec::with_capacity(vectors.len()); - for vector in vectors.iter_deny_null() { - result.push(self.input(BorrowedVector::Vecf32(vector.as_borrowed()))); - } - result - } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { - let vectors = - unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; - let mut result = Vec::with_capacity(vectors.len()); - for vector in vectors.iter_deny_null() { - result.push(self.input(BorrowedVector::Vecf16(vector.as_borrowed()))); - } - result - } - }; - Some(vectors) - } - pub fn output(self, x: Distance) -> f32 { - match self { - Self::VectorCosine | Self::HalfvecCosine => x.to_f32() + 1.0f32, - Self::VectorL2 | Self::HalfvecL2 => x.to_f32().sqrt(), - Self::VectorIp | Self::HalfvecIp | Self::VectorMaxsim | Self::HalfvecMaxsim => { - x.to_f32() - } - } - } - pub const fn distance_kind(self) -> DistanceKind { - match self { - Self::VectorL2 | Self::HalfvecL2 => DistanceKind::L2, - Self::VectorIp - | Self::HalfvecIp - | Self::VectorCosine - | Self::HalfvecCosine - | Self::VectorMaxsim - | Self::HalfvecMaxsim => DistanceKind::Dot, - } - } - pub const fn vector_kind(self) -> VectorKind { - match self { - Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { - VectorKind::Vecf32 - } - Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { - VectorKind::Vecf16 - } - } - } -} - -pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { - use pgrx::pg_sys::Oid; - - let proc = unsafe { pgrx::pg_sys::index_getprocid(index_relation, 1, 1) }; - - if proc == Oid::INVALID { - pgrx::error!("support function 1 is not found"); - } - - let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); - - unsafe { - pgrx::pg_sys::fmgr_info(proc, &mut flinfo); - } - - let fn_addr = flinfo.fn_addr.expect("null function pointer"); - - let mut fcinfo = unsafe { std::mem::zeroed::() }; - fcinfo.flinfo = &mut flinfo; - fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; - fcinfo.context = std::ptr::null_mut(); - fcinfo.resultinfo = std::ptr::null_mut(); - fcinfo.isnull = true; - fcinfo.nargs = 0; - - let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; - - let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; - - let result_string = result_option.expect("null return value"); - - let result = match result_string.as_str() { - "vector_l2_ops" => Opfamily::VectorL2, - "vector_ip_ops" => Opfamily::VectorIp, - "vector_cosine_ops" => Opfamily::VectorCosine, - "halfvec_l2_ops" => Opfamily::HalfvecL2, - "halfvec_ip_ops" => Opfamily::HalfvecIp, - "halfvec_cosine_ops" => Opfamily::HalfvecCosine, - "vector_maxsim_ops" => Opfamily::VectorMaxsim, - "halfvec_maxsim_ops" => Opfamily::HalfvecMaxsim, - _ => pgrx::error!("unknown operator class"), - }; - - unsafe { - pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); - } - - result -} diff --git a/src/index/storage.rs b/src/index/storage.rs index ef88e2b0..0cc09fd4 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -12,53 +12,41 @@ // // Copyright (c) 2025 TensorChord Inc. -use algorithm::*; +use algo::{Opaque, *}; use std::collections::VecDeque; use std::iter::{Chain, Flatten}; +use std::marker::PhantomData; use std::mem::{MaybeUninit, offset_of}; use std::ops::{Deref, DerefMut}; use std::ptr::NonNull; -const _: () = assert!( - offset_of!(pgrx::pg_sys::PageHeaderData, pd_linp) % pgrx::pg_sys::MAXIMUM_ALIGNOF as usize == 0 -); - -const fn size_of_contents() -> usize { - use pgrx::pg_sys::{BLCKSZ, PageHeaderData}; - let size_of_page = BLCKSZ as usize; - let size_of_header = offset_of!(PageHeaderData, pd_linp); - let size_of_opaque = size_of::(); - size_of_page - size_of_header - size_of_opaque -} - #[repr(C, align(8))] #[derive(Debug)] -pub struct PostgresPage { +pub struct PostgresPage { header: pgrx::pg_sys::PageHeaderData, - content: [u8; size_of_contents()], - opaque: Opaque, + content: [u8; pgrx::pg_sys::BLCKSZ as usize - size_of::()], + _opaque: PhantomData O>, } -const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); -const _: () = assert!(size_of::() == pgrx::pg_sys::BLCKSZ as usize); - -impl PostgresPage { - fn init_mut(this: &mut MaybeUninit) -> &mut Self { - unsafe { - pgrx::pg_sys::PageInit( - this.as_mut_ptr() as pgrx::pg_sys::Page, - pgrx::pg_sys::BLCKSZ as usize, - size_of::(), - ); - (&raw mut (*this.as_mut_ptr()).opaque).write(Opaque { - next: u32::MAX, - skip: u32::MAX, - }); - } - let this = unsafe { MaybeUninit::assume_init_mut(this) }; - assert_eq!(offset_of!(Self, opaque), this.header.pd_special as usize); - this - } +// It is a non-guaranteed detection. +// If `PageHeaderData` contains padding bytes, const-eval probably fails. +const _: () = { + use pgrx::pg_sys::PageHeaderData as T; + use std::mem::{transmute, zeroed}; + const _ZERO: &[u8; size_of::()] = unsafe { transmute(&zeroed::()) }; +}; + +// Layout checks of header. +const _: () = { + use pgrx::pg_sys::{MAXIMUM_ALIGNOF, PageHeaderData as T}; + assert!(size_of::() == offset_of!(T, pd_linp)); + assert!(size_of::() % MAXIMUM_ALIGNOF as usize == 0); +}; + +const _: () = assert!(align_of::>() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); +const _: () = assert!(size_of::>() == pgrx::pg_sys::BLCKSZ as usize); + +impl PostgresPage { pub fn clone_into_boxed(&self) -> Box { let mut result = Box::new_uninit(); unsafe { @@ -68,12 +56,15 @@ impl PostgresPage { } } -impl Page for PostgresPage { - fn get_opaque(&self) -> &Opaque { - &self.opaque +impl Page for PostgresPage { + type Opaque = O; + fn get_opaque(&self) -> &O { + assert!(self.header.pd_special as usize + size_of::() == size_of::()); + unsafe { &*((self as *const _ as *const O).byte_add(self.header.pd_special as _)) } } - fn get_opaque_mut(&mut self) -> &mut Opaque { - &mut self.opaque + fn get_opaque_mut(&mut self) -> &mut O { + assert!(self.header.pd_special as usize + size_of::() == size_of::()); + unsafe { &mut *((self as *mut _ as *mut O).byte_add(self.header.pd_special as _)) } } fn len(&self) -> u16 { use pgrx::pg_sys::{ItemIdData, PageHeaderData}; @@ -155,45 +146,48 @@ impl Page for PostgresPage { fn freespace(&self) -> u16 { unsafe { pgrx::pg_sys::PageGetFreeSpace((self as *const Self).cast_mut().cast()) as u16 } } - fn clear(&mut self) { + fn clear(&mut self, opaque: O) { unsafe { - pgrx::pg_sys::PageInit( - (self as *mut PostgresPage as pgrx::pg_sys::Page).cast(), - pgrx::pg_sys::BLCKSZ as usize, - size_of::(), - ); - (&raw mut self.opaque).write(Opaque { - next: u32::MAX, - skip: u32::MAX, - }); + page_init(self, opaque); } - assert_eq!(offset_of!(Self, opaque), self.header.pd_special as usize); } } -const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); +unsafe fn page_init(this: *mut PostgresPage, opaque: O) { + unsafe { + use pgrx::pg_sys::{BLCKSZ, PageHeaderData, PageInit}; + PageInit(this.cast(), BLCKSZ as usize, size_of::()); + assert_eq!( + (*this.cast::()).pd_special as usize + size_of::(), + size_of::>() + ); + this.cast::() + .byte_add(size_of::>() - size_of::()) + .write(opaque); + } +} -pub struct PostgresBufferReadGuard { +pub struct PostgresBufferReadGuard { buf: i32, - page: NonNull, + page: NonNull>, id: u32, } -impl PageGuard for PostgresBufferReadGuard { +impl PageGuard for PostgresBufferReadGuard { fn id(&self) -> u32 { self.id } } -impl Deref for PostgresBufferReadGuard { - type Target = PostgresPage; +impl Deref for PostgresBufferReadGuard { + type Target = PostgresPage; - fn deref(&self) -> &PostgresPage { + fn deref(&self) -> &PostgresPage { unsafe { self.page.as_ref() } } } -impl Drop for PostgresBufferReadGuard { +impl Drop for PostgresBufferReadGuard { fn drop(&mut self) { unsafe { pgrx::pg_sys::UnlockReleaseBuffer(self.buf); @@ -201,36 +195,36 @@ impl Drop for PostgresBufferReadGuard { } } -pub struct PostgresBufferWriteGuard { +pub struct PostgresBufferWriteGuard { raw: pgrx::pg_sys::Relation, buf: i32, - page: NonNull, + page: NonNull>, state: *mut pgrx::pg_sys::GenericXLogState, id: u32, tracking_freespace: bool, } -impl PageGuard for PostgresBufferWriteGuard { +impl PageGuard for PostgresBufferWriteGuard { fn id(&self) -> u32 { self.id } } -impl Deref for PostgresBufferWriteGuard { - type Target = PostgresPage; +impl Deref for PostgresBufferWriteGuard { + type Target = PostgresPage; - fn deref(&self) -> &PostgresPage { + fn deref(&self) -> &PostgresPage { unsafe { self.page.as_ref() } } } -impl DerefMut for PostgresBufferWriteGuard { - fn deref_mut(&mut self) -> &mut PostgresPage { +impl DerefMut for PostgresBufferWriteGuard { + fn deref_mut(&mut self) -> &mut PostgresPage { unsafe { self.page.as_mut() } } } -impl Drop for PostgresBufferWriteGuard { +impl Drop for PostgresBufferWriteGuard { fn drop(&mut self) { unsafe { if std::thread::panicking() { @@ -248,23 +242,29 @@ impl Drop for PostgresBufferWriteGuard { } #[derive(Debug, Clone)] -pub struct PostgresRelation { +pub struct PostgresRelation { raw: pgrx::pg_sys::Relation, + _phantom: PhantomData Opaque>, } -impl PostgresRelation { +impl PostgresRelation { pub unsafe fn new(raw: pgrx::pg_sys::Relation) -> Self { - Self { raw } + Self { + raw, + _phantom: PhantomData, + } } } -impl Relation for PostgresRelation { - type Page = PostgresPage; +impl Relation for PostgresRelation { + type Page = PostgresPage; } -impl RelationRead for PostgresRelation { - type ReadGuard<'a> = PostgresBufferReadGuard; +impl RelationReadTypes for PostgresRelation { + type ReadGuard<'a> = PostgresBufferReadGuard; +} +impl RelationRead for PostgresRelation { fn read(&self, id: u32) -> Self::ReadGuard<'_> { assert!(id != u32::MAX, "no such page"); unsafe { @@ -285,10 +285,12 @@ impl RelationRead for PostgresRelation { } } -impl RelationWrite for PostgresRelation { - type WriteGuard<'a> = PostgresBufferWriteGuard; +impl RelationWriteTypes for PostgresRelation { + type WriteGuard<'a> = PostgresBufferWriteGuard; +} - fn write(&self, id: u32, tracking_freespace: bool) -> PostgresBufferWriteGuard { +impl RelationWrite for PostgresRelation { + fn write(&self, id: u32, tracking_freespace: bool) -> PostgresBufferWriteGuard { assert!(id != u32::MAX, "no such page"); unsafe { use pgrx::pg_sys::{ @@ -307,7 +309,7 @@ impl RelationWrite for PostgresRelation { let state = GenericXLogStart(self.raw); let page = NonNull::new( GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) - .cast::>(), + .cast::>>(), ) .expect("failed to get page"); PostgresBufferWriteGuard { @@ -320,7 +322,11 @@ impl RelationWrite for PostgresRelation { } } } - fn extend(&self, tracking_freespace: bool) -> PostgresBufferWriteGuard { + fn extend( + &self, + opaque: ::Opaque, + tracking_freespace: bool, + ) -> PostgresBufferWriteGuard { unsafe { use pgrx::pg_sys::{ BUFFER_LOCK_EXCLUSIVE, ExclusiveLock, ForkNumber, GENERIC_XLOG_FULL_IMAGE, @@ -340,10 +346,10 @@ impl RelationWrite for PostgresRelation { let state = GenericXLogStart(self.raw); let mut page = NonNull::new( GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) - .cast::>(), + .cast::>>(), ) .expect("failed to get page"); - PostgresPage::init_mut(page.as_mut()); + page_init(page.as_mut().as_mut_ptr(), opaque); PostgresBufferWriteGuard { raw: self.raw, buf, @@ -354,7 +360,7 @@ impl RelationWrite for PostgresRelation { } } } - fn search(&self, freespace: usize) -> Option { + fn search(&self, freespace: usize) -> Option> { unsafe { loop { let id = pgrx::pg_sys::GetPageWithFreeSpace(self.raw, freespace); @@ -374,7 +380,7 @@ impl RelationWrite for PostgresRelation { } } -impl RelationPrefetch for PostgresRelation { +impl RelationPrefetch for PostgresRelation { fn prefetch(&self, id: u32) { assert!(id != u32::MAX, "no such page"); unsafe { @@ -445,40 +451,76 @@ where } } -pub struct PostgresReadStream { - #[cfg(feature = "pg17")] +pub struct PostgresReadStream { + #[cfg(any(feature = "pg17", feature = "pg18"))] raw: *mut pgrx::pg_sys::ReadStream, // Because of `Box`'s special alias rules, `Box` cannot be used here. cache: NonNull>, + _phantom: PhantomData O>, } -#[cfg(feature = "pg17")] -impl PostgresReadStream +pub struct PostgresReadStreamGuards { + #[cfg(any(feature = "pg17", feature = "pg18"))] + raw: *mut pgrx::pg_sys::ReadStream, + #[cfg_attr( + any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"), + expect(dead_code) + )] + list: L, + _phantom: PhantomData (O, I)>, +} + +impl Iterator for PostgresReadStreamGuards where - I::Item: Fetch, + L: Iterator, { - fn read(&mut self, fetch: &[u32]) -> impl Iterator { - fetch.iter().map(|_| unsafe { + type Item = PostgresBufferReadGuard; + + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + fn next(&mut self) -> Option { + unimplemented!() + } + + #[cfg(any(feature = "pg17", feature = "pg18"))] + fn next(&mut self) -> Option { + let _id = self.list.next()?; + unsafe { use pgrx::pg_sys::{ BUFFER_LOCK_SHARE, BufferGetPage, LockBuffer, read_stream_next_buffer, }; let buf = read_stream_next_buffer(self.raw, core::ptr::null_mut()); LockBuffer(buf, BUFFER_LOCK_SHARE as _); let page = NonNull::new(BufferGetPage(buf).cast()).expect("failed to get page"); - PostgresBufferReadGuard { + Some(PostgresBufferReadGuard { buf, page, id: pgrx::pg_sys::BufferGetBlockNumber(buf), - } - }) + }) + } } } -impl<'r, I: Iterator> ReadStream<'r> for PostgresReadStream +#[cfg(any(feature = "pg17", feature = "pg18"))] +impl PostgresReadStream where I::Item: Fetch, { - type Relation = PostgresRelation; + fn read>(&mut self, fetch: L) -> PostgresReadStreamGuards { + PostgresReadStreamGuards { + raw: self.raw, + list: fetch, + _phantom: PhantomData, + } + } +} + +impl<'r, O: Opaque, I: Iterator> ReadStream<'r> for PostgresReadStream +where + I::Item: Fetch, +{ + type Relation = PostgresRelation; + + type Guards = PostgresReadStreamGuards>; type Item = I::Item; @@ -486,35 +528,35 @@ where Chain, Flatten>>; #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] - fn next(&mut self) -> Option<(I::Item, Vec)> { + fn next(&mut self) -> Option<(I::Item, Self::Guards)> { panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); } #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] - fn next_if( + fn next_if bool>( &mut self, - _predicate: impl FnOnce(&I::Item) -> bool, - ) -> Option<(I::Item, Vec)> { + _predicate: P, + ) -> Option<(I::Item, Self::Guards)> { panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); } - #[cfg(feature = "pg17")] - fn next(&mut self) -> Option<(I::Item, Vec)> { + #[cfg(any(feature = "pg17", feature = "pg18"))] + fn next(&mut self) -> Option<(I::Item, Self::Guards)> { if let Some(e) = unsafe { self.cache.as_mut().pop_item() } { - let list = self.read(e.fetch()).collect(); + let list = self.read(e.fetch().into_iter()); Some((e, list)) } else { None } } - #[cfg(feature = "pg17")] - fn next_if( + #[cfg(any(feature = "pg17", feature = "pg18"))] + fn next_if bool>( &mut self, - predicate: impl FnOnce(&I::Item) -> bool, - ) -> Option<(I::Item, Vec)> { + predicate: P, + ) -> Option<(I::Item, Self::Guards)> { if let Some(e) = unsafe { self.cache.as_mut().pop_item_if(predicate) } { - let list = self.read(e.fetch()).collect(); + let list = self.read(e.fetch().into_iter()); Some((e, list)) } else { None @@ -530,11 +572,11 @@ where } } -impl Drop for PostgresReadStream { +impl Drop for PostgresReadStream { fn drop(&mut self) { unsafe { let _ = std::mem::take(self.cache.as_mut()); - #[cfg(feature = "pg17")] + #[cfg(any(feature = "pg17", feature = "pg18"))] if !std::thread::panicking() && pgrx::pg_sys::IsTransactionState() { pgrx::pg_sys::read_stream_end(self.raw); } @@ -543,12 +585,14 @@ impl Drop for PostgresReadStream { } } -impl RelationReadStream for PostgresRelation { - type ReadStream<'s, I: Iterator> - = PostgresReadStream +impl RelationReadStreamTypes for PostgresRelation { + type ReadStream<'r, I: Iterator> + = PostgresReadStream where I::Item: Fetch; +} +impl RelationReadStream for PostgresRelation { #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] fn read_stream(&self, _iter: I, _hints: Hints) -> Self::ReadStream<'_, I> where @@ -557,7 +601,7 @@ impl RelationReadStream for PostgresRelation { panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); } - #[cfg(feature = "pg17")] + #[cfg(any(feature = "pg17", feature = "pg18"))] fn read_stream(&self, iter: I, hints: Hints) -> Self::ReadStream<'_, I> where I::Item: Fetch, @@ -600,7 +644,18 @@ impl RelationReadStream for PostgresRelation { 0, ) }; - PostgresReadStream { raw, cache } + PostgresReadStream { + raw, + cache, + _phantom: PhantomData, + } + } +} + +impl RelationLength for PostgresRelation { + fn len(&self) -> u32 { + use pgrx::pg_sys::{ForkNumber, RelationGetNumberOfBlocksInFork}; + unsafe { RelationGetNumberOfBlocksInFork(self.raw, ForkNumber::MAIN_FORKNUM) } } } diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs new file mode 100644 index 00000000..89a0d1f3 --- /dev/null +++ b/src/index/vchordg/algo.rs @@ -0,0 +1,260 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::index::vchordg::opclass::Opfamily; +use algo::accessor::{Dot, L2S}; +use algo::prefetcher::*; +use algo::*; +use half::f16; +use std::num::NonZero; +use vchordg::Opaque; +use vchordg::operator::Op; +use vchordg::types::*; +use vector::VectorOwned; +use vector::vect::{VectBorrowed, VectOwned}; + +pub fn bulkdelete( + opfamily: Opfamily, + index: &R, + check: impl Fn(), + callback: impl Fn(NonZero) -> bool, +) where + R: RelationRead + RelationWrite + RelationLength, + R::Page: Page, +{ + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordg::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordg::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordg::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + } + } +} + +pub fn maintain(opfamily: Opfamily, index: &R, check: impl Fn()) +where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + let bump = bumpalo::Bump::new(); + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordg::maintain::<_, Op, L2S>>(index, &check, &bump); + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::maintain::<_, Op, L2S>>(index, &check, &bump); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordg::maintain::<_, Op, Dot>>(index, &check, &bump); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordg::maintain::<_, Op, Dot>>(index, &check, &bump); + } + } +} + +pub fn build(vector_options: VectorOptions, vchordg_options: VamanaIndexOptions, index: &R) +where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + match (vector_options.v, vector_options.d) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordg::build::<_, Op, L2S>>(vector_options, vchordg_options, index) + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::build::<_, Op, L2S>>(vector_options, vchordg_options, index) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordg::build::<_, Op, Dot>>(vector_options, vchordg_options, index) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordg::build::<_, Op, Dot>>(vector_options, vchordg_options, index) + } + } +} + +pub fn insert(opfamily: Opfamily, index: &R, payload: NonZero, vector: OwnedVector) +where + R: RelationRead + RelationWrite + RelationReadStream, + R::Page: Page, +{ + let bump = bumpalo::Bump::new(); + let make_vertex_plain_prefetcher = MakePlainPrefetcher { index }; + let make_vector_plain_prefetcher = MakePlainPrefetcher { index }; + match (vector, opfamily.distance_kind()) { + (OwnedVector::Vecf32(unprojected), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + let projected = RandomProject::project(unprojected.as_borrowed()); + vchordg::insert::<_, Op, L2S>>( + index, + projected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Vecf16(unprojected), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + let projected = RandomProject::project(unprojected.as_borrowed()); + vchordg::insert::<_, Op, L2S>>( + index, + projected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Vecf32(unprojected), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + let projected = RandomProject::project(unprojected.as_borrowed()); + vchordg::insert::<_, Op, Dot>>( + index, + projected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Vecf16(unprojected), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + let projected = RandomProject::project(unprojected.as_borrowed()); + vchordg::insert::<_, Op, Dot>>( + index, + projected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + } +} + +pub trait RandomProject { + type Output; + fn project(self) -> Self::Output; +} + +impl RandomProject for VectBorrowed<'_, f32> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use rabitq::rotate::rotate; + let input = self.slice(); + VectOwned::new(rotate(input)) + } +} + +impl RandomProject for VectBorrowed<'_, f16> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use rabitq::rotate::rotate; + use simd::Floating; + let input = f16::vector_to_f32(self.slice()); + VectOwned::new(f16::vector_from_f32(&rotate(&input))) + } +} + +#[derive(Debug)] +pub struct MakePlainPrefetcher<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakePlainPrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakePlainPrefetcher<'r, R> { + type P + = PlainPrefetcher<'r, R, S> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch, + { + PlainPrefetcher::new(self.index, seq) + } +} + +#[derive(Debug)] +pub struct MakeSimplePrefetcher<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakeSimplePrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> + for MakeSimplePrefetcher<'r, R> +{ + type P + = SimplePrefetcher<'r, R, S> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch, + { + SimplePrefetcher::new(self.index, seq) + } +} + +#[derive(Debug)] +pub struct MakeStreamPrefetcher<'r, R> { + pub index: &'r R, + pub hints: Hints, +} + +impl<'r, R> Clone for MakeStreamPrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { + index: self.index, + hints: self.hints.clone(), + } + } +} + +impl<'r, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'r, R> + for MakeStreamPrefetcher<'r, R> +{ + type P + = StreamPrefetcher<'r, R, S> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch, + { + StreamPrefetcher::new(self.index, seq, self.hints.clone()) + } +} diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs new file mode 100644 index 00000000..a076c8ad --- /dev/null +++ b/src/index/vchordg/am/am_build.rs @@ -0,0 +1,704 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::datatype::typmod::Typmod; +use crate::index::storage::PostgresRelation; +use crate::index::vchordg::am::{Reloption, ctid_to_key, kv_to_pointer}; +use crate::index::vchordg::opclass::{Opfamily, opfamily}; +use crate::index::vchordg::types::VamanaIndexingOptions; +use pgrx::pg_sys::{Datum, ItemPointerData}; +use std::ffi::CStr; +use vchordg::types::*; + +#[derive(Debug, Clone, Copy)] +#[repr(u16)] +pub enum BuildPhaseCode { + Initializing = 0, + Inserting = 1, +} + +pub struct BuildPhase(BuildPhaseCode, u16); + +impl BuildPhase { + pub const fn new(code: BuildPhaseCode, k: u16) -> Option { + match (code, k) { + (BuildPhaseCode::Initializing, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::Inserting, 0) => Some(BuildPhase(code, k)), + _ => None, + } + } + pub const fn name(self) -> &'static CStr { + match self { + BuildPhase(BuildPhaseCode::Initializing, k) => { + static RAW: [&CStr; 1] = [c"initializing"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::Inserting, k) => { + static RAW: [&CStr; 1] = [c"inserting tuples from table to index"]; + RAW[k as usize] + } + } + } + pub const fn from_code(code: BuildPhaseCode) -> Self { + Self(code, 0) + } + pub const fn from_value(value: u32) -> Option { + const INITIALIZING: u16 = BuildPhaseCode::Initializing as _; + const INSERTING: u16 = BuildPhaseCode::Inserting as _; + let k = value as u16; + match (value >> 16) as u16 { + INITIALIZING => Self::new(BuildPhaseCode::Initializing, k), + INSERTING => Self::new(BuildPhaseCode::Inserting, k), + _ => None, + } + } + pub const fn into_value(self) -> u32 { + (self.0 as u32) << 16 | (self.1 as u32) + } +} + +#[pgrx::pg_guard] +pub extern "C-unwind" fn ambuildphasename(x: i64) -> *mut core::ffi::c_char { + if let Ok(x) = u32::try_from(x.wrapping_sub(1)) { + if let Some(x) = BuildPhase::from_value(x) { + x.name().as_ptr().cast_mut() + } else { + std::ptr::null_mut() + } + } else { + std::ptr::null_mut() + } +} + +#[derive(Debug, Clone)] +struct Heap { + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + opfamily: Opfamily, + scan: *mut pgrx::pg_sys::TableScanDescData, +} + +impl Heap { + fn traverse))>( + &self, + progress: bool, + callback: F, + ) { + pub struct State<'a, F> { + pub this: &'a Heap, + pub callback: F, + } + #[pgrx::pg_guard] + unsafe extern "C-unwind" fn call( + _index_relation: pgrx::pg_sys::Relation, + ctid: pgrx::pg_sys::ItemPointer, + values: *mut Datum, + is_null: *mut bool, + _tuple_is_alive: bool, + state: *mut core::ffi::c_void, + ) where + F: FnMut((ItemPointerData, Vec<(OwnedVector, u16)>)), + { + let state = unsafe { &mut *state.cast::>() }; + let opfamily = state.this.opfamily; + let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; + let ctid = unsafe { ctid.read() }; + if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { + (state.callback)((ctid, store)); + } + } + let table_am = unsafe { &*(*self.heap_relation).rd_tableam }; + let mut state = State { + this: self, + callback, + }; + unsafe { + table_am.index_build_range_scan.unwrap()( + self.heap_relation, + self.index_relation, + self.index_info, + true, + false, + progress, + 0, + pgrx::pg_sys::InvalidBlockNumber, + Some(call::), + (&mut state) as *mut State as *mut _, + self.scan, + ); + } + } +} + +#[derive(Debug, Clone)] +struct PostgresReporter {} + +impl PostgresReporter { + fn phase(&mut self, phase: BuildPhase) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_SUBPHASE as _, + (phase.into_value() as i64) + 1, + ); + } + } + fn tuples_total(&mut self, tuples_total: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, + tuples_total as _, + ); + } + } + fn tuples_done(&mut self, tuples_done: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, + tuples_done as _, + ); + } + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambuild( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, +) -> *mut pgrx::pg_sys::IndexBuildResult { + use validator::Validate; + let (vector_options, vchordg_options) = unsafe { options(index_relation) }; + if let Err(errors) = Validate::validate(&vector_options) { + pgrx::error!("error while validating options: {}", errors); + } + if let Err(errors) = Validate::validate(&vchordg_options) { + pgrx::error!("error while validating options: {}", errors); + } + if !matches!(vector_options.d, DistanceKind::L2S) && vchordg_options.index.alpha != 1.0 { + pgrx::error!("alpha must be 1 if L2 metric is not used"); + } + let opfamily = unsafe { opfamily(index_relation) }; + let index = unsafe { PostgresRelation::new(index_relation) }; + let heap = Heap { + heap_relation, + index_relation, + index_info, + opfamily, + scan: std::ptr::null_mut(), + }; + let mut reporter = PostgresReporter {}; + crate::index::vchordg::algo::build(vector_options, vchordg_options.index, &index); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); + let cache = vchordg_cached::VamanaCached::_0 {}; + if let Some(leader) = unsafe { + VamanaLeader::enter( + heap_relation, + index_relation, + (*index_info).ii_Concurrent, + cache, + ) + } { + unsafe { + parallel_build( + index_relation, + heap_relation, + index_info, + leader.tablescandesc, + leader.vchordgshared, + leader.vchordgcached, + |indtuples| { + reporter.tuples_done(indtuples); + }, + ); + leader.wait(); + let nparticipants = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordgshared).mutex); + if (*leader.vchordgshared).nparticipantsdone == nparticipants { + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordgshared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordgshared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*leader.vchordgshared).workersdonecv, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + reporter.tuples_done((*leader.vchordgshared).indtuples); + reporter.tuples_total((*leader.vchordgshared).indtuples); + pgrx::pg_sys::ConditionVariableCancelSleep(); + } + } else { + let mut indtuples = 0; + reporter.tuples_done(indtuples); + heap.traverse(true, |(ctid, store)| { + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + crate::index::vchordg::algo::insert(opfamily, &index, payload, vector); + } + indtuples += 1; + reporter.tuples_done(indtuples); + }); + reporter.tuples_total(indtuples); + } + unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } +} + +struct VamanaShared { + /* Immutable state */ + heaprelid: pgrx::pg_sys::Oid, + indexrelid: pgrx::pg_sys::Oid, + isconcurrent: bool, + + /* Worker progress */ + workersdonecv: pgrx::pg_sys::ConditionVariable, + + /* Mutex for mutable state */ + mutex: pgrx::pg_sys::slock_t, + + /* Mutable state */ + nparticipantsdone: i32, + indtuples: u64, +} + +fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { + matches!( + unsafe { (*snapshot).snapshot_type }, + pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + | pgrx::pg_sys::SnapshotType::SNAPSHOT_HISTORIC_MVCC + ) +} + +struct VamanaLeader { + pcxt: *mut pgrx::pg_sys::ParallelContext, + nparticipants: i32, + snapshot: pgrx::pg_sys::Snapshot, + vchordgshared: *mut VamanaShared, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + vchordgcached: *const u8, +} + +impl VamanaLeader { + pub unsafe fn enter( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + isconcurrent: bool, + cache: vchordg_cached::VamanaCached, + ) -> Option { + let _cache = cache.serialize(); + #[expect(clippy::drop_non_drop)] + drop(cache); + let cache = _cache; + + unsafe fn compute_parallel_workers( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + ) -> i32 { + unsafe { + if pgrx::pg_sys::plan_create_index_workers( + (*heap_relation).rd_id, + (*index_relation).rd_id, + ) == 0 + { + return 0; + } + if !(*heap_relation).rd_options.is_null() { + let std_options = (*heap_relation) + .rd_options + .cast::(); + std::cmp::min( + (*std_options).parallel_workers, + pgrx::pg_sys::max_parallel_maintenance_workers, + ) + } else { + pgrx::pg_sys::max_parallel_maintenance_workers + } + } + } + + let request = unsafe { compute_parallel_workers(heap_relation, index_relation) }; + if request <= 0 { + return None; + } + + unsafe { + pgrx::pg_sys::EnterParallelMode(); + } + let pcxt = unsafe { + pgrx::pg_sys::CreateParallelContext( + c"vchord".as_ptr(), + c"vchordg_parallel_build_main".as_ptr(), + request, + ) + }; + + let snapshot = if isconcurrent { + unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } + } else { + &raw mut pgrx::pg_sys::SnapshotAnyData + }; + + fn estimate_chunk(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.space_for_chunks += x.next_multiple_of(pgrx::pg_sys::ALIGNOF_BUFFER as _); + } + fn estimate_keys(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.number_of_keys += x; + } + let est_tablescandesc = + unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap_relation, snapshot) }; + unsafe { + estimate_chunk(&mut (*pcxt).estimator, size_of::()); + estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); + estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, 8 + cache.len()); + estimate_keys(&mut (*pcxt).estimator, 1); + } + + unsafe { + pgrx::pg_sys::InitializeParallelDSM(pcxt); + if (*pcxt).seg.is_null() { + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + let vchordgshared = unsafe { + let vchordgshared = + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) + .cast::(); + vchordgshared.write(VamanaShared { + heaprelid: (*heap_relation).rd_id, + indexrelid: (*index_relation).rd_id, + isconcurrent, + workersdonecv: std::mem::zeroed(), + mutex: std::mem::zeroed(), + nparticipantsdone: 0, + indtuples: 0, + }); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordgshared).workersdonecv); + pgrx::pg_sys::SpinLockInit(&raw mut (*vchordgshared).mutex); + vchordgshared + }; + + let tablescandesc = unsafe { + let tablescandesc = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, est_tablescandesc) + .cast::(); + pgrx::pg_sys::table_parallelscan_initialize(heap_relation, tablescandesc, snapshot); + tablescandesc + }; + + let vchordgcached = unsafe { + let x = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, 8 + cache.len()).cast::(); + (x as *mut u64).write_unaligned(cache.len() as _); + std::ptr::copy(cache.as_ptr(), x.add(8), cache.len()); + x + }; + + unsafe { + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000001, vchordgshared.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000003, vchordgcached.cast()); + } + + unsafe { + pgrx::pg_sys::LaunchParallelWorkers(pcxt); + } + + let nworkers_launched = unsafe { (*pcxt).nworkers_launched }; + + unsafe { + if nworkers_launched == 0 { + pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + Some(Self { + pcxt, + nparticipants: nworkers_launched + 1, + snapshot, + vchordgshared, + tablescandesc, + vchordgcached, + }) + } + + pub fn wait(&self) { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToAttach(self.pcxt); + } + } +} + +impl Drop for VamanaLeader { + fn drop(&mut self) { + if !std::thread::panicking() { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); + if is_mvcc_snapshot(self.snapshot) { + pgrx::pg_sys::UnregisterSnapshot(self.snapshot); + } + pgrx::pg_sys::DestroyParallelContext(self.pcxt); + pgrx::pg_sys::ExitParallelMode(); + } + } + } +} + +#[pgrx::pg_guard] +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vchordg_parallel_build_main( + _seg: *mut pgrx::pg_sys::dsm_segment, + toc: *mut pgrx::pg_sys::shm_toc, +) { + let vchordgshared = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() + }; + let tablescandesc = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) + .cast::() + }; + let vchordgcached = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000003, false) + .cast::() + .cast_const() + }; + let heap_lockmode; + let index_lockmode; + if unsafe { !(*vchordgshared).isconcurrent } { + heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; + index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; + } else { + heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; + index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; + } + let heap = unsafe { pgrx::pg_sys::table_open((*vchordgshared).heaprelid, heap_lockmode) }; + let index = unsafe { pgrx::pg_sys::index_open((*vchordgshared).indexrelid, index_lockmode) }; + let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; + unsafe { + (*index_info).ii_Concurrent = (*vchordgshared).isconcurrent; + } + + unsafe { + parallel_build( + index, + heap, + index_info, + tablescandesc, + vchordgshared, + vchordgcached, + |_| (), + ); + } + + unsafe { + pgrx::pg_sys::index_close(index, index_lockmode); + pgrx::pg_sys::table_close(heap, heap_lockmode); + } +} + +unsafe fn parallel_build( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + vchordgshared: *mut VamanaShared, + vchordgcached: *const u8, + mut callback: impl FnMut(u64), +) { + use vchordg_cached::VamanaCachedReader; + let cached = VamanaCachedReader::deserialize_ref(unsafe { + let bytes = (vchordgcached as *const u64).read_unaligned(); + std::slice::from_raw_parts(vchordgcached.add(8), bytes as _) + }); + + let index = unsafe { PostgresRelation::new(index_relation) }; + + let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap_relation, tablescandesc) }; + let opfamily = unsafe { opfamily(index_relation) }; + let heap = Heap { + heap_relation, + index_relation, + index_info, + opfamily, + scan, + }; + match cached { + VamanaCachedReader::_0(_) => { + heap.traverse(true, move |(ctid, store)| { + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + crate::index::vchordg::algo::insert(opfamily, &index, payload, vector); + } + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordgshared).mutex); + (*vchordgshared).indtuples += 1; + indtuples = (*vchordgshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordgshared).mutex); + } + callback(indtuples); + } + }); + } + } + unsafe { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordgshared).mutex); + (*vchordgshared).nparticipantsdone += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordgshared).mutex); + pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordgshared).workersdonecv); + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambuildempty(_index_relation: pgrx::pg_sys::Relation) { + pgrx::error!("Unlogged indexes are not supported."); +} + +unsafe fn options( + index_relation: pgrx::pg_sys::Relation, +) -> (VectorOptions, VamanaIndexingOptions) { + let att = unsafe { &mut *(*index_relation).rd_att }; + #[cfg(any( + feature = "pg13", + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17" + ))] + let atts = unsafe { att.attrs.as_slice(att.natts as _) }; + #[cfg(feature = "pg18")] + let atts = unsafe { + let ptr = att + .compact_attrs + .as_ptr() + .add(att.natts as _) + .cast::(); + std::slice::from_raw_parts(ptr, att.natts as _) + }; + if atts.is_empty() { + pgrx::error!("indexing on no columns is not supported"); + } + if atts.len() != 1 { + pgrx::error!("multicolumn index is not supported"); + } + // get dims + let typmod = Typmod::parse_from_i32(atts[0].atttypmod).unwrap(); + let dims = if let Some(dims) = typmod.dims() { + dims.get() + } else { + pgrx::error!( + "Dimensions type modifier of a vector column is needed for building the index." + ); + }; + // get v, d + let opfamily = unsafe { opfamily(index_relation) }; + let vector = VectorOptions { + dims, + v: opfamily.vector_kind(), + d: opfamily.distance_kind(), + }; + // get indexing, segment, optimizing + let rabitq = 'rabitq: { + let reloption = unsafe { (*index_relation).rd_options as *const Reloption }; + if reloption.is_null() || unsafe { (*reloption).options == 0 } { + break 'rabitq Default::default(); + } + let s = unsafe { Reloption::options(reloption) }.to_string_lossy(); + match toml::from_str::(&s) { + Ok(p) => p, + Err(e) => pgrx::error!("failed to parse options: {}", e), + } + }; + (vector, rabitq) +} + +mod vchordg_cached { + pub type Tag = u64; + + use algo::tuples::RefChecker; + use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + + #[repr(C, align(8))] + #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct VamanaCachedHeader0 {} + + #[repr(C, align(8))] + #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct VamanaCachedHeader1 { + mapping_s: usize, + mapping_e: usize, + pages_s: usize, + pages_e: usize, + } + + pub enum VamanaCached { + _0 {}, + } + + impl VamanaCached { + pub fn serialize(&self) -> Vec { + let mut buffer = Vec::new(); + match self { + VamanaCached::_0 {} => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + buffer[size_of::()..][..size_of::()] + .copy_from_slice(VamanaCachedHeader0 {}.as_bytes()); + } + } + buffer + } + } + + #[derive(Debug, Clone, Copy)] + pub enum VamanaCachedReader<'a> { + #[allow(dead_code)] + _0(VamanaCachedReader0<'a>), + } + + #[derive(Debug, Clone, Copy)] + pub struct VamanaCachedReader0<'a> { + #[allow(dead_code)] + header: &'a VamanaCachedHeader0, + } + + impl<'a> VamanaCachedReader<'a> { + pub fn deserialize_ref(source: &'a [u8]) -> Self { + let tag = u64::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &VamanaCachedHeader0 = checker.prefix(size_of::()); + Self::_0(VamanaCachedReader0 { header }) + } + _ => panic!("bad bytes"), + } + } + } +} diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs new file mode 100644 index 00000000..76bd4bff --- /dev/null +++ b/src/index/vchordg/am/mod.rs @@ -0,0 +1,438 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub mod am_build; + +use crate::index::fetcher::*; +use crate::index::gucs; +use crate::index::storage::PostgresRelation; +use crate::index::vchordg::opclass::opfamily; +use crate::index::vchordg::scanners::*; +use pgrx::datum::Internal; +use pgrx::pg_sys::Datum; +use std::cell::LazyCell; +use std::ffi::CStr; +use std::num::NonZero; +use std::ops::DerefMut; +use std::ptr::NonNull; +use std::sync::OnceLock; + +#[repr(C)] +struct Reloption { + vl_len_: i32, + pub options: i32, +} + +impl Reloption { + unsafe fn options<'a>(this: *const Self) -> &'a CStr { + unsafe { + let ptr = this + .cast::() + .add((&raw const (*this).options).read() as _); + CStr::from_ptr(ptr.cast()) + } + } +} + +const TABLE: &[pgrx::pg_sys::relopt_parse_elt] = &[pgrx::pg_sys::relopt_parse_elt { + optname: c"options".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, + offset: std::mem::offset_of!(Reloption, options) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, +}]; + +static RELOPT_KIND: OnceLock = OnceLock::new(); + +pub fn init() { + RELOPT_KIND.get_or_init(|| { + let kind; + unsafe { + kind = pgrx::pg_sys::add_reloption_kind(); + pgrx::pg_sys::add_string_reloption( + kind as _, + c"options".as_ptr(), + c"Vector index options, represented as a TOML string.".as_ptr(), + c"".as_ptr(), + None, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + } + kind + }); +} + +#[pgrx::pg_extern(sql = "")] +fn _vchordg_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { + type T = pgrx::pg_sys::IndexAmRoutine; + unsafe { + let index_am_routine = pgrx::pg_sys::palloc0(size_of::()) as *mut T; + index_am_routine.write(AM_HANDLER); + Internal::from(Some(Datum::from(index_am_routine))) + } +} + +const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { + let mut am_routine = unsafe { std::mem::zeroed::() }; + + am_routine.type_ = pgrx::pg_sys::NodeTag::T_IndexAmRoutine; + + am_routine.amsupport = 1; + am_routine.amcanorderbyop = true; + + #[cfg(any(feature = "pg17", feature = "pg18"))] + { + am_routine.amcanbuildparallel = true; + } + + // Index access methods that set `amoptionalkey` to `false` + // must index all tuples, even if the first column is `NULL`. + // However, PostgreSQL does not generate a path if there is no + // index clauses, even if there is a `ORDER BY` clause. + // So we have to set it to `true` and set costs of every path + // for vector index scans without `ORDER BY` clauses a large number + // and throw errors if someone really wants such a path. + am_routine.amoptionalkey = true; + + am_routine.amvalidate = Some(amvalidate); + am_routine.amoptions = Some(amoptions); + am_routine.amcostestimate = Some(amcostestimate); + + am_routine.ambuildphasename = Some(am_build::ambuildphasename); + am_routine.ambuild = Some(am_build::ambuild); + am_routine.ambuildempty = Some(am_build::ambuildempty); + am_routine.aminsert = Some(aminsert); + am_routine.ambulkdelete = Some(ambulkdelete); + am_routine.amvacuumcleanup = Some(amvacuumcleanup); + + am_routine.ambeginscan = Some(ambeginscan); + am_routine.amrescan = Some(amrescan); + am_routine.amgettuple = Some(amgettuple); + am_routine.amendscan = Some(amendscan); + + am_routine +}; + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { + true +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amoptions( + reloptions: Datum, + validate: bool, +) -> *mut pgrx::pg_sys::bytea { + let relopt_kind = RELOPT_KIND.get().copied().expect("init is not called"); + let rdopts = unsafe { + pgrx::pg_sys::build_reloptions( + reloptions, + validate, + relopt_kind, + size_of::(), + TABLE.as_ptr(), + TABLE.len() as _, + ) + }; + rdopts as *mut pgrx::pg_sys::bytea +} + +#[allow(clippy::too_many_arguments)] +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amcostestimate( + _root: *mut pgrx::pg_sys::PlannerInfo, + _path: *mut pgrx::pg_sys::IndexPath, + _loop_count: f64, + index_startup_cost: *mut pgrx::pg_sys::Cost, + index_total_cost: *mut pgrx::pg_sys::Cost, + index_selectivity: *mut pgrx::pg_sys::Selectivity, + index_correlation: *mut f64, + index_pages: *mut f64, +) { + unsafe { + *index_startup_cost = 0.0; + *index_total_cost = 0.0; + *index_selectivity = 1.0; + *index_correlation = 0.0; + *index_pages = 1.0; + } +} + +#[cfg(feature = "pg13")] +#[allow(clippy::too_many_arguments)] +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn aminsert( + index_relation: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, + heap_relation: pgrx::pg_sys::Relation, + _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, + _index_info: *mut pgrx::pg_sys::IndexInfo, +) -> bool { + unsafe { aminsertinner(index_relation, heap_relation, values, is_null, heap_tid) } +} + +#[cfg(any( + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17", + feature = "pg18" +))] +#[allow(clippy::too_many_arguments)] +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn aminsert( + index_relation: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, + heap_relation: pgrx::pg_sys::Relation, + _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, + _index_unchanged: bool, + _index_info: *mut pgrx::pg_sys::IndexInfo, +) -> bool { + unsafe { aminsertinner(index_relation, heap_relation, values, is_null, heap_tid) } +} + +unsafe fn aminsertinner( + index_relation: pgrx::pg_sys::Relation, + _heap_relation: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + ctid: pgrx::pg_sys::ItemPointer, +) -> bool { + let opfamily = unsafe { opfamily(index_relation) }; + let index = unsafe { PostgresRelation::new(index_relation) }; + let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; + let ctid = unsafe { ctid.read() }; + if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + crate::index::vchordg::algo::insert(opfamily, &index, payload, vector); + } + } + false +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambulkdelete( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, + callback: pgrx::pg_sys::IndexBulkDeleteCallback, + callback_state: *mut std::os::raw::c_void, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let opfamily = unsafe { opfamily((*info).index) }; + let index = unsafe { PostgresRelation::new((*info).index) }; + let check = || unsafe { + #[cfg(any( + feature = "pg13", + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17" + ))] + pgrx::pg_sys::vacuum_delay_point(); + #[cfg(feature = "pg18")] + pgrx::pg_sys::vacuum_delay_point(false); + }; + let callback = callback.expect("null function pointer"); + let callback = |pointer: NonZero| { + let (key, _) = pointer_to_kv(pointer); + let mut ctid = key_to_ctid(key); + unsafe { callback(&mut ctid, callback_state) } + }; + crate::index::vchordg::algo::bulkdelete(opfamily, &index, check, callback); + stats +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amvacuumcleanup( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let opfamily = unsafe { opfamily((*info).index) }; + let index = unsafe { PostgresRelation::new((*info).index) }; + let check = || unsafe { + #[cfg(any( + feature = "pg13", + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17" + ))] + pgrx::pg_sys::vacuum_delay_point(); + #[cfg(feature = "pg18")] + pgrx::pg_sys::vacuum_delay_point(false); + }; + crate::index::vchordg::algo::maintain(opfamily, &index, check); + stats +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambeginscan( + index_relation: pgrx::pg_sys::Relation, + n_keys: std::os::raw::c_int, + n_orderbys: std::os::raw::c_int, +) -> pgrx::pg_sys::IndexScanDesc { + use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; + + let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index_relation, n_keys, n_orderbys) }; + let scanner: Scanner = Scanner { + hack: None, + scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), + bump: Box::new(bumpalo::Bump::new()), + }; + unsafe { + (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); + } + scan +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amrescan( + scan: pgrx::pg_sys::IndexScanDesc, + keys: pgrx::pg_sys::ScanKey, + _n_keys: std::os::raw::c_int, + orderbys: pgrx::pg_sys::ScanKey, + _n_orderbys: std::os::raw::c_int, +) { + unsafe { + use crate::index::vchordg::opclass::Opfamily; + if !keys.is_null() && (*scan).numberOfKeys > 0 { + std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); + } + if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { + std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); + } + if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { + pgrx::error!( + "vector search with no WHERE clause and no ORDER BY clause is not supported" + ); + } + let scanner = &mut *(*scan).opaque.cast::(); + scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + scanner.bump.reset(); + let opfamily = opfamily((*scan).indexRelation); + let index = PostgresRelation::new((*scan).indexRelation); + let options = SearchOptions { + ef_search: gucs::vchordg_ef_search(), + beam_search: gucs::vchordg_beam_search(), + max_scan_tuples: gucs::vchordg_max_scan_tuples(), + io_search: gucs::vchordg_io_search(), + io_rerank: gucs::vchordg_io_rerank(), + }; + let fetcher = { + let hack = scanner.hack; + LazyCell::new(move || { + HeapFetcher::new( + (*scan).indexRelation, + (*scan).heapRelation, + (*scan).xs_snapshot, + if let Some(hack) = hack { + hack.as_ptr() + } else { + std::ptr::null_mut() + }, + ) + }) + }; + // PAY ATTENTATION: `scanning` references `bump`, so `scanning` must be dropped before `bump`. + let bump = scanner.bump.as_ref(); + scanner.scanning = match opfamily { + Opfamily::VectorL2 + | Opfamily::VectorCosine + | Opfamily::HalfvecL2 + | Opfamily::HalfvecCosine + | Opfamily::VectorIp + | Opfamily::HalfvecIp => { + let mut builder = DefaultBuilder::new(opfamily); + for i in 0..(*scan).numberOfOrderBys { + let data = (*scan).orderByData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + for i in 0..(*scan).numberOfKeys { + let data = (*scan).keyData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + LazyCell::new(Box::new(move || { + // only do this since `PostgresRelation` has no destructor + let index = bump.alloc(index.clone()); + builder.build(index, options, fetcher, bump) + })) + } + }; + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amgettuple( + scan: pgrx::pg_sys::IndexScanDesc, + direction: pgrx::pg_sys::ScanDirection::Type, +) -> bool { + if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { + pgrx::error!("vector search without a forward scan direction is not supported"); + } + // https://www.postgresql.org/docs/current/index-locking.html + // If heap entries referenced physical pointers are deleted before + // they are consumed by PostgreSQL, PostgreSQL will received wrong + // physical pointers: no rows or irreverent rows are referenced. + if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + { + pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); + } + let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; + if let Some((_, key, recheck)) = scanner.scanning.deref_mut().next() { + unsafe { + (*scan).xs_heaptid = key_to_ctid(key); + (*scan).xs_recheck = recheck; + (*scan).xs_recheckorderby = false; + } + true + } else { + false + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { + let scanner = unsafe { &mut *(*scan).opaque.cast::() }; + scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + scanner.bump.reset(); +} + +type Iter = Box>; + +pub struct Scanner { + pub hack: Option>, + scanning: LazyCell Iter>>, + bump: Box, +} diff --git a/src/index/vchordg/mod.rs b/src/index/vchordg/mod.rs new file mode 100644 index 00000000..d897f7d1 --- /dev/null +++ b/src/index/vchordg/mod.rs @@ -0,0 +1,19 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub mod algo; +pub mod am; +pub mod opclass; +pub mod scanners; +pub mod types; diff --git a/src/index/vchordg/opclass.rs b/src/index/vchordg/opclass.rs new file mode 100644 index 00000000..6a83dd3f --- /dev/null +++ b/src/index/vchordg/opclass.rs @@ -0,0 +1,168 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; +use crate::index::opclass::Sphere; +use distance::Distance; +use pgrx::datum::FromDatum; +use pgrx::heap_tuple::PgHeapTuple; +use pgrx::pg_sys::Datum; +use std::num::NonZero; +use vchordg::types::*; +use vector::VectorBorrowed; + +#[derive(Debug, Clone, Copy)] +pub enum Opfamily { + VectorL2, + VectorCosine, + VectorIp, + HalfvecL2, + HalfvecCosine, + HalfvecIp, +} + +impl Opfamily { + fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { + use {BorrowedVector as B, OwnedVector as O}; + match (vector, self) { + (B::Vecf32(x), Self::VectorL2) => O::Vecf32(x.own()), + (B::Vecf32(x), Self::VectorCosine) => O::Vecf32(x.function_normalize()), + (B::Vecf32(_), _) => unreachable!(), + (B::Vecf16(x), Self::HalfvecL2) => O::Vecf16(x.own()), + (B::Vecf16(x), Self::HalfvecCosine) => O::Vecf16(x.function_normalize()), + (B::Vecf16(_), _) => unreachable!(), + } + } + pub unsafe fn store(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let store = match self { + Self::VectorL2 | Self::VectorCosine | Self::VectorIp => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Vecf32(vector.as_borrowed())), 0)] + } + Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] + } + }; + Some(store) + } + pub unsafe fn input_sphere(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let attno_1 = NonZero::new(1_usize).unwrap(); + let attno_2 = NonZero::new(2_usize).unwrap(); + let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; + let center = match self { + Self::VectorL2 | Self::VectorCosine | Self::VectorIp => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + }; + let radius = tuple.get_by_index::(attno_2).unwrap()?; + Some(Sphere { center, radius }) + } + pub unsafe fn input_vector(self, datum: Datum) -> Option { + if datum.is_null() { + return None; + } + let vector = match self { + Self::VectorL2 | Self::VectorCosine | Self::VectorIp => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + }; + Some(vector) + } + pub fn output(self, x: Distance) -> f32 { + match self { + Self::VectorCosine | Self::HalfvecCosine => x.to_f32() * 0.5, + Self::VectorL2 | Self::HalfvecL2 => x.to_f32().sqrt(), + Self::VectorIp | Self::HalfvecIp => x.to_f32(), + } + } + pub const fn distance_kind(self) -> DistanceKind { + match self { + Self::VectorL2 | Self::HalfvecL2 => DistanceKind::L2S, + Self::VectorCosine | Self::HalfvecCosine => DistanceKind::L2S, + Self::VectorIp | Self::HalfvecIp => DistanceKind::Dot, + } + } + pub const fn vector_kind(self) -> VectorKind { + match self { + Self::VectorL2 | Self::VectorCosine | Self::VectorIp => VectorKind::Vecf32, + Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => VectorKind::Vecf16, + } + } +} + +pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { + use pgrx::pg_sys::Oid; + + let proc = unsafe { pgrx::pg_sys::index_getprocid(index_relation, 1, 1) }; + + if proc == Oid::INVALID { + pgrx::error!("support function 1 is not found"); + } + + let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); + + unsafe { + pgrx::pg_sys::fmgr_info(proc, &mut flinfo); + } + + let fn_addr = flinfo.fn_addr.expect("null function pointer"); + + let mut fcinfo = unsafe { std::mem::zeroed::() }; + fcinfo.flinfo = &mut flinfo; + fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; + fcinfo.context = std::ptr::null_mut(); + fcinfo.resultinfo = std::ptr::null_mut(); + fcinfo.isnull = true; + fcinfo.nargs = 0; + + let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; + + let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; + + let result_string = result_option.expect("null return value"); + + let result = match result_string.as_str() { + "vchordg_vector_l2_ops" => Opfamily::VectorL2, + "vchordg_vector_cosine_ops" => Opfamily::VectorCosine, + "vchordg_halfvec_l2_ops" => Opfamily::HalfvecL2, + "vchordg_halfvec_cosine_ops" => Opfamily::HalfvecCosine, + "vchordg_vector_ip_ops" => Opfamily::VectorIp, + "vchordg_halfvec_ip_ops" => Opfamily::HalfvecIp, + _ => pgrx::error!("unknown operator class"), + }; + + unsafe { + pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); + } + + result +} diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs new file mode 100644 index 00000000..f3d597f9 --- /dev/null +++ b/src/index/vchordg/scanners/default.rs @@ -0,0 +1,501 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use super::{Fetcher, Io, SearchBuilder, SearchOptions}; +use crate::index::fetcher::pointer_to_kv; +use crate::index::opclass::Sphere; +use crate::index::vchordg::algo::*; +use crate::index::vchordg::opclass::Opfamily; +use algo::accessor::{Dot, L2S}; +use algo::*; +use distance::Distance; +use half::f16; +use std::num::NonZero; +use vchordg::operator::{self}; +use vchordg::types::{DistanceKind, OwnedVector, VectorKind}; +use vchordg::*; +use vector::VectorOwned; +use vector::vect::VectOwned; + +pub struct DefaultBuilder { + opfamily: Opfamily, + orderbys: Vec>, + spheres: Vec>>, +} + +impl SearchBuilder for DefaultBuilder { + type Opaque = vchordg::Opaque; + + fn new(opfamily: Opfamily) -> Self { + assert!(matches!( + opfamily, + Opfamily::HalfvecCosine + | Opfamily::HalfvecL2 + | Opfamily::VectorCosine + | Opfamily::VectorL2 + | Opfamily::VectorIp + | Opfamily::HalfvecIp + )); + Self { + opfamily, + orderbys: Vec::new(), + spheres: Vec::new(), + } + } + + unsafe fn add(&mut self, strategy: u16, datum: Option) { + match strategy { + 1 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_vector(x)) }; + self.orderbys.push(x); + } + 2 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_sphere(x)) }; + self.spheres.push(x); + } + _ => unreachable!(), + } + } + + fn build<'a, R>( + self, + index: &'a R, + options: SearchOptions, + _fetcher: impl Fetcher + 'a, + bump: &'a impl Bump, + ) -> Box + 'a> + where + R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page, + { + let mut vector = None; + let mut threshold = None; + let mut recheck = false; + for orderby_vector in self.orderbys.into_iter().flatten() { + if vector.is_none() { + vector = Some(orderby_vector); + } else { + pgrx::error!("vector search with multiple vectors is not supported"); + } + } + for Sphere { center, radius } in self.spheres.into_iter().flatten() { + if vector.is_none() { + (vector, threshold) = (Some(center), Some(radius)); + } else { + recheck = true; + } + } + let opfamily = self.opfamily; + let Some(vector) = vector else { + return Box::new(std::iter::empty()) as Box>; + }; + let make_vertex_plain_prefetcher = MakePlainPrefetcher { index }; + let make_vertex_simple_prefetcher = MakeSimplePrefetcher { index }; + let make_vertex_stream_prefetcher = MakeStreamPrefetcher { + index, + hints: Hints::default().full(true), + }; + let make_vector_plain_prefetcher = MakePlainPrefetcher { index }; + let make_vector_simple_prefetcher = MakeSimplePrefetcher { index }; + let make_vector_stream_prefetcher = MakeStreamPrefetcher { + index, + hints: Hints::default().full(true), + }; + let iter: Box)>> = + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + type Op = operator::Op, L2S>; + let unprojected = if let OwnedVector::Vecf32(vector) = vector { + bump.alloc(vector) + } else { + unreachable!() + }; + let projected = bump.alloc(RandomProject::project(unprojected.as_borrowed())); + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + type Op = operator::Op, L2S>; + let unprojected = if let OwnedVector::Vecf16(vector) = vector { + bump.alloc(vector) + } else { + unreachable!() + }; + let projected = bump.alloc(RandomProject::project(unprojected.as_borrowed())); + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + type Op = operator::Op, Dot>; + let unprojected = if let OwnedVector::Vecf32(vector) = vector { + bump.alloc(vector) + } else { + unreachable!() + }; + let projected = bump.alloc(RandomProject::project(unprojected.as_borrowed())); + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + type Op = operator::Op, Dot>; + let unprojected = if let OwnedVector::Vecf16(vector) = vector { + bump.alloc(vector) + } else { + unreachable!() + }; + let projected = bump.alloc(RandomProject::project(unprojected.as_borrowed())); + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + projected.as_borrowed(), + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + }; + let iter = if let Some(threshold) = threshold { + Box::new(iter.take_while(move |(distance, _)| distance.to_f32() < threshold)) + } else { + iter + }; + let iter = if let Some(max_scan_tuples) = options.max_scan_tuples { + Box::new(iter.take(max_scan_tuples as _)) + } else { + iter + }; + Box::new(iter.map(move |(distance, pointer)| { + let (key, _) = pointer_to_kv(pointer); + (opfamily.output(distance), key, recheck) + })) + } +} diff --git a/src/index/vchordg/scanners/mod.rs b/src/index/vchordg/scanners/mod.rs new file mode 100644 index 00000000..264c29f2 --- /dev/null +++ b/src/index/vchordg/scanners/mod.rs @@ -0,0 +1,61 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +mod default; + +use super::opclass::Opfamily; +use crate::index::fetcher::Fetcher; +use algo::{Bump, Page, RelationPrefetch, RelationRead, RelationReadStream}; +use pgrx::pg_sys::Datum; + +pub use default::DefaultBuilder; + +#[derive(Debug, Clone, Copy)] +pub enum Io { + Plain, + Simple, + #[cfg_attr( + any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"), + expect(dead_code) + )] + Stream, +} + +#[derive(Debug)] +pub struct SearchOptions { + pub ef_search: u32, + pub beam_search: u32, + pub max_scan_tuples: Option, + pub io_search: Io, + pub io_rerank: Io, +} + +pub trait SearchBuilder: 'static { + type Opaque: Copy; + + fn new(opfamily: Opfamily) -> Self; + + unsafe fn add(&mut self, strategy: u16, datum: Option); + + fn build<'a, R>( + self, + relation: &'a R, + options: SearchOptions, + fetcher: impl Fetcher + 'a, + bump: &'a impl Bump, + ) -> Box + 'a> + where + R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page; +} diff --git a/src/index/vchordg/types.rs b/src/index/vchordg/types.rs new file mode 100644 index 00000000..c1511d0f --- /dev/null +++ b/src/index/vchordg/types.rs @@ -0,0 +1,24 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use serde::{Deserialize, Serialize}; +use validator::Validate; +use vchordg::types::VamanaIndexOptions; + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VamanaIndexingOptions { + #[serde(flatten)] + pub index: VamanaIndexOptions, +} diff --git a/src/index/vchordrq/algo.rs b/src/index/vchordrq/algo.rs new file mode 100644 index 00000000..82d87f59 --- /dev/null +++ b/src/index/vchordrq/algo.rs @@ -0,0 +1,421 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::index::vchordrq::am::am_build::InternalBuild; +use crate::index::vchordrq::opclass::Opfamily; +use algo::accessor::{Dot, L2S}; +use algo::prefetcher::*; +use algo::*; +use half::f16; +use std::collections::BinaryHeap; +use std::num::NonZero; +use vchordrq::operator::Op; +use vchordrq::types::*; +use vchordrq::{FastHeap, Opaque}; +use vector::VectorOwned; +use vector::vect::{VectBorrowed, VectOwned}; + +pub fn prewarm(opfamily: Opfamily, index: &R, height: i32) -> String +where + R: RelationRead, + R::Page: Page, +{ + let bump = bumpalo::Bump::new(); + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => vchordrq::prewarm::<_, Op, L2S>>( + index, + height, + &bump, + make_h0_plain_prefetcher, + ), + (VectorKind::Vecf32, DistanceKind::Dot) => vchordrq::prewarm::<_, Op, Dot>>( + index, + height, + &bump, + make_h0_plain_prefetcher, + ), + (VectorKind::Vecf16, DistanceKind::L2S) => vchordrq::prewarm::<_, Op, L2S>>( + index, + height, + &bump, + make_h0_plain_prefetcher, + ), + (VectorKind::Vecf16, DistanceKind::Dot) => vchordrq::prewarm::<_, Op, Dot>>( + index, + height, + &bump, + make_h0_plain_prefetcher, + ), + } +} + +pub fn bulkdelete( + opfamily: Opfamily, + index: &R, + check: impl Fn(), + callback: impl Fn(NonZero) -> bool, +) where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + } + } +} + +pub fn maintain(opfamily: Opfamily, index: &R, check: impl Fn()) +where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; + let maintain = match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordrq::maintain::<_, Op, L2S>>(index, make_h0_plain_prefetcher, check) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordrq::maintain::<_, Op, Dot>>(index, make_h0_plain_prefetcher, check) + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordrq::maintain::<_, Op, L2S>>(index, make_h0_plain_prefetcher, check) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordrq::maintain::<_, Op, Dot>>(index, make_h0_plain_prefetcher, check) + } + }; + pgrx::info!( + "maintain: number_of_formerly_allocated_pages = {}", + maintain.number_of_formerly_allocated_pages + ); + pgrx::info!( + "maintain: number_of_freshly_allocated_pages = {}", + maintain.number_of_freshly_allocated_pages + ); + pgrx::info!( + "maintain: number_of_freed_pages = {}", + maintain.number_of_freed_pages + ); +} + +pub fn build( + vector_options: VectorOptions, + vchordrq_options: VchordrqIndexOptions, + index: &R, + structures: Vec>>, +) where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + match (vector_options.v, vector_options.d) { + (VectorKind::Vecf32, DistanceKind::L2S) => vchordrq::build::<_, Op, L2S>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + (VectorKind::Vecf32, DistanceKind::Dot) => vchordrq::build::<_, Op, Dot>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + (VectorKind::Vecf16, DistanceKind::L2S) => vchordrq::build::<_, Op, L2S>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + (VectorKind::Vecf16, DistanceKind::Dot) => vchordrq::build::<_, Op, Dot>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + ), + } +} + +pub fn insert( + opfamily: Opfamily, + index: &R, + payload: NonZero, + vector: OwnedVector, + skip_freespaces: bool, +) where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + let bump = bumpalo::Bump::new(); + let make_h1_plain_prefetcher = MakeH1PlainPrefetcherForInsertion { index }; + match (vector, opfamily.distance_kind()) { + (OwnedVector::Vecf32(vector), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + let projected = RandomProject::project(vector.as_borrowed()); + let key = vchordrq::insert_vector::<_, Op, L2S>>( + index, + payload, + vector.as_borrowed(), + ); + vchordrq::insert::<_, Op, L2S>>( + index, + payload, + projected.as_borrowed(), + key, + &bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Vecf32(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + let projected = RandomProject::project(vector.as_borrowed()); + let key = vchordrq::insert_vector::<_, Op, Dot>>( + index, + payload, + vector.as_borrowed(), + ); + vchordrq::insert::<_, Op, Dot>>( + index, + payload, + projected.as_borrowed(), + key, + &bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Vecf16(vector), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + let projected = RandomProject::project(vector.as_borrowed()); + let key = vchordrq::insert_vector::<_, Op, L2S>>( + index, + payload, + vector.as_borrowed(), + ); + vchordrq::insert::<_, Op, L2S>>( + index, + payload, + projected.as_borrowed(), + key, + &bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Vecf16(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + let projected = RandomProject::project(vector.as_borrowed()); + let key = vchordrq::insert_vector::<_, Op, Dot>>( + index, + payload, + vector.as_borrowed(), + ); + vchordrq::insert::<_, Op, Dot>>( + index, + payload, + projected.as_borrowed(), + key, + &bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + } +} + +fn map_structures(x: Vec>, f: impl Fn(T) -> U + Copy) -> Vec> { + x.into_iter() + .map( + |Structure { + centroids, + children, + }| Structure { + centroids: centroids.into_iter().map(f).collect(), + children, + }, + ) + .collect() +} + +pub trait RandomProject { + type Output; + fn project(self) -> Self::Output; +} + +impl RandomProject for VectBorrowed<'_, f32> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use rabitq::rotate::rotate; + let input = self.slice(); + VectOwned::new(rotate(input)) + } +} + +impl RandomProject for VectBorrowed<'_, f16> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use rabitq::rotate::rotate; + use simd::Floating; + let input = f16::vector_to_f32(self.slice()); + VectOwned::new(f16::vector_from_f32(&rotate(&input))) + } +} + +#[derive(Debug)] +pub struct MakeH1PlainPrefetcherForInsertion<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakeH1PlainPrefetcherForInsertion<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcherForInsertion<'r, R> { + type P + = PlainPrefetcher<'r, R, FastHeap> + where + T: Ord + Fetch + 'r; + + fn prefetch(&mut self, seq: Vec) -> Self::P + where + T: Ord + Fetch + 'r, + { + PlainPrefetcher::new(self.index, FastHeap::from(seq)) + } +} + +#[derive(Debug)] +pub struct MakeH1PlainPrefetcher<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakeH1PlainPrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcher<'r, R> { + type P + = PlainPrefetcher<'r, R, BinaryHeap> + where + T: Ord + Fetch + 'r; + + fn prefetch(&mut self, seq: Vec) -> Self::P + where + T: Ord + Fetch + 'r, + { + PlainPrefetcher::new(self.index, BinaryHeap::from(seq)) + } +} + +#[derive(Debug)] +pub struct MakeH0PlainPrefetcher<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakeH0PlainPrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakeH0PlainPrefetcher<'r, R> { + type P + = PlainPrefetcher<'r, R, S> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch, + { + PlainPrefetcher::new(self.index, seq) + } +} + +#[derive(Debug)] +pub struct MakeH0SimplePrefetcher<'r, R> { + pub index: &'r R, +} + +impl<'r, R> Clone for MakeH0SimplePrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> + for MakeH0SimplePrefetcher<'r, R> +{ + type P + = SimplePrefetcher<'r, R, S> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch, + { + SimplePrefetcher::new(self.index, seq) + } +} + +#[derive(Debug)] +pub struct MakeH0StreamPrefetcher<'r, R> { + pub index: &'r R, + pub hints: Hints, +} + +impl<'r, R> Clone for MakeH0StreamPrefetcher<'r, R> { + fn clone(&self) -> Self { + Self { + index: self.index, + hints: self.hints.clone(), + } + } +} + +impl<'r, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'r, R> + for MakeH0StreamPrefetcher<'r, R> +{ + type P + = StreamPrefetcher<'r, R, S> + where + S::Item: Fetch; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch, + { + StreamPrefetcher::new(self.index, seq, self.hints.clone()) + } +} diff --git a/src/index/am/am_build.rs b/src/index/vchordrq/am/am_build.rs similarity index 93% rename from src/index/am/am_build.rs rename to src/index/vchordrq/am/am_build.rs index 7ac85117..38afa216 100644 --- a/src/index/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -13,18 +13,21 @@ // Copyright (c) 2025 TensorChord Inc. use crate::datatype::typmod::Typmod; -use crate::index::am::{Reloption, ctid_to_key, kv_to_pointer}; -use crate::index::opclass::{Opfamily, opfamily}; +use crate::index::fetcher::*; use crate::index::storage::{PostgresPage, PostgresRelation}; -use crate::index::types::*; -use algorithm::types::*; -use algorithm::{PageGuard, Relation, RelationRead, RelationWrite}; +use crate::index::vchordrq::am::Reloption; +use crate::index::vchordrq::opclass::{Opfamily, opfamily}; +use crate::index::vchordrq::types::*; +use algo::{ + Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, +}; use half::f16; use pgrx::pg_sys::{Datum, ItemPointerData}; use rand::Rng; use simd::Floating; use std::ffi::CStr; use std::ops::Deref; +use vchordrq::types::*; use vector::VectorOwned; use vector::vect::VectOwned; @@ -323,15 +326,15 @@ pub unsafe extern "C-unwind" fn ambuild( } }; reporter.phase(BuildPhase::from_code(BuildPhaseCode::Build)); - crate::index::algorithm::build(vector_options, vchordrq_options.index, &index, structures); + crate::index::vchordrq::algo::build(vector_options, vchordrq_options.index, &index, structures); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); let cache = if vchordrq_options.build.pin { - let mut trace = algorithm::cache(&index); + let mut trace = vchordrq::cache(&index); trace.sort(); trace.dedup(); if let Some(max) = trace.last().copied() { let mut mapping = vec![u32::MAX; 1 + max as usize]; - let mut pages = Vec::>::with_capacity(trace.len()); + let mut pages = Vec::>>::with_capacity(trace.len()); for id in trace { mapping[id as usize] = pages.len() as u32; pages.push(index.read(id).clone_into_boxed()); @@ -388,7 +391,7 @@ pub unsafe extern "C-unwind" fn ambuild( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, &index, payload, vector, true); + crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); } indtuples += 1; reporter.tuples_done(indtuples); @@ -399,7 +402,7 @@ pub unsafe extern "C-unwind" fn ambuild( pgrx::check_for_interrupts!(); }; reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); - crate::index::algorithm::maintain(opfamily, &index, check); + crate::index::vchordrq::algo::maintain(opfamily, &index, check); unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -425,6 +428,7 @@ mod vchordrq_cached { pub type Tag = u64; use crate::index::storage::PostgresPage; + use algo::tuples::RefChecker; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] @@ -444,7 +448,7 @@ mod vchordrq_cached { _0 {}, _1 { mapping: Vec, - pages: Vec>, + pages: Vec>>, }, } @@ -469,7 +473,9 @@ mod vchordrq_cached { } let pages_s = buffer.len(); buffer.extend(pages.iter().flat_map(|x| unsafe { - std::mem::transmute::<&PostgresPage, &[u8; 8192]>(x.as_ref()) + std::mem::transmute::<&PostgresPage, &[u8; 8192]>( + x.as_ref(), + ) })); let pages_e = buffer.len(); while buffer.len() % ALIGN != 0 { @@ -509,11 +515,11 @@ mod vchordrq_cached { #[allow(dead_code)] header: &'a VchordrqCachedHeader1, mapping: &'a [u32], - pages: &'a [PostgresPage], + pages: &'a [PostgresPage], } impl<'a> VchordrqCachedReader1<'a> { - pub fn get(&self, id: u32) -> Option<&'a PostgresPage> { + pub fn get(&self, id: u32) -> Option<&'a PostgresPage> { let index = *self.mapping.get(id as usize)?; if index == u32::MAX { return None; @@ -547,50 +553,6 @@ mod vchordrq_cached { } } } - - pub struct RefChecker<'a> { - bytes: &'a [u8], - } - - impl<'a> RefChecker<'a> { - pub fn new(bytes: &'a [u8]) -> Self { - Self { bytes } - } - pub fn prefix( - &self, - s: usize, - ) -> &'a T { - let start = s; - let end = s + size_of::(); - let bytes = &self.bytes[start..end]; - FromBytes::ref_from_bytes(bytes).expect("bad bytes") - } - pub fn bytes( - &self, - s: usize, - e: usize, - ) -> &'a T { - let start = s; - let end = e; - let bytes = &self.bytes[start..end]; - FromBytes::ref_from_bytes(bytes).expect("bad bytes") - } - pub unsafe fn bytes_slice_unchecked(&self, s: usize, e: usize) -> &'a [T] { - let start = s; - let end = e; - let bytes = &self.bytes[start..end]; - if size_of::() == 0 || bytes.len() % size_of::() == 0 { - let ptr = bytes as *const [u8] as *const T; - if ptr.is_aligned() { - unsafe { std::slice::from_raw_parts(ptr, bytes.len() / size_of::()) } - } else { - panic!("bad bytes") - } - } else { - panic!("bad bytes") - } - } - } } fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { @@ -870,7 +832,7 @@ unsafe fn parallel_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, &index, payload, vector, true); + crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); } unsafe { let indtuples; @@ -893,7 +855,7 @@ unsafe fn parallel_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, &index, payload, vector, true); + crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); } unsafe { let indtuples; @@ -925,7 +887,23 @@ unsafe fn options( index_relation: pgrx::pg_sys::Relation, ) -> (VectorOptions, VchordrqIndexingOptions) { let att = unsafe { &mut *(*index_relation).rd_att }; + #[cfg(any( + feature = "pg13", + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17" + ))] let atts = unsafe { att.attrs.as_slice(att.natts as _) }; + #[cfg(feature = "pg18")] + let atts = unsafe { + let ptr = att + .compact_attrs + .as_ptr() + .add(att.natts as _) + .cast::(); + std::slice::from_raw_parts(ptr, att.natts as _) + }; if atts.is_empty() { pgrx::error!("indexing on no columns is not supported"); } @@ -933,7 +911,7 @@ unsafe fn options( pgrx::error!("multicolumn index is not supported"); } // get dims - let typmod = Typmod::parse_from_i32(atts[0].type_mod()).unwrap(); + let typmod = Typmod::parse_from_i32(atts[0].atttypmod).unwrap(); let dims = if let Some(dims) = typmod.dims() { dims.get() } else { @@ -1288,12 +1266,15 @@ impl Relation for CachingRelation<'_, R> { type Page = R::Page; } -impl> RelationRead for CachingRelation<'_, R> { - type ReadGuard<'a> - = CachingRelationReadGuard<'a, R::ReadGuard<'a>> - where - Self: 'a; +impl>> RelationReadTypes + for CachingRelation<'_, R> +{ + type ReadGuard<'a> = CachingRelationReadGuard<'a, R::ReadGuard<'a>>; +} +impl>> RelationRead + for CachingRelation<'_, R> +{ fn read(&self, id: u32) -> Self::ReadGuard<'_> { if let Some(x) = self.cache.get(id) { CachingRelationReadGuard::Cached(id, x) @@ -1303,18 +1284,25 @@ impl> RelationRead for CachingRelation<'_, } } -impl> RelationWrite for CachingRelation<'_, R> { - type WriteGuard<'a> - = R::WriteGuard<'a> - where - Self: 'a; +impl>> RelationWriteTypes + for CachingRelation<'_, R> +{ + type WriteGuard<'a> = R::WriteGuard<'a>; +} +impl>> RelationWrite + for CachingRelation<'_, R> +{ fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { self.relation.write(id, tracking_freespace) } - fn extend(&self, tracking_freespace: bool) -> Self::WriteGuard<'_> { - self.relation.extend(tracking_freespace) + fn extend( + &self, + opaque: ::Opaque, + tracking_freespace: bool, + ) -> Self::WriteGuard<'_> { + self.relation.extend(opaque, tracking_freespace) } fn search(&self, freespace: usize) -> Option> { diff --git a/src/index/am/mod.rs b/src/index/vchordrq/am/mod.rs similarity index 75% rename from src/index/am/mod.rs rename to src/index/vchordrq/am/mod.rs index 80c8da9a..4ca9dbff 100644 --- a/src/index/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -14,14 +14,13 @@ pub mod am_build; -use super::algorithm::BumpAlloc; +use crate::index::fetcher::*; use crate::index::gucs; -use crate::index::opclass::{Opfamily, opfamily}; -use crate::index::scanners::*; use crate::index::storage::PostgresRelation; -use algorithm::Bump; +use crate::index::vchordrq::opclass::{Opfamily, opfamily}; +use crate::index::vchordrq::scanners::*; use pgrx::datum::Internal; -use pgrx::pg_sys::{BlockIdData, Datum, ItemPointerData}; +use pgrx::pg_sys::Datum; use std::cell::LazyCell; use std::ffi::CStr; use std::num::NonZero; @@ -50,6 +49,8 @@ const TABLE: &[pgrx::pg_sys::relopt_parse_elt] = &[pgrx::pg_sys::relopt_parse_el optname: c"options".as_ptr(), opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, offset: std::mem::offset_of!(Reloption, options) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, }]; static RELOPT_KIND: OnceLock = OnceLock::new(); @@ -90,7 +91,7 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { am_routine.amsupport = 1; am_routine.amcanorderbyop = true; - #[cfg(feature = "pg17")] + #[cfg(any(feature = "pg17", feature = "pg18"))] { am_routine.amcanbuildparallel = true; } @@ -206,9 +207,9 @@ pub unsafe extern "C-unwind" fn amcostestimate( *index_pages = 1.0; return; } - let index = PostgresRelation::new(relation.raw()); - let probes = gucs::probes(); - let cost = algorithm::cost(&index); + let index = PostgresRelation::::new(relation.raw()); + let probes = gucs::vchordrq_probes(); + let cost = vchordrq::cost(&index); if cost.cells.len() != 1 + probes.len() { panic!( "need {} probes, but {} probes provided", @@ -271,7 +272,13 @@ pub unsafe extern "C-unwind" fn aminsert( unsafe { aminsertinner(index_relation, values, is_null, heap_tid) } } -#[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] +#[cfg(any( + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17", + feature = "pg18" +))] #[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn aminsert( @@ -301,7 +308,7 @@ unsafe fn aminsertinner( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::algorithm::insert(opfamily, &index, payload, vector, false); + crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, false); } } false @@ -323,7 +330,16 @@ pub unsafe extern "C-unwind" fn ambulkdelete( let opfamily = unsafe { opfamily((*info).index) }; let index = unsafe { PostgresRelation::new((*info).index) }; let check = || unsafe { + #[cfg(any( + feature = "pg13", + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17" + ))] pgrx::pg_sys::vacuum_delay_point(); + #[cfg(feature = "pg18")] + pgrx::pg_sys::vacuum_delay_point(false); }; let callback = callback.expect("null function pointer"); let callback = |pointer: NonZero| { @@ -331,7 +347,7 @@ pub unsafe extern "C-unwind" fn ambulkdelete( let mut ctid = key_to_ctid(key); unsafe { callback(&mut ctid, callback_state) } }; - crate::index::algorithm::bulkdelete(opfamily, &index, check, callback); + crate::index::vchordrq::algo::bulkdelete(opfamily, &index, check, callback); stats } @@ -349,9 +365,18 @@ pub unsafe extern "C-unwind" fn amvacuumcleanup( let opfamily = unsafe { opfamily((*info).index) }; let index = unsafe { PostgresRelation::new((*info).index) }; let check = || unsafe { + #[cfg(any( + feature = "pg13", + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17" + ))] pgrx::pg_sys::vacuum_delay_point(); + #[cfg(feature = "pg18")] + pgrx::pg_sys::vacuum_delay_point(false); }; - crate::index::algorithm::maintain(opfamily, &index, check); + crate::index::vchordrq::algo::maintain(opfamily, &index, check); stats } @@ -367,7 +392,7 @@ pub unsafe extern "C-unwind" fn ambeginscan( let scanner: Scanner = Scanner { hack: None, scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), - bump: Box::new(BumpAlloc::new()), + bump: Box::new(bumpalo::Bump::new()), }; unsafe { (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); @@ -384,7 +409,7 @@ pub unsafe extern "C-unwind" fn amrescan( _n_orderbys: std::os::raw::c_int, ) { unsafe { - use crate::index::opclass::Opfamily; + use crate::index::vchordrq::opclass::Opfamily; if !keys.is_null() && (*scan).numberOfKeys > 0 { std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); } @@ -402,14 +427,14 @@ pub unsafe extern "C-unwind" fn amrescan( let opfamily = opfamily((*scan).indexRelation); let index = PostgresRelation::new((*scan).indexRelation); let options = SearchOptions { - epsilon: gucs::epsilon(), - probes: gucs::probes(), - max_scan_tuples: gucs::max_scan_tuples(), - maxsim_refine: gucs::maxsim_refine(), - maxsim_threshold: gucs::maxsim_threshold(), - io_search: gucs::io_search(), - io_rerank: gucs::io_rerank(), - prefilter: gucs::prefilter(), + epsilon: gucs::vchordrq_epsilon(), + probes: gucs::vchordrq_probes(), + max_scan_tuples: gucs::vchordrq_max_scan_tuples(), + maxsim_refine: gucs::vchordrq_maxsim_refine(), + maxsim_threshold: gucs::vchordrq_maxsim_threshold(), + io_search: gucs::vchordrq_io_search(), + io_rerank: gucs::vchordrq_io_rerank(), + prefilter: gucs::vchordrq_prefilter(), }; let fetcher = { let hack = scanner.hack; @@ -519,126 +544,7 @@ type Iter = Box>; pub struct Scanner { pub hack: Option>, scanning: LazyCell Iter>>, - bump: Box, -} - -struct HeapFetcher { - index_info: *mut pgrx::pg_sys::IndexInfo, - estate: *mut pgrx::pg_sys::EState, - econtext: *mut pgrx::pg_sys::ExprContext, - heap_relation: pgrx::pg_sys::Relation, - snapshot: pgrx::pg_sys::Snapshot, - slot: *mut pgrx::pg_sys::TupleTableSlot, - values: [Datum; 32], - is_nulls: [bool; 32], - hack: *mut pgrx::pg_sys::IndexScanState, -} - -impl HeapFetcher { - unsafe fn new( - index_relation: pgrx::pg_sys::Relation, - heap_relation: pgrx::pg_sys::Relation, - snapshot: pgrx::pg_sys::Snapshot, - hack: *mut pgrx::pg_sys::IndexScanState, - ) -> Self { - unsafe { - let index_info = pgrx::pg_sys::BuildIndexInfo(index_relation); - let estate = pgrx::pg_sys::CreateExecutorState(); - let econtext = pgrx::pg_sys::MakePerTupleExprContext(estate); - Self { - index_info, - estate, - econtext, - heap_relation, - snapshot, - slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), - values: [Datum::null(); 32], - is_nulls: [true; 32], - hack, - } - } - } -} - -impl Drop for HeapFetcher { - fn drop(&mut self) { - unsafe { - pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); - // free common resources - pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); - pgrx::pg_sys::FreeExecutorState(self.estate); - } - } -} - -impl Fetcher for HeapFetcher { - type Tuple<'a> = HeapTuple<'a>; - - fn fetch(&mut self, key: [u16; 3]) -> Option> { - unsafe { - let mut ctid = key_to_ctid(key); - let table_am = (*self.heap_relation).rd_tableam; - let fetch_row_version = (*table_am) - .tuple_fetch_row_version - .expect("unsupported heap access method"); - if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { - return None; - } - Some(HeapTuple { this: self }) - } - } -} - -pub struct HeapTuple<'a> { - this: &'a mut HeapFetcher, -} - -impl Tuple for HeapTuple<'_> { - fn build(&mut self) -> (&[Datum; 32], &[bool; 32]) { - unsafe { - let this = &mut self.this; - (*this.econtext).ecxt_scantuple = this.slot; - pgrx::pg_sys::MemoryContextReset((*this.econtext).ecxt_per_tuple_memory); - pgrx::pg_sys::FormIndexDatum( - this.index_info, - this.slot, - this.estate, - this.values.as_mut_ptr(), - this.is_nulls.as_mut_ptr(), - ); - (&this.values, &this.is_nulls) - } - } - - #[allow(clippy::collapsible_if)] - fn filter(&mut self) -> bool { - unsafe { - let this = &mut self.this; - if !this.hack.is_null() { - if let Some(qual) = NonNull::new((*this.hack).ss.ps.qual) { - use pgrx::datum::FromDatum; - use pgrx::memcxt::PgMemoryContexts; - assert!(qual.as_ref().flags & pgrx::pg_sys::EEO_FLAG_IS_QUAL as u8 != 0); - let evalfunc = qual.as_ref().evalfunc.expect("no evalfunc for qual"); - if !(*this.hack).ss.ps.ps_ExprContext.is_null() { - let econtext = (*this.hack).ss.ps.ps_ExprContext; - (*econtext).ecxt_scantuple = this.slot; - pgrx::pg_sys::MemoryContextReset((*econtext).ecxt_per_tuple_memory); - let result = PgMemoryContexts::For((*econtext).ecxt_per_tuple_memory) - .switch_to(|_| { - let mut is_null = true; - let datum = evalfunc(qual.as_ptr(), econtext, &mut is_null); - bool::from_datum(datum, is_null) - }); - if result != Some(true) { - return false; - } - } - } - } - true - } - } + bump: Box, } struct Index { @@ -665,33 +571,3 @@ impl Drop for Index { } } } - -pub const fn ctid_to_key( - ItemPointerData { - ip_blkid: BlockIdData { bi_hi, bi_lo }, - ip_posid, - }: ItemPointerData, -) -> [u16; 3] { - [bi_hi, bi_lo, ip_posid] -} - -pub const fn key_to_ctid([bi_hi, bi_lo, ip_posid]: [u16; 3]) -> ItemPointerData { - ItemPointerData { - ip_blkid: BlockIdData { bi_hi, bi_lo }, - ip_posid, - } -} - -pub const fn pointer_to_kv(pointer: NonZero) -> ([u16; 3], u16) { - let value = pointer.get(); - let bi_hi = ((value >> 48) & 0xffff) as u16; - let bi_lo = ((value >> 32) & 0xffff) as u16; - let ip_posid = ((value >> 16) & 0xffff) as u16; - let extra = value as u16; - ([bi_hi, bi_lo, ip_posid], extra) -} - -pub const fn kv_to_pointer((key, value): ([u16; 3], u16)) -> NonZero { - let x = (key[0] as u64) << 48 | (key[1] as u64) << 32 | (key[2] as u64) << 16 | value as u64; - NonZero::new(x).expect("invalid key") -} diff --git a/src/index/vchordrq/mod.rs b/src/index/vchordrq/mod.rs new file mode 100644 index 00000000..d897f7d1 --- /dev/null +++ b/src/index/vchordrq/mod.rs @@ -0,0 +1,19 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub mod algo; +pub mod am; +pub mod opclass; +pub mod scanners; +pub mod types; diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs new file mode 100644 index 00000000..1fa0f2e2 --- /dev/null +++ b/src/index/vchordrq/opclass.rs @@ -0,0 +1,234 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; +use crate::index::opclass::Sphere; +use distance::Distance; +use pgrx::datum::FromDatum; +use pgrx::heap_tuple::PgHeapTuple; +use pgrx::pg_sys::Datum; +use std::num::NonZero; +use vchordrq::types::*; +use vector::VectorBorrowed; + +#[derive(Debug, Clone, Copy)] +pub enum Opfamily { + VectorL2, + VectorIp, + VectorCosine, + HalfvecL2, + HalfvecIp, + HalfvecCosine, + VectorMaxsim, + HalfvecMaxsim, +} + +impl Opfamily { + fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { + use {BorrowedVector as B, OwnedVector as O}; + match (vector, self) { + (B::Vecf32(x), Self::VectorL2) => O::Vecf32(x.own()), + (B::Vecf32(x), Self::VectorIp | Self::VectorMaxsim) => O::Vecf32(x.own()), + (B::Vecf32(x), Self::VectorCosine) => O::Vecf32(x.function_normalize()), + (B::Vecf32(_), _) => unreachable!(), + (B::Vecf16(x), Self::HalfvecL2) => O::Vecf16(x.own()), + (B::Vecf16(x), Self::HalfvecIp | Self::HalfvecMaxsim) => O::Vecf16(x.own()), + (B::Vecf16(x), Self::HalfvecCosine) => O::Vecf16(x.function_normalize()), + (B::Vecf16(_), _) => unreachable!(), + } + } + pub unsafe fn store(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let store = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Vecf32(vector.as_borrowed())), 0)] + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] + } + Self::VectorMaxsim => { + let vectors = + unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Vecf32(vector.as_borrowed())), + i as u16, + )); + } + result + } + Self::HalfvecMaxsim => { + let vectors = + unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Vecf16(vector.as_borrowed())), + i as u16, + )); + } + result + } + }; + Some(store) + } + pub unsafe fn input_sphere(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let attno_1 = NonZero::new(1_usize).unwrap(); + let attno_2 = NonZero::new(2_usize).unwrap(); + let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; + let center = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + }; + let radius = tuple.get_by_index::(attno_2).unwrap()?; + Some(Sphere { center, radius }) + } + pub unsafe fn input_vector(self, datum: Datum) -> Option { + if datum.is_null() { + return None; + } + let vector = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + }; + Some(vector) + } + pub unsafe fn input_vectors(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let vectors = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { + let vectors = + unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Vecf32(vector.as_borrowed()))); + } + result + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { + let vectors = + unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Vecf16(vector.as_borrowed()))); + } + result + } + }; + Some(vectors) + } + pub fn output(self, x: Distance) -> f32 { + match self { + Self::VectorCosine | Self::HalfvecCosine => x.to_f32() + 1.0f32, + Self::VectorL2 | Self::HalfvecL2 => x.to_f32().sqrt(), + Self::VectorIp | Self::HalfvecIp | Self::VectorMaxsim | Self::HalfvecMaxsim => { + x.to_f32() + } + } + } + pub const fn distance_kind(self) -> DistanceKind { + match self { + Self::VectorL2 | Self::HalfvecL2 => DistanceKind::L2S, + Self::VectorIp + | Self::HalfvecIp + | Self::VectorCosine + | Self::HalfvecCosine + | Self::VectorMaxsim + | Self::HalfvecMaxsim => DistanceKind::Dot, + } + } + pub const fn vector_kind(self) -> VectorKind { + match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { + VectorKind::Vecf32 + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { + VectorKind::Vecf16 + } + } + } +} + +pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { + use pgrx::pg_sys::Oid; + + let proc = unsafe { pgrx::pg_sys::index_getprocid(index_relation, 1, 1) }; + + if proc == Oid::INVALID { + pgrx::error!("support function 1 is not found"); + } + + let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); + + unsafe { + pgrx::pg_sys::fmgr_info(proc, &mut flinfo); + } + + let fn_addr = flinfo.fn_addr.expect("null function pointer"); + + let mut fcinfo = unsafe { std::mem::zeroed::() }; + fcinfo.flinfo = &mut flinfo; + fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; + fcinfo.context = std::ptr::null_mut(); + fcinfo.resultinfo = std::ptr::null_mut(); + fcinfo.isnull = true; + fcinfo.nargs = 0; + + let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; + + let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; + + let result_string = result_option.expect("null return value"); + + let result = match result_string.as_str() { + "vchordrq_vector_l2_ops" => Opfamily::VectorL2, + "vchordrq_vector_ip_ops" => Opfamily::VectorIp, + "vchordrq_vector_cosine_ops" => Opfamily::VectorCosine, + "vchordrq_halfvec_l2_ops" => Opfamily::HalfvecL2, + "vchordrq_halfvec_ip_ops" => Opfamily::HalfvecIp, + "vchordrq_halfvec_cosine_ops" => Opfamily::HalfvecCosine, + "vchordrq_vector_maxsim_ops" => Opfamily::VectorMaxsim, + "vchordrq_halfvec_maxsim_ops" => Opfamily::HalfvecMaxsim, + _ => pgrx::error!("unknown operator class"), + }; + + unsafe { + pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); + } + + result +} diff --git a/src/index/scanners/default.rs b/src/index/vchordrq/scanners/default.rs similarity index 97% rename from src/index/scanners/default.rs rename to src/index/vchordrq/scanners/default.rs index 411a51bc..1f4e6296 100644 --- a/src/index/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -12,18 +12,21 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::{Fetcher, Io, SearchBuilder, SearchOptions, filter}; -use crate::index::algorithm::*; -use crate::index::am::pointer_to_kv; -use crate::index::opclass::{Opfamily, Sphere}; -use crate::index::scanners::Tuple; -use algorithm::operator::{self, Dot, L2}; -use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; -use algorithm::*; +use super::{Io, SearchBuilder, SearchOptions, filter}; +use crate::index::fetcher::*; +use crate::index::opclass::Sphere; +use crate::index::vchordrq::algo::*; +use crate::index::vchordrq::opclass::Opfamily; +use algo::accessor::{Dot, L2S}; +use algo::prefetcher::*; +use algo::*; use always_equal::AlwaysEqual; use half::f16; use std::collections::BinaryHeap; use std::num::NonZero; +use vchordrq::operator::{self}; +use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; +use vchordrq::*; use vector::VectorOwned; use vector::vect::VectOwned; @@ -34,6 +37,8 @@ pub struct DefaultBuilder { } impl SearchBuilder for DefaultBuilder { + type Opaque = vchordrq::Opaque; + fn new(opfamily: Opfamily) -> Self { assert!(matches!( opfamily, @@ -74,6 +79,7 @@ impl SearchBuilder for DefaultBuilder { ) -> Box + 'a> where R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page, { let mut vector = None; let mut threshold = None; @@ -106,8 +112,8 @@ impl SearchBuilder for DefaultBuilder { let f = move |(distance, payload)| (opfamily.output(distance), payload); let iter: Box)>> = match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2) => { - type Op = operator::Op, L2>; + (VectorKind::Vecf32, DistanceKind::L2S) => { + type Op = operator::Op, L2S>; let unprojected = if let OwnedVector::Vecf32(vector) = vector { vector } else { @@ -372,8 +378,8 @@ impl SearchBuilder for DefaultBuilder { } } } - (VectorKind::Vecf16, DistanceKind::L2) => { - type Op = operator::Op, L2>; + (VectorKind::Vecf16, DistanceKind::L2S) => { + type Op = operator::Op, L2S>; let unprojected = if let OwnedVector::Vecf16(vector) = vector { vector } else { diff --git a/src/index/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs similarity index 98% rename from src/index/scanners/maxsim.rs rename to src/index/vchordrq/scanners/maxsim.rs index f70b8570..dd3c8426 100644 --- a/src/index/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -12,20 +12,22 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::{Fetcher, SearchBuilder, SearchOptions}; -use crate::index::algorithm::*; -use crate::index::am::pointer_to_kv; -use crate::index::opclass::Opfamily; -use crate::index::scanners::{Io, Tuple, filter}; -use algorithm::operator::Dot; -use algorithm::types::{DistanceKind, OwnedVector, VectorKind}; -use algorithm::*; +use super::{SearchBuilder, SearchOptions}; +use crate::index::fetcher::*; +use crate::index::vchordrq::algo::*; +use crate::index::vchordrq::opclass::Opfamily; +use crate::index::vchordrq::scanners::{Io, filter}; +use algo::accessor::Dot; +use algo::prefetcher::*; +use algo::*; use always_equal::AlwaysEqual; use distance::Distance; use half::f16; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZero; +use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; +use vchordrq::*; use vector::VectorOwned; use vector::vect::VectOwned; @@ -35,6 +37,8 @@ pub struct MaxsimBuilder { } impl SearchBuilder for MaxsimBuilder { + type Opaque = vchordrq::Opaque; + fn new(opfamily: Opfamily) -> Self { assert!(matches!( opfamily, @@ -65,6 +69,7 @@ impl SearchBuilder for MaxsimBuilder { ) -> Box + 'a> where R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page, { let mut vectors = None; for orderby_vectors in self.orderbys.into_iter().flatten() { diff --git a/src/index/scanners/mod.rs b/src/index/vchordrq/scanners/mod.rs similarity index 77% rename from src/index/scanners/mod.rs rename to src/index/vchordrq/scanners/mod.rs index 5a74485a..a0340a3e 100644 --- a/src/index/scanners/mod.rs +++ b/src/index/vchordrq/scanners/mod.rs @@ -15,11 +15,11 @@ mod default; mod maxsim; +use crate::index::fetcher::Fetcher; + use super::opclass::Opfamily; -use algorithm::{Bump, RelationPrefetch, RelationRead, RelationReadStream, Sequence}; +use algo::{Bump, Page, RelationPrefetch, RelationRead, RelationReadStream, Sequence}; use pgrx::pg_sys::Datum; -use std::cell::LazyCell; -use std::ops::DerefMut; pub use default::DefaultBuilder; pub use maxsim::MaxsimBuilder; @@ -48,6 +48,8 @@ pub struct SearchOptions { } pub trait SearchBuilder: 'static { + type Opaque: Copy; + fn new(opfamily: Opfamily) -> Self; unsafe fn add(&mut self, strategy: u16, datum: Option); @@ -60,31 +62,8 @@ pub trait SearchBuilder: 'static { bump: &'a impl Bump, ) -> Box + 'a> where - R: RelationRead + RelationPrefetch + RelationReadStream; -} - -pub trait Fetcher { - type Tuple<'a>: Tuple - where - Self: 'a; - - fn fetch(&mut self, key: [u16; 3]) -> Option>; -} - -pub trait Tuple { - fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); - fn filter(&mut self) -> bool; -} - -impl T> Fetcher for LazyCell { - type Tuple<'a> - = T::Tuple<'a> - where - Self: 'a; - - fn fetch(&mut self, key: [u16; 3]) -> Option> { - self.deref_mut().fetch(key) - } + R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page; } pub struct Filter { diff --git a/src/index/types.rs b/src/index/vchordrq/types.rs similarity index 98% rename from src/index/types.rs rename to src/index/vchordrq/types.rs index 4e6e286f..a25b4782 100644 --- a/src/index/types.rs +++ b/src/index/vchordrq/types.rs @@ -12,9 +12,9 @@ // // Copyright (c) 2025 TensorChord Inc. -use algorithm::types::VchordrqIndexOptions; use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError, ValidationErrors}; +use vchordrq::types::VchordrqIndexOptions; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] diff --git a/src/lib.rs b/src/lib.rs index 89b7a2a5..64912e7e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. #![allow(unsafe_code)] +#![allow(clippy::type_complexity)] mod datatype; mod index; @@ -31,7 +32,7 @@ extern "C-unwind" fn _PG_init() { unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); - #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))] + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] pgrx::pg_sys::MarkGUCPrefixReserved(c"vchord".as_ptr()); } } diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index c97f6c46..6200b4bb 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -148,9 +148,13 @@ IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhan CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + -- List of access methods CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; -- List of operator families @@ -162,6 +166,12 @@ CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; -- List of operator classes @@ -210,3 +220,39 @@ CREATE OPERATOR CLASS halfvec_maxsim_ops FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); diff --git a/vchord.control b/vchord.control index 5ce41ec2..603cc3ff 100644 --- a/vchord.control +++ b/vchord.control @@ -1,6 +1,6 @@ comment = 'vchord: Vector database plugin for Postgres, written in Rust, specifically designed for LLM' default_version = '0.0.0' -module_pathname = '$libdir/vchord' +module_pathname = 'vchord' relocatable = true superuser = true requires = 'vector' From 9b8889a24844cf785b18c4e23fefe0195ae4ec86 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 1 Jul 2025 12:36:22 +0800 Subject: [PATCH 172/324] feat: add `vchordg_prewarm` (#282) Signed-off-by: usamoi --- crates/vchordg/src/lib.rs | 2 ++ crates/vchordg/src/prewarm.rs | 46 +++++++++++++++++++++++++++++++++++ src/index/functions.rs | 22 +++++++++++++++++ src/index/vchordg/algo.rs | 21 ++++++++++++++++ src/sql/finalize.sql | 3 +++ 5 files changed, 94 insertions(+) create mode 100644 crates/vchordg/src/prewarm.rs diff --git a/crates/vchordg/src/lib.rs b/crates/vchordg/src/lib.rs index 43550656..ad0f767e 100644 --- a/crates/vchordg/src/lib.rs +++ b/crates/vchordg/src/lib.rs @@ -5,6 +5,7 @@ mod bulkdelete; mod candidates; mod insert; mod maintain; +mod prewarm; mod prune; mod results; mod search; @@ -19,6 +20,7 @@ pub use build::build; pub use bulkdelete::bulkdelete; pub use insert::insert; pub use maintain::maintain; +pub use prewarm::prewarm; pub use search::search; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; diff --git a/crates/vchordg/src/prewarm.rs b/crates/vchordg/src/prewarm.rs new file mode 100644 index 00000000..8024be2a --- /dev/null +++ b/crates/vchordg/src/prewarm.rs @@ -0,0 +1,46 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::Opaque; +use crate::operator::Operator; +use algo::{Page, RelationRead}; + +pub fn prewarm(index: &R) -> String +where + R::Page: Page, +{ + use std::fmt::Write; + let meta_guard = index.read(0); + let link = meta_guard.get_opaque().link; + drop(meta_guard); + let mut number_of_vertex_pages = 0_u64; + let mut number_of_vertices = 0_u64; + { + let mut current = link; + while current != u32::MAX { + let vertex_guard = index.read(current); + number_of_vertex_pages += 1; + for i in 1..=vertex_guard.len() { + if vertex_guard.get(i).is_some() { + number_of_vertices += 1; + } + } + current = vertex_guard.get_opaque().next; + } + } + let mut message = String::new(); + writeln!(message, "number of vertex pages: {number_of_vertex_pages}").unwrap(); + writeln!(message, "number of vertices: {number_of_vertices}").unwrap(); + message +} diff --git a/src/index/functions.rs b/src/index/functions.rs index b2dbe03b..3711fac4 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -16,6 +16,28 @@ use crate::index::storage::PostgresRelation; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass, PgClassRelkind}; +#[pgrx::pg_extern(sql = "")] +fn _vchordg_prewarm(indexrelid: Oid) -> String { + let pg_am = PgAm::search_amname(c"vchordg").unwrap(); + let Some(pg_am) = pg_am.get() else { + pgrx::error!("vchord is not installed"); + }; + let pg_class = PgClass::search_reloid(indexrelid).unwrap(); + let Some(pg_class) = pg_class.get() else { + pgrx::error!("the relation does not exist"); + }; + if pg_class.relkind() != PgClassRelkind::Index { + pgrx::error!("the relation {:?} is not an index", pg_class.relname()); + } + if pg_class.relam() != pg_am.oid() { + pgrx::error!("the index {:?} is not a vchordg index", pg_class.relname()); + } + let relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); + let opfamily = unsafe { crate::index::vchordg::opclass::opfamily(relation.raw()) }; + let index = unsafe { PostgresRelation::new(relation.raw()) }; + crate::index::vchordg::algo::prewarm(opfamily, &index) +} + #[pgrx::pg_extern(sql = "")] fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { let pg_am = PgAm::search_amname(c"vchordrq").unwrap(); diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs index 89a0d1f3..f12a6a40 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/algo.rs @@ -24,6 +24,27 @@ use vchordg::types::*; use vector::VectorOwned; use vector::vect::{VectBorrowed, VectOwned}; +pub fn prewarm(opfamily: Opfamily, index: &R) -> String +where + R: RelationRead, + R::Page: Page, +{ + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordg::prewarm::<_, Op, L2S>>(index) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordg::prewarm::<_, Op, Dot>>(index) + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::prewarm::<_, Op, L2S>>(index) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordg::prewarm::<_, Op, Dot>>(index) + } + } +} + pub fn bulkdelete( opfamily: Opfamily, index: &R, diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 6200b4bb..594e3b8a 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -151,6 +151,9 @@ STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + -- List of access methods CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; From ac29775184912253eec386ca908688ca35739c41 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 8 Jul 2025 20:52:12 +0800 Subject: [PATCH 173/324] refactor: simpify vchordg (#283) contains bugfix job: +psql_macos job: +psql_windows Signed-off-by: usamoi --- .github/workflows/check.yml | 38 +- Cargo.lock | 129 +---- Cargo.toml | 17 +- build.rs | 12 + crates/algo/Cargo.toml | 5 +- crates/algo/src/lib.rs | 4 +- crates/algo/src/prefetcher.rs | 14 +- crates/make/Cargo.toml | 6 +- crates/make/src/main.rs | 460 ++++++++------- crates/rabitq/Cargo.toml | 1 - crates/simd/Cargo.toml | 3 +- crates/simd/src/emulate.rs | 10 +- crates/simd/src/u8.rs | 1 - crates/vchordg/Cargo.toml | 3 - crates/vchordg/src/build.rs | 7 +- crates/vchordg/src/bulkdelete.rs | 30 +- crates/vchordg/src/insert.rs | 537 ++++++------------ crates/vchordg/src/lib.rs | 6 +- crates/vchordg/src/maintain.rs | 284 ++++----- crates/vchordg/src/search.rs | 81 +-- crates/vchordg/src/tuples.rs | 12 +- crates/vchordg/src/types.rs | 7 +- crates/vchordg/src/vectors.rs | 168 ++++++ crates/vchordrq/Cargo.toml | 6 +- crates/vchordrq/src/lib.rs | 2 - crates/vchordrq/src/search.rs | 2 - crates/vchordrq/src/tuples.rs | 1 - crates/vchordrq/src/types.rs | 5 +- crates/vector/Cargo.toml | 2 - docker/Dockerfile | 2 +- src/bin/pgrx_embed.rs | 30 + src/index/gucs.rs | 2 +- src/index/storage.rs | 34 +- src/index/vchordg/algo.rs | 9 +- src/index/vchordg/am/mod.rs | 3 - src/index/vchordg/opclass.rs | 20 +- src/index/vchordrq/am/mod.rs | 3 - src/lib.rs | 28 +- tests/README.md | 14 - tests/{general => fail}/null.fail | 0 tests/general/{issue427.slt => issue_427.slt} | 3 + tests/vchordg/partition.slt | 51 ++ tests/vchordg/reindex.slt | 33 ++ .../{general => vchordrq}/external_build.slt | 0 .../filter_rerank_in_index.slt | 0 .../filter_rerank_in_table.slt | 0 tests/{general => vchordrq}/index.slt | 0 tests/{general => vchordrq}/multivector.slt | 0 tests/{general => vchordrq}/partition.slt | 0 .../pg17/filter_rerank_in_index.slt | 0 .../pg17/filter_rerank_in_table.slt | 0 tests/{general => vchordrq}/pin.slt | 0 tests/{general => vchordrq}/pushdown_plan.slt | 0 .../{general => vchordrq}/pushdown_range.slt | 0 tests/{general => vchordrq}/reindex.slt | 0 .../{general => vchordrq}/rerank_in_index.slt | 0 .../{general => vchordrq}/rerank_in_table.slt | 0 57 files changed, 1023 insertions(+), 1062 deletions(-) create mode 100644 crates/vchordg/src/vectors.rs delete mode 100644 tests/README.md rename tests/{general => fail}/null.fail (100%) rename tests/general/{issue427.slt => issue_427.slt} (84%) create mode 100644 tests/vchordg/partition.slt create mode 100644 tests/vchordg/reindex.slt rename tests/{general => vchordrq}/external_build.slt (100%) rename tests/{general => vchordrq}/filter_rerank_in_index.slt (100%) rename tests/{general => vchordrq}/filter_rerank_in_table.slt (100%) rename tests/{general => vchordrq}/index.slt (100%) rename tests/{general => vchordrq}/multivector.slt (100%) rename tests/{general => vchordrq}/partition.slt (100%) rename tests/{ => vchordrq}/pg17/filter_rerank_in_index.slt (100%) rename tests/{ => vchordrq}/pg17/filter_rerank_in_table.slt (100%) rename tests/{general => vchordrq}/pin.slt (100%) rename tests/{general => vchordrq}/pushdown_plan.slt (100%) rename tests/{general => vchordrq}/pushdown_range.slt (100%) rename tests/{general => vchordrq}/reindex.slt (100%) rename tests/{general => vchordrq}/rerank_in_index.slt (100%) rename tests/{general => vchordrq}/rerank_in_table.slt (100%) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 3d1cd725..230143f6 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -192,7 +192,7 @@ jobs: echo "host all all ::1/128 trust" | sudo tee -a /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf sudo -iu postgres createuser -s -r $USER sudo -iu postgres createdb -O $USER $USER - sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.so"' + sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo systemctl stop postgresql curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest @@ -208,7 +208,7 @@ jobs: - name: Install run: | - cargo make build -o ./build/raw + cargo make build -o ./build/raw --profile dev sudo make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install - name: Service @@ -217,11 +217,13 @@ jobs: psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - name: Sqllogictest - run: sqllogictest --db $USER --user $USER './tests/general/*.slt' - - - name: Sqllogictest (PostgreSQL 17) - if: matrix.version == '17' - run: sqllogictest --db $USER --user $USER './tests/pg17/*.slt' + run: | + sqllogictest --db $USER --user $USER './tests/general/*.slt' + sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' + sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' + if [ "${{ matrix.version }}" = "17" ]; then + sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' + fi - name: Package env: @@ -290,11 +292,7 @@ jobs: for i in {1..60}; do [ -S /tmp/.s.PGSQL.5432 ] && echo "PostgreSQL ready" && break || sleep 1; done [ -S /tmp/.s.PGSQL.5432 ] || echo "PostgreSQL socket not found after 60 seconds" $(brew --prefix postgresql@${{ matrix.version }})/bin/createdb -O $USER $USER - if [[ ${{ matrix.version }} != 13 && ${{ matrix.version }} != 14 && ${{ matrix.version }} != 15 ]]; then - $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.dylib"' - else - $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord.so"' - fi + $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' brew services stop postgresql@${{ matrix.version }} echo PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV @@ -312,7 +310,7 @@ jobs: - name: Install run: | - cargo make build -o ./build/raw + cargo make build -o ./build/raw --profile dev sudo make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install - name: Service @@ -323,11 +321,13 @@ jobs: $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - name: Sqllogictest - run: sqllogictest --db $USER --user $USER './tests/general/*.slt' - - - name: Sqllogictest (PostgreSQL 17) - if: matrix.version == '17' - run: sqllogictest --db $USER --user $USER './tests/pg17/*.slt' + run: | + sqllogictest --db $USER --user $USER './tests/general/*.slt' + sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' + sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' + if [ "${{ matrix.version }}" = "17" ]; then + sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' + fi - name: Package env: @@ -392,5 +392,5 @@ jobs: - name: Install run: | - cargo make build -o ./build/raw + cargo make build -o ./build/raw --profile dev make PG_CONFIG="$env:PGRX_PG_CONFIG_PATH" install diff --git a/Cargo.lock b/Cargo.lock index 33cf64de..2fc4e5d7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -86,7 +86,7 @@ version = "1.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6c8bdeb6047d8983be085bab0ba1472e6dc604e7041dbf6fcd5e71523014fae9" dependencies = [ - "windows-sys 0.59.0", + "windows-sys", ] [[package]] @@ -97,7 +97,7 @@ checksum = "403f75924867bb1033c59fbf0797484329750cfbe3c4325cd33127941fabc882" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.59.0", + "windows-sys", ] [[package]] @@ -412,28 +412,12 @@ dependencies = [ "syn", ] -[[package]] -name = "env_home" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe" - [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" -[[package]] -name = "errno" -version = "0.3.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" -dependencies = [ - "libc", - "windows-sys 0.60.2", -] - [[package]] name = "eyre" version = "0.6.12" @@ -444,12 +428,6 @@ dependencies = [ "once_cell", ] -[[package]] -name = "fastrand" -version = "2.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" - [[package]] name = "fixedbitset" version = "0.5.7" @@ -576,7 +554,7 @@ version = "0.5.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589533453244b0995c858700322199b2becb13b627df2851f64a2775d024abcf" dependencies = [ - "windows-sys 0.59.0", + "windows-sys", ] [[package]] @@ -716,7 +694,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys", ] [[package]] @@ -812,12 +790,6 @@ dependencies = [ "libc", ] -[[package]] -name = "linux-raw-sys" -version = "0.9.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd945864f07fe9f5371a27ad7b52a172b4b499999f1d97574c9fa68373937e12" - [[package]] name = "litemap" version = "0.8.0" @@ -834,11 +806,9 @@ checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" name = "make" version = "0.0.0" dependencies = [ - "anyhow", "clap", "object", - "tempfile", - "which", + "target-triple", ] [[package]] @@ -889,9 +859,9 @@ dependencies = [ [[package]] name = "object" -version = "0.36.7" +version = "0.37.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +checksum = "03fd943161069e1768b4b3d050890ba48730e590f57e56d4aa04e7e090e61b4a" dependencies = [ "crc32fast", "flate2", @@ -1163,7 +1133,6 @@ checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" name = "rabitq" version = "0.0.0" dependencies = [ - "distance", "rand", "rand_chacha", "simd", @@ -1269,19 +1238,6 @@ dependencies = [ "semver", ] -[[package]] -name = "rustix" -version = "1.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c71e83d6afe7ff64890ec6b71d6a69bb8a610ab78ce364b3352876bb4c801266" -dependencies = [ - "bitflags", - "errno", - "libc", - "linux-raw-sys", - "windows-sys 0.59.0", -] - [[package]] name = "rustversion" version = "1.0.21" @@ -1290,9 +1246,9 @@ checksum = "8a0d197bd2c9dc6e53b84da9556a69ba4cdfab8619eb41a8bd1cc2027a0f6b1d" [[package]] name = "ruzstd" -version = "0.7.3" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fad02996bfc73da3e301efe90b1837be9ed8f4a462b6ed410aa35d00381de89f" +checksum = "3640bec8aad418d7d03c72ea2de10d5c646a598f9883c7babc160d91e3c1b26c" dependencies = [ "twox-hash", ] @@ -1393,7 +1349,6 @@ version = "0.0.0" dependencies = [ "cc", "half 2.6.0", - "paste", "rand", "seq-macro", "simd_macros", @@ -1427,12 +1382,6 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" -[[package]] -name = "static_assertions" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" - [[package]] name = "strsim" version = "0.11.1" @@ -1487,17 +1436,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" [[package]] -name = "tempfile" -version = "3.20.0" +name = "target-triple" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8a64e3985349f2441a1a9ef0b853f869006c3855f2cda6862a94d26ebb9d6a1" -dependencies = [ - "fastrand", - "getrandom", - "once_cell", - "rustix", - "windows-sys 0.59.0", -] +checksum = "1ac9aa371f599d22256307c24a9d748c041e548cbf599f35d890f9d365361790" [[package]] name = "thiserror" @@ -1572,13 +1514,9 @@ checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" [[package]] name = "twox-hash" -version = "1.6.3" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97fee6b57c6a41524a810daee9286c02d7752c4253064d0b05472833a438f675" -dependencies = [ - "cfg-if", - "static_assertions", -] +checksum = "8b907da542cbced5261bd3256de1b3a1bf340a3d37f93425a07362a1d687de56" [[package]] name = "unescape" @@ -1704,10 +1642,7 @@ dependencies = [ "always_equal", "distance", "half 2.6.0", - "k_means", "min-max-heap", - "paste", - "pin-project", "rabitq", "rand", "serde", @@ -1725,8 +1660,6 @@ dependencies = [ "always_equal", "distance", "half 2.6.0", - "k_means", - "paste", "pin-project", "rabitq", "rand", @@ -1742,7 +1675,6 @@ name = "vector" version = "0.0.0" dependencies = [ "distance", - "half 2.6.0", "simd", ] @@ -1825,25 +1757,13 @@ dependencies = [ [[package]] name = "wasmparser" -version = "0.222.1" +version = "0.234.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa210fd1788e6b37a1d1930f3389c48e1d6ebd1a013d34fa4b7f9e3e3bf03146" +checksum = "be22e5a8f600afce671dd53c8d2dd26b4b7aa810fd18ae27dfc49737f3e02fc5" dependencies = [ "bitflags", ] -[[package]] -name = "which" -version = "7.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "24d643ce3fd3e5b54854602a080f34fb10ab75e0b813ee32d00ca2b44fa74762" -dependencies = [ - "either", - "env_home", - "rustix", - "winsafe", -] - [[package]] name = "winapi" version = "0.3.9" @@ -1866,7 +1786,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys", ] [[package]] @@ -1884,15 +1804,6 @@ dependencies = [ "windows-targets 0.52.6", ] -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets 0.53.2", -] - [[package]] name = "windows-targets" version = "0.52.6" @@ -2030,12 +1941,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "winsafe" -version = "0.0.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" - [[package]] name = "wit-bindgen-rt" version = "0.39.0" diff --git a/Cargo.toml b/Cargo.toml index 61e7071a..1c4995a6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -72,15 +72,30 @@ validator = { version = "0.20.0", features = ["derive"] } zerocopy = { version = "0.8.26", features = ["derive"] } [workspace.lints] +# complexity clippy.identity_op = "allow" clippy.int_plus_one = "allow" -clippy.needless_range_loop = "allow" clippy.nonminimal_bool = "allow" +clippy.too_many_arguments = "allow" +clippy.type_complexity = "allow" +# style +clippy.just_underscores_and_digits = "allow" +clippy.needless_range_loop = "allow" +# unsafe rust.unsafe_code = "deny" rust.unsafe_op_in_unsafe_fn = "deny" +# unused +rust.unused_crate_dependencies = "warn" +rust.unused_extern_crates = "warn" +rust.unused_import_braces = "warn" rust.unused_lifetimes = "warn" +rust.unused_macro_rules = "warn" rust.unused_qualifications = "warn" +[profile.dev] +lto = false +opt-level = 1 + [profile.release] codegen-units = 1 debug = true diff --git a/build.rs b/build.rs index fdf80e9d..d6138350 100644 --- a/build.rs +++ b/build.rs @@ -40,5 +40,17 @@ fn main() -> Result<(), Box> { println!("cargo::rustc-link-arg-cdylib=-Wl,-undefined,dynamic_lookup"); } } + let version = 'version: { + for line in std::fs::read_to_string("./vchord.control")?.lines() { + if let Some(prefix_stripped) = line.strip_prefix("default_version = '") + && let Some(stripped) = prefix_stripped.strip_suffix("'") + { + eprintln!("VectorChord version: {stripped}"); + break 'version stripped.to_string(); + } + } + return Err("VectorChord version is not defined.".into()); + }; + println!("cargo::rustc-env=VCHORD_VERSION={version}"); Ok(()) } diff --git a/crates/algo/Cargo.toml b/crates/algo/Cargo.toml index 44ce2ed5..25568150 100644 --- a/crates/algo/Cargo.toml +++ b/crates/algo/Cargo.toml @@ -6,11 +6,12 @@ publish = false [dependencies] always_equal = { path = "../always_equal" } -bumpalo.workspace = true distance = { path = "../distance" } -half.workspace = true simd = { path = "../simd" } vector = { path = "../vector" } + +bumpalo.workspace = true +half.workspace = true zerocopy.workspace = true [lints] diff --git a/crates/algo/src/lib.rs b/crates/algo/src/lib.rs index 8e5de051..057cefcd 100644 --- a/crates/algo/src/lib.rs +++ b/crates/algo/src/lib.rs @@ -12,8 +12,6 @@ // // Copyright (c) 2025 TensorChord Inc. -#![allow(clippy::type_complexity)] - pub mod accessor; pub mod prefetcher; pub mod tuples; @@ -55,7 +53,7 @@ pub trait PageGuard { pub trait ReadStream<'r> { type Relation: RelationReadTypes; - type Guards: Iterator::ReadGuard<'r>>; + type Guards: ExactSizeIterator::ReadGuard<'r>>; type Item; type Inner: Iterator; fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; diff --git a/crates/algo/src/prefetcher.rs b/crates/algo/src/prefetcher.rs index d0eb5915..c36998bc 100644 --- a/crates/algo/src/prefetcher.rs +++ b/crates/algo/src/prefetcher.rs @@ -27,7 +27,7 @@ where Self::Item: Fetch, { type R: RelationRead + 'r; - type Guards: Iterator::ReadGuard<'r>>; + type Guards: ExactSizeIterator::ReadGuard<'r>>; #[must_use] fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; @@ -109,8 +109,14 @@ impl<'r, R: RelationRead> Iterator for PlainPrefetcherGuards<'r, R> { let id = self.list.next()?; Some(self.relation.read(id)) } + + fn size_hint(&self) -> (usize, Option) { + self.list.size_hint() + } } +impl<'r, R: RelationRead> ExactSizeIterator for PlainPrefetcherGuards<'r, R> {} + pub struct SimplePrefetcher<'r, R, S: Sequence> { relation: &'r R, window: VecDeque, @@ -200,8 +206,14 @@ impl<'r, R: RelationRead> Iterator for SimplePrefetcherGuards<'r, R> { let id = self.list.next()?; Some(self.relation.read(id)) } + + fn size_hint(&self) -> (usize, Option) { + self.list.size_hint() + } } +impl<'r, R: RelationRead> ExactSizeIterator for SimplePrefetcherGuards<'r, R> {} + pub struct StreamPrefetcherSequence(S); impl Iterator for StreamPrefetcherSequence { diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index 7e5d20bc..781e910f 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -5,8 +5,6 @@ edition.workspace = true publish = false [dependencies] -anyhow = "1.0.98" clap = { version = "4.5.38", features = ["derive"] } -object = { version = "0.36.7", features = ["all"] } -tempfile = "3.20.0" -which = "7.0.3" +object = { version = "0.37.1", features = ["all"] } +target-triple = "0.1.4" diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index cf17d576..dec02835 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -12,15 +12,13 @@ // // Copyright (c) 2025 TensorChord Inc. -use anyhow::{Context, Result, bail}; use clap::{Args, Parser, Subcommand}; use object::Object; -use std::collections::HashMap; -use std::env::consts::{DLL_PREFIX, DLL_SUFFIX, OS}; +use std::collections::{HashMap, HashSet}; use std::env::var_os; +use std::error::Error; use std::fs::read_dir; -use std::io::Write; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::process::{Command, Stdio}; #[derive(Parser)] @@ -38,220 +36,237 @@ enum Commands { struct BuildArgs { #[arg(short, long)] output: String, + #[arg(long, default_value = "release")] + profile: String, + #[arg(long, default_value = target_triple::TARGET)] + target: String, } -fn main() -> Result<()> { - let cli = Cli::parse(); - if !std::fs::exists("vchord.control")? { - bail!("The script must be run from the VectorChord source directory.") +fn pg_config(pg_config: impl AsRef) -> Result, Box> { + let mut command = Command::new(pg_config.as_ref()); + command.stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_output = command.output()?; + let contents = String::from_utf8(command_output.stdout)?; + let mut result = HashMap::new(); + for line in contents.lines() { + if let Some((key, value)) = line.split_once(" = ") { + result.insert(key.to_string(), value.to_string()); + } } - let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { - eprintln!("Environment variable `PGRX_PG_CONFIG_PATH`: {value:#?}"); - PathBuf::from(value) - } else { - if let Ok(path) = which::which("pg_config") { - eprintln!("Found executable `pg_config` in PATH: {path:#?}"); - eprintln!("Hint: set environment variable `PGRX_PG_CONFIG_PATH` to {path:#?}"); + Ok(result) +} + +fn control_file(path: impl AsRef) -> Result, Box> { + let path = path.as_ref(); + eprintln!("Reading {path:?}"); + let contents = std::fs::read_to_string(path)?; + let mut result = HashMap::new(); + for line in contents.lines() { + if let Some((key, prefix_stripped)) = line.split_once(" = '") + && let Some(value) = prefix_stripped.strip_suffix("'") + { + result.insert(key.to_string(), value.to_string()); } - bail!("Environment variable `PGRX_PG_CONFIG_PATH` is not set.") + } + Ok(result) +} + +fn cfgs(target: &str) -> Result, Box> { + let mut command = Command::new("rustc"); + command + .args(["--print", "cfg"]) + .args(["--target", target]) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_output = command.output()?; + let contents = String::from_utf8(command_output.stdout)?; + let mut result = HashSet::new(); + for line in contents.lines() { + result.insert(line.to_string()); + } + Ok(result) +} + +fn build( + pg_config: impl AsRef, + pg_version: &str, + cfgs: &HashSet, + profile: &str, + target: &str, +) -> Result> { + let mut command = Command::new("cargo"); + command + .args(["build", "-p", "vchord", "--lib"]) + .args(["--profile", profile]) + .args(["--target", target]) + .args(["--features", pg_version]) + .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_status = command.spawn()?.wait()?; + if !command_status.success() { + return Err(format!("Cargo build failed: {command_status}").into()); + } + let mut result = PathBuf::from("./target"); + result.push(target); + result.push(if profile != "dev" { profile } else { "debug" }); + let is_unix = cfgs.contains("target_family=\"unix\""); + let is_macos = cfgs.contains("target_os=\"macos\""); + let is_windows = cfgs.contains("target_family=\"windows\""); + let is_wasm = cfgs.contains("target_family=\"wasm\""); + let prefix = if is_unix { + "lib" + } else if is_windows || is_wasm { + "" + } else { + return Err("unknown operating system".into()); }; - let map = { - let mut command = Command::new(&path); - command.stderr(Stdio::inherit()); - let command_output = command.output()?; - let command_stdout = String::from_utf8(command_output.stdout)?; - let mut map = HashMap::new(); - for line in command_stdout.lines() { - if let Some((key, value)) = line.split_once(" = ") { - map.insert(key.to_string(), value.to_string()); - eprintln!("Config `{key}`: {value}"); - } - } - map + let suffix = if is_unix && !is_macos { + ".so" + } else if is_macos { + ".dylib" + } else if is_windows { + ".dll" + } else if is_wasm { + ".wasm" + } else { + return Err("unknown operating system".into()); }; - let fork = { - let version = map["VERSION"].clone(); - let fork = if let Some(prefix_stripped) = version.strip_prefix("PostgreSQL ") { - if let Some((stripped, _)) = prefix_stripped.split_once(|c: char| !c.is_ascii_digit()) { - format!("pg{stripped}",) + result.push(format!("{prefix}vchord{suffix}")); + Ok(result) +} + +fn parse(cfgs: &HashSet, obj: impl AsRef) -> Result, Box> { + let obj = obj.as_ref(); + eprintln!("Reading {obj:?}"); + let contents = std::fs::read(obj)?; + let object = object::File::parse(contents.as_slice())?; + let is_macos = cfgs.contains("target_os=\"macos\""); + let exports = object + .exports()? + .into_iter() + .flat_map(|x| str::from_utf8(x.name())) + .flat_map(|x| { + if !is_macos { + Some(x) } else { - format!("pg{prefix_stripped}",) - } - } else { - bail!("PostgreSQL version is invalid.") - }; - eprintln!("Fork: {fork}"); - fork - }; - #[allow(clippy::collapsible_if)] - let version = 'version: { - for line in std::fs::read_to_string("./vchord.control")?.lines() { - if let Some(prefix_stripped) = line.strip_prefix("default_version = '") { - if let Some(stripped) = prefix_stripped.strip_suffix("'") { - eprintln!("VectorChord version: {stripped}"); - break 'version stripped.to_string(); - } + x.strip_prefix("_") } - } - bail!("VectorChord version is not defined.") - }; - let build = || { - let mut command = Command::new("cargo"); - command - .args(["build", "--release"]) - .args(["-p", "vchord", "--lib"]) - .args(["--features", fork.as_str()]) - .env("PGRX_PG_CONFIG_PATH", &path) - .stdout(Stdio::inherit()) - .stderr(Stdio::inherit()); - let debug = format!("{command:?}"); - let status = command.spawn()?.wait()?; - if !status.success() { - bail!("Cargo build failed: {debug}"); - } - Ok(()) - }; - let schema = || { - let object = std::fs::read(format!("./target/release/{DLL_PREFIX}vchord{DLL_SUFFIX}"))?; - let object = object::File::parse(object.as_slice())?; - let exports = object - .exports()? - .into_iter() - .flat_map(|x| str::from_utf8(x.name())); - let exports = if matches!(object.format(), object::BinaryFormat::MachO) { - exports - .flat_map(|x| x.strip_prefix("_")) - .filter(|x| x.starts_with("__pgrx_internals")) - .collect::>() - } else { - exports - .filter(|x| x.starts_with("__pgrx_internals")) - .collect::>() - }; - let pushes = exports - .into_iter() - .map(|x| { - format!( - r#" - entities.push(unsafe {{ - unsafe extern "Rust" {{ - fn {x}() -> ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity; - }} - {x}() - }}); - "# - ) - }) - .collect::>() - .join("\n"); - let code = format!( - r#" - pub fn main() {{ - extern crate vchord as _; - - let mut entities = Vec::new(); - let control_file = std::fs::read_to_string("./vchord.control").unwrap(); - let control_file = ::pgrx::pgrx_sql_entity_graph::ControlFile::try_from(control_file.as_str()).expect(".control file should properly formatted"); - let control_file_entity = ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity::ExtensionRoot(control_file); + }) + .filter(|x| x.starts_with("__pgrx_internals")) + .map(str::to_string) + .collect(); + Ok(exports) +} - entities.push(control_file_entity); +fn generate( + pg_config: impl AsRef, + pg_version: &str, + profile: &str, + target: &str, + exports: Vec, +) -> Result> { + let pgrx_embed = std::env::temp_dir().join("VCHORD_PGRX_EMBED"); + eprintln!("Writing {pgrx_embed:?}"); + std::fs::write( + &pgrx_embed, + format!("crate::schema_generation!({});", exports.join(" ")), + )?; + let mut command = Command::new("cargo"); + command + .args(["rustc", "-p", "vchord", "--bin", "pgrx_embed_vchord"]) + .args(["--profile", profile]) + .args(["--target", target]) + .args(["--features", pg_version]) + .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) + .args(["--", "--cfg", "pgrx_embed"]) + .env("PGRX_EMBED", &pgrx_embed) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_status = command.spawn()?.wait()?; + if !command_status.success() { + return Err(format!("Cargo build failed: {command_status}").into()); + } + let mut result = PathBuf::from("./target"); + result.push(target); + result.push(if profile != "dev" { profile } else { "debug" }); + result.push("pgrx_embed_vchord"); + let mut command = Command::new(result); + command.stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_output = command.output()?; + let command_stdout = String::from_utf8(command_output.stdout)?.replace("\t", " "); + Ok(command_stdout) +} - {pushes} +fn install_by_copying( + src: impl AsRef, + dst: impl AsRef, + #[cfg_attr(not(target_family = "unix"), expect(unused_variables))] is_executable: bool, +) -> Result<(), Box> { + std::fs::copy(src, &dst)?; + #[cfg(target_family = "unix")] + { + use std::fs::Permissions; + use std::os::unix::fs::PermissionsExt; + let perm = Permissions::from_mode(if !is_executable { 0o644 } else { 0o755 }); + std::fs::set_permissions(dst, perm)?; + } + Ok(()) +} - let pgrx_sql = ::pgrx::pgrx_sql_entity_graph::PgrxSql::build( - entities.into_iter(), - "vchord".to_string(), - false, - ) - .expect("SQL generation error"); +fn install_by_writing( + contents: impl AsRef<[u8]>, + dst: impl AsRef, + #[cfg_attr(not(target_family = "unix"), expect(unused_variables))] is_executable: bool, +) -> Result<(), Box> { + std::fs::write(&dst, contents)?; + #[cfg(target_family = "unix")] + { + use std::fs::Permissions; + use std::os::unix::fs::PermissionsExt; + let perm = Permissions::from_mode(if !is_executable { 0o644 } else { 0o755 }); + std::fs::set_permissions(dst, perm)?; + } + Ok(()) +} - pgrx_sql - .write(&mut std::io::stdout()) - .expect("Could not write SQL to stdout"); - }} - "# - ); - let mut file = tempfile::NamedTempFile::new()?; - file.write_all(code.as_bytes())?; - let pgrx_embed_path = file.into_temp_path(); - let mut command = Command::new("cargo"); - command - .args(["rustc"]) - .args(["-p", "vchord", "--bin", "pgrx_embed_vchord"]) - .args(["--features", fork.as_str()]) - .args(["--", "--cfg", "pgrx_embed"]) - .env("PGRX_EMBED", &pgrx_embed_path) - .env("PGRX_PG_CONFIG_PATH", &path) - .stdout(Stdio::inherit()) - .stderr(Stdio::inherit()); - let debug = format!("{command:?}"); - let status = command.spawn()?.wait()?; - if !status.success() { - bail!("Cargo build failed: {debug}"); - } - let mut command = Command::new("./target/debug/pgrx_embed_vchord"); - command.stderr(Stdio::inherit()); - let command_output = command.output()?; - let command_stdout = String::from_utf8(command_output.stdout)?.replace("\t", " "); - Ok(command_stdout) - }; - let uilts_install = |permission: &str, src: &str, dst: &str| { - let mut command = Command::new("install"); - command.args(["-m", permission, src, dst]); - let debug = format!("{command:?}"); - let status = command.spawn()?.wait()?; - if !status.success() { - bail!("Command execution failed: {debug}"); - } - Ok(()) - }; - let dll_suffix = if OS == "macos" && matches!(fork.as_str(), "pg13" | "pg14" | "pg15") { - ".so" +fn main() -> Result<(), Box> { + let cli = Cli::parse(); + if !std::fs::exists("vchord.control")? { + return Err("The script must be run from the VectorChord source directory.".into()); + } + let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { + eprintln!("Environment variable `PGRX_PG_CONFIG_PATH`: {value:#?}"); + PathBuf::from(value) } else { - DLL_SUFFIX + return Err("Environment variable `PGRX_PG_CONFIG_PATH` is not set.".into()); }; - let install = |pkglibdir, sharedir_extension| -> Result<()> { - uilts_install( - "755", - &format!("./target/release/{DLL_PREFIX}vchord{DLL_SUFFIX}"), - &format!("{pkglibdir}/vchord{dll_suffix}"), - )?; - uilts_install( - "644", - "./vchord.control", - &format!("{sharedir_extension}/vchord.control"), - )?; - if version != "0.0.0" { - for maybe_entry in read_dir("./sql/upgrade")? { - let path = maybe_entry?.path(); - let name = path.file_name().context("broken assets")?; - uilts_install( - "644", - &format!("{}", path.display()), - &format!("{sharedir_extension}/{}", name.display()), - )?; + let pg_config = pg_config(&path)?; + let pg_version = { + let version = pg_config["VERSION"].clone(); + if let Some(prefix_stripped) = version.strip_prefix("PostgreSQL ") { + if let Some((stripped, _)) = prefix_stripped.split_once(|c: char| !c.is_ascii_digit()) { + format!("pg{stripped}",) + } else { + format!("pg{prefix_stripped}",) } - uilts_install( - "644", - &format!("./sql/install/vchord--{version}.sql"), - &format!("{sharedir_extension}/vchord--{version}.sql"), - )?; } else { - let contents = schema()?; - let mut file = tempfile::NamedTempFile::new()?; - file.write_all(contents.as_bytes())?; - let path = file.into_temp_path(); - uilts_install( - "644", - &format!("{}", path.display()), - &format!("{sharedir_extension}/vchord--0.0.0.sql"), - )?; + return Err("PostgreSQL version is invalid.".into()); } - Ok(()) }; + let vchord_version = control_file("./vchord.control")?["default_version"].clone(); match cli.command { - Commands::Build(BuildArgs { output }) => { - build()?; + Commands::Build(BuildArgs { + output, + profile, + target, + }) => { + let cfgs = cfgs(&target)?; + let obj = build(&path, &pg_version, &cfgs, &profile, &target)?; let pkglibdir = format!("{output}/pkglibdir"); let sharedir = format!("{output}/sharedir"); let sharedir_extension = format!("{sharedir}/extension"); @@ -262,7 +277,50 @@ fn main() -> Result<()> { std::fs::create_dir_all(&pkglibdir)?; std::fs::create_dir_all(&sharedir)?; std::fs::create_dir_all(&sharedir_extension)?; - install(pkglibdir, sharedir_extension)?; + let is_unix = cfgs.contains("target_family=\"unix\""); + let is_macos = cfgs.contains("target_os=\"macos\""); + let is_windows = cfgs.contains("target_family=\"windows\""); + let is_wasm = cfgs.contains("target_family=\"wasm\""); + let suffix = if is_unix + && (!is_macos || matches!(pg_version.as_str(), "pg13" | "pg14" | "pg15")) + { + ".so" + } else if is_macos { + ".dylib" + } else if is_windows { + ".dll" + } else if is_wasm { + ".wasm" + } else { + return Err("unknown operating system".into()); + }; + install_by_copying(&obj, format!("{pkglibdir}/vchord{suffix}"), true)?; + install_by_copying( + "vchord.control", + format!("{sharedir}/extension/vchord.control"), + false, + )?; + if vchord_version != "0.0.0" { + for e in read_dir("./sql/upgrade")?.collect::, _>>()? { + install_by_copying( + e.path(), + format!("{sharedir}/extension/{}", e.file_name().display()), + false, + )?; + } + install_by_copying( + format!("./sql/install/vchord--{vchord_version}.sql"), + format!("{sharedir}/extension/vchord--{vchord_version}.sql"), + false, + )?; + } else { + let exports = parse(&cfgs, obj)?; + install_by_writing( + generate(&path, &pg_version, &profile, &target, exports)?, + format!("{sharedir_extension}/vchord--0.0.0.sql"), + false, + )?; + } } } Ok(()) diff --git a/crates/rabitq/Cargo.toml b/crates/rabitq/Cargo.toml index 6dff2f62..0d358764 100644 --- a/crates/rabitq/Cargo.toml +++ b/crates/rabitq/Cargo.toml @@ -5,7 +5,6 @@ edition.workspace = true publish = false [dependencies] -distance = { path = "../distance" } simd = { path = "../simd" } zerocopy.workspace = true diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 32bc9779..bdb7eafe 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -8,7 +8,6 @@ publish = false simd_macros = { path = "../simd_macros" } half.workspace = true -paste.workspace = true seq-macro.workspace = true zerocopy.workspace = true @@ -16,7 +15,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.22" +cc = "1.2.27" [lints] workspace = true diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index a7680168..cf6cf1a0 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -96,8 +96,10 @@ pub fn emulate_mm_reduce_add_ps(mut x: std::arch::x86_64::__m128) -> f32 { #[crate::target_cpu(enable = "v4")] pub fn emulate_mm512_reduce_add_epi16(x: std::arch::x86_64::__m512i) -> i16 { use std::arch::x86_64::*; - _mm256_reduce_add_epi16(_mm512_castsi512_si256(x)) - + _mm256_reduce_add_epi16(_mm512_extracti32x8_epi32(x, 1)) + i16::wrapping_add( + _mm256_reduce_add_epi16(_mm512_castsi512_si256(x)), + _mm256_reduce_add_epi16(_mm512_extracti32x8_epi32(x, 1)), + ) } #[inline] @@ -109,7 +111,7 @@ pub fn emulate_mm256_reduce_add_epi16(mut x: std::arch::x86_64::__m256i) -> i16 x = _mm256_hadd_epi16(x, x); x = _mm256_hadd_epi16(x, x); let x = _mm256_cvtsi256_si32(x); - (x as i16) + ((x >> 16) as i16) + i16::wrapping_add(x as i16, (x >> 16) as i16) } #[inline] @@ -120,7 +122,7 @@ pub fn emulate_mm_reduce_add_epi16(mut x: std::arch::x86_64::__m128i) -> i16 { x = _mm_hadd_epi16(x, x); x = _mm_hadd_epi16(x, x); let x = _mm_cvtsi128_si32(x); - (x as i16) + ((x >> 16) as i16) + i16::wrapping_add(x as i16, (x >> 16) as i16) } #[inline] diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index e7490d2f..51b55435 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -12,7 +12,6 @@ // // Copyright (c) 2025 TensorChord Inc. -#![allow(clippy::just_underscores_and_digits)] mod reduce_sum_of_x_as_u32_y_as_u32 { #[inline] #[cfg(target_arch = "x86_64")] diff --git a/crates/vchordg/Cargo.toml b/crates/vchordg/Cargo.toml index a31d3068..0ea18542 100644 --- a/crates/vchordg/Cargo.toml +++ b/crates/vchordg/Cargo.toml @@ -8,15 +8,12 @@ publish = false algo = { path = "../algo" } always_equal = { path = "../always_equal" } distance = { path = "../distance" } -k_means = { path = "../k_means" } rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } half.workspace = true min-max-heap = "1.3.0" -paste.workspace = true -pin-project = "1" rand.workspace = true serde.workspace = true validator.workspace = true diff --git a/crates/vchordg/src/build.rs b/crates/vchordg/src/build.rs index 571575e9..f838bfc7 100644 --- a/crates/vchordg/src/build.rs +++ b/crates/vchordg/src/build.rs @@ -28,9 +28,7 @@ pub fn build( let mut meta_guard = index.extend( Opaque { next: u32::MAX, - skip: u32::MAX, link: 1, - _padding_0: Default::default(), }, false, ); @@ -38,9 +36,7 @@ pub fn build( let vertex_guard = index.extend( Opaque { next: u32::MAX, - skip: 1, link: 2, - _padding_0: Default::default(), }, true, ); @@ -49,9 +45,7 @@ pub fn build( let vector_guard = index.extend( Opaque { next: u32::MAX, - skip: u32::MAX, link: u32::MAX, - _padding_0: Default::default(), }, false, ); @@ -64,6 +58,7 @@ pub fn build( ef_construction: index_options.ef_construction, beam_construction: index_options.beam_construction, start: OptionPointer::NONE, + skip: 1, }); let i = meta_guard .alloc(&serialized) diff --git a/crates/vchordg/src/bulkdelete.rs b/crates/vchordg/src/bulkdelete.rs index ab84ce3f..ab3f921f 100644 --- a/crates/vchordg/src/bulkdelete.rs +++ b/crates/vchordg/src/bulkdelete.rs @@ -14,7 +14,7 @@ use crate::Opaque; use crate::operator::Operator; -use crate::tuples::{VertexTuple, WithReader, WithWriter}; +use crate::tuples::{MetaTuple, VertexTuple, WithReader, WithWriter}; use algo::{Page, RelationLength, RelationRead, RelationWrite}; use std::num::NonZero; @@ -26,19 +26,26 @@ pub fn bulkdelete R::Page: Page, { let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let start = meta_tuple.start(); let link = meta_guard.get_opaque().link; drop(meta_guard); + let Some(_) = start.into_inner() else { + return; + }; let mut current = link; while current != u32::MAX { check(); let read = index.read(current); let flag = 'flag: { for i in 1..=read.len() { - let bytes = read.get(i).expect("data corruption"); - let tuple = VertexTuple::deserialize_ref(bytes); - let p = tuple.payload(); - if Some(true) == p.map(&callback) { - break 'flag true; + if let Some(bytes) = read.get(i) { + let tuple = VertexTuple::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + break 'flag true; + } } } false @@ -47,11 +54,12 @@ pub fn bulkdelete drop(read); let mut write = index.write(current, false); for i in 1..=write.len() { - let bytes = write.get_mut(i).expect("data corruption"); - let mut tuple = VertexTuple::deserialize_mut(bytes); - let p = tuple.payload(); - if Some(true) == p.map(&callback) { - *p = None; + if let Some(bytes) = write.get_mut(i) { + let mut tuple = VertexTuple::deserialize_mut(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + *p = None; + } } } current = write.get_opaque().next; diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index 7d3539cd..753b8ad9 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -17,14 +17,14 @@ use crate::candidates::Candidates; use crate::operator::{CloneAccessor, Operator, Vector}; use crate::results::Results; use crate::tuples::*; +use crate::vectors::{by_prefetch, by_read, copy_all, copy_nothing, copy_outs, update}; use crate::visited::Visited; +use algo::accessor::LAccess; use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use algo::{Bump, Page, PageGuard, RelationRead, RelationWrite}; use always_equal::AlwaysEqual; -use distance::Distance; use std::cmp::Reverse; use std::collections::VecDeque; -use std::iter::{repeat, zip}; use std::num::{NonZero, Wrapping}; use vector::{VectorBorrowed, VectorOwned}; @@ -48,15 +48,79 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let alpha = meta_tuple.alpha(); let ef = meta_tuple.ef_construction(); let beam = meta_tuple.beam_construction(); + let skip = meta_tuple.skip(); drop(meta_guard); - let (vector_pointers, vertex_pointer) = insert_tuples::(index, vector, payload, m); + let version_t = Wrapping(rand::random()); + let (pointers_t, t) = { + let list_of_vector_bytes = { + let (left, right) = O::Vector::split(vector, m as _); + let left = left.into_iter().enumerate().map(|(index, elements)| { + VectorTuple::serialize(&VectorTuple::::_1 { + payload: Some(payload), + elements: elements.to_vec(), + index: index as u32, + }) + }); + let right = VectorTuple::serialize(&VectorTuple::::_0 { + payload: Some(payload), + elements: right.0.to_vec(), + metadata: right.1, + neighbours: vec![OptionNeighbour::NONE; m as usize], + version: version_t, + }); + left.chain(std::iter::once(right)).collect::>() + }; + let mut vertex_bytes = { + let code = O::Vector::code(vector); + VertexTuple::serialize(&VertexTuple { + metadata: code.0.into_array(), + payload: Some(payload), + elements: rabitq::original::binary::pack_code(&code.1), + pointers: vec![Pointer::new((u32::MAX, 0)); list_of_vector_bytes.len()], // a sentinel value + }) + }; + let mut vertex_guard = if let Some(guard) = index.search(vertex_bytes.len()) { + guard + } else { + append_vertex_tuple(index, skip, vertex_bytes.len()) + }; + let link = vertex_guard.get_opaque().link; + let pointers_t = list_of_vector_bytes + .iter() + .map(|vector_bytes| { + let mut vector_guard = append_vector_tuple(index, link, vector_bytes.len()); + let i = vector_guard + .alloc(vector_bytes) + .expect("implementation: a free page cannot accommodate a single tuple"); + Pointer::new((vector_guard.id(), i)) + }) + .collect::>(); + { + let mut vertex_tuple = VertexTuple::deserialize_mut(&mut vertex_bytes); + vertex_tuple.pointers().copy_from_slice(&pointers_t); + } + let t = { + let i = vertex_guard + .alloc(&vertex_bytes) + .expect("implementation: a free page cannot accommodate a single tuple"); + (vertex_guard.id(), i) + }; + drop(vertex_guard); + if skip < t.0 { + let mut meta_guard = index.write(0, false); + let meta_bytes = meta_guard.get_mut(1).expect("data corruption"); + let mut meta_tuple = MetaTuple::deserialize_mut(meta_bytes); + *meta_tuple.skip() = (*meta_tuple.skip()).max(t.0); + } + (pointers_t, t) + }; let start = if start.into_inner().is_none() { let mut meta_guard = index.write(0, false); let meta_bytes = meta_guard.get_mut(1).expect("data corruption"); let mut meta_tuple = MetaTuple::deserialize_mut(meta_bytes); let dst = meta_tuple.start(); if dst.into_inner().is_none() { - *dst = OptionPointer::some(vertex_pointer); + *dst = OptionPointer::some(t); return; } else { *dst @@ -84,13 +148,20 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( } let mut iter = std::iter::from_fn(|| { while let Some(((_, AlwaysEqual((pointers_u, u))), guards)) = candidates.pop() { - let Ok((dis_u, _payload_u, outs_u)) = - rerank::(guards, pointers_u, vector, &mut visited) - else { + let Ok((dis_u, outs_u, _, _)) = crate::vectors::read::( + by_prefetch::(guards, pointers_u.iter().copied()), + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + copy_outs, + ) else { // the link is broken continue; }; - let mut iterator = prefetch_vertices.prefetch(outs_u); + let mut iterator = prefetch_vertices.prefetch( + outs_u + .into_iter() + .filter(|&x| !visited.contains(x)) + .collect::>(), + ); while let Some((v, guards)) = iterator.next() { visited.insert(v); let vertex_guard = { @@ -124,226 +195,105 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( break; } } - let trace = results - .into_inner() - .0 - .into_iter() - .map(|(dis_v, x)| (Reverse(dis_v), x)) - .collect::>(); - let outs = crate::prune::robust_prune( - |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), - (bump.alloc_slice(&vector_pointers), vertex_pointer), - trace.into_iter().flat_map(|item| { - use algo::accessor::Accessor1; - let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; - let m = strict_sub(pointers_u.len(), 1); - let mut accessor = CloneAccessor::::default(); - for i in 0..m { - let vector_guard = index.read(pointers_u[i].into_inner().0); - let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { - // the link is broken - return None; - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_1(segment) = vector_tuple else { - // the link is broken - return None; - }; - if segment.index() as usize != i { - // the link is broken - return None; - } - accessor.push(segment.elements()); - } - let vector_u; - { - let vector_guard = index.read(pointers_u[m].into_inner().0); - let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { - // the link is broken - return None; - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_0(segment) = vector_tuple else { + let trace = { + let (left, right) = results.into_inner(); + left.into_iter() + .map(|(dis_v, v)| (Reverse(dis_v), v)) + .chain(right) + .flat_map(|item| { + let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; + let Ok((vector_u, _, _, _)) = crate::vectors::read::( + by_read(index, pointers_u.iter().copied()), + CloneAccessor::::default(), + copy_nothing, + ) else { // the link is broken return None; }; - accessor.push(segment.elements()); - vector_u = accessor.finish(*segment.metadata()); - } - Some(((Reverse(dis_u), AlwaysEqual((pointers_u, u))), vector_u)) - }), + Some(((Reverse(dis_u), AlwaysEqual((pointers_u, u))), vector_u)) + }) + .collect::>() + }; + let outs = crate::prune::robust_prune( + |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), + (bump.alloc_slice(&pointers_t), t), + trace.into_iter(), m, alpha, |(_, u)| *u, ); - { - let mut vector_guard = index.write(vector_pointers.last().unwrap().into_inner().0, false); - let Some(vector_bytes) = - vector_guard.get_mut(vector_pointers.last().unwrap().into_inner().1) - else { - // the link is broken - return; - }; - let vector_tuple = VectorTuple::::deserialize_mut(vector_bytes); - let VectorTupleWriter::_0(mut vector_tuple) = vector_tuple else { - // the link is broken - return; - }; - let iterator = outs - .iter() - .map(|&(Reverse(dis_v), AlwaysEqual((_, v)))| OptionNeighbour::some(v, dis_v)) - .chain(repeat(OptionNeighbour::NONE)); - for (hole, fill) in zip(vector_tuple.neighbours().iter_mut(), iterator) { - *hole = fill; - } - } + let _ = update::( + (index, pointers_t.as_slice()), + (version_t, VecDeque::new()), + outs.iter() + .map(|&(Reverse(dis_u), AlwaysEqual((_, u)))| (u, dis_u)), + ); for (Reverse(dis_t), AlwaysEqual((pointers_u, u))) in outs { - while add_link::( - index, - (pointers_u, u), - (vertex_pointer, dis_t), - m, - alpha, - bump, - ) == Ok(false) - {} - } -} - -fn rerank<'r, R: RelationRead, O: Operator>( - mut guards: impl Iterator>, - pointers_u: &mut [Pointer], - vector: ::Borrowed<'_>, - visited: &mut Visited, -) -> Result<(Distance, Option>, VecDeque<(u32, u16)>), ()> { - use algo::accessor::{Accessor1, LAccess}; - let m = strict_sub(pointers_u.len(), 1); - let mut accessor = LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()); - for i in 0..m { - let vector_guard = guards.next().expect("internal"); - let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { - // the link is broken - return Err(()); - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_1(segment) = vector_tuple else { - // the link is broken - return Err(()); - }; - if segment.index() as usize != i { - // the link is broken - return Err(()); + 'occ: loop { + let Ok((neighbours_u, _, version)) = + crate::vectors::read_without_accessor::((index, pointers_u), copy_all) + else { + // the link is broken + break 'occ; + }; + let trace = neighbours_u + .iter() + .copied() + .chain(std::iter::once((t, dis_t))) + .map(|(v, dis_v)| (Reverse(dis_v), AlwaysEqual(v))) + .flat_map(|item| { + let (Reverse(dis_u), AlwaysEqual(u)) = item; + let vertex_guard = index.read(u.0); + let Some(vertex_bytes) = vertex_guard.get(u.1) else { + // the link is broken + return None; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_u = vertex_tuple.pointers().to_vec(); + Some((Reverse(dis_u), AlwaysEqual((pointers_u, u)))) + }) + .flat_map(|item| { + let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; + let Ok((vector_u, _, _, _)) = crate::vectors::read::( + by_read(index, pointers_u.iter().copied()), + CloneAccessor::::default(), + copy_nothing, + ) else { + // the link is broken + return None; + }; + Some(((Reverse(dis_u), AlwaysEqual((pointers_u, u))), vector_u)) + }) + .collect::>(); + let outs = crate::prune::robust_prune( + |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), + (pointers_u.to_vec(), u), + trace.into_iter(), + m, + alpha, + |(_, u)| *u, + ); + if update::( + (index, pointers_u), + (version, neighbours_u), + outs.iter() + .map(|&(Reverse(dis_u), AlwaysEqual((_, u)))| (u, dis_u)), + ) != Ok(false) + { + break 'occ; + } } - accessor.push(segment.elements()); - } - let dis_u; - let payload_u; - let neighbours_u; - { - let vector_guard = guards.next().expect("internal"); - let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { - // the link is broken - return Err(()); - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_0(segment) = vector_tuple else { - // the link is broken - return Err(()); - }; - accessor.push(segment.elements()); - dis_u = accessor.finish(*segment.metadata()); - payload_u = segment.payload(); - neighbours_u = segment - .neighbours() - .iter() - .flat_map(|neighbour| neighbour.into_inner()) - .map(|(v, _)| v) - .filter(|&v| !visited.contains(v)) - .collect::>(); } - Ok((dis_u, payload_u, neighbours_u)) } -fn insert_tuples<'b, R: RelationRead + RelationWrite, O: Operator>( - index: &'b R, - vector: ::Borrowed<'b>, - payload: NonZero, - m: u32, -) -> (Vec, (u32, u16)) -where - R::Page: Page, -{ - let vector_bytes = { - let (left, right) = O::Vector::split(vector, m as _); - let left = left.into_iter().enumerate().map(|(index, elements)| { - VectorTuple::serialize(&VectorTuple::::_1 { - payload: Some(payload), - elements: elements.to_vec(), - index: index as u32, - }) - }); - let right = VectorTuple::serialize(&VectorTuple::::_0 { - payload: Some(payload), - elements: right.0.to_vec(), - metadata: right.1, - neighbours: vec![OptionNeighbour::NONE; m as usize], - version: Wrapping(rand::random()), - }); - left.chain(std::iter::once(right)).collect::>() - }; - let mut vertex_bytes = { - let code = O::Vector::code(vector); - VertexTuple::serialize(&VertexTuple { - metadata: code.0.into_array(), - payload: Some(payload), - elements: rabitq::original::binary::pack_code(&code.1), - pointers: vec![Pointer::new((u32::MAX, 0)); vector_bytes.len()], // a sentinel value - }) - }; - append_vertex_tuple(index, 1, vertex_bytes.len(), |mut vertex_guard| { - let vector_pointers = vector_bytes - .iter() - .map(|vector_bytes| { - append_vector_tuple( - index, - vertex_guard.get_opaque().link, - vector_bytes.len(), - |mut vector_guard| { - let i = vector_guard.alloc(vector_bytes).expect( - "implementation: a free page cannot accommodate a single tuple", - ); - Pointer::new((vector_guard.id(), i)) - }, - ) - }) - .collect::>(); - { - let mut vertex_tuple = VertexTuple::deserialize_mut(&mut vertex_bytes); - vertex_tuple.pointers().copy_from_slice(&vector_pointers); - } - let vertex_pointer = { - let i = vertex_guard - .alloc(&vertex_bytes) - .expect("implementation: a free page cannot accommodate a single tuple"); - (vertex_guard.id(), i) - }; - (vector_pointers, vertex_pointer) - }) -} - -#[allow(clippy::collapsible_else_if)] -fn append_vertex_tuple<'r, R: RelationRead + RelationWrite, T>( +fn append_vertex_tuple<'r, R: RelationRead + RelationWrite>( index: &'r R, first: u32, size: usize, - f: impl FnOnce(R::WriteGuard<'r>) -> T, -) -> T +) -> R::WriteGuard<'r> where R::Page: Page, { - if let Some(guard) = index.search(size) { - return f(guard); - } assert!(first != u32::MAX); let mut current = first; loop { @@ -352,16 +302,14 @@ where drop(read); let mut write = index.write(current, true); if write.freespace() as usize >= size { - return f(write); + return write; } if write.get_opaque().next == u32::MAX { let link = index .extend( Opaque { next: u32::MAX, - skip: u32::MAX, link: u32::MAX, - _padding_0: Default::default(), }, false, ) @@ -369,44 +317,29 @@ where let extend = index.extend( Opaque { next: u32::MAX, - skip: u32::MAX, link, - _padding_0: Default::default(), }, true, ); { write }.get_opaque_mut().next = extend.id(); if extend.freespace() as usize >= size { - let fresh = extend.id(); - let result = f(extend); - let mut past = index.write(first, true); - past.get_opaque_mut().skip = fresh.max(past.get_opaque().skip); - return result; + return extend; } else { panic!("implementation: a clear page cannot accommodate a single tuple"); } } - if current == first && write.get_opaque().skip != first { - current = write.get_opaque().skip; - } else { - current = write.get_opaque().next; - } + current = write.get_opaque().next; } else { - if current == first && read.get_opaque().skip != first { - current = read.get_opaque().skip; - } else { - current = read.get_opaque().next; - } + current = read.get_opaque().next; } } } -fn append_vector_tuple<'r, R: RelationRead + RelationWrite, T>( +fn append_vector_tuple<'r, R: RelationRead + RelationWrite>( index: &'r R, first: u32, size: usize, - f: impl FnOnce(R::WriteGuard<'r>) -> T, -) -> T +) -> R::WriteGuard<'r> where R::Page: Page, { @@ -418,21 +351,19 @@ where drop(read); let mut write = index.write(current, false); if write.freespace() as usize >= size { - return f(write); + return write; } if write.get_opaque().next == u32::MAX { let extend = index.extend( Opaque { next: u32::MAX, - skip: u32::MAX, link: u32::MAX, - _padding_0: Default::default(), }, false, ); { write }.get_opaque_mut().next = extend.id(); if extend.freespace() as usize >= size { - return f(extend); + return extend; } else { panic!("implementation: a clear page cannot accommodate a single tuple"); } @@ -443,135 +374,3 @@ where } } } - -fn add_link( - index: &R, - (pointers_u, u): (&mut [Pointer], (u32, u16)), - new: ((u32, u16), Distance), - m: u32, - alpha: f32, - bump: &impl Bump, -) -> Result { - let vector_guard = index.read(pointers_u.last().unwrap().into_inner().0); - let Some(vector_bytes) = vector_guard.get(pointers_u.last().unwrap().into_inner().1) else { - // the link is broken - return Err(()); - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_0(vector_tuple) = vector_tuple else { - // the link is broken - return Err(()); - }; - let check = vector_tuple - .neighbours() - .iter() - .flat_map(|neighbour| neighbour.into_inner()) - .map(|(v, dis_v)| (Reverse(dis_v), AlwaysEqual(v))) - .collect::>(); - let trace = check - .iter() - .copied() - .chain(std::iter::once((Reverse(new.1), AlwaysEqual(new.0)))) - .collect::>(); - let version = vector_tuple.version(); - drop(vector_guard); - let outs = crate::prune::robust_prune( - |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), - (bump.alloc_slice(pointers_u), u), - trace - .into_iter() - .flat_map(|item| { - let (Reverse(dis_u), AlwaysEqual(u)) = item; - let vertex_guard = index.read(u.0); - let Some(vertex_bytes) = vertex_guard.get(u.1) else { - // the link is broken - return None; - }; - let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); - let pointers_u = bump.alloc_slice(vertex_tuple.pointers()); - Some((Reverse(dis_u), AlwaysEqual((pointers_u, u)))) - }) - .flat_map(|item| { - use algo::accessor::Accessor1; - let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; - let m = strict_sub(pointers_u.len(), 1); - let mut accessor = CloneAccessor::::default(); - for i in 0..m { - let vector_guard = index.read(pointers_u[i].into_inner().0); - let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { - // the link is broken - return None; - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_1(segment) = vector_tuple else { - // the link is broken - return None; - }; - if segment.index() as usize != i { - // the link is broken - return None; - } - accessor.push(segment.elements()); - } - let vector_u; - { - let vector_guard = index.read(pointers_u[m].into_inner().0); - let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { - // the link is broken - return None; - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_0(segment) = vector_tuple else { - // the link is broken - return None; - }; - accessor.push(segment.elements()); - vector_u = accessor.finish(*segment.metadata()); - } - Some(((Reverse(dis_u), AlwaysEqual((pointers_u, u))), vector_u)) - }), - m, - alpha, - |(_, u)| *u, - ); - // fast path - if outs - .iter() - .map(|&(x, AlwaysEqual((_, v)))| (x, AlwaysEqual(v))) - .eq(check) - { - return Ok(true); - } - let mut vector_guard = index.write(pointers_u.last().unwrap().into_inner().0, false); - let Some(vector_bytes) = vector_guard.get_mut(pointers_u.last().unwrap().into_inner().1) else { - // the link is broken - return Err(()); - }; - let vector_tuple = VectorTuple::::deserialize_mut(vector_bytes); - let VectorTupleWriter::_0(mut vector_tuple) = vector_tuple else { - // the link is broken - return Err(()); - }; - if *vector_tuple.version() != version { - return Ok(false); - } else { - *vector_tuple.version() += 1; - } - let filling = outs - .iter() - .map(|&(Reverse(dis_v), AlwaysEqual((_, v)))| OptionNeighbour::some(v, dis_v)) - .chain(repeat(OptionNeighbour::NONE)); - for (hole, fill) in zip(vector_tuple.neighbours().iter_mut(), filling) { - *hole = fill; - } - Ok(true) -} - -// Emulate unstable library feature `strict_overflow_ops`. -// See https://github.com/rust-lang/rust/issues/118260. - -#[inline] -pub const fn strict_sub(lhs: usize, rhs: usize) -> usize { - let (a, b) = lhs.overflowing_sub(rhs); - if b { panic!() } else { a } -} diff --git a/crates/vchordg/src/lib.rs b/crates/vchordg/src/lib.rs index ad0f767e..142905af 100644 --- a/crates/vchordg/src/lib.rs +++ b/crates/vchordg/src/lib.rs @@ -1,5 +1,3 @@ -#![allow(clippy::type_complexity)] - mod build; mod bulkdelete; mod candidates; @@ -10,12 +8,12 @@ mod prune; mod results; mod search; mod tuples; +mod vectors; mod visited; pub mod operator; pub mod types; -use algo::tuples::Padding; pub use build::build; pub use bulkdelete::bulkdelete; pub use insert::insert; @@ -29,9 +27,7 @@ use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] pub struct Opaque { pub next: u32, - pub skip: u32, pub link: u32, - pub _padding_0: [Padding; 4], } #[allow(unsafe_code)] diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs index 61799850..7b8eeabc 100644 --- a/crates/vchordg/src/maintain.rs +++ b/crates/vchordg/src/maintain.rs @@ -14,24 +14,15 @@ use crate::Opaque; use crate::operator::{CloneAccessor, Operator}; -use crate::tuples::{ - MetaTuple, OptionNeighbour, Pointer, VectorTuple, VectorTupleReader, VectorTupleWriter, - VertexTuple, WithReader, WithWriter, -}; -use algo::{Bump, Page, PageGuard, RelationRead, RelationWrite}; +use crate::tuples::{MetaTuple, VectorTuple, VertexTuple, WithReader}; +use crate::vectors::{by_read, copy_all, copy_nothing, copy_outs, update}; +use algo::{Page, PageGuard, RelationRead, RelationWrite}; use always_equal::AlwaysEqual; -use distance::Distance; use std::cmp::Reverse; -use std::collections::VecDeque; -use std::iter::{repeat, zip}; -use std::num::Wrapping; use vector::VectorOwned; -pub fn maintain( - index: &R, - check: impl Fn(), - bump: &impl Bump, -) where +pub fn maintain(index: &R, check: impl Fn()) +where R::Page: Page, { let meta_guard = index.read(0); @@ -39,9 +30,12 @@ pub fn maintain( let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let m = meta_tuple.m(); let alpha = meta_tuple.alpha(); - let start = meta_tuple.start().into_inner(); + let start = meta_tuple.start(); let link = meta_guard.get_opaque().link; drop(meta_guard); + let Some(s) = start.into_inner() else { + return; + }; // do it's best to remove broken edges { let mut current = link; @@ -53,14 +47,99 @@ pub fn maintain( if let Some(vertex_bytes) = vertex_guard.get(i) { let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); if vertex_tuple.payload().is_some() { - let pointers = bump.alloc_slice(vertex_tuple.pointers()); - members.push((pointers, (vertex_guard.id(), i))); + let pointers_u = vertex_tuple.pointers().to_vec(); + members.push((pointers_u, (vertex_guard.id(), i))); } } } let next = { vertex_guard }.get_opaque().next; - for member in members { - while fix_link::(index, (member.0, member.1), m, alpha, bump) == Ok(false) {} + for (pointers_u, u) in members { + 'occ: loop { + let Ok((vector_u, neighbours_u, _, version)) = crate::vectors::read::( + by_read::(index, pointers_u.iter().copied()), + CloneAccessor::::default(), + copy_all, + ) else { + // the link is broken + break 'occ; + }; + let trace = { + let mut trace = Vec::new(); + let mut extend = Vec::new(); + for &(v, dis_v) in neighbours_u.iter() { + let vertex_guard = index.read(v.0); + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v = vertex_tuple.pointers().to_vec(); + let payload_v = vertex_tuple.payload(); + drop(vertex_guard); + let Ok((vector_v, outs_v, _, _)) = crate::vectors::read::( + by_read::(index, pointers_v.iter().copied()), + CloneAccessor::::default(), + copy_outs, + ) else { + // the link is broken + continue; + }; + if payload_v.is_some() { + trace.push(( + (Reverse(dis_v), AlwaysEqual((pointers_v, v))), + vector_v, + )); + } else { + extend.extend(outs_v.iter().copied()); + } + } + for v in extend { + let vertex_guard = index.read(v.0); + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v = vertex_tuple.pointers().to_vec(); + let payload_v = vertex_tuple.payload(); + drop(vertex_guard); + let Ok((vector_v, _, _, _)) = crate::vectors::read::( + by_read::(index, pointers_v.iter().copied()), + CloneAccessor::::default(), + copy_nothing, + ) else { + // the link is broken + continue; + }; + if payload_v.is_some() { + let dis_v = + O::distance(vector_u.as_borrowed(), vector_v.as_borrowed()); + trace.push(( + (Reverse(dis_v), AlwaysEqual((pointers_v, v))), + vector_v, + )); + } + } + trace + }; + let outs = crate::prune::robust_prune( + |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), + (pointers_u.to_vec(), u), + trace.into_iter(), + m, + alpha, + |(_, u)| *u, + ); + if update::( + (index, pointers_u.as_slice()), + (version, neighbours_u), + outs.iter() + .map(|&(Reverse(dis_u), AlwaysEqual((_, u)))| (u, dis_u)), + ) != Ok(false) + { + break 'occ; + } + } } current = next; } @@ -75,7 +154,7 @@ pub fn maintain( if let Some(bytes) = vertex_guard.get(i) { let tuple = VertexTuple::deserialize_ref(bytes); let p = tuple.payload(); - if p.is_none() && Some((current, i)) != start { + if p.is_none() && (current, i) != s { vertex_guard.free(i); } }; @@ -85,7 +164,7 @@ pub fn maintain( let mut current = { vertex_guard }.get_opaque().link; while current != u32::MAX { check(); - let mut vector_guard = index.write(current, true); + let mut vector_guard = index.write(current, false); for i in 1..=vector_guard.len() { if let Some(bytes) = vector_guard.get(i) { use crate::tuples::VectorTupleReader; @@ -93,13 +172,13 @@ pub fn maintain( match tuple { VectorTupleReader::_0(tuple) => { let p = tuple.payload(); - if p.is_none() && Some((current, i)) != start { + if p.is_none() && (current, i) != s { vector_guard.free(i); } } VectorTupleReader::_1(tuple) => { let p = tuple.payload(); - if p.is_none() && Some((current, i)) != start { + if p.is_none() && (current, i) != s { vector_guard.free(i); } } @@ -112,162 +191,3 @@ pub fn maintain( } } } - -fn read_tuple( - index: &R, - pointers_u: &mut [Pointer], -) -> Result<(O::Vector, VecDeque<((u32, u16), Distance)>, Wrapping), ()> { - use algo::accessor::Accessor1; - let m = strict_sub(pointers_u.len(), 1); - let mut accessor = CloneAccessor::::default(); - for i in 0..m { - let vector_guard = index.read(pointers_u[i].into_inner().0); - let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { - // the link is broken - return Err(()); - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_1(segment) = vector_tuple else { - // the link is broken - return Err(()); - }; - if segment.index() as usize != i { - // the link is broken - return Err(()); - } - accessor.push(segment.elements()); - } - let vector_u; - let neighbours_u; - let version; - { - let vector_guard = index.read(pointers_u[m].into_inner().0); - let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { - // the link is broken - return Err(()); - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_0(segment) = vector_tuple else { - // the link is broken - return Err(()); - }; - accessor.push(segment.elements()); - vector_u = accessor.finish(*segment.metadata()); - neighbours_u = segment - .neighbours() - .iter() - .flat_map(|neighbour| neighbour.into_inner()) - .collect::>(); - version = segment.version(); - } - Ok((vector_u, neighbours_u, version)) -} - -fn fix_link( - index: &R, - (pointers_u, u): (&mut [Pointer], (u32, u16)), - m: u32, - alpha: f32, - bump: &impl Bump, -) -> Result { - let Ok((vector_u, neighbours_u, version)) = read_tuple::(index, pointers_u) else { - // the link is broken - return Err(()); - }; - let mut trace = Vec::new(); - let mut extend = Vec::new(); - for &(v, dis_v) in neighbours_u.iter() { - let vertex_guard = index.read(v.0); - let Some(vertex_bytes) = vertex_guard.get(v.1) else { - // the link is broken - continue; - }; - let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); - let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); - let payload_v = vertex_tuple.payload(); - drop(vertex_guard); - let Ok((vector_v, neighbours_v, _)) = read_tuple::(index, pointers_v) else { - // the link is broken - continue; - }; - if payload_v.is_some() { - trace.push((v, dis_v, vector_v)); - } else { - let outs_v = neighbours_v.iter().map(|(v, _)| *v).collect::>(); - extend.extend(outs_v); - } - } - // fast path - if extend.is_empty() { - return Ok(true); - } - for v in extend { - let vertex_guard = index.read(v.0); - let Some(vertex_bytes) = vertex_guard.get(v.1) else { - // the link is broken - continue; - }; - let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); - let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); - let payload_v = vertex_tuple.payload(); - drop(vertex_guard); - let Ok((vector_v, _, _)) = read_tuple::(index, pointers_v) else { - // the link is broken - continue; - }; - if payload_v.is_some() { - let dis_v = O::distance(vector_u.as_borrowed(), vector_v.as_borrowed()); - trace.push((v, dis_v, vector_v)); - } - } - let outs = crate::prune::robust_prune( - |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), - (bump.alloc_slice(pointers_u), u), - trace.into_iter().flat_map(|item| { - let (u, dis_u, vector_u) = item; - let vertex_guard = index.read(u.0); - let Some(vertex_bytes) = vertex_guard.get(u.1) else { - // the link is broken - return None; - }; - let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); - let pointer_u = bump.alloc_slice(vertex_tuple.pointers()); - Some(((Reverse(dis_u), AlwaysEqual((pointer_u, u))), vector_u)) - }), - m, - alpha, - |(_, u)| *u, - ); - let mut vector_guard = index.write(pointers_u.last().unwrap().into_inner().0, false); - let Some(vector_bytes) = vector_guard.get_mut(pointers_u.last().unwrap().into_inner().1) else { - // the link is broken - return Err(()); - }; - let vector_tuple = VectorTuple::::deserialize_mut(vector_bytes); - let VectorTupleWriter::_0(mut vector_tuple) = vector_tuple else { - // the link is broken - return Err(()); - }; - if *vector_tuple.version() != version { - return Ok(false); - } else { - *vector_tuple.version() += 1; - } - let iterator = outs - .iter() - .map(|&(Reverse(dis_v), AlwaysEqual((_, v)))| OptionNeighbour::some(v, dis_v)) - .chain(repeat(OptionNeighbour::NONE)); - for (hole, fill) in zip(vector_tuple.neighbours().iter_mut(), iterator) { - *hole = fill; - } - Ok(true) -} - -// Emulate unstable library feature `strict_overflow_ops`. -// See https://github.com/rust-lang/rust/issues/118260. - -#[inline] -pub const fn strict_sub(lhs: usize, rhs: usize) -> usize { - let (a, b) = lhs.overflowing_sub(rhs); - if b { panic!() } else { a } -} diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs index 8d59246a..540101d1 100644 --- a/crates/vchordg/src/search.rs +++ b/crates/vchordg/src/search.rs @@ -17,7 +17,9 @@ use crate::candidates::Candidates; use crate::operator::{Operator, Vector}; use crate::results::Results; use crate::tuples::*; +use crate::vectors::{by_prefetch, copy_outs}; use crate::visited::Visited; +use algo::accessor::LAccess; use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use algo::{Bump, Page, RelationRead}; use always_equal::AlwaysEqual; @@ -39,7 +41,6 @@ pub fn search<'b, R: RelationRead, O: Operator>( where R::Page: Page, { - let _ = bump; let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); @@ -69,13 +70,20 @@ where } let mut iter = std::iter::from_fn(move || { while let Some(((_, AlwaysEqual(pointers_u)), guards)) = candidates.pop() { - let Ok((dis_u, payload_u, outs_u)) = - rerank::(guards, pointers_u, vector, &mut visited) - else { + let Ok((dis_u, outs_u, payload_u, _)) = crate::vectors::read::( + by_prefetch::(guards, pointers_u.iter().copied()), + LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + copy_outs, + ) else { // the link is broken continue; }; - let mut iterator = prefetch_vertices.prefetch(outs_u); + let mut iterator = prefetch_vertices.prefetch( + outs_u + .into_iter() + .filter(|&x| !visited.contains(x)) + .collect::>(), + ); while let Some((v, guards)) = iterator.next() { visited.insert(v); let vertex_guard = { @@ -114,66 +122,3 @@ where }); Box::new(search.filter_map(|(dis_u, AlwaysEqual(payload_u))| Some((dis_u, payload_u?)))) } - -fn rerank<'r, R: RelationRead, O: Operator>( - mut guards: impl Iterator>, - pointers_u: &mut [Pointer], - vector: ::Borrowed<'_>, - visited: &mut Visited, -) -> Result<(Distance, Option>, VecDeque<(u32, u16)>), ()> { - use algo::accessor::{Accessor1, LAccess}; - let m = strict_sub(pointers_u.len(), 1); - let mut accessor = LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()); - for i in 0..m { - let vector_guard = guards.next().expect("internal"); - let Some(vector_bytes) = vector_guard.get(pointers_u[i].into_inner().1) else { - // the link is broken - return Err(()); - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_1(segment) = vector_tuple else { - // the link is broken - return Err(()); - }; - if segment.index() as usize != i { - // the link is broken - return Err(()); - } - accessor.push(segment.elements()); - } - let dis_u; - let payload_u; - let neighbours_u; - { - let vector_guard = guards.next().expect("internal"); - let Some(vector_bytes) = vector_guard.get(pointers_u[m].into_inner().1) else { - // the link is broken - return Err(()); - }; - let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); - let VectorTupleReader::_0(segment) = vector_tuple else { - // the link is broken - return Err(()); - }; - accessor.push(segment.elements()); - dis_u = accessor.finish(*segment.metadata()); - payload_u = segment.payload(); - neighbours_u = segment - .neighbours() - .iter() - .flat_map(|neighbour| neighbour.into_inner()) - .map(|(v, _)| v) - .filter(|&v| !visited.contains(v)) - .collect::>(); - } - Ok((dis_u, payload_u, neighbours_u)) -} - -// Emulate unstable library feature `strict_overflow_ops`. -// See https://github.com/rust-lang/rust/issues/118260. - -#[inline] -pub const fn strict_sub(lhs: usize, rhs: usize) -> usize { - let (a, b) = lhs.overflowing_sub(rhs); - if b { panic!() } else { a } -} diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 2d82a130..47f5d640 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -46,8 +46,8 @@ struct MetaTupleHeader { alpha: f32, ef_construction: u32, beam_construction: u32, - _padding_0: [Padding; 4], start: OptionPointer, + skip: u32, } pub struct MetaTuple { @@ -57,6 +57,7 @@ pub struct MetaTuple { pub ef_construction: u32, pub beam_construction: u32, pub start: OptionPointer, + pub skip: u32, } impl Tuple for MetaTuple { @@ -71,6 +72,7 @@ impl Tuple for MetaTuple { ef_construction, beam_construction, start, + skip, } => { buffer.extend((MAGIC as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); @@ -84,7 +86,7 @@ impl Tuple for MetaTuple { ef_construction: *ef_construction, beam_construction: *beam_construction, start: *start, - _padding_0: Default::default(), + skip: *skip, } .as_bytes(), ); @@ -136,6 +138,9 @@ impl<'a> MetaTupleReader<'a> { pub fn start(self) -> OptionPointer { self.header.start } + pub fn skip(self) -> u32 { + self.header.skip + } } impl WithWriter for MetaTuple { @@ -166,6 +171,9 @@ impl<'a> MetaTupleWriter<'a> { pub fn start(&mut self) -> &mut OptionPointer { &mut self.header.start } + pub fn skip(&mut self) -> &mut u32 { + &mut self.header.skip + } } #[repr(C, align(8))] diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index 04e01ddd..44ae10ee 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -42,10 +42,10 @@ impl VamanaIndexOptions { 1.0 } fn default_ef_construction() -> u32 { - 500 + 64 } fn default_beam_construction() -> u32 { - 1 + 8 } } @@ -111,8 +111,7 @@ pub struct VectorOptions { impl VectorOptions { pub fn validate_self(&self) -> Result<(), ValidationError> { match (self.v, self.d, self.dims) { - (VectorKind::Vecf32, DistanceKind::L2S, 1..=60000) => Ok(()), - (VectorKind::Vecf16, DistanceKind::L2S, 1..=60000) => Ok(()), + (_, _, 1..=60000) => Ok(()), _ => Err(ValidationError::new("invalid vector options")), } } diff --git a/crates/vchordg/src/vectors.rs b/crates/vchordg/src/vectors.rs new file mode 100644 index 00000000..4af30d88 --- /dev/null +++ b/crates/vchordg/src/vectors.rs @@ -0,0 +1,168 @@ +use crate::operator::{Operator, Vector}; +use crate::tuples::{ + OptionNeighbour, Pointer, VectorTuple, VectorTupleReader, VectorTupleWriter, WithReader, + WithWriter, +}; +use algo::accessor::Accessor1; +use algo::{Page, RelationRead, RelationWrite}; +use distance::Distance; +use std::collections::VecDeque; +use std::num::{NonZero, Wrapping}; + +pub fn by_prefetch<'r, R: RelationRead>( + guards: impl ExactSizeIterator>, + pointers_u: impl ExactSizeIterator, +) -> impl ExactSizeIterator, u16)> { + assert!(guards.len() == pointers_u.len() && pointers_u.len() > 0); + pointers_u + .map(Pointer::into_inner) + .zip(guards) + .map(|(pointer, guard)| (guard, pointer.1)) +} + +pub fn by_read<'r, R: RelationRead>( + index: &'r R, + pointers_u: impl ExactSizeIterator, +) -> impl ExactSizeIterator, u16)> { + assert!(pointers_u.len() > 0); + pointers_u + .map(Pointer::into_inner) + .map(|pointer| (pointer, index.read(pointer.0))) + .map(|(pointer, guard)| (guard, pointer.1)) +} + +pub fn copy_nothing(_: &[OptionNeighbour]) {} + +pub fn copy_outs(x: &[OptionNeighbour]) -> VecDeque<(u32, u16)> { + x.iter() + .flat_map(|neighbour| neighbour.into_inner()) + .map(|(v, _)| v) + .collect::>() +} + +pub fn copy_all(x: &[OptionNeighbour]) -> VecDeque<((u32, u16), Distance)> { + x.iter() + .flat_map(|neighbour| neighbour.into_inner()) + .collect::>() +} + +pub fn read< + 'r, + R: RelationRead, + O: Operator, + A: Accessor1<::Element, ::Metadata>, + Output, +>( + mut iterator: impl ExactSizeIterator, u16)>, + accessor: A, + copy: impl FnOnce(&[OptionNeighbour]) -> Output, +) -> Result<(A::Output, Output, Option>, Wrapping), ()> { + let m = strict_sub(iterator.len(), 1); + let mut result = accessor; + for index in 0..m { + let (vector_guard, i) = iterator.next().expect("internal: bad size"); + let Some(vector_bytes) = vector_guard.get(i) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_1(vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + if vector_tuple.index() as usize != index { + // the link is broken + return Err(()); + } + result.push(vector_tuple.elements()); + } + let value_u; + let payload_u; + let neighbours_u; + let version_u; + { + let (vector_guard, i) = iterator.next().expect("internal: bad size"); + let Some(vector_bytes) = vector_guard.get(i) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + result.push(vector_tuple.elements()); + value_u = result.finish(*vector_tuple.metadata()); + neighbours_u = copy(vector_tuple.neighbours()); + payload_u = vector_tuple.payload(); + version_u = vector_tuple.version(); + } + Ok((value_u, neighbours_u, payload_u, version_u)) +} + +pub fn read_without_accessor( + (index, pointers_u): (&R, &[Pointer]), + copy: impl FnOnce(&[OptionNeighbour]) -> Output, +) -> Result<(Output, Option>, Wrapping), ()> { + let payload_u; + let neighbours_u; + let version_u; + { + let (id, i) = pointers_u.last().expect("internal: bad size").into_inner(); + let vector_guard = index.read(id); + let Some(vector_bytes) = vector_guard.get(i) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + neighbours_u = copy(vector_tuple.neighbours()); + payload_u = vector_tuple.payload(); + version_u = vector_tuple.version(); + } + Ok((neighbours_u, payload_u, version_u)) +} + +pub fn update( + (index, pointers_u): (&R, &[Pointer]), + (version, neighbours_u): (Wrapping, VecDeque<((u32, u16), Distance)>), + outs: impl Iterator + Clone, +) -> Result { + if outs.clone().eq(neighbours_u) { + return Ok(true); + } + let mut vector_guard = index.write(pointers_u.last().unwrap().into_inner().0, false); + let Some(vector_bytes) = vector_guard.get_mut(pointers_u.last().unwrap().into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_mut(vector_bytes); + let VectorTupleWriter::_0(mut vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + if *vector_tuple.version() != version { + return Ok(false); + } else { + *vector_tuple.version() += 1; + } + let filling = outs + .map(|(v, dis_v)| OptionNeighbour::some(v, dis_v)) + .chain(std::iter::repeat(OptionNeighbour::NONE)); + for (hole, fill) in std::iter::zip(vector_tuple.neighbours().iter_mut(), filling) { + *hole = fill; + } + Ok(true) +} + +// Emulate unstable library feature `strict_overflow_ops`. +// See https://github.com/rust-lang/rust/issues/118260. + +#[inline] +const fn strict_sub(lhs: usize, rhs: usize) -> usize { + let (a, b) = lhs.overflowing_sub(rhs); + if b { panic!() } else { a } +} diff --git a/crates/vchordrq/Cargo.toml b/crates/vchordrq/Cargo.toml index d65cae4b..a860c744 100644 --- a/crates/vchordrq/Cargo.toml +++ b/crates/vchordrq/Cargo.toml @@ -8,18 +8,18 @@ publish = false algo = { path = "../algo" } always_equal = { path = "../always_equal" } distance = { path = "../distance" } -k_means = { path = "../k_means" } rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } half.workspace = true -paste.workspace = true pin-project = "1" -rand.workspace = true serde.workspace = true validator.workspace = true zerocopy.workspace = true +[dev-dependencies] +rand.workspace = true + [lints] workspace = true diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index 8e7dbd46..040bb3a9 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -12,8 +12,6 @@ // // Copyright (c) 2025 TensorChord Inc. -#![allow(clippy::type_complexity)] - mod build; mod bulkdelete; mod cache; diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index 0ab91141..16624aff 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -30,7 +30,6 @@ use vector::{VectorBorrowed, VectorOwned}; type Extra<'b> = &'b mut (NonZero, u16, &'b mut [u32]); -#[allow(clippy::too_many_arguments)] pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( index: &'r R, vector: ::Borrowed<'_>, @@ -156,7 +155,6 @@ where results.into_vec() } -#[allow(clippy::too_many_arguments)] pub fn maxsim_search<'r, 'b: 'r, R: RelationRead, O: Operator>( index: &'r R, vector: ::Borrowed<'_>, diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index fe3b6fdb..a09033bc 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -437,7 +437,6 @@ struct DirectoryTupleHeader1 { _padding_0: [Padding; 4], } -#[allow(clippy::large_enum_variant)] #[derive(Debug, Clone, PartialEq)] pub enum DirectoryTuple { _0 { elements: Vec }, diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index 941b76c2..3edd066c 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -95,10 +95,7 @@ pub struct VectorOptions { impl VectorOptions { pub fn validate_self(&self) -> Result<(), ValidationError> { match (self.v, self.d, self.dims) { - (VectorKind::Vecf32, DistanceKind::L2S, 1..=60000) => Ok(()), - (VectorKind::Vecf32, DistanceKind::Dot, 1..=60000) => Ok(()), - (VectorKind::Vecf16, DistanceKind::L2S, 1..=60000) => Ok(()), - (VectorKind::Vecf16, DistanceKind::Dot, 1..=60000) => Ok(()), + (_, _, 1..=60000) => Ok(()), _ => Err(ValidationError::new("invalid vector options")), } } diff --git a/crates/vector/Cargo.toml b/crates/vector/Cargo.toml index d5dba499..41d7d173 100644 --- a/crates/vector/Cargo.toml +++ b/crates/vector/Cargo.toml @@ -8,7 +8,5 @@ publish = false distance = { path = "../distance" } simd = { path = "../simd" } -half.workspace = true - [lints] workspace = true diff --git a/docker/Dockerfile b/docker/Dockerfile index 48576eec..1a7966c0 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -11,4 +11,4 @@ RUN echo ${PG_VERSION} COPY ./build/postgresql-${PG_VERSION}-vchord_${SEMVER}-1_${TARGETARCH}.deb /tmp/vchord.deb RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb -CMD ["postgres", "-c" ,"shared_preload_libraries=vchord.so"] +CMD ["postgres", "-c" ,"shared_preload_libraries=vchord"] diff --git a/src/bin/pgrx_embed.rs b/src/bin/pgrx_embed.rs index c7b6102c..505d1b69 100644 --- a/src/bin/pgrx_embed.rs +++ b/src/bin/pgrx_embed.rs @@ -13,4 +13,34 @@ // Copyright (c) 2025 TensorChord Inc. #![allow(unsafe_code)] +#![allow(unused_crate_dependencies)] + ::pgrx::pgrx_embed!(); + +#[macro_export] +macro_rules! schema_generation { + ($($symbol:ident)*) => { + pub fn main() -> Result<(), Box> { + extern crate vchord as _; + + use ::pgrx::pgrx_sql_entity_graph::ControlFile; + use ::pgrx::pgrx_sql_entity_graph::PgrxSql; + use ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity; + + let p = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/vchord.control")); + let control_file = ControlFile::try_from(p)?; + + unsafe extern "Rust" { + $(safe fn $symbol() -> SqlGraphEntity;)* + } + + let mut e = vec![SqlGraphEntity::ExtensionRoot(control_file)]; + $(e.push($symbol());)* + + let pgrx_sql = PgrxSql::build(e.into_iter(), "vchord".to_string(), false)?; + pgrx_sql.write(&mut std::io::stdout())?; + + Ok(()) + } + }; +} diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 080dbe84..5029f462 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -38,7 +38,7 @@ pub enum PostgresIo { static VCHORDG_EF_SEARCH: GucSetting = GucSetting::::new(64); -static VCHORDG_BEAM_SIZE: GucSetting = GucSetting::::new(1); +static VCHORDG_BEAM_SIZE: GucSetting = GucSetting::::new(8); static VCHORDG_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); diff --git a/src/index/storage.rs b/src/index/storage.rs index 0cc09fd4..2409c96a 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -96,7 +96,8 @@ impl Page for PostgresPage { pgrx::pg_sys::LP_DEAD => unimplemented!(), _ => unreachable!(), } - assert!(offset_of!(PageHeaderData, pd_linp) <= lp_off && lp_off <= size_of::()); + assert!(offset_of!(PageHeaderData, pd_linp) <= lp_off); + assert!(lp_off <= size_of::()); assert!(lp_len <= size_of::()); assert!(lp_off + lp_len <= size_of::()); unsafe { @@ -118,7 +119,15 @@ impl Page for PostgresPage { let iid = unsafe { self.header.pd_linp.as_ptr().add((i - 1) as _).read() }; let lp_off = iid.lp_off() as usize; let lp_len = iid.lp_len() as usize; - assert!(offset_of!(PageHeaderData, pd_linp) <= lp_off && lp_off <= size_of::()); + match iid.lp_flags() { + pgrx::pg_sys::LP_UNUSED => return None, + pgrx::pg_sys::LP_NORMAL => (), + pgrx::pg_sys::LP_REDIRECT => unimplemented!(), + pgrx::pg_sys::LP_DEAD => unimplemented!(), + _ => unreachable!(), + } + assert!(offset_of!(PageHeaderData, pd_linp) <= lp_off); + assert!(lp_off <= size_of::()); assert!(lp_len <= size_of::()); assert!(lp_off + lp_len <= size_of::()); unsafe { @@ -294,9 +303,8 @@ impl RelationWrite for PostgresRelation { assert!(id != u32::MAX, "no such page"); unsafe { use pgrx::pg_sys::{ - BUFFER_LOCK_EXCLUSIVE, ForkNumber, GENERIC_XLOG_FULL_IMAGE, - GenericXLogRegisterBuffer, GenericXLogStart, LockBuffer, ReadBufferExtended, - ReadBufferMode, + BUFFER_LOCK_EXCLUSIVE, ForkNumber, GenericXLogRegisterBuffer, GenericXLogStart, + LockBuffer, ReadBufferExtended, ReadBufferMode, }; let buf = ReadBufferExtended( self.raw, @@ -308,8 +316,7 @@ impl RelationWrite for PostgresRelation { LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); let state = GenericXLogStart(self.raw); let page = NonNull::new( - GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) - .cast::>>(), + GenericXLogRegisterBuffer(state, buf, 0).cast::>>(), ) .expect("failed to get page"); PostgresBufferWriteGuard { @@ -462,10 +469,6 @@ pub struct PostgresReadStream { pub struct PostgresReadStreamGuards { #[cfg(any(feature = "pg17", feature = "pg18"))] raw: *mut pgrx::pg_sys::ReadStream, - #[cfg_attr( - any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"), - expect(dead_code) - )] list: L, _phantom: PhantomData (O, I)>, } @@ -498,6 +501,15 @@ where }) } } + + fn size_hint(&self) -> (usize, Option) { + self.list.size_hint() + } +} + +impl ExactSizeIterator for PostgresReadStreamGuards where + L: ExactSizeIterator +{ } #[cfg(any(feature = "pg17", feature = "pg18"))] diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs index f12a6a40..5f9fd2ba 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/algo.rs @@ -75,19 +75,18 @@ where R: RelationRead + RelationWrite, R::Page: Page, { - let bump = bumpalo::Bump::new(); match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { - vchordg::maintain::<_, Op, L2S>>(index, &check, &bump); + vchordg::maintain::<_, Op, L2S>>(index, &check); } (VectorKind::Vecf16, DistanceKind::L2S) => { - vchordg::maintain::<_, Op, L2S>>(index, &check, &bump); + vchordg::maintain::<_, Op, L2S>>(index, &check); } (VectorKind::Vecf32, DistanceKind::Dot) => { - vchordg::maintain::<_, Op, Dot>>(index, &check, &bump); + vchordg::maintain::<_, Op, Dot>>(index, &check); } (VectorKind::Vecf16, DistanceKind::Dot) => { - vchordg::maintain::<_, Op, Dot>>(index, &check, &bump); + vchordg::maintain::<_, Op, Dot>>(index, &check); } } } diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index 76bd4bff..c598d3fe 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -148,7 +148,6 @@ pub unsafe extern "C-unwind" fn amoptions( rdopts as *mut pgrx::pg_sys::bytea } -#[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn amcostestimate( _root: *mut pgrx::pg_sys::PlannerInfo, @@ -170,7 +169,6 @@ pub unsafe extern "C-unwind" fn amcostestimate( } #[cfg(feature = "pg13")] -#[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn aminsert( index_relation: pgrx::pg_sys::Relation, @@ -191,7 +189,6 @@ pub unsafe extern "C-unwind" fn aminsert( feature = "pg17", feature = "pg18" ))] -#[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn aminsert( index_relation: pgrx::pg_sys::Relation, diff --git a/src/index/vchordg/opclass.rs b/src/index/vchordg/opclass.rs index 6a83dd3f..a216265c 100644 --- a/src/index/vchordg/opclass.rs +++ b/src/index/vchordg/opclass.rs @@ -36,13 +36,19 @@ pub enum Opfamily { impl Opfamily { fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { use {BorrowedVector as B, OwnedVector as O}; - match (vector, self) { - (B::Vecf32(x), Self::VectorL2) => O::Vecf32(x.own()), - (B::Vecf32(x), Self::VectorCosine) => O::Vecf32(x.function_normalize()), - (B::Vecf32(_), _) => unreachable!(), - (B::Vecf16(x), Self::HalfvecL2) => O::Vecf16(x.own()), - (B::Vecf16(x), Self::HalfvecCosine) => O::Vecf16(x.function_normalize()), - (B::Vecf16(_), _) => unreachable!(), + match (self, vector) { + (Self::VectorL2, B::Vecf32(x)) => O::Vecf32(x.own()), + (Self::VectorL2, _) => unreachable!(), + (Self::VectorCosine, B::Vecf32(x)) => O::Vecf32(x.function_normalize()), + (Self::VectorCosine, _) => unreachable!(), + (Self::VectorIp, B::Vecf32(x)) => O::Vecf32(x.own()), + (Self::VectorIp, _) => unreachable!(), + (Self::HalfvecL2, B::Vecf16(x)) => O::Vecf16(x.own()), + (Self::HalfvecL2, _) => unreachable!(), + (Self::HalfvecCosine, B::Vecf16(x)) => O::Vecf16(x.function_normalize()), + (Self::HalfvecCosine, _) => unreachable!(), + (Self::HalfvecIp, B::Vecf16(x)) => O::Vecf16(x.own()), + (Self::HalfvecIp, _) => unreachable!(), } } pub unsafe fn store(self, datum: Datum) -> Option> { diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 4ca9dbff..13039ffe 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -148,7 +148,6 @@ pub unsafe extern "C-unwind" fn amoptions( rdopts as *mut pgrx::pg_sys::bytea } -#[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn amcostestimate( root: *mut pgrx::pg_sys::PlannerInfo, @@ -258,7 +257,6 @@ pub unsafe extern "C-unwind" fn amcostestimate( } #[cfg(feature = "pg13")] -#[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn aminsert( index_relation: pgrx::pg_sys::Relation, @@ -279,7 +277,6 @@ pub unsafe extern "C-unwind" fn aminsert( feature = "pg17", feature = "pg18" ))] -#[allow(clippy::too_many_arguments)] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn aminsert( index_relation: pgrx::pg_sys::Relation, diff --git a/src/lib.rs b/src/lib.rs index 64912e7e..1cfd9147 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,13 +13,34 @@ // Copyright (c) 2025 TensorChord Inc. #![allow(unsafe_code)] -#![allow(clippy::type_complexity)] +#![allow(unused_crate_dependencies)] mod datatype; mod index; mod upgrade; -pgrx::pg_module_magic!(); +pgrx::pg_module_magic!( + name = c"vchord", + version = { + const RAW: &str = env!("VCHORD_VERSION"); + const BUFFER: [u8; RAW.len() + 1] = { + let mut buffer = [0u8; RAW.len() + 1]; + let mut i = 0_usize; + while i < RAW.len() { + buffer[i] = RAW.as_bytes()[i]; + i += 1; + } + buffer + }; + const STR: &::core::ffi::CStr = + if let Ok(s) = ::core::ffi::CStr::from_bytes_with_nul(&BUFFER) { + s + } else { + panic!("there are null characters in VCHORD_VERSION") + }; + const { STR } + } +); pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap); pgrx::extension_sql_file!("./sql/finalize.sql", finalize); @@ -37,6 +58,9 @@ extern "C-unwind" fn _PG_init() { } } +#[cfg(not(panic = "unwind"))] +compile_error!("This crate must be compiled with `-Cpanic=unwind`."); + #[cfg(not(target_endian = "little"))] compile_error!("Target architecture is not supported."); diff --git a/tests/README.md b/tests/README.md deleted file mode 100644 index d99133fb..00000000 --- a/tests/README.md +++ /dev/null @@ -1,14 +0,0 @@ -# SQL logic test - -Requires `sqllogictest` to be installed: - -```bash -cargo install sqllogictest-bin -``` - -To run all the tests: - -```bash -PGPASSWORD=postgres psql -h localhost -U postgres -d postgres -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' -sqllogictest './tests/**/*.slt' -``` diff --git a/tests/general/null.fail b/tests/fail/null.fail similarity index 100% rename from tests/general/null.fail rename to tests/fail/null.fail diff --git a/tests/general/issue427.slt b/tests/general/issue_427.slt similarity index 84% rename from tests/general/issue427.slt rename to tests/general/issue_427.slt index de79cb87..2a06a7d4 100644 --- a/tests/general/issue427.slt +++ b/tests/general/issue_427.slt @@ -6,6 +6,9 @@ CREATE TABLE t (val vector(3)); statement ok INSERT INTO t (val) SELECT NULL::vector FROM generate_series(1, 1000); +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + statement ok CREATE INDEX ON t USING vchordrq (val vector_l2_ops); diff --git a/tests/vchordg/partition.slt b/tests/vchordg/partition.slt new file mode 100644 index 00000000..d42ceeb8 --- /dev/null +++ b/tests/vchordg/partition.slt @@ -0,0 +1,51 @@ +# partition table +statement ok +CREATE TABLE t (val vector(3), category_id int) PARTITION BY LIST(category_id); + +statement ok +CREATE TABLE id_123 PARTITION OF t FOR VALUES IN (1, 2, 3); + +statement ok +CREATE TABLE id_456 PARTITION OF t FOR VALUES IN (4, 5, 6); + +statement ok +CREATE TABLE id_789 PARTITION OF t FOR VALUES IN (7, 8, 9); + +statement ok +INSERT INTO t (val, category_id) +SELECT + ARRAY[random(), random(), random()]::real[], + (random() * 6 + 1)::int +FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING vchordg (val public.vector_l2_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement error the relation "t_val_idx" is not an index +select vchordg_prewarm('t_val_idx'); + +statement ok +CREATE INDEX ON id_123 USING vchordg (val vector_cosine_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +# partial index +statement ok +CREATE INDEX ON t USING vchordg (val public.vector_ip_ops) WHERE (category_id = 1); + +query I +SELECT COUNT(1) FROM +(SELECT 1 FROM t WHERE (category_id = 1) ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE id_789, id_456, id_123, t; diff --git a/tests/vchordg/reindex.slt b/tests/vchordg/reindex.slt new file mode 100644 index 00000000..77da3a0b --- /dev/null +++ b/tests/vchordg/reindex.slt @@ -0,0 +1,33 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 700); + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +statement ok +REINDEX INDEX t_val_idx; + +statement ok +INSERT INTO t (val) VALUES ('[0.6,0.6,0.6]'); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +REINDEX INDEX CONCURRENTLY t_val_idx; + +statement ok +INSERT INTO t (val) VALUES ('[0.6,0.6,0.6]'); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE t; diff --git a/tests/general/external_build.slt b/tests/vchordrq/external_build.slt similarity index 100% rename from tests/general/external_build.slt rename to tests/vchordrq/external_build.slt diff --git a/tests/general/filter_rerank_in_index.slt b/tests/vchordrq/filter_rerank_in_index.slt similarity index 100% rename from tests/general/filter_rerank_in_index.slt rename to tests/vchordrq/filter_rerank_in_index.slt diff --git a/tests/general/filter_rerank_in_table.slt b/tests/vchordrq/filter_rerank_in_table.slt similarity index 100% rename from tests/general/filter_rerank_in_table.slt rename to tests/vchordrq/filter_rerank_in_table.slt diff --git a/tests/general/index.slt b/tests/vchordrq/index.slt similarity index 100% rename from tests/general/index.slt rename to tests/vchordrq/index.slt diff --git a/tests/general/multivector.slt b/tests/vchordrq/multivector.slt similarity index 100% rename from tests/general/multivector.slt rename to tests/vchordrq/multivector.slt diff --git a/tests/general/partition.slt b/tests/vchordrq/partition.slt similarity index 100% rename from tests/general/partition.slt rename to tests/vchordrq/partition.slt diff --git a/tests/pg17/filter_rerank_in_index.slt b/tests/vchordrq/pg17/filter_rerank_in_index.slt similarity index 100% rename from tests/pg17/filter_rerank_in_index.slt rename to tests/vchordrq/pg17/filter_rerank_in_index.slt diff --git a/tests/pg17/filter_rerank_in_table.slt b/tests/vchordrq/pg17/filter_rerank_in_table.slt similarity index 100% rename from tests/pg17/filter_rerank_in_table.slt rename to tests/vchordrq/pg17/filter_rerank_in_table.slt diff --git a/tests/general/pin.slt b/tests/vchordrq/pin.slt similarity index 100% rename from tests/general/pin.slt rename to tests/vchordrq/pin.slt diff --git a/tests/general/pushdown_plan.slt b/tests/vchordrq/pushdown_plan.slt similarity index 100% rename from tests/general/pushdown_plan.slt rename to tests/vchordrq/pushdown_plan.slt diff --git a/tests/general/pushdown_range.slt b/tests/vchordrq/pushdown_range.slt similarity index 100% rename from tests/general/pushdown_range.slt rename to tests/vchordrq/pushdown_range.slt diff --git a/tests/general/reindex.slt b/tests/vchordrq/reindex.slt similarity index 100% rename from tests/general/reindex.slt rename to tests/vchordrq/reindex.slt diff --git a/tests/general/rerank_in_index.slt b/tests/vchordrq/rerank_in_index.slt similarity index 100% rename from tests/general/rerank_in_index.slt rename to tests/vchordrq/rerank_in_index.slt diff --git a/tests/general/rerank_in_table.slt b/tests/vchordrq/rerank_in_table.slt similarity index 100% rename from tests/general/rerank_in_table.slt rename to tests/vchordrq/rerank_in_table.slt From c010c36be18aa23ac61eda89c08c296940871e4a Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 10 Jul 2025 17:05:41 +0800 Subject: [PATCH 174/324] feat: improve prune (#285) disables `beam` by default job: +psql_macos job: +psql_windows --------- Signed-off-by: usamoi --- .cargo/config.toml | 3 - Cargo.lock | 107 ++++++++++---------- build.rs | 28 +++--- crates/make/Cargo.toml | 5 +- crates/make/src/main.rs | 165 +++++++++++++++++++------------ crates/simd/Cargo.toml | 1 + crates/simd/build.rs | 32 ++++++ crates/vchordg/src/build.rs | 2 +- crates/vchordg/src/insert.rs | 13 ++- crates/vchordg/src/maintain.rs | 8 +- crates/vchordg/src/operator.rs | 11 +++ crates/vchordg/src/prune.rs | 35 +++++-- crates/vchordg/src/tuples.rs | 12 +-- crates/vchordg/src/types.rs | 12 +-- src/index/gucs.rs | 6 +- src/index/vchordg/am/am_build.rs | 3 - 16 files changed, 276 insertions(+), 167 deletions(-) diff --git a/.cargo/config.toml b/.cargo/config.toml index 63524241..b3550a8f 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,5 +1,2 @@ [alias] make = "run -p make --" - -[env] -CC = "clang" diff --git a/Cargo.lock b/Cargo.lock index 2fc4e5d7..8d4672c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,12 +2,6 @@ # It is not intended for manual editing. version = 4 -[[package]] -name = "adler2" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" - [[package]] name = "aho-corasick" version = "1.1.3" @@ -284,15 +278,6 @@ dependencies = [ "unicode-segmentation", ] -[[package]] -name = "crc32fast" -version = "1.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3" -dependencies = [ - "cfg-if", -] - [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -412,12 +397,28 @@ dependencies = [ "syn", ] +[[package]] +name = "env_home" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe" + [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.3.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" +dependencies = [ + "libc", + "windows-sys", +] + [[package]] name = "eyre" version = "0.6.12" @@ -434,16 +435,6 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" -[[package]] -name = "flate2" -version = "1.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a3d7db9596fecd151c5f638c0ee5d5bd487b6e0ea232e5dc96d5250f6f94b1d" -dependencies = [ - "crc32fast", - "miniz_oxide", -] - [[package]] name = "fnv" version = "1.0.7" @@ -790,6 +781,12 @@ dependencies = [ "libc", ] +[[package]] +name = "linux-raw-sys" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd945864f07fe9f5371a27ad7b52a172b4b499999f1d97574c9fa68373937e12" + [[package]] name = "litemap" version = "0.8.0" @@ -838,15 +835,6 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" -[[package]] -name = "miniz_oxide" -version = "0.8.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" -dependencies = [ - "adler2", -] - [[package]] name = "nom" version = "7.1.3" @@ -863,12 +851,7 @@ version = "0.37.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03fd943161069e1768b4b3d050890ba48730e590f57e56d4aa04e7e090e61b4a" dependencies = [ - "crc32fast", - "flate2", - "hashbrown", - "indexmap", "memchr", - "ruzstd", "wasmparser", ] @@ -1239,19 +1222,23 @@ dependencies = [ ] [[package]] -name = "rustversion" -version = "1.0.21" +name = "rustix" +version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a0d197bd2c9dc6e53b84da9556a69ba4cdfab8619eb41a8bd1cc2027a0f6b1d" +checksum = "c71e83d6afe7ff64890ec6b71d6a69bb8a610ab78ce364b3352876bb4c801266" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] [[package]] -name = "ruzstd" -version = "0.8.1" +name = "rustversion" +version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3640bec8aad418d7d03c72ea2de10d5c646a598f9883c7babc160d91e3c1b26c" -dependencies = [ - "twox-hash", -] +checksum = "8a0d197bd2c9dc6e53b84da9556a69ba4cdfab8619eb41a8bd1cc2027a0f6b1d" [[package]] name = "ryu" @@ -1352,6 +1339,7 @@ dependencies = [ "rand", "seq-macro", "simd_macros", + "which", "zerocopy", ] @@ -1512,12 +1500,6 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" -[[package]] -name = "twox-hash" -version = "2.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b907da542cbced5261bd3256de1b3a1bf340a3d37f93425a07362a1d687de56" - [[package]] name = "unescape" version = "0.1.0" @@ -1764,6 +1746,17 @@ dependencies = [ "bitflags", ] +[[package]] +name = "which" +version = "8.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fabb953106c3c8eea8306e4393700d7657561cb43122571b172bbfb7c7ba1d" +dependencies = [ + "env_home", + "rustix", + "winsafe", +] + [[package]] name = "winapi" version = "0.3.9" @@ -1941,6 +1934,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "winsafe" +version = "0.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" + [[package]] name = "wit-bindgen-rt" version = "0.39.0" diff --git a/build.rs b/build.rs index d6138350..7917f182 100644 --- a/build.rs +++ b/build.rs @@ -18,6 +18,18 @@ use std::error::Error; use std::process::{Command, Stdio}; fn main() -> Result<(), Box> { + let version = 'version: { + for line in std::fs::read_to_string("./vchord.control")?.lines() { + if let Some(prefix_stripped) = line.strip_prefix("default_version = '") + && let Some(stripped) = prefix_stripped.strip_suffix("'") + { + eprintln!("VectorChord version: {stripped}"); + break 'version stripped.to_string(); + } + } + return Err("VectorChord version is not defined.".into()); + }; + println!("cargo::rustc-env=VCHORD_VERSION={version}"); if var("CARGO_CFG_TARGET_OS")? == "macos" { if let Some(path) = var_os("PGRX_PG_CONFIG_PATH") { let map = { @@ -40,17 +52,9 @@ fn main() -> Result<(), Box> { println!("cargo::rustc-link-arg-cdylib=-Wl,-undefined,dynamic_lookup"); } } - let version = 'version: { - for line in std::fs::read_to_string("./vchord.control")?.lines() { - if let Some(prefix_stripped) = line.strip_prefix("default_version = '") - && let Some(stripped) = prefix_stripped.strip_suffix("'") - { - eprintln!("VectorChord version: {stripped}"); - break 'version stripped.to_string(); - } - } - return Err("VectorChord version is not defined.".into()); - }; - println!("cargo::rustc-env=VCHORD_VERSION={version}"); + if var("CARGO_CFG_TARGET_OS")? == "emscripten" { + println!("cargo::rustc-link-arg-cdylib=-sSIDE_MODULE=2"); + println!("cargo::rustc-link-arg-bins=-sEXPORTED_FUNCTIONS=[_main]"); + } Ok(()) } diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index 781e910f..ec2b10eb 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -6,5 +6,8 @@ publish = false [dependencies] clap = { version = "4.5.38", features = ["derive"] } -object = { version = "0.37.1", features = ["all"] } +object = { version = "0.37.1", default-features = false, features = [ + "read", + "wasm", +] } target-triple = "0.1.4" diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index dec02835..9afdd6de 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use clap::{Args, Parser, Subcommand}; -use object::Object; +use object::{Object, ObjectSymbol}; use std::collections::{HashMap, HashSet}; use std::env::var_os; use std::error::Error; @@ -42,6 +42,68 @@ struct BuildArgs { target: String, } +struct TargetSpecificInformation { + is_macos: bool, + is_windows: bool, + is_emscripten: bool, + is_unix: bool, +} + +impl TargetSpecificInformation { + fn dll_prefix(&self) -> Result<&'static str, Box> { + if self.is_macos { + Ok("lib") + } else if self.is_windows || self.is_emscripten { + Ok("") + } else if self.is_unix { + Ok("lib") + } else { + Err("unknown operating system".into()) + } + } + fn dll_suffix(&self) -> Result<&'static str, Box> { + if self.is_macos { + Ok(".dylib") + } else if self.is_windows { + Ok(".dll") + } else if self.is_emscripten { + Ok(".wasm") + } else if self.is_unix { + Ok(".so") + } else { + Err("unknown operating system".into()) + } + } + fn exe_suffix(&self) -> Result<&'static str, Box> { + if self.is_macos { + Ok("") + } else if self.is_windows { + Ok(".exe") + } else if self.is_emscripten { + Ok(".js") + } else if self.is_unix { + Ok("") + } else { + Err("unknown operating system".into()) + } + } + fn ext_suffix(&self, fork: &str) -> Result<&'static str, Box> { + if self.is_macos { + Ok(if matches!(fork, "pg13" | "pg14" | "pg15") { + ".so" + } else { + ".dylib" + }) + } else if self.is_windows { + Ok(".dll") + } else if self.is_emscripten || self.is_unix { + Ok(".so") + } else { + Err("unknown operating system".into()) + } + } +} + fn pg_config(pg_config: impl AsRef) -> Result, Box> { let mut command = Command::new(pg_config.as_ref()); command.stderr(Stdio::inherit()); @@ -72,7 +134,7 @@ fn control_file(path: impl AsRef) -> Result, Box Result, Box> { +fn target_specific_information(target: &str) -> Result> { let mut command = Command::new("rustc"); command .args(["--print", "cfg"]) @@ -81,17 +143,22 @@ fn cfgs(target: &str) -> Result, Box> { eprintln!("Running {command:?}"); let command_output = command.output()?; let contents = String::from_utf8(command_output.stdout)?; - let mut result = HashSet::new(); + let mut cfgs = HashSet::new(); for line in contents.lines() { - result.insert(line.to_string()); + cfgs.insert(line.to_string()); } - Ok(result) + Ok(TargetSpecificInformation { + is_macos: cfgs.contains("target_os=\"macos\""), + is_unix: cfgs.contains("target_family=\"unix\""), + is_emscripten: cfgs.contains("target_os=\"emscripten\""), + is_windows: cfgs.contains("target_os=\"windows\""), + }) } fn build( pg_config: impl AsRef, pg_version: &str, - cfgs: &HashSet, + tsi: &TargetSpecificInformation, profile: &str, target: &str, ) -> Result> { @@ -112,44 +179,23 @@ fn build( let mut result = PathBuf::from("./target"); result.push(target); result.push(if profile != "dev" { profile } else { "debug" }); - let is_unix = cfgs.contains("target_family=\"unix\""); - let is_macos = cfgs.contains("target_os=\"macos\""); - let is_windows = cfgs.contains("target_family=\"windows\""); - let is_wasm = cfgs.contains("target_family=\"wasm\""); - let prefix = if is_unix { - "lib" - } else if is_windows || is_wasm { - "" - } else { - return Err("unknown operating system".into()); - }; - let suffix = if is_unix && !is_macos { - ".so" - } else if is_macos { - ".dylib" - } else if is_windows { - ".dll" - } else if is_wasm { - ".wasm" - } else { - return Err("unknown operating system".into()); - }; - result.push(format!("{prefix}vchord{suffix}")); + result.push(format!("{}vchord{}", tsi.dll_prefix()?, tsi.dll_suffix()?)); Ok(result) } -fn parse(cfgs: &HashSet, obj: impl AsRef) -> Result, Box> { +fn parse( + tsi: &TargetSpecificInformation, + obj: impl AsRef, +) -> Result, Box> { let obj = obj.as_ref(); eprintln!("Reading {obj:?}"); let contents = std::fs::read(obj)?; let object = object::File::parse(contents.as_slice())?; - let is_macos = cfgs.contains("target_os=\"macos\""); let exports = object - .exports()? - .into_iter() - .flat_map(|x| str::from_utf8(x.name())) + .symbols() + .flat_map(|x| x.name().ok()) .flat_map(|x| { - if !is_macos { + if !tsi.is_macos { Some(x) } else { x.strip_prefix("_") @@ -164,6 +210,7 @@ fn parse(cfgs: &HashSet, obj: impl AsRef) -> Result, B fn generate( pg_config: impl AsRef, pg_version: &str, + tsi: &TargetSpecificInformation, profile: &str, target: &str, exports: Vec, @@ -193,8 +240,14 @@ fn generate( let mut result = PathBuf::from("./target"); result.push(target); result.push(if profile != "dev" { profile } else { "debug" }); - result.push("pgrx_embed_vchord"); - let mut command = Command::new(result); + result.push(format!("pgrx_embed_vchord{}", tsi.exe_suffix()?)); + let mut command; + if !(tsi.is_unix && tsi.is_emscripten) { + command = Command::new(result); + } else { + command = Command::new("node"); + command.arg(result); + } command.stderr(Stdio::inherit()); eprintln!("Running {command:?}"); let command_output = command.output()?; @@ -207,6 +260,7 @@ fn install_by_copying( dst: impl AsRef, #[cfg_attr(not(target_family = "unix"), expect(unused_variables))] is_executable: bool, ) -> Result<(), Box> { + eprintln!("Copying {:?} to {:?}", src.as_ref(), dst.as_ref()); std::fs::copy(src, &dst)?; #[cfg(target_family = "unix")] { @@ -223,6 +277,7 @@ fn install_by_writing( dst: impl AsRef, #[cfg_attr(not(target_family = "unix"), expect(unused_variables))] is_executable: bool, ) -> Result<(), Box> { + eprintln!("Writing {:?}", dst.as_ref()); std::fs::write(&dst, contents)?; #[cfg(target_family = "unix")] { @@ -236,11 +291,10 @@ fn install_by_writing( fn main() -> Result<(), Box> { let cli = Cli::parse(); - if !std::fs::exists("vchord.control")? { + if !std::fs::exists("./vchord.control")? { return Err("The script must be run from the VectorChord source directory.".into()); } let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { - eprintln!("Environment variable `PGRX_PG_CONFIG_PATH`: {value:#?}"); PathBuf::from(value) } else { return Err("Environment variable `PGRX_PG_CONFIG_PATH` is not set.".into()); @@ -265,8 +319,8 @@ fn main() -> Result<(), Box> { profile, target, }) => { - let cfgs = cfgs(&target)?; - let obj = build(&path, &pg_version, &cfgs, &profile, &target)?; + let tsi = target_specific_information(&target)?; + let obj = build(&path, &pg_version, &tsi, &profile, &target)?; let pkglibdir = format!("{output}/pkglibdir"); let sharedir = format!("{output}/sharedir"); let sharedir_extension = format!("{sharedir}/extension"); @@ -277,26 +331,13 @@ fn main() -> Result<(), Box> { std::fs::create_dir_all(&pkglibdir)?; std::fs::create_dir_all(&sharedir)?; std::fs::create_dir_all(&sharedir_extension)?; - let is_unix = cfgs.contains("target_family=\"unix\""); - let is_macos = cfgs.contains("target_os=\"macos\""); - let is_windows = cfgs.contains("target_family=\"windows\""); - let is_wasm = cfgs.contains("target_family=\"wasm\""); - let suffix = if is_unix - && (!is_macos || matches!(pg_version.as_str(), "pg13" | "pg14" | "pg15")) - { - ".so" - } else if is_macos { - ".dylib" - } else if is_windows { - ".dll" - } else if is_wasm { - ".wasm" - } else { - return Err("unknown operating system".into()); - }; - install_by_copying(&obj, format!("{pkglibdir}/vchord{suffix}"), true)?; install_by_copying( - "vchord.control", + &obj, + format!("{pkglibdir}/vchord{}", tsi.ext_suffix(&pg_version)?), + true, + )?; + install_by_copying( + "./vchord.control", format!("{sharedir}/extension/vchord.control"), false, )?; @@ -314,9 +355,9 @@ fn main() -> Result<(), Box> { false, )?; } else { - let exports = parse(&cfgs, obj)?; + let exports = parse(&tsi, obj)?; install_by_writing( - generate(&path, &pg_version, &profile, &target, exports)?, + generate(&path, &pg_version, &tsi, &profile, &target, exports)?, format!("{sharedir_extension}/vchord--0.0.0.sql"), false, )?; diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index bdb7eafe..f633a5fb 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -16,6 +16,7 @@ rand.workspace = true [build-dependencies] cc = "1.2.27" +which = "8.0.0" [lints] workspace = true diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 12ce1981..7b5350a5 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -14,19 +14,51 @@ use std::env::var; use std::error::Error; +use std::ffi::OsString; + +fn compiler(host: &str, target: &str) -> Option { + if let Some(cc) = std::env::var_os("CC") { + return Some(cc); + } + if let Some(cc) = std::env::var_os(format!("CC_{target}")) { + return Some(cc); + } + if let Some(cc) = std::env::var_os(format!("CC_{}", target.replace("-", "_"))) { + return Some(cc); + } + if host == target { + if let Ok(cc) = which::which("clang") { + return Some(cc.into()); + } + for i in (16..=99).rev() { + if let Ok(cc) = which::which(format!("clang-{i}")) { + return Some(cc.into()); + } + } + } + None +} fn main() -> Result<(), Box> { println!("cargo::rerun-if-changed=cshim"); + let host = var("HOST")?; + let target = var("TARGET")?; let target_arch = var("CARGO_CFG_TARGET_ARCH")?; match target_arch.as_str() { "aarch64" => { let mut build = cc::Build::new(); + if let Some(compiler) = compiler(&host, &target) { + build.compiler(compiler); + } build.file("./cshim/aarch64.c"); build.opt_level(3); build.compile("simd_cshim"); } "x86_64" => { let mut build = cc::Build::new(); + if let Some(compiler) = compiler(&host, &target) { + build.compiler(compiler); + } build.file("./cshim/x86_64.c"); build.opt_level(3); build.compile("simd_cshim"); diff --git a/crates/vchordg/src/build.rs b/crates/vchordg/src/build.rs index f838bfc7..53c6b773 100644 --- a/crates/vchordg/src/build.rs +++ b/crates/vchordg/src/build.rs @@ -54,7 +54,7 @@ pub fn build( let serialized = MetaTuple::serialize(&MetaTuple { dims: vector_options.dims, m: index_options.m, - alpha: index_options.alpha, + max_alpha: index_options.max_alpha, ef_construction: index_options.ef_construction, beam_construction: index_options.beam_construction, start: OptionPointer::NONE, diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index 753b8ad9..fca0ed73 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -17,6 +17,7 @@ use crate::candidates::Candidates; use crate::operator::{CloneAccessor, Operator, Vector}; use crate::results::Results; use crate::tuples::*; +use crate::types::DistanceKind; use crate::vectors::{by_prefetch, by_read, copy_all, copy_nothing, copy_outs, update}; use crate::visited::Visited; use algo::accessor::LAccess; @@ -45,7 +46,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( assert_eq!(dims, vector.dims(), "unmatched dimensions"); let start = meta_tuple.start(); let m = meta_tuple.m(); - let alpha = meta_tuple.alpha(); + let max_alpha = meta_tuple.max_alpha(); let ef = meta_tuple.ef_construction(); let beam = meta_tuple.beam_construction(); let skip = meta_tuple.skip(); @@ -214,13 +215,14 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( }) .collect::>() }; - let outs = crate::prune::robust_prune( + let outs = crate::prune::prune( |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), (bump.alloc_slice(&pointers_t), t), trace.into_iter(), m, - alpha, + max_alpha, |(_, u)| *u, + O::DISTANCE == DistanceKind::L2S, ); let _ = update::( (index, pointers_t.as_slice()), @@ -265,13 +267,14 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( Some(((Reverse(dis_u), AlwaysEqual((pointers_u, u))), vector_u)) }) .collect::>(); - let outs = crate::prune::robust_prune( + let outs = crate::prune::prune( |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), (pointers_u.to_vec(), u), trace.into_iter(), m, - alpha, + max_alpha, |(_, u)| *u, + O::DISTANCE == DistanceKind::L2S, ); if update::( (index, pointers_u), diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs index 7b8eeabc..fdde458b 100644 --- a/crates/vchordg/src/maintain.rs +++ b/crates/vchordg/src/maintain.rs @@ -15,6 +15,7 @@ use crate::Opaque; use crate::operator::{CloneAccessor, Operator}; use crate::tuples::{MetaTuple, VectorTuple, VertexTuple, WithReader}; +use crate::types::DistanceKind; use crate::vectors::{by_read, copy_all, copy_nothing, copy_outs, update}; use algo::{Page, PageGuard, RelationRead, RelationWrite}; use always_equal::AlwaysEqual; @@ -29,7 +30,7 @@ where let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let m = meta_tuple.m(); - let alpha = meta_tuple.alpha(); + let max_alpha = meta_tuple.max_alpha(); let start = meta_tuple.start(); let link = meta_guard.get_opaque().link; drop(meta_guard); @@ -122,13 +123,14 @@ where } trace }; - let outs = crate::prune::robust_prune( + let outs = crate::prune::prune( |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), (pointers_u.to_vec(), u), trace.into_iter(), m, - alpha, + max_alpha, |(_, u)| *u, + O::DISTANCE == DistanceKind::L2S, ); if update::( (index, pointers_u.as_slice()), diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs index 63776c22..7fa46b51 100644 --- a/crates/vchordg/src/operator.rs +++ b/crates/vchordg/src/operator.rs @@ -12,6 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. +use crate::types::DistanceKind; use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S}; use distance::Distance; use half::f16; @@ -118,6 +119,8 @@ impl Vector for VectOwned { } pub trait Operator: 'static + Debug + Copy { + const DISTANCE: DistanceKind; + type Vector: Vector; type DistanceAccessor: Default @@ -148,6 +151,8 @@ impl Clone for Op { impl Copy for Op {} impl Operator for Op, L2S> { + const DISTANCE: DistanceKind = DistanceKind::L2S; + type Vector = VectOwned; type DistanceAccessor = DistanceAccessor, L2S>; @@ -169,6 +174,8 @@ impl Operator for Op, L2S> { } impl Operator for Op, Dot> { + const DISTANCE: DistanceKind = DistanceKind::Dot; + type Vector = VectOwned; type DistanceAccessor = DistanceAccessor, Dot>; @@ -190,6 +197,8 @@ impl Operator for Op, Dot> { } impl Operator for Op, L2S> { + const DISTANCE: DistanceKind = DistanceKind::L2S; + type Vector = VectOwned; type DistanceAccessor = DistanceAccessor, L2S>; @@ -211,6 +220,8 @@ impl Operator for Op, L2S> { } impl Operator for Op, Dot> { + const DISTANCE: DistanceKind = DistanceKind::Dot; + type Vector = VectOwned; type DistanceAccessor = DistanceAccessor, Dot>; diff --git a/crates/vchordg/src/prune.rs b/crates/vchordg/src/prune.rs index 935120b5..ca1f7128 100644 --- a/crates/vchordg/src/prune.rs +++ b/crates/vchordg/src/prune.rs @@ -16,13 +16,14 @@ use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; -pub fn robust_prune( +pub fn prune( mut d: impl FnMut(&V, &V) -> Distance, u: T, outs: impl Iterator, AlwaysEqual), V)>, m: u32, - alpha: f32, + max_alpha: f32, key: impl Fn(&T) -> K, + is_l2s: bool, ) -> Vec<(Reverse, AlwaysEqual)> { // V ← (V ∪ Nout(p)) \ {p} let mut trace = outs.collect::>(); @@ -31,7 +32,30 @@ pub fn robust_prune( trace.retain(|((_, AlwaysEqual(v)), _)| key(v) != key(&u)); trace.sort_by_key(|&((Reverse(d), _), _)| d); // Nout(p) ← ∅ + let max_alpha = if is_l2s { max_alpha } else { 1.0 }; + let mut alpha = 1.0; let mut result = Vec::new(); + while alpha <= max_alpha { + if result.len() == m as usize { + break; + } + trace = robust_prune(&mut d, m, alpha, &mut result, trace); + alpha *= 1.2; + } + if !(result.len() == m as usize) { + result.extend(trace); + result.truncate(m as usize); + } + result.into_iter().map(|(x, _)| x).collect() +} + +fn robust_prune( + mut d: impl FnMut(&V, &V) -> Distance, + m: u32, + alpha: f32, + result: &mut Vec<((Reverse, AlwaysEqual), V)>, + trace: Vec<((Reverse, AlwaysEqual), V)>, +) -> Vec<((Reverse, AlwaysEqual), V)> { let mut pruned = Vec::new(); for ((Reverse(dis_u), AlwaysEqual(u)), vector_u) in trace { if result.len() == m as usize { @@ -47,10 +71,5 @@ pub fn robust_prune( pruned.push(((Reverse(dis_u), AlwaysEqual(u)), vector_u)); } } - result - .into_iter() - .chain(pruned) - .map(|(x, _)| x) - .take(m as _) - .collect() + pruned } diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 47f5d640..43c95bba 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -43,7 +43,7 @@ struct MetaTupleHeader { version: u64, dims: u32, m: u32, - alpha: f32, + max_alpha: f32, ef_construction: u32, beam_construction: u32, start: OptionPointer, @@ -53,7 +53,7 @@ struct MetaTupleHeader { pub struct MetaTuple { pub dims: u32, pub m: u32, - pub alpha: f32, + pub max_alpha: f32, pub ef_construction: u32, pub beam_construction: u32, pub start: OptionPointer, @@ -68,7 +68,7 @@ impl Tuple for MetaTuple { MetaTuple { dims, m, - alpha, + max_alpha, ef_construction, beam_construction, start, @@ -82,7 +82,7 @@ impl Tuple for MetaTuple { version: VERSION, dims: *dims, m: *m, - alpha: *alpha, + max_alpha: *max_alpha, ef_construction: *ef_construction, beam_construction: *beam_construction, start: *start, @@ -126,8 +126,8 @@ impl<'a> MetaTupleReader<'a> { pub fn m(self) -> u32 { self.header.m } - pub fn alpha(self) -> f32 { - self.header.alpha + pub fn max_alpha(self) -> f32 { + self.header.max_alpha } pub fn ef_construction(self) -> u32 { self.header.ef_construction diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index 44ae10ee..4a6d670d 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -23,9 +23,9 @@ pub struct VamanaIndexOptions { #[serde(default = "VamanaIndexOptions::default_m")] #[validate(range(min = 1, max = 512))] pub m: u32, - #[serde(default = "VamanaIndexOptions::default_alpha")] + #[serde(default = "VamanaIndexOptions::default_max_alpha")] #[validate(range(min = 1.0, max = 2.0))] - pub alpha: f32, + pub max_alpha: f32, #[serde(default = "VamanaIndexOptions::default_ef_construction")] #[validate(range(min = 1, max = 65535))] pub ef_construction: u32, @@ -38,14 +38,14 @@ impl VamanaIndexOptions { fn default_m() -> u32 { 32 } - fn default_alpha() -> f32 { - 1.0 + fn default_max_alpha() -> f32 { + 1.2 } fn default_ef_construction() -> u32 { 64 } fn default_beam_construction() -> u32 { - 8 + 1 } } @@ -53,7 +53,7 @@ impl Default for VamanaIndexOptions { fn default() -> Self { Self { m: Self::default_m(), - alpha: Self::default_alpha(), + max_alpha: Self::default_max_alpha(), ef_construction: Self::default_ef_construction(), beam_construction: Self::default_beam_construction(), } diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 5029f462..c8cb8287 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -38,7 +38,7 @@ pub enum PostgresIo { static VCHORDG_EF_SEARCH: GucSetting = GucSetting::::new(64); -static VCHORDG_BEAM_SIZE: GucSetting = GucSetting::::new(8); +static VCHORDG_BEAM_SEARCH: GucSetting = GucSetting::::new(1); static VCHORDG_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); @@ -175,7 +175,7 @@ pub fn init() { c"vchordg.beam_search", c"`beam_search` argument of vchordg.", c"`beam_search` argument of vchordg.", - &VCHORDG_BEAM_SIZE, + &VCHORDG_BEAM_SEARCH, 1, 65535, GucContext::Userset, @@ -220,7 +220,7 @@ pub fn vchordg_ef_search() -> u32 { } pub fn vchordg_beam_search() -> u32 { - VCHORDG_BEAM_SIZE.get() as u32 + VCHORDG_BEAM_SEARCH.get() as u32 } pub fn vchordg_max_scan_tuples() -> Option { diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index a076c8ad..ac9f09fb 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -186,9 +186,6 @@ pub unsafe extern "C-unwind" fn ambuild( if let Err(errors) = Validate::validate(&vchordg_options) { pgrx::error!("error while validating options: {}", errors); } - if !matches!(vector_options.d, DistanceKind::L2S) && vchordg_options.index.alpha != 1.0 { - pgrx::error!("alpha must be 1 if L2 metric is not used"); - } let opfamily = unsafe { opfamily(index_relation) }; let index = unsafe { PostgresRelation::new(index_relation) }; let heap = Heap { From c31beaa2d0170847b3c8779f935ecc6ac258f488 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Mon, 14 Jul 2025 20:24:31 +0800 Subject: [PATCH 175/324] feat: query recall eval func (#284) ## Usages ```sql > SELECT vchordrq_evaluate_query_recall(query_text =>$$SELECT ctid FROM test ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); vchordrq_evaluate_query_recall -------------------------------- 1 (1 row) > SELECT vchordrq_evaluate_query_recall(query_text => $$SELECT ctid FROM test ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 100$$, exact_search=>true); vchordrq_evaluate_query_recall -------------------------------- 1 > SELECT vchordrq_evaluate_query_recall(query_text => $$SELECT ctid FROM test WHERE FALSE' LIMIT 100$$, exact_search=>true); vchordrq_evaluate_query_recall -------------------------------- NaN (1 row) ``` Signed-off-by: cutecutecat --- src/index/gucs.rs | 28 +++++++++++++++ src/index/vchordg/am/mod.rs | 9 +++++ src/index/vchordrq/am/mod.rs | 4 ++- src/sql/finalize.sql | 69 ++++++++++++++++++++++++++++++++++++ tests/vchordrq/recall.slt | 53 +++++++++++++++++++++++++++ 5 files changed, 162 insertions(+), 1 deletion(-) create mode 100644 tests/vchordrq/recall.slt diff --git a/src/index/gucs.rs b/src/index/gucs.rs index c8cb8287..3d0515b0 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -36,6 +36,8 @@ pub enum PostgresIo { ReadStream, } +static VCHORDG_ENABLE_SCAN: GucSetting = GucSetting::::new(true); + static VCHORDG_EF_SEARCH: GucSetting = GucSetting::::new(64); static VCHORDG_BEAM_SEARCH: GucSetting = GucSetting::::new(1); @@ -56,6 +58,8 @@ static VCHORDG_IO_RERANK: GucSetting = GucSetting::::new PostgresIo::ReadStream, ); +static VCHORDRQ_ENABLE_SCAN: GucSetting = GucSetting::::new(true); + static VCHORDRQ_PROBES: GucSetting> = GucSetting::>::new(Some(c"")); static VCHORDRQ_EPSILON: GucSetting = GucSetting::::new(1.9); @@ -83,6 +87,14 @@ static VCHORDRQ_IO_RERANK: GucSetting = GucSetting::::ne ); pub fn init() { + GucRegistry::define_bool_guc( + c"vchordrq.enable_scan", + c"`enable_scan` argument of vchordrq.", + c"`enable_scan` argument of vchordrq.", + &VCHORDRQ_ENABLE_SCAN, + GucContext::Userset, + GucFlags::default(), + ); GucRegistry::define_string_guc( c"vchordrq.probes", c"`probes` argument of vchordrq.", @@ -161,6 +173,14 @@ pub fn init() { #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrq".as_ptr()); } + GucRegistry::define_bool_guc( + c"vchordg.enable_scan", + c"`enable_scan` argument of vchordg.", + c"`enable_scan` argument of vchordg.", + &VCHORDG_ENABLE_SCAN, + GucContext::Userset, + GucFlags::default(), + ); GucRegistry::define_int_guc( c"vchordg.ef_search", c"`ef_search` argument of vchordg.", @@ -215,6 +235,10 @@ pub fn init() { } } +pub fn vchordg_enable_scan() -> bool { + VCHORDG_ENABLE_SCAN.get() +} + pub fn vchordg_ef_search() -> u32 { VCHORDG_EF_SEARCH.get() as u32 } @@ -248,6 +272,10 @@ pub fn vchordg_io_rerank() -> crate::index::vchordg::scanners::Io { } } +pub fn vchordrq_enable_scan() -> bool { + VCHORDRQ_ENABLE_SCAN.get() +} + pub fn vchordrq_probes() -> Vec { match VCHORDRQ_PROBES.get() { None => Vec::new(), diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index c598d3fe..20a09678 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -160,6 +160,15 @@ pub unsafe extern "C-unwind" fn amcostestimate( index_pages: *mut f64, ) { unsafe { + use pgrx::pg_sys::disable_cost; + if !gucs::vchordg_enable_scan() { + *index_startup_cost = disable_cost; + *index_total_cost = disable_cost; + *index_selectivity = 0.0; + *index_correlation = 0.0; + *index_pages = 1.0; + return; + } *index_startup_cost = 0.0; *index_total_cost = 0.0; *index_selectivity = 1.0; diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 13039ffe..61d28a9c 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -163,7 +163,9 @@ pub unsafe extern "C-unwind" fn amcostestimate( use pgrx::pg_sys::disable_cost; let index_opt_info = (*path).indexinfo; // do not use index, if there are no orderbys or clauses - if (*path).indexorderbys.is_null() && (*path).indexclauses.is_null() { + if ((*path).indexorderbys.is_null() && (*path).indexclauses.is_null()) + || !gucs::vchordrq_enable_scan() + { *index_startup_cost = disable_cost; *index_total_cost = disable_cost; *index_selectivity = 0.0; diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 594e3b8a..7333f209 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -148,6 +148,75 @@ IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhan CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +STRICT +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; + accu_probes text; +BEGIN + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; diff --git a/tests/vchordrq/recall.slt b/tests/vchordrq/recall.slt new file mode 100644 index 00000000..3c1b1bb5 --- /dev/null +++ b/tests/vchordrq/recall.slt @@ -0,0 +1,53 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) +SELECT ARRAY[i * 0.0001, i * 0.00005, i * 0.0002]::vector(3) FROM generate_series(1, 10000) as s(i); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops); + +statement ok +SET vchordrq.epsilon = 0.8; + +statement ok +SET vchordrq.probes = '1'; + +statement error MaxSim operator cannot be used for estimated recall +SELECT * from vchordrq_evaluate_query_recall(query=>'@#'); + +statement error Error executing ANN query +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); + +statement ok +SET vchordrq.probes = ''; + +statement error Error executing ANN query +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT val FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); + +statement error Error executing ANN query +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT * FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); + +query I +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); +---- +1 + +query I +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$, exact_search=>true); +---- +1 + +query I +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t WHERE FALSE ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); +---- +NaN + +query I +SHOW vchordrq.epsilon; +---- +0.8 + +statement ok +DROP TABLE t; \ No newline at end of file From 26d71376c2a25468af68df79cf2625532a66602c Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 16 Jul 2025 13:57:03 +0800 Subject: [PATCH 176/324] fix: add missing license header (#289) Signed-off-by: usamoi --- .github/workflows/check.yml | 23 ++++++++++++++++++----- crates/algo/src/accessor.rs | 14 ++++++++++++++ crates/rabitq/src/b1.rs | 14 ++++++++++++++ crates/simd/src/f32.rs | 20 ++++++++++---------- crates/vchordg/src/lib.rs | 14 ++++++++++++++ crates/vchordg/src/vectors.rs | 19 +++++++++++++++---- src/index/opclass.rs | 14 ++++++++++++++ 7 files changed, 99 insertions(+), 19 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 230143f6..b82424fa 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -74,24 +74,33 @@ jobs: C_FILES=$(find . -type f -name "*.c" -not -path "./target/*") PY_FILES=$(find . -type f -name "*.py" -not -path "./target/*") + FLAG="0" + for p in $RS_FILES; do if ! head -n "$COUNT" "$p" | cmp -s - <(echo "$RS_HEADER"); then echo "license header mismatch in file $p" + FLAG="1" fi done for p in $C_FILES; do if ! head -n "$COUNT" "$p" | cmp -s - <(echo "$C_HEADER"); then echo "license header mismatch in file $p" + FLAG="1" fi done for p in $PY_FILES; do if ! head -n "$COUNT" "$p" | cmp -s - <(echo "$PY_HEADER"); then echo "license header mismatch in file $p" + FLAG="1" fi done + if [ "$FLAG" -eq "1" ]; then + exit 1 + fi + - name: SQL run: | ! grep -P '\t' -r ./sql @@ -138,17 +147,21 @@ jobs: - name: Cargo Test (simd) run: | - cargo test -p simd --release -- --nocapture - if [ "$(uname -m)" == "x86_64" ]; then cargo \ --config 'target.'\''cfg(all())'\''.runner = ["/opt/sde/sde64", "-spr", "--"]' \ - test -p simd --release -- --nocapture + test -p simd -- --nocapture fi if [ "$(uname -m)" == "aarch64" ]; then cargo \ - --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max"]' \ - test -p simd --release -- --nocapture + --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=16"]' \ + test -p simd -- --nocapture + cargo \ + --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=32"]' \ + test -p simd -- --nocapture + cargo \ + --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=64"]' \ + test -p simd -- --nocapture fi psql: diff --git a/crates/algo/src/accessor.rs b/crates/algo/src/accessor.rs index bf94e8b9..f89d6045 100644 --- a/crates/algo/src/accessor.rs +++ b/crates/algo/src/accessor.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use distance::Distance; use half::f16; use simd::Floating; diff --git a/crates/rabitq/src/b1.rs b/crates/rabitq/src/b1.rs index 25fb8fd4..2433ac56 100644 --- a/crates/rabitq/src/b1.rs +++ b/crates/rabitq/src/b1.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use simd::Floating; #[derive(Debug, Clone, Copy)] diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index 70ba66bd..159bb949 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -452,7 +452,7 @@ mod reduce_sum_of_abs_x { #[test] fn reduce_sum_of_abs_x_v4_test() { use rand::Rng; - const EPSILON: f32 = 0.008; + const EPSILON: f32 = 0.009; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; @@ -515,7 +515,7 @@ mod reduce_sum_of_abs_x { #[test] fn reduce_sum_of_abs_x_v3_test() { use rand::Rng; - const EPSILON: f32 = 0.008; + const EPSILON: f32 = 0.009; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; @@ -571,7 +571,7 @@ mod reduce_sum_of_abs_x { #[test] fn reduce_sum_of_abs_x_v2_test() { use rand::Rng; - const EPSILON: f32 = 0.008; + const EPSILON: f32 = 0.009; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); return; @@ -625,7 +625,7 @@ mod reduce_sum_of_abs_x { #[test] fn reduce_sum_of_abs_x_a2_test() { use rand::Rng; - const EPSILON: f32 = 0.008; + const EPSILON: f32 = 0.009; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); return; @@ -665,7 +665,7 @@ mod reduce_sum_of_abs_x { #[test] fn reduce_sum_of_abs_x_a3_256_test() { use rand::Rng; - const EPSILON: f32 = 0.008; + const EPSILON: f32 = 0.009; if !crate::is_cpu_detected!("a3.256") { println!("test {} ... skipped (a3.256)", module_path!()); return; @@ -726,7 +726,7 @@ mod reduce_sum_of_x2 { #[test] fn reduce_sum_of_x2_v4_test() { use rand::Rng; - const EPSILON: f32 = 0.006; + const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; @@ -785,7 +785,7 @@ mod reduce_sum_of_x2 { #[test] fn reduce_sum_of_x2_v3_test() { use rand::Rng; - const EPSILON: f32 = 0.006; + const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; @@ -839,7 +839,7 @@ mod reduce_sum_of_x2 { #[test] fn reduce_sum_of_x2_v2_fma_test() { use rand::Rng; - const EPSILON: f32 = 0.006; + const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { println!("test {} ... skipped (v2:fma)", module_path!()); return; @@ -891,7 +891,7 @@ mod reduce_sum_of_x2 { #[test] fn reduce_sum_of_x2_a2_test() { use rand::Rng; - const EPSILON: f32 = 0.006; + const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); return; @@ -931,7 +931,7 @@ mod reduce_sum_of_x2 { #[test] fn reduce_sum_of_x2_a3_256_test() { use rand::Rng; - const EPSILON: f32 = 0.006; + const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("a3.256") { println!("test {} ... skipped (a3.256)", module_path!()); return; diff --git a/crates/vchordg/src/lib.rs b/crates/vchordg/src/lib.rs index 142905af..aaf7bbbb 100644 --- a/crates/vchordg/src/lib.rs +++ b/crates/vchordg/src/lib.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + mod build; mod bulkdelete; mod candidates; diff --git a/crates/vchordg/src/vectors.rs b/crates/vchordg/src/vectors.rs index 4af30d88..dec945bd 100644 --- a/crates/vchordg/src/vectors.rs +++ b/crates/vchordg/src/vectors.rs @@ -1,8 +1,19 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + use crate::operator::{Operator, Vector}; -use crate::tuples::{ - OptionNeighbour, Pointer, VectorTuple, VectorTupleReader, VectorTupleWriter, WithReader, - WithWriter, -}; +use crate::tuples::*; use algo::accessor::Accessor1; use algo::{Page, RelationRead, RelationWrite}; use distance::Distance; diff --git a/src/index/opclass.rs b/src/index/opclass.rs index 6ae30711..0b182bf4 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -1,3 +1,17 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordg_support_vector_l2_ops() -> String { "vchordg_vector_l2_ops".to_string() From 88fd1e7f65b764a1ec8bea86cbd0a7162738d6f3 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 18 Jul 2025 13:14:54 +0800 Subject: [PATCH 177/324] feat: use 2-bit quantization for vchordg (#288) Signed-off-by: usamoi --- .cargo/config.toml | 2 - .github/workflows/check.yml | 6 +- .github/workflows/release.yml | 2 +- .gitignore | 1 + Cargo.lock | 62 +++++++++-- Cargo.toml | 5 +- Makefile | 2 +- build.rs | 3 + crates/rabitq/Cargo.toml | 4 +- crates/rabitq/src/b1.rs | 184 +++++---------------------------- crates/rabitq/src/b2.rs | 26 ++--- crates/rabitq/src/b4.rs | 24 ++--- crates/rabitq/src/b8.rs | 6 +- crates/rabitq/src/extended.rs | 32 ++++++ crates/simd/Cargo.toml | 3 + crates/simd/src/lib.rs | 1 - crates/simd/src/packed.rs | 92 ----------------- crates/simd_macros/src/lib.rs | 64 +++++++++--- crates/vchordg/src/insert.rs | 14 ++- crates/vchordg/src/operator.rs | 54 +++++----- crates/vchordg/src/search.rs | 12 ++- crates/vchordg/src/types.rs | 2 +- src/lib.rs | 3 +- 23 files changed, 251 insertions(+), 353 deletions(-) delete mode 100644 .cargo/config.toml delete mode 100644 crates/simd/src/packed.rs diff --git a/.cargo/config.toml b/.cargo/config.toml deleted file mode 100644 index b3550a8f..00000000 --- a/.cargo/config.toml +++ /dev/null @@ -1,2 +0,0 @@ -[alias] -make = "run -p make --" diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index b82424fa..ffecc957 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -221,7 +221,7 @@ jobs: - name: Install run: | - cargo make build -o ./build/raw --profile dev + cargo run -p make -- build -o ./build/raw --profile dev sudo make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install - name: Service @@ -323,7 +323,7 @@ jobs: - name: Install run: | - cargo make build -o ./build/raw --profile dev + cargo run -p make -- build -o ./build/raw --profile dev sudo make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install - name: Service @@ -405,5 +405,5 @@ jobs: - name: Install run: | - cargo make build -o ./build/raw --profile dev + cargo run -p make -- build -o ./build/raw --profile dev make PG_CONFIG="$env:PGRX_PG_CONFIG_PATH" install diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 7d663b9a..10adb4bc 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -88,7 +88,7 @@ jobs: ARCH: ${{ matrix.arch }} PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} run: | - cargo make build -o ./build/raw + cargo run -p make -- build -o ./build/raw (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) diff --git a/.gitignore b/.gitignore index 85d1619b..47ad9c00 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ .vscode .ignore build +.cargo/config.toml diff --git a/Cargo.lock b/Cargo.lock index 8d4672c5..bc47266e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -166,7 +166,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02260d489095346e5cafd04dea8e8cb54d1d74fcd759022a9b72986ebe9a1257" dependencies = [ "serde", - "toml", + "toml 0.8.23", ] [[package]] @@ -992,7 +992,7 @@ dependencies = [ "serde", "serde_json", "thiserror", - "toml", + "toml 0.8.23", "url", "winapi", ] @@ -1324,6 +1324,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_spanned" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40734c41988f7306bb04f0ecf60ec0f3f1caa34290e4e8ea471dcd3346483b83" +dependencies = [ + "serde", +] + [[package]] name = "shlex" version = "1.3.0" @@ -1466,11 +1475,26 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" dependencies = [ "serde", - "serde_spanned", - "toml_datetime", + "serde_spanned 0.6.9", + "toml_datetime 0.6.11", "toml_edit", ] +[[package]] +name = "toml" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed0aee96c12fa71097902e0bb061a5e1ebd766a6636bb605ba401c45c1650eac" +dependencies = [ + "indexmap", + "serde", + "serde_spanned 1.0.0", + "toml_datetime 0.7.0", + "toml_parser", + "toml_writer", + "winnow", +] + [[package]] name = "toml_datetime" version = "0.6.11" @@ -1480,6 +1504,15 @@ dependencies = [ "serde", ] +[[package]] +name = "toml_datetime" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bade1c3e902f58d73d3f294cd7f20391c1cb2fbcb643b73566bc773971df91e3" +dependencies = [ + "serde", +] + [[package]] name = "toml_edit" version = "0.22.27" @@ -1488,18 +1521,33 @@ checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" dependencies = [ "indexmap", "serde", - "serde_spanned", - "toml_datetime", + "serde_spanned 0.6.9", + "toml_datetime 0.6.11", "toml_write", "winnow", ] +[[package]] +name = "toml_parser" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97200572db069e74c512a14117b296ba0a80a30123fbbb5aa1f4a348f639ca30" +dependencies = [ + "winnow", +] + [[package]] name = "toml_write" version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" +[[package]] +name = "toml_writer" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcc842091f2def52017664b53082ecbbeb5c7731092bad69d2c63050401dfd64" + [[package]] name = "unescape" version = "0.1.0" @@ -1608,7 +1656,7 @@ dependencies = [ "seq-macro", "serde", "simd", - "toml", + "toml 0.9.2", "validator", "vchordg", "vchordrq", diff --git a/Cargo.toml b/Cargo.toml index 1c4995a6..1ff28765 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,7 @@ name = "pgrx_embed_vchord" path = "./src/bin/pgrx_embed.rs" [features] -default = [] +default = ["simd/init"] pg13 = ["pgrx/pg13", "pgrx-catalog/pg13"] pg14 = ["pgrx/pg14", "pgrx-catalog/pg14"] pg15 = ["pgrx/pg15", "pgrx-catalog/pg15"] @@ -40,7 +40,7 @@ pgrx-catalog = "0.3.0" rand.workspace = true seq-macro.workspace = true serde.workspace = true -toml = "0.8.23" +toml = "0.9.2" validator.workspace = true zerocopy.workspace = true @@ -66,6 +66,7 @@ bumpalo = "3.19.0" half = { version = "2.6.0", features = ["zerocopy"] } paste = "1" rand = "0.9.1" +rand_chacha = "0.9.0" seq-macro = "0.3.6" serde = { version = "1", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } diff --git a/Makefile b/Makefile index 3e10bbd6..b842fa9a 100644 --- a/Makefile +++ b/Makefile @@ -5,7 +5,7 @@ PG_CONFIG ?= pg_config all: build build: - PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo make build -o ./build/raw + PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p make -- build -o ./build/raw install: cp -r ./build/raw/pkglibdir/. $(shell $(PG_CONFIG) --pkglibdir) diff --git a/build.rs b/build.rs index 7917f182..cf775640 100644 --- a/build.rs +++ b/build.rs @@ -30,6 +30,9 @@ fn main() -> Result<(), Box> { return Err("VectorChord version is not defined.".into()); }; println!("cargo::rustc-env=VCHORD_VERSION={version}"); + if var("CARGO_CFG_TARGET_OS")? == "linux" { + println!("cargo::rustc-link-arg-cdylib=-Wl,-Bsymbolic"); + } if var("CARGO_CFG_TARGET_OS")? == "macos" { if let Some(path) = var_os("PGRX_PG_CONFIG_PATH") { let map = { diff --git a/crates/rabitq/Cargo.toml b/crates/rabitq/Cargo.toml index 0d358764..72e65fc6 100644 --- a/crates/rabitq/Cargo.toml +++ b/crates/rabitq/Cargo.toml @@ -10,8 +10,8 @@ simd = { path = "../simd" } zerocopy.workspace = true [build-dependencies] -rand = "0.9.1" -rand_chacha = "0.9.0" +rand.workspace = true +rand_chacha.workspace = true [lints] workspace = true diff --git a/crates/rabitq/src/b1.rs b/crates/rabitq/src/b1.rs index 2433ac56..9109d3b7 100644 --- a/crates/rabitq/src/b1.rs +++ b/crates/rabitq/src/b1.rs @@ -12,184 +12,54 @@ // // Copyright (c) 2025 TensorChord Inc. -use simd::Floating; - -#[derive(Debug, Clone, Copy)] -pub struct CodeMetadata { - pub dis_u_2: f32, - pub factor_cnt: f32, - pub factor_ip: f32, -} - -impl CodeMetadata { - pub fn into_tuple(self) -> (f32, f32, f32) { - (self.dis_u_2, self.factor_cnt, self.factor_ip) - } - pub fn into_array(self) -> [f32; 3] { - [self.dis_u_2, self.factor_cnt, self.factor_ip] - } - pub fn from_tuple((dis_u_2, factor_cnt, factor_ip): (f32, f32, f32)) -> Self { - Self { - dis_u_2, - factor_cnt, - factor_ip, - } - } - pub fn from_array([dis_u_2, factor_cnt, factor_ip]: [f32; 3]) -> Self { - Self { - dis_u_2, - factor_cnt, - factor_ip, - } - } -} - -pub type Code = (CodeMetadata, Vec); +pub use crate::extended::{Code, CodeMetadata}; pub fn code(vector: &[f32]) -> Code { - let n = vector.len(); - let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); - let sum_of_x_2 = f32::reduce_sum_of_x2(vector); - ( - CodeMetadata { - dis_u_2: sum_of_x_2, - factor_cnt: { - let cnt_pos = vector - .iter() - .map(|x| x.is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.is_sign_negative() as i32) - .sum::(); - (cnt_pos - cnt_neg) as f32 - }, - factor_ip: sum_of_x_2 / sum_of_abs_x, - }, - { - let mut signs = Vec::new(); - for i in 0..n { - signs.push(vector[i].is_sign_positive()); - } - signs - }, - ) + crate::extended::code::<1>(vector) } pub mod binary { - pub fn pack_code(input: &[bool]) -> Vec { - let f = |t: &[bool; 64]| { - let mut result = 0_u64; - for i in 0..64 { - result |= (t[i] as u64) << i; - } - result - }; - let (arrays, remainder) = input.as_chunks::<64>(); - let mut buffer = [false; 64]; - let tailing = if !remainder.is_empty() { - buffer[..remainder.len()].copy_from_slice(remainder); - Some(&buffer) - } else { - None - }; - arrays.iter().chain(tailing).map(f).collect() + pub fn pack_code(input: &[u8]) -> Vec { + crate::extended::pack_code::<1>(input) + .into_iter() + .flatten() + .collect() } - use super::CodeMetadata; - use simd::Floating; + use crate::extended::CodeMetadata; - const BITS: usize = 6; - - #[derive(Debug, Clone, Copy)] - pub struct BinaryLutMetadata { - dis_v_2: f32, - b: f32, - k: f32, - qvector_sum: f32, - } + const BITS: usize = 4; + pub type BinaryLutMetadata = CodeMetadata; pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); - pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u64]); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); pub fn preprocess(vector: &[f32]) -> BinaryLut { - let dis_v_2 = f32::reduce_sum_of_x2(vector); - preprocess_with_distance(vector, dis_v_2) - } - - pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BinaryLut { - let (k, b, qvector) = simd::quantize::quantize(vector, ((1 << BITS) - 1) as f32); - let qvector_sum = if vector.len() <= (65535_usize / ((1 << BITS) - 1)) { - simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 - } else { - simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 - }; - ( - BinaryLutMetadata { - dis_v_2, - b, - k, - qvector_sum, - }, - binarize(&qvector), - ) + let (metadata, elements) = crate::extended::code::(vector); + (metadata, crate::extended::pack_code::<4>(&elements)) } - pub fn accumulate(x: &[u64], y: &[impl AsRef<[u64]>; BITS]) -> u32 { - let mut result = 0_u32; - for i in 0..BITS { - result += simd::bit::reduce_sum_of_and(x, y[i].as_ref()) << i; - } - result + pub fn accumulate(lhs: &[u64], rhs: &[Vec; 4]) -> u32 { + crate::extended::accumulate::<1, 4>(lhs, rhs) } - pub fn half_process_l2( - sum: u32, - CodeMetadata { - dis_u_2, - factor_cnt, - factor_ip, - }: CodeMetadata, - BinaryLutMetadata { - dis_v_2, - b, - k, - qvector_sum, - }: BinaryLutMetadata, + pub fn half_process_dot( + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, ) -> (f32,) { - let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; - let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; + let rough = crate::extended::half_process_dot::<1, BITS>(n, value, code, lut); (rough,) } - pub fn half_process_dot( - sum: u32, - CodeMetadata { - dis_u_2: _, - factor_cnt, - factor_ip, - }: CodeMetadata, - BinaryLutMetadata { - dis_v_2: _, - b, - k, - qvector_sum, - }: BinaryLutMetadata, + pub fn half_process_l2( + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, ) -> (f32,) { - let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; - let rough = -e * factor_ip; + let rough = crate::extended::half_process_l2::<1, BITS>(n, value, code, lut); (rough,) } - - pub(crate) fn binarize(vector: &[u8]) -> [Vec; BITS] { - let n = vector.len(); - let mut t: [_; BITS] = std::array::from_fn(|_| vec![0_u64; n.div_ceil(64)]); - for i in 0..BITS { - for j in 0..n { - let bit = (vector[j] >> i) & 1; - t[i][j / 64] |= (bit as u64) << (j % 64); - } - } - t - } } diff --git a/crates/rabitq/src/b2.rs b/crates/rabitq/src/b2.rs index a133bfbf..cd13e269 100644 --- a/crates/rabitq/src/b2.rs +++ b/crates/rabitq/src/b2.rs @@ -19,34 +19,28 @@ pub fn code(vector: &[f32]) -> Code { } pub mod binary { - pub fn pack_code(input: &[u8]) -> Vec { - let f = |&[t_0, t_1, t_2, t_3]: &[u8; 4]| t_0 | (t_1 << 2) | (t_2 << 4) | (t_3 << 6); - let (arrays, remainder) = input.as_chunks::<4>(); - let mut buffer = [0u8; 4]; - let tailing = if !remainder.is_empty() { - buffer[..remainder.len()].copy_from_slice(remainder); - Some(&buffer) - } else { - None - }; - arrays.iter().chain(tailing).map(f).collect() + pub fn pack_code(input: &[u8]) -> Vec { + crate::extended::pack_code::<2>(input) + .into_iter() + .flatten() + .collect() } use crate::extended::CodeMetadata; - const BITS: usize = 8; + const BITS: usize = 4; pub type BinaryLutMetadata = CodeMetadata; - pub type BinaryLut = (BinaryLutMetadata, Vec); + pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); pub fn preprocess(vector: &[f32]) -> BinaryLut { let (metadata, elements) = crate::extended::code::(vector); - (metadata, crate::b8::binary::pack_code(&elements)) + (metadata, crate::extended::pack_code::<4>(&elements)) } - pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { - simd::packed::u2_u8_reduce_sum_of_xy(x, y) + pub fn accumulate(lhs: &[u64], rhs: &[Vec; 4]) -> u32 { + crate::extended::accumulate::<2, 4>(lhs, rhs) } pub fn half_process_dot( diff --git a/crates/rabitq/src/b4.rs b/crates/rabitq/src/b4.rs index 6aa4236e..b0048304 100644 --- a/crates/rabitq/src/b4.rs +++ b/crates/rabitq/src/b4.rs @@ -19,17 +19,11 @@ pub fn code(vector: &[f32]) -> Code { } pub mod binary { - pub fn pack_code(input: &[u8]) -> Vec { - let f = |&[t_0, t_1]: &[u8; 2]| t_0 | (t_1 << 4); - let (arrays, remainder) = input.as_chunks::<2>(); - let mut buffer = [0u8; 2]; - let tailing = if !remainder.is_empty() { - buffer[..remainder.len()].copy_from_slice(remainder); - Some(&buffer) - } else { - None - }; - arrays.iter().chain(tailing).map(f).collect() + pub fn pack_code(input: &[u8]) -> Vec { + crate::extended::pack_code::<4>(input) + .into_iter() + .flatten() + .collect() } use crate::extended::CodeMetadata; @@ -37,16 +31,16 @@ pub mod binary { const BITS: usize = 4; pub type BinaryLutMetadata = CodeMetadata; - pub type BinaryLut = (BinaryLutMetadata, Vec); + pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); pub fn preprocess(vector: &[f32]) -> BinaryLut { let (metadata, elements) = crate::extended::code::(vector); - (metadata, pack_code(&elements)) + (metadata, crate::extended::pack_code::<4>(&elements)) } - pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { - simd::packed::u4_u4_reduce_sum_of_xy(x, y) + pub fn accumulate(lhs: &[u64], rhs: &[Vec; 4]) -> u32 { + crate::extended::accumulate::<4, 4>(lhs, rhs) } pub fn half_process_dot( diff --git a/crates/rabitq/src/b8.rs b/crates/rabitq/src/b8.rs index 3b138890..43a6a688 100644 --- a/crates/rabitq/src/b8.rs +++ b/crates/rabitq/src/b8.rs @@ -19,10 +19,6 @@ pub fn code(vector: &[f32]) -> Code { } pub mod binary { - pub fn pack_code(input: &[u8]) -> Vec { - input.to_vec() - } - use crate::extended::CodeMetadata; const BITS: usize = 8; @@ -33,7 +29,7 @@ pub mod binary { pub fn preprocess(vector: &[f32]) -> BinaryLut { let (metadata, elements) = crate::extended::code::(vector); - (metadata, pack_code(&elements)) + (metadata, elements) } pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { diff --git a/crates/rabitq/src/extended.rs b/crates/rabitq/src/extended.rs index 2a6c2814..2e3c014f 100644 --- a/crates/rabitq/src/extended.rs +++ b/crates/rabitq/src/extended.rs @@ -179,3 +179,35 @@ fn find_scale(o: &[f32]) -> f64 { x_m } + +pub fn pack_code(input: &[u8]) -> [Vec; BITS] { + #[inline(always)] + fn f(array: &[u8; 64], bit: usize) -> u64 { + let mut result = 0_u64; + for i in 0..64 { + result |= ((array[i] as u64 >> bit) & 1) << i; + } + result + } + let (arrays, remainder) = input.as_chunks::<64>(); + let mut buffer = [0_u8; 64]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + std::array::from_fn(|bit| arrays.iter().chain(tailing).map(|t| f(t, bit)).collect()) +} + +pub fn accumulate(lhs: &[u64], rhs: &[Vec; Y]) -> u32 { + assert!(lhs.len().is_multiple_of(X)); + let d = lhs.len() / X; + let mut result = 0_u32; + for i in 0..X { + for j in 0..Y { + result += simd::bit::reduce_sum_of_and(&lhs[i * d..][..d], &rhs[j]) << (i + j); + } + } + result +} diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index f633a5fb..3d298418 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -4,6 +4,9 @@ version.workspace = true edition.workspace = true publish = false +[features] +init = [] + [dependencies] simd_macros = { path = "../simd_macros" } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index f96490ce..16f747cd 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -22,7 +22,6 @@ mod f32; pub mod bit; pub mod fast_scan; pub mod fht; -pub mod packed; pub mod quantize; pub mod rotate; pub mod u8; diff --git a/crates/simd/src/packed.rs b/crates/simd/src/packed.rs deleted file mode 100644 index 645fb8df..00000000 --- a/crates/simd/src/packed.rs +++ /dev/null @@ -1,92 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -pub fn u2_u2_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { - u2_u2_reduce_sum_of_xy::u2_u2_reduce_sum_of_xy(s, t) -} - -mod u2_u2_reduce_sum_of_xy { - #[crate::multiversion("v4", "v3", "v2", "a2")] - pub fn u2_u2_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { - assert_eq!(s.len(), t.len()); - let n = s.len(); - let mut result_0 = 0_u32; - let mut result_1 = 0_u32; - let mut result_2 = 0_u32; - let mut result_3 = 0_u32; - for i in 0..n { - let (s, t) = (s[i], t[i]); - result_0 += (((s >> 0) & 3) * ((t >> 0) & 3)) as u32; - result_1 += (((s >> 2) & 3) * ((t >> 2) & 3)) as u32; - result_2 += (((s >> 4) & 3) * ((t >> 4) & 3)) as u32; - result_3 += (((s >> 6) & 3) * ((t >> 6) & 3)) as u32; - } - result_0 + result_1 + result_2 + result_3 - } -} - -pub fn u4_u4_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { - u4_u4_reduce_sum_of_xy::u4_u4_reduce_sum_of_xy(s, t) -} - -mod u4_u4_reduce_sum_of_xy { - #[crate::multiversion("v4", "v3", "v2", "a2")] - pub fn u4_u4_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { - assert_eq!(s.len(), t.len()); - let n = s.len(); - let mut result_0 = 0; - let mut result_1 = 0; - for i in 0..n { - let (s, t) = (s[i], t[i]); - result_0 += (((s >> 0) & 15) * ((t >> 0) & 15)) as u32; - result_1 += (((s >> 4) & 15) * ((t >> 4) & 15)) as u32; - } - result_0 + result_1 - } -} - -pub fn u2_u8_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { - u2_u8_reduce_sum_of_xy::u2_u8_reduce_sum_of_xy(s, t) -} - -mod u2_u8_reduce_sum_of_xy { - #[crate::multiversion("v4", "v3", "v2", "a2")] - pub fn u2_u8_reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { - assert_eq!(s.len(), t.len().div_ceil(4)); - let mut result_0 = 0_u32; - let mut result_1 = 0_u32; - let mut result_2 = 0_u32; - let mut result_3 = 0_u32; - let (arrays, remainder) = t.as_chunks::<4>(); - assert_eq!(s.len(), arrays.len()); - let n = arrays.len(); - for i in 0..n { - let (s, t) = (s[i], arrays[i]); - result_0 += ((s >> 0) & 3) as u32 * t[0] as u32; - result_1 += ((s >> 2) & 3) as u32 * t[1] as u32; - result_2 += ((s >> 4) & 3) as u32 * t[2] as u32; - result_3 += ((s >> 6) & 3) as u32 * t[3] as u32; - } - let mut buffer = [0u8; 4]; - if !remainder.is_empty() { - buffer[..remainder.len()].copy_from_slice(remainder); - let (s, t) = (s[n], buffer); - result_0 += ((s >> 0) & 3) as u32 * t[0] as u32; - result_1 += ((s >> 2) & 3) as u32 * t[1] as u32; - result_2 += ((s >> 4) & 3) as u32 * t[2] as u32; - result_3 += ((s >> 6) & 3) as u32 * t[3] as u32; - } - result_0 + result_1 + result_2 + result_3 - } -} diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs index cd0eba81..8b903578 100644 --- a/crates/simd_macros/src/lib.rs +++ b/crates/simd_macros/src/lib.rs @@ -99,7 +99,7 @@ pub fn multiversion( } let output = sig.output.clone(); let mut versions = quote::quote! {}; - let mut branches = quote::quote! {}; + let mut cold = quote::quote! {}; for version in attr.versions { let target = version.target.clone(); let name = syn::Ident::new( @@ -123,30 +123,64 @@ pub fn multiversion( fn #name < #generics_params > (#inputs) #output #generics_where { #block } }); } - branches.extend(quote::quote! { + cold.extend(quote::quote! { #[cfg(target_arch = #target_arch)] if crate::is_cpu_detected!(#target_cpu) #(&& crate::is_feature_detected!(#additional_target_features))* { - let _multiversion_internal: unsafe fn(#inputs) #output = #name; - CACHE.store(_multiversion_internal as *mut (), core::sync::atomic::Ordering::Relaxed); - return unsafe { _multiversion_internal(#(#arguments,)*) }; + let ptr = unsafe { std::mem::transmute::(#name) }; + CACHE.store(ptr as *mut (), core::sync::atomic::Ordering::Relaxed); + return ptr; } }); } + cold.extend(quote::quote! { + let ptr = unsafe { std::mem::transmute::(fallback)} ; + CACHE.store(ptr as *mut (), core::sync::atomic::Ordering::Relaxed); + ptr + }); quote::quote! { #versions fn fallback < #generics_params > (#inputs) #output #generics_where { #block } - #[inline(always)] - #(#attrs)* #vis #sig { + #[must_use] + pub(crate) fn pointer() -> fn(#inputs) #output { static CACHE: core::sync::atomic::AtomicPtr<()> = core::sync::atomic::AtomicPtr::new(core::ptr::null_mut()); - let cache = CACHE.load(core::sync::atomic::Ordering::Relaxed); - if !cache.is_null() { - let f = unsafe { core::mem::transmute::<*mut (), unsafe fn(#inputs) #output>(cache as _) }; - return unsafe { f(#(#arguments,)*) }; + #[must_use] + #[cold] + fn cold() -> fn(#inputs) #output { + #cold } - #branches - let _multiversion_internal: unsafe fn(#inputs) #output = fallback; - CACHE.store(_multiversion_internal as *mut (), core::sync::atomic::Ordering::Relaxed); - unsafe { _multiversion_internal(#(#arguments,)*) } + #[cfg(feature = "init")] + #[cfg(target_os = "linux")] + { + #[used] + #[unsafe(link_section = ".init_array")] + static INIT: extern "C" fn() -> usize = { + #[unsafe(link_section = ".text.startup")] + extern "C" fn f() -> usize { + let _ = cold(); + 0 + } + f + }; + let cache = unsafe { CACHE.as_ptr().read() }; + if !cache.is_null() { + let ptr = unsafe { std::mem::transmute::<*mut (), fn(#inputs) #output>(cache) }; + return ptr; + } + panic!("feature `init` is not supported on this platform") + } + #[allow(unreachable_code)] + { + let cache = CACHE.load(core::sync::atomic::Ordering::Relaxed); + if !cache.is_null() { + let ptr = unsafe { std::mem::transmute::<*mut (), fn(#inputs) #output>(cache) }; + return ptr; + } + cold() + } + } + #[inline(always)] + #(#attrs)* #vis #sig { + pointer()(#(#arguments,)*) } } .into() diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index fca0ed73..77f677fb 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -76,7 +76,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( VertexTuple::serialize(&VertexTuple { metadata: code.0.into_array(), payload: Some(payload), - elements: rabitq::original::binary::pack_code(&code.1), + elements: rabitq::b2::binary::pack_code(&code.1), pointers: vec![Pointer::new((u32::MAX, 0)); list_of_vector_bytes.len()], // a sentinel value }) }; @@ -144,7 +144,11 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( }; let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); let pointers_s = bump.alloc_slice(vertex_tuple.pointers()); - let score_s = O::process((vertex_tuple.metadata(), vertex_tuple.elements()), &lut); + let score_s = O::process( + dims, + (vertex_tuple.metadata(), vertex_tuple.elements()), + &lut, + ); candidates.push((Reverse(score_s), AlwaysEqual((pointers_s, s)))); } let mut iter = std::iter::from_fn(|| { @@ -178,7 +182,11 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( }; let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); - let score_v = O::process((vertex_tuple.metadata(), vertex_tuple.elements()), &lut); + let score_v = O::process( + dims, + (vertex_tuple.metadata(), vertex_tuple.elements()), + &lut, + ); candidates.push((Reverse(score_v), AlwaysEqual((pointers_v, v)))); } return Some((Reverse(dis_u), AlwaysEqual((pointers_u, u)))); diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs index 7fa46b51..16dcaeb8 100644 --- a/crates/vchordg/src/operator.rs +++ b/crates/vchordg/src/operator.rs @@ -34,8 +34,8 @@ pub trait Vector: VectorOwned { ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)); fn pack(elements: Vec, metadata: Self::Metadata) -> Self; - fn code(vector: Self::Borrowed<'_>) -> rabitq::b1::Code; - fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b1::binary::BinaryLut; + fn code(vector: Self::Borrowed<'_>) -> rabitq::b2::Code; + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b2::binary::BinaryLut; } impl Vector for VectOwned { @@ -69,12 +69,12 @@ impl Vector for VectOwned { VectOwned::new(elements) } - fn code(vector: Self::Borrowed<'_>) -> rabitq::b1::Code { - rabitq::b1::code(vector.slice()) + fn code(vector: Self::Borrowed<'_>) -> rabitq::b2::Code { + rabitq::b2::code(vector.slice()) } - fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b1::binary::BinaryLut { - rabitq::b1::binary::preprocess(vector.slice()) + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b2::binary::BinaryLut { + rabitq::b2::binary::preprocess(vector.slice()) } } @@ -109,12 +109,12 @@ impl Vector for VectOwned { VectOwned::new(elements) } - fn code(vector: Self::Borrowed<'_>) -> rabitq::b1::Code { - rabitq::b1::code(&f16::vector_to_f32(vector.slice())) + fn code(vector: Self::Borrowed<'_>) -> rabitq::b2::Code { + rabitq::b2::code(&f16::vector_to_f32(vector.slice())) } - fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b1::binary::BinaryLut { - rabitq::b1::binary::preprocess(&f16::vector_to_f32(vector.slice())) + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b2::binary::BinaryLut { + rabitq::b2::binary::preprocess(&f16::vector_to_f32(vector.slice())) } } @@ -132,7 +132,7 @@ pub trait Operator: 'static + Debug + Copy { Output = Distance, >; - fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance; + fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance; fn distance( lhs: ::Borrowed<'_>, rhs: ::Borrowed<'_>, @@ -157,11 +157,11 @@ impl Operator for Op, L2S> { type DistanceAccessor = DistanceAccessor, L2S>; - fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance { - use rabitq::b1::CodeMetadata; - let value = rabitq::b1::binary::accumulate(code.1, &lut.1); + fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance { + use rabitq::b2::CodeMetadata; + let value = rabitq::b2::binary::accumulate(code.1, &lut.1); let (distance,) = - rabitq::b1::binary::half_process_l2(value, CodeMetadata::from_array(code.0), lut.0); + rabitq::b2::binary::half_process_l2(n, value, CodeMetadata::from_array(code.0), lut.0); Distance::from_f32(distance) } @@ -180,11 +180,11 @@ impl Operator for Op, Dot> { type DistanceAccessor = DistanceAccessor, Dot>; - fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance { - use rabitq::b1::CodeMetadata; - let value = rabitq::b1::binary::accumulate(code.1, &lut.1); + fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance { + use rabitq::b2::CodeMetadata; + let value = rabitq::b2::binary::accumulate(code.1, &lut.1); let (distance,) = - rabitq::b1::binary::half_process_dot(value, CodeMetadata::from_array(code.0), lut.0); + rabitq::b2::binary::half_process_dot(n, value, CodeMetadata::from_array(code.0), lut.0); Distance::from_f32(distance) } @@ -203,11 +203,11 @@ impl Operator for Op, L2S> { type DistanceAccessor = DistanceAccessor, L2S>; - fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance { - use rabitq::b1::CodeMetadata; - let value = rabitq::b1::binary::accumulate(code.1, &lut.1); + fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance { + use rabitq::b2::CodeMetadata; + let value = rabitq::b2::binary::accumulate(code.1, &lut.1); let (distance,) = - rabitq::b1::binary::half_process_l2(value, CodeMetadata::from_array(code.0), lut.0); + rabitq::b2::binary::half_process_l2(n, value, CodeMetadata::from_array(code.0), lut.0); Distance::from_f32(distance) } @@ -226,11 +226,11 @@ impl Operator for Op, Dot> { type DistanceAccessor = DistanceAccessor, Dot>; - fn process(code: ([f32; 3], &[u64]), lut: &rabitq::b1::binary::BinaryLut) -> Distance { - use rabitq::b1::CodeMetadata; - let value = rabitq::b1::binary::accumulate(code.1, &lut.1); + fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance { + use rabitq::b2::CodeMetadata; + let value = rabitq::b2::binary::accumulate(code.1, &lut.1); let (distance,) = - rabitq::b1::binary::half_process_dot(value, CodeMetadata::from_array(code.0), lut.0); + rabitq::b2::binary::half_process_dot(n, value, CodeMetadata::from_array(code.0), lut.0); Distance::from_f32(distance) } diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs index 540101d1..9bb55319 100644 --- a/crates/vchordg/src/search.rs +++ b/crates/vchordg/src/search.rs @@ -65,7 +65,11 @@ where }; let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); let pointers_s = bump.alloc_slice(vertex_tuple.pointers()); - let score_s = O::process((vertex_tuple.metadata(), vertex_tuple.elements()), &lut); + let score_s = O::process( + dims, + (vertex_tuple.metadata(), vertex_tuple.elements()), + &lut, + ); candidates.push((Reverse(score_s), AlwaysEqual(pointers_s))); } let mut iter = std::iter::from_fn(move || { @@ -99,7 +103,11 @@ where }; let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); - let score_v = O::process((vertex_tuple.metadata(), vertex_tuple.elements()), &lut); + let score_v = O::process( + dims, + (vertex_tuple.metadata(), vertex_tuple.elements()), + &lut, + ); candidates.push((Reverse(score_v), AlwaysEqual(pointers_v))); } return Some((Reverse(dis_u), AlwaysEqual(payload_u))); diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index 4a6d670d..541f458e 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -39,7 +39,7 @@ impl VamanaIndexOptions { 32 } fn default_max_alpha() -> f32 { - 1.2 + 1.0 } fn default_ef_construction() -> u32 { 64 diff --git a/src/lib.rs b/src/lib.rs index 1cfd9147..d12a1233 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -45,7 +45,8 @@ pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap); pgrx::extension_sql_file!("./sql/finalize.sql", finalize); #[pgrx::pg_guard] -extern "C-unwind" fn _PG_init() { +#[unsafe(export_name = "_PG_init")] +unsafe extern "C-unwind" fn _pg_init() { if unsafe { pgrx::pg_sys::IsUnderPostmaster } { pgrx::error!("vchord must be loaded via shared_preload_libraries."); } From cb440dd2232dc2ca2799d7c9332697aac27b18a0 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 18 Jul 2025 14:00:06 +0800 Subject: [PATCH 178/324] feat: add option bits to control vchordg quantized bits (#290) Signed-off-by: usamoi --- crates/k_means/src/lib.rs | 8 +- crates/rabitq/src/b1.rs | 65 ------------- crates/rabitq/src/b2.rs | 65 ------------- crates/rabitq/src/b4.rs | 65 ------------- crates/rabitq/src/{original.rs => bit.rs} | 0 crates/rabitq/src/bits.rs | 105 +++++++++++++++++++++ crates/rabitq/src/{b8.rs => byte.rs} | 0 crates/rabitq/src/lib.rs | 8 +- crates/vchordg/src/build.rs | 5 +- crates/vchordg/src/insert.rs | 8 +- crates/vchordg/src/operator.rs | 108 +++++++++++++++------- crates/vchordg/src/search.rs | 4 + crates/vchordg/src/tuples.rs | 9 ++ crates/vchordg/src/types.rs | 23 +++-- crates/vchordrq/src/insert.rs | 2 +- crates/vchordrq/src/lib.rs | 2 +- crates/vchordrq/src/maintain.rs | 6 +- crates/vchordrq/src/operator.rs | 78 ++++++++-------- src/index/vchordg/algo.rs | 2 +- src/index/vchordg/am/am_build.rs | 70 +++++++------- src/index/vchordg/types.rs | 6 +- 21 files changed, 307 insertions(+), 332 deletions(-) delete mode 100644 crates/rabitq/src/b1.rs delete mode 100644 crates/rabitq/src/b2.rs delete mode 100644 crates/rabitq/src/b4.rs rename crates/rabitq/src/{original.rs => bit.rs} (100%) create mode 100644 crates/rabitq/src/bits.rs rename crates/rabitq/src/{b8.rs => byte.rs} (100%) diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index 91e1d678..508fb7d9 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use rabitq::original::block::BlockCode; +use rabitq::bit::block::BlockCode; use rabitq::packing::{any_pack, padding_pack}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; @@ -120,7 +120,7 @@ fn rabitq_index(n: usize, c: usize, samples: &[Vec], centroids: &[Vec] let branches = { let mut branches = Vec::new(); for centroid in centroids { - let code = rabitq::original::code(centroid); + let code = rabitq::bit::code(centroid); branches.push(code); } branches @@ -156,10 +156,10 @@ fn rabitq_index(n: usize, c: usize, samples: &[Vec], centroids: &[Vec] (0..n) .into_par_iter() .map(|i| { - let lut = rabitq::original::block::preprocess(&samples[i]); + let lut = rabitq::bit::block::preprocess(&samples[i]); let mut result = (f32::INFINITY, 0); for block in 0..c.div_ceil(32) { - let returns = rabitq::original::block::full_process_l2(blocks[block].code(), &lut); + let returns = rabitq::bit::block::full_process_l2(blocks[block].code(), &lut); let lowerbound = returns.map(|(rough, err)| rough - err * 1.9); for j in block * 32..std::cmp::min(block * 32 + 32, c) { if lowerbound[j - block * 32] < result.0 { diff --git a/crates/rabitq/src/b1.rs b/crates/rabitq/src/b1.rs deleted file mode 100644 index 9109d3b7..00000000 --- a/crates/rabitq/src/b1.rs +++ /dev/null @@ -1,65 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -pub use crate::extended::{Code, CodeMetadata}; - -pub fn code(vector: &[f32]) -> Code { - crate::extended::code::<1>(vector) -} - -pub mod binary { - pub fn pack_code(input: &[u8]) -> Vec { - crate::extended::pack_code::<1>(input) - .into_iter() - .flatten() - .collect() - } - - use crate::extended::CodeMetadata; - - const BITS: usize = 4; - - pub type BinaryLutMetadata = CodeMetadata; - pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); - pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); - - pub fn preprocess(vector: &[f32]) -> BinaryLut { - let (metadata, elements) = crate::extended::code::(vector); - (metadata, crate::extended::pack_code::<4>(&elements)) - } - - pub fn accumulate(lhs: &[u64], rhs: &[Vec; 4]) -> u32 { - crate::extended::accumulate::<1, 4>(lhs, rhs) - } - - pub fn half_process_dot( - n: u32, - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, - ) -> (f32,) { - let rough = crate::extended::half_process_dot::<1, BITS>(n, value, code, lut); - (rough,) - } - - pub fn half_process_l2( - n: u32, - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, - ) -> (f32,) { - let rough = crate::extended::half_process_l2::<1, BITS>(n, value, code, lut); - (rough,) - } -} diff --git a/crates/rabitq/src/b2.rs b/crates/rabitq/src/b2.rs deleted file mode 100644 index cd13e269..00000000 --- a/crates/rabitq/src/b2.rs +++ /dev/null @@ -1,65 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -pub use crate::extended::{Code, CodeMetadata}; - -pub fn code(vector: &[f32]) -> Code { - crate::extended::code::<2>(vector) -} - -pub mod binary { - pub fn pack_code(input: &[u8]) -> Vec { - crate::extended::pack_code::<2>(input) - .into_iter() - .flatten() - .collect() - } - - use crate::extended::CodeMetadata; - - const BITS: usize = 4; - - pub type BinaryLutMetadata = CodeMetadata; - pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); - pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); - - pub fn preprocess(vector: &[f32]) -> BinaryLut { - let (metadata, elements) = crate::extended::code::(vector); - (metadata, crate::extended::pack_code::<4>(&elements)) - } - - pub fn accumulate(lhs: &[u64], rhs: &[Vec; 4]) -> u32 { - crate::extended::accumulate::<2, 4>(lhs, rhs) - } - - pub fn half_process_dot( - n: u32, - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, - ) -> (f32,) { - let rough = crate::extended::half_process_dot::<2, BITS>(n, value, code, lut); - (rough,) - } - - pub fn half_process_l2( - n: u32, - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, - ) -> (f32,) { - let rough = crate::extended::half_process_l2::<2, BITS>(n, value, code, lut); - (rough,) - } -} diff --git a/crates/rabitq/src/b4.rs b/crates/rabitq/src/b4.rs deleted file mode 100644 index b0048304..00000000 --- a/crates/rabitq/src/b4.rs +++ /dev/null @@ -1,65 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -pub use crate::extended::{Code, CodeMetadata}; - -pub fn code(vector: &[f32]) -> Code { - crate::extended::code::<4>(vector) -} - -pub mod binary { - pub fn pack_code(input: &[u8]) -> Vec { - crate::extended::pack_code::<4>(input) - .into_iter() - .flatten() - .collect() - } - - use crate::extended::CodeMetadata; - - const BITS: usize = 4; - - pub type BinaryLutMetadata = CodeMetadata; - pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); - pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); - - pub fn preprocess(vector: &[f32]) -> BinaryLut { - let (metadata, elements) = crate::extended::code::(vector); - (metadata, crate::extended::pack_code::<4>(&elements)) - } - - pub fn accumulate(lhs: &[u64], rhs: &[Vec; 4]) -> u32 { - crate::extended::accumulate::<4, 4>(lhs, rhs) - } - - pub fn half_process_dot( - n: u32, - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, - ) -> (f32,) { - let rough = crate::extended::half_process_dot::<4, BITS>(n, value, code, lut); - (rough,) - } - - pub fn half_process_l2( - n: u32, - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, - ) -> (f32,) { - let rough = crate::extended::half_process_l2::<4, BITS>(n, value, code, lut); - (rough,) - } -} diff --git a/crates/rabitq/src/original.rs b/crates/rabitq/src/bit.rs similarity index 100% rename from crates/rabitq/src/original.rs rename to crates/rabitq/src/bit.rs diff --git a/crates/rabitq/src/bits.rs b/crates/rabitq/src/bits.rs new file mode 100644 index 00000000..47ab0eaa --- /dev/null +++ b/crates/rabitq/src/bits.rs @@ -0,0 +1,105 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub use crate::extended::{Code, CodeMetadata}; + +#[derive(Debug, Clone, Copy)] +#[repr(u8)] +pub enum Bits { + _1 = 1, + _2 = 2, +} + +impl TryFrom for Bits { + type Error = (); + + fn try_from(value: u8) -> Result { + match value { + 1 => Ok(Self::_1), + 2 => Ok(Self::_2), + _ => Err(()), + } + } +} + +pub fn code(bits: Bits, vector: &[f32]) -> Code { + match bits { + Bits::_1 => crate::extended::code::<1>(vector), + Bits::_2 => crate::extended::code::<2>(vector), + } +} + +pub fn pack_code(bits: Bits, input: &[u8]) -> Vec { + match bits { + Bits::_1 => crate::extended::pack_code::<1>(input) + .into_iter() + .flatten() + .collect(), + Bits::_2 => crate::extended::pack_code::<2>(input) + .into_iter() + .flatten() + .collect(), + } +} + +pub mod binary { + use crate::bits::Bits; + use crate::extended::CodeMetadata; + + const BITS: usize = 4; + + pub type BinaryLutMetadata = CodeMetadata; + pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); + + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::code::(vector); + (metadata, crate::extended::pack_code::(&elements)) + } + + pub fn accumulate(bits: Bits, lhs: &[u64], rhs: &[Vec; BITS]) -> u32 { + match bits { + Bits::_1 => crate::extended::accumulate::<1, BITS>(lhs, rhs), + Bits::_2 => crate::extended::accumulate::<2, BITS>(lhs, rhs), + } + } + + pub fn half_process_dot( + bits: Bits, + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = match bits { + Bits::_1 => crate::extended::half_process_dot::<1, BITS>(n, value, code, lut), + Bits::_2 => crate::extended::half_process_dot::<2, BITS>(n, value, code, lut), + }; + (rough,) + } + + pub fn half_process_l2( + bits: Bits, + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = match bits { + Bits::_1 => crate::extended::half_process_l2::<1, BITS>(n, value, code, lut), + Bits::_2 => crate::extended::half_process_l2::<2, BITS>(n, value, code, lut), + }; + (rough,) + } +} diff --git a/crates/rabitq/src/b8.rs b/crates/rabitq/src/byte.rs similarity index 100% rename from crates/rabitq/src/b8.rs rename to crates/rabitq/src/byte.rs diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index 0a1925f1..1925af5b 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -12,11 +12,9 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod b1; -pub mod b2; -pub mod b4; -pub mod b8; +pub mod bit; +pub mod bits; +pub mod byte; mod extended; -pub mod original; pub mod packing; pub mod rotate; diff --git a/crates/vchordg/src/build.rs b/crates/vchordg/src/build.rs index 53c6b773..ef455fc9 100644 --- a/crates/vchordg/src/build.rs +++ b/crates/vchordg/src/build.rs @@ -15,12 +15,12 @@ use crate::Opaque; use crate::operator::Operator; use crate::tuples::{MetaTuple, OptionPointer, Tuple}; -use crate::types::{VamanaIndexOptions, VectorOptions}; +use crate::types::{VchordgIndexOptions, VectorOptions}; use algo::{Page, PageGuard, RelationWrite}; pub fn build( vector_options: VectorOptions, - index_options: VamanaIndexOptions, + index_options: VchordgIndexOptions, index: &R, ) where R::Page: Page, @@ -53,6 +53,7 @@ pub fn build( drop(vector_guard); let serialized = MetaTuple::serialize(&MetaTuple { dims: vector_options.dims, + bits: index_options.bits, m: index_options.m, max_alpha: index_options.max_alpha, ef_construction: index_options.ef_construction, diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index 77f677fb..0e40a74d 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -24,6 +24,7 @@ use algo::accessor::LAccess; use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use algo::{Bump, Page, PageGuard, RelationRead, RelationWrite}; use always_equal::AlwaysEqual; +use rabitq::bits::Bits; use std::cmp::Reverse; use std::collections::VecDeque; use std::num::{NonZero, Wrapping}; @@ -45,6 +46,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let dims = meta_tuple.dims(); assert_eq!(dims, vector.dims(), "unmatched dimensions"); let start = meta_tuple.start(); + let bits = Bits::try_from(meta_tuple.bits()).expect("data corruption"); let m = meta_tuple.m(); let max_alpha = meta_tuple.max_alpha(); let ef = meta_tuple.ef_construction(); @@ -72,11 +74,11 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( left.chain(std::iter::once(right)).collect::>() }; let mut vertex_bytes = { - let code = O::Vector::code(vector); + let code = O::Vector::code(bits, vector); VertexTuple::serialize(&VertexTuple { metadata: code.0.into_array(), payload: Some(payload), - elements: rabitq::b2::binary::pack_code(&code.1), + elements: rabitq::bits::pack_code(bits, &code.1), pointers: vec![Pointer::new((u32::MAX, 0)); list_of_vector_bytes.len()], // a sentinel value }) }; @@ -145,6 +147,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); let pointers_s = bump.alloc_slice(vertex_tuple.pointers()); let score_s = O::process( + bits, dims, (vertex_tuple.metadata(), vertex_tuple.elements()), &lut, @@ -183,6 +186,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); let score_v = O::process( + bits, dims, (vertex_tuple.metadata(), vertex_tuple.elements()), &lut, diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs index 16dcaeb8..a4ad88ac 100644 --- a/crates/vchordg/src/operator.rs +++ b/crates/vchordg/src/operator.rs @@ -16,6 +16,7 @@ use crate::types::DistanceKind; use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S}; use distance::Distance; use half::f16; +use rabitq::bits::Bits; use simd::Floating; use std::fmt::Debug; use std::marker::PhantomData; @@ -34,8 +35,8 @@ pub trait Vector: VectorOwned { ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)); fn pack(elements: Vec, metadata: Self::Metadata) -> Self; - fn code(vector: Self::Borrowed<'_>) -> rabitq::b2::Code; - fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b2::binary::BinaryLut; + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code; + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut; } impl Vector for VectOwned { @@ -69,12 +70,12 @@ impl Vector for VectOwned { VectOwned::new(elements) } - fn code(vector: Self::Borrowed<'_>) -> rabitq::b2::Code { - rabitq::b2::code(vector.slice()) + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code { + rabitq::bits::code(bits, vector.slice()) } - fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b2::binary::BinaryLut { - rabitq::b2::binary::preprocess(vector.slice()) + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut { + rabitq::bits::binary::preprocess(vector.slice()) } } @@ -109,12 +110,12 @@ impl Vector for VectOwned { VectOwned::new(elements) } - fn code(vector: Self::Borrowed<'_>) -> rabitq::b2::Code { - rabitq::b2::code(&f16::vector_to_f32(vector.slice())) + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code { + rabitq::bits::code(bits, &f16::vector_to_f32(vector.slice())) } - fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::b2::binary::BinaryLut { - rabitq::b2::binary::preprocess(&f16::vector_to_f32(vector.slice())) + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut { + rabitq::bits::binary::preprocess(&f16::vector_to_f32(vector.slice())) } } @@ -132,7 +133,12 @@ pub trait Operator: 'static + Debug + Copy { Output = Distance, >; - fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance; + fn process( + bits: Bits, + n: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance; fn distance( lhs: ::Borrowed<'_>, rhs: ::Borrowed<'_>, @@ -157,11 +163,21 @@ impl Operator for Op, L2S> { type DistanceAccessor = DistanceAccessor, L2S>; - fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance { - use rabitq::b2::CodeMetadata; - let value = rabitq::b2::binary::accumulate(code.1, &lut.1); - let (distance,) = - rabitq::b2::binary::half_process_l2(n, value, CodeMetadata::from_array(code.0), lut.0); + fn process( + bits: Bits, + n: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2( + bits, + n, + value, + CodeMetadata::from_array(code.0), + lut.0, + ); Distance::from_f32(distance) } @@ -180,11 +196,21 @@ impl Operator for Op, Dot> { type DistanceAccessor = DistanceAccessor, Dot>; - fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance { - use rabitq::b2::CodeMetadata; - let value = rabitq::b2::binary::accumulate(code.1, &lut.1); - let (distance,) = - rabitq::b2::binary::half_process_dot(n, value, CodeMetadata::from_array(code.0), lut.0); + fn process( + bits: Bits, + n: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_dot( + bits, + n, + value, + CodeMetadata::from_array(code.0), + lut.0, + ); Distance::from_f32(distance) } @@ -203,11 +229,21 @@ impl Operator for Op, L2S> { type DistanceAccessor = DistanceAccessor, L2S>; - fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance { - use rabitq::b2::CodeMetadata; - let value = rabitq::b2::binary::accumulate(code.1, &lut.1); - let (distance,) = - rabitq::b2::binary::half_process_l2(n, value, CodeMetadata::from_array(code.0), lut.0); + fn process( + bits: Bits, + n: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2( + bits, + n, + value, + CodeMetadata::from_array(code.0), + lut.0, + ); Distance::from_f32(distance) } @@ -226,11 +262,21 @@ impl Operator for Op, Dot> { type DistanceAccessor = DistanceAccessor, Dot>; - fn process(n: u32, code: ([f32; 3], &[u64]), lut: &rabitq::b2::binary::BinaryLut) -> Distance { - use rabitq::b2::CodeMetadata; - let value = rabitq::b2::binary::accumulate(code.1, &lut.1); - let (distance,) = - rabitq::b2::binary::half_process_dot(n, value, CodeMetadata::from_array(code.0), lut.0); + fn process( + bits: Bits, + n: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_dot( + bits, + n, + value, + CodeMetadata::from_array(code.0), + lut.0, + ); Distance::from_f32(distance) } diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs index 9bb55319..78117bb4 100644 --- a/crates/vchordg/src/search.rs +++ b/crates/vchordg/src/search.rs @@ -24,6 +24,7 @@ use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use algo::{Bump, Page, RelationRead}; use always_equal::AlwaysEqual; use distance::Distance; +use rabitq::bits::Bits; use std::cmp::Reverse; use std::collections::VecDeque; use std::num::NonZero; @@ -46,6 +47,7 @@ where let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dims = meta_tuple.dims(); let start = meta_tuple.start(); + let bits = Bits::try_from(meta_tuple.bits()).expect("data corruption"); assert_eq!(dims, vector.dims(), "unmatched dimensions"); let ef = ef_search; let beam = beam_search; @@ -66,6 +68,7 @@ where let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); let pointers_s = bump.alloc_slice(vertex_tuple.pointers()); let score_s = O::process( + bits, dims, (vertex_tuple.metadata(), vertex_tuple.elements()), &lut, @@ -104,6 +107,7 @@ where let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); let score_v = O::process( + bits, dims, (vertex_tuple.metadata(), vertex_tuple.elements()), &lut, diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 43c95bba..6d38389e 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -42,6 +42,8 @@ pub trait WithWriter: Tuple { struct MetaTupleHeader { version: u64, dims: u32, + bits: u8, + _padding: [Padding; 7], m: u32, max_alpha: f32, ef_construction: u32, @@ -52,6 +54,7 @@ struct MetaTupleHeader { pub struct MetaTuple { pub dims: u32, + pub bits: u8, pub m: u32, pub max_alpha: f32, pub ef_construction: u32, @@ -67,6 +70,7 @@ impl Tuple for MetaTuple { match self { MetaTuple { dims, + bits, m, max_alpha, ef_construction, @@ -81,12 +85,14 @@ impl Tuple for MetaTuple { MetaTupleHeader { version: VERSION, dims: *dims, + bits: *bits, m: *m, max_alpha: *max_alpha, ef_construction: *ef_construction, beam_construction: *beam_construction, start: *start, skip: *skip, + _padding: Default::default(), } .as_bytes(), ); @@ -123,6 +129,9 @@ impl<'a> MetaTupleReader<'a> { pub fn dims(self) -> u32 { self.header.dims } + pub fn bits(self) -> u8 { + self.header.bits + } pub fn m(self) -> u32 { self.header.m } diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index 541f458e..f139d939 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -19,22 +19,28 @@ use vector::vect::{VectBorrowed, VectOwned}; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] -pub struct VamanaIndexOptions { - #[serde(default = "VamanaIndexOptions::default_m")] +pub struct VchordgIndexOptions { + #[serde(default = "VchordgIndexOptions::default_bits")] + #[validate(range(min = 1, max = 2))] + pub bits: u8, + #[serde(default = "VchordgIndexOptions::default_m")] #[validate(range(min = 1, max = 512))] pub m: u32, - #[serde(default = "VamanaIndexOptions::default_max_alpha")] + #[serde(default = "VchordgIndexOptions::default_max_alpha")] #[validate(range(min = 1.0, max = 2.0))] pub max_alpha: f32, - #[serde(default = "VamanaIndexOptions::default_ef_construction")] + #[serde(default = "VchordgIndexOptions::default_ef_construction")] #[validate(range(min = 1, max = 65535))] pub ef_construction: u32, - #[serde(default = "VamanaIndexOptions::default_beam_construction")] + #[serde(default = "VchordgIndexOptions::default_beam_construction")] #[validate(range(min = 1, max = 65535))] pub beam_construction: u32, } -impl VamanaIndexOptions { +impl VchordgIndexOptions { + fn default_bits() -> u8 { + 2 + } fn default_m() -> u32 { 32 } @@ -49,9 +55,10 @@ impl VamanaIndexOptions { } } -impl Default for VamanaIndexOptions { +impl Default for VchordgIndexOptions { fn default() -> Self { Self { + bits: Self::default_bits(), m: Self::default_m(), max_alpha: Self::default_max_alpha(), ef_construction: Self::default_ef_construction(), @@ -111,7 +118,7 @@ pub struct VectorOptions { impl VectorOptions { pub fn validate_self(&self) -> Result<(), ValidationError> { match (self.v, self.d, self.dims) { - (_, _, 1..=60000) => Ok(()), + (_, _, 1..=30000) => Ok(()), _ => Err(ValidationError::new("invalid vector options")), } } diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index 563fcb12..ab43a90a 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -174,7 +174,7 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( payload: Some(payload), prefetch, head, - elements: rabitq::original::binary::pack_code(&code.1), + elements: rabitq::bit::binary::pack_code(&code.1), }); tape::append( diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index 040bb3a9..a842e0f5 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -58,7 +58,7 @@ pub struct Opaque { unsafe impl algo::Opaque for Opaque {} pub(crate) struct Branch { - pub code: rabitq::original::Code, + pub code: rabitq::bit::Code, pub delta: f32, pub prefetch: Vec, pub head: u16, diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs index e2b52186..74d621b1 100644 --- a/crates/vchordrq/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -149,7 +149,7 @@ where let signs = unpacked[i].iter().flat_map(f).collect::>(); ( ( - rabitq::original::CodeMetadata { + rabitq::bit::CodeMetadata { dis_u_2: metadata[0][i], factor_cnt: metadata[1][i], factor_ip: metadata[2][i], @@ -178,7 +178,7 @@ where .collect::>(); ( ( - rabitq::original::CodeMetadata { + rabitq::bit::CodeMetadata { dis_u_2: metadata[0], factor_cnt: metadata[1], factor_ip: metadata[2], @@ -204,7 +204,7 @@ where branch.code.0.factor_ip, branch.code.0.factor_err, ], - elements: rabitq::original::binary::pack_code(&branch.code.1), + elements: rabitq::bit::binary::pack_code(&branch.code.1), delta: branch.delta, prefetch: branch.prefetch, head: branch.head, diff --git a/crates/vchordrq/src/operator.rs b/crates/vchordrq/src/operator.rs index 04b0bd07..0ee9ac86 100644 --- a/crates/vchordrq/src/operator.rs +++ b/crates/vchordrq/src/operator.rs @@ -15,9 +15,9 @@ use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S, RAccess}; use distance::Distance; use half::f16; -use rabitq::original::CodeMetadata; -use rabitq::original::binary::{BinaryLut, BinaryLutMetadata}; -use rabitq::original::block::{BlockLut, BlockLutMetadata, STEP}; +use rabitq::bit::CodeMetadata; +use rabitq::bit::binary::{BinaryLut, BinaryLutMetadata}; +use rabitq::bit::block::{BlockLut, BlockLutMetadata, STEP}; use simd::Floating; use std::fmt::Debug; use std::marker::PhantomData; @@ -102,7 +102,7 @@ pub trait Vector: VectorOwned { fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut); - fn code(vector: Self::Borrowed<'_>) -> rabitq::original::Code; + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code; fn squared_norm(vector: Self::Borrowed<'_>) -> f32; } @@ -143,15 +143,15 @@ impl Vector for VectOwned { } fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { - rabitq::original::block::preprocess(vector.slice()) + rabitq::bit::block::preprocess(vector.slice()) } fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { - rabitq::original::preprocess(vector.slice()) + rabitq::bit::preprocess(vector.slice()) } - fn code(vector: Self::Borrowed<'_>) -> rabitq::original::Code { - rabitq::original::code(vector.slice()) + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code { + rabitq::bit::code(vector.slice()) } fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { @@ -195,15 +195,15 @@ impl Vector for VectOwned { } fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { - rabitq::original::block::preprocess(&f16::vector_to_f32(vector.slice())) + rabitq::bit::block::preprocess(&f16::vector_to_f32(vector.slice())) } fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { - rabitq::original::preprocess(&f16::vector_to_f32(vector.slice())) + rabitq::bit::preprocess(&f16::vector_to_f32(vector.slice())) } - fn code(vector: Self::Borrowed<'_>) -> rabitq::original::Code { - rabitq::original::code(&f16::vector_to_f32(vector.slice())) + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code { + rabitq::bit::code(&f16::vector_to_f32(vector.slice())) } fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { @@ -256,7 +256,7 @@ pub trait Operator: 'static + Debug + Copy { fn build( vector: ::Borrowed<'_>, centroid: Option, - ) -> (rabitq::original::Code, f32); + ) -> (rabitq::bit::Code, f32); } #[derive(Debug)] @@ -288,7 +288,7 @@ impl Operator for Op, L2S> { mut f: F, ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::original::binary::accumulate(elements, &lut.1); + let value = rabitq::bit::binary::accumulate(elements, &lut.1); f.call( value, CodeMetadata { @@ -313,9 +313,9 @@ impl Operator for Op, L2S> { _: f32, ) -> (f32, f32) { if !is_residual { - rabitq::original::block::half_process_l2(value, code, lut) + rabitq::bit::block::half_process_l2(value, code, lut) } else { - rabitq::original::block::half_process_l2_residual(value, code, lut, dis_f, delta) + rabitq::bit::block::half_process_l2_residual(value, code, lut, dis_f, delta) } } @@ -329,16 +329,16 @@ impl Operator for Op, L2S> { _: f32, ) -> (f32, f32) { if !is_residual { - rabitq::original::binary::half_process_l2(value, code, lut) + rabitq::bit::binary::half_process_l2(value, code, lut) } else { - rabitq::original::binary::half_process_l2_residual(value, code, lut, dis_f, delta) + rabitq::bit::binary::half_process_l2_residual(value, code, lut, dis_f, delta) } } fn build( vector: VectBorrowed<'_, f32>, centroid: Option, - ) -> (rabitq::original::Code, f32) { + ) -> (rabitq::bit::Code, f32) { if let Some(centroid) = centroid { let residual = VectOwned::new(f32::vector_sub(vector.slice(), centroid.slice())); let code = Self::Vector::code(residual.as_borrowed()); @@ -382,7 +382,7 @@ impl Operator for Op, Dot> { mut f: F, ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::original::binary::accumulate(elements, &lut.1); + let value = rabitq::bit::binary::accumulate(elements, &lut.1); f.call( value, CodeMetadata { @@ -407,9 +407,9 @@ impl Operator for Op, Dot> { norm: f32, ) -> (f32, f32) { if !is_residual { - rabitq::original::block::half_process_dot(value, code, lut) + rabitq::bit::block::half_process_dot(value, code, lut) } else { - rabitq::original::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + rabitq::bit::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) } } @@ -423,18 +423,16 @@ impl Operator for Op, Dot> { norm: f32, ) -> (f32, f32) { if !is_residual { - rabitq::original::binary::half_process_dot(value, code, lut) + rabitq::bit::binary::half_process_dot(value, code, lut) } else { - rabitq::original::binary::half_process_dot_residual( - value, code, lut, dis_f, delta, norm, - ) + rabitq::bit::binary::half_process_dot_residual(value, code, lut, dis_f, delta, norm) } } fn build( vector: VectBorrowed<'_, f32>, centroid: Option, - ) -> (rabitq::original::Code, f32) { + ) -> (rabitq::bit::Code, f32) { if let Some(centroid) = centroid { let residual = VectOwned::new(f32::vector_sub(vector.slice(), centroid.slice())); let code = Self::Vector::code(residual.as_borrowed()); @@ -478,7 +476,7 @@ impl Operator for Op, L2S> { mut f: F, ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::original::binary::accumulate(elements, &lut.1); + let value = rabitq::bit::binary::accumulate(elements, &lut.1); f.call( value, CodeMetadata { @@ -503,9 +501,9 @@ impl Operator for Op, L2S> { _: f32, ) -> (f32, f32) { if !is_residual { - rabitq::original::block::half_process_l2(value, code, lut) + rabitq::bit::block::half_process_l2(value, code, lut) } else { - rabitq::original::block::half_process_l2_residual(value, code, lut, dis_f, delta) + rabitq::bit::block::half_process_l2_residual(value, code, lut, dis_f, delta) } } @@ -519,16 +517,16 @@ impl Operator for Op, L2S> { _: f32, ) -> (f32, f32) { if !is_residual { - rabitq::original::binary::half_process_l2(value, code, lut) + rabitq::bit::binary::half_process_l2(value, code, lut) } else { - rabitq::original::binary::half_process_l2_residual(value, code, lut, dis_f, delta) + rabitq::bit::binary::half_process_l2_residual(value, code, lut, dis_f, delta) } } fn build( vector: VectBorrowed<'_, f16>, centroid: Option, - ) -> (rabitq::original::Code, f32) { + ) -> (rabitq::bit::Code, f32) { if let Some(centroid) = centroid { let residual = VectOwned::new(f16::vector_sub(vector.slice(), centroid.slice())); let code = Self::Vector::code(residual.as_borrowed()); @@ -572,7 +570,7 @@ impl Operator for Op, Dot> { mut f: F, ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::original::binary::accumulate(elements, &lut.1); + let value = rabitq::bit::binary::accumulate(elements, &lut.1); f.call( value, CodeMetadata { @@ -597,9 +595,9 @@ impl Operator for Op, Dot> { norm: f32, ) -> (f32, f32) { if !is_residual { - rabitq::original::block::half_process_dot(value, code, lut) + rabitq::bit::block::half_process_dot(value, code, lut) } else { - rabitq::original::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + rabitq::bit::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) } } @@ -613,18 +611,16 @@ impl Operator for Op, Dot> { norm: f32, ) -> (f32, f32) { if !is_residual { - rabitq::original::binary::half_process_dot(value, code, lut) + rabitq::bit::binary::half_process_dot(value, code, lut) } else { - rabitq::original::binary::half_process_dot_residual( - value, code, lut, dis_f, delta, norm, - ) + rabitq::bit::binary::half_process_dot_residual(value, code, lut, dis_f, delta, norm) } } fn build( vector: VectBorrowed<'_, f16>, centroid: Option, - ) -> (rabitq::original::Code, f32) { + ) -> (rabitq::bit::Code, f32) { if let Some(centroid) = centroid { let residual = VectOwned::new(f16::vector_sub(vector.slice(), centroid.slice())); let code = Self::Vector::code(residual.as_borrowed()); diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs index 5f9fd2ba..9c6e409d 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/algo.rs @@ -91,7 +91,7 @@ where } } -pub fn build(vector_options: VectorOptions, vchordg_options: VamanaIndexOptions, index: &R) +pub fn build(vector_options: VectorOptions, vchordg_options: VchordgIndexOptions, index: &R) where R: RelationRead + RelationWrite, R::Page: Page, diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index ac9f09fb..d32e5a42 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -16,7 +16,7 @@ use crate::datatype::typmod::Typmod; use crate::index::storage::PostgresRelation; use crate::index::vchordg::am::{Reloption, ctid_to_key, kv_to_pointer}; use crate::index::vchordg::opclass::{Opfamily, opfamily}; -use crate::index::vchordg::types::VamanaIndexingOptions; +use crate::index::vchordg::types::VchordgIndexingOptions; use pgrx::pg_sys::{Datum, ItemPointerData}; use std::ffi::CStr; use vchordg::types::*; @@ -198,9 +198,9 @@ pub unsafe extern "C-unwind" fn ambuild( let mut reporter = PostgresReporter {}; crate::index::vchordg::algo::build(vector_options, vchordg_options.index, &index); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); - let cache = vchordg_cached::VamanaCached::_0 {}; + let cache = vchordg_cached::VchordgCached::_0 {}; if let Some(leader) = unsafe { - VamanaLeader::enter( + VchordgLeader::enter( heap_relation, index_relation, (*index_info).ii_Concurrent, @@ -254,7 +254,7 @@ pub unsafe extern "C-unwind" fn ambuild( unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } -struct VamanaShared { +struct VchordgShared { /* Immutable state */ heaprelid: pgrx::pg_sys::Oid, indexrelid: pgrx::pg_sys::Oid, @@ -279,21 +279,21 @@ fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { ) } -struct VamanaLeader { +struct VchordgLeader { pcxt: *mut pgrx::pg_sys::ParallelContext, nparticipants: i32, snapshot: pgrx::pg_sys::Snapshot, - vchordgshared: *mut VamanaShared, + vchordgshared: *mut VchordgShared, tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, vchordgcached: *const u8, } -impl VamanaLeader { +impl VchordgLeader { pub unsafe fn enter( heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, isconcurrent: bool, - cache: vchordg_cached::VamanaCached, + cache: vchordg_cached::VchordgCached, ) -> Option { let _cache = cache.serialize(); #[expect(clippy::drop_non_drop)] @@ -357,7 +357,7 @@ impl VamanaLeader { let est_tablescandesc = unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap_relation, snapshot) }; unsafe { - estimate_chunk(&mut (*pcxt).estimator, size_of::()); + estimate_chunk(&mut (*pcxt).estimator, size_of::()); estimate_keys(&mut (*pcxt).estimator, 1); estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); estimate_keys(&mut (*pcxt).estimator, 1); @@ -379,9 +379,9 @@ impl VamanaLeader { let vchordgshared = unsafe { let vchordgshared = - pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) - .cast::(); - vchordgshared.write(VamanaShared { + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) + .cast::(); + vchordgshared.write(VchordgShared { heaprelid: (*heap_relation).rd_id, indexrelid: (*index_relation).rd_id, isconcurrent, @@ -450,7 +450,7 @@ impl VamanaLeader { } } -impl Drop for VamanaLeader { +impl Drop for VchordgLeader { fn drop(&mut self) { if !std::thread::panicking() { unsafe { @@ -472,7 +472,7 @@ pub unsafe extern "C-unwind" fn vchordg_parallel_build_main( toc: *mut pgrx::pg_sys::shm_toc, ) { let vchordgshared = unsafe { - pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() }; let tablescandesc = unsafe { pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) @@ -522,12 +522,12 @@ unsafe fn parallel_build( heap_relation: pgrx::pg_sys::Relation, index_info: *mut pgrx::pg_sys::IndexInfo, tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, - vchordgshared: *mut VamanaShared, + vchordgshared: *mut VchordgShared, vchordgcached: *const u8, mut callback: impl FnMut(u64), ) { - use vchordg_cached::VamanaCachedReader; - let cached = VamanaCachedReader::deserialize_ref(unsafe { + use vchordg_cached::VchordgCachedReader; + let cached = VchordgCachedReader::deserialize_ref(unsafe { let bytes = (vchordgcached as *const u64).read_unaligned(); std::slice::from_raw_parts(vchordgcached.add(8), bytes as _) }); @@ -544,7 +544,7 @@ unsafe fn parallel_build( scan, }; match cached { - VamanaCachedReader::_0(_) => { + VchordgCachedReader::_0(_) => { heap.traverse(true, move |(ctid, store)| { for (vector, extra) in store { let key = ctid_to_key(ctid); @@ -579,7 +579,7 @@ pub unsafe extern "C-unwind" fn ambuildempty(_index_relation: pgrx::pg_sys::Rela unsafe fn options( index_relation: pgrx::pg_sys::Relation, -) -> (VectorOptions, VamanaIndexingOptions) { +) -> (VectorOptions, VchordgIndexingOptions) { let att = unsafe { &mut *(*index_relation).rd_att }; #[cfg(any( feature = "pg13", @@ -627,7 +627,7 @@ unsafe fn options( break 'rabitq Default::default(); } let s = unsafe { Reloption::options(reloption) }.to_string_lossy(); - match toml::from_str::(&s) { + match toml::from_str::(&s) { Ok(p) => p, Err(e) => pgrx::error!("failed to parse options: {}", e), } @@ -643,30 +643,30 @@ mod vchordg_cached { #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] - struct VamanaCachedHeader0 {} + struct VchordgCachedHeader0 {} #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] - struct VamanaCachedHeader1 { + struct VchordgCachedHeader1 { mapping_s: usize, mapping_e: usize, pages_s: usize, pages_e: usize, } - pub enum VamanaCached { + pub enum VchordgCached { _0 {}, } - impl VamanaCached { + impl VchordgCached { pub fn serialize(&self) -> Vec { let mut buffer = Vec::new(); match self { - VamanaCached::_0 {} => { + VchordgCached::_0 {} => { buffer.extend((0 as Tag).to_ne_bytes()); - buffer.extend(std::iter::repeat_n(0, size_of::())); - buffer[size_of::()..][..size_of::()] - .copy_from_slice(VamanaCachedHeader0 {}.as_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + buffer[size_of::()..][..size_of::()] + .copy_from_slice(VchordgCachedHeader0 {}.as_bytes()); } } buffer @@ -674,25 +674,25 @@ mod vchordg_cached { } #[derive(Debug, Clone, Copy)] - pub enum VamanaCachedReader<'a> { + pub enum VchordgCachedReader<'a> { #[allow(dead_code)] - _0(VamanaCachedReader0<'a>), + _0(VchordgCachedReader0<'a>), } #[derive(Debug, Clone, Copy)] - pub struct VamanaCachedReader0<'a> { + pub struct VchordgCachedReader0<'a> { #[allow(dead_code)] - header: &'a VamanaCachedHeader0, + header: &'a VchordgCachedHeader0, } - impl<'a> VamanaCachedReader<'a> { + impl<'a> VchordgCachedReader<'a> { pub fn deserialize_ref(source: &'a [u8]) -> Self { let tag = u64::from_ne_bytes(std::array::from_fn(|i| source[i])); match tag { 0 => { let checker = RefChecker::new(source); - let header: &VamanaCachedHeader0 = checker.prefix(size_of::()); - Self::_0(VamanaCachedReader0 { header }) + let header: &VchordgCachedHeader0 = checker.prefix(size_of::()); + Self::_0(VchordgCachedReader0 { header }) } _ => panic!("bad bytes"), } diff --git a/src/index/vchordg/types.rs b/src/index/vchordg/types.rs index c1511d0f..b9ff9c27 100644 --- a/src/index/vchordg/types.rs +++ b/src/index/vchordg/types.rs @@ -14,11 +14,11 @@ use serde::{Deserialize, Serialize}; use validator::Validate; -use vchordg::types::VamanaIndexOptions; +use vchordg::types::VchordgIndexOptions; #[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] -pub struct VamanaIndexingOptions { +pub struct VchordgIndexingOptions { #[serde(flatten)] - pub index: VamanaIndexOptions, + pub index: VchordgIndexOptions, } From 33bbd619aece58283c05add9af34466bc34a1475 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 18 Jul 2025 17:24:01 +0800 Subject: [PATCH 179/324] feat: add `v4`, `v3` version of `simd::bit` (#291) Signed-off-by: usamoi --- crates/simd/src/bit.rs | 762 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 739 insertions(+), 23 deletions(-) diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index e547ec4c..e0c0e378 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -18,8 +18,6 @@ pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { } mod reduce_sum_of_and { - // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 - #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -64,7 +62,146 @@ mod reduce_sum_of_and { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_and_v4(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::emulate_mm512_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mask_1 = _mm512_set1_epi16(0x00ff); + let mut and_0 = _mm512_setzero_si512(); + let mut and_1 = _mm512_setzero_si512(); + let mut and_2 = _mm512_setzero_si512(); + let mut and_3 = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + // + let and = _mm512_and_si512(x, y); + let and_lo = _mm512_and_si512(and, mask_0); + let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); + and_0 = _mm512_add_epi16(and_0, _mm512_and_si512(and_res_lo, mask_1)); + and_1 = _mm512_add_epi16(and_1, _mm512_srli_epi16(and_res_lo, 8)); + let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); + and_2 = _mm512_add_epi16(and_2, _mm512_and_si512(and_res_hi, mask_1)); + and_3 = _mm512_add_epi16(and_3, _mm512_srli_epi16(and_res_hi, 8)); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + // + let and = _mm512_and_si512(x, y); + let and_lo = _mm512_and_si512(and, mask_0); + let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); + and_0 = _mm512_add_epi16(and_0, _mm512_and_si512(and_res_lo, mask_1)); + and_1 = _mm512_add_epi16(and_1, _mm512_srli_epi16(and_res_lo, 8)); + let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); + and_2 = _mm512_add_epi16(and_2, _mm512_and_si512(and_res_hi, mask_1)); + and_3 = _mm512_add_epi16(and_3, _mm512_srli_epi16(and_res_hi, 8)); + } + emulate_mm512_reduce_add_epi16(_mm512_add_epi16( + _mm512_add_epi16(and_0, and_1), + _mm512_add_epi16(and_2, and_3), + )) as u32 + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_and_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_v4(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_and_v3(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::emulate_mm256_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mask_1 = _mm256_set1_epi16(0x00ff); + let mut and_0 = _mm256_setzero_si256(); + let mut and_1 = _mm256_setzero_si256(); + let mut and_2 = _mm256_setzero_si256(); + let mut and_3 = _mm256_setzero_si256(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + // + let and = _mm256_and_si256(x, y); + let and_lo = _mm256_and_si256(and, mask_0); + let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); + and_0 = _mm256_add_epi16(and_0, _mm256_and_si256(and_res_lo, mask_1)); + and_1 = _mm256_add_epi16(and_1, _mm256_srli_epi16(and_res_lo, 8)); + let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); + and_2 = _mm256_add_epi16(and_2, _mm256_and_si256(and_res_hi, mask_1)); + and_3 = _mm256_add_epi16(and_3, _mm256_srli_epi16(and_res_hi, 8)); + } + let mut and = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( + _mm256_add_epi16(and_0, and_1), + _mm256_add_epi16(and_2, and_3), + )) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + and += (x & y).count_ones(); + } + and + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_and_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_v3(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -82,8 +219,6 @@ pub fn reduce_sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { } mod reduce_sum_of_or { - // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 - #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -128,7 +263,146 @@ mod reduce_sum_of_or { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_or_v4(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::emulate_mm512_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mask_1 = _mm512_set1_epi16(0x00ff); + let mut or_0 = _mm512_setzero_si512(); + let mut or_1 = _mm512_setzero_si512(); + let mut or_2 = _mm512_setzero_si512(); + let mut or_3 = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + // + let or = _mm512_or_si512(x, y); + let or_lo = _mm512_and_si512(or, mask_0); + let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); + or_0 = _mm512_add_epi16(or_0, _mm512_and_si512(or_res_lo, mask_1)); + or_1 = _mm512_add_epi16(or_1, _mm512_srli_epi16(or_res_lo, 8)); + let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); + or_2 = _mm512_add_epi16(or_2, _mm512_and_si512(or_res_hi, mask_1)); + or_3 = _mm512_add_epi16(or_3, _mm512_srli_epi16(or_res_hi, 8)); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + // + let or = _mm512_or_si512(x, y); + let or_lo = _mm512_and_si512(or, mask_0); + let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); + or_0 = _mm512_add_epi16(or_0, _mm512_and_si512(or_res_lo, mask_1)); + or_1 = _mm512_add_epi16(or_1, _mm512_srli_epi16(or_res_lo, 8)); + let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); + or_2 = _mm512_add_epi16(or_2, _mm512_and_si512(or_res_hi, mask_1)); + or_3 = _mm512_add_epi16(or_3, _mm512_srli_epi16(or_res_hi, 8)); + } + emulate_mm512_reduce_add_epi16(_mm512_add_epi16( + _mm512_add_epi16(or_0, or_1), + _mm512_add_epi16(or_2, or_3), + )) as u32 + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_or_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_or_v4(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_or_v3(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::emulate_mm256_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mask_1 = _mm256_set1_epi16(0x00ff); + let mut or_0 = _mm256_setzero_si256(); + let mut or_1 = _mm256_setzero_si256(); + let mut or_2 = _mm256_setzero_si256(); + let mut or_3 = _mm256_setzero_si256(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + // + let or = _mm256_or_si256(x, y); + let or_lo = _mm256_and_si256(or, mask_0); + let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); + or_0 = _mm256_add_epi16(or_0, _mm256_and_si256(or_res_lo, mask_1)); + or_1 = _mm256_add_epi16(or_1, _mm256_srli_epi16(or_res_lo, 8)); + let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); + or_2 = _mm256_add_epi16(or_2, _mm256_and_si256(or_res_hi, mask_1)); + or_3 = _mm256_add_epi16(or_3, _mm256_srli_epi16(or_res_hi, 8)); + } + let mut or = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( + _mm256_add_epi16(or_0, or_1), + _mm256_add_epi16(or_2, or_3), + )) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + or += (x | y).count_ones(); + } + or + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_or_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_or_v3(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] pub fn reduce_sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -146,8 +420,6 @@ pub fn reduce_sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { } mod reduce_sum_of_xor { - // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 - #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -192,7 +464,146 @@ mod reduce_sum_of_xor { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xor_v4(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::emulate_mm512_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mask_1 = _mm512_set1_epi16(0x00ff); + let mut xor_0 = _mm512_setzero_si512(); + let mut xor_1 = _mm512_setzero_si512(); + let mut xor_2 = _mm512_setzero_si512(); + let mut xor_3 = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + // + let xor = _mm512_xor_si512(x, y); + let xor_lo = _mm512_and_si512(xor, mask_0); + let xor_hi = _mm512_and_si512(_mm512_srli_epi16(xor, 4), mask_0); + let xor_res_lo = _mm512_shuffle_epi8(lut, xor_lo); + xor_0 = _mm512_add_epi16(xor_0, _mm512_and_si512(xor_res_lo, mask_1)); + xor_1 = _mm512_add_epi16(xor_1, _mm512_srli_epi16(xor_res_lo, 8)); + let xor_res_hi = _mm512_shuffle_epi8(lut, xor_hi); + xor_2 = _mm512_add_epi16(xor_2, _mm512_and_si512(xor_res_hi, mask_1)); + xor_3 = _mm512_add_epi16(xor_3, _mm512_srli_epi16(xor_res_hi, 8)); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + // + let xor = _mm512_xor_si512(x, y); + let xor_lo = _mm512_and_si512(xor, mask_0); + let xor_hi = _mm512_and_si512(_mm512_srli_epi16(xor, 4), mask_0); + let xor_res_lo = _mm512_shuffle_epi8(lut, xor_lo); + xor_0 = _mm512_add_epi16(xor_0, _mm512_and_si512(xor_res_lo, mask_1)); + xor_1 = _mm512_add_epi16(xor_1, _mm512_srli_epi16(xor_res_lo, 8)); + let xor_res_hi = _mm512_shuffle_epi8(lut, xor_hi); + xor_2 = _mm512_add_epi16(xor_2, _mm512_and_si512(xor_res_hi, mask_1)); + xor_3 = _mm512_add_epi16(xor_3, _mm512_srli_epi16(xor_res_hi, 8)); + } + emulate_mm512_reduce_add_epi16(_mm512_add_epi16( + _mm512_add_epi16(xor_0, xor_1), + _mm512_add_epi16(xor_2, xor_3), + )) as u32 + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_xor_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_xor_v4(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_xor_v3(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::emulate_mm256_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mask_1 = _mm256_set1_epi16(0x00ff); + let mut xor_0 = _mm256_setzero_si256(); + let mut xor_1 = _mm256_setzero_si256(); + let mut xor_2 = _mm256_setzero_si256(); + let mut xor_3 = _mm256_setzero_si256(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + // + let xor = _mm256_xor_si256(x, y); + let xor_lo = _mm256_and_si256(xor, mask_0); + let xor_hi = _mm256_and_si256(_mm256_srli_epi16(xor, 4), mask_0); + let xor_res_lo = _mm256_shuffle_epi8(lut, xor_lo); + xor_0 = _mm256_add_epi16(xor_0, _mm256_and_si256(xor_res_lo, mask_1)); + xor_1 = _mm256_add_epi16(xor_1, _mm256_srli_epi16(xor_res_lo, 8)); + let xor_res_hi = _mm256_shuffle_epi8(lut, xor_hi); + xor_2 = _mm256_add_epi16(xor_2, _mm256_and_si256(xor_res_hi, mask_1)); + xor_3 = _mm256_add_epi16(xor_3, _mm256_srli_epi16(xor_res_hi, 8)); + } + let mut xor = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( + _mm256_add_epi16(xor_0, xor_1), + _mm256_add_epi16(xor_2, xor_3), + )) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + xor += (x ^ y).count_ones(); + } + xor + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_xor_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_xor_v3(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] pub fn reduce_sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -210,8 +621,6 @@ pub fn reduce_sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { } mod reduce_sum_of_and_or { - // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 - #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -248,7 +657,7 @@ mod reduce_sum_of_and_or { #[cfg(all(target_arch = "x86_64", test, not(miri)))] #[test] - fn reduce_sum_of_xor_v4_avx512vpopcntdq_test() { + fn reduce_sum_of_and_or_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); return; @@ -262,7 +671,195 @@ mod reduce_sum_of_and_or { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_and_or_v4(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + assert!(lhs.len() == rhs.len()); + use crate::emulate::emulate_mm512_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mask_1 = _mm512_set1_epi16(0x00ff); + let mut and_0 = _mm512_setzero_si512(); + let mut and_1 = _mm512_setzero_si512(); + let mut and_2 = _mm512_setzero_si512(); + let mut and_3 = _mm512_setzero_si512(); + let mut or_0 = _mm512_setzero_si512(); + let mut or_1 = _mm512_setzero_si512(); + let mut or_2 = _mm512_setzero_si512(); + let mut or_3 = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + a = unsafe { a.add(8) }; + b = unsafe { b.add(8) }; + n -= 8; + // + let and = _mm512_and_si512(x, y); + let and_lo = _mm512_and_si512(and, mask_0); + let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); + and_0 = _mm512_add_epi16(and_0, _mm512_and_si512(and_res_lo, mask_1)); + and_1 = _mm512_add_epi16(and_1, _mm512_srli_epi16(and_res_lo, 8)); + let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); + and_2 = _mm512_add_epi16(and_2, _mm512_and_si512(and_res_hi, mask_1)); + and_3 = _mm512_add_epi16(and_3, _mm512_srli_epi16(and_res_hi, 8)); + // + let or = _mm512_or_si512(x, y); + let or_lo = _mm512_and_si512(or, mask_0); + let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); + or_0 = _mm512_add_epi16(or_0, _mm512_and_si512(or_res_lo, mask_1)); + or_1 = _mm512_add_epi16(or_1, _mm512_srli_epi16(or_res_lo, 8)); + let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); + or_2 = _mm512_add_epi16(or_2, _mm512_and_si512(or_res_hi, mask_1)); + or_3 = _mm512_add_epi16(or_3, _mm512_srli_epi16(or_res_hi, 8)); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + // + let and = _mm512_and_si512(x, y); + let and_lo = _mm512_and_si512(and, mask_0); + let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); + and_0 = _mm512_add_epi16(and_0, _mm512_and_si512(and_res_lo, mask_1)); + and_1 = _mm512_add_epi16(and_1, _mm512_srli_epi16(and_res_lo, 8)); + let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); + and_2 = _mm512_add_epi16(and_2, _mm512_and_si512(and_res_hi, mask_1)); + and_3 = _mm512_add_epi16(and_3, _mm512_srli_epi16(and_res_hi, 8)); + // + let or = _mm512_or_si512(x, y); + let or_lo = _mm512_and_si512(or, mask_0); + let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); + or_0 = _mm512_add_epi16(or_0, _mm512_and_si512(or_res_lo, mask_1)); + or_1 = _mm512_add_epi16(or_1, _mm512_srli_epi16(or_res_lo, 8)); + let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); + or_2 = _mm512_add_epi16(or_2, _mm512_and_si512(or_res_hi, mask_1)); + or_3 = _mm512_add_epi16(or_3, _mm512_srli_epi16(or_res_hi, 8)); + } + ( + emulate_mm512_reduce_add_epi16(_mm512_add_epi16( + _mm512_add_epi16(and_0, and_1), + _mm512_add_epi16(and_2, and_3), + )) as u32, + emulate_mm512_reduce_add_epi16(_mm512_add_epi16( + _mm512_add_epi16(or_0, or_1), + _mm512_add_epi16(or_2, or_3), + )) as u32, + ) + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_and_or_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_or_v4(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_and_or_v3(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + assert!(lhs.len() == rhs.len()); + use crate::emulate::emulate_mm256_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mask_1 = _mm256_set1_epi16(0x00ff); + let mut and_0 = _mm256_setzero_si256(); + let mut and_1 = _mm256_setzero_si256(); + let mut and_2 = _mm256_setzero_si256(); + let mut and_3 = _mm256_setzero_si256(); + let mut or_0 = _mm256_setzero_si256(); + let mut or_1 = _mm256_setzero_si256(); + let mut or_2 = _mm256_setzero_si256(); + let mut or_3 = _mm256_setzero_si256(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + a = unsafe { a.add(4) }; + b = unsafe { b.add(4) }; + n -= 4; + // + let and = _mm256_and_si256(x, y); + let and_lo = _mm256_and_si256(and, mask_0); + let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); + and_0 = _mm256_add_epi16(and_0, _mm256_and_si256(and_res_lo, mask_1)); + and_1 = _mm256_add_epi16(and_1, _mm256_srli_epi16(and_res_lo, 8)); + let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); + and_2 = _mm256_add_epi16(and_2, _mm256_and_si256(and_res_hi, mask_1)); + and_3 = _mm256_add_epi16(and_3, _mm256_srli_epi16(and_res_hi, 8)); + // + let or = _mm256_or_si256(x, y); + let or_lo = _mm256_and_si256(or, mask_0); + let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); + or_0 = _mm256_add_epi16(or_0, _mm256_and_si256(or_res_lo, mask_1)); + or_1 = _mm256_add_epi16(or_1, _mm256_srli_epi16(or_res_lo, 8)); + let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); + or_2 = _mm256_add_epi16(or_2, _mm256_and_si256(or_res_hi, mask_1)); + or_3 = _mm256_add_epi16(or_3, _mm256_srli_epi16(or_res_hi, 8)); + } + let mut and = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( + _mm256_add_epi16(and_0, and_1), + _mm256_add_epi16(and_2, and_3), + )) as u32; + let mut or = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( + _mm256_add_epi16(or_0, or_1), + _mm256_add_epi16(or_2, or_3), + )) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + and += (x & y).count_ones(); + or += (x | y).count_ones(); + } + (and, or) + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_and_or_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_or_v3(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] pub fn reduce_sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -282,29 +879,27 @@ pub fn reduce_sum_of_x(this: &[u64]) -> u32 { } mod reduce_sum_of_x { - // FIXME: add manually-implemented SIMD version for AVX512 and AVX2 - #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] fn reduce_sum_of_x_v4_avx512vpopcntdq(this: &[u64]) -> u32 { use std::arch::x86_64::*; - let mut and = _mm512_setzero_si512(); + let mut sum = _mm512_setzero_si512(); let mut a = this.as_ptr(); let mut n = this.len(); while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; a = unsafe { a.add(8) }; n -= 8; - and = _mm512_add_epi64(and, _mm512_popcnt_epi64(x)); + sum = _mm512_add_epi64(sum, _mm512_popcnt_epi64(x)); } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; - and = _mm512_add_epi64(and, _mm512_popcnt_epi64(x)); + sum = _mm512_add_epi64(sum, _mm512_popcnt_epi64(x)); } - _mm512_reduce_add_epi64(and) as u32 + _mm512_reduce_add_epi64(sum) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -322,14 +917,135 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", "v4", "v3", "v2", "a2")] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[u64]) -> u32 { + use crate::emulate::emulate_mm512_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mask_1 = _mm512_set1_epi16(0x00ff); + let mut accu_0 = _mm512_setzero_si512(); + let mut accu_1 = _mm512_setzero_si512(); + let mut accu_2 = _mm512_setzero_si512(); + let mut accu_3 = _mm512_setzero_si512(); + let mut a = this.as_ptr(); + let mut n = this.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + a = unsafe { a.add(8) }; + n -= 8; + let xlo = _mm512_and_si512(x, mask_0); + let xhi = _mm512_and_si512(_mm512_srli_epi16(x, 4), mask_0); + let res_lo = _mm512_shuffle_epi8(lut, xlo); + accu_0 = _mm512_add_epi16(accu_0, _mm512_and_si512(res_lo, mask_1)); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_shuffle_epi8(lut, xhi); + accu_2 = _mm512_add_epi16(accu_2, _mm512_and_si512(res_hi, mask_1)); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let xlo = _mm512_and_si512(x, mask_0); + let xhi = _mm512_and_si512(_mm512_srli_epi16(x, 4), mask_0); + let res_lo = _mm512_shuffle_epi8(lut, xlo); + accu_0 = _mm512_add_epi16(accu_0, _mm512_and_si512(res_lo, mask_1)); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_shuffle_epi8(lut, xhi); + accu_2 = _mm512_add_epi16(accu_2, _mm512_and_si512(res_hi, mask_1)); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + } + emulate_mm512_reduce_add_epi16(_mm512_add_epi16( + _mm512_add_epi16(accu_0, accu_1), + _mm512_add_epi16(accu_2, accu_3), + )) as u32 + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let this = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_x_v4(&this) }; + let fallback = fallback(&this); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_v3(this: &[u64]) -> u32 { + use crate::emulate::emulate_mm256_reduce_add_epi16; + use std::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mask_1 = _mm256_set1_epi16(0x00ff); + let mut accu_0 = _mm256_setzero_si256(); + let mut accu_1 = _mm256_setzero_si256(); + let mut accu_2 = _mm256_setzero_si256(); + let mut accu_3 = _mm256_setzero_si256(); + let mut a = this.as_ptr(); + let mut n = this.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + a = unsafe { a.add(4) }; + n -= 4; + let xlo = _mm256_and_si256(x, mask_0); + let xhi = _mm256_and_si256(_mm256_srli_epi16(x, 4), mask_0); + let res_lo = _mm256_shuffle_epi8(lut, xlo); + accu_0 = _mm256_add_epi16(accu_0, _mm256_and_si256(res_lo, mask_1)); + accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); + let res_hi = _mm256_shuffle_epi8(lut, xhi); + accu_2 = _mm256_add_epi16(accu_2, _mm256_and_si256(res_hi, mask_1)); + accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); + } + let mut sum = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( + _mm256_add_epi16(accu_0, accu_1), + _mm256_add_epi16(accu_2, accu_3), + )) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + a = unsafe { a.add(1) }; + n -= 1; + let p = x.count_ones(); + sum += p; + } + sum + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let this = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_x_v3(&this) }; + let fallback = fallback(&this); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] pub fn reduce_sum_of_x(this: &[u64]) -> u32 { let n = this.len(); - let mut and = 0; + let mut sum = 0; for i in 0..n { - and += this[i].count_ones(); + sum += this[i].count_ones(); } - and + sum } } From eb04e14fb650ab715c51a47f185afa224f50f31c Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 18 Jul 2025 19:31:42 +0800 Subject: [PATCH 180/324] feat: use `sad` in intel intrinsics (#292) Signed-off-by: usamoi --- crates/simd/src/bit.rs | 301 ++++++++++++------------------------- crates/simd/src/emulate.rs | 9 ++ 2 files changed, 107 insertions(+), 203 deletions(-) diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index e0c0e378..65f342e2 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -67,16 +67,11 @@ mod reduce_sum_of_and { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_and_v4(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm512_reduce_add_epi16; use std::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); - let mask_1 = _mm512_set1_epi16(0x00ff); - let mut and_0 = _mm512_setzero_si512(); - let mut and_1 = _mm512_setzero_si512(); - let mut and_2 = _mm512_setzero_si512(); - let mut and_3 = _mm512_setzero_si512(); + let mut sum_and = _mm512_setzero_si512(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut n = lhs.len(); @@ -91,11 +86,10 @@ mod reduce_sum_of_and { let and_lo = _mm512_and_si512(and, mask_0); let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); - and_0 = _mm512_add_epi16(and_0, _mm512_and_si512(and_res_lo, mask_1)); - and_1 = _mm512_add_epi16(and_1, _mm512_srli_epi16(and_res_lo, 8)); let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); - and_2 = _mm512_add_epi16(and_2, _mm512_and_si512(and_res_hi, mask_1)); - and_3 = _mm512_add_epi16(and_3, _mm512_srli_epi16(and_res_hi, 8)); + let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); + sum_and = _mm512_add_epi64(sum_and, and_sad); } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -106,16 +100,12 @@ mod reduce_sum_of_and { let and_lo = _mm512_and_si512(and, mask_0); let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); - and_0 = _mm512_add_epi16(and_0, _mm512_and_si512(and_res_lo, mask_1)); - and_1 = _mm512_add_epi16(and_1, _mm512_srli_epi16(and_res_lo, 8)); let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); - and_2 = _mm512_add_epi16(and_2, _mm512_and_si512(and_res_hi, mask_1)); - and_3 = _mm512_add_epi16(and_3, _mm512_srli_epi16(and_res_hi, 8)); + let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); + sum_and = _mm512_add_epi64(sum_and, and_sad); } - emulate_mm512_reduce_add_epi16(_mm512_add_epi16( - _mm512_add_epi16(and_0, and_1), - _mm512_add_epi16(and_2, and_3), - )) as u32 + _mm512_reduce_add_epi64(sum_and) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -139,16 +129,12 @@ mod reduce_sum_of_and { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_and_v3(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm256_reduce_add_epi16; + use crate::emulate::emulate_mm256_reduce_add_epi64; use std::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); - let mask_1 = _mm256_set1_epi16(0x00ff); - let mut and_0 = _mm256_setzero_si256(); - let mut and_1 = _mm256_setzero_si256(); - let mut and_2 = _mm256_setzero_si256(); - let mut and_3 = _mm256_setzero_si256(); + let mut sum_and = _mm256_setzero_si256(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut n = lhs.len(); @@ -163,16 +149,12 @@ mod reduce_sum_of_and { let and_lo = _mm256_and_si256(and, mask_0); let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); - and_0 = _mm256_add_epi16(and_0, _mm256_and_si256(and_res_lo, mask_1)); - and_1 = _mm256_add_epi16(and_1, _mm256_srli_epi16(and_res_lo, 8)); let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); - and_2 = _mm256_add_epi16(and_2, _mm256_and_si256(and_res_hi, mask_1)); - and_3 = _mm256_add_epi16(and_3, _mm256_srli_epi16(and_res_hi, 8)); + let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); + sum_and = _mm256_add_epi64(sum_and, and_sad); } - let mut and = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( - _mm256_add_epi16(and_0, and_1), - _mm256_add_epi16(and_2, and_3), - )) as u32; + let mut and = emulate_mm256_reduce_add_epi64(sum_and) as u32; // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { let x = unsafe { a.read() }; @@ -268,16 +250,11 @@ mod reduce_sum_of_or { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_or_v4(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm512_reduce_add_epi16; use std::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); - let mask_1 = _mm512_set1_epi16(0x00ff); - let mut or_0 = _mm512_setzero_si512(); - let mut or_1 = _mm512_setzero_si512(); - let mut or_2 = _mm512_setzero_si512(); - let mut or_3 = _mm512_setzero_si512(); + let mut sum_or = _mm512_setzero_si512(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut n = lhs.len(); @@ -292,11 +269,10 @@ mod reduce_sum_of_or { let or_lo = _mm512_and_si512(or, mask_0); let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); - or_0 = _mm512_add_epi16(or_0, _mm512_and_si512(or_res_lo, mask_1)); - or_1 = _mm512_add_epi16(or_1, _mm512_srli_epi16(or_res_lo, 8)); let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); - or_2 = _mm512_add_epi16(or_2, _mm512_and_si512(or_res_hi, mask_1)); - or_3 = _mm512_add_epi16(or_3, _mm512_srli_epi16(or_res_hi, 8)); + let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); + sum_or = _mm512_add_epi64(sum_or, or_sad); } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -307,16 +283,12 @@ mod reduce_sum_of_or { let or_lo = _mm512_and_si512(or, mask_0); let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); - or_0 = _mm512_add_epi16(or_0, _mm512_and_si512(or_res_lo, mask_1)); - or_1 = _mm512_add_epi16(or_1, _mm512_srli_epi16(or_res_lo, 8)); let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); - or_2 = _mm512_add_epi16(or_2, _mm512_and_si512(or_res_hi, mask_1)); - or_3 = _mm512_add_epi16(or_3, _mm512_srli_epi16(or_res_hi, 8)); + let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); + sum_or = _mm512_add_epi64(sum_or, or_sad); } - emulate_mm512_reduce_add_epi16(_mm512_add_epi16( - _mm512_add_epi16(or_0, or_1), - _mm512_add_epi16(or_2, or_3), - )) as u32 + _mm512_reduce_add_epi64(sum_or) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -340,16 +312,12 @@ mod reduce_sum_of_or { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_or_v3(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm256_reduce_add_epi16; + use crate::emulate::emulate_mm256_reduce_add_epi64; use std::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); - let mask_1 = _mm256_set1_epi16(0x00ff); - let mut or_0 = _mm256_setzero_si256(); - let mut or_1 = _mm256_setzero_si256(); - let mut or_2 = _mm256_setzero_si256(); - let mut or_3 = _mm256_setzero_si256(); + let mut sum_or = _mm256_setzero_si256(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut n = lhs.len(); @@ -364,16 +332,12 @@ mod reduce_sum_of_or { let or_lo = _mm256_and_si256(or, mask_0); let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); - or_0 = _mm256_add_epi16(or_0, _mm256_and_si256(or_res_lo, mask_1)); - or_1 = _mm256_add_epi16(or_1, _mm256_srli_epi16(or_res_lo, 8)); let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); - or_2 = _mm256_add_epi16(or_2, _mm256_and_si256(or_res_hi, mask_1)); - or_3 = _mm256_add_epi16(or_3, _mm256_srli_epi16(or_res_hi, 8)); + let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); + sum_or = _mm256_add_epi64(sum_or, or_sad); } - let mut or = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( - _mm256_add_epi16(or_0, or_1), - _mm256_add_epi16(or_2, or_3), - )) as u32; + let mut or = emulate_mm256_reduce_add_epi64(sum_or) as u32; // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { let x = unsafe { a.read() }; @@ -469,16 +433,11 @@ mod reduce_sum_of_xor { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_xor_v4(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm512_reduce_add_epi16; use std::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); - let mask_1 = _mm512_set1_epi16(0x00ff); - let mut xor_0 = _mm512_setzero_si512(); - let mut xor_1 = _mm512_setzero_si512(); - let mut xor_2 = _mm512_setzero_si512(); - let mut xor_3 = _mm512_setzero_si512(); + let mut sum_xor = _mm512_setzero_si512(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut n = lhs.len(); @@ -493,11 +452,10 @@ mod reduce_sum_of_xor { let xor_lo = _mm512_and_si512(xor, mask_0); let xor_hi = _mm512_and_si512(_mm512_srli_epi16(xor, 4), mask_0); let xor_res_lo = _mm512_shuffle_epi8(lut, xor_lo); - xor_0 = _mm512_add_epi16(xor_0, _mm512_and_si512(xor_res_lo, mask_1)); - xor_1 = _mm512_add_epi16(xor_1, _mm512_srli_epi16(xor_res_lo, 8)); let xor_res_hi = _mm512_shuffle_epi8(lut, xor_hi); - xor_2 = _mm512_add_epi16(xor_2, _mm512_and_si512(xor_res_hi, mask_1)); - xor_3 = _mm512_add_epi16(xor_3, _mm512_srli_epi16(xor_res_hi, 8)); + let xor_res = _mm512_add_epi8(xor_res_lo, xor_res_hi); + let xor_sad = _mm512_sad_epu8(xor_res, _mm512_setzero_si512()); + sum_xor = _mm512_add_epi64(sum_xor, xor_sad); } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -508,16 +466,12 @@ mod reduce_sum_of_xor { let xor_lo = _mm512_and_si512(xor, mask_0); let xor_hi = _mm512_and_si512(_mm512_srli_epi16(xor, 4), mask_0); let xor_res_lo = _mm512_shuffle_epi8(lut, xor_lo); - xor_0 = _mm512_add_epi16(xor_0, _mm512_and_si512(xor_res_lo, mask_1)); - xor_1 = _mm512_add_epi16(xor_1, _mm512_srli_epi16(xor_res_lo, 8)); let xor_res_hi = _mm512_shuffle_epi8(lut, xor_hi); - xor_2 = _mm512_add_epi16(xor_2, _mm512_and_si512(xor_res_hi, mask_1)); - xor_3 = _mm512_add_epi16(xor_3, _mm512_srli_epi16(xor_res_hi, 8)); + let xor_res = _mm512_add_epi8(xor_res_lo, xor_res_hi); + let xor_sad = _mm512_sad_epu8(xor_res, _mm512_setzero_si512()); + sum_xor = _mm512_add_epi64(sum_xor, xor_sad); } - emulate_mm512_reduce_add_epi16(_mm512_add_epi16( - _mm512_add_epi16(xor_0, xor_1), - _mm512_add_epi16(xor_2, xor_3), - )) as u32 + _mm512_reduce_add_epi64(sum_xor) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -541,16 +495,12 @@ mod reduce_sum_of_xor { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_xor_v3(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm256_reduce_add_epi16; + use crate::emulate::emulate_mm256_reduce_add_epi64; use std::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); - let mask_1 = _mm256_set1_epi16(0x00ff); - let mut xor_0 = _mm256_setzero_si256(); - let mut xor_1 = _mm256_setzero_si256(); - let mut xor_2 = _mm256_setzero_si256(); - let mut xor_3 = _mm256_setzero_si256(); + let mut sum_xor = _mm256_setzero_si256(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut n = lhs.len(); @@ -565,16 +515,12 @@ mod reduce_sum_of_xor { let xor_lo = _mm256_and_si256(xor, mask_0); let xor_hi = _mm256_and_si256(_mm256_srli_epi16(xor, 4), mask_0); let xor_res_lo = _mm256_shuffle_epi8(lut, xor_lo); - xor_0 = _mm256_add_epi16(xor_0, _mm256_and_si256(xor_res_lo, mask_1)); - xor_1 = _mm256_add_epi16(xor_1, _mm256_srli_epi16(xor_res_lo, 8)); let xor_res_hi = _mm256_shuffle_epi8(lut, xor_hi); - xor_2 = _mm256_add_epi16(xor_2, _mm256_and_si256(xor_res_hi, mask_1)); - xor_3 = _mm256_add_epi16(xor_3, _mm256_srli_epi16(xor_res_hi, 8)); + let xor_res = _mm256_add_epi8(xor_res_lo, xor_res_hi); + let xor_sad = _mm256_sad_epu8(xor_res, _mm256_setzero_si256()); + sum_xor = _mm256_add_epi64(sum_xor, xor_sad); } - let mut xor = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( - _mm256_add_epi16(xor_0, xor_1), - _mm256_add_epi16(xor_2, xor_3), - )) as u32; + let mut xor = emulate_mm256_reduce_add_epi64(sum_xor) as u32; // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { let x = unsafe { a.read() }; @@ -676,20 +622,12 @@ mod reduce_sum_of_and_or { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_and_or_v4(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm512_reduce_add_epi16; use std::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); - let mask_1 = _mm512_set1_epi16(0x00ff); - let mut and_0 = _mm512_setzero_si512(); - let mut and_1 = _mm512_setzero_si512(); - let mut and_2 = _mm512_setzero_si512(); - let mut and_3 = _mm512_setzero_si512(); - let mut or_0 = _mm512_setzero_si512(); - let mut or_1 = _mm512_setzero_si512(); - let mut or_2 = _mm512_setzero_si512(); - let mut or_3 = _mm512_setzero_si512(); + let mut sum_and = _mm512_setzero_si512(); + let mut sum_or = _mm512_setzero_si512(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut n = lhs.len(); @@ -704,21 +642,19 @@ mod reduce_sum_of_and_or { let and_lo = _mm512_and_si512(and, mask_0); let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); - and_0 = _mm512_add_epi16(and_0, _mm512_and_si512(and_res_lo, mask_1)); - and_1 = _mm512_add_epi16(and_1, _mm512_srli_epi16(and_res_lo, 8)); let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); - and_2 = _mm512_add_epi16(and_2, _mm512_and_si512(and_res_hi, mask_1)); - and_3 = _mm512_add_epi16(and_3, _mm512_srli_epi16(and_res_hi, 8)); + let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); + sum_and = _mm512_add_epi64(sum_and, and_sad); // let or = _mm512_or_si512(x, y); let or_lo = _mm512_and_si512(or, mask_0); let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); - or_0 = _mm512_add_epi16(or_0, _mm512_and_si512(or_res_lo, mask_1)); - or_1 = _mm512_add_epi16(or_1, _mm512_srli_epi16(or_res_lo, 8)); let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); - or_2 = _mm512_add_epi16(or_2, _mm512_and_si512(or_res_hi, mask_1)); - or_3 = _mm512_add_epi16(or_3, _mm512_srli_epi16(or_res_hi, 8)); + let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); + sum_or = _mm512_add_epi64(sum_or, or_sad); } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -729,31 +665,23 @@ mod reduce_sum_of_and_or { let and_lo = _mm512_and_si512(and, mask_0); let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); - and_0 = _mm512_add_epi16(and_0, _mm512_and_si512(and_res_lo, mask_1)); - and_1 = _mm512_add_epi16(and_1, _mm512_srli_epi16(and_res_lo, 8)); let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); - and_2 = _mm512_add_epi16(and_2, _mm512_and_si512(and_res_hi, mask_1)); - and_3 = _mm512_add_epi16(and_3, _mm512_srli_epi16(and_res_hi, 8)); + let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); + sum_and = _mm512_add_epi64(sum_and, and_sad); // let or = _mm512_or_si512(x, y); let or_lo = _mm512_and_si512(or, mask_0); let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); - or_0 = _mm512_add_epi16(or_0, _mm512_and_si512(or_res_lo, mask_1)); - or_1 = _mm512_add_epi16(or_1, _mm512_srli_epi16(or_res_lo, 8)); let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); - or_2 = _mm512_add_epi16(or_2, _mm512_and_si512(or_res_hi, mask_1)); - or_3 = _mm512_add_epi16(or_3, _mm512_srli_epi16(or_res_hi, 8)); + let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); + sum_or = _mm512_add_epi64(sum_or, or_sad); } ( - emulate_mm512_reduce_add_epi16(_mm512_add_epi16( - _mm512_add_epi16(and_0, and_1), - _mm512_add_epi16(and_2, and_3), - )) as u32, - emulate_mm512_reduce_add_epi16(_mm512_add_epi16( - _mm512_add_epi16(or_0, or_1), - _mm512_add_epi16(or_2, or_3), - )) as u32, + _mm512_reduce_add_epi64(sum_and) as u32, + _mm512_reduce_add_epi64(sum_or) as u32, ) } @@ -778,20 +706,13 @@ mod reduce_sum_of_and_or { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_and_or_v3(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm256_reduce_add_epi16; + use crate::emulate::emulate_mm256_reduce_add_epi64; use std::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); - let mask_1 = _mm256_set1_epi16(0x00ff); - let mut and_0 = _mm256_setzero_si256(); - let mut and_1 = _mm256_setzero_si256(); - let mut and_2 = _mm256_setzero_si256(); - let mut and_3 = _mm256_setzero_si256(); - let mut or_0 = _mm256_setzero_si256(); - let mut or_1 = _mm256_setzero_si256(); - let mut or_2 = _mm256_setzero_si256(); - let mut or_3 = _mm256_setzero_si256(); + let mut sum_and = _mm256_setzero_si256(); + let mut sum_or = _mm256_setzero_si256(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); let mut n = lhs.len(); @@ -806,30 +727,22 @@ mod reduce_sum_of_and_or { let and_lo = _mm256_and_si256(and, mask_0); let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); - and_0 = _mm256_add_epi16(and_0, _mm256_and_si256(and_res_lo, mask_1)); - and_1 = _mm256_add_epi16(and_1, _mm256_srli_epi16(and_res_lo, 8)); let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); - and_2 = _mm256_add_epi16(and_2, _mm256_and_si256(and_res_hi, mask_1)); - and_3 = _mm256_add_epi16(and_3, _mm256_srli_epi16(and_res_hi, 8)); + let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); + sum_and = _mm256_add_epi64(sum_and, and_sad); // let or = _mm256_or_si256(x, y); let or_lo = _mm256_and_si256(or, mask_0); let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); - or_0 = _mm256_add_epi16(or_0, _mm256_and_si256(or_res_lo, mask_1)); - or_1 = _mm256_add_epi16(or_1, _mm256_srli_epi16(or_res_lo, 8)); let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); - or_2 = _mm256_add_epi16(or_2, _mm256_and_si256(or_res_hi, mask_1)); - or_3 = _mm256_add_epi16(or_3, _mm256_srli_epi16(or_res_hi, 8)); - } - let mut and = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( - _mm256_add_epi16(and_0, and_1), - _mm256_add_epi16(and_2, and_3), - )) as u32; - let mut or = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( - _mm256_add_epi16(or_0, or_1), - _mm256_add_epi16(or_2, or_3), - )) as u32; + let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); + sum_or = _mm256_add_epi64(sum_or, or_sad); + } + let mut and = emulate_mm256_reduce_add_epi64(sum_and) as u32; + let mut or = emulate_mm256_reduce_add_epi64(sum_or) as u32; // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { let x = unsafe { a.read() }; @@ -921,47 +834,37 @@ mod reduce_sum_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_x_v4(this: &[u64]) -> u32 { - use crate::emulate::emulate_mm512_reduce_add_epi16; use std::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); - let mask_1 = _mm512_set1_epi16(0x00ff); - let mut accu_0 = _mm512_setzero_si512(); - let mut accu_1 = _mm512_setzero_si512(); - let mut accu_2 = _mm512_setzero_si512(); - let mut accu_3 = _mm512_setzero_si512(); + let mut sum = _mm512_setzero_si512(); let mut a = this.as_ptr(); let mut n = this.len(); while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; a = unsafe { a.add(8) }; n -= 8; - let xlo = _mm512_and_si512(x, mask_0); - let xhi = _mm512_and_si512(_mm512_srli_epi16(x, 4), mask_0); - let res_lo = _mm512_shuffle_epi8(lut, xlo); - accu_0 = _mm512_add_epi16(accu_0, _mm512_and_si512(res_lo, mask_1)); - accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); - let res_hi = _mm512_shuffle_epi8(lut, xhi); - accu_2 = _mm512_add_epi16(accu_2, _mm512_and_si512(res_hi, mask_1)); - accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + let lo = _mm512_and_si512(x, mask_0); + let hi = _mm512_and_si512(_mm512_srli_epi16(x, 4), mask_0); + let res_lo = _mm512_shuffle_epi8(lut, lo); + let res_hi = _mm512_shuffle_epi8(lut, hi); + let res = _mm512_add_epi8(res_lo, res_hi); + let sad = _mm512_sad_epu8(res, _mm512_setzero_si512()); + sum = _mm512_add_epi64(sum, sad); } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; - let xlo = _mm512_and_si512(x, mask_0); - let xhi = _mm512_and_si512(_mm512_srli_epi16(x, 4), mask_0); - let res_lo = _mm512_shuffle_epi8(lut, xlo); - accu_0 = _mm512_add_epi16(accu_0, _mm512_and_si512(res_lo, mask_1)); - accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); - let res_hi = _mm512_shuffle_epi8(lut, xhi); - accu_2 = _mm512_add_epi16(accu_2, _mm512_and_si512(res_hi, mask_1)); - accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); - } - emulate_mm512_reduce_add_epi16(_mm512_add_epi16( - _mm512_add_epi16(accu_0, accu_1), - _mm512_add_epi16(accu_2, accu_3), - )) as u32 + let lo = _mm512_and_si512(x, mask_0); + let hi = _mm512_and_si512(_mm512_srli_epi16(x, 4), mask_0); + let res_lo = _mm512_shuffle_epi8(lut, lo); + let res_hi = _mm512_shuffle_epi8(lut, hi); + let res = _mm512_add_epi8(res_lo, res_hi); + let sad = _mm512_sad_epu8(res, _mm512_setzero_si512()); + sum = _mm512_add_epi64(sum, sad); + } + _mm512_reduce_add_epi64(sum) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -983,35 +886,27 @@ mod reduce_sum_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_x_v3(this: &[u64]) -> u32 { - use crate::emulate::emulate_mm256_reduce_add_epi16; + use crate::emulate::emulate_mm256_reduce_add_epi64; use std::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); - let mask_1 = _mm256_set1_epi16(0x00ff); - let mut accu_0 = _mm256_setzero_si256(); - let mut accu_1 = _mm256_setzero_si256(); - let mut accu_2 = _mm256_setzero_si256(); - let mut accu_3 = _mm256_setzero_si256(); + let mut sum = _mm256_setzero_si256(); let mut a = this.as_ptr(); let mut n = this.len(); while n >= 4 { let x = unsafe { _mm256_loadu_si256(a.cast()) }; a = unsafe { a.add(4) }; n -= 4; - let xlo = _mm256_and_si256(x, mask_0); - let xhi = _mm256_and_si256(_mm256_srli_epi16(x, 4), mask_0); - let res_lo = _mm256_shuffle_epi8(lut, xlo); - accu_0 = _mm256_add_epi16(accu_0, _mm256_and_si256(res_lo, mask_1)); - accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); - let res_hi = _mm256_shuffle_epi8(lut, xhi); - accu_2 = _mm256_add_epi16(accu_2, _mm256_and_si256(res_hi, mask_1)); - accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); - } - let mut sum = emulate_mm256_reduce_add_epi16(_mm256_add_epi16( - _mm256_add_epi16(accu_0, accu_1), - _mm256_add_epi16(accu_2, accu_3), - )) as u32; + let lo = _mm256_and_si256(x, mask_0); + let hi = _mm256_and_si256(_mm256_srli_epi16(x, 4), mask_0); + let res_lo = _mm256_shuffle_epi8(lut, lo); + let res_hi = _mm256_shuffle_epi8(lut, hi); + let res = _mm256_add_epi8(res_lo, res_hi); + let sad = _mm256_sad_epu8(res, _mm256_setzero_si256()); + sum = _mm256_add_epi64(sum, sad); + } + let mut sum = emulate_mm256_reduce_add_epi64(sum) as u32; // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { let x = unsafe { a.read() }; diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index cf6cf1a0..4f298d47 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -205,3 +205,12 @@ pub fn emulate_mm_reduce_max_ps(x: std::arch::x86_64::__m128) -> f32 { } f32::max(f32::max(x.0[0], x.0[1]), f32::max(x.0[2], x.0[3])) } + +#[inline] +#[cfg(target_arch = "x86_64")] +#[crate::target_cpu(enable = "v3")] +pub fn emulate_mm256_reduce_add_epi64(mut x: std::arch::x86_64::__m256i) -> i64 { + use std::arch::x86_64::*; + x = _mm256_add_epi64(x, _mm256_permute2f128_si256(x, x, 1)); + _mm256_extract_epi64(x, 0) + _mm256_extract_epi64(x, 1) +} From cba3ef0db459d39d625359874e33c2b3888449ee Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 22 Jul 2025 09:50:46 +0800 Subject: [PATCH 181/324] feat: arbitrary choice `alpha` (#293) Signed-off-by: usamoi --- crates/vchordg/src/build.rs | 2 +- crates/vchordg/src/insert.rs | 6 +++--- crates/vchordg/src/maintain.rs | 4 ++-- crates/vchordg/src/prune.rs | 7 ++----- crates/vchordg/src/tuples.rs | 25 ++++++++++++++++++------- crates/vchordg/src/types.rs | 24 ++++++++++++++++++------ src/index/vchordg/types.rs | 1 + src/index/vchordrq/types.rs | 3 +++ 8 files changed, 48 insertions(+), 24 deletions(-) diff --git a/crates/vchordg/src/build.rs b/crates/vchordg/src/build.rs index ef455fc9..38259f56 100644 --- a/crates/vchordg/src/build.rs +++ b/crates/vchordg/src/build.rs @@ -55,7 +55,7 @@ pub fn build( dims: vector_options.dims, bits: index_options.bits, m: index_options.m, - max_alpha: index_options.max_alpha, + alpha: index_options.alpha, ef_construction: index_options.ef_construction, beam_construction: index_options.beam_construction, start: OptionPointer::NONE, diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index 0e40a74d..648e098b 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -48,7 +48,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let start = meta_tuple.start(); let bits = Bits::try_from(meta_tuple.bits()).expect("data corruption"); let m = meta_tuple.m(); - let max_alpha = meta_tuple.max_alpha(); + let alpha = meta_tuple.alpha().to_vec(); let ef = meta_tuple.ef_construction(); let beam = meta_tuple.beam_construction(); let skip = meta_tuple.skip(); @@ -232,7 +232,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( (bump.alloc_slice(&pointers_t), t), trace.into_iter(), m, - max_alpha, + &alpha, |(_, u)| *u, O::DISTANCE == DistanceKind::L2S, ); @@ -284,7 +284,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( (pointers_u.to_vec(), u), trace.into_iter(), m, - max_alpha, + &alpha, |(_, u)| *u, O::DISTANCE == DistanceKind::L2S, ); diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs index fdde458b..58d39580 100644 --- a/crates/vchordg/src/maintain.rs +++ b/crates/vchordg/src/maintain.rs @@ -30,7 +30,7 @@ where let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let m = meta_tuple.m(); - let max_alpha = meta_tuple.max_alpha(); + let alpha = meta_tuple.alpha().to_vec(); let start = meta_tuple.start(); let link = meta_guard.get_opaque().link; drop(meta_guard); @@ -128,7 +128,7 @@ where (pointers_u.to_vec(), u), trace.into_iter(), m, - max_alpha, + &alpha, |(_, u)| *u, O::DISTANCE == DistanceKind::L2S, ); diff --git a/crates/vchordg/src/prune.rs b/crates/vchordg/src/prune.rs index ca1f7128..b3f1b552 100644 --- a/crates/vchordg/src/prune.rs +++ b/crates/vchordg/src/prune.rs @@ -21,7 +21,7 @@ pub fn prune( u: T, outs: impl Iterator, AlwaysEqual), V)>, m: u32, - max_alpha: f32, + alpha: &[f32], key: impl Fn(&T) -> K, is_l2s: bool, ) -> Vec<(Reverse, AlwaysEqual)> { @@ -32,15 +32,12 @@ pub fn prune( trace.retain(|((_, AlwaysEqual(v)), _)| key(v) != key(&u)); trace.sort_by_key(|&((Reverse(d), _), _)| d); // Nout(p) ← ∅ - let max_alpha = if is_l2s { max_alpha } else { 1.0 }; - let mut alpha = 1.0; let mut result = Vec::new(); - while alpha <= max_alpha { + for &alpha in if is_l2s { alpha } else { &[1.0] } { if result.len() == m as usize { break; } trace = robust_prune(&mut d, m, alpha, &mut result, trace); - alpha *= 1.2; } if !(result.len() == m as usize) { result.extend(trace); diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 6d38389e..05e6fd46 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -45,7 +45,8 @@ struct MetaTupleHeader { bits: u8, _padding: [Padding; 7], m: u32, - max_alpha: f32, + alpha_s: u16, + alpha_e: u16, ef_construction: u32, beam_construction: u32, start: OptionPointer, @@ -56,7 +57,7 @@ pub struct MetaTuple { pub dims: u32, pub bits: u8, pub m: u32, - pub max_alpha: f32, + pub alpha: Vec, pub ef_construction: u32, pub beam_construction: u32, pub start: OptionPointer, @@ -72,7 +73,7 @@ impl Tuple for MetaTuple { dims, bits, m, - max_alpha, + alpha, ef_construction, beam_construction, start, @@ -80,6 +81,13 @@ impl Tuple for MetaTuple { } => { buffer.extend((MAGIC as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let alpha_s = buffer.len() as u16; + buffer.extend(alpha.as_bytes()); + let alpha_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } // header buffer[size_of::()..][..size_of::()].copy_from_slice( MetaTupleHeader { @@ -87,7 +95,8 @@ impl Tuple for MetaTuple { dims: *dims, bits: *bits, m: *m, - max_alpha: *max_alpha, + alpha_s, + alpha_e, ef_construction: *ef_construction, beam_construction: *beam_construction, start: *start, @@ -113,7 +122,8 @@ impl WithReader for MetaTuple { panic!("deserialization: bad version number"); } let header: &MetaTupleHeader = checker.prefix(size_of::()); - MetaTupleReader { header } + let alpha = checker.bytes(header.alpha_s, header.alpha_e); + MetaTupleReader { header, alpha } } _ => panic!("deserialization: bad magic number"), } @@ -123,6 +133,7 @@ impl WithReader for MetaTuple { #[derive(Debug, Clone, Copy)] pub struct MetaTupleReader<'a> { header: &'a MetaTupleHeader, + alpha: &'a [f32], } impl<'a> MetaTupleReader<'a> { @@ -135,8 +146,8 @@ impl<'a> MetaTupleReader<'a> { pub fn m(self) -> u32 { self.header.m } - pub fn max_alpha(self) -> f32 { - self.header.max_alpha + pub fn alpha(self) -> &'a [f32] { + self.alpha } pub fn ef_construction(self) -> u32 { self.header.ef_construction diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index f139d939..d3df660c 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -26,9 +26,9 @@ pub struct VchordgIndexOptions { #[serde(default = "VchordgIndexOptions::default_m")] #[validate(range(min = 1, max = 512))] pub m: u32, - #[serde(default = "VchordgIndexOptions::default_max_alpha")] - #[validate(range(min = 1.0, max = 2.0))] - pub max_alpha: f32, + #[serde(default = "VchordgIndexOptions::default_alpha")] + #[validate(length(min = 1, max = 8), custom(function = VchordgIndexOptions::validate_alpha))] + pub alpha: Vec, #[serde(default = "VchordgIndexOptions::default_ef_construction")] #[validate(range(min = 1, max = 65535))] pub ef_construction: u32, @@ -44,8 +44,20 @@ impl VchordgIndexOptions { fn default_m() -> u32 { 32 } - fn default_max_alpha() -> f32 { - 1.0 + fn default_alpha() -> Vec { + vec![1.0, 1.2] + } + fn validate_alpha(alpha: &[f32]) -> Result<(), ValidationError> { + if !alpha.is_sorted() { + return Err(ValidationError::new("`alpha` should be in ascending order")); + } + if !alpha.iter().all(|x| (1.0..2.0).contains(x)) { + return Err(ValidationError::new("alpha is too large or too small")); + } + if alpha[0] != 1.0 { + return Err(ValidationError::new("`alpha` should contain `1.0`")); + } + Ok(()) } fn default_ef_construction() -> u32 { 64 @@ -60,7 +72,7 @@ impl Default for VchordgIndexOptions { Self { bits: Self::default_bits(), m: Self::default_m(), - max_alpha: Self::default_max_alpha(), + alpha: Self::default_alpha(), ef_construction: Self::default_ef_construction(), beam_construction: Self::default_beam_construction(), } diff --git a/src/index/vchordg/types.rs b/src/index/vchordg/types.rs index b9ff9c27..84f88fb2 100644 --- a/src/index/vchordg/types.rs +++ b/src/index/vchordg/types.rs @@ -20,5 +20,6 @@ use vchordg::types::VchordgIndexOptions; #[serde(deny_unknown_fields)] pub struct VchordgIndexingOptions { #[serde(flatten)] + #[validate(nested)] pub index: VchordgIndexOptions, } diff --git a/src/index/vchordrq/types.rs b/src/index/vchordrq/types.rs index a25b4782..1a1b21a0 100644 --- a/src/index/vchordrq/types.rs +++ b/src/index/vchordrq/types.rs @@ -109,6 +109,7 @@ impl Validate for VchordrqBuildSourceOptions { #[serde(rename_all = "snake_case")] pub struct VchordrqBuildOptions { #[serde(flatten)] + #[validate(nested)] pub source: VchordrqBuildSourceOptions, #[serde(default = "VchordrqBuildOptions::default_pin")] pub pin: bool, @@ -124,6 +125,8 @@ impl VchordrqBuildOptions { #[serde(deny_unknown_fields)] pub struct VchordrqIndexingOptions { #[serde(flatten)] + #[validate(nested)] pub index: VchordrqIndexOptions, + #[validate(nested)] pub build: VchordrqBuildOptions, } From 583cd54831c9292e7732ed99b53ec2ced2dc439e Mon Sep 17 00:00:00 2001 From: usamoi Date: Sat, 26 Jul 2025 13:43:54 +0800 Subject: [PATCH 182/324] readme: fix links (#297) Signed-off-by: usamoi --- Cargo.lock | 21 --------------------- Cargo.toml | 5 +---- README.md | 14 +++++++++++--- src/index/vchordg/am/am_build.rs | 6 ++++++ src/lib.rs | 8 ++------ 5 files changed, 20 insertions(+), 34 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bc47266e..d45a8a45 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -715,26 +715,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" -[[package]] -name = "jemalloc-sys" -version = "0.5.4+5.3.0-patched" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac6c1946e1cea1788cbfde01c993b52a10e2da07f4bac608228d1bed20bfebf2" -dependencies = [ - "cc", - "libc", -] - -[[package]] -name = "jemallocator" -version = "0.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0de374a9f8e63150e6f5e8a60cc14c668226d7a347d8aee1a45766e3c4dd3bc" -dependencies = [ - "jemalloc-sys", - "libc", -] - [[package]] name = "js-sys" version = "0.3.77" @@ -1645,7 +1625,6 @@ dependencies = [ "bumpalo", "distance", "half 2.6.0", - "jemallocator", "k_means", "mimalloc", "paste", diff --git a/Cargo.toml b/Cargo.toml index 1ff28765..91d5a3cb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,10 +44,7 @@ toml = "0.9.2" validator.workspace = true zerocopy.workspace = true -[target.'cfg(all(target_arch = "x86_64", target_os = "linux"))'.dependencies] -jemallocator = { version = "0.5.4", features = ["disable_initial_exec_tls"] } - -[target.'cfg(all(target_arch = "aarch64", target_os = "linux"))'.dependencies] +[target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), target_os = "linux"))'.dependencies] mimalloc = { version = "0.1.47", features = ["local_dynamic_tls"] } [lints] diff --git a/README.md b/README.md index 49f13564..3b109a33 100644 --- a/README.md +++ b/README.md @@ -112,9 +112,17 @@ SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5; For more usage, please read: -* [Indexing](https://docs.vectorchord.ai/vectorchord/usage/indexing.html) -* [Performance Tuning](https://docs.vectorchord.ai/vectorchord/usage/performance-tuning.html) -* [Advanced Features](https://docs.vectorchord.ai/vectorchord/usage/advanced-features.html) +- [Indexing](https://docs.vectorchord.ai/vectorchord/usage/indexing.html) +- [Search](https://docs.vectorchord.ai/vectorchord/usage/search.html) +- [Multi-Vector Retrieval](https://docs.vectorchord.ai/vectorchord/usage/indexing-with-maxsim-operators.html) +- [Similarity Filter](https://docs.vectorchord.ai/vectorchord/usage/range-query.html) +- [Performance Tuning](https://docs.vectorchord.ai/vectorchord/usage/performance-tuning.html) +- [Monitoring](https://docs.vectorchord.ai/vectorchord/usage/monitoring.html) +- [Prewarm](https://docs.vectorchord.ai/vectorchord/usage/prewarm.html) +- [Prefilter](https://docs.vectorchord.ai/vectorchord/usage/prefilter.html) +- [Prefetch](https://docs.vectorchord.ai/vectorchord/usage/prefetch.html) +- [Rerank In Table](https://docs.vectorchord.ai/vectorchord/usage/rerank-in-table.html) +- [External Build](https://docs.vectorchord.ai/vectorchord/usage/external-index-precomputation.html) > [!NOTE] > The partition parameter, `lists`, should be configured based on the number of rows. The following table provides guidance for this selection. When searching, remember to set [`vchordrq.probes`](https://docs.vectorchord.ai/vectorchord/usage/search.html#vchordrq-probes) based on the value of lists. diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index d32e5a42..5b0af1e8 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -186,6 +186,12 @@ pub unsafe extern "C-unwind" fn ambuild( if let Err(errors) = Validate::validate(&vchordg_options) { pgrx::error!("error while validating options: {}", errors); } + if vector_options.d != DistanceKind::L2S + && (vchordg_options.index.alpha != [1.0] && vchordg_options.index.alpha != [1.0, 1.2]) + { + let errors = "alpha not equal to `1.0` are only applicable to l2 and cosine distance."; + pgrx::warning!("warning while validating options: {errors}"); + } let opfamily = unsafe { opfamily(index_relation) }; let index = unsafe { PostgresRelation::new(index_relation) }; let heap = Heap { diff --git a/src/lib.rs b/src/lib.rs index d12a1233..f83eacf7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -66,11 +66,7 @@ compile_error!("This crate must be compiled with `-Cpanic=unwind`."); compile_error!("Target architecture is not supported."); #[cfg(not(miri))] -#[cfg(all(target_arch = "x86_64", target_os = "linux"))] -#[global_allocator] -static GLOBAL_ALLOCATOR: jemallocator::Jemalloc = jemallocator::Jemalloc; - -#[cfg(not(miri))] -#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] +#[cfg(target_os = "linux")] #[global_allocator] static GLOBAL_ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; From fa4cc957a50feec1a5141f3fbe86d2a18b9fffe8 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 28 Jul 2025 10:19:10 +0800 Subject: [PATCH 183/324] fix: disable debuginfo on profile release (#298) Profile `release` includes debug information for profiling, while it unnecessarily increased the binary size for users. Now disable debuginfo for profile `release` and add a profile `prof` for profiling. The size of the release binary has been reduced from 96.9MB to 4.4MB. Signed-off-by: usamoi --- Cargo.lock | 61 ++++++++++++++++++++--------------- Cargo.toml | 20 +++++++++--- crates/make/Cargo.toml | 2 +- crates/make/src/main.rs | 12 +++++-- crates/simd/Cargo.toml | 2 +- crates/simd_macros/Cargo.toml | 6 ++-- crates/vchordrq/Cargo.toml | 2 +- taplo.toml | 4 +-- 8 files changed, 67 insertions(+), 42 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d45a8a45..55691e19 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -80,7 +80,7 @@ version = "1.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6c8bdeb6047d8983be085bab0ba1472e6dc604e7041dbf6fcd5e71523014fae9" dependencies = [ - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -91,7 +91,7 @@ checksum = "403f75924867bb1033c59fbf0797484329750cfbe3c4325cd33127941fabc882" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -161,19 +161,19 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "cargo_toml" -version = "0.22.1" +version = "0.22.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02260d489095346e5cafd04dea8e8cb54d1d74fcd759022a9b72986ebe9a1257" +checksum = "374b7c592d9c00c1f4972ea58390ac6b18cbb6ab79011f3bdc90a0b82ca06b77" dependencies = [ "serde", - "toml 0.8.23", + "toml 0.9.2", ] [[package]] name = "cc" -version = "1.2.27" +version = "1.2.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d487aa071b5f64da6f19a3e848e3578944b726ee5a4854b82172f02aa876bfdc" +checksum = "deec109607ca693028562ed836a5f1c4b8bd77755c4e132fc5ce11b0b6211ae7" dependencies = [ "shlex", ] @@ -216,9 +216,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.40" +version = "4.5.41" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40b6887a1d8685cebccf115538db5c0efe625ccac9696ad45c409d96566e910f" +checksum = "be92d32e80243a54711e5d7ce823c35c41c9d929dc4ab58e1276f625841aadf9" dependencies = [ "clap_builder", "clap_derive", @@ -226,9 +226,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.40" +version = "4.5.41" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e0c66c08ce9f0c698cbce5c0279d0bb6ac936d8674174fe48f736533b964f59e" +checksum = "707eab41e9622f9139419d573eca0900137718000c517d47da73045f54331c3d" dependencies = [ "anstream", "anstyle", @@ -238,9 +238,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.5.40" +version = "4.5.41" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2c7947ae4cc3d851207c1adb5b5e260ff0cca11446b1d6d1423788e442257ce" +checksum = "ef4f52386a59ca4c860f7393bcf8abd8dfd91ecccc0f774635ff68e92eeef491" dependencies = [ "heck", "proc-macro2", @@ -416,7 +416,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.60.2", ] [[package]] @@ -545,7 +545,7 @@ version = "0.5.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589533453244b0995c858700322199b2becb13b627df2851f64a2775d024abcf" dependencies = [ - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -685,7 +685,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -1110,9 +1110,9 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "rand" -version = "0.9.1" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" dependencies = [ "rand_chacha", "rand_core", @@ -1203,15 +1203,15 @@ dependencies = [ [[package]] name = "rustix" -version = "1.0.7" +version = "1.0.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c71e83d6afe7ff64890ec6b71d6a69bb8a610ab78ce364b3352876bb4c801266" +checksum = "11181fbabf243db407ef8df94a6ce0b2f9a733bd8be4ad02b4eda9602296cac8" dependencies = [ "bitflags", "errno", "libc", "linux-raw-sys", - "windows-sys", + "windows-sys 0.60.2", ] [[package]] @@ -1285,9 +1285,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.140" +version = "1.0.141" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373" +checksum = "30b9eff21ebe718216c6ec64e1d9ac57087aad11efc64e32002bce4a0d4c03d3" dependencies = [ "itoa", "memchr", @@ -1806,7 +1806,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -1824,6 +1824,15 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.2", +] + [[package]] name = "windows-targets" version = "0.52.6" @@ -1954,9 +1963,9 @@ checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486" [[package]] name = "winnow" -version = "0.7.11" +version = "0.7.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74c7b26e3480b707944fc872477815d29a8e429d2f93a1ce000f5fa84a15cbcd" +checksum = "f3edebf492c8125044983378ecb5766203ad3b4c2f7a922bd7dd207f6d443e95" dependencies = [ "memchr", ] diff --git a/Cargo.toml b/Cargo.toml index 91d5a3cb..7a61bd65 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,11 +61,11 @@ edition = "2024" [workspace.dependencies] bumpalo = "3.19.0" half = { version = "2.6.0", features = ["zerocopy"] } -paste = "1" -rand = "0.9.1" +paste = "1.0.15" +rand = "0.9.2" rand_chacha = "0.9.0" seq-macro = "0.3.6" -serde = { version = "1", features = ["derive"] } +serde = { version = "1.0", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } zerocopy = { version = "0.8.26", features = ["derive"] } @@ -91,10 +91,20 @@ rust.unused_macro_rules = "warn" rust.unused_qualifications = "warn" [profile.dev] -lto = false +codegen-units = 256 +lto = "off" opt-level = 1 +debug = "full" +strip = "none" [profile.release] codegen-units = 1 -debug = true lto = "fat" +opt-level = 3 +debug = "none" +strip = "debuginfo" + +[profile.prof] +inherits = "release" +debug = "full" +strip = "none" diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index ec2b10eb..602c9bdb 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -5,7 +5,7 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.38", features = ["derive"] } +clap = { version = "4.5.41", features = ["derive"] } object = { version = "0.37.1", default-features = false, features = [ "read", "wasm", diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index 9afdd6de..133e6dcd 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -178,7 +178,11 @@ fn build( } let mut result = PathBuf::from("./target"); result.push(target); - result.push(if profile != "dev" { profile } else { "debug" }); + result.push(match profile { + "dev" | "test" => "debug", + "release" | "bench" => "release", + profile => profile, + }); result.push(format!("{}vchord{}", tsi.dll_prefix()?, tsi.dll_suffix()?)); Ok(result) } @@ -239,7 +243,11 @@ fn generate( } let mut result = PathBuf::from("./target"); result.push(target); - result.push(if profile != "dev" { profile } else { "debug" }); + result.push(match profile { + "dev" | "test" => "debug", + "release" | "bench" => "release", + profile => profile, + }); result.push(format!("pgrx_embed_vchord{}", tsi.exe_suffix()?)); let mut command; if !(tsi.is_unix && tsi.is_emscripten) { diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 3d298418..ae3b286e 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -18,7 +18,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.27" +cc = "1.2.30" which = "8.0.0" [lints] diff --git a/crates/simd_macros/Cargo.toml b/crates/simd_macros/Cargo.toml index f206bff0..6924ac1c 100644 --- a/crates/simd_macros/Cargo.toml +++ b/crates/simd_macros/Cargo.toml @@ -8,9 +8,9 @@ publish = false proc-macro = true [dependencies] -proc-macro2 = { version = "1", features = ["proc-macro"] } -quote = "1" -syn = { version = "2", default-features = false, features = [ +proc-macro2 = { version = "1.0", features = ["proc-macro"] } +quote = "1.0" +syn = { version = "2.0", default-features = false, features = [ "clone-impls", "full", "parsing", diff --git a/crates/vchordrq/Cargo.toml b/crates/vchordrq/Cargo.toml index a860c744..6f920599 100644 --- a/crates/vchordrq/Cargo.toml +++ b/crates/vchordrq/Cargo.toml @@ -13,7 +13,7 @@ simd = { path = "../simd" } vector = { path = "../vector" } half.workspace = true -pin-project = "1" +pin-project = "1.1" serde.workspace = true validator.workspace = true zerocopy.workspace = true diff --git a/taplo.toml b/taplo.toml index 41af6e14..b913433c 100644 --- a/taplo.toml +++ b/taplo.toml @@ -9,9 +9,7 @@ keys = [ "build-dependencies", "target.*.dependencies", "lints", - "patch.*", - "profile.*", "workspace.dependencies", - "lints.dependencies", + "workspace.lints", ] formatting = { reorder_keys = true, reorder_arrays = true, align_comments = true } From 128a5f48cf8ccc13a8f68b9be4239a0affdf9121 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 29 Jul 2025 17:27:43 +0800 Subject: [PATCH 184/324] readme: sync with docs (#300) Signed-off-by: usamoi --- README.md | 23 ++--------------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 3b109a33..380690d9 100644 --- a/README.md +++ b/README.md @@ -96,27 +96,21 @@ INSERT INTO items (embedding) SELECT ARRAY[random(), random(), random()]::real[] With VectorChord, you can create `vchordrq` indexes. ```SQL -CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops) WITH (options = $$ -residual_quantization = true -[build.internal] -lists = [] -$$); +CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops); ``` And then perform a vector search using `SELECT ... ORDER BY ... LIMIT ...`. ```SQL -SET vchordrq.probes TO ''; SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5; ``` For more usage, please read: - [Indexing](https://docs.vectorchord.ai/vectorchord/usage/indexing.html) -- [Search](https://docs.vectorchord.ai/vectorchord/usage/search.html) - [Multi-Vector Retrieval](https://docs.vectorchord.ai/vectorchord/usage/indexing-with-maxsim-operators.html) - [Similarity Filter](https://docs.vectorchord.ai/vectorchord/usage/range-query.html) -- [Performance Tuning](https://docs.vectorchord.ai/vectorchord/usage/performance-tuning.html) +- [PostgreSQL Tuning](https://docs.vectorchord.ai/vectorchord/usage/performance-tuning.html) - [Monitoring](https://docs.vectorchord.ai/vectorchord/usage/monitoring.html) - [Prewarm](https://docs.vectorchord.ai/vectorchord/usage/prewarm.html) - [Prefilter](https://docs.vectorchord.ai/vectorchord/usage/prefilter.html) @@ -124,19 +118,6 @@ For more usage, please read: - [Rerank In Table](https://docs.vectorchord.ai/vectorchord/usage/rerank-in-table.html) - [External Build](https://docs.vectorchord.ai/vectorchord/usage/external-index-precomputation.html) -> [!NOTE] -> The partition parameter, `lists`, should be configured based on the number of rows. The following table provides guidance for this selection. When searching, remember to set [`vchordrq.probes`](https://docs.vectorchord.ai/vectorchord/usage/search.html#vchordrq-probes) based on the value of lists. - - - -| Number of Rows $N$ | Recommended Number of Partitions $L$ | Example `lists` | -| -------------------------------------- | ------------------------------------ | --------------- | -| $N \in [0, 10^5)$ | N/A | `[]` | -| $N \in [10^5, 2 \times 10^6)$ | $L = \frac{N}{500}$ | `[2000]` | -| $N \in [2 \times 10^6, 5 \times 10^7)$ | $L \in [4 \sqrt{N}, 8 \sqrt{N}]$ | `[10000]` | -| $N \in [5 \times 10^7, \infty)$ | $L \in [8 \sqrt{N}, 16\sqrt{N}]$ | `[80000]` | - - ## License This software is licensed under a dual license model: From 676b13af5b6764436e55b6da4c2fc71325963c42 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 30 Jul 2025 15:00:00 +0800 Subject: [PATCH 185/324] ci: enable x86_64 windows and x86_64 macos (#301) job: +psql job: +psql_macos job: +psql_windows Signed-off-by: usamoi --- .github/workflows/check.yml | 87 +++++++++++++++--- .github/workflows/release.yml | 2 +- .gitignore | 1 - Cargo.lock | 31 ++++--- Cargo.toml | 2 +- crates/make/Cargo.toml | 2 +- crates/make/src/main.rs | 39 ++++++--- scripts/README.md | 76 ---------------- src/lib.rs | 2 +- tools/dev.md | 160 ---------------------------------- 10 files changed, 126 insertions(+), 276 deletions(-) delete mode 100644 tools/dev.md diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index ffecc957..61413a29 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -259,7 +259,7 @@ jobs: Maintainer: Tensorchord Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM Homepage: https://vectorchord.ai/ - License: AGPL-3 or Elastic-2" \ + License: AGPL-3.0-only or Elastic-2.0" \ > ./build/deb/DEBIAN/control (cd ./build/deb && find usr -type f -print0 | xargs -0 md5sum) > ./build/deb/DEBIAN/md5sums dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb @@ -282,10 +282,10 @@ jobs: strategy: matrix: - version: ["14", "17"] - arch: ["aarch64"] + version: ["13", "14", "15", "16", "17"] + arch: ["aarch64", "x86_64"] - runs-on: "macos-15" + runs-on: ${{ matrix.arch == 'aarch64' && 'macos-15' || 'macos-13' }} env: RUSTC_WRAPPER: sccache @@ -297,10 +297,10 @@ jobs: rustup set profile rustup default 1.89-beta + HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install llvm@18 echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install postgresql@${{ matrix.version }} - HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install pgvector brew services start postgresql@${{ matrix.version }} for i in {1..60}; do [ -S /tmp/.s.PGSQL.5432 ] && echo "PostgreSQL ready" && break || sleep 1; done [ -S /tmp/.s.PGSQL.5432 ] || echo "PostgreSQL socket not found after 60 seconds" @@ -309,6 +309,11 @@ jobs: brew services stop postgresql@${{ matrix.version }} echo PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV + mkdir ~/pgvector-install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.0.tar.gz | tar -xz -C ~/pgvector-install + make -C ~/pgvector-install/pgvector-0.8.0 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config + sudo make -C ~/pgvector-install/pgvector-0.8.0 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config install + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest - name: Set up Sccache @@ -367,7 +372,7 @@ jobs: strategy: matrix: - version: ["16"] + version: ["13", "14", "15", "16", "17"] arch: ["x86_64"] runs-on: "windows-2022" @@ -379,15 +384,44 @@ jobs: steps: - name: Set up Environment run: | + 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } + rustup set profile rustup default 1.89-beta - Invoke-WebRequest -Uri "https://get.enterprisedb.com/postgresql/postgresql-16.9-1-windows-x64-binaries.zip" -OutFile "$env:TEMP\postgresql-install.zip" + & 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise\Common7\Tools\Launch-VsDevShell.ps1' -HostArch amd64 -Arch amd64 + + if ( "${{ matrix.version }}" -eq "13" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-13.21-1-windows-x64-binaries.zip" + } + if ( "${{ matrix.version }}" -eq "14" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-14.18-1-windows-x64-binaries.zip" + } + if ( "${{ matrix.version }}" -eq "15" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-15.13-1-windows-x64-binaries.zip" + } + if ( "${{ matrix.version }}" -eq "16" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-16.9-1-windows-x64-binaries.zip" + } + if ( "${{ matrix.version }}" -eq "17" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-17.5-1-windows-x64-binaries.zip" + } + Invoke-WebRequest -Uri $postgresql_url -OutFile "$env:TEMP\postgresql-install.zip" Expand-Archive -Path "$env:TEMP\postgresql-install.zip" -DestinationPath "D:\postgresql-install" -Force Add-Content -Path $env:GITHUB_ENV -Value "PGRX_PG_CONFIG_PATH=D:\postgresql-install\pgsql\bin\pg_config.exe" - Add-Content -Path $env:GITHUB_ENV -Value "PGBIN=D:\postgresql-install\pgsql\bin" - Add-Content -Path $env:GITHUB_ENV -Value "PGDATA=D:\postgresql-install\pgsql\data" - Add-Content -Path $env:GITHUB_ENV -Value "PGROOT=D:\postgresql-install\pgsql" + D:\postgresql-install\pgsql\bin\initdb.exe -D D:\postgresql-install\pgsql\data -U postgres + D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data + D:\postgresql-install\pgsql\bin\createuser.exe -U postgres -s -r $env:USERNAME + D:\postgresql-install\pgsql\bin\createdb.exe -O $env:USERNAME $env:USERNAME + D:\postgresql-install\pgsql\bin\psql.exe -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + D:\postgresql-install\pgsql\bin\pg_ctl.exe stop -D D:\postgresql-install\pgsql\data + + Invoke-WebRequest -Uri "https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.0.zip" -OutFile "$env:TEMP\pgvector-install.zip" + Expand-Archive -Path "$env:TEMP\pgvector-install.zip" -DestinationPath "D:\pgvector-install" -Force + Push-Location -Path D:\pgvector-install\pgvector-0.8.0 + nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" + nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" install + Pop-Location Invoke-WebRequest -Uri https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-${{ matrix.arch }}-pc-windows-msvc.zip -OutFile "$env:TEMP\sqllogictest-install.zip" Expand-Archive -Path "$env:TEMP\sqllogictest-install.zip" -DestinationPath "D:\sqllogictest-install" -Force @@ -407,3 +441,36 @@ jobs: run: | cargo run -p make -- build -o ./build/raw --profile dev make PG_CONFIG="$env:PGRX_PG_CONFIG_PATH" install + + - name: Service + run: | + 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } + + D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data + D:\postgresql-install\pgsql\bin\psql.exe -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/general/*.slt' + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordg/*.slt' + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/*.slt' + if ( "${{ matrix.version }}" -eq "17" ) { + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/pg17/*.slt' + } + + - name: Package + env: + SEMVER: "0.0.0" + VERSION: ${{ matrix.version }} + ARCH: ${{ matrix.arch }} + run: | + Compress-Archive -Path ./build/raw/* -DestinationPath ./build/postgresql-$($env:VERSION)-vchord_$($env:SEMVER)_$($env:ARCH)-pc-windows-msvc.zip + + - name: Upload Artifacts + uses: actions/upload-artifact@v4 + with: + name: artifacts-psql_windows-${{ matrix.version }}-${{ matrix.arch }} + path: | + ./build/postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc.zip + compression-level: 9 + retention-days: 14 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 10adb4bc..175c0027 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -104,7 +104,7 @@ jobs: Maintainer: Tensorchord Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM Homepage: https://vectorchord.ai/ - License: AGPL-3 or Elastic-2" \ + License: AGPL-3.0-only or Elastic-2.0" \ > ./build/deb/DEBIAN/control (cd ./build/deb && find usr -type f -print0 | xargs -0 md5sum) > ./build/deb/DEBIAN/md5sums dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb diff --git a/.gitignore b/.gitignore index 47ad9c00..85d1619b 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,3 @@ .vscode .ignore build -.cargo/config.toml diff --git a/Cargo.lock b/Cargo.lock index 55691e19..a1fba94d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -166,7 +166,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "374b7c592d9c00c1f4972ea58390ac6b18cbb6ab79011f3bdc90a0b82ca06b77" dependencies = [ "serde", - "toml 0.9.2", + "toml 0.9.4", ] [[package]] @@ -216,9 +216,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.41" +version = "4.5.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be92d32e80243a54711e5d7ce823c35c41c9d929dc4ab58e1276f625841aadf9" +checksum = "ed87a9d530bb41a67537289bafcac159cb3ee28460e0a4571123d2a778a6a882" dependencies = [ "clap_builder", "clap_derive", @@ -226,9 +226,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.41" +version = "4.5.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "707eab41e9622f9139419d573eca0900137718000c517d47da73045f54331c3d" +checksum = "64f4f3f3c77c94aff3c7e9aac9a2ca1974a5adf392a8bb751e827d6d127ab966" dependencies = [ "anstream", "anstyle", @@ -748,7 +748,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.53.2", + "windows-targets 0.53.3", ] [[package]] @@ -1462,9 +1462,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.9.2" +version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed0aee96c12fa71097902e0bb061a5e1ebd766a6636bb605ba401c45c1650eac" +checksum = "41ae868b5a0f67631c14589f7e250c1ea2c574ee5ba21c6c8dd4b1485705a5a1" dependencies = [ "indexmap", "serde", @@ -1635,7 +1635,7 @@ dependencies = [ "seq-macro", "serde", "simd", - "toml 0.9.2", + "toml 0.9.4", "validator", "vchordg", "vchordrq", @@ -1815,6 +1815,12 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-link" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" + [[package]] name = "windows-sys" version = "0.59.0" @@ -1830,7 +1836,7 @@ version = "0.60.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" dependencies = [ - "windows-targets 0.53.2", + "windows-targets 0.53.3", ] [[package]] @@ -1851,10 +1857,11 @@ dependencies = [ [[package]] name = "windows-targets" -version = "0.53.2" +version = "0.53.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c66f69fcc9ce11da9966ddb31a40968cad001c5bedeb5c2b82ede4253ab48aef" +checksum = "d5fe6031c4041849d7c496a8ded650796e7b6ecc19df1a431c1a363342e5dc91" dependencies = [ + "windows-link", "windows_aarch64_gnullvm 0.53.0", "windows_aarch64_msvc 0.53.0", "windows_i686_gnu 0.53.0", diff --git a/Cargo.toml b/Cargo.toml index 7a61bd65..40c2358f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,7 +40,7 @@ pgrx-catalog = "0.3.0" rand.workspace = true seq-macro.workspace = true serde.workspace = true -toml = "0.9.2" +toml = "0.9.4" validator.workspace = true zerocopy.workspace = true diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index 602c9bdb..4f75b1c6 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -5,7 +5,7 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.41", features = ["derive"] } +clap = { version = "4.5.42", features = ["derive"] } object = { version = "0.37.1", default-features = false, features = [ "read", "wasm", diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index 133e6dcd..b9c426cf 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -195,19 +195,32 @@ fn parse( eprintln!("Reading {obj:?}"); let contents = std::fs::read(obj)?; let object = object::File::parse(contents.as_slice())?; - let exports = object - .symbols() - .flat_map(|x| x.name().ok()) - .flat_map(|x| { - if !tsi.is_macos { - Some(x) - } else { - x.strip_prefix("_") - } - }) - .filter(|x| x.starts_with("__pgrx_internals")) - .map(str::to_string) - .collect(); + let exports; + if tsi.is_macos { + exports = object + .exports()? + .into_iter() + .flat_map(|x| std::str::from_utf8(x.name())) + .flat_map(|x| x.strip_prefix("_")) + .filter(|x| x.starts_with("__pgrx_internals")) + .map(str::to_string) + .collect(); + } else if tsi.is_emscripten { + exports = object + .symbols() + .flat_map(|x| x.name().ok()) + .filter(|x| x.starts_with("__pgrx_internals")) + .map(str::to_string) + .collect(); + } else { + exports = object + .exports()? + .into_iter() + .flat_map(|x| std::str::from_utf8(x.name())) + .filter(|x| x.starts_with("__pgrx_internals")) + .map(str::to_string) + .collect(); + } Ok(exports) } diff --git a/scripts/README.md b/scripts/README.md index 07ca423b..37119811 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -1,79 +1,3 @@ -## Build Docker - -```shell -SEMVER='0.2.1' -VERSION='17' -ARCH='x86_64' -PLATFORM='amd64' - -git clone https://github.com/tensorchord/VectorChord.git -cd VectorChord -git checkout $SEMVER - -sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config - -sudo apt-get install -y postgresql-common -sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y -sudo apt-get install -y postgresql-server-dev-$VERSION -sudo apt-get install -y postgresql-$VERSION -sudo apt-get install -y postgresql-$VERSION-pgvector - -curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 -sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 - -curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y - -cargo install cargo-pgrx@$(sed -n 's/.*pgrx = { version = "\(=.*\)",.*/\1/p' Cargo.toml) --locked -cargo pgrx init --pg$VERSION=$(which pg_config) - -cargo build --lib --features pg$VERSION --release - -mkdir -p ./build/zip -cp -a ./sql/upgrade/. ./build/zip/ -cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql -cp ./vchord.control ./build/zip/vchord.control -cp ./target/release/libvchord.so ./build/zip/vchord.so -zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* - -mkdir -p ./build/deb -mkdir -p ./build/deb/DEBIAN -mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ -mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ -for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file -done -for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file -done -for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file -done -echo "Package: postgresql-${VERSION}-vchord -Version: ${SEMVER}-1 -Section: database -Priority: optional -Architecture: ${PLATFORM} -Maintainer: Tensorchord -Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM -Homepage: https://vectorchord.ai/ -License: AGPL-3 or Elastic-2" \ -> ./build/deb/DEBIAN/control -(cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums -dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb - -ls ./build - -docker build -t vchord:pg$VERSION-latest --build-arg PG_VERSION=$VERSION -f ./docker/Dockerfile . -``` - -## Run Instance - -```shell -VERSION='17' - -docker run --name vchord -e POSTGRES_PASSWORD=123 -p 5432:5432 -d vchord:pg$VERSION-latest -``` - ## Run External Index Precomputation Toolkit 1. Install requirements diff --git a/src/lib.rs b/src/lib.rs index f83eacf7..cbd6f4c1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -47,7 +47,7 @@ pgrx::extension_sql_file!("./sql/finalize.sql", finalize); #[pgrx::pg_guard] #[unsafe(export_name = "_PG_init")] unsafe extern "C-unwind" fn _pg_init() { - if unsafe { pgrx::pg_sys::IsUnderPostmaster } { + if !unsafe { pgrx::pg_sys::process_shared_preload_libraries_in_progress } { pgrx::error!("vchord must be loaded via shared_preload_libraries."); } index::init(); diff --git a/tools/dev.md b/tools/dev.md deleted file mode 100644 index 7ce13007..00000000 --- a/tools/dev.md +++ /dev/null @@ -1,160 +0,0 @@ -## Release a new version - -### Pre-requisite - -```shell -SEMVER='0.2.1' -VERSION='17' - -git clone https://github.com/tensorchord/VectorChord.git -cd VectorChord -git checkout $SEMVER - -sudo apt install -y build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc ccache pkg-config - -sudo apt-get install -y postgresql-common -sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y -sudo apt-get install -y postgresql-server-dev-$VERSION -sudo apt-get install -y postgresql-$VERSION -sudo apt-get install -y postgresql-$VERSION-pgvector - -curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 -sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 - -curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y - -cargo install cargo-pgrx@$(sed -n 's/.*pgrx = { version = "\(=.*\)",.*/\1/p' Cargo.toml) --locked -cargo pgrx init --pg$VERSION=$(which pg_config) -``` - -### Step 1 - generate install schema - -```shell -SEMVER='0.2.1' -VERSION='17' - -cargo build --lib --features pg$VERSION --release -cargo pgrx schema --features pg$VERSION | expand -t 4 > ./sql/install/vchord--$SEMVER.sql -``` - -### Step 2 - generate upgrade schema - -```shell -PREV='0.2.0' -SEMVER='0.2.1' - -diff -u ./sql/install/vchord--$PREV.sql ./sql/install/vchord--$SEMVER.sql | awk ' - /^\+/ && !/^+++/ { - print substr($0, 2) > "target/upgrade.sql" - next - } - /^-/ && !/^---/ || /^@/ { print } - { next } -' -cp target/upgrade.sql ./sql/upgrade/vchord--$PREV--$SEMVER.sql -``` - -### Step 3 - validate - -```shell -PREV='0.2.0' -SEMVER='0.2.1' - -sudo -u postgres createdb vchord - -sudo -u postgres ./tools/dump.sh $PREV $SEMVER > target/upgrade.sql -sudo -u postgres ./tools/dump.sh $SEMVER > target/install.sql -code --diff target/upgrade.sql target/install.sql - -sudo -u postgres dropdb vchord -``` - -### Step 4 - package and download - -Edit `./vchord.control` and bump the version. - -```shell -SEMVER='0.2.1' -VERSION='17' -ARCH='x86_64' -PLATFORM='amd64' - -cargo build --lib --features pg$VERSION --release - -mkdir -p ./build/zip -cp -a ./sql/upgrade/. ./build/zip/ -cp ./sql/install/vchord--$SEMVER.sql ./build/zip/vchord--$SEMVER.sql -cp ./vchord.control ./build/zip/vchord.control -cp ./target/release/libvchord.so ./build/zip/vchord.so -zip ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip -j ./build/zip/* - -mkdir -p ./build/deb -mkdir -p ./build/deb/DEBIAN -mkdir -p ./build/deb/usr/share/postgresql/$VERSION/extension/ -mkdir -p ./build/deb/usr/lib/postgresql/$VERSION/lib/ -for file in $(ls ./build/zip/*.sql | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file -done -for file in $(ls ./build/zip/*.control | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/share/postgresql/$VERSION/extension/$file -done -for file in $(ls ./build/zip/*.so | xargs -n 1 basename); do - cp ./build/zip/$file ./build/deb/usr/lib/postgresql/$VERSION/lib/$file -done -echo "Package: postgresql-${VERSION}-vchord -Version: ${SEMVER}-1 -Section: database -Priority: optional -Architecture: ${PLATFORM} -Maintainer: Tensorchord -Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM -Homepage: https://vectorchord.ai/ -License: AGPL-3 or Elastic-2" \ -> ./build/deb/DEBIAN/control -(cd ./build/deb && md5sum usr/share/postgresql/$VERSION/extension/* usr/lib/postgresql/$VERSION/lib/*) > ./build/deb/DEBIAN/md5sums -dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb - -ls ./build - -wget https://github.com/tensorchord/VectorChord/releases/download/${PREV}/postgresql-${VERSION}-vchord_${PREV}-1_${PLATFORM}.deb -O ./build/postgresql-${VERSION}-vchord_${PREV}-1_${PLATFORM}.deb -``` - -### Step 5 - further test - -```shell -PREV='0.2.0' -SEMVER='0.2.1' -VERSION='17' -ARCH='x86_64' -PLATFORM='amd64' - -cargo install sqllogictest-bin - -# upgrade test - -sudo apt install -y ./build/postgresql-${VERSION}-vchord_${PREV}-1_${PLATFORM}.deb -sudo -u postgres psql -d vchord -c "ALTER SYSTEM SET SHARED_PRELOAD_LIBRARIES='vchord.so';" -sudo systemctl restart postgresql@$VERSION-main.service -sudo -u postgres psql -d vchord -c 'CREATE EXTENSION vchord CASCADE;' -sudo -u postgres psql -d vchord -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord';" -sudo -u postgres sqllogictest -d vchord '../tests/**/*.slt' - -sudo apt install -y ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb -sudo systemctl restart postgresql@$VERSION-main.service -sudo -u postgres psql -d vchord -c 'ALTER EXTENSION vchord UPDATE;' -sudo -u postgres psql -d vchord -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord';" -sudo -u postgres sqllogictest -d vchord '../tests/**/*.slt' - -sudo apt remove -y postgresql-${VERSION}-vchord - -# install test - -sudo apt install -y ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb -sudo -u postgres psql -d vchord -c "ALTER SYSTEM SET SHARED_PRELOAD_LIBRARIES='vchord.so';" -sudo systemctl restart postgresql@$VERSION-main.service -sudo -u postgres psql -d vchord -c 'CREATE EXTENSION vchord CASCADE;' -sudo -u postgres psql -d vchord -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord';" -sudo -u postgres sqllogictest -d vchord '../tests/**/*.slt' - -sudo apt remove -y postgresql-${VERSION}-vchord -``` From 9fdcc5c54e0eff970915644095029da7c5c9b3ee Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 1 Aug 2025 15:40:35 +0800 Subject: [PATCH 186/324] fix: add pg_guard_ffi_boundary properly (#302) Signed-off-by: usamoi --- src/index/fetcher.rs | 15 +++++++-- src/index/hook.rs | 51 ++++++++----------------------- src/index/vchordg/am/am_build.rs | 31 +++++++++++-------- src/index/vchordg/am/mod.rs | 6 +++- src/index/vchordg/opclass.rs | 5 ++- src/index/vchordrq/am/am_build.rs | 31 +++++++++++-------- src/index/vchordrq/am/mod.rs | 6 +++- src/index/vchordrq/opclass.rs | 5 ++- src/lib.rs | 11 +++++++ 9 files changed, 90 insertions(+), 71 deletions(-) diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs index 823a9bf4..6726ce0c 100644 --- a/src/index/fetcher.rs +++ b/src/index/fetcher.rs @@ -96,12 +96,16 @@ impl Fetcher for HeapFetcher { fn fetch(&mut self, key: [u16; 3]) -> Option> { unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; let mut ctid = key_to_ctid(key); let table_am = (*self.heap_relation).rd_tableam; let fetch_row_version = (*table_am) .tuple_fetch_row_version .expect("unsupported heap access method"); - if !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + if pg_guard_ffi_boundary(|| { + !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) + }) { return None; } Some(HeapTuple { this: self }) @@ -133,6 +137,7 @@ impl Tuple for HeapTuple<'_> { #[allow(clippy::collapsible_if)] fn filter(&mut self) -> bool { unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; let this = &mut self.this; if !this.hack.is_null() { if let Some(qual) = NonNull::new((*this.hack).ss.ps.qual) { @@ -147,7 +152,13 @@ impl Tuple for HeapTuple<'_> { let result = PgMemoryContexts::For((*econtext).ecxt_per_tuple_memory) .switch_to(|_| { let mut is_null = true; - let datum = evalfunc(qual.as_ptr(), econtext, &mut is_null); + #[allow( + ffi_unwind_calls, + reason = "protected by pg_guard_ffi_boundary" + )] + let datum = pg_guard_ffi_boundary(|| { + evalfunc(qual.as_ptr(), econtext, &mut is_null) + }); bool::from_datum(datum, is_null) }); if result != Some(true) { diff --git a/src/index/hook.rs b/src/index/hook.rs index 636c4c5d..bfada43a 100644 --- a/src/index/hook.rs +++ b/src/index/hook.rs @@ -12,8 +12,6 @@ // // Copyright (c) 2025 TensorChord Inc. -use std::sync::atomic::AtomicPtr; - #[pgrx::pg_guard] unsafe extern "C-unwind" fn rewrite_plan_state( node: *mut pgrx::pg_sys::PlanState, @@ -118,55 +116,30 @@ unsafe extern "C-unwind" fn rewrite_plan_state( } } -static PREV_EXECUTOR_START: AtomicPtr<()> = AtomicPtr::new(core::ptr::null_mut()); - -#[cfg(any( - feature = "pg13", - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17" -))] -type ExecutorStartReturnType = (); - -#[cfg(feature = "pg18")] -type ExecutorStartReturnType = bool; +static mut PREV_EXECUTOR_START: pgrx::pg_sys::ExecutorStart_hook_type = None; #[pgrx::pg_guard] unsafe extern "C-unwind" fn executor_start( query_desc: *mut pgrx::pg_sys::QueryDesc, - eflags: ::std::os::raw::c_int, -) -> ExecutorStartReturnType { + eflags: core::ffi::c_int, +) { unsafe { - use core::mem::transmute; - use pgrx::pg_sys::ExecutorStart_hook_type; - use std::sync::atomic::Ordering; - let value = transmute::<*mut (), ExecutorStart_hook_type>( - PREV_EXECUTOR_START.load(Ordering::Relaxed), - ); - #[allow(clippy::let_unit_value)] - let result = if let Some(prev_executor_start) = value { - prev_executor_start(query_desc, eflags) + use core::ptr::null_mut; + use pgrx::pg_sys::submodules::ffi::pg_guard_ffi_boundary; + if let Some(prev_executor_start) = PREV_EXECUTOR_START { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| prev_executor_start(query_desc, eflags)) } else { pgrx::pg_sys::standard_ExecutorStart(query_desc, eflags) - }; - let planstate = (*query_desc).planstate; - let context = core::ptr::null_mut(); - rewrite_plan_state(planstate, context); - result + } + pg_guard_ffi_boundary(|| rewrite_plan_state((*query_desc).planstate, null_mut())); } } pub fn init() { + assert!(crate::is_main()); unsafe { - use core::mem::transmute; - use std::sync::atomic::Ordering; - PREV_EXECUTOR_START.store( - transmute:: _>, *mut ()>( - pgrx::pg_sys::ExecutorStart_hook, - ), - Ordering::Relaxed, - ); + PREV_EXECUTOR_START = pgrx::pg_sys::ExecutorStart_hook; pgrx::pg_sys::ExecutorStart_hook = Some(executor_start); } } diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index 5b0af1e8..d777635d 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -96,6 +96,7 @@ impl Heap { progress: bool, callback: F, ) { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; pub struct State<'a, F> { pub this: &'a Heap, pub callback: F, @@ -124,20 +125,24 @@ impl Heap { this: self, callback, }; + let index_build_range_scan = table_am.index_build_range_scan.expect("bad table"); unsafe { - table_am.index_build_range_scan.unwrap()( - self.heap_relation, - self.index_relation, - self.index_info, - true, - false, - progress, - 0, - pgrx::pg_sys::InvalidBlockNumber, - Some(call::), - (&mut state) as *mut State as *mut _, - self.scan, - ); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + index_build_range_scan( + self.heap_relation, + self.index_relation, + self.index_info, + true, + false, + progress, + 0, + pgrx::pg_sys::InvalidBlockNumber, + Some(call::), + (&mut state) as *mut State as *mut _, + self.scan, + ) + }); } } } diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index 20a09678..ea71d05c 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -240,6 +240,7 @@ pub unsafe extern "C-unwind" fn ambulkdelete( callback: pgrx::pg_sys::IndexBulkDeleteCallback, callback_state: *mut std::os::raw::c_void, ) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; let mut stats = stats; if stats.is_null() { stats = unsafe { @@ -264,7 +265,10 @@ pub unsafe extern "C-unwind" fn ambulkdelete( let callback = |pointer: NonZero| { let (key, _) = pointer_to_kv(pointer); let mut ctid = key_to_ctid(key); - unsafe { callback(&mut ctid, callback_state) } + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + unsafe { + pg_guard_ffi_boundary(|| callback(&mut ctid, callback_state)) + } }; crate::index::vchordg::algo::bulkdelete(opfamily, &index, check, callback); stats diff --git a/src/index/vchordg/opclass.rs b/src/index/vchordg/opclass.rs index a216265c..298687b2 100644 --- a/src/index/vchordg/opclass.rs +++ b/src/index/vchordg/opclass.rs @@ -150,7 +150,10 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { fcinfo.isnull = true; fcinfo.nargs = 0; - let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; + let result_datum = unsafe { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) + }; let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 38afa216..be3b0b79 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -172,6 +172,7 @@ impl Heap { progress: bool, callback: F, ) { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; pub struct State<'a, F> { pub this: &'a Heap, pub callback: F, @@ -200,20 +201,24 @@ impl Heap { this: self, callback, }; + let index_build_range_scan = table_am.index_build_range_scan.expect("bad table"); unsafe { - table_am.index_build_range_scan.unwrap()( - self.heap_relation, - self.index_relation, - self.index_info, - true, - false, - progress, - 0, - pgrx::pg_sys::InvalidBlockNumber, - Some(call::), - (&mut state) as *mut State as *mut _, - self.scan, - ); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + index_build_range_scan( + self.heap_relation, + self.index_relation, + self.index_info, + true, + false, + progress, + 0, + pgrx::pg_sys::InvalidBlockNumber, + Some(call::), + (&mut state) as *mut State as *mut _, + self.scan, + ) + }); } } } diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 61d28a9c..58e57423 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -320,6 +320,7 @@ pub unsafe extern "C-unwind" fn ambulkdelete( callback: pgrx::pg_sys::IndexBulkDeleteCallback, callback_state: *mut std::os::raw::c_void, ) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; let mut stats = stats; if stats.is_null() { stats = unsafe { @@ -344,7 +345,10 @@ pub unsafe extern "C-unwind" fn ambulkdelete( let callback = |pointer: NonZero| { let (key, _) = pointer_to_kv(pointer); let mut ctid = key_to_ctid(key); - unsafe { callback(&mut ctid, callback_state) } + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + unsafe { + pg_guard_ffi_boundary(|| callback(&mut ctid, callback_state)) + } }; crate::index::vchordrq::algo::bulkdelete(opfamily, &index, check, callback); stats diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs index 1fa0f2e2..f70eee1c 100644 --- a/src/index/vchordrq/opclass.rs +++ b/src/index/vchordrq/opclass.rs @@ -208,7 +208,10 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { fcinfo.isnull = true; fcinfo.nargs = 0; - let result_datum = unsafe { pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) }; + let result_datum = unsafe { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) + }; let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; diff --git a/src/lib.rs b/src/lib.rs index cbd6f4c1..8aa7ab8f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,6 +14,7 @@ #![allow(unsafe_code)] #![allow(unused_crate_dependencies)] +#![warn(ffi_unwind_calls)] mod datatype; mod index; @@ -50,6 +51,7 @@ unsafe extern "C-unwind" fn _pg_init() { if !unsafe { pgrx::pg_sys::process_shared_preload_libraries_in_progress } { pgrx::error!("vchord must be loaded via shared_preload_libraries."); } + IS_MAIN.set(true); index::init(); unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] @@ -59,6 +61,15 @@ unsafe extern "C-unwind" fn _pg_init() { } } +std::thread_local! { + static IS_MAIN: core::cell::Cell = const { core::cell::Cell::new(false) }; +} + +#[must_use] +fn is_main() -> bool { + IS_MAIN.get() +} + #[cfg(not(panic = "unwind"))] compile_error!("This crate must be compiled with `-Cpanic=unwind`."); From 49b6642fea5b934f679a0eb2d72ac35617640640 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 4 Aug 2025 16:57:12 +0800 Subject: [PATCH 187/324] fix: fast path for prefetch-disabled cases (#304) Signed-off-by: usamoi --- .github/workflows/check.yml | 215 +++++++++++++++++++++--------- .github/workflows/release.yml | 3 +- Cargo.lock | 96 +++---------- Cargo.toml | 6 +- crates/algo/src/prefetcher.rs | 4 + crates/make/Cargo.toml | 2 +- crates/simd/Cargo.toml | 2 +- crates/simd/build.rs | 46 ++++--- crates/simd_macros/src/lib.rs | 2 +- crates/vchordrq/src/build.rs | 15 ++- crates/vchordrq/src/maintain.rs | 1 + crates/vchordrq/src/prewarm.rs | 28 ++-- crates/vchordrq/src/search.rs | 44 ++++-- crates/vchordrq/src/tuples.rs | 13 +- src/datatype/operators_halfvec.rs | 2 +- src/datatype/operators_vector.rs | 2 +- src/index/gucs.rs | 14 +- src/index/vchordg/algo.rs | 12 ++ src/index/vchordrq/algo.rs | 20 +++ src/index/vchordrq/opclass.rs | 14 +- src/lib.rs | 2 +- 21 files changed, 333 insertions(+), 210 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 61413a29..120d9a7d 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -9,11 +9,6 @@ concurrency: group: ${{ github.ref }}-${{ github.workflow }} cancel-in-progress: true -env: - CARGO_TERM_COLOR: always - RUST_BACKTRACE: 1 - RUSTFLAGS: "-Dwarnings" - jobs: style: runs-on: "ubuntu-latest" @@ -21,8 +16,8 @@ jobs: steps: - name: Set up Environment run: | - rustup set profile rustup default 1.89-beta + rustup component add rustfmt clippy curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo @@ -114,12 +109,15 @@ jobs: env: RUSTC_WRAPPER: sccache SCCACHE_GHA_ENABLED: true + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 + RUSTFLAGS: "-Dwarnings" steps: - name: Set up Environment run: | - rustup set profile rustup default 1.89-beta + rustup component add rustfmt clippy sudo apt-get update @@ -179,12 +177,14 @@ jobs: env: RUSTC_WRAPPER: sccache SCCACHE_GHA_ENABLED: true + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 steps: - name: Set up Environment run: | - rustup set profile rustup default 1.89-beta + rustup component add rustfmt clippy sudo apt-get update @@ -216,8 +216,17 @@ jobs: - name: Checkout uses: actions/checkout@v4 + - name: Patch + run: | + mkdir ./.cargo && touch ./.cargo/config.toml + echo "unstable.host-config = true" >> ./.cargo/config.toml + echo "unstable.target-applies-to-host = true" >> ./.cargo/config.toml + echo "host.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml + echo "build.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + - name: Clippy - run: cargo clippy -p vchord --features pg${{ matrix.version }} -- --no-deps + run: cargo clippy --target ${{ matrix.arch }}-unknown-linux-gnu -p vchord --features pg${{ matrix.version }} --no-deps - name: Install run: | @@ -238,39 +247,11 @@ jobs: sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' fi - - name: Package - env: - SEMVER: "0.0.0" - VERSION: ${{ matrix.version }} - ARCH: ${{ matrix.arch }} - PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} - run: | - (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) - - mkdir -p ./build/deb - mkdir -p ./build/deb/DEBIAN - mkdir -p ./build/deb$(pg_config --pkglibdir) && cp -r ./build/raw/pkglibdir/. ./build/deb$(pg_config --pkglibdir) - mkdir -p ./build/deb$(pg_config --sharedir) && cp -r ./build/raw/sharedir/. ./build/deb$(pg_config --sharedir) - echo "Package: postgresql-${VERSION}-vchord - Version: ${SEMVER}-1 - Section: database - Priority: optional - Architecture: ${PLATFORM} - Maintainer: Tensorchord - Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM - Homepage: https://vectorchord.ai/ - License: AGPL-3.0-only or Elastic-2.0" \ - > ./build/deb/DEBIAN/control - (cd ./build/deb && find usr -type f -print0 | xargs -0 md5sum) > ./build/deb/DEBIAN/md5sums - dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb - - name: Upload Artifacts uses: actions/upload-artifact@v4 with: - name: artifacts-psql-${{ matrix.version }}-${{ matrix.arch }} - path: | - ./build/postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu.zip - ./build/postgresql-${{ matrix.version }}-vchord_0.0.0-1_${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }}.deb + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu + path: ./build/raw compression-level: 9 retention-days: 14 @@ -290,12 +271,14 @@ jobs: env: RUSTC_WRAPPER: sccache SCCACHE_GHA_ENABLED: true + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 steps: - name: Set up Environment run: | - rustup set profile rustup default 1.89-beta + rustup component add rustfmt clippy HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install llvm@18 echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV @@ -322,9 +305,18 @@ jobs: - name: Checkout uses: actions/checkout@v4 + - name: Patch + run: | + mkdir ./.cargo && touch ./.cargo/config.toml + echo "unstable.host-config = true" >> ./.cargo/config.toml + echo "unstable.target-applies-to-host = true" >> ./.cargo/config.toml + echo "host.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml + echo "build.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + - name: Clippy run: | - cargo clippy -p vchord --features pg${{ matrix.version }} -- --no-deps + cargo clippy -p vchord --target ${{ matrix.arch }}-apple-darwin --features pg${{ matrix.version }} --no-deps - name: Install run: | @@ -347,20 +339,11 @@ jobs: sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' fi - - name: Package - env: - SEMVER: "0.0.0" - VERSION: ${{ matrix.version }} - ARCH: ${{ matrix.arch }} - run: | - (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-apple-darwin.zip .) - - name: Upload Artifacts uses: actions/upload-artifact@v4 with: - name: artifacts-psql_macos-${{ matrix.version }}-${{ matrix.arch }} - path: | - ./build/postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin.zip + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin + path: ./build/raw compression-level: 9 retention-days: 14 @@ -380,14 +363,16 @@ jobs: env: RUSTC_WRAPPER: sccache SCCACHE_GHA_ENABLED: true + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 steps: - name: Set up Environment run: | 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } - rustup set profile rustup default 1.89-beta + rustup component add rustfmt clippy & 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise\Common7\Tools\Launch-VsDevShell.ps1' -HostArch amd64 -Arch amd64 @@ -433,9 +418,19 @@ jobs: - name: Checkout uses: actions/checkout@v4 + - name: Patch + shell: bash + run: | + mkdir ./.cargo && touch ./.cargo/config.toml + echo "unstable.host-config = true" >> ./.cargo/config.toml + echo "unstable.target-applies-to-host = true" >> ./.cargo/config.toml + echo "host.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml + echo "build.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + - name: Clippy run: | - cargo clippy -p vchord --features pg${{ matrix.version }} -- --no-deps + cargo clippy --target ${{ matrix.arch }}-pc-windows-msvc -p vchord --features pg${{ matrix.version }} --no-deps - name: Install run: | @@ -458,19 +453,113 @@ jobs: sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/pg17/*.slt' } - - name: Package - env: - SEMVER: "0.0.0" - VERSION: ${{ matrix.version }} - ARCH: ${{ matrix.arch }} + - name: Upload Artifacts + uses: actions/upload-artifact@v4 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc + path: ./build/raw + compression-level: 9 + retention-days: 14 + + psql_alpine: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_alpine')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +psql_alpine')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + version: ["15", "16", "17"] + arch: ["x86_64", "aarch64"] + runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-24.04' || 'ubuntu-24.04-arm' }} + container: + image: alpine:3.22 + volumes: + - /opt:/opt:rw,rshared + - /opt:/__e/node20:ro,rshared + + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: true + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 + USER: root + + steps: + - name: Set up Environment run: | - Compress-Archive -Path ./build/raw/* -DestinationPath ./build/postgresql-$($env:VERSION)-vchord_$($env:SEMVER)_$($env:ARCH)-pc-windows-msvc.zip + sed -i "/^ID=/s/alpine/NotpineForGHA/" /etc/os-release + apk add nodejs --update-cache + mkdir /opt/bin + ln -s /usr/bin/node /opt/bin/node + + apk add --no-cache sudo curl make zip clang18-dev postgresql${{ matrix.version }} postgresql${{ matrix.version }}-dev + + curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain 1.89-beta + + echo PGRX_PG_CONFIG_PATH=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV + + install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql + install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql/data + install --verbose --directory --owner postgres --group postgres --mode 3777 /run/postgresql + sudo -iu postgres initdb -D /var/lib/postgresql/data + sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data + sudo -iu postgres createuser -s -r $USER + sudo -iu postgres createdb -O $USER $USER + sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + sudo -iu postgres pg_ctl stop -D /var/lib/postgresql/data + + mkdir ~/pgvector-install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.0.tar.gz | tar -xz -C ~/pgvector-install + make -C ~/pgvector-install/pgvector-0.8.0 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config + make -C ~/pgvector-install/pgvector-0.8.0 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config install + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Checkout + uses: actions/checkout@v4 + + - name: Patch + run: | + mkdir ./.cargo && touch ./.cargo/config.toml + echo "unstable.host-config = true" >> ./.cargo/config.toml + echo "unstable.target-applies-to-host = true" >> ./.cargo/config.toml + echo "host.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml + echo "build.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + + - name: Clippy + run: | + . "$HOME/.cargo/env" + cargo clippy --target ${{ matrix.arch }}-unknown-linux-musl -p vchord --features pg${{ matrix.version }} --no-deps + + - name: Install + run: | + . "$HOME/.cargo/env" + cargo run -p make -- build -o ./build/raw --profile dev + make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install + + - name: Service + run: | + sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data + psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $USER --user $USER './tests/general/*.slt' + sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' + sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' + if [ "${{ matrix.version }}" = "17" ]; then + sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' + fi - name: Upload Artifacts uses: actions/upload-artifact@v4 with: - name: artifacts-psql_windows-${{ matrix.version }}-${{ matrix.arch }} - path: | - ./build/postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc.zip + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-musl + path: ./build/raw compression-level: 9 retention-days: 14 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 175c0027..6846127b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -46,13 +46,12 @@ jobs: env: CARGO_TERM_COLOR: always RUST_BACKTRACE: 1 - RUSTFLAGS: "-Dwarnings" steps: - name: Set up Environment run: | - rustup set profile rustup default 1.89-beta + rustup component add rustfmt clippy sudo apt-get update diff --git a/Cargo.lock b/Cargo.lock index a1fba94d..491e18bf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -100,16 +100,6 @@ version = "1.0.98" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" -[[package]] -name = "atomic-traits" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "707f750b93bd1b739cf9ddf85f8fe7c97a4a62c60ccf8b6f232514bd9103bedc" -dependencies = [ - "cfg-if", - "rustc_version", -] - [[package]] name = "bindgen" version = "0.71.1" @@ -153,12 +143,6 @@ version = "3.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43" -[[package]] -name = "byteorder" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" - [[package]] name = "cargo_toml" version = "0.22.3" @@ -171,9 +155,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.30" +version = "1.2.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "deec109607ca693028562ed836a5f1c4b8bd77755c4e132fc5ce11b0b6211ae7" +checksum = "c3a42d84bb6b69d3a8b3eaacf0d88f179e1929695e1ad012b6cf64d9caaa5fd2" dependencies = [ "shlex", ] @@ -497,15 +481,6 @@ dependencies = [ "zerocopy", ] -[[package]] -name = "hash32" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" -dependencies = [ - "byteorder", -] - [[package]] name = "hashbrown" version = "0.15.4" @@ -517,16 +492,6 @@ dependencies = [ "foldhash", ] -[[package]] -name = "heapless" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" -dependencies = [ - "hash32", - "stable_deref_trait", -] - [[package]] name = "heck" version = "0.5.0" @@ -827,9 +792,9 @@ dependencies = [ [[package]] name = "object" -version = "0.37.1" +version = "0.37.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03fd943161069e1768b4b3d050890ba48730e590f57e56d4aa04e7e090e61b4a" +checksum = "b3e3d0a7419f081f4a808147e845310313a39f322d7ae1f996b7f001d6cbed04" dependencies = [ "memchr", "wasmparser", @@ -893,15 +858,13 @@ dependencies = [ [[package]] name = "pgrx" -version = "0.15.0" +version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bab5bc1d60d3bc3c966d307a3c7313b1ebfb49a0ec183be3f1a057df0bcc9988" +checksum = "2c898b37587b29f8fa7ca20bfaaf57103e10fa81e6caed2cc7b742f96088d851" dependencies = [ - "atomic-traits", "bitflags", "bitvec", "enum-map", - "heapless", "libc", "once_cell", "pgrx-macros", @@ -917,9 +880,9 @@ dependencies = [ [[package]] name = "pgrx-bindgen" -version = "0.15.0" +version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9804b74c211a9edd550cd974718f8cc407dec50d8e9cafb906e0b042ba434af0" +checksum = "35439a65a6c0cc17e9758005da0557cff583a3e4da0f92a4a744494ac21c190e" dependencies = [ "bindgen", "cc", @@ -936,9 +899,9 @@ dependencies = [ [[package]] name = "pgrx-catalog" -version = "0.3.0" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fad3b930e3f5abd77c689fd9f64dba3bbb97993c493700daa4c1242c27fc141c" +checksum = "3342c1e98f08dc57ceb8a44835f362ae3de4422bada1fbbccfcd74a4b55dc978" dependencies = [ "paste", "pgrx", @@ -946,9 +909,9 @@ dependencies = [ [[package]] name = "pgrx-macros" -version = "0.15.0" +version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f230769493bf567f137de23264d604d267dd72b8a77c596528e43cf423c6208e" +checksum = "3586274c7d7237e68bebb871cf9244444f1ad78bb395679ad4ef1015db302454" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", @@ -958,9 +921,9 @@ dependencies = [ [[package]] name = "pgrx-pg-config" -version = "0.15.0" +version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49b64c071c2a46a19ab4521120a25b02b598f4abf6e9b4b1769a7922edeee3de" +checksum = "0cc24f6663377dca8797ff3561d9ecdaef47e518ca2e2438fd85ef3c20f421ea" dependencies = [ "cargo_toml", "codepage", @@ -979,9 +942,9 @@ dependencies = [ [[package]] name = "pgrx-pg-sys" -version = "0.15.0" +version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fcbfa98ec7a90252d13a78ac666541173dbb01a2fc1ba20131db6490c0711125" +checksum = "b44a35477c0752c00cab38c2d8b7ea2d2d3554e3385870d7ad43e3ae5ade66a9" dependencies = [ "cee-scape", "libc", @@ -994,9 +957,9 @@ dependencies = [ [[package]] name = "pgrx-sql-entity-graph" -version = "0.15.0" +version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e79bbf5a33cff6cfdc6dda3a976cd931c995eaa2c073a7c59b8f8fe8f6faa073" +checksum = "59036a5849c5b0bca0f33df1fdd72206ceb74712959756b3fecba56a8e19d52a" dependencies = [ "convert_case", "eyre", @@ -1192,15 +1155,6 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" -[[package]] -name = "rustc_version" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" -dependencies = [ - "semver", -] - [[package]] name = "rustix" version = "1.0.8" @@ -1241,12 +1195,6 @@ version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" -[[package]] -name = "semver" -version = "1.0.26" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56e6fa9c48d24d85fb3de5ad847117517440f6beceb7798af16b4a87d616b8d0" - [[package]] name = "seq-macro" version = "0.3.6" @@ -1285,9 +1233,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.141" +version = "1.0.142" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30b9eff21ebe718216c6ec64e1d9ac57087aad11efc64e32002bce4a0d4c03d3" +checksum = "030fedb782600dcbd6f02d479bf0d817ac3bb40d644745b769d6a96bc3afc5a7" dependencies = [ "itoa", "memchr", @@ -1766,9 +1714,9 @@ dependencies = [ [[package]] name = "wasmparser" -version = "0.234.0" +version = "0.236.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be22e5a8f600afce671dd53c8d2dd26b4b7aa810fd18ae27dfc49737f3e02fc5" +checksum = "16d1eee846a705f6f3cb9d7b9f79b54583810f1fb57a1e3aea76d1742db2e3d2" dependencies = [ "bitflags", ] diff --git a/Cargo.toml b/Cargo.toml index 40c2358f..a2c6ae94 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,8 +35,8 @@ vector = { path = "./crates/vector" } bumpalo.workspace = true half.workspace = true paste.workspace = true -pgrx = { version = "=0.15.0", default-features = false, features = ["cshim"] } -pgrx-catalog = "0.3.0" +pgrx = { version = "=0.16.0", default-features = false, features = ["cshim"] } +pgrx-catalog = "0.3.1" rand.workspace = true seq-macro.workspace = true serde.workspace = true @@ -44,7 +44,7 @@ toml = "0.9.4" validator.workspace = true zerocopy.workspace = true -[target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), target_os = "linux"))'.dependencies] +[target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos")))'.dependencies] mimalloc = { version = "0.1.47", features = ["local_dynamic_tls"] } [lints] diff --git a/crates/algo/src/prefetcher.rs b/crates/algo/src/prefetcher.rs index c36998bc..2150c887 100644 --- a/crates/algo/src/prefetcher.rs +++ b/crates/algo/src/prefetcher.rs @@ -279,6 +279,8 @@ pub trait PrefetcherSequenceFamily<'r, R> { fn prefetch(&mut self, seq: S) -> Self::P where S::Item: Fetch; + + fn is_not_plain(&self) -> bool; } pub trait PrefetcherHeapFamily<'r, R> { @@ -289,6 +291,8 @@ pub trait PrefetcherHeapFamily<'r, R> { fn prefetch(&mut self, seq: Vec) -> Self::P where T: Ord + Fetch + 'r; + + fn is_not_plain(&self) -> bool; } // Emulate unstable library feature `vec_deque_pop_if`. diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index 4f75b1c6..0d7d6414 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -6,7 +6,7 @@ publish = false [dependencies] clap = { version = "4.5.42", features = ["derive"] } -object = { version = "0.37.1", default-features = false, features = [ +object = { version = "0.37.2", default-features = false, features = [ "read", "wasm", ] } diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index ae3b286e..69ed8684 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -18,7 +18,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.30" +cc = "1.2.31" which = "8.0.0" [lints] diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 7b5350a5..e0d9ffd5 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -15,25 +15,41 @@ use std::env::var; use std::error::Error; use std::ffi::OsString; +use std::path::Path; -fn compiler(host: &str, target: &str) -> Option { - if let Some(cc) = std::env::var_os("CC") { - return Some(cc); +fn compiler_version(cc: impl AsRef) -> Option { + let cc = cc.as_ref(); + if let Ok(r) = std::process::Command::new(cc).arg("-dumpversion").output() + && r.status.success() + && let Some(major) = r.stdout.split(|c| !c.is_ascii_digit()).next() + && let Ok(major) = std::str::from_utf8(major) + && let Ok(major) = major.parse::() + { + return Some(major); } - if let Some(cc) = std::env::var_os(format!("CC_{target}")) { - return Some(cc); - } - if let Some(cc) = std::env::var_os(format!("CC_{}", target.replace("-", "_"))) { - return Some(cc); + None +} + +fn compiler(host: &str, target: &str, clang_version: u16, gcc_version: u16) -> Option { + let keys = [ + &format!("CC_{target}"), + &format!("CC_{}", target.replace("-", "_")), + "TARGET_CC", + "CC", + ]; + if keys.iter().any(|key| std::env::var_os(key).is_some()) { + return None; } if host == target { - if let Ok(cc) = which::which("clang") { + if let Ok(cc) = which::which("clang") + && compiler_version(&cc) >= Some(clang_version) + { return Some(cc.into()); } - for i in (16..=99).rev() { - if let Ok(cc) = which::which(format!("clang-{i}")) { - return Some(cc.into()); - } + if let Ok(cc) = which::which("gcc") + && compiler_version(&cc) >= Some(gcc_version) + { + return Some(cc.into()); } } None @@ -47,7 +63,7 @@ fn main() -> Result<(), Box> { match target_arch.as_str() { "aarch64" => { let mut build = cc::Build::new(); - if let Some(compiler) = compiler(&host, &target) { + if let Some(compiler) = compiler(&host, &target, 16, 14) { build.compiler(compiler); } build.file("./cshim/aarch64.c"); @@ -56,7 +72,7 @@ fn main() -> Result<(), Box> { } "x86_64" => { let mut build = cc::Build::new(); - if let Some(compiler) = compiler(&host, &target) { + if let Some(compiler) = compiler(&host, &target, 16, 12) { build.compiler(compiler); } build.file("./cshim/x86_64.c"); diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs index 8b903578..a085260c 100644 --- a/crates/simd_macros/src/lib.rs +++ b/crates/simd_macros/src/lib.rs @@ -149,7 +149,7 @@ pub fn multiversion( #cold } #[cfg(feature = "init")] - #[cfg(target_os = "linux")] + #[cfg(all(target_os = "linux", target_env = "gnu"))] { #[used] #[unsafe(link_section = ".init_array")] diff --git a/crates/vchordrq/src/build.rs b/crates/vchordrq/src/build.rs index c53499e5..abde0caa 100644 --- a/crates/vchordrq/src/build.rs +++ b/crates/vchordrq/src/build.rs @@ -14,7 +14,7 @@ use crate::operator::{Operator, Vector}; use crate::tape::TapeWriter; -use crate::tape_writer::H1TapeWriter; +use crate::tape_writer::{DirectoryTapeWriter, H1TapeWriter}; use crate::tuples::*; use crate::types::*; use crate::{Branch, Opaque}; @@ -73,12 +73,19 @@ pub fn build( let mut level = Vec::new(); for j in 0..structures[i].len() { if i == 0 { - let directory_tape = TapeWriter::<_, DirectoryTuple>::create(index, false); + let frozen_tape = TapeWriter::<_, FrozenTuple>::create(index, false); let appendable_tape = TapeWriter::<_, AppendableTuple>::create(index, false); + let frozen_first = { frozen_tape }.first(); + + let mut directory_tape = DirectoryTapeWriter::create(index, false); + directory_tape.push(&[frozen_first]); + let directory_tape = directory_tape.into_inner(); + let mut jump = TapeWriter::<_, JumpTuple>::create(index, false); jump.push(JumpTuple { - directory_first: directory_tape.first(), - appendable_first: appendable_tape.first(), + directory_first: { directory_tape }.first(), + frozen_first, + appendable_first: { appendable_tape }.first(), centroid_prefetch: pointer_of_centroids[i][j].0.clone(), centroid_head: pointer_of_centroids[i][j].1, tuples: 0, diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs index 74d621b1..de0b64b6 100644 --- a/crates/vchordrq/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -224,6 +224,7 @@ where let directory_tape = directory_tape.into_inner(); *jump_tuple.directory_first() = { directory_tape }.first(); + *jump_tuple.frozen_first() = frozen_first; *jump_tuple.appendable_first() = { appendable_tape }.first(); *jump_tuple.tuples() = tuples; diff --git a/crates/vchordrq/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs index 24a8e8d9..b46bbb92 100644 --- a/crates/vchordrq/src/prewarm.rs +++ b/crates/vchordrq/src/prewarm.rs @@ -93,15 +93,25 @@ where let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let directory = - tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); - tape::read_frozen_tape::( - by_directory(&mut prefetch_h0_tuples, directory).inspect(|_| counter += 1), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), - id_2(|_, _, _, _| { - results.push(()); - }), - ); + if prefetch_h0_tuples.is_not_plain() { + let directory = + tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory).inspect(|_| counter += 1), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + id_2(|_, _, _, _| { + results.push(()); + }), + ); + } else { + tape::read_frozen_tape::( + by_next(index, jump_tuple.frozen_first()).inspect(|_| counter += 1), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + id_2(|_, _, _, _| { + results.push(()); + }), + ); + } tape::read_appendable_tape::( by_next(index, jump_tuple.appendable_first()).inspect(|_| counter += 1), |_, _, _| (), diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index 16624aff..fa8e02f2 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -132,8 +132,6 @@ where let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let directory = - tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); let mut callback = id_2(|(rough, err), head, payload, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( @@ -141,11 +139,21 @@ where AlwaysEqual(bump.alloc((payload, head, bump.alloc_slice(prefetch)))), )); }); - tape::read_frozen_tape::( - by_directory(&mut prefetch_h0_tuples, directory), - || O::block_access(&lut.0, block_process), - &mut callback, - ); + if prefetch_h0_tuples.is_not_plain() { + let directory = + tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory), + || O::block_access(&lut.0, block_process), + &mut callback, + ); + } else { + tape::read_frozen_tape::( + by_next(index, jump_tuple.frozen_first()), + || O::block_access(&lut.0, block_process), + &mut callback, + ); + } tape::read_appendable_tape::( by_next(index, jump_tuple.appendable_first()), O::binary_access(&lut.1, binary_process), @@ -266,8 +274,6 @@ where let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); - let directory = - tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); let mut callback = id_2(|(rough, err), head, payload, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); let rough = Distance::from_f32(rough); @@ -276,11 +282,21 @@ where AlwaysEqual(bump.alloc((payload, head, bump.alloc_slice(prefetch)))), )); }); - tape::read_frozen_tape::( - by_directory(&mut prefetch_h0_tuples, directory), - || O::block_access(&lut.0, block_process), - &mut callback, - ); + if prefetch_h0_tuples.is_not_plain() { + let directory = + tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory), + || O::block_access(&lut.0, block_process), + &mut callback, + ); + } else { + tape::read_frozen_tape::( + by_next(index, jump_tuple.frozen_first()), + || O::block_access(&lut.0, block_process), + &mut callback, + ); + } tape::read_appendable_tape::( by_next(index, jump_tuple.appendable_first()), O::binary_access(&lut.1, binary_process), diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index a09033bc..d8d5ac83 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -20,7 +20,7 @@ use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 10; +const VERSION: u64 = 11; pub trait Tuple: 'static { fn serialize(&self) -> Vec; @@ -795,8 +795,9 @@ struct JumpTupleHeader { centroid_prefetch_s: u16, centroid_prefetch_e: u16, centroid_head: u16, - _padding_0: [Padding; 2], + _padding_0: [Padding; 6], directory_first: u32, + frozen_first: u32, appendable_first: u32, tuples: u64, } @@ -806,6 +807,7 @@ pub struct JumpTuple { pub centroid_prefetch: Vec, pub centroid_head: u16, pub directory_first: u32, + pub frozen_first: u32, pub appendable_first: u32, pub tuples: u64, } @@ -828,6 +830,7 @@ impl Tuple for JumpTuple { centroid_prefetch_e, centroid_head: self.centroid_head, directory_first: self.directory_first, + frozen_first: self.frozen_first, appendable_first: self.appendable_first, tuples: self.tuples, _padding_0: Default::default(), @@ -877,6 +880,9 @@ impl<'a> JumpTupleReader<'a> { pub fn directory_first(self) -> u32 { self.header.directory_first } + pub fn frozen_first(self) -> u32 { + self.header.frozen_first + } pub fn appendable_first(self) -> u32 { self.header.appendable_first } @@ -894,6 +900,9 @@ impl JumpTupleWriter<'_> { pub fn directory_first(&mut self) -> &mut u32 { &mut self.header.directory_first } + pub fn frozen_first(&mut self) -> &mut u32 { + &mut self.header.frozen_first + } pub fn appendable_first(&mut self) -> &mut u32 { &mut self.header.appendable_first } diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index 25727123..41ff8568 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; -use pgrx::Array; +use pgrx::datum::Array; use std::num::NonZero; use vector::VectorBorrowed; use vector::vect::VectBorrowed; diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index eb96ac19..c648bfa7 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::datatype::memory_vector::{VectorInput, VectorOutput}; -use pgrx::Array; +use pgrx::datum::Array; use std::num::NonZero; use vector::VectorBorrowed; use vector::vect::VectBorrowed; diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 3d0515b0..855d6b82 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -12,26 +12,16 @@ // // Copyright (c) 2025 TensorChord Inc. -use pgrx::PostgresGucEnum; -use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; +use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting, PostgresGucEnum}; use std::ffi::CString; -#[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] -#[derive(Debug, Clone, Copy, PostgresGucEnum)] -pub enum PostgresIo { - #[name = c"read_buffer"] - ReadBuffer, - #[name = c"prefetch_buffer"] - PrefetchBuffer, -} - -#[cfg(any(feature = "pg17", feature = "pg18"))] #[derive(Debug, Clone, Copy, PostgresGucEnum)] pub enum PostgresIo { #[name = c"read_buffer"] ReadBuffer, #[name = c"prefetch_buffer"] PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] #[name = c"read_stream"] ReadStream, } diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs index 9c6e409d..31f0ade2 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/algo.rs @@ -219,6 +219,10 @@ impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakePlainPrefetche { PlainPrefetcher::new(self.index, seq) } + + fn is_not_plain(&self) -> bool { + false + } } #[derive(Debug)] @@ -246,6 +250,10 @@ impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> { SimplePrefetcher::new(self.index, seq) } + + fn is_not_plain(&self) -> bool { + true + } } #[derive(Debug)] @@ -277,4 +285,8 @@ impl<'r, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'r, R> { StreamPrefetcher::new(self.index, seq, self.hints.clone()) } + + fn is_not_plain(&self) -> bool { + true + } } diff --git a/src/index/vchordrq/algo.rs b/src/index/vchordrq/algo.rs index 82d87f59..a8e42c27 100644 --- a/src/index/vchordrq/algo.rs +++ b/src/index/vchordrq/algo.rs @@ -310,6 +310,10 @@ impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcherF { PlainPrefetcher::new(self.index, FastHeap::from(seq)) } + + fn is_not_plain(&self) -> bool { + false + } } #[derive(Debug)] @@ -335,6 +339,10 @@ impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcher< { PlainPrefetcher::new(self.index, BinaryHeap::from(seq)) } + + fn is_not_plain(&self) -> bool { + false + } } #[derive(Debug)] @@ -360,6 +368,10 @@ impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakeH0PlainPrefetc { PlainPrefetcher::new(self.index, seq) } + + fn is_not_plain(&self) -> bool { + false + } } #[derive(Debug)] @@ -387,6 +399,10 @@ impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> { SimplePrefetcher::new(self.index, seq) } + + fn is_not_plain(&self) -> bool { + true + } } #[derive(Debug)] @@ -418,4 +434,8 @@ impl<'r, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'r, R> { StreamPrefetcher::new(self.index, seq, self.hints.clone()) } + + fn is_not_plain(&self) -> bool { + true + } } diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs index f70eee1c..ae02dc6a 100644 --- a/src/index/vchordrq/opclass.rs +++ b/src/index/vchordrq/opclass.rs @@ -64,7 +64,7 @@ impl Opfamily { } Self::VectorMaxsim => { let vectors = - unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -75,8 +75,9 @@ impl Opfamily { result } Self::HalfvecMaxsim => { - let vectors = - unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -132,7 +133,7 @@ impl Opfamily { let vectors = match self { Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { let vectors = - unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Vecf32(vector.as_borrowed()))); @@ -140,8 +141,9 @@ impl Opfamily { result } Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { - let vectors = - unsafe { pgrx::Array::::from_datum(datum, false).unwrap() }; + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Vecf16(vector.as_borrowed()))); diff --git a/src/lib.rs b/src/lib.rs index 8aa7ab8f..4447d88b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -78,6 +78,6 @@ compile_error!("Target architecture is not supported."); #[cfg(not(miri))] #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] -#[cfg(target_os = "linux")] +#[cfg(any(target_os = "linux", target_os = "macos"))] #[global_allocator] static GLOBAL_ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; From efa126a3c6ec77dd6eb4dafa2d86ecf4e1a8531f Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 6 Aug 2025 17:21:40 +0800 Subject: [PATCH 188/324] perf: inlining more functions (#305) Signed-off-by: usamoi --- crates/algo/src/accessor.rs | 38 ++++ crates/rabitq/src/bit.rs | 18 ++ crates/vchordrq/src/insert.rs | 5 +- crates/vchordrq/src/operator.rs | 382 +++++++++++++------------------- crates/vchordrq/src/search.rs | 34 +-- 5 files changed, 214 insertions(+), 263 deletions(-) diff --git a/crates/algo/src/accessor.rs b/crates/algo/src/accessor.rs index f89d6045..17ce29ab 100644 --- a/crates/algo/src/accessor.rs +++ b/crates/algo/src/accessor.rs @@ -33,18 +33,22 @@ pub trait Accessor2 { impl Accessor2 for () { type Output = (); + #[inline(always)] fn push(&mut self, _: &[E0], _: &[E1]) {} + #[inline(always)] fn finish(self, _: M0, _: M1) -> Self::Output {} } impl> Accessor2 for (A,) { type Output = (A::Output,); + #[inline(always)] fn push(&mut self, input: &[E0], target: &[E1]) { self.0.push(input, target); } + #[inline(always)] fn finish(self, input: M0, target: M1) -> Self::Output { (self.0.finish(input, target),) } @@ -55,11 +59,13 @@ impl, B: Accessor2 Self::Output { (self.0.finish(input, target), self.1.finish(input, target)) } @@ -74,8 +80,10 @@ pub trait Accessor1 { impl Accessor1 for () { type Output = (); + #[inline(always)] fn push(&mut self, _: &[E]) {} + #[inline(always)] fn finish(self, _: M) -> Self::Output {} } @@ -85,10 +93,12 @@ where { type Output = (A::Output,); + #[inline(always)] fn push(&mut self, input: &[E]) { self.0.push(input); } + #[inline(always)] fn finish(self, input: M) -> Self::Output { (self.0.finish(input),) } @@ -101,11 +111,13 @@ where { type Output = (A::Output, B::Output); + #[inline(always)] fn push(&mut self, input: &[E]) { self.0.push(input); self.1.push(input); } + #[inline(always)] fn finish(self, input: M) -> Self::Output { (self.0.finish(input), self.1.finish(input)) } @@ -130,10 +142,12 @@ where { type Output = R; + #[inline(always)] fn push(&mut self, input: &[E]) { (self.p)(&mut self.data, input); } + #[inline(always)] fn finish(self, input: M) -> Self::Output { (self.f)(self.data, input) } @@ -146,6 +160,7 @@ pub struct LAccess<'a, E, M, A> { } impl<'a, E, M, A> LAccess<'a, E, M, A> { + #[inline(always)] pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { Self { elements, @@ -158,12 +173,14 @@ impl<'a, E, M, A> LAccess<'a, E, M, A> { impl> Accessor1 for LAccess<'_, E0, M0, A> { type Output = A::Output; + #[inline(always)] fn push(&mut self, rhs: &[E1]) { let (lhs, elements) = self.elements.split_at(rhs.len()); self.accessor.push(lhs, rhs); self.elements = elements; } + #[inline(always)] fn finish(self, rhs: M1) -> Self::Output { assert!(self.elements.is_empty(), "goal is shorter than expected"); self.accessor.finish(self.metadata, rhs) @@ -177,6 +194,7 @@ pub struct RAccess<'a, E, M, A> { } impl<'a, E, M, A> RAccess<'a, E, M, A> { + #[inline(always)] pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { Self { elements, @@ -189,12 +207,14 @@ impl<'a, E, M, A> RAccess<'a, E, M, A> { impl> Accessor1 for RAccess<'_, E1, M1, A> { type Output = A::Output; + #[inline(always)] fn push(&mut self, lhs: &[E0]) { let (rhs, elements) = self.elements.split_at(lhs.len()); self.accessor.push(lhs, rhs); self.elements = elements; } + #[inline(always)] fn finish(self, lhs: M0) -> Self::Output { assert!(self.elements.is_empty(), "goal is shorter than expected"); self.accessor.finish(lhs, self.metadata) @@ -212,10 +232,12 @@ pub trait TryAccessor1: Sized { impl TryAccessor1 for () { type Output = (); + #[inline(always)] fn push(&mut self, _: &[E]) -> Option<()> { Some(()) } + #[inline(always)] fn finish(self, _: M) -> Option { Some(()) } @@ -227,11 +249,13 @@ where { type Output = (A::Output,); + #[inline(always)] fn push(&mut self, input: &[E]) -> Option<()> { self.0.push(input)?; Some(()) } + #[inline(always)] fn finish(self, input: M) -> Option { Some((self.0.finish(input)?,)) } @@ -244,12 +268,14 @@ where { type Output = (A::Output, B::Output); + #[inline(always)] fn push(&mut self, input: &[E]) -> Option<()> { self.0.push(input)?; self.1.push(input)?; Some(()) } + #[inline(always)] fn finish(self, input: M) -> Option { Some((self.0.finish(input)?, self.1.finish(input)?)) } @@ -262,6 +288,7 @@ pub struct LTryAccess<'a, E, M, A> { } impl<'a, E, M, A> LTryAccess<'a, E, M, A> { + #[inline(always)] pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { Self { elements, @@ -276,6 +303,7 @@ impl> TryAccessor1 { type Output = A::Output; + #[inline(always)] fn push(&mut self, rhs: &[E1]) -> Option<()> { let (lhs, elements) = self.elements.split_at_checked(rhs.len())?; self.accessor.push(lhs, rhs); @@ -283,6 +311,7 @@ impl> TryAccessor1 Some(()) } + #[inline(always)] fn finish(self, rhs: M1) -> Option { if !self.elements.is_empty() { return None; @@ -295,6 +324,7 @@ impl> TryAccessor1 pub struct DistanceAccessor(f32, PhantomData V>, PhantomData D>); impl Default for DistanceAccessor { + #[inline(always)] fn default() -> Self { Self(0.0, PhantomData, PhantomData) } @@ -303,10 +333,12 @@ impl Default for DistanceAccessor { impl Accessor2 for DistanceAccessor, L2S> { type Output = Distance; + #[inline(always)] fn push(&mut self, target: &[f32], input: &[f32]) { self.0 += f32::reduce_sum_of_d2(target, input) } + #[inline(always)] fn finish(self, (): (), (): ()) -> Self::Output { Distance::from_f32(self.0) } @@ -315,10 +347,12 @@ impl Accessor2 for DistanceAccessor, L2S> { impl Accessor2 for DistanceAccessor, Dot> { type Output = Distance; + #[inline(always)] fn push(&mut self, target: &[f32], input: &[f32]) { self.0 += f32::reduce_sum_of_xy(target, input) } + #[inline(always)] fn finish(self, (): (), (): ()) -> Self::Output { Distance::from_f32(-self.0) } @@ -327,10 +361,12 @@ impl Accessor2 for DistanceAccessor, Dot> { impl Accessor2 for DistanceAccessor, L2S> { type Output = Distance; + #[inline(always)] fn push(&mut self, target: &[f16], input: &[f16]) { self.0 += f16::reduce_sum_of_d2(target, input) } + #[inline(always)] fn finish(self, (): (), (): ()) -> Self::Output { Distance::from_f32(self.0) } @@ -339,10 +375,12 @@ impl Accessor2 for DistanceAccessor, L2S> { impl Accessor2 for DistanceAccessor, Dot> { type Output = Distance; + #[inline(always)] fn push(&mut self, target: &[f16], input: &[f16]) { self.0 += f16::reduce_sum_of_xy(target, input) } + #[inline(always)] fn finish(self, (): (), (): ()) -> Self::Output { Distance::from_f32(-self.0) } diff --git a/crates/rabitq/src/bit.rs b/crates/rabitq/src/bit.rs index a24c3111..3d8fe9e2 100644 --- a/crates/rabitq/src/bit.rs +++ b/crates/rabitq/src/bit.rs @@ -25,6 +25,7 @@ pub struct CodeMetadata { } impl CodeMetadata { + #[inline(always)] pub fn into_tuple(self) -> (f32, f32, f32, f32) { ( self.dis_u_2, @@ -33,6 +34,7 @@ impl CodeMetadata { self.factor_err, ) } + #[inline(always)] pub fn into_array(self) -> [f32; 4] { [ self.dis_u_2, @@ -41,6 +43,7 @@ impl CodeMetadata { self.factor_err, ] } + #[inline(always)] pub fn from_tuple((dis_u_2, factor_cnt, factor_ip, factor_err): (f32, f32, f32, f32)) -> Self { Self { dis_u_2, @@ -49,6 +52,7 @@ impl CodeMetadata { factor_err, } } + #[inline(always)] pub fn from_array([dis_u_2, factor_cnt, factor_ip, factor_err]: [f32; 4]) -> Self { Self { dis_u_2, @@ -140,6 +144,7 @@ pub mod binary { pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u64]); + #[inline(always)] pub fn preprocess(vector: &[f32]) -> BinaryLut { let dis_v_2 = f32::reduce_sum_of_x2(vector); preprocess_with_distance(vector, dis_v_2) @@ -163,6 +168,7 @@ pub mod binary { ) } + #[inline(always)] pub fn accumulate(x: &[u64], y: &[impl AsRef<[u64]>; BITS]) -> u32 { let mut result = 0_u32; for i in 0..BITS { @@ -171,6 +177,7 @@ pub mod binary { result } + #[inline(always)] pub fn half_process_l2( sum: u32, CodeMetadata { @@ -192,6 +199,7 @@ pub mod binary { (rough, err) } + #[inline(always)] pub fn half_process_l2_residual( sum: u32, CodeMetadata { @@ -215,6 +223,7 @@ pub mod binary { (rough, err) } + #[inline(always)] pub fn half_process_dot( sum: u32, CodeMetadata { @@ -236,6 +245,7 @@ pub mod binary { (rough, err) } + #[inline(always)] pub fn half_process_dot_residual( sum: u32, CodeMetadata { @@ -260,6 +270,7 @@ pub mod binary { (rough, err) } + #[inline(always)] pub(crate) fn binarize(vector: &[u8]) -> [Vec; BITS] { let n = vector.len(); let mut t: [_; BITS] = std::array::from_fn(|_| vec![0_u64; n.div_ceil(64)]); @@ -296,6 +307,7 @@ pub mod block { &'a [[u8; 16]], ); + #[inline(always)] pub fn preprocess(vector: &[f32]) -> BlockLut { let dis_v_2 = f32::reduce_sum_of_x2(vector); preprocess_with_distance(vector, dis_v_2) @@ -315,6 +327,7 @@ pub mod block { ) } + #[inline(always)] pub fn full_process_l2( (dis_u_2, _, factor_ip, factor_err, t): BlockCode<'_>, lut: &BlockLut, @@ -334,6 +347,7 @@ pub mod block { }) } + #[inline(always)] pub fn full_process_dot( (_, _, factor_ip, factor_err, t): BlockCode<'_>, lut: &BlockLut, @@ -353,6 +367,7 @@ pub mod block { }) } + #[inline(always)] pub fn half_process_l2( sum: u32, CodeMetadata { @@ -369,6 +384,7 @@ pub mod block { (rough, err) } + #[inline(always)] pub fn half_process_l2_residual( sum: u32, CodeMetadata { @@ -387,6 +403,7 @@ pub mod block { (rough, err) } + #[inline(always)] pub fn half_process_dot( sum: u32, CodeMetadata { @@ -403,6 +420,7 @@ pub mod block { (rough, err) } + #[inline(always)] pub fn half_process_dot_residual( sum: u32, CodeMetadata { diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index ab43a90a..1ec50219 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -105,12 +105,9 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); { let (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) = state; - let process = |value, code, delta, lut| { - O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) - }; tape::read_h1_tape::( by_next(index, first), - || O::block_access(&lut.0, process), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), |(rough, err), head, norm, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( diff --git a/crates/vchordrq/src/operator.rs b/crates/vchordrq/src/operator.rs index 0ee9ac86..b26542f9 100644 --- a/crates/vchordrq/src/operator.rs +++ b/crates/vchordrq/src/operator.rs @@ -16,8 +16,8 @@ use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S, RAccess}; use distance::Distance; use half::f16; use rabitq::bit::CodeMetadata; -use rabitq::bit::binary::{BinaryLut, BinaryLutMetadata}; -use rabitq::bit::block::{BlockLut, BlockLutMetadata, STEP}; +use rabitq::bit::binary::BinaryLut; +use rabitq::bit::block::{BlockLut, STEP}; use simd::Floating; use std::fmt::Debug; use std::marker::PhantomData; @@ -28,12 +28,12 @@ use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[derive(Debug)] pub struct BlockAccessor([u32; 32], F); -impl> - Accessor2<[u8; 16], [u8; 16], (&[[f32; 32]; 4], &[f32; 32]), BlockLutMetadata> - for BlockAccessor +impl> + Accessor2<[u8; 16], [u8; 16], (&[[f32; 32]; 4], &[f32; 32]), ()> for BlockAccessor { type Output = [F::Output; 32]; + #[inline(always)] fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { use std::iter::zip; @@ -43,11 +43,8 @@ impl> } } - fn finish( - mut self, - (metadata, delta): (&[[f32; 32]; 4], &[f32; 32]), - lut: BlockLutMetadata, - ) -> Self::Output { + #[inline(always)] + fn finish(mut self, (metadata, delta): (&[[f32; 32]; 4], &[f32; 32]), (): ()) -> Self::Output { std::array::from_fn(|i| { (self.1).call( self.0[i], @@ -58,7 +55,6 @@ impl> factor_err: metadata[3][i], }, delta[i], - lut, ) }) } @@ -68,6 +64,7 @@ impl> pub struct CloneAccessor(Vec); impl Default for CloneAccessor { + #[inline(always)] fn default() -> Self { Self(Vec::new()) } @@ -76,10 +73,12 @@ impl Default for CloneAccessor { impl Accessor1 for CloneAccessor { type Output = V; + #[inline(always)] fn push(&mut self, input: &[V::Element]) { self.0.extend(input); } + #[inline(always)] fn finish(self, metadata: V::Metadata) -> Self::Output { V::pack(self.0, metadata) } @@ -223,35 +222,19 @@ pub trait Operator: 'static + Debug + Copy { Output = Distance, >; - fn block_access>( + fn block_access( lut: &BlockLut, - f: F, - ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]>; - - fn binary_access>( - lut: &BinaryLut, - f: F, - ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output; - - fn block_process( - value: u32, - code: CodeMetadata, - lut: BlockLutMetadata, is_residual: bool, dis_f: f32, - delta: f32, norm: f32, - ) -> (f32, f32); + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]>; - fn binary_process( - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, + fn binary_access( + lut: &BinaryLut, is_residual: bool, dis_f: f32, - delta: f32, norm: f32, - ) -> (f32, f32); + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32); fn build( vector: ::Borrowed<'_>, @@ -275,63 +258,44 @@ impl Operator for Op, L2S> { type DistanceAccessor = DistanceAccessor, L2S>; - fn block_access>( + fn block_access( lut: &BlockLut, - f: F, - ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]> - { - RAccess::new((&lut.1, lut.0), BlockAccessor([0u32; 32], f)) - } - - fn binary_access>( - lut: &BinaryLut, - mut f: F, - ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { - move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::bit::binary::accumulate(elements, &lut.1); - f.call( - value, - CodeMetadata { - dis_u_2: metadata[0], - factor_cnt: metadata[1], - factor_ip: metadata[2], - factor_err: metadata[3], - }, - delta, - lut.0, - ) - } - } - - fn block_process( - value: u32, - code: CodeMetadata, - lut: BlockLutMetadata, is_residual: bool, dis_f: f32, - delta: f32, - _: f32, - ) -> (f32, f32) { - if !is_residual { - rabitq::bit::block::half_process_l2(value, code, lut) - } else { - rabitq::bit::block::half_process_l2_residual(value, code, lut, dis_f, delta) - } + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, delta| { + if !is_residual { + rabitq::bit::block::half_process_l2(value, code, lut.0) + } else { + rabitq::bit::block::half_process_l2_residual(value, code, lut.0, dis_f, delta) + } + }), + ) } - fn binary_process( - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, + fn binary_access( + lut: &BinaryLut, is_residual: bool, dis_f: f32, - delta: f32, - _: f32, - ) -> (f32, f32) { - if !is_residual { - rabitq::bit::binary::half_process_l2(value, code, lut) - } else { - rabitq::bit::binary::half_process_l2_residual(value, code, lut, dis_f, delta) + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + if !is_residual { + rabitq::bit::binary::half_process_l2(value, code, lut.0) + } else { + rabitq::bit::binary::half_process_l2_residual(value, code, lut.0, dis_f, delta) + } } } @@ -369,63 +333,48 @@ impl Operator for Op, Dot> { type DistanceAccessor = DistanceAccessor, Dot>; - fn block_access>( + fn block_access( lut: &BlockLut, - f: F, - ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]> - { - RAccess::new((&lut.1, lut.0), BlockAccessor([0u32; 32], f)) - } - - fn binary_access>( - lut: &BinaryLut, - mut f: F, - ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { - move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::bit::binary::accumulate(elements, &lut.1); - f.call( - value, - CodeMetadata { - dis_u_2: metadata[0], - factor_cnt: metadata[1], - factor_ip: metadata[2], - factor_err: metadata[3], - }, - delta, - lut.0, - ) - } - } - - fn block_process( - value: u32, - code: CodeMetadata, - lut: BlockLutMetadata, is_residual: bool, dis_f: f32, - delta: f32, norm: f32, - ) -> (f32, f32) { - if !is_residual { - rabitq::bit::block::half_process_dot(value, code, lut) - } else { - rabitq::bit::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) - } + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, delta| { + if !is_residual { + rabitq::bit::block::half_process_dot(value, code, lut.0) + } else { + rabitq::bit::block::half_process_dot_residual( + value, code, lut.0, dis_f, delta, norm, + ) + } + }), + ) } - fn binary_process( - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, + fn binary_access( + lut: &BinaryLut, is_residual: bool, dis_f: f32, - delta: f32, norm: f32, - ) -> (f32, f32) { - if !is_residual { - rabitq::bit::binary::half_process_dot(value, code, lut) - } else { - rabitq::bit::binary::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + if !is_residual { + rabitq::bit::binary::half_process_dot(value, code, lut.0) + } else { + rabitq::bit::binary::half_process_dot_residual( + value, code, lut.0, dis_f, delta, norm, + ) + } } } @@ -463,63 +412,44 @@ impl Operator for Op, L2S> { type DistanceAccessor = DistanceAccessor, L2S>; - fn block_access>( + fn block_access( lut: &BlockLut, - f: F, - ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]> - { - RAccess::new((&lut.1, lut.0), BlockAccessor([0u32; 32], f)) - } - - fn binary_access>( - lut: &BinaryLut, - mut f: F, - ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { - move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::bit::binary::accumulate(elements, &lut.1); - f.call( - value, - CodeMetadata { - dis_u_2: metadata[0], - factor_cnt: metadata[1], - factor_ip: metadata[2], - factor_err: metadata[3], - }, - delta, - lut.0, - ) - } - } - - fn block_process( - value: u32, - code: CodeMetadata, - lut: BlockLutMetadata, is_residual: bool, dis_f: f32, - delta: f32, - _: f32, - ) -> (f32, f32) { - if !is_residual { - rabitq::bit::block::half_process_l2(value, code, lut) - } else { - rabitq::bit::block::half_process_l2_residual(value, code, lut, dis_f, delta) - } + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, delta| { + if !is_residual { + rabitq::bit::block::half_process_l2(value, code, lut.0) + } else { + rabitq::bit::block::half_process_l2_residual(value, code, lut.0, dis_f, delta) + } + }), + ) } - fn binary_process( - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, + fn binary_access( + lut: &BinaryLut, is_residual: bool, dis_f: f32, - delta: f32, - _: f32, - ) -> (f32, f32) { - if !is_residual { - rabitq::bit::binary::half_process_l2(value, code, lut) - } else { - rabitq::bit::binary::half_process_l2_residual(value, code, lut, dis_f, delta) + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + if !is_residual { + rabitq::bit::binary::half_process_l2(value, code, lut.0) + } else { + rabitq::bit::binary::half_process_l2_residual(value, code, lut.0, dis_f, delta) + } } } @@ -557,63 +487,48 @@ impl Operator for Op, Dot> { type DistanceAccessor = DistanceAccessor, Dot>; - fn block_access>( + fn block_access( lut: &BlockLut, - f: F, - ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [F::Output; 32]> - { - RAccess::new((&lut.1, lut.0), BlockAccessor([0u32; 32], f)) - } - - fn binary_access>( - lut: &BinaryLut, - mut f: F, - ) -> impl FnMut([f32; 4], &[u64], f32) -> F::Output { - move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::bit::binary::accumulate(elements, &lut.1); - f.call( - value, - CodeMetadata { - dis_u_2: metadata[0], - factor_cnt: metadata[1], - factor_ip: metadata[2], - factor_err: metadata[3], - }, - delta, - lut.0, - ) - } - } - - fn block_process( - value: u32, - code: CodeMetadata, - lut: BlockLutMetadata, is_residual: bool, dis_f: f32, - delta: f32, norm: f32, - ) -> (f32, f32) { - if !is_residual { - rabitq::bit::block::half_process_dot(value, code, lut) - } else { - rabitq::bit::block::half_process_dot_residual(value, code, lut, dis_f, delta, norm) - } + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, delta| { + if !is_residual { + rabitq::bit::block::half_process_dot(value, code, lut.0) + } else { + rabitq::bit::block::half_process_dot_residual( + value, code, lut.0, dis_f, delta, norm, + ) + } + }), + ) } - fn binary_process( - value: u32, - code: CodeMetadata, - lut: BinaryLutMetadata, + fn binary_access( + lut: &BinaryLut, is_residual: bool, dis_f: f32, - delta: f32, norm: f32, - ) -> (f32, f32) { - if !is_residual { - rabitq::bit::binary::half_process_dot(value, code, lut) - } else { - rabitq::bit::binary::half_process_dot_residual(value, code, lut, dis_f, delta, norm) + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + if !is_residual { + rabitq::bit::binary::half_process_dot(value, code, lut.0) + } else { + rabitq::bit::binary::half_process_dot_residual( + value, code, lut.0, dis_f, delta, norm, + ) + } } } @@ -646,16 +561,17 @@ impl Operator for Op, Dot> { } } -pub trait Call { +pub trait Call { type Output; - fn call(&mut self, a: A, b: B, c: C, d: D) -> Self::Output; + fn call(&mut self, a: A, b: B, c: C) -> Self::Output; } -impl R, R> Call for F { +impl R, R> Call for F { type Output = R; - fn call(&mut self, a: A, b: B, c: C, d: D) -> R { - (self)(a, b, c, d) + #[inline(always)] + fn call(&mut self, a: A, b: B, c: C) -> R { + (self)(a, b, c) } } diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index fa8e02f2..015b6aa6 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -85,12 +85,9 @@ where let mut step = |state: State| { let mut results = LinkedVec::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { - let process = |value, code, delta, lut| { - O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) - }; tape::read_h1_tape::( by_next(index, first), - || O::block_access(&lut.0, process), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), |(rough, err), head, norm, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( @@ -123,12 +120,6 @@ where let mut results = LinkedVec::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { - let block_process = |value, code, delta, lut| { - O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) - }; - let binary_process = |value, code, delta, lut| { - O::binary_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) - }; let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); @@ -144,19 +135,19 @@ where tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); tape::read_frozen_tape::( by_directory(&mut prefetch_h0_tuples, directory), - || O::block_access(&lut.0, block_process), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), &mut callback, ); } else { tape::read_frozen_tape::( by_next(index, jump_tuple.frozen_first()), - || O::block_access(&lut.0, block_process), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), &mut callback, ); } tape::read_appendable_tape::( by_next(index, jump_tuple.appendable_first()), - O::binary_access(&lut.1, binary_process), + O::binary_access(&lut.1, is_residual, dis_f.to_f32(), norm), &mut callback, ); } @@ -225,12 +216,9 @@ where let mut step = |state: State| { let mut results = LinkedVec::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { - let process = |value, code, delta, lut| { - O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) - }; tape::read_h1_tape::( by_next(index, first), - || O::block_access(&lut.0, process), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), |(rough, err), head, norm, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( @@ -265,12 +253,6 @@ where let mut results = LinkedVec::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { - let block_process = |value, code, delta, lut| { - O::block_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) - }; - let binary_process = |value, code, delta, lut| { - O::binary_process(value, code, lut, is_residual, dis_f.to_f32(), delta, norm) - }; let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); @@ -287,19 +269,19 @@ where tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); tape::read_frozen_tape::( by_directory(&mut prefetch_h0_tuples, directory), - || O::block_access(&lut.0, block_process), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), &mut callback, ); } else { tape::read_frozen_tape::( by_next(index, jump_tuple.frozen_first()), - || O::block_access(&lut.0, block_process), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), &mut callback, ); } tape::read_appendable_tape::( by_next(index, jump_tuple.appendable_first()), - O::binary_access(&lut.1, binary_process), + O::binary_access(&lut.1, is_residual, dis_f.to_f32(), norm), &mut callback, ); threshold = threshold.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as _); From f2268621f2c6a4e66a4766d92f514d5d400a3b39 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 7 Aug 2025 17:44:19 +0800 Subject: [PATCH 189/324] perf: use stack-allocated vectors (#306) Signed-off-by: usamoi --- Cargo.lock | 18 ++- Cargo.toml | 3 +- crates/algo/Cargo.toml | 1 + crates/algo/src/accessor.rs | 1 + crates/algo/src/lib.rs | 67 ++++++++--- crates/algo/src/prefetcher.rs | 38 ++++-- crates/algo/src/tuples.rs | 4 + crates/always_equal/src/lib.rs | 9 +- crates/sbsii/Cargo.toml | 10 ++ crates/sbsii/src/borrowed.rs | 119 +++++++++++++++++++ crates/sbsii/src/lib.rs | 19 +++ crates/sbsii/src/owned.rs | 156 +++++++++++++++++++++++++ crates/sbsii/src/stack.rs | 61 ++++++++++ crates/vchordrq/src/insert.rs | 19 +-- crates/vchordrq/src/prewarm.rs | 9 +- crates/vchordrq/src/rerank.rs | 8 +- crates/vchordrq/src/search.rs | 50 +++++--- src/index/storage.rs | 12 +- src/index/vchordg/algo.rs | 2 +- src/index/vchordg/am/mod.rs | 2 +- src/index/vchordrq/algo.rs | 39 +++---- src/index/vchordrq/am/mod.rs | 2 +- src/index/vchordrq/scanners/default.rs | 2 +- src/index/vchordrq/scanners/maxsim.rs | 2 +- 24 files changed, 546 insertions(+), 107 deletions(-) create mode 100644 crates/sbsii/Cargo.toml create mode 100644 crates/sbsii/src/borrowed.rs create mode 100644 crates/sbsii/src/lib.rs create mode 100644 crates/sbsii/src/owned.rs create mode 100644 crates/sbsii/src/stack.rs diff --git a/Cargo.lock b/Cargo.lock index 491e18bf..df8d3d28 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -19,6 +19,7 @@ dependencies = [ "bumpalo", "distance", "half 2.6.0", + "sbsii", "simd", "vector", "zerocopy", @@ -150,7 +151,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "374b7c592d9c00c1f4972ea58390ac6b18cbb6ab79011f3bdc90a0b82ca06b77" dependencies = [ "serde", - "toml 0.9.4", + "toml 0.9.5", ] [[package]] @@ -1189,6 +1190,10 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "sbsii" +version = "0.0.0" + [[package]] name = "seahash" version = "4.1.0" @@ -1410,9 +1415,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.9.4" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41ae868b5a0f67631c14589f7e250c1ea2c574ee5ba21c6c8dd4b1485705a5a1" +checksum = "75129e1dc5000bfbaa9fee9d1b21f974f9fbad9daec557a521ee6e080825f6e8" dependencies = [ "indexmap", "serde", @@ -1457,9 +1462,9 @@ dependencies = [ [[package]] name = "toml_parser" -version = "1.0.1" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97200572db069e74c512a14117b296ba0a80a30123fbbb5aa1f4a348f639ca30" +checksum = "b551886f449aa90d4fe2bdaa9f4a2577ad2dde302c61ecf262d80b116db95c10" dependencies = [ "winnow", ] @@ -1580,10 +1585,11 @@ dependencies = [ "pgrx-catalog", "rabitq", "rand", + "sbsii", "seq-macro", "serde", "simd", - "toml 0.9.4", + "toml 0.9.5", "validator", "vchordg", "vchordrq", diff --git a/Cargo.toml b/Cargo.toml index a2c6ae94..9bcada1e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ always_equal = { path = "./crates/always_equal" } distance = { path = "./crates/distance" } k_means = { path = "./crates/k_means" } rabitq = { path = "./crates/rabitq" } +sbsii = { path = "./crates/sbsii" } simd = { path = "./crates/simd" } vchordg = { path = "./crates/vchordg" } vchordrq = { path = "./crates/vchordrq" } @@ -40,7 +41,7 @@ pgrx-catalog = "0.3.1" rand.workspace = true seq-macro.workspace = true serde.workspace = true -toml = "0.9.4" +toml = "0.9.5" validator.workspace = true zerocopy.workspace = true diff --git a/crates/algo/Cargo.toml b/crates/algo/Cargo.toml index 25568150..198acfca 100644 --- a/crates/algo/Cargo.toml +++ b/crates/algo/Cargo.toml @@ -7,6 +7,7 @@ publish = false [dependencies] always_equal = { path = "../always_equal" } distance = { path = "../distance" } +sbsii = { path = "../sbsii" } simd = { path = "../simd" } vector = { path = "../vector" } diff --git a/crates/algo/src/accessor.rs b/crates/algo/src/accessor.rs index 17ce29ab..9a361481 100644 --- a/crates/algo/src/accessor.rs +++ b/crates/algo/src/accessor.rs @@ -130,6 +130,7 @@ pub struct FunctionalAccessor { } impl FunctionalAccessor { + #[inline(always)] pub fn new(data: T, p: P, f: F) -> Self { Self { data, p, f } } diff --git a/crates/algo/src/lib.rs b/crates/algo/src/lib.rs index 057cefcd..a80c187c 100644 --- a/crates/algo/src/lib.rs +++ b/crates/algo/src/lib.rs @@ -32,6 +32,7 @@ pub trait Page: Sized + 'static { #[must_use] fn len(&self) -> u16; #[must_use] + #[inline] fn is_empty(&self) -> bool { self.len() == 0 } @@ -92,6 +93,7 @@ pub trait RelationWrite: RelationWriteTypes { pub trait RelationLength: Relation { fn len(&self) -> u32; + #[inline] fn is_empty(&self) -> bool { self.len() == 0 } @@ -108,6 +110,7 @@ pub struct Hints { impl Hints { #[allow(clippy::needless_update)] + #[inline] pub fn full(self, full: bool) -> Self { Self { full, ..self } } @@ -140,60 +143,73 @@ pub trait Bump: 'static { } impl Bump for bumpalo::Bump { + #[inline] fn alloc(&self, value: T) -> &mut T { self.alloc(value) } + #[inline] fn alloc_slice(&self, slice: &[T]) -> &mut [T] { self.alloc_slice_copy(slice) } + #[inline] fn reset(&mut self) { self.reset(); } } +pub type BorrowedIter<'b> = sbsii::borrowed::IntoIter<'b, u32, 1>; +pub type OwnedIter = sbsii::owned::IntoIter; + pub trait Fetch { - fn fetch(&self) -> Vec; + #[must_use] + fn fetch(&self) -> OwnedIter; } impl Fetch for u32 { - fn fetch(&self) -> Vec { - std::slice::from_ref(self).to_vec() + #[inline(always)] + fn fetch(&self) -> OwnedIter { + OwnedIter::from_slice(std::slice::from_ref(self)) } } -impl<'b, T, A, B> Fetch for (T, AlwaysEqual<&'b mut (A, B, &'b mut [u32])>) { - fn fetch(&self) -> Vec { +impl Fetch for (T, AlwaysEqual<&mut (A, B, OwnedIter)>) { + #[inline(always)] + fn fetch(&self) -> OwnedIter { let (_, AlwaysEqual((.., list))) = self; - list.to_vec() + list.clone() } } -impl<'b, T, A, B, C> Fetch for (T, AlwaysEqual<&'b mut (A, B, C, &'b mut [u32])>) { - fn fetch(&self) -> Vec { +impl Fetch for (T, AlwaysEqual<&mut (A, B, C, OwnedIter)>) { + #[inline(always)] + fn fetch(&self) -> OwnedIter { let (_, AlwaysEqual((.., list))) = self; - list.to_vec() + list.clone() } } impl Fetch for (T, AlwaysEqual<(u32, u16)>) { - fn fetch(&self) -> Vec { + #[inline(always)] + fn fetch(&self) -> OwnedIter { let (_, AlwaysEqual((x, _))) = self; - std::slice::from_ref(x).to_vec() + OwnedIter::from_slice(std::slice::from_ref(x)) } } impl Fetch for (u32, u16) { - fn fetch(&self) -> Vec { - std::slice::from_ref(&self.0).to_vec() + #[inline(always)] + fn fetch(&self) -> OwnedIter { + OwnedIter::from_slice(std::slice::from_ref(&self.0)) } } impl Fetch for (T, AlwaysEqual<((u32, u16), (u32, u16))>) { - fn fetch(&self) -> Vec { + #[inline(always)] + fn fetch(&self) -> OwnedIter { let (_, AlwaysEqual(((x, _), _))) = self; - std::slice::from_ref(x).to_vec() + OwnedIter::from_slice(std::slice::from_ref(x)) } } @@ -202,14 +218,18 @@ pub trait Fetch1 { } impl Fetch for (T, AlwaysEqual<&mut [F]>) { - fn fetch(&self) -> Vec { - self.1.0.iter().map(|x| x.fetch_1()).collect() + #[inline(always)] + fn fetch(&self) -> OwnedIter { + let vec = self.1.0.iter().map(|x| x.fetch_1()).collect::>(); + OwnedIter::from_slice(vec.as_slice()) } } impl Fetch for (T, AlwaysEqual<(&mut [F], U)>) { - fn fetch(&self) -> Vec { - self.1.0.0.iter().map(|x| x.fetch_1()).collect() + #[inline(always)] + fn fetch(&self) -> OwnedIter { + let vec = self.1.0.0.iter().map(|x| x.fetch_1()).collect::>(); + OwnedIter::from_slice(vec.as_slice()) } } @@ -227,12 +247,15 @@ pub trait Sequence { impl Sequence for BinaryHeap { type Item = T; type Inner = std::vec::IntoIter; + #[inline] fn next(&mut self) -> Option { self.pop() } + #[inline] fn peek(&mut self) -> Option<&T> { (self as &Self).peek() } + #[inline] fn into_inner(self) -> Self::Inner { self.into_vec().into_iter() } @@ -241,12 +264,15 @@ impl Sequence for BinaryHeap { impl Sequence for Peekable { type Item = I::Item; type Inner = Peekable; + #[inline] fn next(&mut self) -> Option { Iterator::next(self) } + #[inline] fn peek(&mut self) -> Option<&I::Item> { self.peek() } + #[inline] fn into_inner(self) -> Self::Inner { self } @@ -255,12 +281,15 @@ impl Sequence for Peekable { impl Sequence for VecDeque { type Item = T; type Inner = std::collections::vec_deque::IntoIter; + #[inline] fn next(&mut self) -> Option { self.pop_front() } + #[inline] fn peek(&mut self) -> Option<&T> { self.front() } + #[inline] fn into_inner(self) -> Self::Inner { self.into_iter() } diff --git a/crates/algo/src/prefetcher.rs b/crates/algo/src/prefetcher.rs index 2150c887..f2159e64 100644 --- a/crates/algo/src/prefetcher.rs +++ b/crates/algo/src/prefetcher.rs @@ -44,6 +44,7 @@ pub struct PlainPrefetcher<'r, R, S: Sequence> { } impl<'r, R, S: Sequence> PlainPrefetcher<'r, R, S> { + #[inline] pub fn new(relation: &'r R, sequence: S) -> Self { Self { relation, sequence } } @@ -54,6 +55,7 @@ impl<'r, R, S: Sequence> IntoIterator for PlainPrefetcher<'r, R, S> { type IntoIter = S::Inner; + #[inline] fn into_iter(self) -> Self::IntoIter { self.sequence.into_inner() } @@ -66,9 +68,10 @@ where type R = R; type Guards = PlainPrefetcherGuards<'r, R>; + #[inline] fn next(&mut self) -> Option<(Self::Item, PlainPrefetcherGuards<'r, R>)> { let e = self.sequence.next()?; - let list = e.fetch().into_iter(); + let list = e.fetch(); Some(( e, PlainPrefetcherGuards { @@ -78,6 +81,7 @@ where )) } + #[inline] fn next_if( &mut self, predicate: impl FnOnce(&Self::Item) -> bool, @@ -86,7 +90,7 @@ where return None; } let e = self.sequence.next()?; - let list = e.fetch().into_iter(); + let list = e.fetch(); Some(( e, PlainPrefetcherGuards { @@ -99,17 +103,19 @@ where pub struct PlainPrefetcherGuards<'r, R> { relation: &'r R, - list: std::vec::IntoIter, + list: crate::OwnedIter, } impl<'r, R: RelationRead> Iterator for PlainPrefetcherGuards<'r, R> { type Item = R::ReadGuard<'r>; + #[inline] fn next(&mut self) -> Option { let id = self.list.next()?; Some(self.relation.read(id)) } + #[inline] fn size_hint(&self) -> (usize, Option) { self.list.size_hint() } @@ -124,6 +130,7 @@ pub struct SimplePrefetcher<'r, R, S: Sequence> { } impl<'r, R, S: Sequence> SimplePrefetcher<'r, R, S> { + #[inline] pub fn new(relation: &'r R, sequence: S) -> Self { Self { relation, @@ -138,6 +145,7 @@ impl<'r, R, S: Sequence> IntoIterator for SimplePrefetcher<'r, R, S> { type IntoIter = Chain, S::Inner>; + #[inline] fn into_iter(self) -> Self::IntoIter { self.window.into_iter().chain(self.sequence.into_inner()) } @@ -151,17 +159,18 @@ where type R = R; type Guards = SimplePrefetcherGuards<'r, R>; + #[inline] fn next(&mut self) -> Option<(Self::Item, SimplePrefetcherGuards<'r, R>)> { while self.window.len() < WINDOW_SIZE && let Some(e) = self.sequence.next() { - for id in e.fetch().into_iter() { + for id in e.fetch() { self.relation.prefetch(id); } self.window.push_back(e); } let e = self.window.pop_front()?; - let list = e.fetch().into_iter(); + let list = e.fetch(); Some(( e, SimplePrefetcherGuards { @@ -170,6 +179,8 @@ where }, )) } + + #[inline] fn next_if( &mut self, predicate: impl FnOnce(&S::Item) -> bool, @@ -177,13 +188,13 @@ where while self.window.len() < WINDOW_SIZE && let Some(e) = self.sequence.next() { - for id in e.fetch().into_iter() { + for id in e.fetch() { self.relation.prefetch(id); } self.window.push_back(e); } let e = vec_deque_pop_front_if(&mut self.window, predicate)?; - let list = e.fetch().into_iter(); + let list = e.fetch(); Some(( e, SimplePrefetcherGuards { @@ -196,17 +207,19 @@ where pub struct SimplePrefetcherGuards<'r, R> { relation: &'r R, - list: std::vec::IntoIter, + list: crate::OwnedIter, } impl<'r, R: RelationRead> Iterator for SimplePrefetcherGuards<'r, R> { type Item = R::ReadGuard<'r>; + #[inline] fn next(&mut self) -> Option { let id = self.list.next()?; Some(self.relation.read(id)) } + #[inline] fn size_hint(&self) -> (usize, Option) { self.list.size_hint() } @@ -218,6 +231,8 @@ pub struct StreamPrefetcherSequence(S); impl Iterator for StreamPrefetcherSequence { type Item = S::Item; + + #[inline] fn next(&mut self) -> Option { self.0.next() } @@ -234,6 +249,7 @@ impl<'r, R: RelationReadStream, S: Sequence> StreamPrefetcher<'r, R, S> where S::Item: Fetch, { + #[inline] pub fn new(relation: &'r R, sequence: S, hints: Hints) -> Self { let stream = relation.read_stream(StreamPrefetcherSequence(sequence), hints); Self { stream } @@ -248,6 +264,7 @@ where type IntoIter = > as ReadStream<'r>>::Inner; + #[inline] fn into_iter(self) -> Self::IntoIter { self.stream.into_inner() } @@ -259,10 +276,15 @@ where S::Item: Fetch, { type R = R; + type Guards = > as ReadStream<'r>>::Guards; + + #[inline] fn next(&mut self) -> Option<(S::Item, Self::Guards)> { self.stream.next() } + + #[inline] fn next_if( &mut self, predicate: impl FnOnce(&Self::Item) -> bool, diff --git a/crates/algo/src/tuples.rs b/crates/algo/src/tuples.rs index fd741f70..8483db05 100644 --- a/crates/algo/src/tuples.rs +++ b/crates/algo/src/tuples.rs @@ -23,6 +23,7 @@ pub struct RefChecker<'a> { } impl<'a> RefChecker<'a> { + #[inline(always)] pub fn new(bytes: &'a [u8]) -> Self { Self { bytes } } @@ -80,6 +81,7 @@ pub struct MutChecker<'a> { } impl<'a> MutChecker<'a> { + #[inline(always)] pub fn new(bytes: &'a mut [u8]) -> Self { Self { flag: 0, @@ -221,12 +223,14 @@ impl Bool { } impl From for bool { + #[inline] fn from(value: Bool) -> Self { value != Bool::FALSE } } impl From for Bool { + #[inline] fn from(value: bool) -> Self { std::hint::select_unpredictable(value, Self::TRUE, Self::FALSE) } diff --git a/crates/always_equal/src/lib.rs b/crates/always_equal/src/lib.rs index c05d10f9..ff4b51fd 100644 --- a/crates/always_equal/src/lib.rs +++ b/crates/always_equal/src/lib.rs @@ -20,6 +20,7 @@ use std::hash::Hash; pub struct AlwaysEqual(pub T); impl PartialEq for AlwaysEqual { + #[inline(always)] fn eq(&self, _: &Self) -> bool { true } @@ -27,18 +28,22 @@ impl PartialEq for AlwaysEqual { impl Eq for AlwaysEqual {} +#[allow(clippy::non_canonical_partial_ord_impl)] impl PartialOrd for AlwaysEqual { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) + #[inline(always)] + fn partial_cmp(&self, _: &Self) -> Option { + Some(Ordering::Equal) } } impl Ord for AlwaysEqual { + #[inline(always)] fn cmp(&self, _: &Self) -> Ordering { Ordering::Equal } } impl Hash for AlwaysEqual { + #[inline(always)] fn hash(&self, _: &mut H) {} } diff --git a/crates/sbsii/Cargo.toml b/crates/sbsii/Cargo.toml new file mode 100644 index 00000000..4af4bd3c --- /dev/null +++ b/crates/sbsii/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "sbsii" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] + +[lints] +workspace = true diff --git a/crates/sbsii/src/borrowed.rs b/crates/sbsii/src/borrowed.rs new file mode 100644 index 00000000..af7eb462 --- /dev/null +++ b/crates/sbsii/src/borrowed.rs @@ -0,0 +1,119 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::stack::StackIntoIter; + +#[derive(Clone)] +pub struct HeapIntoIter<'a, T> { + inner: std::iter::Copied>, +} + +impl<'a, T: Copy> HeapIntoIter<'a, T> { + #[cold] + pub(crate) fn from_slice(slice: &[T], alloc: impl Fn(&[T]) -> &'a [T]) -> Self { + assert!(slice.len() <= 65535_usize); + Self { + inner: alloc(slice).iter().copied(), + } + } +} + +impl<'a, T: Copy> Iterator for HeapIntoIter<'a, T> { + type Item = T; + + #[inline(always)] + fn next(&mut self) -> Option { + self.inner.next() + } +} + +#[derive(Clone)] +pub enum IntoIter<'a, T: Copy, const N: usize> { + Stack(StackIntoIter), + Heap(HeapIntoIter<'a, T>), +} + +impl<'a, T: Copy, const N: usize> IntoIter<'a, T, N> { + #[inline(always)] + pub fn from_slice(slice: &[T], alloc: impl Fn(&[T]) -> &'a [T]) -> Self { + assert!(slice.len() <= 65535); + if slice.len() <= N && N <= 65535 { + IntoIter::Stack(StackIntoIter::from_slice(slice)) + } else { + IntoIter::Heap(HeapIntoIter::from_slice(slice, alloc)) + } + } +} + +impl<'a, T: Copy, const N: usize> Iterator for IntoIter<'a, T, N> { + type Item = T; + + #[inline(always)] + fn next(&mut self) -> Option { + match self { + IntoIter::Stack(iter) => iter.next(), + IntoIter::Heap(iter) => iter.next(), + } + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + match self { + IntoIter::Stack(iter) => iter.size_hint(), + IntoIter::Heap(iter) => iter.size_hint(), + } + } +} + +impl<'a, T: Copy, const N: usize> ExactSizeIterator for IntoIter<'a, T, N> {} + +#[cfg(target_pointer_width = "64")] +const _: () = { + assert!(size_of::>() == 16); +}; + +#[test] +fn tests() { + let alloc = |slice: &[u32]| -> &'static [u32] { + // make miri happy + static GLOBAL: std::sync::Mutex> = std::sync::Mutex::new(Vec::new()); + let pointer = Vec::leak(slice.to_vec()); + GLOBAL.lock().expect("failed to lock").push(pointer); + pointer + }; + for x in IntoIter::::from_slice(&[1; 0], alloc) { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 1], alloc) { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 2], alloc) { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 3], alloc) { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 4], alloc) { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 5], alloc) { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 6], alloc) { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 7], alloc) { + assert_eq!(x, 1); + } +} diff --git a/crates/sbsii/src/lib.rs b/crates/sbsii/src/lib.rs new file mode 100644 index 00000000..304b9354 --- /dev/null +++ b/crates/sbsii/src/lib.rs @@ -0,0 +1,19 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +// sbsii, small boxed slice into iter + +pub mod borrowed; +pub mod owned; +pub mod stack; diff --git a/crates/sbsii/src/owned.rs b/crates/sbsii/src/owned.rs new file mode 100644 index 00000000..477d0215 --- /dev/null +++ b/crates/sbsii/src/owned.rs @@ -0,0 +1,156 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::stack::StackIntoIter; +use std::ptr::NonNull; + +pub struct HeapIntoIter { + pointer: NonNull, + off: u16, + len: u16, +} + +impl HeapIntoIter { + #[cold] + pub(crate) fn from_slice(slice: &[T]) -> Self { + assert!(slice.len() <= 65535_usize); + #[allow(unsafe_code)] + unsafe { + let c = NonNull::new_unchecked(Box::<[T]>::into_raw(Box::from(slice)).cast::()); + Self { + pointer: c, + off: 0, + len: slice.len() as u16, + } + } + } +} + +impl Iterator for HeapIntoIter { + type Item = T; + + #[inline(always)] + fn next(&mut self) -> Option { + #[allow(unsafe_code)] + unsafe { + if self.off < self.len { + let r = self.pointer.as_ptr().add(self.off as _).read(); + self.off += 1; + Some(r) + } else { + None + } + } + } +} + +impl Clone for HeapIntoIter { + #[inline(always)] + fn clone(&self) -> Self { + #[allow(unsafe_code)] + unsafe { + let p = std::slice::from_raw_parts_mut(self.pointer.as_ptr(), self.len as _); + let c = NonNull::new_unchecked(Box::<[T]>::into_raw(Box::from(p)).cast::()); + Self { + pointer: c, + off: self.off, + len: self.len, + } + } + } +} + +impl Drop for HeapIntoIter { + #[inline(always)] + fn drop(&mut self) { + #[allow(unsafe_code)] + unsafe { + let p = std::slice::from_raw_parts_mut(self.pointer.as_ptr(), self.len as _); + let _ = Box::<[T]>::from_raw(p); + } + } +} + +#[derive(Clone)] +pub enum IntoIter { + Stack(StackIntoIter), + Heap(HeapIntoIter), +} + +impl IntoIter { + #[inline(always)] + pub fn from_slice(slice: &[T]) -> Self { + assert!(slice.len() <= 65535); + if slice.len() <= N && N <= 65535 { + IntoIter::Stack(StackIntoIter::from_slice(slice)) + } else { + IntoIter::Heap(HeapIntoIter::from_slice(slice)) + } + } +} + +impl Iterator for IntoIter { + type Item = T; + + #[inline(always)] + fn next(&mut self) -> Option { + match self { + IntoIter::Stack(iter) => iter.next(), + IntoIter::Heap(iter) => iter.next(), + } + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + match self { + IntoIter::Stack(iter) => iter.size_hint(), + IntoIter::Heap(iter) => iter.size_hint(), + } + } +} + +impl ExactSizeIterator for IntoIter {} + +#[cfg(target_pointer_width = "64")] +const _: () = { + assert!(size_of::>() == 16); +}; + +#[test] +fn tests() { + for x in IntoIter::::from_slice(&[1; 0]).clone() { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 1]).clone() { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 2]).clone() { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 3]).clone() { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 4]).clone() { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 5]).clone() { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 6]).clone() { + assert_eq!(x, 1); + } + for x in IntoIter::::from_slice(&[1; 7]).clone() { + assert_eq!(x, 1); + } +} diff --git a/crates/sbsii/src/stack.rs b/crates/sbsii/src/stack.rs new file mode 100644 index 00000000..13218694 --- /dev/null +++ b/crates/sbsii/src/stack.rs @@ -0,0 +1,61 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use std::mem::MaybeUninit; + +#[derive(Clone)] +pub struct StackIntoIter { + start: u16, + end: u16, + buffer: [MaybeUninit; N], +} + +impl StackIntoIter { + #[inline(always)] + pub(crate) fn from_slice(slice: &[T]) -> Self { + assert!(slice.len() <= N && N <= 65535); + Self { + start: 0, + end: slice.len() as u16, + buffer: { + let mut buffer = [const { MaybeUninit::uninit() }; N]; + for i in 0..slice.len() { + buffer[i] = MaybeUninit::new(slice[i]); + } + buffer + }, + } + } +} + +impl Iterator for StackIntoIter { + type Item = T; + + #[inline(always)] + fn next(&mut self) -> Option { + if self.start == self.end { + return None; + } + #[allow(unsafe_code)] + let result = unsafe { self.buffer[self.start as usize].assume_init() }; + self.start += 1; + Some(result) + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + let size = (self.end - self.start) as usize; + (size, Some(size)) + } +} diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index 1ec50219..29cf2689 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -20,7 +20,7 @@ use crate::vectors::{self}; use crate::{Opaque, Page, tape}; use algo::accessor::{FunctionalAccessor, LAccess}; use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily}; -use algo::{Bump, RelationRead, RelationWrite}; +use algo::{Bump, OwnedIter, RelationRead, RelationWrite}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -28,7 +28,7 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -type Extra<'b> = &'b mut (u32, u16, f32, &'b mut [u32]); +type Extra1<'b> = &'b mut (u32, f32, u16, OwnedIter); pub fn insert_vector( index: &R, @@ -86,12 +86,12 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( AlwaysEqual(first), ) } else { - let prefetch = bump.alloc_slice(meta_tuple.centroid_prefetch()); + let prefetch = OwnedIter::from_slice(meta_tuple.centroid_prefetch()); let head = meta_tuple.centroid_head(); let norm = meta_tuple.centroid_norm(); let first = meta_tuple.first(); let distance = vectors::read_for_h1_tuple::( - prefetch.iter().map(|&id| index.read(id)), + prefetch.map(|id| index.read(id)), head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); @@ -102,7 +102,7 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( let lut = (O::Vector::block_preprocess(vector),); let mut step = |state: State| { - let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); + let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); { let (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) = state; tape::read_h1_tape::( @@ -112,7 +112,12 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( Reverse(lowerbound), - AlwaysEqual(bump.alloc((first, head, norm, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc(( + first, + norm, + head, + OwnedIter::from_slice(prefetch), + ))), )); }, ); @@ -120,7 +125,7 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); let mut cache = BinaryHeap::<(_, _, _)>::new(); { - while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, norm, ..))), prefetch)) = + while let Some(((Reverse(_), AlwaysEqual(&mut (first, norm, head, ..))), prefetch)) = heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { let distance = vectors::read_for_h1_tuple::( diff --git a/crates/vchordrq/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs index b46bbb92..4a9de383 100644 --- a/crates/vchordrq/src/prewarm.rs +++ b/crates/vchordrq/src/prewarm.rs @@ -17,15 +17,14 @@ use crate::operator::Operator; use crate::tape::{by_directory, by_next}; use crate::tuples::*; use crate::{Opaque, Page, tape, vectors}; +use algo::RelationRead; use algo::accessor::FunctionalAccessor; use algo::prefetcher::PrefetcherSequenceFamily; -use algo::{Bump, RelationRead}; use std::fmt::Write; -pub fn prewarm<'r, 'b: 'r, R: RelationRead, O: Operator>( +pub fn prewarm<'r, R: RelationRead, O: Operator>( index: &'r R, height: i32, - bump: &'b impl Bump, mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, ) -> String where @@ -46,10 +45,10 @@ where type State = Vec; let mut state: State = { let mut results = Vec::new(); - let prefetch = bump.alloc_slice(meta_tuple.centroid_prefetch()); + let prefetch = algo::OwnedIter::from_slice(meta_tuple.centroid_prefetch()); let head = meta_tuple.centroid_head(); let first = meta_tuple.first(); - vectors::read_for_h1_tuple::(prefetch.iter().map(|&id| index.read(id)), head, ()); + vectors::read_for_h1_tuple::(prefetch.map(|id| index.read(id)), head, ()); results.push(first); writeln!(message, "------------------------").unwrap(); writeln!(message, "number of nodes: {}", results.len()).unwrap(); diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs index 4247e622..8e4688af 100644 --- a/crates/vchordrq/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -27,7 +27,7 @@ use std::marker::PhantomData; use std::num::NonZero; use vector::VectorOwned; -type Extra<'b> = &'b mut (NonZero, u16, &'b mut [u32]); +type Extra0<'b> = &'b mut (NonZero, u16, algo::OwnedIter); type Result = (Reverse, AlwaysEqual>); @@ -53,7 +53,7 @@ pub struct Reranker { impl<'r, 'b, T, F, P> Iterator for Reranker where F: FnMut(NonZero, P::Guards, u16) -> Option, - P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, { type Item = (Distance, NonZero); @@ -82,7 +82,7 @@ pub fn rerank_index< 'b, O: Operator, T, - P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, >( vector: O::Vector, prefetcher: P, @@ -110,7 +110,7 @@ pub fn rerank_heap< 'b, O: Operator, T, - P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, >( vector: O::Vector, prefetcher: P, diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index 015b6aa6..c898af0f 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -20,7 +20,7 @@ use crate::tuples::*; use crate::{Opaque, Page, tape, vectors}; use algo::accessor::LAccess; use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; -use algo::{Bump, RelationRead}; +use algo::{Bump, OwnedIter, RelationRead}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -28,7 +28,8 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -type Extra<'b> = &'b mut (NonZero, u16, &'b mut [u32]); +type Extra1<'b> = &'b mut (u32, f32, u16, OwnedIter); +type Extra0<'b> = &'b mut (NonZero, u16, OwnedIter); pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( index: &'r R, @@ -38,7 +39,10 @@ pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( bump: &'b impl Bump, mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, -) -> Vec<((Reverse, AlwaysEqual<()>), AlwaysEqual>)> +) -> Vec<( + (Reverse, AlwaysEqual<()>), + AlwaysEqual>, +)> where R::Page: Page, { @@ -67,12 +71,12 @@ where AlwaysEqual(first), )] } else { - let prefetch = bump.alloc_slice(meta_tuple.centroid_prefetch()); + let prefetch = OwnedIter::from_slice(meta_tuple.centroid_prefetch()); let head = meta_tuple.centroid_head(); let norm = meta_tuple.centroid_norm(); let first = meta_tuple.first(); let distance = vectors::read_for_h1_tuple::( - prefetch.iter().map(|&id| index.read(id)), + prefetch.map(|id| index.read(id)), head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); @@ -83,7 +87,7 @@ where let lut = O::Vector::preprocess(vector); let mut step = |state: State| { - let mut results = LinkedVec::new(); + let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { tape::read_h1_tape::( by_next(index, first), @@ -92,7 +96,12 @@ where let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( Reverse(lowerbound), - AlwaysEqual(bump.alloc((first, head, norm, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc(( + first, + norm, + head, + OwnedIter::from_slice(prefetch), + ))), )); }, ); @@ -100,7 +109,7 @@ where let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); let mut cache = BinaryHeap::<(_, _, _)>::new(); std::iter::from_fn(move || { - while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, norm, ..))), prefetch)) = + while let Some(((Reverse(_), AlwaysEqual(&mut (first, norm, head, ..))), prefetch)) = heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { let distance = vectors::read_for_h1_tuple::( @@ -118,7 +127,7 @@ where state = step(state).take(probes[i as usize - 1] as _).collect(); } - let mut results = LinkedVec::new(); + let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); @@ -127,7 +136,7 @@ where let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( (Reverse(lowerbound), AlwaysEqual(())), - AlwaysEqual(bump.alloc((payload, head, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc((payload, head, OwnedIter::from_slice(prefetch)))), )); }); if prefetch_h0_tuples.is_not_plain() { @@ -166,7 +175,7 @@ pub fn maxsim_search<'r, 'b: 'r, R: RelationRead, O: Operator>( ) -> ( Vec<( (Reverse, AlwaysEqual), - AlwaysEqual>, + AlwaysEqual>, )>, Distance, ) @@ -198,12 +207,12 @@ where AlwaysEqual(first), )] } else { - let prefetch = bump.alloc_slice(meta_tuple.centroid_prefetch()); + let prefetch = OwnedIter::from_slice(meta_tuple.centroid_prefetch()); let head = meta_tuple.centroid_head(); let norm = meta_tuple.centroid_norm(); let first = meta_tuple.first(); let distance = vectors::read_for_h1_tuple::( - prefetch.iter().map(|&id| index.read(id)), + prefetch.map(|id| index.read(id)), head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); @@ -214,7 +223,7 @@ where let lut = O::Vector::preprocess(vector); let mut step = |state: State| { - let mut results = LinkedVec::new(); + let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { tape::read_h1_tape::( by_next(index, first), @@ -223,7 +232,12 @@ where let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( Reverse(lowerbound), - AlwaysEqual(bump.alloc((first, head, norm, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc(( + first, + norm, + head, + OwnedIter::from_slice(prefetch), + ))), )); }, ); @@ -231,7 +245,7 @@ where let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); let mut cache = BinaryHeap::<(_, _, _)>::new(); std::iter::from_fn(move || { - while let Some(((Reverse(_), AlwaysEqual(&mut (first, head, norm, ..))), prefetch)) = + while let Some(((Reverse(_), AlwaysEqual(&mut (first, norm, head, ..))), prefetch)) = heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { let distance = vectors::read_for_h1_tuple::( @@ -251,7 +265,7 @@ where state = it.take(probes[i as usize - 1] as _).collect(); } - let mut results = LinkedVec::new(); + let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); @@ -261,7 +275,7 @@ where let rough = Distance::from_f32(rough); results.push(( (Reverse(lowerbound), AlwaysEqual(rough)), - AlwaysEqual(bump.alloc((payload, head, bump.alloc_slice(prefetch)))), + AlwaysEqual(bump.alloc((payload, head, OwnedIter::from_slice(prefetch)))), )); }); if prefetch_h0_tuples.is_not_plain() { diff --git a/src/index/storage.rs b/src/index/storage.rs index 2409c96a..6047981a 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -423,7 +423,7 @@ where && let Some(iter) = self.iter.as_mut() && let Some(e) = iter.next() { - for id in e.fetch().iter().copied() { + for id in e.fetch() { self.tail.push_back(id); } self.window.push_back(e); @@ -436,7 +436,7 @@ where && let Some(iter) = self.iter.as_mut() && let Some(e) = iter.next() { - for id in e.fetch().iter().copied() { + for id in e.fetch() { self.tail.push_back(id); } self.window.push_back(e); @@ -449,7 +449,7 @@ where && let Some(iter) = self.iter.as_mut() && let Some(e) = iter.next() { - for id in e.fetch().iter().copied() { + for id in e.fetch() { self.tail.push_back(id); } self.window.push_back(e); @@ -532,7 +532,7 @@ where { type Relation = PostgresRelation; - type Guards = PostgresReadStreamGuards>; + type Guards = PostgresReadStreamGuards; type Item = I::Item; @@ -555,7 +555,7 @@ where #[cfg(any(feature = "pg17", feature = "pg18"))] fn next(&mut self) -> Option<(I::Item, Self::Guards)> { if let Some(e) = unsafe { self.cache.as_mut().pop_item() } { - let list = self.read(e.fetch().into_iter()); + let list = self.read(e.fetch()); Some((e, list)) } else { None @@ -568,7 +568,7 @@ where predicate: P, ) -> Option<(I::Item, Self::Guards)> { if let Some(e) = unsafe { self.cache.as_mut().pop_item_if(predicate) } { - let list = self.read(e.fetch().into_iter()); + let list = self.read(e.fetch()); Some((e, list)) } else { None diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs index 31f0ade2..5114555b 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/algo.rs @@ -117,7 +117,7 @@ where R: RelationRead + RelationWrite + RelationReadStream, R::Page: Page, { - let bump = bumpalo::Bump::new(); + let bump = bumpalo::Bump::with_capacity(2 << 20); let make_vertex_plain_prefetcher = MakePlainPrefetcher { index }; let make_vector_plain_prefetcher = MakePlainPrefetcher { index }; match (vector, opfamily.distance_kind()) { diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index ea71d05c..3eb31a04 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -315,7 +315,7 @@ pub unsafe extern "C-unwind" fn ambeginscan( let scanner: Scanner = Scanner { hack: None, scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), - bump: Box::new(bumpalo::Bump::new()), + bump: Box::new(bumpalo::Bump::with_capacity(2 << 20)), }; unsafe { (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); diff --git a/src/index/vchordrq/algo.rs b/src/index/vchordrq/algo.rs index a8e42c27..2ac6c538 100644 --- a/src/index/vchordrq/algo.rs +++ b/src/index/vchordrq/algo.rs @@ -31,33 +31,20 @@ where R: RelationRead, R::Page: Page, { - let bump = bumpalo::Bump::new(); let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; match (opfamily.vector_kind(), opfamily.distance_kind()) { - (VectorKind::Vecf32, DistanceKind::L2S) => vchordrq::prewarm::<_, Op, L2S>>( - index, - height, - &bump, - make_h0_plain_prefetcher, - ), - (VectorKind::Vecf32, DistanceKind::Dot) => vchordrq::prewarm::<_, Op, Dot>>( - index, - height, - &bump, - make_h0_plain_prefetcher, - ), - (VectorKind::Vecf16, DistanceKind::L2S) => vchordrq::prewarm::<_, Op, L2S>>( - index, - height, - &bump, - make_h0_plain_prefetcher, - ), - (VectorKind::Vecf16, DistanceKind::Dot) => vchordrq::prewarm::<_, Op, Dot>>( - index, - height, - &bump, - make_h0_plain_prefetcher, - ), + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordrq::prewarm::<_, Op, L2S>>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordrq::prewarm::<_, Op, Dot>>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordrq::prewarm::<_, Op, L2S>>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordrq::prewarm::<_, Op, Dot>>(index, height, make_h0_plain_prefetcher) + } } } @@ -171,7 +158,7 @@ pub fn insert( R: RelationRead + RelationWrite, R::Page: Page, { - let bump = bumpalo::Bump::new(); + let bump = bumpalo::Bump::with_capacity(2 << 20); let make_h1_plain_prefetcher = MakeH1PlainPrefetcherForInsertion { index }; match (vector, opfamily.distance_kind()) { (OwnedVector::Vecf32(vector), DistanceKind::L2S) => { diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 58e57423..5708f167 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -395,7 +395,7 @@ pub unsafe extern "C-unwind" fn ambeginscan( let scanner: Scanner = Scanner { hack: None, scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), - bump: Box::new(bumpalo::Bump::new()), + bump: Box::new(bumpalo::Bump::with_capacity(2 << 20)), }; unsafe { (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index 1f4e6296..f6b63d5c 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -665,7 +665,7 @@ impl SearchBuilder for DefaultBuilder { #[inline(always)] pub fn id_0(f: F) -> F where - F: for<'a> FnMut(&(A, AlwaysEqual<&'a mut (B, C, &'a mut D)>)) -> R, + F: for<'a> FnMut(&(A, AlwaysEqual<&'a mut (B, C, D)>)) -> R, { f } diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index dd3c8426..787bbb1d 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -486,7 +486,7 @@ impl std::iter::FusedIterator for IntoIterSorted {} #[inline(always)] pub fn id_0(f: F) -> F where - F: for<'a> FnMut(&(A, AlwaysEqual<&'a mut (B, C, &'a mut D)>)) -> R, + F: for<'a> FnMut(&(A, AlwaysEqual<&'a mut (B, C, D)>)) -> R, { f } From 5e8b04d352b2484a0950eb7c5fc86403b3179984 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 8 Aug 2025 09:29:22 +0800 Subject: [PATCH 190/324] chore: update Rust to 1.89 (#307) Signed-off-by: usamoi --- .github/workflows/check.yml | 50 ++++++----------------------- .github/workflows/release.yml | 6 ++-- crates/algo/src/lib.rs | 10 +++--- crates/sbsii/src/borrowed.rs | 59 ++++++++++++++++++++++++++++++----- crates/sbsii/src/owned.rs | 50 ++++++++++++++++++++++++----- crates/sbsii/src/stack.rs | 28 +++++++++++++---- crates/vchordrq/src/insert.rs | 9 +++--- crates/vchordrq/src/rerank.rs | 2 +- crates/vchordrq/src/search.rs | 28 +++++++++++------ rust-toolchain.toml | 2 -- 10 files changed, 160 insertions(+), 84 deletions(-) delete mode 100644 rust-toolchain.toml diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 120d9a7d..2525bbf1 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -16,8 +16,7 @@ jobs: steps: - name: Set up Environment run: | - rustup default 1.89-beta - rustup component add rustfmt clippy + rustup update curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo @@ -116,8 +115,7 @@ jobs: steps: - name: Set up Environment run: | - rustup default 1.89-beta - rustup component add rustfmt clippy + rustup update sudo apt-get update @@ -179,12 +177,12 @@ jobs: SCCACHE_GHA_ENABLED: true CARGO_TERM_COLOR: always RUST_BACKTRACE: 1 + RUSTFLAGS: -Dwarnings steps: - name: Set up Environment run: | - rustup default 1.89-beta - rustup component add rustfmt clippy + rustup update sudo apt-get update @@ -216,15 +214,6 @@ jobs: - name: Checkout uses: actions/checkout@v4 - - name: Patch - run: | - mkdir ./.cargo && touch ./.cargo/config.toml - echo "unstable.host-config = true" >> ./.cargo/config.toml - echo "unstable.target-applies-to-host = true" >> ./.cargo/config.toml - echo "host.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml - echo "build.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml - echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV - - name: Clippy run: cargo clippy --target ${{ matrix.arch }}-unknown-linux-gnu -p vchord --features pg${{ matrix.version }} --no-deps @@ -264,21 +253,21 @@ jobs: strategy: matrix: version: ["13", "14", "15", "16", "17"] - arch: ["aarch64", "x86_64"] + arch: ["aarch64"] - runs-on: ${{ matrix.arch == 'aarch64' && 'macos-15' || 'macos-13' }} + runs-on: "macos-15" env: RUSTC_WRAPPER: sccache SCCACHE_GHA_ENABLED: true CARGO_TERM_COLOR: always RUST_BACKTRACE: 1 + RUSTFLAGS: -Dwarnings steps: - name: Set up Environment run: | - rustup default 1.89-beta - rustup component add rustfmt clippy + rustup update HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install llvm@18 echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV @@ -305,15 +294,6 @@ jobs: - name: Checkout uses: actions/checkout@v4 - - name: Patch - run: | - mkdir ./.cargo && touch ./.cargo/config.toml - echo "unstable.host-config = true" >> ./.cargo/config.toml - echo "unstable.target-applies-to-host = true" >> ./.cargo/config.toml - echo "host.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml - echo "build.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml - echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV - - name: Clippy run: | cargo clippy -p vchord --target ${{ matrix.arch }}-apple-darwin --features pg${{ matrix.version }} --no-deps @@ -365,14 +345,14 @@ jobs: SCCACHE_GHA_ENABLED: true CARGO_TERM_COLOR: always RUST_BACKTRACE: 1 + RUSTFLAGS: -Dwarnings steps: - name: Set up Environment run: | 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } - rustup default 1.89-beta - rustup component add rustfmt clippy + rustup update & 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise\Common7\Tools\Launch-VsDevShell.ps1' -HostArch amd64 -Arch amd64 @@ -418,16 +398,6 @@ jobs: - name: Checkout uses: actions/checkout@v4 - - name: Patch - shell: bash - run: | - mkdir ./.cargo && touch ./.cargo/config.toml - echo "unstable.host-config = true" >> ./.cargo/config.toml - echo "unstable.target-applies-to-host = true" >> ./.cargo/config.toml - echo "host.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml - echo "build.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml - echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV - - name: Clippy run: | cargo clippy --target ${{ matrix.arch }}-pc-windows-msvc -p vchord --features pg${{ matrix.version }} --no-deps diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 6846127b..0b3bd44b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -46,12 +46,12 @@ jobs: env: CARGO_TERM_COLOR: always RUST_BACKTRACE: 1 + RUSTFLAGS: "-Dwarnings" steps: - name: Set up Environment run: | - rustup default 1.89-beta - rustup component add rustfmt clippy + rustup update sudo apt-get update @@ -87,7 +87,7 @@ jobs: ARCH: ${{ matrix.arch }} PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} run: | - cargo run -p make -- build -o ./build/raw + cargo run -p make -- build -o ./build/raw --profile release (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) diff --git a/crates/algo/src/lib.rs b/crates/algo/src/lib.rs index a80c187c..6ce00364 100644 --- a/crates/algo/src/lib.rs +++ b/crates/algo/src/lib.rs @@ -160,7 +160,7 @@ impl Bump for bumpalo::Bump { } pub type BorrowedIter<'b> = sbsii::borrowed::IntoIter<'b, u32, 1>; -pub type OwnedIter = sbsii::owned::IntoIter; +pub type OwnedIter = sbsii::owned::IntoIter; pub trait Fetch { #[must_use] @@ -174,19 +174,19 @@ impl Fetch for u32 { } } -impl Fetch for (T, AlwaysEqual<&mut (A, B, OwnedIter)>) { +impl Fetch for (T, AlwaysEqual<&mut (A, B, BorrowedIter<'_>)>) { #[inline(always)] fn fetch(&self) -> OwnedIter { let (_, AlwaysEqual((.., list))) = self; - list.clone() + OwnedIter::from_slice(list.as_slice()) } } -impl Fetch for (T, AlwaysEqual<&mut (A, B, C, OwnedIter)>) { +impl Fetch for (T, AlwaysEqual<&mut (A, B, C, BorrowedIter<'_>)>) { #[inline(always)] fn fetch(&self) -> OwnedIter { let (_, AlwaysEqual((.., list))) = self; - list.clone() + OwnedIter::from_slice(list.as_slice()) } } diff --git a/crates/sbsii/src/borrowed.rs b/crates/sbsii/src/borrowed.rs index af7eb462..4cffdc6f 100644 --- a/crates/sbsii/src/borrowed.rs +++ b/crates/sbsii/src/borrowed.rs @@ -13,18 +13,38 @@ // Copyright (c) 2025 TensorChord Inc. use crate::stack::StackIntoIter; +use std::marker::PhantomData; +use std::ptr::NonNull; #[derive(Clone)] -pub struct HeapIntoIter<'a, T> { - inner: std::iter::Copied>, +pub struct HeapIntoIter<'a, T: Copy> { + pointer: NonNull, + off: u16, + len: u16, + _phantom: PhantomData<&'a ()>, } impl<'a, T: Copy> HeapIntoIter<'a, T> { #[cold] - pub(crate) fn from_slice(slice: &[T], alloc: impl Fn(&[T]) -> &'a [T]) -> Self { + pub(crate) fn from_slice(slice: &'a [T]) -> Self { assert!(slice.len() <= 65535_usize); + let c = NonNull::from_ref(slice).cast::(); Self { - inner: alloc(slice).iter().copied(), + pointer: c, + off: 0, + len: slice.len() as u16, + _phantom: PhantomData, + } + } + + #[inline(always)] + pub(crate) fn as_slice(&self) -> &[T] { + #[allow(unsafe_code)] + unsafe { + std::slice::from_raw_parts( + self.pointer.as_ptr().add(self.off as _), + (self.len - self.off) as _, + ) } } } @@ -34,24 +54,49 @@ impl<'a, T: Copy> Iterator for HeapIntoIter<'a, T> { #[inline(always)] fn next(&mut self) -> Option { - self.inner.next() + #[allow(unsafe_code)] + unsafe { + if self.off < self.len { + let r = self.pointer.as_ptr().add(self.off as _).read(); + self.off += 1; + Some(r) + } else { + None + } + } + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + let size = (self.len - self.off) as usize; + (size, Some(size)) } } +impl<'a, T: Copy> ExactSizeIterator for HeapIntoIter<'a, T> {} + #[derive(Clone)] pub enum IntoIter<'a, T: Copy, const N: usize> { Stack(StackIntoIter), Heap(HeapIntoIter<'a, T>), } -impl<'a, T: Copy, const N: usize> IntoIter<'a, T, N> { +impl<'a, T: Copy + 'a, const N: usize> IntoIter<'a, T, N> { #[inline(always)] pub fn from_slice(slice: &[T], alloc: impl Fn(&[T]) -> &'a [T]) -> Self { assert!(slice.len() <= 65535); if slice.len() <= N && N <= 65535 { IntoIter::Stack(StackIntoIter::from_slice(slice)) } else { - IntoIter::Heap(HeapIntoIter::from_slice(slice, alloc)) + IntoIter::Heap(HeapIntoIter::from_slice(alloc(slice))) + } + } + + #[inline(always)] + pub fn as_slice(&self) -> &[T] { + match self { + IntoIter::Stack(x) => x.as_slice(), + IntoIter::Heap(x) => x.as_slice(), } } } diff --git a/crates/sbsii/src/owned.rs b/crates/sbsii/src/owned.rs index 477d0215..e2a163cf 100644 --- a/crates/sbsii/src/owned.rs +++ b/crates/sbsii/src/owned.rs @@ -25,14 +25,22 @@ impl HeapIntoIter { #[cold] pub(crate) fn from_slice(slice: &[T]) -> Self { assert!(slice.len() <= 65535_usize); + let c = box_into_non_null::<[T]>(Box::from(slice)).cast::(); + Self { + pointer: c, + off: 0, + len: slice.len() as u16, + } + } + + #[inline(always)] + pub(crate) fn as_slice(&self) -> &[T] { #[allow(unsafe_code)] unsafe { - let c = NonNull::new_unchecked(Box::<[T]>::into_raw(Box::from(slice)).cast::()); - Self { - pointer: c, - off: 0, - len: slice.len() as u16, - } + std::slice::from_raw_parts( + self.pointer.as_ptr().add(self.off as _), + (self.len - self.off) as _, + ) } } } @@ -53,8 +61,16 @@ impl Iterator for HeapIntoIter { } } } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + let size = (self.len - self.off) as usize; + (size, Some(size)) + } } +impl ExactSizeIterator for HeapIntoIter {} + impl Clone for HeapIntoIter { #[inline(always)] fn clone(&self) -> Self { @@ -98,6 +114,14 @@ impl IntoIter { IntoIter::Heap(HeapIntoIter::from_slice(slice)) } } + + #[inline(always)] + pub fn as_slice(&self) -> &[T] { + match self { + IntoIter::Stack(x) => x.as_slice(), + IntoIter::Heap(x) => x.as_slice(), + } + } } impl Iterator for IntoIter { @@ -124,9 +148,21 @@ impl ExactSizeIterator for IntoIter {} #[cfg(target_pointer_width = "64")] const _: () = { - assert!(size_of::>() == 16); + assert!(size_of::>() == 24); }; +// Emulate unstable library feature `box_vec_non_null`. +// See https://github.com/rust-lang/rust/issues/130364. + +#[allow(dead_code)] +#[must_use] +fn box_into_non_null(b: Box) -> NonNull { + #[allow(unsafe_code)] + unsafe { + NonNull::new_unchecked(Box::into_raw(b)) + } +} + #[test] fn tests() { for x in IntoIter::::from_slice(&[1; 0]).clone() { diff --git a/crates/sbsii/src/stack.rs b/crates/sbsii/src/stack.rs index 13218694..2351db07 100644 --- a/crates/sbsii/src/stack.rs +++ b/crates/sbsii/src/stack.rs @@ -37,6 +37,17 @@ impl StackIntoIter { }, } } + + #[inline(always)] + pub(crate) fn as_slice(&self) -> &[T] { + #[allow(unsafe_code)] + unsafe { + std::slice::from_raw_parts( + self.buffer.as_ptr().add(self.start as _).cast(), + (self.end - self.start) as _, + ) + } + } } impl Iterator for StackIntoIter { @@ -44,13 +55,16 @@ impl Iterator for StackIntoIter { #[inline(always)] fn next(&mut self) -> Option { - if self.start == self.end { - return None; - } #[allow(unsafe_code)] - let result = unsafe { self.buffer[self.start as usize].assume_init() }; - self.start += 1; - Some(result) + unsafe { + if self.start < self.end { + let r = self.buffer[self.start as usize].assume_init(); + self.start += 1; + Some(r) + } else { + None + } + } } #[inline(always)] @@ -59,3 +73,5 @@ impl Iterator for StackIntoIter { (size, Some(size)) } } + +impl ExactSizeIterator for StackIntoIter {} diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index 29cf2689..c615ddf4 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -20,7 +20,7 @@ use crate::vectors::{self}; use crate::{Opaque, Page, tape}; use algo::accessor::{FunctionalAccessor, LAccess}; use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily}; -use algo::{Bump, OwnedIter, RelationRead, RelationWrite}; +use algo::{BorrowedIter, Bump, RelationRead, RelationWrite}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -28,7 +28,7 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -type Extra1<'b> = &'b mut (u32, f32, u16, OwnedIter); +type Extra1<'b> = &'b mut (u32, f32, u16, BorrowedIter<'b>); pub fn insert_vector( index: &R, @@ -86,7 +86,8 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( AlwaysEqual(first), ) } else { - let prefetch = OwnedIter::from_slice(meta_tuple.centroid_prefetch()); + let prefetch = + BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); let head = meta_tuple.centroid_head(); let norm = meta_tuple.centroid_norm(); let first = meta_tuple.first(); @@ -116,7 +117,7 @@ pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( first, norm, head, - OwnedIter::from_slice(prefetch), + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), ))), )); }, diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs index 8e4688af..7d3cc032 100644 --- a/crates/vchordrq/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -27,7 +27,7 @@ use std::marker::PhantomData; use std::num::NonZero; use vector::VectorOwned; -type Extra0<'b> = &'b mut (NonZero, u16, algo::OwnedIter); +type Extra0<'b> = &'b mut (NonZero, u16, algo::BorrowedIter<'b>); type Result = (Reverse, AlwaysEqual>); diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index c898af0f..c3acee32 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -20,7 +20,7 @@ use crate::tuples::*; use crate::{Opaque, Page, tape, vectors}; use algo::accessor::LAccess; use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; -use algo::{Bump, OwnedIter, RelationRead}; +use algo::{BorrowedIter, Bump, RelationRead}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -28,8 +28,8 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; -type Extra1<'b> = &'b mut (u32, f32, u16, OwnedIter); -type Extra0<'b> = &'b mut (NonZero, u16, OwnedIter); +type Extra1<'b> = &'b mut (u32, f32, u16, BorrowedIter<'b>); +type Extra0<'b> = &'b mut (NonZero, u16, BorrowedIter<'b>); pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( index: &'r R, @@ -71,7 +71,8 @@ where AlwaysEqual(first), )] } else { - let prefetch = OwnedIter::from_slice(meta_tuple.centroid_prefetch()); + let prefetch = + BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); let head = meta_tuple.centroid_head(); let norm = meta_tuple.centroid_norm(); let first = meta_tuple.first(); @@ -100,7 +101,7 @@ where first, norm, head, - OwnedIter::from_slice(prefetch), + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), ))), )); }, @@ -136,7 +137,11 @@ where let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( (Reverse(lowerbound), AlwaysEqual(())), - AlwaysEqual(bump.alloc((payload, head, OwnedIter::from_slice(prefetch)))), + AlwaysEqual(bump.alloc(( + payload, + head, + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), + ))), )); }); if prefetch_h0_tuples.is_not_plain() { @@ -207,7 +212,8 @@ where AlwaysEqual(first), )] } else { - let prefetch = OwnedIter::from_slice(meta_tuple.centroid_prefetch()); + let prefetch = + BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); let head = meta_tuple.centroid_head(); let norm = meta_tuple.centroid_norm(); let first = meta_tuple.first(); @@ -236,7 +242,7 @@ where first, norm, head, - OwnedIter::from_slice(prefetch), + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), ))), )); }, @@ -275,7 +281,11 @@ where let rough = Distance::from_f32(rough); results.push(( (Reverse(lowerbound), AlwaysEqual(rough)), - AlwaysEqual(bump.alloc((payload, head, OwnedIter::from_slice(prefetch)))), + AlwaysEqual(bump.alloc(( + payload, + head, + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), + ))), )); }); if prefetch_h0_tuples.is_not_plain() { diff --git a/rust-toolchain.toml b/rust-toolchain.toml deleted file mode 100644 index 55827dfe..00000000 --- a/rust-toolchain.toml +++ /dev/null @@ -1,2 +0,0 @@ -[toolchain] -channel = "1.89-beta" From fb5d3d4c160639d46d78e15df6694ec367448a20 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Mon, 11 Aug 2025 14:18:04 +0800 Subject: [PATCH 191/324] chore: add accu_probes to recall function (#308) Signed-off-by: cutecutecat --- src/sql/finalize.sql | 20 ++++++++++++-------- tests/vchordrq/recall.slt | 9 ++++++--- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 7333f209..d6c186ff 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -151,11 +151,11 @@ STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; CREATE FUNCTION vchordrq_evaluate_query_recall( query text, exact_search boolean default false, + accu_probes TEXT default NULL, accu_epsilon real default 1.9 ) RETURNS real LANGUAGE plpgsql -STRICT AS $$ DECLARE rough tid[]; @@ -164,8 +164,10 @@ DECLARE accu_k integer; recall real; rough_probes text; - accu_probes text; BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; IF query LIKE '%@#%' AND NOT exact_search THEN RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; END IF; @@ -188,12 +190,14 @@ BEGIN IF exact_search THEN SET LOCAL vchordrq.enable_scan = off; ELSE - IF rough_probes = '' THEN - accu_probes := ''; - ELSIF position(',' in rough_probes) > 0 THEN - accu_probes := '65535,65535'; - ELSE - accu_probes := '65535'; + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; END IF; EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); diff --git a/tests/vchordrq/recall.slt b/tests/vchordrq/recall.slt index 3c1b1bb5..04584fd0 100644 --- a/tests/vchordrq/recall.slt +++ b/tests/vchordrq/recall.slt @@ -17,18 +17,21 @@ SET vchordrq.probes = '1'; statement error MaxSim operator cannot be used for estimated recall SELECT * from vchordrq_evaluate_query_recall(query=>'@#'); -statement error Error executing ANN query +statement error Error executing ANN query (.+) need 0 probes, but 1 probes provided SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); statement ok SET vchordrq.probes = ''; -statement error Error executing ANN query +statement error Error executing ANN query (.+) could not convert type SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT val FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); -statement error Error executing ANN query +statement error Error executing ANN query (.+) could not convert type SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT * FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); +statement error Error executing Ground Truth query (.+) need 0 probes, but 1 probes provided +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$, accu_probes=>'1'); + query I SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); ---- From cfbeb513b2552db5f4cf73c0828ed63b21566682 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 11 Aug 2025 17:21:48 +0800 Subject: [PATCH 192/324] refactor: unify `'b` and `'r` (#309) Signed-off-by: usamoi --- Cargo.lock | 12 +- Cargo.toml | 2 +- crates/algo/Cargo.toml | 2 +- crates/algo/src/lib.rs | 121 +++++++---- crates/algo/src/prefetcher.rs | 138 +++++++------ crates/sbsii/src/owned.rs | 192 ------------------ crates/{sbsii => small_iter}/Cargo.toml | 2 +- crates/{sbsii => small_iter}/src/borrowed.rs | 81 +++----- crates/{sbsii => small_iter}/src/lib.rs | 3 - crates/{sbsii => small_iter}/src/stack.rs | 37 ++-- crates/vchordg/src/candidates.rs | 4 +- crates/vchordg/src/insert.rs | 18 +- crates/vchordg/src/search.rs | 2 +- crates/vchordg/src/vectors.rs | 16 +- .../vchordrq/src/closure_lifetime_binder.rs | 6 +- crates/vchordrq/src/insert.rs | 6 +- crates/vchordrq/src/maintain.rs | 18 +- crates/vchordrq/src/prewarm.rs | 8 +- crates/vchordrq/src/rerank.rs | 10 +- crates/vchordrq/src/search.rs | 16 +- crates/vchordrq/src/tape.rs | 48 ++--- crates/vector/src/vect.rs | 2 +- src/datatype/binary_scalar8.rs | 2 +- src/datatype/text_scalar8.rs | 3 +- src/index/gucs.rs | 13 +- src/index/mod.rs | 1 + src/index/scanners.rs | 51 +++++ src/index/storage.rs | 45 ++-- src/index/vchordg/algo.rs | 46 ++--- src/index/vchordg/am/mod.rs | 1 + src/index/vchordg/scanners/default.rs | 39 ++-- src/index/vchordg/scanners/mod.rs | 39 +--- src/index/vchordrq/algo.rs | 74 +++---- src/index/vchordrq/am/mod.rs | 1 + src/index/vchordrq/filter.rs | 55 +++++ src/index/vchordrq/mod.rs | 1 + src/index/vchordrq/scanners/default.rs | 18 +- src/index/vchordrq/scanners/maxsim.rs | 19 +- src/index/vchordrq/scanners/mod.rs | 76 +------ 39 files changed, 539 insertions(+), 689 deletions(-) delete mode 100644 crates/sbsii/src/owned.rs rename crates/{sbsii => small_iter}/Cargo.toml (85%) rename crates/{sbsii => small_iter}/src/borrowed.rs (56%) rename crates/{sbsii => small_iter}/src/lib.rs (91%) rename crates/{sbsii => small_iter}/src/stack.rs (61%) create mode 100644 src/index/scanners.rs create mode 100644 src/index/vchordrq/filter.rs diff --git a/Cargo.lock b/Cargo.lock index df8d3d28..59eff7f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -19,8 +19,8 @@ dependencies = [ "bumpalo", "distance", "half 2.6.0", - "sbsii", "simd", + "small_iter", "vector", "zerocopy", ] @@ -1190,10 +1190,6 @@ dependencies = [ "winapi-util", ] -[[package]] -name = "sbsii" -version = "0.0.0" - [[package]] name = "seahash" version = "4.1.0" @@ -1294,6 +1290,10 @@ dependencies = [ "syn", ] +[[package]] +name = "small_iter" +version = "0.0.0" + [[package]] name = "smallvec" version = "1.15.1" @@ -1585,10 +1585,10 @@ dependencies = [ "pgrx-catalog", "rabitq", "rand", - "sbsii", "seq-macro", "serde", "simd", + "small_iter", "toml 0.9.5", "validator", "vchordg", diff --git a/Cargo.toml b/Cargo.toml index 9bcada1e..ab874f21 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,8 +27,8 @@ always_equal = { path = "./crates/always_equal" } distance = { path = "./crates/distance" } k_means = { path = "./crates/k_means" } rabitq = { path = "./crates/rabitq" } -sbsii = { path = "./crates/sbsii" } simd = { path = "./crates/simd" } +small_iter = { path = "./crates/small_iter" } vchordg = { path = "./crates/vchordg" } vchordrq = { path = "./crates/vchordrq" } vector = { path = "./crates/vector" } diff --git a/crates/algo/Cargo.toml b/crates/algo/Cargo.toml index 198acfca..166a1e88 100644 --- a/crates/algo/Cargo.toml +++ b/crates/algo/Cargo.toml @@ -7,8 +7,8 @@ publish = false [dependencies] always_equal = { path = "../always_equal" } distance = { path = "../distance" } -sbsii = { path = "../sbsii" } simd = { path = "../simd" } +small_iter = { path = "../small_iter" } vector = { path = "../vector" } bumpalo.workspace = true diff --git a/crates/algo/src/lib.rs b/crates/algo/src/lib.rs index 6ce00364..e28c423f 100644 --- a/crates/algo/src/lib.rs +++ b/crates/algo/src/lib.rs @@ -52,9 +52,9 @@ pub trait PageGuard { fn id(&self) -> u32; } -pub trait ReadStream<'r> { +pub trait ReadStream<'b> { type Relation: RelationReadTypes; - type Guards: ExactSizeIterator::ReadGuard<'r>>; + type Guards: ExactSizeIterator::ReadGuard<'b>>; type Item; type Inner: Iterator; fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; @@ -70,7 +70,7 @@ pub trait Relation { } pub trait RelationReadTypes: Relation { - type ReadGuard<'r>: PageGuard + Deref; + type ReadGuard<'b>: PageGuard + Deref; } pub trait RelationRead: RelationReadTypes { @@ -78,7 +78,7 @@ pub trait RelationRead: RelationReadTypes { } pub trait RelationWriteTypes: Relation { - type WriteGuard<'r>: PageGuard + DerefMut; + type WriteGuard<'b>: PageGuard + DerefMut; } pub trait RelationWrite: RelationWriteTypes { @@ -117,15 +117,15 @@ impl Hints { } pub trait RelationReadStreamTypes: RelationReadTypes { - type ReadStream<'r, I: Iterator>: ReadStream<'r, Item = I::Item, Relation = Self> + type ReadStream<'b, I: Iterator>: ReadStream<'b, Item = I::Item, Relation = Self> where - I::Item: Fetch; + I::Item: Fetch<'b>; } pub trait RelationReadStream: RelationReadStreamTypes { - fn read_stream<'r, I: Iterator>(&'r self, iter: I, hints: Hints) -> Self::ReadStream<'r, I> + fn read_stream<'b, I: Iterator>(&'b self, iter: I, hints: Hints) -> Self::ReadStream<'b, I> where - I::Item: Fetch; + I::Item: Fetch<'b>; } #[derive(Debug, Clone, Copy)] @@ -136,15 +136,16 @@ pub enum RerankMethod { pub trait Bump: 'static { #[allow(clippy::mut_from_ref)] - fn alloc(&self, value: T) -> &mut T; + fn alloc(&self, value: T) -> &mut T; #[allow(clippy::mut_from_ref)] fn alloc_slice(&self, slice: &[T]) -> &mut [T]; - fn reset(&mut self); + #[allow(clippy::mut_from_ref)] + fn alloc_any(&self, value: T) -> &mut T; } impl Bump for bumpalo::Bump { #[inline] - fn alloc(&self, value: T) -> &mut T { + fn alloc(&self, value: T) -> &mut T { self.alloc(value) } @@ -154,62 +155,68 @@ impl Bump for bumpalo::Bump { } #[inline] - fn reset(&mut self) { - self.reset(); + fn alloc_any(&self, value: T) -> &mut T { + self.alloc(value) } } -pub type BorrowedIter<'b> = sbsii::borrowed::IntoIter<'b, u32, 1>; -pub type OwnedIter = sbsii::owned::IntoIter; +pub type BorrowedIter<'b> = small_iter::borrowed::Iter<'b, u32, 1>; -pub trait Fetch { +pub trait Fetch<'b> { + type Iter: ExactSizeIterator + 'b; #[must_use] - fn fetch(&self) -> OwnedIter; + fn fetch(&self) -> Self::Iter; } -impl Fetch for u32 { +impl Fetch<'_> for u32 { + type Iter = std::iter::Once; #[inline(always)] - fn fetch(&self) -> OwnedIter { - OwnedIter::from_slice(std::slice::from_ref(self)) + fn fetch(&self) -> std::iter::Once { + std::iter::once(*self) } } -impl Fetch for (T, AlwaysEqual<&mut (A, B, BorrowedIter<'_>)>) { +impl<'b, T, A, B> Fetch<'b> for (T, AlwaysEqual<&mut (A, B, BorrowedIter<'b>)>) { + type Iter = BorrowedIter<'b>; #[inline(always)] - fn fetch(&self) -> OwnedIter { + fn fetch(&self) -> BorrowedIter<'b> { let (_, AlwaysEqual((.., list))) = self; - OwnedIter::from_slice(list.as_slice()) + *list } } -impl Fetch for (T, AlwaysEqual<&mut (A, B, C, BorrowedIter<'_>)>) { +impl<'b, T, A, B, C> Fetch<'b> for (T, AlwaysEqual<&mut (A, B, C, BorrowedIter<'b>)>) { + type Iter = BorrowedIter<'b>; #[inline(always)] - fn fetch(&self) -> OwnedIter { + fn fetch(&self) -> BorrowedIter<'b> { let (_, AlwaysEqual((.., list))) = self; - OwnedIter::from_slice(list.as_slice()) + *list } } -impl Fetch for (T, AlwaysEqual<(u32, u16)>) { +impl Fetch<'_> for (T, AlwaysEqual<(u32, u16)>) { + type Iter = std::iter::Once; #[inline(always)] - fn fetch(&self) -> OwnedIter { + fn fetch(&self) -> std::iter::Once { let (_, AlwaysEqual((x, _))) = self; - OwnedIter::from_slice(std::slice::from_ref(x)) + std::iter::once(*x) } } -impl Fetch for (u32, u16) { +impl Fetch<'_> for (u32, u16) { + type Iter = std::iter::Once; #[inline(always)] - fn fetch(&self) -> OwnedIter { - OwnedIter::from_slice(std::slice::from_ref(&self.0)) + fn fetch(&self) -> std::iter::Once { + std::iter::once(self.0) } } -impl Fetch for (T, AlwaysEqual<((u32, u16), (u32, u16))>) { +impl Fetch<'_> for (T, AlwaysEqual<((u32, u16), (u32, u16))>) { + type Iter = std::iter::Once; #[inline(always)] - fn fetch(&self) -> OwnedIter { + fn fetch(&self) -> std::iter::Once { let (_, AlwaysEqual(((x, _), _))) = self; - OwnedIter::from_slice(std::slice::from_ref(x)) + std::iter::once(*x) } } @@ -217,19 +224,47 @@ pub trait Fetch1 { fn fetch_1(&self) -> u32; } -impl Fetch for (T, AlwaysEqual<&mut [F]>) { +#[repr(transparent)] +pub struct Fetch1Iter { + iter: I, +} + +impl Iterator for Fetch1Iter +where + I::Item: Fetch1, +{ + type Item = u32; + + #[inline(always)] + fn next(&mut self) -> Option { + self.iter.next().map(|x| x.fetch_1()) + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + self.iter.size_hint() + } +} + +impl ExactSizeIterator for Fetch1Iter where I::Item: Fetch1 {} + +impl<'b, T, F: Fetch1 + Copy> Fetch<'b> for (T, AlwaysEqual<&'b [F]>) { + type Iter = Fetch1Iter>>; #[inline(always)] - fn fetch(&self) -> OwnedIter { - let vec = self.1.0.iter().map(|x| x.fetch_1()).collect::>(); - OwnedIter::from_slice(vec.as_slice()) + fn fetch(&self) -> Self::Iter { + Fetch1Iter { + iter: self.1.0.iter().copied(), + } } } -impl Fetch for (T, AlwaysEqual<(&mut [F], U)>) { +impl<'b, T, U, F: Fetch1 + Copy> Fetch<'b> for (T, AlwaysEqual<(&'b [F], U)>) { + type Iter = Fetch1Iter>>; #[inline(always)] - fn fetch(&self) -> OwnedIter { - let vec = self.1.0.0.iter().map(|x| x.fetch_1()).collect::>(); - OwnedIter::from_slice(vec.as_slice()) + fn fetch(&self) -> Self::Iter { + Fetch1Iter { + iter: self.1.0.0.iter().copied(), + } } } diff --git a/crates/algo/src/prefetcher.rs b/crates/algo/src/prefetcher.rs index f2159e64..ca140047 100644 --- a/crates/algo/src/prefetcher.rs +++ b/crates/algo/src/prefetcher.rs @@ -22,12 +22,12 @@ use std::iter::Chain; pub const WINDOW_SIZE: usize = 32; const _: () = assert!(WINDOW_SIZE > 0); -pub trait Prefetcher<'r>: IntoIterator +pub trait Prefetcher<'b>: IntoIterator where - Self::Item: Fetch, + Self::Item: Fetch<'b>, { - type R: RelationRead + 'r; - type Guards: ExactSizeIterator::ReadGuard<'r>>; + type R: RelationRead + 'b; + type Guards: ExactSizeIterator::ReadGuard<'b>>; #[must_use] fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; @@ -38,19 +38,19 @@ where ) -> Option<(Self::Item, Self::Guards)>; } -pub struct PlainPrefetcher<'r, R, S: Sequence> { - relation: &'r R, +pub struct PlainPrefetcher<'b, R, S: Sequence> { + relation: &'b R, sequence: S, } -impl<'r, R, S: Sequence> PlainPrefetcher<'r, R, S> { +impl<'b, R, S: Sequence> PlainPrefetcher<'b, R, S> { #[inline] - pub fn new(relation: &'r R, sequence: S) -> Self { + pub fn new(relation: &'b R, sequence: S) -> Self { Self { relation, sequence } } } -impl<'r, R, S: Sequence> IntoIterator for PlainPrefetcher<'r, R, S> { +impl<'b, R, S: Sequence> IntoIterator for PlainPrefetcher<'b, R, S> { type Item = S::Item; type IntoIter = S::Inner; @@ -61,15 +61,20 @@ impl<'r, R, S: Sequence> IntoIterator for PlainPrefetcher<'r, R, S> { } } -impl<'r, R: RelationRead, S: Sequence> Prefetcher<'r> for PlainPrefetcher<'r, R, S> +impl<'b, R: RelationRead, S: Sequence> Prefetcher<'b> for PlainPrefetcher<'b, R, S> where - S::Item: Fetch, + S::Item: Fetch<'b>, { type R = R; - type Guards = PlainPrefetcherGuards<'r, R>; + type Guards = PlainPrefetcherGuards<'b, R, >::Iter>; #[inline] - fn next(&mut self) -> Option<(Self::Item, PlainPrefetcherGuards<'r, R>)> { + fn next( + &mut self, + ) -> Option<( + Self::Item, + PlainPrefetcherGuards<'b, R, >::Iter>, + )> { let e = self.sequence.next()?; let list = e.fetch(); Some(( @@ -85,7 +90,10 @@ where fn next_if( &mut self, predicate: impl FnOnce(&Self::Item) -> bool, - ) -> Option<(S::Item, PlainPrefetcherGuards<'r, R>)> { + ) -> Option<( + S::Item, + PlainPrefetcherGuards<'b, R, >::Iter>, + )> { if !predicate(self.sequence.peek()?) { return None; } @@ -101,13 +109,13 @@ where } } -pub struct PlainPrefetcherGuards<'r, R> { - relation: &'r R, - list: crate::OwnedIter, +pub struct PlainPrefetcherGuards<'b, R, L> { + relation: &'b R, + list: L, } -impl<'r, R: RelationRead> Iterator for PlainPrefetcherGuards<'r, R> { - type Item = R::ReadGuard<'r>; +impl<'b, R: RelationRead, L: Iterator> Iterator for PlainPrefetcherGuards<'b, R, L> { + type Item = R::ReadGuard<'b>; #[inline] fn next(&mut self) -> Option { @@ -121,17 +129,20 @@ impl<'r, R: RelationRead> Iterator for PlainPrefetcherGuards<'r, R> { } } -impl<'r, R: RelationRead> ExactSizeIterator for PlainPrefetcherGuards<'r, R> {} +impl<'b, R: RelationRead, L: Iterator> ExactSizeIterator + for PlainPrefetcherGuards<'b, R, L> +{ +} -pub struct SimplePrefetcher<'r, R, S: Sequence> { - relation: &'r R, +pub struct SimplePrefetcher<'b, R, S: Sequence> { + relation: &'b R, window: VecDeque, sequence: S, } -impl<'r, R, S: Sequence> SimplePrefetcher<'r, R, S> { +impl<'b, R, S: Sequence> SimplePrefetcher<'b, R, S> { #[inline] - pub fn new(relation: &'r R, sequence: S) -> Self { + pub fn new(relation: &'b R, sequence: S) -> Self { Self { relation, window: VecDeque::new(), @@ -140,7 +151,7 @@ impl<'r, R, S: Sequence> SimplePrefetcher<'r, R, S> { } } -impl<'r, R, S: Sequence> IntoIterator for SimplePrefetcher<'r, R, S> { +impl<'b, R, S: Sequence> IntoIterator for SimplePrefetcher<'b, R, S> { type Item = S::Item; type IntoIter = Chain, S::Inner>; @@ -151,16 +162,21 @@ impl<'r, R, S: Sequence> IntoIterator for SimplePrefetcher<'r, R, S> { } } -impl<'r, R: RelationRead + RelationPrefetch, S: Sequence> Prefetcher<'r> - for SimplePrefetcher<'r, R, S> +impl<'b, R: RelationRead + RelationPrefetch, S: Sequence> Prefetcher<'b> + for SimplePrefetcher<'b, R, S> where - S::Item: Fetch, + S::Item: Fetch<'b>, { type R = R; - type Guards = SimplePrefetcherGuards<'r, R>; + type Guards = SimplePrefetcherGuards<'b, R, >::Iter>; #[inline] - fn next(&mut self) -> Option<(Self::Item, SimplePrefetcherGuards<'r, R>)> { + fn next( + &mut self, + ) -> Option<( + Self::Item, + SimplePrefetcherGuards<'b, R, >::Iter>, + )> { while self.window.len() < WINDOW_SIZE && let Some(e) = self.sequence.next() { @@ -184,7 +200,10 @@ where fn next_if( &mut self, predicate: impl FnOnce(&S::Item) -> bool, - ) -> Option<(S::Item, SimplePrefetcherGuards<'r, R>)> { + ) -> Option<( + S::Item, + SimplePrefetcherGuards<'b, R, >::Iter>, + )> { while self.window.len() < WINDOW_SIZE && let Some(e) = self.sequence.next() { @@ -205,13 +224,13 @@ where } } -pub struct SimplePrefetcherGuards<'r, R> { - relation: &'r R, - list: crate::OwnedIter, +pub struct SimplePrefetcherGuards<'b, R, L> { + relation: &'b R, + list: L, } -impl<'r, R: RelationRead> Iterator for SimplePrefetcherGuards<'r, R> { - type Item = R::ReadGuard<'r>; +impl<'b, R: RelationRead, L: Iterator> Iterator for SimplePrefetcherGuards<'b, R, L> { + type Item = R::ReadGuard<'b>; #[inline] fn next(&mut self) -> Option { @@ -225,7 +244,10 @@ impl<'r, R: RelationRead> Iterator for SimplePrefetcherGuards<'r, R> { } } -impl<'r, R: RelationRead> ExactSizeIterator for SimplePrefetcherGuards<'r, R> {} +impl<'b, R: RelationRead, L: ExactSizeIterator> ExactSizeIterator + for SimplePrefetcherGuards<'b, R, L> +{ +} pub struct StreamPrefetcherSequence(S); @@ -238,31 +260,31 @@ impl Iterator for StreamPrefetcherSequence { } } -pub struct StreamPrefetcher<'r, R: RelationReadStream + 'r, S: Sequence> +pub struct StreamPrefetcher<'b, R: RelationReadStream + 'b, S: Sequence> where - S::Item: Fetch, + S::Item: Fetch<'b>, { - stream: R::ReadStream<'r, StreamPrefetcherSequence>, + stream: R::ReadStream<'b, StreamPrefetcherSequence>, } -impl<'r, R: RelationReadStream, S: Sequence> StreamPrefetcher<'r, R, S> +impl<'b, R: RelationReadStream, S: Sequence> StreamPrefetcher<'b, R, S> where - S::Item: Fetch, + S::Item: Fetch<'b>, { #[inline] - pub fn new(relation: &'r R, sequence: S, hints: Hints) -> Self { + pub fn new(relation: &'b R, sequence: S, hints: Hints) -> Self { let stream = relation.read_stream(StreamPrefetcherSequence(sequence), hints); Self { stream } } } -impl<'r, R: RelationReadStream, S: Sequence> IntoIterator for StreamPrefetcher<'r, R, S> +impl<'b, R: RelationReadStream, S: Sequence> IntoIterator for StreamPrefetcher<'b, R, S> where - S::Item: Fetch, + S::Item: Fetch<'b>, { type Item = S::Item; - type IntoIter = > as ReadStream<'r>>::Inner; + type IntoIter = > as ReadStream<'b>>::Inner; #[inline] fn into_iter(self) -> Self::IntoIter { @@ -270,14 +292,14 @@ where } } -impl<'r, R: RelationRead + RelationReadStream, S: Sequence> Prefetcher<'r> - for StreamPrefetcher<'r, R, S> +impl<'b, R: RelationRead + RelationReadStream, S: Sequence> Prefetcher<'b> + for StreamPrefetcher<'b, R, S> where - S::Item: Fetch, + S::Item: Fetch<'b>, { type R = R; - type Guards = > as ReadStream<'r>>::Guards; + type Guards = > as ReadStream<'b>>::Guards; #[inline] fn next(&mut self) -> Option<(S::Item, Self::Guards)> { @@ -293,26 +315,26 @@ where } } -pub trait PrefetcherSequenceFamily<'r, R> { - type P: Prefetcher<'r, R = R, Item = S::Item> +pub trait PrefetcherSequenceFamily<'b, R> { + type P: Prefetcher<'b, R = R, Item = S::Item> where - S::Item: Fetch; + S::Item: Fetch<'b>; fn prefetch(&mut self, seq: S) -> Self::P where - S::Item: Fetch; + S::Item: Fetch<'b>; fn is_not_plain(&self) -> bool; } -pub trait PrefetcherHeapFamily<'r, R> { - type P: Prefetcher<'r, R = R, Item = T> +pub trait PrefetcherHeapFamily<'b, R> { + type P: Prefetcher<'b, R = R, Item = T> where - T: Ord + Fetch + 'r; + T: Ord + Fetch<'b>; fn prefetch(&mut self, seq: Vec) -> Self::P where - T: Ord + Fetch + 'r; + T: Ord + Fetch<'b>; fn is_not_plain(&self) -> bool; } diff --git a/crates/sbsii/src/owned.rs b/crates/sbsii/src/owned.rs deleted file mode 100644 index e2a163cf..00000000 --- a/crates/sbsii/src/owned.rs +++ /dev/null @@ -1,192 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -use crate::stack::StackIntoIter; -use std::ptr::NonNull; - -pub struct HeapIntoIter { - pointer: NonNull, - off: u16, - len: u16, -} - -impl HeapIntoIter { - #[cold] - pub(crate) fn from_slice(slice: &[T]) -> Self { - assert!(slice.len() <= 65535_usize); - let c = box_into_non_null::<[T]>(Box::from(slice)).cast::(); - Self { - pointer: c, - off: 0, - len: slice.len() as u16, - } - } - - #[inline(always)] - pub(crate) fn as_slice(&self) -> &[T] { - #[allow(unsafe_code)] - unsafe { - std::slice::from_raw_parts( - self.pointer.as_ptr().add(self.off as _), - (self.len - self.off) as _, - ) - } - } -} - -impl Iterator for HeapIntoIter { - type Item = T; - - #[inline(always)] - fn next(&mut self) -> Option { - #[allow(unsafe_code)] - unsafe { - if self.off < self.len { - let r = self.pointer.as_ptr().add(self.off as _).read(); - self.off += 1; - Some(r) - } else { - None - } - } - } - - #[inline(always)] - fn size_hint(&self) -> (usize, Option) { - let size = (self.len - self.off) as usize; - (size, Some(size)) - } -} - -impl ExactSizeIterator for HeapIntoIter {} - -impl Clone for HeapIntoIter { - #[inline(always)] - fn clone(&self) -> Self { - #[allow(unsafe_code)] - unsafe { - let p = std::slice::from_raw_parts_mut(self.pointer.as_ptr(), self.len as _); - let c = NonNull::new_unchecked(Box::<[T]>::into_raw(Box::from(p)).cast::()); - Self { - pointer: c, - off: self.off, - len: self.len, - } - } - } -} - -impl Drop for HeapIntoIter { - #[inline(always)] - fn drop(&mut self) { - #[allow(unsafe_code)] - unsafe { - let p = std::slice::from_raw_parts_mut(self.pointer.as_ptr(), self.len as _); - let _ = Box::<[T]>::from_raw(p); - } - } -} - -#[derive(Clone)] -pub enum IntoIter { - Stack(StackIntoIter), - Heap(HeapIntoIter), -} - -impl IntoIter { - #[inline(always)] - pub fn from_slice(slice: &[T]) -> Self { - assert!(slice.len() <= 65535); - if slice.len() <= N && N <= 65535 { - IntoIter::Stack(StackIntoIter::from_slice(slice)) - } else { - IntoIter::Heap(HeapIntoIter::from_slice(slice)) - } - } - - #[inline(always)] - pub fn as_slice(&self) -> &[T] { - match self { - IntoIter::Stack(x) => x.as_slice(), - IntoIter::Heap(x) => x.as_slice(), - } - } -} - -impl Iterator for IntoIter { - type Item = T; - - #[inline(always)] - fn next(&mut self) -> Option { - match self { - IntoIter::Stack(iter) => iter.next(), - IntoIter::Heap(iter) => iter.next(), - } - } - - #[inline(always)] - fn size_hint(&self) -> (usize, Option) { - match self { - IntoIter::Stack(iter) => iter.size_hint(), - IntoIter::Heap(iter) => iter.size_hint(), - } - } -} - -impl ExactSizeIterator for IntoIter {} - -#[cfg(target_pointer_width = "64")] -const _: () = { - assert!(size_of::>() == 24); -}; - -// Emulate unstable library feature `box_vec_non_null`. -// See https://github.com/rust-lang/rust/issues/130364. - -#[allow(dead_code)] -#[must_use] -fn box_into_non_null(b: Box) -> NonNull { - #[allow(unsafe_code)] - unsafe { - NonNull::new_unchecked(Box::into_raw(b)) - } -} - -#[test] -fn tests() { - for x in IntoIter::::from_slice(&[1; 0]).clone() { - assert_eq!(x, 1); - } - for x in IntoIter::::from_slice(&[1; 1]).clone() { - assert_eq!(x, 1); - } - for x in IntoIter::::from_slice(&[1; 2]).clone() { - assert_eq!(x, 1); - } - for x in IntoIter::::from_slice(&[1; 3]).clone() { - assert_eq!(x, 1); - } - for x in IntoIter::::from_slice(&[1; 4]).clone() { - assert_eq!(x, 1); - } - for x in IntoIter::::from_slice(&[1; 5]).clone() { - assert_eq!(x, 1); - } - for x in IntoIter::::from_slice(&[1; 6]).clone() { - assert_eq!(x, 1); - } - for x in IntoIter::::from_slice(&[1; 7]).clone() { - assert_eq!(x, 1); - } -} diff --git a/crates/sbsii/Cargo.toml b/crates/small_iter/Cargo.toml similarity index 85% rename from crates/sbsii/Cargo.toml rename to crates/small_iter/Cargo.toml index 4af4bd3c..011d90b9 100644 --- a/crates/sbsii/Cargo.toml +++ b/crates/small_iter/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "sbsii" +name = "small_iter" version.workspace = true edition.workspace = true publish = false diff --git a/crates/sbsii/src/borrowed.rs b/crates/small_iter/src/borrowed.rs similarity index 56% rename from crates/sbsii/src/borrowed.rs rename to crates/small_iter/src/borrowed.rs index 4cffdc6f..3e76300e 100644 --- a/crates/sbsii/src/borrowed.rs +++ b/crates/small_iter/src/borrowed.rs @@ -12,44 +12,33 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::stack::StackIntoIter; +use crate::stack::StackIter; use std::marker::PhantomData; use std::ptr::NonNull; -#[derive(Clone)] -pub struct HeapIntoIter<'a, T: Copy> { - pointer: NonNull, +#[derive(Clone, Copy)] +pub struct HeapIter<'a, T: Copy> { off: u16, len: u16, - _phantom: PhantomData<&'a ()>, + pointer: NonNull, + _phantom: PhantomData<&'a T>, } -impl<'a, T: Copy> HeapIntoIter<'a, T> { +impl<'a, T: Copy> HeapIter<'a, T> { #[cold] pub(crate) fn from_slice(slice: &'a [T]) -> Self { assert!(slice.len() <= 65535_usize); let c = NonNull::from_ref(slice).cast::(); Self { - pointer: c, off: 0, len: slice.len() as u16, + pointer: c, _phantom: PhantomData, } } - - #[inline(always)] - pub(crate) fn as_slice(&self) -> &[T] { - #[allow(unsafe_code)] - unsafe { - std::slice::from_raw_parts( - self.pointer.as_ptr().add(self.off as _), - (self.len - self.off) as _, - ) - } - } } -impl<'a, T: Copy> Iterator for HeapIntoIter<'a, T> { +impl<'a, T: Copy> Iterator for HeapIter<'a, T> { type Item = T; #[inline(always)] @@ -73,59 +62,51 @@ impl<'a, T: Copy> Iterator for HeapIntoIter<'a, T> { } } -impl<'a, T: Copy> ExactSizeIterator for HeapIntoIter<'a, T> {} +impl<'a, T: Copy> ExactSizeIterator for HeapIter<'a, T> {} -#[derive(Clone)] -pub enum IntoIter<'a, T: Copy, const N: usize> { - Stack(StackIntoIter), - Heap(HeapIntoIter<'a, T>), +#[derive(Clone, Copy)] +pub enum Iter<'a, T: Copy, const N: usize> { + Stack(StackIter), + Heap(HeapIter<'a, T>), } -impl<'a, T: Copy + 'a, const N: usize> IntoIter<'a, T, N> { +impl<'a, T: Copy, const N: usize> Iter<'a, T, N> { #[inline(always)] pub fn from_slice(slice: &[T], alloc: impl Fn(&[T]) -> &'a [T]) -> Self { assert!(slice.len() <= 65535); if slice.len() <= N && N <= 65535 { - IntoIter::Stack(StackIntoIter::from_slice(slice)) + Iter::Stack(StackIter::from_slice(slice)) } else { - IntoIter::Heap(HeapIntoIter::from_slice(alloc(slice))) - } - } - - #[inline(always)] - pub fn as_slice(&self) -> &[T] { - match self { - IntoIter::Stack(x) => x.as_slice(), - IntoIter::Heap(x) => x.as_slice(), + Iter::Heap(HeapIter::from_slice(alloc(slice))) } } } -impl<'a, T: Copy, const N: usize> Iterator for IntoIter<'a, T, N> { +impl<'a, T: Copy, const N: usize> Iterator for Iter<'a, T, N> { type Item = T; #[inline(always)] fn next(&mut self) -> Option { match self { - IntoIter::Stack(iter) => iter.next(), - IntoIter::Heap(iter) => iter.next(), + Iter::Stack(iter) => iter.next(), + Iter::Heap(iter) => iter.next(), } } #[inline(always)] fn size_hint(&self) -> (usize, Option) { match self { - IntoIter::Stack(iter) => iter.size_hint(), - IntoIter::Heap(iter) => iter.size_hint(), + Iter::Stack(iter) => iter.size_hint(), + Iter::Heap(iter) => iter.size_hint(), } } } -impl<'a, T: Copy, const N: usize> ExactSizeIterator for IntoIter<'a, T, N> {} +impl<'a, T: Copy, const N: usize> ExactSizeIterator for Iter<'a, T, N> {} #[cfg(target_pointer_width = "64")] const _: () = { - assert!(size_of::>() == 16); + assert!(size_of::>() == 16); }; #[test] @@ -137,28 +118,28 @@ fn tests() { GLOBAL.lock().expect("failed to lock").push(pointer); pointer }; - for x in IntoIter::::from_slice(&[1; 0], alloc) { + for x in Iter::::from_slice(&[1; 0], alloc) { assert_eq!(x, 1); } - for x in IntoIter::::from_slice(&[1; 1], alloc) { + for x in Iter::::from_slice(&[1; 1], alloc) { assert_eq!(x, 1); } - for x in IntoIter::::from_slice(&[1; 2], alloc) { + for x in Iter::::from_slice(&[1; 2], alloc) { assert_eq!(x, 1); } - for x in IntoIter::::from_slice(&[1; 3], alloc) { + for x in Iter::::from_slice(&[1; 3], alloc) { assert_eq!(x, 1); } - for x in IntoIter::::from_slice(&[1; 4], alloc) { + for x in Iter::::from_slice(&[1; 4], alloc) { assert_eq!(x, 1); } - for x in IntoIter::::from_slice(&[1; 5], alloc) { + for x in Iter::::from_slice(&[1; 5], alloc) { assert_eq!(x, 1); } - for x in IntoIter::::from_slice(&[1; 6], alloc) { + for x in Iter::::from_slice(&[1; 6], alloc) { assert_eq!(x, 1); } - for x in IntoIter::::from_slice(&[1; 7], alloc) { + for x in Iter::::from_slice(&[1; 7], alloc) { assert_eq!(x, 1); } } diff --git a/crates/sbsii/src/lib.rs b/crates/small_iter/src/lib.rs similarity index 91% rename from crates/sbsii/src/lib.rs rename to crates/small_iter/src/lib.rs index 304b9354..886062e5 100644 --- a/crates/sbsii/src/lib.rs +++ b/crates/small_iter/src/lib.rs @@ -12,8 +12,5 @@ // // Copyright (c) 2025 TensorChord Inc. -// sbsii, small boxed slice into iter - pub mod borrowed; -pub mod owned; pub mod stack; diff --git a/crates/sbsii/src/stack.rs b/crates/small_iter/src/stack.rs similarity index 61% rename from crates/sbsii/src/stack.rs rename to crates/small_iter/src/stack.rs index 2351db07..6e78ecd1 100644 --- a/crates/sbsii/src/stack.rs +++ b/crates/small_iter/src/stack.rs @@ -14,20 +14,20 @@ use std::mem::MaybeUninit; -#[derive(Clone)] -pub struct StackIntoIter { - start: u16, - end: u16, +#[derive(Clone, Copy)] +pub struct StackIter { + off: u16, + len: u16, buffer: [MaybeUninit; N], } -impl StackIntoIter { +impl StackIter { #[inline(always)] pub(crate) fn from_slice(slice: &[T]) -> Self { assert!(slice.len() <= N && N <= 65535); Self { - start: 0, - end: slice.len() as u16, + off: 0, + len: slice.len() as u16, buffer: { let mut buffer = [const { MaybeUninit::uninit() }; N]; for i in 0..slice.len() { @@ -37,29 +37,18 @@ impl StackIntoIter { }, } } - - #[inline(always)] - pub(crate) fn as_slice(&self) -> &[T] { - #[allow(unsafe_code)] - unsafe { - std::slice::from_raw_parts( - self.buffer.as_ptr().add(self.start as _).cast(), - (self.end - self.start) as _, - ) - } - } } -impl Iterator for StackIntoIter { +impl Iterator for StackIter { type Item = T; #[inline(always)] fn next(&mut self) -> Option { #[allow(unsafe_code)] unsafe { - if self.start < self.end { - let r = self.buffer[self.start as usize].assume_init(); - self.start += 1; + if self.off < self.len { + let r = self.buffer[self.off as usize].assume_init(); + self.off += 1; Some(r) } else { None @@ -69,9 +58,9 @@ impl Iterator for StackIntoIter { #[inline(always)] fn size_hint(&self) -> (usize, Option) { - let size = (self.end - self.start) as usize; + let size = (self.len - self.off) as usize; (size, Some(size)) } } -impl ExactSizeIterator for StackIntoIter {} +impl ExactSizeIterator for StackIter {} diff --git a/crates/vchordg/src/candidates.rs b/crates/vchordg/src/candidates.rs index fb77a012..fd901e2f 100644 --- a/crates/vchordg/src/candidates.rs +++ b/crates/vchordg/src/candidates.rs @@ -20,7 +20,7 @@ use std::marker::PhantomData; pub struct Candidates<'a, P, F, R> where P: Prefetcher<'a>, - P::Item: Fetch, + P::Item: Fetch<'a>, { beam: usize, front: Option

, @@ -32,7 +32,7 @@ where impl<'a, P, F, R> Candidates<'a, P, F, R> where P: Prefetcher<'a, R = R>, - P::Item: Fetch + Ord, + P::Item: Fetch<'a> + Ord, F: PrefetcherSequenceFamily<'a, R, P> = P>, R: RelationRead, { diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index 648e098b..90f68566 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -145,7 +145,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( return; }; let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); - let pointers_s = bump.alloc_slice(vertex_tuple.pointers()); + let pointers_s: &[_] = bump.alloc_slice(vertex_tuple.pointers()); let score_s = O::process( bits, dims, @@ -184,7 +184,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( continue; }; let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); - let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); + let pointers_v: &[_] = bump.alloc_slice(vertex_tuple.pointers()); let score_v = O::process( bits, dims, @@ -229,7 +229,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( }; let outs = crate::prune::prune( |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), - (bump.alloc_slice(&pointers_t), t), + (bump.alloc_slice(&pointers_t) as &[_], t), trace.into_iter(), m, &alpha, @@ -301,11 +301,11 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( } } -fn append_vertex_tuple<'r, R: RelationRead + RelationWrite>( - index: &'r R, +fn append_vertex_tuple<'b, R: RelationRead + RelationWrite>( + index: &'b R, first: u32, size: usize, -) -> R::WriteGuard<'r> +) -> R::WriteGuard<'b> where R::Page: Page, { @@ -350,11 +350,11 @@ where } } -fn append_vector_tuple<'r, R: RelationRead + RelationWrite>( - index: &'r R, +fn append_vector_tuple<'b, R: RelationRead + RelationWrite>( + index: &'b R, first: u32, size: usize, -) -> R::WriteGuard<'r> +) -> R::WriteGuard<'b> where R::Page: Page, { diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs index 78117bb4..af8d6721 100644 --- a/crates/vchordg/src/search.rs +++ b/crates/vchordg/src/search.rs @@ -66,7 +66,7 @@ where return Box::new(std::iter::empty()); }; let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); - let pointers_s = bump.alloc_slice(vertex_tuple.pointers()); + let pointers_s: &[_] = bump.alloc_slice(vertex_tuple.pointers()); let score_s = O::process( bits, dims, diff --git a/crates/vchordg/src/vectors.rs b/crates/vchordg/src/vectors.rs index dec945bd..94a43aa0 100644 --- a/crates/vchordg/src/vectors.rs +++ b/crates/vchordg/src/vectors.rs @@ -20,10 +20,10 @@ use distance::Distance; use std::collections::VecDeque; use std::num::{NonZero, Wrapping}; -pub fn by_prefetch<'r, R: RelationRead>( - guards: impl ExactSizeIterator>, +pub fn by_prefetch<'b, R: RelationRead>( + guards: impl ExactSizeIterator>, pointers_u: impl ExactSizeIterator, -) -> impl ExactSizeIterator, u16)> { +) -> impl ExactSizeIterator, u16)> { assert!(guards.len() == pointers_u.len() && pointers_u.len() > 0); pointers_u .map(Pointer::into_inner) @@ -31,10 +31,10 @@ pub fn by_prefetch<'r, R: RelationRead>( .map(|(pointer, guard)| (guard, pointer.1)) } -pub fn by_read<'r, R: RelationRead>( - index: &'r R, +pub fn by_read<'b, R: RelationRead>( + index: &'b R, pointers_u: impl ExactSizeIterator, -) -> impl ExactSizeIterator, u16)> { +) -> impl ExactSizeIterator, u16)> { assert!(pointers_u.len() > 0); pointers_u .map(Pointer::into_inner) @@ -58,13 +58,13 @@ pub fn copy_all(x: &[OptionNeighbour]) -> VecDeque<((u32, u16), Distance)> { } pub fn read< - 'r, + 'b, R: RelationRead, O: Operator, A: Accessor1<::Element, ::Metadata>, Output, >( - mut iterator: impl ExactSizeIterator, u16)>, + mut iterator: impl ExactSizeIterator, u16)>, accessor: A, copy: impl FnOnce(&[OptionNeighbour]) -> Output, ) -> Result<(A::Output, Output, Option>, Wrapping), ()> { diff --git a/crates/vchordrq/src/closure_lifetime_binder.rs b/crates/vchordrq/src/closure_lifetime_binder.rs index 4eb2d063..a7472d30 100644 --- a/crates/vchordrq/src/closure_lifetime_binder.rs +++ b/crates/vchordrq/src/closure_lifetime_binder.rs @@ -49,10 +49,10 @@ where } #[inline(always)] -pub fn id_4<'r, F, P, A, B, R: ?Sized>(f: F) -> F +pub fn id_4<'b, F, P, A, B, R: ?Sized>(f: F) -> F where - P: algo::prefetcher::Prefetcher<'r>, -

::Item: algo::Fetch, + P: algo::prefetcher::Prefetcher<'b>, +

::Item: algo::Fetch<'b>, F: FnMut(A, P::Guards, B) -> R, { f diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index c615ddf4..3d6b337f 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -55,13 +55,13 @@ where } } -pub fn insert<'r, 'b: 'r, R: RelationRead + RelationWrite, O: Operator>( - index: &'r R, +pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( + index: &'b R, payload: NonZero, vector: ::Borrowed<'_>, key: (Vec, u16), bump: &'b impl Bump, - mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, + mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'b, R>, skip_freespaces: bool, ) where R::Page: Page, diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs index de0b64b6..3d559c50 100644 --- a/crates/vchordrq/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -32,9 +32,9 @@ pub struct Maintain { pub number_of_freed_pages: usize, } -pub fn maintain<'r, R: RelationRead + RelationWrite, O: Operator>( - index: &'r R, - mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, +pub fn maintain<'b, R: RelationRead + RelationWrite, O: Operator>( + index: &'b R, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, check: impl Fn(), ) -> Maintain where @@ -249,9 +249,9 @@ where } #[derive(Clone)] -struct RelationHooked<'r, R, E>(&'r R, E); +struct RelationHooked<'b, R, E>(&'b R, E); -impl<'r, R, E> Relation for RelationHooked<'r, R, E> +impl<'b, R, E> Relation for RelationHooked<'b, R, E> where R: Relation, E: Clone, @@ -259,7 +259,7 @@ where type Page = R::Page; } -impl<'r, R, E> RelationReadTypes for RelationHooked<'r, R, E> +impl<'b, R, E> RelationReadTypes for RelationHooked<'b, R, E> where R: RelationRead, E: Clone, @@ -267,7 +267,7 @@ where type ReadGuard<'a> = R::ReadGuard<'a>; } -impl<'r, R, E> RelationRead for RelationHooked<'r, R, E> +impl<'b, R, E> RelationRead for RelationHooked<'b, R, E> where R: RelationRead, E: Clone, @@ -277,7 +277,7 @@ where } } -impl<'r, R, E> RelationWriteTypes for RelationHooked<'r, R, E> +impl<'b, R, E> RelationWriteTypes for RelationHooked<'b, R, E> where R: RelationWrite, E: Clone + for<'a> Fn(&'a R, ::Opaque, bool) -> R::WriteGuard<'a>, @@ -285,7 +285,7 @@ where type WriteGuard<'a> = R::WriteGuard<'a>; } -impl<'r, R, E> RelationWrite for RelationHooked<'r, R, E> +impl<'b, R, E> RelationWrite for RelationHooked<'b, R, E> where R: RelationWrite, E: Clone + for<'a> Fn(&'a R, ::Opaque, bool) -> R::WriteGuard<'a>, diff --git a/crates/vchordrq/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs index 4a9de383..d8066f69 100644 --- a/crates/vchordrq/src/prewarm.rs +++ b/crates/vchordrq/src/prewarm.rs @@ -22,10 +22,10 @@ use algo::accessor::FunctionalAccessor; use algo::prefetcher::PrefetcherSequenceFamily; use std::fmt::Write; -pub fn prewarm<'r, R: RelationRead, O: Operator>( - index: &'r R, +pub fn prewarm<'b, R: RelationRead, O: Operator>( + index: &'b R, height: i32, - mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, ) -> String where R::Page: Page, @@ -45,7 +45,7 @@ where type State = Vec; let mut state: State = { let mut results = Vec::new(); - let prefetch = algo::OwnedIter::from_slice(meta_tuple.centroid_prefetch()); + let prefetch = meta_tuple.centroid_prefetch().to_vec().into_iter(); let head = meta_tuple.centroid_head(); let first = meta_tuple.first(); vectors::read_for_h1_tuple::(prefetch.map(|id| index.read(id)), head, ()); diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs index 7d3cc032..62e49006 100644 --- a/crates/vchordrq/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -50,10 +50,10 @@ pub struct Reranker { _phantom: PhantomData T>, } -impl<'r, 'b, T, F, P> Iterator for Reranker +impl<'b, T, F, P> Iterator for Reranker where F: FnMut(NonZero, P::Guards, u16) -> Option, - P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, { type Item = (Distance, NonZero); @@ -78,11 +78,10 @@ impl Reranker { } pub fn rerank_index< - 'r, 'b, O: Operator, T, - P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, >( vector: O::Vector, prefetcher: P, @@ -106,11 +105,10 @@ pub fn rerank_index< } pub fn rerank_heap< - 'r, 'b, O: Operator, T, - P: Prefetcher<'r, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, >( vector: O::Vector, prefetcher: P, diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index c3acee32..c737c702 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -31,14 +31,14 @@ use vector::{VectorBorrowed, VectorOwned}; type Extra1<'b> = &'b mut (u32, f32, u16, BorrowedIter<'b>); type Extra0<'b> = &'b mut (NonZero, u16, BorrowedIter<'b>); -pub fn default_search<'r, 'b: 'r, R: RelationRead, O: Operator>( - index: &'r R, +pub fn default_search<'b, R: RelationRead, O: Operator>( + index: &'b R, vector: ::Borrowed<'_>, probes: Vec, epsilon: f32, bump: &'b impl Bump, - mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, - mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, + mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'b, R>, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, ) -> Vec<( (Reverse, AlwaysEqual<()>), AlwaysEqual>, @@ -168,15 +168,15 @@ where results.into_vec() } -pub fn maxsim_search<'r, 'b: 'r, R: RelationRead, O: Operator>( - index: &'r R, +pub fn maxsim_search<'b, R: RelationRead, O: Operator>( + index: &'b R, vector: ::Borrowed<'_>, probes: Vec, epsilon: f32, mut threshold: u32, bump: &'b impl Bump, - mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'r, R>, - mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'r, R>, + mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'b, R>, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, ) -> ( Vec<( (Reverse, AlwaysEqual), diff --git a/crates/vchordrq/src/tape.rs b/crates/vchordrq/src/tape.rs index 2de5e546..16f4e54a 100644 --- a/crates/vchordrq/src/tape.rs +++ b/crates/vchordrq/src/tape.rs @@ -113,24 +113,24 @@ where } } -pub fn read_directory_tape<'r, R>( - iter: impl Iterator>, +pub fn read_directory_tape<'b, R>( + iter: impl Iterator>, ) -> impl Iterator where - R: RelationRead + 'r, + R: RelationRead + 'b, { use std::pin::Pin; use std::ptr::NonNull; #[pin_project::pin_project] - struct State<'r, R: RelationRead + 'r, I> { + struct State<'b, R: RelationRead + 'b, I> { slice: NonNull<[u32]>, #[pin] - now: Option<(R::ReadGuard<'r>, u16)>, + now: Option<(R::ReadGuard<'b>, u16)>, iter: I, } - impl<'r, R: RelationRead + 'r, I: Iterator>> State<'r, R, I> { + impl<'b, R: RelationRead + 'b, I: Iterator>> State<'b, R, I> { fn init(self: Pin<&mut Self>) { let mut this = self.project(); let now = this.iter.next().map(|guard| (guard, 0)); @@ -172,14 +172,14 @@ where } } - let mut state = Box::pin(State::<'r, R, _> { + let mut state = Box::pin(State::<'b, R, _> { slice: NonNull::from(&mut []), now: None, iter, }); - impl<'r, R: RelationRead + 'r, I: Iterator>> Iterator - for Pin>> + impl<'b, R: RelationRead + 'b, I: Iterator>> Iterator + for Pin>> { type Item = u32; @@ -193,12 +193,12 @@ where state } -pub fn by_directory<'r, R>( - p: &mut impl PrefetcherSequenceFamily<'r, R>, +pub fn by_directory<'b, R>( + p: &mut impl PrefetcherSequenceFamily<'b, R>, iter: impl Iterator, -) -> impl Iterator> +) -> impl Iterator> where - R: RelationRead + 'r, + R: RelationRead + 'b, { let mut t = p.prefetch(iter.peekable()); std::iter::from_fn(move || { @@ -209,9 +209,9 @@ where }) } -pub fn by_next<'r, R>(index: &'r R, first: u32) -> impl Iterator> +pub fn by_next<'b, R>(index: &'b R, first: u32) -> impl Iterator> where - R: RelationRead + 'r, + R: RelationRead + 'b, R::Page: Page, { let mut current = first; @@ -226,12 +226,12 @@ where }) } -pub fn read_h1_tape<'r, R, A, T>( - iter: impl Iterator>, +pub fn read_h1_tape<'b, R, A, T>( + iter: impl Iterator>, accessor: impl Fn() -> A, mut callback: impl for<'a> FnMut(T, u16, f32, u32, &'a [u32]), ) where - R: RelationRead + 'r, + R: RelationRead + 'b, A: for<'a> Accessor1<[u8; 16], (&'a [[f32; 32]; 4], &'a [f32; 32]), Output = [T; 32]>, { let mut x = None; @@ -265,12 +265,12 @@ pub fn read_h1_tape<'r, R, A, T>( } } -pub fn read_frozen_tape<'r, R, A, T>( - iter: impl Iterator>, +pub fn read_frozen_tape<'b, R, A, T>( + iter: impl Iterator>, accessor: impl Fn() -> A, mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), ) where - R: RelationRead + 'r, + R: RelationRead + 'b, A: for<'a> Accessor1<[u8; 16], (&'a [[f32; 32]; 4], &'a [f32; 32]), Output = [T; 32]>, { let mut x = None; @@ -298,12 +298,12 @@ pub fn read_frozen_tape<'r, R, A, T>( } } -pub fn read_appendable_tape<'r, R, T>( - iter: impl Iterator>, +pub fn read_appendable_tape<'b, R, T>( + iter: impl Iterator>, mut access: impl for<'a> FnMut([f32; 4], &'a [u64], f32) -> T, mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), ) where - R: RelationRead + 'r, + R: RelationRead + 'b, { for guard in iter { for i in 1..=guard.len() { diff --git a/crates/vector/src/vect.rs b/crates/vector/src/vect.rs index 8854a055..8f40cf2b 100644 --- a/crates/vector/src/vect.rs +++ b/crates/vector/src/vect.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::{VectorBorrowed, VectorOwned}; +use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; use simd::Floating; use std::cmp::Ordering; diff --git a/src/datatype/binary_scalar8.rs b/src/datatype/binary_scalar8.rs index 57f510f5..e1c35efe 100644 --- a/src/datatype/binary_scalar8.rs +++ b/src/datatype/binary_scalar8.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::memory_scalar8::{Scalar8Input, Scalar8Output}; +use crate::datatype::memory_scalar8::{Scalar8Input, Scalar8Output}; use pgrx::datum::Internal; use pgrx::pg_sys::Oid; use vector::VectorBorrowed; diff --git a/src/datatype/text_scalar8.rs b/src/datatype/text_scalar8.rs index 0b519589..e123df51 100644 --- a/src/datatype/text_scalar8.rs +++ b/src/datatype/text_scalar8.rs @@ -12,8 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::memory_scalar8::Scalar8Output; -use crate::datatype::memory_scalar8::Scalar8Input; +use crate::datatype::memory_scalar8::{Scalar8Input, Scalar8Output}; use pgrx::pg_sys::Oid; use std::ffi::{CStr, CString}; use vector::scalar8::Scalar8Borrowed; diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 855d6b82..2e039c14 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -12,6 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. +use crate::index::scanners::Io; use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting, PostgresGucEnum}; use std::ffi::CString; @@ -242,8 +243,7 @@ pub fn vchordg_max_scan_tuples() -> Option { if x < 0 { None } else { Some(x as u32) } } -pub fn vchordg_io_search() -> crate::index::vchordg::scanners::Io { - use crate::index::vchordg::scanners::Io; +pub fn vchordg_io_search() -> Io { match VCHORDG_IO_SEARCH.get() { PostgresIo::ReadBuffer => Io::Plain, PostgresIo::PrefetchBuffer => Io::Simple, @@ -252,8 +252,7 @@ pub fn vchordg_io_search() -> crate::index::vchordg::scanners::Io { } } -pub fn vchordg_io_rerank() -> crate::index::vchordg::scanners::Io { - use crate::index::vchordg::scanners::Io; +pub fn vchordg_io_rerank() -> Io { match VCHORDG_IO_RERANK.get() { PostgresIo::ReadBuffer => Io::Plain, PostgresIo::PrefetchBuffer => Io::Simple, @@ -315,8 +314,7 @@ pub fn vchordrq_prefilter() -> bool { VCHORDRQ_PREFILTER.get() } -pub fn vchordrq_io_search() -> crate::index::vchordrq::scanners::Io { - use crate::index::vchordrq::scanners::Io; +pub fn vchordrq_io_search() -> Io { match VCHORDRQ_IO_SEARCH.get() { PostgresIo::ReadBuffer => Io::Plain, PostgresIo::PrefetchBuffer => Io::Simple, @@ -325,8 +323,7 @@ pub fn vchordrq_io_search() -> crate::index::vchordrq::scanners::Io { } } -pub fn vchordrq_io_rerank() -> crate::index::vchordrq::scanners::Io { - use crate::index::vchordrq::scanners::Io; +pub fn vchordrq_io_rerank() -> Io { match VCHORDRQ_IO_RERANK.get() { PostgresIo::ReadBuffer => Io::Plain, PostgresIo::PrefetchBuffer => Io::Simple, diff --git a/src/index/mod.rs b/src/index/mod.rs index 55b8d86c..4cc178cd 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -17,6 +17,7 @@ pub mod functions; pub mod gucs; pub mod hook; pub mod opclass; +pub mod scanners; pub mod storage; pub mod vchordg; pub mod vchordrq; diff --git a/src/index/scanners.rs b/src/index/scanners.rs new file mode 100644 index 00000000..b4203b97 --- /dev/null +++ b/src/index/scanners.rs @@ -0,0 +1,51 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::index::fetcher::Fetcher; +use algo::{Bump, Page, RelationPrefetch, RelationRead, RelationReadStream}; +use pgrx::pg_sys::Datum; + +#[derive(Debug, Clone, Copy)] +pub enum Io { + Plain, + Simple, + #[cfg_attr( + any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"), + expect(dead_code) + )] + Stream, +} + +pub trait SearchBuilder: 'static { + type Options; + + type Opfamily; + + type Opaque: Copy; + + fn new(opfamily: Self::Opfamily) -> Self; + + unsafe fn add(&mut self, strategy: u16, datum: Option); + + fn build<'b, R>( + self, + relation: &'b R, + options: Self::Options, + fetcher: impl Fetcher + 'b, + bump: &'b impl Bump, + ) -> Box + 'b> + where + R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page; +} diff --git a/src/index/storage.rs b/src/index/storage.rs index 6047981a..34027948 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -397,25 +397,27 @@ impl RelationPrefetch for PostgresRelation { } } -pub struct Cache { +pub struct Cache<'b, I: Iterator> { window: VecDeque, tail: VecDeque, iter: Option, + _phantom: PhantomData<&'b mut ()>, } -impl Default for Cache { +impl<'b, I: Iterator> Default for Cache<'b, I> { fn default() -> Self { Self { window: Default::default(), tail: Default::default(), iter: Default::default(), + _phantom: PhantomData, } } } -impl Cache +impl<'b, I: Iterator> Cache<'b, I> where - I::Item: Fetch, + I::Item: Fetch<'b>, { #[allow(dead_code)] pub fn pop_id(&mut self) -> Option { @@ -458,11 +460,11 @@ where } } -pub struct PostgresReadStream { +pub struct PostgresReadStream<'b, O: Opaque, I: Iterator> { #[cfg(any(feature = "pg17", feature = "pg18"))] raw: *mut pgrx::pg_sys::ReadStream, // Because of `Box`'s special alias rules, `Box` cannot be used here. - cache: NonNull>, + cache: NonNull>, _phantom: PhantomData O>, } @@ -513,9 +515,9 @@ impl ExactSizeIterator for PostgresReadStreamGuards PostgresReadStream +impl<'b, O: Opaque, I: Iterator> PostgresReadStream<'b, O, I> where - I::Item: Fetch, + I::Item: Fetch<'b>, { fn read>(&mut self, fetch: L) -> PostgresReadStreamGuards { PostgresReadStreamGuards { @@ -526,13 +528,13 @@ where } } -impl<'r, O: Opaque, I: Iterator> ReadStream<'r> for PostgresReadStream +impl<'b, O: Opaque, I: Iterator> ReadStream<'b> for PostgresReadStream<'b, O, I> where - I::Item: Fetch, + I::Item: Fetch<'b>, { type Relation = PostgresRelation; - type Guards = PostgresReadStreamGuards; + type Guards = PostgresReadStreamGuards>::Iter>; type Item = I::Item; @@ -584,7 +586,7 @@ where } } -impl Drop for PostgresReadStream { +impl<'b, O: Opaque, I: Iterator> Drop for PostgresReadStream<'b, O, I> { fn drop(&mut self) { unsafe { let _ = std::mem::take(self.cache.as_mut()); @@ -598,34 +600,34 @@ impl Drop for PostgresReadStream { } impl RelationReadStreamTypes for PostgresRelation { - type ReadStream<'r, I: Iterator> - = PostgresReadStream + type ReadStream<'b, I: Iterator> + = PostgresReadStream<'b, O, I> where - I::Item: Fetch; + I::Item: Fetch<'b>; } impl RelationReadStream for PostgresRelation { #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] - fn read_stream(&self, _iter: I, _hints: Hints) -> Self::ReadStream<'_, I> + fn read_stream<'b, I: Iterator>(&'b self, _iter: I, _hints: Hints) -> Self::ReadStream<'b, I> where - I::Item: Fetch, + I::Item: Fetch<'b>, { panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); } #[cfg(any(feature = "pg17", feature = "pg18"))] - fn read_stream(&self, iter: I, hints: Hints) -> Self::ReadStream<'_, I> + fn read_stream<'b, I: Iterator>(&'b self, iter: I, hints: Hints) -> Self::ReadStream<'b, I> where - I::Item: Fetch, + I::Item: Fetch<'b>, { #[pgrx::pg_guard] - unsafe extern "C-unwind" fn callback( + unsafe extern "C-unwind" fn callback<'b, I: Iterator>( _stream: *mut pgrx::pg_sys::ReadStream, callback_private_data: *mut core::ffi::c_void, _per_buffer_data: *mut core::ffi::c_void, ) -> pgrx::pg_sys::BlockNumber where - I::Item: Fetch, + I::Item: Fetch<'b>, { unsafe { use pgrx::pg_sys::InvalidBlockNumber; @@ -637,6 +639,7 @@ impl RelationReadStream for PostgresRelation { window: VecDeque::new(), tail: VecDeque::new(), iter: Some(iter), + _phantom: PhantomData, })); let raw = unsafe { use pgrx::pg_sys::{ diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs index 5114555b..ab9dff83 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/algo.rs @@ -197,25 +197,25 @@ impl RandomProject for VectBorrowed<'_, f16> { } #[derive(Debug)] -pub struct MakePlainPrefetcher<'r, R> { - pub index: &'r R, +pub struct MakePlainPrefetcher<'b, R> { + pub index: &'b R, } -impl<'r, R> Clone for MakePlainPrefetcher<'r, R> { +impl<'b, R> Clone for MakePlainPrefetcher<'b, R> { fn clone(&self) -> Self { Self { index: self.index } } } -impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakePlainPrefetcher<'r, R> { +impl<'b, R: RelationRead> PrefetcherSequenceFamily<'b, R> for MakePlainPrefetcher<'b, R> { type P - = PlainPrefetcher<'r, R, S> + = PlainPrefetcher<'b, R, S> where - S::Item: Fetch; + S::Item: Fetch<'b>; fn prefetch(&mut self, seq: S) -> Self::P where - S::Item: Fetch, + S::Item: Fetch<'b>, { PlainPrefetcher::new(self.index, seq) } @@ -226,27 +226,27 @@ impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakePlainPrefetche } #[derive(Debug)] -pub struct MakeSimplePrefetcher<'r, R> { - pub index: &'r R, +pub struct MakeSimplePrefetcher<'b, R> { + pub index: &'b R, } -impl<'r, R> Clone for MakeSimplePrefetcher<'r, R> { +impl<'b, R> Clone for MakeSimplePrefetcher<'b, R> { fn clone(&self) -> Self { Self { index: self.index } } } -impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> - for MakeSimplePrefetcher<'r, R> +impl<'b, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'b, R> + for MakeSimplePrefetcher<'b, R> { type P - = SimplePrefetcher<'r, R, S> + = SimplePrefetcher<'b, R, S> where - S::Item: Fetch; + S::Item: Fetch<'b>; fn prefetch(&mut self, seq: S) -> Self::P where - S::Item: Fetch, + S::Item: Fetch<'b>, { SimplePrefetcher::new(self.index, seq) } @@ -257,12 +257,12 @@ impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> } #[derive(Debug)] -pub struct MakeStreamPrefetcher<'r, R> { - pub index: &'r R, +pub struct MakeStreamPrefetcher<'b, R> { + pub index: &'b R, pub hints: Hints, } -impl<'r, R> Clone for MakeStreamPrefetcher<'r, R> { +impl<'b, R> Clone for MakeStreamPrefetcher<'b, R> { fn clone(&self) -> Self { Self { index: self.index, @@ -271,17 +271,17 @@ impl<'r, R> Clone for MakeStreamPrefetcher<'r, R> { } } -impl<'r, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'r, R> - for MakeStreamPrefetcher<'r, R> +impl<'b, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'b, R> + for MakeStreamPrefetcher<'b, R> { type P - = StreamPrefetcher<'r, R, S> + = StreamPrefetcher<'b, R, S> where - S::Item: Fetch; + S::Item: Fetch<'b>; fn prefetch(&mut self, seq: S) -> Self::P where - S::Item: Fetch, + S::Item: Fetch<'b>, { StreamPrefetcher::new(self.index, seq, self.hints.clone()) } diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index 3eb31a04..a4b023c9 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -16,6 +16,7 @@ pub mod am_build; use crate::index::fetcher::*; use crate::index::gucs; +use crate::index::scanners::SearchBuilder; use crate::index::storage::PostgresRelation; use crate::index::vchordg::opclass::opfamily; use crate::index::vchordg::scanners::*; diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs index f3d597f9..a0124f71 100644 --- a/src/index/vchordg/scanners/default.rs +++ b/src/index/vchordg/scanners/default.rs @@ -12,11 +12,12 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::{Fetcher, Io, SearchBuilder, SearchOptions}; -use crate::index::fetcher::pointer_to_kv; +use crate::index::fetcher::{Fetcher, pointer_to_kv}; use crate::index::opclass::Sphere; +use crate::index::scanners::{Io, SearchBuilder}; use crate::index::vchordg::algo::*; use crate::index::vchordg::opclass::Opfamily; +use crate::index::vchordg::scanners::SearchOptions; use algo::accessor::{Dot, L2S}; use algo::*; use distance::Distance; @@ -35,6 +36,10 @@ pub struct DefaultBuilder { } impl SearchBuilder for DefaultBuilder { + type Options = SearchOptions; + + type Opfamily = Opfamily; + type Opaque = vchordg::Opaque; fn new(opfamily: Opfamily) -> Self { @@ -68,13 +73,13 @@ impl SearchBuilder for DefaultBuilder { } } - fn build<'a, R>( + fn build<'b, R>( self, - index: &'a R, + index: &'b R, options: SearchOptions, - _fetcher: impl Fetcher + 'a, - bump: &'a impl Bump, - ) -> Box + 'a> + _fetcher: impl Fetcher + 'b, + bump: &'b impl Bump, + ) -> Box + 'b> where R: RelationRead + RelationPrefetch + RelationReadStream, R::Page: Page, @@ -117,11 +122,12 @@ impl SearchBuilder for DefaultBuilder { (VectorKind::Vecf32, DistanceKind::L2S) => { type Op = operator::Op, L2S>; let unprojected = if let OwnedVector::Vecf32(vector) = vector { - bump.alloc(vector) + bump.alloc_any(vector) } else { unreachable!() }; - let projected = bump.alloc(RandomProject::project(unprojected.as_borrowed())); + let projected = + bump.alloc_any(RandomProject::project(unprojected.as_borrowed())); match (options.io_search, options.io_rerank) { (Io::Plain, Io::Plain) => search::<_, Op>( index, @@ -209,11 +215,12 @@ impl SearchBuilder for DefaultBuilder { (VectorKind::Vecf16, DistanceKind::L2S) => { type Op = operator::Op, L2S>; let unprojected = if let OwnedVector::Vecf16(vector) = vector { - bump.alloc(vector) + bump.alloc_any(vector) } else { unreachable!() }; - let projected = bump.alloc(RandomProject::project(unprojected.as_borrowed())); + let projected = + bump.alloc_any(RandomProject::project(unprojected.as_borrowed())); match (options.io_search, options.io_rerank) { (Io::Plain, Io::Plain) => search::<_, Op>( index, @@ -301,11 +308,12 @@ impl SearchBuilder for DefaultBuilder { (VectorKind::Vecf32, DistanceKind::Dot) => { type Op = operator::Op, Dot>; let unprojected = if let OwnedVector::Vecf32(vector) = vector { - bump.alloc(vector) + bump.alloc_any(vector) } else { unreachable!() }; - let projected = bump.alloc(RandomProject::project(unprojected.as_borrowed())); + let projected = + bump.alloc_any(RandomProject::project(unprojected.as_borrowed())); match (options.io_search, options.io_rerank) { (Io::Plain, Io::Plain) => search::<_, Op>( index, @@ -393,11 +401,12 @@ impl SearchBuilder for DefaultBuilder { (VectorKind::Vecf16, DistanceKind::Dot) => { type Op = operator::Op, Dot>; let unprojected = if let OwnedVector::Vecf16(vector) = vector { - bump.alloc(vector) + bump.alloc_any(vector) } else { unreachable!() }; - let projected = bump.alloc(RandomProject::project(unprojected.as_borrowed())); + let projected = + bump.alloc_any(RandomProject::project(unprojected.as_borrowed())); match (options.io_search, options.io_rerank) { (Io::Plain, Io::Plain) => search::<_, Op>( index, diff --git a/src/index/vchordg/scanners/mod.rs b/src/index/vchordg/scanners/mod.rs index 264c29f2..a9545ccc 100644 --- a/src/index/vchordg/scanners/mod.rs +++ b/src/index/vchordg/scanners/mod.rs @@ -14,48 +14,13 @@ mod default; -use super::opclass::Opfamily; -use crate::index::fetcher::Fetcher; -use algo::{Bump, Page, RelationPrefetch, RelationRead, RelationReadStream}; -use pgrx::pg_sys::Datum; - pub use default::DefaultBuilder; -#[derive(Debug, Clone, Copy)] -pub enum Io { - Plain, - Simple, - #[cfg_attr( - any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"), - expect(dead_code) - )] - Stream, -} - #[derive(Debug)] pub struct SearchOptions { pub ef_search: u32, pub beam_search: u32, pub max_scan_tuples: Option, - pub io_search: Io, - pub io_rerank: Io, -} - -pub trait SearchBuilder: 'static { - type Opaque: Copy; - - fn new(opfamily: Opfamily) -> Self; - - unsafe fn add(&mut self, strategy: u16, datum: Option); - - fn build<'a, R>( - self, - relation: &'a R, - options: SearchOptions, - fetcher: impl Fetcher + 'a, - bump: &'a impl Bump, - ) -> Box + 'a> - where - R: RelationRead + RelationPrefetch + RelationReadStream, - R::Page: Page; + pub io_search: crate::index::scanners::Io, + pub io_rerank: crate::index::scanners::Io, } diff --git a/src/index/vchordrq/algo.rs b/src/index/vchordrq/algo.rs index 2ac6c538..aaf030b4 100644 --- a/src/index/vchordrq/algo.rs +++ b/src/index/vchordrq/algo.rs @@ -275,25 +275,25 @@ impl RandomProject for VectBorrowed<'_, f16> { } #[derive(Debug)] -pub struct MakeH1PlainPrefetcherForInsertion<'r, R> { - pub index: &'r R, +pub struct MakeH1PlainPrefetcherForInsertion<'b, R> { + pub index: &'b R, } -impl<'r, R> Clone for MakeH1PlainPrefetcherForInsertion<'r, R> { +impl<'b, R> Clone for MakeH1PlainPrefetcherForInsertion<'b, R> { fn clone(&self) -> Self { Self { index: self.index } } } -impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcherForInsertion<'r, R> { +impl<'b, R: RelationRead> PrefetcherHeapFamily<'b, R> for MakeH1PlainPrefetcherForInsertion<'b, R> { type P - = PlainPrefetcher<'r, R, FastHeap> + = PlainPrefetcher<'b, R, FastHeap> where - T: Ord + Fetch + 'r; + T: Ord + Fetch<'b>; fn prefetch(&mut self, seq: Vec) -> Self::P where - T: Ord + Fetch + 'r, + T: Ord + Fetch<'b>, { PlainPrefetcher::new(self.index, FastHeap::from(seq)) } @@ -304,25 +304,25 @@ impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcherF } #[derive(Debug)] -pub struct MakeH1PlainPrefetcher<'r, R> { - pub index: &'r R, +pub struct MakeH1PlainPrefetcher<'b, R> { + pub index: &'b R, } -impl<'r, R> Clone for MakeH1PlainPrefetcher<'r, R> { +impl<'b, R> Clone for MakeH1PlainPrefetcher<'b, R> { fn clone(&self) -> Self { Self { index: self.index } } } -impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcher<'r, R> { +impl<'b, R: RelationRead> PrefetcherHeapFamily<'b, R> for MakeH1PlainPrefetcher<'b, R> { type P - = PlainPrefetcher<'r, R, BinaryHeap> + = PlainPrefetcher<'b, R, BinaryHeap> where - T: Ord + Fetch + 'r; + T: Ord + Fetch<'b>; fn prefetch(&mut self, seq: Vec) -> Self::P where - T: Ord + Fetch + 'r, + T: Ord + Fetch<'b>, { PlainPrefetcher::new(self.index, BinaryHeap::from(seq)) } @@ -333,25 +333,25 @@ impl<'r, R: RelationRead> PrefetcherHeapFamily<'r, R> for MakeH1PlainPrefetcher< } #[derive(Debug)] -pub struct MakeH0PlainPrefetcher<'r, R> { - pub index: &'r R, +pub struct MakeH0PlainPrefetcher<'b, R> { + pub index: &'b R, } -impl<'r, R> Clone for MakeH0PlainPrefetcher<'r, R> { +impl<'b, R> Clone for MakeH0PlainPrefetcher<'b, R> { fn clone(&self) -> Self { Self { index: self.index } } } -impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakeH0PlainPrefetcher<'r, R> { +impl<'b, R: RelationRead> PrefetcherSequenceFamily<'b, R> for MakeH0PlainPrefetcher<'b, R> { type P - = PlainPrefetcher<'r, R, S> + = PlainPrefetcher<'b, R, S> where - S::Item: Fetch; + S::Item: Fetch<'b>; fn prefetch(&mut self, seq: S) -> Self::P where - S::Item: Fetch, + S::Item: Fetch<'b>, { PlainPrefetcher::new(self.index, seq) } @@ -362,27 +362,27 @@ impl<'r, R: RelationRead> PrefetcherSequenceFamily<'r, R> for MakeH0PlainPrefetc } #[derive(Debug)] -pub struct MakeH0SimplePrefetcher<'r, R> { - pub index: &'r R, +pub struct MakeH0SimplePrefetcher<'b, R> { + pub index: &'b R, } -impl<'r, R> Clone for MakeH0SimplePrefetcher<'r, R> { +impl<'b, R> Clone for MakeH0SimplePrefetcher<'b, R> { fn clone(&self) -> Self { Self { index: self.index } } } -impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> - for MakeH0SimplePrefetcher<'r, R> +impl<'b, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'b, R> + for MakeH0SimplePrefetcher<'b, R> { type P - = SimplePrefetcher<'r, R, S> + = SimplePrefetcher<'b, R, S> where - S::Item: Fetch; + S::Item: Fetch<'b>; fn prefetch(&mut self, seq: S) -> Self::P where - S::Item: Fetch, + S::Item: Fetch<'b>, { SimplePrefetcher::new(self.index, seq) } @@ -393,12 +393,12 @@ impl<'r, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'r, R> } #[derive(Debug)] -pub struct MakeH0StreamPrefetcher<'r, R> { - pub index: &'r R, +pub struct MakeH0StreamPrefetcher<'b, R> { + pub index: &'b R, pub hints: Hints, } -impl<'r, R> Clone for MakeH0StreamPrefetcher<'r, R> { +impl<'b, R> Clone for MakeH0StreamPrefetcher<'b, R> { fn clone(&self) -> Self { Self { index: self.index, @@ -407,17 +407,17 @@ impl<'r, R> Clone for MakeH0StreamPrefetcher<'r, R> { } } -impl<'r, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'r, R> - for MakeH0StreamPrefetcher<'r, R> +impl<'b, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'b, R> + for MakeH0StreamPrefetcher<'b, R> { type P - = StreamPrefetcher<'r, R, S> + = StreamPrefetcher<'b, R, S> where - S::Item: Fetch; + S::Item: Fetch<'b>; fn prefetch(&mut self, seq: S) -> Self::P where - S::Item: Fetch, + S::Item: Fetch<'b>, { StreamPrefetcher::new(self.index, seq, self.hints.clone()) } diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 5708f167..b1c3d58b 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -16,6 +16,7 @@ pub mod am_build; use crate::index::fetcher::*; use crate::index::gucs; +use crate::index::scanners::SearchBuilder; use crate::index::storage::PostgresRelation; use crate::index::vchordrq::opclass::{Opfamily, opfamily}; use crate::index::vchordrq::scanners::*; diff --git a/src/index/vchordrq/filter.rs b/src/index/vchordrq/filter.rs new file mode 100644 index 00000000..1fec85bb --- /dev/null +++ b/src/index/vchordrq/filter.rs @@ -0,0 +1,55 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use algo::Sequence; + +pub struct Filter { + sequence: S, + predicate: P, +} + +impl Sequence for Filter +where + S: Sequence, + P: FnMut(&S::Item) -> bool, +{ + type Item = S::Item; + + type Inner = S::Inner; + + fn next(&mut self) -> Option { + while !(self.predicate)(self.sequence.peek()?) { + let _ = self.sequence.next(); + } + self.sequence.next() + } + + fn peek(&mut self) -> Option<&Self::Item> { + while !(self.predicate)(self.sequence.peek()?) { + let _ = self.sequence.next(); + } + self.sequence.peek() + } + + fn into_inner(self) -> Self::Inner { + self.sequence.into_inner() + } +} + +pub fn filter(sequence: S, predicate: P) -> Filter { + Filter { + sequence, + predicate, + } +} diff --git a/src/index/vchordrq/mod.rs b/src/index/vchordrq/mod.rs index d897f7d1..78d81435 100644 --- a/src/index/vchordrq/mod.rs +++ b/src/index/vchordrq/mod.rs @@ -14,6 +14,7 @@ pub mod algo; pub mod am; +pub mod filter; pub mod opclass; pub mod scanners; pub mod types; diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index f6b63d5c..2d4739bf 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -12,11 +12,13 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::{Io, SearchBuilder, SearchOptions, filter}; use crate::index::fetcher::*; use crate::index::opclass::Sphere; +use crate::index::scanners::{Io, SearchBuilder}; use crate::index::vchordrq::algo::*; +use crate::index::vchordrq::filter::filter; use crate::index::vchordrq::opclass::Opfamily; +use crate::index::vchordrq::scanners::SearchOptions; use algo::accessor::{Dot, L2S}; use algo::prefetcher::*; use algo::*; @@ -37,6 +39,10 @@ pub struct DefaultBuilder { } impl SearchBuilder for DefaultBuilder { + type Options = SearchOptions; + + type Opfamily = Opfamily; + type Opaque = vchordrq::Opaque; fn new(opfamily: Opfamily) -> Self { @@ -70,13 +76,13 @@ impl SearchBuilder for DefaultBuilder { } } - fn build<'a, R>( + fn build<'b, R>( self, - index: &'a R, + index: &'b R, options: SearchOptions, - mut fetcher: impl Fetcher + 'a, - bump: &'a impl Bump, - ) -> Box + 'a> + mut fetcher: impl Fetcher + 'b, + bump: &'b impl Bump, + ) -> Box + 'b> where R: RelationRead + RelationPrefetch + RelationReadStream, R::Page: Page, diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 787bbb1d..3000bc06 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -12,11 +12,12 @@ // // Copyright (c) 2025 TensorChord Inc. -use super::{SearchBuilder, SearchOptions}; use crate::index::fetcher::*; +use crate::index::scanners::{Io, SearchBuilder}; use crate::index::vchordrq::algo::*; +use crate::index::vchordrq::filter::filter; use crate::index::vchordrq::opclass::Opfamily; -use crate::index::vchordrq::scanners::{Io, filter}; +use crate::index::vchordrq::scanners::SearchOptions; use algo::accessor::Dot; use algo::prefetcher::*; use algo::*; @@ -37,6 +38,10 @@ pub struct MaxsimBuilder { } impl SearchBuilder for MaxsimBuilder { + type Options = SearchOptions; + + type Opfamily = Opfamily; + type Opaque = vchordrq::Opaque; fn new(opfamily: Opfamily) -> Self { @@ -60,13 +65,13 @@ impl SearchBuilder for MaxsimBuilder { } } - fn build<'a, R>( + fn build<'b, R>( self, - index: &'a R, + index: &'b R, options: SearchOptions, - mut fetcher: impl Fetcher + 'a, - bump: &'a impl Bump, - ) -> Box + 'a> + mut fetcher: impl Fetcher + 'b, + bump: &'b impl Bump, + ) -> Box + 'b> where R: RelationRead + RelationPrefetch + RelationReadStream, R::Page: Page, diff --git a/src/index/vchordrq/scanners/mod.rs b/src/index/vchordrq/scanners/mod.rs index a0340a3e..a1c7e8f1 100644 --- a/src/index/vchordrq/scanners/mod.rs +++ b/src/index/vchordrq/scanners/mod.rs @@ -15,26 +15,11 @@ mod default; mod maxsim; -use crate::index::fetcher::Fetcher; - -use super::opclass::Opfamily; -use algo::{Bump, Page, RelationPrefetch, RelationRead, RelationReadStream, Sequence}; -use pgrx::pg_sys::Datum; +use crate::index::scanners::Io; pub use default::DefaultBuilder; pub use maxsim::MaxsimBuilder; -#[derive(Debug, Clone, Copy)] -pub enum Io { - Plain, - Simple, - #[cfg_attr( - any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"), - expect(dead_code) - )] - Stream, -} - #[derive(Debug)] pub struct SearchOptions { pub epsilon: f32, @@ -46,62 +31,3 @@ pub struct SearchOptions { pub io_rerank: Io, pub prefilter: bool, } - -pub trait SearchBuilder: 'static { - type Opaque: Copy; - - fn new(opfamily: Opfamily) -> Self; - - unsafe fn add(&mut self, strategy: u16, datum: Option); - - fn build<'a, R>( - self, - relation: &'a R, - options: SearchOptions, - fetcher: impl Fetcher + 'a, - bump: &'a impl Bump, - ) -> Box + 'a> - where - R: RelationRead + RelationPrefetch + RelationReadStream, - R::Page: Page; -} - -pub struct Filter { - sequence: S, - predicate: P, -} - -impl Sequence for Filter -where - S: Sequence, - P: FnMut(&S::Item) -> bool, -{ - type Item = S::Item; - - type Inner = S::Inner; - - fn next(&mut self) -> Option { - while !(self.predicate)(self.sequence.peek()?) { - let _ = self.sequence.next(); - } - self.sequence.next() - } - - fn peek(&mut self) -> Option<&Self::Item> { - while !(self.predicate)(self.sequence.peek()?) { - let _ = self.sequence.next(); - } - self.sequence.peek() - } - - fn into_inner(self) -> Self::Inner { - self.sequence.into_inner() - } -} - -pub fn filter(sequence: S, predicate: P) -> Filter { - Filter { - sequence, - predicate, - } -} From c4f7ccc95fe15e11840675e1e54662bc632f745f Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 11 Aug 2025 17:44:04 +0800 Subject: [PATCH 193/324] fix: free memory in maxsim search (#310) Signed-off-by: usamoi --- crates/algo/src/lib.rs | 7 -- src/index/vchordg/algo.rs | 2 +- src/index/vchordg/am/mod.rs | 2 +- src/index/vchordg/scanners/default.rs | 107 ++++++++++++++------------ src/index/vchordrq/algo.rs | 2 +- src/index/vchordrq/am/mod.rs | 2 +- 6 files changed, 61 insertions(+), 61 deletions(-) diff --git a/crates/algo/src/lib.rs b/crates/algo/src/lib.rs index e28c423f..39cfc215 100644 --- a/crates/algo/src/lib.rs +++ b/crates/algo/src/lib.rs @@ -139,8 +139,6 @@ pub trait Bump: 'static { fn alloc(&self, value: T) -> &mut T; #[allow(clippy::mut_from_ref)] fn alloc_slice(&self, slice: &[T]) -> &mut [T]; - #[allow(clippy::mut_from_ref)] - fn alloc_any(&self, value: T) -> &mut T; } impl Bump for bumpalo::Bump { @@ -153,11 +151,6 @@ impl Bump for bumpalo::Bump { fn alloc_slice(&self, slice: &[T]) -> &mut [T] { self.alloc_slice_copy(slice) } - - #[inline] - fn alloc_any(&self, value: T) -> &mut T { - self.alloc(value) - } } pub type BorrowedIter<'b> = small_iter::borrowed::Iter<'b, u32, 1>; diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs index ab9dff83..528a64dc 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/algo.rs @@ -117,7 +117,7 @@ where R: RelationRead + RelationWrite + RelationReadStream, R::Page: Page, { - let bump = bumpalo::Bump::with_capacity(2 << 20); + let bump = bumpalo::Bump::new(); let make_vertex_plain_prefetcher = MakePlainPrefetcher { index }; let make_vector_plain_prefetcher = MakePlainPrefetcher { index }; match (vector, opfamily.distance_kind()) { diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index a4b023c9..dc6a233e 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -316,7 +316,7 @@ pub unsafe extern "C-unwind" fn ambeginscan( let scanner: Scanner = Scanner { hack: None, scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), - bump: Box::new(bumpalo::Bump::with_capacity(2 << 20)), + bump: Box::new(bumpalo::Bump::new()), }; unsafe { (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs index a0124f71..67562fad 100644 --- a/src/index/vchordg/scanners/default.rs +++ b/src/index/vchordg/scanners/default.rs @@ -26,8 +26,7 @@ use std::num::NonZero; use vchordg::operator::{self}; use vchordg::types::{DistanceKind, OwnedVector, VectorKind}; use vchordg::*; -use vector::VectorOwned; -use vector::vect::VectOwned; +use vector::vect::{VectBorrowed, VectOwned}; pub struct DefaultBuilder { opfamily: Opfamily, @@ -122,16 +121,18 @@ impl SearchBuilder for DefaultBuilder { (VectorKind::Vecf32, DistanceKind::L2S) => { type Op = operator::Op, L2S>; let unprojected = if let OwnedVector::Vecf32(vector) = vector { - bump.alloc_any(vector) + VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() }; - let projected = - bump.alloc_any(RandomProject::project(unprojected.as_borrowed())); + let projected = { + let projected = RandomProject::project(unprojected); + VectBorrowed::new(bump.alloc_slice(projected.slice())) + }; match (options.io_search, options.io_rerank) { (Io::Plain, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -140,7 +141,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Plain, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -149,7 +150,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Plain, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -158,7 +159,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -167,7 +168,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -176,7 +177,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -185,7 +186,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -194,7 +195,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -203,7 +204,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -215,16 +216,18 @@ impl SearchBuilder for DefaultBuilder { (VectorKind::Vecf16, DistanceKind::L2S) => { type Op = operator::Op, L2S>; let unprojected = if let OwnedVector::Vecf16(vector) = vector { - bump.alloc_any(vector) + VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() }; - let projected = - bump.alloc_any(RandomProject::project(unprojected.as_borrowed())); + let projected = { + let projected = RandomProject::project(unprojected); + VectBorrowed::new(bump.alloc_slice(projected.slice())) + }; match (options.io_search, options.io_rerank) { (Io::Plain, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -233,7 +236,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Plain, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -242,7 +245,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Plain, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -251,7 +254,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -260,7 +263,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -269,7 +272,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -278,7 +281,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -287,7 +290,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -296,7 +299,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -308,16 +311,18 @@ impl SearchBuilder for DefaultBuilder { (VectorKind::Vecf32, DistanceKind::Dot) => { type Op = operator::Op, Dot>; let unprojected = if let OwnedVector::Vecf32(vector) = vector { - bump.alloc_any(vector) + VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() }; - let projected = - bump.alloc_any(RandomProject::project(unprojected.as_borrowed())); + let projected = { + let projected = RandomProject::project(unprojected); + VectBorrowed::new(bump.alloc_slice(projected.slice())) + }; match (options.io_search, options.io_rerank) { (Io::Plain, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -326,7 +331,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Plain, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -335,7 +340,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Plain, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -344,7 +349,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -353,7 +358,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -362,7 +367,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -371,7 +376,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -380,7 +385,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -389,7 +394,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -401,16 +406,18 @@ impl SearchBuilder for DefaultBuilder { (VectorKind::Vecf16, DistanceKind::Dot) => { type Op = operator::Op, Dot>; let unprojected = if let OwnedVector::Vecf16(vector) = vector { - bump.alloc_any(vector) + VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() }; - let projected = - bump.alloc_any(RandomProject::project(unprojected.as_borrowed())); + let projected = { + let projected = RandomProject::project(unprojected); + VectBorrowed::new(bump.alloc_slice(projected.slice())) + }; match (options.io_search, options.io_rerank) { (Io::Plain, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -419,7 +426,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Plain, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -428,7 +435,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Plain, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -437,7 +444,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -446,7 +453,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -455,7 +462,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Simple, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -464,7 +471,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Plain) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -473,7 +480,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Simple) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, @@ -482,7 +489,7 @@ impl SearchBuilder for DefaultBuilder { ), (Io::Stream, Io::Stream) => search::<_, Op>( index, - projected.as_borrowed(), + projected, options.ef_search, options.beam_search, bump, diff --git a/src/index/vchordrq/algo.rs b/src/index/vchordrq/algo.rs index aaf030b4..508537fb 100644 --- a/src/index/vchordrq/algo.rs +++ b/src/index/vchordrq/algo.rs @@ -158,7 +158,7 @@ pub fn insert( R: RelationRead + RelationWrite, R::Page: Page, { - let bump = bumpalo::Bump::with_capacity(2 << 20); + let bump = bumpalo::Bump::new(); let make_h1_plain_prefetcher = MakeH1PlainPrefetcherForInsertion { index }; match (vector, opfamily.distance_kind()) { (OwnedVector::Vecf32(vector), DistanceKind::L2S) => { diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index b1c3d58b..b8603c35 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -396,7 +396,7 @@ pub unsafe extern "C-unwind" fn ambeginscan( let scanner: Scanner = Scanner { hack: None, scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), - bump: Box::new(bumpalo::Bump::with_capacity(2 << 20)), + bump: Box::new(bumpalo::Bump::new()), }; unsafe { (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); From 1dc4b35247f789bf1c7ca1ecb408b32fc08f8c26 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 12 Aug 2025 13:39:36 +0800 Subject: [PATCH 194/324] chore: upload 0.5.0 schema scripts (#311) job: +psql_macos job: +psql_windows job: +psql_alpine Signed-off-by: usamoi --- .github/workflows/check.yml | 20 +- .github/workflows/release.yml | 8 +- Makefile | 16 +- sql/install/vchord--0.5.0.sql | 752 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.4.3--0.5.0.sql | 151 ++++++ tools/dump.sh | 9 +- 6 files changed, 934 insertions(+), 22 deletions(-) create mode 100644 sql/install/vchord--0.5.0.sql create mode 100644 sql/upgrade/vchord--0.4.3--0.5.0.sql diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 2525bbf1..962c8bdb 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -196,6 +196,7 @@ jobs: sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} echo PGRX_PG_CONFIG_PATH=pg_config >> $GITHUB_ENV + echo PG_CONFIG=pg_config >> $GITHUB_ENV sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector echo "local all all trust" | sudo tee /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf @@ -219,8 +220,8 @@ jobs: - name: Install run: | - cargo run -p make -- build -o ./build/raw --profile dev - sudo make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install + make PG_CONFIG=$PG_CONFIG PROFILE=dev build + sudo make PG_CONFIG=$PG_CONFIG install - name: Service run: | @@ -280,6 +281,7 @@ jobs: $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' brew services stop postgresql@${{ matrix.version }} echo PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV + echo PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.0.tar.gz | tar -xz -C ~/pgvector-install @@ -300,8 +302,8 @@ jobs: - name: Install run: | - cargo run -p make -- build -o ./build/raw --profile dev - sudo make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install + make PG_CONFIG=$PG_CONFIG PROFILE=dev build + sudo make PG_CONFIG=$PG_CONFIG install - name: Service run: | @@ -374,6 +376,7 @@ jobs: Invoke-WebRequest -Uri $postgresql_url -OutFile "$env:TEMP\postgresql-install.zip" Expand-Archive -Path "$env:TEMP\postgresql-install.zip" -DestinationPath "D:\postgresql-install" -Force Add-Content -Path $env:GITHUB_ENV -Value "PGRX_PG_CONFIG_PATH=D:\postgresql-install\pgsql\bin\pg_config.exe" + Add-Content -Path $env:GITHUB_ENV -Value "PG_CONFIG=D:\postgresql-install\pgsql\bin\pg_config.exe" D:\postgresql-install\pgsql\bin\initdb.exe -D D:\postgresql-install\pgsql\data -U postgres D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data D:\postgresql-install\pgsql\bin\createuser.exe -U postgres -s -r $env:USERNAME @@ -404,8 +407,8 @@ jobs: - name: Install run: | - cargo run -p make -- build -o ./build/raw --profile dev - make PG_CONFIG="$env:PGRX_PG_CONFIG_PATH" install + make PG_CONFIG=$env:PG_CONFIG PROFILE=dev build + make PG_CONFIG=$env:PG_CONFIG install - name: Service run: | @@ -468,6 +471,7 @@ jobs: curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain 1.89-beta echo PGRX_PG_CONFIG_PATH=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV + echo PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql/data @@ -509,8 +513,8 @@ jobs: - name: Install run: | . "$HOME/.cargo/env" - cargo run -p make -- build -o ./build/raw --profile dev - make PG_CONFIG=${PGRX_PG_CONFIG_PATH} install + make PG_CONFIG=$PG_CONFIG PROFILE=dev build + sudo make PG_CONFIG=$PG_CONFIG install - name: Service run: | diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 0b3bd44b..2011987c 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -65,6 +65,7 @@ jobs: sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} echo PGRX_PG_CONFIG_PATH=pg_config >> $GITHUB_ENV + echo PG_CONFIG=pg_config >> $GITHUB_ENV sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector @@ -87,16 +88,15 @@ jobs: ARCH: ${{ matrix.arch }} PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} run: | - cargo run -p make -- build -o ./build/raw --profile release + make PG_CONFIG=$PG_CONFIG build (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) - mkdir -p ./build/deb + make DESTDIR="./build/deb" install mkdir -p ./build/deb/DEBIAN - mkdir -p ./build/deb$(pg_config --pkglibdir) && cp -r ./build/raw/pkglibdir/. ./build/deb$(pg_config --pkglibdir) - mkdir -p ./build/deb$(pg_config --sharedir) && cp -r ./build/raw/sharedir/. ./build/deb$(pg_config --sharedir) echo "Package: postgresql-${VERSION}-vchord Version: ${SEMVER}-1 + Depends: postgresql-${VERSION}, libgcc-s1, libc6 (>= 2.35) Section: database Priority: optional Architecture: ${PLATFORM} diff --git a/Makefile b/Makefile index b842fa9a..920c6c30 100644 --- a/Makefile +++ b/Makefile @@ -1,17 +1,21 @@ PG_CONFIG ?= pg_config +PKGLIBDIR := $(shell $(PG_CONFIG) --pkglibdir) +SHAREDIR := $(shell $(PG_CONFIG) --sharedir) +PROFILE ?= release +MKDIR ?= mkdir +CP ?= cp .PHONY: all build install uninstall all: build build: - PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p make -- build -o ./build/raw + PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p make -- build --output ./build/raw --profile $(PROFILE) install: - cp -r ./build/raw/pkglibdir/. $(shell $(PG_CONFIG) --pkglibdir) - cp -r ./build/raw/sharedir/. $(shell $(PG_CONFIG) --sharedir) + $(MKDIR) -p $(DESTDIR)$(PKGLIBDIR) $(DESTDIR)$(SHAREDIR) && \ + $(CP) -r ./build/raw/pkglibdir/. $(DESTDIR)$(PKGLIBDIR) && \ + $(CP) -r ./build/raw/sharedir/. $(DESTDIR)$(SHAREDIR) uninstall: - rm -f $(shell find $(shell $(PG_CONFIG) --pkglibdir) -type f -name 'vchord.*') - rm -f $(shell find $(shell $(PG_CONFIG) --sharedir)/extension -type f -name 'vchord.*') - rm -f $(shell find $(shell $(PG_CONFIG) --sharedir)/extension -type f -name 'vchord--*.sql') + $(RM) $(wildcard $(DESTDIR)$(PKGLIBDIR)/vchord.*) $(wildcard $(DESTDIR)$(SHAREDIR)/extension/vchord.*) $(wildcard $(DESTDIR)$(SHAREDIR)/extension/vchord--*.sql) diff --git a/sql/install/vchord--0.5.0.sql b/sql/install/vchord--0.5.0.sql new file mode 100644 index 00000000..173d6807 --- /dev/null +++ b/sql/install/vchord--0.5.0.sql @@ -0,0 +1,752 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:45 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:77 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:77 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:41 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:46 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.4.3--0.5.0.sql b/sql/upgrade/vchord--0.4.3--0.5.0.sql new file mode 100644 index 00000000..ee5bdbcf --- /dev/null +++ b/sql/upgrade/vchord--0.4.3--0.5.0.sql @@ -0,0 +1,151 @@ +-- internal changes + +CREATE FUNCTION _vchordg_support_halfvec_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_halfvec_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_halfvec_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_vector_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_vector_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_vector_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; + +-- List of functions + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); diff --git a/tools/dump.sh b/tools/dump.sh index 110884fd..04e28daf 100755 --- a/tools/dump.sh +++ b/tools/dump.sh @@ -31,14 +31,15 @@ $(dirname "$0")/dump-codegen.sh | gcc -I $(pg_config --includedir-server) -fPIC sql=$(mktemp) echo "BEGIN;" >> $sql -echo "CREATE SCHEMA public;" >> $sql -echo "CREATE EXTENSION vector;" >> $sql +echo "CREATE SCHEMA vchord;" >> $sql +echo "SET LOCAL search_path TO vchord,public;" >> $sql cat ${f[@]} \ | grep -v '^\\' \ + | sed "s|@extschema@|vchord|g" \ | sed "s|MODULE_PATHNAME|$so|g" \ >> $sql echo "END;" >> $sql psql -d vchord -f $sql 1>&2 -pg_dump -d vchord -psql -d vchord -c "DROP SCHEMA IF EXISTS public CASCADE;" 1>&2 \ No newline at end of file +pg_dump -s -d vchord +psql -d vchord -c "DROP SCHEMA IF EXISTS vchord CASCADE;" 1>&2 From 3eb09c37c186b0deb8dff2336eea09a67841910b Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 13 Aug 2025 10:35:13 +0800 Subject: [PATCH 195/324] docker: download vchord with wget in build (#312) Signed-off-by: usamoi --- .github/workflows/release.yml | 21 ++------------------- docker/Dockerfile | 21 ++++++++++++--------- docker/binary.Dockerfile | 8 -------- 3 files changed, 14 insertions(+), 36 deletions(-) delete mode 100644 docker/binary.Dockerfile diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 2011987c..b03b5499 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -147,18 +147,6 @@ jobs: with: username: ${{ secrets.DOCKERIO_USERNAME }} password: ${{ secrets.DOCKERIO_TOKEN }} - - name: Push binary release to Docker Registry - uses: docker/build-push-action@v6 - with: - context: . - push: true - platforms: ${{ matrix.runner == 'ubuntu-22.04' && 'linux/amd64' || 'linux/arm64' }} - file: ./docker/binary.Dockerfile - provenance: false - tags: tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} - build-args: | - PG_VERSION=${{ matrix.version }} - SEMVER=${{ env.SEMVER }} - name: Login to GHCR uses: docker/login-action@v3 with: @@ -177,8 +165,8 @@ jobs: tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} build-args: | - PG_VERSION=${{ matrix.version }} - SEMVER=${{ env.SEMVER }} + BASE=${{ matrix.version }}-bookworm + VCHORD=${{ env.SEMVER }} PGVECTOR=0.8.0 create-manifests: @@ -207,11 +195,6 @@ jobs: password: ${{ secrets.GITHUB_TOKEN }} - name: Create manifest and push run: | - docker manifest create \ - tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }} \ - --amend tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ - --amend tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 - docker manifest push tensorchord/vchord-binary:pg${{ matrix.version }}-v${{ env.SEMVER }} docker manifest create \ tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} \ --amend tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ diff --git a/docker/Dockerfile b/docker/Dockerfile index 1a7966c0..04c15c8e 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,14 +1,17 @@ -ARG PG_VERSION=17 -ARG PGVECTOR=0.8.0 +ARG BASE=17-bookworm -FROM pgvector/pgvector:${PGVECTOR}-pg${PG_VERSION} +FROM postgres:${BASE} -ARG PG_VERSION -ARG SEMVER -ARG TARGETARCH +ARG VCHORD +ARG PGVECTOR -RUN echo ${PG_VERSION} -COPY ./build/postgresql-${PG_VERSION}-vchord_${SEMVER}-1_${TARGETARCH}.deb /tmp/vchord.deb -RUN apt-get install -y /tmp/vchord.deb && rm -f /tmp/vchord.deb +RUN apt-get update && \ + apt-get install -y --no-install-recommends wget ca-certificates && \ + wget https://github.com/tensorchord/VectorChord/releases/download/${VCHORD}/postgresql-${PG_MAJOR}-vchord_${VCHORD}-1_$(dpkg --print-architecture).deb -P /tmp && \ + apt-get install -y postgresql-${PG_MAJOR}-pgvector=${PGVECTOR}-* && \ + apt-get install -y /tmp/postgresql-${PG_MAJOR}-vchord_${VCHORD}-1_$(dpkg --print-architecture).deb && \ + apt-get remove -y wget ca-certificates && \ + apt-get purge -y --auto-remove && \ + rm -rf /tmp/* /var/lib/apt/lists/* CMD ["postgres", "-c" ,"shared_preload_libraries=vchord"] diff --git a/docker/binary.Dockerfile b/docker/binary.Dockerfile deleted file mode 100644 index 438031d7..00000000 --- a/docker/binary.Dockerfile +++ /dev/null @@ -1,8 +0,0 @@ -FROM scratch - -ARG SEMVER -ARG PG_VERSION -ARG TARGETARCH - -WORKDIR /workspace -COPY ./build/postgresql-${PG_VERSION}-vchord_${SEMVER}-1_${TARGETARCH}.deb /workspace/ From e7021c47cd735dac0631887bfc2e9dc2a2729dd2 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 15 Aug 2025 15:20:25 +0800 Subject: [PATCH 196/324] readme: sync with docs (#313) [Rendered](https://github.com/usamoi/VectorChord/blob/0.5.0-readme/README.md) Signed-off-by: usamoi --- README.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 380690d9..db627121 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: **⚡ Enhanced Performance**: Delivering optimized operations with up to 5x faster queries, 16x higher insert throughput, and 16x quicker[^1] index building compared to pgvector's HNSW implementation. -[^1]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/#/) with 768-dimensional vectors and 95% recall. Please checkout our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. +[^1]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/) with 768-dimensional vectors and 95% recall. Please checkout our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. **💰 Affordable Vector Search**: Query 100M 768-dimensional vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%, helping you keep infrastructure costs down while maintaining high search quality. @@ -56,7 +56,7 @@ VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: [^4]: Please check our [blog post](https://blog.vectorchord.ai/vector-search-at-10000-qps-in-postgresql-with-vectorchord) for more details, the PostgreSQL scalability is powered by [CloudNative-PG](https://github.com/cloudnative-pg/cloudnative-pg). -**🏭 Production Proven**: Deployed in mission-critical environments, VectorChord ​​reliably handles 3B+ vectors​​ in production with consistent performance. Please check out [3B vectors in PostgresQL to Protect the Earth](https://blog.vectorchord.ai/3-billion-vectors-in-postgresql-to-protect-the-earth). +**🏭 Production Proven**: Deployed in mission-critical environments, VectorChord ​​reliably handles 3B+ vectors​​ in production with consistent performance. Please check out [3B vectors in PostgreSQL to Protect the Earth](https://blog.vectorchord.ai/3-billion-vectors-in-postgresql-to-protect-the-earth). ## Quick Start @@ -67,7 +67,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.4.3 + -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.5.0 ``` > [!NOTE] > In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). @@ -109,9 +109,11 @@ For more usage, please read: - [Indexing](https://docs.vectorchord.ai/vectorchord/usage/indexing.html) - [Multi-Vector Retrieval](https://docs.vectorchord.ai/vectorchord/usage/indexing-with-maxsim-operators.html) +- [Graph Index](https://docs.vectorchord.ai/vectorchord/usage/graph-index.html) - [Similarity Filter](https://docs.vectorchord.ai/vectorchord/usage/range-query.html) - [PostgreSQL Tuning](https://docs.vectorchord.ai/vectorchord/usage/performance-tuning.html) - [Monitoring](https://docs.vectorchord.ai/vectorchord/usage/monitoring.html) +- [Measure Recall](https://docs.vectorchord.ai/vectorchord/usage/measure-recall.html) - [Prewarm](https://docs.vectorchord.ai/vectorchord/usage/prewarm.html) - [Prefilter](https://docs.vectorchord.ai/vectorchord/usage/prefilter.html) - [Prefetch](https://docs.vectorchord.ai/vectorchord/usage/prefetch.html) From fb00f48570c7e75138777224935eb9dec4684a9f Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 18 Aug 2025 13:03:39 +0800 Subject: [PATCH 197/324] perf: heap optimizations (#314) Signed-off-by: usamoi --- Cargo.lock | 8 + Cargo.toml | 2 + crates/algo/Cargo.toml | 1 + crates/algo/src/lib.rs | 72 ++++++- crates/vchordg/src/tuples.rs | 17 +- crates/vchordrq/src/rerank.rs | 28 +-- crates/vchordrq/src/search.rs | 19 +- crates/vchordrq/src/tape.rs | 31 ++- crates/vchordrq/src/tape_writer.rs | 45 ++-- crates/vchordrq/src/tuples.rs | 21 +- src/index/storage.rs | 10 +- src/index/vchordg/am/am_build.rs | 9 - src/index/vchordrq/scanners/default.rs | 272 ++++++++++++++----------- src/index/vchordrq/scanners/maxsim.rs | 130 ++++++------ tests/vchordrq/index.slt | 6 +- 15 files changed, 409 insertions(+), 262 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 59eff7f1..858b22cd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,7 @@ version = "0.0.0" dependencies = [ "always_equal", "bumpalo", + "dary_heap", "distance", "half 2.6.0", "simd", @@ -329,6 +330,12 @@ dependencies = [ "syn", ] +[[package]] +name = "dary_heap" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04d2cd9c18b9f454ed67da600630b021a8a80bf33f8c95896ab33aaf1c26b728" + [[package]] name = "displaydoc" version = "0.2.5" @@ -1576,6 +1583,7 @@ dependencies = [ "algo", "always_equal", "bumpalo", + "dary_heap", "distance", "half 2.6.0", "k_means", diff --git a/Cargo.toml b/Cargo.toml index ab874f21..7902d33a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,7 @@ vchordrq = { path = "./crates/vchordrq" } vector = { path = "./crates/vector" } bumpalo.workspace = true +dary_heap.workspace = true half.workspace = true paste.workspace = true pgrx = { version = "=0.16.0", default-features = false, features = ["cshim"] } @@ -61,6 +62,7 @@ edition = "2024" [workspace.dependencies] bumpalo = "3.19.0" +dary_heap = "0.3.7" half = { version = "2.6.0", features = ["zerocopy"] } paste = "1.0.15" rand = "0.9.2" diff --git a/crates/algo/Cargo.toml b/crates/algo/Cargo.toml index 166a1e88..2951aad5 100644 --- a/crates/algo/Cargo.toml +++ b/crates/algo/Cargo.toml @@ -12,6 +12,7 @@ small_iter = { path = "../small_iter" } vector = { path = "../vector" } bumpalo.workspace = true +dary_heap.workspace = true half.workspace = true zerocopy.workspace = true diff --git a/crates/algo/src/lib.rs b/crates/algo/src/lib.rs index 39cfc215..baedf91c 100644 --- a/crates/algo/src/lib.rs +++ b/crates/algo/src/lib.rs @@ -17,6 +17,7 @@ pub mod prefetcher; pub mod tuples; use always_equal::AlwaysEqual; +use dary_heap::DaryHeap; use std::collections::{BinaryHeap, VecDeque}; use std::iter::Peekable; use std::ops::{Deref, DerefMut}; @@ -169,11 +170,13 @@ impl Fetch<'_> for u32 { } } -impl<'b, T, A, B> Fetch<'b> for (T, AlwaysEqual<&mut (A, B, BorrowedIter<'b>)>) { +impl<'b, T, A, B, W: 'b + PackedRefMut)>> Fetch<'b> + for (T, AlwaysEqual) +{ type Iter = BorrowedIter<'b>; #[inline(always)] fn fetch(&self) -> BorrowedIter<'b> { - let (_, AlwaysEqual((.., list))) = self; + let (.., list) = self.1.0.get(); *list } } @@ -289,6 +292,23 @@ impl Sequence for BinaryHeap { } } +impl Sequence for DaryHeap { + type Item = T; + type Inner = std::vec::IntoIter; + #[inline] + fn next(&mut self) -> Option { + self.pop() + } + #[inline] + fn peek(&mut self) -> Option<&T> { + (self as &Self).peek() + } + #[inline] + fn into_inner(self) -> Self::Inner { + self.into_vec().into_iter() + } +} + impl Sequence for Peekable { type Item = I::Item; type Inner = Peekable; @@ -328,3 +348,51 @@ impl Sequence for VecDeque { /// * `Opaque` must aligned to 8 bytes. #[allow(unsafe_code)] pub unsafe trait Opaque: Copy + Send + Sync + FromBytes + IntoBytes + 'static {} + +pub trait PackedRefMut { + type T; + fn get(&self) -> &Self::T; + fn get_mut(&mut self) -> &mut Self::T; +} + +impl PackedRefMut for &mut T { + type T = T; + #[inline(always)] + fn get(&self) -> &T { + self + } + #[inline(always)] + fn get_mut(&mut self) -> &mut T { + self + } +} + +#[repr(Rust, packed(4))] +pub struct PackedRefMut4<'b, T>(pub &'b mut T); + +impl<'b, T> PackedRefMut for PackedRefMut4<'b, T> { + type T = T; + #[inline(always)] + fn get(&self) -> &T { + self.0 + } + #[inline(always)] + fn get_mut(&mut self) -> &mut T { + self.0 + } +} + +#[repr(Rust, packed(8))] +pub struct PackedRefMut8<'b, T>(pub &'b mut T); + +impl<'a, T> PackedRefMut for PackedRefMut8<'a, T> { + type T = T; + #[inline(always)] + fn get(&self) -> &T { + self.0 + } + #[inline(always)] + fn get_mut(&mut self) -> &mut T { + self.0 + } +} diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 05e6fd46..a33fc134 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -23,6 +23,15 @@ pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordg\0"); const VERSION: u64 = 10; +#[inline(always)] +fn tag(source: &[u8]) -> Tag { + assert!(source.len() >= size_of::()); + #[allow(unsafe_code)] + unsafe { + source.as_ptr().cast::().read_unaligned() + } +} + pub trait Tuple: 'static { fn serialize(&self) -> Vec; } @@ -114,7 +123,7 @@ impl Tuple for MetaTuple { impl WithReader for MetaTuple { type Reader<'a> = MetaTupleReader<'a>; fn deserialize_ref(source: &[u8]) -> MetaTupleReader<'_> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { MAGIC => { let checker = RefChecker::new(source); @@ -166,7 +175,7 @@ impl<'a> MetaTupleReader<'a> { impl WithWriter for MetaTuple { type Writer<'a> = MetaTupleWriter<'a>; fn deserialize_mut(source: &mut [u8]) -> MetaTupleWriter<'_> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { MAGIC => { let checker = RefChecker::new(source); @@ -449,7 +458,7 @@ impl WithReader for VectorTuple { type Reader<'a> = VectorTupleReader<'a, V>; fn deserialize_ref(source: &[u8]) -> VectorTupleReader<'_, V> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { 0 => { let checker = RefChecker::new(source); @@ -479,7 +488,7 @@ impl WithWriter for VectorTuple { type Writer<'a> = VectorTupleWriter<'a, V>; fn deserialize_mut(source: &mut [u8]) -> VectorTupleWriter<'_, V> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { 0 => { let mut checker = MutChecker::new(source); diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs index 62e49006..bb3948d2 100644 --- a/crates/vchordrq/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -18,7 +18,7 @@ use crate::tuples::{MetaTuple, WithReader}; use crate::{Page, vectors}; use algo::accessor::{Accessor2, LTryAccess}; use algo::prefetcher::Prefetcher; -use algo::{RelationRead, RerankMethod}; +use algo::{PackedRefMut, RelationRead, RerankMethod}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -27,8 +27,6 @@ use std::marker::PhantomData; use std::num::NonZero; use vector::VectorOwned; -type Extra0<'b> = &'b mut (NonZero, u16, algo::BorrowedIter<'b>); - type Result = (Reverse, AlwaysEqual>); pub fn how(index: &impl RelationRead) -> RerankMethod { @@ -43,25 +41,27 @@ pub fn how(index: &impl RelationRead) -> RerankMethod { } } -pub struct Reranker { +pub struct Reranker { prefetcher: P, cache: BinaryHeap, f: F, - _phantom: PhantomData T>, + _phantom: PhantomData (T, W)>, } -impl<'b, T, F, P> Iterator for Reranker +impl<'b, T, F, P, W> Iterator for Reranker where F: FnMut(NonZero, P::Guards, u16) -> Option, - P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual)>, + W: 'b + PackedRefMut, u16, algo::BorrowedIter<'b>)>, { type Item = (Distance, NonZero); fn next(&mut self) -> Option { - while let Some(((_, AlwaysEqual(&mut (payload, head, ..))), prefetch)) = self + while let Some(((_, AlwaysEqual(mut w)), prefetch)) = self .prefetcher .next_if(|((d, _), ..)| Some(*d) > self.cache.peek().map(|(d, ..)| *d)) { + let &mut (payload, head, ..) = w.get_mut(); if let Some(distance) = (self.f)(payload, prefetch, head) { self.cache.push((Reverse(distance), AlwaysEqual(payload))); }; @@ -71,7 +71,7 @@ where } } -impl Reranker { +impl Reranker { pub fn finish(self) -> (P, impl Iterator) { (self.prefetcher, self.cache.into_iter()) } @@ -81,11 +81,12 @@ pub fn rerank_index< 'b, O: Operator, T, - P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual)>, + W: 'b + PackedRefMut, u16, algo::BorrowedIter<'b>)>, >( vector: O::Vector, prefetcher: P, -) -> Reranker, P::Guards, u16) -> Option, P> { +) -> Reranker, P::Guards, u16) -> Option, P, W> { Reranker { prefetcher, cache: BinaryHeap::new(), @@ -108,12 +109,13 @@ pub fn rerank_heap< 'b, O: Operator, T, - P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual>)>, + P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual)>, + W: 'b + PackedRefMut, u16, algo::BorrowedIter<'b>)>, >( vector: O::Vector, prefetcher: P, mut fetch: impl FnMut(NonZero) -> Option + 'b, -) -> Reranker, P::Guards, u16) -> Option, P> { +) -> Reranker, P::Guards, u16) -> Option, P, W> { Reranker { prefetcher, cache: BinaryHeap::new(), diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index c737c702..4d18db9b 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -20,7 +20,7 @@ use crate::tuples::*; use crate::{Opaque, Page, tape, vectors}; use algo::accessor::LAccess; use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; -use algo::{BorrowedIter, Bump, RelationRead}; +use algo::{BorrowedIter, Bump, PackedRefMut4, PackedRefMut8, RelationRead}; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; @@ -29,7 +29,6 @@ use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; type Extra1<'b> = &'b mut (u32, f32, u16, BorrowedIter<'b>); -type Extra0<'b> = &'b mut (NonZero, u16, BorrowedIter<'b>); pub fn default_search<'b, R: RelationRead, O: Operator>( index: &'b R, @@ -41,7 +40,7 @@ pub fn default_search<'b, R: RelationRead, O: Operator>( mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, ) -> Vec<( (Reverse, AlwaysEqual<()>), - AlwaysEqual>, + AlwaysEqual, u16, BorrowedIter<'b>)>>, )> where R::Page: Page, @@ -128,7 +127,7 @@ where state = step(state).take(probes[i as usize - 1] as _).collect(); } - let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); + let mut results = LinkedVec::<(_, AlwaysEqual<_>)>::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); @@ -137,11 +136,11 @@ where let lowerbound = Distance::from_f32(rough - err * epsilon); results.push(( (Reverse(lowerbound), AlwaysEqual(())), - AlwaysEqual(bump.alloc(( + AlwaysEqual(PackedRefMut4(bump.alloc(( payload, head, BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), - ))), + )))), )); }); if prefetch_h0_tuples.is_not_plain() { @@ -180,7 +179,7 @@ pub fn maxsim_search<'b, R: RelationRead, O: Operator>( ) -> ( Vec<( (Reverse, AlwaysEqual), - AlwaysEqual>, + AlwaysEqual, u16, BorrowedIter<'b>)>>, )>, Distance, ) @@ -271,7 +270,7 @@ where state = it.take(probes[i as usize - 1] as _).collect(); } - let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); + let mut results = LinkedVec::<(_, AlwaysEqual<_>)>::new(); for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { let jump_guard = index.read(first); let jump_bytes = jump_guard.get(1).expect("data corruption"); @@ -281,11 +280,11 @@ where let rough = Distance::from_f32(rough); results.push(( (Reverse(lowerbound), AlwaysEqual(rough)), - AlwaysEqual(bump.alloc(( + AlwaysEqual(PackedRefMut8(bump.alloc(( payload, head, BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), - ))), + )))), )); }); if prefetch_h0_tuples.is_not_plain() { diff --git a/crates/vchordrq/src/tape.rs b/crates/vchordrq/src/tape.rs index 16f4e54a..755a8f00 100644 --- a/crates/vchordrq/src/tape.rs +++ b/crates/vchordrq/src/tape.rs @@ -244,7 +244,9 @@ pub fn read_h1_tape<'b, R, A, T>( let mut x = x.take().unwrap_or_else(&accessor); x.push(tuple.elements()); let values = x.finish((tuple.metadata(), tuple.delta())); - let prefetch: [_; 32] = fix_0(tuple.prefetch()); + let prefetch = tuple.prefetch(); + let flattened = prefetch.as_flattened(); + let step = prefetch.len(); for (j, value) in values.into_iter().enumerate() { if j < tuple.len() as usize { callback( @@ -252,7 +254,7 @@ pub fn read_h1_tape<'b, R, A, T>( tuple.head()[j], tuple.norm()[j], tuple.first()[j], - fix_1(prefetch[j]), + &flattened[j * step..][..step], ); } } @@ -283,10 +285,17 @@ pub fn read_frozen_tape<'b, R, A, T>( let mut x = x.take().unwrap_or_else(&accessor); x.push(tuple.elements()); let values = x.finish((tuple.metadata(), tuple.delta())); - let prefetch: [_; 32] = fix_0(tuple.prefetch()); + let prefetch = tuple.prefetch(); + let flattened = prefetch.as_flattened(); + let step = prefetch.len(); for (j, value) in values.into_iter().enumerate() { if let Some(payload) = tuple.payload()[j] { - callback(value, tuple.head()[j], payload, fix_1(prefetch[j])); + callback( + value, + tuple.head()[j], + payload, + &flattened[j * step..][..step], + ); } } } @@ -393,17 +402,3 @@ where } } } - -fn fix_0(x: &[[T; 32]]) -> [&[T]; 32] { - let step = x.len(); - let flat = x.as_flattened(); - std::array::from_fn(|i| &flat[i * step..][..step]) -} - -fn fix_1(x: &[u32]) -> &[u32] { - if let Some(i) = x.iter().position(|&x| x == u32::MAX) { - &x[..i] - } else { - x - } -} diff --git a/crates/vchordrq/src/tape_writer.rs b/crates/vchordrq/src/tape_writer.rs index 19caab88..3078ea10 100644 --- a/crates/vchordrq/src/tape_writer.rs +++ b/crates/vchordrq/src/tape_writer.rs @@ -97,7 +97,7 @@ where chunk.each_ref().map(|x| x.code.0.factor_err), ], delta: chunk.each_ref().map(|x| x.delta), - prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), + prefetch: fix_good(chunk.each_ref().map(|x| x.prefetch.as_slice())), norm: chunk.each_ref().map(|x| x.norm), head: chunk.each_ref().map(|x| x.head), first: chunk.each_ref().map(|x| x.extra), @@ -138,7 +138,7 @@ where any_pack(chunk.iter().map(|x| x.code.0.factor_err)), ], delta: any_pack(chunk.iter().map(|x| x.delta)), - prefetch: fix(chunk.iter().map(|x| x.prefetch.as_slice())), + prefetch: fix_bad(chunk.iter().map(|x| x.prefetch.as_slice())), head: any_pack(chunk.iter().map(|x| x.head)), norm: any_pack(chunk.iter().map(|x| x.norm)), first: any_pack(chunk.iter().map(|x| x.extra)), @@ -197,7 +197,7 @@ where chunk.each_ref().map(|x| x.code.0.factor_err), ], delta: chunk.each_ref().map(|x| x.delta), - prefetch: fix(chunk.each_ref().map(|x| x.prefetch.as_slice())), + prefetch: fix_good(chunk.each_ref().map(|x| x.prefetch.as_slice())), head: chunk.each_ref().map(|x| x.head), payload: chunk.each_ref().map(|x| Some(x.extra)), elements: remain, @@ -222,15 +222,36 @@ where } } -fn fix<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { - use std::array::from_fn; - let mut iter = into_iter.into_iter(); - let mut array: [_; 32] = from_fn(|_| iter.next().map(<[u32]>::to_vec).unwrap_or_default()); - if iter.next().is_some() { +fn fix_good<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { + let slices = into_iter.into_iter().collect::>(); + if slices.len() != 32 { + panic!("too many or too few slices"); + } + let min = slices.iter().map(|x| x.len()).min().unwrap_or_default(); + let max = slices.iter().map(|x| x.len()).max().unwrap_or_default(); + if min != max { + panic!("the number of pages for prefetching is not a constant"); + } + let flattened = slices.into_iter().flatten().copied().collect::>(); + let (arrays, remainder) = flattened.as_chunks::<32>(); + assert!(remainder.is_empty()); + arrays.to_vec() +} + +fn fix_bad<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { + let mut slices = into_iter.into_iter().collect::>(); + if slices.len() > 32 { panic!("too many slices"); } - let step = array.iter().map(Vec::len).max().unwrap_or_default(); - array.iter_mut().for_each(|x| x.resize(step, u32::MAX)); - let flat = array.into_iter().flatten().collect::>(); - (0..step).map(|i| from_fn(|j| flat[i * 32 + j])).collect() + let min = slices.iter().map(|x| x.len()).min().unwrap_or_default(); + let max = slices.iter().map(|x| x.len()).max().unwrap_or_default(); + if min != max { + panic!("the number of pages for prefetching is not a constant"); + } + let temp = vec![u32::MAX; max]; + slices.resize(32, temp.as_slice()); + let flattened = slices.into_iter().flatten().copied().collect::>(); + let (arrays, remainder) = flattened.as_chunks::<32>(); + assert!(remainder.is_empty()); + arrays.to_vec() } diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index d8d5ac83..a1018151 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -22,6 +22,15 @@ pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); const VERSION: u64 = 11; +#[inline(always)] +fn tag(source: &[u8]) -> Tag { + assert!(source.len() >= size_of::()); + #[allow(unsafe_code)] + unsafe { + source.as_ptr().cast::().read_unaligned() + } +} + pub trait Tuple: 'static { fn serialize(&self) -> Vec; } @@ -137,7 +146,7 @@ impl Tuple for MetaTuple { impl WithReader for MetaTuple { type Reader<'a> = MetaTupleReader<'a>; fn deserialize_ref(source: &[u8]) -> MetaTupleReader<'_> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { MAGIC => { let checker = RefChecker::new(source); @@ -351,7 +360,7 @@ impl WithReader for VectorTuple { type Reader<'a> = VectorTupleReader<'a, V>; fn deserialize_ref(source: &[u8]) -> VectorTupleReader<'_, V> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { 0 => { let checker = RefChecker::new(source); @@ -519,7 +528,7 @@ impl WithReader for DirectoryTuple { type Reader<'a> = DirectoryTupleReader<'a>; fn deserialize_ref(source: &[u8]) -> DirectoryTupleReader<'_> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { 0 => { let checker = RefChecker::new(source); @@ -713,7 +722,7 @@ impl WithReader for H1Tuple { type Reader<'a> = H1TupleReader<'a>; fn deserialize_ref(source: &[u8]) -> H1TupleReader<'_> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { 0 => { let checker = RefChecker::new(source); @@ -1047,7 +1056,7 @@ impl WithReader for FrozenTuple { type Reader<'a> = FrozenTupleReader<'a>; fn deserialize_ref(source: &[u8]) -> FrozenTupleReader<'_> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { 0 => { let checker = RefChecker::new(source); @@ -1075,7 +1084,7 @@ impl WithWriter for FrozenTuple { type Writer<'a> = FrozenTupleWriter<'a>; fn deserialize_mut(source: &mut [u8]) -> FrozenTupleWriter<'_> { - let tag = Tag::from_ne_bytes(std::array::from_fn(|i| source[i])); + let tag = tag(source); match tag { 0 => { let mut checker = MutChecker::new(source); diff --git a/src/index/storage.rs b/src/index/storage.rs index 34027948..8b7849a7 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -89,7 +89,7 @@ impl Page for PostgresPage { let iid = unsafe { self.header.pd_linp.as_ptr().add((i - 1) as _).read() }; let lp_off = iid.lp_off() as usize; let lp_len = iid.lp_len() as usize; - match iid.lp_flags() { + match lp_flags(iid) { pgrx::pg_sys::LP_UNUSED => return None, pgrx::pg_sys::LP_NORMAL => (), pgrx::pg_sys::LP_REDIRECT => unimplemented!(), @@ -119,7 +119,7 @@ impl Page for PostgresPage { let iid = unsafe { self.header.pd_linp.as_ptr().add((i - 1) as _).read() }; let lp_off = iid.lp_off() as usize; let lp_len = iid.lp_len() as usize; - match iid.lp_flags() { + match lp_flags(iid) { pgrx::pg_sys::LP_UNUSED => return None, pgrx::pg_sys::LP_NORMAL => (), pgrx::pg_sys::LP_REDIRECT => unimplemented!(), @@ -674,6 +674,12 @@ impl RelationLength for PostgresRelation { } } +#[inline(always)] +fn lp_flags(x: pgrx::pg_sys::ItemIdData) -> u32 { + let x: u32 = unsafe { std::mem::transmute(x) }; + (x >> 15) & 0b11 +} + // Emulate unstable library feature `vec_deque_pop_if`. // See https://github.com/rust-lang/rust/issues/135889. diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index d777635d..a8cb518e 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -656,15 +656,6 @@ mod vchordg_cached { #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VchordgCachedHeader0 {} - #[repr(C, align(8))] - #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] - struct VchordgCachedHeader1 { - mapping_s: usize, - mapping_e: usize, - pages_s: usize, - pages_e: usize, - } - pub enum VchordgCached { _0 {}, } diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index 2d4739bf..a5fa5b30 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -23,8 +23,8 @@ use algo::accessor::{Dot, L2S}; use algo::prefetcher::*; use algo::*; use always_equal::AlwaysEqual; +use dary_heap::QuaternaryHeap as Heap; use half::f16; -use std::collections::BinaryHeap; use std::num::NonZero; use vchordrq::operator::{self}; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; @@ -156,57 +156,60 @@ impl SearchBuilder for DefaultBuilder { ), }; let method = how(index); - let sequence = BinaryHeap::from(results); + let sequence = Heap::from(results); match (method, options.io_rerank, options.prefilter) { (RerankMethod::Index, Io::Plain, false) => { let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Plain, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, false) => { let prefetcher = SimplePrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = SimplePrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, false) => { let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Heap, _, false) => { let fetch = move |payload| { @@ -225,7 +228,9 @@ impl SearchBuilder for DefaultBuilder { Some(raw) }; let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) } (RerankMethod::Heap, _, true) => { let fetch = move |payload| { @@ -247,7 +252,9 @@ impl SearchBuilder for DefaultBuilder { Some(raw) }; let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) } } } @@ -289,57 +296,60 @@ impl SearchBuilder for DefaultBuilder { ), }; let method = how(index); - let sequence = BinaryHeap::from(results); + let sequence = Heap::from(results); match (method, options.io_rerank, options.prefilter) { (RerankMethod::Index, Io::Plain, false) => { let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Plain, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, false) => { let prefetcher = SimplePrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = SimplePrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, false) => { let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Heap, _, false) => { let fetch = move |payload| { @@ -358,7 +368,9 @@ impl SearchBuilder for DefaultBuilder { Some(raw) }; let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) } (RerankMethod::Heap, _, true) => { let fetch = move |payload| { @@ -380,7 +392,9 @@ impl SearchBuilder for DefaultBuilder { Some(raw) }; let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) } } } @@ -422,57 +436,60 @@ impl SearchBuilder for DefaultBuilder { ), }; let method = how(index); - let sequence = BinaryHeap::from(results); + let sequence = Heap::from(results); match (method, options.io_rerank, options.prefilter) { (RerankMethod::Index, Io::Plain, false) => { let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Plain, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, false) => { let prefetcher = SimplePrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = SimplePrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, false) => { let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Heap, _, false) => { let fetch = move |payload| { @@ -491,7 +508,9 @@ impl SearchBuilder for DefaultBuilder { Some(raw) }; let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) } (RerankMethod::Heap, _, true) => { let fetch = move |payload| { @@ -513,7 +532,9 @@ impl SearchBuilder for DefaultBuilder { Some(raw) }; let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) } } } @@ -555,57 +576,60 @@ impl SearchBuilder for DefaultBuilder { ), }; let method = how(index); - let sequence = BinaryHeap::from(results); + let sequence = Heap::from(results); match (method, options.io_rerank, options.prefilter) { (RerankMethod::Index, Io::Plain, false) => { let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Plain, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, false) => { let prefetcher = SimplePrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Simple, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = SimplePrefetcher::new(index, sequence); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, false) => { let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { - let predicate = id_0(move |(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); - Box::new(rerank_index::(unprojected, prefetcher).map(f)) + Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Heap, _, false) => { let fetch = move |payload| { @@ -624,7 +648,9 @@ impl SearchBuilder for DefaultBuilder { Some(raw) }; let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) } (RerankMethod::Heap, _, true) => { let fetch = move |payload| { @@ -646,7 +672,9 @@ impl SearchBuilder for DefaultBuilder { Some(raw) }; let prefetcher = PlainPrefetcher::new(index, sequence); - Box::new(rerank_heap::(unprojected, prefetcher, fetch).map(f)) + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) } } } @@ -671,7 +699,7 @@ impl SearchBuilder for DefaultBuilder { #[inline(always)] pub fn id_0(f: F) -> F where - F: for<'a> FnMut(&(A, AlwaysEqual<&'a mut (B, C, D)>)) -> R, + F: for<'a> FnMut(&(A, AlwaysEqual>)) -> R, { f } diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 3000bc06..f01710ec 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -22,6 +22,7 @@ use algo::accessor::Dot; use algo::prefetcher::*; use algo::*; use always_equal::AlwaysEqual; +use dary_heap::QuaternaryHeap as Heap; use distance::Distance; use half::f16; use std::cmp::Reverse; @@ -107,10 +108,11 @@ impl SearchBuilder for MaxsimBuilder { }; let n = vectors.len(); let accu_map = |(Reverse(distance), AlwaysEqual(payload))| (distance, payload); - let rough_map = |((_, AlwaysEqual(rough)), AlwaysEqual(&mut (payload, ..))): ( - _, - AlwaysEqual<&mut (NonZero, _, _)>, - )| (rough, payload); + let rough_map = + |((_, AlwaysEqual(rough)), AlwaysEqual(PackedRefMut8(&mut (payload, ..)))): ( + _, + AlwaysEqual, _, _)>>, + )| (rough, payload); let iter: Box> = match opfamily.vector_kind() { VectorKind::Vecf32 => { type Op = operator::Op, Dot>; @@ -163,29 +165,30 @@ impl SearchBuilder for MaxsimBuilder { }; let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { - let sequence = BinaryHeap::from(results); + let sequence = Heap::from(results); match (options.io_rerank, options.prefilter) { (Io::Plain, false) => { let prefetcher = PlainPrefetcher::new(index, sequence); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } (Io::Plain, true) => { - let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = PlainPrefetcher::new(index, sequence); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); @@ -194,24 +197,25 @@ impl SearchBuilder for MaxsimBuilder { (Io::Simple, false) => { let prefetcher = SimplePrefetcher::new(index, sequence); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } (Io::Simple, true) => { - let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = SimplePrefetcher::new(index, sequence); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); @@ -221,25 +225,26 @@ impl SearchBuilder for MaxsimBuilder { let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } (Io::Stream, true) => { - let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); @@ -304,29 +309,30 @@ impl SearchBuilder for MaxsimBuilder { }; let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { - let sequence = BinaryHeap::from(results); + let sequence = Heap::from(results); match (options.io_rerank, options.prefilter) { (Io::Plain, false) => { let prefetcher = PlainPrefetcher::new(index, sequence); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } (Io::Plain, true) => { - let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = PlainPrefetcher::new(index, sequence); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); @@ -335,24 +341,25 @@ impl SearchBuilder for MaxsimBuilder { (Io::Simple, false) => { let prefetcher = SimplePrefetcher::new(index, sequence); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } (Io::Simple, true) => { - let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = SimplePrefetcher::new(index, sequence); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); @@ -362,25 +369,26 @@ impl SearchBuilder for MaxsimBuilder { let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); rough_set.extend(rough_iter.into_iter().map(rough_map)); } (Io::Stream, true) => { - let predicate = id_0(|(_, AlwaysEqual((pointer, _, _)))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); let sequence = filter(sequence, predicate); let prefetcher = StreamPrefetcher::new(index, sequence, Hints::default()); let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); + rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); let (rough_iter, accu_iter) = reranker.finish(); accu_set.extend(accu_iter.map(accu_map)); @@ -491,7 +499,7 @@ impl std::iter::FusedIterator for IntoIterSorted {} #[inline(always)] pub fn id_0(f: F) -> F where - F: for<'a> FnMut(&(A, AlwaysEqual<&'a mut (B, C, D)>)) -> R, + F: for<'a> FnMut(&(A, AlwaysEqual>)) -> R, { f } diff --git a/tests/vchordrq/index.slt b/tests/vchordrq/index.slt index f2688752..0c95e8c8 100644 --- a/tests/vchordrq/index.slt +++ b/tests/vchordrq/index.slt @@ -16,7 +16,7 @@ CREATE INDEX ON t USING vchordrq (val vector_l2_ops) WITH (options = $$ residual_quantization = true [build.internal] -lists = [32] +lists = [33] spherical_centroids = false $$); @@ -25,7 +25,7 @@ CREATE INDEX ON t USING vchordrq (val vector_ip_ops) WITH (options = $$ residual_quantization = false [build.internal] -lists = [32] +lists = [33] spherical_centroids = true $$); @@ -34,7 +34,7 @@ CREATE INDEX ON t USING vchordrq (val vector_cosine_ops) WITH (options = $$ residual_quantization = false [build.internal] -lists = [32] +lists = [33] spherical_centroids = true $$); From 9eae0e7c4bf7a2d3083e51be7c183d5f9b410244 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 18 Aug 2025 16:01:29 +0800 Subject: [PATCH 198/324] chore: don't use attribute target arch (#316) https://github.com/llvm/llvm-project/issues/153199 Signed-off-by: usamoi --- crates/simd/cshim/x86_64.c | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/crates/simd/cshim/x86_64.c b/crates/simd/cshim/x86_64.c index 61a93421..0f469dd6 100644 --- a/crates/simd/cshim/x86_64.c +++ b/crates/simd/cshim/x86_64.c @@ -58,8 +58,10 @@ typedef _Float16 f16; typedef float f32; -__attribute__((target("arch=x86-64-v4,avx512fp16"))) float -fp16_reduce_sum_of_xy_v4_avx512fp16(f16 *restrict a, f16 *restrict b, size_t n) { +__attribute__((target("avx512bw,avx512cd,avx512dq,avx512vl,bmi,bmi2,lzcnt," + "movbe,popcnt,avx512fp16"))) float +fp16_reduce_sum_of_xy_v4_avx512fp16(f16 *restrict a, f16 *restrict b, + size_t n) { __m512h xy = _mm512_setzero_ph(); while (n >= 32) { __m512h x = _mm512_loadu_ph(a); @@ -78,8 +80,10 @@ fp16_reduce_sum_of_xy_v4_avx512fp16(f16 *restrict a, f16 *restrict b, size_t n) return _mm512_reduce_add_ph(xy); } -__attribute__((target("arch=x86-64-v4,avx512fp16"))) float -fp16_reduce_sum_of_d2_v4_avx512fp16(f16 *restrict a, f16 *restrict b, size_t n) { +__attribute__((target("avx512bw,avx512cd,avx512dq,avx512vl,bmi,bmi2,lzcnt," + "movbe,popcnt,avx512fp16"))) float +fp16_reduce_sum_of_d2_v4_avx512fp16(f16 *restrict a, f16 *restrict b, + size_t n) { __m512h d2 = _mm512_setzero_ph(); while (n >= 32) { __m512h x = _mm512_loadu_ph(a); From dfff814932b6f8071c29c33a081328a3b4750e5e Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 21 Aug 2025 17:49:13 +0800 Subject: [PATCH 199/324] fix compilation on aarch64_be (#318) not tested Signed-off-by: usamoi --- Cargo.lock | 131 +++++++++++++++++----------------- Cargo.toml | 2 +- crates/k_means/Cargo.toml | 2 +- crates/make/Cargo.toml | 5 +- crates/make/src/main.rs | 64 +++++++++++------ crates/simd/Cargo.toml | 2 +- crates/simd/cshim/aarch64.c | 18 +++++ crates/simd/src/f16.rs | 18 +++-- crates/simd/src/f32.rs | 54 +++++++------- crates/simd/src/fast_scan.rs | 7 +- crates/simd/src/quantize.rs | 6 ++ crates/simd_macros/src/lib.rs | 34 +++++---- src/lib.rs | 3 - 13 files changed, 194 insertions(+), 152 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 858b22cd..f0b41cd7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -48,9 +48,9 @@ dependencies = [ [[package]] name = "anstream" -version = "0.6.19" +version = "0.6.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "301af1932e46185686725e0fad2f8f2aa7da69dd70bf6ecc44d6b703844a3933" +checksum = "3ae563653d1938f79b1ab1b5e668c87c76a9930414574a6583a7b7e11a8e6192" dependencies = [ "anstyle", "anstyle-parse", @@ -78,29 +78,29 @@ dependencies = [ [[package]] name = "anstyle-query" -version = "1.1.3" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c8bdeb6047d8983be085bab0ba1472e6dc604e7041dbf6fcd5e71523014fae9" +checksum = "9e231f6134f61b71076a3eab506c379d4f36122f2af15a9ff04415ea4c3339e2" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] name = "anstyle-wincon" -version = "3.0.9" +version = "3.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "403f75924867bb1033c59fbf0797484329750cfbe3c4325cd33127941fabc882" +checksum = "3e0633414522a32ffaac8ac6cc8f748e090c5717661fddeea04219e2344f5f2a" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] name = "anyhow" -version = "1.0.98" +version = "1.0.99" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" +checksum = "b0674a1ddeecb70197781e945de4b3b8ffb61fa939a5597bcf48503737663100" [[package]] name = "bindgen" @@ -123,9 +123,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.9.1" +version = "2.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967" +checksum = "6a65b545ab31d687cff52899d4890855fec459eb6afe0da6417b8a18da87aa29" [[package]] name = "bitvec" @@ -157,9 +157,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.31" +version = "1.2.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3a42d84bb6b69d3a8b3eaacf0d88f179e1929695e1ad012b6cf64d9caaa5fd2" +checksum = "3ee0f8803222ba5a7e2777dd72ca451868909b1ac410621b676adf07280e9b5f" dependencies = [ "shlex", ] @@ -185,9 +185,9 @@ dependencies = [ [[package]] name = "cfg-if" -version = "1.0.1" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268" +checksum = "2fd1289c04a9ea8cb22300a459a72a385d7c73d3259e2ed7dcb2af674838cfa9" [[package]] name = "clang-sys" @@ -202,9 +202,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.42" +version = "4.5.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed87a9d530bb41a67537289bafcac159cb3ee28460e0a4571123d2a778a6a882" +checksum = "1fc0e74a703892159f5ae7d3aac52c8e6c392f5ae5f359c70b5881d60aaac318" dependencies = [ "clap_builder", "clap_derive", @@ -212,9 +212,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.42" +version = "4.5.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64f4f3f3c77c94aff3c7e9aac9a2ca1974a5adf392a8bb751e827d6d127ab966" +checksum = "b3e7f4214277f3c7aa526a59dd3fbe306a370daee1f8b7b8c987069cd8e888a8" dependencies = [ "anstream", "anstyle", @@ -224,9 +224,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.5.41" +version = "4.5.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef4f52386a59ca4c860f7393bcf8abd8dfd91ecccc0f774635ff68e92eeef491" +checksum = "14cb31bb0a7d536caef2639baa7fad459e15c3144efefa6dbd1c84562c4739f6" dependencies = [ "heck", "proc-macro2", @@ -441,9 +441,9 @@ checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" [[package]] name = "form_urlencoded" -version = "1.2.1" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" dependencies = [ "percent-encoding", ] @@ -468,9 +468,9 @@ dependencies = [ [[package]] name = "glob" -version = "0.3.2" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" [[package]] name = "half" @@ -491,9 +491,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.15.4" +version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5971ac85611da7067dbfcabef3c70ebb5606018acd9e2a3903a0da507521e0d5" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ "allocator-api2", "equivalent", @@ -615,9 +615,9 @@ checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" [[package]] name = "idna" -version = "1.0.3" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" dependencies = [ "idna_adapter", "smallvec", @@ -636,9 +636,9 @@ dependencies = [ [[package]] name = "indenter" -version = "0.3.3" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" +checksum = "964de6e86d545b246d84badc0fef527924ace5134f30641c203ef52ba83f58d5" [[package]] name = "indexmap" @@ -710,9 +710,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.174" +version = "0.2.175" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776" +checksum = "6a82ae493e598baaea5209805c49bbf2ea7de956d50d7da0da1164f9c6d28543" [[package]] name = "libloading" @@ -758,6 +758,7 @@ version = "0.0.0" dependencies = [ "clap", "object", + "shlex", "target-triple", ] @@ -800,9 +801,9 @@ dependencies = [ [[package]] name = "object" -version = "0.37.2" +version = "0.37.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3e3d0a7419f081f4a808147e845310313a39f322d7ae1f996b7f001d6cbed04" +checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" dependencies = [ "memchr", "wasmparser", @@ -848,9 +849,9 @@ dependencies = [ [[package]] name = "percent-encoding" -version = "2.3.1" +version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" [[package]] name = "petgraph" @@ -1041,9 +1042,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.95" +version = "1.0.101" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02b3e5e68a3a1a02aad3ec490a98007cbc13c37cbe84a3cd7b8e406d76e7f778" +checksum = "89ae43fd86e4158d6db51ad8e2b80f313af9cc74f5c0e03ccb87de09998732de" dependencies = [ "unicode-ident", ] @@ -1110,9 +1111,9 @@ dependencies = [ [[package]] name = "rayon" -version = "1.10.0" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" dependencies = [ "either", "rayon-core", @@ -1120,9 +1121,9 @@ dependencies = [ [[package]] name = "rayon-core" -version = "1.12.1" +version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" dependencies = [ "crossbeam-deque", "crossbeam-utils", @@ -1178,9 +1179,9 @@ dependencies = [ [[package]] name = "rustversion" -version = "1.0.21" +version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a0d197bd2c9dc6e53b84da9556a69ba4cdfab8619eb41a8bd1cc2027a0f6b1d" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] name = "ryu" @@ -1241,9 +1242,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.142" +version = "1.0.143" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "030fedb782600dcbd6f02d479bf0d817ac3bb40d644745b769d6a96bc3afc5a7" +checksum = "d401abef1d108fbd9cbaebc3e46611f4b1021f714a0597a71f41ee463f5f4a5a" dependencies = [ "itoa", "memchr", @@ -1346,9 +1347,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.104" +version = "2.0.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17b6f705963418cdb9927482fa304bc562ece2fdd4f616084c50b7023b435a40" +checksum = "ede7c438028d4436d71104916910f5bb611972c5cfd7f89b8300a8186e6fada6" dependencies = [ "proc-macro2", "quote", @@ -1380,18 +1381,18 @@ checksum = "1ac9aa371f599d22256307c24a9d748c041e548cbf599f35d890f9d365361790" [[package]] name = "thiserror" -version = "2.0.12" +version = "2.0.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" +checksum = "3467d614147380f2e4e374161426ff399c91084acd2363eaf549172b3d5e60c0" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.12" +version = "2.0.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d" +checksum = "6c5e1be1c48b9172ee610da68fd9cd2770e7a4056cb3fc98710ee6906f0c7960" dependencies = [ "proc-macro2", "quote", @@ -1514,9 +1515,9 @@ checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" [[package]] name = "url" -version = "2.5.4" +version = "2.5.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32f8b686cadd1473f4bd0117a5d28d36b1ade384ea9b5069a1c40aefed7fda60" +checksum = "ec961601b32b6f5d14ae8dabd35ff2ff2e2c6cb4c0e6641845ff105abe96d958" dependencies = [ "form_urlencoded", "idna", @@ -1537,9 +1538,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.17.0" +version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3cf4199d1e5d15ddd86a694e4d0dffa9c323ce759fea589f00fef9d81cc1931d" +checksum = "f33196643e165781c20a5ead5582283a7dacbb87855d867fbc2df3f81eddc1be" dependencies = [ "getrandom", "js-sys", @@ -1728,9 +1729,9 @@ dependencies = [ [[package]] name = "wasmparser" -version = "0.236.0" +version = "0.236.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16d1eee846a705f6f3cb9d7b9f79b54583810f1fb57a1e3aea76d1742db2e3d2" +checksum = "a9b1e81f3eb254cf7404a82cee6926a4a3ccc5aad80cc3d43608a070c67aa1d7" dependencies = [ "bitflags", ] @@ -1764,11 +1765,11 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] name = "winapi-util" -version = "0.1.9" +version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" +checksum = "0978bf7171b3d90bac376700cb56d606feb40f251a475a5d6634613564460b22" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -2047,9 +2048,9 @@ dependencies = [ [[package]] name = "zerovec" -version = "0.11.2" +version = "0.11.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a05eb080e015ba39cc9e23bbe5e7fb04d5fb040350f99f34e338d5fdd294428" +checksum = "e7aa2bd55086f1ab526693ecbe444205da57e25f4489879da80635a46d90e73b" dependencies = [ "yoke", "zerofrom", diff --git a/Cargo.toml b/Cargo.toml index 7902d33a..15e8d202 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -53,7 +53,7 @@ mimalloc = { version = "0.1.47", features = ["local_dynamic_tls"] } workspace = true [workspace] -resolver = "2" +resolver = "3" members = ["crates/*"] [workspace.package] diff --git a/crates/k_means/Cargo.toml b/crates/k_means/Cargo.toml index fd3b43a7..6475cd6e 100644 --- a/crates/k_means/Cargo.toml +++ b/crates/k_means/Cargo.toml @@ -9,7 +9,7 @@ rabitq = { path = "../rabitq" } simd = { path = "../simd" } rand.workspace = true -rayon = "1.10.0" +rayon = "1.11.0" [lints] workspace = true diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index 0d7d6414..d4565466 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -5,9 +5,10 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.42", features = ["derive"] } -object = { version = "0.37.2", default-features = false, features = [ +clap = { version = "4.5.45", features = ["derive"] } +object = { version = "0.37.3", default-features = false, features = [ "read", "wasm", ] } +shlex = "1.3.0" target-triple = "0.1.4" diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index b9c426cf..8e942920 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -40,6 +40,8 @@ struct BuildArgs { profile: String, #[arg(long, default_value = target_triple::TARGET)] target: String, + #[arg(long)] + runner: Option, } struct TargetSpecificInformation { @@ -225,6 +227,7 @@ fn parse( } fn generate( + runner: &Option>, pg_config: impl AsRef, pg_version: &str, tsi: &TargetSpecificInformation, @@ -263,11 +266,14 @@ fn generate( }); result.push(format!("pgrx_embed_vchord{}", tsi.exe_suffix()?)); let mut command; - if !(tsi.is_unix && tsi.is_emscripten) { - command = Command::new(result); - } else { - command = Command::new("node"); + if let Some(runner) = runner { + command = Command::new(&runner[0]); + for arg in runner[1..].iter() { + command.arg(arg); + } command.arg(result); + } else { + command = Command::new(result); } command.stderr(Stdio::inherit()); eprintln!("Running {command:?}"); @@ -315,31 +321,35 @@ fn main() -> Result<(), Box> { if !std::fs::exists("./vchord.control")? { return Err("The script must be run from the VectorChord source directory.".into()); } - let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { - PathBuf::from(value) - } else { - return Err("Environment variable `PGRX_PG_CONFIG_PATH` is not set.".into()); - }; - let pg_config = pg_config(&path)?; - let pg_version = { - let version = pg_config["VERSION"].clone(); - if let Some(prefix_stripped) = version.strip_prefix("PostgreSQL ") { - if let Some((stripped, _)) = prefix_stripped.split_once(|c: char| !c.is_ascii_digit()) { - format!("pg{stripped}",) - } else { - format!("pg{prefix_stripped}",) - } - } else { - return Err("PostgreSQL version is invalid.".into()); - } - }; let vchord_version = control_file("./vchord.control")?["default_version"].clone(); match cli.command { Commands::Build(BuildArgs { output, profile, target, + runner, }) => { + let runner = runner.and_then(|runner| shlex::split(&runner)); + let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { + PathBuf::from(value) + } else { + return Err("Environment variable `PGRX_PG_CONFIG_PATH` is not set.".into()); + }; + let pg_config = pg_config(&path)?; + let pg_version = { + let version = pg_config["VERSION"].clone(); + if let Some(prefix_stripped) = version.strip_prefix("PostgreSQL ") { + if let Some((stripped, _)) = + prefix_stripped.split_once(|c: char| !c.is_ascii_digit()) + { + format!("pg{stripped}",) + } else { + format!("pg{prefix_stripped}",) + } + } else { + return Err("PostgreSQL version is invalid.".into()); + } + }; let tsi = target_specific_information(&target)?; let obj = build(&path, &pg_version, &tsi, &profile, &target)?; let pkglibdir = format!("{output}/pkglibdir"); @@ -378,7 +388,15 @@ fn main() -> Result<(), Box> { } else { let exports = parse(&tsi, obj)?; install_by_writing( - generate(&path, &pg_version, &tsi, &profile, &target, exports)?, + generate( + &runner, + &path, + &pg_version, + &tsi, + &profile, + &target, + exports, + )?, format!("{sharedir_extension}/vchord--0.0.0.sql"), false, )?; diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 69ed8684..98a9bd5c 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -18,7 +18,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.31" +cc = "1.2.33" which = "8.0.0" [lints] diff --git a/crates/simd/cshim/aarch64.c b/crates/simd/cshim/aarch64.c index 8e2d9336..5af96c3b 100644 --- a/crates/simd/cshim/aarch64.c +++ b/crates/simd/cshim/aarch64.c @@ -25,7 +25,9 @@ #endif #include +#if defined(__AARCH64EL__) #include +#endif #include #include @@ -81,6 +83,7 @@ fp16_reduce_sum_of_xy_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { return vaddvq_f32(lo) + vaddvq_f32(hi); } +#if defined(__AARCH64EL__) __attribute__((target("+sve"))) float fp16_reduce_sum_of_xy_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svfloat16_t xy = svdup_f16(0.0); @@ -92,6 +95,7 @@ fp16_reduce_sum_of_xy_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { } return svaddv_f16(svptrue_b16(), xy); } +#endif __attribute__((target("+fp16"))) float fp16_reduce_sum_of_d2_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { @@ -150,6 +154,7 @@ fp16_reduce_sum_of_d2_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { return vaddvq_f32(lo) + vaddvq_f32(hi); } +#if defined(__AARCH64EL__) __attribute__((target("+sve"))) float fp16_reduce_sum_of_d2_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svfloat16_t d2 = svdup_f16(0.0); @@ -162,7 +167,9 @@ fp16_reduce_sum_of_d2_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { } return svaddv_f16(svptrue_b16(), d2); } +#endif +#if defined(__AARCH64EL__) __attribute__((target("+sve"))) float fp32_reduce_sum_of_x_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); @@ -173,7 +180,9 @@ fp32_reduce_sum_of_x_a3_256(float *restrict this, size_t n) { } return svaddv_f32(svptrue_b32(), sum); } +#endif +#if defined(__AARCH64EL__) __attribute__((target("+sve"))) float fp32_reduce_sum_of_abs_x_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); @@ -184,7 +193,9 @@ fp32_reduce_sum_of_abs_x_a3_256(float *restrict this, size_t n) { } return svaddv_f32(svptrue_b32(), sum); } +#endif +#if defined(__AARCH64EL__) __attribute__((target("+sve"))) float fp32_reduce_sum_of_x2_a3_256(float *restrict this, size_t n) { svfloat32_t sum = svdup_f32(0.0); @@ -195,7 +206,9 @@ fp32_reduce_sum_of_x2_a3_256(float *restrict this, size_t n) { } return svaddv_f32(svptrue_b32(), sum); } +#endif +#if defined(__AARCH64EL__) __attribute__((target("+sve"))) void fp32_reduce_min_max_of_x_a3_256(float *restrict this, size_t n, float *out_min, float *out_max) { @@ -210,7 +223,9 @@ fp32_reduce_min_max_of_x_a3_256(float *restrict this, size_t n, float *out_min, *out_min = svminv_f32(svptrue_b32(), min); *out_max = svmaxv_f32(svptrue_b32(), max); } +#endif +#if defined(__AARCH64EL__) __attribute__((target("+sve"))) float fp32_reduce_sum_of_xy_a3_256(float *restrict lhs, float *restrict rhs, size_t n) { @@ -223,7 +238,9 @@ fp32_reduce_sum_of_xy_a3_256(float *restrict lhs, float *restrict rhs, } return svaddv_f32(svptrue_b32(), sum); } +#endif +#if defined(__AARCH64EL__) __attribute__((target("+sve"))) float fp32_reduce_sum_of_d2_a3_256(float *restrict lhs, float *restrict rhs, size_t n) { @@ -237,3 +254,4 @@ fp32_reduce_sum_of_d2_a3_256(float *restrict lhs, float *restrict rhs, } return svaddv_f32(svptrue_b32(), sum); } +#endif diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index e6828907..95686e0c 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -458,9 +458,8 @@ mod reduce_sum_of_xy { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "sve")] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.512")] pub fn reduce_sum_of_xy_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { @@ -472,7 +471,7 @@ mod reduce_sum_of_xy { } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] #[test] fn reduce_sum_of_xy_a3_512_test() { use rand::Rng; @@ -503,7 +502,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -742,9 +741,8 @@ mod reduce_sum_of_d2 { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "sve")] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.512")] pub fn reduce_sum_of_d2_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { @@ -756,7 +754,7 @@ mod reduce_sum_of_d2 { } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] #[test] fn reduce_sum_of_d2_a3_512_test() { use rand::Rng; @@ -787,7 +785,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16")] pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index 159bb949..7e534922 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -373,9 +373,8 @@ mod reduce_sum_of_x { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "sve")] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] fn reduce_sum_of_x_a3_256(this: &[f32]) -> f32 { unsafe { unsafe extern "C" { @@ -385,7 +384,7 @@ mod reduce_sum_of_x { } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] #[test] fn reduce_sum_of_x_a3_256_test() { use rand::Rng; @@ -412,7 +411,7 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2")] pub fn reduce_sum_of_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -649,9 +648,8 @@ mod reduce_sum_of_abs_x { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "sve")] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] fn reduce_sum_of_abs_x_a3_256(this: &[f32]) -> f32 { unsafe { unsafe extern "C" { @@ -661,7 +659,7 @@ mod reduce_sum_of_abs_x { } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] #[test] fn reduce_sum_of_abs_x_a3_256_test() { use rand::Rng; @@ -688,7 +686,7 @@ mod reduce_sum_of_abs_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2")] pub fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -915,9 +913,8 @@ mod reduce_sum_of_x2 { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "sve")] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] fn reduce_sum_of_x2_a3_256(this: &[f32]) -> f32 { unsafe { unsafe extern "C" { @@ -927,7 +924,7 @@ mod reduce_sum_of_x2 { } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] #[test] fn reduce_sum_of_x2_a3_256_test() { use rand::Rng; @@ -954,7 +951,7 @@ mod reduce_sum_of_x2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2")] pub fn reduce_sum_of_x2(this: &[f32]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; @@ -1182,9 +1179,8 @@ mod reduce_min_max_of_x { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "sve")] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] fn reduce_min_max_of_x_a3_256(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = -f32::INFINITY; @@ -1202,7 +1198,7 @@ mod reduce_min_max_of_x { (min, max) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] #[test] fn reduce_min_max_of_x_a3_256_test() { use rand::Rng; @@ -1226,7 +1222,7 @@ mod reduce_min_max_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2")] pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -1496,9 +1492,8 @@ mod reduce_sum_of_xy { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "sve")] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] fn reduce_sum_of_xy_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { @@ -1513,7 +1508,7 @@ mod reduce_sum_of_xy { } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] #[test] fn reduce_sum_of_xy_a3_256_test() { use rand::Rng; @@ -1544,7 +1539,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2")] pub fn reduce_sum_of_xy(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -1822,9 +1817,8 @@ mod reduce_sum_of_d2 { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "sve")] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] fn reduce_sum_of_d2_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); unsafe { @@ -1839,7 +1833,7 @@ mod reduce_sum_of_d2 { } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] #[test] fn reduce_sum_of_d2_a3_256_test() { use rand::Rng; @@ -1870,7 +1864,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2")] pub fn reduce_sum_of_d2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index 43673381..c766cda5 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -435,6 +435,9 @@ mod scan { fn shuffle(a: [u8; 16], b: [u8; 16]) -> [u8; 16] { std::array::from_fn(|i| a[b[i] as usize]) } + fn transmute(x: [u8; 16]) -> [u16; 8] { + std::array::from_fn(|i| u16::from_le_bytes([x[i << 1 | 0], x[i << 1 | 1]])) + } assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -451,10 +454,10 @@ mod scan { let chi = code.map(|x| x >> 4); let lut = lut[i]; - let res_lo = zerocopy::transmute!(shuffle(lut, clo)); + let res_lo = transmute(shuffle(lut, clo)); a_0 = binary(u16::wrapping_add, a_0, res_lo); a_1 = binary(u16::wrapping_add, a_1, res_lo.map(|x| x >> 8)); - let res_hi = zerocopy::transmute!(shuffle(lut, chi)); + let res_hi = transmute(shuffle(lut, chi)); a_2 = binary(u16::wrapping_add, a_2, res_hi); a_3 = binary(u16::wrapping_add, a_3, res_hi.map(|x| x >> 8)); } diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index cb363df0..874f66f7 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -216,10 +216,16 @@ mod mul_add_round { fn mul_add_round_a2(this: &[f32], k: f32, b: f32) -> Vec { let mut r = Vec::::with_capacity(this.len()); use std::arch::aarch64::*; + #[cfg(target_endian = "little")] const CONS: [u8; 16] = [ 0, 4, 8, 12, 0xff, 0xff, 0xff, 0xff, // 0..8 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // 8..15 ]; + #[cfg(target_endian = "big")] + const CONS: [u8; 16] = [ + 12, 8, 4, 0, 0xff, 0xff, 0xff, 0xff, // 0..8 + 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // 8..15 + ]; let cons = unsafe { vld1q_u8(CONS.as_ptr()) }; let lk = vdupq_n_f32(k); let lb = vdupq_n_f32(b); diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs index a085260c..75f7b12d 100644 --- a/crates/simd_macros/src/lib.rs +++ b/crates/simd_macros/src/lib.rs @@ -15,27 +15,30 @@ mod target; struct MultiversionVersion { + attrs: Vec, target: String, import: bool, } impl syn::parse::Parse for MultiversionVersion { fn parse(input: syn::parse::ParseStream) -> syn::Result { - let lookahead1 = input.lookahead1(); - if lookahead1.peek(syn::Token![@]) { + let attrs = if input.lookahead1().peek(syn::Token![#]) { + syn::Attribute::parse_outer(input)? + } else { + Vec::new() + }; + let import = if input.lookahead1().peek(syn::Token![@]) { let _: syn::Token![@] = input.parse()?; - let target: syn::LitStr = input.parse()?; - Ok(Self { - target: target.value(), - import: true, - }) + true } else { - let target: syn::LitStr = input.parse()?; - Ok(Self { - target: target.value(), - import: false, - }) - } + false + }; + let target = input.parse::()?.value(); + Ok(Self { + attrs, + target, + import, + }) } } @@ -101,6 +104,7 @@ pub fn multiversion( let mut versions = quote::quote! {}; let mut cold = quote::quote! {}; for version in attr.versions { + let attrs = version.attrs.clone(); let target = version.target.clone(); let name = syn::Ident::new( &format!("{name}_{}", target.replace(":", "_").replace(".", "_")), @@ -116,14 +120,16 @@ pub fn multiversion( let target_cpu = target_cpu.target_cpu; if !version.import { versions.extend(quote::quote! { + #(#attrs)* #[inline] - #[cfg(any(target_arch = #target_arch))] + #[cfg(target_arch = #target_arch)] #[crate::target_cpu(enable = #target_cpu)] #(#[target_feature(enable = #additional_target_features)])* fn #name < #generics_params > (#inputs) #output #generics_where { #block } }); } cold.extend(quote::quote! { + #(#attrs)* #[cfg(target_arch = #target_arch)] if crate::is_cpu_detected!(#target_cpu) #(&& crate::is_feature_detected!(#additional_target_features))* { let ptr = unsafe { std::mem::transmute::(#name) }; diff --git a/src/lib.rs b/src/lib.rs index 4447d88b..b458d51c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -73,9 +73,6 @@ fn is_main() -> bool { #[cfg(not(panic = "unwind"))] compile_error!("This crate must be compiled with `-Cpanic=unwind`."); -#[cfg(not(target_endian = "little"))] -compile_error!("Target architecture is not supported."); - #[cfg(not(miri))] #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] #[cfg(any(target_os = "linux", target_os = "macos"))] From c464f4bea6fbd988347764ac66017faabfca1f8c Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 22 Aug 2025 13:49:12 +0800 Subject: [PATCH 200/324] fix compilation on riscv64 (#320) Signed-off-by: usamoi --- Cargo.toml | 2 +- crates/simd/src/lib.rs | 19 +++---------------- crates/simd_macros/src/lib.rs | 4 +++- 3 files changed, 7 insertions(+), 18 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 15e8d202..ff4d1bea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -86,7 +86,7 @@ clippy.needless_range_loop = "allow" rust.unsafe_code = "deny" rust.unsafe_op_in_unsafe_fn = "deny" # unused -rust.unused_crate_dependencies = "warn" +rust.unused_crate_dependencies = "allow" rust.unused_extern_crates = "warn" rust.unused_import_braces = "warn" rust.unused_lifetimes = "warn" diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 16f747cd..c6baba23 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -67,9 +67,6 @@ mod internal { #[cfg(target_arch = "aarch64")] simd_macros::define_is_cpu_detected!("aarch64"); - #[cfg(target_arch = "riscv64")] - simd_macros::define_is_cpu_detected!("riscv64"); - #[cfg(target_arch = "x86_64")] #[allow(unused_imports)] pub use is_x86_64_cpu_detected; @@ -78,10 +75,6 @@ mod internal { #[allow(unused_imports)] pub use is_aarch64_cpu_detected; - #[cfg(target_arch = "riscv64")] - #[allow(unused_imports)] - pub use is_riscv64_cpu_detected; - #[cfg(target_arch = "x86_64")] pub fn is_v4_detected() -> bool { std::arch::is_x86_feature_detected!("avx512bw") @@ -114,6 +107,7 @@ mod internal { } #[cfg(target_arch = "aarch64")] + #[cfg_attr(target_endian = "big", expect(dead_code))] pub fn is_a3_512_detected() -> bool { #[target_feature(enable = "sve")] fn is_512_detected() -> bool { @@ -130,6 +124,7 @@ mod internal { } #[cfg(target_arch = "aarch64")] + #[cfg_attr(target_endian = "big", expect(dead_code))] pub fn is_a3_256_detected() -> bool { #[target_feature(enable = "sve")] fn is_256_detected() -> bool { @@ -145,8 +140,8 @@ mod internal { std::arch::is_aarch64_feature_detected!("sve") && unsafe { is_256_detected() } } - #[expect(dead_code)] #[cfg(target_arch = "aarch64")] + #[expect(dead_code)] pub fn is_a3_128_detected() -> bool { std::arch::is_aarch64_feature_detected!("sve") } @@ -167,10 +162,6 @@ pub use std::arch::is_x86_feature_detected as is_feature_detected; #[allow(unused_imports)] pub use std::arch::is_aarch64_feature_detected as is_feature_detected; -#[cfg(target_arch = "riscv64")] -#[allow(unused_imports)] -pub use std::arch::is_riscv_feature_detected as is_feature_detected; - #[cfg(target_arch = "x86_64")] #[allow(unused_imports)] pub use internal::is_x86_64_cpu_detected as is_cpu_detected; @@ -178,7 +169,3 @@ pub use internal::is_x86_64_cpu_detected as is_cpu_detected; #[cfg(target_arch = "aarch64")] #[allow(unused_imports)] pub use internal::is_aarch64_cpu_detected as is_cpu_detected; - -#[cfg(target_arch = "riscv64")] -#[allow(unused_imports)] -pub use internal::is_riscv64_cpu_detected as is_cpu_detected; diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs index 75f7b12d..6b5b67d1 100644 --- a/crates/simd_macros/src/lib.rs +++ b/crates/simd_macros/src/lib.rs @@ -231,7 +231,9 @@ pub fn target_cpu( #[proc_macro] pub fn define_is_cpu_detected(input: proc_macro::TokenStream) -> proc_macro::TokenStream { let target_arch = syn::parse_macro_input!(input as syn::LitStr).value(); - let mut arms = quote::quote! {}; + let mut arms = quote::quote! { + () => { compile_error!("cpu is not provided") }; + }; for target_cpu in target::TARGET_CPUS { if target_cpu.target_arch != target_arch { continue; From c6bab25972adf63ab8f4a41fe211ee99ebc5c2ad Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 27 Aug 2025 11:51:46 +0800 Subject: [PATCH 201/324] fix: never use _x suffix intrinsics in SVE (#322) Signed-off-by: usamoi --- crates/simd/cshim/aarch64.c | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/crates/simd/cshim/aarch64.c b/crates/simd/cshim/aarch64.c index 5af96c3b..5b8307f3 100644 --- a/crates/simd/cshim/aarch64.c +++ b/crates/simd/cshim/aarch64.c @@ -91,7 +91,7 @@ fp16_reduce_sum_of_xy_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); svfloat16_t x = svld1_f16(mask, a + i); svfloat16_t y = svld1_f16(mask, b + i); - xy = svmla_f16_x(mask, xy, x, y); + xy = svmla_f16_m(mask, xy, x, y); } return svaddv_f16(svptrue_b16(), xy); } @@ -162,8 +162,8 @@ fp16_reduce_sum_of_d2_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); svfloat16_t x = svld1_f16(mask, a + i); svfloat16_t y = svld1_f16(mask, b + i); - svfloat16_t d = svsub_f16_x(mask, x, y); - d2 = svmla_f16_x(mask, d2, d, d); + svfloat16_t d = svsub_f16_z(mask, x, y); + d2 = svmla_f16_m(mask, d2, d, d); } return svaddv_f16(svptrue_b16(), d2); } @@ -176,7 +176,7 @@ fp32_reduce_sum_of_x_a3_256(float *restrict this, size_t n) { for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, this + i); - sum = svadd_f32_x(mask, sum, x); + sum = svadd_f32_m(mask, sum, x); } return svaddv_f32(svptrue_b32(), sum); } @@ -189,7 +189,7 @@ fp32_reduce_sum_of_abs_x_a3_256(float *restrict this, size_t n) { for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, this + i); - sum = svadd_f32_x(mask, sum, svabs_f32_x(mask, x)); + sum = svadd_f32_m(mask, sum, svabs_f32_z(mask, x)); } return svaddv_f32(svptrue_b32(), sum); } @@ -202,7 +202,7 @@ fp32_reduce_sum_of_x2_a3_256(float *restrict this, size_t n) { for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, this + i); - sum = svmla_f32_x(mask, sum, x, x); + sum = svmla_f32_m(mask, sum, x, x); } return svaddv_f32(svptrue_b32(), sum); } @@ -217,8 +217,8 @@ fp32_reduce_min_max_of_x_a3_256(float *restrict this, size_t n, float *out_min, for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, this + i); - min = svmin_f32_x(mask, min, x); - max = svmax_f32_x(mask, max, x); + min = svmin_f32_m(mask, min, x); + max = svmax_f32_m(mask, max, x); } *out_min = svminv_f32(svptrue_b32(), min); *out_max = svmaxv_f32(svptrue_b32(), max); @@ -234,7 +234,7 @@ fp32_reduce_sum_of_xy_a3_256(float *restrict lhs, float *restrict rhs, svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, lhs + i); svfloat32_t y = svld1_f32(mask, rhs + i); - sum = svmla_f32_x(mask, sum, x, y); + sum = svmla_f32_m(mask, sum, x, y); } return svaddv_f32(svptrue_b32(), sum); } @@ -249,8 +249,8 @@ fp32_reduce_sum_of_d2_a3_256(float *restrict lhs, float *restrict rhs, svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, lhs + i); svfloat32_t y = svld1_f32(mask, rhs + i); - svfloat32_t d = svsub_f32_x(mask, x, y); - sum = svmla_f32_x(mask, sum, d, d); + svfloat32_t d = svsub_f32_z(mask, x, y); + sum = svmla_f32_m(mask, sum, d, d); } return svaddv_f32(svptrue_b32(), sum); } From d9c859cfcfa4f4529b7c1a334b1a3333a5dbe734 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 27 Aug 2025 19:39:49 +0800 Subject: [PATCH 202/324] feat: s390x simd (#323) Signed-off-by: usamoi --- .github/workflows/check.yml | 54 ++++++------ .github/workflows/release.yml | 5 +- Cargo.lock | 49 +++++------ Cargo.toml | 2 +- crates/make/Cargo.toml | 2 +- crates/simd/Cargo.toml | 2 +- crates/simd/src/bit.rs | 16 ++-- crates/simd/src/f16.rs | 41 +++++---- crates/simd/src/f32.rs | 32 +++---- crates/simd/src/fast_scan.rs | 138 +++++++++++++++++++++++++------ crates/simd/src/fht.rs | 6 +- crates/simd/src/lib.rs | 71 ++++++++++++++++ crates/simd/src/quantize.rs | 2 +- crates/simd/src/rotate.rs | 6 +- crates/simd/src/u8.rs | 6 +- crates/simd_macros/src/target.rs | 73 +++++++++++++--- 16 files changed, 362 insertions(+), 143 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 962c8bdb..e2bc9c53 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -106,11 +106,12 @@ jobs: runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-24.04' || 'ubuntu-24.04-arm' }} env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: true - CARGO_TERM_COLOR: always - RUST_BACKTRACE: 1 + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" steps: - name: Set up Environment @@ -173,11 +174,12 @@ jobs: runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: true - CARGO_TERM_COLOR: always - RUST_BACKTRACE: 1 - RUSTFLAGS: -Dwarnings + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" steps: - name: Set up Environment @@ -259,11 +261,12 @@ jobs: runs-on: "macos-15" env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: true - CARGO_TERM_COLOR: always - RUST_BACKTRACE: 1 - RUSTFLAGS: -Dwarnings + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" steps: - name: Set up Environment @@ -343,11 +346,12 @@ jobs: runs-on: "windows-2022" env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: true - CARGO_TERM_COLOR: always - RUST_BACKTRACE: 1 - RUSTFLAGS: -Dwarnings + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" steps: - name: Set up Environment @@ -452,11 +456,13 @@ jobs: - /opt:/__e/node20:ro,rshared env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: true - CARGO_TERM_COLOR: always - RUST_BACKTRACE: 1 - USER: root + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + # RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + USER: "root" steps: - name: Set up Environment diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index b03b5499..6e0d0ef0 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -44,9 +44,10 @@ jobs: runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} env: - CARGO_TERM_COLOR: always - RUST_BACKTRACE: 1 + RUSTUP_AUTO_INSTALL: "0" RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" steps: - name: Set up Environment diff --git a/Cargo.lock b/Cargo.lock index f0b41cd7..1115fc9c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -123,9 +123,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.9.2" +version = "2.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a65b545ab31d687cff52899d4890855fec459eb6afe0da6417b8a18da87aa29" +checksum = "34efbcccd345379ca2868b2b2c9d3782e9cc58ba87bc7d79d5b53d9c9ae6f25d" [[package]] name = "bitvec" @@ -157,9 +157,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.33" +version = "1.2.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ee0f8803222ba5a7e2777dd72ca451868909b1ac410621b676adf07280e9b5f" +checksum = "42bc4aea80032b7bf409b0bc7ccad88853858911b7713a8062fdc0623867bedc" dependencies = [ "shlex", ] @@ -202,9 +202,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.45" +version = "4.5.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fc0e74a703892159f5ae7d3aac52c8e6c392f5ae5f359c70b5881d60aaac318" +checksum = "2c5e4fcf9c21d2e544ca1ee9d8552de13019a42aa7dbf32747fa7aaf1df76e57" dependencies = [ "clap_builder", "clap_derive", @@ -212,9 +212,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.44" +version = "4.5.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3e7f4214277f3c7aa526a59dd3fbe306a370daee1f8b7b8c987069cd8e888a8" +checksum = "fecb53a0e6fcfb055f686001bc2e2592fa527efaf38dbe81a6a9563562e57d41" dependencies = [ "anstream", "anstyle", @@ -642,9 +642,9 @@ checksum = "964de6e86d545b246d84badc0fef527924ace5134f30641c203ef52ba83f58d5" [[package]] name = "indexmap" -version = "2.10.0" +version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe4cd85333e22411419a0bcae1297d25e58c9443848b11dc6a86fefe8c78a661" +checksum = "f2481980430f9f78649238835720ddccc57e52df14ffce1c6f37391d61b563e9" dependencies = [ "equivalent", "hashbrown", @@ -726,9 +726,9 @@ dependencies = [ [[package]] name = "libmimalloc-sys" -version = "0.1.43" +version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf88cd67e9de251c1781dbe2f641a1a3ad66eaae831b8a2c38fbdc5ddae16d4d" +checksum = "667f4fec20f29dfc6bc7357c582d91796c169ad7e2fce709468aefeb2c099870" dependencies = [ "cc", "libc", @@ -770,9 +770,9 @@ checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" [[package]] name = "mimalloc" -version = "0.1.47" +version = "0.1.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1791cbe101e95af5764f06f20f6760521f7158f69dbf9d6baf941ee1bf6bc40" +checksum = "e1ee66a4b64c74f4ef288bcbb9192ad9c3feaad75193129ac8509af543894fd8" dependencies = [ "libmimalloc-sys", ] @@ -1131,9 +1131,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.11.1" +version = "1.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191" +checksum = "23d7fd106d8c02486a8d64e778353d1cffe08ce79ac2e82f540c86d0facf6912" dependencies = [ "aho-corasick", "memchr", @@ -1143,9 +1143,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.9" +version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "809e8dc61f6de73b46c85f4c96486310fe304c434cfa43669d7b40f711150908" +checksum = "6b9458fa0bfeeac22b5ca447c63aaf45f28439a709ccd244698632f9aa6394d6" dependencies = [ "aho-corasick", "memchr", @@ -1154,9 +1154,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.5" +version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" +checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" [[package]] name = "rustc-hash" @@ -1515,13 +1515,14 @@ checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" [[package]] name = "url" -version = "2.5.5" +version = "2.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec961601b32b6f5d14ae8dabd35ff2ff2e2c6cb4c0e6641845ff105abe96d958" +checksum = "08bc136a29a3d1758e07a9cca267be308aeebf5cfd5a10f3f67ab2097683ef5b" dependencies = [ "form_urlencoded", "idna", "percent-encoding", + "serde", ] [[package]] @@ -1933,9 +1934,9 @@ checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486" [[package]] name = "winnow" -version = "0.7.12" +version = "0.7.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3edebf492c8125044983378ecb5766203ad3b4c2f7a922bd7dd207f6d443e95" +checksum = "21a0236b59786fed61e2a80582dd500fe61f18b5dca67a4a067d0bc9039339cf" dependencies = [ "memchr", ] diff --git a/Cargo.toml b/Cargo.toml index ff4d1bea..2be9b7af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,7 +47,7 @@ validator.workspace = true zerocopy.workspace = true [target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos")))'.dependencies] -mimalloc = { version = "0.1.47", features = ["local_dynamic_tls"] } +mimalloc = { version = "0.1.48", features = ["local_dynamic_tls"] } [lints] workspace = true diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index d4565466..eb92030f 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -5,7 +5,7 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.45", features = ["derive"] } +clap = { version = "4.5.46", features = ["derive"] } object = { version = "0.37.3", default-features = false, features = [ "read", "wasm", diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 98a9bd5c..764836ba 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -18,7 +18,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.33" +cc = "1.2.34" which = "8.0.0" [lints] diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index 65f342e2..7538f7f8 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -183,7 +183,7 @@ mod reduce_sum_of_and { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -366,7 +366,7 @@ mod reduce_sum_of_or { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -549,7 +549,7 @@ mod reduce_sum_of_xor { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -772,7 +772,7 @@ mod reduce_sum_of_and_or { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -933,7 +933,7 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_x(this: &[u64]) -> u32 { let n = this.len(); let mut sum = 0; @@ -950,7 +950,7 @@ pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_and { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -973,7 +973,7 @@ pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_or { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -996,7 +996,7 @@ pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_xor { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index 95686e0c..c34f12aa 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -14,7 +14,6 @@ use crate::{Floating, f32}; use half::f16; -use zerocopy::FromZeros; trait AsF32 { #[allow(clippy::wrong_self_convention)] @@ -30,7 +29,7 @@ impl AsF32 for f16 { impl Floating for f16 { #[inline(always)] fn zero() -> Self { - FromZeros::new_zeroed() + f16::ZERO } #[inline(always)] @@ -162,10 +161,10 @@ impl Floating for f16 { mod reduce_or_of_is_zero_x { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { for &x in this { - if x == FromZeros::new_zeroed() { + if x == f16::ZERO { return true; } } @@ -178,7 +177,7 @@ mod reduce_sum_of_x { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; @@ -194,7 +193,7 @@ mod reduce_sum_of_abs_x { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; @@ -210,7 +209,7 @@ mod reduce_sum_of_x2 { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_x2(this: &[f16]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; @@ -226,7 +225,7 @@ mod reduce_min_max_of_x { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -502,7 +501,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -785,7 +784,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -804,7 +803,7 @@ mod reduce_sum_of_xy_sparse { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -837,7 +836,7 @@ mod reduce_sum_of_d2_sparse { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -876,7 +875,7 @@ mod reduce_sum_of_d2_sparse { mod vector_add { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_add(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -896,7 +895,7 @@ mod vector_add { mod vector_add_inplace { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_add_inplace(lhs: &mut [f16], rhs: &[f16]) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -909,7 +908,7 @@ mod vector_add_inplace { mod vector_sub { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_sub(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -929,7 +928,7 @@ mod vector_sub { mod vector_mul { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_mul(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -949,7 +948,7 @@ mod vector_mul { mod vector_mul_scalar { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_mul_scalar(lhs: &[f16], rhs: f32) -> Vec { let rhs = f16::from_f32(rhs); let n = lhs.len(); @@ -969,7 +968,7 @@ mod vector_mul_scalar { mod vector_mul_scalar_inplace { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_mul_scalar_inplace(lhs: &mut [f16], rhs: f32) { let rhs = f16::from_f32(rhs); let n = lhs.len(); @@ -982,7 +981,7 @@ mod vector_mul_scalar_inplace { mod vector_abs_inplace { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_abs_inplace(this: &mut [f16]) { let n = this.len(); for i in 0..n { @@ -994,7 +993,7 @@ mod vector_abs_inplace { mod vector_from_f32 { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_from_f32(this: &[f32]) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); @@ -1013,7 +1012,7 @@ mod vector_from_f32 { mod vector_to_f32 { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_to_f32(this: &[f16]) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index 7e534922..01136328 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -147,7 +147,7 @@ impl Floating for f32 { } mod reduce_or_of_is_zero_x { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { for &x in this { if x == 0.0f32 { @@ -411,7 +411,7 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -686,7 +686,7 @@ mod reduce_sum_of_abs_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -951,7 +951,7 @@ mod reduce_sum_of_x2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_x2(this: &[f32]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; @@ -1222,7 +1222,7 @@ mod reduce_min_max_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -1539,7 +1539,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_xy(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -1864,7 +1864,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_d2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -1966,7 +1966,7 @@ mod reduce_sum_of_xy_sparse { } } - #[crate::multiversion(@"v4", "v3", "v2", "a2")] + #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -2116,7 +2116,7 @@ mod reduce_sum_of_d2_sparse { } } - #[crate::multiversion(@"v4", "v3", "v2", "a2")] + #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -2153,7 +2153,7 @@ mod reduce_sum_of_d2_sparse { } mod vector_add { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2171,7 +2171,7 @@ mod vector_add { } mod vector_add_inplace { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2182,7 +2182,7 @@ mod vector_add_inplace { } mod vector_sub { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2200,7 +2200,7 @@ mod vector_sub { } mod vector_mul { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2218,7 +2218,7 @@ mod vector_mul { } mod vector_mul_scalar { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { let n = lhs.len(); let mut r = Vec::::with_capacity(n); @@ -2235,7 +2235,7 @@ mod vector_mul_scalar { } mod vector_mul_scalar_inplace { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { let n = lhs.len(); for i in 0..n { @@ -2245,7 +2245,7 @@ mod vector_mul_scalar_inplace { } mod vector_abs_inplace { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn vector_abs_inplace(this: &mut [f32]) { let n = this.len(); for i in 0..n { diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index c766cda5..2a3032a2 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -427,45 +427,133 @@ mod scan { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] - pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { - fn binary(op: impl Fn(u16, u16) -> u16, a: [u16; 8], b: [u16; 8]) -> [u16; 8] { - std::array::from_fn(|i| op(a[i], b[i])) + #[cfg(target_arch = "s390x")] + #[crate::target_cpu(enable = "z13")] + fn scan_z13(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + unsafe { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + use std::arch::s390x::*; + use std::mem::transmute; + use {vector_unsigned_char as u8x16, vector_unsigned_short as u16x8}; + + let _0001_u16x8 = vec_splat_u16::<0x0001>(); + let _00ff_u16x8 = vec_splat_u16::<0x00ff>(); + let _ff00_u16x8 = vec_splat_u16::<{ 0xff00u16 as i16 }>(); + + let mut accu_0 = vec_splat_u16::<0>(); + let mut accu_1 = vec_splat_u16::<0>(); + let mut accu_2 = vec_splat_u16::<0>(); + let mut accu_3 = vec_splat_u16::<0>(); + + let mut i = 0_usize; + while i < n { + let code: u8x16 = vec_xl((i as isize) * 16, code.as_ptr().cast()); + + let clo = vec_and(code, vec_splat_u8::<0xf>()); + let chi = vec_srl(code, vec_splat_u8::<4>()); + + let lut: u8x16 = vec_xl((i as isize) * 16, lut.as_ptr().cast()); + let res_lo = vec_revb(transmute::(vec_perm(lut, lut, clo))); + accu_0 = vec_add(accu_0, res_lo); + accu_1 = vec_add(accu_1, vec_and(vec_rli(res_lo, 8), _00ff_u16x8)); + let res_hi = vec_revb(transmute::(vec_perm(lut, lut, chi))); + accu_2 = vec_add(accu_2, res_hi); + accu_3 = vec_add(accu_3, vec_and(vec_rli(res_hi, 8), _00ff_u16x8)); + + i += 1; + } + debug_assert_eq!(i, n); + + let mut result = [0_u16; 32]; + + accu_0 = vec_sub(accu_0, vec_and(vec_rli(accu_1, 120), _ff00_u16x8)); + vec_xst(accu_0, 0, result.as_mut_ptr().cast()); + vec_xst(accu_1, 16, result.as_mut_ptr().cast()); + + accu_2 = vec_sub(accu_2, vec_and(vec_rli(accu_3, 120), _ff00_u16x8)); + vec_xst(accu_2, 32, result.as_mut_ptr().cast()); + vec_xst(accu_3, 48, result.as_mut_ptr().cast()); + + result } - fn shuffle(a: [u8; 16], b: [u8; 16]) -> [u8; 16] { - std::array::from_fn(|i| a[b[i] as usize]) + } + + #[cfg(all(target_arch = "s390x", test, not(miri)))] + #[test] + fn scan_z13_test() { + if !crate::is_cpu_detected!("z13") { + println!("test {} ... skipped (z13)", module_path!()); + return; } - fn transmute(x: [u8; 16]) -> [u16; 8] { - std::array::from_fn(|i| u16::from_le_bytes([x[i << 1 | 0], x[i << 1 | 1]])) + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + for n in 90..110 { + let code = (0..n) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..n) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + unsafe { + assert_eq!(scan_z13(&code, &lut), fallback(&code, &lut)); + } + } } + } + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", @"z13")] + pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { assert_eq!(code.len(), lut.len()); let n = code.len(); - let mut a_0 = [0u16; 8]; - let mut a_1 = [0u16; 8]; - let mut a_2 = [0u16; 8]; - let mut a_3 = [0u16; 8]; + let mut result = [0u16; 32]; for i in 0..n { let code = code[i]; - let clo = code.map(|x| x & 0xf); let chi = code.map(|x| x >> 4); - let lut = lut[i]; - let res_lo = transmute(shuffle(lut, clo)); - a_0 = binary(u16::wrapping_add, a_0, res_lo); - a_1 = binary(u16::wrapping_add, a_1, res_lo.map(|x| x >> 8)); - let res_hi = transmute(shuffle(lut, chi)); - a_2 = binary(u16::wrapping_add, a_2, res_hi); - a_3 = binary(u16::wrapping_add, a_3, res_hi.map(|x| x >> 8)); - } - a_0 = binary(u16::wrapping_sub, a_0, a_1.map(|x| x << 8)); - a_2 = binary(u16::wrapping_sub, a_2, a_3.map(|x| x << 8)); + let res_lo: [u8; 16] = std::array::from_fn(|i| lut[clo[i] as usize]); + let res_hi: [u8; 16] = std::array::from_fn(|i| lut[chi[i] as usize]); + + result[0x00] += res_lo[0x0] as u16; + result[0x01] += res_lo[0x2] as u16; + result[0x02] += res_lo[0x4] as u16; + result[0x03] += res_lo[0x6] as u16; + result[0x04] += res_lo[0x8] as u16; + result[0x05] += res_lo[0xa] as u16; + result[0x06] += res_lo[0xc] as u16; + result[0x07] += res_lo[0xe] as u16; + result[0x08] += res_lo[0x1] as u16; + result[0x09] += res_lo[0x3] as u16; + result[0x0a] += res_lo[0x5] as u16; + result[0x0b] += res_lo[0x7] as u16; + result[0x0c] += res_lo[0x9] as u16; + result[0x0d] += res_lo[0xb] as u16; + result[0x0e] += res_lo[0xd] as u16; + result[0x0f] += res_lo[0xf] as u16; + result[0x10] += res_hi[0x0] as u16; + result[0x11] += res_hi[0x2] as u16; + result[0x12] += res_hi[0x4] as u16; + result[0x13] += res_hi[0x6] as u16; + result[0x14] += res_hi[0x8] as u16; + result[0x15] += res_hi[0xa] as u16; + result[0x16] += res_hi[0xc] as u16; + result[0x17] += res_hi[0xe] as u16; + result[0x18] += res_hi[0x1] as u16; + result[0x19] += res_hi[0x3] as u16; + result[0x1a] += res_hi[0x5] as u16; + result[0x1b] += res_hi[0x7] as u16; + result[0x1c] += res_hi[0x9] as u16; + result[0x1d] += res_hi[0xb] as u16; + result[0x1e] += res_hi[0xd] as u16; + result[0x1f] += res_hi[0xf] as u16; + } - zerocopy::transmute!([a_0, a_1, a_2, a_3]) + result } } @@ -475,7 +563,7 @@ pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { } mod accu { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn accu(sum: &mut [u32; 32], delta: &[u16; 32]) { for i in 0..32 { sum[i] += delta[i] as u32; diff --git a/crates/simd/src/fht.rs b/crates/simd/src/fht.rs index be277473..b45b13c8 100644 --- a/crates/simd/src/fht.rs +++ b/crates/simd/src/fht.rs @@ -33,7 +33,7 @@ mod step_1 { seq_macro::seq!( Q in 0..16 { mod dispatch_~Q { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn f(x: &mut [f32]) { crate::fht::basic_1::(x); } @@ -48,7 +48,7 @@ mod step_2 { seq_macro::seq!( Q in 0..16 { mod dispatch_~Q { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn f(x: &mut [f32]) { crate::fht::basic_2::(x); } @@ -62,7 +62,7 @@ mod step_2 { macro_rules! fht { ($p:literal, 0) => { { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] fn walk(x: &mut [f32]) { assert!(x.len() == (1 << $p)); seq_macro::seq!( diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index c6baba23..90f2db6c 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -13,6 +13,9 @@ // Copyright (c) 2025 TensorChord Inc. #![allow(unsafe_code)] +#![cfg_attr(target_arch = "s390x", feature(stdarch_s390x_feature_detection))] +#![cfg_attr(target_arch = "s390x", feature(s390x_target_feature))] +#![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] mod aligned; mod emulate; @@ -67,6 +70,9 @@ mod internal { #[cfg(target_arch = "aarch64")] simd_macros::define_is_cpu_detected!("aarch64"); + #[cfg(target_arch = "s390x")] + simd_macros::define_is_cpu_detected!("s390x"); + #[cfg(target_arch = "x86_64")] #[allow(unused_imports)] pub use is_x86_64_cpu_detected; @@ -75,6 +81,10 @@ mod internal { #[allow(unused_imports)] pub use is_aarch64_cpu_detected; + #[cfg(target_arch = "s390x")] + #[allow(unused_imports)] + pub use is_s390x_cpu_detected; + #[cfg(target_arch = "x86_64")] pub fn is_v4_detected() -> bool { std::arch::is_x86_feature_detected!("avx512bw") @@ -150,6 +160,59 @@ mod internal { pub fn is_a2_detected() -> bool { std::arch::is_aarch64_feature_detected!("neon") } + + #[cfg(target_arch = "s390x")] + pub fn is_z17_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + && std::arch::is_s390x_feature_detected!("vector-enhancements-1") + && std::arch::is_s390x_feature_detected!("vector-enhancements-2") + && std::arch::is_s390x_feature_detected!("vector-enhancements-3") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement-2") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement-3") + && std::arch::is_s390x_feature_detected!("nnp-assist") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-2") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-3") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-4") + } + + #[cfg(target_arch = "s390x")] + pub fn is_z16_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + && std::arch::is_s390x_feature_detected!("vector-enhancements-1") + && std::arch::is_s390x_feature_detected!("vector-enhancements-2") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement-2") + && std::arch::is_s390x_feature_detected!("nnp-assist") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-2") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-3") + } + + #[cfg(target_arch = "s390x")] + pub fn is_z15_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + && std::arch::is_s390x_feature_detected!("vector-enhancements-1") + && std::arch::is_s390x_feature_detected!("vector-enhancements-2") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-2") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-3") + } + + #[cfg(target_arch = "s390x")] + pub fn is_z14_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + && std::arch::is_s390x_feature_detected!("vector-enhancements-1") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-2") + } + + #[cfg(target_arch = "s390x")] + pub fn is_z13_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + } } pub use simd_macros::{multiversion, target_cpu}; @@ -162,6 +225,10 @@ pub use std::arch::is_x86_feature_detected as is_feature_detected; #[allow(unused_imports)] pub use std::arch::is_aarch64_feature_detected as is_feature_detected; +#[cfg(target_arch = "s390x")] +#[allow(unused_imports)] +pub use std::arch::is_s390x_feature_detected as is_feature_detected; + #[cfg(target_arch = "x86_64")] #[allow(unused_imports)] pub use internal::is_x86_64_cpu_detected as is_cpu_detected; @@ -169,3 +236,7 @@ pub use internal::is_x86_64_cpu_detected as is_cpu_detected; #[cfg(target_arch = "aarch64")] #[allow(unused_imports)] pub use internal::is_aarch64_cpu_detected as is_cpu_detected; + +#[cfg(target_arch = "s390x")] +#[allow(unused_imports)] +pub use internal::is_s390x_cpu_detected as is_cpu_detected; diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index 874f66f7..eb87baaf 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -283,7 +283,7 @@ mod mul_add_round { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn mul_add_round(this: &[f32], k: f32, b: f32) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); diff --git a/crates/simd/src/rotate.rs b/crates/simd/src/rotate.rs index 00642bd3..9bf99ca7 100644 --- a/crates/simd/src/rotate.rs +++ b/crates/simd/src/rotate.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. pub mod givens { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -29,10 +29,10 @@ pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { } pub mod flip { - #[crate::multiversion("v4", "v3", "v2", "a2")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] pub fn flip(bits: &[u64; 1024], result: &mut [f32]) { use std::hint::select_unpredictable; - let result: &mut [u32] = unsafe { std::mem::transmute(result) }; + let result: &mut [u32] = zerocopy::transmute_mut!(result); let (arrays, remainder) = result.as_chunks_mut::<64>(); let n = arrays.len(); assert!(n <= 1024); diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index 51b55435..b21145bd 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -311,7 +311,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); @@ -517,7 +517,7 @@ mod reduce_sum_of_x_as_u16 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_x_as_u16(this: &[u8]) -> u16 { let n = this.len(); let mut sum = 0; @@ -722,7 +722,7 @@ mod reduce_sum_of_x_as_u32 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13")] pub fn reduce_sum_of_x_as_u32(this: &[u8]) -> u32 { let n = this.len(); let mut sum = 0; diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs index 720abc6c..4030c5f4 100644 --- a/crates/simd_macros/src/target.rs +++ b/crates/simd_macros/src/target.rs @@ -46,29 +46,82 @@ pub const TARGET_CPUS: &[TargetCpu] = &[ TargetCpu { target_cpu: "a3.512", target_arch: "aarch64", - target_features: &[ - "sve", // simd - ], + target_features: &["sve"], }, TargetCpu { target_cpu: "a3.256", target_arch: "aarch64", - target_features: &[ - "sve", // simd - ], + target_features: &["sve"], }, TargetCpu { target_cpu: "a3.128", target_arch: "aarch64", - target_features: &[ - "sve", // simd - ], + target_features: &["sve"], }, TargetCpu { target_cpu: "a2", target_arch: "aarch64", + target_features: &["neon"], + }, + TargetCpu { + target_cpu: "z17", + target_arch: "s390x", + target_features: &[ + "vector", + "vector-enhancements-1", + "vector-enhancements-2", + "vector-enhancements-3", + "vector-packed-decimal", + "vector-packed-decimal-enhancement", + "vector-packed-decimal-enhancement-2", + "vector-packed-decimal-enhancement-3", + "nnp-assist", + "miscellaneous-extensions-2", + "miscellaneous-extensions-3", + "miscellaneous-extensions-4", + ], + }, + TargetCpu { + target_cpu: "z16", + target_arch: "s390x", + target_features: &[ + "vector", + "vector-enhancements-1", + "vector-enhancements-2", + "vector-packed-decimal", + "vector-packed-decimal-enhancement", + "vector-packed-decimal-enhancement-2", + "nnp-assist", + "miscellaneous-extensions-2", + "miscellaneous-extensions-3", + ], + }, + TargetCpu { + target_cpu: "z15", + target_arch: "s390x", target_features: &[ - "neon", // simd + "vector", + "vector-enhancements-1", + "vector-enhancements-2", + "vector-packed-decimal", + "vector-packed-decimal-enhancement", + "miscellaneous-extensions-2", + "miscellaneous-extensions-3", ], }, + TargetCpu { + target_cpu: "z14", + target_arch: "s390x", + target_features: &[ + "vector", + "vector-enhancements-1", + "vector-packed-decimal", + "miscellaneous-extensions-2", + ], + }, + TargetCpu { + target_cpu: "z13", + target_arch: "s390x", + target_features: &["vector"], + }, ]; From ab26411d7f9ec7d64d17286c7ab4c7b6dbc303a1 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 28 Aug 2025 11:14:14 +0800 Subject: [PATCH 203/324] feat: ppc64le simd (#324) Signed-off-by: usamoi --- crates/simd/src/bit.rs | 22 +++++--- crates/simd/src/f16.rs | 68 ++++++++++++++++------ crates/simd/src/f32.rs | 48 ++++++++++------ crates/simd/src/fast_scan.rs | 96 +++++++++++++++++++++++++++++--- crates/simd/src/fht.rs | 6 +- crates/simd/src/lib.rs | 44 +++++++++++++++ crates/simd/src/quantize.rs | 2 +- crates/simd/src/rotate.rs | 8 ++- crates/simd/src/u8.rs | 6 +- crates/simd_macros/src/target.rs | 29 ++++++++++ 10 files changed, 270 insertions(+), 59 deletions(-) diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index 7538f7f8..7bc4305a 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -183,7 +183,7 @@ mod reduce_sum_of_and { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -366,7 +366,7 @@ mod reduce_sum_of_or { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -549,7 +549,7 @@ mod reduce_sum_of_xor { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -772,7 +772,7 @@ mod reduce_sum_of_and_or { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -933,7 +933,7 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_x(this: &[u64]) -> u32 { let n = this.len(); let mut sum = 0; @@ -950,7 +950,9 @@ pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_and { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -973,7 +975,9 @@ pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_or { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -996,7 +1000,9 @@ pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { } mod vector_xor { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); diff --git a/crates/simd/src/f16.rs b/crates/simd/src/f16.rs index c34f12aa..eeaef4ff 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/f16.rs @@ -161,7 +161,9 @@ impl Floating for f16 { mod reduce_or_of_is_zero_x { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { for &x in this { if x == f16::ZERO { @@ -177,7 +179,9 @@ mod reduce_sum_of_x { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn reduce_sum_of_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; @@ -193,7 +197,9 @@ mod reduce_sum_of_abs_x { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; @@ -209,7 +215,9 @@ mod reduce_sum_of_x2 { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn reduce_sum_of_x2(this: &[f16]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; @@ -225,7 +233,9 @@ mod reduce_min_max_of_x { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -501,7 +511,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -784,7 +794,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -803,7 +813,9 @@ mod reduce_sum_of_xy_sparse { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -836,7 +848,9 @@ mod reduce_sum_of_d2_sparse { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -875,7 +889,9 @@ mod reduce_sum_of_d2_sparse { mod vector_add { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_add(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -895,7 +911,9 @@ mod vector_add { mod vector_add_inplace { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_add_inplace(lhs: &mut [f16], rhs: &[f16]) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -908,7 +926,9 @@ mod vector_add_inplace { mod vector_sub { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_sub(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -928,7 +948,9 @@ mod vector_sub { mod vector_mul { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_mul(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -948,7 +970,9 @@ mod vector_mul { mod vector_mul_scalar { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_mul_scalar(lhs: &[f16], rhs: f32) -> Vec { let rhs = f16::from_f32(rhs); let n = lhs.len(); @@ -968,7 +992,9 @@ mod vector_mul_scalar { mod vector_mul_scalar_inplace { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_mul_scalar_inplace(lhs: &mut [f16], rhs: f32) { let rhs = f16::from_f32(rhs); let n = lhs.len(); @@ -981,7 +1007,9 @@ mod vector_mul_scalar_inplace { mod vector_abs_inplace { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_abs_inplace(this: &mut [f16]) { let n = this.len(); for i in 0..n { @@ -993,7 +1021,9 @@ mod vector_abs_inplace { mod vector_from_f32 { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_from_f32(this: &[f32]) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); @@ -1012,7 +1042,9 @@ mod vector_from_f32 { mod vector_to_f32 { use super::*; - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_to_f32(this: &[f16]) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); diff --git a/crates/simd/src/f32.rs b/crates/simd/src/f32.rs index 01136328..13aae3da 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/f32.rs @@ -147,7 +147,9 @@ impl Floating for f32 { } mod reduce_or_of_is_zero_x { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { for &x in this { if x == 0.0f32 { @@ -411,7 +413,7 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -686,7 +688,7 @@ mod reduce_sum_of_abs_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -951,7 +953,7 @@ mod reduce_sum_of_x2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_x2(this: &[f32]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; @@ -1222,7 +1224,7 @@ mod reduce_min_max_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -1539,7 +1541,7 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_xy(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -1864,7 +1866,7 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_d2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -1966,7 +1968,7 @@ mod reduce_sum_of_xy_sparse { } } - #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -2116,7 +2118,7 @@ mod reduce_sum_of_d2_sparse { } } - #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -2153,7 +2155,9 @@ mod reduce_sum_of_d2_sparse { } mod vector_add { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2171,7 +2175,9 @@ mod vector_add { } mod vector_add_inplace { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2182,7 +2188,9 @@ mod vector_add_inplace { } mod vector_sub { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2200,7 +2208,9 @@ mod vector_sub { } mod vector_mul { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -2218,7 +2228,9 @@ mod vector_mul { } mod vector_mul_scalar { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { let n = lhs.len(); let mut r = Vec::::with_capacity(n); @@ -2235,7 +2247,9 @@ mod vector_mul_scalar { } mod vector_mul_scalar_inplace { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { let n = lhs.len(); for i in 0..n { @@ -2245,7 +2259,9 @@ mod vector_mul_scalar_inplace { } mod vector_abs_inplace { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn vector_abs_inplace(this: &mut [f32]) { let n = this.len(); for i in 0..n { diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index 2a3032a2..59040fa9 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -456,12 +456,14 @@ mod scan { let chi = vec_srl(code, vec_splat_u8::<4>()); let lut: u8x16 = vec_xl((i as isize) * 16, lut.as_ptr().cast()); - let res_lo = vec_revb(transmute::(vec_perm(lut, lut, clo))); + let res_lo_r = transmute::(vec_perm(lut, lut, clo)); + let res_lo = vec_revb(res_lo_r); accu_0 = vec_add(accu_0, res_lo); - accu_1 = vec_add(accu_1, vec_and(vec_rli(res_lo, 8), _00ff_u16x8)); - let res_hi = vec_revb(transmute::(vec_perm(lut, lut, chi))); + accu_1 = vec_add(accu_1, vec_and(res_lo_r, _00ff_u16x8)); + let res_hi_r = transmute::(vec_perm(lut, lut, chi)); + let res_hi = vec_revb(res_hi_r); accu_2 = vec_add(accu_2, res_hi); - accu_3 = vec_add(accu_3, vec_and(vec_rli(res_hi, 8), _00ff_u16x8)); + accu_3 = vec_add(accu_3, vec_and(res_hi_r, _00ff_u16x8)); i += 1; } @@ -469,11 +471,11 @@ mod scan { let mut result = [0_u16; 32]; - accu_0 = vec_sub(accu_0, vec_and(vec_rli(accu_1, 120), _ff00_u16x8)); + accu_0 = vec_sub(accu_0, vec_and(vec_revb(accu_1), _ff00_u16x8)); vec_xst(accu_0, 0, result.as_mut_ptr().cast()); vec_xst(accu_1, 16, result.as_mut_ptr().cast()); - accu_2 = vec_sub(accu_2, vec_and(vec_rli(accu_3, 120), _ff00_u16x8)); + accu_2 = vec_sub(accu_2, vec_and(vec_revb(accu_3), _ff00_u16x8)); vec_xst(accu_2, 32, result.as_mut_ptr().cast()); vec_xst(accu_3, 48, result.as_mut_ptr().cast()); @@ -503,7 +505,83 @@ mod scan { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", @"z13")] + #[cfg(target_arch = "powerpc64")] + #[crate::target_cpu(enable = "p7")] + fn scan_p7(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + unsafe { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + use std::arch::powerpc64::*; + use std::mem::transmute; + use {vector_unsigned_char as u8x16, vector_unsigned_short as u16x8}; + + let _0008_u16x8 = vec_splat_u16::<0x0008>(); + let _00ff_u16x8 = vec_splat_u16::<{ 0x00ffu8 as i8 }>(); + let _ff00_u16x8 = vec_splats(0xff00u16); + + let mut accu_0 = vec_splat_u16::<0>(); + let mut accu_1 = vec_splat_u16::<0>(); + let mut accu_2 = vec_splat_u16::<0>(); + let mut accu_3 = vec_splat_u16::<0>(); + + let mut i = 0_usize; + while i < n { + let code: u8x16 = vec_xl((i as isize) * 16, code.as_ptr().cast::()); + + let clo = vec_and(code, vec_splat_u8::<0xf>()); + let chi = vec_srl(code, vec_splat_u8::<4>()); + + let lut: u8x16 = vec_xl((i as isize) * 16, lut.as_ptr().cast::()); + let res_lo = transmute::(vec_perm(lut, lut, clo)); + accu_0 = vec_add(accu_0, res_lo); + accu_1 = vec_add(accu_1, vec_sr(res_lo, _0008_u16x8)); + let res_hi = transmute::(vec_perm(lut, lut, chi)); + accu_2 = vec_add(accu_2, res_hi); + accu_3 = vec_add(accu_3, vec_sr(res_hi, _0008_u16x8)); + + i += 1; + } + debug_assert_eq!(i, n); + + let mut result = [0_u16; 32]; + + accu_0 = vec_sub(accu_0, vec_sl(accu_1, _0008_u16x8)); + vec_xst(accu_0, 0, result.as_mut_ptr().cast()); + vec_xst(accu_1, 16, result.as_mut_ptr().cast()); + + accu_2 = vec_sub(accu_2, vec_sl(accu_3, _0008_u16x8)); + vec_xst(accu_2, 32, result.as_mut_ptr().cast()); + vec_xst(accu_3, 48, result.as_mut_ptr().cast()); + + result + } + } + + #[cfg(all(target_arch = "powerpc64", test, not(miri)))] + #[test] + fn scan_p7_test() { + if !crate::is_cpu_detected!("p7") { + println!("test {} ... skipped (p7)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + for n in 90..110 { + let code = (0..n) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..n) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + unsafe { + assert_eq!(scan_p7(&code, &lut), fallback(&code, &lut)); + } + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", @"z13", @"p7")] pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { assert_eq!(code.len(), lut.len()); let n = code.len(); @@ -563,7 +641,9 @@ pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { } mod accu { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn accu(sum: &mut [u32; 32], delta: &[u16; 32]) { for i in 0..32 { sum[i] += delta[i] as u32; diff --git a/crates/simd/src/fht.rs b/crates/simd/src/fht.rs index b45b13c8..84238b48 100644 --- a/crates/simd/src/fht.rs +++ b/crates/simd/src/fht.rs @@ -33,7 +33,7 @@ mod step_1 { seq_macro::seq!( Q in 0..16 { mod dispatch_~Q { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn f(x: &mut [f32]) { crate::fht::basic_1::(x); } @@ -48,7 +48,7 @@ mod step_2 { seq_macro::seq!( Q in 0..16 { mod dispatch_~Q { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn f(x: &mut [f32]) { crate::fht::basic_2::(x); } @@ -62,7 +62,7 @@ mod step_2 { macro_rules! fht { ($p:literal, 0) => { { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] fn walk(x: &mut [f32]) { assert!(x.len() == (1 << $p)); seq_macro::seq!( diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 90f2db6c..47600f7a 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -16,6 +16,9 @@ #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x_feature_detection))] #![cfg_attr(target_arch = "s390x", feature(s390x_target_feature))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] +#![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc_feature_detection))] +#![cfg_attr(target_arch = "powerpc64", feature(powerpc_target_feature))] +#![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc))] mod aligned; mod emulate; @@ -73,6 +76,9 @@ mod internal { #[cfg(target_arch = "s390x")] simd_macros::define_is_cpu_detected!("s390x"); + #[cfg(target_arch = "powerpc64")] + simd_macros::define_is_cpu_detected!("powerpc64"); + #[cfg(target_arch = "x86_64")] #[allow(unused_imports)] pub use is_x86_64_cpu_detected; @@ -85,6 +91,10 @@ mod internal { #[allow(unused_imports)] pub use is_s390x_cpu_detected; + #[cfg(target_arch = "powerpc64")] + #[allow(unused_imports)] + pub use is_powerpc64_cpu_detected; + #[cfg(target_arch = "x86_64")] pub fn is_v4_detected() -> bool { std::arch::is_x86_feature_detected!("avx512bw") @@ -213,6 +223,32 @@ mod internal { pub fn is_z13_detected() -> bool { std::arch::is_s390x_feature_detected!("vector") } + + #[cfg(target_arch = "powerpc64")] + pub fn is_p9_detected() -> bool { + std::arch::is_powerpc64_feature_detected!("altivec") + && std::arch::is_powerpc64_feature_detected!("vsx") + && std::arch::is_powerpc64_feature_detected!("power8-altivec") + && std::arch::is_powerpc64_feature_detected!("power8-crypto") + && std::arch::is_powerpc64_feature_detected!("power8-vector") + && std::arch::is_powerpc64_feature_detected!("power9-altivec") + && std::arch::is_powerpc64_feature_detected!("power9-vector") + } + + #[cfg(target_arch = "powerpc64")] + pub fn is_p8_detected() -> bool { + std::arch::is_powerpc64_feature_detected!("altivec") + && std::arch::is_powerpc64_feature_detected!("vsx") + && std::arch::is_powerpc64_feature_detected!("power8-altivec") + && std::arch::is_powerpc64_feature_detected!("power8-crypto") + && std::arch::is_powerpc64_feature_detected!("power8-vector") + } + + #[cfg(target_arch = "powerpc64")] + pub fn is_p7_detected() -> bool { + std::arch::is_powerpc64_feature_detected!("altivec") + && std::arch::is_powerpc64_feature_detected!("vsx") + } } pub use simd_macros::{multiversion, target_cpu}; @@ -229,6 +265,10 @@ pub use std::arch::is_aarch64_feature_detected as is_feature_detected; #[allow(unused_imports)] pub use std::arch::is_s390x_feature_detected as is_feature_detected; +#[cfg(target_arch = "powerpc64")] +#[allow(unused_imports)] +pub use std::arch::is_powerpc64_feature_detected as is_feature_detected; + #[cfg(target_arch = "x86_64")] #[allow(unused_imports)] pub use internal::is_x86_64_cpu_detected as is_cpu_detected; @@ -240,3 +280,7 @@ pub use internal::is_aarch64_cpu_detected as is_cpu_detected; #[cfg(target_arch = "s390x")] #[allow(unused_imports)] pub use internal::is_s390x_cpu_detected as is_cpu_detected; + +#[cfg(target_arch = "powerpc64")] +#[allow(unused_imports)] +pub use internal::is_powerpc64_cpu_detected as is_cpu_detected; diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index eb87baaf..24e99830 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -283,7 +283,7 @@ mod mul_add_round { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn mul_add_round(this: &[f32], k: f32, b: f32) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); diff --git a/crates/simd/src/rotate.rs b/crates/simd/src/rotate.rs index 9bf99ca7..795ea28c 100644 --- a/crates/simd/src/rotate.rs +++ b/crates/simd/src/rotate.rs @@ -13,7 +13,9 @@ // Copyright (c) 2025 TensorChord Inc. pub mod givens { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -29,7 +31,9 @@ pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { } pub mod flip { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + )] pub fn flip(bits: &[u64; 1024], result: &mut [f32]) { use std::hint::select_unpredictable; let result: &mut [u32] = zerocopy::transmute_mut!(result); diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index b21145bd..628bedf3 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -311,7 +311,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); @@ -517,7 +517,7 @@ mod reduce_sum_of_x_as_u16 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_x_as_u16(this: &[u8]) -> u16 { let n = this.len(); let mut sum = 0; @@ -722,7 +722,7 @@ mod reduce_sum_of_x_as_u32 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] pub fn reduce_sum_of_x_as_u32(this: &[u8]) -> u32 { let n = this.len(); let mut sum = 0; diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs index 4030c5f4..ba94cf60 100644 --- a/crates/simd_macros/src/target.rs +++ b/crates/simd_macros/src/target.rs @@ -124,4 +124,33 @@ pub const TARGET_CPUS: &[TargetCpu] = &[ target_arch: "s390x", target_features: &["vector"], }, + TargetCpu { + target_cpu: "p9", + target_arch: "powerpc64", + target_features: &[ + "altivec", + "vsx", + "power8-altivec", + "power8-crypto", + "power8-vector", + "power9-altivec", + "power9-vector", + ], + }, + TargetCpu { + target_cpu: "p8", + target_arch: "powerpc64", + target_features: &[ + "altivec", + "vsx", + "power8-altivec", + "power8-crypto", + "power8-vector", + ], + }, + TargetCpu { + target_cpu: "p7", + target_arch: "powerpc64", + target_features: &["altivec", "vsx"], + }, ]; From 2104191f50d4be5b818a535881e7e5e2eccdd577 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 29 Aug 2025 10:18:38 +0800 Subject: [PATCH 204/324] fix: SET_VARSIZE_4B on big endian (#326) Signed-off-by: usamoi --- src/datatype/memory_halfvec.rs | 4 ++++ src/datatype/memory_scalar8.rs | 4 ++++ src/datatype/memory_vector.rs | 4 ++++ 3 files changed, 12 insertions(+) diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs index 95afee4d..5fd46fb1 100644 --- a/src/datatype/memory_halfvec.rs +++ b/src/datatype/memory_halfvec.rs @@ -87,6 +87,10 @@ impl HalfvecOutput { let size = HalfvecHeader::size_of(slice.len()); let ptr = pgrx::pg_sys::palloc0(size) as *mut HalfvecHeader; + // SET_VARSIZE_4B + #[cfg(target_endian = "big")] + (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); + #[cfg(target_endian = "little")] (&raw mut (*ptr).varlena).write((size << 2) as u32); (&raw mut (*ptr).dims).write(vector.dims() as _); (&raw mut (*ptr).unused).write(0); diff --git a/src/datatype/memory_scalar8.rs b/src/datatype/memory_scalar8.rs index 6fa927ef..a57c7db4 100644 --- a/src/datatype/memory_scalar8.rs +++ b/src/datatype/memory_scalar8.rs @@ -95,6 +95,10 @@ impl Scalar8Output { let size = Scalar8Header::size_of(code.len()); let ptr = pgrx::pg_sys::palloc0(size) as *mut Scalar8Header; + // SET_VARSIZE_4B + #[cfg(target_endian = "big")] + (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); + #[cfg(target_endian = "little")] (&raw mut (*ptr).varlena).write((size << 2) as u32); (&raw mut (*ptr).dims).write(vector.dims() as _); (&raw mut (*ptr).unused).write(0); diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs index 5b3b1b3f..bbf818d3 100644 --- a/src/datatype/memory_vector.rs +++ b/src/datatype/memory_vector.rs @@ -86,6 +86,10 @@ impl VectorOutput { let size = VectorHeader::size_of(slice.len()); let ptr = pgrx::pg_sys::palloc0(size) as *mut VectorHeader; + // SET_VARSIZE_4B + #[cfg(target_endian = "big")] + (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); + #[cfg(target_endian = "little")] (&raw mut (*ptr).varlena).write((size << 2) as u32); (&raw mut (*ptr).dims).write(vector.dims() as _); (&raw mut (*ptr).unused).write(0); From dd741e5f6a8548515080642d56286be3d7175085 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 29 Aug 2025 10:48:55 +0800 Subject: [PATCH 205/324] feat: enable fast math in simd (#325) Signed-off-by: usamoi --- Cargo.lock | 4 - Cargo.toml | 2 - crates/algo/Cargo.toml | 1 - crates/algo/src/accessor.rs | 3 +- crates/make/Cargo.toml | 5 +- crates/make/src/main.rs | 24 ++- crates/simd/Cargo.toml | 3 +- crates/simd/build.rs | 25 ++- crates/simd/src/{f16.rs => floating_f16.rs} | 178 +++++++++++--------- crates/simd/src/{f32.rs => floating_f32.rs} | 48 +++++- crates/simd/src/lib.rs | 45 ++++- crates/vchordg/Cargo.toml | 1 - crates/vchordg/src/operator.rs | 3 +- crates/vchordg/src/types.rs | 2 +- crates/vchordrq/Cargo.toml | 1 - crates/vchordrq/src/operator.rs | 9 +- crates/vchordrq/src/types.rs | 2 +- src/datatype/functions_scalar8.rs | 3 +- src/datatype/memory_halfvec.rs | 2 +- src/index/vchordg/algo.rs | 2 +- src/index/vchordg/scanners/default.rs | 2 +- src/index/vchordrq/algo.rs | 2 +- src/index/vchordrq/am/am_build.rs | 3 +- src/index/vchordrq/scanners/default.rs | 2 +- src/index/vchordrq/scanners/maxsim.rs | 2 +- 25 files changed, 253 insertions(+), 121 deletions(-) rename crates/simd/src/{f16.rs => floating_f16.rs} (87%) rename crates/simd/src/{f32.rs => floating_f32.rs} (98%) diff --git a/Cargo.lock b/Cargo.lock index 1115fc9c..ab9b7d02 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -19,7 +19,6 @@ dependencies = [ "bumpalo", "dary_heap", "distance", - "half 2.6.0", "simd", "small_iter", "vector", @@ -1587,7 +1586,6 @@ dependencies = [ "bumpalo", "dary_heap", "distance", - "half 2.6.0", "k_means", "mimalloc", "paste", @@ -1614,7 +1612,6 @@ dependencies = [ "algo", "always_equal", "distance", - "half 2.6.0", "min-max-heap", "rabitq", "rand", @@ -1632,7 +1629,6 @@ dependencies = [ "algo", "always_equal", "distance", - "half 2.6.0", "pin-project", "rabitq", "rand", diff --git a/Cargo.toml b/Cargo.toml index 2be9b7af..a3e84be6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,7 +35,6 @@ vector = { path = "./crates/vector" } bumpalo.workspace = true dary_heap.workspace = true -half.workspace = true paste.workspace = true pgrx = { version = "=0.16.0", default-features = false, features = ["cshim"] } pgrx-catalog = "0.3.1" @@ -63,7 +62,6 @@ edition = "2024" [workspace.dependencies] bumpalo = "3.19.0" dary_heap = "0.3.7" -half = { version = "2.6.0", features = ["zerocopy"] } paste = "1.0.15" rand = "0.9.2" rand_chacha = "0.9.0" diff --git a/crates/algo/Cargo.toml b/crates/algo/Cargo.toml index 2951aad5..3f11dd14 100644 --- a/crates/algo/Cargo.toml +++ b/crates/algo/Cargo.toml @@ -13,7 +13,6 @@ vector = { path = "../vector" } bumpalo.workspace = true dary_heap.workspace = true -half.workspace = true zerocopy.workspace = true [lints] diff --git a/crates/algo/src/accessor.rs b/crates/algo/src/accessor.rs index 9a361481..d411eeb5 100644 --- a/crates/algo/src/accessor.rs +++ b/crates/algo/src/accessor.rs @@ -13,8 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use distance::Distance; -use half::f16; -use simd::Floating; +use simd::{Floating, f16}; use std::marker::PhantomData; use vector::vect::VectOwned; diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index eb92030f..66a7f525 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -5,10 +5,13 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.46", features = ["derive"] } +clap = { version = "4.5.46", features = ["derive", "env"] } object = { version = "0.37.3", default-features = false, features = [ "read", "wasm", ] } shlex = "1.3.0" target-triple = "0.1.4" + +[lints] +workspace = true diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index 8e942920..243cc5a7 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -42,6 +42,8 @@ struct BuildArgs { target: String, #[arg(long)] runner: Option, + #[arg(long, action = clap::ArgAction::SetTrue, env = "EXPERIMENTAL", value_parser = clap::builder::FalseyValueParser::new())] + experimental: bool, } struct TargetSpecificInformation { @@ -163,13 +165,20 @@ fn build( tsi: &TargetSpecificInformation, profile: &str, target: &str, + experimental: bool, ) -> Result> { let mut command = Command::new("cargo"); command .args(["build", "-p", "vchord", "--lib"]) .args(["--profile", profile]) .args(["--target", target]) - .args(["--features", pg_version]) + .args(["--features".into(), { + let mut features = vec![pg_version]; + if experimental { + features.push("simd/experimental"); + } + features.join(",") + }]) .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) .stdout(Stdio::inherit()) .stderr(Stdio::inherit()); @@ -234,6 +243,7 @@ fn generate( profile: &str, target: &str, exports: Vec, + experimental: bool, ) -> Result> { let pgrx_embed = std::env::temp_dir().join("VCHORD_PGRX_EMBED"); eprintln!("Writing {pgrx_embed:?}"); @@ -246,7 +256,13 @@ fn generate( .args(["rustc", "-p", "vchord", "--bin", "pgrx_embed_vchord"]) .args(["--profile", profile]) .args(["--target", target]) - .args(["--features", pg_version]) + .args(["--features".into(), { + let mut features = vec![pg_version]; + if experimental { + features.push("simd/experimental"); + } + features.join(",") + }]) .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) .args(["--", "--cfg", "pgrx_embed"]) .env("PGRX_EMBED", &pgrx_embed) @@ -328,6 +344,7 @@ fn main() -> Result<(), Box> { profile, target, runner, + experimental, }) => { let runner = runner.and_then(|runner| shlex::split(&runner)); let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { @@ -351,7 +368,7 @@ fn main() -> Result<(), Box> { } }; let tsi = target_specific_information(&target)?; - let obj = build(&path, &pg_version, &tsi, &profile, &target)?; + let obj = build(&path, &pg_version, &tsi, &profile, &target, experimental)?; let pkglibdir = format!("{output}/pkglibdir"); let sharedir = format!("{output}/sharedir"); let sharedir_extension = format!("{sharedir}/extension"); @@ -396,6 +413,7 @@ fn main() -> Result<(), Box> { &profile, &target, exports, + experimental, )?, format!("{sharedir_extension}/vchord--0.0.0.sql"), false, diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 764836ba..33774431 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -6,11 +6,12 @@ publish = false [features] init = [] +experimental = ["zerocopy/float-nightly"] [dependencies] simd_macros = { path = "../simd_macros" } -half.workspace = true +half = { version = "2.6.0", features = ["zerocopy"] } seq-macro.workspace = true zerocopy.workspace = true diff --git a/crates/simd/build.rs b/crates/simd/build.rs index e0d9ffd5..4bbcf142 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use std::env::var; +use std::env::{VarError, var}; use std::error::Error; use std::ffi::OsString; use std::path::Path; @@ -70,6 +70,16 @@ fn main() -> Result<(), Box> { build.opt_level(3); build.compile("simd_cshim"); } + "powerpc64" => { + if let Err(VarError::NotPresent) = var("CARGO_FEATURE_EXPERIMENTAL") { + println!("cargo::error=`experimental` should be enabled on this platform"); + } + } + "s390x" => { + if let Err(VarError::NotPresent) = var("CARGO_FEATURE_EXPERIMENTAL") { + println!("cargo::error=`experimental` should be enabled on this platform"); + } + } "x86_64" => { let mut build = cc::Build::new(); if let Some(compiler) = compiler(&host, &target, 16, 12) { @@ -79,7 +89,18 @@ fn main() -> Result<(), Box> { build.opt_level(3); build.compile("simd_cshim"); } - _ => (), + _ => { + if let Err(VarError::NotPresent) = var("CARGO_FEATURE_EXPERIMENTAL") { + println!("cargo::error=`experimental` should be enabled on this platform"); + } + let messages = [ + "This platform has poor SIMD implementation.", + "Please submit a feature request on https://github.com/tensorchord/VectorChord/issues.", + ]; + for message in messages { + println!("cargo::warning={message}"); + } + } } Ok(()) } diff --git a/crates/simd/src/f16.rs b/crates/simd/src/floating_f16.rs similarity index 87% rename from crates/simd/src/f16.rs rename to crates/simd/src/floating_f16.rs index eeaef4ff..4d9dff3b 100644 --- a/crates/simd/src/f16.rs +++ b/crates/simd/src/floating_f16.rs @@ -12,24 +12,12 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::{Floating, f32}; -use half::f16; - -trait AsF32 { - #[allow(clippy::wrong_self_convention)] - fn as_f32(self) -> f32; -} - -impl AsF32 for f16 { - fn as_f32(self) -> f32 { - self.into() - } -} +use crate::{F16, Floating, f16}; impl Floating for f16 { #[inline(always)] fn zero() -> Self { - f16::ZERO + f16::_ZERO } #[inline(always)] @@ -159,14 +147,14 @@ impl Floating for f16 { } mod reduce_or_of_is_zero_x { - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" )] pub fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { for &x in this { - if x == f16::ZERO { + if x == f16::_ZERO { return true; } } @@ -177,7 +165,7 @@ mod reduce_or_of_is_zero_x { mod reduce_sum_of_x { // FIXME: add manually-implemented SIMD version - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -186,7 +174,14 @@ mod reduce_sum_of_x { let n = this.len(); let mut x = 0.0f32; for i in 0..n { - x += this[i].as_f32(); + #[cfg(not(feature = "experimental"))] + { + x += this[i]._to_f32(); + } + #[cfg(feature = "experimental")] + { + x = x.algebraic_add(this[i]._to_f32()); + } } x } @@ -195,7 +190,7 @@ mod reduce_sum_of_x { mod reduce_sum_of_abs_x { // FIXME: add manually-implemented SIMD version - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -204,7 +199,14 @@ mod reduce_sum_of_abs_x { let n = this.len(); let mut x = 0.0f32; for i in 0..n { - x += (this[i].as_f32()).abs(); + #[cfg(not(feature = "experimental"))] + { + x += this[i]._to_f32().abs(); + } + #[cfg(feature = "experimental")] + { + x = x.algebraic_add(this[i]._to_f32().abs()); + } } x } @@ -213,7 +215,7 @@ mod reduce_sum_of_abs_x { mod reduce_sum_of_x2 { // FIXME: add manually-implemented SIMD version - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -222,7 +224,14 @@ mod reduce_sum_of_x2 { let n = this.len(); let mut x2 = 0.0f32; for i in 0..n { - x2 += this[i].as_f32() * this[i].as_f32(); + #[cfg(not(feature = "experimental"))] + { + x2 += this[i]._to_f32() * this[i]._to_f32(); + } + #[cfg(feature = "experimental")] + { + x2 = x2.algebraic_add(this[i]._to_f32().algebraic_mul(this[i]._to_f32())); + } } x2 } @@ -231,7 +240,7 @@ mod reduce_sum_of_x2 { mod reduce_min_max_of_x { // FIXME: add manually-implemented SIMD version - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -241,15 +250,15 @@ mod reduce_min_max_of_x { let mut max = f32::NEG_INFINITY; let n = this.len(); for i in 0..n { - min = min.min(this[i].as_f32()); - max = max.max(this[i].as_f32()); + min = min.min(this[i]._to_f32()); + max = max.max(this[i]._to_f32()); } (min, max) } } mod reduce_sum_of_xy { - use super::*; + use crate::{F16, f16}; #[inline] #[cfg(target_arch = "x86_64")] @@ -282,10 +291,10 @@ mod reduce_sum_of_xy { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -340,10 +349,10 @@ mod reduce_sum_of_xy { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let specialized = unsafe { reduce_sum_of_xy_v4(&lhs, &rhs) }; let fallback = fallback(&lhs, &rhs); @@ -376,8 +385,8 @@ mod reduce_sum_of_xy { let mut xy = emulate_mm256_reduce_add_ps(xy); // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { - let x = unsafe { a.read().as_f32() }; - let y = unsafe { b.read().as_f32() }; + let x = unsafe { a.read()._to_f32() }; + let y = unsafe { b.read()._to_f32() }; a = unsafe { a.add(1) }; b = unsafe { b.add(1) }; n -= 1; @@ -399,10 +408,10 @@ mod reduce_sum_of_xy { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -448,10 +457,10 @@ mod reduce_sum_of_xy { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -493,10 +502,10 @@ mod reduce_sum_of_xy { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -517,14 +526,21 @@ mod reduce_sum_of_xy { let n = lhs.len(); let mut xy = 0.0f32; for i in 0..n { - xy += lhs[i].as_f32() * rhs[i].as_f32(); + #[cfg(not(feature = "experimental"))] + { + xy += lhs[i]._to_f32() * rhs[i]._to_f32(); + } + #[cfg(feature = "experimental")] + { + xy = xy.algebraic_add(lhs[i]._to_f32().algebraic_mul(rhs[i]._to_f32())); + } } xy } } mod reduce_sum_of_d2 { - use super::*; + use crate::{F16, f16}; #[inline] #[cfg(target_arch = "x86_64")] @@ -557,10 +573,10 @@ mod reduce_sum_of_d2 { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -617,10 +633,10 @@ mod reduce_sum_of_d2 { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -658,8 +674,8 @@ mod reduce_sum_of_d2 { let mut d2 = emulate_mm256_reduce_add_ps(d2); // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { - let x = unsafe { a.read().as_f32() }; - let y = unsafe { b.read().as_f32() }; + let x = unsafe { a.read()._to_f32() }; + let y = unsafe { b.read()._to_f32() }; a = unsafe { a.add(1) }; b = unsafe { b.add(1) }; n -= 1; @@ -682,10 +698,10 @@ mod reduce_sum_of_d2 { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -731,10 +747,10 @@ mod reduce_sum_of_d2 { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -776,10 +792,10 @@ mod reduce_sum_of_d2 { for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; let lhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); let rhs = (0..n) - .map(|_| f16::from_f32(rng.random_range(-1.0..=1.0))) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) .collect::>(); for z in 3984..4016 { let lhs = &lhs[..z]; @@ -798,10 +814,18 @@ mod reduce_sum_of_d2 { pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); - let mut d2 = 0.0; + let mut d2 = 0.0_f32; for i in 0..n { - let d = lhs[i].as_f32() - rhs[i].as_f32(); - d2 += d * d; + #[cfg(not(feature = "experimental"))] + { + let d = lhs[i]._to_f32() - rhs[i]._to_f32(); + d2 += d * d; + } + #[cfg(feature = "experimental")] + { + let d = lhs[i]._to_f32().algebraic_sub(rhs[i]._to_f32()); + d2 = d2.algebraic_add(d.algebraic_mul(d)); + } } d2 } @@ -811,7 +835,7 @@ mod reduce_sum_of_xy_sparse { // There is no manually-implemented SIMD version. // Add it if `svecf16` is supported. - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -826,7 +850,7 @@ mod reduce_sum_of_xy_sparse { while lp < ln && rp < rn { match Ord::cmp(&lidx[lp], &ridx[rp]) { Ordering::Equal => { - xy += lval[lp].as_f32() * rval[rp].as_f32(); + xy += lval[lp]._to_f32() * rval[rp]._to_f32(); lp += 1; rp += 1; } @@ -846,7 +870,7 @@ mod reduce_sum_of_d2_sparse { // There is no manually-implemented SIMD version. // Add it if `svecf16` is supported. - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -861,33 +885,33 @@ mod reduce_sum_of_d2_sparse { while lp < ln && rp < rn { match Ord::cmp(&lidx[lp], &ridx[rp]) { Ordering::Equal => { - let d = lval[lp].as_f32() - rval[rp].as_f32(); + let d = lval[lp]._to_f32() - rval[rp]._to_f32(); d2 += d * d; lp += 1; rp += 1; } Ordering::Less => { - d2 += lval[lp].as_f32() * lval[lp].as_f32(); + d2 += lval[lp]._to_f32() * lval[lp]._to_f32(); lp += 1; } Ordering::Greater => { - d2 += rval[rp].as_f32() * rval[rp].as_f32(); + d2 += rval[rp]._to_f32() * rval[rp]._to_f32(); rp += 1; } } } for i in lp..ln { - d2 += lval[i].as_f32() * lval[i].as_f32(); + d2 += lval[i]._to_f32() * lval[i]._to_f32(); } for i in rp..rn { - d2 += rval[i].as_f32() * rval[i].as_f32(); + d2 += rval[i]._to_f32() * rval[i]._to_f32(); } d2 } } mod vector_add { - use super::*; + use crate::f16; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -909,7 +933,7 @@ mod vector_add { } mod vector_add_inplace { - use super::*; + use crate::f16; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -924,7 +948,7 @@ mod vector_add_inplace { } mod vector_sub { - use super::*; + use crate::f16; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -946,7 +970,7 @@ mod vector_sub { } mod vector_mul { - use super::*; + use crate::f16; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -968,13 +992,13 @@ mod vector_mul { } mod vector_mul_scalar { - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" )] pub fn vector_mul_scalar(lhs: &[f16], rhs: f32) -> Vec { - let rhs = f16::from_f32(rhs); + let rhs = f16::_from_f32(rhs); let n = lhs.len(); let mut r = Vec::::with_capacity(n); for i in 0..n { @@ -990,13 +1014,13 @@ mod vector_mul_scalar { } mod vector_mul_scalar_inplace { - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" )] pub fn vector_mul_scalar_inplace(lhs: &mut [f16], rhs: f32) { - let rhs = f16::from_f32(rhs); + let rhs = f16::_from_f32(rhs); let n = lhs.len(); for i in 0..n { lhs[i] *= rhs; @@ -1005,7 +1029,7 @@ mod vector_mul_scalar_inplace { } mod vector_abs_inplace { - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -1013,13 +1037,13 @@ mod vector_abs_inplace { pub fn vector_abs_inplace(this: &mut [f16]) { let n = this.len(); for i in 0..n { - this[i] = f16::from_f32(this[i].as_f32().abs()); + this[i] = f16::_from_f32(this[i]._to_f32().abs()); } } } mod vector_from_f32 { - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -1029,7 +1053,7 @@ mod vector_from_f32 { let mut r = Vec::::with_capacity(n); for i in 0..n { unsafe { - r.as_mut_ptr().add(i).write(f16::from_f32(this[i])); + r.as_mut_ptr().add(i).write(f16::_from_f32(this[i])); } } unsafe { @@ -1040,7 +1064,7 @@ mod vector_from_f32 { } mod vector_to_f32 { - use super::*; + use crate::{F16, f16}; #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" @@ -1050,7 +1074,7 @@ mod vector_to_f32 { let mut r = Vec::::with_capacity(n); for i in 0..n { unsafe { - r.as_mut_ptr().add(i).write(this[i].as_f32()); + r.as_mut_ptr().add(i).write(this[i]._to_f32()); } } unsafe { diff --git a/crates/simd/src/f32.rs b/crates/simd/src/floating_f32.rs similarity index 98% rename from crates/simd/src/f32.rs rename to crates/simd/src/floating_f32.rs index 13aae3da..f1ec2e12 100644 --- a/crates/simd/src/f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -418,7 +418,14 @@ mod reduce_sum_of_x { let n = this.len(); let mut sum = 0.0f32; for i in 0..n { - sum += this[i]; + #[cfg(not(feature = "experimental"))] + { + sum += this[i]; + } + #[cfg(feature = "experimental")] + { + sum = sum.algebraic_add(this[i]); + } } sum } @@ -693,7 +700,14 @@ mod reduce_sum_of_abs_x { let n = this.len(); let mut sum = 0.0f32; for i in 0..n { - sum += this[i].abs(); + #[cfg(not(feature = "experimental"))] + { + sum += this[i].abs(); + } + #[cfg(feature = "experimental")] + { + sum = sum.algebraic_add(this[i].abs()); + } } sum } @@ -958,7 +972,14 @@ mod reduce_sum_of_x2 { let n = this.len(); let mut x2 = 0.0f32; for i in 0..n { - x2 += this[i] * this[i]; + #[cfg(not(feature = "experimental"))] + { + x2 += this[i] * this[i]; + } + #[cfg(feature = "experimental")] + { + x2 = x2.algebraic_add(this[i].algebraic_mul(this[i])); + } } x2 } @@ -1547,7 +1568,14 @@ mod reduce_sum_of_xy { let n = lhs.len(); let mut xy = 0.0f32; for i in 0..n { - xy += lhs[i] * rhs[i]; + #[cfg(not(feature = "experimental"))] + { + xy += lhs[i] * rhs[i]; + } + #[cfg(feature = "experimental")] + { + xy = xy.algebraic_add(lhs[i].algebraic_mul(rhs[i])); + } } xy } @@ -1872,8 +1900,16 @@ mod reduce_sum_of_d2 { let n = lhs.len(); let mut d2 = 0.0f32; for i in 0..n { - let d = lhs[i] - rhs[i]; - d2 += d * d; + #[cfg(not(feature = "experimental"))] + { + let d = lhs[i] - rhs[i]; + d2 += d * d; + } + #[cfg(feature = "experimental")] + { + let d = lhs[i].algebraic_sub(rhs[i]); + d2 = d2.algebraic_add(d.algebraic_mul(d)); + } } d2 } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 47600f7a..a8585a37 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. #![allow(unsafe_code)] +#![cfg_attr(feature = "experimental", feature(float_algebraic, f16))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x_feature_detection))] #![cfg_attr(target_arch = "s390x", feature(s390x_target_feature))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] @@ -22,8 +23,8 @@ mod aligned; mod emulate; -mod f16; -mod f32; +mod floating_f16; +mod floating_f32; pub mod bit; pub mod fast_scan; @@ -32,6 +33,46 @@ pub mod quantize; pub mod rotate; pub mod u8; +#[cfg(not(feature = "experimental"))] +pub use half::f16; + +#[cfg(feature = "experimental")] +pub use f16; + +pub trait F16: Sized { + const _ZERO: Self; + + fn _from_f32(x: f32) -> Self; + + fn _to_f32(self) -> f32; +} + +#[cfg(not(feature = "experimental"))] +impl F16 for f16 { + const _ZERO: Self = f16::ZERO; + + fn _from_f32(x: f32) -> Self { + f16::from_f32(x) + } + + fn _to_f32(self) -> f32 { + self.to_f32() + } +} + +#[cfg(feature = "experimental")] +impl F16 for f16 { + const _ZERO: Self = 0.0; + + fn _from_f32(x: f32) -> Self { + x as _ + } + + fn _to_f32(self) -> f32 { + self as _ + } +} + pub trait Floating: Copy + Send + Sync + std::fmt::Debug + Default + 'static + PartialEq + PartialOrd { diff --git a/crates/vchordg/Cargo.toml b/crates/vchordg/Cargo.toml index 0ea18542..a9aa0b54 100644 --- a/crates/vchordg/Cargo.toml +++ b/crates/vchordg/Cargo.toml @@ -12,7 +12,6 @@ rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } -half.workspace = true min-max-heap = "1.3.0" rand.workspace = true serde.workspace = true diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs index a4ad88ac..9cb40636 100644 --- a/crates/vchordg/src/operator.rs +++ b/crates/vchordg/src/operator.rs @@ -15,9 +15,8 @@ use crate::types::DistanceKind; use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S}; use distance::Distance; -use half::f16; use rabitq::bits::Bits; -use simd::Floating; +use simd::{Floating, f16}; use std::fmt::Debug; use std::marker::PhantomData; use vector::vect::VectOwned; diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index d3df660c..cf5d6f55 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -12,8 +12,8 @@ // // Copyright (c) 2025 TensorChord Inc. -use half::f16; use serde::{Deserialize, Serialize}; +use simd::f16; use validator::{Validate, ValidationError}; use vector::vect::{VectBorrowed, VectOwned}; diff --git a/crates/vchordrq/Cargo.toml b/crates/vchordrq/Cargo.toml index 6f920599..e453f0fd 100644 --- a/crates/vchordrq/Cargo.toml +++ b/crates/vchordrq/Cargo.toml @@ -12,7 +12,6 @@ rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } -half.workspace = true pin-project = "1.1" serde.workspace = true validator.workspace = true diff --git a/crates/vchordrq/src/operator.rs b/crates/vchordrq/src/operator.rs index b26542f9..627e8aeb 100644 --- a/crates/vchordrq/src/operator.rs +++ b/crates/vchordrq/src/operator.rs @@ -14,11 +14,10 @@ use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S, RAccess}; use distance::Distance; -use half::f16; use rabitq::bit::CodeMetadata; use rabitq::bit::binary::BinaryLut; use rabitq::bit::block::{BlockLut, STEP}; -use simd::Floating; +use simd::{Floating, f16}; use std::fmt::Debug; use std::marker::PhantomData; use vector::vect::{VectBorrowed, VectOwned}; @@ -464,7 +463,8 @@ impl Operator for Op, L2S> { use std::iter::zip; let dims = vector.dims(); let t = zip(&code.1, centroid.slice()) - .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num).to_f32()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) + .map(simd::F16::_to_f32) .sum::() / (dims as f32).sqrt(); let sum_of_x_2 = code.0.dis_u_2; @@ -543,7 +543,8 @@ impl Operator for Op, Dot> { use std::iter::zip; let dims = vector.dims(); let t = zip(&code.1, centroid.slice()) - .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num).to_f32()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) + .map(simd::F16::_to_f32) .sum::() / (dims as f32).sqrt(); let sum_of_x_2 = code.0.dis_u_2; diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index 3edd066c..0f759bd1 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -12,8 +12,8 @@ // // Copyright (c) 2025 TensorChord Inc. -use half::f16; use serde::{Deserialize, Serialize}; +use simd::f16; use validator::{Validate, ValidationError}; use vector::vect::{VectBorrowed, VectOwned}; diff --git a/src/datatype/functions_scalar8.rs b/src/datatype/functions_scalar8.rs index d37d81b9..99990b15 100644 --- a/src/datatype/functions_scalar8.rs +++ b/src/datatype/functions_scalar8.rs @@ -15,8 +15,7 @@ use crate::datatype::memory_halfvec::HalfvecInput; use crate::datatype::memory_scalar8::Scalar8Output; use crate::datatype::memory_vector::VectorInput; -use half::f16; -use simd::Floating; +use simd::{Floating, f16}; use vector::scalar8::Scalar8Borrowed; #[pgrx::pg_extern(sql = "")] diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs index 5fd46fb1..16ad34a5 100644 --- a/src/datatype/memory_halfvec.rs +++ b/src/datatype/memory_halfvec.rs @@ -12,10 +12,10 @@ // // Copyright (c) 2025 TensorChord Inc. -use half::f16; use pgrx::datum::{FromDatum, IntoDatum}; use pgrx::pg_sys::{Datum, Oid}; use pgrx::pgrx_sql_entity_graph::metadata::*; +use simd::f16; use std::marker::PhantomData; use std::ptr::NonNull; use vector::VectorBorrowed; diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs index 528a64dc..b3e3f7f9 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/algo.rs @@ -16,7 +16,7 @@ use crate::index::vchordg::opclass::Opfamily; use algo::accessor::{Dot, L2S}; use algo::prefetcher::*; use algo::*; -use half::f16; +use simd::f16; use std::num::NonZero; use vchordg::Opaque; use vchordg::operator::Op; diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs index 67562fad..30683d8e 100644 --- a/src/index/vchordg/scanners/default.rs +++ b/src/index/vchordg/scanners/default.rs @@ -21,7 +21,7 @@ use crate::index::vchordg::scanners::SearchOptions; use algo::accessor::{Dot, L2S}; use algo::*; use distance::Distance; -use half::f16; +use simd::f16; use std::num::NonZero; use vchordg::operator::{self}; use vchordg::types::{DistanceKind, OwnedVector, VectorKind}; diff --git a/src/index/vchordrq/algo.rs b/src/index/vchordrq/algo.rs index 508537fb..55f094bc 100644 --- a/src/index/vchordrq/algo.rs +++ b/src/index/vchordrq/algo.rs @@ -17,7 +17,7 @@ use crate::index::vchordrq::opclass::Opfamily; use algo::accessor::{Dot, L2S}; use algo::prefetcher::*; use algo::*; -use half::f16; +use simd::f16; use std::collections::BinaryHeap; use std::num::NonZero; use vchordrq::operator::Op; diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index be3b0b79..8bf65eff 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -21,10 +21,9 @@ use crate::index::vchordrq::types::*; use algo::{ Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, }; -use half::f16; use pgrx::pg_sys::{Datum, ItemPointerData}; use rand::Rng; -use simd::Floating; +use simd::{Floating, f16}; use std::ffi::CStr; use std::ops::Deref; use vchordrq::types::*; diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index a5fa5b30..8ac6af34 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -24,7 +24,7 @@ use algo::prefetcher::*; use algo::*; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; -use half::f16; +use simd::f16; use std::num::NonZero; use vchordrq::operator::{self}; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index f01710ec..eecf8158 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -24,7 +24,7 @@ use algo::*; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; use distance::Distance; -use half::f16; +use simd::f16; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZero; From 8c701ce992caff5ff740d2023baf60558eca74e4 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 29 Aug 2025 16:11:23 +0800 Subject: [PATCH 206/324] feat: ppc64 BE scan (#327) The compiler refuses to inline the intrinsics in `stdarch` unless the baseline is raised with target-cpu. This is strange. Signed-off-by: usamoi --- crates/simd/src/fast_scan.rs | 48 +++++++++++++++++++++++++++++------- crates/simd/src/lib.rs | 3 ++- 2 files changed, 41 insertions(+), 10 deletions(-) diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index 59040fa9..a2827f98 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -514,11 +514,13 @@ mod scan { let n = code.len(); use std::arch::powerpc64::*; + #[cfg(target_endian = "big")] + use std::intrinsics::simd::simd_bswap as vec_revb; use std::mem::transmute; use {vector_unsigned_char as u8x16, vector_unsigned_short as u16x8}; let _0008_u16x8 = vec_splat_u16::<0x0008>(); - let _00ff_u16x8 = vec_splat_u16::<{ 0x00ffu8 as i8 }>(); + let _00ff_u16x8 = vec_splats(0x00ffu16); let _ff00_u16x8 = vec_splats(0xff00u16); let mut accu_0 = vec_splat_u16::<0>(); @@ -534,12 +536,26 @@ mod scan { let chi = vec_srl(code, vec_splat_u8::<4>()); let lut: u8x16 = vec_xl((i as isize) * 16, lut.as_ptr().cast::()); - let res_lo = transmute::(vec_perm(lut, lut, clo)); - accu_0 = vec_add(accu_0, res_lo); - accu_1 = vec_add(accu_1, vec_sr(res_lo, _0008_u16x8)); - let res_hi = transmute::(vec_perm(lut, lut, chi)); - accu_2 = vec_add(accu_2, res_hi); - accu_3 = vec_add(accu_3, vec_sr(res_hi, _0008_u16x8)); + #[cfg(target_endian = "big")] + { + let res_lo_r = transmute::(vec_perm(lut, lut, clo)); + let res_lo = vec_revb(res_lo_r); + accu_0 = vec_add(accu_0, res_lo); + accu_1 = vec_add(accu_1, vec_and(res_lo_r, _00ff_u16x8)); + let res_hi_r = transmute::(vec_perm(lut, lut, chi)); + let res_hi = vec_revb(res_hi_r); + accu_2 = vec_add(accu_2, res_hi); + accu_3 = vec_add(accu_3, vec_and(res_hi_r, _00ff_u16x8)); + } + #[cfg(target_endian = "little")] + { + let res_lo = transmute::(vec_perm(lut, lut, clo)); + accu_0 = vec_add(accu_0, res_lo); + accu_1 = vec_add(accu_1, vec_sr(res_lo, _0008_u16x8)); + let res_hi = transmute::(vec_perm(lut, lut, chi)); + accu_2 = vec_add(accu_2, res_hi); + accu_3 = vec_add(accu_3, vec_sr(res_hi, _0008_u16x8)); + } i += 1; } @@ -547,11 +563,25 @@ mod scan { let mut result = [0_u16; 32]; - accu_0 = vec_sub(accu_0, vec_sl(accu_1, _0008_u16x8)); + #[cfg(target_endian = "big")] + { + accu_0 = vec_sub(accu_0, vec_and(vec_revb(accu_1), _ff00_u16x8)); + } + #[cfg(target_endian = "little")] + { + accu_0 = vec_sub(accu_0, vec_sl(accu_1, _0008_u16x8)); + } vec_xst(accu_0, 0, result.as_mut_ptr().cast()); vec_xst(accu_1, 16, result.as_mut_ptr().cast()); - accu_2 = vec_sub(accu_2, vec_sl(accu_3, _0008_u16x8)); + #[cfg(target_endian = "big")] + { + accu_2 = vec_sub(accu_2, vec_and(vec_revb(accu_3), _ff00_u16x8)); + } + #[cfg(target_endian = "little")] + { + accu_2 = vec_sub(accu_2, vec_sl(accu_3, _0008_u16x8)); + } vec_xst(accu_2, 32, result.as_mut_ptr().cast()); vec_xst(accu_3, 48, result.as_mut_ptr().cast()); diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index a8585a37..5e2c91a4 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -12,8 +12,9 @@ // // Copyright (c) 2025 TensorChord Inc. -#![allow(unsafe_code)] +#![allow(unsafe_code, internal_features)] #![cfg_attr(feature = "experimental", feature(float_algebraic, f16))] +#![cfg_attr(feature = "experimental", feature(core_intrinsics))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x_feature_detection))] #![cfg_attr(target_arch = "s390x", feature(s390x_target_feature))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] From 692e17bc4211c5eeca96a53bccc7a4e9da1d95cb Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 1 Sep 2025 07:39:20 +0800 Subject: [PATCH 207/324] ci: s390x, powerpc64le, riscv64 (#328) job: +psql_macos job: +psql_windows job: +psql_alpine job: +check_debian Signed-off-by: usamoi --- .github/workflows/check.yml | 265 +++++++++++++++++++++++++++---- Cargo.toml | 2 + Makefile | 3 +- crates/make/src/main.rs | 60 +++---- crates/simd/Cargo.toml | 3 +- crates/simd/build.rs | 19 +-- crates/simd/src/bit.rs | 16 +- crates/simd/src/fast_scan.rs | 2 +- crates/simd/src/fht.rs | 6 +- crates/simd/src/floating_f16.rs | 56 +++---- crates/simd/src/floating_f32.rs | 52 +++--- crates/simd/src/lib.rs | 31 +++- crates/simd/src/quantize.rs | 2 +- crates/simd/src/rotate.rs | 4 +- crates/simd/src/u8.rs | 6 +- crates/simd_macros/src/target.rs | 5 + src/bin/pgrx_embed.rs | 11 +- 17 files changed, 389 insertions(+), 154 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index e2bc9c53..cf55169e 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -147,18 +147,18 @@ jobs: if [ "$(uname -m)" == "x86_64" ]; then cargo \ --config 'target.'\''cfg(all())'\''.runner = ["/opt/sde/sde64", "-spr", "--"]' \ - test -p simd -- --nocapture + test -p simd -- --no-capture fi if [ "$(uname -m)" == "aarch64" ]; then cargo \ --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=16"]' \ - test -p simd -- --nocapture + test -p simd -- --no-capture cargo \ --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=32"]' \ - test -p simd -- --nocapture + test -p simd -- --no-capture cargo \ --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=64"]' \ - test -p simd -- --nocapture + test -p simd -- --no-capture fi psql: @@ -225,6 +225,14 @@ jobs: make PG_CONFIG=$PG_CONFIG PROFILE=dev build sudo make PG_CONFIG=$PG_CONFIG install + - name: Upload Artifacts + uses: actions/upload-artifact@v4 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu + path: ./build/raw + compression-level: 9 + retention-days: 14 + - name: Service run: | sudo systemctl start postgresql @@ -239,14 +247,6 @@ jobs: sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' fi - - name: Upload Artifacts - uses: actions/upload-artifact@v4 - with: - name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu - path: ./build/raw - compression-level: 9 - retention-days: 14 - psql_macos: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_macos')) || @@ -308,6 +308,14 @@ jobs: make PG_CONFIG=$PG_CONFIG PROFILE=dev build sudo make PG_CONFIG=$PG_CONFIG install + - name: Upload Artifacts + uses: actions/upload-artifact@v4 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin + path: ./build/raw + compression-level: 9 + retention-days: 14 + - name: Service run: | brew services start postgresql@${{ matrix.version }} @@ -324,14 +332,6 @@ jobs: sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' fi - - name: Upload Artifacts - uses: actions/upload-artifact@v4 - with: - name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin - path: ./build/raw - compression-level: 9 - retention-days: 14 - psql_windows: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_windows')) || @@ -414,6 +414,14 @@ jobs: make PG_CONFIG=$env:PG_CONFIG PROFILE=dev build make PG_CONFIG=$env:PG_CONFIG install + - name: Upload Artifacts + uses: actions/upload-artifact@v4 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc + path: ./build/raw + compression-level: 9 + retention-days: 14 + - name: Service run: | 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } @@ -430,14 +438,6 @@ jobs: sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/pg17/*.slt' } - - name: Upload Artifacts - uses: actions/upload-artifact@v4 - with: - name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc - path: ./build/raw - compression-level: 9 - retention-days: 14 - psql_alpine: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_alpine')) || @@ -486,7 +486,7 @@ jobs: sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data sudo -iu postgres createuser -s -r $USER sudo -iu postgres createdb -O $USER $USER - sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo -iu postgres pg_ctl stop -D /var/lib/postgresql/data mkdir ~/pgvector-install @@ -522,10 +522,18 @@ jobs: make PG_CONFIG=$PG_CONFIG PROFILE=dev build sudo make PG_CONFIG=$PG_CONFIG install + - name: Upload Artifacts + uses: actions/upload-artifact@v4 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-musl + path: ./build/raw + compression-level: 9 + retention-days: 14 + - name: Service run: | sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data - psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + psql -d $USER -U $USER -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - name: Sqllogictest run: | @@ -536,10 +544,205 @@ jobs: sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' fi + check_debian: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +check_debian')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +check_debian')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + include: + - version: "15" + platform: "s390x" + clang_triple: "s390x-linux-gnu" + rust_triple: "s390x-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "17" + platform: "s390x" + clang_triple: "s390x-linux-gnu" + rust_triple: "s390x-unknown-linux-gnu" + gcc_version: "14" + debian_version: "trixie" + - version: "13" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "14" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "15" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "16" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "17" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "17" + platform: "riscv64" + clang_triple: "riscv64-linux-gnu" + rust_triple: "riscv64gc-unknown-linux-gnu" + gcc_version: "14" + debian_version: "trixie" + runs-on: "ubuntu-24.04" + + env: + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + # RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + sudo apt-get update + sudo apt-get install -y qemu-user-static clang lld + sudo systemctl enable --now systemd-binfmt + sudo mkdir /sysroot + curl -fsSL https://github.com/debuerreotype/docker-debian-artifacts/raw/refs/heads/dist-${{ matrix.platform }}/${{ matrix.debian_version }}/oci/blobs/rootfs.tar.gz | sudo tar -xz -C /sysroot + sudo mount --bind /dev /sysroot/dev + sudo mount --bind /dev/pts /sysroot/dev/pts + sudo mount --bind /etc/resolv.conf /sysroot/etc/resolv.conf + sudo mount --bind /proc /sysroot/proc + sudo mount --bind /sys /sysroot/sys + sudo mount --bind /tmp /sysroot/tmp + sudo chroot /sysroot apt-get update + sudo chroot /sysroot apt-get install --no-install-recommends -y libc6-dev libgcc-${{ matrix.gcc_version }}-dev + sudo chroot /sysroot apt-get install --no-install-recommends -y ca-certificates sudo + QEMU_LD_PREFIX=/sysroot + QEMU_CPU=max + if [ "${{ matrix.platform }}" = "ppc64le" ]; then + sudo rm -f /sysroot/lib64/ld64.so.2 + sudo ln /sysroot/usr/lib/powerpc64le-linux-gnu/ld64.so.2 /sysroot/lib64/ld64.so.2 + QEMU_CPU=power10 + fi + export QEMU_LD_PREFIX + export QEMU_CPU + echo QEMU_LD_PREFIX=$QEMU_LD_PREFIX >> $GITHUB_ENV + echo QEMU_CPU=$QEMU_CPU >> $GITHUB_ENV + echo "Defaults env_keep += \"QEMU_CPU\"" | sudo tee -a /etc/sudoers + + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' + sudo chroot /sysroot apt-get update + sudo chroot /sysroot apt-get install --no-install-recommends -y postgresql-common + sudo chroot /sysroot /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo chroot /sysroot apt-get install --no-install-recommends -y postgresql-${{ matrix.version }} postgresql-server-dev-${{ matrix.version }} + sudo touch /usr/bin/pg_config + echo "#!/usr/bin/env bash" | sudo tee -a /usr/bin/pg_config + echo "sudo chroot /sysroot pg_config \"\$@\" \\" | sudo tee -a /usr/bin/pg_config + echo " | sed -E 's|^/(.*)$|/sysroot/\1|' \\" | sudo tee -a /usr/bin/pg_config + echo " | sed -E 's|^([A-Z]+) = /(.*)$|\1 = /sysroot/\2|'" | sudo tee -a /usr/bin/pg_config + sudo chmod 755 /usr/bin/pg_config + pg_config + + echo "local all all trust" | sudo tee /sysroot/etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + echo "host all all 127.0.0.1/32 trust" | sudo tee -a /sysroot/etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + echo "host all all ::1/128 trust" | sudo tee -a /sysroot/etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main start + sudo chroot /sysroot sudo -iu postgres createuser -s -r $USER + sudo chroot /sysroot sudo -iu postgres createdb -O $USER $USER + sudo chroot /sysroot sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main stop + + mkdir ~/pgvector-install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.0.tar.gz | tar -xz -C ~/pgvector-install + make -C ~/pgvector-install/pgvector-0.8.0 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" + sudo make -C ~/pgvector-install/pgvector-0.8.0 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Checkout + uses: actions/checkout@v4 + + - name: Patch + run: | + cat << EOF > rust-toolchain.toml + [toolchain] + channel = "nightly-2025-08-29" + components = ["rustfmt", "clippy"] + targets = ["${{ matrix.rust_triple }}"] + profile = "minimal" + EOF + mkdir -p .cargo + cat << EOF > .cargo/config.toml + [target.${{ matrix.rust_triple }}] + runner = ["qemu-${{ matrix.platform }}-static"] + linker = "clang" + rustflags = [ + "-Clink-arg=-fuse-ld=lld", + "-Clink-arg=--target=${{ matrix.clang_triple }}", + "-Clink-arg=--sysroot=/sysroot", + "-Dwarnings", + ] + [env] + CC_${{ matrix.rust_triple }} = "clang" + CFLAGS_${{ matrix.rust_triple }} = "--target=${{ matrix.clang_triple }} --sysroot=/sysroot" + BINDGEN_EXTRA_CLANG_ARGS_${{ matrix.rust_triple }} = "--sysroot=/sysroot" + [patch.crates-io] + pgrx = { git = "https://github.com/tensorchord/pgrx.git", branch = "big-endian" } + EOF + rustup toolchain install + + - name: Clippy & Test + run: | + PGRX_PG_CONFIG_PATH=pg_config \ + cargo clippy --target ${{ matrix.rust_triple }} \ + --workspace --features pg${{ matrix.version }} + PGRX_PG_CONFIG_PATH=pg_config \ + cargo test --target ${{ matrix.rust_triple }} \ + --workspace --exclude vchord --no-fail-fast \ + -- --no-capture + + - name: Install + run: | + make \ + TARGET=${{ matrix.rust_triple }} \ + PROFILE=dev \ + RUNNER='qemu-${{ matrix.platform }}-static' \ + build + sudo make install + - name: Upload Artifacts uses: actions/upload-artifact@v4 with: - name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-musl + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.clang_triple }} path: ./build/raw compression-level: 9 retention-days: 14 + + - name: Service + run: | + sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main start + sudo chroot /sysroot psql -d $USER -U $USER -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $USER --user $USER './tests/general/*.slt' + sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' + sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' + if [ "${{ matrix.version }}" = "17" ]; then + sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' + fi diff --git a/Cargo.toml b/Cargo.toml index a3e84be6..9a087514 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,9 +71,11 @@ validator = { version = "0.20.0", features = ["derive"] } zerocopy = { version = "0.8.26", features = ["derive"] } [workspace.lints] +rust.unknown_lints = "allow" # complexity clippy.identity_op = "allow" clippy.int_plus_one = "allow" +clippy.manual_is_multiple_of = "allow" clippy.nonminimal_bool = "allow" clippy.too_many_arguments = "allow" clippy.type_complexity = "allow" diff --git a/Makefile b/Makefile index 920c6c30..8f7390a4 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,6 @@ PG_CONFIG ?= pg_config PKGLIBDIR := $(shell $(PG_CONFIG) --pkglibdir) SHAREDIR := $(shell $(PG_CONFIG) --sharedir) -PROFILE ?= release MKDIR ?= mkdir CP ?= cp @@ -10,7 +9,7 @@ CP ?= cp all: build build: - PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p make -- build --output ./build/raw --profile $(PROFILE) + PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p make -- build --output ./build/raw install: $(MKDIR) -p $(DESTDIR)$(PKGLIBDIR) $(DESTDIR)$(SHAREDIR) && \ diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index 243cc5a7..51d217cb 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -36,14 +36,12 @@ enum Commands { struct BuildArgs { #[arg(short, long)] output: String, - #[arg(long, default_value = "release")] - profile: String, - #[arg(long, default_value = target_triple::TARGET)] + #[arg(long, default_value = target_triple::TARGET, env = "TARGET")] target: String, - #[arg(long)] + #[arg(long, default_value = "release", env = "PROFILE")] + profile: String, + #[arg(long, env = "RUNNER")] runner: Option, - #[arg(long, action = clap::ArgAction::SetTrue, env = "EXPERIMENTAL", value_parser = clap::builder::FalseyValueParser::new())] - experimental: bool, } struct TargetSpecificInformation { @@ -51,6 +49,7 @@ struct TargetSpecificInformation { is_windows: bool, is_emscripten: bool, is_unix: bool, + is_powerpc64: bool, } impl TargetSpecificInformation { @@ -156,6 +155,7 @@ fn target_specific_information(target: &str) -> Result Result> { let mut command = Command::new("cargo"); command .args(["build", "-p", "vchord", "--lib"]) .args(["--profile", profile]) .args(["--target", target]) - .args(["--features".into(), { - let mut features = vec![pg_version]; - if experimental { - features.push("simd/experimental"); - } - features.join(",") - }]) + .args(["--features", pg_version]) .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) .stdout(Stdio::inherit()) .stderr(Stdio::inherit()); @@ -243,26 +236,39 @@ fn generate( profile: &str, target: &str, exports: Vec, - experimental: bool, + postmaster: impl AsRef, ) -> Result> { + let imports = if tsi.is_powerpc64 { + let postmaster = postmaster.as_ref(); + eprintln!("Reading {postmaster:?}"); + let contents = std::fs::read(postmaster)?; + let object = object::File::parse(contents.as_slice())?; + object + .exports()? + .into_iter() + .flat_map(|x| std::str::from_utf8(x.name())) + .filter(|x| !["_start", "_IO_stdin_used", "main"].contains(x)) + .map(str::to_string) + .collect::>() + } else { + Vec::new() + }; let pgrx_embed = std::env::temp_dir().join("VCHORD_PGRX_EMBED"); eprintln!("Writing {pgrx_embed:?}"); std::fs::write( &pgrx_embed, - format!("crate::schema_generation!({});", exports.join(" ")), + format!( + "crate::schema_generation!({}; {});", + exports.join(" "), + imports.join(" ") + ), )?; let mut command = Command::new("cargo"); command .args(["rustc", "-p", "vchord", "--bin", "pgrx_embed_vchord"]) .args(["--profile", profile]) .args(["--target", target]) - .args(["--features".into(), { - let mut features = vec![pg_version]; - if experimental { - features.push("simd/experimental"); - } - features.join(",") - }]) + .args(["--features", pg_version]) .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) .args(["--", "--cfg", "pgrx_embed"]) .env("PGRX_EMBED", &pgrx_embed) @@ -341,10 +347,9 @@ fn main() -> Result<(), Box> { match cli.command { Commands::Build(BuildArgs { output, - profile, target, + profile, runner, - experimental, }) => { let runner = runner.and_then(|runner| shlex::split(&runner)); let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { @@ -367,8 +372,9 @@ fn main() -> Result<(), Box> { return Err("PostgreSQL version is invalid.".into()); } }; + let postmaster = format!("{}/postgres", pg_config["BINDIR"]); let tsi = target_specific_information(&target)?; - let obj = build(&path, &pg_version, &tsi, &profile, &target, experimental)?; + let obj = build(&path, &pg_version, &tsi, &profile, &target)?; let pkglibdir = format!("{output}/pkglibdir"); let sharedir = format!("{output}/sharedir"); let sharedir_extension = format!("{sharedir}/extension"); @@ -413,7 +419,7 @@ fn main() -> Result<(), Box> { &profile, &target, exports, - experimental, + postmaster, )?, format!("{sharedir_extension}/vchord--0.0.0.sql"), false, diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 33774431..5157fed0 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -6,7 +6,8 @@ publish = false [features] init = [] -experimental = ["zerocopy/float-nightly"] +experimental_f16 = ["zerocopy/float-nightly"] +experimental_math = [] [dependencies] simd_macros = { path = "../simd_macros" } diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 4bbcf142..36e35d29 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use std::env::{VarError, var}; +use std::env::var; use std::error::Error; use std::ffi::OsString; use std::path::Path; @@ -70,16 +70,8 @@ fn main() -> Result<(), Box> { build.opt_level(3); build.compile("simd_cshim"); } - "powerpc64" => { - if let Err(VarError::NotPresent) = var("CARGO_FEATURE_EXPERIMENTAL") { - println!("cargo::error=`experimental` should be enabled on this platform"); - } - } - "s390x" => { - if let Err(VarError::NotPresent) = var("CARGO_FEATURE_EXPERIMENTAL") { - println!("cargo::error=`experimental` should be enabled on this platform"); - } - } + "powerpc64" => {} + "s390x" => {} "x86_64" => { let mut build = cc::Build::new(); if let Some(compiler) = compiler(&host, &target, 16, 12) { @@ -90,9 +82,7 @@ fn main() -> Result<(), Box> { build.compile("simd_cshim"); } _ => { - if let Err(VarError::NotPresent) = var("CARGO_FEATURE_EXPERIMENTAL") { - println!("cargo::error=`experimental` should be enabled on this platform"); - } + /* let messages = [ "This platform has poor SIMD implementation.", "Please submit a feature request on https://github.com/tensorchord/VectorChord/issues.", @@ -100,6 +90,7 @@ fn main() -> Result<(), Box> { for message in messages { println!("cargo::warning={message}"); } + */ } } Ok(()) diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index 7bc4305a..e372f860 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -183,7 +183,7 @@ mod reduce_sum_of_and { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -366,7 +366,7 @@ mod reduce_sum_of_or { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -549,7 +549,7 @@ mod reduce_sum_of_xor { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -772,7 +772,7 @@ mod reduce_sum_of_and_or { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert_eq!(lhs.len(), rhs.len()); let n = lhs.len(); @@ -933,7 +933,7 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_x(this: &[u64]) -> u32 { let n = this.len(); let mut sum = 0; @@ -951,7 +951,7 @@ pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { mod vector_and { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); @@ -976,7 +976,7 @@ pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { mod vector_or { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); @@ -1001,7 +1001,7 @@ pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { mod vector_xor { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { assert_eq!(lhs.len(), rhs.len()); diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index a2827f98..47a6c451 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -672,7 +672,7 @@ pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { mod accu { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn accu(sum: &mut [u32; 32], delta: &[u16; 32]) { for i in 0..32 { diff --git a/crates/simd/src/fht.rs b/crates/simd/src/fht.rs index 84238b48..e7effbab 100644 --- a/crates/simd/src/fht.rs +++ b/crates/simd/src/fht.rs @@ -33,7 +33,7 @@ mod step_1 { seq_macro::seq!( Q in 0..16 { mod dispatch_~Q { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn f(x: &mut [f32]) { crate::fht::basic_1::(x); } @@ -48,7 +48,7 @@ mod step_2 { seq_macro::seq!( Q in 0..16 { mod dispatch_~Q { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn f(x: &mut [f32]) { crate::fht::basic_2::(x); } @@ -62,7 +62,7 @@ mod step_2 { macro_rules! fht { ($p:literal, 0) => { { - #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] fn walk(x: &mut [f32]) { assert!(x.len() == (1 << $p)); seq_macro::seq!( diff --git a/crates/simd/src/floating_f16.rs b/crates/simd/src/floating_f16.rs index 4d9dff3b..0316394e 100644 --- a/crates/simd/src/floating_f16.rs +++ b/crates/simd/src/floating_f16.rs @@ -150,7 +150,7 @@ mod reduce_or_of_is_zero_x { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { for &x in this { @@ -168,17 +168,17 @@ mod reduce_sum_of_x { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] { x += this[i]._to_f32(); } - #[cfg(feature = "experimental")] + #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] { x = x.algebraic_add(this[i]._to_f32()); } @@ -193,17 +193,17 @@ mod reduce_sum_of_abs_x { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { let n = this.len(); let mut x = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] { x += this[i]._to_f32().abs(); } - #[cfg(feature = "experimental")] + #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] { x = x.algebraic_add(this[i]._to_f32().abs()); } @@ -218,17 +218,17 @@ mod reduce_sum_of_x2 { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_x2(this: &[f16]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] { x2 += this[i]._to_f32() * this[i]._to_f32(); } - #[cfg(feature = "experimental")] + #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] { x2 = x2.algebraic_add(this[i]._to_f32().algebraic_mul(this[i]._to_f32())); } @@ -243,7 +243,7 @@ mod reduce_min_max_of_x { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { let mut min = f32::INFINITY; @@ -520,17 +520,17 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); let mut xy = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] { xy += lhs[i]._to_f32() * rhs[i]._to_f32(); } - #[cfg(feature = "experimental")] + #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] { xy = xy.algebraic_add(lhs[i]._to_f32().algebraic_mul(rhs[i]._to_f32())); } @@ -810,18 +810,18 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); let mut d2 = 0.0_f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] { let d = lhs[i]._to_f32() - rhs[i]._to_f32(); d2 += d * d; } - #[cfg(feature = "experimental")] + #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] { let d = lhs[i]._to_f32().algebraic_sub(rhs[i]._to_f32()); d2 = d2.algebraic_add(d.algebraic_mul(d)); @@ -838,7 +838,7 @@ mod reduce_sum_of_xy_sparse { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; @@ -873,7 +873,7 @@ mod reduce_sum_of_d2_sparse { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { use std::cmp::Ordering; @@ -914,7 +914,7 @@ mod vector_add { use crate::f16; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_add(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); @@ -936,7 +936,7 @@ mod vector_add_inplace { use crate::f16; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_add_inplace(lhs: &mut [f16], rhs: &[f16]) { assert_eq!(lhs.len(), rhs.len()); @@ -951,7 +951,7 @@ mod vector_sub { use crate::f16; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_sub(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); @@ -973,7 +973,7 @@ mod vector_mul { use crate::f16; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_mul(lhs: &[f16], rhs: &[f16]) -> Vec { assert_eq!(lhs.len(), rhs.len()); @@ -995,7 +995,7 @@ mod vector_mul_scalar { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_mul_scalar(lhs: &[f16], rhs: f32) -> Vec { let rhs = f16::_from_f32(rhs); @@ -1017,7 +1017,7 @@ mod vector_mul_scalar_inplace { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_mul_scalar_inplace(lhs: &mut [f16], rhs: f32) { let rhs = f16::_from_f32(rhs); @@ -1032,7 +1032,7 @@ mod vector_abs_inplace { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_abs_inplace(this: &mut [f16]) { let n = this.len(); @@ -1046,7 +1046,7 @@ mod vector_from_f32 { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_from_f32(this: &[f32]) -> Vec { let n = this.len(); @@ -1067,7 +1067,7 @@ mod vector_to_f32 { use crate::{F16, f16}; #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_to_f32(this: &[f16]) -> Vec { let n = this.len(); diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs index f1ec2e12..4fe48642 100644 --- a/crates/simd/src/floating_f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -148,7 +148,7 @@ impl Floating for f32 { mod reduce_or_of_is_zero_x { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { for &x in this { @@ -413,16 +413,16 @@ mod reduce_sum_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(feature = "experimental_math"))] { sum += this[i]; } - #[cfg(feature = "experimental")] + #[cfg(feature = "experimental_math")] { sum = sum.algebraic_add(this[i]); } @@ -695,16 +695,16 @@ mod reduce_sum_of_abs_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(feature = "experimental_math"))] { sum += this[i].abs(); } - #[cfg(feature = "experimental")] + #[cfg(feature = "experimental_math")] { sum = sum.algebraic_add(this[i].abs()); } @@ -967,16 +967,16 @@ mod reduce_sum_of_x2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_x2(this: &[f32]) -> f32 { let n = this.len(); let mut x2 = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(feature = "experimental_math"))] { x2 += this[i] * this[i]; } - #[cfg(feature = "experimental")] + #[cfg(feature = "experimental_math")] { x2 = x2.algebraic_add(this[i].algebraic_mul(this[i])); } @@ -1245,7 +1245,7 @@ mod reduce_min_max_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -1562,17 +1562,17 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_xy(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); let mut xy = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(feature = "experimental_math"))] { xy += lhs[i] * rhs[i]; } - #[cfg(feature = "experimental")] + #[cfg(feature = "experimental_math")] { xy = xy.algebraic_add(lhs[i].algebraic_mul(rhs[i])); } @@ -1894,18 +1894,18 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_d2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); let mut d2 = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental"))] + #[cfg(not(feature = "experimental_math"))] { let d = lhs[i] - rhs[i]; d2 += d * d; } - #[cfg(feature = "experimental")] + #[cfg(feature = "experimental_math")] { let d = lhs[i].algebraic_sub(rhs[i]); d2 = d2.algebraic_add(d.algebraic_mul(d)); @@ -2004,7 +2004,7 @@ mod reduce_sum_of_xy_sparse { } } - #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -2154,7 +2154,7 @@ mod reduce_sum_of_d2_sparse { } } - #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { use std::cmp::Ordering; assert_eq!(lidx.len(), lval.len()); @@ -2192,7 +2192,7 @@ mod reduce_sum_of_d2_sparse { mod vector_add { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); @@ -2212,7 +2212,7 @@ mod vector_add { mod vector_add_inplace { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { assert_eq!(lhs.len(), rhs.len()); @@ -2225,7 +2225,7 @@ mod vector_add_inplace { mod vector_sub { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); @@ -2245,7 +2245,7 @@ mod vector_sub { mod vector_mul { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { assert_eq!(lhs.len(), rhs.len()); @@ -2265,7 +2265,7 @@ mod vector_mul { mod vector_mul_scalar { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { let n = lhs.len(); @@ -2284,7 +2284,7 @@ mod vector_mul_scalar { mod vector_mul_scalar_inplace { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { let n = lhs.len(); @@ -2296,7 +2296,7 @@ mod vector_mul_scalar_inplace { mod vector_abs_inplace { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn vector_abs_inplace(this: &mut [f32]) { let n = this.len(); diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 5e2c91a4..7422f795 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -13,14 +13,17 @@ // Copyright (c) 2025 TensorChord Inc. #![allow(unsafe_code, internal_features)] -#![cfg_attr(feature = "experimental", feature(float_algebraic, f16))] -#![cfg_attr(feature = "experimental", feature(core_intrinsics))] +#![cfg_attr(feature = "experimental_f16", feature(f16))] +#![cfg_attr(feature = "experimental_math", feature(float_algebraic))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x_feature_detection))] #![cfg_attr(target_arch = "s390x", feature(s390x_target_feature))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] #![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc_feature_detection))] #![cfg_attr(target_arch = "powerpc64", feature(powerpc_target_feature))] #![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc))] +#![cfg_attr(target_arch = "powerpc64", feature(core_intrinsics))] +#![cfg_attr(target_arch = "riscv64", feature(stdarch_riscv_feature_detection))] +#![cfg_attr(target_arch = "riscv64", feature(riscv_target_feature))] mod aligned; mod emulate; @@ -34,10 +37,10 @@ pub mod quantize; pub mod rotate; pub mod u8; -#[cfg(not(feature = "experimental"))] +#[cfg(not(feature = "experimental_f16"))] pub use half::f16; -#[cfg(feature = "experimental")] +#[cfg(feature = "experimental_f16")] pub use f16; pub trait F16: Sized { @@ -48,7 +51,7 @@ pub trait F16: Sized { fn _to_f32(self) -> f32; } -#[cfg(not(feature = "experimental"))] +#[cfg(not(feature = "experimental_f16"))] impl F16 for f16 { const _ZERO: Self = f16::ZERO; @@ -61,7 +64,7 @@ impl F16 for f16 { } } -#[cfg(feature = "experimental")] +#[cfg(feature = "experimental_f16")] impl F16 for f16 { const _ZERO: Self = 0.0; @@ -121,6 +124,9 @@ mod internal { #[cfg(target_arch = "powerpc64")] simd_macros::define_is_cpu_detected!("powerpc64"); + #[cfg(target_arch = "riscv64")] + simd_macros::define_is_cpu_detected!("riscv64"); + #[cfg(target_arch = "x86_64")] #[allow(unused_imports)] pub use is_x86_64_cpu_detected; @@ -137,6 +143,10 @@ mod internal { #[allow(unused_imports)] pub use is_powerpc64_cpu_detected; + #[cfg(target_arch = "riscv64")] + #[allow(unused_imports)] + pub use is_riscv64_cpu_detected; + #[cfg(target_arch = "x86_64")] pub fn is_v4_detected() -> bool { std::arch::is_x86_feature_detected!("avx512bw") @@ -291,6 +301,11 @@ mod internal { std::arch::is_powerpc64_feature_detected!("altivec") && std::arch::is_powerpc64_feature_detected!("vsx") } + + #[cfg(target_arch = "riscv64")] + pub fn is_r1_detected() -> bool { + std::arch::is_riscv_feature_detected!("v") + } } pub use simd_macros::{multiversion, target_cpu}; @@ -326,3 +341,7 @@ pub use internal::is_s390x_cpu_detected as is_cpu_detected; #[cfg(target_arch = "powerpc64")] #[allow(unused_imports)] pub use internal::is_powerpc64_cpu_detected as is_cpu_detected; + +#[cfg(target_arch = "riscv64")] +#[allow(unused_imports)] +pub use internal::is_riscv64_cpu_detected as is_cpu_detected; diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index 24e99830..90c88fcd 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -283,7 +283,7 @@ mod mul_add_round { } } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn mul_add_round(this: &[f32], k: f32, b: f32) -> Vec { let n = this.len(); let mut r = Vec::::with_capacity(n); diff --git a/crates/simd/src/rotate.rs b/crates/simd/src/rotate.rs index 795ea28c..15916229 100644 --- a/crates/simd/src/rotate.rs +++ b/crates/simd/src/rotate.rs @@ -14,7 +14,7 @@ pub mod givens { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { assert!(lhs.len() == rhs.len()); @@ -32,7 +32,7 @@ pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { pub mod flip { #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7" + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn flip(bits: &[u64; 1024], result: &mut [f32]) { use std::hint::select_unpredictable; diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index 628bedf3..f0cfa839 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -311,7 +311,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); @@ -517,7 +517,7 @@ mod reduce_sum_of_x_as_u16 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_x_as_u16(this: &[u8]) -> u16 { let n = this.len(); let mut sum = 0; @@ -722,7 +722,7 @@ mod reduce_sum_of_x_as_u32 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_x_as_u32(this: &[u8]) -> u32 { let n = this.len(); let mut sum = 0; diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs index ba94cf60..5903f72a 100644 --- a/crates/simd_macros/src/target.rs +++ b/crates/simd_macros/src/target.rs @@ -153,4 +153,9 @@ pub const TARGET_CPUS: &[TargetCpu] = &[ target_arch: "powerpc64", target_features: &["altivec", "vsx"], }, + TargetCpu { + target_cpu: "r1", + target_arch: "riscv64", + target_features: &["v"], + }, ]; diff --git a/src/bin/pgrx_embed.rs b/src/bin/pgrx_embed.rs index 505d1b69..36ae772e 100644 --- a/src/bin/pgrx_embed.rs +++ b/src/bin/pgrx_embed.rs @@ -19,8 +19,17 @@ #[macro_export] macro_rules! schema_generation { - ($($symbol:ident)*) => { + ($($symbol:ident)*; $($import:ident)*) => { pub fn main() -> Result<(), Box> { + $( + const _: () = { + #[unsafe(no_mangle)] + unsafe extern "C" fn $import() { + panic!("{} is called unexpectedly.", stringify!($import)); + } + }; + )* + extern crate vchord as _; use ::pgrx::pgrx_sql_entity_graph::ControlFile; From 9dbe9b6abaaf08b2002a1d9b86f4d628bd9ce99a Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 1 Sep 2025 08:40:47 +0800 Subject: [PATCH 208/324] ci: fix libm.so linking, powerpc64le (#329) https://github.com/tensorchord/VectorChord/actions/runs/17363793655/job/49287408204 Signed-off-by: usamoi --- .github/workflows/check.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index cf55169e..afd34a42 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -633,6 +633,8 @@ jobs: if [ "${{ matrix.platform }}" = "ppc64le" ]; then sudo rm -f /sysroot/lib64/ld64.so.2 sudo ln /sysroot/usr/lib/powerpc64le-linux-gnu/ld64.so.2 /sysroot/lib64/ld64.so.2 + sudo rm -f /sysroot/usr/lib/powerpc64le-linux-gnu/libm.so + sudo ln /sysroot/lib/powerpc64le-linux-gnu/libm.so.6 /sysroot/usr/lib/powerpc64le-linux-gnu/libm.so QEMU_CPU=power10 fi export QEMU_LD_PREFIX From 7c9420fb509c9f7f3e9b01d7197d05bebd7415cb Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 1 Sep 2025 16:04:02 +0800 Subject: [PATCH 209/324] fix: apply same rotation on big endian (#330) job: +check_debian Signed-off-by: usamoi --- crates/rabitq/build.rs | 7 ++++++- crates/rabitq/src/rotate.rs | 31 +++++++++++++++++++++++-------- 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/crates/rabitq/build.rs b/crates/rabitq/build.rs index a56f6a95..54e24743 100644 --- a/crates/rabitq/build.rs +++ b/crates/rabitq/build.rs @@ -20,7 +20,12 @@ fn main() -> Result<(), Box> { use rand::{Rng, SeedableRng}; use rand_chacha::ChaCha12Rng; let mut rng = ChaCha12Rng::from_seed([7; 32]); - let bits = (0..262144).map(|_| rng.random::()).collect::>(); + let mut bits = (0..262144).map(|_| rng.random::()).collect::>(); + if var("CARGO_CFG_TARGET_ENDIAN")?.as_str() == "big" { + for chunk in bits.as_chunks_mut::<8>().0 { + chunk.reverse(); + } + } let out_dir = var("OUT_DIR")?; std::fs::write(PathBuf::from(out_dir).join("bits"), bits)?; Ok(()) diff --git a/crates/rabitq/src/rotate.rs b/crates/rabitq/src/rotate.rs index 130d7dfa..1be5eb3f 100644 --- a/crates/rabitq/src/rotate.rs +++ b/crates/rabitq/src/rotate.rs @@ -14,14 +14,29 @@ const BITS: &[u8; 262144] = include_bytes!(concat!(env!("OUT_DIR"), "/bits")); -const _: () = assert!(BITS[0] == 246); -const _: () = assert!(BITS[1] == 133); -const _: () = assert!(BITS[2] == 163); -const _: () = assert!(BITS[3] == 106); -const _: () = assert!(BITS[4] == 54); -const _: () = assert!(BITS[5] == 126); -const _: () = assert!(BITS[6] == 9); -const _: () = assert!(BITS[7] == 115); +#[cfg(target_endian = "little")] +const _: () = { + assert!(BITS[0] == 246); + assert!(BITS[1] == 133); + assert!(BITS[2] == 163); + assert!(BITS[3] == 106); + assert!(BITS[4] == 54); + assert!(BITS[5] == 126); + assert!(BITS[6] == 9); + assert!(BITS[7] == 115); +}; + +#[cfg(target_endian = "big")] +const _: () = { + assert!(BITS[7] == 246); + assert!(BITS[6] == 133); + assert!(BITS[5] == 163); + assert!(BITS[4] == 106); + assert!(BITS[3] == 54); + assert!(BITS[2] == 126); + assert!(BITS[1] == 9); + assert!(BITS[0] == 115); +}; static BITS_0: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[0]); static BITS_1: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[1]); From 14d21fa1943ede9b7bd2cd911211e86e05653266 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 2 Sep 2025 15:04:09 +0800 Subject: [PATCH 210/324] chore: upload 0.5.1 schema scripts (#331) Signed-off-by: usamoi --- sql/install/vchord--0.5.1.sql | 752 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.5.0--0.5.1.sql | 1 + 2 files changed, 753 insertions(+) create mode 100644 sql/install/vchord--0.5.1.sql create mode 100644 sql/upgrade/vchord--0.5.0--0.5.1.sql diff --git a/sql/install/vchord--0.5.1.sql b/sql/install/vchord--0.5.1.sql new file mode 100644 index 00000000..dcaf9b98 --- /dev/null +++ b/sql/install/vchord--0.5.1.sql @@ -0,0 +1,752 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:45 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:31 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:21 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:77 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:77 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:41 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:46 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/upgrade/vchord--0.5.0--0.5.1.sql b/sql/upgrade/vchord--0.5.0--0.5.1.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.5.0--0.5.1.sql @@ -0,0 +1 @@ +/* nothing to do */ From 29437d2f1b70f76fe5a3865b0c62a6aa252c2581 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 2 Sep 2025 15:46:10 +0800 Subject: [PATCH 211/324] readme: sync with docs (#333) https://github.com/tensorchord/pgvecto.rs-docs/pull/140 Signed-off-by: usamoi --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index db627121..e0996700 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.5.0 + -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.5.1 ``` > [!NOTE] > In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). From bbcf054f460af95a575ab484c1fe3c7737677154 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Thu, 4 Sep 2025 20:12:10 +0800 Subject: [PATCH 212/324] feat: record queries by sqlite (#321) Real-time sampling of SQL queries for recall evaluation and monitoring. Signed-off-by: cutecutecat --- Cargo.lock | 59 ++++++++++ Cargo.toml | 1 + scripts/train.py | 2 +- src/index/functions.rs | 24 ++++ src/index/gucs.rs | 46 ++++++++ src/index/scanners.rs | 2 + src/index/vchordg/am/mod.rs | 10 +- src/index/vchordg/scanners/default.rs | 21 +++- src/index/vchordrq/am/mod.rs | 15 ++- src/index/vchordrq/scanners/default.rs | 20 +++- src/index/vchordrq/scanners/maxsim.rs | 2 + src/lib.rs | 2 + src/recorder/hook.rs | 53 +++++++++ src/recorder/mod.rs | 26 +++++ src/recorder/text.rs | 40 +++++++ src/recorder/types.rs | 62 +++++++++++ src/recorder/worker.rs | 146 +++++++++++++++++++++++++ src/sql/finalize.sql | 96 ++++++++++++++++ tests/vchordrq/recall.slt | 132 +++++++++++++++++++++- 19 files changed, 744 insertions(+), 15 deletions(-) create mode 100644 src/recorder/hook.rs create mode 100644 src/recorder/mod.rs create mode 100644 src/recorder/text.rs create mode 100644 src/recorder/types.rs create mode 100644 src/recorder/worker.rs diff --git a/Cargo.lock b/Cargo.lock index ab9b7d02..d178abb2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -420,6 +420,18 @@ dependencies = [ "once_cell", ] +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "fixedbitset" version = "0.5.7" @@ -499,6 +511,15 @@ dependencies = [ "foldhash", ] +[[package]] +name = "hashlink" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7382cf6263419f2d8df38c55d7da83da5c18aef87fc7a7fc1fb1e344edfe14c1" +dependencies = [ + "hashbrown", +] + [[package]] name = "heck" version = "0.5.0" @@ -733,6 +754,17 @@ dependencies = [ "libc", ] +[[package]] +name = "libsqlite3-sys" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "133c182a6a2c87864fe97778797e46c7e999672690dc9fa3ee8e241aa4a9c13f" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + [[package]] name = "linux-raw-sys" version = "0.9.4" @@ -999,6 +1031,12 @@ dependencies = [ "syn", ] +[[package]] +name = "pkg-config" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" + [[package]] name = "potential_utf" version = "0.1.2" @@ -1157,6 +1195,20 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" +[[package]] +name = "rusqlite" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "165ca6e57b20e1351573e3729b958bc62f0e48025386970b6e4d29e7a7e71f3f" +dependencies = [ + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", +] + [[package]] name = "rustc-hash" version = "2.1.1" @@ -1593,6 +1645,7 @@ dependencies = [ "pgrx-catalog", "rabitq", "rand", + "rusqlite", "seq-macro", "serde", "simd", @@ -1639,6 +1692,12 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + [[package]] name = "vector" version = "0.0.0" diff --git a/Cargo.toml b/Cargo.toml index 9a087514..67ea8192 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,7 @@ paste.workspace = true pgrx = { version = "=0.16.0", default-features = false, features = ["cshim"] } pgrx-catalog = "0.3.1" rand.workspace = true +rusqlite = { version = "0.37.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true toml = "0.9.5" diff --git a/scripts/train.py b/scripts/train.py index e5f38925..54f86c28 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -29,7 +29,7 @@ import numpy as np DEFAULT_LISTS = 4096 -N_ITER = 25 +N_ITER = 10 CHUNKS = 10 SEED = 42 MAX_POINTS_PER_CLUSTER = 256 diff --git a/src/index/functions.rs b/src/index/functions.rs index 3711fac4..17cf6741 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -13,6 +13,8 @@ // Copyright (c) 2025 TensorChord Inc. use crate::index::storage::PostgresRelation; +use crate::recorder::dump; +use pgrx::iter::SetOfIterator; use pgrx::pg_sys::Oid; use pgrx_catalog::{PgAm, PgClass, PgClassRelkind}; @@ -84,3 +86,25 @@ impl Drop for Index { } } } + +#[pgrx::pg_extern(sql = "")] +fn _vchordrq_sampled_vectors(indexrelid: Oid) -> SetOfIterator<'static, String> { + let pg_am = PgAm::search_amname(c"vchordrq").unwrap(); + let Some(pg_am) = pg_am.get() else { + pgrx::error!("vchord is not installed"); + }; + let pg_class = PgClass::search_reloid(indexrelid).unwrap(); + let Some(pg_class) = pg_class.get() else { + pgrx::error!("the relation does not exist"); + }; + if pg_class.relkind() != PgClassRelkind::Index { + pgrx::error!("the relation {:?} is not an index", pg_class.relname()); + } + if pg_class.relam() != pg_am.oid() { + pgrx::error!("the index {:?} is not a vchordrq index", pg_class.relname()); + } + // The user must have access to the index, if not, raise an error from Postgres. + let _relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); + let queries = dump(indexrelid.to_u32()); + SetOfIterator::new(queries) +} diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 2e039c14..ed501a52 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -27,6 +27,12 @@ pub enum PostgresIo { ReadStream, } +static VCHORDRQ_QUERY_SAMPLING_ENABLE: GucSetting = GucSetting::::new(false); + +static VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS: GucSetting = GucSetting::::new(0); + +static VCHORDRQ_QUERY_SAMPLING_RATE: GucSetting = GucSetting::::new(0.0); + static VCHORDG_ENABLE_SCAN: GucSetting = GucSetting::::new(true); static VCHORDG_EF_SEARCH: GucSetting = GucSetting::::new(64); @@ -158,6 +164,34 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_bool_guc( + c"vchordrq.query_sampling_enable", + c"`query_sampling_enable` argument of vchordrq.", + c"`query_sampling_enable` argument of vchordrq.", + &VCHORDRQ_QUERY_SAMPLING_ENABLE, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.query_sampling_max_records", + c"`query_sampling_max_records` argument of vchordrq.", + c"`query_sampling_max_records` argument of vchordrq.", + &VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS, + 0, + 10000, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_float_guc( + c"vchordrq.query_sampling_rate", + c"`query_sampling_rate` argument of vchordrq.", + c"`query_sampling_rate` argument of vchordrq.", + &VCHORDRQ_QUERY_SAMPLING_RATE, + 0.0, + 1.0, + GucContext::Userset, + GucFlags::default(), + ); unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); @@ -331,3 +365,15 @@ pub fn vchordrq_io_rerank() -> Io { PostgresIo::ReadStream => Io::Stream, } } + +pub fn vchordrq_query_sampling_enable() -> bool { + VCHORDRQ_QUERY_SAMPLING_ENABLE.get() +} + +pub fn vchordrq_query_sampling_max_records() -> u32 { + VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS.get() as u32 +} + +pub fn vchordrq_query_sampling_rate() -> f64 { + VCHORDRQ_QUERY_SAMPLING_RATE.get() +} diff --git a/src/index/scanners.rs b/src/index/scanners.rs index b4203b97..97d8b1b4 100644 --- a/src/index/scanners.rs +++ b/src/index/scanners.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::index::fetcher::Fetcher; +use crate::recorder::Recorder; use algo::{Bump, Page, RelationPrefetch, RelationRead, RelationReadStream}; use pgrx::pg_sys::Datum; @@ -44,6 +45,7 @@ pub trait SearchBuilder: 'static { options: Self::Options, fetcher: impl Fetcher + 'b, bump: &'b impl Bump, + recorder: impl Recorder, ) -> Box + 'b> where R: RelationRead + RelationPrefetch + RelationReadStream, diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index dc6a233e..d3f5dca5 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -20,6 +20,7 @@ use crate::index::scanners::SearchBuilder; use crate::index::storage::PostgresRelation; use crate::index::vchordg::opclass::opfamily; use crate::index::vchordg::scanners::*; +use crate::recorder::DefaultRecorder; use pgrx::datum::Internal; use pgrx::pg_sys::Datum; use std::cell::LazyCell; @@ -372,6 +373,13 @@ pub unsafe extern "C-unwind" fn amrescan( ) }) }; + // Query recorde is disable for vchordg indexes for now. + let recorder = DefaultRecorder { + enable: false, + rate: None, + max_records: 0, + index: (*(*scan).indexRelation).rd_id.to_u32(), + }; // PAY ATTENTATION: `scanning` references `bump`, so `scanning` must be dropped before `bump`. let bump = scanner.bump.as_ref(); scanner.scanning = match opfamily { @@ -397,7 +405,7 @@ pub unsafe extern "C-unwind" fn amrescan( LazyCell::new(Box::new(move || { // only do this since `PostgresRelation` has no destructor let index = bump.alloc(index.clone()); - builder.build(index, options, fetcher, bump) + builder.build(index, options, fetcher, bump, recorder) })) } }; diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs index 30683d8e..54a2c560 100644 --- a/src/index/vchordg/scanners/default.rs +++ b/src/index/vchordg/scanners/default.rs @@ -18,6 +18,7 @@ use crate::index::scanners::{Io, SearchBuilder}; use crate::index::vchordg::algo::*; use crate::index::vchordg::opclass::Opfamily; use crate::index::vchordg::scanners::SearchOptions; +use crate::recorder::{Recorder, halfvec_out, vector_out}; use algo::accessor::{Dot, L2S}; use algo::*; use distance::Distance; @@ -26,6 +27,7 @@ use std::num::NonZero; use vchordg::operator::{self}; use vchordg::types::{DistanceKind, OwnedVector, VectorKind}; use vchordg::*; +use vector::VectorOwned; use vector::vect::{VectBorrowed, VectOwned}; pub struct DefaultBuilder { @@ -78,6 +80,7 @@ impl SearchBuilder for DefaultBuilder { options: SearchOptions, _fetcher: impl Fetcher + 'b, bump: &'b impl Bump, + recorder: impl Recorder, ) -> Box + 'b> where R: RelationRead + RelationPrefetch + RelationReadStream, @@ -120,7 +123,7 @@ impl SearchBuilder for DefaultBuilder { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { type Op = operator::Op, L2S>; - let unprojected = if let OwnedVector::Vecf32(vector) = vector { + let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() @@ -215,7 +218,7 @@ impl SearchBuilder for DefaultBuilder { } (VectorKind::Vecf16, DistanceKind::L2S) => { type Op = operator::Op, L2S>; - let unprojected = if let OwnedVector::Vecf16(vector) = vector { + let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() @@ -310,7 +313,7 @@ impl SearchBuilder for DefaultBuilder { } (VectorKind::Vecf32, DistanceKind::Dot) => { type Op = operator::Op, Dot>; - let unprojected = if let OwnedVector::Vecf32(vector) = vector { + let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() @@ -405,7 +408,7 @@ impl SearchBuilder for DefaultBuilder { } (VectorKind::Vecf16, DistanceKind::Dot) => { type Op = operator::Op, Dot>; - let unprojected = if let OwnedVector::Vecf16(vector) = vector { + let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() @@ -509,6 +512,16 @@ impl SearchBuilder for DefaultBuilder { } else { iter }; + if recorder.is_enabled() { + match &vector { + OwnedVector::Vecf32(v) => { + recorder.send(&vector_out(v.as_borrowed())); + } + OwnedVector::Vecf16(v) => { + recorder.send(&halfvec_out(v.as_borrowed())); + } + } + } Box::new(iter.map(move |(distance, pointer)| { let (key, _) = pointer_to_kv(pointer); (opfamily.output(distance), key, recheck) diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index b8603c35..45e774b4 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -20,6 +20,7 @@ use crate::index::scanners::SearchBuilder; use crate::index::storage::PostgresRelation; use crate::index::vchordrq::opclass::{Opfamily, opfamily}; use crate::index::vchordrq::scanners::*; +use crate::recorder::DefaultRecorder; use pgrx::datum::Internal; use pgrx::pg_sys::Datum; use std::cell::LazyCell; @@ -455,6 +456,16 @@ pub unsafe extern "C-unwind" fn amrescan( ) }) }; + let rate = match gucs::vchordrq_query_sampling_rate() { + 0.0 => None, + rate => Some(rate), + }; + let recorder = DefaultRecorder { + enable: gucs::vchordrq_query_sampling_enable(), + rate, + max_records: gucs::vchordrq_query_sampling_max_records(), + index: (*(*scan).indexRelation).rd_id.to_u32(), + }; // PAY ATTENTATION: `scanning` references `bump`, so `scanning` must be dropped before `bump`. let bump = scanner.bump.as_ref(); scanner.scanning = match opfamily { @@ -480,7 +491,7 @@ pub unsafe extern "C-unwind" fn amrescan( LazyCell::new(Box::new(move || { // only do this since `PostgresRelation` has no destructor let index = bump.alloc(index.clone()); - builder.build(index, options, fetcher, bump) + builder.build(index, options, fetcher, bump, recorder) })) } Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim => { @@ -500,7 +511,7 @@ pub unsafe extern "C-unwind" fn amrescan( LazyCell::new(Box::new(move || { // only do this since `PostgresRelation` has no destructor let index = bump.alloc(index.clone()); - builder.build(index, options, fetcher, bump) + builder.build(index, options, fetcher, bump, recorder) })) } }; diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index 8ac6af34..c61d0bab 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -19,6 +19,7 @@ use crate::index::vchordrq::algo::*; use crate::index::vchordrq::filter::filter; use crate::index::vchordrq::opclass::Opfamily; use crate::index::vchordrq::scanners::SearchOptions; +use crate::recorder::{Recorder, halfvec_out, vector_out}; use algo::accessor::{Dot, L2S}; use algo::prefetcher::*; use algo::*; @@ -82,6 +83,7 @@ impl SearchBuilder for DefaultBuilder { options: SearchOptions, mut fetcher: impl Fetcher + 'b, bump: &'b impl Bump, + recorder: impl Recorder, ) -> Box + 'b> where R: RelationRead + RelationPrefetch + RelationReadStream, @@ -120,7 +122,7 @@ impl SearchBuilder for DefaultBuilder { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { type Op = operator::Op, L2S>; - let unprojected = if let OwnedVector::Vecf32(vector) = vector { + let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { vector } else { unreachable!() @@ -260,7 +262,7 @@ impl SearchBuilder for DefaultBuilder { } (VectorKind::Vecf32, DistanceKind::Dot) => { type Op = operator::Op, Dot>; - let unprojected = if let OwnedVector::Vecf32(vector) = vector { + let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { vector } else { unreachable!() @@ -400,7 +402,7 @@ impl SearchBuilder for DefaultBuilder { } (VectorKind::Vecf16, DistanceKind::L2S) => { type Op = operator::Op, L2S>; - let unprojected = if let OwnedVector::Vecf16(vector) = vector { + let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { vector } else { unreachable!() @@ -540,7 +542,7 @@ impl SearchBuilder for DefaultBuilder { } (VectorKind::Vecf16, DistanceKind::Dot) => { type Op = operator::Op, Dot>; - let unprojected = if let OwnedVector::Vecf16(vector) = vector { + let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { vector } else { unreachable!() @@ -689,6 +691,16 @@ impl SearchBuilder for DefaultBuilder { } else { iter }; + if recorder.is_enabled() { + match &vector { + OwnedVector::Vecf32(v) => { + recorder.send(&vector_out(v.as_borrowed())); + } + OwnedVector::Vecf16(v) => { + recorder.send(&halfvec_out(v.as_borrowed())); + } + } + } Box::new(iter.map(move |(distance, pointer)| { let (key, _) = pointer_to_kv(pointer); (distance, key, recheck) diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index eecf8158..65603995 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -18,6 +18,7 @@ use crate::index::vchordrq::algo::*; use crate::index::vchordrq::filter::filter; use crate::index::vchordrq::opclass::Opfamily; use crate::index::vchordrq::scanners::SearchOptions; +use crate::recorder::Recorder; use algo::accessor::Dot; use algo::prefetcher::*; use algo::*; @@ -72,6 +73,7 @@ impl SearchBuilder for MaxsimBuilder { options: SearchOptions, mut fetcher: impl Fetcher + 'b, bump: &'b impl Bump, + _sender: impl Recorder, ) -> Box + 'b> where R: RelationRead + RelationPrefetch + RelationReadStream, diff --git a/src/lib.rs b/src/lib.rs index b458d51c..25ff506d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,6 +18,7 @@ mod datatype; mod index; +mod recorder; mod upgrade; pgrx::pg_module_magic!( @@ -53,6 +54,7 @@ unsafe extern "C-unwind" fn _pg_init() { } IS_MAIN.set(true); index::init(); + recorder::init(); unsafe { #[cfg(any(feature = "pg13", feature = "pg14"))] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); diff --git a/src/recorder/hook.rs b/src/recorder/hook.rs new file mode 100644 index 00000000..8dadee99 --- /dev/null +++ b/src/recorder/hook.rs @@ -0,0 +1,53 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::recorder::worker::{delete_database, delete_index}; + +static mut PREV_OBJECT_ACCESS: pgrx::pg_sys::object_access_hook_type = None; + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn recorder_object_access( + access: pgrx::pg_sys::ObjectAccessType::Type, + class_id: pgrx::pg_sys::Oid, + object_id: pgrx::pg_sys::Oid, + sub_id: ::std::os::raw::c_int, + arg: *mut ::std::os::raw::c_void, +) { + unsafe { + use pgrx::pg_sys::submodules::ffi::pg_guard_ffi_boundary; + if let Some(prev_object_access_hook) = PREV_OBJECT_ACCESS { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + prev_object_access_hook(access, class_id, object_id, sub_id, arg) + }); + } + if access == pgrx::pg_sys::ObjectAccessType::OAT_DROP + && class_id == pgrx::pg_sys::DatabaseRelationId + { + delete_database(object_id.to_u32()); + } else if access == pgrx::pg_sys::ObjectAccessType::OAT_DROP + && class_id == pgrx::pg_sys::RelationRelationId + { + delete_index(object_id.to_u32()); + } + } +} + +pub fn init() { + assert!(crate::is_main()); + unsafe { + PREV_OBJECT_ACCESS = pgrx::pg_sys::object_access_hook; + pgrx::pg_sys::object_access_hook = Some(recorder_object_access); + } +} diff --git a/src/recorder/mod.rs b/src/recorder/mod.rs new file mode 100644 index 00000000..77a09d4b --- /dev/null +++ b/src/recorder/mod.rs @@ -0,0 +1,26 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub use text::{halfvec_out, vector_out}; +pub use types::{DefaultRecorder, Recorder}; +pub use worker::dump; + +mod hook; +mod text; +mod types; +mod worker; + +pub fn init() { + hook::init(); +} diff --git a/src/recorder/text.rs b/src/recorder/text.rs new file mode 100644 index 00000000..321599af --- /dev/null +++ b/src/recorder/text.rs @@ -0,0 +1,40 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use simd::f16; +use vector::vect::VectBorrowed; + +pub fn vector_out(vector: VectBorrowed<'_, f32>) -> String { + let mut result = String::from("["); + for x in vector.slice() { + if !result.ends_with('[') { + result.push(','); + } + result.push_str(&x.to_string()); + } + result.push(']'); + result +} + +pub fn halfvec_out(vector: VectBorrowed<'_, f16>) -> String { + let mut result = String::from("["); + for x in vector.slice() { + if !result.ends_with('[') { + result.push(','); + } + result.push_str(&x.to_string()); + } + result.push(']'); + result +} diff --git a/src/recorder/types.rs b/src/recorder/types.rs new file mode 100644 index 00000000..39ec9a97 --- /dev/null +++ b/src/recorder/types.rs @@ -0,0 +1,62 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::recorder::worker::push; +use rand::Rng; +use std::cell::RefMut; + +pub trait Recorder { + fn is_enabled(&self) -> bool; + fn send(&self, sample: &str); +} + +#[derive(Debug)] +pub struct DefaultRecorder { + pub enable: bool, + pub rate: Option, + pub max_records: u32, + pub index: u32, +} + +pub struct PgRefCell(std::cell::RefCell); + +unsafe impl Send for PgRefCell {} +unsafe impl Sync for PgRefCell {} + +impl PgRefCell { + pub const fn new(x: T) -> Self { + Self(std::cell::RefCell::new(x)) + } + pub fn borrow_mut(&self) -> RefMut<'_, T> { + assert!( + crate::is_main(), + "cannot borrow the value outside main thread" + ); + self.0.borrow_mut() + } +} + +impl Recorder for DefaultRecorder { + fn is_enabled(&self) -> bool { + self.enable + } + fn send(&self, sample: &str) { + if let Some(rate) = self.rate { + let mut rng = rand::rng(); + if rng.random_bool(rate) { + push(self.index, sample, self.max_records); + } + } + } +} diff --git a/src/recorder/worker.rs b/src/recorder/worker.rs new file mode 100644 index 00000000..a6da90f0 --- /dev/null +++ b/src/recorder/worker.rs @@ -0,0 +1,146 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. +use crate::recorder::types::PgRefCell; +use std::cell::RefMut; +use std::fs; +use std::path::Path; + +// Safety: The directory name must start with "pgsql_tmp" to be excluded by pg_basebackup +const RECORDER_DIR: &str = "pgsql_tmp_vchord_sampling"; +const RECORDER_VERSION: u32 = 1; + +static CONNECTION: PgRefCell> = + PgRefCell::>::new(None); + +fn get<'a>() -> Option> { + if unsafe { !pgrx::pg_sys::IsBackendPid(pgrx::pg_sys::MyProcPid) } { + return None; + } + let database_oid = unsafe { pgrx::pg_sys::MyDatabaseId.to_u32() }; + if database_oid == 0 { + return None; + } + let mut connection = CONNECTION.borrow_mut(); + if connection.is_none() + && let Err(err) = || -> rusqlite::Result<()> { + if !Path::new(RECORDER_DIR).exists() { + let _ = fs::create_dir_all(RECORDER_DIR); + } + let p = format!("{RECORDER_DIR}/database_{database_oid}.sqlite"); + let mut conn = rusqlite::Connection::open(&p)?; + conn.pragma_update(Some("main"), "journal_mode", "WAL")?; + conn.pragma_update(Some("main"), "synchronous", "NORMAL")?; + let tx = conn.transaction()?; + let version: u32 = tx + .pragma_query_value(Some("main"), "user_version", |row| row.get(0)) + .unwrap_or(RECORDER_VERSION); + if version != RECORDER_VERSION && version != 0 { + let mut statement = tx.prepare( + "SELECT name FROM sqlite_master WHERE type='table' AND name LIKE 'index_%';", + )?; + let tables = statement.query_map((), |row| row.get::(0))?; + for name in tables.into_iter().flatten() { + let drop_statement = format!("DROP TABLE IF EXISTS {name}"); + tx.execute(&drop_statement, ())?; + } + } + tx.pragma_update(Some("main"), "user_version", RECORDER_VERSION)?; + tx.commit()?; + let _ = connection.insert(conn); + Ok(()) + }() + { + if err.sqlite_error_code() == Some(rusqlite::ErrorCode::DatabaseCorrupt) { + delete_database(database_oid); + } + pgrx::debug1!("Recorder: Error initializing database: {}", err); + return None; + } + RefMut::filter_map(connection, |c| c.as_mut()).ok() +} + +pub fn push(index: u32, sample: &str, max_records: u32) { + let mut connection = match get() { + Some(c) => c, + None => return, + }; + let init_statement = format!( + " + CREATE TABLE IF NOT EXISTS index_{index} (sample TEXT, create_at REAL); + CREATE INDEX IF NOT EXISTS i ON index_{index} (create_at); + " + ); + let insert_statement = + format!("INSERT INTO index_{index} (sample, create_at) VALUES (?1, unixepoch('subsec'))"); + let count_statement = format!("SELECT COUNT(create_at) FROM index_{index}"); + let maintain_statement = format!( + "DELETE FROM index_{index} WHERE rowid = ( + SELECT rowid FROM index_{index} ORDER BY create_at ASC LIMIT ?1);" + ); + if let Err(err) = || -> rusqlite::Result<()> { + let tx = connection.transaction()?; + tx.execute_batch(&init_statement)?; + tx.prepare_cached(&insert_statement)?.execute((sample,))?; + let records = tx.query_one(&count_statement, (), |row| row.get::(0))?; + if records > max_records { + tx.execute(&maintain_statement, (records - max_records,))?; + } + tx.commit()?; + Ok(()) + }() { + pgrx::debug1!("Recorder: Error pushing sample: {}", err); + } +} + +pub fn delete_index(index: u32) { + let connection = match get() { + Some(c) => c, + None => return, + }; + let drop_statement = format!("DROP TABLE IF EXISTS index_{index}"); + if let Err(e) = connection.execute(&drop_statement, ()) { + pgrx::debug1!("Recorder: Error deleting index table: {}", e); + }; +} + +pub fn delete_database(database_oid: u32) { + let _ = fs::remove_file(format!("{RECORDER_DIR}/database_{database_oid}.sqlite")); + let _ = fs::remove_file(format!("{RECORDER_DIR}/database_{database_oid}.sqlite-shm")); + let _ = fs::remove_file(format!("{RECORDER_DIR}/database_{database_oid}.sqlite-wal")); +} + +pub fn dump(index: u32) -> Vec { + let connection = match get() { + Some(c) => c, + None => return Vec::new(), + }; + let load_statement = format!("SELECT sample FROM index_{index} ORDER BY create_at DESC"); + match || -> rusqlite::Result> { + let mut stmt = connection.prepare(&load_statement)?; + let mut rows = stmt.query(())?; + let mut result = Vec::new(); + while let Some(row) = rows.next()? { + if let Ok(sample) = row.get::(0) { + result.push(sample); + } + } + Ok(result) + }() { + Ok(v) => v, + Err(e) => { + pgrx::debug1!("Recorder: Error loading samples: {}", e); + Vec::new() + } + } +} diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index d6c186ff..a520ba39 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -142,6 +142,102 @@ IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_ CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; +CREATE FUNCTION vchordrq_sampled_vectors(regclass) +RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_vectors_wrapper'; + +CREATE OR REPLACE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + table_name NAME, + column_name NAME, + operator TEXT, + vector_text TEXT +) +LANGUAGE plpgsql +STRICT AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + C.relname AS table_name, + PA.attname AS column_name, + CASE + WHEN OP.opcname LIKE '%%l2%%' THEN '<->' + WHEN OP.opcname LIKE '%%ip%%' THEN '<#>' + WHEN OP.opcname LIKE '%%cosine%%' THEN '<=>' + ELSE '' + END AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OP ON OP.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.table_name, + im.column_name, + im.operator, + s.vector_text + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_vectors(%1$s) AS s(vector_text); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.table_name, + record.column_name, + record.operator, + record.vector_text +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; + CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; diff --git a/tests/vchordrq/recall.slt b/tests/vchordrq/recall.slt index 04584fd0..ca3fb095 100644 --- a/tests/vchordrq/recall.slt +++ b/tests/vchordrq/recall.slt @@ -1,12 +1,12 @@ statement ok -CREATE TABLE t (val vector(3)); +CREATE TABLE t (id SERIAL PRIMARY KEY, val vector(3)); statement ok INSERT INTO t (val) SELECT ARRAY[i * 0.0001, i * 0.00005, i * 0.0002]::vector(3) FROM generate_series(1, 10000) as s(i); statement ok -CREATE INDEX ON t USING vchordrq (val vector_l2_ops); +CREATE INDEX idx1 ON t USING vchordrq (val vector_l2_ops); statement ok SET vchordrq.epsilon = 0.8; @@ -53,4 +53,130 @@ SHOW vchordrq.epsilon; 0.8 statement ok -DROP TABLE t; \ No newline at end of file +CREATE TABLE t_dim4 (val vector(4), id SERIAL PRIMARY KEY); + +statement ok +INSERT INTO t_dim4 (val) +SELECT ARRAY[i * 0.0001, i * 0.00005, i * 0.0002, i * 0.001]::vector(4) FROM generate_series(1, 10000) as s(i); + +statement ok +CREATE INDEX idx2 ON t_dim4 USING vchordrq (val vector_l2_ops); + +statement ok +ALTER SYSTEM SET vchordrq.query_sampling_max_records = 1; + +statement ok +ALTER SYSTEM SET vchordrq.query_sampling_rate = 1; + +statement ok +ALTER SYSTEM SET vchordrq.query_sampling_enable = on; + +statement ok +SELECT pg_reload_conf(); + +query I retry 5 backoff 1s +SHOW vchordrq.query_sampling_enable; +---- +on + +statement ok +SELECT * from t ORDER BY val <-> '[0.50, 0.25, 1.00]'; + +statement ok +SELECT * from t_dim4 ORDER BY val <-> '[1.00, 0.50, 0.25, 0]'; + +query I +SELECT vector_text from vchordrq_sampled_queries('idx1'); +---- +[0.5,0.25,1] + +query I +SELECT vector_text from vchordrq_sampled_queries('idx2'); +---- +[1,0.5,0.25,0] + +query I +SELECT COUNT(*) from vchordrq_sampled_queries; +---- +2 + +statement ok +SELECT * from t_dim4 ORDER BY val <-> '[2.1, 0.3, 0.7, 0.9]'; + +query I +SELECT * from vchordrq_sampled_queries('idx2'); +---- +public t_dim4 val <-> [2.1,0.3,0.7,0.9] + +query I +SELECT AVG(recall_value) +FROM ( + SELECT + vchordrq_evaluate_query_recall( + query => format( + 'SELECT ctid FROM %I.%I ORDER BY %I %s ''%s'' LIMIT 10', + lq.schema_name, + lq.table_name, + lq.column_name, + lq.operator, + lq.vector_text + ) + ) AS recall_value + FROM + vchordrq_sampled_queries('idx2') AS lq +) AS eval_results; +---- +1 + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX idx3 ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops); + +statement ok +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.0000, 0.5000, 0.2500]' limit 1; + +query I +SELECT column_name from vchordrq_sampled_queries('idx3'); +---- +NULL + +query I +SELECT vector_text from vchordrq_sampled_queries('idx3'); +---- +[1,0.5,0.25] + +statement ok +SET search_path='@'; + +query I +SELECT vector_text from public.vchordrq_sampled_queries('public.idx3'); +---- +[1,0.5,0.25] + +query I +SELECT COUNT(*) from public.vchordrq_sampled_queries; +---- +3 + +statement ok +RESET search_path; + +statement ok +ALTER SYSTEM RESET vchordrq.query_sampling_enable; + +statement ok +ALTER SYSTEM RESET vchordrq.query_sampling_max_records; + +statement ok +ALTER SYSTEM RESET vchordrq.query_sampling_rate; + +statement ok +SELECT pg_reload_conf(); + +statement ok +DROP TABLE t, t_dim4, t_expr; \ No newline at end of file From a748e767dd6b51bb9b48557fee124c0aff9addf2 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Tue, 9 Sep 2025 12:45:53 +0800 Subject: [PATCH 213/324] chord: add index_name into recorder (#334) Signed-off-by: cutecutecat --- src/recorder/worker.rs | 1 + src/sql/finalize.sql | 12 ++++++++---- tests/vchordrq/recall.slt | 12 ++++++------ 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/src/recorder/worker.rs b/src/recorder/worker.rs index a6da90f0..cbef6b11 100644 --- a/src/recorder/worker.rs +++ b/src/recorder/worker.rs @@ -11,6 +11,7 @@ // vectorchord-inquiry@tensorchord.ai // // Copyright (c) 2025 TensorChord Inc. + use crate::recorder::types::PgRefCell; use std::cell::RefMut; use std::fs; diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index a520ba39..c5844c44 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -149,10 +149,11 @@ STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_vectors_wrapper'; CREATE OR REPLACE FUNCTION vchordrq_sampled_queries(regclass) RETURNS TABLE( schema_name NAME, + index_name NAME, table_name NAME, column_name NAME, operator TEXT, - vector_text TEXT + value TEXT ) LANGUAGE plpgsql STRICT AS $$ @@ -175,6 +176,7 @@ BEGIN WITH index_metadata AS ( SELECT NS.nspname AS schema_name, + I.relname AS index_name, C.relname AS table_name, PA.attname AS column_name, CASE @@ -205,13 +207,14 @@ BEGIN ) SELECT im.schema_name, + im.index_name, im.table_name, im.column_name, im.operator, - s.vector_text + s.value FROM index_metadata im, - LATERAL %2$I.vchordrq_sampled_vectors(%1$s) AS s(vector_text); + LATERAL %2$I.vchordrq_sampled_vectors(%1$s) AS s(value); $q$, $1::oid, ext_schema @@ -223,10 +226,11 @@ $$; CREATE VIEW vchordrq_sampled_queries AS SELECT record.schema_name, + record.index_name, record.table_name, record.column_name, record.operator, - record.vector_text + record.value FROM ( SELECT i.oid diff --git a/tests/vchordrq/recall.slt b/tests/vchordrq/recall.slt index ca3fb095..d75068cc 100644 --- a/tests/vchordrq/recall.slt +++ b/tests/vchordrq/recall.slt @@ -86,12 +86,12 @@ statement ok SELECT * from t_dim4 ORDER BY val <-> '[1.00, 0.50, 0.25, 0]'; query I -SELECT vector_text from vchordrq_sampled_queries('idx1'); +SELECT value from vchordrq_sampled_queries('idx1'); ---- [0.5,0.25,1] query I -SELECT vector_text from vchordrq_sampled_queries('idx2'); +SELECT value from vchordrq_sampled_queries('idx2'); ---- [1,0.5,0.25,0] @@ -106,7 +106,7 @@ SELECT * from t_dim4 ORDER BY val <-> '[2.1, 0.3, 0.7, 0.9]'; query I SELECT * from vchordrq_sampled_queries('idx2'); ---- -public t_dim4 val <-> [2.1,0.3,0.7,0.9] +public idx2 t_dim4 val <-> [2.1,0.3,0.7,0.9] query I SELECT AVG(recall_value) @@ -119,7 +119,7 @@ FROM ( lq.table_name, lq.column_name, lq.operator, - lq.vector_text + lq.value ) ) AS recall_value FROM @@ -146,7 +146,7 @@ SELECT column_name from vchordrq_sampled_queries('idx3'); NULL query I -SELECT vector_text from vchordrq_sampled_queries('idx3'); +SELECT value from vchordrq_sampled_queries('idx3'); ---- [1,0.5,0.25] @@ -154,7 +154,7 @@ statement ok SET search_path='@'; query I -SELECT vector_text from public.vchordrq_sampled_queries('public.idx3'); +SELECT value from public.vchordrq_sampled_queries('public.idx3'); ---- [1,0.5,0.25] From c9c1eb3cccb965d9347f40f4b9f77157f721afec Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 11 Sep 2025 17:30:31 +0800 Subject: [PATCH 214/324] fix: check mvcc in heap fetch (#336) https://github.com/tensorchord/VectorChord/issues/335 --------- Signed-off-by: usamoi --- .github/workflows/check.yml | 24 ++++++++--------- crates/make/src/main.rs | 6 +++-- src/index/fetcher.rs | 52 ++++++++++++++++++++++++++++++------ src/index/vchordg/am/mod.rs | 1 + src/index/vchordrq/am/mod.rs | 1 + tests/vchordrq/issue_335.slt | 37 +++++++++++++++++++++++++ 6 files changed, 99 insertions(+), 22 deletions(-) create mode 100644 tests/vchordrq/issue_335.slt diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index afd34a42..b3b62f3d 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -137,28 +137,28 @@ jobs: uses: actions/checkout@v4 - name: Clippy - run: cargo clippy --workspace --exclude vchord + run: cargo clippy --locked --workspace --exclude vchord - name: Cargo Test - run: cargo test --workspace --exclude vchord --exclude simd --no-fail-fast + run: cargo test --locked --workspace --exclude vchord --exclude simd --no-fail-fast - name: Cargo Test (simd) run: | if [ "$(uname -m)" == "x86_64" ]; then cargo \ --config 'target.'\''cfg(all())'\''.runner = ["/opt/sde/sde64", "-spr", "--"]' \ - test -p simd -- --no-capture + test --locked -p simd -- --no-capture fi if [ "$(uname -m)" == "aarch64" ]; then cargo \ --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=16"]' \ - test -p simd -- --no-capture + test --locked -p simd -- --no-capture cargo \ --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=32"]' \ - test -p simd -- --no-capture + test --locked -p simd -- --no-capture cargo \ --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=64"]' \ - test -p simd -- --no-capture + test --locked -p simd -- --no-capture fi psql: @@ -218,7 +218,7 @@ jobs: uses: actions/checkout@v4 - name: Clippy - run: cargo clippy --target ${{ matrix.arch }}-unknown-linux-gnu -p vchord --features pg${{ matrix.version }} --no-deps + run: cargo clippy --locked --target ${{ matrix.arch }}-unknown-linux-gnu -p vchord --features pg${{ matrix.version }} --no-deps - name: Install run: | @@ -301,7 +301,7 @@ jobs: - name: Clippy run: | - cargo clippy -p vchord --target ${{ matrix.arch }}-apple-darwin --features pg${{ matrix.version }} --no-deps + cargo clippy --locked -p vchord --target ${{ matrix.arch }}-apple-darwin --features pg${{ matrix.version }} --no-deps - name: Install run: | @@ -407,7 +407,7 @@ jobs: - name: Clippy run: | - cargo clippy --target ${{ matrix.arch }}-pc-windows-msvc -p vchord --features pg${{ matrix.version }} --no-deps + cargo clippy --locked --target ${{ matrix.arch }}-pc-windows-msvc -p vchord --features pg${{ matrix.version }} --no-deps - name: Install run: | @@ -514,7 +514,7 @@ jobs: - name: Clippy run: | . "$HOME/.cargo/env" - cargo clippy --target ${{ matrix.arch }}-unknown-linux-musl -p vchord --features pg${{ matrix.version }} --no-deps + cargo clippy --locked --target ${{ matrix.arch }}-unknown-linux-musl -p vchord --features pg${{ matrix.version }} --no-deps - name: Install run: | @@ -711,10 +711,10 @@ jobs: - name: Clippy & Test run: | PGRX_PG_CONFIG_PATH=pg_config \ - cargo clippy --target ${{ matrix.rust_triple }} \ + cargo clippy --locked --target ${{ matrix.rust_triple }} \ --workspace --features pg${{ matrix.version }} PGRX_PG_CONFIG_PATH=pg_config \ - cargo test --target ${{ matrix.rust_triple }} \ + cargo test --locked --target ${{ matrix.rust_triple }} \ --workspace --exclude vchord --no-fail-fast \ -- --no-capture diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index 51d217cb..9e980e0b 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -168,7 +168,8 @@ fn build( ) -> Result> { let mut command = Command::new("cargo"); command - .args(["build", "-p", "vchord", "--lib"]) + .args(["build", "--locked"]) + .args(["-p", "vchord", "--lib"]) .args(["--profile", profile]) .args(["--target", target]) .args(["--features", pg_version]) @@ -265,7 +266,8 @@ fn generate( )?; let mut command = Command::new("cargo"); command - .args(["rustc", "-p", "vchord", "--bin", "pgrx_embed_vchord"]) + .args(["rustc", "--locked"]) + .args(["-p", "vchord", "--bin", "pgrx_embed_vchord"]) .args(["--profile", profile]) .args(["--target", target]) .args(["--features", pg_version]) diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs index 6726ce0c..f13eaac9 100644 --- a/src/index/fetcher.rs +++ b/src/index/fetcher.rs @@ -48,6 +48,7 @@ pub struct HeapFetcher { econtext: *mut pgrx::pg_sys::ExprContext, heap_relation: pgrx::pg_sys::Relation, snapshot: pgrx::pg_sys::Snapshot, + heapfetch: *mut pgrx::pg_sys::IndexFetchTableData, slot: *mut pgrx::pg_sys::TupleTableSlot, values: [Datum; 32], is_nulls: [bool; 32], @@ -59,6 +60,7 @@ impl HeapFetcher { index_relation: pgrx::pg_sys::Relation, heap_relation: pgrx::pg_sys::Relation, snapshot: pgrx::pg_sys::Snapshot, + heapfetch: *mut pgrx::pg_sys::IndexFetchTableData, hack: *mut pgrx::pg_sys::IndexScanState, ) -> Self { unsafe { @@ -71,6 +73,7 @@ impl HeapFetcher { econtext, heap_relation, snapshot, + heapfetch, slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), values: [Datum::null(); 32], is_nulls: [true; 32], @@ -99,16 +102,49 @@ impl Fetcher for HeapFetcher { use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; let mut ctid = key_to_ctid(key); let table_am = (*self.heap_relation).rd_tableam; - let fetch_row_version = (*table_am) - .tuple_fetch_row_version + let index_fetch_tuple = (*table_am) + .index_fetch_tuple .expect("unsupported heap access method"); - #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] - if pg_guard_ffi_boundary(|| { - !fetch_row_version(self.heap_relation, &mut ctid, self.snapshot, self.slot) - }) { - return None; + let found = 'a: { + let mut call_again = false; + let mut all_dead = false; + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + let found = pg_guard_ffi_boundary(|| { + index_fetch_tuple( + self.heapfetch, + &mut ctid, + self.snapshot, + self.slot, + &mut call_again, + &mut all_dead, + ) + }); + if found { + break 'a true; + } + while call_again { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + let found = pg_guard_ffi_boundary(|| { + index_fetch_tuple( + self.heapfetch, + &mut ctid, + self.snapshot, + self.slot, + &mut call_again, + &mut all_dead, + ) + }); + if found { + break 'a true; + } + } + false + }; + if found { + Some(HeapTuple { this: self }) + } else { + None } - Some(HeapTuple { this: self }) } } } diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index d3f5dca5..fc2f2e1d 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -365,6 +365,7 @@ pub unsafe extern "C-unwind" fn amrescan( (*scan).indexRelation, (*scan).heapRelation, (*scan).xs_snapshot, + (*scan).xs_heapfetch, if let Some(hack) = hack { hack.as_ptr() } else { diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 45e774b4..bb24d316 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -448,6 +448,7 @@ pub unsafe extern "C-unwind" fn amrescan( (*scan).indexRelation, (*scan).heapRelation, (*scan).xs_snapshot, + (*scan).xs_heapfetch, if let Some(hack) = hack { hack.as_ptr() } else { diff --git a/tests/vchordrq/issue_335.slt b/tests/vchordrq/issue_335.slt new file mode 100644 index 00000000..7a5dfd9e --- /dev/null +++ b/tests/vchordrq/issue_335.slt @@ -0,0 +1,37 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3), visible boolean default true); + +statement ok +INSERT INTO items (embedding) SELECT ARRAY[i * 0.001, i * 0.001, i * 0.001]::real[] FROM generate_series(1, 1000) s(i); + +statement ok +CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops); + +statement ok +SET vchordrq.prefilter = on; + +query I +SELECT id FROM items WHERE visible = true ORDER BY embedding <-> '[0.0031,0.0031,0.0031]' LIMIT 3; +---- +3 +4 +2 + +statement ok +UPDATE items SET visible = false WHERE id = 3; + +statement ok +UPDATE items SET visible = true WHERE id = 3; + +query I +SELECT id FROM items WHERE visible = true ORDER BY embedding <-> '[0.0031,0.0031,0.0031]' LIMIT 3; +---- +3 +4 +2 + +statement ok +DROP TABLE items; From 346364716867a3329a65721e7f329363bd6fae8b Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Mon, 15 Sep 2025 09:45:58 +0800 Subject: [PATCH 215/324] fix: permission and sql for query sample (#337) - Add permission(superuser-only) for 3 query sampling GUCs - Refactor SQL with better compatibility Signed-off-by: cutecutecat --- src/index/gucs.rs | 6 +++--- src/sql/finalize.sql | 16 ++++++++-------- tests/vchordrq/recall.slt | 2 +- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/index/gucs.rs b/src/index/gucs.rs index ed501a52..5b74bc2b 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -169,7 +169,7 @@ pub fn init() { c"`query_sampling_enable` argument of vchordrq.", c"`query_sampling_enable` argument of vchordrq.", &VCHORDRQ_QUERY_SAMPLING_ENABLE, - GucContext::Userset, + GucContext::Suset, GucFlags::default(), ); GucRegistry::define_int_guc( @@ -179,7 +179,7 @@ pub fn init() { &VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS, 0, 10000, - GucContext::Userset, + GucContext::Suset, GucFlags::default(), ); GucRegistry::define_float_guc( @@ -189,7 +189,7 @@ pub fn init() { &VCHORDRQ_QUERY_SAMPLING_RATE, 0.0, 1.0, - GucContext::Userset, + GucContext::Suset, GucFlags::default(), ); unsafe { diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index c5844c44..7396fb74 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -152,7 +152,7 @@ RETURNS TABLE( index_name NAME, table_name NAME, column_name NAME, - operator TEXT, + operator NAME, value TEXT ) LANGUAGE plpgsql @@ -179,12 +179,7 @@ BEGIN I.relname AS index_name, C.relname AS table_name, PA.attname AS column_name, - CASE - WHEN OP.opcname LIKE '%%l2%%' THEN '<->' - WHEN OP.opcname LIKE '%%ip%%' THEN '<#>' - WHEN OP.opcname LIKE '%%cosine%%' THEN '<=>' - ELSE '' - END AS operator + OP.oprname AS operator FROM pg_catalog.pg_index X JOIN @@ -196,11 +191,16 @@ BEGIN JOIN pg_catalog.pg_am A ON A.oid = I.relam LEFT JOIN - pg_catalog.pg_opclass AS OP ON OP.oid = X.indclass[0] + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr LEFT JOIN pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] WHERE A.amname = 'vchordrq' + AND AO.amopstrategy = 1 AND C.relkind = 'r' AND X.indnatts = 1 AND X.indexrelid = %1$s diff --git a/tests/vchordrq/recall.slt b/tests/vchordrq/recall.slt index d75068cc..475046fa 100644 --- a/tests/vchordrq/recall.slt +++ b/tests/vchordrq/recall.slt @@ -114,7 +114,7 @@ FROM ( SELECT vchordrq_evaluate_query_recall( query => format( - 'SELECT ctid FROM %I.%I ORDER BY %I %s ''%s'' LIMIT 10', + 'SELECT ctid FROM %I.%I ORDER BY %I OPERATOR(%s) %L LIMIT 10', lq.schema_name, lq.table_name, lq.column_name, From fef6806174b03ee0157ea756b4c5794a377cdd20 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 15 Sep 2025 10:04:05 +0800 Subject: [PATCH 216/324] fix: schema (#338) Signed-off-by: usamoi --- src/index/functions.rs | 2 +- src/sql/finalize.sql | 52 +++++++++++++++++++++--------------------- 2 files changed, 27 insertions(+), 27 deletions(-) diff --git a/src/index/functions.rs b/src/index/functions.rs index 17cf6741..65c71c78 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -88,7 +88,7 @@ impl Drop for Index { } #[pgrx::pg_extern(sql = "")] -fn _vchordrq_sampled_vectors(indexrelid: Oid) -> SetOfIterator<'static, String> { +fn _vchordrq_sampled_values(indexrelid: Oid) -> SetOfIterator<'static, String> { let pg_am = PgAm::search_amname(c"vchordrq").unwrap(); let Some(pg_am) = pg_am.get() else { pgrx::error!("vchord is not installed"); diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 7396fb74..a1aadcdc 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -142,11 +142,10 @@ IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_ CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; -CREATE FUNCTION vchordrq_sampled_vectors(regclass) -RETURNS SETOF TEXT -STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_vectors_wrapper'; +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; -CREATE OR REPLACE FUNCTION vchordrq_sampled_queries(regclass) +CREATE FUNCTION vchordrq_sampled_queries(regclass) RETURNS TABLE( schema_name NAME, index_name NAME, @@ -155,8 +154,7 @@ RETURNS TABLE( operator NAME, value TEXT ) -LANGUAGE plpgsql -STRICT AS $$ +STRICT LANGUAGE plpgsql AS $$ DECLARE ext_schema TEXT; query_text TEXT; @@ -214,7 +212,7 @@ BEGIN s.value FROM index_metadata im, - LATERAL %2$I.vchordrq_sampled_vectors(%1$s) AS s(value); + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); $q$, $1::oid, ext_schema @@ -223,25 +221,6 @@ BEGIN END; $$; -CREATE VIEW vchordrq_sampled_queries AS -SELECT - record.schema_name, - record.index_name, - record.table_name, - record.column_name, - record.operator, - record.value -FROM - ( - SELECT i.oid - FROM pg_catalog.pg_class AS i - JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid - JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid - JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid - WHERE am.amname = 'vchordrq' - ) AS index_oids -CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; - CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; @@ -432,3 +411,24 @@ CREATE OPERATOR CLASS halfvec_cosine_ops OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; From 1cf6068e40a7b839a97866fd4a333b17a6ef6c40 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 15 Sep 2025 10:26:12 +0800 Subject: [PATCH 217/324] chore: upload 0.5.2 schema scripts (#339) Signed-off-by: usamoi --- {tools => devtools}/dump-codegen.sh | 0 {tools => devtools}/dump.sh | 0 sql/install/vchord--0.5.2.sql | 857 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.5.1--0.5.2.sql | 101 ++++ 4 files changed, 958 insertions(+) rename {tools => devtools}/dump-codegen.sh (100%) rename {tools => devtools}/dump.sh (100%) create mode 100644 sql/install/vchord--0.5.2.sql create mode 100644 sql/upgrade/vchord--0.5.1--0.5.2.sql diff --git a/tools/dump-codegen.sh b/devtools/dump-codegen.sh similarity index 100% rename from tools/dump-codegen.sh rename to devtools/dump-codegen.sh diff --git a/tools/dump.sh b/devtools/dump.sh similarity index 100% rename from tools/dump.sh rename to devtools/dump.sh diff --git a/sql/install/vchord--0.5.2.sql b/sql/install/vchord--0.5.2.sql new file mode 100644 index 00000000..0e7e473e --- /dev/null +++ b/sql/install/vchord--0.5.2.sql @@ -0,0 +1,857 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:46 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:31 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:21 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:78 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:78 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:47 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/upgrade/vchord--0.5.1--0.5.2.sql b/sql/upgrade/vchord--0.5.1--0.5.2.sql new file mode 100644 index 00000000..a28eefcd --- /dev/null +++ b/sql/upgrade/vchord--0.5.1--0.5.2.sql @@ -0,0 +1,101 @@ +-- List of functions + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; From 5be691fb5b67d356dd97d7c2a9e1c36b3697fc4c Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 15 Sep 2025 11:23:37 +0800 Subject: [PATCH 218/324] build: do not use --locked (#341) job: +check_debian Signed-off-by: usamoi --- .github/workflows/check.yml | 4 ++-- crates/make/src/main.rs | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index b3b62f3d..7291b49e 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -711,10 +711,10 @@ jobs: - name: Clippy & Test run: | PGRX_PG_CONFIG_PATH=pg_config \ - cargo clippy --locked --target ${{ matrix.rust_triple }} \ + cargo clippy --target ${{ matrix.rust_triple }} \ --workspace --features pg${{ matrix.version }} PGRX_PG_CONFIG_PATH=pg_config \ - cargo test --locked --target ${{ matrix.rust_triple }} \ + cargo test --target ${{ matrix.rust_triple }} \ --workspace --exclude vchord --no-fail-fast \ -- --no-capture diff --git a/crates/make/src/main.rs b/crates/make/src/main.rs index 9e980e0b..572774d1 100644 --- a/crates/make/src/main.rs +++ b/crates/make/src/main.rs @@ -168,7 +168,7 @@ fn build( ) -> Result> { let mut command = Command::new("cargo"); command - .args(["build", "--locked"]) + .args(["build"]) .args(["-p", "vchord", "--lib"]) .args(["--profile", profile]) .args(["--target", target]) @@ -266,7 +266,7 @@ fn generate( )?; let mut command = Command::new("cargo"); command - .args(["rustc", "--locked"]) + .args(["rustc"]) .args(["-p", "vchord", "--bin", "pgrx_embed_vchord"]) .args(["--profile", profile]) .args(["--target", target]) From b43d9bb6c6f0a3d4cb5483aabed15c51a4ef3ea7 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 15 Sep 2025 14:47:25 +0800 Subject: [PATCH 219/324] fix: order of probes (#342) Signed-off-by: usamoi --- crates/vchordrq/src/search.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index 4d18db9b..89e21362 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -123,7 +123,7 @@ where }) }; - for i in (1..height_of_root).rev() { + for i in 1..height_of_root { state = step(state).take(probes[i as usize - 1] as _).collect(); } @@ -265,7 +265,7 @@ where }; let mut it = None; - for i in (1..height_of_root).rev() { + for i in 1..height_of_root { let it = it.insert(step(state)); state = it.take(probes[i as usize - 1] as _).collect(); } From 9d1f2ec8817386efac3b7e7047face39496e39af Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 15 Sep 2025 17:44:22 +0800 Subject: [PATCH 220/324] readme: sync with docs (#340) blocked by https://github.com/tensorchord/pgvecto.rs-docs/pull/143 Signed-off-by: usamoi --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index e0996700..56f13ecf 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.5.1 + -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.5.2 ``` > [!NOTE] > In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). From 0021aa7f57529a8ec588dfdfb0c8b9a01836fa02 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 15 Sep 2025 17:46:29 +0800 Subject: [PATCH 221/324] chore: remove dockerfile (#343) Since `docker build` makes releasing difficult, `Dockerfile` has been moved to [VectorChord-images](https://github.com/tensorchord/VectorChord-images) now. Signed-off-by: usamoi --- .github/workflows/release.yml | 87 ----------------------------------- docker/Dockerfile | 17 ------- 2 files changed, 104 deletions(-) delete mode 100644 docker/Dockerfile diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 6e0d0ef0..77f67650 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -120,93 +120,6 @@ jobs: gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip - docker: - runs-on: ${{ matrix.runner }} - needs: ["semver", "build"] - strategy: - matrix: - version: ["14", "15", "16", "17"] - runner: ["ubuntu-22.04", "ubuntu-22.04-arm"] - env: - SEMVER: ${{ needs.semver.outputs.SEMVER }} - PLATFORM: ${{ matrix.runner == 'ubuntu-22.04' && 'amd64' || 'arm64' }} - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Download - env: - GH_TOKEN: ${{ github.token }} - run: | - mkdir -p build - gh release download $SEMVER --pattern "postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${PLATFORM}.deb" --output ./build/postgresql-${{ matrix.version }}-vchord_${SEMVER}-1_${PLATFORM}.deb - - name: Set up QEMU - uses: docker/setup-qemu-action@v3 - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 - - name: Login to Docker Hub - uses: docker/login-action@v3 - with: - username: ${{ secrets.DOCKERIO_USERNAME }} - password: ${{ secrets.DOCKERIO_TOKEN }} - - name: Login to GHCR - uses: docker/login-action@v3 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} - - name: Push PostgreSQL release to Docker Registry - uses: docker/build-push-action@v6 - with: - context: . - push: true - platforms: ${{ matrix.runner == 'ubuntu-22.04' && 'linux/amd64' || 'linux/arm64' }} - file: ./docker/Dockerfile - provenance: false - tags: | - tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} - ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-${{ env.PLATFORM }} - build-args: | - BASE=${{ matrix.version }}-bookworm - VCHORD=${{ env.SEMVER }} - PGVECTOR=0.8.0 - - create-manifests: - runs-on: ubuntu-latest - needs: ["semver", "build", "docker"] - strategy: - matrix: - version: ["14", "15", "16", "17"] - env: - SEMVER: ${{ needs.semver.outputs.SEMVER }} - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Set up QEMU - uses: docker/setup-qemu-action@v3 - - name: Login to Docker Hub - uses: docker/login-action@v3 - with: - username: ${{ secrets.DOCKERIO_USERNAME }} - password: ${{ secrets.DOCKERIO_TOKEN }} - - name: Login to GHCR - uses: docker/login-action@v3 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} - - name: Create manifest and push - run: | - docker manifest create \ - tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} \ - --amend tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ - --amend tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 - docker manifest push tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} - docker manifest create \ - ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} \ - --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-amd64 \ - --amend ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }}-arm64 - docker manifest push ghcr.io/tensorchord/vchord-postgres:pg${{ matrix.version }}-v${{ env.SEMVER }} - pgxn: runs-on: "ubuntu-latest" needs: ["semver", "build"] diff --git a/docker/Dockerfile b/docker/Dockerfile deleted file mode 100644 index 04c15c8e..00000000 --- a/docker/Dockerfile +++ /dev/null @@ -1,17 +0,0 @@ -ARG BASE=17-bookworm - -FROM postgres:${BASE} - -ARG VCHORD -ARG PGVECTOR - -RUN apt-get update && \ - apt-get install -y --no-install-recommends wget ca-certificates && \ - wget https://github.com/tensorchord/VectorChord/releases/download/${VCHORD}/postgresql-${PG_MAJOR}-vchord_${VCHORD}-1_$(dpkg --print-architecture).deb -P /tmp && \ - apt-get install -y postgresql-${PG_MAJOR}-pgvector=${PGVECTOR}-* && \ - apt-get install -y /tmp/postgresql-${PG_MAJOR}-vchord_${VCHORD}-1_$(dpkg --print-architecture).deb && \ - apt-get remove -y wget ca-certificates && \ - apt-get purge -y --auto-remove && \ - rm -rf /tmp/* /var/lib/apt/lists/* - -CMD ["postgres", "-c" ,"shared_preload_libraries=vchord"] From d4b592b1a6fcade9b95effe5f6b6d22aedecc7a0 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 16 Sep 2025 17:11:53 +0800 Subject: [PATCH 222/324] ci: add beta rust and beta postgresql (#345) Signed-off-by: usamoi --- .github/workflows/check.yml | 66 ++++++++++++++++++-------------- tests/vchordrq/pushdown_plan.slt | 20 ++++++++++ 2 files changed, 57 insertions(+), 29 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 7291b49e..0ed4a000 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -11,7 +11,7 @@ concurrency: jobs: style: - runs-on: "ubuntu-latest" + runs-on: "ubuntu-24.04" steps: - name: Set up Environment @@ -169,7 +169,7 @@ jobs: strategy: matrix: - version: ["13", "14", "15", "16", "17"] + version: ["13", "14", "15", "16", "17", "18"] arch: ["x86_64", "aarch64"] runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} @@ -185,8 +185,10 @@ jobs: - name: Set up Environment run: | rustup update - - sudo apt-get update + if [ "${{ matrix.version }}" = "18" ]; then + rustup default beta + rustup component add clippy rustfmt + fi sudo apt-get remove -y '^postgres.*' '^libpq.*' sudo apt-get purge -y '^postgres.*' '^libpq.*' @@ -194,8 +196,14 @@ jobs: curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 + sudo apt-get update sudo apt-get install -y postgresql-common sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + if [ "${{ matrix.version }}" = "18" ]; then + sudo add-apt-repository "deb https://apt.postgresql.org/pub/repos/apt/ $(lsb_release -s -c)-pgdg main 18" + fi + + sudo apt-get update sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} echo PGRX_PG_CONFIG_PATH=pg_config >> $GITHUB_ENV echo PG_CONFIG=pg_config >> $GITHUB_ENV @@ -240,11 +248,11 @@ jobs: - name: Sqllogictest run: | - sqllogictest --db $USER --user $USER './tests/general/*.slt' - sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' - sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' - if [ "${{ matrix.version }}" = "17" ]; then - sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' + sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi psql_macos: @@ -325,11 +333,11 @@ jobs: - name: Sqllogictest run: | - sqllogictest --db $USER --user $USER './tests/general/*.slt' - sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' - sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' - if [ "${{ matrix.version }}" = "17" ]; then - sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' + sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi psql_windows: @@ -431,11 +439,11 @@ jobs: - name: Sqllogictest run: | - sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/general/*.slt' - sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordg/*.slt' - sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/*.slt' - if ( "${{ matrix.version }}" -eq "17" ) { - sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/pg17/*.slt' + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if ( "${{ matrix.version }}" -ge "17" ) { + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} } psql_alpine: @@ -537,11 +545,11 @@ jobs: - name: Sqllogictest run: | - sqllogictest --db $USER --user $USER './tests/general/*.slt' - sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' - sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' - if [ "${{ matrix.version }}" = "17" ]; then - sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' + sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi check_debian: @@ -742,9 +750,9 @@ jobs: - name: Sqllogictest run: | - sqllogictest --db $USER --user $USER './tests/general/*.slt' - sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' - sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' - if [ "${{ matrix.version }}" = "17" ]; then - sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' + sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi diff --git a/tests/vchordrq/pushdown_plan.slt b/tests/vchordrq/pushdown_plan.slt index f6aecc48..d732d3bc 100644 --- a/tests/vchordrq/pushdown_plan.slt +++ b/tests/vchordrq/pushdown_plan.slt @@ -28,11 +28,22 @@ SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) ORDER BY val0 Order By: (val0 <-> '[0,0,0]'::vector) # 1 vector key + 0 order_by key + original style + +skipif pg18 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +onlyif pg18 query I EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; ---- Seq Scan on t + Disabled: true Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) # 1 vector key + 0 order_by key + sphere style @@ -135,10 +146,19 @@ ORDER BY val0 <-> '[0, 0, 0]'; # Filter: (val0 <<->> '("[0,0,0]",1)'::sphere_vector) # 0 vector key + 0 order_by key(variable) + +skipif pg18 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +onlyif pg18 query I EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; ---- Seq Scan on t + Disabled: true statement ok DROP TABLE t; From 5bdb993d755a14e781f1247012bd5f40188e090d Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 18 Sep 2025 16:08:55 +0800 Subject: [PATCH 223/324] chore: multiple fixes and fast paths (#346) 1. fixes `build.pin` does not work if the build is not parallelized 2. adds fast path if `lists <= probes && !residual_quantization` 3. fixes cost estimation if `lists < probes` 4. adds `build.default` to reduce special handling 5. rotate vectors after all kinds of `build` to reduce special handling Signed-off-by: usamoi --- crates/k_means/src/lib.rs | 38 +----- crates/rabitq/src/bit.rs | 12 +- crates/rabitq/src/packing.rs | 21 ++-- crates/vchordrq/src/bulkdelete.rs | 2 +- crates/vchordrq/src/cache.rs | 2 +- crates/vchordrq/src/insert.rs | 20 ++-- crates/vchordrq/src/maintain.rs | 2 +- crates/vchordrq/src/prewarm.rs | 6 +- crates/vchordrq/src/search.rs | 95 ++++++++++----- crates/vchordrq/src/types.rs | 2 +- src/index/vchordg/am/am_build.rs | 89 +++++++++----- src/index/vchordrq/am/am_build.rs | 174 ++++++++++++++++++++-------- src/index/vchordrq/am/mod.rs | 2 +- src/index/vchordrq/types.rs | 15 ++- tests/vchordg/sequential_build.slt | 17 +++ tests/vchordrq/sequential_build.slt | 24 ++++ 16 files changed, 341 insertions(+), 180 deletions(-) create mode 100644 tests/vchordg/sequential_build.slt create mode 100644 tests/vchordrq/sequential_build.slt diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index 508fb7d9..aab50292 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -16,23 +16,9 @@ use rabitq::bit::block::BlockCode; use rabitq::packing::{any_pack, padding_pack}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; -use rayon::iter::{IntoParallelIterator, IntoParallelRefMutIterator, ParallelIterator}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; use simd::Floating; -pub fn preprocess(num_threads: usize, x: &mut [T], f: impl Fn(&mut T) + Sync) { - rayon::ThreadPoolBuilder::new() - .num_threads(num_threads) - .build_scoped( - |thread| thread.run(), - move |pool| { - pool.install(|| { - x.par_iter_mut().for_each(&f); - }); - }, - ) - .expect("failed to build thread pool") -} - pub fn k_means( num_threads: usize, mut check: impl FnMut(usize), @@ -57,7 +43,7 @@ pub fn k_means( if n >= 1024 && c >= 1024 { rabitq_index(n, c, samples, centroids) } else { - flat_index(dims, n, c, samples, centroids) + flat_index(n, c, samples, centroids) } }; let mut lloyd_k_means = @@ -75,18 +61,6 @@ pub fn k_means( } } -pub fn k_means_lookup(vector: &[f32], centroids: &[Vec]) -> usize { - assert_ne!(centroids.len(), 0); - let mut result = (f32::INFINITY, 0); - for i in 0..centroids.len() { - let dis = f32::reduce_sum_of_d2(vector, ¢roids[i]); - if dis <= result.0 { - result = (dis, i); - } - } - result.1 -} - fn quick_centers( c: usize, dims: usize, @@ -175,13 +149,7 @@ fn rabitq_index(n: usize, c: usize, samples: &[Vec], centroids: &[Vec] .collect::>() } -fn flat_index( - _dims: usize, - n: usize, - c: usize, - samples: &[Vec], - centroids: &[Vec], -) -> Vec { +fn flat_index(n: usize, c: usize, samples: &[Vec], centroids: &[Vec]) -> Vec { (0..n) .into_par_iter() .map(|i| { diff --git a/crates/rabitq/src/bit.rs b/crates/rabitq/src/bit.rs index 3d8fe9e2..27dd5237 100644 --- a/crates/rabitq/src/bit.rs +++ b/crates/rabitq/src/bit.rs @@ -73,15 +73,9 @@ pub fn code(vector: &[f32]) -> Code { CodeMetadata { dis_u_2: sum_of_x_2, factor_cnt: { - let cnt_pos = vector - .iter() - .map(|x| x.is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.is_sign_negative() as i32) - .sum::(); - (cnt_pos - cnt_neg) as f32 + let cnt_pos = vector.iter().filter(|x| x.is_sign_positive()).count(); + let cnt_neg = vector.iter().filter(|x| x.is_sign_negative()).count(); + cnt_pos as f32 - cnt_neg as f32 }, factor_ip: sum_of_x_2 / sum_of_abs_x, factor_err: { diff --git a/crates/rabitq/src/packing.rs b/crates/rabitq/src/packing.rs index 008140dd..79c8e2a4 100644 --- a/crates/rabitq/src/packing.rs +++ b/crates/rabitq/src/packing.rs @@ -101,14 +101,15 @@ pub fn any_pack(mut x: impl Iterator) -> [T; 32] { std::array::from_fn(|_| x.next()).map(|x| x.unwrap_or_default()) } -pub fn pack_to_u4(signs: &[bool]) -> Vec { - fn f(x: [bool; 4]) -> u8 { - x[0] as u8 | (x[1] as u8) << 1 | (x[2] as u8) << 2 | (x[3] as u8) << 3 - } - let mut result = Vec::with_capacity(signs.len().div_ceil(4)); - for i in 0..signs.len().div_ceil(4) { - let x = std::array::from_fn(|j| signs.get(i * 4 + j).copied().unwrap_or_default()); - result.push(f(x)); - } - result +pub fn pack_to_u4(input: &[bool]) -> Vec { + let f = |t: &[bool; 4]| t[0] as u8 | (t[1] as u8) << 1 | (t[2] as u8) << 2 | (t[3] as u8) << 3; + let (arrays, remainder) = input.as_chunks::<4>(); + let mut buffer = [false; 4]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() } diff --git a/crates/vchordrq/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs index 00a446ff..3f2556fd 100644 --- a/crates/vchordrq/src/bulkdelete.rs +++ b/crates/vchordrq/src/bulkdelete.rs @@ -43,7 +43,7 @@ pub fn bulkdelete( for first in state { tape::read_h1_tape::( by_next(index, first).inspect(|_| check()), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), |(), _, _, first, _| results.push(first), ); } diff --git a/crates/vchordrq/src/cache.rs b/crates/vchordrq/src/cache.rs index 0a21f3cd..eecc7275 100644 --- a/crates/vchordrq/src/cache.rs +++ b/crates/vchordrq/src/cache.rs @@ -39,7 +39,7 @@ where for first in state { tape::read_h1_tape::( by_next(index, first).inspect(|guard| trace.push(guard.id())), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), |(), _, _, first, _| { results.push(first); }, diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index 3d6b337f..77d1c448 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -77,25 +77,23 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let epsilon = 1.9; type State = (Reverse, AlwaysEqual, AlwaysEqual); - let mut state: State = if !is_residual { - let first = meta_tuple.first(); - // it's safe to leave it a fake value - ( - Reverse(Distance::ZERO), - AlwaysEqual(0.0), - AlwaysEqual(first), - ) - } else { + let mut state: State = if is_residual { let prefetch = BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); let head = meta_tuple.centroid_head(); - let norm = meta_tuple.centroid_norm(); - let first = meta_tuple.first(); let distance = vectors::read_for_h1_tuple::( prefetch.map(|id| index.read(id)), head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + (Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first)) + } else { + // fast path + let distance = Distance::ZERO; + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); (Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first)) }; diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs index 3d559c50..4b7b1866 100644 --- a/crates/vchordrq/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -57,7 +57,7 @@ where for first in state { tape::read_h1_tape::( by_next(index, first).inspect(|_| check()), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), |(), _, _, first, _| results.push(first), ); } diff --git a/crates/vchordrq/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs index d8066f69..6b33a165 100644 --- a/crates/vchordrq/src/prewarm.rs +++ b/crates/vchordrq/src/prewarm.rs @@ -64,7 +64,7 @@ where for first in state { tape::read_h1_tape::( by_next(index, first).inspect(|_| counter += 1), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), |(), head, _, first, prefetch| { vectors::read_for_h1_tuple::( prefetch.iter().map(|&id| index.read(id)), @@ -97,7 +97,7 @@ where tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); tape::read_frozen_tape::( by_directory(&mut prefetch_h0_tuples, directory).inspect(|_| counter += 1), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), id_2(|_, _, _, _| { results.push(()); }), @@ -105,7 +105,7 @@ where } else { tape::read_frozen_tape::( by_next(index, jump_tuple.frozen_first()).inspect(|_| counter += 1), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); 32])), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), id_2(|_, _, _, _| { results.push(()); }), diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index 89e21362..4b3bd2cd 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -12,13 +12,13 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::closure_lifetime_binder::id_2; +use crate::closure_lifetime_binder::{id_0, id_1, id_2}; use crate::linked_vec::LinkedVec; use crate::operator::*; use crate::tape::{by_directory, by_next}; use crate::tuples::*; use crate::{Opaque, Page, tape, vectors}; -use algo::accessor::LAccess; +use algo::accessor::{FunctionalAccessor, LAccess}; use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; use algo::{BorrowedIter, Bump, PackedRefMut4, PackedRefMut8, RelationRead}; use always_equal::AlwaysEqual; @@ -51,6 +51,7 @@ where let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); let height_of_root = meta_tuple.height_of_root(); + let cells = meta_tuple.cells().to_vec(); assert_eq!(dims, vector.dims(), "unmatched dimensions"); if height_of_root as usize != 1 + probes.len() { panic!( @@ -59,27 +60,26 @@ where probes.len() ); } + debug_assert_eq!(cells[(height_of_root - 1) as usize], 1); type State = Vec<(Reverse, AlwaysEqual, AlwaysEqual)>; - let mut state: State = if !is_residual { - let first = meta_tuple.first(); - // it's safe to leave it a fake value - vec![( - Reverse(Distance::ZERO), - AlwaysEqual(0.0), - AlwaysEqual(first), - )] - } else { + let mut state: State = if is_residual { let prefetch = BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); let head = meta_tuple.centroid_head(); - let norm = meta_tuple.centroid_norm(); - let first = meta_tuple.first(); let distance = vectors::read_for_h1_tuple::( prefetch.map(|id| index.read(id)), head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] + } else { + // fast path + let distance = Distance::ZERO; + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] }; @@ -124,7 +124,27 @@ where }; for i in 1..height_of_root { - state = step(state).take(probes[i as usize - 1] as _).collect(); + let partial_scan = probes[i as usize - 1] < cells[(height_of_root - 1 - i) as usize]; + if partial_scan || is_residual { + state = step(state).take(probes[i as usize - 1] as _).collect(); + } else { + // fast path + let mut results = LinkedVec::new(); + for (Reverse(_), AlwaysEqual(_), AlwaysEqual(first)) in state { + tape::read_h1_tape::( + by_next(index, first), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + |(), _, norm, first, _| { + results.push(( + Reverse(Distance::ZERO), + AlwaysEqual(norm), + AlwaysEqual(first), + )); + }, + ); + } + state = results.into_vec(); + } } let mut results = LinkedVec::<(_, AlwaysEqual<_>)>::new(); @@ -192,6 +212,7 @@ where let dims = meta_tuple.dims(); let is_residual = meta_tuple.is_residual(); let height_of_root = meta_tuple.height_of_root(); + let cells = meta_tuple.cells().to_vec(); assert_eq!(dims, vector.dims(), "unmatched dimensions"); if height_of_root as usize != 1 + probes.len() { panic!( @@ -200,27 +221,26 @@ where probes.len() ); } + debug_assert_eq!(cells[(height_of_root - 1) as usize], 1); type State = Vec<(Reverse, AlwaysEqual, AlwaysEqual)>; - let mut state: State = if !is_residual { - let first = meta_tuple.first(); - // it's safe to leave it a fake value - vec![( - Reverse(Distance::ZERO), - AlwaysEqual(0.0), - AlwaysEqual(first), - )] - } else { + let mut state: State = if is_residual { let prefetch = BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); let head = meta_tuple.centroid_head(); - let norm = meta_tuple.centroid_norm(); - let first = meta_tuple.first(); let distance = vectors::read_for_h1_tuple::( prefetch.map(|id| index.read(id)), head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), ); + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] + } else { + // fast path + let distance = Distance::ZERO; + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] }; @@ -266,8 +286,29 @@ where let mut it = None; for i in 1..height_of_root { - let it = it.insert(step(state)); - state = it.take(probes[i as usize - 1] as _).collect(); + let partial_scan = probes[i as usize - 1] < cells[(height_of_root - 1 - i) as usize]; + let needs_sort = i + 1 == height_of_root && threshold != 0; + if partial_scan || is_residual || needs_sort { + let it = it.insert(step(state)); + state = it.take(probes[i as usize - 1] as _).collect(); + } else { + // fast path + let mut results = LinkedVec::new(); + for (Reverse(_), AlwaysEqual(_), AlwaysEqual(first)) in state { + tape::read_h1_tape::( + by_next(index, first), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + |(), _, norm, first, _| { + results.push(( + Reverse(Distance::ZERO), + AlwaysEqual(norm), + AlwaysEqual(first), + )); + }, + ); + } + state = results.into_vec(); + } } let mut results = LinkedVec::<(_, AlwaysEqual<_>)>::new(); diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index 0f759bd1..9296dfd5 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -79,7 +79,7 @@ impl VectorKind { } } -#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] #[validate(schema(function = "Self::validate_self"))] pub struct VectorOptions { diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index a8cb518e..fbc63f56 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -197,27 +197,20 @@ pub unsafe extern "C-unwind" fn ambuild( let errors = "alpha not equal to `1.0` are only applicable to l2 and cosine distance."; pgrx::warning!("warning while validating options: {errors}"); } - let opfamily = unsafe { opfamily(index_relation) }; let index = unsafe { PostgresRelation::new(index_relation) }; - let heap = Heap { - heap_relation, - index_relation, - index_info, - opfamily, - scan: std::ptr::null_mut(), - }; let mut reporter = PostgresReporter {}; crate::index::vchordg::algo::build(vector_options, vchordg_options.index, &index); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); - let cache = vchordg_cached::VchordgCached::_0 {}; + let cached = vchordg_cached::VchordgCached::_0 {}.serialize(); if let Some(leader) = unsafe { VchordgLeader::enter( heap_relation, index_relation, (*index_info).ii_Concurrent, - cache, + &cached, ) } { + drop(cached); unsafe { parallel_build( index_relation, @@ -249,18 +242,18 @@ pub unsafe extern "C-unwind" fn ambuild( pgrx::pg_sys::ConditionVariableCancelSleep(); } } else { - let mut indtuples = 0; - reporter.tuples_done(indtuples); - heap.traverse(true, |(ctid, store)| { - for (vector, extra) in store { - let key = ctid_to_key(ctid); - let payload = kv_to_pointer((key, extra)); - crate::index::vchordg::algo::insert(opfamily, &index, payload, vector); - } - indtuples += 1; - reporter.tuples_done(indtuples); - }); - reporter.tuples_total(indtuples); + unsafe { + let indtuples = sequential_build( + index_relation, + heap_relation, + index_info, + &cached, + |indtuples| { + reporter.tuples_done(indtuples); + }, + ); + reporter.tuples_total(indtuples); + } } unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -304,13 +297,8 @@ impl VchordgLeader { heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, isconcurrent: bool, - cache: vchordg_cached::VchordgCached, + cached: &[u8], ) -> Option { - let _cache = cache.serialize(); - #[expect(clippy::drop_non_drop)] - drop(cache); - let cache = _cache; - unsafe fn compute_parallel_workers( heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, @@ -372,7 +360,7 @@ impl VchordgLeader { estimate_keys(&mut (*pcxt).estimator, 1); estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); estimate_keys(&mut (*pcxt).estimator, 1); - estimate_chunk(&mut (*pcxt).estimator, 8 + cache.len()); + estimate_chunk(&mut (*pcxt).estimator, 8 + cached.len()); estimate_keys(&mut (*pcxt).estimator, 1); } @@ -414,9 +402,9 @@ impl VchordgLeader { }; let vchordgcached = unsafe { - let x = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, 8 + cache.len()).cast::(); - (x as *mut u64).write_unaligned(cache.len() as _); - std::ptr::copy(cache.as_ptr(), x.add(8), cache.len()); + let x = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, 8 + cached.len()).cast::(); + (x as *mut u64).write_unaligned(cached.len() as _); + std::ptr::copy(cached.as_ptr(), x.add(8), cached.len()); x }; @@ -583,6 +571,43 @@ unsafe fn parallel_build( } } +unsafe fn sequential_build( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + vchordgcached: &[u8], + mut callback: impl FnMut(u64), +) -> u64 { + use vchordg_cached::VchordgCachedReader; + let cached = VchordgCachedReader::deserialize_ref(vchordgcached); + let index = unsafe { PostgresRelation::new(index_relation) }; + + let opfamily = unsafe { opfamily(index_relation) }; + let heap = Heap { + heap_relation, + index_relation, + index_info, + opfamily, + scan: std::ptr::null_mut(), + }; + + let mut indtuples = 0; + match cached { + VchordgCachedReader::_0(_) => { + heap.traverse(true, |(ctid, store)| { + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + crate::index::vchordg::algo::insert(opfamily, &index, payload, vector); + } + indtuples += 1; + callback(indtuples); + }); + } + } + indtuples +} + #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn ambuildempty(_index_relation: pgrx::pg_sys::Relation) { pgrx::error!("Unlogged indexes are not supported."); diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 8bf65eff..df476cb7 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -34,11 +34,12 @@ use vector::vect::VectOwned; #[repr(u16)] pub enum BuildPhaseCode { Initializing = 0, - InternalBuild = 1, - ExternalBuild = 2, - Build = 3, - Inserting = 4, - Compacting = 5, + DefaultBuild = 1, + InternalBuild = 2, + ExternalBuild = 3, + Build = 4, + Inserting = 5, + Compacting = 6, } pub struct BuildPhase(BuildPhaseCode, u16); @@ -47,6 +48,7 @@ impl BuildPhase { pub const fn new(code: BuildPhaseCode, k: u16) -> Option { match (code, k) { (BuildPhaseCode::Initializing, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::DefaultBuild, 0) => Some(BuildPhase(code, k)), (BuildPhaseCode::InternalBuild, 0..102) => Some(BuildPhase(code, k)), (BuildPhaseCode::ExternalBuild, 0) => Some(BuildPhase(code, k)), (BuildPhaseCode::Build, 0) => Some(BuildPhase(code, k)), @@ -61,6 +63,10 @@ impl BuildPhase { static RAW: [&CStr; 1] = [c"initializing"]; RAW[k as usize] } + BuildPhase(BuildPhaseCode::DefaultBuild, k) => { + static RAW: [&CStr; 1] = [c"initializing index, by default build"]; + RAW[k as usize] + } BuildPhase(BuildPhaseCode::InternalBuild, k) => { static RAWS: [&[&CStr]; 2] = [ &[c"initializing index, by internal build"], @@ -122,6 +128,7 @@ impl BuildPhase { } pub const fn from_value(value: u32) -> Option { const INITIALIZING: u16 = BuildPhaseCode::Initializing as _; + const DEFAULT_BUILD: u16 = BuildPhaseCode::DefaultBuild as _; const INTERNAL_BUILD: u16 = BuildPhaseCode::InternalBuild as _; const EXTERNAL_BUILD: u16 = BuildPhaseCode::ExternalBuild as _; const BUILD: u16 = BuildPhaseCode::Build as _; @@ -130,6 +137,7 @@ impl BuildPhase { let k = value as u16; match (value >> 16) as u16 { INITIALIZING => Self::new(BuildPhaseCode::Initializing, k), + DEFAULT_BUILD => Self::new(BuildPhaseCode::DefaultBuild, k), INTERNAL_BUILD => Self::new(BuildPhaseCode::InternalBuild, k), EXTERNAL_BUILD => Self::new(BuildPhaseCode::ExternalBuild, k), BUILD => Self::new(BuildPhaseCode::Build, k), @@ -276,12 +284,11 @@ pub unsafe extern "C-unwind" fn ambuild( scan: std::ptr::null_mut(), }; let mut reporter = PostgresReporter {}; - let structures = match vchordrq_options.build.source.clone() { - VchordrqBuildSourceOptions::External(external_build) => { - reporter.phase(BuildPhase::from_code(BuildPhaseCode::ExternalBuild)); - let reltuples = unsafe { (*(*index_relation).rd_rel).reltuples }; - reporter.tuples_total(reltuples as u64); - make_external_build(vector_options.clone(), opfamily, external_build.clone()) + reporter.tuples_total(unsafe { (*(*index_relation).rd_rel).reltuples as u64 }); + let mut structures = match vchordrq_options.build.source.clone() { + VchordrqBuildSourceOptions::Default(default_build) => { + reporter.phase(BuildPhase::from_code(BuildPhaseCode::DefaultBuild)); + make_default_build(vector_options, default_build) } VchordrqBuildSourceOptions::Internal(internal_build) => { reporter.phase(BuildPhase::from_code(BuildPhaseCode::InternalBuild)); @@ -321,18 +328,22 @@ pub unsafe extern "C-unwind" fn ambuild( samples }; reporter.tuples_total(tuples_total); - make_internal_build( - vector_options.clone(), - internal_build.clone(), - samples, - &mut reporter, - ) + make_internal_build(vector_options, internal_build, samples, &mut reporter) + } + VchordrqBuildSourceOptions::External(external_build) => { + reporter.phase(BuildPhase::from_code(BuildPhaseCode::ExternalBuild)); + make_external_build(vector_options, opfamily, external_build) } }; + for structure in structures.iter_mut() { + for centroid in structure.centroids.iter_mut() { + *centroid = rabitq::rotate::rotate(centroid); + } + } reporter.phase(BuildPhase::from_code(BuildPhaseCode::Build)); crate::index::vchordrq::algo::build(vector_options, vchordrq_options.index, &index, structures); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); - let cache = if vchordrq_options.build.pin { + let cached = if vchordrq_options.build.pin { let mut trace = vchordrq::cache(&index); trace.sort(); trace.dedup(); @@ -349,15 +360,17 @@ pub unsafe extern "C-unwind" fn ambuild( } } else { vchordrq_cached::VchordrqCached::_0 {} - }; + } + .serialize(); if let Some(leader) = unsafe { VchordrqLeader::enter( heap_relation, index_relation, (*index_info).ii_Concurrent, - cache, + &cached, ) } { + drop(cached); unsafe { parallel_build( index_relation, @@ -389,18 +402,18 @@ pub unsafe extern "C-unwind" fn ambuild( pgrx::pg_sys::ConditionVariableCancelSleep(); } } else { - let mut indtuples = 0; - reporter.tuples_done(indtuples); - heap.traverse(true, |(ctid, store)| { - for (vector, extra) in store { - let key = ctid_to_key(ctid); - let payload = kv_to_pointer((key, extra)); - crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); - } - indtuples += 1; - reporter.tuples_done(indtuples); - }); - reporter.tuples_total(indtuples); + unsafe { + let indtuples = sequential_build( + index_relation, + heap_relation, + index_info, + &cached, + |indtuples| { + reporter.tuples_done(indtuples); + }, + ); + reporter.tuples_total(indtuples); + } } let check = || { pgrx::check_for_interrupts!(); @@ -581,12 +594,8 @@ impl VchordrqLeader { heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, isconcurrent: bool, - cache: vchordrq_cached::VchordrqCached, + vchordrq_cached: &[u8], ) -> Option { - let _cache = cache.serialize(); - drop(cache); - let cache = _cache; - unsafe fn compute_parallel_workers( heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, @@ -648,7 +657,7 @@ impl VchordrqLeader { estimate_keys(&mut (*pcxt).estimator, 1); estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); estimate_keys(&mut (*pcxt).estimator, 1); - estimate_chunk(&mut (*pcxt).estimator, 8 + cache.len()); + estimate_chunk(&mut (*pcxt).estimator, 8 + vchordrq_cached.len()); estimate_keys(&mut (*pcxt).estimator, 1); } @@ -690,9 +699,10 @@ impl VchordrqLeader { }; let vchordrqcached = unsafe { - let x = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, 8 + cache.len()).cast::(); - (x as *mut u64).write_unaligned(cache.len() as _); - std::ptr::copy(cache.as_ptr(), x.add(8), cache.len()); + let x = + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, 8 + vchordrq_cached.len()).cast::(); + (x as *mut u64).write_unaligned(vchordrq_cached.len() as _); + std::ptr::copy(vchordrq_cached.as_ptr(), x.add(8), vchordrq_cached.len()); x }; @@ -882,6 +892,58 @@ unsafe fn parallel_build( } } +unsafe fn sequential_build( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + vchordrqcached: &[u8], + mut callback: impl FnMut(u64), +) -> u64 { + use vchordrq_cached::VchordrqCachedReader; + let cached = VchordrqCachedReader::deserialize_ref(vchordrqcached); + let index = unsafe { PostgresRelation::new(index_relation) }; + + let opfamily = unsafe { opfamily(index_relation) }; + let heap = Heap { + heap_relation, + index_relation, + index_info, + opfamily, + scan: std::ptr::null_mut(), + }; + + let mut indtuples = 0; + match cached { + VchordrqCachedReader::_0(_) => { + heap.traverse(true, |(ctid, store)| { + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); + } + indtuples += 1; + callback(indtuples); + }); + } + VchordrqCachedReader::_1(cached) => { + let index = CachingRelation { + cache: cached, + relation: index, + }; + heap.traverse(true, |(ctid, store)| { + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); + } + indtuples += 1; + callback(indtuples); + }); + } + } + indtuples +} + #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn ambuildempty(_index_relation: pgrx::pg_sys::Relation) { pgrx::error!("Unlogged indexes are not supported."); @@ -945,16 +1007,23 @@ unsafe fn options( (vector, rabitq) } +fn make_default_build( + vector_options: VectorOptions, + _default_build: VchordrqDefaultBuildOptions, +) -> Vec>> { + vec![Structure::> { + centroids: vec![vec![0.0f32; vector_options.dims as usize]], + children: vec![vec![]], + }] +} + fn make_internal_build( vector_options: VectorOptions, internal_build: VchordrqInternalBuildOptions, - mut samples: Vec>, + samples: Vec>, reporter: &mut PostgresReporter, ) -> Vec>> { use std::iter::once; - k_means::preprocess(internal_build.build_threads as _, &mut samples, |sample| { - *sample = rabitq::rotate::rotate(sample) - }); let mut result = Vec::>>::new(); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { let input = if let Some(structure) = result.last() { @@ -1014,7 +1083,18 @@ fn make_internal_build( if let Some(structure) = result.last() { let mut children = vec![Vec::new(); centroids.len()]; for i in 0..structure.len() as u32 { - let target = k_means::k_means_lookup(&structure.centroids[i as usize], ¢roids); + pub fn k_means_lookup(vector: &[f32], centroids: &[Vec]) -> usize { + assert_ne!(centroids.len(), 0); + let mut result = (f32::INFINITY, 0); + for i in 0..centroids.len() { + let dis = f32::reduce_sum_of_d2(vector, ¢roids[i]); + if dis <= result.0 { + result = (dis, i); + } + } + result.1 + } + let target = k_means_lookup(&structure.centroids[i as usize], ¢roids); children[target].push(i); } let (centroids, children) = std::iter::zip(centroids, children) @@ -1078,7 +1158,7 @@ fn make_external_build( if vector_options.dims != vector.as_borrowed().dims() { pgrx::error!("external build: incorrect dimension, id = {id}"); } - vectors.insert(id, rabitq::rotate::rotate(vector.as_borrowed().slice())); + vectors.insert(id, vector.as_borrowed().slice().to_vec()); } }); if parents.len() >= 2 && parents.values().all(|x| x.is_none()) { diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index bb24d316..03bb274c 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -228,7 +228,7 @@ pub unsafe extern "C-unwind" fn amcostestimate( let denumerator = r.clone(); let scale = r.skip(1).chain(std::iter::once(tuples)); for (scale, (numerator, denumerator)) in scale.zip(numerator.zip(denumerator)) { - count += (scale as f64) * ((numerator as f64) / (denumerator as f64)); + count += (scale as f64) * 1.0f64.min((numerator as f64) / (denumerator as f64)); } count }; diff --git a/src/index/vchordrq/types.rs b/src/index/vchordrq/types.rs index 1a1b21a0..44fb0ec9 100644 --- a/src/index/vchordrq/types.rs +++ b/src/index/vchordrq/types.rs @@ -16,6 +16,17 @@ use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError, ValidationErrors}; use vchordrq::types::VchordrqIndexOptions; +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqDefaultBuildOptions {} + +#[allow(clippy::derivable_impls)] +impl Default for VchordrqDefaultBuildOptions { + fn default() -> Self { + Self {} + } +} + #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] pub struct VchordrqInternalBuildOptions { @@ -84,13 +95,14 @@ pub struct VchordrqExternalBuildOptions { #[serde(deny_unknown_fields)] #[serde(rename_all = "snake_case")] pub enum VchordrqBuildSourceOptions { + Default(VchordrqDefaultBuildOptions), Internal(VchordrqInternalBuildOptions), External(VchordrqExternalBuildOptions), } impl Default for VchordrqBuildSourceOptions { fn default() -> Self { - Self::Internal(Default::default()) + Self::Default(Default::default()) } } @@ -98,6 +110,7 @@ impl Validate for VchordrqBuildSourceOptions { fn validate(&self) -> Result<(), ValidationErrors> { use VchordrqBuildSourceOptions::*; match self { + Default(default_build) => default_build.validate(), Internal(internal_build) => internal_build.validate(), External(external_build) => external_build.validate(), } diff --git a/tests/vchordg/sequential_build.slt b/tests/vchordg/sequential_build.slt new file mode 100644 index 00000000..826de42f --- /dev/null +++ b/tests/vchordg/sequential_build.slt @@ -0,0 +1,17 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SET max_parallel_workers = 0; + +statement ok +SET max_parallel_maintenance_workers = 0; + +statement ok +CREATE INDEX ON t USING vchordg (val vector_ip_ops); + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/sequential_build.slt b/tests/vchordrq/sequential_build.slt new file mode 100644 index 00000000..b214b64f --- /dev/null +++ b/tests/vchordrq/sequential_build.slt @@ -0,0 +1,24 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SET max_parallel_workers = 0; + +statement ok +SET max_parallel_maintenance_workers = 0; + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_ip_ops) +WITH (options = $$ +residual_quantization = false +build.pin = true +build.internal.lists = [32] +build.internal.build_threads = 2 +build.internal.spherical_centroids = true +$$); + +statement ok +DROP TABLE t; From 6bc0fdb04d900abfd5490ad92c14491cd9f87317 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 22 Sep 2025 13:22:56 +0800 Subject: [PATCH 224/324] ci: add macos x86_64 (#348) job: -psql job: +psql_macos Signed-off-by: usamoi --- .github/workflows/check.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 0ed4a000..ae8cf035 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -264,9 +264,9 @@ jobs: strategy: matrix: version: ["13", "14", "15", "16", "17"] - arch: ["aarch64"] + arch: ["aarch64", "x86_64"] - runs-on: "macos-15" + runs-on: ${{ matrix.arch == 'aarch64' && 'macos-15' || 'macos-15-intel' }} env: SCCACHE_GHA_ENABLED: "true" From 22e7c4557ad626d12b2fcb3ce307dbdf4f676632 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 25 Sep 2025 15:23:45 +0800 Subject: [PATCH 225/324] perf: reduce lock contention (#347) ```sql create index on laion1m768_train using vchordrq (embedding vector_ip_ops) with (options = $$ build.pin = 2 build.internal.lists = [256, 4096] build.internal.build_threads = 20 build.internal.spherical_centroids = true degree_of_parallelism = 40 $$); ``` `degree_of_parallelism` is 32 by default. It reduces lock contention during insertion, especially, in a build where `lists.len() != 1`. `build.pin` is `-1` by default. `build.pin = 2` caches almost everything in memory, which speeds build but spends more memory. `build.pin = false` and `build.pin = true` still works, and they're mapped to `build.pin = -1` and `build.pin = 1`. Signed-off-by: usamoi --- Cargo.lock | 230 ++++++++++++++++++------------ crates/algo/src/lib.rs | 8 -- crates/vchordg/src/bulkdelete.rs | 4 +- crates/vchordg/src/types.rs | 10 +- crates/vchordrq/src/build.rs | 16 ++- crates/vchordrq/src/bulkdelete.rs | 60 ++++---- crates/vchordrq/src/cache.rs | 54 ++++--- crates/vchordrq/src/centroids.rs | 44 ++++++ crates/vchordrq/src/insert.rs | 19 ++- crates/vchordrq/src/lib.rs | 6 + crates/vchordrq/src/prewarm.rs | 10 +- crates/vchordrq/src/rerank.rs | 2 +- crates/vchordrq/src/search.rs | 10 +- crates/vchordrq/src/tuples.rs | 195 +++++++++++++++++++++++-- crates/vchordrq/src/types.rs | 17 ++- crates/vchordrq/src/vectors.rs | 44 ++---- src/datatype/typmod.rs | 3 +- src/index/storage.rs | 48 ++++--- src/index/vchordg/algo.rs | 2 +- src/index/vchordg/am/am_build.rs | 27 ++-- src/index/vchordrq/algo.rs | 22 ++- src/index/vchordrq/am/am_build.rs | 114 ++++++++++++--- src/index/vchordrq/am/mod.rs | 22 ++- src/index/vchordrq/types.rs | 24 +++- 24 files changed, 692 insertions(+), 299 deletions(-) create mode 100644 crates/vchordrq/src/centroids.rs diff --git a/Cargo.lock b/Cargo.lock index d178abb2..dff15fa8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -97,9 +97,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.99" +version = "1.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0674a1ddeecb70197781e945de4b3b8ffb61fa939a5597bcf48503737663100" +checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" [[package]] name = "bindgen" @@ -122,9 +122,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.9.3" +version = "2.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34efbcccd345379ca2868b2b2c9d3782e9cc58ba87bc7d79d5b53d9c9ae6f25d" +checksum = "2261d10cca569e4643e526d8dc2e62e433cc8aba21ab764233731f8d369bf394" [[package]] name = "bitvec" @@ -151,15 +151,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "374b7c592d9c00c1f4972ea58390ac6b18cbb6ab79011f3bdc90a0b82ca06b77" dependencies = [ "serde", - "toml 0.9.5", + "toml 0.9.7", ] [[package]] name = "cc" -version = "1.2.34" +version = "1.2.38" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42bc4aea80032b7bf409b0bc7ccad88853858911b7713a8062fdc0623867bedc" +checksum = "80f41ae168f955c12fb8960b057d70d0ca153fb83182b57d86380443527be7e9" dependencies = [ + "find-msvc-tools", "shlex", ] @@ -201,9 +202,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.46" +version = "4.5.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c5e4fcf9c21d2e544ca1ee9d8552de13019a42aa7dbf32747fa7aaf1df76e57" +checksum = "e2134bb3ea021b78629caa971416385309e0131b351b25e01dc16fb54e1b5fae" dependencies = [ "clap_builder", "clap_derive", @@ -211,9 +212,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.46" +version = "4.5.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fecb53a0e6fcfb055f686001bc2e2592fa527efaf38dbe81a6a9563562e57d41" +checksum = "c2ba64afa3c0a6df7fa517765e31314e983f51dda798ffba27b988194fb65dc9" dependencies = [ "anstream", "anstyle", @@ -223,9 +224,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.5.45" +version = "4.5.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14cb31bb0a7d536caef2639baa7fad459e15c3144efefa6dbd1c84562c4739f6" +checksum = "bbfd7eae0b0f1a6e63d4b13c9c478de77c2eb546fba158ad50b4203dc24b9f9c" dependencies = [ "heck", "proc-macro2", @@ -331,9 +332,9 @@ dependencies = [ [[package]] name = "dary_heap" -version = "0.3.7" +version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04d2cd9c18b9f454ed67da600630b021a8a80bf33f8c95896ab33aaf1c26b728" +checksum = "06d2e3287df1c007e74221c49ca10a95d557349e54b3a75dc2fb14712c751f04" [[package]] name = "displaydoc" @@ -402,12 +403,12 @@ checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "errno" -version = "0.3.13" +version = "0.3.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.0", ] [[package]] @@ -432,6 +433,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "find-msvc-tools" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ced73b1dacfc750a6db6c0a0c3a3853c8b41997e2e2c563dc90804ae6867959" + [[package]] name = "fixedbitset" version = "0.5.7" @@ -511,13 +518,19 @@ dependencies = [ "foldhash", ] +[[package]] +name = "hashbrown" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" + [[package]] name = "hashlink" version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7382cf6263419f2d8df38c55d7da83da5c18aef87fc7a7fc1fb1e344edfe14c1" dependencies = [ - "hashbrown", + "hashbrown 0.15.5", ] [[package]] @@ -662,12 +675,12 @@ checksum = "964de6e86d545b246d84badc0fef527924ace5134f30641c203ef52ba83f58d5" [[package]] name = "indexmap" -version = "2.11.0" +version = "2.11.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2481980430f9f78649238835720ddccc57e52df14ffce1c6f37391d61b563e9" +checksum = "4b0f83760fb341a774ed326568e19f5a863af4a952def8c39f9ab92fd95b88e5" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.16.0", ] [[package]] @@ -710,9 +723,9 @@ checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" [[package]] name = "js-sys" -version = "0.3.77" +version = "0.3.80" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1cfaf33c695fc6e08064efbc1f72ec937429614f25eef83af942d0e227c3a28f" +checksum = "852f13bec5eba4ba9afbeb93fd7c13fe56147f055939ae21c43a29a0ecb2702e" dependencies = [ "once_cell", "wasm-bindgen", @@ -730,18 +743,18 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.175" +version = "0.2.176" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a82ae493e598baaea5209805c49bbf2ea7de956d50d7da0da1164f9c6d28543" +checksum = "58f929b4d672ea937a23a1ab494143d968337a5f47e56d0815df1e0890ddf174" [[package]] name = "libloading" -version = "0.8.8" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" dependencies = [ "cfg-if", - "windows-targets 0.53.3", + "windows-link 0.2.0", ] [[package]] @@ -767,9 +780,9 @@ dependencies = [ [[package]] name = "linux-raw-sys" -version = "0.9.4" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd945864f07fe9f5371a27ad7b52a172b4b499999f1d97574c9fa68373937e12" +checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039" [[package]] name = "litemap" @@ -779,9 +792,9 @@ checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956" [[package]] name = "log" -version = "0.4.27" +version = "0.4.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" +checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432" [[package]] name = "make" @@ -891,7 +904,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "54acf3a685220b533e437e264e4d932cfbdc4cc7ec0cd232ed73c08d03b8a7ca" dependencies = [ "fixedbitset", - "hashbrown", + "hashbrown 0.15.5", "indexmap", "serde", ] @@ -1039,9 +1052,9 @@ checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" [[package]] name = "potential_utf" -version = "0.1.2" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5a7c30837279ca13e7c867e9e40053bc68740f988cb07f7ca6df43cc734b585" +checksum = "84df19adbe5b5a0782edcab45899906947ab039ccf4573713735ee7de1e6b08a" dependencies = [ "zerovec", ] @@ -1217,15 +1230,15 @@ checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" [[package]] name = "rustix" -version = "1.0.8" +version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11181fbabf243db407ef8df94a6ce0b2f9a733bd8be4ad02b4eda9602296cac8" +checksum = "cd15f8a2c5551a84d56efdc1cd049089e409ac19a3072d5037a17fd70719ff3e" dependencies = [ "bitflags", "errno", "libc", "linux-raw-sys", - "windows-sys 0.60.2", + "windows-sys 0.61.0", ] [[package]] @@ -1263,10 +1276,11 @@ checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" -version = "1.0.219" +version = "1.0.226" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6" +checksum = "0dca6411025b24b60bfa7ec1fe1f8e710ac09782dca409ee8237ba74b51295fd" dependencies = [ + "serde_core", "serde_derive", ] @@ -1280,11 +1294,20 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_core" +version = "1.0.226" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba2ba63999edb9dac981fb34b3e5c0d111a69b0924e253ed29d83f7c99e966a4" +dependencies = [ + "serde_derive", +] + [[package]] name = "serde_derive" -version = "1.0.219" +version = "1.0.226" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00" +checksum = "8db53ae22f34573731bafa1db20f04027b2d25e02d8205921b569171699cdb33" dependencies = [ "proc-macro2", "quote", @@ -1293,14 +1316,15 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.143" +version = "1.0.145" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d401abef1d108fbd9cbaebc3e46611f4b1021f714a0597a71f41ee463f5f4a5a" +checksum = "402a6f66d8c709116cf22f558eab210f5a50187f702eb4d7e5ef38d9a7f1c79c" dependencies = [ "itoa", "memchr", "ryu", "serde", + "serde_core", ] [[package]] @@ -1314,11 +1338,11 @@ dependencies = [ [[package]] name = "serde_spanned" -version = "1.0.0" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40734c41988f7306bb04f0ecf60ec0f3f1caa34290e4e8ea471dcd3346483b83" +checksum = "5417783452c2be558477e104686f7de5dae53dba813c28435e0e70f82d9b04ee" dependencies = [ - "serde", + "serde_core", ] [[package]] @@ -1474,14 +1498,14 @@ dependencies = [ [[package]] name = "toml" -version = "0.9.5" +version = "0.9.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75129e1dc5000bfbaa9fee9d1b21f974f9fbad9daec557a521ee6e080825f6e8" +checksum = "00e5e5d9bf2475ac9d4f0d9edab68cc573dc2fd644b0dba36b0c30a92dd9eaa0" dependencies = [ "indexmap", - "serde", - "serde_spanned 1.0.0", - "toml_datetime 0.7.0", + "serde_core", + "serde_spanned 1.0.2", + "toml_datetime 0.7.2", "toml_parser", "toml_writer", "winnow", @@ -1498,11 +1522,11 @@ dependencies = [ [[package]] name = "toml_datetime" -version = "0.7.0" +version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bade1c3e902f58d73d3f294cd7f20391c1cb2fbcb643b73566bc773971df91e3" +checksum = "32f1085dec27c2b6632b04c80b3bb1b4300d6495d1e129693bdda7d91e72eec1" dependencies = [ - "serde", + "serde_core", ] [[package]] @@ -1521,9 +1545,9 @@ dependencies = [ [[package]] name = "toml_parser" -version = "1.0.2" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b551886f449aa90d4fe2bdaa9f4a2577ad2dde302c61ecf262d80b116db95c10" +checksum = "4cf893c33be71572e0e9aa6dd15e6677937abd686b066eac3f8cd3531688a627" dependencies = [ "winnow", ] @@ -1536,9 +1560,9 @@ checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" [[package]] name = "toml_writer" -version = "1.0.2" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fcc842091f2def52017664b53082ecbbeb5c7731092bad69d2c63050401dfd64" +checksum = "d163a63c116ce562a22cda521fcc4d79152e7aba014456fb5eb442f6d6a10109" [[package]] name = "unescape" @@ -1548,9 +1572,9 @@ checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] name = "unicode-ident" -version = "1.0.18" +version = "1.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" +checksum = "f63a545481291138910575129486daeaf8ac54aee4387fe7906919f7830c7d9d" [[package]] name = "unicode-segmentation" @@ -1590,9 +1614,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.18.0" +version = "1.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f33196643e165781c20a5ead5582283a7dacbb87855d867fbc2df3f81eddc1be" +checksum = "2f87b8aa10b915a06587d0dec516c282ff295b475d94abf425d62b57710070a2" dependencies = [ "getrandom", "js-sys", @@ -1650,7 +1674,7 @@ dependencies = [ "serde", "simd", "small_iter", - "toml 0.9.5", + "toml 0.9.7", "validator", "vchordg", "vchordrq", @@ -1718,30 +1742,40 @@ dependencies = [ [[package]] name = "wasi" -version = "0.14.2+wasi-0.2.4" +version = "0.14.7+wasi-0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" dependencies = [ - "wit-bindgen-rt", + "wasip2", +] + +[[package]] +name = "wasip2" +version = "1.0.1+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" +dependencies = [ + "wit-bindgen", ] [[package]] name = "wasm-bindgen" -version = "0.2.100" +version = "0.2.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1edc8929d7499fc4e8f0be2262a241556cfc54a0bea223790e71446f2aab1ef5" +checksum = "ab10a69fbd0a177f5f649ad4d8d3305499c42bab9aef2f7ff592d0ec8f833819" dependencies = [ "cfg-if", "once_cell", "rustversion", "wasm-bindgen-macro", + "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-backend" -version = "0.2.100" +version = "0.2.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f0a0651a5c2bc21487bde11ee802ccaf4c51935d0d3d42a6101f98161700bc6" +checksum = "0bb702423545a6007bbc368fde243ba47ca275e549c8a28617f56f6ba53b1d1c" dependencies = [ "bumpalo", "log", @@ -1753,9 +1787,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.100" +version = "0.2.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fe63fc6d09ed3792bd0897b314f53de8e16568c2b3f7982f468c0bf9bd0b407" +checksum = "fc65f4f411d91494355917b605e1480033152658d71f722a90647f56a70c88a0" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1763,9 +1797,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.100" +version = "0.2.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" +checksum = "ffc003a991398a8ee604a401e194b6b3a39677b3173d6e74495eb51b82e99a32" dependencies = [ "proc-macro2", "quote", @@ -1776,9 +1810,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.100" +version = "0.2.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a05d73b933a847d6cccdda8f838a22ff101ad9bf93e33684f39c1f5f0eece3d" +checksum = "293c37f4efa430ca14db3721dfbe48d8c33308096bd44d80ebaa775ab71ba1cf" dependencies = [ "unicode-ident", ] @@ -1821,11 +1855,11 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] name = "winapi-util" -version = "0.1.10" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0978bf7171b3d90bac376700cb56d606feb40f251a475a5d6634613564460b22" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.0", ] [[package]] @@ -1840,6 +1874,12 @@ version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" +[[package]] +name = "windows-link" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45e46c0661abb7180e7b9c281db115305d49ca1709ab8242adf09666d2173c65" + [[package]] name = "windows-sys" version = "0.59.0" @@ -1858,6 +1898,15 @@ dependencies = [ "windows-targets 0.53.3", ] +[[package]] +name = "windows-sys" +version = "0.61.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e201184e40b2ede64bc2ea34968b28e33622acdbbf37104f0e4a33f7abe657aa" +dependencies = [ + "windows-link 0.2.0", +] + [[package]] name = "windows-targets" version = "0.52.6" @@ -1880,7 +1929,7 @@ version = "0.53.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d5fe6031c4041849d7c496a8ded650796e7b6ecc19df1a431c1a363342e5dc91" dependencies = [ - "windows-link", + "windows-link 0.1.3", "windows_aarch64_gnullvm 0.53.0", "windows_aarch64_msvc 0.53.0", "windows_i686_gnu 0.53.0", @@ -2003,13 +2052,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" [[package]] -name = "wit-bindgen-rt" -version = "0.39.0" +name = "wit-bindgen" +version = "0.46.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" -dependencies = [ - "bitflags", -] +checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" [[package]] name = "writeable" @@ -2052,18 +2098,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.26" +version = "0.8.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1039dd0d3c310cf05de012d8a39ff557cb0d23087fd44cad61df08fc31907a2f" +checksum = "0894878a5fa3edfd6da3f88c4805f4c8558e2b996227a3d864f47fe11e38282c" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.26" +version = "0.8.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ecf5b4cc5364572d7f4c329661bcc82724222973f2cab6f050a4e5c22f75181" +checksum = "88d2b8d9c68ad2b9e4340d7832716a4d21a22a1154777ad56ea55c51a9cf3831" dependencies = [ "proc-macro2", "quote", diff --git a/crates/algo/src/lib.rs b/crates/algo/src/lib.rs index baedf91c..995afd88 100644 --- a/crates/algo/src/lib.rs +++ b/crates/algo/src/lib.rs @@ -92,14 +92,6 @@ pub trait RelationWrite: RelationWriteTypes { fn search(&self, freespace: usize) -> Option>; } -pub trait RelationLength: Relation { - fn len(&self) -> u32; - #[inline] - fn is_empty(&self) -> bool { - self.len() == 0 - } -} - pub trait RelationPrefetch: Relation { fn prefetch(&self, id: u32); } diff --git a/crates/vchordg/src/bulkdelete.rs b/crates/vchordg/src/bulkdelete.rs index ab3f921f..a373559a 100644 --- a/crates/vchordg/src/bulkdelete.rs +++ b/crates/vchordg/src/bulkdelete.rs @@ -15,10 +15,10 @@ use crate::Opaque; use crate::operator::Operator; use crate::tuples::{MetaTuple, VertexTuple, WithReader, WithWriter}; -use algo::{Page, RelationLength, RelationRead, RelationWrite}; +use algo::{Page, RelationRead, RelationWrite}; use std::num::NonZero; -pub fn bulkdelete( +pub fn bulkdelete( index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index cf5d6f55..d262587b 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -92,14 +92,14 @@ pub enum BorrowedVector<'a> { } #[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum DistanceKind { L2S, Dot, } #[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum VectorKind { Vecf32, Vecf16, @@ -114,16 +114,12 @@ impl VectorKind { } } -#[derive(Debug, Clone, Serialize, Deserialize, Validate)] -#[serde(deny_unknown_fields)] +#[derive(Debug, Clone, Validate)] #[validate(schema(function = "Self::validate_self"))] pub struct VectorOptions { #[validate(range(min = 1))] - #[serde(rename = "dimensions")] pub dims: u32, - #[serde(rename = "vector")] pub v: VectorKind, - #[serde(rename = "distance")] pub d: DistanceKind, } diff --git a/crates/vchordrq/src/build.rs b/crates/vchordrq/src/build.rs index abde0caa..5bcbe177 100644 --- a/crates/vchordrq/src/build.rs +++ b/crates/vchordrq/src/build.rs @@ -35,7 +35,10 @@ pub fn build( assert_eq!(meta.first(), 0); let mut freepages = TapeWriter::<_, FreepagesTuple>::create(index, false); freepages.push(FreepagesTuple {}); - let mut vectors = TapeWriter::<_, VectorTuple>::create(index, true); + let mut centroids = TapeWriter::<_, CentroidTuple>::create(index, false); + let vectors = (0..vchordrq_options.degree_of_parallelism) + .map(|_| TapeWriter::<_, VectorTuple>::create(index, true).first()) + .collect::>(); let mut pointer_of_centroids = Vec::, u16)>>::new(); for i in 0..structures.len() { let mut level = Vec::new(); @@ -45,14 +48,12 @@ pub fn build( let mut chain = Ok(metadata); let mut prefetch = Vec::new(); for i in (0..slices.len()).rev() { - let (id, head) = vectors.push(match chain { - Ok(metadata) => VectorTuple::_0 { - payload: None, + let (id, head) = centroids.push(match chain { + Ok(metadata) => CentroidTuple::_0 { elements: slices[i].to_vec(), metadata, }, - Err(head) => VectorTuple::_1 { - payload: None, + Err(head) => CentroidTuple::_1 { elements: slices[i].to_vec(), head, }, @@ -118,7 +119,8 @@ pub fn build( height_of_root: structures.len() as u32, is_residual, rerank_in_heap: vchordrq_options.rerank_in_table, - vectors_first: vectors.first(), + centroids_first: centroids.first(), + vectors_first: vectors, centroid_prefetch: pointer_of_centroids .last() .expect("internal error: empty structure")[0] diff --git a/crates/vchordrq/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs index 3f2556fd..f93751fb 100644 --- a/crates/vchordrq/src/bulkdelete.rs +++ b/crates/vchordrq/src/bulkdelete.rs @@ -142,41 +142,43 @@ pub fn bulkdelete_vectors( let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); - let vectors_first = meta_tuple.vectors_first(); + let vectors_first = meta_tuple.vectors_first().to_vec(); drop(meta_guard); - let mut current = vectors_first; - while current != u32::MAX { - check(); - let read = index.read(current); - let flag = 'flag: { - for i in 1..=read.len() { - if let Some(bytes) = read.get(i) { - let tuple = VectorTuple::::deserialize_ref(bytes); - let p = tuple.payload(); - if Some(true) == p.map(&callback) { - break 'flag true; + for vectors_first in vectors_first { + let mut current = vectors_first; + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + if let Some(bytes) = read.get(i) { + let tuple = VectorTuple::::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + break 'flag true; + } } } + false + }; + if flag { + drop(read); + let mut write = index.write(current, true); + for i in 1..=write.len() { + if let Some(bytes) = write.get(i) { + let tuple = VectorTuple::::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + write.free(i); + } + }; + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; } - false - }; - if flag { - drop(read); - let mut write = index.write(current, true); - for i in 1..=write.len() { - if let Some(bytes) = write.get(i) { - let tuple = VectorTuple::::deserialize_ref(bytes); - let p = tuple.payload(); - if Some(true) == p.map(&callback) { - write.free(i); - } - }; - } - current = write.get_opaque().next; - } else { - current = read.get_opaque().next; } } } diff --git a/crates/vchordrq/src/cache.rs b/crates/vchordrq/src/cache.rs index eecc7275..307446d4 100644 --- a/crates/vchordrq/src/cache.rs +++ b/crates/vchordrq/src/cache.rs @@ -19,41 +19,59 @@ use crate::{Opaque, Page, PageGuard, tape}; use algo::RelationRead; use algo::accessor::FunctionalAccessor; -pub fn cache(index: &R) -> Vec +pub fn cache(index: &R, level: i32) -> Vec where R::Page: Page, { - let mut trace = vec![0_u32]; let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let height_of_root = meta_tuple.height_of_root(); + let centroids_first = meta_tuple.centroids_first(); type State = Vec; let mut state: State = vec![meta_tuple.first()]; drop(meta_guard); - let mut step = |state: State| { - let mut results = Vec::new(); - for first in state { - tape::read_h1_tape::( - by_next(index, first).inspect(|guard| trace.push(guard.id())), - || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), - |(), _, _, first, _| { - results.push(first); - }, - ); + let mut trace = Vec::new(); + + if level >= 0 { + // meta tuple + trace.push(0); + } + + if level >= 1 { + let mut step = |state: State| { + let mut results = Vec::new(); + for first in state { + tape::read_h1_tape::( + by_next(index, first).inspect(|guard| trace.push(guard.id())), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + |(), _, _, first, _| { + results.push(first); + }, + ); + } + results + }; + + // h1 tuples + for _ in (1..height_of_root).rev() { + state = step(state); } - results - }; - for _ in (1..height_of_root).rev() { - state = step(state); + // jump tuples + for first in state { + trace.push(first); + } } - for first in state { - trace.push(first); + if level >= 2 { + // centroid tuples + for guard in by_next(index, centroids_first) { + trace.push(guard.id()); + } } trace diff --git a/crates/vchordrq/src/centroids.rs b/crates/vchordrq/src/centroids.rs new file mode 100644 index 00000000..a7490e1c --- /dev/null +++ b/crates/vchordrq/src/centroids.rs @@ -0,0 +1,44 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::Page; +use crate::operator::*; +use crate::tuples::*; +use algo::RelationRead; +use algo::accessor::Accessor1; + +pub fn read< + 'a, + R: RelationRead + 'a, + O: Operator, + A: Accessor1<::Element, ::Metadata>, +>( + mut prefetch: impl Iterator>, + head: u16, + accessor: A, +) -> A::Output { + let mut cursor = Err(head); + let mut result = accessor; + while let Err(head) = cursor { + let guard = prefetch.next().expect("data corruption"); + let bytes = guard.get(head).expect("data corruption"); + let tuple = CentroidTuple::::deserialize_ref(bytes); + result.push(tuple.elements()); + cursor = tuple.metadata_or_head(); + } + if prefetch.next().is_some() { + panic!("data corruption"); + } + result.finish(cursor.expect("data corruption")) +} diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index 77d1c448..c3aca5c7 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -17,7 +17,7 @@ use crate::operator::*; use crate::tape::by_next; use crate::tuples::*; use crate::vectors::{self}; -use crate::{Opaque, Page, tape}; +use crate::{Chooser, Opaque, Page, centroids, tape}; use algo::accessor::{FunctionalAccessor, LAccess}; use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily}; use algo::{BorrowedIter, Bump, RelationRead, RelationWrite}; @@ -34,6 +34,8 @@ pub fn insert_vector( index: &R, payload: NonZero, vector: ::Borrowed<'_>, + chooser: &mut impl Chooser, + skip_search: bool, ) -> (Vec, u16) where R::Page: Page, @@ -44,12 +46,17 @@ where let dims = meta_tuple.dims(); let rerank_in_heap = meta_tuple.rerank_in_heap(); assert_eq!(dims, vector.dims(), "unmatched dimensions"); - let vectors_first = meta_tuple.vectors_first(); + let vectors_first = { + let l = meta_tuple.vectors_first(); + let n = NonZero::new(l.len()).expect("data corruption"); + let i = chooser.choose(n); + l[i] + }; drop(meta_guard); if !rerank_in_heap { - vectors::append::(index, vectors_first, vector, payload) + vectors::append::(index, vectors_first, vector, payload, skip_search) } else { (Vec::new(), 0) } @@ -81,7 +88,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let prefetch = BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); let head = meta_tuple.centroid_head(); - let distance = vectors::read_for_h1_tuple::( + let distance = centroids::read::( prefetch.map(|id| index.read(id)), head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), @@ -127,7 +134,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( while let Some(((Reverse(_), AlwaysEqual(&mut (first, norm, head, ..))), prefetch)) = heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { - let distance = vectors::read_for_h1_tuple::( + let distance = centroids::read::( prefetch, head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), @@ -152,7 +159,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let (code, delta) = O::build( vector, is_residual.then(|| { - vectors::read_for_h1_tuple::( + centroids::read::( jump_tuple .centroid_prefetch() .iter() diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index a842e0f5..771efcf1 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -15,6 +15,7 @@ mod build; mod bulkdelete; mod cache; +mod centroids; mod closure_lifetime_binder; mod cost; mod fast_heap; @@ -44,6 +45,7 @@ pub use maintain::maintain; pub use prewarm::prewarm; pub use rerank::{how, rerank_heap, rerank_index}; pub use search::{default_search, maxsim_search}; +use std::num::NonZero; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -65,3 +67,7 @@ pub(crate) struct Branch { pub norm: f32, pub extra: T, } + +pub trait Chooser { + fn choose(&mut self, n: NonZero) -> usize; +} diff --git a/crates/vchordrq/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs index 6b33a165..fa4c0f9b 100644 --- a/crates/vchordrq/src/prewarm.rs +++ b/crates/vchordrq/src/prewarm.rs @@ -16,7 +16,7 @@ use crate::closure_lifetime_binder::{id_0, id_1, id_2}; use crate::operator::Operator; use crate::tape::{by_directory, by_next}; use crate::tuples::*; -use crate::{Opaque, Page, tape, vectors}; +use crate::{Opaque, Page, centroids, tape}; use algo::RelationRead; use algo::accessor::FunctionalAccessor; use algo::prefetcher::PrefetcherSequenceFamily; @@ -48,7 +48,7 @@ where let prefetch = meta_tuple.centroid_prefetch().to_vec().into_iter(); let head = meta_tuple.centroid_head(); let first = meta_tuple.first(); - vectors::read_for_h1_tuple::(prefetch.map(|id| index.read(id)), head, ()); + centroids::read::(prefetch.map(|id| index.read(id)), head, ()); results.push(first); writeln!(message, "------------------------").unwrap(); writeln!(message, "number of nodes: {}", results.len()).unwrap(); @@ -66,11 +66,7 @@ where by_next(index, first).inspect(|_| counter += 1), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), |(), head, _, first, prefetch| { - vectors::read_for_h1_tuple::( - prefetch.iter().map(|&id| index.read(id)), - head, - (), - ); + centroids::read::(prefetch.iter().map(|&id| index.read(id)), head, ()); results.push(first); }, ); diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs index bb3948d2..0c828d92 100644 --- a/crates/vchordrq/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -91,7 +91,7 @@ pub fn rerank_index< prefetcher, cache: BinaryHeap::new(), f: id_4::<_, P, _, _, _>(move |payload, prefetch, head| { - vectors::read_for_h0_tuple::( + vectors::read::( prefetch, head, payload, diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index 4b3bd2cd..eaf503e0 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -17,7 +17,7 @@ use crate::linked_vec::LinkedVec; use crate::operator::*; use crate::tape::{by_directory, by_next}; use crate::tuples::*; -use crate::{Opaque, Page, tape, vectors}; +use crate::{Opaque, Page, centroids, tape}; use algo::accessor::{FunctionalAccessor, LAccess}; use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; use algo::{BorrowedIter, Bump, PackedRefMut4, PackedRefMut8, RelationRead}; @@ -67,7 +67,7 @@ where let prefetch = BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); let head = meta_tuple.centroid_head(); - let distance = vectors::read_for_h1_tuple::( + let distance = centroids::read::( prefetch.map(|id| index.read(id)), head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), @@ -112,7 +112,7 @@ where while let Some(((Reverse(_), AlwaysEqual(&mut (first, norm, head, ..))), prefetch)) = heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { - let distance = vectors::read_for_h1_tuple::( + let distance = centroids::read::( prefetch, head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), @@ -228,7 +228,7 @@ where let prefetch = BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); let head = meta_tuple.centroid_head(); - let distance = vectors::read_for_h1_tuple::( + let distance = centroids::read::( prefetch.map(|id| index.read(id)), head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), @@ -273,7 +273,7 @@ where while let Some(((Reverse(_), AlwaysEqual(&mut (first, norm, head, ..))), prefetch)) = heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) { - let distance = vectors::read_for_h1_tuple::( + let distance = centroids::read::( prefetch, head, LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index a1018151..054a8c5e 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -20,7 +20,7 @@ use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 11; +const VERSION: u64 = 12; #[inline(always)] fn tag(source: &[u8]) -> Tag { @@ -56,9 +56,11 @@ struct MetaTupleHeader { cells_s: u16, cells_e: u16, _padding_0: [Padding; 2], - vectors_first: u32, + centroids_first: u32, + vectors_first_s: u16, + vectors_first_e: u16, freepages_first: u32, - _padding_1: [Padding; 2], + _padding_1: [Padding; 6], // tree centroid_prefetch_s: u16, centroid_prefetch_e: u16, @@ -73,7 +75,8 @@ pub struct MetaTuple { pub is_residual: bool, pub rerank_in_heap: bool, pub cells: Vec, - pub vectors_first: u32, + pub centroids_first: u32, + pub vectors_first: Vec, pub freepages_first: u32, pub centroid_prefetch: Vec, pub centroid_head: u16, @@ -92,6 +95,7 @@ impl Tuple for MetaTuple { is_residual, rerank_in_heap, cells, + centroids_first, vectors_first, freepages_first, centroid_prefetch, @@ -108,6 +112,13 @@ impl Tuple for MetaTuple { while buffer.len() % ALIGN != 0 { buffer.push(0); } + // vectors_first + let vectors_first_s = buffer.len() as u16; + buffer.extend(vectors_first.as_bytes()); + let vectors_first_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } // centroid_prefetch let centroid_prefetch_s = buffer.len() as u16; buffer.extend(centroid_prefetch.as_bytes()); @@ -125,7 +136,9 @@ impl Tuple for MetaTuple { rerank_in_heap: (*rerank_in_heap).into(), cells_s, cells_e, - vectors_first: *vectors_first, + centroids_first: *centroids_first, + vectors_first_s, + vectors_first_e, freepages_first: *freepages_first, centroid_prefetch_s, centroid_prefetch_e, @@ -157,13 +170,15 @@ impl WithReader for MetaTuple { ); } let header: &MetaTupleHeader = checker.prefix(size_of::()); + let cells = checker.bytes(header.cells_s, header.cells_e); + let vectors_first = checker.bytes(header.vectors_first_s, header.vectors_first_e); let centroid_prefetch = checker.bytes(header.centroid_prefetch_s, header.centroid_prefetch_e); - let cells = checker.bytes(header.cells_s, header.cells_e); MetaTupleReader { header, - centroid_prefetch, cells, + vectors_first, + centroid_prefetch, } } _ => panic!("deserialization: bad magic number"), @@ -174,8 +189,9 @@ impl WithReader for MetaTuple { #[derive(Debug, Clone, Copy)] pub struct MetaTupleReader<'a> { header: &'a MetaTupleHeader, - centroid_prefetch: &'a [u32], cells: &'a [u32], + vectors_first: &'a [u32], + centroid_prefetch: &'a [u32], } impl<'a> MetaTupleReader<'a> { @@ -194,8 +210,11 @@ impl<'a> MetaTupleReader<'a> { pub fn cells(self) -> &'a [u32] { self.cells } - pub fn vectors_first(self) -> u32 { - self.header.vectors_first + pub fn centroids_first(self) -> u32 { + self.header.centroids_first + } + pub fn vectors_first(self) -> &'a [u32] { + self.vectors_first } pub fn freepages_first(self) -> u32 { self.header.freepages_first @@ -255,6 +274,162 @@ impl FreepagesTupleWriter<'_> { } } +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct CentroidTupleHeader0 { + metadata_s: u16, + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 2], +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct CentroidTupleHeader1 { + head: u16, + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 2], +} + +#[derive(Debug, Clone, PartialEq)] +pub enum CentroidTuple { + _0 { + metadata: V::Metadata, + elements: Vec, + }, + _1 { + head: u16, + elements: Vec, + }, +} + +impl Tuple for CentroidTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + CentroidTuple::_0 { metadata, elements } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // metadata + let metadata_s = buffer.len() as u16; + buffer.extend(metadata.as_bytes()); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + CentroidTupleHeader0 { + metadata_s, + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + CentroidTuple::_1 { head, elements } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + CentroidTupleHeader1 { + head: *head, + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for CentroidTuple { + type Reader<'a> = CentroidTupleReader<'a, V>; + + fn deserialize_ref(source: &[u8]) -> CentroidTupleReader<'_, V> { + let tag = tag(source); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &CentroidTupleHeader0 = checker.prefix(size_of::()); + let metadata = checker.prefix(header.metadata_s); + let elements = checker.bytes(header.elements_s, header.elements_e); + CentroidTupleReader::_0(CentroidTupleReader0 { + header, + elements, + metadata, + }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &CentroidTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + CentroidTupleReader::_1(CentroidTupleReader1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +#[derive(Clone)] +pub struct CentroidTupleReader0<'a, V: Vector> { + #[allow(dead_code)] + header: &'a CentroidTupleHeader0, + metadata: &'a V::Metadata, + elements: &'a [V::Element], +} + +impl Copy for CentroidTupleReader0<'_, V> {} + +#[derive(Clone)] +pub struct CentroidTupleReader1<'a, V: Vector> { + header: &'a CentroidTupleHeader1, + elements: &'a [V::Element], +} + +impl Copy for CentroidTupleReader1<'_, V> {} + +#[derive(Clone)] +pub enum CentroidTupleReader<'a, V: Vector> { + _0(CentroidTupleReader0<'a, V>), + _1(CentroidTupleReader1<'a, V>), +} + +impl Copy for CentroidTupleReader<'_, V> {} + +impl<'a, V: Vector> CentroidTupleReader<'a, V> { + pub fn elements(self) -> &'a [::Element] { + match self { + CentroidTupleReader::_0(this) => this.elements, + CentroidTupleReader::_1(this) => this.elements, + } + } + pub fn metadata_or_head(self) -> Result { + match self { + CentroidTupleReader::_0(this) => Ok(*this.metadata), + CentroidTupleReader::_1(this) => Err(this.header.head), + } + } +} + #[repr(C, align(8))] #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader0 { diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index 9296dfd5..c42edecd 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -24,6 +24,9 @@ pub struct VchordrqIndexOptions { pub residual_quantization: bool, #[serde(default = "VchordrqIndexOptions::default_rerank_in_table")] pub rerank_in_table: bool, + #[serde(default = "VchordrqIndexOptions::default_degree_of_parallelism")] + #[validate(range(min = 1, max = 256))] + pub degree_of_parallelism: u32, } impl VchordrqIndexOptions { @@ -33,6 +36,9 @@ impl VchordrqIndexOptions { fn default_rerank_in_table() -> bool { false } + fn default_degree_of_parallelism() -> u32 { + 32 + } } impl Default for VchordrqIndexOptions { @@ -40,6 +46,7 @@ impl Default for VchordrqIndexOptions { Self { residual_quantization: Self::default_residual_quantization(), rerank_in_table: Self::default_rerank_in_table(), + degree_of_parallelism: Self::default_degree_of_parallelism(), } } } @@ -57,14 +64,14 @@ pub enum BorrowedVector<'a> { } #[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum DistanceKind { L2S, Dot, } #[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum VectorKind { Vecf32, Vecf16, @@ -79,16 +86,12 @@ impl VectorKind { } } -#[derive(Debug, Clone, Copy, Serialize, Deserialize, Validate)] -#[serde(deny_unknown_fields)] +#[derive(Debug, Clone, Copy, Validate)] #[validate(schema(function = "Self::validate_self"))] pub struct VectorOptions { #[validate(range(min = 1))] - #[serde(rename = "dimensions")] pub dims: u32, - #[serde(rename = "vector")] pub v: VectorKind, - #[serde(rename = "distance")] pub d: DistanceKind, } diff --git a/crates/vchordrq/src/vectors.rs b/crates/vchordrq/src/vectors.rs index b50a71f2..183fb2e0 100644 --- a/crates/vchordrq/src/vectors.rs +++ b/crates/vchordrq/src/vectors.rs @@ -15,40 +15,12 @@ use crate::operator::*; use crate::tuples::*; use crate::{Opaque, Page, PageGuard, tape}; -use algo::accessor::{Accessor1, TryAccessor1}; +use algo::accessor::TryAccessor1; use algo::{RelationRead, RelationWrite}; use std::num::NonZero; use vector::VectorOwned; -pub fn read_for_h1_tuple< - 'a, - R: RelationRead + 'a, - O: Operator, - A: Accessor1<::Element, ::Metadata>, ->( - mut prefetch: impl Iterator>, - head: u16, - accessor: A, -) -> A::Output { - let mut cursor = Err(head); - let mut result = accessor; - while let Err(head) = cursor { - let guard = prefetch.next().expect("data corruption"); - let bytes = guard.get(head).expect("data corruption"); - let tuple = VectorTuple::::deserialize_ref(bytes); - if tuple.payload().is_some() { - panic!("data corruption"); - } - result.push(tuple.elements()); - cursor = tuple.metadata_or_head(); - } - if prefetch.next().is_some() { - panic!("data corruption"); - } - result.finish(cursor.expect("data corruption")) -} - -pub fn read_for_h0_tuple< +pub fn read< 'a, R: RelationRead + 'a, O: Operator, @@ -85,15 +57,21 @@ pub fn append( vectors_first: u32, vector: ::Borrowed<'_>, payload: NonZero, + skip_search: bool, ) -> (Vec, u16) where R::Page: Page, { - fn append(index: &R, first: u32, bytes: &[u8]) -> (u32, u16) + fn append( + index: &R, + first: u32, + bytes: &[u8], + skip_search: bool, + ) -> (u32, u16) where R::Page: Page, { - if let Some(mut write) = index.search(bytes.len()) { + if !skip_search && let Some(mut write) = index.search(bytes.len()) { let i = write .alloc(bytes) .expect("implementation: a free page cannot accommodate a single tuple"); @@ -117,7 +95,7 @@ where head, }, }); - let (id, head) = append(index, vectors_first, &bytes); + let (id, head) = append(index, vectors_first, &bytes, skip_search); chain = Err(head); prefetch.push(id); } diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index 14bbcd94..1fd60055 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -12,11 +12,10 @@ // // Copyright (c) 2025 TensorChord Inc. -use serde::{Deserialize, Serialize}; use std::ffi::{CStr, CString}; use std::num::NonZero; -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy)] pub enum Typmod { Any, Dims(NonZero), diff --git a/src/index/storage.rs b/src/index/storage.rs index 8b7849a7..6bda0a4e 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -336,19 +336,36 @@ impl RelationWrite for PostgresRelation { ) -> PostgresBufferWriteGuard { unsafe { use pgrx::pg_sys::{ - BUFFER_LOCK_EXCLUSIVE, ExclusiveLock, ForkNumber, GENERIC_XLOG_FULL_IMAGE, - GenericXLogRegisterBuffer, GenericXLogStart, LockBuffer, LockRelationForExtension, - ReadBufferExtended, ReadBufferMode, UnlockRelationForExtension, + BUFFER_LOCK_EXCLUSIVE, GENERIC_XLOG_FULL_IMAGE, GenericXLogRegisterBuffer, + GenericXLogStart, LockBuffer, }; - LockRelationForExtension(self.raw, ExclusiveLock as _); - let buf = ReadBufferExtended( - self.raw, - ForkNumber::MAIN_FORKNUM, - u32::MAX, - ReadBufferMode::RBM_NORMAL, - std::ptr::null_mut(), - ); - UnlockRelationForExtension(self.raw, ExclusiveLock as _); + let buf; + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15"))] + { + use pgrx::pg_sys::{ + ExclusiveLock, ForkNumber, LockRelationForExtension, ReadBufferExtended, + ReadBufferMode, UnlockRelationForExtension, + }; + LockRelationForExtension(self.raw, ExclusiveLock as _); + buf = ReadBufferExtended( + self.raw, + ForkNumber::MAIN_FORKNUM, + u32::MAX, + ReadBufferMode::RBM_NORMAL, + std::ptr::null_mut(), + ); + UnlockRelationForExtension(self.raw, ExclusiveLock as _); + } + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + { + use pgrx::pg_sys::{BufferManagerRelation, ExtendBufferedRel, ForkNumber}; + let bmr = BufferManagerRelation { + rel: self.raw, + smgr: std::ptr::null_mut(), + relpersistence: 0, + }; + buf = ExtendBufferedRel(bmr, ForkNumber::MAIN_FORKNUM, std::ptr::null_mut(), 0); + } LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); let state = GenericXLogStart(self.raw); let mut page = NonNull::new( @@ -667,13 +684,6 @@ impl RelationReadStream for PostgresRelation { } } -impl RelationLength for PostgresRelation { - fn len(&self) -> u32 { - use pgrx::pg_sys::{ForkNumber, RelationGetNumberOfBlocksInFork}; - unsafe { RelationGetNumberOfBlocksInFork(self.raw, ForkNumber::MAIN_FORKNUM) } - } -} - #[inline(always)] fn lp_flags(x: pgrx::pg_sys::ItemIdData) -> u32 { let x: u32 = unsafe { std::mem::transmute(x) }; diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/algo.rs index b3e3f7f9..52416c01 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/algo.rs @@ -51,7 +51,7 @@ pub fn bulkdelete( check: impl Fn(), callback: impl Fn(NonZero) -> bool, ) where - R: RelationRead + RelationWrite + RelationLength, + R: RelationRead + RelationWrite, R::Page: Page, { match (opfamily.vector_kind(), opfamily.distance_kind()) { diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index fbc63f56..f660e7ec 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -227,13 +227,13 @@ pub unsafe extern "C-unwind" fn ambuild( let nparticipants = leader.nparticipants; loop { pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordgshared).mutex); - if (*leader.vchordgshared).nparticipantsdone == nparticipants { + if (*leader.vchordgshared).workers_done == nparticipants { pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordgshared).mutex); break; } pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordgshared).mutex); pgrx::pg_sys::ConditionVariableSleep( - &raw mut (*leader.vchordgshared).workersdonecv, + &raw mut (*leader.vchordgshared).condvar_workers_done, pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, ); } @@ -259,19 +259,17 @@ pub unsafe extern "C-unwind" fn ambuild( } struct VchordgShared { - /* Immutable state */ + /* immutable state */ heaprelid: pgrx::pg_sys::Oid, indexrelid: pgrx::pg_sys::Oid, isconcurrent: bool, - /* Worker progress */ - workersdonecv: pgrx::pg_sys::ConditionVariable, - - /* Mutex for mutable state */ + /* locking */ mutex: pgrx::pg_sys::slock_t, + condvar_workers_done: pgrx::pg_sys::ConditionVariable, - /* Mutable state */ - nparticipantsdone: i32, + /* mutable state */ + workers_done: i32, indtuples: u64, } @@ -384,12 +382,12 @@ impl VchordgLeader { heaprelid: (*heap_relation).rd_id, indexrelid: (*index_relation).rd_id, isconcurrent, - workersdonecv: std::mem::zeroed(), + condvar_workers_done: std::mem::zeroed(), mutex: std::mem::zeroed(), - nparticipantsdone: 0, + workers_done: 0, indtuples: 0, }); - pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordgshared).workersdonecv); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordgshared).condvar_workers_done); pgrx::pg_sys::SpinLockInit(&raw mut (*vchordgshared).mutex); vchordgshared }; @@ -470,6 +468,7 @@ pub unsafe extern "C-unwind" fn vchordg_parallel_build_main( _seg: *mut pgrx::pg_sys::dsm_segment, toc: *mut pgrx::pg_sys::shm_toc, ) { + let _ = rand::rng().reseed(); let vchordgshared = unsafe { pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() }; @@ -565,9 +564,9 @@ unsafe fn parallel_build( } unsafe { pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordgshared).mutex); - (*vchordgshared).nparticipantsdone += 1; + (*vchordgshared).workers_done += 1; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordgshared).mutex); - pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordgshared).workersdonecv); + pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordgshared).condvar_workers_done); } } diff --git a/src/index/vchordrq/algo.rs b/src/index/vchordrq/algo.rs index 55f094bc..15802c51 100644 --- a/src/index/vchordrq/algo.rs +++ b/src/index/vchordrq/algo.rs @@ -22,7 +22,7 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vchordrq::operator::Op; use vchordrq::types::*; -use vchordrq::{FastHeap, Opaque}; +use vchordrq::{Chooser, FastHeap, Opaque}; use vector::VectorOwned; use vector::vect::{VectBorrowed, VectOwned}; @@ -154,11 +154,13 @@ pub fn insert( payload: NonZero, vector: OwnedVector, skip_freespaces: bool, + skip_search: bool, + chooser: &mut impl Chooser, + bump: &impl Bump, ) where R: RelationRead + RelationWrite, R::Page: Page, { - let bump = bumpalo::Bump::new(); let make_h1_plain_prefetcher = MakeH1PlainPrefetcherForInsertion { index }; match (vector, opfamily.distance_kind()) { (OwnedVector::Vecf32(vector), DistanceKind::L2S) => { @@ -168,13 +170,15 @@ pub fn insert( index, payload, vector.as_borrowed(), + chooser, + skip_search, ); vchordrq::insert::<_, Op, L2S>>( index, payload, projected.as_borrowed(), key, - &bump, + bump, make_h1_plain_prefetcher, skip_freespaces, ) @@ -186,13 +190,15 @@ pub fn insert( index, payload, vector.as_borrowed(), + chooser, + skip_search, ); vchordrq::insert::<_, Op, Dot>>( index, payload, projected.as_borrowed(), key, - &bump, + bump, make_h1_plain_prefetcher, skip_freespaces, ) @@ -204,13 +210,15 @@ pub fn insert( index, payload, vector.as_borrowed(), + chooser, + skip_search, ); vchordrq::insert::<_, Op, L2S>>( index, payload, projected.as_borrowed(), key, - &bump, + bump, make_h1_plain_prefetcher, skip_freespaces, ) @@ -222,13 +230,15 @@ pub fn insert( index, payload, vector.as_borrowed(), + chooser, + skip_search, ); vchordrq::insert::<_, Op, Dot>>( index, payload, projected.as_borrowed(), key, - &bump, + bump, make_h1_plain_prefetcher, skip_freespaces, ) diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index df476cb7..42347bc0 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -25,7 +25,9 @@ use pgrx::pg_sys::{Datum, ItemPointerData}; use rand::Rng; use simd::{Floating, f16}; use std::ffi::CStr; +use std::num::NonZero; use std::ops::Deref; +use vchordrq::Chooser; use vchordrq::types::*; use vector::VectorOwned; use vector::vect::VectOwned; @@ -343,8 +345,8 @@ pub unsafe extern "C-unwind" fn ambuild( reporter.phase(BuildPhase::from_code(BuildPhaseCode::Build)); crate::index::vchordrq::algo::build(vector_options, vchordrq_options.index, &index, structures); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); - let cached = if vchordrq_options.build.pin { - let mut trace = vchordrq::cache(&index); + let cached = if vchordrq_options.build.pin >= 0 { + let mut trace = vchordrq::cache(&index, vchordrq_options.build.pin); trace.sort(); trace.dedup(); if let Some(max) = trace.last().copied() { @@ -387,13 +389,13 @@ pub unsafe extern "C-unwind" fn ambuild( let nparticipants = leader.nparticipants; loop { pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordrqshared).mutex); - if (*leader.vchordrqshared).nparticipantsdone == nparticipants { + if (*leader.vchordrqshared).workers_done == nparticipants { pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); break; } pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); pgrx::pg_sys::ConditionVariableSleep( - &raw mut (*leader.vchordrqshared).workersdonecv, + &raw mut (*leader.vchordrqshared).condvar_workers_done, pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, ); } @@ -424,19 +426,18 @@ pub unsafe extern "C-unwind" fn ambuild( } struct VchordrqShared { - /* Immutable state */ + /* immutable state */ heaprelid: pgrx::pg_sys::Oid, indexrelid: pgrx::pg_sys::Oid, isconcurrent: bool, - /* Worker progress */ - workersdonecv: pgrx::pg_sys::ConditionVariable, - - /* Mutex for mutable state */ + /* locking */ mutex: pgrx::pg_sys::slock_t, + condvar_workers_done: pgrx::pg_sys::ConditionVariable, - /* Mutable state */ - nparticipantsdone: i32, + /* mutable state */ + workers_ready: i32, + workers_done: i32, indtuples: u64, } @@ -681,12 +682,13 @@ impl VchordrqLeader { heaprelid: (*heap_relation).rd_id, indexrelid: (*index_relation).rd_id, isconcurrent, - workersdonecv: std::mem::zeroed(), + condvar_workers_done: std::mem::zeroed(), mutex: std::mem::zeroed(), - nparticipantsdone: 0, + workers_ready: 0, + workers_done: 0, indtuples: 0, }); - pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).workersdonecv); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_workers_done); pgrx::pg_sys::SpinLockInit(&raw mut (*vchordrqshared).mutex); vchordrqshared }; @@ -768,6 +770,7 @@ pub unsafe extern "C-unwind" fn vchordrq_parallel_build_main( _seg: *mut pgrx::pg_sys::dsm_segment, toc: *mut pgrx::pg_sys::shm_toc, ) { + let _ = rand::rng().reseed(); let vchordrqshared = unsafe { pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() }; @@ -824,6 +827,16 @@ unsafe fn parallel_build( mut callback: impl FnMut(u64), ) { use vchordrq_cached::VchordrqCachedReader; + + let wid; + + unsafe { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + wid = (*vchordrqshared).workers_ready as u32; + (*vchordrqshared).workers_ready += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + let cached = VchordrqCachedReader::deserialize_ref(unsafe { let bytes = (vchordrqcached as *const u64).read_unaligned(); std::slice::from_raw_parts(vchordrqcached.add(8), bytes as _) @@ -840,13 +853,32 @@ unsafe fn parallel_build( opfamily, scan, }; + + struct IdChooser(u32); + impl Chooser for IdChooser { + fn choose(&mut self, n: NonZero) -> usize { + self.0 as usize % n.get() + } + } + match cached { VchordrqCachedReader::_0(_) => { heap.traverse(true, |(ctid, store)| { for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); + let mut chooser = IdChooser(wid); + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::algo::insert( + opfamily, + &index, + payload, + vector, + true, + true, + &mut chooser, + &bump, + ); } unsafe { let indtuples; @@ -869,7 +901,18 @@ unsafe fn parallel_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); + let mut chooser = IdChooser(wid); + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::algo::insert( + opfamily, + &index, + payload, + vector, + true, + true, + &mut chooser, + &bump, + ); } unsafe { let indtuples; @@ -886,9 +929,9 @@ unsafe fn parallel_build( } unsafe { pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).nparticipantsdone += 1; + (*vchordrqshared).workers_done += 1; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordrqshared).workersdonecv); + pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordrqshared).condvar_workers_done); } } @@ -900,7 +943,9 @@ unsafe fn sequential_build( mut callback: impl FnMut(u64), ) -> u64 { use vchordrq_cached::VchordrqCachedReader; + let cached = VchordrqCachedReader::deserialize_ref(vchordrqcached); + let index = unsafe { PostgresRelation::new(index_relation) }; let opfamily = unsafe { opfamily(index_relation) }; @@ -912,6 +957,13 @@ unsafe fn sequential_build( scan: std::ptr::null_mut(), }; + struct ChooseZero; + impl Chooser for ChooseZero { + fn choose(&mut self, _: NonZero) -> usize { + 0 + } + } + let mut indtuples = 0; match cached { VchordrqCachedReader::_0(_) => { @@ -919,7 +971,18 @@ unsafe fn sequential_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); + let mut chooser = ChooseZero; + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::algo::insert( + opfamily, + &index, + payload, + vector, + true, + true, + &mut chooser, + &bump, + ); } indtuples += 1; callback(indtuples); @@ -934,7 +997,18 @@ unsafe fn sequential_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, true); + let mut chooser = ChooseZero; + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::algo::insert( + opfamily, + &index, + payload, + vector, + true, + true, + &mut chooser, + &bump, + ); } indtuples += 1; callback(indtuples); diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 03bb274c..dd898b6f 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -23,12 +23,14 @@ use crate::index::vchordrq::scanners::*; use crate::recorder::DefaultRecorder; use pgrx::datum::Internal; use pgrx::pg_sys::Datum; +use rand::RngCore; use std::cell::LazyCell; use std::ffi::CStr; use std::num::NonZero; use std::ops::DerefMut; use std::ptr::NonNull; use std::sync::OnceLock; +use vchordrq::Chooser; #[repr(C)] struct Reloption { @@ -301,6 +303,13 @@ unsafe fn aminsertinner( is_null: *mut bool, ctid: pgrx::pg_sys::ItemPointer, ) -> bool { + struct RngChooser(T); + impl Chooser for RngChooser { + fn choose(&mut self, n: NonZero) -> usize { + rand::Rng::random_range(&mut self.0, 0..n.get()) + } + } + let opfamily = unsafe { opfamily(index_relation) }; let index = unsafe { PostgresRelation::new(index_relation) }; let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; @@ -309,7 +318,18 @@ unsafe fn aminsertinner( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::vchordrq::algo::insert(opfamily, &index, payload, vector, false); + let mut chooser = RngChooser(rand::rng()); + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::algo::insert( + opfamily, + &index, + payload, + vector, + false, + false, + &mut chooser, + &bump, + ); } } false diff --git a/src/index/vchordrq/types.rs b/src/index/vchordrq/types.rs index 44fb0ec9..dd2c166b 100644 --- a/src/index/vchordrq/types.rs +++ b/src/index/vchordrq/types.rs @@ -119,18 +119,34 @@ impl Validate for VchordrqBuildSourceOptions { #[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] -#[serde(rename_all = "snake_case")] pub struct VchordrqBuildOptions { #[serde(flatten)] #[validate(nested)] pub source: VchordrqBuildSourceOptions, + #[serde(deserialize_with = "VchordrqBuildOptions::deserialize_pin")] #[serde(default = "VchordrqBuildOptions::default_pin")] - pub pin: bool, + #[validate(range(min = -1, max = 2))] + pub pin: i32, } impl VchordrqBuildOptions { - pub fn default_pin() -> bool { - false + pub fn deserialize_pin<'de, D: serde::Deserializer<'de>>( + deserializer: D, + ) -> Result { + #[derive(Deserialize)] + #[serde(untagged)] + enum Untagged { + Bool(bool), + I32(i32), + } + + match Untagged::deserialize(deserializer)? { + Untagged::Bool(b) => Ok(if b { 1 } else { -1 }), + Untagged::I32(i) => Ok(i), + } + } + pub fn default_pin() -> i32 { + -1 } } From 07fd78517e3371f95ffb51ef6d60754ea0540ee8 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 26 Sep 2025 11:36:13 +0800 Subject: [PATCH 226/324] chore: pg18 updates (#349) job: +psql_macos job: +psql_windows job: +psql_alpine job: +check_debian Signed-off-by: usamoi --- .github/workflows/check.yml | 53 +++--- .github/workflows/release.yml | 2 +- Cargo.lock | 300 +++++++--------------------------- Cargo.toml | 8 +- crates/make/Cargo.toml | 2 +- crates/simd/Cargo.toml | 2 +- 6 files changed, 98 insertions(+), 269 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index ae8cf035..611a14ca 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -185,7 +185,7 @@ jobs: - name: Set up Environment run: | rustup update - if [ "${{ matrix.version }}" = "18" ]; then + if [ "${{ matrix.version }}" = "13" ]; then rustup default beta rustup component add clippy rustfmt fi @@ -199,9 +199,6 @@ jobs: sudo apt-get update sudo apt-get install -y postgresql-common sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y - if [ "${{ matrix.version }}" = "18" ]; then - sudo add-apt-repository "deb https://apt.postgresql.org/pub/repos/apt/ $(lsb_release -s -c)-pgdg main 18" - fi sudo apt-get update sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} @@ -263,7 +260,7 @@ jobs: strategy: matrix: - version: ["13", "14", "15", "16", "17"] + version: ["13", "14", "15", "16", "17", "18"] arch: ["aarch64", "x86_64"] runs-on: ${{ matrix.arch == 'aarch64' && 'macos-15' || 'macos-15-intel' }} @@ -281,7 +278,8 @@ jobs: run: | rustup update - HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install llvm@18 + brew update + echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install postgresql@${{ matrix.version }} @@ -295,9 +293,9 @@ jobs: echo PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV mkdir ~/pgvector-install - curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.0.tar.gz | tar -xz -C ~/pgvector-install - make -C ~/pgvector-install/pgvector-0.8.0 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config - sudo make -C ~/pgvector-install/pgvector-0.8.0 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config + sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config install curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest @@ -348,7 +346,7 @@ jobs: strategy: matrix: - version: ["13", "14", "15", "16", "17"] + version: ["13", "14", "15", "16", "17", "18"] arch: ["x86_64"] runs-on: "windows-2022" @@ -371,19 +369,22 @@ jobs: & 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise\Common7\Tools\Launch-VsDevShell.ps1' -HostArch amd64 -Arch amd64 if ( "${{ matrix.version }}" -eq "13" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-13.21-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-13.22-1-windows-x64-binaries.zip" } if ( "${{ matrix.version }}" -eq "14" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-14.18-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-14.19-1-windows-x64-binaries.zip" } if ( "${{ matrix.version }}" -eq "15" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-15.13-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-15.14-1-windows-x64-binaries.zip" } if ( "${{ matrix.version }}" -eq "16" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-16.9-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-16.10-1-windows-x64-binaries.zip" } if ( "${{ matrix.version }}" -eq "17" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-17.5-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-17.6-1-windows-x64-binaries.zip" + } + if ( "${{ matrix.version }}" -eq "18" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-18.0-1-windows-x64-binaries.zip" } Invoke-WebRequest -Uri $postgresql_url -OutFile "$env:TEMP\postgresql-install.zip" Expand-Archive -Path "$env:TEMP\postgresql-install.zip" -DestinationPath "D:\postgresql-install" -Force @@ -396,9 +397,9 @@ jobs: D:\postgresql-install\pgsql\bin\psql.exe -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' D:\postgresql-install\pgsql\bin\pg_ctl.exe stop -D D:\postgresql-install\pgsql\data - Invoke-WebRequest -Uri "https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.0.zip" -OutFile "$env:TEMP\pgvector-install.zip" + Invoke-WebRequest -Uri "https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.zip" -OutFile "$env:TEMP\pgvector-install.zip" Expand-Archive -Path "$env:TEMP\pgvector-install.zip" -DestinationPath "D:\pgvector-install" -Force - Push-Location -Path D:\pgvector-install\pgvector-0.8.0 + Push-Location -Path D:\pgvector-install\pgvector-0.8.1 nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" install Pop-Location @@ -498,9 +499,9 @@ jobs: sudo -iu postgres pg_ctl stop -D /var/lib/postgresql/data mkdir ~/pgvector-install - curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.0.tar.gz | tar -xz -C ~/pgvector-install - make -C ~/pgvector-install/pgvector-0.8.0 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config - make -C ~/pgvector-install/pgvector-0.8.0 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config install curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest @@ -603,6 +604,12 @@ jobs: rust_triple: "powerpc64le-unknown-linux-gnu" gcc_version: "12" debian_version: "bookworm" + - version: "18" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" - version: "17" platform: "riscv64" clang_triple: "riscv64-linux-gnu" @@ -675,9 +682,9 @@ jobs: sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main stop mkdir ~/pgvector-install - curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.0.tar.gz | tar -xz -C ~/pgvector-install - make -C ~/pgvector-install/pgvector-0.8.0 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" - sudo make -C ~/pgvector-install/pgvector-0.8.0 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install + make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" + sudo make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 77f67650..d6223e81 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -39,7 +39,7 @@ jobs: needs: ["semver"] strategy: matrix: - version: ["13", "14", "15", "16", "17"] + version: ["13", "14", "15", "16", "17", "18"] arch: ["x86_64", "aarch64"] runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} diff --git a/Cargo.lock b/Cargo.lock index dff15fa8..28b37649 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -151,7 +151,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "374b7c592d9c00c1f4972ea58390ac6b18cbb6ab79011f3bdc90a0b82ca06b77" dependencies = [ "serde", - "toml 0.9.7", + "toml", ] [[package]] @@ -408,7 +408,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.0", + "windows-sys 0.61.1", ] [[package]] @@ -539,21 +539,6 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" -[[package]] -name = "hermit-abi" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" - -[[package]] -name = "home" -version = "0.5.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "589533453244b0995c858700322199b2becb13b627df2851f64a2775d024abcf" -dependencies = [ - "windows-sys 0.59.0", -] - [[package]] name = "icu_collections" version = "2.0.0" @@ -683,17 +668,6 @@ dependencies = [ "hashbrown 0.16.0", ] -[[package]] -name = "is-terminal" -version = "0.4.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" -dependencies = [ - "hermit-abi", - "libc", - "windows-sys 0.59.0", -] - [[package]] name = "is_ci" version = "1.2.0" @@ -723,9 +697,9 @@ checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" [[package]] name = "js-sys" -version = "0.3.80" +version = "0.3.81" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "852f13bec5eba4ba9afbeb93fd7c13fe56147f055939ae21c43a29a0ecb2702e" +checksum = "ec48937a97411dcb524a265206ccd4c90bb711fca92b2792c407f268825b9305" dependencies = [ "once_cell", "wasm-bindgen", @@ -754,7 +728,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" dependencies = [ "cfg-if", - "windows-link 0.2.0", + "windows-link", ] [[package]] @@ -808,9 +782,9 @@ dependencies = [ [[package]] name = "memchr" -version = "2.7.5" +version = "2.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" +checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" [[package]] name = "mimalloc" @@ -871,8 +845,7 @@ version = "4.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48dd4f4a2c8405440fd0462561f0e5806bd0f77e86f51c761481bdd4018b545e" dependencies = [ - "supports-color 2.1.0", - "supports-color 3.0.2", + "supports-color", ] [[package]] @@ -911,15 +884,14 @@ dependencies = [ [[package]] name = "pgrx" -version = "0.16.0" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c898b37587b29f8fa7ca20bfaaf57103e10fa81e6caed2cc7b742f96088d851" +checksum = "fdcfb88f7fa9ba42b4ea9d1f85a1d968bbb407cc30308b35f73bdfe6c966f64b" dependencies = [ "bitflags", "bitvec", "enum-map", "libc", - "once_cell", "pgrx-macros", "pgrx-pg-sys", "pgrx-sql-entity-graph", @@ -933,9 +905,9 @@ dependencies = [ [[package]] name = "pgrx-bindgen" -version = "0.16.0" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35439a65a6c0cc17e9758005da0557cff583a3e4da0f92a4a744494ac21c190e" +checksum = "00e35193b7e71e2f612d336cecd00db0f049f4cc609f2b1c9a34755b5ec559d7" dependencies = [ "bindgen", "cc", @@ -962,9 +934,9 @@ dependencies = [ [[package]] name = "pgrx-macros" -version = "0.16.0" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3586274c7d7237e68bebb871cf9244444f1ad78bb395679ad4ef1015db302454" +checksum = "dab542dd4041773874f90cd8e3448195749548dc3fb1daf501e7e11ebfb1dd22" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", @@ -974,30 +946,29 @@ dependencies = [ [[package]] name = "pgrx-pg-config" -version = "0.16.0" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cc24f6663377dca8797ff3561d9ecdaef47e518ca2e2438fd85ef3c20f421ea" +checksum = "eff9b29df94c3f9fcb0cde220f92eea6975ed05962784a98fb557754ad663501" dependencies = [ "cargo_toml", "codepage", "encoding_rs", "eyre", - "home", "owo-colors", "pathsearch", "serde", "serde_json", "thiserror", - "toml 0.8.23", + "toml", "url", "winapi", ] [[package]] name = "pgrx-pg-sys" -version = "0.16.0" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b44a35477c0752c00cab38c2d8b7ea2d2d3554e3385870d7ad43e3ae5ade66a9" +checksum = "934f2536953ccb6722bef2cfdfb1f8d6d3cd4a4f2c508d56ec85b649c5680c2b" dependencies = [ "cee-scape", "libc", @@ -1005,14 +976,13 @@ dependencies = [ "pgrx-macros", "pgrx-sql-entity-graph", "serde", - "sptr", ] [[package]] name = "pgrx-sql-entity-graph" -version = "0.16.0" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59036a5849c5b0bca0f33df1fdd72206ceb74712959756b3fecba56a8e19d52a" +checksum = "07a767cb9faa612f1ba7f13718136d4006950d6f253b414ef487a03e85c47a94" dependencies = [ "convert_case", "eyre", @@ -1181,9 +1151,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.11.2" +version = "1.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23d7fd106d8c02486a8d64e778353d1cffe08ce79ac2e82f540c86d0facf6912" +checksum = "8b5288124840bee7b386bc413c487869b360b2b4ec421ea56425128692f2a82c" dependencies = [ "aho-corasick", "memchr", @@ -1193,9 +1163,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.10" +version = "0.4.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b9458fa0bfeeac22b5ca447c63aaf45f28439a709ccd244698632f9aa6394d6" +checksum = "833eb9ce86d40ef33cb1306d8accf7bc8ec2bfea4355cbdebb3df68b40925cad" dependencies = [ "aho-corasick", "memchr", @@ -1238,7 +1208,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.0", + "windows-sys 0.61.1", ] [[package]] @@ -1276,9 +1246,9 @@ checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" -version = "1.0.226" +version = "1.0.227" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0dca6411025b24b60bfa7ec1fe1f8e710ac09782dca409ee8237ba74b51295fd" +checksum = "80ece43fc6fbed4eb5392ab50c07334d3e577cbf40997ee896fe7af40bba4245" dependencies = [ "serde_core", "serde_derive", @@ -1296,18 +1266,18 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.226" +version = "1.0.227" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba2ba63999edb9dac981fb34b3e5c0d111a69b0924e253ed29d83f7c99e966a4" +checksum = "7a576275b607a2c86ea29e410193df32bc680303c82f31e275bbfcafe8b33be5" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.226" +version = "1.0.227" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8db53ae22f34573731bafa1db20f04027b2d25e02d8205921b569171699cdb33" +checksum = "51e694923b8824cf0e9b382adf0f60d4e05f348f357b38833a3fa5ed7c2ede04" dependencies = [ "proc-macro2", "quote", @@ -1327,15 +1297,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "serde_spanned" -version = "0.6.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3" -dependencies = [ - "serde", -] - [[package]] name = "serde_spanned" version = "1.0.2" @@ -1383,12 +1344,6 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" -[[package]] -name = "sptr" -version = "0.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b9b39299b249ad65f3b7e96443bad61c02ca5cd3589f46cb6d610a0fd6c0d6a" - [[package]] name = "stable_deref_trait" version = "1.2.0" @@ -1401,16 +1356,6 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" -[[package]] -name = "supports-color" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6398cde53adc3c4557306a96ce67b302968513830a77a95b2b17305d9719a89" -dependencies = [ - "is-terminal", - "is_ci", -] - [[package]] name = "supports-color" version = "3.0.2" @@ -1484,18 +1429,6 @@ dependencies = [ "zerovec", ] -[[package]] -name = "toml" -version = "0.8.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" -dependencies = [ - "serde", - "serde_spanned 0.6.9", - "toml_datetime 0.6.11", - "toml_edit", -] - [[package]] name = "toml" version = "0.9.7" @@ -1504,22 +1437,13 @@ checksum = "00e5e5d9bf2475ac9d4f0d9edab68cc573dc2fd644b0dba36b0c30a92dd9eaa0" dependencies = [ "indexmap", "serde_core", - "serde_spanned 1.0.2", - "toml_datetime 0.7.2", + "serde_spanned", + "toml_datetime", "toml_parser", "toml_writer", "winnow", ] -[[package]] -name = "toml_datetime" -version = "0.6.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c" -dependencies = [ - "serde", -] - [[package]] name = "toml_datetime" version = "0.7.2" @@ -1529,20 +1453,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "toml_edit" -version = "0.22.27" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" -dependencies = [ - "indexmap", - "serde", - "serde_spanned 0.6.9", - "toml_datetime 0.6.11", - "toml_write", - "winnow", -] - [[package]] name = "toml_parser" version = "1.0.3" @@ -1552,12 +1462,6 @@ dependencies = [ "winnow", ] -[[package]] -name = "toml_write" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" - [[package]] name = "toml_writer" version = "1.0.3" @@ -1674,7 +1578,7 @@ dependencies = [ "serde", "simd", "small_iter", - "toml 0.9.7", + "toml", "validator", "vchordg", "vchordrq", @@ -1760,9 +1664,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.103" +version = "0.2.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab10a69fbd0a177f5f649ad4d8d3305499c42bab9aef2f7ff592d0ec8f833819" +checksum = "c1da10c01ae9f1ae40cbfac0bac3b1e724b320abfcf52229f80b547c0d250e2d" dependencies = [ "cfg-if", "once_cell", @@ -1773,9 +1677,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-backend" -version = "0.2.103" +version = "0.2.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bb702423545a6007bbc368fde243ba47ca275e549c8a28617f56f6ba53b1d1c" +checksum = "671c9a5a66f49d8a47345ab942e2cb93c7d1d0339065d4f8139c486121b43b19" dependencies = [ "bumpalo", "log", @@ -1787,9 +1691,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.103" +version = "0.2.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc65f4f411d91494355917b605e1480033152658d71f722a90647f56a70c88a0" +checksum = "7ca60477e4c59f5f2986c50191cd972e3a50d8a95603bc9434501cf156a9a119" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1797,9 +1701,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.103" +version = "0.2.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ffc003a991398a8ee604a401e194b6b3a39677b3173d6e74495eb51b82e99a32" +checksum = "9f07d2f20d4da7b26400c9f4a0511e6e0345b040694e8a75bd41d578fa4421d7" dependencies = [ "proc-macro2", "quote", @@ -1810,9 +1714,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.103" +version = "0.2.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "293c37f4efa430ca14db3721dfbe48d8c33308096bd44d80ebaa775ab71ba1cf" +checksum = "bad67dc8b2a1a6e5448428adec4c3e84c43e561d8c9ee8a9e5aabeb193ec41d1" dependencies = [ "unicode-ident", ] @@ -1859,7 +1763,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.0", + "windows-sys 0.61.1", ] [[package]] @@ -1868,168 +1772,89 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" -[[package]] -name = "windows-link" -version = "0.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" - [[package]] name = "windows-link" version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "45e46c0661abb7180e7b9c281db115305d49ca1709ab8242adf09666d2173c65" -[[package]] -name = "windows-sys" -version = "0.59.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" -dependencies = [ - "windows-targets 0.52.6", -] - [[package]] name = "windows-sys" version = "0.60.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" dependencies = [ - "windows-targets 0.53.3", + "windows-targets", ] [[package]] name = "windows-sys" -version = "0.61.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e201184e40b2ede64bc2ea34968b28e33622acdbbf37104f0e4a33f7abe657aa" -dependencies = [ - "windows-link 0.2.0", -] - -[[package]] -name = "windows-targets" -version = "0.52.6" +version = "0.61.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +checksum = "6f109e41dd4a3c848907eb83d5a42ea98b3769495597450cf6d153507b166f0f" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm 0.52.6", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", + "windows-link", ] [[package]] name = "windows-targets" -version = "0.53.3" +version = "0.53.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5fe6031c4041849d7c496a8ded650796e7b6ecc19df1a431c1a363342e5dc91" +checksum = "2d42b7b7f66d2a06854650af09cfdf8713e427a439c97ad65a6375318033ac4b" dependencies = [ - "windows-link 0.1.3", - "windows_aarch64_gnullvm 0.53.0", - "windows_aarch64_msvc 0.53.0", - "windows_i686_gnu 0.53.0", - "windows_i686_gnullvm 0.53.0", - "windows_i686_msvc 0.53.0", - "windows_x86_64_gnu 0.53.0", - "windows_x86_64_gnullvm 0.53.0", - "windows_x86_64_msvc 0.53.0", + "windows-link", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" - [[package]] name = "windows_aarch64_gnullvm" version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "86b8d5f90ddd19cb4a147a5fa63ca848db3df085e25fee3cc10b39b6eebae764" -[[package]] -name = "windows_aarch64_msvc" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" - [[package]] name = "windows_aarch64_msvc" version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7651a1f62a11b8cbd5e0d42526e55f2c99886c77e007179efff86c2b137e66c" -[[package]] -name = "windows_i686_gnu" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" - [[package]] name = "windows_i686_gnu" version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1dc67659d35f387f5f6c479dc4e28f1d4bb90ddd1a5d3da2e5d97b42d6272c3" -[[package]] -name = "windows_i686_gnullvm" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" - [[package]] name = "windows_i686_gnullvm" version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ce6ccbdedbf6d6354471319e781c0dfef054c81fbc7cf83f338a4296c0cae11" -[[package]] -name = "windows_i686_msvc" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" - [[package]] name = "windows_i686_msvc" version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "581fee95406bb13382d2f65cd4a908ca7b1e4c2f1917f143ba16efe98a589b5d" -[[package]] -name = "windows_x86_64_gnu" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" - [[package]] name = "windows_x86_64_gnu" version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2e55b5ac9ea33f2fc1716d1742db15574fd6fc8dadc51caab1c16a3d3b4190ba" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" - [[package]] name = "windows_x86_64_gnullvm" version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0a6e035dd0599267ce1ee132e51c27dd29437f63325753051e71dd9e42406c57" -[[package]] -name = "windows_x86_64_msvc" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" - [[package]] name = "windows_x86_64_msvc" version = "0.53.0" @@ -2041,9 +1866,6 @@ name = "winnow" version = "0.7.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "21a0236b59786fed61e2a80582dd500fe61f18b5dca67a4a067d0bc9039339cf" -dependencies = [ - "memchr", -] [[package]] name = "winsafe" diff --git a/Cargo.toml b/Cargo.toml index 67ea8192..adc8fc2a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,13 +36,13 @@ vector = { path = "./crates/vector" } bumpalo.workspace = true dary_heap.workspace = true paste.workspace = true -pgrx = { version = "=0.16.0", default-features = false, features = ["cshim"] } +pgrx = { version = "=0.16.1", default-features = false, features = ["cshim"] } pgrx-catalog = "0.3.1" rand.workspace = true rusqlite = { version = "0.37.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true -toml = "0.9.5" +toml = "0.9.7" validator.workspace = true zerocopy.workspace = true @@ -62,14 +62,14 @@ edition = "2024" [workspace.dependencies] bumpalo = "3.19.0" -dary_heap = "0.3.7" +dary_heap = "0.3.8" paste = "1.0.15" rand = "0.9.2" rand_chacha = "0.9.0" seq-macro = "0.3.6" serde = { version = "1.0", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } -zerocopy = { version = "0.8.26", features = ["derive"] } +zerocopy = { version = "0.8.27", features = ["derive"] } [workspace.lints] rust.unknown_lints = "allow" diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index 66a7f525..9c1e7182 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -5,7 +5,7 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.46", features = ["derive", "env"] } +clap = { version = "4.5.48", features = ["derive", "env"] } object = { version = "0.37.3", default-features = false, features = [ "read", "wasm", diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 5157fed0..40e8597e 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -20,7 +20,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.34" +cc = "1.2.38" which = "8.0.0" [lints] From 9c02f219b883a3e25b99103ff2c7ce50da08659f Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 26 Sep 2025 18:58:30 +0800 Subject: [PATCH 227/324] chore: upload 0.5.3 schema scripts (#350) Signed-off-by: usamoi --- sql/install/vchord--0.5.3.sql | 857 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.5.2--0.5.3.sql | 1 + 2 files changed, 858 insertions(+) create mode 100644 sql/install/vchord--0.5.3.sql create mode 100644 sql/upgrade/vchord--0.5.2--0.5.3.sql diff --git a/sql/install/vchord--0.5.3.sql b/sql/install/vchord--0.5.3.sql new file mode 100644 index 00000000..9c428f43 --- /dev/null +++ b/sql/install/vchord--0.5.3.sql @@ -0,0 +1,857 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:46 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:31 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:58 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:76 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:21 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:78 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:80 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:47 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/upgrade/vchord--0.5.2--0.5.3.sql b/sql/upgrade/vchord--0.5.2--0.5.3.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.5.2--0.5.3.sql @@ -0,0 +1 @@ +/* nothing to do */ From 9f65baadad8e30ff9f653bfa8f0044201f04d057 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 26 Sep 2025 20:08:01 +0800 Subject: [PATCH 228/324] readme: sync with docs (#351) https://github.com/tensorchord/pgvecto.rs-docs/pull/145 https://github.com/tensorchord/pgvecto.rs-docs/pull/146 Signed-off-by: usamoi --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 56f13ecf..7a398145 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg17-v0.5.2 + -d ghcr.io/tensorchord/vchord-postgres:pg18-v0.5.3 ``` > [!NOTE] > In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). From d67771896df6b2394f943e9bd2908b1016d9412b Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 9 Oct 2025 11:28:27 +0800 Subject: [PATCH 229/324] fix: check varlena length (#355) job: +check_debian Signed-off-by: usamoi --- .github/workflows/check.yml | 2 -- crates/vector/src/bvect.rs | 2 +- src/datatype/memory_halfvec.rs | 28 +++++++++++++++++++++++++--- src/datatype/memory_scalar8.rs | 24 +++++++++++++++++++++++- src/datatype/memory_vector.rs | 28 +++++++++++++++++++++++++--- 5 files changed, 74 insertions(+), 10 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 611a14ca..42937cd1 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -718,8 +718,6 @@ jobs: CC_${{ matrix.rust_triple }} = "clang" CFLAGS_${{ matrix.rust_triple }} = "--target=${{ matrix.clang_triple }} --sysroot=/sysroot" BINDGEN_EXTRA_CLANG_ARGS_${{ matrix.rust_triple }} = "--sysroot=/sysroot" - [patch.crates-io] - pgrx = { git = "https://github.com/tensorchord/pgrx.git", branch = "big-endian" } EOF rustup toolchain install diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs index 74b76239..90551703 100644 --- a/crates/vector/src/bvect.rs +++ b/crates/vector/src/bvect.rs @@ -210,7 +210,7 @@ impl VectorBorrowed for BVectBorrowed<'_> { fn operator_and(&self, rhs: Self) -> Self::Owned { assert_eq!(self.dims, rhs.dims); - let data = simd::bit::vector_and(self.data, self.data); + let data = simd::bit::vector_and(self.data, rhs.data); BVectOwned::new(self.dims, data) } diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs index 16ad34a5..89d94538 100644 --- a/src/datatype/memory_halfvec.rs +++ b/src/datatype/memory_halfvec.rs @@ -21,7 +21,7 @@ use std::ptr::NonNull; use vector::VectorBorrowed; use vector::vect::VectBorrowed; -#[repr(C, align(8))] +#[repr(C)] struct HalfvecHeader { varlena: u32, dims: u16, @@ -34,7 +34,7 @@ impl HalfvecHeader { if len > 65535 { panic!("vector is too large"); } - (size_of::() + size_of::() * len).next_multiple_of(8) + size_of::() + size_of::() * len } unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f16> { unsafe { @@ -54,6 +54,17 @@ impl HalfvecInput<'_> { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dims = q.byte_add(4).cast::().read(); + assert_eq!(HalfvecHeader::size_of(dims as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } HalfvecInput(q, PhantomData, p != q) } pub fn as_borrowed(&self) -> VectBorrowed<'_, f16> { @@ -78,9 +89,20 @@ impl HalfvecOutput { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dims = q.byte_add(4).cast::().read(); + assert_eq!(HalfvecHeader::size_of(dims as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } Self(q) } - #[allow(dead_code)] + #[expect(dead_code)] pub fn new(vector: VectBorrowed<'_, f16>) -> Self { unsafe { let slice = vector.slice(); diff --git a/src/datatype/memory_scalar8.rs b/src/datatype/memory_scalar8.rs index a57c7db4..f030b865 100644 --- a/src/datatype/memory_scalar8.rs +++ b/src/datatype/memory_scalar8.rs @@ -20,7 +20,7 @@ use std::ptr::NonNull; use vector::VectorBorrowed; use vector::scalar8::Scalar8Borrowed; -#[repr(C, align(8))] +#[repr(C)] struct Scalar8Header { varlena: u32, dims: u16, @@ -63,6 +63,17 @@ impl Scalar8Input<'_> { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dims = q.byte_add(4).cast::().read(); + assert_eq!(Scalar8Header::size_of(dims as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } Scalar8Input(q, PhantomData, p != q) } pub fn as_borrowed(&self) -> Scalar8Borrowed<'_> { @@ -87,6 +98,17 @@ impl Scalar8Output { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dims = q.byte_add(4).cast::().read(); + assert_eq!(Scalar8Header::size_of(dims as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } Self(q) } pub fn new(vector: Scalar8Borrowed<'_>) -> Self { diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs index bbf818d3..1de17baa 100644 --- a/src/datatype/memory_vector.rs +++ b/src/datatype/memory_vector.rs @@ -20,7 +20,7 @@ use std::ptr::NonNull; use vector::VectorBorrowed; use vector::vect::VectBorrowed; -#[repr(C, align(8))] +#[repr(C)] struct VectorHeader { varlena: u32, dims: u16, @@ -33,7 +33,7 @@ impl VectorHeader { if len > 65535 { panic!("vector is too large"); } - (size_of::() + size_of::() * len).next_multiple_of(8) + size_of::() + size_of::() * len } unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f32> { unsafe { @@ -53,6 +53,17 @@ impl VectorInput<'_> { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dims = q.byte_add(4).cast::().read(); + assert_eq!(VectorHeader::size_of(dims as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } VectorInput(q, PhantomData, p != q) } pub fn as_borrowed(&self) -> VectBorrowed<'_, f32> { @@ -77,9 +88,20 @@ impl VectorOutput { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dims = q.byte_add(4).cast::().read(); + assert_eq!(VectorHeader::size_of(dims as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } Self(q) } - #[allow(dead_code)] + #[expect(dead_code)] pub fn new(vector: VectBorrowed<'_, f32>) -> Self { unsafe { let slice = vector.slice(); From 710c4c33304895f868cf2636ce86726fc5d80b91 Mon Sep 17 00:00:00 2001 From: Maxime Schoemans Date: Fri, 10 Oct 2025 09:44:23 +0200 Subject: [PATCH 230/324] Update scripts for ease of use (#354) I tried using the scripts to generate centroids from a dataset and I hit a few issues, so I though I would send a PR with some improvements. The specific changes are detailed in the commit messages. Let me know if you are interested in these changes. I can also revert individual commits if you do not want these specific changes. --------- Signed-off-by: Maxime Schoemans --- scripts/README.md | 12 +++++------ scripts/dump.py | 17 +++++++-------- scripts/index.py | 54 +++++++++++++++++++++++++---------------------- 3 files changed, 43 insertions(+), 40 deletions(-) diff --git a/scripts/README.md b/scripts/README.md index 37119811..41115ab3 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -17,7 +17,7 @@ conda install conda-forge::pgvector-python numpy pytorch::faiss-gpu conda-forge: - If you already have your vectors stored in `PostgreSQL` using `pgvector`, you can export them to a local file by: ```shell - python script/dump.py -n [table name] -c [column name] -d [dim] -o export.hdf5 + python scripts/dump.py -n [table name] -c [column name] -d [dim] -o export.hdf5 --url postgresql://USERNAME:PASSWORD@localhost:5432/DBNAME ``` - If you don't have any data, but would like to give it a try, you can choose one of these datasets: @@ -34,9 +34,9 @@ conda install conda-forge::pgvector-python numpy pytorch::faiss-gpu conda-forge: ```shell # For small dataset size from 1M to 5M - python script/train.py -i [dataset file(export.hdf5)] -o [centroid filename(centroid.npy)] -lists [lists] -m [metric(l2/cos/dot)] + python scripts/train.py -i [dataset file(export.hdf5)] -o [centroid filename(centroid.npy)] --lists [lists] -m [metric(l2/cos/dot)] # For large datasets size, 5M to 100M in size, use GPU and mmap chunks - python script/train.py -i [dataset file(export.hdf5)] -o [centroid filename(centroid.npy)] --lists [lists] -m [metric(l2/cos/dot)] -g --mmap + python scripts/train.py -i [dataset file(export.hdf5)] -o [centroid filename(centroid.npy)] --lists [lists] -m [metric(l2/cos/dot)] -g --mmap ``` `lists` is the number of centroids for clustering, and a typical value for large datasets(>5M) could range from: @@ -48,11 +48,11 @@ conda install conda-forge::pgvector-python numpy pytorch::faiss-gpu conda-forge: 3. To insert vectors and centroids into the database, and then create an index ```shell - python script/index.py -n [table name] -i [dataset file(export.hdf5)] -c [centroid filename(centroid.npy)] -m [metric(l2/cos/dot)] -d [dim] --url postgresql://postgres:123@localhost:5432/postgres + python scripts/index.py -n [table name] -i [dataset file(export.hdf5)] -c [centroid filename(centroid.npy)] -m [metric(l2/cos/dot)] -d [dim] --url postgresql://postgres:123@localhost:5432/postgres ``` 4. Let's start our tour to check the benchmark result of VectorChord ```shell - python script/bench.py -n [table name] -i [dataset file(export.hdf5)] -m [metric(l2/cos/dot)] --nprob 100 --epsilon 1.0 --url postgresql://postgres:123@localhost:5432/postgres - ``` \ No newline at end of file + python scripts/bench.py -n [table name] -i [dataset file(export.hdf5)] -m [metric(l2/cos/dot)] --nprob 100 --epsilon 1.0 --url postgresql://postgres:123@localhost:5432/postgres + ``` diff --git a/scripts/dump.py b/scripts/dump.py index abdcdd3e..68814a99 100644 --- a/scripts/dump.py +++ b/scripts/dump.py @@ -23,17 +23,19 @@ def build_arg_parse(): parser = argparse.ArgumentParser(description="Dump embeddings to a local file") - parser.add_argument("-n", "--name", help="table name", required=True) parser.add_argument( - "-p", "--password", help="Database password", default="password" + "--url", + help="url, like `postgresql://USERNAME:PASSWORD@localhost:5432/DBNAME`", + default="postgresql://postgres:postgres@localhost:5432/postgres" ) - parser.add_argument("-o", "--output", help="Output filepath", required=True) + parser.add_argument("-n", "--name", help="Table name", required=True) parser.add_argument("-c", "--column", help="Column name", default="embedding") parser.add_argument("-d", "--dim", help="Dimension", type=int, required=True) + parser.add_argument("-o", "--output", help="Output filepath", required=True) return parser -def create_connection(password): +def create_connection(url): keepalive_kwargs = { "keepalives": 1, "keepalives_idle": 30, @@ -41,13 +43,10 @@ def create_connection(password): "keepalives_count": 5, } conn = psycopg.connect( - conninfo=f"postgresql://postgres:{password}@localhost:5432/postgres", - dbname="postgres", + conninfo=url, autocommit=True, **keepalive_kwargs, ) - conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") register_vector(conn) return conn @@ -78,5 +77,5 @@ def write_to_h5(filepath, vecs, dim): args = parser.parse_args() print(args) - conn = create_connection(args.password) + conn = create_connection(args.url) write_to_h5(args.output, extract_vectors(conn, args.name, args.column), args.dim) diff --git a/scripts/index.py b/scripts/index.py index db238f9b..39aa4710 100644 --- a/scripts/index.py +++ b/scripts/index.py @@ -44,7 +44,7 @@ def build_arg_parse(): choices=["l2", "cos", "dot"], ) parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) - parser.add_argument("-i", "--input", help="Input filepath", required=True) + parser.add_argument("-i", "--input", help="Input filepath", required=False) parser.add_argument( "--url", help="url, like `postgresql://postgres:123@localhost:5432/postgres`", required=True ) @@ -70,6 +70,8 @@ def build_arg_parse(): ) # Internal build parser.add_argument("--lists", help="Number of centroids", type=int, required=False) + # Do not build index + parser.add_argument("--noindex", help="Do not build index", action='store_true', required=False) return parser @@ -119,12 +121,10 @@ def get_ivf_ops_config(metric, workers, k=None, name=None): async def create_connection(url): conn = await psycopg.AsyncConnection.connect( conninfo=url, - dbname="postgres", autocommit=True, **KEEPALIVE_KWARGS, ) await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - await conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") await register_vector_async(conn) return conn @@ -176,6 +176,7 @@ async def add_embeddings(conn, name, dim, train, chunks, workers): async def build_index( conn, name, workers, metric_ops, ivf_config, finish: asyncio.Event ): + await conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") start_time = perf_counter() await conn.execute(f"SET max_parallel_maintenance_workers TO {workers}") await conn.execute(f"SET max_parallel_workers TO {workers}") @@ -208,32 +209,35 @@ async def monitor_index_build(conn, finish: asyncio.Event): async def main(dataset): - dataset = h5py.File(Path(args.input), "r") conn = await create_connection(args.url) if args.centroids: centroids = np.load(args.centroids, allow_pickle=False) await add_centroids(conn, args.name, centroids) - metric_ops, ivf_config = get_ivf_ops_config( - args.metric, args.workers, args.lists, args.name if args.centroids else None - ) - await add_embeddings(conn, args.name, args.dim, dataset["train"], args.chunks, args.workers) - - index_finish = asyncio.Event() - # Need a separate connection for monitor process - monitor_conn = await create_connection(args.url) - monitor_task = monitor_index_build( - monitor_conn, - index_finish, - ) - index_task = build_index( - conn, - args.name, - args.workers, - metric_ops, - ivf_config, - index_finish, - ) - await asyncio.gather(index_task, monitor_task) + + if args.input: + dataset = h5py.File(Path(args.input), "r") + await add_embeddings(conn, args.name, args.dim, dataset["train"], args.chunks, args.workers) + + if not args.noindex: + metric_ops, ivf_config = get_ivf_ops_config( + args.metric, args.workers, args.lists, args.name if args.centroids else None + ) + index_finish = asyncio.Event() + # Need a separate connection for monitor process + monitor_conn = await create_connection(args.url) + monitor_task = monitor_index_build( + monitor_conn, + index_finish, + ) + index_task = build_index( + conn, + args.name, + args.workers, + metric_ops, + ivf_config, + index_finish, + ) + await asyncio.gather(index_task, monitor_task) if __name__ == "__main__": From 94526f7b384fc5c7342130fbd9351da5f0937968 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 10 Oct 2025 16:57:30 +0800 Subject: [PATCH 231/324] chore: code cleanup (#356) Signed-off-by: usamoi --- .github/workflows/check.yml | 10 +- Cargo.lock | 178 ++++---- Cargo.toml | 6 +- crates/algo/src/lib.rs | 390 ------------------ crates/{algo => index}/Cargo.toml | 2 +- crates/{algo => index}/src/accessor.rs | 0 crates/index/src/bump.rs | 32 ++ crates/index/src/fetch.rs | 128 ++++++ crates/index/src/lib.rs | 21 + crates/index/src/packed.rs | 61 +++ crates/{algo => index}/src/prefetcher.rs | 91 +++- crates/index/src/relation.rs | 136 ++++++ crates/{algo => index}/src/tuples.rs | 0 crates/simd/Cargo.toml | 4 +- crates/simd/build.rs | 12 +- crates/simd/src/lib.rs | 5 +- crates/vchordg/Cargo.toml | 2 +- crates/vchordg/src/build.rs | 2 +- crates/vchordg/src/bulkdelete.rs | 2 +- crates/vchordg/src/candidates.rs | 5 +- crates/vchordg/src/insert.rs | 7 +- crates/vchordg/src/lib.rs | 4 +- crates/vchordg/src/maintain.rs | 2 +- crates/vchordg/src/operator.rs | 2 +- crates/vchordg/src/prewarm.rs | 2 +- crates/vchordg/src/search.rs | 7 +- crates/vchordg/src/tuples.rs | 2 +- crates/vchordg/src/vectors.rs | 4 +- crates/vchordrq/Cargo.toml | 2 +- crates/vchordrq/src/build.rs | 2 +- crates/vchordrq/src/bulkdelete.rs | 6 +- crates/vchordrq/src/cache.rs | 11 +- crates/vchordrq/src/centroids.rs | 5 +- .../vchordrq/src/closure_lifetime_binder.rs | 6 +- crates/vchordrq/src/cost.rs | 2 +- crates/vchordrq/src/fast_heap.rs | 2 +- crates/vchordrq/src/freepages.rs | 4 +- crates/vchordrq/src/insert.rs | 14 +- crates/vchordrq/src/lib.rs | 11 +- crates/vchordrq/src/maintain.rs | 23 +- crates/vchordrq/src/operator.rs | 2 +- crates/vchordrq/src/prewarm.rs | 8 +- crates/vchordrq/src/rerank.rs | 16 +- crates/vchordrq/src/search.rs | 11 +- crates/vchordrq/src/tape.rs | 8 +- crates/vchordrq/src/tape_writer.rs | 2 +- crates/vchordrq/src/tuples.rs | 2 +- crates/vchordrq/src/vectors.rs | 6 +- src/bin/pgrx_embed.rs | 1 - src/datatype/mod.rs | 12 +- src/index/functions.rs | 4 +- src/index/mod.rs | 18 +- src/index/scanners.rs | 3 +- src/index/storage.rs | 22 +- src/index/vchordg/am/am_build.rs | 8 +- src/index/vchordg/am/mod.rs | 8 +- src/index/vchordg/{algo.rs => dispatch.rs} | 24 +- src/index/vchordg/mod.rs | 4 +- src/index/vchordg/scanners/default.rs | 15 +- src/index/vchordrq/am/am_build.rs | 73 ++-- src/index/vchordrq/am/mod.rs | 8 +- src/index/vchordrq/build.rs | 44 ++ src/index/vchordrq/{algo.rs => dispatch.rs} | 38 +- src/index/vchordrq/filter.rs | 2 +- src/index/vchordrq/mod.rs | 7 +- src/index/vchordrq/scanners/default.rs | 49 +-- src/index/vchordrq/scanners/maxsim.rs | 28 +- src/lib.rs | 3 +- 68 files changed, 864 insertions(+), 767 deletions(-) delete mode 100644 crates/algo/src/lib.rs rename crates/{algo => index}/Cargo.toml (96%) rename crates/{algo => index}/src/accessor.rs (100%) create mode 100644 crates/index/src/bump.rs create mode 100644 crates/index/src/fetch.rs create mode 100644 crates/index/src/lib.rs create mode 100644 crates/index/src/packed.rs rename crates/{algo => index}/src/prefetcher.rs (80%) create mode 100644 crates/index/src/relation.rs rename crates/{algo => index}/src/tuples.rs (100%) rename src/index/vchordg/{algo.rs => dispatch.rs} (94%) create mode 100644 src/index/vchordrq/build.rs rename src/index/vchordrq/{algo.rs => dispatch.rs} (93%) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 42937cd1..ac317acf 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -214,7 +214,7 @@ jobs: sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo systemctl stop postgresql - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -297,7 +297,7 @@ jobs: make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config install - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -404,7 +404,7 @@ jobs: nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" install Pop-Location - Invoke-WebRequest -Uri https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-${{ matrix.arch }}-pc-windows-msvc.zip -OutFile "$env:TEMP\sqllogictest-install.zip" + Invoke-WebRequest -Uri https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-${{ matrix.arch }}-pc-windows-msvc.zip -OutFile "$env:TEMP\sqllogictest-install.zip" Expand-Archive -Path "$env:TEMP\sqllogictest-install.zip" -DestinationPath "D:\sqllogictest-install" -Force Add-Content -Path $env:GITHUB_PATH -Value "D:\sqllogictest-install" @@ -503,7 +503,7 @@ jobs: make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config install - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -686,7 +686,7 @@ jobs: make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" sudo make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.26.4/sqllogictest-bin-v0.26.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 diff --git a/Cargo.lock b/Cargo.lock index 28b37649..5bad0f56 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,26 +11,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "algo" -version = "0.0.0" -dependencies = [ - "always_equal", - "bumpalo", - "dary_heap", - "distance", - "simd", - "small_iter", - "vector", - "zerocopy", -] - -[[package]] -name = "allocator-api2" -version = "0.2.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" - [[package]] name = "always_equal" version = "0.0.0" @@ -47,9 +27,9 @@ dependencies = [ [[package]] name = "anstream" -version = "0.6.20" +version = "0.6.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ae563653d1938f79b1ab1b5e668c87c76a9930414574a6583a7b7e11a8e6192" +checksum = "43d5b281e737544384e969a5ccad3f1cdd24b48086a0fc1b2a5262a26b8f4f4a" dependencies = [ "anstyle", "anstyle-parse", @@ -62,9 +42,9 @@ dependencies = [ [[package]] name = "anstyle" -version = "1.0.11" +version = "1.0.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "862ed96ca487e809f1c8e5a8447f6ee2cf102f846893800b20cebdf541fc6bbd" +checksum = "5192cca8006f1fd4f7237516f40fa183bb07f8fbdfedaa0036de5ea9b0b45e78" [[package]] name = "anstyle-parse" @@ -156,9 +136,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.38" +version = "1.2.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "80f41ae168f955c12fb8960b057d70d0ca153fb83182b57d86380443527be7e9" +checksum = "e1d05d92f4b1fd76aad469d46cdd858ca761576082cd37df81416691e50199fb" dependencies = [ "find-msvc-tools", "shlex", @@ -408,7 +388,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.1", + "windows-sys 0.61.2", ] [[package]] @@ -435,9 +415,9 @@ checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" [[package]] name = "find-msvc-tools" -version = "0.1.2" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ced73b1dacfc750a6db6c0a0c3a3853c8b41997e2e2c563dc90804ae6867959" +checksum = "0399f9d26e5191ce32c498bebd31e7a3ceabc2745f0ac54af3f335126c3f24b3" [[package]] name = "fixedbitset" @@ -498,9 +478,9 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403" [[package]] name = "half" -version = "2.6.0" +version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "459196ed295495a68f7d7fe1d84f6c4b7ff0e21fe3017b2f283c6fac3ad803c9" +checksum = "e54c115d4f30f52c67202f079c5f9d8b49db4691f460fdb0b4c2e838261b2ba5" dependencies = [ "cfg-if", "crunchy", @@ -513,8 +493,6 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ - "allocator-api2", - "equivalent", "foldhash", ] @@ -658,6 +636,20 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "964de6e86d545b246d84badc0fef527924ace5134f30641c203ef52ba83f58d5" +[[package]] +name = "index" +version = "0.0.0" +dependencies = [ + "always_equal", + "bumpalo", + "dary_heap", + "distance", + "simd", + "small_iter", + "vector", + "zerocopy", +] + [[package]] name = "indexmap" version = "2.11.4" @@ -717,9 +709,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.176" +version = "0.2.177" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "58f929b4d672ea937a23a1ab494143d968337a5f47e56d0815df1e0890ddf174" +checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976" [[package]] name = "libloading" @@ -841,9 +833,9 @@ checksum = "a4895175b425cb1f87721b59f0f286c2092bd4af812243672510e1ac53e2e0ad" [[package]] name = "owo-colors" -version = "4.2.2" +version = "4.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48dd4f4a2c8405440fd0462561f0e5806bd0f77e86f51c761481bdd4018b545e" +checksum = "9c6901729fa79e91a0913333229e9ca5dc725089d1c363b2f4b4760709dc4a52" dependencies = [ "supports-color", ] @@ -872,9 +864,9 @@ checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" [[package]] name = "petgraph" -version = "0.8.2" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54acf3a685220b533e437e264e4d932cfbdc4cc7ec0cd232ed73c08d03b8a7ca" +checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" dependencies = [ "fixedbitset", "hashbrown 0.15.5", @@ -1071,9 +1063,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.40" +version = "1.0.41" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" +checksum = "ce25767e7b499d1b604768e7cde645d14cc8584231ea6b295e9c9eb22c02e1d1" dependencies = [ "proc-macro2", ] @@ -1208,7 +1200,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.1", + "windows-sys 0.61.2", ] [[package]] @@ -1246,9 +1238,9 @@ checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" -version = "1.0.227" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "80ece43fc6fbed4eb5392ab50c07334d3e577cbf40997ee896fe7af40bba4245" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" dependencies = [ "serde_core", "serde_derive", @@ -1266,18 +1258,18 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.227" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a576275b607a2c86ea29e410193df32bc680303c82f31e275bbfcafe8b33be5" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.227" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51e694923b8824cf0e9b382adf0f60d4e05f348f357b38833a3fa5ed7c2ede04" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", @@ -1299,9 +1291,9 @@ dependencies = [ [[package]] name = "serde_spanned" -version = "1.0.2" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5417783452c2be558477e104686f7de5dae53dba813c28435e0e70f82d9b04ee" +checksum = "e24345aa0fe688594e73770a5f6d1b216508b4f93484c0026d521acd30134392" dependencies = [ "serde_core", ] @@ -1317,7 +1309,7 @@ name = "simd" version = "0.0.0" dependencies = [ "cc", - "half 2.6.0", + "half 2.7.0", "rand", "seq-macro", "simd_macros", @@ -1346,9 +1338,9 @@ checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" [[package]] name = "stable_deref_trait" -version = "1.2.0" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "strsim" @@ -1401,18 +1393,18 @@ checksum = "1ac9aa371f599d22256307c24a9d748c041e548cbf599f35d890f9d365361790" [[package]] name = "thiserror" -version = "2.0.16" +version = "2.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3467d614147380f2e4e374161426ff399c91084acd2363eaf549172b3d5e60c0" +checksum = "f63587ca0f12b72a0600bcba1d40081f830876000bb46dd2337a3051618f4fc8" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.16" +version = "2.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c5e1be1c48b9172ee610da68fd9cd2770e7a4056cb3fc98710ee6906f0c7960" +checksum = "3ff15c8ecd7de3849db632e14d18d2571fa09dfc5ed93479bc4485c7a517c913" dependencies = [ "proc-macro2", "quote", @@ -1431,9 +1423,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.9.7" +version = "0.9.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00e5e5d9bf2475ac9d4f0d9edab68cc573dc2fd644b0dba36b0c30a92dd9eaa0" +checksum = "f0dc8b1fb61449e27716ec0e1bdf0f6b8f3e8f6b05391e8497b8b6d7804ea6d8" dependencies = [ "indexmap", "serde_core", @@ -1446,27 +1438,27 @@ dependencies = [ [[package]] name = "toml_datetime" -version = "0.7.2" +version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32f1085dec27c2b6632b04c80b3bb1b4300d6495d1e129693bdda7d91e72eec1" +checksum = "f2cdb639ebbc97961c51720f858597f7f24c4fc295327923af55b74c3c724533" dependencies = [ "serde_core", ] [[package]] name = "toml_parser" -version = "1.0.3" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cf893c33be71572e0e9aa6dd15e6677937abd686b066eac3f8cd3531688a627" +checksum = "c0cbe268d35bdb4bb5a56a2de88d0ad0eb70af5384a99d648cd4b3d04039800e" dependencies = [ "winnow", ] [[package]] name = "toml_writer" -version = "1.0.3" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d163a63c116ce562a22cda521fcc4d79152e7aba014456fb5eb442f6d6a10109" +checksum = "df8b2b54733674ad286d16267dcfc7a71ed5c776e4ac7aa3c3e2561f7c637bf2" [[package]] name = "unescape" @@ -1488,9 +1480,9 @@ checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" [[package]] name = "unicode-width" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" [[package]] name = "url" @@ -1561,11 +1553,11 @@ dependencies = [ name = "vchord" version = "0.0.0" dependencies = [ - "algo", "always_equal", "bumpalo", "dary_heap", "distance", + "index", "k_means", "mimalloc", "paste", @@ -1590,9 +1582,9 @@ dependencies = [ name = "vchordg" version = "0.0.0" dependencies = [ - "algo", "always_equal", "distance", + "index", "min-max-heap", "rabitq", "rand", @@ -1607,9 +1599,9 @@ dependencies = [ name = "vchordrq" version = "0.0.0" dependencies = [ - "algo", "always_equal", "distance", + "index", "pin-project", "rabitq", "rand", @@ -1763,7 +1755,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.1", + "windows-sys 0.61.2", ] [[package]] @@ -1774,9 +1766,9 @@ checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" [[package]] name = "windows-link" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45e46c0661abb7180e7b9c281db115305d49ca1709ab8242adf09666d2173c65" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" [[package]] name = "windows-sys" @@ -1789,18 +1781,18 @@ dependencies = [ [[package]] name = "windows-sys" -version = "0.61.1" +version = "0.61.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f109e41dd4a3c848907eb83d5a42ea98b3769495597450cf6d153507b166f0f" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" dependencies = [ "windows-link", ] [[package]] name = "windows-targets" -version = "0.53.4" +version = "0.53.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d42b7b7f66d2a06854650af09cfdf8713e427a439c97ad65a6375318033ac4b" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" dependencies = [ "windows-link", "windows_aarch64_gnullvm", @@ -1815,51 +1807,51 @@ dependencies = [ [[package]] name = "windows_aarch64_gnullvm" -version = "0.53.0" +version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "86b8d5f90ddd19cb4a147a5fa63ca848db3df085e25fee3cc10b39b6eebae764" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" [[package]] name = "windows_aarch64_msvc" -version = "0.53.0" +version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7651a1f62a11b8cbd5e0d42526e55f2c99886c77e007179efff86c2b137e66c" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" [[package]] name = "windows_i686_gnu" -version = "0.53.0" +version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1dc67659d35f387f5f6c479dc4e28f1d4bb90ddd1a5d3da2e5d97b42d6272c3" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" [[package]] name = "windows_i686_gnullvm" -version = "0.53.0" +version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ce6ccbdedbf6d6354471319e781c0dfef054c81fbc7cf83f338a4296c0cae11" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" [[package]] name = "windows_i686_msvc" -version = "0.53.0" +version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "581fee95406bb13382d2f65cd4a908ca7b1e4c2f1917f143ba16efe98a589b5d" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" [[package]] name = "windows_x86_64_gnu" -version = "0.53.0" +version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e55b5ac9ea33f2fc1716d1742db15574fd6fc8dadc51caab1c16a3d3b4190ba" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" [[package]] name = "windows_x86_64_gnullvm" -version = "0.53.0" +version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0a6e035dd0599267ce1ee132e51c27dd29437f63325753051e71dd9e42406c57" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" [[package]] name = "windows_x86_64_msvc" -version = "0.53.0" +version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" [[package]] name = "winnow" diff --git a/Cargo.toml b/Cargo.toml index adc8fc2a..8fd543c4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,9 +22,9 @@ pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] pg18 = ["pgrx/pg18", "pgrx-catalog/pg18"] [dependencies] -algo = { path = "./crates/algo" } always_equal = { path = "./crates/always_equal" } distance = { path = "./crates/distance" } +index = { path = "./crates/index" } k_means = { path = "./crates/k_means" } rabitq = { path = "./crates/rabitq" } simd = { path = "./crates/simd" } @@ -42,7 +42,7 @@ rand.workspace = true rusqlite = { version = "0.37.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true -toml = "0.9.7" +toml = "0.9.8" validator.workspace = true zerocopy.workspace = true @@ -72,7 +72,6 @@ validator = { version = "0.20.0", features = ["derive"] } zerocopy = { version = "0.8.27", features = ["derive"] } [workspace.lints] -rust.unknown_lints = "allow" # complexity clippy.identity_op = "allow" clippy.int_plus_one = "allow" @@ -87,7 +86,6 @@ clippy.needless_range_loop = "allow" rust.unsafe_code = "deny" rust.unsafe_op_in_unsafe_fn = "deny" # unused -rust.unused_crate_dependencies = "allow" rust.unused_extern_crates = "warn" rust.unused_import_braces = "warn" rust.unused_lifetimes = "warn" diff --git a/crates/algo/src/lib.rs b/crates/algo/src/lib.rs deleted file mode 100644 index 995afd88..00000000 --- a/crates/algo/src/lib.rs +++ /dev/null @@ -1,390 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -pub mod accessor; -pub mod prefetcher; -pub mod tuples; - -use always_equal::AlwaysEqual; -use dary_heap::DaryHeap; -use std::collections::{BinaryHeap, VecDeque}; -use std::iter::Peekable; -use std::ops::{Deref, DerefMut}; -use zerocopy::{FromBytes, IntoBytes}; - -pub trait Page: Sized + 'static { - type Opaque: Opaque; - - #[must_use] - fn get_opaque(&self) -> &Self::Opaque; - #[must_use] - fn get_opaque_mut(&mut self) -> &mut Self::Opaque; - #[must_use] - fn len(&self) -> u16; - #[must_use] - #[inline] - fn is_empty(&self) -> bool { - self.len() == 0 - } - #[must_use] - fn get(&self, i: u16) -> Option<&[u8]>; - #[must_use] - fn get_mut(&mut self, i: u16) -> Option<&mut [u8]>; - #[must_use] - fn alloc(&mut self, data: &[u8]) -> Option; - fn free(&mut self, i: u16); - #[must_use] - fn freespace(&self) -> u16; - fn clear(&mut self, opaque: Self::Opaque); -} - -pub trait PageGuard { - fn id(&self) -> u32; -} - -pub trait ReadStream<'b> { - type Relation: RelationReadTypes; - type Guards: ExactSizeIterator::ReadGuard<'b>>; - type Item; - type Inner: Iterator; - fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; - fn next_if bool>( - &mut self, - predicate: P, - ) -> Option<(Self::Item, Self::Guards)>; - fn into_inner(self) -> Self::Inner; -} - -pub trait Relation { - type Page: Page; -} - -pub trait RelationReadTypes: Relation { - type ReadGuard<'b>: PageGuard + Deref; -} - -pub trait RelationRead: RelationReadTypes { - fn read(&self, id: u32) -> Self::ReadGuard<'_>; -} - -pub trait RelationWriteTypes: Relation { - type WriteGuard<'b>: PageGuard + DerefMut; -} - -pub trait RelationWrite: RelationWriteTypes { - fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_>; - fn extend( - &self, - opaque: ::Opaque, - tracking_freespace: bool, - ) -> Self::WriteGuard<'_>; - fn search(&self, freespace: usize) -> Option>; -} - -pub trait RelationPrefetch: Relation { - fn prefetch(&self, id: u32); -} - -#[derive(Debug, Default, Clone)] -pub struct Hints { - pub full: bool, -} - -impl Hints { - #[allow(clippy::needless_update)] - #[inline] - pub fn full(self, full: bool) -> Self { - Self { full, ..self } - } -} - -pub trait RelationReadStreamTypes: RelationReadTypes { - type ReadStream<'b, I: Iterator>: ReadStream<'b, Item = I::Item, Relation = Self> - where - I::Item: Fetch<'b>; -} - -pub trait RelationReadStream: RelationReadStreamTypes { - fn read_stream<'b, I: Iterator>(&'b self, iter: I, hints: Hints) -> Self::ReadStream<'b, I> - where - I::Item: Fetch<'b>; -} - -#[derive(Debug, Clone, Copy)] -pub enum RerankMethod { - Index, - Heap, -} - -pub trait Bump: 'static { - #[allow(clippy::mut_from_ref)] - fn alloc(&self, value: T) -> &mut T; - #[allow(clippy::mut_from_ref)] - fn alloc_slice(&self, slice: &[T]) -> &mut [T]; -} - -impl Bump for bumpalo::Bump { - #[inline] - fn alloc(&self, value: T) -> &mut T { - self.alloc(value) - } - - #[inline] - fn alloc_slice(&self, slice: &[T]) -> &mut [T] { - self.alloc_slice_copy(slice) - } -} - -pub type BorrowedIter<'b> = small_iter::borrowed::Iter<'b, u32, 1>; - -pub trait Fetch<'b> { - type Iter: ExactSizeIterator + 'b; - #[must_use] - fn fetch(&self) -> Self::Iter; -} - -impl Fetch<'_> for u32 { - type Iter = std::iter::Once; - #[inline(always)] - fn fetch(&self) -> std::iter::Once { - std::iter::once(*self) - } -} - -impl<'b, T, A, B, W: 'b + PackedRefMut)>> Fetch<'b> - for (T, AlwaysEqual) -{ - type Iter = BorrowedIter<'b>; - #[inline(always)] - fn fetch(&self) -> BorrowedIter<'b> { - let (.., list) = self.1.0.get(); - *list - } -} - -impl<'b, T, A, B, C> Fetch<'b> for (T, AlwaysEqual<&mut (A, B, C, BorrowedIter<'b>)>) { - type Iter = BorrowedIter<'b>; - #[inline(always)] - fn fetch(&self) -> BorrowedIter<'b> { - let (_, AlwaysEqual((.., list))) = self; - *list - } -} - -impl Fetch<'_> for (T, AlwaysEqual<(u32, u16)>) { - type Iter = std::iter::Once; - #[inline(always)] - fn fetch(&self) -> std::iter::Once { - let (_, AlwaysEqual((x, _))) = self; - std::iter::once(*x) - } -} - -impl Fetch<'_> for (u32, u16) { - type Iter = std::iter::Once; - #[inline(always)] - fn fetch(&self) -> std::iter::Once { - std::iter::once(self.0) - } -} - -impl Fetch<'_> for (T, AlwaysEqual<((u32, u16), (u32, u16))>) { - type Iter = std::iter::Once; - #[inline(always)] - fn fetch(&self) -> std::iter::Once { - let (_, AlwaysEqual(((x, _), _))) = self; - std::iter::once(*x) - } -} - -pub trait Fetch1 { - fn fetch_1(&self) -> u32; -} - -#[repr(transparent)] -pub struct Fetch1Iter { - iter: I, -} - -impl Iterator for Fetch1Iter -where - I::Item: Fetch1, -{ - type Item = u32; - - #[inline(always)] - fn next(&mut self) -> Option { - self.iter.next().map(|x| x.fetch_1()) - } - - #[inline(always)] - fn size_hint(&self) -> (usize, Option) { - self.iter.size_hint() - } -} - -impl ExactSizeIterator for Fetch1Iter where I::Item: Fetch1 {} - -impl<'b, T, F: Fetch1 + Copy> Fetch<'b> for (T, AlwaysEqual<&'b [F]>) { - type Iter = Fetch1Iter>>; - #[inline(always)] - fn fetch(&self) -> Self::Iter { - Fetch1Iter { - iter: self.1.0.iter().copied(), - } - } -} - -impl<'b, T, U, F: Fetch1 + Copy> Fetch<'b> for (T, AlwaysEqual<(&'b [F], U)>) { - type Iter = Fetch1Iter>>; - #[inline(always)] - fn fetch(&self) -> Self::Iter { - Fetch1Iter { - iter: self.1.0.0.iter().copied(), - } - } -} - -pub trait Sequence { - type Item; - type Inner: Iterator; - #[must_use] - fn next(&mut self) -> Option; - #[must_use] - fn peek(&mut self) -> Option<&Self::Item>; - #[must_use] - fn into_inner(self) -> Self::Inner; -} - -impl Sequence for BinaryHeap { - type Item = T; - type Inner = std::vec::IntoIter; - #[inline] - fn next(&mut self) -> Option { - self.pop() - } - #[inline] - fn peek(&mut self) -> Option<&T> { - (self as &Self).peek() - } - #[inline] - fn into_inner(self) -> Self::Inner { - self.into_vec().into_iter() - } -} - -impl Sequence for DaryHeap { - type Item = T; - type Inner = std::vec::IntoIter; - #[inline] - fn next(&mut self) -> Option { - self.pop() - } - #[inline] - fn peek(&mut self) -> Option<&T> { - (self as &Self).peek() - } - #[inline] - fn into_inner(self) -> Self::Inner { - self.into_vec().into_iter() - } -} - -impl Sequence for Peekable { - type Item = I::Item; - type Inner = Peekable; - #[inline] - fn next(&mut self) -> Option { - Iterator::next(self) - } - #[inline] - fn peek(&mut self) -> Option<&I::Item> { - self.peek() - } - #[inline] - fn into_inner(self) -> Self::Inner { - self - } -} - -impl Sequence for VecDeque { - type Item = T; - type Inner = std::collections::vec_deque::IntoIter; - #[inline] - fn next(&mut self) -> Option { - self.pop_front() - } - #[inline] - fn peek(&mut self) -> Option<&T> { - self.front() - } - #[inline] - fn into_inner(self) -> Self::Inner { - self.into_iter() - } -} - -/// # Safety -/// -/// * `Opaque` must aligned to 8 bytes. -#[allow(unsafe_code)] -pub unsafe trait Opaque: Copy + Send + Sync + FromBytes + IntoBytes + 'static {} - -pub trait PackedRefMut { - type T; - fn get(&self) -> &Self::T; - fn get_mut(&mut self) -> &mut Self::T; -} - -impl PackedRefMut for &mut T { - type T = T; - #[inline(always)] - fn get(&self) -> &T { - self - } - #[inline(always)] - fn get_mut(&mut self) -> &mut T { - self - } -} - -#[repr(Rust, packed(4))] -pub struct PackedRefMut4<'b, T>(pub &'b mut T); - -impl<'b, T> PackedRefMut for PackedRefMut4<'b, T> { - type T = T; - #[inline(always)] - fn get(&self) -> &T { - self.0 - } - #[inline(always)] - fn get_mut(&mut self) -> &mut T { - self.0 - } -} - -#[repr(Rust, packed(8))] -pub struct PackedRefMut8<'b, T>(pub &'b mut T); - -impl<'a, T> PackedRefMut for PackedRefMut8<'a, T> { - type T = T; - #[inline(always)] - fn get(&self) -> &T { - self.0 - } - #[inline(always)] - fn get_mut(&mut self) -> &mut T { - self.0 - } -} diff --git a/crates/algo/Cargo.toml b/crates/index/Cargo.toml similarity index 96% rename from crates/algo/Cargo.toml rename to crates/index/Cargo.toml index 3f11dd14..f6fe282e 100644 --- a/crates/algo/Cargo.toml +++ b/crates/index/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "algo" +name = "index" version.workspace = true edition.workspace = true publish = false diff --git a/crates/algo/src/accessor.rs b/crates/index/src/accessor.rs similarity index 100% rename from crates/algo/src/accessor.rs rename to crates/index/src/accessor.rs diff --git a/crates/index/src/bump.rs b/crates/index/src/bump.rs new file mode 100644 index 00000000..e80e242e --- /dev/null +++ b/crates/index/src/bump.rs @@ -0,0 +1,32 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub trait Bump: 'static { + #[allow(clippy::mut_from_ref)] + fn alloc(&self, value: T) -> &mut T; + #[allow(clippy::mut_from_ref)] + fn alloc_slice(&self, slice: &[T]) -> &mut [T]; +} + +impl Bump for bumpalo::Bump { + #[inline] + fn alloc(&self, value: T) -> &mut T { + self.alloc(value) + } + + #[inline] + fn alloc_slice(&self, slice: &[T]) -> &mut [T] { + self.alloc_slice_copy(slice) + } +} diff --git a/crates/index/src/fetch.rs b/crates/index/src/fetch.rs new file mode 100644 index 00000000..685c90dc --- /dev/null +++ b/crates/index/src/fetch.rs @@ -0,0 +1,128 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::packed::PackedRefMut; +use always_equal::AlwaysEqual; + +pub type BorrowedIter<'b> = small_iter::borrowed::Iter<'b, u32, 1>; + +pub trait Fetch<'b> { + type Iter: ExactSizeIterator + 'b; + #[must_use] + fn fetch(&self) -> Self::Iter; +} + +impl Fetch<'_> for u32 { + type Iter = std::iter::Once; + #[inline(always)] + fn fetch(&self) -> std::iter::Once { + std::iter::once(*self) + } +} + +impl<'b, T, A, B, W: 'b + PackedRefMut)>> Fetch<'b> + for (T, AlwaysEqual) +{ + type Iter = BorrowedIter<'b>; + #[inline(always)] + fn fetch(&self) -> BorrowedIter<'b> { + let (.., list) = self.1.0.get(); + *list + } +} + +impl<'b, T, A, B, C> Fetch<'b> for (T, AlwaysEqual<&mut (A, B, C, BorrowedIter<'b>)>) { + type Iter = BorrowedIter<'b>; + #[inline(always)] + fn fetch(&self) -> BorrowedIter<'b> { + let (_, AlwaysEqual((.., list))) = self; + *list + } +} + +impl Fetch<'_> for (T, AlwaysEqual<(u32, u16)>) { + type Iter = std::iter::Once; + #[inline(always)] + fn fetch(&self) -> std::iter::Once { + let (_, AlwaysEqual((x, _))) = self; + std::iter::once(*x) + } +} + +impl Fetch<'_> for (u32, u16) { + type Iter = std::iter::Once; + #[inline(always)] + fn fetch(&self) -> std::iter::Once { + std::iter::once(self.0) + } +} + +impl Fetch<'_> for (T, AlwaysEqual<((u32, u16), (u32, u16))>) { + type Iter = std::iter::Once; + #[inline(always)] + fn fetch(&self) -> std::iter::Once { + let (_, AlwaysEqual(((x, _), _))) = self; + std::iter::once(*x) + } +} + +pub trait Fetch1 { + fn fetch_1(&self) -> u32; +} + +#[repr(transparent)] +pub struct Fetch1Iter { + iter: I, +} + +impl Iterator for Fetch1Iter +where + I::Item: Fetch1, +{ + type Item = u32; + + #[inline(always)] + fn next(&mut self) -> Option { + self.iter.next().map(|x| x.fetch_1()) + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + self.iter.size_hint() + } +} + +impl ExactSizeIterator for Fetch1Iter where I::Item: Fetch1 {} + +impl<'b, T, F: Fetch1 + Copy> Fetch<'b> for (T, AlwaysEqual<&'b [F]>) { + type Iter = Fetch1Iter>>; + + #[inline(always)] + fn fetch(&self) -> Self::Iter { + Fetch1Iter { + iter: self.1.0.iter().copied(), + } + } +} + +impl<'b, T, U, F: Fetch1 + Copy> Fetch<'b> for (T, AlwaysEqual<(&'b [F], U)>) { + type Iter = Fetch1Iter>>; + + #[inline(always)] + fn fetch(&self) -> Self::Iter { + Fetch1Iter { + iter: self.1.0.0.iter().copied(), + } + } +} diff --git a/crates/index/src/lib.rs b/crates/index/src/lib.rs new file mode 100644 index 00000000..1ba2ed4e --- /dev/null +++ b/crates/index/src/lib.rs @@ -0,0 +1,21 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub mod accessor; +pub mod bump; +pub mod fetch; +pub mod packed; +pub mod prefetcher; +pub mod relation; +pub mod tuples; diff --git a/crates/index/src/packed.rs b/crates/index/src/packed.rs new file mode 100644 index 00000000..65d93db7 --- /dev/null +++ b/crates/index/src/packed.rs @@ -0,0 +1,61 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub trait PackedRefMut { + type T; + fn get(&self) -> &Self::T; + fn get_mut(&mut self) -> &mut Self::T; +} + +impl PackedRefMut for &mut T { + type T = T; + #[inline(always)] + fn get(&self) -> &T { + self + } + #[inline(always)] + fn get_mut(&mut self) -> &mut T { + self + } +} + +#[repr(Rust, packed(4))] +pub struct PackedRefMut4<'b, T>(pub &'b mut T); + +impl<'b, T> PackedRefMut for PackedRefMut4<'b, T> { + type T = T; + #[inline(always)] + fn get(&self) -> &T { + self.0 + } + #[inline(always)] + fn get_mut(&mut self) -> &mut T { + self.0 + } +} + +#[repr(Rust, packed(8))] +pub struct PackedRefMut8<'b, T>(pub &'b mut T); + +impl<'a, T> PackedRefMut for PackedRefMut8<'a, T> { + type T = T; + #[inline(always)] + fn get(&self) -> &T { + self.0 + } + #[inline(always)] + fn get_mut(&mut self) -> &mut T { + self.0 + } +} diff --git a/crates/algo/src/prefetcher.rs b/crates/index/src/prefetcher.rs similarity index 80% rename from crates/algo/src/prefetcher.rs rename to crates/index/src/prefetcher.rs index ca140047..cd53a263 100644 --- a/crates/algo/src/prefetcher.rs +++ b/crates/index/src/prefetcher.rs @@ -12,16 +12,95 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::{ - Fetch, Hints, ReadStream, RelationPrefetch, RelationRead, RelationReadStream, - RelationReadTypes, Sequence, +use crate::fetch::Fetch; +use crate::relation::{ + Hints, ReadStream, RelationPrefetch, RelationRead, RelationReadStream, RelationReadTypes, }; -use std::collections::{VecDeque, vec_deque}; -use std::iter::Chain; +use dary_heap::DaryHeap; +use std::collections::{BinaryHeap, VecDeque}; pub const WINDOW_SIZE: usize = 32; const _: () = assert!(WINDOW_SIZE > 0); +pub trait Sequence { + type Item; + type Inner: Iterator; + #[must_use] + fn next(&mut self) -> Option; + #[must_use] + fn peek(&mut self) -> Option<&Self::Item>; + #[must_use] + fn into_inner(self) -> Self::Inner; +} + +impl Sequence for BinaryHeap { + type Item = T; + type Inner = std::vec::IntoIter; + #[inline] + fn next(&mut self) -> Option { + self.pop() + } + #[inline] + fn peek(&mut self) -> Option<&T> { + (self as &Self).peek() + } + #[inline] + fn into_inner(self) -> Self::Inner { + self.into_vec().into_iter() + } +} + +impl Sequence for DaryHeap { + type Item = T; + type Inner = std::vec::IntoIter; + #[inline] + fn next(&mut self) -> Option { + self.pop() + } + #[inline] + fn peek(&mut self) -> Option<&T> { + (self as &Self).peek() + } + #[inline] + fn into_inner(self) -> Self::Inner { + self.into_vec().into_iter() + } +} + +impl Sequence for std::iter::Peekable { + type Item = I::Item; + type Inner = std::iter::Peekable; + #[inline] + fn next(&mut self) -> Option { + Iterator::next(self) + } + #[inline] + fn peek(&mut self) -> Option<&I::Item> { + self.peek() + } + #[inline] + fn into_inner(self) -> Self::Inner { + self + } +} + +impl Sequence for VecDeque { + type Item = T; + type Inner = std::collections::vec_deque::IntoIter; + #[inline] + fn next(&mut self) -> Option { + self.pop_front() + } + #[inline] + fn peek(&mut self) -> Option<&T> { + self.front() + } + #[inline] + fn into_inner(self) -> Self::Inner { + self.into_iter() + } +} + pub trait Prefetcher<'b>: IntoIterator where Self::Item: Fetch<'b>, @@ -154,7 +233,7 @@ impl<'b, R, S: Sequence> SimplePrefetcher<'b, R, S> { impl<'b, R, S: Sequence> IntoIterator for SimplePrefetcher<'b, R, S> { type Item = S::Item; - type IntoIter = Chain, S::Inner>; + type IntoIter = std::iter::Chain, S::Inner>; #[inline] fn into_iter(self) -> Self::IntoIter { diff --git a/crates/index/src/relation.rs b/crates/index/src/relation.rs new file mode 100644 index 00000000..e6023542 --- /dev/null +++ b/crates/index/src/relation.rs @@ -0,0 +1,136 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::fetch::Fetch; +use std::ops::{Deref, DerefMut}; +use zerocopy::{FromBytes, IntoBytes}; + +/// # Safety +/// +/// * `Opaque` must aligned to 8 bytes. +#[allow(unsafe_code)] +pub unsafe trait Opaque: Copy + Send + Sync + FromBytes + IntoBytes + 'static {} + +pub trait Page: Sized + 'static { + type Opaque: Opaque; + + #[must_use] + fn get_opaque(&self) -> &Self::Opaque; + #[must_use] + fn get_opaque_mut(&mut self) -> &mut Self::Opaque; + #[must_use] + fn len(&self) -> u16; + #[must_use] + #[inline] + fn is_empty(&self) -> bool { + self.len() == 0 + } + #[must_use] + fn get(&self, i: u16) -> Option<&[u8]>; + #[must_use] + fn get_mut(&mut self, i: u16) -> Option<&mut [u8]>; + #[must_use] + fn alloc(&mut self, data: &[u8]) -> Option; + fn free(&mut self, i: u16); + #[must_use] + fn freespace(&self) -> u16; + fn clear(&mut self, opaque: Self::Opaque); +} + +pub trait PageGuard { + fn id(&self) -> u32; +} + +pub trait ReadStream<'b> { + type Relation: RelationReadTypes; + type Guards: ExactSizeIterator::ReadGuard<'b>>; + type Item; + type Inner: Iterator; + fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; + fn next_if bool>( + &mut self, + predicate: P, + ) -> Option<(Self::Item, Self::Guards)>; + fn into_inner(self) -> Self::Inner; +} + +pub trait Relation { + type Page: Page; +} + +pub trait RelationReadTypes: Relation { + type ReadGuard<'b>: PageGuard + Deref; +} + +pub trait RelationRead: RelationReadTypes { + fn read(&self, id: u32) -> Self::ReadGuard<'_>; +} + +pub trait RelationWriteTypes: Relation { + type WriteGuard<'b>: PageGuard + DerefMut; +} + +pub trait RelationWrite: RelationWriteTypes { + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_>; + fn extend( + &self, + opaque: ::Opaque, + tracking_freespace: bool, + ) -> Self::WriteGuard<'_>; + fn search(&self, freespace: usize) -> Option>; +} + +pub trait RelationPrefetch: Relation { + fn prefetch(&self, id: u32); +} + +#[derive(Debug, Clone, Copy)] +#[non_exhaustive] +pub struct Hints { + pub full: bool, + pub batch: bool, +} + +impl Default for Hints { + #[inline] + fn default() -> Self { + Self { + full: false, + batch: false, + } + } +} + +impl Hints { + #[inline] + pub fn full(self, full: bool) -> Self { + Self { full, ..self } + } + #[inline] + pub fn batch(self, batch: bool) -> Self { + Self { batch, ..self } + } +} + +pub trait RelationReadStreamTypes: RelationReadTypes { + type ReadStream<'b, I: Iterator>: ReadStream<'b, Item = I::Item, Relation = Self> + where + I::Item: Fetch<'b>; +} + +pub trait RelationReadStream: RelationReadStreamTypes { + fn read_stream<'b, I: Iterator>(&'b self, iter: I, hints: Hints) -> Self::ReadStream<'b, I> + where + I::Item: Fetch<'b>; +} diff --git a/crates/algo/src/tuples.rs b/crates/index/src/tuples.rs similarity index 100% rename from crates/algo/src/tuples.rs rename to crates/index/src/tuples.rs diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 40e8597e..f4c9fc7e 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -12,7 +12,7 @@ experimental_math = [] [dependencies] simd_macros = { path = "../simd_macros" } -half = { version = "2.6.0", features = ["zerocopy"] } +half = { version = "2.7.0", features = ["zerocopy"] } seq-macro.workspace = true zerocopy.workspace = true @@ -20,7 +20,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.38" +cc = "1.2.40" which = "8.0.0" [lints] diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 36e35d29..bba723b8 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -81,17 +81,7 @@ fn main() -> Result<(), Box> { build.opt_level(3); build.compile("simd_cshim"); } - _ => { - /* - let messages = [ - "This platform has poor SIMD implementation.", - "Please submit a feature request on https://github.com/tensorchord/VectorChord/issues.", - ]; - for message in messages { - println!("cargo::warning={message}"); - } - */ - } + _ => {} } Ok(()) } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 7422f795..90e8aeb3 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -21,7 +21,10 @@ #![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc_feature_detection))] #![cfg_attr(target_arch = "powerpc64", feature(powerpc_target_feature))] #![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc))] -#![cfg_attr(target_arch = "powerpc64", feature(core_intrinsics))] +#![cfg_attr( + all(target_arch = "powerpc64", target_endian = "big"), + feature(core_intrinsics) +)] #![cfg_attr(target_arch = "riscv64", feature(stdarch_riscv_feature_detection))] #![cfg_attr(target_arch = "riscv64", feature(riscv_target_feature))] diff --git a/crates/vchordg/Cargo.toml b/crates/vchordg/Cargo.toml index a9aa0b54..7b5a1658 100644 --- a/crates/vchordg/Cargo.toml +++ b/crates/vchordg/Cargo.toml @@ -5,9 +5,9 @@ edition.workspace = true publish = false [dependencies] -algo = { path = "../algo" } always_equal = { path = "../always_equal" } distance = { path = "../distance" } +index = { path = "../index" } rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } diff --git a/crates/vchordg/src/build.rs b/crates/vchordg/src/build.rs index 38259f56..2f13f734 100644 --- a/crates/vchordg/src/build.rs +++ b/crates/vchordg/src/build.rs @@ -16,7 +16,7 @@ use crate::Opaque; use crate::operator::Operator; use crate::tuples::{MetaTuple, OptionPointer, Tuple}; use crate::types::{VchordgIndexOptions, VectorOptions}; -use algo::{Page, PageGuard, RelationWrite}; +use index::relation::{Page, PageGuard, RelationWrite}; pub fn build( vector_options: VectorOptions, diff --git a/crates/vchordg/src/bulkdelete.rs b/crates/vchordg/src/bulkdelete.rs index a373559a..3cbfffbb 100644 --- a/crates/vchordg/src/bulkdelete.rs +++ b/crates/vchordg/src/bulkdelete.rs @@ -15,7 +15,7 @@ use crate::Opaque; use crate::operator::Operator; use crate::tuples::{MetaTuple, VertexTuple, WithReader, WithWriter}; -use algo::{Page, RelationRead, RelationWrite}; +use index::relation::{Page, RelationRead, RelationWrite}; use std::num::NonZero; pub fn bulkdelete( diff --git a/crates/vchordg/src/candidates.rs b/crates/vchordg/src/candidates.rs index fd901e2f..2a30b5a9 100644 --- a/crates/vchordg/src/candidates.rs +++ b/crates/vchordg/src/candidates.rs @@ -12,8 +12,9 @@ // // Copyright (c) 2025 TensorChord Inc. -use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; -use algo::{Fetch, RelationRead}; +use index::fetch::Fetch; +use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use index::relation::RelationRead; use std::collections::{BinaryHeap, VecDeque}; use std::marker::PhantomData; diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index 90f68566..4fd7b76a 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -20,10 +20,11 @@ use crate::tuples::*; use crate::types::DistanceKind; use crate::vectors::{by_prefetch, by_read, copy_all, copy_nothing, copy_outs, update}; use crate::visited::Visited; -use algo::accessor::LAccess; -use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; -use algo::{Bump, Page, PageGuard, RelationRead, RelationWrite}; use always_equal::AlwaysEqual; +use index::accessor::LAccess; +use index::bump::Bump; +use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; use rabitq::bits::Bits; use std::cmp::Reverse; use std::collections::VecDeque; diff --git a/crates/vchordg/src/lib.rs b/crates/vchordg/src/lib.rs index aaf7bbbb..fbedf22d 100644 --- a/crates/vchordg/src/lib.rs +++ b/crates/vchordg/src/lib.rs @@ -45,11 +45,11 @@ pub struct Opaque { } #[allow(unsafe_code)] -unsafe impl algo::Opaque for Opaque {} +unsafe impl index::relation::Opaque for Opaque {} pub type Id = (u32, u16); -impl algo::Fetch1 for tuples::Pointer { +impl index::fetch::Fetch1 for tuples::Pointer { fn fetch_1(&self) -> u32 { self.into_inner().0 } diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs index 58d39580..1f36e52e 100644 --- a/crates/vchordg/src/maintain.rs +++ b/crates/vchordg/src/maintain.rs @@ -17,8 +17,8 @@ use crate::operator::{CloneAccessor, Operator}; use crate::tuples::{MetaTuple, VectorTuple, VertexTuple, WithReader}; use crate::types::DistanceKind; use crate::vectors::{by_read, copy_all, copy_nothing, copy_outs, update}; -use algo::{Page, PageGuard, RelationRead, RelationWrite}; use always_equal::AlwaysEqual; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; use std::cmp::Reverse; use vector::VectorOwned; diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs index 9cb40636..1193de48 100644 --- a/crates/vchordg/src/operator.rs +++ b/crates/vchordg/src/operator.rs @@ -13,8 +13,8 @@ // Copyright (c) 2025 TensorChord Inc. use crate::types::DistanceKind; -use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S}; use distance::Distance; +use index::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S}; use rabitq::bits::Bits; use simd::{Floating, f16}; use std::fmt::Debug; diff --git a/crates/vchordg/src/prewarm.rs b/crates/vchordg/src/prewarm.rs index 8024be2a..6c0cda0b 100644 --- a/crates/vchordg/src/prewarm.rs +++ b/crates/vchordg/src/prewarm.rs @@ -14,7 +14,7 @@ use crate::Opaque; use crate::operator::Operator; -use algo::{Page, RelationRead}; +use index::relation::{Page, RelationRead}; pub fn prewarm(index: &R) -> String where diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs index af8d6721..5e260492 100644 --- a/crates/vchordg/src/search.rs +++ b/crates/vchordg/src/search.rs @@ -19,11 +19,12 @@ use crate::results::Results; use crate::tuples::*; use crate::vectors::{by_prefetch, copy_outs}; use crate::visited::Visited; -use algo::accessor::LAccess; -use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; -use algo::{Bump, Page, RelationRead}; use always_equal::AlwaysEqual; use distance::Distance; +use index::accessor::LAccess; +use index::bump::Bump; +use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use index::relation::{Page, RelationRead}; use rabitq::bits::Bits; use std::cmp::Reverse; use std::collections::VecDeque; diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index a33fc134..1bc1ead2 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -13,8 +13,8 @@ // Copyright (c) 2025 TensorChord Inc. use crate::operator::Vector; -use algo::tuples::{Bool, MutChecker, Padding, RefChecker}; use distance::Distance; +use index::tuples::{Bool, MutChecker, Padding, RefChecker}; use std::num::{NonZero, Wrapping}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; diff --git a/crates/vchordg/src/vectors.rs b/crates/vchordg/src/vectors.rs index 94a43aa0..ab37b2af 100644 --- a/crates/vchordg/src/vectors.rs +++ b/crates/vchordg/src/vectors.rs @@ -14,9 +14,9 @@ use crate::operator::{Operator, Vector}; use crate::tuples::*; -use algo::accessor::Accessor1; -use algo::{Page, RelationRead, RelationWrite}; use distance::Distance; +use index::accessor::Accessor1; +use index::relation::{Page, RelationRead, RelationWrite}; use std::collections::VecDeque; use std::num::{NonZero, Wrapping}; diff --git a/crates/vchordrq/Cargo.toml b/crates/vchordrq/Cargo.toml index e453f0fd..acd9d932 100644 --- a/crates/vchordrq/Cargo.toml +++ b/crates/vchordrq/Cargo.toml @@ -5,9 +5,9 @@ edition.workspace = true publish = false [dependencies] -algo = { path = "../algo" } always_equal = { path = "../always_equal" } distance = { path = "../distance" } +index = { path = "../index" } rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } diff --git a/crates/vchordrq/src/build.rs b/crates/vchordrq/src/build.rs index 5bcbe177..f0905df0 100644 --- a/crates/vchordrq/src/build.rs +++ b/crates/vchordrq/src/build.rs @@ -18,7 +18,7 @@ use crate::tape_writer::{DirectoryTapeWriter, H1TapeWriter}; use crate::tuples::*; use crate::types::*; use crate::{Branch, Opaque}; -use algo::{Page, RelationWrite}; +use index::relation::{Page, RelationWrite}; use vector::{VectorBorrowed, VectorOwned}; pub fn build( diff --git a/crates/vchordrq/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs index f93751fb..6bdbe530 100644 --- a/crates/vchordrq/src/bulkdelete.rs +++ b/crates/vchordrq/src/bulkdelete.rs @@ -16,9 +16,9 @@ use crate::closure_lifetime_binder::{id_0, id_1}; use crate::operator::Operator; use crate::tape::by_next; use crate::tuples::*; -use crate::{Opaque, Page, tape}; -use algo::accessor::FunctionalAccessor; -use algo::{RelationRead, RelationWrite}; +use crate::{Opaque, tape}; +use index::accessor::FunctionalAccessor; +use index::relation::{Page, RelationRead, RelationWrite}; use std::num::NonZero; pub fn bulkdelete( diff --git a/crates/vchordrq/src/cache.rs b/crates/vchordrq/src/cache.rs index 307446d4..89c9178e 100644 --- a/crates/vchordrq/src/cache.rs +++ b/crates/vchordrq/src/cache.rs @@ -13,11 +13,10 @@ // Copyright (c) 2025 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1}; -use crate::tape::by_next; use crate::tuples::{MetaTuple, WithReader}; -use crate::{Opaque, Page, PageGuard, tape}; -use algo::RelationRead; -use algo::accessor::FunctionalAccessor; +use crate::{Opaque, tape}; +use index::accessor::FunctionalAccessor; +use index::relation::{Page, PageGuard, RelationRead}; pub fn cache(index: &R, level: i32) -> Vec where @@ -46,7 +45,7 @@ where let mut results = Vec::new(); for first in state { tape::read_h1_tape::( - by_next(index, first).inspect(|guard| trace.push(guard.id())), + tape::by_next(index, first).inspect(|guard| trace.push(guard.id())), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), |(), _, _, first, _| { results.push(first); @@ -69,7 +68,7 @@ where if level >= 2 { // centroid tuples - for guard in by_next(index, centroids_first) { + for guard in tape::by_next(index, centroids_first) { trace.push(guard.id()); } } diff --git a/crates/vchordrq/src/centroids.rs b/crates/vchordrq/src/centroids.rs index a7490e1c..44221a86 100644 --- a/crates/vchordrq/src/centroids.rs +++ b/crates/vchordrq/src/centroids.rs @@ -12,11 +12,10 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::Page; use crate::operator::*; use crate::tuples::*; -use algo::RelationRead; -use algo::accessor::Accessor1; +use index::accessor::Accessor1; +use index::relation::{Page, RelationRead}; pub fn read< 'a, diff --git a/crates/vchordrq/src/closure_lifetime_binder.rs b/crates/vchordrq/src/closure_lifetime_binder.rs index a7472d30..f247e9e2 100644 --- a/crates/vchordrq/src/closure_lifetime_binder.rs +++ b/crates/vchordrq/src/closure_lifetime_binder.rs @@ -42,7 +42,7 @@ where #[inline(always)] pub fn id_3(f: F) -> F where - T: algo::RelationWrite, + T: index::relation::RelationWrite, F: for<'a> Fn(&'a T, A, B) -> T::WriteGuard<'a>, { f @@ -51,8 +51,8 @@ where #[inline(always)] pub fn id_4<'b, F, P, A, B, R: ?Sized>(f: F) -> F where - P: algo::prefetcher::Prefetcher<'b>, -

::Item: algo::Fetch<'b>, + P: index::prefetcher::Prefetcher<'b>, +

::Item: index::fetch::Fetch<'b>, F: FnMut(A, P::Guards, B) -> R, { f diff --git a/crates/vchordrq/src/cost.rs b/crates/vchordrq/src/cost.rs index c4eefecd..5eb2dc6d 100644 --- a/crates/vchordrq/src/cost.rs +++ b/crates/vchordrq/src/cost.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::tuples::{MetaTuple, WithReader}; -use algo::{Page, RelationRead}; +use index::relation::{Page, RelationRead}; pub struct Cost { pub dims: u32, diff --git a/crates/vchordrq/src/fast_heap.rs b/crates/vchordrq/src/fast_heap.rs index d989ddfe..27149082 100644 --- a/crates/vchordrq/src/fast_heap.rs +++ b/crates/vchordrq/src/fast_heap.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use algo::Sequence; +use index::prefetcher::Sequence; use std::collections::BinaryHeap; use std::num::NonZero; diff --git a/crates/vchordrq/src/freepages.rs b/crates/vchordrq/src/freepages.rs index b5420355..7657f6e0 100644 --- a/crates/vchordrq/src/freepages.rs +++ b/crates/vchordrq/src/freepages.rs @@ -12,9 +12,9 @@ // // Copyright (c) 2025 TensorChord Inc. +use crate::Opaque; use crate::tuples::*; -use crate::{Opaque, Page}; -use algo::RelationWrite; +use index::relation::{Page, RelationWrite}; pub fn alloc(index: &R, freepages_first: u32) -> Option> where diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index c3aca5c7..7d4abfa0 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -14,15 +14,15 @@ use crate::linked_vec::LinkedVec; use crate::operator::*; -use crate::tape::by_next; use crate::tuples::*; -use crate::vectors::{self}; -use crate::{Chooser, Opaque, Page, centroids, tape}; -use algo::accessor::{FunctionalAccessor, LAccess}; -use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily}; -use algo::{BorrowedIter, Bump, RelationRead, RelationWrite}; +use crate::{Chooser, Opaque, centroids, tape, vectors}; use always_equal::AlwaysEqual; use distance::Distance; +use index::accessor::{FunctionalAccessor, LAccess}; +use index::bump::Bump; +use index::fetch::BorrowedIter; +use index::prefetcher::{Prefetcher, PrefetcherHeapFamily}; +use index::relation::{Page, RelationRead, RelationWrite}; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZero; @@ -112,7 +112,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( { let (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) = state; tape::read_h1_tape::( - by_next(index, first), + tape::by_next(index, first), || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), |(rough, err), head, norm, first, prefetch| { let lowerbound = Distance::from_f32(rough - err * epsilon); diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index 771efcf1..f176809e 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -34,7 +34,6 @@ mod vectors; pub mod operator; pub mod types; -use algo::{Page, PageGuard}; pub use build::build; pub use bulkdelete::{bulkdelete, bulkdelete_vectors}; pub use cache::cache; @@ -45,8 +44,8 @@ pub use maintain::maintain; pub use prewarm::prewarm; pub use rerank::{how, rerank_heap, rerank_index}; pub use search::{default_search, maxsim_search}; -use std::num::NonZero; +use std::num::NonZero; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] @@ -57,7 +56,7 @@ pub struct Opaque { } #[allow(unsafe_code)] -unsafe impl algo::Opaque for Opaque {} +unsafe impl index::relation::Opaque for Opaque {} pub(crate) struct Branch { pub code: rabitq::bit::Code, @@ -71,3 +70,9 @@ pub(crate) struct Branch { pub trait Chooser { fn choose(&mut self, n: NonZero) -> usize; } + +#[derive(Debug, Clone, Copy)] +pub enum RerankMethod { + Index, + Heap, +} diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs index 4b7b1866..f7ee7e01 100644 --- a/crates/vchordrq/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -14,14 +14,13 @@ use crate::closure_lifetime_binder::{id_0, id_1, id_2, id_3}; use crate::operator::{Operator, Vector}; -use crate::tape::{self, TapeWriter, by_directory, by_next}; use crate::tape_writer::{DirectoryTapeWriter, FrozenTapeWriter}; use crate::tuples::*; -use crate::{Branch, Opaque, Page, freepages}; -use algo::accessor::FunctionalAccessor; -use algo::prefetcher::PrefetcherSequenceFamily; -use algo::{ - PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, +use crate::{Branch, Opaque, freepages, tape}; +use index::accessor::FunctionalAccessor; +use index::prefetcher::PrefetcherSequenceFamily; +use index::relation::{ + Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, }; use rabitq::packing::unpack; use std::cell::RefCell; @@ -56,7 +55,7 @@ where let mut results = Vec::new(); for first in state { tape::read_h1_tape::( - by_next(index, first).inspect(|_| check()), + tape::by_next(index, first).inspect(|_| check()), || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), |(), _, _, first, _| results.push(first), ); @@ -129,12 +128,12 @@ where tuples += 1; }); let directory = tape::read_directory_tape::( - by_next(index, *jump_tuple.directory_first()) + tape::by_next(index, *jump_tuple.directory_first()) .inspect(|_| check()) .inspect(|guard| trace_directory.push(guard.id())), ); tape::read_frozen_tape::( - by_directory(&mut prefetch_h0_tuples, directory) + tape::by_directory(&mut prefetch_h0_tuples, directory) .inspect(|_| check()) .inspect(|guard| trace_forzen.push(guard.id())), || { @@ -167,7 +166,7 @@ where &mut callback, ); tape::read_appendable_tape::( - by_next(index, *jump_tuple.appendable_first()) + tape::by_next(index, *jump_tuple.appendable_first()) .inspect(|_| check()) .inspect(|guard| trace_appendable.push(guard.id())), |metadata, elements, delta| { @@ -194,7 +193,7 @@ where let (frozen_tape, branches) = tape.into_inner(); - let mut appendable_tape = TapeWriter::create(&hooked_index, false); + let mut appendable_tape = tape::TapeWriter::create(&hooked_index, false); for branch in branches { appendable_tape.push(AppendableTuple { @@ -214,7 +213,7 @@ where let frozen_first = { frozen_tape }.first(); - let directory = by_next(index, frozen_first) + let directory = tape::by_next(index, frozen_first) .inspect(|_| check()) .map(|guard| guard.id()) .collect::>(); diff --git a/crates/vchordrq/src/operator.rs b/crates/vchordrq/src/operator.rs index 627e8aeb..a27bf41f 100644 --- a/crates/vchordrq/src/operator.rs +++ b/crates/vchordrq/src/operator.rs @@ -12,8 +12,8 @@ // // Copyright (c) 2025 TensorChord Inc. -use algo::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S, RAccess}; use distance::Distance; +use index::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S, RAccess}; use rabitq::bit::CodeMetadata; use rabitq::bit::binary::BinaryLut; use rabitq::bit::block::{BlockLut, STEP}; diff --git a/crates/vchordrq/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs index fa4c0f9b..d56314e4 100644 --- a/crates/vchordrq/src/prewarm.rs +++ b/crates/vchordrq/src/prewarm.rs @@ -16,10 +16,10 @@ use crate::closure_lifetime_binder::{id_0, id_1, id_2}; use crate::operator::Operator; use crate::tape::{by_directory, by_next}; use crate::tuples::*; -use crate::{Opaque, Page, centroids, tape}; -use algo::RelationRead; -use algo::accessor::FunctionalAccessor; -use algo::prefetcher::PrefetcherSequenceFamily; +use crate::{Opaque, centroids, tape}; +use index::accessor::FunctionalAccessor; +use index::prefetcher::PrefetcherSequenceFamily; +use index::relation::{Page, RelationRead}; use std::fmt::Write; pub fn prewarm<'b, R: RelationRead, O: Operator>( diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs index 0c828d92..446630c7 100644 --- a/crates/vchordrq/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -15,12 +15,14 @@ use crate::closure_lifetime_binder::id_4; use crate::operator::*; use crate::tuples::{MetaTuple, WithReader}; -use crate::{Page, vectors}; -use algo::accessor::{Accessor2, LTryAccess}; -use algo::prefetcher::Prefetcher; -use algo::{PackedRefMut, RelationRead, RerankMethod}; +use crate::{RerankMethod, vectors}; use always_equal::AlwaysEqual; use distance::Distance; +use index::accessor::{Accessor2, LTryAccess}; +use index::fetch::BorrowedIter; +use index::packed::PackedRefMut; +use index::prefetcher::Prefetcher; +use index::relation::{Page, RelationRead}; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::marker::PhantomData; @@ -52,7 +54,7 @@ impl<'b, T, F, P, W> Iterator for Reranker where F: FnMut(NonZero, P::Guards, u16) -> Option, P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual)>, - W: 'b + PackedRefMut, u16, algo::BorrowedIter<'b>)>, + W: 'b + PackedRefMut, u16, BorrowedIter<'b>)>, { type Item = (Distance, NonZero); @@ -82,7 +84,7 @@ pub fn rerank_index< O: Operator, T, P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual)>, - W: 'b + PackedRefMut, u16, algo::BorrowedIter<'b>)>, + W: 'b + PackedRefMut, u16, BorrowedIter<'b>)>, >( vector: O::Vector, prefetcher: P, @@ -110,7 +112,7 @@ pub fn rerank_heap< O: Operator, T, P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual)>, - W: 'b + PackedRefMut, u16, algo::BorrowedIter<'b>)>, + W: 'b + PackedRefMut, u16, BorrowedIter<'b>)>, >( vector: O::Vector, prefetcher: P, diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index eaf503e0..583569a2 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -17,12 +17,15 @@ use crate::linked_vec::LinkedVec; use crate::operator::*; use crate::tape::{by_directory, by_next}; use crate::tuples::*; -use crate::{Opaque, Page, centroids, tape}; -use algo::accessor::{FunctionalAccessor, LAccess}; -use algo::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; -use algo::{BorrowedIter, Bump, PackedRefMut4, PackedRefMut8, RelationRead}; +use crate::{Opaque, centroids, tape}; use always_equal::AlwaysEqual; use distance::Distance; +use index::accessor::{FunctionalAccessor, LAccess}; +use index::bump::Bump; +use index::fetch::BorrowedIter; +use index::packed::{PackedRefMut4, PackedRefMut8}; +use index::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; +use index::relation::{Page, RelationRead}; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZero; diff --git a/crates/vchordrq/src/tape.rs b/crates/vchordrq/src/tape.rs index 755a8f00..aa135dd2 100644 --- a/crates/vchordrq/src/tape.rs +++ b/crates/vchordrq/src/tape.rs @@ -13,10 +13,10 @@ // Copyright (c) 2025 TensorChord Inc. use crate::tuples::*; -use crate::{Opaque, Page, PageGuard, freepages}; -use algo::accessor::Accessor1; -use algo::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; -use algo::{RelationRead, RelationWrite}; +use crate::{Opaque, freepages}; +use index::accessor::Accessor1; +use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; use std::marker::PhantomData; use std::num::NonZero; diff --git a/crates/vchordrq/src/tape_writer.rs b/crates/vchordrq/src/tape_writer.rs index 3078ea10..258ead51 100644 --- a/crates/vchordrq/src/tape_writer.rs +++ b/crates/vchordrq/src/tape_writer.rs @@ -15,7 +15,7 @@ use crate::tape::TapeWriter; use crate::tuples::*; use crate::{Branch, Opaque}; -use algo::{Page, RelationWrite}; +use index::relation::{Page, RelationWrite}; use rabitq::packing::{any_pack, padding_pack}; use std::num::NonZero; diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index 054a8c5e..36b16427 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::operator::Vector; -use algo::tuples::{Bool, MutChecker, Padding, RefChecker}; +use index::tuples::{Bool, MutChecker, Padding, RefChecker}; use std::num::NonZero; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; diff --git a/crates/vchordrq/src/vectors.rs b/crates/vchordrq/src/vectors.rs index 183fb2e0..0767425c 100644 --- a/crates/vchordrq/src/vectors.rs +++ b/crates/vchordrq/src/vectors.rs @@ -14,9 +14,9 @@ use crate::operator::*; use crate::tuples::*; -use crate::{Opaque, Page, PageGuard, tape}; -use algo::accessor::TryAccessor1; -use algo::{RelationRead, RelationWrite}; +use crate::{Opaque, tape}; +use index::accessor::TryAccessor1; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; use std::num::NonZero; use vector::VectorOwned; diff --git a/src/bin/pgrx_embed.rs b/src/bin/pgrx_embed.rs index 36ae772e..6e0b06da 100644 --- a/src/bin/pgrx_embed.rs +++ b/src/bin/pgrx_embed.rs @@ -13,7 +13,6 @@ // Copyright (c) 2025 TensorChord Inc. #![allow(unsafe_code)] -#![allow(unused_crate_dependencies)] ::pgrx::pgrx_embed!(); diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 4ce1c38f..8da4ee16 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -12,13 +12,13 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod binary_scalar8; -pub mod functions_scalar8; +mod binary_scalar8; +mod functions_scalar8; pub mod memory_halfvec; pub mod memory_scalar8; pub mod memory_vector; -pub mod operators_halfvec; -pub mod operators_scalar8; -pub mod operators_vector; -pub mod text_scalar8; +mod operators_halfvec; +mod operators_scalar8; +mod operators_vector; +mod text_scalar8; pub mod typmod; diff --git a/src/index/functions.rs b/src/index/functions.rs index 65c71c78..922e46db 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -37,7 +37,7 @@ fn _vchordg_prewarm(indexrelid: Oid) -> String { let relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); let opfamily = unsafe { crate::index::vchordg::opclass::opfamily(relation.raw()) }; let index = unsafe { PostgresRelation::new(relation.raw()) }; - crate::index::vchordg::algo::prewarm(opfamily, &index) + crate::index::vchordg::dispatch::prewarm(opfamily, &index) } #[pgrx::pg_extern(sql = "")] @@ -59,7 +59,7 @@ fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { let relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); let opfamily = unsafe { crate::index::vchordrq::opclass::opfamily(relation.raw()) }; let index = unsafe { PostgresRelation::new(relation.raw()) }; - crate::index::vchordrq::algo::prewarm(opfamily, &index, height) + crate::index::vchordrq::dispatch::prewarm(opfamily, &index, height) } struct Index { diff --git a/src/index/mod.rs b/src/index/mod.rs index 4cc178cd..b68eb4e8 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -12,15 +12,15 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod fetcher; -pub mod functions; -pub mod gucs; -pub mod hook; -pub mod opclass; -pub mod scanners; -pub mod storage; -pub mod vchordg; -pub mod vchordrq; +mod fetcher; +mod functions; +mod gucs; +mod hook; +mod opclass; +mod scanners; +mod storage; +mod vchordg; +mod vchordrq; pub fn init() { gucs::init(); diff --git a/src/index/scanners.rs b/src/index/scanners.rs index 97d8b1b4..e0d11c7a 100644 --- a/src/index/scanners.rs +++ b/src/index/scanners.rs @@ -14,7 +14,8 @@ use crate::index::fetcher::Fetcher; use crate::recorder::Recorder; -use algo::{Bump, Page, RelationPrefetch, RelationRead, RelationReadStream}; +use index::bump::Bump; +use index::relation::{Page, RelationPrefetch, RelationRead, RelationReadStream}; use pgrx::pg_sys::Datum; #[derive(Debug, Clone, Copy)] diff --git a/src/index/storage.rs b/src/index/storage.rs index 6bda0a4e..2299c95c 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -12,7 +12,12 @@ // // Copyright (c) 2025 TensorChord Inc. -use algo::{Opaque, *}; +use index::fetch::Fetch; +use index::relation::{ + Hints, Opaque, Page, PageGuard, ReadStream, Relation, RelationPrefetch, RelationRead, + RelationReadStream, RelationReadStreamTypes, RelationReadTypes, RelationWrite, + RelationWriteTypes, +}; use std::collections::VecDeque; use std::iter::{Chain, Flatten}; use std::marker::PhantomData; @@ -659,18 +664,19 @@ impl RelationReadStream for PostgresRelation { _phantom: PhantomData, })); let raw = unsafe { - use pgrx::pg_sys::{ - ForkNumber, READ_STREAM_DEFAULT, READ_STREAM_FULL, read_stream_begin_relation, - }; - let mut flags = READ_STREAM_DEFAULT; + let mut flags = pgrx::pg_sys::READ_STREAM_DEFAULT; if hints.full { - flags |= READ_STREAM_FULL; + flags |= pgrx::pg_sys::READ_STREAM_FULL; } - read_stream_begin_relation( + #[cfg(feature = "pg18")] + if hints.batch { + flags |= pgrx::pg_sys::READ_STREAM_USE_BATCHING; + } + pgrx::pg_sys::read_stream_begin_relation( flags as i32, core::ptr::null_mut(), self.raw, - ForkNumber::MAIN_FORKNUM, + pgrx::pg_sys::ForkNumber::MAIN_FORKNUM, Some(callback::), cache.as_ptr().cast(), 0, diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index f660e7ec..c2ca4888 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -199,7 +199,7 @@ pub unsafe extern "C-unwind" fn ambuild( } let index = unsafe { PostgresRelation::new(index_relation) }; let mut reporter = PostgresReporter {}; - crate::index::vchordg::algo::build(vector_options, vchordg_options.index, &index); + crate::index::vchordg::dispatch::build(vector_options, vchordg_options.index, &index); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); let cached = vchordg_cached::VchordgCached::_0 {}.serialize(); if let Some(leader) = unsafe { @@ -547,7 +547,7 @@ unsafe fn parallel_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::vchordg::algo::insert(opfamily, &index, payload, vector); + crate::index::vchordg::dispatch::insert(opfamily, &index, payload, vector); } unsafe { let indtuples; @@ -597,7 +597,7 @@ unsafe fn sequential_build( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::vchordg::algo::insert(opfamily, &index, payload, vector); + crate::index::vchordg::dispatch::insert(opfamily, &index, payload, vector); } indtuples += 1; callback(indtuples); @@ -673,7 +673,7 @@ unsafe fn options( mod vchordg_cached { pub type Tag = u64; - use algo::tuples::RefChecker; + use index::tuples::RefChecker; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index fc2f2e1d..c7ef9f8c 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod am_build; +mod am_build; use crate::index::fetcher::*; use crate::index::gucs; @@ -229,7 +229,7 @@ unsafe fn aminsertinner( for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - crate::index::vchordg::algo::insert(opfamily, &index, payload, vector); + crate::index::vchordg::dispatch::insert(opfamily, &index, payload, vector); } } false @@ -272,7 +272,7 @@ pub unsafe extern "C-unwind" fn ambulkdelete( pg_guard_ffi_boundary(|| callback(&mut ctid, callback_state)) } }; - crate::index::vchordg::algo::bulkdelete(opfamily, &index, check, callback); + crate::index::vchordg::dispatch::bulkdelete(opfamily, &index, check, callback); stats } @@ -301,7 +301,7 @@ pub unsafe extern "C-unwind" fn amvacuumcleanup( #[cfg(feature = "pg18")] pgrx::pg_sys::vacuum_delay_point(false); }; - crate::index::vchordg::algo::maintain(opfamily, &index, check); + crate::index::vchordg::dispatch::maintain(opfamily, &index, check); stats } diff --git a/src/index/vchordg/algo.rs b/src/index/vchordg/dispatch.rs similarity index 94% rename from src/index/vchordg/algo.rs rename to src/index/vchordg/dispatch.rs index 52416c01..2c186a7c 100644 --- a/src/index/vchordg/algo.rs +++ b/src/index/vchordg/dispatch.rs @@ -13,12 +13,14 @@ // Copyright (c) 2025 TensorChord Inc. use crate::index::vchordg::opclass::Opfamily; -use algo::accessor::{Dot, L2S}; -use algo::prefetcher::*; -use algo::*; +use index::accessor::{Dot, L2S}; +use index::fetch::Fetch; +use index::prefetcher::*; +use index::relation::{ + Hints, Page, RelationPrefetch, RelationRead, RelationReadStream, RelationWrite, +}; use simd::f16; use std::num::NonZero; -use vchordg::Opaque; use vchordg::operator::Op; use vchordg::types::*; use vector::VectorOwned; @@ -27,7 +29,7 @@ use vector::vect::{VectBorrowed, VectOwned}; pub fn prewarm(opfamily: Opfamily, index: &R) -> String where R: RelationRead, - R::Page: Page, + R::Page: Page, { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { @@ -52,7 +54,7 @@ pub fn bulkdelete( callback: impl Fn(NonZero) -> bool, ) where R: RelationRead + RelationWrite, - R::Page: Page, + R::Page: Page, { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { @@ -73,7 +75,7 @@ pub fn bulkdelete( pub fn maintain(opfamily: Opfamily, index: &R, check: impl Fn()) where R: RelationRead + RelationWrite, - R::Page: Page, + R::Page: Page, { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { @@ -94,7 +96,7 @@ where pub fn build(vector_options: VectorOptions, vchordg_options: VchordgIndexOptions, index: &R) where R: RelationRead + RelationWrite, - R::Page: Page, + R::Page: Page, { match (vector_options.v, vector_options.d) { (VectorKind::Vecf32, DistanceKind::L2S) => { @@ -115,7 +117,7 @@ where pub fn insert(opfamily: Opfamily, index: &R, payload: NonZero, vector: OwnedVector) where R: RelationRead + RelationWrite + RelationReadStream, - R::Page: Page, + R::Page: Page, { let bump = bumpalo::Bump::new(); let make_vertex_plain_prefetcher = MakePlainPrefetcher { index }; @@ -266,7 +268,7 @@ impl<'b, R> Clone for MakeStreamPrefetcher<'b, R> { fn clone(&self) -> Self { Self { index: self.index, - hints: self.hints.clone(), + hints: self.hints, } } } @@ -283,7 +285,7 @@ impl<'b, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'b, R> where S::Item: Fetch<'b>, { - StreamPrefetcher::new(self.index, seq, self.hints.clone()) + StreamPrefetcher::new(self.index, seq, self.hints) } fn is_not_plain(&self) -> bool { diff --git a/src/index/vchordg/mod.rs b/src/index/vchordg/mod.rs index d897f7d1..157200da 100644 --- a/src/index/vchordg/mod.rs +++ b/src/index/vchordg/mod.rs @@ -12,8 +12,8 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod algo; pub mod am; +pub mod dispatch; pub mod opclass; -pub mod scanners; +mod scanners; pub mod types; diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs index 54a2c560..b0945bf4 100644 --- a/src/index/vchordg/scanners/default.rs +++ b/src/index/vchordg/scanners/default.rs @@ -15,18 +15,19 @@ use crate::index::fetcher::{Fetcher, pointer_to_kv}; use crate::index::opclass::Sphere; use crate::index::scanners::{Io, SearchBuilder}; -use crate::index::vchordg::algo::*; +use crate::index::vchordg::dispatch::*; use crate::index::vchordg::opclass::Opfamily; use crate::index::vchordg::scanners::SearchOptions; use crate::recorder::{Recorder, halfvec_out, vector_out}; -use algo::accessor::{Dot, L2S}; -use algo::*; use distance::Distance; +use index::accessor::{Dot, L2S}; +use index::bump::Bump; +use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadStream}; use simd::f16; use std::num::NonZero; use vchordg::operator::{self}; +use vchordg::search; use vchordg::types::{DistanceKind, OwnedVector, VectorKind}; -use vchordg::*; use vector::VectorOwned; use vector::vect::{VectBorrowed, VectOwned}; @@ -107,17 +108,19 @@ impl SearchBuilder for DefaultBuilder { let Some(vector) = vector else { return Box::new(std::iter::empty()) as Box>; }; + let vertex_hints = Hints::default().full(true); + let vector_hints = Hints::default().full(true); let make_vertex_plain_prefetcher = MakePlainPrefetcher { index }; let make_vertex_simple_prefetcher = MakeSimplePrefetcher { index }; let make_vertex_stream_prefetcher = MakeStreamPrefetcher { index, - hints: Hints::default().full(true), + hints: vertex_hints, }; let make_vector_plain_prefetcher = MakePlainPrefetcher { index }; let make_vector_simple_prefetcher = MakeSimplePrefetcher { index }; let make_vector_stream_prefetcher = MakeStreamPrefetcher { index, - hints: Hints::default().full(true), + hints: vector_hints, }; let iter: Box)>> = match (opfamily.vector_kind(), opfamily.distance_kind()) { diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 42347bc0..c9aaa0b7 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -16,20 +16,20 @@ use crate::datatype::typmod::Typmod; use crate::index::fetcher::*; use crate::index::storage::{PostgresPage, PostgresRelation}; use crate::index::vchordrq::am::Reloption; +use crate::index::vchordrq::build::{Normalize, Normalized}; use crate::index::vchordrq::opclass::{Opfamily, opfamily}; use crate::index::vchordrq::types::*; -use algo::{ +use index::relation::{ Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, }; use pgrx::pg_sys::{Datum, ItemPointerData}; use rand::Rng; -use simd::{Floating, f16}; +use simd::Floating; use std::ffi::CStr; use std::num::NonZero; use std::ops::Deref; use vchordrq::Chooser; use vchordrq::types::*; -use vector::VectorOwned; use vector::vect::VectOwned; #[derive(Debug, Clone, Copy)] @@ -309,8 +309,8 @@ pub unsafe extern "C-unwind" fn ambuild( heap.traverse(false, |(_, store)| { for (vector, _) in store { let x = match vector { - OwnedVector::Vecf32(x) => VectOwned::build_to_vecf32(x.as_borrowed()), - OwnedVector::Vecf16(x) => VectOwned::build_to_vecf32(x.as_borrowed()), + OwnedVector::Vecf32(x) => VectOwned::normalize(x), + OwnedVector::Vecf16(x) => VectOwned::normalize(x), }; assert_eq!( vector_options.dims, @@ -343,7 +343,12 @@ pub unsafe extern "C-unwind" fn ambuild( } } reporter.phase(BuildPhase::from_code(BuildPhaseCode::Build)); - crate::index::vchordrq::algo::build(vector_options, vchordrq_options.index, &index, structures); + crate::index::vchordrq::dispatch::build( + vector_options, + vchordrq_options.index, + &index, + structures, + ); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); let cached = if vchordrq_options.build.pin >= 0 { let mut trace = vchordrq::cache(&index, vchordrq_options.build.pin); @@ -421,7 +426,7 @@ pub unsafe extern "C-unwind" fn ambuild( pgrx::check_for_interrupts!(); }; reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); - crate::index::vchordrq::algo::maintain(opfamily, &index, check); + crate::index::vchordrq::dispatch::maintain(opfamily, &index, check); unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -446,7 +451,7 @@ mod vchordrq_cached { pub type Tag = u64; use crate::index::storage::PostgresPage; - use algo::tuples::RefChecker; + use index::tuples::RefChecker; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] @@ -869,7 +874,7 @@ unsafe fn parallel_build( let payload = kv_to_pointer((key, extra)); let mut chooser = IdChooser(wid); let bump = bumpalo::Bump::new(); - crate::index::vchordrq::algo::insert( + crate::index::vchordrq::dispatch::insert( opfamily, &index, payload, @@ -903,7 +908,7 @@ unsafe fn parallel_build( let payload = kv_to_pointer((key, extra)); let mut chooser = IdChooser(wid); let bump = bumpalo::Bump::new(); - crate::index::vchordrq::algo::insert( + crate::index::vchordrq::dispatch::insert( opfamily, &index, payload, @@ -973,7 +978,7 @@ unsafe fn sequential_build( let payload = kv_to_pointer((key, extra)); let mut chooser = ChooseZero; let bump = bumpalo::Bump::new(); - crate::index::vchordrq::algo::insert( + crate::index::vchordrq::dispatch::insert( opfamily, &index, payload, @@ -999,7 +1004,7 @@ unsafe fn sequential_build( let payload = kv_to_pointer((key, extra)); let mut chooser = ChooseZero; let bump = bumpalo::Bump::new(); - crate::index::vchordrq::algo::insert( + crate::index::vchordrq::dispatch::insert( opfamily, &index, payload, @@ -1084,8 +1089,8 @@ unsafe fn options( fn make_default_build( vector_options: VectorOptions, _default_build: VchordrqDefaultBuildOptions, -) -> Vec>> { - vec![Structure::> { +) -> Vec> { + vec![Structure:: { centroids: vec![vec![0.0f32; vector_options.dims as usize]], children: vec![vec![]], }] @@ -1094,11 +1099,11 @@ fn make_default_build( fn make_internal_build( vector_options: VectorOptions, internal_build: VchordrqInternalBuildOptions, - samples: Vec>, + samples: Vec, reporter: &mut PostgresReporter, -) -> Vec>> { +) -> Vec> { use std::iter::once; - let mut result = Vec::>>::new(); + let mut result = Vec::>::new(); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { let input = if let Some(structure) = result.last() { &structure.centroids @@ -1157,7 +1162,7 @@ fn make_internal_build( if let Some(structure) = result.last() { let mut children = vec![Vec::new(); centroids.len()]; for i in 0..structure.len() as u32 { - pub fn k_means_lookup(vector: &[f32], centroids: &[Vec]) -> usize { + pub fn k_means_lookup(vector: &[f32], centroids: &[Normalized]) -> usize { assert_ne!(centroids.len(), 0); let mut result = (f32::INFINITY, 0); for i in 0..centroids.len() { @@ -1194,7 +1199,7 @@ fn make_external_build( vector_options: VectorOptions, _opfamily: Opfamily, external_build: VchordrqExternalBuildOptions, -) -> Vec>> { +) -> Vec> { use std::collections::BTreeMap; let VchordrqExternalBuildOptions { table } = external_build; let mut parents = BTreeMap::new(); @@ -1330,9 +1335,9 @@ fn make_external_build( fn extract( height: u32, labels: &BTreeMap, - vectors: &BTreeMap>, + vectors: &BTreeMap, children: &BTreeMap>, - ) -> (Vec>, Vec>) { + ) -> (Vec, Vec>) { labels .iter() .filter(|(_, (h, _))| *h == height) @@ -1355,32 +1360,6 @@ fn make_external_build( result } -pub trait InternalBuild: VectorOwned { - fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec; - - fn build_from_vecf32(x: &[f32]) -> Self; -} - -impl InternalBuild for VectOwned { - fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { - vector.slice().to_vec() - } - - fn build_from_vecf32(x: &[f32]) -> Self { - Self::new(x.to_vec()) - } -} - -impl InternalBuild for VectOwned { - fn build_to_vecf32(vector: Self::Borrowed<'_>) -> Vec { - f16::vector_to_f32(vector.slice()) - } - - fn build_from_vecf32(x: &[f32]) -> Self { - Self::new(f16::vector_from_f32(x)) - } -} - struct CachingRelation<'a, R> { cache: vchordrq_cached::VchordrqCachedReader1<'a>, relation: R, diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index dd898b6f..d4fd8087 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod am_build; +mod am_build; use crate::index::fetcher::*; use crate::index::gucs; @@ -320,7 +320,7 @@ unsafe fn aminsertinner( let payload = kv_to_pointer((key, extra)); let mut chooser = RngChooser(rand::rng()); let bump = bumpalo::Bump::new(); - crate::index::vchordrq::algo::insert( + crate::index::vchordrq::dispatch::insert( opfamily, &index, payload, @@ -372,7 +372,7 @@ pub unsafe extern "C-unwind" fn ambulkdelete( pg_guard_ffi_boundary(|| callback(&mut ctid, callback_state)) } }; - crate::index::vchordrq::algo::bulkdelete(opfamily, &index, check, callback); + crate::index::vchordrq::dispatch::bulkdelete(opfamily, &index, check, callback); stats } @@ -401,7 +401,7 @@ pub unsafe extern "C-unwind" fn amvacuumcleanup( #[cfg(feature = "pg18")] pgrx::pg_sys::vacuum_delay_point(false); }; - crate::index::vchordrq::algo::maintain(opfamily, &index, check); + crate::index::vchordrq::dispatch::maintain(opfamily, &index, check); stats } diff --git a/src/index/vchordrq/build.rs b/src/index/vchordrq/build.rs new file mode 100644 index 00000000..007fc9cb --- /dev/null +++ b/src/index/vchordrq/build.rs @@ -0,0 +1,44 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use simd::{Floating, f16}; +use vector::VectorOwned; +use vector::vect::VectOwned; + +pub type Normalized = Vec; + +pub trait Normalize: VectorOwned { + fn normalize(vector: Self) -> Normalized; + fn denormalize(x: Normalized) -> Self; +} + +impl Normalize for VectOwned { + fn normalize(vector: Self) -> Normalized { + vector.into_vec() + } + + fn denormalize(x: Normalized) -> Self { + Self::new(x) + } +} + +impl Normalize for VectOwned { + fn normalize(vector: Self) -> Normalized { + f16::vector_to_f32(vector.slice()) + } + + fn denormalize(x: Normalized) -> Self { + Self::new(f16::vector_from_f32(&x)) + } +} diff --git a/src/index/vchordrq/algo.rs b/src/index/vchordrq/dispatch.rs similarity index 93% rename from src/index/vchordrq/algo.rs rename to src/index/vchordrq/dispatch.rs index 15802c51..209b4812 100644 --- a/src/index/vchordrq/algo.rs +++ b/src/index/vchordrq/dispatch.rs @@ -12,24 +12,28 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::index::vchordrq::am::am_build::InternalBuild; +use crate::index::vchordrq::build::{Normalize, Normalized}; use crate::index::vchordrq::opclass::Opfamily; -use algo::accessor::{Dot, L2S}; -use algo::prefetcher::*; -use algo::*; +use index::accessor::{Dot, L2S}; +use index::bump::Bump; +use index::fetch::Fetch; +use index::prefetcher::*; +use index::relation::{ + Hints, Page, RelationPrefetch, RelationRead, RelationReadStream, RelationWrite, +}; use simd::f16; use std::collections::BinaryHeap; use std::num::NonZero; use vchordrq::operator::Op; use vchordrq::types::*; -use vchordrq::{Chooser, FastHeap, Opaque}; +use vchordrq::{Chooser, FastHeap}; use vector::VectorOwned; use vector::vect::{VectBorrowed, VectOwned}; pub fn prewarm(opfamily: Opfamily, index: &R, height: i32) -> String where R: RelationRead, - R::Page: Page, + R::Page: Page, { let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; match (opfamily.vector_kind(), opfamily.distance_kind()) { @@ -55,7 +59,7 @@ pub fn bulkdelete( callback: impl Fn(NonZero) -> bool, ) where R: RelationRead + RelationWrite, - R::Page: Page, + R::Page: Page, { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { @@ -80,7 +84,7 @@ pub fn bulkdelete( pub fn maintain(opfamily: Opfamily, index: &R, check: impl Fn()) where R: RelationRead + RelationWrite, - R::Page: Page, + R::Page: Page, { let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; let maintain = match (opfamily.vector_kind(), opfamily.distance_kind()) { @@ -115,35 +119,35 @@ pub fn build( vector_options: VectorOptions, vchordrq_options: VchordrqIndexOptions, index: &R, - structures: Vec>>, + structures: Vec>, ) where R: RelationRead + RelationWrite, - R::Page: Page, + R::Page: Page, { match (vector_options.v, vector_options.d) { (VectorKind::Vecf32, DistanceKind::L2S) => vchordrq::build::<_, Op, L2S>>( vector_options, vchordrq_options, index, - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + map_structures(structures, Normalize::denormalize), ), (VectorKind::Vecf32, DistanceKind::Dot) => vchordrq::build::<_, Op, Dot>>( vector_options, vchordrq_options, index, - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + map_structures(structures, Normalize::denormalize), ), (VectorKind::Vecf16, DistanceKind::L2S) => vchordrq::build::<_, Op, L2S>>( vector_options, vchordrq_options, index, - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + map_structures(structures, Normalize::denormalize), ), (VectorKind::Vecf16, DistanceKind::Dot) => vchordrq::build::<_, Op, Dot>>( vector_options, vchordrq_options, index, - map_structures(structures, |x| InternalBuild::build_from_vecf32(&x)), + map_structures(structures, Normalize::denormalize), ), } } @@ -159,7 +163,7 @@ pub fn insert( bump: &impl Bump, ) where R: RelationRead + RelationWrite, - R::Page: Page, + R::Page: Page, { let make_h1_plain_prefetcher = MakeH1PlainPrefetcherForInsertion { index }; match (vector, opfamily.distance_kind()) { @@ -412,7 +416,7 @@ impl<'b, R> Clone for MakeH0StreamPrefetcher<'b, R> { fn clone(&self) -> Self { Self { index: self.index, - hints: self.hints.clone(), + hints: self.hints, } } } @@ -429,7 +433,7 @@ impl<'b, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'b, R> where S::Item: Fetch<'b>, { - StreamPrefetcher::new(self.index, seq, self.hints.clone()) + StreamPrefetcher::new(self.index, seq, self.hints) } fn is_not_plain(&self) -> bool { diff --git a/src/index/vchordrq/filter.rs b/src/index/vchordrq/filter.rs index 1fec85bb..f63c7729 100644 --- a/src/index/vchordrq/filter.rs +++ b/src/index/vchordrq/filter.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use algo::Sequence; +use index::prefetcher::Sequence; pub struct Filter { sequence: S, diff --git a/src/index/vchordrq/mod.rs b/src/index/vchordrq/mod.rs index 78d81435..400542cf 100644 --- a/src/index/vchordrq/mod.rs +++ b/src/index/vchordrq/mod.rs @@ -12,9 +12,10 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod algo; pub mod am; -pub mod filter; +mod build; +pub mod dispatch; +mod filter; pub mod opclass; -pub mod scanners; +mod scanners; pub mod types; diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index c61d0bab..388396f2 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -15,21 +15,22 @@ use crate::index::fetcher::*; use crate::index::opclass::Sphere; use crate::index::scanners::{Io, SearchBuilder}; -use crate::index::vchordrq::algo::*; +use crate::index::vchordrq::dispatch::*; use crate::index::vchordrq::filter::filter; use crate::index::vchordrq::opclass::Opfamily; use crate::index::vchordrq::scanners::SearchOptions; use crate::recorder::{Recorder, halfvec_out, vector_out}; -use algo::accessor::{Dot, L2S}; -use algo::prefetcher::*; -use algo::*; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; +use index::accessor::{Dot, L2S}; +use index::bump::Bump; +use index::packed::PackedRefMut4; +use index::prefetcher::*; +use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadStream}; use simd::f16; use std::num::NonZero; -use vchordrq::operator::{self}; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; -use vchordrq::*; +use vchordrq::{RerankMethod, default_search, how, rerank_heap, rerank_index}; use vector::VectorOwned; use vector::vect::VectOwned; @@ -110,18 +111,20 @@ impl SearchBuilder for DefaultBuilder { let Some(vector) = vector else { return Box::new(std::iter::empty()) as Box>; }; + let search_hints = Hints::default().full(true); + let rerank_hints = Hints::default().full(false); let make_h1_plain_prefetcher = MakeH1PlainPrefetcher { index }; let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; let make_h0_simple_prefetcher = MakeH0SimplePrefetcher { index }; let make_h0_stream_prefetcher = MakeH0StreamPrefetcher { index, - hints: Hints::default().full(true), + hints: search_hints, }; let f = move |(distance, payload)| (opfamily.output(distance), payload); let iter: Box)>> = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { - type Op = operator::Op, L2S>; + type Op = vchordrq::operator::Op, L2S>; let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { vector } else { @@ -195,8 +198,7 @@ impl SearchBuilder for DefaultBuilder { Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { @@ -209,8 +211,7 @@ impl SearchBuilder for DefaultBuilder { tuple.filter() }); let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Heap, _, false) => { @@ -261,7 +262,7 @@ impl SearchBuilder for DefaultBuilder { } } (VectorKind::Vecf32, DistanceKind::Dot) => { - type Op = operator::Op, Dot>; + type Op = vchordrq::operator::Op, Dot>; let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { vector } else { @@ -335,8 +336,7 @@ impl SearchBuilder for DefaultBuilder { Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { @@ -349,8 +349,7 @@ impl SearchBuilder for DefaultBuilder { tuple.filter() }); let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Heap, _, false) => { @@ -401,7 +400,7 @@ impl SearchBuilder for DefaultBuilder { } } (VectorKind::Vecf16, DistanceKind::L2S) => { - type Op = operator::Op, L2S>; + type Op = vchordrq::operator::Op, L2S>; let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { vector } else { @@ -475,8 +474,7 @@ impl SearchBuilder for DefaultBuilder { Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { @@ -489,8 +487,7 @@ impl SearchBuilder for DefaultBuilder { tuple.filter() }); let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Heap, _, false) => { @@ -541,7 +538,7 @@ impl SearchBuilder for DefaultBuilder { } } (VectorKind::Vecf16, DistanceKind::Dot) => { - type Op = operator::Op, Dot>; + type Op = vchordrq::operator::Op, Dot>; let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { vector } else { @@ -615,8 +612,7 @@ impl SearchBuilder for DefaultBuilder { Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Index, Io::Stream, true) => { @@ -629,8 +625,7 @@ impl SearchBuilder for DefaultBuilder { tuple.filter() }); let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); Box::new(rerank_index::(unprojected, prefetcher).map(f)) } (RerankMethod::Heap, _, false) => { diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 65603995..934a2e94 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -14,23 +14,25 @@ use crate::index::fetcher::*; use crate::index::scanners::{Io, SearchBuilder}; -use crate::index::vchordrq::algo::*; +use crate::index::vchordrq::dispatch::*; use crate::index::vchordrq::filter::filter; use crate::index::vchordrq::opclass::Opfamily; use crate::index::vchordrq::scanners::SearchOptions; use crate::recorder::Recorder; -use algo::accessor::Dot; -use algo::prefetcher::*; -use algo::*; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; use distance::Distance; +use index::accessor::Dot; +use index::bump::Bump; +use index::packed::PackedRefMut8; +use index::prefetcher::*; +use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadStream}; use simd::f16; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZero; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; -use vchordrq::*; +use vchordrq::{RerankMethod, how, maxsim_search, rerank_index}; use vector::VectorOwned; use vector::vect::VectOwned; @@ -101,12 +103,14 @@ impl SearchBuilder for MaxsimBuilder { pgrx::error!("maxsim search with rerank_in_table is not supported"); } assert!(matches!(opfamily.distance_kind(), DistanceKind::Dot)); + let search_hints = Hints::default().full(true); + let rerank_hints = Hints::default().full(false); let make_h1_plain_prefetcher = MakeH1PlainPrefetcher { index }; let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; let make_h0_simple_prefetcher = MakeH0SimplePrefetcher { index }; let make_h0_stream_prefetcher = MakeH0StreamPrefetcher { index, - hints: Hints::default().full(true), + hints: search_hints, }; let n = vectors.len(); let accu_map = |(Reverse(distance), AlwaysEqual(payload))| (distance, payload); @@ -117,7 +121,7 @@ impl SearchBuilder for MaxsimBuilder { )| (rough, payload); let iter: Box> = match opfamily.vector_kind() { VectorKind::Vecf32 => { - type Op = operator::Op, Dot>; + type Op = vchordrq::operator::Op, Dot>; let unprojected = vectors .into_iter() .map(|vector| { @@ -225,7 +229,7 @@ impl SearchBuilder for MaxsimBuilder { } (Io::Stream, false) => { let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + StreamPrefetcher::new(index, sequence, rerank_hints); let mut reranker = rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); @@ -244,7 +248,7 @@ impl SearchBuilder for MaxsimBuilder { }); let sequence = filter(sequence, predicate); let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + StreamPrefetcher::new(index, sequence, rerank_hints); let mut reranker = rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); @@ -261,7 +265,7 @@ impl SearchBuilder for MaxsimBuilder { })) } VectorKind::Vecf16 => { - type Op = operator::Op, Dot>; + type Op = vchordrq::operator::Op, Dot>; let unprojected = vectors .into_iter() .map(|vector| { @@ -369,7 +373,7 @@ impl SearchBuilder for MaxsimBuilder { } (Io::Stream, false) => { let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + StreamPrefetcher::new(index, sequence, rerank_hints); let mut reranker = rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); @@ -388,7 +392,7 @@ impl SearchBuilder for MaxsimBuilder { }); let sequence = filter(sequence, predicate); let prefetcher = - StreamPrefetcher::new(index, sequence, Hints::default()); + StreamPrefetcher::new(index, sequence, rerank_hints); let mut reranker = rerank_index::(unprojected[i].clone(), prefetcher); accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); diff --git a/src/lib.rs b/src/lib.rs index 25ff506d..594d7824 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,8 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. #![allow(unsafe_code)] -#![allow(unused_crate_dependencies)] -#![warn(ffi_unwind_calls)] +#![deny(ffi_unwind_calls)] mod datatype; mod index; From 1487d647bd8fef2ac74611a264478303d209dd95 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 14 Oct 2025 18:18:32 +0800 Subject: [PATCH 232/324] feat: parallel maintain (#357) PostgreSQL does not support nested parallelism, so this patch only takes effect during the build process. --------- Signed-off-by: usamoi --- .github/workflows/check.yml | 59 +++- crates/vchordrq/src/insert.rs | 8 +- crates/vchordrq/src/lib.rs | 9 +- crates/vchordrq/src/maintain.rs | 11 +- src/index/vchordg/am/am_build.rs | 23 +- src/index/vchordg/am/mod.rs | 3 + src/index/vchordrq/am/am_build.rs | 364 ++++++++++++++++++---- src/index/vchordrq/am/am_vacuumcleanup.rs | 70 +++++ src/index/vchordrq/am/mod.rs | 39 +-- src/index/vchordrq/dispatch.rs | 40 ++- tests/vchordg/vacuum.slt | 77 +++++ tests/vchordg/vacuum_parallel.slt | 53 ++++ tests/vchordrq/pushdown_plan.slt | 60 +++- tests/vchordrq/vacuum.slt | 80 +++++ tests/vchordrq/vacuum_parallel.slt | 65 ++++ 15 files changed, 827 insertions(+), 134 deletions(-) create mode 100644 src/index/vchordrq/am/am_vacuumcleanup.rs create mode 100644 tests/vchordg/vacuum.slt create mode 100644 tests/vchordg/vacuum_parallel.slt create mode 100644 tests/vchordrq/vacuum.slt create mode 100644 tests/vchordrq/vacuum_parallel.slt diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index ac317acf..11d515e0 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -209,8 +209,16 @@ jobs: echo "local all all trust" | sudo tee /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf echo "host all all 127.0.0.1/32 trust" | sudo tee -a /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf echo "host all all ::1/128 trust" | sudo tee -a /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + sudo mkdir -p /etc/systemd/system/postgresql@.service.d/ + echo "[Service]" | sudo tee /etc/systemd/system/postgresql@.service.d/limit.conf + echo "LimitNOFILE=infinity" | sudo tee -a /etc/systemd/system/postgresql@.service.d/limit.conf + echo "LimitMEMLOCK=infinity" | sudo tee -a /etc/systemd/system/postgresql@.service.d/limit.conf sudo -iu postgres createuser -s -r $USER sudo -iu postgres createdb -O $USER $USER + if [ "${{ matrix.version }}" -ge "18" ]; then + sudo -iu postgres psql -c 'ALTER SYSTEM SET io_method = io_uring' + fi + sudo -iu postgres psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo systemctl stop postgresql @@ -252,6 +260,11 @@ jobs: sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi + - name: Logging + if: always() + run: | + cat /var/log/postgresql/postgresql-${{ matrix.version }}-main.log + psql_macos: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_macos')) || @@ -287,6 +300,7 @@ jobs: for i in {1..60}; do [ -S /tmp/.s.PGSQL.5432 ] && echo "PostgreSQL ready" && break || sleep 1; done [ -S /tmp/.s.PGSQL.5432 ] || echo "PostgreSQL socket not found after 60 seconds" $(brew --prefix postgresql@${{ matrix.version }})/bin/createdb -O $USER $USER + $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' brew services stop postgresql@${{ matrix.version }} echo PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV @@ -338,6 +352,11 @@ jobs: sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi + - name: Logging + if: always() + run: | + cat $(brew --prefix)/var/log/postgresql@${{ matrix.version }}.log + psql_windows: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_windows')) || @@ -391,9 +410,10 @@ jobs: Add-Content -Path $env:GITHUB_ENV -Value "PGRX_PG_CONFIG_PATH=D:\postgresql-install\pgsql\bin\pg_config.exe" Add-Content -Path $env:GITHUB_ENV -Value "PG_CONFIG=D:\postgresql-install\pgsql\bin\pg_config.exe" D:\postgresql-install\pgsql\bin\initdb.exe -D D:\postgresql-install\pgsql\data -U postgres - D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data + D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data -l D:\postgresql-install\postgresql.log D:\postgresql-install\pgsql\bin\createuser.exe -U postgres -s -r $env:USERNAME D:\postgresql-install\pgsql\bin\createdb.exe -O $env:USERNAME $env:USERNAME + D:\postgresql-install\pgsql\bin\psql.exe -c 'ALTER SYSTEM SET max_worker_processes = 1024' D:\postgresql-install\pgsql\bin\psql.exe -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' D:\postgresql-install\pgsql\bin\pg_ctl.exe stop -D D:\postgresql-install\pgsql\data @@ -435,7 +455,7 @@ jobs: run: | 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } - D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data + D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data -l D:\postgresql-install\postgresql.log D:\postgresql-install\pgsql\bin\psql.exe -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - name: Sqllogictest @@ -447,6 +467,11 @@ jobs: sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} } + - name: Logging + if: always() + run: | + Get-Content -Path 'D:\postgresql-install\postgresql.log' + psql_alpine: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_alpine')) || @@ -488,13 +513,20 @@ jobs: echo PGRX_PG_CONFIG_PATH=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV echo PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV + sudo touch /var/log/postgresql.log + sudo chmod 644 /var/log/postgresql.log + sudo chown postgres /var/log/postgresql.log install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql/data install --verbose --directory --owner postgres --group postgres --mode 3777 /run/postgresql sudo -iu postgres initdb -D /var/lib/postgresql/data - sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data + sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data -l /var/log/postgresql.log sudo -iu postgres createuser -s -r $USER sudo -iu postgres createdb -O $USER $USER + if [ "${{ matrix.version }}" -ge "18" ]; then + sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET io_method = io_uring' + fi + sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET max_worker_processes = 1024' sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo -iu postgres pg_ctl stop -D /var/lib/postgresql/data @@ -541,7 +573,7 @@ jobs: - name: Service run: | - sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data + sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data -l /var/log/postgresql.log psql -d $USER -U $USER -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - name: Sqllogictest @@ -553,6 +585,11 @@ jobs: sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi + - name: Logging + if: always() + run: | + cat /var/log/postgresql.log + check_debian: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +check_debian')) || @@ -678,6 +715,15 @@ jobs: sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main start sudo chroot /sysroot sudo -iu postgres createuser -s -r $USER sudo chroot /sysroot sudo -iu postgres createdb -O $USER $USER + sudo mkdir -p /sysroot/etc/systemd/system/postgresql@.service.d/ + echo "[Service]" | sudo tee /sysroot/etc/systemd/system/postgresql@.service.d/limit.conf + echo "LimitNOFILE=infinity" | sudo tee -a /sysroot/etc/systemd/system/postgresql@.service.d/limit.conf + echo "LimitMEMLOCK=infinity" | sudo tee -a /sysroot/etc/systemd/system/postgresql@.service.d/limit.conf + if [ "${{ matrix.version }}" -ge "18" ]; then + # io_uring is not supported on qemu-user + sudo chroot /sysroot sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET io_method = worker' + fi + sudo chroot /sysroot sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET max_worker_processes = 1024' sudo chroot /sysroot sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main stop @@ -761,3 +807,8 @@ jobs: if [ "${{ matrix.version }}" -ge "17" ]; then sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi + + - name: Logging + if: always() + run: | + cat /sysroot/var/log/postgresql/postgresql-${{ matrix.version }}-main.log diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index 7d4abfa0..bd30f62f 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -15,7 +15,7 @@ use crate::linked_vec::LinkedVec; use crate::operator::*; use crate::tuples::*; -use crate::{Chooser, Opaque, centroids, tape, vectors}; +use crate::{Opaque, centroids, tape, vectors}; use always_equal::AlwaysEqual; use distance::Distance; use index::accessor::{FunctionalAccessor, LAccess}; @@ -28,13 +28,17 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vector::{VectorBorrowed, VectorOwned}; +pub trait InsertChooser { + fn choose(&mut self, n: NonZero) -> usize; +} + type Extra1<'b> = &'b mut (u32, f32, u16, BorrowedIter<'b>); pub fn insert_vector( index: &R, payload: NonZero, vector: ::Borrowed<'_>, - chooser: &mut impl Chooser, + chooser: &mut impl InsertChooser, skip_search: bool, ) -> (Vec, u16) where diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index f176809e..50c9c2b6 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -39,13 +39,12 @@ pub use bulkdelete::{bulkdelete, bulkdelete_vectors}; pub use cache::cache; pub use cost::cost; pub use fast_heap::FastHeap; -pub use insert::{insert, insert_vector}; -pub use maintain::maintain; +pub use insert::{InsertChooser, insert, insert_vector}; +pub use maintain::{MaintainChooser, maintain}; pub use prewarm::prewarm; pub use rerank::{how, rerank_heap, rerank_index}; pub use search::{default_search, maxsim_search}; -use std::num::NonZero; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] @@ -67,10 +66,6 @@ pub(crate) struct Branch { pub extra: T, } -pub trait Chooser { - fn choose(&mut self, n: NonZero) -> usize; -} - #[derive(Debug, Clone, Copy)] pub enum RerankMethod { Index, diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs index f7ee7e01..e9492679 100644 --- a/crates/vchordrq/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -25,6 +25,10 @@ use index::relation::{ use rabitq::packing::unpack; use std::cell::RefCell; +pub trait MaintainChooser { + fn choose(&mut self, i: usize) -> bool; +} + pub struct Maintain { pub number_of_formerly_allocated_pages: usize, pub number_of_freshly_allocated_pages: usize, @@ -34,6 +38,7 @@ pub struct Maintain { pub fn maintain<'b, R: RelationRead + RelationWrite, O: Operator>( index: &'b R, mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, + chooser: &mut impl MaintainChooser, check: impl Fn(), ) -> Maintain where @@ -79,7 +84,11 @@ where number_of_freshly_allocated_pages: 0, }); - for first in state { + for (idx, first) in state.into_iter().enumerate() { + if !chooser.choose(idx) { + continue; + } + let mut jump_guard = index.write(first, false); let jump_bytes = jump_guard.get_mut(1).expect("data corruption"); let mut jump_tuple = JumpTuple::deserialize_mut(jump_bytes); diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index c2ca4888..84498277 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -19,6 +19,7 @@ use crate::index::vchordg::opclass::{Opfamily, opfamily}; use crate::index::vchordg::types::VchordgIndexingOptions; use pgrx::pg_sys::{Datum, ItemPointerData}; use std::ffi::CStr; +use std::marker::PhantomData; use vchordg::types::*; #[derive(Debug, Clone, Copy)] @@ -148,10 +149,12 @@ impl Heap { } #[derive(Debug, Clone)] -struct PostgresReporter {} +struct PostgresReporter { + _phantom: PhantomData<*mut ()>, +} impl PostgresReporter { - fn phase(&mut self, phase: BuildPhase) { + fn phase(&self, phase: BuildPhase) { unsafe { pgrx::pg_sys::pgstat_progress_update_param( pgrx::pg_sys::PROGRESS_CREATEIDX_SUBPHASE as _, @@ -159,7 +162,7 @@ impl PostgresReporter { ); } } - fn tuples_total(&mut self, tuples_total: u64) { + fn tuples_total(&self, tuples_total: u64) { unsafe { pgrx::pg_sys::pgstat_progress_update_param( pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, @@ -167,7 +170,7 @@ impl PostgresReporter { ); } } - fn tuples_done(&mut self, tuples_done: u64) { + fn tuples_done(&self, tuples_done: u64) { unsafe { pgrx::pg_sys::pgstat_progress_update_param( pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, @@ -198,12 +201,15 @@ pub unsafe extern "C-unwind" fn ambuild( pgrx::warning!("warning while validating options: {errors}"); } let index = unsafe { PostgresRelation::new(index_relation) }; - let mut reporter = PostgresReporter {}; + let reporter = PostgresReporter { + _phantom: PhantomData, + }; crate::index::vchordg::dispatch::build(vector_options, vchordg_options.index, &index); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); let cached = vchordg_cached::VchordgCached::_0 {}.serialize(); if let Some(leader) = unsafe { VchordgLeader::enter( + c"vchordg_parallel_build_main", heap_relation, index_relation, (*index_info).ii_Concurrent, @@ -292,6 +298,7 @@ struct VchordgLeader { impl VchordgLeader { pub unsafe fn enter( + main: &'static CStr, heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, isconcurrent: bool, @@ -332,11 +339,7 @@ impl VchordgLeader { pgrx::pg_sys::EnterParallelMode(); } let pcxt = unsafe { - pgrx::pg_sys::CreateParallelContext( - c"vchord".as_ptr(), - c"vchordg_parallel_build_main".as_ptr(), - request, - ) + pgrx::pg_sys::CreateParallelContext(c"vchord".as_ptr(), main.as_ptr(), request) }; let snapshot = if isconcurrent { diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index c7ef9f8c..f7ae5984 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -123,6 +123,9 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { am_routine.amgettuple = Some(amgettuple); am_routine.amendscan = Some(amendscan); + am_routine.amparallelvacuumoptions = pgrx::pg_sys::VACUUM_OPTION_PARALLEL_BULKDEL as u8 + | pgrx::pg_sys::VACUUM_OPTION_PARALLEL_CLEANUP as u8; + am_routine }; diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index c9aaa0b7..8d11b38b 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -26,10 +26,11 @@ use pgrx::pg_sys::{Datum, ItemPointerData}; use rand::Rng; use simd::Floating; use std::ffi::CStr; +use std::marker::PhantomData; use std::num::NonZero; use std::ops::Deref; -use vchordrq::Chooser; use vchordrq::types::*; +use vchordrq::{InsertChooser, MaintainChooser}; use vector::vect::VectOwned; #[derive(Debug, Clone, Copy)] @@ -233,10 +234,12 @@ impl Heap { } #[derive(Debug, Clone)] -struct PostgresReporter {} +struct PostgresReporter { + _phantom: PhantomData<*mut ()>, +} impl PostgresReporter { - fn phase(&mut self, phase: BuildPhase) { + fn phase(&self, phase: BuildPhase) { unsafe { pgrx::pg_sys::pgstat_progress_update_param( pgrx::pg_sys::PROGRESS_CREATEIDX_SUBPHASE as _, @@ -244,7 +247,7 @@ impl PostgresReporter { ); } } - fn tuples_total(&mut self, tuples_total: u64) { + fn tuples_total(&self, tuples_total: u64) { unsafe { pgrx::pg_sys::pgstat_progress_update_param( pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, @@ -252,7 +255,7 @@ impl PostgresReporter { ); } } - fn tuples_done(&mut self, tuples_done: u64) { + fn tuples_done(&self, tuples_done: u64) { unsafe { pgrx::pg_sys::pgstat_progress_update_param( pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, @@ -285,7 +288,9 @@ pub unsafe extern "C-unwind" fn ambuild( opfamily, scan: std::ptr::null_mut(), }; - let mut reporter = PostgresReporter {}; + let reporter = PostgresReporter { + _phantom: PhantomData, + }; reporter.tuples_total(unsafe { (*(*index_relation).rd_rel).reltuples as u64 }); let mut structures = match vchordrq_options.build.source.clone() { VchordrqBuildSourceOptions::Default(default_build) => { @@ -330,7 +335,7 @@ pub unsafe extern "C-unwind" fn ambuild( samples }; reporter.tuples_total(tuples_total); - make_internal_build(vector_options, internal_build, samples, &mut reporter) + make_internal_build(vector_options, internal_build, samples, &reporter) } VchordrqBuildSourceOptions::External(external_build) => { reporter.phase(BuildPhase::from_code(BuildPhaseCode::ExternalBuild)); @@ -349,7 +354,6 @@ pub unsafe extern "C-unwind" fn ambuild( &index, structures, ); - reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); let cached = if vchordrq_options.build.pin >= 0 { let mut trace = vchordrq::cache(&index, vchordrq_options.build.pin); trace.sort(); @@ -371,6 +375,7 @@ pub unsafe extern "C-unwind" fn ambuild( .serialize(); if let Some(leader) = unsafe { VchordrqLeader::enter( + c"vchordrq_parallel_build_main", heap_relation, index_relation, (*index_info).ii_Concurrent, @@ -379,6 +384,7 @@ pub unsafe extern "C-unwind" fn ambuild( } { drop(cached); unsafe { + leader.wait(); parallel_build( index_relation, heap_relation, @@ -389,28 +395,113 @@ pub unsafe extern "C-unwind" fn ambuild( |indtuples| { reporter.tuples_done(indtuples); }, + || { + #[allow(clippy::needless_late_init)] + let order; + // enter the barrier + let shared = leader.vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).nparticipants = leader.nparticipants as u32; + order = (*shared).barrier_enter_0 as u32; + (*shared).barrier_enter_0 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_0, + ); + // leave the barrier + let total = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_enter_0 == total { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_enter_0, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_leave_0 = true; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_leave_0, + ); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); + order + }, + |indtuples| { + reporter.tuples_done(indtuples); + reporter.tuples_total(indtuples); + // enter the barrier + let shared = leader.vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_enter_1 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_1, + ); + // leave the barrier + let total = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_enter_1 == total { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_enter_1, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_leave_1 = true; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_leave_1, + ); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); + }, + || { + // enter the barrier + let shared = leader.vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_enter_2 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_2, + ); + // leave the barrier + let total = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_enter_2 == total { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_enter_2, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_leave_2 = true; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_leave_2, + ); + }, ); - leader.wait(); - let nparticipants = leader.nparticipants; - loop { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordrqshared).mutex); - if (*leader.vchordrqshared).workers_done == nparticipants { - pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); - break; - } - pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordrqshared).mutex); - pgrx::pg_sys::ConditionVariableSleep( - &raw mut (*leader.vchordrqshared).condvar_workers_done, - pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, - ); - } - reporter.tuples_done((*leader.vchordrqshared).indtuples); - reporter.tuples_total((*leader.vchordrqshared).indtuples); - pgrx::pg_sys::ConditionVariableCancelSleep(); } } else { unsafe { - let indtuples = sequential_build( + sequential_build( index_relation, heap_relation, index_info, @@ -418,15 +509,18 @@ pub unsafe extern "C-unwind" fn ambuild( |indtuples| { reporter.tuples_done(indtuples); }, + || { + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); + }, + |indtuples| { + reporter.tuples_done(indtuples); + reporter.tuples_total(indtuples); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); + }, + || {}, ); - reporter.tuples_total(indtuples); } } - let check = || { - pgrx::check_for_interrupts!(); - }; - reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); - crate::index::vchordrq::dispatch::maintain(opfamily, &index, check); unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } } @@ -438,12 +532,22 @@ struct VchordrqShared { /* locking */ mutex: pgrx::pg_sys::slock_t, - condvar_workers_done: pgrx::pg_sys::ConditionVariable, + condvar_barrier_enter_0: pgrx::pg_sys::ConditionVariable, + condvar_barrier_leave_0: pgrx::pg_sys::ConditionVariable, + condvar_barrier_enter_1: pgrx::pg_sys::ConditionVariable, + condvar_barrier_leave_1: pgrx::pg_sys::ConditionVariable, + condvar_barrier_enter_2: pgrx::pg_sys::ConditionVariable, + condvar_barrier_leave_2: pgrx::pg_sys::ConditionVariable, /* mutable state */ - workers_ready: i32, - workers_done: i32, + barrier_enter_0: i32, + nparticipants: u32, indtuples: u64, + barrier_leave_0: bool, + barrier_enter_1: i32, + barrier_leave_1: bool, + barrier_enter_2: i32, + barrier_leave_2: bool, } mod vchordrq_cached { @@ -597,6 +701,7 @@ struct VchordrqLeader { impl VchordrqLeader { pub unsafe fn enter( + main: &'static CStr, heap_relation: pgrx::pg_sys::Relation, index_relation: pgrx::pg_sys::Relation, isconcurrent: bool, @@ -637,11 +742,7 @@ impl VchordrqLeader { pgrx::pg_sys::EnterParallelMode(); } let pcxt = unsafe { - pgrx::pg_sys::CreateParallelContext( - c"vchord".as_ptr(), - c"vchordrq_parallel_build_main".as_ptr(), - request, - ) + pgrx::pg_sys::CreateParallelContext(c"vchord".as_ptr(), main.as_ptr(), request) }; let snapshot = if isconcurrent { @@ -687,13 +788,28 @@ impl VchordrqLeader { heaprelid: (*heap_relation).rd_id, indexrelid: (*index_relation).rd_id, isconcurrent, - condvar_workers_done: std::mem::zeroed(), + nparticipants: 0, + condvar_barrier_enter_0: std::mem::zeroed(), + condvar_barrier_leave_0: std::mem::zeroed(), + condvar_barrier_enter_1: std::mem::zeroed(), + condvar_barrier_leave_1: std::mem::zeroed(), + condvar_barrier_enter_2: std::mem::zeroed(), + condvar_barrier_leave_2: std::mem::zeroed(), + barrier_enter_0: 0, + barrier_leave_0: false, + barrier_enter_1: 0, + barrier_leave_1: false, + barrier_enter_2: 0, + barrier_leave_2: false, mutex: std::mem::zeroed(), - workers_ready: 0, - workers_done: 0, indtuples: 0, }); - pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_workers_done); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_enter_0); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_leave_0); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_enter_1); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_leave_1); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_enter_2); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_leave_2); pgrx::pg_sys::SpinLockInit(&raw mut (*vchordrqshared).mutex); vchordrqshared }; @@ -813,6 +929,82 @@ pub unsafe extern "C-unwind" fn vchordrq_parallel_build_main( vchordrqshared, vchordrqcached, |_| (), + || { + #[allow(clippy::needless_late_init)] + let order; + // enter the barrier + let shared = vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + order = (*shared).barrier_enter_0 as u32; + (*shared).barrier_enter_0 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_0, + ); + // leave the barrier + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_leave_0 { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_leave_0, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + order + }, + |_| { + // enter the barrier + let shared = vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_enter_1 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_1, + ); + // leave the barrier + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_leave_1 { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_leave_1, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + }, + || { + // enter the barrier + let shared = vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_enter_2 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_2, + ); + // leave the barrier + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_leave_2 { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_leave_2, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + }, ); } @@ -830,18 +1022,12 @@ unsafe fn parallel_build( vchordrqshared: *mut VchordrqShared, vchordrqcached: *const u8, mut callback: impl FnMut(u64), + sync_0: impl FnOnce() -> u32, + sync_1: impl FnOnce(u64), + sync_2: impl FnOnce(), ) { use vchordrq_cached::VchordrqCachedReader; - let wid; - - unsafe { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - wid = (*vchordrqshared).workers_ready as u32; - (*vchordrqshared).workers_ready += 1; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - } - let cached = VchordrqCachedReader::deserialize_ref(unsafe { let bytes = (vchordrqcached as *const u64).read_unaligned(); std::slice::from_raw_parts(vchordrqcached.add(8), bytes as _) @@ -860,19 +1046,35 @@ unsafe fn parallel_build( }; struct IdChooser(u32); - impl Chooser for IdChooser { + impl InsertChooser for IdChooser { fn choose(&mut self, n: NonZero) -> usize { self.0 as usize % n.get() } } + struct ChooseSome { + n: usize, + k: usize, + } + impl MaintainChooser for ChooseSome { + fn choose(&mut self, i: usize) -> bool { + i % self.n == self.k + } + } + + let check = || { + pgrx::check_for_interrupts!(); + }; + + let order = sync_0(); + match cached { VchordrqCachedReader::_0(_) => { heap.traverse(true, |(ctid, store)| { for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - let mut chooser = IdChooser(wid); + let mut chooser = IdChooser(order); let bump = bumpalo::Bump::new(); crate::index::vchordrq::dispatch::insert( opfamily, @@ -900,13 +1102,13 @@ unsafe fn parallel_build( VchordrqCachedReader::_1(cached) => { let index = CachingRelation { cache: cached, - relation: index, + relation: index.clone(), }; heap.traverse(true, |(ctid, store)| { for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); - let mut chooser = IdChooser(wid); + let mut chooser = IdChooser(order); let bump = bumpalo::Bump::new(); crate::index::vchordrq::dispatch::insert( opfamily, @@ -932,12 +1134,16 @@ unsafe fn parallel_build( }); } } - unsafe { - pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); - (*vchordrqshared).workers_done += 1; - pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); - pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordrqshared).condvar_workers_done); - } + + sync_1(unsafe { (*vchordrqshared).indtuples }); + + let mut chooser = ChooseSome { + n: unsafe { (*vchordrqshared).nparticipants as usize }, + k: order as usize, + }; + crate::index::vchordrq::dispatch::maintain(opfamily, &index, &mut chooser, check); + + sync_2(); } unsafe fn sequential_build( @@ -946,7 +1152,10 @@ unsafe fn sequential_build( index_info: *mut pgrx::pg_sys::IndexInfo, vchordrqcached: &[u8], mut callback: impl FnMut(u64), -) -> u64 { + sync_0: impl FnOnce(), + sync_1: impl FnOnce(u64), + sync_2: impl FnOnce(), +) { use vchordrq_cached::VchordrqCachedReader; let cached = VchordrqCachedReader::deserialize_ref(vchordrqcached); @@ -963,12 +1172,25 @@ unsafe fn sequential_build( }; struct ChooseZero; - impl Chooser for ChooseZero { + impl InsertChooser for ChooseZero { fn choose(&mut self, _: NonZero) -> usize { 0 } } + struct ChooseAll; + impl MaintainChooser for ChooseAll { + fn choose(&mut self, _: usize) -> bool { + true + } + } + + let check = || { + pgrx::check_for_interrupts!(); + }; + + sync_0(); + let mut indtuples = 0; match cached { VchordrqCachedReader::_0(_) => { @@ -996,7 +1218,7 @@ unsafe fn sequential_build( VchordrqCachedReader::_1(cached) => { let index = CachingRelation { cache: cached, - relation: index, + relation: index.clone(), }; heap.traverse(true, |(ctid, store)| { for (vector, extra) in store { @@ -1020,7 +1242,13 @@ unsafe fn sequential_build( }); } } - indtuples + + sync_1(indtuples); + + let mut chooser = ChooseAll; + crate::index::vchordrq::dispatch::maintain(opfamily, &index, &mut chooser, check); + + sync_2(); } #[pgrx::pg_guard] @@ -1100,7 +1328,7 @@ fn make_internal_build( vector_options: VectorOptions, internal_build: VchordrqInternalBuildOptions, samples: Vec, - reporter: &mut PostgresReporter, + reporter: &PostgresReporter, ) -> Vec> { use std::iter::once; let mut result = Vec::>::new(); diff --git a/src/index/vchordrq/am/am_vacuumcleanup.rs b/src/index/vchordrq/am/am_vacuumcleanup.rs new file mode 100644 index 00000000..c7e4f835 --- /dev/null +++ b/src/index/vchordrq/am/am_vacuumcleanup.rs @@ -0,0 +1,70 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::index::vchordrq::am::PostgresRelation; +use crate::index::vchordrq::opclass::opfamily; +use vchordrq::MaintainChooser; + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amvacuumcleanup( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let index_relation = unsafe { (*info).index }; + unsafe { + sequential_vacuumcleanup(index_relation, || (), || ()); + } + stats +} + +unsafe fn sequential_vacuumcleanup( + index_relation: pgrx::pg_sys::Relation, + sync_0: impl FnOnce(), + sync_1: impl FnOnce(), +) { + struct ChooseAll; + impl MaintainChooser for ChooseAll { + fn choose(&mut self, _: usize) -> bool { + true + } + } + + let opfamily = unsafe { opfamily(index_relation) }; + let index = unsafe { PostgresRelation::new(index_relation) }; + let check = || unsafe { + #[cfg(any( + feature = "pg13", + feature = "pg14", + feature = "pg15", + feature = "pg16", + feature = "pg17" + ))] + pgrx::pg_sys::vacuum_delay_point(); + #[cfg(feature = "pg18")] + pgrx::pg_sys::vacuum_delay_point(false); + }; + + sync_0(); + + let mut chooser = ChooseAll; + crate::index::vchordrq::dispatch::maintain(opfamily, &index, &mut chooser, check); + + sync_1(); +} diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index d4fd8087..69641f13 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. mod am_build; +mod am_vacuumcleanup; use crate::index::fetcher::*; use crate::index::gucs; @@ -30,7 +31,7 @@ use std::num::NonZero; use std::ops::DerefMut; use std::ptr::NonNull; use std::sync::OnceLock; -use vchordrq::Chooser; +use vchordrq::InsertChooser; #[repr(C)] struct Reloption { @@ -118,13 +119,16 @@ const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { am_routine.ambuildempty = Some(am_build::ambuildempty); am_routine.aminsert = Some(aminsert); am_routine.ambulkdelete = Some(ambulkdelete); - am_routine.amvacuumcleanup = Some(amvacuumcleanup); + am_routine.amvacuumcleanup = Some(am_vacuumcleanup::amvacuumcleanup); am_routine.ambeginscan = Some(ambeginscan); am_routine.amrescan = Some(amrescan); am_routine.amgettuple = Some(amgettuple); am_routine.amendscan = Some(amendscan); + am_routine.amparallelvacuumoptions = pgrx::pg_sys::VACUUM_OPTION_PARALLEL_BULKDEL as u8 + | pgrx::pg_sys::VACUUM_OPTION_PARALLEL_CLEANUP as u8; + am_routine }; @@ -304,7 +308,7 @@ unsafe fn aminsertinner( ctid: pgrx::pg_sys::ItemPointer, ) -> bool { struct RngChooser(T); - impl Chooser for RngChooser { + impl InsertChooser for RngChooser { fn choose(&mut self, n: NonZero) -> usize { rand::Rng::random_range(&mut self.0, 0..n.get()) } @@ -376,35 +380,6 @@ pub unsafe extern "C-unwind" fn ambulkdelete( stats } -#[pgrx::pg_guard] -pub unsafe extern "C-unwind" fn amvacuumcleanup( - info: *mut pgrx::pg_sys::IndexVacuumInfo, - stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, -) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { - let mut stats = stats; - if stats.is_null() { - stats = unsafe { - pgrx::pg_sys::palloc0(size_of::()).cast() - }; - } - let opfamily = unsafe { opfamily((*info).index) }; - let index = unsafe { PostgresRelation::new((*info).index) }; - let check = || unsafe { - #[cfg(any( - feature = "pg13", - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17" - ))] - pgrx::pg_sys::vacuum_delay_point(); - #[cfg(feature = "pg18")] - pgrx::pg_sys::vacuum_delay_point(false); - }; - crate::index::vchordrq::dispatch::maintain(opfamily, &index, check); - stats -} - #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn ambeginscan( index_relation: pgrx::pg_sys::Relation, diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs index 209b4812..cdbb0608 100644 --- a/src/index/vchordrq/dispatch.rs +++ b/src/index/vchordrq/dispatch.rs @@ -26,7 +26,7 @@ use std::collections::BinaryHeap; use std::num::NonZero; use vchordrq::operator::Op; use vchordrq::types::*; -use vchordrq::{Chooser, FastHeap}; +use vchordrq::{FastHeap, InsertChooser, MaintainChooser}; use vector::VectorOwned; use vector::vect::{VectBorrowed, VectOwned}; @@ -81,24 +81,48 @@ pub fn bulkdelete( } } -pub fn maintain(opfamily: Opfamily, index: &R, check: impl Fn()) -where +pub fn maintain( + opfamily: Opfamily, + index: &R, + chooser: &mut impl MaintainChooser, + check: impl Fn(), +) where R: RelationRead + RelationWrite, R::Page: Page, { let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; let maintain = match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { - vchordrq::maintain::<_, Op, L2S>>(index, make_h0_plain_prefetcher, check) + vchordrq::maintain::<_, Op, L2S>>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ) } (VectorKind::Vecf32, DistanceKind::Dot) => { - vchordrq::maintain::<_, Op, Dot>>(index, make_h0_plain_prefetcher, check) + vchordrq::maintain::<_, Op, Dot>>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ) } (VectorKind::Vecf16, DistanceKind::L2S) => { - vchordrq::maintain::<_, Op, L2S>>(index, make_h0_plain_prefetcher, check) + vchordrq::maintain::<_, Op, L2S>>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ) } (VectorKind::Vecf16, DistanceKind::Dot) => { - vchordrq::maintain::<_, Op, Dot>>(index, make_h0_plain_prefetcher, check) + vchordrq::maintain::<_, Op, Dot>>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ) } }; pgrx::info!( @@ -159,7 +183,7 @@ pub fn insert( vector: OwnedVector, skip_freespaces: bool, skip_search: bool, - chooser: &mut impl Chooser, + chooser: &mut impl InsertChooser, bump: &impl Bump, ) where R: RelationRead + RelationWrite, diff --git a/tests/vchordg/vacuum.slt b/tests/vchordg/vacuum.slt new file mode 100644 index 00000000..4b1a548c --- /dev/null +++ b/tests/vchordg/vacuum.slt @@ -0,0 +1,77 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.667; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DROP TABLE t; diff --git a/tests/vchordg/vacuum_parallel.slt b/tests/vchordg/vacuum_parallel.slt new file mode 100644 index 00000000..56d1d67e --- /dev/null +++ b/tests/vchordg/vacuum_parallel.slt @@ -0,0 +1,53 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SET min_parallel_index_scan_size = 1; + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.334; + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.334; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.667; + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/pushdown_plan.slt b/tests/vchordrq/pushdown_plan.slt index d732d3bc..b586edf2 100644 --- a/tests/vchordrq/pushdown_plan.slt +++ b/tests/vchordrq/pushdown_plan.slt @@ -29,7 +29,39 @@ SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) ORDER BY val0 # 1 vector key + 0 order_by key + original style -skipif pg18 +onlyif pg13 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +onlyif pg14 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +onlyif pg15 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +onlyif pg16 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +onlyif pg17 query I EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; @@ -147,7 +179,31 @@ ORDER BY val0 <-> '[0, 0, 0]'; # 0 vector key + 0 order_by key(variable) -skipif pg18 +onlyif pg13 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +onlyif pg14 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +onlyif pg15 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +onlyif pg16 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +onlyif pg17 query I EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; ---- diff --git a/tests/vchordrq/vacuum.slt b/tests/vchordrq/vacuum.slt new file mode 100644 index 00000000..4cbdc08a --- /dev/null +++ b/tests/vchordrq/vacuum.slt @@ -0,0 +1,80 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +# statement ok +# SET vchordrq.max_parallel_vacuum_workers = 0; + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.667; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/vacuum_parallel.slt b/tests/vchordrq/vacuum_parallel.slt new file mode 100644 index 00000000..427ea608 --- /dev/null +++ b/tests/vchordrq/vacuum_parallel.slt @@ -0,0 +1,65 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +# statement ok +# SET vchordrq.max_parallel_vacuum_workers = 32; + +statement ok +SET min_parallel_index_scan_size = 1; + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ build.internal.lists = [31] $$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ build.internal.lists = [33] $$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ build.internal.lists = [67] $$); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.334; + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.334; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.667; + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DROP TABLE t; From afb7f7909a56f37c0622853de43a7c4b2181c53f Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 15 Oct 2025 18:13:59 +0800 Subject: [PATCH 233/324] ci: archlinux, enable-debug, enable-cassert, valgrind, second latest clang, beta rust (#359) Adds a new job: * archlinux * PostgreSQL, debug build, valgrind enabled * second latest clang * beta rust It can help us assert assertions in PostgreSQL, detect memory unsafety, and anticipate changes introduced by future compiler versions and environment. job: -psql job: +check_arch Signed-off-by: usamoi --- .github/workflows/check.yml | 161 +++++++++++++++++++++++++++++- tests/vchordg/vacuum_parallel.slt | 12 +-- 2 files changed, 163 insertions(+), 10 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 11d515e0..00dc1474 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -185,10 +185,6 @@ jobs: - name: Set up Environment run: | rustup update - if [ "${{ matrix.version }}" = "13" ]; then - rustup default beta - rustup component add clippy rustfmt - fi sudo apt-get remove -y '^postgres.*' '^libpq.*' sudo apt-get purge -y '^postgres.*' '^libpq.*' @@ -812,3 +808,160 @@ jobs: if: always() run: | cat /sysroot/var/log/postgresql/postgresql-${{ matrix.version }}-main.log + + check_arch: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +check_arch')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +check_arch')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + version: ["17"] + + runs-on: "ubuntu-24.04" + + container: + image: archlinux/archlinux:base-devel + volumes: + - /etc/passwd:/etc/passwd:rw,rshared + - /etc/group:/etc/group:rw,rshared + - /etc/shadow:/etc/shadow:rw,rshared + - /etc/gshadow:/etc/gshadow:rw,rshared + - /etc/sudoers.d/runner:/etc/sudoers.d/runner:rw,rshared + options: --user 1001 --privileged + + env: + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + USER: "runner" + + steps: + - name: Set up Environment + run: | + sudo sed -i 's/^DownloadUser = alpm$/DownloadUser = runner/' /etc/pacman.conf + sudo pacman -Syu git github-cli nodejs --ignore systemd,systemd-libs,systemd-sysvcompat --noconfirm + + # fix home directory + sudo mkdir -p /var/lib/postgresql + sudo chmod 700 /var/lib/postgresql + sudo chown postgres:postgres /var/lib/postgresql + + sudo pacman -S python clang valgrind --noconfirm + curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain beta + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | sudo install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Install PostgreSQL + run: | + mkdir ~/postgresql-install + curl -fsSL https://ftp.postgresql.org/pub/source/v17.6/postgresql-17.6.tar.bz2 | tar -xj -C ~/postgresql-install + pushd ~/postgresql-install/postgresql-17.6 + ./configure \ + --prefix=/usr \ + --sysconfdir=/etc \ + --mandir=/usr/share/man \ + --datadir=/usr/share/postgresql \ + --disable-rpath \ + --enable-nls \ + --with-gssapi \ + --with-icu \ + --with-ldap \ + --with-lz4 \ + --with-openssl \ + --with-python \ + --with-readline \ + --with-system-tzdata=/usr/share/zoneinfo \ + --with-uuid=e2fs \ + --with-zstd \ + --enable-debug \ + --enable-cassert \ + CC='sccache clang' \ + CFLAGS='-O0 -g -fno-omit-frame-pointer -mno-omit-leaf-frame-pointer -fasynchronous-unwind-tables' \ + CPPFLAGS='-DUSE_ASSERT_CHECKING=1 -DRANDOMIZE_ALLOCATED_MEMORY=1 -DUSE_VALGRIND=1' + make all -j$(nproc) + sudo make install + popd + + sudo touch /usr/bin/postmaster + echo '#!/usr/bin/sh' | sudo tee -a /usr/bin/postmaster + echo 'export DEBUGINFOD_URLS=https://debuginfod.archlinux.org' | sudo tee -a /usr/bin/postmaster + echo 'exec \' | sudo tee -a /usr/bin/postmaster + echo ' valgrind \' | sudo tee -a /usr/bin/postmaster + echo ' --leak-check=no \' | sudo tee -a /usr/bin/postmaster + echo ' --gen-suppressions=all \' | sudo tee -a /usr/bin/postmaster + echo ' --time-stamp=yes \' | sudo tee -a /usr/bin/postmaster + echo ' --error-markers=VALGRINDERROR-BEGIN,VALGRINDERROR-END \' | sudo tee -a /usr/bin/postmaster + echo ' --trace-children=yes \' | sudo tee -a /usr/bin/postmaster + echo ' /usr/bin/postgres "$@"' | sudo tee -a /usr/bin/postmaster + sudo chmod 755 /usr/bin/postmaster + + sudo mkdir -p /var/lib/postgres + sudo chmod 700 /var/lib/postgres + sudo chown postgres:postgres /var/lib/postgres + sudo touch /var/log/postgresql.log + sudo chmod 700 /var/log/postgresql.log + sudo chown postgres:postgres /var/log/postgresql.log + + sudo -iu postgres initdb -D /var/lib/postgres/data + + sudo -iu postgres pg_ctl start -D /var/lib/postgres/data -l /var/log/postgresql.log -p /usr/bin/postmaster + sudo -iu postgres createuser -s -r $USER + sudo -iu postgres createdb -O $USER $USER + if [ "${{ matrix.version }}" -ge "18" ]; then + sudo -iu postgres psql -c 'ALTER SYSTEM SET io_method = io_uring' + fi + sudo -iu postgres psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' + sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + sudo -iu postgres pg_ctl stop -D /var/lib/postgres/data + + - name: Install pgvector + run: | + mkdir ~/pgvector-install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install + pushd ~/pgvector-install/pgvector-0.8.1 + make -j$(nproc) + sudo make install + popd + + - name: Checkout + uses: actions/checkout@v4 + + - name: Clippy + run: | + . "$HOME/.cargo/env" + PGRX_PG_CONFIG_PATH=pg_config \ + cargo clippy --locked --target x86_64-unknown-linux-gnu \ + --workspace --features pg${{ matrix.version }} + + - name: Install + run: | + . "$HOME/.cargo/env" + make PROFILE=dev build + sudo make install + + - name: Service + run: | + sudo -iu postgres pg_ctl start -D /var/lib/postgres/data -l /var/log/postgresql.log -p /usr/bin/postmaster + psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + fi + + - name: Logging + if: always() + run: | + sudo cat /var/log/postgresql.log diff --git a/tests/vchordg/vacuum_parallel.slt b/tests/vchordg/vacuum_parallel.slt index 56d1d67e..2bb62cad 100644 --- a/tests/vchordg/vacuum_parallel.slt +++ b/tests/vchordg/vacuum_parallel.slt @@ -20,31 +20,31 @@ statement ok VACUUM (PARALLEL 8) t; statement ok -INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); statement ok VACUUM (PARALLEL 8) t; statement ok -DELETE FROM t WHERE vector_norm(val) > 0.334; +DELETE FROM t WHERE vector_norm(val) < 0.333; statement ok VACUUM (PARALLEL 8) t; statement ok -DELETE FROM t WHERE vector_norm(val) > 0.334; +DELETE FROM t WHERE vector_norm(val) < 0.333; statement ok -INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); statement ok VACUUM (PARALLEL 8) t; statement ok -INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); statement ok -DELETE FROM t WHERE vector_norm(val) < 0.667; +DELETE FROM t WHERE vector_norm(val) > 0.667; statement ok VACUUM (PARALLEL 8) t; From ea9bbd11de1f71ac91efdbde21dfc077d249777d Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 20 Oct 2025 16:04:32 +0800 Subject: [PATCH 234/324] feat: fast sample (#360) Signed-off-by: usamoi --- .github/workflows/check.yml | 28 ++- Cargo.lock | 19 ++ Cargo.toml | 3 + crates/feistel/Cargo.toml | 14 ++ crates/feistel/src/lib.rs | 118 +++++++++++ src/index/fetcher.rs | 20 +- src/index/mod.rs | 2 + src/index/sample.rs | 334 ++++++++++++++++++++++++++++++ src/index/traverse.rs | 127 ++++++++++++ src/index/vchordg/am/am_build.rs | 107 +++------- src/index/vchordrq/am/am_build.rs | 193 +++++++---------- 11 files changed, 757 insertions(+), 208 deletions(-) create mode 100644 crates/feistel/Cargo.toml create mode 100644 crates/feistel/src/lib.rs create mode 100644 src/index/sample.rs create mode 100644 src/index/traverse.rs diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 00dc1474..3232c0b5 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -11,6 +11,11 @@ concurrency: jobs: style: + if: | + (github.event_name == 'push' && !contains(github.event.head_commit.message, 'job: -style')) || + (github.event_name == 'pull_request' && !contains(github.event.pull_request.body, 'job: -style')) || + github.event_name == 'workflow_dispatch' + runs-on: "ubuntu-24.04" steps: @@ -100,6 +105,11 @@ jobs: ! grep -P '\t' -r ./sql lint: + if: | + (github.event_name == 'push' && !contains(github.event.head_commit.message, 'job: -lint')) || + (github.event_name == 'pull_request' && !contains(github.event.pull_request.body, 'job: -lint')) || + github.event_name == 'workflow_dispatch' + strategy: matrix: arch: ["x86_64", "aarch64"] @@ -817,7 +827,7 @@ jobs: strategy: matrix: - version: ["17"] + version: ["17", "18"] runs-on: "ubuntu-24.04" @@ -851,7 +861,7 @@ jobs: sudo chmod 700 /var/lib/postgresql sudo chown postgres:postgres /var/lib/postgresql - sudo pacman -S python clang valgrind --noconfirm + sudo pacman -S python clang valgrind liburing numactl --noconfirm curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain beta curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | sudo install -m 755 /dev/stdin /usr/local/bin/sqllogictest @@ -861,9 +871,16 @@ jobs: - name: Install PostgreSQL run: | + if [ "${{ matrix.version }}" = "17" ]; then + VER=17.6 + fi + if [ "${{ matrix.version }}" = "18" ]; then + VER=18.0 + fi + mkdir ~/postgresql-install - curl -fsSL https://ftp.postgresql.org/pub/source/v17.6/postgresql-17.6.tar.bz2 | tar -xj -C ~/postgresql-install - pushd ~/postgresql-install/postgresql-17.6 + curl -fsSL https://ftp.postgresql.org/pub/source/v$VER/postgresql-$VER.tar.bz2 | tar -xj -C ~/postgresql-install + pushd ~/postgresql-install/postgresql-$VER ./configure \ --prefix=/usr \ --sysconfdir=/etc \ @@ -881,6 +898,9 @@ jobs: --with-system-tzdata=/usr/share/zoneinfo \ --with-uuid=e2fs \ --with-zstd \ + ${{ matrix.version >= 18 && '--with-libcurl' || '' }} \ + ${{ matrix.version >= 18 && '--with-libnuma' || '' }} \ + ${{ matrix.version >= 18 && '--with-liburing' || '' }} \ --enable-debug \ --enable-cassert \ CC='sccache clang' \ diff --git a/Cargo.lock b/Cargo.lock index 5bad0f56..5941d3bd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -413,6 +413,14 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "feistel" +version = "0.0.0" +dependencies = [ + "rand", + "wyhash", +] + [[package]] name = "find-msvc-tools" version = "0.1.3" @@ -1557,6 +1565,7 @@ dependencies = [ "bumpalo", "dary_heap", "distance", + "feistel", "index", "k_means", "mimalloc", @@ -1575,6 +1584,7 @@ dependencies = [ "vchordg", "vchordrq", "vector", + "wyhash", "zerocopy", ] @@ -1877,6 +1887,15 @@ version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" +[[package]] +name = "wyhash" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca4d373340c479fd1e779f7a763acee85da3e423b1a9a9acccf97babcc92edbb" +dependencies = [ + "rand_core", +] + [[package]] name = "wyz" version = "0.5.1" diff --git a/Cargo.toml b/Cargo.toml index 8fd543c4..4345a1a9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ pg18 = ["pgrx/pg18", "pgrx-catalog/pg18"] [dependencies] always_equal = { path = "./crates/always_equal" } distance = { path = "./crates/distance" } +feistel = { path = "./crates/feistel" } index = { path = "./crates/index" } k_means = { path = "./crates/k_means" } rabitq = { path = "./crates/rabitq" } @@ -44,6 +45,7 @@ seq-macro.workspace = true serde.workspace = true toml = "0.9.8" validator.workspace = true +wyhash.workspace = true zerocopy.workspace = true [target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos")))'.dependencies] @@ -69,6 +71,7 @@ rand_chacha = "0.9.0" seq-macro = "0.3.6" serde = { version = "1.0", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } +wyhash = "0.6.0" zerocopy = { version = "0.8.27", features = ["derive"] } [workspace.lints] diff --git a/crates/feistel/Cargo.toml b/crates/feistel/Cargo.toml new file mode 100644 index 00000000..826e327f --- /dev/null +++ b/crates/feistel/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "feistel" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] + +[dev-dependencies] +rand.workspace = true +wyhash.workspace = true + +[lints] +workspace = true diff --git a/crates/feistel/src/lib.rs b/crates/feistel/src/lib.rs new file mode 100644 index 00000000..afb184be --- /dev/null +++ b/crates/feistel/src/lib.rs @@ -0,0 +1,118 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub fn feistel(width: u32, x: I, round: u32, secret: impl Fn(u32, I) -> I) -> I +where + I: Copy, + I: Ord, + I: std::ops::Add, + I: std::ops::Sub, + I: std::ops::BitAnd, + I: std::ops::BitOr, + I: std::ops::BitXor, + I: std::ops::Shl, + I: std::ops::Shr, + I: Zero, + I: One, + I: std::fmt::Debug, +{ + assert_eq!(width % 2, 0); + assert_eq!(x >> width, I::zero()); + + let half_width = width >> 1; + let half_mask = (I::one() << half_width) - I::one(); + + let mut left = (x >> half_width) & half_mask; + let mut right = x & half_mask; + + for i in 0..round { + (left, right) = (right, left ^ (secret(i, right) & half_mask)); + } + + (left << half_width) | right +} + +pub trait Zero { + fn zero() -> Self; +} + +pub trait One { + fn one() -> Self; +} + +macro_rules! impl_traits { + ($t:ty) => { + impl Zero for $t { + fn zero() -> Self { + 0 + } + } + + impl One for $t { + fn one() -> Self { + 1 + } + } + }; +} + +impl_traits!(u8); +impl_traits!(u16); +impl_traits!(u32); +impl_traits!(u64); +impl_traits!(u128); + +// This is a standalone crate, simply because we want to run tests for it. + +#[test] +fn is_a_permutation() { + let key_0 = [7u8; _]; + let key_1 = [8u8; _]; + let secret = move |round: u32, x: u32| { + let buffer = [round.to_le_bytes(), x.to_le_bytes(), key_0, key_1]; + wyhash::wyhash(buffer.as_flattened(), 0) as u32 + }; + for width in (0..20).step_by(2) { + let mut y = Vec::new(); + for x in 0..1 << width { + y.push(feistel::(width, x, 8, secret)); + } + if width <= 8 { + eprintln!("feistel({width}, _, 8, *) = {y:?}"); + } + y.sort_unstable(); + for x in 0..1 << width { + assert_eq!(y[x as usize], x); + } + } +} + +#[test] +fn sample() { + let n = 6370_u32; + let width = (n.ilog2() + 1).next_multiple_of(2); + let key_0 = rand::Rng::random(&mut rand::rng()); + let key_1 = rand::Rng::random(&mut rand::rng()); + let secret = move |round: u32, x: u32| { + let buffer = [round.to_le_bytes(), x.to_le_bytes(), key_0, key_1]; + wyhash::wyhash(buffer.as_flattened(), 0) as u32 + }; + let mut permutation = (0..1 << width) + .map(move |i| feistel(width, i, 8, secret)) + .filter(move |&x| x < n); + for _ in 0..n { + assert!(permutation.next().is_some()); + } + assert_eq!(permutation.next(), None); +} diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs index f13eaac9..e9c779bc 100644 --- a/src/index/fetcher.rs +++ b/src/index/fetcher.rs @@ -18,19 +18,22 @@ use std::num::NonZero; use std::ops::DerefMut; use std::ptr::NonNull; +pub trait FilterableTuple: Tuple { + fn filter(&mut self) -> bool; +} + +pub trait Tuple { + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); +} + pub trait Fetcher { - type Tuple<'a>: Tuple + type Tuple<'a>: FilterableTuple where Self: 'a; fn fetch(&mut self, key: [u16; 3]) -> Option>; } -pub trait Tuple { - fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); - fn filter(&mut self) -> bool; -} - impl T> Fetcher for LazyCell { type Tuple<'a> = T::Tuple<'a> @@ -102,6 +105,9 @@ impl Fetcher for HeapFetcher { use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; let mut ctid = key_to_ctid(key); let table_am = (*self.heap_relation).rd_tableam; + if table_am.is_null() { + panic!("unknown heap access method"); + } let index_fetch_tuple = (*table_am) .index_fetch_tuple .expect("unsupported heap access method"); @@ -169,7 +175,9 @@ impl Tuple for HeapTuple<'_> { (&this.values, &this.is_nulls) } } +} +impl FilterableTuple for HeapTuple<'_> { #[allow(clippy::collapsible_if)] fn filter(&mut self) -> bool { unsafe { diff --git a/src/index/mod.rs b/src/index/mod.rs index b68eb4e8..f4cd19d4 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -17,8 +17,10 @@ mod functions; mod gucs; mod hook; mod opclass; +mod sample; mod scanners; mod storage; +mod traverse; mod vchordg; mod vchordrq; diff --git a/src/index/sample.rs b/src/index/sample.rs new file mode 100644 index 00000000..bbdbc3f5 --- /dev/null +++ b/src/index/sample.rs @@ -0,0 +1,334 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use pgrx::pg_sys::{Datum, ItemPointerData}; +use std::ptr::NonNull; + +pub trait Tuple { + #[expect(dead_code)] + fn id(&mut self) -> ItemPointerData; + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); +} + +pub trait Sample { + type Tuple<'a>: Tuple + where + Self: 'a; + + fn next(&mut self) -> Option>; +} + +pub trait Sampler { + type Sample: Sample; + + fn sample(&self) -> Self::Sample; +} + +pub struct HeapSampler { + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, +} + +impl HeapSampler { + pub unsafe fn new( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + ) -> Self { + Self { + heap_relation, + index_relation, + snapshot, + } + } +} + +impl Drop for HeapSampler { + fn drop(&mut self) {} +} + +impl Sampler for HeapSampler { + type Sample = HeapSample; + + fn sample(&self) -> Self::Sample { + unsafe { + let state = NonNull::new_unchecked(Box::into_raw(Box::new(State { + blocks: None, + tuples: None, + }))); + let sample_scan_state = + NonNull::new_unchecked(Box::into_raw(Box::new(pgrx::pg_sys::SampleScanState { + tsmroutine: (&raw const TSM.0).cast_mut().cast(), + tsm_state: state.as_ptr().cast(), + ..core::mem::zeroed() + }))); + let table_scan_desc = pgrx::pg_sys::table_beginscan_sampling( + self.heap_relation, + self.snapshot, + 0, + std::ptr::null_mut(), + true, + false, + true, + ); + let index_info = pgrx::pg_sys::BuildIndexInfo(self.index_relation); + let estate = pgrx::pg_sys::CreateExecutorState(); + let econtext = pgrx::pg_sys::MakePerTupleExprContext(estate); + HeapSample { + index_info, + estate, + econtext, + slot: pgrx::pg_sys::table_slot_create(self.heap_relation, std::ptr::null_mut()), + values: [Datum::null(); 32], + is_nulls: [true; 32], + state, + sample_scan_state, + table_scan_desc, + done: false, + have_block: false, + } + } + } +} + +pub struct HeapSample { + index_info: *mut pgrx::pg_sys::IndexInfo, + estate: *mut pgrx::pg_sys::EState, + econtext: *mut pgrx::pg_sys::ExprContext, + slot: *mut pgrx::pg_sys::TupleTableSlot, + values: [Datum; 32], + is_nulls: [bool; 32], + state: NonNull, + sample_scan_state: NonNull, + table_scan_desc: pgrx::pg_sys::TableScanDesc, + done: bool, + have_block: bool, +} + +impl Drop for HeapSample { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); + // free common resources + pgrx::pg_sys::table_endscan(self.table_scan_desc); + let _ = Box::from_raw(self.sample_scan_state.as_ptr()); + let _ = Box::from_raw(self.state.as_ptr()); + pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); + pgrx::pg_sys::FreeExecutorState(self.estate); + } + } +} + +impl Sample for HeapSample { + type Tuple<'a> = HeapTuple<'a>; + + fn next(&mut self) -> Option> { + unsafe { + use pgrx::pg_sys::{table_scan_sample_next_block, table_scan_sample_next_tuple}; + + if self.done { + return None; + } + + loop { + if !self.have_block { + if !table_scan_sample_next_block( + self.table_scan_desc, + self.sample_scan_state.as_ptr(), + ) { + self.have_block = false; + self.done = true; + return None; + } + + self.have_block = true; + } + + if !table_scan_sample_next_tuple( + self.table_scan_desc, + self.sample_scan_state.as_ptr(), + self.slot, + ) { + self.have_block = false; + continue; + } + + break; + } + + Some(HeapTuple { this: self }) + } + } +} + +pub struct HeapTuple<'a> { + this: &'a mut HeapSample, +} + +impl Tuple for HeapTuple<'_> { + fn id(&mut self) -> ItemPointerData { + unsafe { + let this = &mut self.this; + (*this.slot).tts_tid + } + } + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]) { + unsafe { + let this = &mut self.this; + (*this.econtext).ecxt_scantuple = this.slot; + pgrx::pg_sys::MemoryContextReset((*this.econtext).ecxt_per_tuple_memory); + pgrx::pg_sys::FormIndexDatum( + this.index_info, + this.slot, + this.estate, + this.values.as_mut_ptr(), + this.is_nulls.as_mut_ptr(), + ); + (&this.values, &this.is_nulls) + } + } +} + +fn sample(n: u32) -> Box> { + let width = (n.ilog2() + 1).next_multiple_of(2); + let key_0 = rand::Rng::random(&mut rand::rng()); + let key_1 = rand::Rng::random(&mut rand::rng()); + let secret = move |round: u32, x: u32| { + let buffer = [round.to_le_bytes(), x.to_le_bytes(), key_0, key_1]; + wyhash::wyhash(buffer.as_flattened(), 0) as u32 + }; + let permutation = (0..1 << width) + .map(move |i| feistel::feistel(width, i, 8, secret)) + .filter(move |&x| x < n); + Box::new(permutation) +} + +pub struct State { + blocks: Option>>, + tuples: Option>, +} + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn feistel_rows_nextsampleblock( + node: *mut pgrx::pg_sys::SampleScanState, + nblocks: pgrx::pg_sys::BlockNumber, +) -> pgrx::pg_sys::BlockNumber { + let state: &mut State = unsafe { &mut *(*node).tsm_state.cast() }; + let iter = state.blocks.get_or_insert_with(|| sample(nblocks)); + if let Some(number) = iter.next() { + number + } else { + pgrx::pg_sys::InvalidBlockNumber + } +} + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn feistel_rows_nextsampletuple( + node: *mut pgrx::pg_sys::SampleScanState, + _blockno: pgrx::pg_sys::BlockNumber, + maxoffset: pgrx::pg_sys::OffsetNumber, +) -> pgrx::pg_sys::OffsetNumber { + let state: &mut State = unsafe { &mut *(*node).tsm_state.cast() }; + let iter = state.tuples.get_or_insert(1..=maxoffset); + if let Some(number) = iter.next() { + number + } else { + state.tuples = None; + pgrx::pg_sys::InvalidOffsetNumber + } +} + +struct AssertSync(T); + +unsafe impl Sync for AssertSync {} + +static TSM: AssertSync = AssertSync(sys::TsmRoutine { + type_: pgrx::pg_sys::NodeTag::T_TsmRoutine, + NextSampleBlock: Some(feistel_rows_nextsampleblock), + NextSampleTuple: Some(feistel_rows_nextsampletuple), + ..unsafe { core::mem::zeroed() } +}); + +#[allow(non_camel_case_types)] +#[allow(non_snake_case)] +mod sys { + #[cfg(not(feature = "pg13"))] + #[cfg(not(feature = "pg14"))] + #[cfg(not(feature = "pg15"))] + #[cfg(not(feature = "pg16"))] + #[cfg(not(feature = "pg17"))] + #[cfg(not(feature = "pg18"))] + compile_error!("bindings are not checked"); + + use core::ffi::c_int; + use pgrx::pg_sys::{ + BlockNumber, Datum, List, NodeTag, OffsetNumber, PlannerInfo, RelOptInfo, SampleScanState, + }; + + pub type SampleScanGetSampleSize_function = Option< + unsafe extern "C-unwind" fn( + root: *mut PlannerInfo, + baserel: *mut RelOptInfo, + paramexprs: *mut List, + pages: *mut BlockNumber, + tuples: *mut f64, + ), + >; + + pub type InitSampleScan_function = + Option; + + pub type BeginSampleScan_function = Option< + unsafe extern "C-unwind" fn( + node: *mut SampleScanState, + params: *mut Datum, + nparams: c_int, + seed: u32, + ), + >; + + pub type NextSampleBlock_function = Option< + unsafe extern "C-unwind" fn( + node: *mut SampleScanState, + nblocks: BlockNumber, + ) -> BlockNumber, + >; + + pub type NextSampleTuple_function = Option< + unsafe extern "C-unwind" fn( + node: *mut SampleScanState, + blockno: BlockNumber, + maxoffset: OffsetNumber, + ) -> OffsetNumber, + >; + + pub type EndSampleScan_function = + Option; + + #[repr(C)] + #[derive(Debug, Copy, Clone)] + pub struct TsmRoutine { + pub type_: NodeTag, + pub parameterTypes: *mut List, + pub repeatable_across_queries: bool, + pub repeatable_across_scans: bool, + pub SampleScanGetSampleSize: SampleScanGetSampleSize_function, + pub InitSampleScan: InitSampleScan_function, + pub BeginSampleScan: BeginSampleScan_function, + pub NextSampleBlock: NextSampleBlock_function, + pub NextSampleTuple: NextSampleTuple_function, + pub EndSampleScan: EndSampleScan_function, + } +} diff --git a/src/index/traverse.rs b/src/index/traverse.rs new file mode 100644 index 00000000..b30ff39b --- /dev/null +++ b/src/index/traverse.rs @@ -0,0 +1,127 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use pgrx::pg_sys::{Datum, ItemPointerData}; + +pub trait Tuple { + fn id(&mut self) -> ItemPointerData; + fn build(&mut self) -> (*const Datum, *const bool); +} + +pub trait Traverse { + fn next(&mut self, tuple: T); +} + +impl Traverse for F { + fn next(&mut self, mut tuple: T) { + self(&mut tuple); + } +} + +pub trait Traverser { + fn traverse(&self, progress: bool, traverse: T); +} + +#[derive(Debug, Clone)] +pub struct HeapTraverser { + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + scan: *mut pgrx::pg_sys::TableScanDescData, +} + +impl HeapTraverser { + pub unsafe fn new( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + scan: *mut pgrx::pg_sys::TableScanDescData, + ) -> Self { + Self { + heap_relation, + index_relation, + index_info, + scan, + } + } +} + +impl Drop for HeapTraverser { + fn drop(&mut self) {} +} + +impl Traverser for HeapTraverser { + fn traverse(&self, progress: bool, mut traverse: T) { + unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + let table_am = (*self.heap_relation).rd_tableam; + if table_am.is_null() { + panic!("unknown heap access method"); + } + let index_build_range_scan = (*table_am) + .index_build_range_scan + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + index_build_range_scan( + self.heap_relation, + self.index_relation, + self.index_info, + true, + false, + progress, + 0, + pgrx::pg_sys::InvalidBlockNumber, + Some(callback::), + (&raw mut traverse).cast(), + self.scan, + ) + }); + } + } +} + +struct HeapTuple { + id: ItemPointerData, + values: *const Datum, + is_nulls: *const bool, +} + +impl Tuple for HeapTuple { + fn id(&mut self) -> ItemPointerData { + self.id + } + + fn build(&mut self) -> (*const Datum, *const bool) { + (self.values, self.is_nulls) + } +} + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn callback( + _index_relation: pgrx::pg_sys::Relation, + ctid: pgrx::pg_sys::ItemPointer, + values: *mut Datum, + is_null: *mut bool, + _tuple_is_alive: bool, + state: *mut core::ffi::c_void, +) { + let state = unsafe { &mut *state.cast::() }; + + state.next(HeapTuple { + id: unsafe { *ctid }, + values: values.cast_const(), + is_nulls: is_null.cast_const(), + }); +} diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index 84498277..602a09a2 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -14,10 +14,10 @@ use crate::datatype::typmod::Typmod; use crate::index::storage::PostgresRelation; +use crate::index::traverse::{HeapTraverser, Traverser}; use crate::index::vchordg::am::{Reloption, ctid_to_key, kv_to_pointer}; -use crate::index::vchordg::opclass::{Opfamily, opfamily}; +use crate::index::vchordg::opclass::opfamily; use crate::index::vchordg::types::VchordgIndexingOptions; -use pgrx::pg_sys::{Datum, ItemPointerData}; use std::ffi::CStr; use std::marker::PhantomData; use vchordg::types::*; @@ -82,72 +82,6 @@ pub extern "C-unwind" fn ambuildphasename(x: i64) -> *mut core::ffi::c_char { } } -#[derive(Debug, Clone)] -struct Heap { - heap_relation: pgrx::pg_sys::Relation, - index_relation: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, - opfamily: Opfamily, - scan: *mut pgrx::pg_sys::TableScanDescData, -} - -impl Heap { - fn traverse))>( - &self, - progress: bool, - callback: F, - ) { - use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; - pub struct State<'a, F> { - pub this: &'a Heap, - pub callback: F, - } - #[pgrx::pg_guard] - unsafe extern "C-unwind" fn call( - _index_relation: pgrx::pg_sys::Relation, - ctid: pgrx::pg_sys::ItemPointer, - values: *mut Datum, - is_null: *mut bool, - _tuple_is_alive: bool, - state: *mut core::ffi::c_void, - ) where - F: FnMut((ItemPointerData, Vec<(OwnedVector, u16)>)), - { - let state = unsafe { &mut *state.cast::>() }; - let opfamily = state.this.opfamily; - let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; - let ctid = unsafe { ctid.read() }; - if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { - (state.callback)((ctid, store)); - } - } - let table_am = unsafe { &*(*self.heap_relation).rd_tableam }; - let mut state = State { - this: self, - callback, - }; - let index_build_range_scan = table_am.index_build_range_scan.expect("bad table"); - unsafe { - #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] - pg_guard_ffi_boundary(|| { - index_build_range_scan( - self.heap_relation, - self.index_relation, - self.index_info, - true, - false, - progress, - 0, - pgrx::pg_sys::InvalidBlockNumber, - Some(call::), - (&mut state) as *mut State as *mut _, - self.scan, - ) - }); - } - } -} - #[derive(Debug, Clone)] struct PostgresReporter { _phantom: PhantomData<*mut ()>, @@ -537,16 +471,16 @@ unsafe fn parallel_build( let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap_relation, tablescandesc) }; let opfamily = unsafe { opfamily(index_relation) }; - let heap = Heap { - heap_relation, - index_relation, - index_info, - opfamily, - scan, - }; + let traverser = unsafe { HeapTraverser::new(heap_relation, index_relation, index_info, scan) }; match cached { VchordgCachedReader::_0(_) => { - heap.traverse(true, move |(ctid, store)| { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -585,18 +519,25 @@ unsafe fn sequential_build( let index = unsafe { PostgresRelation::new(index_relation) }; let opfamily = unsafe { opfamily(index_relation) }; - let heap = Heap { - heap_relation, - index_relation, - index_info, - opfamily, - scan: std::ptr::null_mut(), + let traverser = unsafe { + HeapTraverser::new( + heap_relation, + index_relation, + index_info, + std::ptr::null_mut(), + ) }; let mut indtuples = 0; match cached { VchordgCachedReader::_0(_) => { - heap.traverse(true, |(ctid, store)| { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 8d11b38b..b1c687da 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -14,7 +14,9 @@ use crate::datatype::typmod::Typmod; use crate::index::fetcher::*; +use crate::index::sample::{HeapSampler, Sample, Sampler, Tuple}; use crate::index::storage::{PostgresPage, PostgresRelation}; +use crate::index::traverse::{HeapTraverser, Traverser}; use crate::index::vchordrq::am::Reloption; use crate::index::vchordrq::build::{Normalize, Normalized}; use crate::index::vchordrq::opclass::{Opfamily, opfamily}; @@ -22,8 +24,6 @@ use crate::index::vchordrq::types::*; use index::relation::{ Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, }; -use pgrx::pg_sys::{Datum, ItemPointerData}; -use rand::Rng; use simd::Floating; use std::ffi::CStr; use std::marker::PhantomData; @@ -167,72 +167,6 @@ pub extern "C-unwind" fn ambuildphasename(x: i64) -> *mut core::ffi::c_char { } } -#[derive(Debug, Clone)] -struct Heap { - heap_relation: pgrx::pg_sys::Relation, - index_relation: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, - opfamily: Opfamily, - scan: *mut pgrx::pg_sys::TableScanDescData, -} - -impl Heap { - fn traverse))>( - &self, - progress: bool, - callback: F, - ) { - use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; - pub struct State<'a, F> { - pub this: &'a Heap, - pub callback: F, - } - #[pgrx::pg_guard] - unsafe extern "C-unwind" fn call( - _index_relation: pgrx::pg_sys::Relation, - ctid: pgrx::pg_sys::ItemPointer, - values: *mut Datum, - is_null: *mut bool, - _tuple_is_alive: bool, - state: *mut core::ffi::c_void, - ) where - F: FnMut((ItemPointerData, Vec<(OwnedVector, u16)>)), - { - let state = unsafe { &mut *state.cast::>() }; - let opfamily = state.this.opfamily; - let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; - let ctid = unsafe { ctid.read() }; - if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { - (state.callback)((ctid, store)); - } - } - let table_am = unsafe { &*(*self.heap_relation).rd_tableam }; - let mut state = State { - this: self, - callback, - }; - let index_build_range_scan = table_am.index_build_range_scan.expect("bad table"); - unsafe { - #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] - pg_guard_ffi_boundary(|| { - index_build_range_scan( - self.heap_relation, - self.index_relation, - self.index_info, - true, - false, - progress, - 0, - pgrx::pg_sys::InvalidBlockNumber, - Some(call::), - (&mut state) as *mut State as *mut _, - self.scan, - ) - }); - } - } -} - #[derive(Debug, Clone)] struct PostgresReporter { _phantom: PhantomData<*mut ()>, @@ -280,14 +214,6 @@ pub unsafe extern "C-unwind" fn ambuild( pgrx::error!("error while validating options: {}", errors); } let opfamily = unsafe { opfamily(index_relation) }; - let index = unsafe { PostgresRelation::new(index_relation) }; - let heap = Heap { - heap_relation, - index_relation, - index_info, - opfamily, - scan: std::ptr::null_mut(), - }; let reporter = PostgresReporter { _phantom: PhantomData, }; @@ -299,9 +225,14 @@ pub unsafe extern "C-unwind" fn ambuild( } VchordrqBuildSourceOptions::Internal(internal_build) => { reporter.phase(BuildPhase::from_code(BuildPhaseCode::InternalBuild)); - let mut tuples_total = 0_u64; + let snapshot = if unsafe { (*index_info).ii_Concurrent } { + unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } + } else { + &raw mut pgrx::pg_sys::SnapshotAnyData + }; + let sampler = unsafe { HeapSampler::new(index_relation, heap_relation, snapshot) }; + let mut sample = sampler.sample(); let samples = 'a: { - let mut rand = rand::rng(); let Some(max_number_of_samples) = internal_build .lists .last() @@ -310,31 +241,43 @@ pub unsafe extern "C-unwind" fn ambuild( break 'a Vec::new(); }; let mut samples = Vec::new(); - let mut number_of_samples = 0_u32; - heap.traverse(false, |(_, store)| { - for (vector, _) in store { - let x = match vector { - OwnedVector::Vecf32(x) => VectOwned::normalize(x), - OwnedVector::Vecf16(x) => VectOwned::normalize(x), - }; - assert_eq!( - vector_options.dims, - x.len() as u32, - "invalid vector dimensions" - ); - if number_of_samples < max_number_of_samples { - samples.push(x); - number_of_samples += 1; + while samples.len() < max_number_of_samples as usize { + if let Some(mut tuple) = sample.next() { + let (values, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(values[0]); + if let Some(datum) = datum { + let vectors = unsafe { opfamily.store(datum) }; + if let Some(vectors) = vectors { + for (vector, _) in vectors { + let x = match vector { + OwnedVector::Vecf32(x) => VectOwned::normalize(x), + OwnedVector::Vecf16(x) => VectOwned::normalize(x), + }; + assert_eq!( + vector_options.dims, + x.len() as u32, + "invalid vector dimensions" + ); + samples.push(x); + } + } else { + continue; + } } else { - let index = rand.random_range(0..max_number_of_samples) as usize; - samples[index] = x; + continue; } + } else { + break; } - tuples_total += 1; - }); + } + samples.truncate(max_number_of_samples as usize); samples }; - reporter.tuples_total(tuples_total); + if is_mvcc_snapshot(snapshot) { + unsafe { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + } make_internal_build(vector_options, internal_build, samples, &reporter) } VchordrqBuildSourceOptions::External(external_build) => { @@ -348,6 +291,7 @@ pub unsafe extern "C-unwind" fn ambuild( } } reporter.phase(BuildPhase::from_code(BuildPhaseCode::Build)); + let index = unsafe { PostgresRelation::new(index_relation) }; crate::index::vchordrq::dispatch::build( vector_options, vchordrq_options.index, @@ -1037,13 +981,7 @@ unsafe fn parallel_build( let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap_relation, tablescandesc) }; let opfamily = unsafe { opfamily(index_relation) }; - let heap = Heap { - heap_relation, - index_relation, - index_info, - opfamily, - scan, - }; + let traverser = unsafe { HeapTraverser::new(heap_relation, index_relation, index_info, scan) }; struct IdChooser(u32); impl InsertChooser for IdChooser { @@ -1070,7 +1008,13 @@ unsafe fn parallel_build( match cached { VchordrqCachedReader::_0(_) => { - heap.traverse(true, |(ctid, store)| { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1104,7 +1048,13 @@ unsafe fn parallel_build( cache: cached, relation: index.clone(), }; - heap.traverse(true, |(ctid, store)| { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1163,12 +1113,13 @@ unsafe fn sequential_build( let index = unsafe { PostgresRelation::new(index_relation) }; let opfamily = unsafe { opfamily(index_relation) }; - let heap = Heap { - heap_relation, - index_relation, - index_info, - opfamily, - scan: std::ptr::null_mut(), + let traverser = unsafe { + HeapTraverser::new( + heap_relation, + index_relation, + index_info, + std::ptr::null_mut(), + ) }; struct ChooseZero; @@ -1194,7 +1145,13 @@ unsafe fn sequential_build( let mut indtuples = 0; match cached { VchordrqCachedReader::_0(_) => { - heap.traverse(true, |(ctid, store)| { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1220,7 +1177,13 @@ unsafe fn sequential_build( cache: cached, relation: index.clone(), }; - heap.traverse(true, |(ctid, store)| { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); From 5574da571f3cca1ae558edb3be445acb04948b51 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 22 Oct 2025 11:05:53 +0800 Subject: [PATCH 235/324] ci: add miri (#362) job: +miri Signed-off-by: usamoi --- .github/workflows/check.yml | 103 ++++++++++++++++++++++--------- Cargo.lock | 85 ++++++++++++------------- crates/feistel/src/lib.rs | 4 +- crates/make/Cargo.toml | 2 +- crates/simd/Cargo.toml | 4 +- crates/simd/src/fht.rs | 2 +- crates/simd/src/u8.rs | 2 +- crates/vchordrq/src/fast_heap.rs | 6 +- 8 files changed, 120 insertions(+), 88 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 3232c0b5..d0b14177 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -40,7 +40,7 @@ jobs: uses: astral-sh/ruff-action@v1 - name: Rustfmt - run: cargo fmt --check -- --config-path /dev/null --config imports_granularity=Module + run: cargo fmt --check - name: Deny run: cargo deny check @@ -104,6 +104,35 @@ jobs: run: | ! grep -P '\t' -r ./sql + miri: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +miri')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +miri')) || + github.event_name == 'workflow_dispatch' + + runs-on: "ubuntu-24.04" + + env: + RUSTUP_AUTO_INSTALL: "0" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + rustup default nightly + rustup component add miri + + - name: Checkout + uses: actions/checkout@v4 + + - name: Cargo Test (Miri) + run: | + RUSTFLAGS="-Ctarget-cpu=sapphirerapids" MIRIFLAGS="-Zmiri-strict-provenance" \ + cargo miri test --locked --target x86_64-unknown-linux-gnu \ + --workspace --exclude vchord --no-fail-fast \ + -- --no-capture + lint: if: | (github.event_name == 'push' && !contains(github.event.head_commit.message, 'job: -lint')) || @@ -146,13 +175,13 @@ jobs: - name: Checkout uses: actions/checkout@v4 - - name: Clippy - run: cargo clippy --locked --workspace --exclude vchord - - name: Cargo Test - run: cargo test --locked --workspace --exclude vchord --exclude simd --no-fail-fast + run: | + cargo test --locked \ + --workspace --exclude vchord --no-fail-fast \ + -- --no-capture - - name: Cargo Test (simd) + - name: Cargo Test (QEMU) run: | if [ "$(uname -m)" == "x86_64" ]; then cargo \ @@ -171,6 +200,9 @@ jobs: test --locked -p simd -- --no-capture fi + - name: Clippy + run: cargo clippy --locked --workspace --exclude vchord + psql: if: | (github.event_name == 'push' && !contains(github.event.head_commit.message, 'job: -psql')) || @@ -502,7 +534,6 @@ jobs: # RUSTFLAGS: "-Dwarnings" CARGO_TERM_COLOR: "always" RUST_BACKTRACE: "1" - USER: "root" steps: - name: Set up Environment @@ -511,10 +542,13 @@ jobs: apk add nodejs --update-cache mkdir /opt/bin ln -s /usr/bin/node /opt/bin/node + export USER=root + echo USER=$USER >> $GITHUB_ENV apk add --no-cache sudo curl make zip clang18-dev postgresql${{ matrix.version }} postgresql${{ matrix.version }}-dev - curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain 1.89-beta + curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain beta + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV echo PGRX_PG_CONFIG_PATH=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV echo PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV @@ -552,11 +586,12 @@ jobs: - name: Patch run: | mkdir ./.cargo && touch ./.cargo/config.toml - echo "unstable.host-config = true" >> ./.cargo/config.toml - echo "unstable.target-applies-to-host = true" >> ./.cargo/config.toml - echo "host.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml - echo "build.rustflags = [\"-Dwarnings\", \"-Ctarget-feature=-crt-static\"]" >> ./.cargo/config.toml - echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + cat << EOF > ./.cargo/config.toml + unstable.host-config = true + unstable.target-applies-to-host = true + host.rustflags = ["-Dwarnings", "-Ctarget-feature=-crt-static"] + build.rustflags = ["-Dwarnings", "-Ctarget-feature=-crt-static"] + EOF - name: Clippy run: | @@ -672,6 +707,10 @@ jobs: steps: - name: Set up Environment run: | + rustup default beta + rustup component add rustfmt clippy + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + sudo apt-get update sudo apt-get install -y qemu-user-static clang lld sudo systemctl enable --now systemd-binfmt @@ -748,15 +787,8 @@ jobs: - name: Patch run: | - cat << EOF > rust-toolchain.toml - [toolchain] - channel = "nightly-2025-08-29" - components = ["rustfmt", "clippy"] - targets = ["${{ matrix.rust_triple }}"] - profile = "minimal" - EOF - mkdir -p .cargo - cat << EOF > .cargo/config.toml + mkdir ./.cargo && touch ./.cargo/config.toml + cat << EOF > ./.cargo/config.toml [target.${{ matrix.rust_triple }}] runner = ["qemu-${{ matrix.platform }}-static"] linker = "clang" @@ -771,17 +803,18 @@ jobs: CFLAGS_${{ matrix.rust_triple }} = "--target=${{ matrix.clang_triple }} --sysroot=/sysroot" BINDGEN_EXTRA_CLANG_ARGS_${{ matrix.rust_triple }} = "--sysroot=/sysroot" EOF - rustup toolchain install - - name: Clippy & Test + - name: Cargo Test + run: | + cargo test --locked --target ${{ matrix.rust_triple }} \ + --workspace --exclude vchord --no-fail-fast \ + -- --no-capture + + - name: Clippy run: | PGRX_PG_CONFIG_PATH=pg_config \ - cargo clippy --target ${{ matrix.rust_triple }} \ + cargo clippy --locked --target ${{ matrix.rust_triple }} \ --workspace --features pg${{ matrix.version }} - PGRX_PG_CONFIG_PATH=pg_config \ - cargo test --target ${{ matrix.rust_triple }} \ - --workspace --exclude vchord --no-fail-fast \ - -- --no-capture - name: Install run: | @@ -848,13 +881,14 @@ jobs: RUSTFLAGS: "-Dwarnings" CARGO_TERM_COLOR: "always" RUST_BACKTRACE: "1" - USER: "runner" steps: - name: Set up Environment run: | sudo sed -i 's/^DownloadUser = alpm$/DownloadUser = runner/' /etc/pacman.conf sudo pacman -Syu git github-cli nodejs --ignore systemd,systemd-libs,systemd-sysvcompat --noconfirm + export USER=runner + echo USER=$USER >> $GITHUB_ENV # fix home directory sudo mkdir -p /var/lib/postgresql @@ -862,7 +896,9 @@ jobs: sudo chown postgres:postgres /var/lib/postgresql sudo pacman -S python clang valgrind liburing numactl --noconfirm + curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain beta + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | sudo install -m 755 /dev/stdin /usr/local/bin/sqllogictest @@ -954,6 +990,13 @@ jobs: - name: Checkout uses: actions/checkout@v4 + - name: Cargo Test + run: | + . "$HOME/.cargo/env" + cargo test --locked --target x86_64-unknown-linux-gnu \ + --workspace --exclude vchord --no-fail-fast \ + -- --no-capture + - name: Clippy run: | . "$HOME/.cargo/env" diff --git a/Cargo.lock b/Cargo.lock index 5941d3bd..4f095bca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -102,9 +102,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.9.4" +version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2261d10cca569e4643e526d8dc2e62e433cc8aba21ab764233731f8d369bf394" +checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" [[package]] name = "bitvec" @@ -136,9 +136,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.40" +version = "1.2.41" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1d05d92f4b1fd76aad469d46cdd858ca761576082cd37df81416691e50199fb" +checksum = "ac9fe6cdbb24b6ade63616c0a0688e45bb56732262c158df3c0c4bea4ca47cb7" dependencies = [ "find-msvc-tools", "shlex", @@ -165,9 +165,9 @@ dependencies = [ [[package]] name = "cfg-if" -version = "1.0.3" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2fd1289c04a9ea8cb22300a459a72a385d7c73d3259e2ed7dcb2af674838cfa9" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "clang-sys" @@ -182,9 +182,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.48" +version = "4.5.50" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2134bb3ea021b78629caa971416385309e0131b351b25e01dc16fb54e1b5fae" +checksum = "0c2cfd7bf8a6017ddaa4e32ffe7403d547790db06bd171c1c53926faab501623" dependencies = [ "clap_builder", "clap_derive", @@ -192,9 +192,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.48" +version = "4.5.50" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2ba64afa3c0a6df7fa517765e31314e983f51dda798ffba27b988194fb65dc9" +checksum = "0a4c05b9e80c5ccd3a7ef080ad7b6ba7d6fc00a985b8b157197075677c82c7a0" dependencies = [ "anstream", "anstyle", @@ -204,9 +204,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.5.47" +version = "4.5.49" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbfd7eae0b0f1a6e63d4b13c9c478de77c2eb546fba158ad50b4203dc24b9f9c" +checksum = "2a0b5487afeab2deb2ff4e03a807ad1a03ac532ff5a2cee5d86884440c7f7671" dependencies = [ "heck", "proc-macro2", @@ -216,9 +216,9 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.7.5" +version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" +checksum = "a1d728cc89cf3aee9ff92b05e62b19ee65a02b5702cff7d5a377e32c6ae29d8d" [[package]] name = "codepage" @@ -423,9 +423,9 @@ dependencies = [ [[package]] name = "find-msvc-tools" -version = "0.1.3" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0399f9d26e5191ce32c498bebd31e7a3ceabc2745f0ac54af3f335126c3f24b3" +checksum = "52051878f80a721bb68ebfbc930e07b65ba72f2da88968ea5c06fd6ca3d3a127" [[package]] name = "fixedbitset" @@ -462,14 +462,14 @@ checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" [[package]] name = "getrandom" -version = "0.3.3" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", "libc", "r-efi", - "wasi", + "wasip2", ] [[package]] @@ -486,9 +486,9 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403" [[package]] name = "half" -version = "2.7.0" +version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e54c115d4f30f52c67202f079c5f9d8b49db4691f460fdb0b4c2e838261b2ba5" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ "cfg-if", "crunchy", @@ -660,9 +660,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.11.4" +version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b0f83760fb341a774ed326568e19f5a863af4a952def8c39f9ab92fd95b88e5" +checksum = "6717a8d2a5a929a1a2eb43a12812498ed141a0bcfb7e8f7844fbdbe4303bba9f" dependencies = [ "equivalent", "hashbrown 0.16.0", @@ -676,9 +676,9 @@ checksum = "7655c9839580ee829dfacba1d1278c2b7883e50a277ff7541299489d6bdfdc45" [[package]] name = "is_terminal_polyfill" -version = "1.70.1" +version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" [[package]] name = "itertools" @@ -835,9 +835,9 @@ checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" [[package]] name = "once_cell_polyfill" -version = "1.70.1" +version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4895175b425cb1f87721b59f0f286c2092bd4af812243672510e1ac53e2e0ad" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] name = "owo-colors" @@ -1151,9 +1151,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.11.3" +version = "1.12.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b5288124840bee7b386bc413c487869b360b2b4ec421ea56425128692f2a82c" +checksum = "843bc0191f75f3e22651ae5f1e72939ab2f72a4bc30fa80a066bd66edefc24d4" dependencies = [ "aho-corasick", "memchr", @@ -1163,9 +1163,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.11" +version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "833eb9ce86d40ef33cb1306d8accf7bc8ec2bfea4355cbdebb3df68b40925cad" +checksum = "5276caf25ac86c8d810222b3dbb938e512c55c6831a10f3e6ed1c93b84041f1c" dependencies = [ "aho-corasick", "memchr", @@ -1174,9 +1174,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.6" +version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" +checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" [[package]] name = "rusqlite" @@ -1317,7 +1317,7 @@ name = "simd" version = "0.0.0" dependencies = [ "cc", - "half 2.7.0", + "half 2.7.1", "rand", "seq-macro", "simd_macros", @@ -1367,9 +1367,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.106" +version = "2.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ede7c438028d4436d71104916910f5bb611972c5cfd7f89b8300a8186e6fada6" +checksum = "2a26dbd934e5451d21ef060c018dae56fc073894c5a7896f882928a76e6d081b" dependencies = [ "proc-macro2", "quote", @@ -1476,9 +1476,9 @@ checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] name = "unicode-ident" -version = "1.0.19" +version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f63a545481291138910575129486daeaf8ac54aee4387fe7906919f7830c7d9d" +checksum = "462eeb75aeb73aea900253ce739c8e18a67423fadf006037cd3ff27e82748a06" [[package]] name = "unicode-segmentation" @@ -1646,15 +1646,6 @@ dependencies = [ "winapi-util", ] -[[package]] -name = "wasi" -version = "0.14.7+wasi-0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" -dependencies = [ - "wasip2", -] - [[package]] name = "wasip2" version = "1.0.1+wasi-0.2.4" diff --git a/crates/feistel/src/lib.rs b/crates/feistel/src/lib.rs index afb184be..54b16380 100644 --- a/crates/feistel/src/lib.rs +++ b/crates/feistel/src/lib.rs @@ -83,7 +83,7 @@ fn is_a_permutation() { let buffer = [round.to_le_bytes(), x.to_le_bytes(), key_0, key_1]; wyhash::wyhash(buffer.as_flattened(), 0) as u32 }; - for width in (0..20).step_by(2) { + for width in (0..if cfg!(not(miri)) { 20 } else { 10 }).step_by(2) { let mut y = Vec::new(); for x in 0..1 << width { y.push(feistel::(width, x, 8, secret)); @@ -100,7 +100,7 @@ fn is_a_permutation() { #[test] fn sample() { - let n = 6370_u32; + let n = if cfg!(not(miri)) { 6370_u32 } else { 637_u32 }; let width = (n.ilog2() + 1).next_multiple_of(2); let key_0 = rand::Rng::random(&mut rand::rng()); let key_1 = rand::Rng::random(&mut rand::rng()); diff --git a/crates/make/Cargo.toml b/crates/make/Cargo.toml index 9c1e7182..49a781e0 100644 --- a/crates/make/Cargo.toml +++ b/crates/make/Cargo.toml @@ -5,7 +5,7 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.48", features = ["derive", "env"] } +clap = { version = "4.5.50", features = ["derive", "env"] } object = { version = "0.37.3", default-features = false, features = [ "read", "wasm", diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index f4c9fc7e..0703c552 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -12,7 +12,7 @@ experimental_math = [] [dependencies] simd_macros = { path = "../simd_macros" } -half = { version = "2.7.0", features = ["zerocopy"] } +half = { version = "2.7.1", features = ["zerocopy"] } seq-macro.workspace = true zerocopy.workspace = true @@ -20,7 +20,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.40" +cc = "1.2.41" which = "8.0.0" [lints] diff --git a/crates/simd/src/fht.rs b/crates/simd/src/fht.rs index e7effbab..1378ff1a 100644 --- a/crates/simd/src/fht.rs +++ b/crates/simd/src/fht.rs @@ -156,7 +156,7 @@ mod tests { const EPSILON: f32 = 1e-6; let mut rng = rand::rng(); let mut n = 1_usize; - while n <= 65536 { + while n <= if cfg!(not(miri)) { 65536 } else { 4096 } { let x = (0..n) .map(|_| rng.random_range(-1.0_f32..=1.0_f32)) .collect::>(); diff --git a/crates/simd/src/u8.rs b/crates/simd/src/u8.rs index f0cfa839..2985413a 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/u8.rs @@ -701,7 +701,7 @@ mod reduce_sum_of_x_as_u32 { sum } - #[cfg(all(target_arch = "aarch64", test))] + #[cfg(all(target_arch = "aarch64", test, not(miri)))] #[test] fn reduce_sum_of_x_as_u32_a2_test() { use rand::Rng; diff --git a/crates/vchordrq/src/fast_heap.rs b/crates/vchordrq/src/fast_heap.rs index 27149082..d4c61711 100644 --- a/crates/vchordrq/src/fast_heap.rs +++ b/crates/vchordrq/src/fast_heap.rs @@ -92,10 +92,8 @@ impl Sequence for FastHeap { #[test] fn test_fast_heap() { - for _ in 0..1000 { - let sequence = (0..10000) - .map(|_| rand::random::()) - .collect::>(); + for _ in 0..if cfg!(not(miri)) { 1000 } else { 1 } { + let sequence = (0..1000).map(|_| rand::random::()).collect::>(); let answer = { let mut x = sequence.clone(); x.sort_by_key(|x| std::cmp::Reverse(*x)); From 49ce09ce3964c79c6389854001cdf59230357d21 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 23 Oct 2025 16:40:27 +0800 Subject: [PATCH 236/324] fix: don't use rayon `fold` and `reduce` (#363) Signed-off-by: usamoi --- .github/workflows/check.yml | 2 +- crates/k_means/src/flat.rs | 166 ++++++++++++++++ crates/k_means/src/lib.rs | 315 ++++-------------------------- crates/k_means/src/quick.rs | 79 ++++++++ crates/k_means/src/rabitq.rs | 205 +++++++++++++++++++ crates/k_means/src/square.rs | 110 +++++++++++ crates/rabitq/src/bit.rs | 42 +++- crates/rabitq/src/rotate.rs | 78 ++++++-- src/index/vchordrq/am/am_build.rs | 75 ++++--- 9 files changed, 752 insertions(+), 320 deletions(-) create mode 100644 crates/k_means/src/flat.rs create mode 100644 crates/k_means/src/quick.rs create mode 100644 crates/k_means/src/rabitq.rs create mode 100644 crates/k_means/src/square.rs diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index d0b14177..a2cd63b8 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -40,7 +40,7 @@ jobs: uses: astral-sh/ruff-action@v1 - name: Rustfmt - run: cargo fmt --check + run: cargo fmt --check -- --config-path /dev/null --config imports_granularity=Module - name: Deny run: cargo deny check diff --git a/crates/k_means/src/flat.rs b/crates/k_means/src/flat.rs new file mode 100644 index 00000000..810f1da1 --- /dev/null +++ b/crates/k_means/src/flat.rs @@ -0,0 +1,166 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::square::Square; +use crate::{KMeans, This}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; +use simd::Floating; + +struct Flat { + this: This, +} + +impl KMeans for Flat { + fn this(&mut self) -> &mut This { + &mut self.this + } + + fn assign(&mut self) { + let this = &mut self.this; + this.pool.install(|| { + this.targets + .par_iter_mut() + .zip(this.samples.into_par_iter()) + .for_each(|(target, sample)| { + let mut result = (f32::INFINITY, 0); + for (j, centroid) in this.centroids.into_iter().enumerate() { + let dis_2 = f32::reduce_sum_of_d2(sample, centroid); + if dis_2 <= result.0 { + result = (dis_2, j); + } + } + *target = result.1; + }); + }); + } + + fn update(&mut self) { + let this = &mut self.this; + this.pool.install(|| { + const DELTA: f32 = 9.7656e-4_f32; + + let d = this.d; + let n = this.samples.len(); + let c = this.c; + + let list = rayon::broadcast({ + |ctx| { + let mut sum = Square::from_zeros(d, c); + let mut count = vec![0.0f32; c]; + for i in (ctx.index()..this.samples.len()).step_by(ctx.num_threads()) { + let target = this.targets[i]; + let sample = &this.samples[i]; + f32::vector_add_inplace(&mut sum[target], sample); + count[target] += 1.0; + } + (sum, count) + } + }); + let mut sum = Square::from_zeros(d, c); + let mut count = vec![0.0f32; c]; + for (sum_1, count_1) in list { + for i in 0..c { + f32::vector_add_inplace(&mut sum[i], &sum_1[i]); + count[i] += count_1[i]; + } + } + + sum.par_iter_mut() + .enumerate() + .for_each(|(i, sum)| f32::vector_mul_scalar_inplace(sum, 1.0 / count[i])); + + this.centroids = sum; + + for i in 0..c { + if count[i] != 0.0f32 { + continue; + } + let mut o = 0; + loop { + let alpha = this.rng.random_range(0.0..1.0f32); + let beta = (count[o] - 1.0) / (n - c) as f32; + if alpha < beta { + break; + } + o = (o + 1) % c; + } + this.centroids.copy_within(o..o + 1, i); + vector_mul_scalars_inplace(&mut this.centroids[i], [1.0 + DELTA, 1.0 - DELTA]); + vector_mul_scalars_inplace(&mut this.centroids[o], [1.0 - DELTA, 1.0 + DELTA]); + count[i] = count[o] / 2.0; + count[o] -= count[i]; + } + }); + } + + fn finish(self: Box) -> Square { + let this = self.this; + this.centroids + } +} + +pub fn new( + d: usize, + samples: Square, + c: usize, + num_threads: usize, + seed: [u8; 32], +) -> Box { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .build() + .expect("failed to build thread pool"); + let mut rng = StdRng::from_seed(seed); + + let mut centroids = Square::new(d); + + for index in rand::seq::index::sample(&mut rng, samples.len(), c.min(samples.len())) { + centroids.push_slice(&samples[index]); + } + + if centroids.is_empty() && c == 1 { + centroids.push_iter(std::iter::repeat_n(0.0, d as _)); + } + + while centroids.len() < c { + centroids.push_iter((0..d).map(|_| rng.random_range(-1.0f32..1.0f32))); + } + + let targets = vec![0; samples.len()]; + + Box::new(Flat { + this: This { + pool, + d, + c, + centroids, + targets, + rng, + samples, + }, + }) +} + +fn vector_mul_scalars_inplace(this: &mut [f32], scalars: [f32; 2]) { + let n: usize = this.len(); + for i in 0..n { + if i % 2 == 0 { + this[i] *= scalars[0]; + } else { + this[i] *= scalars[1]; + } + } +} diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index aab50292..bc70a8cb 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -12,291 +12,56 @@ // // Copyright (c) 2025 TensorChord Inc. -use rabitq::bit::block::BlockCode; -use rabitq::packing::{any_pack, padding_pack}; -use rand::rngs::StdRng; -use rand::{Rng, SeedableRng}; -use rayon::iter::{IntoParallelIterator, ParallelIterator}; -use simd::Floating; - -pub fn k_means( - num_threads: usize, - mut check: impl FnMut(usize), - c: usize, - dims: usize, - samples: &[Vec], - is_spherical: bool, - iterations: usize, -) -> Vec> { - assert!(c > 0); - assert!(dims > 0); - let n = samples.len(); - if n <= c { - quick_centers(c, dims, samples.to_vec(), is_spherical) - } else { - rayon::ThreadPoolBuilder::new() - .num_threads(num_threads) - .build_scoped( - |thread| thread.run(), - move |pool| { - let compute = |centroids: &[Vec]| { - if n >= 1024 && c >= 1024 { - rabitq_index(n, c, samples, centroids) - } else { - flat_index(n, c, samples, centroids) - } - }; - let mut lloyd_k_means = - pool.install(|| LloydKMeans::new(c, dims, samples, is_spherical, compute)); - for i in 0..iterations { - check(i); - if pool.install(|| lloyd_k_means.iterate()) { - break; - } - } - pool.install(|| lloyd_k_means.finish()) - }, - ) - .expect("failed to build thread pool") - } -} - -fn quick_centers( - c: usize, - dims: usize, - samples: Vec>, - is_spherical: bool, -) -> Vec> { - let n = samples.len(); - assert!(c >= n); - if c == 1 && n == 0 { - return vec![vec![0.0; dims]]; - } - let mut rng = rand::rng(); - let mut centroids = samples; - for _ in n..c { - let r = (0..dims) - .map(|_| rng.random_range(-1.0f32..1.0f32)) - .collect(); - centroids.push(r); - } - if is_spherical { - for i in 0..c { - let centroid = &mut centroids[i]; - let l = f32::reduce_sum_of_x2(centroid).sqrt(); - f32::vector_mul_scalar_inplace(centroid, 1.0 / l); - } - } - centroids -} +pub mod flat; +pub mod quick; +pub mod rabitq; +pub mod square; -fn rabitq_index(n: usize, c: usize, samples: &[Vec], centroids: &[Vec]) -> Vec { - let branches = { - let mut branches = Vec::new(); - for centroid in centroids { - let code = rabitq::bit::code(centroid); - branches.push(code); - } - branches - }; - struct Block { - dis_u_2: [f32; 32], - factor_cnt: [f32; 32], - factor_ip: [f32; 32], - factor_err: [f32; 32], - elements: Vec<[u8; 16]>, - } - impl Block { - fn code(&self) -> BlockCode<'_> { - ( - &self.dis_u_2, - &self.factor_cnt, - &self.factor_ip, - &self.factor_err, - &self.elements, - ) - } - } - let mut blocks = Vec::new(); - for chunk in branches.chunks(32) { - blocks.push(Block { - dis_u_2: any_pack(chunk.iter().map(|x| x.0.dis_u_2)), - factor_cnt: any_pack(chunk.iter().map(|x| x.0.factor_cnt)), - factor_ip: any_pack(chunk.iter().map(|x| x.0.factor_ip)), - factor_err: any_pack(chunk.iter().map(|x| x.0.factor_err)), - elements: padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.1))), - }); - } - (0..n) - .into_par_iter() - .map(|i| { - let lut = rabitq::bit::block::preprocess(&samples[i]); - let mut result = (f32::INFINITY, 0); - for block in 0..c.div_ceil(32) { - let returns = rabitq::bit::block::full_process_l2(blocks[block].code(), &lut); - let lowerbound = returns.map(|(rough, err)| rough - err * 1.9); - for j in block * 32..std::cmp::min(block * 32 + 32, c) { - if lowerbound[j - block * 32] < result.0 { - let dis = f32::reduce_sum_of_d2(&samples[i], ¢roids[j]); - if dis <= result.0 { - result = (dis, j); - } - } - } - } - result.1 - }) - .collect::>() -} - -fn flat_index(n: usize, c: usize, samples: &[Vec], centroids: &[Vec]) -> Vec { - (0..n) - .into_par_iter() - .map(|i| { - let mut result = (f32::INFINITY, 0); - for j in 0..c { - let dis_2 = f32::reduce_sum_of_d2(&samples[i], ¢roids[j]); - if dis_2 <= result.0 { - result = (dis_2, j); - } - } - result.1 - }) - .collect::>() -} +use crate::square::Square; +use rand::rngs::StdRng; -struct LloydKMeans<'a, F> { - dims: usize, - c: usize, - is_spherical: bool, - centroids: Vec>, - assign: Vec, +pub struct This { + pool: rayon::ThreadPool, rng: StdRng, - samples: &'a [Vec], - compute: F, + d: usize, + samples: Square, + c: usize, + centroids: Square, + targets: Vec, } -const DELTA: f32 = 9.7656e-4_f32; - -impl<'a, F: Fn(&[Vec]) -> Vec> LloydKMeans<'a, F> { - fn new(c: usize, dims: usize, samples: &'a [Vec], is_spherical: bool, compute: F) -> Self { - let n = samples.len(); - - let mut rng = StdRng::from_seed([7; 32]); - let mut centroids = Vec::with_capacity(c); - - for index in rand::seq::index::sample(&mut rng, n, c).into_iter() { - centroids.push(samples[index].clone()); - } - - let assign = (0..n) - .into_par_iter() - .map(|i| { - let mut result = (f32::INFINITY, 0); - for j in 0..c { - let dis_2 = f32::reduce_sum_of_d2(&samples[i], ¢roids[j]); - if dis_2 <= result.0 { - result = (dis_2, j); - } - } - result.1 - }) - .collect::>(); - - Self { - dims, - c, - is_spherical, - centroids, - assign, - rng, - samples, - compute, - } - } - - fn iterate(&mut self) -> bool { - let dims = self.dims; - let c = self.c; - let rand = &mut self.rng; - let samples = self.samples; - let n = samples.len(); - - let (sum, mut count) = (0..n) - .into_par_iter() - .fold( - || (vec![vec![f32::zero(); dims]; c], vec![0.0f32; c]), - |(mut sum, mut count), i| { - f32::vector_add_inplace(&mut sum[self.assign[i]], &samples[i]); - count[self.assign[i]] += 1.0; - (sum, count) - }, - ) - .reduce( - || (vec![vec![f32::zero(); dims]; c], vec![0.0f32; c]), - |(mut sum, mut count), (sum_1, count_1)| { - for i in 0..c { - f32::vector_add_inplace(&mut sum[i], &sum_1[i]); - count[i] += count_1[i]; - } - (sum, count) - }, - ); - - let mut centroids = (0..c) - .into_par_iter() - .map(|i| f32::vector_mul_scalar(&sum[i], 1.0 / count[i])) - .collect::>(); - - for i in 0..c { - if count[i] != 0.0f32 { - continue; - } - let mut o = 0; - loop { - let alpha = rand.random_range(0.0..1.0f32); - let beta = (count[o] - 1.0) / (n - c) as f32; - if alpha < beta { - break; - } - o = (o + 1) % c; - } - centroids[i] = centroids[o].clone(); - vector_mul_scalars_inplace(&mut centroids[i], [1.0 + DELTA, 1.0 - DELTA]); - vector_mul_scalars_inplace(&mut centroids[o], [1.0 - DELTA, 1.0 + DELTA]); - count[i] = count[o] / 2.0; - count[o] -= count[i]; - } - - if self.is_spherical { - (&mut centroids).into_par_iter().for_each(|centroid| { +pub trait KMeans { + fn this(&mut self) -> &mut This; + fn assign(&mut self); + fn update(&mut self); + fn sphericalize(&mut self) { + use rayon::iter::{IntoParallelIterator, ParallelIterator}; + use simd::Floating; + let this = self.this(); + this.pool.install(|| { + (&mut this.centroids).into_par_iter().for_each(|centroid| { let l = f32::reduce_sum_of_x2(centroid).sqrt(); f32::vector_mul_scalar_inplace(centroid, 1.0 / l); }); - } - - let assign = (self.compute)(¢roids); - - let result = (0..n).all(|i| assign[i] == self.assign[i]); - - self.centroids = centroids; - self.assign = assign; - - result - } - - fn finish(self) -> Vec> { - self.centroids + }); } + fn finish(self: Box) -> Square; } -fn vector_mul_scalars_inplace(this: &mut [f32], scalars: [f32; 2]) { - let n: usize = this.len(); - for i in 0..n { - if i % 2 == 0 { - this[i] *= scalars[0]; - } else { - this[i] *= scalars[1]; - } +pub fn k_means( + d: usize, + samples: Square, + c: usize, + num_threads: usize, + seed: [u8; 32], +) -> Box { + assert!(d > 0 && c > 0 && num_threads > 0); + let n = samples.len(); + if n <= c { + quick::new(d, samples, c, num_threads, seed) + } else if n <= c * 2 { + flat::new(d, samples, c, num_threads, seed) + } else { + rabitq::new(d, samples, c, num_threads, seed) } } diff --git a/crates/k_means/src/quick.rs b/crates/k_means/src/quick.rs new file mode 100644 index 00000000..45a31e77 --- /dev/null +++ b/crates/k_means/src/quick.rs @@ -0,0 +1,79 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::square::Square; +use crate::{KMeans, This}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +struct Quick { + this: This, +} + +impl KMeans for Quick { + fn this(&mut self) -> &mut This { + &mut self.this + } + + fn assign(&mut self) {} + + fn update(&mut self) {} + + fn finish(self: Box) -> Square { + let this = self.this; + this.centroids + } +} + +pub fn new( + d: usize, + samples: Square, + c: usize, + num_threads: usize, + seed: [u8; 32], +) -> Box { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .build() + .expect("failed to build thread pool"); + let mut rng = StdRng::from_seed(seed); + + let mut centroids = Square::new(d); + + for index in rand::seq::index::sample(&mut rng, samples.len(), c.min(samples.len())) { + centroids.push_slice(&samples[index]); + } + + if centroids.is_empty() && c == 1 { + centroids.push_iter(std::iter::repeat_n(0.0, d as _)); + } + + while centroids.len() < c { + centroids.push_iter((0..d).map(|_| rng.random_range(-1.0f32..1.0f32))); + } + + let targets = vec![0; samples.len()]; + + Box::new(Quick { + this: This { + pool, + rng, + d, + samples, + c, + centroids, + targets, + }, + }) +} diff --git a/crates/k_means/src/rabitq.rs b/crates/k_means/src/rabitq.rs new file mode 100644 index 00000000..5749222c --- /dev/null +++ b/crates/k_means/src/rabitq.rs @@ -0,0 +1,205 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::square::Square; +use crate::{KMeans, This}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; +use simd::Floating; + +struct RaBitQ { + this: This, +} + +impl KMeans for RaBitQ { + fn this(&mut self) -> &mut This { + &mut self.this + } + + fn assign(&mut self) { + let this = &mut self.this; + this.pool.install(|| { + use rabitq::packing::{pack_to_u4, padding_pack}; + + let metadata = this + .centroids + .par_iter() + .map(rabitq::bit::code_metadata) + .collect::>(); + + let blocks = this + .centroids + .par_iter() + .chunks(32) + .map(|chunk| { + let f = |x: &&_| pack_to_u4(&rabitq::bit::code_elements(x)); + padding_pack(chunk.iter().map(f)) + }) + .collect::>(); + + this.targets + .par_iter_mut() + .zip(this.samples.into_par_iter()) + .for_each(|(target, sample)| { + let lut = rabitq::bit::block::preprocess(sample); + let mut result = (f32::INFINITY, 0); + let mut sum = [0u32; 32]; + for (j, centroid) in this.centroids.into_iter().enumerate() { + if j % 32 == 0 { + sum = rabitq::bit::block::accumulate(&blocks[j / 32], &lut.1); + } + let (rough, err) = + rabitq::bit::block::half_process_l2(sum[j % 32], metadata[j], lut.0); + let lowerbound = rough - err * 1.9; + if lowerbound < result.0 { + let dis_2 = f32::reduce_sum_of_d2(sample, centroid); + if dis_2 <= result.0 { + result = (dis_2, j); + } + } + } + *target = result.1; + }); + }); + } + + fn update(&mut self) { + let this = &mut self.this; + this.pool.install(|| { + const DELTA: f32 = 9.7656e-4_f32; + + let d = this.d; + let n = this.samples.len(); + let c = this.c; + + let list = rayon::broadcast({ + |ctx| { + let mut sum = Square::from_zeros(d, c); + let mut count = vec![0.0f32; c]; + for i in (ctx.index()..this.samples.len()).step_by(ctx.num_threads()) { + let target = this.targets[i]; + let sample = &this.samples[i]; + f32::vector_add_inplace(&mut sum[target], sample); + count[target] += 1.0; + } + (sum, count) + } + }); + let mut sum = Square::from_zeros(d, c); + let mut count = vec![0.0f32; c]; + for (sum_1, count_1) in list { + for i in 0..c { + f32::vector_add_inplace(&mut sum[i], &sum_1[i]); + count[i] += count_1[i]; + } + } + + sum.par_iter_mut() + .enumerate() + .for_each(|(i, sum)| f32::vector_mul_scalar_inplace(sum, 1.0 / count[i])); + + this.centroids = sum; + + for i in 0..c { + if count[i] != 0.0f32 { + continue; + } + let mut o = 0; + loop { + let alpha = this.rng.random_range(0.0..1.0f32); + let beta = (count[o] - 1.0) / (n - c) as f32; + if alpha < beta { + break; + } + o = (o + 1) % c; + } + this.centroids.copy_within(o..o + 1, i); + vector_mul_scalars_inplace(&mut this.centroids[i], [1.0 + DELTA, 1.0 - DELTA]); + vector_mul_scalars_inplace(&mut this.centroids[o], [1.0 - DELTA, 1.0 + DELTA]); + count[i] = count[o] / 2.0; + count[o] -= count[i]; + } + }); + } + + fn finish(self: Box) -> Square { + let mut this = self.this; + this.pool.install(|| { + this.centroids.par_iter_mut().for_each(|centroid| { + rabitq::rotate::rotate_reversed_inplace(centroid); + }); + }); + this.centroids + } +} + +pub fn new( + d: usize, + mut samples: Square, + c: usize, + num_threads: usize, + seed: [u8; 32], +) -> Box { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .build() + .expect("failed to build thread pool"); + let mut rng = StdRng::from_seed(seed); + + pool.install(|| { + samples.par_iter_mut().for_each(|sample| { + rabitq::rotate::rotate_inplace(sample); + }); + }); + + let mut centroids = Square::new(d); + + for index in rand::seq::index::sample(&mut rng, samples.len(), c.min(samples.len())) { + centroids.push_slice(&samples[index]); + } + + if centroids.is_empty() && c == 1 { + centroids.push_iter(std::iter::repeat_n(0.0, d as _)); + } + + while centroids.len() < c { + centroids.push_iter((0..d).map(|_| rng.random_range(-1.0f32..1.0f32))); + } + + let targets = vec![0; samples.len()]; + + Box::new(RaBitQ { + this: This { + pool, + d, + c, + centroids, + targets, + rng, + samples, + }, + }) +} + +fn vector_mul_scalars_inplace(this: &mut [f32], scalars: [f32; 2]) { + let n: usize = this.len(); + for i in 0..n { + if i % 2 == 0 { + this[i] *= scalars[0]; + } else { + this[i] *= scalars[1]; + } + } +} diff --git a/crates/k_means/src/square.rs b/crates/k_means/src/square.rs new file mode 100644 index 00000000..40692261 --- /dev/null +++ b/crates/k_means/src/square.rs @@ -0,0 +1,110 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +#[derive(Debug, Clone)] +pub struct Square { + d: usize, + p: Vec, +} + +impl Square { + pub fn d(&self) -> usize { + self.d + } + pub fn len(&self) -> usize { + self.p.len() / self.d + } + pub fn is_empty(&self) -> bool { + self.p.is_empty() + } + pub fn new(d: usize) -> Self { + Self { d, p: Vec::new() } + } + pub fn push_slice(&mut self, slice: &[f32]) { + assert_eq!(slice.len(), self.d); + self.p.extend_from_slice(slice); + } + pub fn push_iter(&mut self, iter: impl ExactSizeIterator) { + assert_eq!(iter.len(), self.d); + self.p.extend(iter); + } + pub fn copy_within>(&mut self, src: R, dest: usize) { + let src_start = src.start_bound().map(|x| self.d * x); + let src_end = src.end_bound().map(|x| self.d * x); + self.p.copy_within((src_start, src_end), self.d * dest); + } + pub fn from_zeros(d: usize, p: usize) -> Self { + Self { + d, + p: vec![0.0; d * p], + } + } + pub fn truncate(&mut self, len: usize) { + self.p.truncate(self.d * len); + } +} + +impl std::ops::Index for Square { + type Output = [f32]; + + fn index(&self, index: usize) -> &Self::Output { + &self.p[self.d * index..][..self.d] + } +} + +impl std::ops::IndexMut for Square { + fn index_mut(&mut self, index: usize) -> &mut Self::Output { + &mut self.p[self.d * index..][..self.d] + } +} + +impl<'a> IntoIterator for &'a Square { + type Item = &'a [f32]; + + type IntoIter = std::slice::ChunksExact<'a, f32>; + + fn into_iter(self) -> Self::IntoIter { + self.p.chunks_exact(self.d) + } +} + +impl<'a> IntoIterator for &'a mut Square { + type Item = &'a mut [f32]; + + type IntoIter = std::slice::ChunksExactMut<'a, f32>; + + fn into_iter(self) -> Self::IntoIter { + self.p.chunks_exact_mut(self.d) + } +} + +impl<'a> rayon::prelude::IntoParallelIterator for &'a Square { + type Item = &'a [f32]; + + type Iter = rayon::slice::ChunksExact<'a, f32>; + + fn into_par_iter(self) -> Self::Iter { + rayon::slice::ParallelSlice::par_chunks_exact(self.p.as_slice(), self.d) + } +} + +impl<'a> rayon::prelude::IntoParallelIterator for &'a mut Square { + type Item = &'a mut [f32]; + + type Iter = rayon::slice::ChunksExactMut<'a, f32>; + + fn into_par_iter(self) -> Self::Iter { + rayon::slice::ParallelSliceMut::par_chunks_exact_mut(self.p.as_mut_slice(), self.d) + } +} diff --git a/crates/rabitq/src/bit.rs b/crates/rabitq/src/bit.rs index 27dd5237..59c97b68 100644 --- a/crates/rabitq/src/bit.rs +++ b/crates/rabitq/src/bit.rs @@ -16,7 +16,7 @@ use binary::BinaryLut; use block::BlockLut; use simd::Floating; -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Default)] pub struct CodeMetadata { pub dis_u_2: f32, pub factor_cnt: f32, @@ -65,6 +65,35 @@ impl CodeMetadata { pub type Code = (CodeMetadata, Vec); +pub fn code_metadata(vector: &[f32]) -> CodeMetadata { + let n = vector.len(); + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x_2 = f32::reduce_sum_of_x2(vector); + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector.iter().filter(|x| x.is_sign_positive()).count(); + let cnt_neg = vector.iter().filter(|x| x.is_sign_negative()).count(); + cnt_pos as f32 - cnt_neg as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + factor_err: { + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (n as f32).sqrt(); + dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (n as f32 - 1.0).sqrt() + }, + } +} + +pub fn code_elements(vector: &[f32]) -> Vec { + let n = vector.len(); + let mut signs = Vec::new(); + for i in 0..n { + signs.push(vector[i].is_sign_positive()); + } + signs +} + pub fn code(vector: &[f32]) -> Code { let n = vector.len(); let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); @@ -321,6 +350,17 @@ pub mod block { ) } + #[inline(always)] + pub fn accumulate(t: &[[u8; 16]], s: &[[u8; 16]]) -> [u32; 32] { + use std::iter::zip; + let mut sum = [0_u32; 32]; + for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { + let delta = simd::fast_scan::scan(t, s); + simd::fast_scan::accu(&mut sum, &delta); + } + sum + } + #[inline(always)] pub fn full_process_l2( (dis_u_2, _, factor_ip, factor_err, t): BlockCode<'_>, diff --git a/crates/rabitq/src/rotate.rs b/crates/rabitq/src/rotate.rs index 1be5eb3f..fb54a96b 100644 --- a/crates/rabitq/src/rotate.rs +++ b/crates/rabitq/src/rotate.rs @@ -52,44 +52,100 @@ fn kacs_walk(result: &mut [f32]) { } pub fn rotate(vector: &[f32]) -> Vec { + let mut vector = vector.to_vec(); + rotate_inplace(&mut vector); + vector +} + +pub fn rotate_inplace(result: &mut [f32]) { use simd::Floating; use std::ops::Bound::{Excluded, Included, Unbounded}; - let mut result = vector.to_vec(); - let n = vector.len(); + let n = result.len(); let base = n.ilog2(); let scale = 1.0 / ((1_usize << base) as f32).sqrt(); let l = (Unbounded, Excluded(1_usize << base)); let r = (Included(n - (1_usize << base)), Unbounded); - simd::rotate::flip(&BITS_0, &mut result); + simd::rotate::flip(&BITS_0, result); simd::fht::fht(&mut result[l]); f32::vector_mul_scalar_inplace(&mut result[l], scale); if n != (1_usize << base) { - kacs_walk(&mut result); + kacs_walk(result); } - simd::rotate::flip(&BITS_1, &mut result); + simd::rotate::flip(&BITS_1, result); simd::fht::fht(&mut result[r]); f32::vector_mul_scalar_inplace(&mut result[r], scale); if n != (1_usize << base) { - kacs_walk(&mut result); + kacs_walk(result); } - simd::rotate::flip(&BITS_2, &mut result); + simd::rotate::flip(&BITS_2, result); simd::fht::fht(&mut result[l]); f32::vector_mul_scalar_inplace(&mut result[l], scale); if n != (1_usize << base) { - kacs_walk(&mut result); + kacs_walk(result); } - simd::rotate::flip(&BITS_3, &mut result); + simd::rotate::flip(&BITS_3, result); simd::fht::fht(&mut result[r]); f32::vector_mul_scalar_inplace(&mut result[r], scale); if n != (1_usize << base) { - kacs_walk(&mut result); + kacs_walk(result); + } +} + +pub fn rotate_reversed_inplace(result: &mut [f32]) { + use simd::Floating; + use std::ops::Bound::{Excluded, Included, Unbounded}; + + let n = result.len(); + let base = n.ilog2(); + let scale = 1.0 / ((1_usize << base) as f32).sqrt(); + + let l = (Unbounded, Excluded(1_usize << base)); + let r = (Included(n - (1_usize << base)), Unbounded); + + if n != (1_usize << base) { + kacs_walk(result); } + f32::vector_mul_scalar_inplace(&mut result[r], scale); + simd::fht::fht(&mut result[r]); + simd::rotate::flip(&BITS_3, result); + + if n != (1_usize << base) { + kacs_walk(result); + } + f32::vector_mul_scalar_inplace(&mut result[l], scale); + simd::fht::fht(&mut result[l]); + simd::rotate::flip(&BITS_2, result); + + if n != (1_usize << base) { + kacs_walk(result); + } + f32::vector_mul_scalar_inplace(&mut result[r], scale); + simd::fht::fht(&mut result[r]); + simd::rotate::flip(&BITS_1, result); + + if n != (1_usize << base) { + kacs_walk(result); + } + f32::vector_mul_scalar_inplace(&mut result[l], scale); + simd::fht::fht(&mut result[l]); + simd::rotate::flip(&BITS_0, result); +} - result +#[test] +fn reverse() { + let mut x = vec![2.0, 3.0, 4.0]; + rotate_inplace(&mut x); + assert!((x[0] - 3.981917).abs() < 1e-6); + assert!((x[1] - 1.8043789).abs() < 1e-6); + assert!((x[2] - 3.1446066).abs() < 1e-6); + rotate_reversed_inplace(&mut x); + assert!((x[0] - 2.0).abs() < 1e-6); + assert!((x[1] - 3.0).abs() < 1e-6); + assert!((x[2] - 4.0).abs() < 1e-6); } diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index b1c687da..bd37bc11 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -24,6 +24,7 @@ use crate::index::vchordrq::types::*; use index::relation::{ Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, }; +use k_means::square::Square; use simd::Floating; use std::ffi::CStr; use std::marker::PhantomData; @@ -233,14 +234,14 @@ pub unsafe extern "C-unwind" fn ambuild( let sampler = unsafe { HeapSampler::new(index_relation, heap_relation, snapshot) }; let mut sample = sampler.sample(); let samples = 'a: { + let mut samples = Square::new(vector_options.dims as _); let Some(max_number_of_samples) = internal_build .lists .last() .map(|x| x.saturating_mul(internal_build.sampling_factor)) else { - break 'a Vec::new(); + break 'a samples; }; - let mut samples = Vec::new(); while samples.len() < max_number_of_samples as usize { if let Some(mut tuple) = sample.next() { let (values, is_nulls) = tuple.build(); @@ -258,7 +259,7 @@ pub unsafe extern "C-unwind" fn ambuild( x.len() as u32, "invalid vector dimensions" ); - samples.push(x); + samples.push_slice(x.as_slice()); } } else { continue; @@ -1290,16 +1291,23 @@ fn make_default_build( fn make_internal_build( vector_options: VectorOptions, internal_build: VchordrqInternalBuildOptions, - samples: Vec, + samples: Square, reporter: &PostgresReporter, ) -> Vec> { use std::iter::once; let mut result = Vec::>::new(); + let mut samples = Some(samples); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { let input = if let Some(structure) = result.last() { - &structure.centroids + let mut input = Square::new(vector_options.dims as _); + for slice in structure.centroids.iter() { + input.push_slice(slice); + } + input + } else if let Some(samples) = samples.take() { + samples } else { - &samples + unreachable!() }; let num_threads = internal_build.build_threads as _; let num_points = input.len(); @@ -1318,28 +1326,30 @@ fn make_internal_build( "clustering: starting, using {num_threads} threads, clustering {num_points} vectors of {num_dims} dimension into {num_lists} clusters, in {num_iterations} iterations" ); } - let centroids = k_means::k_means( - num_threads, - |i| { - pgrx::check_for_interrupts!(); - if result.is_empty() { - let percentage = - ((i as f64 / num_iterations as f64) * 100.0).clamp(0.0, 100.0) as u16; - let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); - let phase = BuildPhase::new(BuildPhaseCode::InternalBuild, 1 + percentage) - .unwrap_or(default); - reporter.phase(phase); - } - if num_lists > 1 { - pgrx::info!("clustering: iteration {}", i + 1); - } - }, - num_lists, - num_dims, - input, - internal_build.spherical_centroids, - num_iterations, - ); + let mut f = k_means::k_means(num_dims, input, num_lists, num_threads, [7; 32]); + if internal_build.spherical_centroids { + f.sphericalize(); + } + for i in 0..num_iterations { + pgrx::check_for_interrupts!(); + if result.is_empty() { + let percentage = + ((i as f64 / num_iterations as f64) * 100.0).clamp(0.0, 100.0) as u16; + let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); + let phase = BuildPhase::new(BuildPhaseCode::InternalBuild, 1 + percentage) + .unwrap_or(default); + reporter.phase(phase); + } + if num_lists > 1 { + pgrx::info!("clustering: iteration {}", i + 1); + } + f.assign(); + f.update(); + if internal_build.spherical_centroids { + f.sphericalize(); + } + } + let centroids = f.finish(); if result.is_empty() { let percentage = 100; let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); @@ -1353,7 +1363,7 @@ fn make_internal_build( if let Some(structure) = result.last() { let mut children = vec![Vec::new(); centroids.len()]; for i in 0..structure.len() as u32 { - pub fn k_means_lookup(vector: &[f32], centroids: &[Normalized]) -> usize { + pub fn k_means_lookup(vector: &[f32], centroids: &Square) -> usize { assert_ne!(centroids.len(), 0); let mut result = (f32::INFINITY, 0); for i in 0..centroids.len() { @@ -1367,8 +1377,9 @@ fn make_internal_build( let target = k_means_lookup(&structure.centroids[i as usize], ¢roids); children[target].push(i); } - let (centroids, children) = std::iter::zip(centroids, children) - .filter(|(_, x)| !x.is_empty()) + let (centroids, children) = std::iter::zip(¢roids, children) + .filter(|(_, children)| !children.is_empty()) + .map(|(centroids, children)| (centroids.to_vec(), children)) .unzip::<_, _, Vec<_>, Vec<_>>(); result.push(Structure { centroids, @@ -1377,7 +1388,7 @@ fn make_internal_build( } else { let children = vec![Vec::new(); centroids.len()]; result.push(Structure { - centroids, + centroids: centroids.into_iter().map(|x| x.to_vec()).collect(), children, }); } From cb99b312a867945373134c515b3b3ee73a19d563 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 24 Oct 2025 14:02:50 +0800 Subject: [PATCH 237/324] chore: add `xtask clippy` (#364) Signed-off-by: usamoi --- .github/workflows/check.yml | 391 +++++++++++++++-------------- .github/workflows/release.yml | 70 +++--- Cargo.lock | 78 +++++- Makefile | 22 +- crates/{make => xtask}/Cargo.toml | 7 +- crates/{make => xtask}/src/main.rs | 71 +++++- devtools/dump.sh | 8 +- src/index/storage.rs | 21 +- 8 files changed, 405 insertions(+), 263 deletions(-) rename crates/{make => xtask}/Cargo.toml (66%) rename crates/{make => xtask}/src/main.rs (84%) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index a2cd63b8..5d89ca94 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -21,12 +21,13 @@ jobs: steps: - name: Set up Environment run: | - rustup update - curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo curl -fsSL https://github.com/EmbarkStudios/cargo-deny/releases/download/0.18.2/cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl/cargo-deny | install -m 755 /dev/stdin /usr/local/bin/cargo-deny + - name: Install Rust + run: rustup update + - name: Checkout uses: actions/checkout@v4 @@ -118,7 +119,7 @@ jobs: RUST_BACKTRACE: "1" steps: - - name: Set up Environment + - name: Install Rust run: | rustup default nightly rustup component add miri @@ -155,8 +156,6 @@ jobs: steps: - name: Set up Environment run: | - rustup update - sudo apt-get update if [ "$(uname -m)" == "x86_64" ]; then @@ -165,10 +164,13 @@ jobs: sudo mv /opt/sde-external-9.48.0-2024-11-25-lin /opt/sde fi - if [ "$(uname -m)" == "aarch64" ]; then + if [ "$(uname -m)" == "aarch64" ]; then sudo apt-get install -y qemu-user-static fi + - name: Install Rust + run: rustup update + - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -176,10 +178,7 @@ jobs: uses: actions/checkout@v4 - name: Cargo Test - run: | - cargo test --locked \ - --workspace --exclude vchord --no-fail-fast \ - -- --no-capture + run: cargo test --locked --workspace --exclude vchord --no-fail-fast -- --no-capture - name: Cargo Test (QEMU) run: | @@ -226,41 +225,51 @@ jobs: steps: - name: Set up Environment run: | - rustup update + sudo apt-get update - sudo apt-get remove -y '^postgres.*' '^libpq.*' - sudo apt-get purge -y '^postgres.*' '^libpq.*' + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + - name: Install Rust + run: rustup update + + - name: Install Clang + run: | curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 - sudo apt-get update - sudo apt-get install -y postgresql-common - sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + - name: Install PostgreSQL & pgvector + run: | + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' - sudo apt-get update - sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} - echo PGRX_PG_CONFIG_PATH=pg_config >> $GITHUB_ENV - echo PG_CONFIG=pg_config >> $GITHUB_ENV + sudo apt-get install -y --no-install-recommends postgresql-common + sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo apt-get install -y --no-install-recommends postgresql-server-dev-${{ matrix.version }} + sudo apt-get install -y --no-install-recommends postgresql-${{ matrix.version }} - sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector echo "local all all trust" | sudo tee /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf echo "host all all 127.0.0.1/32 trust" | sudo tee -a /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf echo "host all all ::1/128 trust" | sudo tee -a /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + sudo mkdir -p /etc/systemd/system/postgresql@.service.d/ echo "[Service]" | sudo tee /etc/systemd/system/postgresql@.service.d/limit.conf echo "LimitNOFILE=infinity" | sudo tee -a /etc/systemd/system/postgresql@.service.d/limit.conf echo "LimitMEMLOCK=infinity" | sudo tee -a /etc/systemd/system/postgresql@.service.d/limit.conf - sudo -iu postgres createuser -s -r $USER - sudo -iu postgres createdb -O $USER $USER + sudo systemctl daemon-reload + + sudo -iu postgres createuser -s -r $(whoami) + sudo -iu postgres createdb -O $(whoami) $(whoami) if [ "${{ matrix.version }}" -ge "18" ]; then - sudo -iu postgres psql -c 'ALTER SYSTEM SET io_method = io_uring' + psql -c 'ALTER SYSTEM SET io_method = io_uring' fi - sudo -iu postgres psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' - sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' + psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo systemctl stop postgresql - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + pg_config + echo PG_CONFIG=pg_config >> $GITHUB_ENV + + sudo apt-get install -y --no-install-recommends postgresql-${{ matrix.version }}-pgvector - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -269,18 +278,19 @@ jobs: uses: actions/checkout@v4 - name: Clippy - run: cargo clippy --locked --target ${{ matrix.arch }}-unknown-linux-gnu -p vchord --features pg${{ matrix.version }} --no-deps + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build - name: Install - run: | - make PG_CONFIG=$PG_CONFIG PROFILE=dev build - sudo make PG_CONFIG=$PG_CONFIG install + run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts uses: actions/upload-artifact@v4 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu - path: ./build/raw + path: ./build compression-level: 9 retention-days: 14 @@ -291,11 +301,11 @@ jobs: - name: Sqllogictest run: | - sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} if [ "${{ matrix.version }}" -ge "17" ]; then - sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi - name: Logging @@ -327,21 +337,29 @@ jobs: steps: - name: Set up Environment run: | - rustup update - brew update - echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest + + - name: Install Rust + run: rustup update - HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1 brew install postgresql@${{ matrix.version }} + - name: Install Clang + run: echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV + + - name: Install PostgreSQL & pgvector + run: | + brew install postgresql@${{ matrix.version }} brew services start postgresql@${{ matrix.version }} for i in {1..60}; do [ -S /tmp/.s.PGSQL.5432 ] && echo "PostgreSQL ready" && break || sleep 1; done [ -S /tmp/.s.PGSQL.5432 ] || echo "PostgreSQL socket not found after 60 seconds" - $(brew --prefix postgresql@${{ matrix.version }})/bin/createdb -O $USER $USER + + $(brew --prefix postgresql@${{ matrix.version }})/bin/createdb -O $(whoami) $(whoami) $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' brew services stop postgresql@${{ matrix.version }} - echo PGRX_PG_CONFIG_PATH=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV + + $(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config echo PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV mkdir ~/pgvector-install @@ -349,8 +367,6 @@ jobs: make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config install - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest - - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -358,19 +374,19 @@ jobs: uses: actions/checkout@v4 - name: Clippy - run: | - cargo clippy --locked -p vchord --target ${{ matrix.arch }}-apple-darwin --features pg${{ matrix.version }} --no-deps + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build - name: Install - run: | - make PG_CONFIG=$PG_CONFIG PROFILE=dev build - sudo make PG_CONFIG=$PG_CONFIG install + run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts uses: actions/upload-artifact@v4 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin - path: ./build/raw + path: ./build compression-level: 9 retention-days: 14 @@ -383,11 +399,11 @@ jobs: - name: Sqllogictest run: | - sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} if [ "${{ matrix.version }}" -ge "17" ]; then - sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi - name: Logging @@ -419,11 +435,16 @@ jobs: steps: - name: Set up Environment run: | - 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } + Invoke-WebRequest -Uri https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-${{ matrix.arch }}-pc-windows-msvc.zip -OutFile "$env:TEMP\sqllogictest-install.zip" + Expand-Archive -Path "$env:TEMP\sqllogictest-install.zip" -DestinationPath "D:\sqllogictest-install" -Force + Add-Content -Path $env:GITHUB_PATH -Value "D:\sqllogictest-install" - rustup update + - name: Install Rust + run: rustup update - & 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise\Common7\Tools\Launch-VsDevShell.ps1' -HostArch amd64 -Arch amd64 + - name: Install PostgreSQL & pgvector + run: | + 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } if ( "${{ matrix.version }}" -eq "13" ) { $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-13.22-1-windows-x64-binaries.zip" @@ -445,27 +466,27 @@ jobs: } Invoke-WebRequest -Uri $postgresql_url -OutFile "$env:TEMP\postgresql-install.zip" Expand-Archive -Path "$env:TEMP\postgresql-install.zip" -DestinationPath "D:\postgresql-install" -Force - Add-Content -Path $env:GITHUB_ENV -Value "PGRX_PG_CONFIG_PATH=D:\postgresql-install\pgsql\bin\pg_config.exe" - Add-Content -Path $env:GITHUB_ENV -Value "PG_CONFIG=D:\postgresql-install\pgsql\bin\pg_config.exe" D:\postgresql-install\pgsql\bin\initdb.exe -D D:\postgresql-install\pgsql\data -U postgres + D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data -l D:\postgresql-install\postgresql.log + D:\postgresql-install\pgsql\bin\createuser.exe -U postgres -s -r $env:USERNAME D:\postgresql-install\pgsql\bin\createdb.exe -O $env:USERNAME $env:USERNAME D:\postgresql-install\pgsql\bin\psql.exe -c 'ALTER SYSTEM SET max_worker_processes = 1024' D:\postgresql-install\pgsql\bin\psql.exe -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' D:\postgresql-install\pgsql\bin\pg_ctl.exe stop -D D:\postgresql-install\pgsql\data + D:\postgresql-install\pgsql\bin\pg_config.exe + Add-Content -Path $env:GITHUB_ENV -Value "PG_CONFIG=D:\postgresql-install\pgsql\bin\pg_config.exe" + Invoke-WebRequest -Uri "https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.zip" -OutFile "$env:TEMP\pgvector-install.zip" Expand-Archive -Path "$env:TEMP\pgvector-install.zip" -DestinationPath "D:\pgvector-install" -Force + & 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise\Common7\Tools\Launch-VsDevShell.ps1' -HostArch amd64 -Arch amd64 Push-Location -Path D:\pgvector-install\pgvector-0.8.1 nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" install Pop-Location - Invoke-WebRequest -Uri https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-${{ matrix.arch }}-pc-windows-msvc.zip -OutFile "$env:TEMP\sqllogictest-install.zip" - Expand-Archive -Path "$env:TEMP\sqllogictest-install.zip" -DestinationPath "D:\sqllogictest-install" -Force - Add-Content -Path $env:GITHUB_PATH -Value "D:\sqllogictest-install" - - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -473,19 +494,19 @@ jobs: uses: actions/checkout@v4 - name: Clippy - run: | - cargo clippy --locked --target ${{ matrix.arch }}-pc-windows-msvc -p vchord --features pg${{ matrix.version }} --no-deps + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build - name: Install - run: | - make PG_CONFIG=$env:PG_CONFIG PROFILE=dev build - make PG_CONFIG=$env:PG_CONFIG install + run: make install - name: Upload Artifacts uses: actions/upload-artifact@v4 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc - path: ./build/raw + path: ./build compression-level: 9 retention-days: 14 @@ -539,44 +560,51 @@ jobs: - name: Set up Environment run: | sed -i "/^ID=/s/alpine/NotpineForGHA/" /etc/os-release - apk add nodejs --update-cache + apk add --update-cache curl git make nodejs sudo zip mkdir /opt/bin ln -s /usr/bin/node /opt/bin/node - export USER=root - echo USER=$USER >> $GITHUB_ENV - apk add --no-cache sudo curl make zip clang18-dev postgresql${{ matrix.version }} postgresql${{ matrix.version }}-dev + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + - name: Install Rust + run: | curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain beta + echo $HOME/.cargo/bin >> $GITHUB_PATH echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV - echo PGRX_PG_CONFIG_PATH=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV - echo PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV + - name: Install Clang + run: sudo apk add --no-cache clang18-dev + - name: Install PostgreSQL & pgvector + run: | + sudo apk add --no-cache postgresql${{ matrix.version }} postgresql${{ matrix.version }}-dev sudo touch /var/log/postgresql.log sudo chmod 644 /var/log/postgresql.log sudo chown postgres /var/log/postgresql.log - install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql - install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql/data - install --verbose --directory --owner postgres --group postgres --mode 3777 /run/postgresql + sudo install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql + sudo install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql/data + sudo install --verbose --directory --owner postgres --group postgres --mode 3777 /run/postgresql sudo -iu postgres initdb -D /var/lib/postgresql/data + sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data -l /var/log/postgresql.log - sudo -iu postgres createuser -s -r $USER - sudo -iu postgres createdb -O $USER $USER + + sudo -iu postgres createuser -s -r $(whoami) + sudo -iu postgres createdb -O $(whoami) $(whoami) if [ "${{ matrix.version }}" -ge "18" ]; then - sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET io_method = io_uring' + psql -c 'ALTER SYSTEM SET io_method = io_uring' fi - sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET max_worker_processes = 1024' - sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' + psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo -iu postgres pg_ctl stop -D /var/lib/postgresql/data + /usr/libexec/postgresql${{ matrix.version }}/pg_config + echo PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV + mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config install - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -594,36 +622,34 @@ jobs: EOF - name: Clippy - run: | - . "$HOME/.cargo/env" - cargo clippy --locked --target ${{ matrix.arch }}-unknown-linux-musl -p vchord --features pg${{ matrix.version }} --no-deps + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build - name: Install - run: | - . "$HOME/.cargo/env" - make PG_CONFIG=$PG_CONFIG PROFILE=dev build - sudo make PG_CONFIG=$PG_CONFIG install + run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts uses: actions/upload-artifact@v4 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-musl - path: ./build/raw + path: ./build compression-level: 9 retention-days: 14 - name: Service run: | sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data -l /var/log/postgresql.log - psql -d $USER -U $USER -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - name: Sqllogictest run: | - sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} if [ "${{ matrix.version }}" -ge "17" ]; then - sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi - name: Logging @@ -707,13 +733,11 @@ jobs: steps: - name: Set up Environment run: | - rustup default beta - rustup component add rustfmt clippy - echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV - sudo apt-get update - sudo apt-get install -y qemu-user-static clang lld + sudo apt-get install -y qemu-user-static sudo systemctl enable --now systemd-binfmt + sudo touch /etc/sudoers.d/qemu + echo "Defaults env_keep += \"QEMU_CPU\"" | sudo tee -a /etc/sudoers.d/qemu sudo mkdir /sysroot curl -fsSL https://github.com/debuerreotype/docker-debian-artifacts/raw/refs/heads/dist-${{ matrix.platform }}/${{ matrix.debian_version }}/oci/blobs/rootfs.tar.gz | sudo tar -xz -C /sysroot sudo mount --bind /dev /sysroot/dev @@ -738,47 +762,59 @@ jobs: export QEMU_CPU echo QEMU_LD_PREFIX=$QEMU_LD_PREFIX >> $GITHUB_ENV echo QEMU_CPU=$QEMU_CPU >> $GITHUB_ENV - echo "Defaults env_keep += \"QEMU_CPU\"" | sudo tee -a /etc/sudoers + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Install Rust + run: | + rustup default beta + rustup component add rustfmt clippy + rustup target add ${{ matrix.rust_triple }} + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + + - name: Install Clang + run: sudo apt-get install -y clang lld + + - name: Install PostgreSQL & pgvector + run: | sudo apt-get remove -y '^postgres.*' '^libpq.*' sudo apt-get purge -y '^postgres.*' '^libpq.*' - sudo chroot /sysroot apt-get update - sudo chroot /sysroot apt-get install --no-install-recommends -y postgresql-common + + sudo chroot /sysroot apt-get install -y --no-install-recommends postgresql-common sudo chroot /sysroot /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y - sudo chroot /sysroot apt-get install --no-install-recommends -y postgresql-${{ matrix.version }} postgresql-server-dev-${{ matrix.version }} - sudo touch /usr/bin/pg_config - echo "#!/usr/bin/env bash" | sudo tee -a /usr/bin/pg_config - echo "sudo chroot /sysroot pg_config \"\$@\" \\" | sudo tee -a /usr/bin/pg_config - echo " | sed -E 's|^/(.*)$|/sysroot/\1|' \\" | sudo tee -a /usr/bin/pg_config - echo " | sed -E 's|^([A-Z]+) = /(.*)$|\1 = /sysroot/\2|'" | sudo tee -a /usr/bin/pg_config - sudo chmod 755 /usr/bin/pg_config - pg_config + sudo chroot /sysroot apt-get install -y --no-install-recommends postgresql-server-dev-${{ matrix.version }} + sudo chroot /sysroot apt-get install -y --no-install-recommends postgresql-${{ matrix.version }} echo "local all all trust" | sudo tee /sysroot/etc/postgresql/${{ matrix.version }}/main/pg_hba.conf echo "host all all 127.0.0.1/32 trust" | sudo tee -a /sysroot/etc/postgresql/${{ matrix.version }}/main/pg_hba.conf echo "host all all ::1/128 trust" | sudo tee -a /sysroot/etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main start - sudo chroot /sysroot sudo -iu postgres createuser -s -r $USER - sudo chroot /sysroot sudo -iu postgres createdb -O $USER $USER - sudo mkdir -p /sysroot/etc/systemd/system/postgresql@.service.d/ - echo "[Service]" | sudo tee /sysroot/etc/systemd/system/postgresql@.service.d/limit.conf - echo "LimitNOFILE=infinity" | sudo tee -a /sysroot/etc/systemd/system/postgresql@.service.d/limit.conf - echo "LimitMEMLOCK=infinity" | sudo tee -a /sysroot/etc/systemd/system/postgresql@.service.d/limit.conf + + sudo chroot /sysroot sudo -iu postgres createuser -s -r $(whoami) + sudo chroot /sysroot sudo -iu postgres createdb -O $(whoami) $(whoami) if [ "${{ matrix.version }}" -ge "18" ]; then - # io_uring is not supported on qemu-user - sudo chroot /sysroot sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET io_method = worker' + sudo chroot /sysroot sudo -iu postgres psql -c 'ALTER SYSTEM SET io_method = worker' fi - sudo chroot /sysroot sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET max_worker_processes = 1024' - sudo chroot /sysroot sudo -iu postgres psql -d $USER -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + sudo chroot /sysroot sudo -iu postgres psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' + sudo chroot /sysroot sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main stop + sudo touch /usr/bin/pg_config + echo "#!/usr/bin/env bash" | sudo tee -a /usr/bin/pg_config + echo "sudo chroot /sysroot pg_config \"\$@\" \\" | sudo tee -a /usr/bin/pg_config + echo " | sed -E 's|^/(.*)$|/sysroot/\1|' \\" | sudo tee -a /usr/bin/pg_config + echo " | sed -E 's|^([A-Z]+) = /(.*)$|\1 = /sysroot/\2|'" | sudo tee -a /usr/bin/pg_config + sudo chmod 755 /usr/bin/pg_config + + pg_config + echo PG_CONFIG=pg_config >> $GITHUB_ENV + mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" sudo make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -805,46 +841,37 @@ jobs: EOF - name: Cargo Test - run: | - cargo test --locked --target ${{ matrix.rust_triple }} \ - --workspace --exclude vchord --no-fail-fast \ - -- --no-capture + run: cargo test --locked --target ${{ matrix.rust_triple }} --workspace --exclude vchord --no-fail-fast -- --no-capture - name: Clippy - run: | - PGRX_PG_CONFIG_PATH=pg_config \ - cargo clippy --locked --target ${{ matrix.rust_triple }} \ - --workspace --features pg${{ matrix.version }} + run: make TARGET=${{ matrix.rust_triple }} PROFILE=dev RUNNER='qemu-${{ matrix.platform }}-static' clippy + + - name: Build + run: make TARGET=${{ matrix.rust_triple }} PROFILE=dev RUNNER='qemu-${{ matrix.platform }}-static' build - name: Install - run: | - make \ - TARGET=${{ matrix.rust_triple }} \ - PROFILE=dev \ - RUNNER='qemu-${{ matrix.platform }}-static' \ - build - sudo make install + run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts uses: actions/upload-artifact@v4 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.clang_triple }} - path: ./build/raw + path: ./build compression-level: 9 retention-days: 14 - name: Service run: | sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main start - sudo chroot /sysroot psql -d $USER -U $USER -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + sudo chroot /sysroot psql -d $(whoami) -U $(whoami) -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' - name: Sqllogictest run: | - sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} if [ "${{ matrix.version }}" -ge "17" ]; then - sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi - name: Logging @@ -886,37 +913,38 @@ jobs: - name: Set up Environment run: | sudo sed -i 's/^DownloadUser = alpm$/DownloadUser = runner/' /etc/pacman.conf - sudo pacman -Syu git github-cli nodejs --ignore systemd,systemd-libs,systemd-sysvcompat --noconfirm - export USER=runner - echo USER=$USER >> $GITHUB_ENV - - # fix home directory + sudo pacman -Syu --ignore systemd,systemd-libs,systemd-sysvcompat --noconfirm git nodejs sudo mkdir -p /var/lib/postgresql sudo chmod 700 /var/lib/postgresql sudo chown postgres:postgres /var/lib/postgresql - sudo pacman -S python clang valgrind liburing numactl --noconfirm + sudo pacman -S --noconfirm liburing numactl python valgrind + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | sudo install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Install Rust + run: | curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain beta + echo $HOME/.cargo/bin >> $GITHUB_PATH echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | sudo install -m 755 /dev/stdin /usr/local/bin/sqllogictest + - name: Install Clang + run: sudo pacman -S --noconfirm clang - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 - - name: Install PostgreSQL + - name: Install PostgreSQL & pgvector run: | + mkdir ~/postgresql-install if [ "${{ matrix.version }}" = "17" ]; then - VER=17.6 + curl -fsSL https://ftp.postgresql.org/pub/source/v17.6/postgresql-17.6.tar.bz2 | tar -xj -C ~/postgresql-install + pushd ~/postgresql-install/postgresql-17.6 fi if [ "${{ matrix.version }}" = "18" ]; then - VER=18.0 + curl -fsSL https://ftp.postgresql.org/pub/source/v18.0/postgresql-18.0.tar.bz2 | tar -xj -C ~/postgresql-install + pushd ~/postgresql-install/postgresql-18.0 fi - - mkdir ~/postgresql-install - curl -fsSL https://ftp.postgresql.org/pub/source/v$VER/postgresql-$VER.tar.bz2 | tar -xj -C ~/postgresql-install - pushd ~/postgresql-install/postgresql-$VER ./configure \ --prefix=/usr \ --sysconfdir=/etc \ @@ -958,19 +986,17 @@ jobs: echo ' --trace-children=yes \' | sudo tee -a /usr/bin/postmaster echo ' /usr/bin/postgres "$@"' | sudo tee -a /usr/bin/postmaster sudo chmod 755 /usr/bin/postmaster - sudo mkdir -p /var/lib/postgres sudo chmod 700 /var/lib/postgres sudo chown postgres:postgres /var/lib/postgres sudo touch /var/log/postgresql.log sudo chmod 700 /var/log/postgresql.log sudo chown postgres:postgres /var/log/postgresql.log - sudo -iu postgres initdb -D /var/lib/postgres/data - sudo -iu postgres pg_ctl start -D /var/lib/postgres/data -l /var/log/postgresql.log -p /usr/bin/postmaster - sudo -iu postgres createuser -s -r $USER - sudo -iu postgres createdb -O $USER $USER + + sudo -iu postgres createuser -s -r $(whoami) + sudo -iu postgres createdb -O $(whoami) $(whoami) if [ "${{ matrix.version }}" -ge "18" ]; then sudo -iu postgres psql -c 'ALTER SYSTEM SET io_method = io_uring' fi @@ -978,8 +1004,9 @@ jobs: sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' sudo -iu postgres pg_ctl stop -D /var/lib/postgres/data - - name: Install pgvector - run: | + pg_config + echo PG_CONFIG=pg_config >> $GITHUB_ENV + mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install pushd ~/pgvector-install/pgvector-0.8.1 @@ -991,24 +1018,16 @@ jobs: uses: actions/checkout@v4 - name: Cargo Test - run: | - . "$HOME/.cargo/env" - cargo test --locked --target x86_64-unknown-linux-gnu \ - --workspace --exclude vchord --no-fail-fast \ - -- --no-capture + run: cargo test --locked --workspace --exclude vchord --no-fail-fast -- --no-capture - name: Clippy - run: | - . "$HOME/.cargo/env" - PGRX_PG_CONFIG_PATH=pg_config \ - cargo clippy --locked --target x86_64-unknown-linux-gnu \ - --workspace --features pg${{ matrix.version }} + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build - name: Install - run: | - . "$HOME/.cargo/env" - make PROFILE=dev build - sudo make install + run: sudo make PG_CONFIG=$PG_CONFIG install - name: Service run: | @@ -1017,11 +1036,11 @@ jobs: - name: Sqllogictest run: | - sqllogictest --db $USER --user $USER './tests/general/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordg/*.slt' --label pg${{ matrix.version }} - sqllogictest --db $USER --user $USER './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} if [ "${{ matrix.version }}" -ge "17" ]; then - sqllogictest --db $USER --user $USER './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} fi - name: Logging diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index d6223e81..17c24812 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -52,49 +52,55 @@ jobs: steps: - name: Set up Environment run: | - rustup update - sudo apt-get update - sudo apt-get remove -y '^postgres.*' '^libpq.*' - sudo apt-get purge -y '^postgres.*' '^libpq.*' + - name: Install Rust + run: rustup update + - name: Install Clang + run: | curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 - sudo apt-get install -y postgresql-common - sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y - sudo apt-get install -y postgresql-server-dev-${{ matrix.version }} - echo PGRX_PG_CONFIG_PATH=pg_config >> $GITHUB_ENV - echo PG_CONFIG=pg_config >> $GITHUB_ENV + - name: Install PostgreSQL + run: | + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' - sudo apt-get install -y postgresql-${{ matrix.version }} postgresql-${{ matrix.version }}-pgvector + sudo apt-get install -y --no-install-recommends postgresql-common + sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo apt-get install -y --no-install-recommends postgresql-server-dev-${{ matrix.version }} - mkdir -p ~/.pgrx - touch ~/.pgrx/config.toml + pg_config - name: Checkout uses: actions/checkout@v4 - - name: Check + - name: Build + run: | + grep -q "default_version = '${{ needs.semver.outputs.SEMVER }}'" vchord.control || exit 1 + make build + + - name: Package (binary package) env: + GH_TOKEN: ${{ github.token }} SEMVER: ${{ needs.semver.outputs.SEMVER }} + VERSION: ${{ matrix.version }} run: | - grep -q "default_version = '${SEMVER}'" vchord.control || exit 1 + ARCH=$(uname -m) + (cd ./build && zip -r ~/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) + gh release upload --clobber $SEMVER ~/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip - - name: Build + - name: Package (debian package) env: + GH_TOKEN: ${{ github.token }} SEMVER: ${{ needs.semver.outputs.SEMVER }} VERSION: ${{ matrix.version }} - ARCH: ${{ matrix.arch }} - PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} run: | - make PG_CONFIG=$PG_CONFIG build - - (cd ./build/raw && zip -r ../postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) - - make DESTDIR="./build/deb" install - mkdir -p ./build/deb/DEBIAN + make DESTDIR="~/debian" install + ARCH=$(uname -m) + PLATFORM=$(dpkg --print-architecture) + mkdir -p ~/debian/DEBIAN echo "Package: postgresql-${VERSION}-vchord Version: ${SEMVER}-1 Depends: postgresql-${VERSION}, libgcc-s1, libc6 (>= 2.35) @@ -105,20 +111,10 @@ jobs: Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM Homepage: https://vectorchord.ai/ License: AGPL-3.0-only or Elastic-2.0" \ - > ./build/deb/DEBIAN/control - (cd ./build/deb && find usr -type f -print0 | xargs -0 md5sum) > ./build/deb/DEBIAN/md5sums - dpkg-deb --root-owner-group -Zxz --build ./build/deb/ ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb - - - name: Upload Artifacts - env: - GH_TOKEN: ${{ github.token }} - SEMVER: ${{ needs.semver.outputs.SEMVER }} - VERSION: ${{ matrix.version }} - ARCH: ${{ matrix.arch }} - PLATFORM: ${{ matrix.arch == 'x86_64' && 'amd64' || 'arm64' }} - run: | - gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb - gh release upload --clobber $SEMVER ./build/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip + > ~/debian/DEBIAN/control + (cd ~/debian && find usr -type f -print0 | xargs -0 md5sum) > ~/debian/DEBIAN/md5sums + dpkg-deb --root-owner-group -Zxz --build ~/debian ~/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb + gh release upload --clobber $SEMVER ~/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb pgxn: runs-on: "ubuntu-latest" diff --git a/Cargo.lock b/Cargo.lock index 4f095bca..a6f0eb23 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "aho-corasick" version = "1.1.3" @@ -244,6 +250,15 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -433,6 +448,16 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" +[[package]] +name = "flate2" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc5a4e564e38c699f2880d3fda590bedc2e69f3f84cd48b457bd892ce61d0aa9" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + [[package]] name = "fnv" version = "1.0.7" @@ -770,16 +795,6 @@ version = "0.4.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432" -[[package]] -name = "make" -version = "0.0.0" -dependencies = [ - "clap", - "object", - "shlex", - "target-triple", -] - [[package]] name = "memchr" version = "2.7.6" @@ -807,6 +822,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "nom" version = "7.1.3" @@ -823,7 +848,9 @@ version = "0.37.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" dependencies = [ + "flate2", "memchr", + "ruzstd", "wasmparser", ] @@ -1217,6 +1244,15 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +[[package]] +name = "ruzstd" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3640bec8aad418d7d03c72ea2de10d5c646a598f9883c7babc160d91e3c1b26c" +dependencies = [ + "twox-hash", +] + [[package]] name = "ryu" version = "1.0.20" @@ -1325,6 +1361,12 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "simd-adler32" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" + [[package]] name = "simd_macros" version = "0.0.0" @@ -1468,6 +1510,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df8b2b54733674ad286d16267dcfc7a71ed5c776e4ac7aa3c3e2561f7c637bf2" +[[package]] +name = "twox-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" + [[package]] name = "unescape" version = "0.1.0" @@ -1896,6 +1944,16 @@ dependencies = [ "tap", ] +[[package]] +name = "xtask" +version = "0.0.0" +dependencies = [ + "clap", + "object", + "shlex", + "target-triple", +] + [[package]] name = "yoke" version = "0.8.0" diff --git a/Makefile b/Makefile index 8f7390a4..fac35b23 100644 --- a/Makefile +++ b/Makefile @@ -1,20 +1,26 @@ PG_CONFIG ?= pg_config PKGLIBDIR := $(shell $(PG_CONFIG) --pkglibdir) SHAREDIR := $(shell $(PG_CONFIG) --sharedir) -MKDIR ?= mkdir -CP ?= cp +CP_R ?= cp -r +MKDIR_P ?= mkdir -p -.PHONY: all build install uninstall +.PHONY: all build clippy install uninstall all: build build: - PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p make -- build --output ./build/raw + PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p xtask -- build + +clippy: + PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p xtask -- clippy install: - $(MKDIR) -p $(DESTDIR)$(PKGLIBDIR) $(DESTDIR)$(SHAREDIR) && \ - $(CP) -r ./build/raw/pkglibdir/. $(DESTDIR)$(PKGLIBDIR) && \ - $(CP) -r ./build/raw/sharedir/. $(DESTDIR)$(SHAREDIR) + $(MKDIR_P) $(DESTDIR)$(PKGLIBDIR) && \ + $(CP_R) ./build/pkglibdir/. $(DESTDIR)$(PKGLIBDIR) && \ + $(MKDIR_P) $(DESTDIR)$(SHAREDIR) && \ + $(CP_R) ./build/sharedir/. $(DESTDIR)$(SHAREDIR) uninstall: - $(RM) $(wildcard $(DESTDIR)$(PKGLIBDIR)/vchord.*) $(wildcard $(DESTDIR)$(SHAREDIR)/extension/vchord.*) $(wildcard $(DESTDIR)$(SHAREDIR)/extension/vchord--*.sql) + $(RM) $(wildcard $(DESTDIR)$(PKGLIBDIR)/vchord.*) && \ + $(RM) $(wildcard $(DESTDIR)$(SHAREDIR)/extension/vchord.*) && \ + $(RM) $(wildcard $(DESTDIR)$(SHAREDIR)/extension/vchord--*.sql) diff --git a/crates/make/Cargo.toml b/crates/xtask/Cargo.toml similarity index 66% rename from crates/make/Cargo.toml rename to crates/xtask/Cargo.toml index 49a781e0..039ba184 100644 --- a/crates/make/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -1,15 +1,12 @@ [package] -name = "make" +name = "xtask" version.workspace = true edition.workspace = true publish = false [dependencies] clap = { version = "4.5.50", features = ["derive", "env"] } -object = { version = "0.37.3", default-features = false, features = [ - "read", - "wasm", -] } +object = { version = "0.37.3", features = ["read", "wasm"] } shlex = "1.3.0" target-triple = "0.1.4" diff --git a/crates/make/src/main.rs b/crates/xtask/src/main.rs similarity index 84% rename from crates/make/src/main.rs rename to crates/xtask/src/main.rs index 572774d1..ddcccd03 100644 --- a/crates/make/src/main.rs +++ b/crates/xtask/src/main.rs @@ -30,11 +30,12 @@ struct Cli { #[derive(Subcommand)] enum Commands { Build(BuildArgs), + Clippy(ClippyArgs), } #[derive(Args)] struct BuildArgs { - #[arg(short, long)] + #[arg(short, long, default_value = "./build")] output: String, #[arg(long, default_value = target_triple::TARGET, env = "TARGET")] target: String, @@ -44,6 +45,14 @@ struct BuildArgs { runner: Option, } +#[derive(Args)] +struct ClippyArgs { + #[arg(long, default_value = target_triple::TARGET, env = "TARGET")] + target: String, + #[arg(long, default_value = "release", env = "PROFILE")] + profile: String, +} + struct TargetSpecificInformation { is_macos: bool, is_windows: bool, @@ -272,7 +281,7 @@ fn generate( .args(["--target", target]) .args(["--features", pg_version]) .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) - .args(["--", "--cfg", "pgrx_embed"]) + .args(["--", "-C", "lto=off", "--cfg", "pgrx_embed"]) .env("PGRX_EMBED", &pgrx_embed) .stdout(Stdio::inherit()) .stderr(Stdio::inherit()); @@ -340,12 +349,32 @@ fn install_by_writing( Ok(()) } +fn clippy( + pg_config: impl AsRef, + pg_version: &str, + profile: &str, + target: &str, +) -> Result<(), Box> { + let mut command = Command::new("cargo"); + command + .args(["clippy"]) + .args(["--workspace"]) + .args(["--profile", profile]) + .args(["--target", target]) + .args(["--features", pg_version]) + .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_status = command.spawn()?.wait()?; + if !command_status.success() { + return Err(format!("Cargo clippy failed: {command_status}").into()); + } + Ok(()) +} + fn main() -> Result<(), Box> { let cli = Cli::parse(); - if !std::fs::exists("./vchord.control")? { - return Err("The script must be run from the VectorChord source directory.".into()); - } - let vchord_version = control_file("./vchord.control")?["default_version"].clone(); match cli.command { Commands::Build(BuildArgs { output, @@ -353,6 +382,10 @@ fn main() -> Result<(), Box> { profile, runner, }) => { + if !std::fs::exists("./vchord.control")? { + return Err("The script must be run from the VectorChord source directory.".into()); + } + let vchord_version = control_file("./vchord.control")?["default_version"].clone(); let runner = runner.and_then(|runner| shlex::split(&runner)); let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { PathBuf::from(value) @@ -428,6 +461,32 @@ fn main() -> Result<(), Box> { )?; } } + Commands::Clippy(ClippyArgs { target, profile }) => { + if !std::fs::exists("./vchord.control")? { + return Err("The script must be run from the VectorChord source directory.".into()); + } + let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { + PathBuf::from(value) + } else { + return Err("Environment variable `PGRX_PG_CONFIG_PATH` is not set.".into()); + }; + let pg_config = pg_config(&path)?; + let pg_version = { + let version = pg_config["VERSION"].clone(); + if let Some(prefix_stripped) = version.strip_prefix("PostgreSQL ") { + if let Some((stripped, _)) = + prefix_stripped.split_once(|c: char| !c.is_ascii_digit()) + { + format!("pg{stripped}",) + } else { + format!("pg{prefix_stripped}",) + } + } else { + return Err("PostgreSQL version is invalid.".into()); + } + }; + clippy(&path, &pg_version, &profile, &target)?; + } } Ok(()) } diff --git a/devtools/dump.sh b/devtools/dump.sh index 04e28daf..63e21672 100755 --- a/devtools/dump.sh +++ b/devtools/dump.sh @@ -1,15 +1,15 @@ #!/usr/bin/env bash set -e -mkdir -p $(dirname "$0")/../target/tools +mkdir -p $(dirname "$0")/../target/devtools if [ $# -eq 0 ] ; then echo "Usage: $0 INITIAL_VERSION VERSIONS.." echo "Dump the extension members. Install INITIAL_VERSION first and upgrade to every version in VERSIONS." echo "The extension members are installed in database \"vchord\", so use \"createdb vchord\" to setup." echo "Examples:" - echo " ./tools/dump.sh 0.1.11 0.2.0 > ./dump_upgrade.sql" - echo " ./tools/dump.sh 0.2.0 > ./dump_install.sql" + echo " ./devtools/dump.sh 0.1.11 0.2.0 > ./dump_upgrade.sql" + echo " ./devtools/dump.sh 0.2.0 > ./dump_install.sql" echo " diff ./dump_upgrade.sql ./dump_install.sql" exit 0 fi @@ -26,7 +26,7 @@ for arg in "$@"; do f+=("$x") done -so=$(realpath $(dirname "$0")/../target/tools/vchord.so) +so=$(realpath $(dirname "$0")/../target/devtools/vchord.so) $(dirname "$0")/dump-codegen.sh | gcc -I $(pg_config --includedir-server) -fPIC -shared -o $so -x c - sql=$(mktemp) diff --git a/src/index/storage.rs b/src/index/storage.rs index 2299c95c..6a3d0ee5 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -341,15 +341,15 @@ impl RelationWrite for PostgresRelation { ) -> PostgresBufferWriteGuard { unsafe { use pgrx::pg_sys::{ - BUFFER_LOCK_EXCLUSIVE, GENERIC_XLOG_FULL_IMAGE, GenericXLogRegisterBuffer, - GenericXLogStart, LockBuffer, + GENERIC_XLOG_FULL_IMAGE, GenericXLogRegisterBuffer, GenericXLogStart, }; let buf; #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15"))] { use pgrx::pg_sys::{ - ExclusiveLock, ForkNumber, LockRelationForExtension, ReadBufferExtended, - ReadBufferMode, UnlockRelationForExtension, + BUFFER_LOCK_EXCLUSIVE, ExclusiveLock, ForkNumber, LockBuffer, + LockRelationForExtension, ReadBufferExtended, ReadBufferMode, + UnlockRelationForExtension, }; LockRelationForExtension(self.raw, ExclusiveLock as _); buf = ReadBufferExtended( @@ -360,18 +360,25 @@ impl RelationWrite for PostgresRelation { std::ptr::null_mut(), ); UnlockRelationForExtension(self.raw, ExclusiveLock as _); + LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); } #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] { - use pgrx::pg_sys::{BufferManagerRelation, ExtendBufferedRel, ForkNumber}; + use pgrx::pg_sys::{ + BufferManagerRelation, ExtendBufferedFlags, ExtendBufferedRel, ForkNumber, + }; let bmr = BufferManagerRelation { rel: self.raw, smgr: std::ptr::null_mut(), relpersistence: 0, }; - buf = ExtendBufferedRel(bmr, ForkNumber::MAIN_FORKNUM, std::ptr::null_mut(), 0); + buf = ExtendBufferedRel( + bmr, + ForkNumber::MAIN_FORKNUM, + std::ptr::null_mut(), + ExtendBufferedFlags::EB_LOCK_FIRST as _, + ); } - LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); let state = GenericXLogStart(self.raw); let mut page = NonNull::new( GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) From bc37f4a7a3a4ed9597c8784ef03481ac0a9a8619 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 27 Oct 2025 12:39:00 +0800 Subject: [PATCH 238/324] feat: postgres buffer extend cache (#365) Signed-off-by: usamoi --- Cargo.lock | 7 ++ Cargo.toml | 1 + build.rs | 1 + crates/vchordrq/src/consume.rs | 41 +++++++++ crates/vchordrq/src/lib.rs | 2 + src/index/storage.rs | 7 +- src/index/storage/buffered.rs | 135 ++++++++++++++++++++++++++++++ src/index/vchordrq/am/am_build.rs | 40 +++++---- 8 files changed, 216 insertions(+), 18 deletions(-) create mode 100644 crates/vchordrq/src/consume.rs create mode 100644 src/index/storage/buffered.rs diff --git a/Cargo.lock b/Cargo.lock index a6f0eb23..56b057d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -87,6 +87,12 @@ version = "1.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" +[[package]] +name = "arrayvec" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" + [[package]] name = "bindgen" version = "0.71.1" @@ -1610,6 +1616,7 @@ name = "vchord" version = "0.0.0" dependencies = [ "always_equal", + "arrayvec", "bumpalo", "dary_heap", "distance", diff --git a/Cargo.toml b/Cargo.toml index 4345a1a9..340a70df 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,7 @@ vchordg = { path = "./crates/vchordg" } vchordrq = { path = "./crates/vchordrq" } vector = { path = "./crates/vector" } +arrayvec = "0.7.6" bumpalo.workspace = true dary_heap.workspace = true paste.workspace = true diff --git a/build.rs b/build.rs index cf775640..bc506af2 100644 --- a/build.rs +++ b/build.rs @@ -18,6 +18,7 @@ use std::error::Error; use std::process::{Command, Stdio}; fn main() -> Result<(), Box> { + println!("cargo::rerun-if-changed=vchord.control"); let version = 'version: { for line in std::fs::read_to_string("./vchord.control")?.lines() { if let Some(prefix_stripped) = line.strip_prefix("default_version = '") diff --git a/crates/vchordrq/src/consume.rs b/crates/vchordrq/src/consume.rs new file mode 100644 index 00000000..97749161 --- /dev/null +++ b/crates/vchordrq/src/consume.rs @@ -0,0 +1,41 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::tuples::*; +use crate::{Opaque, freepages}; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; + +pub fn consume(index: &R) +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let freepages_first = meta_tuple.freepages_first(); + + drop(meta_guard); + + let id = index + .extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + false, + ) + .id(); + + freepages::free(index, freepages_first, id); +} diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index 50c9c2b6..e2938dab 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -17,6 +17,7 @@ mod bulkdelete; mod cache; mod centroids; mod closure_lifetime_binder; +mod consume; mod cost; mod fast_heap; mod freepages; @@ -37,6 +38,7 @@ pub mod types; pub use build::build; pub use bulkdelete::{bulkdelete, bulkdelete_vectors}; pub use cache::cache; +pub use consume::consume; pub use cost::cost; pub use fast_heap::FastHeap; pub use insert::{InsertChooser, insert, insert_vector}; diff --git a/src/index/storage.rs b/src/index/storage.rs index 6a3d0ee5..a0b36950 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -12,6 +12,8 @@ // // Copyright (c) 2025 TensorChord Inc. +pub mod buffered; + use index::fetch::Fetch; use index::relation::{ Hints, Opaque, Page, PageGuard, ReadStream, Relation, RelationPrefetch, RelationRead, @@ -283,11 +285,12 @@ impl RelationRead for PostgresRelation { assert!(id != u32::MAX, "no such page"); unsafe { use pgrx::pg_sys::{ - BUFFER_LOCK_SHARE, BufferGetPage, LockBuffer, ReadBufferExtended, ReadBufferMode, + BUFFER_LOCK_SHARE, BufferGetPage, ForkNumber, LockBuffer, ReadBufferExtended, + ReadBufferMode, }; let buf = ReadBufferExtended( self.raw, - 0, + ForkNumber::MAIN_FORKNUM, id, ReadBufferMode::RBM_NORMAL, std::ptr::null_mut(), diff --git a/src/index/storage/buffered.rs b/src/index/storage/buffered.rs new file mode 100644 index 00000000..151987a3 --- /dev/null +++ b/src/index/storage/buffered.rs @@ -0,0 +1,135 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::index::storage::{ + PostgresBufferReadGuard, PostgresBufferWriteGuard, PostgresPage, PostgresRelation, +}; +use index::relation::{ + Opaque, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, +}; + +#[derive(Debug)] +pub struct BufferedPostgresRelation { + postgres: PostgresRelation, + list: std::cell::RefCell>, +} + +impl BufferedPostgresRelation { + pub unsafe fn new(raw: pgrx::pg_sys::Relation) -> Self { + Self { + postgres: unsafe { PostgresRelation::new(raw) }, + list: std::cell::RefCell::new(arrayvec::ArrayVec::new()), + } + } + pub fn release(&self) -> bool { + let list = self.list.borrow(); + !list.is_empty() + } +} + +impl Drop for BufferedPostgresRelation { + fn drop(&mut self) { + let free = self.list.get_mut(); + #[cfg(debug_assertions)] + assert!(free.is_empty()); + for &mut buf in free { + unsafe { + pgrx::pg_sys::ReleaseBuffer(buf); + } + } + } +} + +impl Relation for BufferedPostgresRelation { + type Page = PostgresPage; +} + +impl RelationReadTypes for BufferedPostgresRelation { + type ReadGuard<'a> = PostgresBufferReadGuard; +} + +impl RelationWriteTypes for BufferedPostgresRelation { + type WriteGuard<'a> = PostgresBufferWriteGuard; +} + +impl RelationRead for BufferedPostgresRelation { + fn read(&self, id: u32) -> Self::ReadGuard<'_> { + self.postgres.read(id) + } +} + +impl RelationWrite for BufferedPostgresRelation { + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.postgres.write(id, tracking_freespace) + } + fn search(&self, freespace: usize) -> Option> { + self.postgres.search(freespace) + } + #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15"))] + fn extend(&self, opaque: O, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.postgres.extend(opaque, tracking_freespace) + } + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + fn extend(&self, opaque: O, tracking_freespace: bool) -> Self::WriteGuard<'_> { + unsafe { + use pgrx::pg_sys::{ + BUFFER_LOCK_EXCLUSIVE, GENERIC_XLOG_FULL_IMAGE, GenericXLogRegisterBuffer, + GenericXLogStart, LockBuffer, + }; + use std::mem::MaybeUninit; + use std::ptr::NonNull; + let buf = { + let mut list = self.list.borrow_mut(); + if list.is_empty() { + use pgrx::pg_sys::{BufferManagerRelation, ExtendBufferedRelBy, ForkNumber}; + let bmr = BufferManagerRelation { + rel: self.postgres.raw, + smgr: std::ptr::null_mut(), + relpersistence: 0, + }; + let mut len = 0_u32; + _ = ExtendBufferedRelBy( + bmr, + ForkNumber::MAIN_FORKNUM, + std::ptr::null_mut(), + 0, + list.capacity() as _, + list.as_mut_ptr(), + &raw mut len, + ); + assert!(1 <= len && len <= list.capacity() as _); + list.set_len(len as _); + list.reverse(); + } + list.pop().expect("number of allocated pages is zero") + }; + LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); + let state = GenericXLogStart(self.postgres.raw); + let mut page = NonNull::new( + GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) + .cast::>>(), + ) + .expect("failed to get page"); + crate::index::storage::page_init(page.as_mut().as_mut_ptr(), opaque); + PostgresBufferWriteGuard { + raw: self.postgres.raw, + buf, + page: page.cast(), + state, + id: pgrx::pg_sys::BufferGetBlockNumber(buf), + tracking_freespace, + } + } + } +} diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index bd37bc11..75f27044 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -15,6 +15,7 @@ use crate::datatype::typmod::Typmod; use crate::index::fetcher::*; use crate::index::sample::{HeapSampler, Sample, Sampler, Tuple}; +use crate::index::storage::buffered::BufferedPostgresRelation; use crate::index::storage::{PostgresPage, PostgresRelation}; use crate::index::traverse::{HeapTraverser, Traverser}; use crate::index::vchordrq::am::Reloption; @@ -978,8 +979,6 @@ unsafe fn parallel_build( std::slice::from_raw_parts(vchordrqcached.add(8), bytes as _) }); - let index = unsafe { PostgresRelation::new(index_relation) }; - let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap_relation, tablescandesc) }; let opfamily = unsafe { opfamily(index_relation) }; let traverser = unsafe { HeapTraverser::new(heap_relation, index_relation, index_info, scan) }; @@ -1007,6 +1006,8 @@ unsafe fn parallel_build( let order = sync_0(); + let index = unsafe { BufferedPostgresRelation::new(index_relation) }; + match cached { VchordrqCachedReader::_0(_) => { traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { @@ -1047,7 +1048,7 @@ unsafe fn parallel_build( VchordrqCachedReader::_1(cached) => { let index = CachingRelation { cache: cached, - relation: index.clone(), + relation: &index, }; traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { let ctid = tuple.id(); @@ -1086,8 +1087,16 @@ unsafe fn parallel_build( } } + while index.release() { + vchordrq::consume(&index); + } + + drop(index); + sync_1(unsafe { (*vchordrqshared).indtuples }); + let index = unsafe { PostgresRelation::new(index_relation) }; + let mut chooser = ChooseSome { n: unsafe { (*vchordrqshared).nparticipants as usize }, k: order as usize, @@ -1111,8 +1120,6 @@ unsafe fn sequential_build( let cached = VchordrqCachedReader::deserialize_ref(vchordrqcached); - let index = unsafe { PostgresRelation::new(index_relation) }; - let opfamily = unsafe { opfamily(index_relation) }; let traverser = unsafe { HeapTraverser::new( @@ -1143,6 +1150,8 @@ unsafe fn sequential_build( sync_0(); + let index = unsafe { BufferedPostgresRelation::new(index_relation) }; + let mut indtuples = 0; match cached { VchordrqCachedReader::_0(_) => { @@ -1176,7 +1185,7 @@ unsafe fn sequential_build( VchordrqCachedReader::_1(cached) => { let index = CachingRelation { cache: cached, - relation: index.clone(), + relation: &index, }; traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { let ctid = tuple.id(); @@ -1207,8 +1216,16 @@ unsafe fn sequential_build( } } + while index.release() { + vchordrq::consume(&index); + } + + drop(index); + sync_1(indtuples); + let index = unsafe { PostgresRelation::new(index_relation) }; + let mut chooser = ChooseAll; crate::index::vchordrq::dispatch::maintain(opfamily, &index, &mut chooser, check); @@ -1564,16 +1581,7 @@ fn make_external_build( struct CachingRelation<'a, R> { cache: vchordrq_cached::VchordrqCachedReader1<'a>, - relation: R, -} - -impl Clone for CachingRelation<'_, R> { - fn clone(&self) -> Self { - Self { - cache: self.cache, - relation: self.relation.clone(), - } - } + relation: &'a R, } enum CachingRelationReadGuard<'a, G: Deref> { From 9fa094591a0fe7776d65b817237be6af4ac6f8f8 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 28 Oct 2025 18:48:36 +0800 Subject: [PATCH 239/324] feat: add a message of estimated memory usage (#367) #299 https://github.com/tensorchord/pgvecto.rs-docs/pull/151 Signed-off-by: usamoi --- Cargo.lock | 16 +++++ Cargo.toml | 1 + crates/k_means/src/flat.rs | 2 +- crates/k_means/src/quick.rs | 2 +- crates/k_means/src/rabitq.rs | 2 +- crates/k_means/src/square.rs | 6 ++ src/index/vchordrq/am/am_build.rs | 105 +++++++++++++++++------------- 7 files changed, 86 insertions(+), 48 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 56b057d9..3ab71796 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -556,6 +556,15 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "humansize" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb51c9a029ddc91b07a787f1d86b53ccfa49b0e86688c946ebe8d3555685dd7" +dependencies = [ + "libm", +] + [[package]] name = "icu_collections" version = "2.0.0" @@ -762,6 +771,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "libm" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" + [[package]] name = "libmimalloc-sys" version = "0.1.44" @@ -1621,6 +1636,7 @@ dependencies = [ "dary_heap", "distance", "feistel", + "humansize", "index", "k_means", "mimalloc", diff --git a/Cargo.toml b/Cargo.toml index 340a70df..eafed4e1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ vector = { path = "./crates/vector" } arrayvec = "0.7.6" bumpalo.workspace = true dary_heap.workspace = true +humansize = "2.1.3" paste.workspace = true pgrx = { version = "=0.16.1", default-features = false, features = ["cshim"] } pgrx-catalog = "0.3.1" diff --git a/crates/k_means/src/flat.rs b/crates/k_means/src/flat.rs index 810f1da1..68c5ac92 100644 --- a/crates/k_means/src/flat.rs +++ b/crates/k_means/src/flat.rs @@ -125,7 +125,7 @@ pub fn new( .expect("failed to build thread pool"); let mut rng = StdRng::from_seed(seed); - let mut centroids = Square::new(d); + let mut centroids = Square::with_capacity(d, c); for index in rand::seq::index::sample(&mut rng, samples.len(), c.min(samples.len())) { centroids.push_slice(&samples[index]); diff --git a/crates/k_means/src/quick.rs b/crates/k_means/src/quick.rs index 45a31e77..3a588915 100644 --- a/crates/k_means/src/quick.rs +++ b/crates/k_means/src/quick.rs @@ -49,7 +49,7 @@ pub fn new( .expect("failed to build thread pool"); let mut rng = StdRng::from_seed(seed); - let mut centroids = Square::new(d); + let mut centroids = Square::with_capacity(d, c); for index in rand::seq::index::sample(&mut rng, samples.len(), c.min(samples.len())) { centroids.push_slice(&samples[index]); diff --git a/crates/k_means/src/rabitq.rs b/crates/k_means/src/rabitq.rs index 5749222c..2940cee3 100644 --- a/crates/k_means/src/rabitq.rs +++ b/crates/k_means/src/rabitq.rs @@ -164,7 +164,7 @@ pub fn new( }); }); - let mut centroids = Square::new(d); + let mut centroids = Square::with_capacity(d, c); for index in rand::seq::index::sample(&mut rng, samples.len(), c.min(samples.len())) { centroids.push_slice(&samples[index]); diff --git a/crates/k_means/src/square.rs b/crates/k_means/src/square.rs index 40692261..0497c765 100644 --- a/crates/k_means/src/square.rs +++ b/crates/k_means/src/square.rs @@ -28,6 +28,12 @@ impl Square { pub fn is_empty(&self) -> bool { self.p.is_empty() } + pub fn with_capacity(d: usize, p: usize) -> Self { + Self { + d, + p: Vec::with_capacity(usize::saturating_mul(d, p)), + } + } pub fn new(d: usize) -> Self { Self { d, p: Vec::new() } } diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 75f27044..2adc8dc4 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -233,54 +233,14 @@ pub unsafe extern "C-unwind" fn ambuild( &raw mut pgrx::pg_sys::SnapshotAnyData }; let sampler = unsafe { HeapSampler::new(index_relation, heap_relation, snapshot) }; - let mut sample = sampler.sample(); - let samples = 'a: { - let mut samples = Square::new(vector_options.dims as _); - let Some(max_number_of_samples) = internal_build - .lists - .last() - .map(|x| x.saturating_mul(internal_build.sampling_factor)) - else { - break 'a samples; - }; - while samples.len() < max_number_of_samples as usize { - if let Some(mut tuple) = sample.next() { - let (values, is_nulls) = tuple.build(); - let datum = (!is_nulls[0]).then_some(values[0]); - if let Some(datum) = datum { - let vectors = unsafe { opfamily.store(datum) }; - if let Some(vectors) = vectors { - for (vector, _) in vectors { - let x = match vector { - OwnedVector::Vecf32(x) => VectOwned::normalize(x), - OwnedVector::Vecf16(x) => VectOwned::normalize(x), - }; - assert_eq!( - vector_options.dims, - x.len() as u32, - "invalid vector dimensions" - ); - samples.push_slice(x.as_slice()); - } - } else { - continue; - } - } else { - continue; - } - } else { - break; - } - } - samples.truncate(max_number_of_samples as usize); - samples - }; + let result = + make_internal_build(vector_options, opfamily, internal_build, sampler, &reporter); if is_mvcc_snapshot(snapshot) { unsafe { pgrx::pg_sys::UnregisterSnapshot(snapshot); } } - make_internal_build(vector_options, internal_build, samples, &reporter) + result } VchordrqBuildSourceOptions::External(external_build) => { reporter.phase(BuildPhase::from_code(BuildPhaseCode::ExternalBuild)); @@ -1307,16 +1267,71 @@ fn make_default_build( fn make_internal_build( vector_options: VectorOptions, + opfamily: Opfamily, internal_build: VchordrqInternalBuildOptions, - samples: Square, + sampler: impl Sampler, reporter: &PostgresReporter, ) -> Vec> { + use humansize::{BINARY, format_size}; use std::iter::once; + { + let d = vector_options.dims as u64; + let c = internal_build.lists.last().copied().unwrap_or_default() as u64; + let f = internal_build.sampling_factor as u64; + let t = internal_build.build_threads as u64; + let estimated_memory_usage = 4 * c * d * (1 + f + t); + pgrx::info!( + "clustering: estimated memory usage is {}", + format_size( + estimated_memory_usage, + BINARY.decimal_places(2).decimal_zeroes(2) + ) + ); + } + let mut sample = sampler.sample(); + let samples = 'a: { + let max_number_of_samples = internal_build + .lists + .last() + .map(|x| x.saturating_mul(internal_build.sampling_factor)) + .unwrap_or_default(); + let mut samples = + Square::with_capacity(vector_options.dims as _, max_number_of_samples as _); + if samples.len() >= max_number_of_samples as usize { + break 'a samples; + } + while let Some(mut tuple) = sample.next() { + let (values, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(values[0]); + if let Some(datum) = datum { + let vectors = unsafe { opfamily.store(datum) }; + if let Some(vectors) = vectors { + for (vector, _) in vectors { + let x = match vector { + OwnedVector::Vecf32(x) => VectOwned::normalize(x), + OwnedVector::Vecf16(x) => VectOwned::normalize(x), + }; + assert_eq!( + vector_options.dims, + x.len() as u32, + "invalid vector dimensions" + ); + samples.push_slice(x.as_slice()); + if samples.len() >= max_number_of_samples as usize { + break; + } + } + } + } + } + samples + }; let mut result = Vec::>::new(); let mut samples = Some(samples); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { let input = if let Some(structure) = result.last() { - let mut input = Square::new(vector_options.dims as _); + let mut input = + Square::with_capacity(vector_options.dims as _, structure.centroids.len()); for slice in structure.centroids.iter() { input.push_slice(slice); } From 0f5fd52092b07c318d1aeb6fe165ebfaa90b816d Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 29 Oct 2025 19:59:08 +0800 Subject: [PATCH 240/324] chore: cleanup simd for ppc64 and remove relaxed arithmetic (#368) job: +check_debian Signed-off-by: usamoi --- Cargo.lock | 136 +++++++++----------- Cargo.toml | 2 +- crates/simd/Cargo.toml | 5 +- crates/simd/src/fast_scan.rs | 213 +++++++++++++++++++++----------- crates/simd/src/floating_f16.rs | 48 +------ crates/simd/src/floating_f32.rs | 48 +------ crates/simd/src/lib.rs | 17 +-- crates/simd_macros/Cargo.toml | 6 +- crates/vchordrq/Cargo.toml | 2 +- crates/xtask/Cargo.toml | 2 +- 10 files changed, 220 insertions(+), 259 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3ab71796..180f19a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10,9 +10,9 @@ checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "aho-corasick" -version = "1.1.3" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" dependencies = [ "memchr", ] @@ -148,9 +148,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.41" +version = "1.2.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac9fe6cdbb24b6ade63616c0a0688e45bb56732262c158df3c0c4bea4ca47cb7" +checksum = "739eb0f94557554b3ca9a86d2d37bebd49c5e6d0c1d2bda35ba5bdac830befc2" dependencies = [ "find-msvc-tools", "shlex", @@ -456,9 +456,9 @@ checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" [[package]] name = "flate2" -version = "1.1.4" +version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc5a4e564e38c699f2880d3fda590bedc2e69f3f84cd48b457bd892ce61d0aa9" +checksum = "bfe33edd8e85a12a67454e37f8c75e730830d83e313556ab9ebf9ee7fbeb3bfb" dependencies = [ "crc32fast", "miniz_oxide", @@ -567,9 +567,9 @@ dependencies = [ [[package]] name = "icu_collections" -version = "2.0.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "200072f5d0e3614556f94a9930d5dc3e0662a652823904c3a75dc3b0af7fee47" +checksum = "4c6b649701667bbe825c3b7e6388cb521c23d88644678e83c0c4d0a621a34b43" dependencies = [ "displaydoc", "potential_utf", @@ -580,9 +580,9 @@ dependencies = [ [[package]] name = "icu_locale_core" -version = "2.0.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cde2700ccaed3872079a65fb1a78f6c0a36c91570f28755dda67bc8f7d9f00a" +checksum = "edba7861004dd3714265b4db54a3c390e880ab658fec5f7db895fae2046b5bb6" dependencies = [ "displaydoc", "litemap", @@ -593,11 +593,10 @@ dependencies = [ [[package]] name = "icu_normalizer" -version = "2.0.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "436880e8e18df4d7bbc06d58432329d6458cc84531f7ac5f024e93deadb37979" +checksum = "5f6c8828b67bf8908d82127b2054ea1b4427ff0230ee9141c54251934ab1b599" dependencies = [ - "displaydoc", "icu_collections", "icu_normalizer_data", "icu_properties", @@ -608,42 +607,38 @@ dependencies = [ [[package]] name = "icu_normalizer_data" -version = "2.0.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00210d6893afc98edb752b664b8890f0ef174c8adbb8d0be9710fa66fbbf72d3" +checksum = "7aedcccd01fc5fe81e6b489c15b247b8b0690feb23304303a9e560f37efc560a" [[package]] name = "icu_properties" -version = "2.0.1" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "016c619c1eeb94efb86809b015c58f479963de65bdb6253345c1a1276f22e32b" +checksum = "e93fcd3157766c0c8da2f8cff6ce651a31f0810eaa1c51ec363ef790bbb5fb99" dependencies = [ - "displaydoc", "icu_collections", "icu_locale_core", "icu_properties_data", "icu_provider", - "potential_utf", "zerotrie", "zerovec", ] [[package]] name = "icu_properties_data" -version = "2.0.1" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "298459143998310acd25ffe6810ed544932242d3f07083eee1084d83a71bd632" +checksum = "02845b3647bb045f1100ecd6480ff52f34c35f82d9880e029d329c21d1054899" [[package]] name = "icu_provider" -version = "2.0.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03c80da27b5f4187909049ee2d72f276f0d9f99a42c306bd0131ecfe04d8e5af" +checksum = "85962cf0ce02e1e0a629cc34e7ca3e373ce20dda4c4d7294bbd0bf1fdb59e614" dependencies = [ "displaydoc", "icu_locale_core", - "stable_deref_trait", - "tinystr", "writeable", "yoke", "zerofrom", @@ -737,9 +732,9 @@ checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" [[package]] name = "js-sys" -version = "0.3.81" +version = "0.3.82" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec48937a97411dcb524a265206ccd4c90bb711fca92b2792c407f268825b9305" +checksum = "b011eec8cc36da2aab2d5cff675ec18454fad408585853910a202391cf9f8e65" dependencies = [ "once_cell", "wasm-bindgen", @@ -806,15 +801,9 @@ checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039" [[package]] name = "litemap" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956" - -[[package]] -name = "log" -version = "0.4.28" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432" +checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" [[package]] name = "memchr" @@ -1070,9 +1059,9 @@ checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" [[package]] name = "potential_utf" -version = "0.1.3" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "84df19adbe5b5a0782edcab45899906947ab039ccf4573713735ee7de1e6b08a" +checksum = "b73949432f5e2a09657003c25bca5e19a0e9c84f8058ca374f49e0ebe605af77" dependencies = [ "zerovec", ] @@ -1110,9 +1099,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.101" +version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89ae43fd86e4158d6db51ad8e2b80f313af9cc74f5c0e03ccb87de09998732de" +checksum = "5ee95bc4ef87b8d5ba32e8b7714ccc834865276eab0aed5c9958d00ec45f49e8" dependencies = [ "unicode-ident", ] @@ -1430,9 +1419,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.107" +version = "2.0.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2a26dbd934e5451d21ef060c018dae56fc073894c5a7896f882928a76e6d081b" +checksum = "da58917d35242480a05c2897064da0a80589a2a0476c9a3f2fdc83b53502e917" dependencies = [ "proc-macro2", "quote", @@ -1458,9 +1447,9 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" [[package]] name = "target-triple" -version = "0.1.4" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ac9aa371f599d22256307c24a9d748c041e548cbf599f35d890f9d365361790" +checksum = "591ef38edfb78ca4771ee32cf494cb8771944bee237a9b91fc9c1424ac4b777b" [[package]] name = "thiserror" @@ -1484,9 +1473,9 @@ dependencies = [ [[package]] name = "tinystr" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d4f6d1145dcb577acf783d4e601bc1d76a13337bb54e6233add580b07344c8b" +checksum = "42d3e9c45c09de15d06dd8acf5f4e0e399e85927b7f00711024eb7ae10fa4869" dependencies = [ "displaydoc", "zerovec", @@ -1728,9 +1717,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.104" +version = "0.2.105" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1da10c01ae9f1ae40cbfac0bac3b1e724b320abfcf52229f80b547c0d250e2d" +checksum = "da95793dfc411fbbd93f5be7715b0578ec61fe87cb1a42b12eb625caa5c5ea60" dependencies = [ "cfg-if", "once_cell", @@ -1739,25 +1728,11 @@ dependencies = [ "wasm-bindgen-shared", ] -[[package]] -name = "wasm-bindgen-backend" -version = "0.2.104" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "671c9a5a66f49d8a47345ab942e2cb93c7d1d0339065d4f8139c486121b43b19" -dependencies = [ - "bumpalo", - "log", - "proc-macro2", - "quote", - "syn", - "wasm-bindgen-shared", -] - [[package]] name = "wasm-bindgen-macro" -version = "0.2.104" +version = "0.2.105" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ca60477e4c59f5f2986c50191cd972e3a50d8a95603bc9434501cf156a9a119" +checksum = "04264334509e04a7bf8690f2384ef5265f05143a4bff3889ab7a3269adab59c2" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1765,22 +1740,22 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.104" +version = "0.2.105" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f07d2f20d4da7b26400c9f4a0511e6e0345b040694e8a75bd41d578fa4421d7" +checksum = "420bc339d9f322e562942d52e115d57e950d12d88983a14c79b86859ee6c7ebc" dependencies = [ + "bumpalo", "proc-macro2", "quote", "syn", - "wasm-bindgen-backend", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.104" +version = "0.2.105" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bad67dc8b2a1a6e5448428adec4c3e84c43e561d8c9ee8a9e5aabeb193ec41d1" +checksum = "76f218a38c84bcb33c25ec7059b07847d465ce0e0a76b995e134a45adcb6af76" dependencies = [ "unicode-ident", ] @@ -1945,9 +1920,9 @@ checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" [[package]] name = "writeable" -version = "0.6.1" +version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" +checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" [[package]] name = "wyhash" @@ -1979,11 +1954,10 @@ dependencies = [ [[package]] name = "yoke" -version = "0.8.0" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f41bb01b8226ef4bfd589436a297c53d118f65921786300e427be8d487695cc" +checksum = "72d6e5c6afb84d73944e5cedb052c4680d5657337201555f9f2a16b7406d4954" dependencies = [ - "serde", "stable_deref_trait", "yoke-derive", "zerofrom", @@ -1991,9 +1965,9 @@ dependencies = [ [[package]] name = "yoke-derive" -version = "0.8.0" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38da3c9736e16c5d3c8c597a9aaa5d1fa565d0532ae05e27c24aa62fb32c0ab6" +checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" dependencies = [ "proc-macro2", "quote", @@ -2044,9 +2018,9 @@ dependencies = [ [[package]] name = "zerotrie" -version = "0.2.2" +version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36f0bbd478583f79edad978b407914f61b2972f5af6fa089686016be8f9af595" +checksum = "2a59c17a5562d507e4b54960e8569ebee33bee890c70aa3fe7b97e85a9fd7851" dependencies = [ "displaydoc", "yoke", @@ -2055,9 +2029,9 @@ dependencies = [ [[package]] name = "zerovec" -version = "0.11.4" +version = "0.11.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7aa2bd55086f1ab526693ecbe444205da57e25f4489879da80635a46d90e73b" +checksum = "6c28719294829477f525be0186d13efa9a3c602f7ec202ca9e353d310fb9a002" dependencies = [ "yoke", "zerofrom", @@ -2066,9 +2040,9 @@ dependencies = [ [[package]] name = "zerovec-derive" -version = "0.11.1" +version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b96237efa0c878c64bd89c436f661be4e46b2f3eff1ebb976f7ef2321d2f58f" +checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index eafed4e1..96a11943 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,7 +71,7 @@ paste = "1.0.15" rand = "0.9.2" rand_chacha = "0.9.0" seq-macro = "0.3.6" -serde = { version = "1.0", features = ["derive"] } +serde = { version = "1.0.228", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } wyhash = "0.6.0" zerocopy = { version = "0.8.27", features = ["derive"] } diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 0703c552..72baa9bb 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -6,8 +6,7 @@ publish = false [features] init = [] -experimental_f16 = ["zerocopy/float-nightly"] -experimental_math = [] +nightly_f16 = ["zerocopy/float-nightly"] [dependencies] simd_macros = { path = "../simd_macros" } @@ -20,7 +19,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.41" +cc = "1.2.43" which = "8.0.0" [lints] diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index 47a6c451..a0221584 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -505,90 +505,155 @@ mod scan { } } + /// The instructions required for byte order swapping differ across CPUs, so + /// different instructions needs to be generated for different CPUs for the + /// same code. #[cfg(target_arch = "powerpc64")] - #[crate::target_cpu(enable = "p7")] - fn scan_p7(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { - unsafe { - // bounds checking is not enforced by compiler, so check it manually - assert_eq!(code.len(), lut.len()); - let n = code.len(); - - use std::arch::powerpc64::*; - #[cfg(target_endian = "big")] - use std::intrinsics::simd::simd_bswap as vec_revb; - use std::mem::transmute; - use {vector_unsigned_char as u8x16, vector_unsigned_short as u16x8}; - - let _0008_u16x8 = vec_splat_u16::<0x0008>(); - let _00ff_u16x8 = vec_splats(0x00ffu16); - let _ff00_u16x8 = vec_splats(0xff00u16); - - let mut accu_0 = vec_splat_u16::<0>(); - let mut accu_1 = vec_splat_u16::<0>(); - let mut accu_2 = vec_splat_u16::<0>(); - let mut accu_3 = vec_splat_u16::<0>(); - - let mut i = 0_usize; - while i < n { - let code: u8x16 = vec_xl((i as isize) * 16, code.as_ptr().cast::()); - - let clo = vec_and(code, vec_splat_u8::<0xf>()); - let chi = vec_srl(code, vec_splat_u8::<4>()); - - let lut: u8x16 = vec_xl((i as isize) * 16, lut.as_ptr().cast::()); - #[cfg(target_endian = "big")] - { - let res_lo_r = transmute::(vec_perm(lut, lut, clo)); - let res_lo = vec_revb(res_lo_r); - accu_0 = vec_add(accu_0, res_lo); - accu_1 = vec_add(accu_1, vec_and(res_lo_r, _00ff_u16x8)); - let res_hi_r = transmute::(vec_perm(lut, lut, chi)); - let res_hi = vec_revb(res_hi_r); - accu_2 = vec_add(accu_2, res_hi); - accu_3 = vec_add(accu_3, vec_and(res_hi_r, _00ff_u16x8)); - } - #[cfg(target_endian = "little")] - { - let res_lo = transmute::(vec_perm(lut, lut, clo)); - accu_0 = vec_add(accu_0, res_lo); - accu_1 = vec_add(accu_1, vec_sr(res_lo, _0008_u16x8)); - let res_hi = transmute::(vec_perm(lut, lut, chi)); - accu_2 = vec_add(accu_2, res_hi); - accu_3 = vec_add(accu_3, vec_sr(res_hi, _0008_u16x8)); + macro_rules! scan_powerpc64 { + ($name:ident, $cpu:literal) => { + #[crate::target_cpu(enable = $cpu)] + fn $name(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + unsafe { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + use std::arch::powerpc64::*; + use std::mem::transmute; + use {vector_unsigned_char as u8x16, vector_unsigned_short as u16x8}; + + #[cfg(target_endian = "big")] + let revb = transmute::<[u8; 16], u8x16>([ + 1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10, 13, 12, 15, 14, + ]); + + let _0008_u16x8 = vec_splat_u16::<0x0008>(); + let _00ff_u16x8 = vec_splats(0x00ffu16); + let _ff00_u16x8 = vec_splats(0xff00u16); + + let mut accu_0 = vec_splat_u16::<0>(); + let mut accu_1 = vec_splat_u16::<0>(); + let mut accu_2 = vec_splat_u16::<0>(); + let mut accu_3 = vec_splat_u16::<0>(); + + let mut i = 0_usize; + while i < n { + let code: u8x16 = vec_xl((i as isize) * 16, code.as_ptr().cast::()); + + let clo = vec_and(code, vec_splat_u8::<0xf>()); + let chi = vec_srl(code, vec_splat_u8::<4>()); + + let lut: u8x16 = vec_xl((i as isize) * 16, lut.as_ptr().cast::()); + #[cfg(target_endian = "big")] + { + let res_lo_r = transmute::(vec_perm(lut, lut, clo)); + let res_lo = vec_perm(res_lo_r, res_lo_r, revb); + accu_0 = vec_add(accu_0, res_lo); + accu_1 = vec_add(accu_1, vec_and(res_lo_r, _00ff_u16x8)); + let res_hi_r = transmute::(vec_perm(lut, lut, chi)); + let res_hi = vec_perm(res_hi_r, res_hi_r, revb); + accu_2 = vec_add(accu_2, res_hi); + accu_3 = vec_add(accu_3, vec_and(res_hi_r, _00ff_u16x8)); + } + #[cfg(target_endian = "little")] + { + let res_lo = transmute::(vec_perm(lut, lut, clo)); + accu_0 = vec_add(accu_0, res_lo); + accu_1 = vec_add(accu_1, vec_sr(res_lo, _0008_u16x8)); + let res_hi = transmute::(vec_perm(lut, lut, chi)); + accu_2 = vec_add(accu_2, res_hi); + accu_3 = vec_add(accu_3, vec_sr(res_hi, _0008_u16x8)); + } + + i += 1; + } + debug_assert_eq!(i, n); + + let mut result = [0_u16; 32]; + + #[cfg(target_endian = "big")] + { + accu_0 = + vec_sub(accu_0, vec_and(vec_perm(accu_1, accu_1, revb), _ff00_u16x8)); + } + #[cfg(target_endian = "little")] + { + accu_0 = vec_sub(accu_0, vec_sl(accu_1, _0008_u16x8)); + } + vec_xst(accu_0, 0, result.as_mut_ptr().cast()); + vec_xst(accu_1, 16, result.as_mut_ptr().cast()); + + #[cfg(target_endian = "big")] + { + accu_2 = + vec_sub(accu_2, vec_and(vec_perm(accu_3, accu_3, revb), _ff00_u16x8)); + } + #[cfg(target_endian = "little")] + { + accu_2 = vec_sub(accu_2, vec_sl(accu_3, _0008_u16x8)); + } + vec_xst(accu_2, 32, result.as_mut_ptr().cast()); + vec_xst(accu_3, 48, result.as_mut_ptr().cast()); + + result } - - i += 1; } - debug_assert_eq!(i, n); + }; + } - let mut result = [0_u16; 32]; + #[cfg(target_arch = "powerpc64")] + scan_powerpc64!(scan_p9, "p9"); - #[cfg(target_endian = "big")] - { - accu_0 = vec_sub(accu_0, vec_and(vec_revb(accu_1), _ff00_u16x8)); - } - #[cfg(target_endian = "little")] - { - accu_0 = vec_sub(accu_0, vec_sl(accu_1, _0008_u16x8)); + #[cfg(all(target_arch = "powerpc64", test, not(miri)))] + #[test] + fn scan_p9_test() { + if !crate::is_cpu_detected!("p9") { + println!("test {} ... skipped (p9)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + for n in 90..110 { + let code = (0..n) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..n) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + unsafe { + assert_eq!(scan_p9(&code, &lut), fallback(&code, &lut)); + } } - vec_xst(accu_0, 0, result.as_mut_ptr().cast()); - vec_xst(accu_1, 16, result.as_mut_ptr().cast()); + } + } - #[cfg(target_endian = "big")] - { - accu_2 = vec_sub(accu_2, vec_and(vec_revb(accu_3), _ff00_u16x8)); - } - #[cfg(target_endian = "little")] - { - accu_2 = vec_sub(accu_2, vec_sl(accu_3, _0008_u16x8)); - } - vec_xst(accu_2, 32, result.as_mut_ptr().cast()); - vec_xst(accu_3, 48, result.as_mut_ptr().cast()); + #[cfg(target_arch = "powerpc64")] + scan_powerpc64!(scan_p8, "p8"); - result + #[cfg(all(target_arch = "powerpc64", test, not(miri)))] + #[test] + fn scan_p8_test() { + if !crate::is_cpu_detected!("p8") { + println!("test {} ... skipped (p8)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + for n in 90..110 { + let code = (0..n) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..n) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + unsafe { + assert_eq!(scan_p8(&code, &lut), fallback(&code, &lut)); + } + } } } + #[cfg(target_arch = "powerpc64")] + scan_powerpc64!(scan_p7, "p7"); + #[cfg(all(target_arch = "powerpc64", test, not(miri)))] #[test] fn scan_p7_test() { @@ -611,7 +676,7 @@ mod scan { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", @"z13", @"p7")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", @"z13", @"p9", @"p8", @"p7")] pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { assert_eq!(code.len(), lut.len()); let n = code.len(); diff --git a/crates/simd/src/floating_f16.rs b/crates/simd/src/floating_f16.rs index 0316394e..136d6ca3 100644 --- a/crates/simd/src/floating_f16.rs +++ b/crates/simd/src/floating_f16.rs @@ -174,14 +174,7 @@ mod reduce_sum_of_x { let n = this.len(); let mut x = 0.0f32; for i in 0..n { - #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] - { - x += this[i]._to_f32(); - } - #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] - { - x = x.algebraic_add(this[i]._to_f32()); - } + x += this[i]._to_f32(); } x } @@ -199,14 +192,7 @@ mod reduce_sum_of_abs_x { let n = this.len(); let mut x = 0.0f32; for i in 0..n { - #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] - { - x += this[i]._to_f32().abs(); - } - #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] - { - x = x.algebraic_add(this[i]._to_f32().abs()); - } + x += this[i]._to_f32().abs(); } x } @@ -224,14 +210,7 @@ mod reduce_sum_of_x2 { let n = this.len(); let mut x2 = 0.0f32; for i in 0..n { - #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] - { - x2 += this[i]._to_f32() * this[i]._to_f32(); - } - #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] - { - x2 = x2.algebraic_add(this[i]._to_f32().algebraic_mul(this[i]._to_f32())); - } + x2 += this[i]._to_f32() * this[i]._to_f32(); } x2 } @@ -526,14 +505,7 @@ mod reduce_sum_of_xy { let n = lhs.len(); let mut xy = 0.0f32; for i in 0..n { - #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] - { - xy += lhs[i]._to_f32() * rhs[i]._to_f32(); - } - #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] - { - xy = xy.algebraic_add(lhs[i]._to_f32().algebraic_mul(rhs[i]._to_f32())); - } + xy += lhs[i]._to_f32() * rhs[i]._to_f32(); } xy } @@ -816,16 +788,8 @@ mod reduce_sum_of_d2 { let n = lhs.len(); let mut d2 = 0.0_f32; for i in 0..n { - #[cfg(not(all(feature = "experimental_math", feature = "experimental_f16")))] - { - let d = lhs[i]._to_f32() - rhs[i]._to_f32(); - d2 += d * d; - } - #[cfg(all(feature = "experimental_math", feature = "experimental_f16"))] - { - let d = lhs[i]._to_f32().algebraic_sub(rhs[i]._to_f32()); - d2 = d2.algebraic_add(d.algebraic_mul(d)); - } + let d = lhs[i]._to_f32() - rhs[i]._to_f32(); + d2 += d * d; } d2 } diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs index 4fe48642..5c291a57 100644 --- a/crates/simd/src/floating_f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -418,14 +418,7 @@ mod reduce_sum_of_x { let n = this.len(); let mut sum = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental_math"))] - { - sum += this[i]; - } - #[cfg(feature = "experimental_math")] - { - sum = sum.algebraic_add(this[i]); - } + sum += this[i]; } sum } @@ -700,14 +693,7 @@ mod reduce_sum_of_abs_x { let n = this.len(); let mut sum = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental_math"))] - { - sum += this[i].abs(); - } - #[cfg(feature = "experimental_math")] - { - sum = sum.algebraic_add(this[i].abs()); - } + sum += this[i].abs(); } sum } @@ -972,14 +958,7 @@ mod reduce_sum_of_x2 { let n = this.len(); let mut x2 = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental_math"))] - { - x2 += this[i] * this[i]; - } - #[cfg(feature = "experimental_math")] - { - x2 = x2.algebraic_add(this[i].algebraic_mul(this[i])); - } + x2 += this[i] * this[i]; } x2 } @@ -1568,14 +1547,7 @@ mod reduce_sum_of_xy { let n = lhs.len(); let mut xy = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental_math"))] - { - xy += lhs[i] * rhs[i]; - } - #[cfg(feature = "experimental_math")] - { - xy = xy.algebraic_add(lhs[i].algebraic_mul(rhs[i])); - } + xy += lhs[i] * rhs[i]; } xy } @@ -1900,16 +1872,8 @@ mod reduce_sum_of_d2 { let n = lhs.len(); let mut d2 = 0.0f32; for i in 0..n { - #[cfg(not(feature = "experimental_math"))] - { - let d = lhs[i] - rhs[i]; - d2 += d * d; - } - #[cfg(feature = "experimental_math")] - { - let d = lhs[i].algebraic_sub(rhs[i]); - d2 = d2.algebraic_add(d.algebraic_mul(d)); - } + let d = lhs[i] - rhs[i]; + d2 += d * d; } d2 } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 90e8aeb3..6798a9b1 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -12,19 +12,14 @@ // // Copyright (c) 2025 TensorChord Inc. -#![allow(unsafe_code, internal_features)] -#![cfg_attr(feature = "experimental_f16", feature(f16))] -#![cfg_attr(feature = "experimental_math", feature(float_algebraic))] +#![allow(unsafe_code)] +#![cfg_attr(feature = "nightly_f16", feature(f16))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x_feature_detection))] #![cfg_attr(target_arch = "s390x", feature(s390x_target_feature))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] #![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc_feature_detection))] #![cfg_attr(target_arch = "powerpc64", feature(powerpc_target_feature))] #![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc))] -#![cfg_attr( - all(target_arch = "powerpc64", target_endian = "big"), - feature(core_intrinsics) -)] #![cfg_attr(target_arch = "riscv64", feature(stdarch_riscv_feature_detection))] #![cfg_attr(target_arch = "riscv64", feature(riscv_target_feature))] @@ -40,10 +35,10 @@ pub mod quantize; pub mod rotate; pub mod u8; -#[cfg(not(feature = "experimental_f16"))] +#[cfg(not(feature = "nightly_f16"))] pub use half::f16; -#[cfg(feature = "experimental_f16")] +#[cfg(feature = "nightly_f16")] pub use f16; pub trait F16: Sized { @@ -54,7 +49,7 @@ pub trait F16: Sized { fn _to_f32(self) -> f32; } -#[cfg(not(feature = "experimental_f16"))] +#[cfg(not(feature = "nightly_f16"))] impl F16 for f16 { const _ZERO: Self = f16::ZERO; @@ -67,7 +62,7 @@ impl F16 for f16 { } } -#[cfg(feature = "experimental_f16")] +#[cfg(feature = "nightly_f16")] impl F16 for f16 { const _ZERO: Self = 0.0; diff --git a/crates/simd_macros/Cargo.toml b/crates/simd_macros/Cargo.toml index 6924ac1c..f206bff0 100644 --- a/crates/simd_macros/Cargo.toml +++ b/crates/simd_macros/Cargo.toml @@ -8,9 +8,9 @@ publish = false proc-macro = true [dependencies] -proc-macro2 = { version = "1.0", features = ["proc-macro"] } -quote = "1.0" -syn = { version = "2.0", default-features = false, features = [ +proc-macro2 = { version = "1", features = ["proc-macro"] } +quote = "1" +syn = { version = "2", default-features = false, features = [ "clone-impls", "full", "parsing", diff --git a/crates/vchordrq/Cargo.toml b/crates/vchordrq/Cargo.toml index acd9d932..89c4c360 100644 --- a/crates/vchordrq/Cargo.toml +++ b/crates/vchordrq/Cargo.toml @@ -12,7 +12,7 @@ rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } -pin-project = "1.1" +pin-project = "1.1.10" serde.workspace = true validator.workspace = true zerocopy.workspace = true diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index 039ba184..3fcfff23 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -8,7 +8,7 @@ publish = false clap = { version = "4.5.50", features = ["derive", "env"] } object = { version = "0.37.3", features = ["read", "wasm"] } shlex = "1.3.0" -target-triple = "0.1.4" +target-triple = "1.0.0" [lints] workspace = true From 56ecc24d868d373a96d3369d81905968dd1502f1 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 31 Oct 2025 13:56:14 +0800 Subject: [PATCH 241/324] chore: fix for cargo target-directory (#369) Signed-off-by: usamoi --- Cargo.lock | 2 + Makefile | 6 +- crates/vchordg/src/vectors.rs | 11 +--- crates/xtask/Cargo.toml | 2 + crates/xtask/src/main.rs | 108 +++++++++++++++++++++++++--------- 5 files changed, 87 insertions(+), 42 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 180f19a4..ac7f1527 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1948,6 +1948,8 @@ version = "0.0.0" dependencies = [ "clap", "object", + "serde", + "serde_json", "shlex", "target-triple", ] diff --git a/Makefile b/Makefile index fac35b23..d453aba9 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -PG_CONFIG ?= pg_config +export PG_CONFIG ?= pg_config PKGLIBDIR := $(shell $(PG_CONFIG) --pkglibdir) SHAREDIR := $(shell $(PG_CONFIG) --sharedir) CP_R ?= cp -r @@ -9,10 +9,10 @@ MKDIR_P ?= mkdir -p all: build build: - PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p xtask -- build + cargo run -p xtask -- build clippy: - PGRX_PG_CONFIG_PATH="$(PG_CONFIG)" cargo run -p xtask -- clippy + cargo run -p xtask -- clippy install: $(MKDIR_P) $(DESTDIR)$(PKGLIBDIR) && \ diff --git a/crates/vchordg/src/vectors.rs b/crates/vchordg/src/vectors.rs index ab37b2af..95b83e25 100644 --- a/crates/vchordg/src/vectors.rs +++ b/crates/vchordg/src/vectors.rs @@ -68,7 +68,7 @@ pub fn read< accessor: A, copy: impl FnOnce(&[OptionNeighbour]) -> Output, ) -> Result<(A::Output, Output, Option>, Wrapping), ()> { - let m = strict_sub(iterator.len(), 1); + let m = iterator.len().strict_sub(1); let mut result = accessor; for index in 0..m { let (vector_guard, i) = iterator.next().expect("internal: bad size"); @@ -168,12 +168,3 @@ pub fn update( } Ok(true) } - -// Emulate unstable library feature `strict_overflow_ops`. -// See https://github.com/rust-lang/rust/issues/118260. - -#[inline] -const fn strict_sub(lhs: usize, rhs: usize) -> usize { - let (a, b) = lhs.overflowing_sub(rhs); - if b { panic!() } else { a } -} diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index 3fcfff23..d334d24e 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -7,6 +7,8 @@ publish = false [dependencies] clap = { version = "4.5.50", features = ["derive", "env"] } object = { version = "0.37.3", features = ["read", "wasm"] } +serde.workspace = true +serde_json = "1.0.145" shlex = "1.3.0" target-triple = "1.0.0" diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index ddcccd03..cfbf936f 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -14,6 +14,7 @@ use clap::{Args, Parser, Subcommand}; use object::{Object, ObjectSymbol}; +use serde::Deserialize; use std::collections::{HashMap, HashSet}; use std::env::var_os; use std::error::Error; @@ -53,7 +54,7 @@ struct ClippyArgs { profile: String, } -struct TargetSpecificInformation { +struct RustcCfg { is_macos: bool, is_windows: bool, is_emscripten: bool, @@ -61,7 +62,7 @@ struct TargetSpecificInformation { is_powerpc64: bool, } -impl TargetSpecificInformation { +impl RustcCfg { fn dll_prefix(&self) -> Result<&'static str, Box> { if self.is_macos { Ok("lib") @@ -116,11 +117,20 @@ impl TargetSpecificInformation { } } +#[derive(Deserialize)] +struct CargoMetadata { + target_directory: String, +} + fn pg_config(pg_config: impl AsRef) -> Result, Box> { let mut command = Command::new(pg_config.as_ref()); command.stderr(Stdio::inherit()); eprintln!("Running {command:?}"); let command_output = command.output()?; + let command_status = command_output.status; + if !command_status.success() { + return Err(format!("PostgreSQL failed: {command_status}").into()); + } let contents = String::from_utf8(command_output.stdout)?; let mut result = HashMap::new(); for line in contents.lines() { @@ -146,7 +156,7 @@ fn control_file(path: impl AsRef) -> Result, Box Result> { +fn rustc_cfg(target: &str) -> Result> { let mut command = Command::new("rustc"); command .args(["--print", "cfg"]) @@ -154,12 +164,16 @@ fn target_specific_information(target: &str) -> Result Result Result> { + let mut command = Command::new("cargo"); + command + .args(["metadata", "--format-version", "1"]) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_output = command.output()?; + let command_status = command_output.status; + if !command_status.success() { + return Err(format!("Cargo failed: {command_status}").into()); + } + let contents = String::from_utf8(command_output.stdout)?; + let cargo_metadata: CargoMetadata = serde_json::from_str(&contents)?; + Ok(cargo_metadata) +} + fn build( pg_config: impl AsRef, pg_version: &str, - tsi: &TargetSpecificInformation, + rustc_cfg: &RustcCfg, + cargo_metadata: &CargoMetadata, profile: &str, target: &str, ) -> Result> { @@ -188,29 +219,30 @@ fn build( eprintln!("Running {command:?}"); let command_status = command.spawn()?.wait()?; if !command_status.success() { - return Err(format!("Cargo build failed: {command_status}").into()); + return Err(format!("Cargo failed: {command_status}").into()); } - let mut result = PathBuf::from("./target"); + let mut result = PathBuf::from(&cargo_metadata.target_directory); result.push(target); result.push(match profile { "dev" | "test" => "debug", "release" | "bench" => "release", profile => profile, }); - result.push(format!("{}vchord{}", tsi.dll_prefix()?, tsi.dll_suffix()?)); + result.push(format!( + "{}vchord{}", + rustc_cfg.dll_prefix()?, + rustc_cfg.dll_suffix()? + )); Ok(result) } -fn parse( - tsi: &TargetSpecificInformation, - obj: impl AsRef, -) -> Result, Box> { +fn parse(rustc_cfg: &RustcCfg, obj: impl AsRef) -> Result, Box> { let obj = obj.as_ref(); eprintln!("Reading {obj:?}"); let contents = std::fs::read(obj)?; let object = object::File::parse(contents.as_slice())?; let exports; - if tsi.is_macos { + if rustc_cfg.is_macos { exports = object .exports()? .into_iter() @@ -219,7 +251,7 @@ fn parse( .filter(|x| x.starts_with("__pgrx_internals")) .map(str::to_string) .collect(); - } else if tsi.is_emscripten { + } else if rustc_cfg.is_emscripten { exports = object .symbols() .flat_map(|x| x.name().ok()) @@ -242,13 +274,14 @@ fn generate( runner: &Option>, pg_config: impl AsRef, pg_version: &str, - tsi: &TargetSpecificInformation, + rustc_cfg: &RustcCfg, + cargo_metadata: &CargoMetadata, profile: &str, target: &str, exports: Vec, postmaster: impl AsRef, ) -> Result> { - let imports = if tsi.is_powerpc64 { + let imports = if rustc_cfg.is_powerpc64 { let postmaster = postmaster.as_ref(); eprintln!("Reading {postmaster:?}"); let contents = std::fs::read(postmaster)?; @@ -288,16 +321,16 @@ fn generate( eprintln!("Running {command:?}"); let command_status = command.spawn()?.wait()?; if !command_status.success() { - return Err(format!("Cargo build failed: {command_status}").into()); + return Err(format!("Cargo failed: {command_status}").into()); } - let mut result = PathBuf::from("./target"); + let mut result = PathBuf::from(&cargo_metadata.target_directory); result.push(target); result.push(match profile { "dev" | "test" => "debug", "release" | "bench" => "release", profile => profile, }); - result.push(format!("pgrx_embed_vchord{}", tsi.exe_suffix()?)); + result.push(format!("pgrx_embed_vchord{}", rustc_cfg.exe_suffix()?)); let mut command; if let Some(runner) = runner { command = Command::new(&runner[0]); @@ -311,6 +344,10 @@ fn generate( command.stderr(Stdio::inherit()); eprintln!("Running {command:?}"); let command_output = command.output()?; + let command_status = command_output.status; + if !command_status.success() { + return Err(format!("Cargo failed: {command_status}").into()); + } let command_stdout = String::from_utf8(command_output.stdout)?.replace("\t", " "); Ok(command_stdout) } @@ -368,12 +405,16 @@ fn clippy( eprintln!("Running {command:?}"); let command_status = command.spawn()?.wait()?; if !command_status.success() { - return Err(format!("Cargo clippy failed: {command_status}").into()); + return Err(format!("Cargo failed: {command_status}").into()); } Ok(()) } fn main() -> Result<(), Box> { + #[allow(unsafe_code)] + unsafe { + std::env::remove_var("PGRX_PG_CONFIG_PATH"); + } let cli = Cli::parse(); match cli.command { Commands::Build(BuildArgs { @@ -387,10 +428,10 @@ fn main() -> Result<(), Box> { } let vchord_version = control_file("./vchord.control")?["default_version"].clone(); let runner = runner.and_then(|runner| shlex::split(&runner)); - let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { + let path = if let Some(value) = var_os("PG_CONFIG") { PathBuf::from(value) } else { - return Err("Environment variable `PGRX_PG_CONFIG_PATH` is not set.".into()); + return Err("Environment variable `PG_CONFIG` is not set.".into()); }; let pg_config = pg_config(&path)?; let pg_version = { @@ -408,8 +449,16 @@ fn main() -> Result<(), Box> { } }; let postmaster = format!("{}/postgres", pg_config["BINDIR"]); - let tsi = target_specific_information(&target)?; - let obj = build(&path, &pg_version, &tsi, &profile, &target)?; + let rustc_cfg = rustc_cfg(&target)?; + let cargo_metadata = cargo_metadata()?; + let obj = build( + &path, + &pg_version, + &rustc_cfg, + &cargo_metadata, + &profile, + &target, + )?; let pkglibdir = format!("{output}/pkglibdir"); let sharedir = format!("{output}/sharedir"); let sharedir_extension = format!("{sharedir}/extension"); @@ -422,7 +471,7 @@ fn main() -> Result<(), Box> { std::fs::create_dir_all(&sharedir_extension)?; install_by_copying( &obj, - format!("{pkglibdir}/vchord{}", tsi.ext_suffix(&pg_version)?), + format!("{pkglibdir}/vchord{}", rustc_cfg.ext_suffix(&pg_version)?), true, )?; install_by_copying( @@ -444,13 +493,14 @@ fn main() -> Result<(), Box> { false, )?; } else { - let exports = parse(&tsi, obj)?; + let exports = parse(&rustc_cfg, obj)?; install_by_writing( generate( &runner, &path, &pg_version, - &tsi, + &rustc_cfg, + &cargo_metadata, &profile, &target, exports, @@ -465,10 +515,10 @@ fn main() -> Result<(), Box> { if !std::fs::exists("./vchord.control")? { return Err("The script must be run from the VectorChord source directory.".into()); } - let path = if let Some(value) = var_os("PGRX_PG_CONFIG_PATH") { + let path = if let Some(value) = var_os("PG_CONFIG") { PathBuf::from(value) } else { - return Err("Environment variable `PGRX_PG_CONFIG_PATH` is not set.".into()); + return Err("Environment variable `PG_CONFIG` is not set.".into()); }; let pg_config = pg_config(&path)?; let pg_version = { From 6375def1e3c85d642970bd869c7f2a7b69030a67 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Fri, 31 Oct 2025 18:23:30 +0800 Subject: [PATCH 242/324] feat: approximate kmeans and dimension reduction (#358) This PR provides a much more faster kmeans algorithm at the cost of less accuracy. Signed-off-by: cutecutecat --- Cargo.lock | 2 + crates/k_means/Cargo.toml | 2 + crates/k_means/src/flat.rs | 54 +++-- crates/k_means/src/hierarchical.rs | 307 +++++++++++++++++++++++++++++ crates/k_means/src/lib.rs | 23 ++- crates/k_means/src/quick.rs | 15 +- crates/k_means/src/rabitq.rs | 36 ++-- crates/k_means/src/square.rs | 87 ++++++++ src/index/vchordrq/am/am_build.rs | 138 +++++++++++-- src/index/vchordrq/types.rs | 21 ++ tests/vchordrq/approximate.slt | 57 ++++++ 11 files changed, 680 insertions(+), 62 deletions(-) create mode 100644 crates/k_means/src/hierarchical.rs create mode 100644 tests/vchordrq/approximate.slt diff --git a/Cargo.lock b/Cargo.lock index ac7f1527..4d55aea1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -744,6 +744,8 @@ dependencies = [ name = "k_means" version = "0.0.0" dependencies = [ + "always_equal", + "distance", "rabitq", "rand", "rayon", diff --git a/crates/k_means/Cargo.toml b/crates/k_means/Cargo.toml index 6475cd6e..9361a3f9 100644 --- a/crates/k_means/Cargo.toml +++ b/crates/k_means/Cargo.toml @@ -7,6 +7,8 @@ publish = false [dependencies] rabitq = { path = "../rabitq" } simd = { path = "../simd" } +always_equal = { path = "../always_equal" } +distance = { path = "../distance" } rand.workspace = true rayon = "1.11.0" diff --git a/crates/k_means/src/flat.rs b/crates/k_means/src/flat.rs index 68c5ac92..270b7f5d 100644 --- a/crates/k_means/src/flat.rs +++ b/crates/k_means/src/flat.rs @@ -12,57 +12,51 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::square::Square; +use crate::square::{Square, SquareMut}; use crate::{KMeans, This}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use rayon::prelude::*; use simd::Floating; -struct Flat { +struct Flat<'a> { this: This, + samples: SquareMut<'a>, } -impl KMeans for Flat { +impl<'a> KMeans for Flat<'a> { fn this(&mut self) -> &mut This { &mut self.this } - fn assign(&mut self) { let this = &mut self.this; this.pool.install(|| { this.targets .par_iter_mut() - .zip(this.samples.into_par_iter()) + .zip(self.samples.par_iter_mut()) .for_each(|(target, sample)| { - let mut result = (f32::INFINITY, 0); - for (j, centroid) in this.centroids.into_iter().enumerate() { - let dis_2 = f32::reduce_sum_of_d2(sample, centroid); - if dis_2 <= result.0 { - result = (dis_2, j); - } - } - *target = result.1; + *target = k_means_lookup(sample, &this.centroids); }); }); } fn update(&mut self) { let this = &mut self.this; + let samples = &mut self.samples; this.pool.install(|| { const DELTA: f32 = 9.7656e-4_f32; let d = this.d; - let n = this.samples.len(); + let n = samples.len(); let c = this.c; let list = rayon::broadcast({ |ctx| { let mut sum = Square::from_zeros(d, c); let mut count = vec![0.0f32; c]; - for i in (ctx.index()..this.samples.len()).step_by(ctx.num_threads()) { + for i in (ctx.index()..samples.len()).step_by(ctx.num_threads()) { let target = this.targets[i]; - let sample = &this.samples[i]; + let sample = &samples[i]; f32::vector_add_inplace(&mut sum[target], sample); count[target] += 1.0; } @@ -106,19 +100,25 @@ impl KMeans for Flat { }); } + fn index(&mut self) -> Box u32> { + let this = self.this(); + let centroids = this.centroids.clone(); + Box::new(move |sample| k_means_lookup(sample, ¢roids) as u32) + } + fn finish(self: Box) -> Square { let this = self.this; this.centroids } } -pub fn new( +pub fn new<'a>( d: usize, - samples: Square, + samples: SquareMut<'a>, c: usize, num_threads: usize, seed: [u8; 32], -) -> Box { +) -> Box { let pool = rayon::ThreadPoolBuilder::new() .num_threads(num_threads) .build() @@ -144,16 +144,28 @@ pub fn new( Box::new(Flat { this: This { pool, + rng, d, c, centroids, targets, - rng, - samples, }, + samples, }) } +pub fn k_means_lookup(vector: &[f32], centroids: &Square) -> usize { + assert_ne!(centroids.len(), 0); + let mut result = (f32::INFINITY, 0); + for i in 0..centroids.len() { + let dis = f32::reduce_sum_of_d2(vector, ¢roids[i]); + if dis <= result.0 { + result = (dis, i); + } + } + result.1 +} + fn vector_mul_scalars_inplace(this: &mut [f32], scalars: [f32; 2]) { let n: usize = this.len(); for i in 0..n { diff --git a/crates/k_means/src/hierarchical.rs b/crates/k_means/src/hierarchical.rs new file mode 100644 index 00000000..ffd4a36c --- /dev/null +++ b/crates/k_means/src/hierarchical.rs @@ -0,0 +1,307 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::square::{Square, SquareMut}; +use crate::{KMeans, This}; +use always_equal::AlwaysEqual; +use distance::Distance; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator}; +use std::collections::BinaryHeap; + +struct Hierarchical<'a> { + this: This, + top_centroids: Square, + partition_start: Vec, + bottom_k_means: Vec>, +} + +impl<'a> KMeans for Hierarchical<'a> { + fn this(&mut self) -> &mut This { + &mut self.this + } + + fn assign(&mut self) { + for k_means in &mut self.bottom_k_means { + k_means.assign(); + } + } + + fn update(&mut self) { + for k_means in &mut self.bottom_k_means { + k_means.update(); + } + } + + fn sphericalize(&mut self) { + for k_means in &mut self.bottom_k_means { + k_means.sphericalize(); + } + } + + fn index(&mut self) -> Box u32 + '_> { + let top_centroids = self.top_centroids.clone(); + let top = move |sample: &[f32]| crate::flat::k_means_lookup(sample, &top_centroids) as u32; + let bottom_centroids: Vec<_> = self + .bottom_k_means + .iter_mut() + .map(|k_means| k_means.index()) + .collect(); + let partition_start = self.partition_start.clone(); + let index = move |sample: &[f32]| { + let top_id = top(sample) as usize; + let bottom_id = bottom_centroids[top_id](sample) as usize; + (partition_start[top_id] + bottom_id) as u32 + }; + Box::new(index) + } + + fn finish(self: Box) -> Square { + let mut ret = Square::new(self.this.d); + for k_means in self.bottom_k_means { + let centroids = k_means.finish(); + for centroid in centroids.into_iter() { + ret.push_slice(centroid); + } + } + ret + } +} + +const LOCAL_SAMPLE_FACTOR: usize = 256; +const LOCAL_NUM_ITERATIONS: usize = 10; + +pub fn new<'a>( + d: usize, + mut samples: SquareMut<'a>, + c: usize, + num_threads: usize, + seed: [u8; 32], + is_spherical: bool, +) -> Box { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .build() + .expect("failed to build thread pool"); + let centroids = Square::new(d); + let targets = vec![0; samples.len()]; + let samples_len = samples.len(); + let top_list = + ((c as f64).sqrt().floor() as u32).clamp(1, (samples_len as f64).sqrt().floor() as u32); + let top_samples_len = LOCAL_SAMPLE_FACTOR * (top_list as usize); + let mut top_samples = Square::new(d); + let mut rng = StdRng::from_seed(seed); + for index in + rand::seq::index::sample(&mut rng, samples.len(), top_samples_len.min(samples.len())) + { + top_samples.push_slice(&samples[index]); + } + let mut f = crate::k_means( + d, + top_samples.as_mut_view(), + top_list as usize, + num_threads, + seed, + ); + if is_spherical { + f.sphericalize(); + } + for _ in 0..LOCAL_NUM_ITERATIONS { + f.assign(); + f.update(); + if is_spherical { + f.sphericalize(); + } + } + let top_centroids = f.finish(); + let mut final_assign = vec![0; samples.len()]; + pool.install(|| { + final_assign + .par_iter_mut() + .zip(samples.par_iter_mut()) + .for_each(|(target, sample)| { + *target = crate::flat::k_means_lookup(sample, &top_centroids); + }); + }); + let alloc = final_assign.into_iter().enumerate().fold( + vec![vec![]; top_centroids.len()], + |mut acc, (i, target)| { + acc[target].push(i); + acc + }, + ); + let alloc_size = alloc.iter().map(|x| x.len() as u32).collect::>(); + let keep_indices: Vec = alloc_size + .iter() + .enumerate() + .filter_map(|(i, size)| if *size > 0 { Some(i) } else { None }) + .collect(); + let alloc: Vec<_> = keep_indices.iter().map(|&i| alloc[i].clone()).collect(); + let alloc_size: Vec<_> = keep_indices.iter().map(|&i| alloc_size[i]).collect(); + let alloc_lists = successive_quotients_allocate(c, alloc_size); + + let mut bottom_k_means = vec![]; + let mut partition_start = vec![]; + let mut offset = 0; + let all_sub_samples = partition_mut(samples, &alloc); + for (sub_samples, nlist) in all_sub_samples.into_iter().zip(alloc_lists) { + partition_start.push(offset); + offset += nlist as usize; + let f = crate::k_means(d, sub_samples, nlist as usize, num_threads, seed); + bottom_k_means.push(f); + } + Box::new(Hierarchical { + this: This { + pool, + d, + c, + centroids, + targets, + rng, + }, + top_centroids, + partition_start, + bottom_k_means, + }) +} + +/// Allocate clusters to different parts according to the given proportions +/// +/// See: https://en.wikipedia.org/wiki/Sainte-Lagu%C3%AB_method +fn successive_quotients_allocate(all_clusters: usize, proportion: Vec) -> Vec { + let mut alloc_lists = vec![1u32; proportion.len()]; + let mut diff = all_clusters as i32 - proportion.len() as i32; + assert!(diff >= 0); + let mut priorities: BinaryHeap<(AlwaysEqual, Distance)> = proportion + .iter() + .enumerate() + .map(|(i, x)| { + ( + AlwaysEqual(i), + Distance::from_f32(*x as f32 / (alloc_lists[i] as f32 + 0.5)), + ) + }) + .collect(); + while diff > 0 { + let index = priorities.pop().unwrap().0.0; + alloc_lists[index] += 1; + priorities.push(( + AlwaysEqual(index), + Distance::from_f32(proportion[index] as f32 / (alloc_lists[index] as f32 + 0.5)), + )); + diff -= 1; + } + alloc_lists +} + +pub fn partition_mut<'a>( + mut a: SquareMut<'a>, + groups: &[impl AsRef<[usize]>], +) -> Vec> { + let n = a.len(); + let permutation = groups + .iter() + .flat_map(AsRef::as_ref) + .copied() + .collect::>(); + let mut marked = vec![false; a.len()]; + let mut buffer = vec![0.0; a.d()]; + for i in 0..n { + if marked[i] { + continue; + } + buffer.copy_from_slice(&a[i]); + let (mut src, mut dst) = (permutation[i], i); + while src != i { + a.copy_within(src..src + 1, dst); + marked[dst] = true; + (src, dst) = (permutation[src], src); + } + a[dst].copy_from_slice(&buffer); + marked[dst] = true; + } + let mut result = Vec::with_capacity(groups.len()); + let (d, mut p) = a.into_inner(); + for group in groups { + let group = group.as_ref(); + let head; + (head, p) = std::mem::take(&mut p).split_at_mut(group.len() * d); + result.push(SquareMut::new(d, head)); + } + result +} + +#[cfg(test)] +mod partition_tests { + use super::*; + use rand::prelude::*; + + fn mk_square_random(d: usize, rows: usize, rng: &mut impl Rng) -> Square { + let mut s = Square::with_capacity(d, rows); + for _ in 0..rows { + let row: Vec = (0..d).map(|_| rng.random_range(-1000.0..1000.0)).collect(); + s.push_slice(&row); + } + s + } + + fn gen_random_alloc(rows: usize, groups: usize, rng: &mut impl Rng) -> Vec> { + let mut idx: Vec = (0..rows).collect(); + idx.shuffle(rng); + let mut alloc = Vec::with_capacity(groups); + let mut start = 0; + for g in 0..groups { + let rem = groups - g; + let take = if rem == 1 { + rows - start + } else { + rng.random_range(1..=(rows - start - (rem - 1))) + }; + alloc.push(idx[start..start + take].to_vec()); + start += take; + } + alloc + } + #[test] + fn random_partition() { + let mut rng = StdRng::seed_from_u64(7); + for trial in 0..1000 { + let d = rng.random_range(1..10); + let rows = rng.random_range(1000..2000); + let groups = rng.random_range(10..=rows.min(20)); + let mut s = mk_square_random(d, rows, &mut rng); + let golden: Vec> = (0..s.len()).map(|i| s[i].to_vec()).collect(); + let alloc = gen_random_alloc(rows, groups, &mut rng); + let views = partition_mut(s.as_mut_view(), &alloc); + assert_eq!(views.len(), alloc.len(), "trial {}", trial); + + for (g, group) in alloc.iter().enumerate() { + let v = &views[g]; + assert_eq!(v.len(), group.len(), "trial {}, group {}", trial, g); + assert_eq!(v.d(), d); + for (r, &row_idx) in group.iter().enumerate() { + assert_eq!( + v.row(r), + &golden[row_idx][..], + "trial {}, group {}, row {}", + trial, + g, + r + ); + } + } + } + } +} diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index bc70a8cb..dae6a10a 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -13,18 +13,18 @@ // Copyright (c) 2025 TensorChord Inc. pub mod flat; +pub mod hierarchical; pub mod quick; pub mod rabitq; pub mod square; -use crate::square::Square; +use crate::square::{Square, SquareMut}; use rand::rngs::StdRng; pub struct This { pool: rayon::ThreadPool, rng: StdRng, d: usize, - samples: Square, c: usize, centroids: Square, targets: Vec, @@ -45,16 +45,17 @@ pub trait KMeans { }); }); } + fn index(&mut self) -> Box u32 + '_>; fn finish(self: Box) -> Square; } -pub fn k_means( +pub fn k_means<'a>( d: usize, - samples: Square, + samples: SquareMut<'a>, c: usize, num_threads: usize, seed: [u8; 32], -) -> Box { +) -> Box { assert!(d > 0 && c > 0 && num_threads > 0); let n = samples.len(); if n <= c { @@ -65,3 +66,15 @@ pub fn k_means( rabitq::new(d, samples, c, num_threads, seed) } } + +pub fn hierarchical_k_means<'a>( + d: usize, + samples: SquareMut<'a>, + c: usize, + num_threads: usize, + seed: [u8; 32], + is_spherical: bool, +) -> Box { + assert!(d > 0 && c > 0 && num_threads > 0); + hierarchical::new(d, samples, c, num_threads, seed, is_spherical) +} diff --git a/crates/k_means/src/quick.rs b/crates/k_means/src/quick.rs index 3a588915..47b8e234 100644 --- a/crates/k_means/src/quick.rs +++ b/crates/k_means/src/quick.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::square::Square; +use crate::square::{Square, SquareMut}; use crate::{KMeans, This}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; @@ -30,19 +30,25 @@ impl KMeans for Quick { fn update(&mut self) {} + fn index(&mut self) -> Box u32> { + let this = self.this(); + let centroids = this.centroids.clone(); + Box::new(move |sample| crate::flat::k_means_lookup(sample, ¢roids) as u32) + } + fn finish(self: Box) -> Square { let this = self.this; this.centroids } } -pub fn new( +pub fn new<'a>( d: usize, - samples: Square, + samples: SquareMut<'a>, c: usize, num_threads: usize, seed: [u8; 32], -) -> Box { +) -> Box { let pool = rayon::ThreadPoolBuilder::new() .num_threads(num_threads) .build() @@ -70,7 +76,6 @@ pub fn new( pool, rng, d, - samples, c, centroids, targets, diff --git a/crates/k_means/src/rabitq.rs b/crates/k_means/src/rabitq.rs index 2940cee3..861e055e 100644 --- a/crates/k_means/src/rabitq.rs +++ b/crates/k_means/src/rabitq.rs @@ -12,24 +12,26 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::square::Square; +use crate::square::{Square, SquareMut}; use crate::{KMeans, This}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use rayon::prelude::*; use simd::Floating; -struct RaBitQ { +struct RaBitQ<'a> { this: This, + samples: SquareMut<'a>, } -impl KMeans for RaBitQ { +impl<'a> KMeans for RaBitQ<'a> { fn this(&mut self) -> &mut This { &mut self.this } fn assign(&mut self) { let this = &mut self.this; + let samples = &mut self.samples; this.pool.install(|| { use rabitq::packing::{pack_to_u4, padding_pack}; @@ -51,7 +53,7 @@ impl KMeans for RaBitQ { this.targets .par_iter_mut() - .zip(this.samples.into_par_iter()) + .zip(samples.par_iter_mut()) .for_each(|(target, sample)| { let lut = rabitq::bit::block::preprocess(sample); let mut result = (f32::INFINITY, 0); @@ -77,20 +79,21 @@ impl KMeans for RaBitQ { fn update(&mut self) { let this = &mut self.this; + let samples = &mut self.samples; this.pool.install(|| { const DELTA: f32 = 9.7656e-4_f32; let d = this.d; - let n = this.samples.len(); + let n = samples.len(); let c = this.c; let list = rayon::broadcast({ |ctx| { let mut sum = Square::from_zeros(d, c); let mut count = vec![0.0f32; c]; - for i in (ctx.index()..this.samples.len()).step_by(ctx.num_threads()) { + for i in (ctx.index()..samples.len()).step_by(ctx.num_threads()) { let target = this.targets[i]; - let sample = &this.samples[i]; + let sample = &samples[i]; f32::vector_add_inplace(&mut sum[target], sample); count[target] += 1.0; } @@ -134,6 +137,17 @@ impl KMeans for RaBitQ { }); } + fn index(&mut self) -> Box u32> { + let this = self.this(); + let mut centroids = this.centroids.clone(); + this.pool.install(|| { + centroids.par_iter_mut().for_each(|centroid| { + rabitq::rotate::rotate_reversed_inplace(centroid); + }); + }); + Box::new(move |sample| crate::flat::k_means_lookup(sample, ¢roids) as u32) + } + fn finish(self: Box) -> Square { let mut this = self.this; this.pool.install(|| { @@ -145,13 +159,13 @@ impl KMeans for RaBitQ { } } -pub fn new( +pub fn new<'a>( d: usize, - mut samples: Square, + mut samples: SquareMut<'a>, c: usize, num_threads: usize, seed: [u8; 32], -) -> Box { +) -> Box { let pool = rayon::ThreadPoolBuilder::new() .num_threads(num_threads) .build() @@ -188,8 +202,8 @@ pub fn new( centroids, targets, rng, - samples, }, + samples, }) } diff --git a/crates/k_means/src/square.rs b/crates/k_means/src/square.rs index 0497c765..0309fd3c 100644 --- a/crates/k_means/src/square.rs +++ b/crates/k_means/src/square.rs @@ -59,6 +59,13 @@ impl Square { pub fn truncate(&mut self, len: usize) { self.p.truncate(self.d * len); } + #[inline] + pub fn as_mut_view(&mut self) -> SquareMut<'_> { + SquareMut { + d: self.d, + p: self.p.as_mut_slice(), + } + } } impl std::ops::Index for Square { @@ -114,3 +121,83 @@ impl<'a> rayon::prelude::IntoParallelIterator for &'a mut Square { rayon::slice::ParallelSliceMut::par_chunks_exact_mut(self.p.as_mut_slice(), self.d) } } + +#[derive(Debug)] +pub struct SquareMut<'a> { + d: usize, + p: &'a mut [f32], +} +impl<'a> SquareMut<'a> { + #[inline] + pub fn d(&self) -> usize { + self.d + } + #[inline] + pub fn len(&self) -> usize { + self.p.len() / self.d + } + #[inline] + pub fn is_empty(&self) -> bool { + self.p.is_empty() + } + pub fn new(d: usize, p: &'a mut [f32]) -> Self { + Self { d, p } + } + #[inline] + pub fn iter_mut<'b>(&'b mut self) -> std::slice::ChunksExactMut<'b, f32> + where + 'a: 'b, + { + self.p.chunks_exact_mut(self.d) + } + #[inline] + pub fn par_iter_mut<'b>(&'b mut self) -> rayon::slice::ChunksExactMut<'b, f32> + where + 'a: 'b, + { + rayon::slice::ParallelSliceMut::par_chunks_exact_mut(self.p, self.d) + } + #[inline] + pub fn row(&self, i: usize) -> &[f32] { + let d = self.d; + &self.p[i * d..(i + 1) * d] + } + #[inline] + pub fn row_mut(&mut self, i: usize) -> &mut [f32] { + let d = self.d; + &mut self.p[i * d..(i + 1) * d] + } + pub fn copy_within>(&mut self, src: R, dest: usize) { + let src_start = src.start_bound().map(|x| self.d * x); + let src_end = src.end_bound().map(|x| self.d * x); + self.p.copy_within((src_start, src_end), self.d * dest); + } + #[inline] + pub fn into_inner(self) -> (usize, &'a mut [f32]) { + (self.d, self.p) + } +} + +impl std::ops::Index for SquareMut<'_> { + type Output = [f32]; + + fn index(&self, index: usize) -> &Self::Output { + &self.p[self.d * index..][..self.d] + } +} + +impl std::ops::IndexMut for SquareMut<'_> { + fn index_mut(&mut self, index: usize) -> &mut Self::Output { + &mut self.p[self.d * index..][..self.d] + } +} + +impl<'a> rayon::prelude::IntoParallelIterator for &'a mut SquareMut<'a> { + type Item = &'a mut [f32]; + + type Iter = rayon::slice::ChunksExactMut<'a, f32>; + + fn into_par_iter(self) -> Self::Iter { + rayon::slice::ParallelSliceMut::par_chunks_exact_mut(self.p, self.d) + } +} diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 2adc8dc4..da58892f 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -25,6 +25,7 @@ use crate::index::vchordrq::types::*; use index::relation::{ Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, }; +use k_means::flat::k_means_lookup; use k_means::square::Square; use simd::Floating; use std::ffi::CStr; @@ -1274,8 +1275,20 @@ fn make_internal_build( ) -> Vec> { use humansize::{BINARY, format_size}; use std::iter::once; + let (reduction, sample_dim) = match internal_build.kmeans_dimension_reduction { + None => (false, vector_options.dims as usize), + Some(d) if d < vector_options.dims => (true, d as usize), + Some(d) => { + pgrx::warning!( + "ignoring `kmeans_dimension_reduction = {}` because it is less than the vector dimension {}", + d, + vector_options.dims + ); + (false, vector_options.dims as usize) + } + }; { - let d = vector_options.dims as u64; + let d = sample_dim as u64; let c = internal_build.lists.last().copied().unwrap_or_default() as u64; let f = internal_build.sampling_factor as u64; let t = internal_build.build_threads as u64; @@ -1295,8 +1308,7 @@ fn make_internal_build( .last() .map(|x| x.saturating_mul(internal_build.sampling_factor)) .unwrap_or_default(); - let mut samples = - Square::with_capacity(vector_options.dims as _, max_number_of_samples as _); + let mut samples = Square::with_capacity(sample_dim, max_number_of_samples as _); if samples.len() >= max_number_of_samples as usize { break 'a samples; } @@ -1307,7 +1319,7 @@ fn make_internal_build( let vectors = unsafe { opfamily.store(datum) }; if let Some(vectors) = vectors { for (vector, _) in vectors { - let x = match vector { + let mut x = match vector { OwnedVector::Vecf32(x) => VectOwned::normalize(x), OwnedVector::Vecf16(x) => VectOwned::normalize(x), }; @@ -1316,9 +1328,13 @@ fn make_internal_build( x.len() as u32, "invalid vector dimensions" ); + if reduction { + rabitq::rotate::rotate_inplace(&mut x); + x.truncate(sample_dim); + } samples.push_slice(x.as_slice()); if samples.len() >= max_number_of_samples as usize { - break; + break 'a samples; } } } @@ -1329,7 +1345,7 @@ fn make_internal_build( let mut result = Vec::>::new(); let mut samples = Some(samples); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { - let input = if let Some(structure) = result.last() { + let mut input = if let Some(structure) = result.last() { let mut input = Square::with_capacity(vector_options.dims as _, structure.centroids.len()); for slice in structure.centroids.iter() { @@ -1343,7 +1359,7 @@ fn make_internal_build( }; let num_threads = internal_build.build_threads as _; let num_points = input.len(); - let num_dims = vector_options.dims as usize; + let num_dims = input.d(); let num_lists = w as usize; let num_iterations = internal_build.kmeans_iterations as _; if result.is_empty() { @@ -1358,7 +1374,23 @@ fn make_internal_build( "clustering: starting, using {num_threads} threads, clustering {num_points} vectors of {num_dims} dimension into {num_lists} clusters, in {num_iterations} iterations" ); } - let mut f = k_means::k_means(num_dims, input, num_lists, num_threads, [7; 32]); + let mut f = { + let view = input.as_mut_view(); + if result.last().is_none() + && let KMeansAlgorithm::Hierarchical {} = internal_build.kmeans_algorithm + { + k_means::hierarchical_k_means( + num_dims, + view, + num_lists, + num_threads, + [7; 32], + internal_build.spherical_centroids, + ) + } else { + k_means::k_means(num_dims, view, num_lists, num_threads, [7; 32]) + } + }; if internal_build.spherical_centroids { f.sphericalize(); } @@ -1381,7 +1413,45 @@ fn make_internal_build( f.sphericalize(); } } - let centroids = f.finish(); + let centroids = if reduction && result.last().is_none() { + let mut sample = sampler.sample(); + let mut next_sample = || -> Option<(Vec, Vec)> { + let mut tuple = sample.next()?; + let (values, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(values[0]); + if let Some(datum) = datum { + let vectors = unsafe { opfamily.store(datum) }; + if let Some(vectors) = vectors { + if let Some((vector, _)) = vectors.into_iter().next() { + let x = match vector { + OwnedVector::Vecf32(x) => VectOwned::normalize(x), + OwnedVector::Vecf16(x) => VectOwned::normalize(x), + }; + let mut y = rabitq::rotate::rotate(&x); + y.truncate(sample_dim); + return Some((x, y)); + } + None + } else { + None + } + } else { + None + } + }; + k_means_dimension_restore( + vector_options.dims as usize, + num_lists, + num_points, + internal_build.spherical_centroids, + [7; 32], + f.index(), + &mut next_sample, + ) + } else { + f.finish() + }; + if result.is_empty() { let percentage = 100; let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); @@ -1395,17 +1465,6 @@ fn make_internal_build( if let Some(structure) = result.last() { let mut children = vec![Vec::new(); centroids.len()]; for i in 0..structure.len() as u32 { - pub fn k_means_lookup(vector: &[f32], centroids: &Square) -> usize { - assert_ne!(centroids.len(), 0); - let mut result = (f32::INFINITY, 0); - for i in 0..centroids.len() { - let dis = f32::reduce_sum_of_d2(vector, ¢roids[i]); - if dis <= result.0 { - result = (dis, i); - } - } - result.1 - } let target = k_means_lookup(&structure.centroids[i as usize], ¢roids); children[target].push(i); } @@ -1428,6 +1487,45 @@ fn make_internal_build( result } +fn k_means_dimension_restore<'a>( + d: usize, + c: usize, + num_points: usize, + is_spherical: bool, + seed: [u8; 32], + index: Box u32 + 'a>, + mut next_sample: impl FnMut() -> Option<(Vec, Vec)>, +) -> Square { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + let mut centroids = Square::from_zeros(d, c); + let mut count = vec![0_u32; c]; + let mut traveled = 0; + let mut rng = StdRng::from_seed(seed); + while let Some((original, reduction)) = next_sample() + && traveled < num_points + { + let cid = index(&reduction) as usize; + f32::vector_add_inplace(&mut centroids[cid], &original); + count[cid] += 1; + traveled += 1; + } + let mut ret = Square::new(d); + for (i, centroid) in centroids.as_mut_view().iter_mut().enumerate() { + if count[i] == 0 { + centroid.fill_with(|| rng.random_range(-1.0f32..1.0f32)); + } else { + f32::vector_mul_scalar_inplace(centroid, 1.0 / count[i] as f32); + } + if is_spherical { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + } + ret.push_slice(centroid); + } + ret +} + #[allow(clippy::collapsible_else_if)] fn make_external_build( vector_options: VectorOptions, diff --git a/src/index/vchordrq/types.rs b/src/index/vchordrq/types.rs index dd2c166b..897e0b53 100644 --- a/src/index/vchordrq/types.rs +++ b/src/index/vchordrq/types.rs @@ -27,6 +27,14 @@ impl Default for VchordrqDefaultBuildOptions { } } +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +#[serde(rename_all = "snake_case")] +pub enum KMeansAlgorithm { + Lloyd {}, + Hierarchical {}, +} + #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] pub struct VchordrqInternalBuildOptions { @@ -44,6 +52,11 @@ pub struct VchordrqInternalBuildOptions { #[serde(default = "VchordrqInternalBuildOptions::default_build_threads")] #[validate(range(min = 1, max = 255))] pub build_threads: u16, + #[serde(default = "VchordrqInternalBuildOptions::default_kmeans_algorithm")] + pub kmeans_algorithm: KMeansAlgorithm, + #[serde(default = "VchordrqInternalBuildOptions::default_kmeans_dimension_reduction")] + #[validate(range(min = 1, max = 16000))] + pub kmeans_dimension_reduction: Option, } impl VchordrqInternalBuildOptions { @@ -71,6 +84,12 @@ impl VchordrqInternalBuildOptions { fn default_build_threads() -> u16 { 1 } + fn default_kmeans_algorithm() -> KMeansAlgorithm { + KMeansAlgorithm::Lloyd {} + } + fn default_kmeans_dimension_reduction() -> Option { + None + } } impl Default for VchordrqInternalBuildOptions { @@ -81,6 +100,8 @@ impl Default for VchordrqInternalBuildOptions { sampling_factor: Self::default_sampling_factor(), kmeans_iterations: Self::default_kmeans_iterations(), build_threads: Self::default_build_threads(), + kmeans_algorithm: Self::default_kmeans_algorithm(), + kmeans_dimension_reduction: Self::default_kmeans_dimension_reduction(), } } } diff --git a/tests/vchordrq/approximate.slt b/tests/vchordrq/approximate.slt new file mode 100644 index 00000000..afb8715e --- /dev/null +++ b/tests/vchordrq/approximate.slt @@ -0,0 +1,57 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.internal] +lists = [33] +spherical_centroids = false +kmeans_dimension_reduction = 2 +kmeans_algorithm.hierarchical = {} +$$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_ip_ops) +WITH (options = $$ +residual_quantization = false +[build.internal] +lists = [33] +spherical_centroids = true +kmeans_algorithm.hierarchical = {} +$$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_cosine_ops) +WITH (options = $$ +residual_quantization = false +[build.internal] +lists = [33] +spherical_centroids = true +kmeans_algorithm.hierarchical = {} +$$); + +statement ok +SET vchordrq.probes = '16'; + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE t; From e3c45f5fbc0abd1429014e8da0532573aab70f49 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 4 Nov 2025 11:08:01 +0800 Subject: [PATCH 243/324] fix: vchordg freespace recycle (#370) job: +check_arch Signed-off-by: usamoi --- Cargo.toml | 1 + crates/vchordg/src/maintain.rs | 34 +++++++++------------- crates/vchordg/src/tuples.rs | 1 + src/index/fetcher.rs | 1 - src/index/storage.rs | 4 ++- tests/vchordg/freespace.slt | 53 ++++++++++++++++++++++++++++++++++ 6 files changed, 71 insertions(+), 23 deletions(-) create mode 100644 tests/vchordg/freespace.slt diff --git a/Cargo.toml b/Cargo.toml index 96a11943..08847046 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -85,6 +85,7 @@ clippy.nonminimal_bool = "allow" clippy.too_many_arguments = "allow" clippy.type_complexity = "allow" # style +clippy.collapsible_if = "allow" clippy.just_underscores_and_digits = "allow" clippy.needless_range_loop = "allow" # unsafe diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs index 1f36e52e..1451fd5b 100644 --- a/crates/vchordg/src/maintain.rs +++ b/crates/vchordg/src/maintain.rs @@ -14,7 +14,7 @@ use crate::Opaque; use crate::operator::{CloneAccessor, Operator}; -use crate::tuples::{MetaTuple, VectorTuple, VertexTuple, WithReader}; +use crate::tuples::{MetaTuple, VertexTuple, WithReader}; use crate::types::DistanceKind; use crate::vectors::{by_read, copy_all, copy_nothing, copy_outs, update}; use always_equal::AlwaysEqual; @@ -152,44 +152,36 @@ where while current != u32::MAX { check(); let mut vertex_guard = index.write(current, true); + let mut reachable_set = Vec::<(u32, u16)>::new(); for i in 1..=vertex_guard.len() { if let Some(bytes) = vertex_guard.get(i) { let tuple = VertexTuple::deserialize_ref(bytes); let p = tuple.payload(); if p.is_none() && (current, i) != s { vertex_guard.free(i); + } else { + let iter = tuple.pointers().iter().map(|pointer| pointer.into_inner()); + reachable_set.extend(iter); } - }; + } } - current = vertex_guard.get_opaque().next; + reachable_set.sort_unstable(); { - let mut current = { vertex_guard }.get_opaque().link; + let mut current = vertex_guard.get_opaque().link; while current != u32::MAX { check(); let mut vector_guard = index.write(current, false); for i in 1..=vector_guard.len() { - if let Some(bytes) = vector_guard.get(i) { - use crate::tuples::VectorTupleReader; - let tuple = VectorTuple::::deserialize_ref(bytes); - match tuple { - VectorTupleReader::_0(tuple) => { - let p = tuple.payload(); - if p.is_none() && (current, i) != s { - vector_guard.free(i); - } - } - VectorTupleReader::_1(tuple) => { - let p = tuple.payload(); - if p.is_none() && (current, i) != s { - vector_guard.free(i); - } - } + if vector_guard.get(i).is_some() { + if reachable_set.binary_search(&(current, i)).is_err() { + vector_guard.free(i); } - }; + } } current = vector_guard.get_opaque().next; } } + current = vertex_guard.get_opaque().next; } } } diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 1bc1ead2..99f3f9d4 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -577,6 +577,7 @@ impl<'a, V: Vector> Clone for VectorTupleReader1<'a, V> { } impl<'a, V: Vector> VectorTupleReader1<'a, V> { + #[allow(dead_code)] pub fn payload(self) -> Option> { self.header.payload } diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs index e9c779bc..bba55206 100644 --- a/src/index/fetcher.rs +++ b/src/index/fetcher.rs @@ -178,7 +178,6 @@ impl Tuple for HeapTuple<'_> { } impl FilterableTuple for HeapTuple<'_> { - #[allow(clippy::collapsible_if)] fn filter(&mut self) -> bool { unsafe { use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; diff --git a/src/index/storage.rs b/src/index/storage.rs index a0b36950..d6896e97 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -79,7 +79,7 @@ impl Page for PostgresPage { assert!(self.header.pd_upper as usize <= size_of::()); let lower = self.header.pd_lower as usize; let upper = self.header.pd_upper as usize; - assert!(lower <= upper); + assert!(offset_of!(PageHeaderData, pd_linp) <= lower && lower <= upper); ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16 } fn get(&self, i: u16) -> Option<&[u8]> { @@ -89,6 +89,7 @@ impl Page for PostgresPage { } assert!(self.header.pd_lower as usize <= size_of::()); let lower = self.header.pd_lower as usize; + assert!(offset_of!(PageHeaderData, pd_linp) <= lower); let n = ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16; if i > n { return None; @@ -119,6 +120,7 @@ impl Page for PostgresPage { } assert!(self.header.pd_lower as usize <= size_of::()); let lower = self.header.pd_lower as usize; + assert!(offset_of!(PageHeaderData, pd_linp) <= lower); let n = ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16; if i > n { return None; diff --git a/tests/vchordg/freespace.slt b/tests/vchordg/freespace.slt new file mode 100644 index 00000000..9f3e1ed9 --- /dev/null +++ b/tests/vchordg/freespace.slt @@ -0,0 +1,53 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +3 + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 18); + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +3 + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1); + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +4 + +statement ok +DELETE FROM t; + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +4 + +statement ok +VACUUM (INDEX_CLEANUP ON) t; + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +4 + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 19); + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +4 + +statement ok +DROP TABLE t; From 55b8e8d2e3b834a060fcc949801b2c70436ccc0c Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 4 Nov 2025 12:20:21 +0800 Subject: [PATCH 244/324] fix: enforce thread pool size upper limit (#371) Signed-off-by: usamoi --- Cargo.lock | 1 + Cargo.toml | 2 ++ crates/k_means/Cargo.toml | 2 +- crates/k_means/src/flat.rs | 14 +++++--------- crates/k_means/src/hierarchical.rs | 26 ++++++++------------------ crates/k_means/src/lib.rs | 28 ++++++++++++++-------------- crates/k_means/src/quick.rs | 16 ++++++---------- crates/k_means/src/rabitq.rs | 14 +++++--------- src/index/vchordrq/am/am_build.rs | 12 ++++++++---- 9 files changed, 50 insertions(+), 65 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4d55aea1..274bf278 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1636,6 +1636,7 @@ dependencies = [ "pgrx-catalog", "rabitq", "rand", + "rayon", "rusqlite", "seq-macro", "serde", diff --git a/Cargo.toml b/Cargo.toml index 08847046..a84582d5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,6 +42,7 @@ paste.workspace = true pgrx = { version = "=0.16.1", default-features = false, features = ["cshim"] } pgrx-catalog = "0.3.1" rand.workspace = true +rayon.workspace = true rusqlite = { version = "0.37.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true @@ -70,6 +71,7 @@ dary_heap = "0.3.8" paste = "1.0.15" rand = "0.9.2" rand_chacha = "0.9.0" +rayon = "1.11.0" seq-macro = "0.3.6" serde = { version = "1.0.228", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } diff --git a/crates/k_means/Cargo.toml b/crates/k_means/Cargo.toml index 9361a3f9..75cf247f 100644 --- a/crates/k_means/Cargo.toml +++ b/crates/k_means/Cargo.toml @@ -11,7 +11,7 @@ always_equal = { path = "../always_equal" } distance = { path = "../distance" } rand.workspace = true -rayon = "1.11.0" +rayon.workspace = true [lints] workspace = true diff --git a/crates/k_means/src/flat.rs b/crates/k_means/src/flat.rs index 270b7f5d..7230809f 100644 --- a/crates/k_means/src/flat.rs +++ b/crates/k_means/src/flat.rs @@ -20,12 +20,12 @@ use rayon::prelude::*; use simd::Floating; struct Flat<'a> { - this: This, + this: This<'a>, samples: SquareMut<'a>, } -impl<'a> KMeans for Flat<'a> { - fn this(&mut self) -> &mut This { +impl<'a> KMeans<'a> for Flat<'a> { + fn this(&mut self) -> &mut This<'a> { &mut self.this } fn assign(&mut self) { @@ -113,16 +113,12 @@ impl<'a> KMeans for Flat<'a> { } pub fn new<'a>( + pool: &'a rayon::ThreadPool, d: usize, samples: SquareMut<'a>, c: usize, - num_threads: usize, seed: [u8; 32], -) -> Box { - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(num_threads) - .build() - .expect("failed to build thread pool"); +) -> Box + 'a> { let mut rng = StdRng::from_seed(seed); let mut centroids = Square::with_capacity(d, c); diff --git a/crates/k_means/src/hierarchical.rs b/crates/k_means/src/hierarchical.rs index ffd4a36c..31beece0 100644 --- a/crates/k_means/src/hierarchical.rs +++ b/crates/k_means/src/hierarchical.rs @@ -22,14 +22,14 @@ use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelI use std::collections::BinaryHeap; struct Hierarchical<'a> { - this: This, + this: This<'a>, top_centroids: Square, partition_start: Vec, - bottom_k_means: Vec>, + bottom_k_means: Vec + 'a>>, } -impl<'a> KMeans for Hierarchical<'a> { - fn this(&mut self) -> &mut This { +impl<'a> KMeans<'a> for Hierarchical<'a> { + fn this(&mut self) -> &mut This<'a> { &mut self.this } @@ -84,17 +84,13 @@ const LOCAL_SAMPLE_FACTOR: usize = 256; const LOCAL_NUM_ITERATIONS: usize = 10; pub fn new<'a>( + pool: &'a rayon::ThreadPool, d: usize, mut samples: SquareMut<'a>, c: usize, - num_threads: usize, seed: [u8; 32], is_spherical: bool, -) -> Box { - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(num_threads) - .build() - .expect("failed to build thread pool"); +) -> Box + 'a> { let centroids = Square::new(d); let targets = vec![0; samples.len()]; let samples_len = samples.len(); @@ -108,13 +104,7 @@ pub fn new<'a>( { top_samples.push_slice(&samples[index]); } - let mut f = crate::k_means( - d, - top_samples.as_mut_view(), - top_list as usize, - num_threads, - seed, - ); + let mut f = crate::k_means(pool, d, top_samples.as_mut_view(), top_list as usize, seed); if is_spherical { f.sphericalize(); } @@ -159,7 +149,7 @@ pub fn new<'a>( for (sub_samples, nlist) in all_sub_samples.into_iter().zip(alloc_lists) { partition_start.push(offset); offset += nlist as usize; - let f = crate::k_means(d, sub_samples, nlist as usize, num_threads, seed); + let f = crate::k_means(pool, d, sub_samples, nlist as usize, seed); bottom_k_means.push(f); } Box::new(Hierarchical { diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index dae6a10a..e8c531be 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -21,8 +21,8 @@ pub mod square; use crate::square::{Square, SquareMut}; use rand::rngs::StdRng; -pub struct This { - pool: rayon::ThreadPool, +pub struct This<'a> { + pool: &'a rayon::ThreadPool, rng: StdRng, d: usize, c: usize, @@ -30,8 +30,8 @@ pub struct This { targets: Vec, } -pub trait KMeans { - fn this(&mut self) -> &mut This; +pub trait KMeans<'a> { + fn this(&mut self) -> &mut This<'a>; fn assign(&mut self); fn update(&mut self); fn sphericalize(&mut self) { @@ -50,31 +50,31 @@ pub trait KMeans { } pub fn k_means<'a>( + pool: &'a rayon::ThreadPool, d: usize, samples: SquareMut<'a>, c: usize, - num_threads: usize, seed: [u8; 32], -) -> Box { - assert!(d > 0 && c > 0 && num_threads > 0); +) -> Box + 'a> { + assert!(d > 0 && c > 0); let n = samples.len(); if n <= c { - quick::new(d, samples, c, num_threads, seed) + quick::new(pool, d, samples, c, seed) } else if n <= c * 2 { - flat::new(d, samples, c, num_threads, seed) + flat::new(pool, d, samples, c, seed) } else { - rabitq::new(d, samples, c, num_threads, seed) + rabitq::new(pool, d, samples, c, seed) } } pub fn hierarchical_k_means<'a>( + pool: &'a rayon::ThreadPool, d: usize, samples: SquareMut<'a>, c: usize, - num_threads: usize, seed: [u8; 32], is_spherical: bool, -) -> Box { - assert!(d > 0 && c > 0 && num_threads > 0); - hierarchical::new(d, samples, c, num_threads, seed, is_spherical) +) -> Box + 'a> { + assert!(d > 0 && c > 0); + hierarchical::new(pool, d, samples, c, seed, is_spherical) } diff --git a/crates/k_means/src/quick.rs b/crates/k_means/src/quick.rs index 47b8e234..5572027f 100644 --- a/crates/k_means/src/quick.rs +++ b/crates/k_means/src/quick.rs @@ -17,12 +17,12 @@ use crate::{KMeans, This}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; -struct Quick { - this: This, +struct Quick<'a> { + this: This<'a>, } -impl KMeans for Quick { - fn this(&mut self) -> &mut This { +impl<'a> KMeans<'a> for Quick<'a> { + fn this(&mut self) -> &mut This<'a> { &mut self.this } @@ -43,16 +43,12 @@ impl KMeans for Quick { } pub fn new<'a>( + pool: &'a rayon::ThreadPool, d: usize, samples: SquareMut<'a>, c: usize, - num_threads: usize, seed: [u8; 32], -) -> Box { - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(num_threads) - .build() - .expect("failed to build thread pool"); +) -> Box + 'a> { let mut rng = StdRng::from_seed(seed); let mut centroids = Square::with_capacity(d, c); diff --git a/crates/k_means/src/rabitq.rs b/crates/k_means/src/rabitq.rs index 861e055e..c30e5c2d 100644 --- a/crates/k_means/src/rabitq.rs +++ b/crates/k_means/src/rabitq.rs @@ -20,12 +20,12 @@ use rayon::prelude::*; use simd::Floating; struct RaBitQ<'a> { - this: This, + this: This<'a>, samples: SquareMut<'a>, } -impl<'a> KMeans for RaBitQ<'a> { - fn this(&mut self) -> &mut This { +impl<'a> KMeans<'a> for RaBitQ<'a> { + fn this(&mut self) -> &mut This<'a> { &mut self.this } @@ -160,16 +160,12 @@ impl<'a> KMeans for RaBitQ<'a> { } pub fn new<'a>( + pool: &'a rayon::ThreadPool, d: usize, mut samples: SquareMut<'a>, c: usize, - num_threads: usize, seed: [u8; 32], -) -> Box { - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(num_threads) - .build() - .expect("failed to build thread pool"); +) -> Box + 'a> { let mut rng = StdRng::from_seed(seed); pool.install(|| { diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index da58892f..9a03cbf3 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -1342,6 +1342,11 @@ fn make_internal_build( } samples }; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(internal_build.build_threads as usize) + .build() + .expect("failed to build thread pool"); + pgrx::info!("clustering: using {} threads", pool.current_num_threads()); let mut result = Vec::>::new(); let mut samples = Some(samples); for w in internal_build.lists.iter().rev().copied().chain(once(1)) { @@ -1357,7 +1362,6 @@ fn make_internal_build( } else { unreachable!() }; - let num_threads = internal_build.build_threads as _; let num_points = input.len(); let num_dims = input.d(); let num_lists = w as usize; @@ -1371,7 +1375,7 @@ fn make_internal_build( } if num_lists > 1 { pgrx::info!( - "clustering: starting, using {num_threads} threads, clustering {num_points} vectors of {num_dims} dimension into {num_lists} clusters, in {num_iterations} iterations" + "clustering: starting, clustering {num_points} vectors of {num_dims} dimension into {num_lists} clusters, in {num_iterations} iterations" ); } let mut f = { @@ -1380,15 +1384,15 @@ fn make_internal_build( && let KMeansAlgorithm::Hierarchical {} = internal_build.kmeans_algorithm { k_means::hierarchical_k_means( + &pool, num_dims, view, num_lists, - num_threads, [7; 32], internal_build.spherical_centroids, ) } else { - k_means::k_means(num_dims, view, num_lists, num_threads, [7; 32]) + k_means::k_means(&pool, num_dims, view, num_lists, [7; 32]) } }; if internal_build.spherical_centroids { From 5d5fee637bba105c1db48bbe38b54cf28f145bb0 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 5 Nov 2025 16:24:38 +0800 Subject: [PATCH 245/324] feat: multithreading dimension expansion (#372) * adds multithreading dimension expansion * don't copy centroids when creating index for K-means * fix `build.internal.kmeans_algorithm.hierarchical = {}` on empty table * fix ci `style`, toml formatting --------- Signed-off-by: usamoi --- .github/workflows/check.yml | 2 +- Cargo.lock | 10 + Cargo.toml | 1 + crates/k_means/Cargo.toml | 4 +- crates/k_means/src/flat.rs | 134 +++------ crates/k_means/src/hierarchical.rs | 360 ++++++++++++----------- crates/k_means/src/index.rs | 156 ++++++++++ crates/k_means/src/lib.rs | 57 ++-- crates/k_means/src/quick.rs | 47 ++- crates/k_means/src/rabitq.rs | 178 ++++------- crates/k_means/src/square.rs | 11 + src/index/vchordrq/am/am_build.rs | 246 +++++++++------- tests/vchordrq/internal_build_kmeans.slt | 19 ++ 13 files changed, 674 insertions(+), 551 deletions(-) create mode 100644 crates/k_means/src/index.rs create mode 100644 tests/vchordrq/internal_build_kmeans.slt diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 5d89ca94..b88a4be8 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -21,7 +21,7 @@ jobs: steps: - name: Set up Environment run: | - curl -fsSL https://github.com/tamasfe/taplo/releases/latest/download/taplo-full-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo + curl -fsSL https://github.com/tamasfe/taplo/releases/download/0.10.0/taplo-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo curl -fsSL https://github.com/EmbarkStudios/cargo-deny/releases/download/0.18.2/cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl/cargo-deny | install -m 755 /dev/stdin /usr/local/bin/cargo-deny diff --git a/Cargo.lock b/Cargo.lock index 274bf278..fd441118 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -265,6 +265,15 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -1624,6 +1633,7 @@ dependencies = [ "always_equal", "arrayvec", "bumpalo", + "crossbeam-channel", "dary_heap", "distance", "feistel", diff --git a/Cargo.toml b/Cargo.toml index a84582d5..2c776853 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ vector = { path = "./crates/vector" } arrayvec = "0.7.6" bumpalo.workspace = true +crossbeam-channel = "0.5.15" dary_heap.workspace = true humansize = "2.1.3" paste.workspace = true diff --git a/crates/k_means/Cargo.toml b/crates/k_means/Cargo.toml index 75cf247f..0ddf105f 100644 --- a/crates/k_means/Cargo.toml +++ b/crates/k_means/Cargo.toml @@ -5,10 +5,10 @@ edition.workspace = true publish = false [dependencies] -rabitq = { path = "../rabitq" } -simd = { path = "../simd" } always_equal = { path = "../always_equal" } distance = { path = "../distance" } +rabitq = { path = "../rabitq" } +simd = { path = "../simd" } rand.workspace = true rayon.workspace = true diff --git a/crates/k_means/src/flat.rs b/crates/k_means/src/flat.rs index 7230809f..f1770616 100644 --- a/crates/k_means/src/flat.rs +++ b/crates/k_means/src/flat.rs @@ -12,103 +12,52 @@ // // Copyright (c) 2025 TensorChord Inc. +use crate::index::{flat_index as prefect_index, flat_index as index}; use crate::square::{Square, SquareMut}; use crate::{KMeans, This}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use rayon::prelude::*; -use simd::Floating; struct Flat<'a> { this: This<'a>, samples: SquareMut<'a>, + centroids: Square, + targets: Vec, } -impl<'a> KMeans<'a> for Flat<'a> { - fn this(&mut self) -> &mut This<'a> { - &mut self.this +impl<'a> KMeans for Flat<'a> { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_> { + Box::new(prefect_index(&self.centroids)) } + fn assign(&mut self) { let this = &mut self.this; + let samples = &mut self.samples; + let centroids = &self.centroids; + let targets = &mut self.targets; + let index = index(centroids); this.pool.install(|| { - this.targets + targets .par_iter_mut() - .zip(self.samples.par_iter_mut()) + .zip(samples.par_iter_mut()) .for_each(|(target, sample)| { - *target = k_means_lookup(sample, &this.centroids); + *target = index(sample).1; }); }); } fn update(&mut self) { - let this = &mut self.this; - let samples = &mut self.samples; - this.pool.install(|| { - const DELTA: f32 = 9.7656e-4_f32; - - let d = this.d; - let n = samples.len(); - let c = this.c; - - let list = rayon::broadcast({ - |ctx| { - let mut sum = Square::from_zeros(d, c); - let mut count = vec![0.0f32; c]; - for i in (ctx.index()..samples.len()).step_by(ctx.num_threads()) { - let target = this.targets[i]; - let sample = &samples[i]; - f32::vector_add_inplace(&mut sum[target], sample); - count[target] += 1.0; - } - (sum, count) - } - }); - let mut sum = Square::from_zeros(d, c); - let mut count = vec![0.0f32; c]; - for (sum_1, count_1) in list { - for i in 0..c { - f32::vector_add_inplace(&mut sum[i], &sum_1[i]); - count[i] += count_1[i]; - } - } - - sum.par_iter_mut() - .enumerate() - .for_each(|(i, sum)| f32::vector_mul_scalar_inplace(sum, 1.0 / count[i])); - - this.centroids = sum; - - for i in 0..c { - if count[i] != 0.0f32 { - continue; - } - let mut o = 0; - loop { - let alpha = this.rng.random_range(0.0..1.0f32); - let beta = (count[o] - 1.0) / (n - c) as f32; - if alpha < beta { - break; - } - o = (o + 1) % c; - } - this.centroids.copy_within(o..o + 1, i); - vector_mul_scalars_inplace(&mut this.centroids[i], [1.0 + DELTA, 1.0 - DELTA]); - vector_mul_scalars_inplace(&mut this.centroids[o], [1.0 - DELTA, 1.0 + DELTA]); - count[i] = count[o] / 2.0; - count[o] -= count[i]; - } - }); - } - - fn index(&mut self) -> Box u32> { - let this = self.this(); - let centroids = this.centroids.clone(); - Box::new(move |sample| k_means_lookup(sample, ¢roids) as u32) + crate::index::update( + &mut self.this, + &self.samples, + &self.targets, + &mut self.centroids, + ); } fn finish(self: Box) -> Square { - let this = self.this; - this.centroids + self.centroids } } @@ -118,7 +67,8 @@ pub fn new<'a>( samples: SquareMut<'a>, c: usize, seed: [u8; 32], -) -> Box + 'a> { + is_spherical: bool, +) -> Box { let mut rng = StdRng::from_seed(seed); let mut centroids = Square::with_capacity(d, c); @@ -135,6 +85,16 @@ pub fn new<'a>( centroids.push_iter((0..d).map(|_| rng.random_range(-1.0f32..1.0f32))); } + pool.install(|| { + if is_spherical { + use simd::Floating; + (&mut centroids).into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + }); + let targets = vec![0; samples.len()]; Box::new(Flat { @@ -143,32 +103,10 @@ pub fn new<'a>( rng, d, c, - centroids, - targets, + is_spherical, }, samples, + centroids, + targets, }) } - -pub fn k_means_lookup(vector: &[f32], centroids: &Square) -> usize { - assert_ne!(centroids.len(), 0); - let mut result = (f32::INFINITY, 0); - for i in 0..centroids.len() { - let dis = f32::reduce_sum_of_d2(vector, ¢roids[i]); - if dis <= result.0 { - result = (dis, i); - } - } - result.1 -} - -fn vector_mul_scalars_inplace(this: &mut [f32], scalars: [f32; 2]) { - let n: usize = this.len(); - for i in 0..n { - if i % 2 == 0 { - this[i] *= scalars[0]; - } else { - this[i] *= scalars[1]; - } - } -} diff --git a/crates/k_means/src/hierarchical.rs b/crates/k_means/src/hierarchical.rs index 31beece0..aa241f38 100644 --- a/crates/k_means/src/hierarchical.rs +++ b/crates/k_means/src/hierarchical.rs @@ -23,65 +23,91 @@ use std::collections::BinaryHeap; struct Hierarchical<'a> { this: This<'a>, - top_centroids: Square, - partition_start: Vec, - bottom_k_means: Vec + 'a>>, + partitions: Vec>, + coarse_centroids: Square, + offsets: Vec, } -impl<'a> KMeans<'a> for Hierarchical<'a> { - fn this(&mut self) -> &mut This<'a> { - &mut self.this +impl<'a> KMeans for Hierarchical<'a> { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_> { + Box::new(prefect_index(&self.partitions, &self.offsets)) + } + + fn index(&self) -> Box (f32, usize) + Sync + '_> { + Box::new(index( + &self.partitions, + &self.coarse_centroids, + &self.offsets, + )) } fn assign(&mut self) { - for k_means in &mut self.bottom_k_means { - k_means.assign(); + for partial in self.partitions.iter_mut() { + partial.assign(); } } fn update(&mut self) { - for k_means in &mut self.bottom_k_means { - k_means.update(); + for partial in self.partitions.iter_mut() { + partial.update(); } } - fn sphericalize(&mut self) { - for k_means in &mut self.bottom_k_means { - k_means.sphericalize(); + fn finish(self: Box) -> Square { + let mut centroids = Square::new(self.this.d); + for k_means in self.partitions { + let partial_centroids = k_means.finish(); + for centroid in partial_centroids.into_iter() { + centroids.push_slice(centroid); + } } + centroids } +} - fn index(&mut self) -> Box u32 + '_> { - let top_centroids = self.top_centroids.clone(); - let top = move |sample: &[f32]| crate::flat::k_means_lookup(sample, &top_centroids) as u32; - let bottom_centroids: Vec<_> = self - .bottom_k_means - .iter_mut() - .map(|k_means| k_means.index()) - .collect(); - let partition_start = self.partition_start.clone(); - let index = move |sample: &[f32]| { - let top_id = top(sample) as usize; - let bottom_id = bottom_centroids[top_id](sample) as usize; - (partition_start[top_id] + bottom_id) as u32 - }; - Box::new(index) +fn prefect_index( + partitions: &[Box], + offsets: &[usize], +) -> impl Fn(&[f32]) -> (f32, usize) + Sync { + let indexes = partitions + .iter() + .map(|p| p.prefect_index()) + .collect::>(); + move |sample| { + let mut result = (f32::INFINITY, 0); + for (id, index) in indexes.iter().enumerate() { + let partial_result = index(sample); + if partial_result.0 <= result.0 { + result = (partial_result.0, offsets[id] + partial_result.1); + } + } + result } +} - fn finish(self: Box) -> Square { - let mut ret = Square::new(self.this.d); - for k_means in self.bottom_k_means { - let centroids = k_means.finish(); - for centroid in centroids.into_iter() { - ret.push_slice(centroid); +fn index( + partitions: &[Box], + coarse_centroids: &Square, + offsets: &[usize], +) -> impl Fn(&[f32]) -> (f32, usize) + Sync { + let indexes = partitions.iter().map(|p| p.index()).collect::>(); + move |sample| { + use simd::Floating; + let mut result = (f32::INFINITY, 0); + for i in 0..coarse_centroids.len() { + let dis = f32::reduce_sum_of_d2(sample, &coarse_centroids[i]); + if dis <= result.0 { + result = (dis, i); } } - ret + let id = result.1; + let result = indexes[id](sample); + (result.0, offsets[id] + result.1) } } -const LOCAL_SAMPLE_FACTOR: usize = 256; -const LOCAL_NUM_ITERATIONS: usize = 10; +const COARSE_SAMPLING_FACTOR: usize = 256; +const COARSE_ITERATIONS: usize = 10; pub fn new<'a>( pool: &'a rayon::ThreadPool, @@ -90,116 +116,111 @@ pub fn new<'a>( c: usize, seed: [u8; 32], is_spherical: bool, -) -> Box + 'a> { - let centroids = Square::new(d); - let targets = vec![0; samples.len()]; - let samples_len = samples.len(); - let top_list = - ((c as f64).sqrt().floor() as u32).clamp(1, (samples_len as f64).sqrt().floor() as u32); - let top_samples_len = LOCAL_SAMPLE_FACTOR * (top_list as usize); - let mut top_samples = Square::new(d); +) -> Box { let mut rng = StdRng::from_seed(seed); - for index in - rand::seq::index::sample(&mut rng, samples.len(), top_samples_len.min(samples.len())) - { - top_samples.push_slice(&samples[index]); - } - let mut f = crate::k_means(pool, d, top_samples.as_mut_view(), top_list as usize, seed); - if is_spherical { - f.sphericalize(); - } - for _ in 0..LOCAL_NUM_ITERATIONS { - f.assign(); - f.update(); - if is_spherical { - f.sphericalize(); + + let mut coarse_samples = { + let mut coarse_samples = Square::new(d); + let s = c.isqrt().saturating_mul(COARSE_SAMPLING_FACTOR); + for index in rand::seq::index::sample(&mut rng, samples.len(), s.min(samples.len())) { + coarse_samples.push_slice(&samples[index]); } - } - let top_centroids = f.finish(); - let mut final_assign = vec![0; samples.len()]; - pool.install(|| { - final_assign - .par_iter_mut() - .zip(samples.par_iter_mut()) - .for_each(|(target, sample)| { - *target = crate::flat::k_means_lookup(sample, &top_centroids); - }); - }); - let alloc = final_assign.into_iter().enumerate().fold( - vec![vec![]; top_centroids.len()], - |mut acc, (i, target)| { - acc[target].push(i); - acc - }, - ); - let alloc_size = alloc.iter().map(|x| x.len() as u32).collect::>(); - let keep_indices: Vec = alloc_size - .iter() - .enumerate() - .filter_map(|(i, size)| if *size > 0 { Some(i) } else { None }) - .collect(); - let alloc: Vec<_> = keep_indices.iter().map(|&i| alloc[i].clone()).collect(); - let alloc_size: Vec<_> = keep_indices.iter().map(|&i| alloc_size[i]).collect(); - let alloc_lists = successive_quotients_allocate(c, alloc_size); + coarse_samples + }; + let coarse_k_means = { + let mut coarse_k_means = crate::lloyd_k_means( + pool, + d, + coarse_samples.as_mut_view(), + c.isqrt(), + seed, + is_spherical, + ); + for _ in 0..COARSE_ITERATIONS { + coarse_k_means.assign(); + coarse_k_means.update(); + } + coarse_k_means + }; + let coarse_assign = { + let coarse_index = coarse_k_means.prefect_index(); + let mut coarse_assign = vec![0; samples.len()]; + pool.install(|| { + coarse_assign + .par_iter_mut() + .zip(samples.par_iter_mut()) + .for_each(|(target, sample)| { + *target = coarse_index(sample).1; + }); + }); + coarse_assign + }; + let coarse_centroids = coarse_k_means.finish(); - let mut bottom_k_means = vec![]; - let mut partition_start = vec![]; + let (weight, groups) = { + let mut weight = vec![0_usize; coarse_centroids.len()]; + let mut groups = vec![vec![]; coarse_centroids.len()]; + for (i, &target) in coarse_assign.iter().enumerate() { + weight[target] += 1; + groups[target].push(i); + } + (weight, groups) + }; + let seats = modified_webster_method(c, &weight); + + let mut partitions = vec![]; + let mut offsets = vec![]; let mut offset = 0; - let all_sub_samples = partition_mut(samples, &alloc); - for (sub_samples, nlist) in all_sub_samples.into_iter().zip(alloc_lists) { - partition_start.push(offset); - offset += nlist as usize; - let f = crate::k_means(pool, d, sub_samples, nlist as usize, seed); - bottom_k_means.push(f); + for (samples, c) in std::iter::zip(partition(samples, &groups), seats) { + partitions.push(crate::lloyd_k_means( + pool, + d, + samples, + c, + seed, + is_spherical, + )); + offsets.push(offset); + offset += c; } + Box::new(Hierarchical { this: This { pool, d, c, - centroids, - targets, rng, + is_spherical, }, - top_centroids, - partition_start, - bottom_k_means, + coarse_centroids, + partitions, + offsets, }) } -/// Allocate clusters to different parts according to the given proportions -/// -/// See: https://en.wikipedia.org/wiki/Sainte-Lagu%C3%AB_method -fn successive_quotients_allocate(all_clusters: usize, proportion: Vec) -> Vec { - let mut alloc_lists = vec![1u32; proportion.len()]; - let mut diff = all_clusters as i32 - proportion.len() as i32; - assert!(diff >= 0); - let mut priorities: BinaryHeap<(AlwaysEqual, Distance)> = proportion - .iter() - .enumerate() - .map(|(i, x)| { - ( - AlwaysEqual(i), - Distance::from_f32(*x as f32 / (alloc_lists[i] as f32 + 0.5)), - ) - }) - .collect(); - while diff > 0 { - let index = priorities.pop().unwrap().0.0; - alloc_lists[index] += 1; - priorities.push(( - AlwaysEqual(index), - Distance::from_f32(proportion[index] as f32 / (alloc_lists[index] as f32 + 0.5)), - )); - diff -= 1; +// https://en.wikipedia.org/wiki/Sainte-Lagu%C3%AB_method +fn modified_webster_method(n: usize, weight: &[usize]) -> Vec { + assert!(n >= weight.len()); + let mut seats = vec![0_usize; weight.len()]; + let mut quotients = Vec::new(); + for index in 0..weight.len() { + seats[index] += 1; + let quotient = weight[index] as f64 / (seats[index] as f64 + 0.5); + quotients.push((Distance::from_f32(quotient as _), AlwaysEqual(index))); + } + let mut quotients = BinaryHeap::<_>::from(quotients); + for _ in weight.len()..n { + let Some((_, AlwaysEqual(index))) = quotients.pop() else { + break; + }; + seats[index] += 1; + let quotient = weight[index] as f64 / (seats[index] as f64 + 0.5); + quotients.push((Distance::from_f32(quotient as _), AlwaysEqual(index))); } - alloc_lists + seats } -pub fn partition_mut<'a>( - mut a: SquareMut<'a>, - groups: &[impl AsRef<[usize]>], -) -> Vec> { +fn partition<'a>(mut a: SquareMut<'a>, groups: &[impl AsRef<[usize]>]) -> Vec> { let n = a.len(); let permutation = groups .iter() @@ -233,20 +254,18 @@ pub fn partition_mut<'a>( result } -#[cfg(test)] -mod partition_tests { - use super::*; - use rand::prelude::*; - - fn mk_square_random(d: usize, rows: usize, rng: &mut impl Rng) -> Square { - let mut s = Square::with_capacity(d, rows); - for _ in 0..rows { - let row: Vec = (0..d).map(|_| rng.random_range(-1000.0..1000.0)).collect(); - s.push_slice(&row); - } - s - } +#[test] +fn test_modified_webster_method() { + let seats = modified_webster_method(51, &[10, 10, 10, 11, 9]); + assert_eq!(seats[0], 10); + assert_eq!(seats[1], 10); + assert_eq!(seats[2], 10); + assert_eq!(seats[3], 12); + assert_eq!(seats[4], 9); +} +#[test] +fn test_partition() { fn gen_random_alloc(rows: usize, groups: usize, rng: &mut impl Rng) -> Vec> { let mut idx: Vec = (0..rows).collect(); idx.shuffle(rng); @@ -264,33 +283,38 @@ mod partition_tests { } alloc } - #[test] - fn random_partition() { - let mut rng = StdRng::seed_from_u64(7); - for trial in 0..1000 { - let d = rng.random_range(1..10); - let rows = rng.random_range(1000..2000); - let groups = rng.random_range(10..=rows.min(20)); - let mut s = mk_square_random(d, rows, &mut rng); - let golden: Vec> = (0..s.len()).map(|i| s[i].to_vec()).collect(); - let alloc = gen_random_alloc(rows, groups, &mut rng); - let views = partition_mut(s.as_mut_view(), &alloc); - assert_eq!(views.len(), alloc.len(), "trial {}", trial); - for (g, group) in alloc.iter().enumerate() { - let v = &views[g]; - assert_eq!(v.len(), group.len(), "trial {}, group {}", trial, g); - assert_eq!(v.d(), d); - for (r, &row_idx) in group.iter().enumerate() { - assert_eq!( - v.row(r), - &golden[row_idx][..], - "trial {}, group {}, row {}", - trial, - g, - r - ); - } + use rand::prelude::*; + let mut rng = StdRng::seed_from_u64(7); + for trial in 0..1000 { + let d = rng.random_range(1..10); + let rows = rng.random_range(1000..2000); + let groups = rng.random_range(10..=rows.min(20)); + let mut s = { + let mut result = Square::with_capacity(d, rows); + for _ in 0..rows { + result.push_iter((0..d).map(|_| rng.random_range(-1000.0..1000.0))); + } + result + }; + let golden: Vec> = (0..s.len()).map(|i| s[i].to_vec()).collect(); + let alloc = gen_random_alloc(rows, groups, &mut rng); + let views = partition(s.as_mut_view(), &alloc); + assert_eq!(views.len(), alloc.len(), "trial {}", trial); + + for (g, group) in alloc.iter().enumerate() { + let v = &views[g]; + assert_eq!(v.len(), group.len(), "trial {}, group {}", trial, g); + assert_eq!(v.d(), d); + for (r, &row_idx) in group.iter().enumerate() { + assert_eq!( + v.row(r), + &golden[row_idx][..], + "trial {}, group {}, row {}", + trial, + g, + r + ); } } } diff --git a/crates/k_means/src/index.rs b/crates/k_means/src/index.rs new file mode 100644 index 00000000..612e3365 --- /dev/null +++ b/crates/k_means/src/index.rs @@ -0,0 +1,156 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::This; +use crate::square::{Square, SquareMut}; +use rand::Rng; +use rayon::prelude::*; +use simd::Floating; + +pub fn flat_index(centroids: &Square) -> impl Fn(&[f32]) -> (f32, usize) + Sync { + move |sample| { + let mut result = (f32::INFINITY, 0); + for (i, centroid) in centroids.into_iter().enumerate() { + let dis = f32::reduce_sum_of_d2(sample, centroid); + if dis <= result.0 { + result = (dis, i); + } + } + result + } +} + +pub fn rabitq_index( + pool: &rayon::ThreadPool, + centroids: &Square, +) -> impl Fn(&[f32]) -> (f32, usize) + Sync { + use rabitq::packing::{pack_to_u4, padding_pack}; + let (metadata, blocks) = pool.install(|| { + let metadata = centroids + .par_iter() + .map(rabitq::bit::code_metadata) + .collect::>(); + + let blocks = centroids + .par_iter() + .chunks(32) + .map(|chunk| { + let f = |x: &&_| pack_to_u4(&rabitq::bit::code_elements(x)); + padding_pack(chunk.iter().map(f)) + }) + .collect::>(); + + (metadata, blocks) + }); + move |sample| { + let lut = rabitq::bit::block::preprocess(sample); + let mut result = (f32::INFINITY, 0); + let mut sum = [0u32; 32]; + for (i, centroid) in centroids.into_iter().enumerate() { + if i % 32 == 0 { + sum = rabitq::bit::block::accumulate(&blocks[i / 32], &lut.1); + } + let (rough, err) = rabitq::bit::block::half_process_l2(sum[i % 32], metadata[i], lut.0); + let lowerbound = rough - err * 1.9; + if lowerbound < result.0 { + let dis = f32::reduce_sum_of_d2(sample, centroid); + if dis <= result.0 { + result = (dis, i); + } + } + } + result + } +} + +pub fn update( + this: &mut This<'_>, + samples: &SquareMut<'_>, + targets: &[usize], + centroids: &mut Square, +) { + this.pool.install(|| { + const DELTA: f32 = 9.7656e-4_f32; + + let d = this.d; + let n = samples.len(); + let c = this.c; + + let list = rayon::broadcast({ + |ctx| { + let mut sum = Square::from_zeros(d, c); + let mut count = vec![0.0f32; c]; + for i in (ctx.index()..samples.len()).step_by(ctx.num_threads()) { + let target = targets[i]; + let sample = &samples[i]; + f32::vector_add_inplace(&mut sum[target], sample); + count[target] += 1.0; + } + (sum, count) + } + }); + let mut sum = Square::from_zeros(d, c); + let mut count = vec![0.0f32; c]; + for (sum_1, count_1) in list { + for i in 0..c { + f32::vector_add_inplace(&mut sum[i], &sum_1[i]); + count[i] += count_1[i]; + } + } + + sum.par_iter_mut() + .zip(count.par_iter()) + .for_each(|(sum, count)| f32::vector_mul_scalar_inplace(sum, 1.0 / count)); + + *centroids = sum; + + for i in 0..c { + if count[i] != 0.0f32 { + continue; + } + let mut o = 0; + loop { + let alpha = this.rng.random_range(0.0..1.0f32); + let beta = (count[o] - 1.0) / (n - c) as f32; + if alpha < beta { + break; + } + o = (o + 1) % c; + } + centroids.copy_within(o..o + 1, i); + vector_mul_scalars_inplace(&mut centroids[i], [1.0 + DELTA, 1.0 - DELTA]); + vector_mul_scalars_inplace(&mut centroids[o], [1.0 - DELTA, 1.0 + DELTA]); + count[i] = count[o] / 2.0; + count[o] -= count[i]; + } + + if this.is_spherical { + centroids.into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + }); +} + +fn vector_mul_scalars_inplace(this: &mut [f32], scalars: [f32; 2]) { + let n: usize = this.len(); + for i in 0..n { + if i % 2 == 0 { + this[i] *= scalars[0]; + } else { + this[i] *= scalars[1]; + } + } +} diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index e8c531be..0c8a5f3d 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -12,10 +12,12 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod flat; -pub mod hierarchical; -pub mod quick; -pub mod rabitq; +mod flat; +mod hierarchical; +mod index; +mod quick; +mod rabitq; + pub mod square; use crate::square::{Square, SquareMut}; @@ -26,44 +28,47 @@ pub struct This<'a> { rng: StdRng, d: usize, c: usize, - centroids: Square, - targets: Vec, + is_spherical: bool, } -pub trait KMeans<'a> { - fn this(&mut self) -> &mut This<'a>; +pub trait KMeans { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_>; + fn index(&self) -> Box (f32, usize) + Sync + '_> { + self.prefect_index() + } fn assign(&mut self); fn update(&mut self); - fn sphericalize(&mut self) { - use rayon::iter::{IntoParallelIterator, ParallelIterator}; - use simd::Floating; - let this = self.this(); - this.pool.install(|| { - (&mut this.centroids).into_par_iter().for_each(|centroid| { - let l = f32::reduce_sum_of_x2(centroid).sqrt(); - f32::vector_mul_scalar_inplace(centroid, 1.0 / l); - }); - }); - } - fn index(&mut self) -> Box u32 + '_>; fn finish(self: Box) -> Square; } -pub fn k_means<'a>( +pub fn k_means_lookup(sample: &[f32], centroids: &Square) -> usize { + use simd::Floating; + let mut result = (f32::INFINITY, 0); + for (i, centroid) in centroids.into_iter().enumerate() { + let dis = f32::reduce_sum_of_d2(sample, centroid); + if dis <= result.0 { + result = (dis, i); + } + } + result.1 +} + +pub fn lloyd_k_means<'a>( pool: &'a rayon::ThreadPool, d: usize, samples: SquareMut<'a>, c: usize, seed: [u8; 32], -) -> Box + 'a> { + is_spherical: bool, +) -> Box { assert!(d > 0 && c > 0); let n = samples.len(); if n <= c { - quick::new(pool, d, samples, c, seed) + quick::new(pool, d, samples, c, seed, is_spherical) } else if n <= c * 2 { - flat::new(pool, d, samples, c, seed) + flat::new(pool, d, samples, c, seed, is_spherical) } else { - rabitq::new(pool, d, samples, c, seed) + rabitq::new(pool, d, samples, c, seed, is_spherical) } } @@ -74,7 +79,7 @@ pub fn hierarchical_k_means<'a>( c: usize, seed: [u8; 32], is_spherical: bool, -) -> Box + 'a> { +) -> Box { assert!(d > 0 && c > 0); hierarchical::new(pool, d, samples, c, seed, is_spherical) } diff --git a/crates/k_means/src/quick.rs b/crates/k_means/src/quick.rs index 5572027f..d07db425 100644 --- a/crates/k_means/src/quick.rs +++ b/crates/k_means/src/quick.rs @@ -12,33 +12,28 @@ // // Copyright (c) 2025 TensorChord Inc. +use crate::KMeans; +use crate::index::flat_index as prefect_index; use crate::square::{Square, SquareMut}; -use crate::{KMeans, This}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; -struct Quick<'a> { - this: This<'a>, +struct Quick { + centroids: Square, } -impl<'a> KMeans<'a> for Quick<'a> { - fn this(&mut self) -> &mut This<'a> { - &mut self.this +impl KMeans for Quick { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_> { + Box::new(prefect_index(&self.centroids)) } fn assign(&mut self) {} fn update(&mut self) {} - fn index(&mut self) -> Box u32> { - let this = self.this(); - let centroids = this.centroids.clone(); - Box::new(move |sample| crate::flat::k_means_lookup(sample, ¢roids) as u32) - } - fn finish(self: Box) -> Square { - let this = self.this; - this.centroids + self.centroids } } @@ -48,7 +43,8 @@ pub fn new<'a>( samples: SquareMut<'a>, c: usize, seed: [u8; 32], -) -> Box + 'a> { + is_spherical: bool, +) -> Box { let mut rng = StdRng::from_seed(seed); let mut centroids = Square::with_capacity(d, c); @@ -65,16 +61,15 @@ pub fn new<'a>( centroids.push_iter((0..d).map(|_| rng.random_range(-1.0f32..1.0f32))); } - let targets = vec![0; samples.len()]; + pool.install(|| { + if is_spherical { + use simd::Floating; + (&mut centroids).into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + }); - Box::new(Quick { - this: This { - pool, - rng, - d, - c, - centroids, - targets, - }, - }) + Box::new(Quick { centroids }) } diff --git a/crates/k_means/src/rabitq.rs b/crates/k_means/src/rabitq.rs index c30e5c2d..a7e567e8 100644 --- a/crates/k_means/src/rabitq.rs +++ b/crates/k_means/src/rabitq.rs @@ -12,150 +12,71 @@ // // Copyright (c) 2025 TensorChord Inc. +use crate::index::{flat_index as prefect_index, rabitq_index as index}; use crate::square::{Square, SquareMut}; use crate::{KMeans, This}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use rayon::prelude::*; -use simd::Floating; struct RaBitQ<'a> { this: This<'a>, samples: SquareMut<'a>, + centroids: Square, + targets: Vec, } -impl<'a> KMeans<'a> for RaBitQ<'a> { - fn this(&mut self) -> &mut This<'a> { - &mut self.this +impl<'a> KMeans for RaBitQ<'a> { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_> { + let index = prefect_index(&self.centroids); + Box::new(move |sample| { + let rotated = rabitq::rotate::rotate(sample); + let sample = rotated.as_slice(); + index(sample) + }) + } + + fn index(&self) -> Box (f32, usize) + Sync + '_> { + let index = index(self.this.pool, &self.centroids); + Box::new(move |sample| { + let rotated = rabitq::rotate::rotate(sample); + let sample = rotated.as_slice(); + index(sample) + }) } fn assign(&mut self) { let this = &mut self.this; let samples = &mut self.samples; + let centroids = &self.centroids; + let targets = &mut self.targets; + let index = index(this.pool, centroids); this.pool.install(|| { - use rabitq::packing::{pack_to_u4, padding_pack}; - - let metadata = this - .centroids - .par_iter() - .map(rabitq::bit::code_metadata) - .collect::>(); - - let blocks = this - .centroids - .par_iter() - .chunks(32) - .map(|chunk| { - let f = |x: &&_| pack_to_u4(&rabitq::bit::code_elements(x)); - padding_pack(chunk.iter().map(f)) - }) - .collect::>(); - - this.targets + targets .par_iter_mut() .zip(samples.par_iter_mut()) .for_each(|(target, sample)| { - let lut = rabitq::bit::block::preprocess(sample); - let mut result = (f32::INFINITY, 0); - let mut sum = [0u32; 32]; - for (j, centroid) in this.centroids.into_iter().enumerate() { - if j % 32 == 0 { - sum = rabitq::bit::block::accumulate(&blocks[j / 32], &lut.1); - } - let (rough, err) = - rabitq::bit::block::half_process_l2(sum[j % 32], metadata[j], lut.0); - let lowerbound = rough - err * 1.9; - if lowerbound < result.0 { - let dis_2 = f32::reduce_sum_of_d2(sample, centroid); - if dis_2 <= result.0 { - result = (dis_2, j); - } - } - } - *target = result.1; + *target = index(sample).1; }); }); } fn update(&mut self) { - let this = &mut self.this; - let samples = &mut self.samples; - this.pool.install(|| { - const DELTA: f32 = 9.7656e-4_f32; - - let d = this.d; - let n = samples.len(); - let c = this.c; - - let list = rayon::broadcast({ - |ctx| { - let mut sum = Square::from_zeros(d, c); - let mut count = vec![0.0f32; c]; - for i in (ctx.index()..samples.len()).step_by(ctx.num_threads()) { - let target = this.targets[i]; - let sample = &samples[i]; - f32::vector_add_inplace(&mut sum[target], sample); - count[target] += 1.0; - } - (sum, count) - } - }); - let mut sum = Square::from_zeros(d, c); - let mut count = vec![0.0f32; c]; - for (sum_1, count_1) in list { - for i in 0..c { - f32::vector_add_inplace(&mut sum[i], &sum_1[i]); - count[i] += count_1[i]; - } - } - - sum.par_iter_mut() - .enumerate() - .for_each(|(i, sum)| f32::vector_mul_scalar_inplace(sum, 1.0 / count[i])); - - this.centroids = sum; - - for i in 0..c { - if count[i] != 0.0f32 { - continue; - } - let mut o = 0; - loop { - let alpha = this.rng.random_range(0.0..1.0f32); - let beta = (count[o] - 1.0) / (n - c) as f32; - if alpha < beta { - break; - } - o = (o + 1) % c; - } - this.centroids.copy_within(o..o + 1, i); - vector_mul_scalars_inplace(&mut this.centroids[i], [1.0 + DELTA, 1.0 - DELTA]); - vector_mul_scalars_inplace(&mut this.centroids[o], [1.0 - DELTA, 1.0 + DELTA]); - count[i] = count[o] / 2.0; - count[o] -= count[i]; - } - }); + crate::index::update( + &mut self.this, + &self.samples, + &self.targets, + &mut self.centroids, + ); } - fn index(&mut self) -> Box u32> { - let this = self.this(); - let mut centroids = this.centroids.clone(); - this.pool.install(|| { - centroids.par_iter_mut().for_each(|centroid| { + fn finish(mut self: Box) -> Square { + self.this.pool.install(|| { + self.centroids.par_iter_mut().for_each(|centroid| { rabitq::rotate::rotate_reversed_inplace(centroid); }); }); - Box::new(move |sample| crate::flat::k_means_lookup(sample, ¢roids) as u32) - } - - fn finish(self: Box) -> Square { - let mut this = self.this; - this.pool.install(|| { - this.centroids.par_iter_mut().for_each(|centroid| { - rabitq::rotate::rotate_reversed_inplace(centroid); - }); - }); - this.centroids + self.centroids } } @@ -165,7 +86,8 @@ pub fn new<'a>( mut samples: SquareMut<'a>, c: usize, seed: [u8; 32], -) -> Box + 'a> { + is_spherical: bool, +) -> Box { let mut rng = StdRng::from_seed(seed); pool.install(|| { @@ -188,6 +110,16 @@ pub fn new<'a>( centroids.push_iter((0..d).map(|_| rng.random_range(-1.0f32..1.0f32))); } + pool.install(|| { + if is_spherical { + use simd::Floating; + (&mut centroids).into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + }); + let targets = vec![0; samples.len()]; Box::new(RaBitQ { @@ -195,21 +127,11 @@ pub fn new<'a>( pool, d, c, - centroids, - targets, rng, + is_spherical, }, samples, + centroids, + targets, }) } - -fn vector_mul_scalars_inplace(this: &mut [f32], scalars: [f32; 2]) { - let n: usize = this.len(); - for i in 0..n { - if i % 2 == 0 { - this[i] *= scalars[0]; - } else { - this[i] *= scalars[1]; - } - } -} diff --git a/crates/k_means/src/square.rs b/crates/k_means/src/square.rs index 0309fd3c..3fa52008 100644 --- a/crates/k_means/src/square.rs +++ b/crates/k_means/src/square.rs @@ -127,6 +127,7 @@ pub struct SquareMut<'a> { d: usize, p: &'a mut [f32], } + impl<'a> SquareMut<'a> { #[inline] pub fn d(&self) -> usize { @@ -192,6 +193,16 @@ impl std::ops::IndexMut for SquareMut<'_> { } } +impl<'a> IntoIterator for &'a mut SquareMut<'a> { + type Item = &'a mut [f32]; + + type IntoIter = std::slice::ChunksExactMut<'a, f32>; + + fn into_iter(self) -> Self::IntoIter { + self.p.chunks_exact_mut(self.d) + } +} + impl<'a> rayon::prelude::IntoParallelIterator for &'a mut SquareMut<'a> { type Item = &'a mut [f32]; diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 9a03cbf3..a6792260 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -25,7 +25,7 @@ use crate::index::vchordrq::types::*; use index::relation::{ Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, }; -use k_means::flat::k_means_lookup; +use k_means::k_means_lookup; use k_means::square::Square; use simd::Floating; use std::ffi::CStr; @@ -1276,23 +1276,24 @@ fn make_internal_build( use humansize::{BINARY, format_size}; use std::iter::once; let (reduction, sample_dim) = match internal_build.kmeans_dimension_reduction { - None => (false, vector_options.dims as usize), - Some(d) if d < vector_options.dims => (true, d as usize), + None => (None, vector_options.dims as usize), + Some(d) if d < vector_options.dims => (Some(d as usize), d as usize), Some(d) => { pgrx::warning!( "ignoring `kmeans_dimension_reduction = {}` because it is less than the vector dimension {}", d, vector_options.dims ); - (false, vector_options.dims as usize) + (None, vector_options.dims as usize) } }; { - let d = sample_dim as u64; + let d = vector_options.dims as u64; + let d_ = sample_dim as u64; let c = internal_build.lists.last().copied().unwrap_or_default() as u64; let f = internal_build.sampling_factor as u64; let t = internal_build.build_threads as u64; - let estimated_memory_usage = 4 * c * d * (1 + f + t); + let estimated_memory_usage = 4 * c * d_ * f + 4 * c * d * (1 + t); pgrx::info!( "clustering: estimated memory usage is {}", format_size( @@ -1301,16 +1302,16 @@ fn make_internal_build( ) ); } + let max_number_of_samples = internal_build + .lists + .last() + .map(|x| x.saturating_mul(internal_build.sampling_factor)) + .unwrap_or_default(); let mut sample = sampler.sample(); - let samples = 'a: { - let max_number_of_samples = internal_build - .lists - .last() - .map(|x| x.saturating_mul(internal_build.sampling_factor)) - .unwrap_or_default(); + let samples = 'samples: { let mut samples = Square::with_capacity(sample_dim, max_number_of_samples as _); if samples.len() >= max_number_of_samples as usize { - break 'a samples; + break 'samples samples; } while let Some(mut tuple) = sample.next() { let (values, is_nulls) = tuple.build(); @@ -1328,13 +1329,13 @@ fn make_internal_build( x.len() as u32, "invalid vector dimensions" ); - if reduction { + if let Some(sample_dim) = reduction { rabitq::rotate::rotate_inplace(&mut x); x.truncate(sample_dim); } samples.push_slice(x.as_slice()); if samples.len() >= max_number_of_samples as usize { - break 'a samples; + break 'samples samples; } } } @@ -1342,6 +1343,7 @@ fn make_internal_build( } samples }; + drop(sample); let pool = rayon::ThreadPoolBuilder::new() .num_threads(internal_build.build_threads as usize) .build() @@ -1378,26 +1380,37 @@ fn make_internal_build( "clustering: starting, clustering {num_points} vectors of {num_dims} dimension into {num_lists} clusters, in {num_iterations} iterations" ); } - let mut f = { + let mut f = 'f: { let view = input.as_mut_view(); - if result.last().is_none() - && let KMeansAlgorithm::Hierarchical {} = internal_build.kmeans_algorithm - { - k_means::hierarchical_k_means( + if result.last().is_some() { + break 'f k_means::lloyd_k_means( &pool, num_dims, view, num_lists, [7; 32], internal_build.spherical_centroids, - ) - } else { - k_means::k_means(&pool, num_dims, view, num_lists, [7; 32]) + ); + }; + match internal_build.kmeans_algorithm { + KMeansAlgorithm::Lloyd {} => k_means::lloyd_k_means( + &pool, + num_dims, + view, + num_lists, + [7; 32], + internal_build.spherical_centroids, + ), + KMeansAlgorithm::Hierarchical {} => k_means::hierarchical_k_means( + &pool, + num_dims, + view, + num_lists, + [7; 32], + internal_build.spherical_centroids, + ), } }; - if internal_build.spherical_centroids { - f.sphericalize(); - } for i in 0..num_iterations { pgrx::check_for_interrupts!(); if result.is_empty() { @@ -1413,47 +1426,115 @@ fn make_internal_build( } f.assign(); f.update(); - if internal_build.spherical_centroids { - f.sphericalize(); - } } - let centroids = if reduction && result.last().is_none() { - let mut sample = sampler.sample(); - let mut next_sample = || -> Option<(Vec, Vec)> { - let mut tuple = sample.next()?; - let (values, is_nulls) = tuple.build(); - let datum = (!is_nulls[0]).then_some(values[0]); - if let Some(datum) = datum { - let vectors = unsafe { opfamily.store(datum) }; - if let Some(vectors) = vectors { - if let Some((vector, _)) = vectors.into_iter().next() { - let x = match vector { - OwnedVector::Vecf32(x) => VectOwned::normalize(x), - OwnedVector::Vecf16(x) => VectOwned::normalize(x), - }; - let mut y = rabitq::rotate::rotate(&x); - y.truncate(sample_dim); - return Some((x, y)); + let centroids = 'centroids: { + if result.last().is_some() { + break 'centroids f.finish(); + } + if let Some(sample_dim) = reduction { + let mut sample = sampler.sample(); + let index = f.index(); + let index = &index; + let is_spherical = internal_build.spherical_centroids; + let num_threads = internal_build.build_threads as usize; + let (tx, rx) = crossbeam_channel::bounded::>(1024); + let list = std::thread::scope(move |scope| { + assert_ne!(num_threads, 0); + let mut handles = Vec::with_capacity(num_threads); + for _ in 0..num_threads { + let rx = rx.clone(); + let mut sum = Square::from_zeros(vector_options.dims as _, num_lists); + let mut count = vec![0.0f32; num_lists]; + handles.push(scope.spawn(move || { + while let Ok(sample) = rx.recv() { + let mut x = rabitq::rotate::rotate(&sample); + x.truncate(sample_dim); + let (_, id) = index(&x); + f32::vector_add_inplace(&mut sum[id], &sample); + count[id] += 1.0; + } + (sum, count) + })); + } + { + drop(rx); + let tx = tx; + let mut samples = 0_usize; + 'samples: { + if samples >= max_number_of_samples as usize { + break 'samples; + } + while let Some(mut tuple) = sample.next() { + let (values, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(values[0]); + if let Some(datum) = datum { + let vectors = unsafe { opfamily.store(datum) }; + if let Some(vectors) = vectors { + for (vector, _) in vectors { + let x = match vector { + OwnedVector::Vecf32(x) => VectOwned::normalize(x), + OwnedVector::Vecf16(x) => VectOwned::normalize(x), + }; + assert_eq!( + vector_options.dims, + x.len() as u32, + "invalid vector dimensions" + ); + if tx.send(x).is_err() { + break 'samples; + } + samples += 1; + if samples >= max_number_of_samples as usize { + break 'samples; + } + } + } + } + } } - None - } else { - None } - } else { - None + handles + .into_iter() + .map(|handle| handle.join().expect("failed to spawn threads")) + .collect::>() + }); + let mut sum = Square::from_zeros(vector_options.dims as _, num_lists); + let mut count = vec![0.0f32; num_lists]; + for (sum_1, count_1) in list { + for i in 0..num_lists { + f32::vector_add_inplace(&mut sum[i], &sum_1[i]); + count[i] += count_1[i]; + } } - }; - k_means_dimension_restore( - vector_options.dims as usize, - num_lists, - num_points, - internal_build.spherical_centroids, - [7; 32], - f.index(), - &mut next_sample, - ) - } else { - f.finish() + pool.install(|| { + use rayon::prelude::*; + + sum.par_iter_mut() + .zip(count.par_iter()) + .for_each(|(sum, count)| f32::vector_mul_scalar_inplace(sum, 1.0 / count)); + + let mut centroids = sum; + + centroids + .par_iter_mut() + .zip(count.par_iter()) + .filter(|&(_, &count)| count == 0.0) + .for_each(|(centroid, _)| { + centroid.fill_with(|| rand::random_range(-1.0..=1.0)) + }); + + if is_spherical { + (&mut centroids).into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + + centroids + }) + } else { + f.finish() + } }; if result.is_empty() { @@ -1491,45 +1572,6 @@ fn make_internal_build( result } -fn k_means_dimension_restore<'a>( - d: usize, - c: usize, - num_points: usize, - is_spherical: bool, - seed: [u8; 32], - index: Box u32 + 'a>, - mut next_sample: impl FnMut() -> Option<(Vec, Vec)>, -) -> Square { - use rand::rngs::StdRng; - use rand::{Rng, SeedableRng}; - let mut centroids = Square::from_zeros(d, c); - let mut count = vec![0_u32; c]; - let mut traveled = 0; - let mut rng = StdRng::from_seed(seed); - while let Some((original, reduction)) = next_sample() - && traveled < num_points - { - let cid = index(&reduction) as usize; - f32::vector_add_inplace(&mut centroids[cid], &original); - count[cid] += 1; - traveled += 1; - } - let mut ret = Square::new(d); - for (i, centroid) in centroids.as_mut_view().iter_mut().enumerate() { - if count[i] == 0 { - centroid.fill_with(|| rng.random_range(-1.0f32..1.0f32)); - } else { - f32::vector_mul_scalar_inplace(centroid, 1.0 / count[i] as f32); - } - if is_spherical { - let l = f32::reduce_sum_of_x2(centroid).sqrt(); - f32::vector_mul_scalar_inplace(centroid, 1.0 / l); - } - ret.push_slice(centroid); - } - ret -} - #[allow(clippy::collapsible_else_if)] fn make_external_build( vector_options: VectorOptions, diff --git a/tests/vchordrq/internal_build_kmeans.slt b/tests/vchordrq/internal_build_kmeans.slt new file mode 100644 index 00000000..5c5a7518 --- /dev/null +++ b/tests/vchordrq/internal_build_kmeans.slt @@ -0,0 +1,19 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +build.internal.lists = [1000] +build.internal.kmeans_algorithm.lloyd = {} +$$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +build.internal.lists = [1000] +build.internal.kmeans_algorithm.hierarchical = {} +$$); + +statement ok +DROP TABLE t; From 15c6a5c923a23ba3f66a7539e2726460052b708f Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 6 Nov 2025 19:34:08 +0800 Subject: [PATCH 246/324] chore: use less threads for dimension expansion (#373) to consume less memory, and keep consistent with https://github.com/tensorchord/pgvecto.rs-docs/pull/151 Signed-off-by: usamoi --- src/index/vchordrq/am/am_build.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index a6792260..0fdae017 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -1288,12 +1288,11 @@ fn make_internal_build( } }; { - let d = vector_options.dims as u64; - let d_ = sample_dim as u64; + let d = sample_dim as u64; let c = internal_build.lists.last().copied().unwrap_or_default() as u64; let f = internal_build.sampling_factor as u64; let t = internal_build.build_threads as u64; - let estimated_memory_usage = 4 * c * d_ * f + 4 * c * d * (1 + t); + let estimated_memory_usage = 4 * c * d * (1 + t + f); pgrx::info!( "clustering: estimated memory usage is {}", format_size( @@ -1436,7 +1435,11 @@ fn make_internal_build( let index = f.index(); let index = &index; let is_spherical = internal_build.spherical_centroids; - let num_threads = internal_build.build_threads as usize; + let num_threads = { + let config = internal_build.build_threads as f64; + let ratio = sample_dim as f64 / vector_options.dims as f64; + (config * ratio).ceil().max(1.0).min(config) as usize + }; let (tx, rx) = crossbeam_channel::bounded::>(1024); let list = std::thread::scope(move |scope| { assert_ne!(num_threads, 0); From 28cac71b66c87d2b6d14276fc35d7ee473d7c4b4 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 7 Nov 2025 14:11:28 +0800 Subject: [PATCH 247/324] fix: use `tempfile` in xtask (#374) This can prevent permission issues when multiple users are present. Signed-off-by: usamoi --- Cargo.lock | 48 ++++++++++++++++++++++++++++------------ crates/simd/Cargo.toml | 2 +- crates/xtask/Cargo.toml | 3 ++- crates/xtask/src/main.rs | 28 ++++++++++------------- 4 files changed, 49 insertions(+), 32 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fd441118..5398538b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -148,9 +148,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.43" +version = "1.2.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "739eb0f94557554b3ca9a86d2d37bebd49c5e6d0c1d2bda35ba5bdac830befc2" +checksum = "37521ac7aabe3d13122dc382493e20c9416f299d2ccd5b3a5340a2570cdeb0f3" dependencies = [ "find-msvc-tools", "shlex", @@ -194,9 +194,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.50" +version = "4.5.51" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c2cfd7bf8a6017ddaa4e32ffe7403d547790db06bd171c1c53926faab501623" +checksum = "4c26d721170e0295f191a69bd9a1f93efcdb0aff38684b61ab5750468972e5f5" dependencies = [ "clap_builder", "clap_derive", @@ -204,9 +204,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.50" +version = "4.5.51" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0a4c05b9e80c5ccd3a7ef080ad7b6ba7d6fc00a985b8b157197075677c82c7a0" +checksum = "75835f0c7bf681bfd05abe44e965760fea999a5286c6eb2d59883634fd02011a" dependencies = [ "anstream", "anstyle", @@ -443,6 +443,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "fastrand" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" + [[package]] name = "feistel" version = "0.0.0" @@ -1119,9 +1125,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.41" +version = "1.0.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce25767e7b499d1b604768e7cde645d14cc8584231ea6b295e9c9eb22c02e1d1" +checksum = "a338cc41d27e6cc6dce6cefc13a0729dfbb81c262b1f519331575dd80ef3067f" dependencies = [ "proc-macro2", ] @@ -1267,9 +1273,9 @@ checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] name = "ruzstd" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3640bec8aad418d7d03c72ea2de10d5c646a598f9883c7babc160d91e3c1b26c" +checksum = "e5ff0cc5e135c8870a775d3320910cd9b564ec036b4dc0b8741629020be63f01" dependencies = [ "twox-hash", ] @@ -1430,9 +1436,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.108" +version = "2.0.109" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da58917d35242480a05c2897064da0a80589a2a0476c9a3f2fdc83b53502e917" +checksum = "2f17c7e013e88258aa9543dcbe81aca68a667a9ac37cd69c9fbc07858bfe0e2f" dependencies = [ "proc-macro2", "quote", @@ -1462,6 +1468,19 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "591ef38edfb78ca4771ee32cf494cb8771944bee237a9b91fc9c1424ac4b777b" +[[package]] +name = "tempfile" +version = "3.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d31c77bdf42a745371d260a26ca7163f1e0924b64afa0b688e61b5a9fa02f16" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "2.0.17" @@ -1545,9 +1564,9 @@ checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] name = "unicode-ident" -version = "1.0.20" +version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "462eeb75aeb73aea900253ce739c8e18a67423fadf006037cd3ff27e82748a06" +checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" [[package]] name = "unicode-segmentation" @@ -1965,6 +1984,7 @@ dependencies = [ "serde_json", "shlex", "target-triple", + "tempfile", ] [[package]] diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 72baa9bb..181661b1 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -19,7 +19,7 @@ zerocopy.workspace = true rand.workspace = true [build-dependencies] -cc = "1.2.43" +cc = "1.2.44" which = "8.0.0" [lints] diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index d334d24e..3e03f42f 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -5,12 +5,13 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.50", features = ["derive", "env"] } +clap = { version = "4.5.51", features = ["derive", "env"] } object = { version = "0.37.3", features = ["read", "wasm"] } serde.workspace = true serde_json = "1.0.145" shlex = "1.3.0" target-triple = "1.0.0" +tempfile = "3.23.0" [lints] workspace = true diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index cfbf936f..59f1d11a 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -16,9 +16,7 @@ use clap::{Args, Parser, Subcommand}; use object::{Object, ObjectSymbol}; use serde::Deserialize; use std::collections::{HashMap, HashSet}; -use std::env::var_os; use std::error::Error; -use std::fs::read_dir; use std::path::{Path, PathBuf}; use std::process::{Command, Stdio}; @@ -296,16 +294,14 @@ fn generate( } else { Vec::new() }; - let pgrx_embed = std::env::temp_dir().join("VCHORD_PGRX_EMBED"); - eprintln!("Writing {pgrx_embed:?}"); - std::fs::write( - &pgrx_embed, - format!( - "crate::schema_generation!({}; {});", - exports.join(" "), - imports.join(" ") - ), - )?; + let mut pgrx_embed = tempfile::NamedTempFile::new()?; + eprintln!("Writing {:?}", pgrx_embed.path()); + let contents = format!( + "crate::schema_generation!({}; {});", + exports.join(" "), + imports.join(" ") + ); + std::io::Write::write_all(pgrx_embed.as_file_mut(), contents.as_bytes())?; let mut command = Command::new("cargo"); command .args(["rustc"]) @@ -315,7 +311,7 @@ fn generate( .args(["--features", pg_version]) .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) .args(["--", "-C", "lto=off", "--cfg", "pgrx_embed"]) - .env("PGRX_EMBED", &pgrx_embed) + .env("PGRX_EMBED", pgrx_embed.path()) .stdout(Stdio::inherit()) .stderr(Stdio::inherit()); eprintln!("Running {command:?}"); @@ -428,7 +424,7 @@ fn main() -> Result<(), Box> { } let vchord_version = control_file("./vchord.control")?["default_version"].clone(); let runner = runner.and_then(|runner| shlex::split(&runner)); - let path = if let Some(value) = var_os("PG_CONFIG") { + let path = if let Some(value) = std::env::var_os("PG_CONFIG") { PathBuf::from(value) } else { return Err("Environment variable `PG_CONFIG` is not set.".into()); @@ -480,7 +476,7 @@ fn main() -> Result<(), Box> { false, )?; if vchord_version != "0.0.0" { - for e in read_dir("./sql/upgrade")?.collect::, _>>()? { + for e in std::fs::read_dir("./sql/upgrade")?.collect::, _>>()? { install_by_copying( e.path(), format!("{sharedir}/extension/{}", e.file_name().display()), @@ -515,7 +511,7 @@ fn main() -> Result<(), Box> { if !std::fs::exists("./vchord.control")? { return Err("The script must be run from the VectorChord source directory.".into()); } - let path = if let Some(value) = var_os("PG_CONFIG") { + let path = if let Some(value) = std::env::var_os("PG_CONFIG") { PathBuf::from(value) } else { return Err("Environment variable `PG_CONFIG` is not set.".into()); From e03ad476911191bf9e95a6b01e9b241ad98af511 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 7 Nov 2025 16:24:14 +0800 Subject: [PATCH 248/324] chore: upload 1.0.0 schema scripts (#375) Signed-off-by: usamoi --- crates/vchordg/src/tuples.rs | 2 +- crates/vchordrq/src/tuples.rs | 2 +- sql/install/vchord--1.0.0.sql | 857 +++++++++++++++++++++++++++ sql/upgrade/vchord--0.5.3--1.0.0.sql | 1 + 4 files changed, 860 insertions(+), 2 deletions(-) create mode 100644 sql/install/vchord--1.0.0.sql create mode 100644 sql/upgrade/vchord--0.5.3--1.0.0.sql diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 99f3f9d4..1555e4c9 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -21,7 +21,7 @@ use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordg\0"); -const VERSION: u64 = 10; +const VERSION: u64 = 1000; #[inline(always)] fn tag(source: &[u8]) -> Tag { diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index 36b16427..29a6ea7a 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -20,7 +20,7 @@ use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 12; +const VERSION: u64 = 1000; #[inline(always)] fn tag(source: &[u8]) -> Tag { diff --git a/sql/install/vchord--1.0.0.sql b/sql/install/vchord--1.0.0.sql new file mode 100644 index 00000000..e50ca725 --- /dev/null +++ b/sql/install/vchord--1.0.0.sql @@ -0,0 +1,857 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:45 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ + "rhs" halfvec[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:31 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:58 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:76 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ + "rhs" vector[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:21 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:78 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:81 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:46 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/upgrade/vchord--0.5.3--1.0.0.sql b/sql/upgrade/vchord--0.5.3--1.0.0.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.5.3--1.0.0.sql @@ -0,0 +1 @@ +/* nothing to do */ From 83649d9734f8d88196227f09b2c8461848f61722 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 7 Nov 2025 17:27:11 +0800 Subject: [PATCH 249/324] chore: rename kmeans_dimension_reduction to kmeans_dimension (#376) Signed-off-by: usamoi --- src/index/vchordrq/am/am_build.rs | 4 ++-- src/index/vchordrq/types.rs | 8 ++++---- tests/vchordrq/approximate.slt | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 0fdae017..0c016eb4 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -1275,12 +1275,12 @@ fn make_internal_build( ) -> Vec> { use humansize::{BINARY, format_size}; use std::iter::once; - let (reduction, sample_dim) = match internal_build.kmeans_dimension_reduction { + let (reduction, sample_dim) = match internal_build.kmeans_dimension { None => (None, vector_options.dims as usize), Some(d) if d < vector_options.dims => (Some(d as usize), d as usize), Some(d) => { pgrx::warning!( - "ignoring `kmeans_dimension_reduction = {}` because it is less than the vector dimension {}", + "ignoring `kmeans_dimension = {}` because it is less than the vector dimension {}", d, vector_options.dims ); diff --git a/src/index/vchordrq/types.rs b/src/index/vchordrq/types.rs index 897e0b53..f16f5bec 100644 --- a/src/index/vchordrq/types.rs +++ b/src/index/vchordrq/types.rs @@ -54,9 +54,9 @@ pub struct VchordrqInternalBuildOptions { pub build_threads: u16, #[serde(default = "VchordrqInternalBuildOptions::default_kmeans_algorithm")] pub kmeans_algorithm: KMeansAlgorithm, - #[serde(default = "VchordrqInternalBuildOptions::default_kmeans_dimension_reduction")] + #[serde(default = "VchordrqInternalBuildOptions::default_kmeans_dimension")] #[validate(range(min = 1, max = 16000))] - pub kmeans_dimension_reduction: Option, + pub kmeans_dimension: Option, } impl VchordrqInternalBuildOptions { @@ -87,7 +87,7 @@ impl VchordrqInternalBuildOptions { fn default_kmeans_algorithm() -> KMeansAlgorithm { KMeansAlgorithm::Lloyd {} } - fn default_kmeans_dimension_reduction() -> Option { + fn default_kmeans_dimension() -> Option { None } } @@ -101,7 +101,7 @@ impl Default for VchordrqInternalBuildOptions { kmeans_iterations: Self::default_kmeans_iterations(), build_threads: Self::default_build_threads(), kmeans_algorithm: Self::default_kmeans_algorithm(), - kmeans_dimension_reduction: Self::default_kmeans_dimension_reduction(), + kmeans_dimension: Self::default_kmeans_dimension(), } } } diff --git a/tests/vchordrq/approximate.slt b/tests/vchordrq/approximate.slt index afb8715e..c19f6bc6 100644 --- a/tests/vchordrq/approximate.slt +++ b/tests/vchordrq/approximate.slt @@ -11,7 +11,7 @@ residual_quantization = true [build.internal] lists = [33] spherical_centroids = false -kmeans_dimension_reduction = 2 +kmeans_dimension = 2 kmeans_algorithm.hierarchical = {} $$); From 5af9a3eb92ca235bebd3abe88cea38e5ab8761b1 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 12 Nov 2025 12:34:15 +0800 Subject: [PATCH 250/324] readme: sync with docs (#377) Signed-off-by: usamoi --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 7a398145..fa7f0c54 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: **⚡ Enhanced Performance**: Delivering optimized operations with up to 5x faster queries, 16x higher insert throughput, and 16x quicker[^1] index building compared to pgvector's HNSW implementation. -[^1]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/) with 768-dimensional vectors and 95% recall. Please checkout our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. +[^1]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/) with 768-dimensional vectors and 95% recall. Please check out our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. **💰 Affordable Vector Search**: Query 100M 768-dimensional vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%, helping you keep infrastructure costs down while maintaining high search quality. @@ -67,7 +67,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg18-v0.5.3 + -d ghcr.io/tensorchord/vchord-postgres:pg18-v1.0.0 ``` > [!NOTE] > In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). @@ -117,7 +117,7 @@ For more usage, please read: - [Prewarm](https://docs.vectorchord.ai/vectorchord/usage/prewarm.html) - [Prefilter](https://docs.vectorchord.ai/vectorchord/usage/prefilter.html) - [Prefetch](https://docs.vectorchord.ai/vectorchord/usage/prefetch.html) -- [Rerank In Table](https://docs.vectorchord.ai/vectorchord/usage/rerank-in-table.html) +- [Rerank in Table](https://docs.vectorchord.ai/vectorchord/usage/rerank-in-table.html) - [External Build](https://docs.vectorchord.ai/vectorchord/usage/external-index-precomputation.html) ## License From 3ed1b92d89822942fe3c4050e17b1dab1c71af98 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 24 Nov 2025 15:25:46 +0800 Subject: [PATCH 251/324] chore: strip symbols in profile release (#383) It reduces the size of the release binary by 16.5% (with `-Csymbol-mangling-version=v0`). It prevent users from getting the error's backtrace. PostgreSQL errors do not show a backtrace by default, so it is not an issue. Signed-off-by: usamoi --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 2c776853..1e952a64 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -113,7 +113,7 @@ codegen-units = 1 lto = "fat" opt-level = 3 debug = "none" -strip = "debuginfo" +strip = "symbols" [profile.prof] inherits = "release" From dd98adf4647b559caa55ce0c1150a80b65af55a8 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 25 Nov 2025 21:59:03 +0800 Subject: [PATCH 252/324] ci: fix miri test (#384) https://github.com/rust-lang/rust/issues/149314 job: -psql job: +miri Signed-off-by: usamoi --- .github/workflows/check.yml | 9 ++++++--- crates/k_means/src/hierarchical.rs | 2 +- crates/simd_macros/src/lib.rs | 2 +- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index b88a4be8..292fbfe4 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -128,11 +128,14 @@ jobs: uses: actions/checkout@v4 - name: Cargo Test (Miri) + env: + # https://github.com/rust-lang/rust/issues/149314 + RUSTFLAGS: "-Ctarget-cpu=sapphirerapids -Zcodegen-backend=llvm" + MIRIFLAGS: "-Zmiri-strict-provenance" run: | - RUSTFLAGS="-Ctarget-cpu=sapphirerapids" MIRIFLAGS="-Zmiri-strict-provenance" \ - cargo miri test --locked --target x86_64-unknown-linux-gnu \ + cargo miri test --locked --target x86_64-unknown-linux-gnu \ --workspace --exclude vchord --no-fail-fast \ - -- --no-capture + -- --no-capture --test-threads=1 lint: if: | diff --git a/crates/k_means/src/hierarchical.rs b/crates/k_means/src/hierarchical.rs index aa241f38..180e0f08 100644 --- a/crates/k_means/src/hierarchical.rs +++ b/crates/k_means/src/hierarchical.rs @@ -286,7 +286,7 @@ fn test_partition() { use rand::prelude::*; let mut rng = StdRng::seed_from_u64(7); - for trial in 0..1000 { + for trial in 0..if cfg!(not(miri)) { 1000 } else { 1 } { let d = rng.random_range(1..10); let rows = rng.random_range(1000..2000); let groups = rng.random_range(10..=rows.min(20)); diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs index 6b5b67d1..094a5066 100644 --- a/crates/simd_macros/src/lib.rs +++ b/crates/simd_macros/src/lib.rs @@ -154,7 +154,7 @@ pub fn multiversion( fn cold() -> fn(#inputs) #output { #cold } - #[cfg(feature = "init")] + #[cfg(all(feature = "init", not(miri)))] #[cfg(all(target_os = "linux", target_env = "gnu"))] { #[used] From d8e880ce7de4629a60b4843df700b3f56721a540 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 3 Dec 2025 17:04:21 +0800 Subject: [PATCH 253/324] feat: rabitq8 (#386) It removes undocumented `scalar8`. Signed-off-by: usamoi --- Cargo.lock | 2 + crates/index/Cargo.toml | 1 + crates/index/src/accessor.rs | 63 ++++ crates/rabitq/Cargo.toml | 2 + crates/rabitq/src/bits.rs | 12 + crates/rabitq/src/byte.rs | 21 ++ crates/rabitq/src/extended.rs | 132 +++++++- crates/vchordg/src/operator.rs | 137 +++++++- crates/vchordg/src/types.rs | 5 + crates/vchordrq/src/operator.rs | 210 +++++++++++- crates/vchordrq/src/types.rs | 5 + crates/vector/Cargo.toml | 1 + crates/vector/src/bvect.rs | 39 --- crates/vector/src/lib.rs | 6 +- crates/vector/src/rabitq8.rs | 320 ++++++++++++++++++ crates/vector/src/scalar8.rs | 305 ----------------- crates/vector/src/svect.rs | 34 -- crates/vector/src/vect.rs | 17 - .../{binary_scalar8.rs => binary_rabitq8.rs} | 24 +- src/datatype/functions_rabitq8.rs | 47 +++ src/datatype/functions_scalar8.rs | 39 --- .../{memory_scalar8.rs => memory_rabitq8.rs} | 104 +++--- src/datatype/mod.rs | 11 +- ...rators_scalar8.rs => operators_rabitq8.rs} | 65 ++-- .../{text_scalar8.rs => text_rabitq8.rs} | 24 +- src/datatype/typmod.rs | 44 +-- src/datatype/typmod_rabitq8.rs | 37 ++ src/index/opclass.rs | 35 ++ src/index/vchordg/am/am_build.rs | 6 +- src/index/vchordg/am/mod.rs | 7 +- src/index/vchordg/dispatch.rs | 77 ++++- src/index/vchordg/opclass.rs | 42 ++- src/index/vchordg/scanners/default.rs | 227 ++++++++++++- src/index/vchordrq/am/am_build.rs | 13 +- src/index/vchordrq/am/mod.rs | 10 +- src/index/vchordrq/build.rs | 28 +- src/index/vchordrq/dispatch.rs | 77 +++++ src/index/vchordrq/opclass.rs | 71 +++- src/index/vchordrq/scanners/default.rs | 287 +++++++++++++++- src/index/vchordrq/scanners/maxsim.rs | 143 +++++++- src/recorder/mod.rs | 4 +- src/recorder/text.rs | 23 ++ src/sql/bootstrap.sql | 4 +- src/sql/finalize.sql | 130 +++++-- tests/general/distance.slt | 30 -- tests/vchordg/index_vector.slt | 255 ++++++++++++++ tests/vchordrq/index_multivector.slt | 104 ++++++ tests/vchordrq/index_vector.slt | 255 ++++++++++++++ 48 files changed, 2825 insertions(+), 710 deletions(-) create mode 100644 crates/vector/src/rabitq8.rs delete mode 100644 crates/vector/src/scalar8.rs rename src/datatype/{binary_scalar8.rs => binary_rabitq8.rs} (83%) create mode 100644 src/datatype/functions_rabitq8.rs delete mode 100644 src/datatype/functions_scalar8.rs rename src/datatype/{memory_scalar8.rs => memory_rabitq8.rs} (68%) rename src/datatype/{operators_scalar8.rs => operators_rabitq8.rs} (63%) rename src/datatype/{text_scalar8.rs => text_rabitq8.rs} (87%) create mode 100644 src/datatype/typmod_rabitq8.rs create mode 100644 tests/vchordg/index_vector.slt create mode 100644 tests/vchordrq/index_multivector.slt create mode 100644 tests/vchordrq/index_vector.slt diff --git a/Cargo.lock b/Cargo.lock index 5398538b..525c9b2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -702,6 +702,7 @@ dependencies = [ "bumpalo", "dary_heap", "distance", + "rabitq", "simd", "small_iter", "vector", @@ -1725,6 +1726,7 @@ name = "vector" version = "0.0.0" dependencies = [ "distance", + "rabitq", "simd", ] diff --git a/crates/index/Cargo.toml b/crates/index/Cargo.toml index f6fe282e..deb0d497 100644 --- a/crates/index/Cargo.toml +++ b/crates/index/Cargo.toml @@ -7,6 +7,7 @@ publish = false [dependencies] always_equal = { path = "../always_equal" } distance = { path = "../distance" } +rabitq = { path = "../rabitq" } simd = { path = "../simd" } small_iter = { path = "../small_iter" } vector = { path = "../vector" } diff --git a/crates/index/src/accessor.rs b/crates/index/src/accessor.rs index d411eeb5..0d1a9eb1 100644 --- a/crates/index/src/accessor.rs +++ b/crates/index/src/accessor.rs @@ -13,8 +13,10 @@ // Copyright (c) 2025 TensorChord Inc. use distance::Distance; +use rabitq::byte::CodeMetadata; use simd::{Floating, f16}; use std::marker::PhantomData; +use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; #[derive(Debug, Clone, Copy)] @@ -385,3 +387,64 @@ impl Accessor2 for DistanceAccessor, Dot> { Distance::from_f32(-self.0) } } + +#[derive(Debug)] +pub struct DimensionDistanceAccessor( + u32, + u32, + PhantomData V>, + PhantomData D>, +); + +impl Default for DimensionDistanceAccessor { + #[inline(always)] + fn default() -> Self { + Self(0, 0, PhantomData, PhantomData) + } +} + +impl Accessor2 for DimensionDistanceAccessor { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[u8], input: &[u8]) { + self.0 += target.len() as u32; + self.1 += simd::u8::reduce_sum_of_x_as_u32_y_as_u32(target, input); + } + + #[inline(always)] + fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { + Distance::from_f32( + rabitq::byte::binary::half_process_l2( + self.0, + self.1, + CodeMetadata::from_array(std::array::from_fn(|i| target[i])), + CodeMetadata::from_array(std::array::from_fn(|i| input[i])), + ) + .0, + ) + } +} + +impl Accessor2 for DimensionDistanceAccessor { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[u8], input: &[u8]) { + self.0 += target.len() as u32; + self.1 += simd::u8::reduce_sum_of_x_as_u32_y_as_u32(target, input); + } + + #[inline(always)] + fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { + Distance::from_f32( + rabitq::byte::binary::half_process_dot( + self.0, + self.1, + CodeMetadata::from_array(std::array::from_fn(|i| target[i])), + CodeMetadata::from_array(std::array::from_fn(|i| input[i])), + ) + .0, + ) + } +} diff --git a/crates/rabitq/Cargo.toml b/crates/rabitq/Cargo.toml index 72e65fc6..dc6e6c91 100644 --- a/crates/rabitq/Cargo.toml +++ b/crates/rabitq/Cargo.toml @@ -7,6 +7,8 @@ publish = false [dependencies] simd = { path = "../simd" } +rand.workspace = true +rand_chacha.workspace = true zerocopy.workspace = true [build-dependencies] diff --git a/crates/rabitq/src/bits.rs b/crates/rabitq/src/bits.rs index 47ab0eaa..66bb62f8 100644 --- a/crates/rabitq/src/bits.rs +++ b/crates/rabitq/src/bits.rs @@ -40,6 +40,13 @@ pub fn code(bits: Bits, vector: &[f32]) -> Code { } } +pub fn ugly_code(bits: Bits, vector: &[f32]) -> Code { + match bits { + Bits::_1 => crate::extended::ugly_code::<1>(vector), + Bits::_2 => crate::extended::ugly_code::<2>(vector), + } +} + pub fn pack_code(bits: Bits, input: &[u8]) -> Vec { match bits { Bits::_1 => crate::extended::pack_code::<1>(input) @@ -68,6 +75,11 @@ pub mod binary { (metadata, crate::extended::pack_code::(&elements)) } + pub fn ugly_preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::ugly_code::(vector); + (metadata, crate::extended::pack_code::(&elements)) + } + pub fn accumulate(bits: Bits, lhs: &[u64], rhs: &[Vec; BITS]) -> u32 { match bits { Bits::_1 => crate::extended::accumulate::<1, BITS>(lhs, rhs), diff --git a/crates/rabitq/src/byte.rs b/crates/rabitq/src/byte.rs index 43a6a688..9f145dbe 100644 --- a/crates/rabitq/src/byte.rs +++ b/crates/rabitq/src/byte.rs @@ -14,10 +14,15 @@ pub use crate::extended::{Code, CodeMetadata}; +#[deprecated] pub fn code(vector: &[f32]) -> Code { crate::extended::code::<8>(vector) } +pub fn ugly_code(vector: &[f32]) -> Code { + crate::extended::ugly_code::<8>(vector) +} + pub mod binary { use crate::extended::CodeMetadata; @@ -27,11 +32,17 @@ pub mod binary { pub type BinaryLut = (BinaryLutMetadata, Vec); pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); + #[deprecated] pub fn preprocess(vector: &[f32]) -> BinaryLut { let (metadata, elements) = crate::extended::code::(vector); (metadata, elements) } + pub fn ugly_preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::ugly_code::(vector); + (metadata, elements) + } + pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { simd::u8::reduce_sum_of_x_as_u32_y_as_u32(x, y) } @@ -55,4 +66,14 @@ pub mod binary { let rough = crate::extended::half_process_l2::<8, BITS>(n, value, code, lut); (rough,) } + + pub fn half_process_cos( + n: u32, + value: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_cos::<8, BITS>(n, value, code, lut); + (rough,) + } } diff --git a/crates/rabitq/src/extended.rs b/crates/rabitq/src/extended.rs index 2e3c014f..53b04212 100644 --- a/crates/rabitq/src/extended.rs +++ b/crates/rabitq/src/extended.rs @@ -62,7 +62,7 @@ pub fn code(vector: &[f32]) -> Code { let scale = { let mut o = normalized_vector.clone(); f32::vector_abs_inplace(&mut o); - find_scale::(&o) as f32 + f32::EPSILON + find_scale::(&o) }; let mut code = Vec::with_capacity(n as _); for i in 0..n { @@ -99,6 +99,60 @@ pub fn code(vector: &[f32]) -> Code { ) } +pub fn ugly_code(vector: &[f32]) -> Code { + assert!((1..=8).contains(&BITS)); + + let n = vector.len(); + let dis_u_2 = f32::reduce_sum_of_x2(vector); + let code = { + let min = -(1 << (BITS - 1)); + let max = (1 << (BITS - 1)) - 1; + let normalized_vector = { + let mut vector = vector.to_vec(); + f32::vector_mul_scalar_inplace(&mut vector, 1.0 / dis_u_2.sqrt()); + vector + }; + let (scale, delta) = { + let mut o = normalized_vector.clone(); + f32::vector_abs_inplace(&mut o); + ugly_find_scale::(&o) + }; + let mut code = Vec::with_capacity(n as _); + for i in 0..n { + let v = scale * normalized_vector[i]; + let v = v + v.signum() * delta[i] as f32; + let c = v.floor().clamp(min as f32, max as f32) as i32; + code.push((c + (1 << (BITS - 1))) as _); + } + code + }; + let norm_of_lattice = { + let base = -0.5 * ((1 << BITS) - 1) as f32; + let mut y = 0.0; + for i in 0..n { + let x = base + code[i] as f32; + y += x * x; + } + y.sqrt() + }; + let sum_of_code = { + let mut y = 0; + for i in 0..n { + let x = code[i] as u32; + y += x; + } + y as f32 + }; + ( + CodeMetadata { + dis_u_2, + norm_of_lattice, + sum_of_code, + }, + code, + ) +} + pub fn half_process_l2( n: u32, value: u32, @@ -135,7 +189,23 @@ pub fn half_process_dot( -ip * (lhs.dis_u_2.sqrt() / lhs.norm_of_lattice) * (rhs.dis_u_2.sqrt() / rhs.norm_of_lattice) } -fn find_scale(o: &[f32]) -> f64 { +pub fn half_process_cos( + n: u32, + value: u32, + lhs: CodeMetadata, + rhs: CodeMetadata, +) -> f32 { + assert!((1..=8).contains(&X)); + assert!((1..=8).contains(&Y)); + + let c_x = ((1 << X) - 1) as f32 * 0.5; + let c_y = ((1 << Y) - 1) as f32 * 0.5; + + let ip = value as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + n as f32 * c_x * c_y; + -ip / lhs.norm_of_lattice / rhs.norm_of_lattice +} + +fn find_scale(o: &[f32]) -> f32 { assert!((1..=8).contains(&B)); let mask = (1_u32 << (B - 1)) - 1; @@ -149,8 +219,9 @@ fn find_scale(o: &[f32]) -> f64 { for i in 0..dims { code.push(0); - numerator += 0.5 * o[i] as f64; - sqr_denominator += 0.5 * 0.5; + let value = 0.5; + numerator += value * o[i] as f64; + sqr_denominator += value * value; } { let x = 0.0; @@ -177,7 +248,58 @@ fn find_scale(o: &[f32]) -> f64 { } } - x_m + x_m as f32 + f32::EPSILON +} + +fn ugly_find_scale(o: &[f32]) -> (f32, Vec) { + assert!((1..=8).contains(&B)); + + let dims = o.len(); + + let mut code = Vec::::with_capacity(dims); + let mut numerator_m = 0.0f64; + let mut sqr_denominator_m = 0.0f64; + + let scale = (1 << (B - 1)) as f32 / f32::reduce_min_max_of_x(o).1; + for i in 0..dims { + code.push((o[i] as f64 * scale as f64) as u8); + let value = code[i] as f64 + 0.5; + numerator_m += value * o[i] as f64; + sqr_denominator_m += value * value; + } + let mut y_m = numerator_m / sqr_denominator_m.sqrt(); + + let mut delta = vec![0_i32; dims]; + for _ in 0..8 { + for i in 0..dims { + if code[i] < (1 << (B - 1)) - 1 { + let numerator = numerator_m + o[i] as f64; + let sqr_denominator = sqr_denominator_m + 2.0 * (code[i] as f64 + 1.0); + let y = numerator / sqr_denominator.sqrt(); + if y > y_m { + y_m = y; + numerator_m = numerator; + sqr_denominator_m = sqr_denominator; + code[i] += 1; + delta[i] += 1; + } + } + if code[i] > 0 { + let numerator = numerator_m - o[i] as f64; + let sqr_denominator = sqr_denominator_m - 2.0 * code[i] as f64; + let y = numerator / sqr_denominator.sqrt(); + if y > y_m { + y_m = y; + numerator_m = numerator; + sqr_denominator_m = sqr_denominator; + code[i] -= 1; + delta[i] -= 1; + } + } + } + } + + (scale, delta) } pub fn pack_code(input: &[u8]) -> [Vec; BITS] { diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs index 1193de48..4cf31f57 100644 --- a/crates/vchordg/src/operator.rs +++ b/crates/vchordg/src/operator.rs @@ -14,11 +14,14 @@ use crate::types::DistanceKind; use distance::Distance; -use index::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S}; +use index::accessor::{ + Accessor1, Accessor2, DimensionDistanceAccessor, DistanceAccessor, Dot, L2S, +}; use rabitq::bits::Bits; use simd::{Floating, f16}; use std::fmt::Debug; use std::marker::PhantomData; +use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; use vector::{VectorBorrowed, VectorOwned}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -118,6 +121,72 @@ impl Vector for VectOwned { } } +impl Vector for Rabitq8Owned { + type Metadata = [f32; 4]; + + type Element = u8; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + ( + vector.code(), + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)) { + let metadata = [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ]; + let slice = vector.code(); + let tailing = (size_of::() * m) + .next_multiple_of(crate::tuples::ALIGN); + assert!(tailing <= 8000); + if slice.len() <= (8000 - tailing) / size_of::() { + return (vec![], (slice, metadata)); + } + let (l, r) = slice.split_at(slice.len() - (8000 - tailing) / size_of::()); + ( + l.chunks(8000 / size_of::()).collect::>(), + (r, metadata), + ) + } + + fn pack(code: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq8Owned::new(_0, _1, _2, _3, code) + } + + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dims() as _); + for c in vector.code().iter().copied() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bits::code(bits, &result) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dims() as _); + for c in vector.code().iter().copied() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bits::binary::preprocess(&result) + } +} + pub trait Operator: 'static + Debug + Copy { const DISTANCE: DistanceKind; @@ -287,6 +356,72 @@ impl Operator for Op, Dot> { } } +impl Operator for Op { + const DISTANCE: DistanceKind = DistanceKind::L2S; + + type Vector = Rabitq8Owned; + + type DistanceAccessor = DimensionDistanceAccessor; + + fn process( + bits: Bits, + n: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2( + bits, + n, + value, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_l2s(rhs) + } +} + +impl Operator for Op { + const DISTANCE: DistanceKind = DistanceKind::Dot; + + type Vector = Rabitq8Owned; + + type DistanceAccessor = DimensionDistanceAccessor; + + fn process( + bits: Bits, + n: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_dot( + bits, + n, + value, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_dot(rhs) + } +} + #[derive(Debug, Clone)] pub struct CloneAccessor(Vec); diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index d262587b..c1d94cf5 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -15,6 +15,7 @@ use serde::{Deserialize, Serialize}; use simd::f16; use validator::{Validate, ValidationError}; +use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] @@ -83,12 +84,14 @@ impl Default for VchordgIndexOptions { pub enum OwnedVector { Vecf32(VectOwned), Vecf16(VectOwned), + Rabitq8(Rabitq8Owned), } #[derive(Debug, Clone, Copy)] pub enum BorrowedVector<'a> { Vecf32(VectBorrowed<'a, f32>), Vecf16(VectBorrowed<'a, f16>), + Rabitq8(Rabitq8Borrowed<'a>), } #[repr(u8)] @@ -103,6 +106,7 @@ pub enum DistanceKind { pub enum VectorKind { Vecf32, Vecf16, + Rabitq8, } impl VectorKind { @@ -110,6 +114,7 @@ impl VectorKind { match self { VectorKind::Vecf32 => size_of::() as _, VectorKind::Vecf16 => size_of::() as _, + VectorKind::Rabitq8 => size_of::() as _, } } } diff --git a/crates/vchordrq/src/operator.rs b/crates/vchordrq/src/operator.rs index a27bf41f..8b0688ce 100644 --- a/crates/vchordrq/src/operator.rs +++ b/crates/vchordrq/src/operator.rs @@ -13,13 +13,16 @@ // Copyright (c) 2025 TensorChord Inc. use distance::Distance; -use index::accessor::{Accessor1, Accessor2, DistanceAccessor, Dot, L2S, RAccess}; +use index::accessor::{ + Accessor1, Accessor2, DimensionDistanceAccessor, DistanceAccessor, Dot, L2S, RAccess, +}; use rabitq::bit::CodeMetadata; use rabitq::bit::binary::BinaryLut; use rabitq::bit::block::{BlockLut, STEP}; use simd::{Floating, f16}; use std::fmt::Debug; use std::marker::PhantomData; +use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; use vector::{VectorBorrowed, VectorOwned}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; @@ -209,6 +212,109 @@ impl Vector for VectOwned { } } +impl Vector for Rabitq8Owned { + type Metadata = [f32; 4]; + + type Element = u8; + + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[u8]>, [f32; 4]) { + ( + match vector.code().len() { + 0 => unreachable!(), + 1..=3840 => vec![vector.code()], + 3841..=5120 => vec![&vector.code()[..1280], &vector.code()[1280..]], + 5121.. => vector.code().chunks(7680).collect(), + }, + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn count(n: usize) -> usize { + match n { + 0 => unreachable!(), + 1..=3840 => 1, + 3841..=5120 => 2, + 5121.. => n.div_ceil(7680), + } + } + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + ( + vector.code(), + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn pack(elements: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq8Owned::new(_0, _1, _2, _3, elements) + } + + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dims() as _); + for c in vector.code().iter().copied() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bit::block::preprocess(&result) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dims() as _); + for c in vector.code().iter().copied() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bit::preprocess(&result) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code { + let n = vector.dims(); + let sum_of_abs_x = vector.sum_of_abs_x(); + let sum_of_x_2 = vector.sum_of_x2(); + ( + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector.code().iter().filter(|&&x| x >= 128).count(); + let cnt_neg = vector.code().iter().filter(|&&x| x <= 127).count(); + cnt_pos as f32 - cnt_neg as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + factor_err: { + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (n as f32).sqrt(); + dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (n as f32 - 1.0).sqrt() + }, + }, + { + let vector = vector.code(); + let n = vector.len(); + let mut signs = Vec::new(); + for i in 0..n { + signs.push(vector[i] >= 128); + } + signs + }, + ) + } + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { + vector.sum_of_x2() + } +} + pub trait Operator: 'static + Debug + Copy { type Vector: Vector; @@ -562,6 +668,108 @@ impl Operator for Op, Dot> { } } +impl Operator for Op { + type Vector = Rabitq8Owned; + + type DistanceAccessor = DimensionDistanceAccessor; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + assert!(!is_residual); + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, _delta| { + rabitq::bit::block::half_process_l2(value, code, lut.0) + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + assert!(!is_residual); + move |metadata: [f32; 4], elements: &[u64], _delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + rabitq::bit::binary::half_process_l2(value, code, lut.0) + } + } + + fn build( + vector: Rabitq8Borrowed<'_>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if centroid.is_some() { + unimplemented!(); + } + (Self::Vector::code(vector), 0.0) + } +} + +impl Operator for Op { + type Vector = Rabitq8Owned; + + type DistanceAccessor = DimensionDistanceAccessor; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + assert!(!is_residual); + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, _delta| { + rabitq::bit::block::half_process_dot(value, code, lut.0) + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + assert!(!is_residual); + move |metadata: [f32; 4], elements: &[u64], _delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + rabitq::bit::binary::half_process_dot(value, code, lut.0) + } + } + + fn build( + vector: Rabitq8Borrowed<'_>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if centroid.is_some() { + unimplemented!(); + } + (Self::Vector::code(vector), 0.0) + } +} + pub trait Call { type Output; diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index c42edecd..23bd43c4 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -15,6 +15,7 @@ use serde::{Deserialize, Serialize}; use simd::f16; use validator::{Validate, ValidationError}; +use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] @@ -55,12 +56,14 @@ impl Default for VchordrqIndexOptions { pub enum OwnedVector { Vecf32(VectOwned), Vecf16(VectOwned), + Rabitq8(Rabitq8Owned), } #[derive(Debug, Clone, Copy)] pub enum BorrowedVector<'a> { Vecf32(VectBorrowed<'a, f32>), Vecf16(VectBorrowed<'a, f16>), + Rabitq8(Rabitq8Borrowed<'a>), } #[repr(u8)] @@ -75,6 +78,7 @@ pub enum DistanceKind { pub enum VectorKind { Vecf32, Vecf16, + Rabitq8, } impl VectorKind { @@ -82,6 +86,7 @@ impl VectorKind { match self { VectorKind::Vecf32 => size_of::() as _, VectorKind::Vecf16 => size_of::() as _, + VectorKind::Rabitq8 => size_of::() as _, } } } diff --git a/crates/vector/Cargo.toml b/crates/vector/Cargo.toml index 41d7d173..d68f1f92 100644 --- a/crates/vector/Cargo.toml +++ b/crates/vector/Cargo.toml @@ -6,6 +6,7 @@ publish = false [dependencies] distance = { path = "../distance" } +rabitq = { path = "../rabitq" } simd = { path = "../simd" } [lints] diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs index 90551703..dd549d34 100644 --- a/crates/vector/src/bvect.rs +++ b/crates/vector/src/bvect.rs @@ -14,7 +14,6 @@ use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; -use std::ops::{Bound, RangeBounds}; pub const BVECTOR_WIDTH: u32 = u64::BITS; @@ -70,11 +69,6 @@ impl VectorOwned for BVectOwned { data: &self.data, } } - - #[inline(always)] - fn zero(dims: u32) -> Self { - Self::new(dims, vec![0; dims.div_ceil(BVECTOR_WIDTH) as usize]) - } } #[derive(Debug, Clone, Copy)] @@ -225,39 +219,6 @@ impl VectorBorrowed for BVectBorrowed<'_> { let data = simd::bit::vector_xor(self.data, rhs.data); BVectOwned::new(self.dims, data) } - - #[inline(always)] - fn subvector(&self, bounds: impl RangeBounds) -> Option { - let start = match bounds.start_bound().cloned() { - Bound::Included(x) => x, - Bound::Excluded(u32::MAX) => return None, - Bound::Excluded(x) => x + 1, - Bound::Unbounded => 0, - }; - let end = match bounds.end_bound().cloned() { - Bound::Included(u32::MAX) => return None, - Bound::Included(x) => x + 1, - Bound::Excluded(x) => x, - Bound::Unbounded => self.dims, - }; - if start >= end || end > self.dims { - return None; - } - let dims = end - start; - let mut data = vec![0_u64; dims.div_ceil(BVECTOR_WIDTH) as _]; - { - let mut i = 0; - let mut j = start; - while j < end { - if self.data[(j / BVECTOR_WIDTH) as usize] & (1 << (j % BVECTOR_WIDTH)) != 0 { - data[(i / BVECTOR_WIDTH) as usize] |= 1 << (i % BVECTOR_WIDTH); - } - i += 1; - j += 1; - } - } - Self::Owned::new_checked(dims, data) - } } impl PartialEq for BVectBorrowed<'_> { diff --git a/crates/vector/src/lib.rs b/crates/vector/src/lib.rs index 5f7df24b..67403990 100644 --- a/crates/vector/src/lib.rs +++ b/crates/vector/src/lib.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. pub mod bvect; -pub mod scalar8; +pub mod rabitq8; pub mod svect; pub mod vect; @@ -21,8 +21,6 @@ pub trait VectorOwned: Clone + 'static { type Borrowed<'a>: VectorBorrowed; fn as_borrowed(&self) -> Self::Borrowed<'_>; - - fn zero(dims: u32) -> Self; } pub trait VectorBorrowed: Copy { @@ -57,6 +55,4 @@ pub trait VectorBorrowed: Copy { fn operator_or(&self, rhs: Self) -> Self::Owned; fn operator_xor(&self, rhs: Self) -> Self::Owned; - - fn subvector(&self, bounds: impl std::ops::RangeBounds) -> Option; } diff --git a/crates/vector/src/rabitq8.rs b/crates/vector/src/rabitq8.rs new file mode 100644 index 00000000..0cbd074f --- /dev/null +++ b/crates/vector/src/rabitq8.rs @@ -0,0 +1,320 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; + +#[derive(Debug, Clone)] +pub struct Rabitq8Owned { + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + code: Vec, +} + +impl Rabitq8Owned { + #[inline(always)] + pub fn new( + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + code: Vec, + ) -> Self { + Self::new_checked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, code) + .expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + code: Vec, + ) -> Option { + if !(1..=65535).contains(&code.len()) { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { + Self::new_unchecked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, code) + }) + } + + /// # Safety + /// + /// * `code.len()` must not be zero. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked( + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + code: Vec, + ) -> Self { + Self { + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + code, + } + } +} + +impl VectorOwned for Rabitq8Owned { + type Borrowed<'a> = Rabitq8Borrowed<'a>; + + #[inline(always)] + fn as_borrowed(&self) -> Rabitq8Borrowed<'_> { + Rabitq8Borrowed { + sum_of_x2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x, + code: self.code.as_slice(), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct Rabitq8Borrowed<'a> { + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + code: &'a [u8], +} + +impl<'a> Rabitq8Borrowed<'a> { + #[inline(always)] + pub fn new( + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + code: &'a [u8], + ) -> Self { + Self::new_checked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, code) + .expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + code: &'a [u8], + ) -> Option { + if !(1..=65535).contains(&code.len()) { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { + Self::new_unchecked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, code) + }) + } + + /// # Safety + /// + /// * `code.len()` must not be zero. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked( + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + code: &'a [u8], + ) -> Self { + Self { + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + code, + } + } + + #[inline(always)] + pub fn sum_of_x2(&self) -> f32 { + self.sum_of_x2 + } + + #[inline(always)] + pub fn norm_of_lattice(&self) -> f32 { + self.norm_of_lattice + } + + #[inline(always)] + pub fn sum_of_code(&self) -> f32 { + self.sum_of_code + } + + #[inline(always)] + pub fn sum_of_abs_x(&self) -> f32 { + self.sum_of_abs_x + } + + #[inline(always)] + pub fn code(&self) -> &'a [u8] { + self.code + } +} + +impl VectorBorrowed for Rabitq8Borrowed<'_> { + type Owned = Rabitq8Owned; + + #[inline(always)] + fn dims(&self) -> u32 { + self.code.len() as u32 + } + + #[inline(always)] + fn own(&self) -> Rabitq8Owned { + Rabitq8Owned { + sum_of_x2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x, + code: self.code.to_owned(), + } + } + + #[inline(always)] + fn norm(&self) -> f32 { + self.sum_of_x2.sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + assert_eq!(self.code.len(), rhs.code.len()); + let n = self.code.len() as u32; + let value = rabitq::byte::binary::accumulate(self.code, rhs.code); + Distance::from_f32( + rabitq::byte::binary::half_process_dot( + n, + value, + rabitq::byte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::byte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_l2s(self, rhs: Self) -> Distance { + assert_eq!(self.code.len(), rhs.code.len()); + let n = self.code.len() as u32; + let value = rabitq::byte::binary::accumulate(self.code, rhs.code); + Distance::from_f32( + rabitq::byte::binary::half_process_l2( + n, + value, + rabitq::byte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::byte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_cos(self, rhs: Self) -> Distance { + assert_eq!(self.code.len(), rhs.code.len()); + let n = self.code.len() as u32; + let value = rabitq::byte::binary::accumulate(self.code, rhs.code); + Distance::from_f32( + rabitq::byte::binary::half_process_cos( + n, + value, + rabitq::byte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::byte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_hamming(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_jaccard(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn function_normalize(&self) -> Rabitq8Owned { + Rabitq8Owned { + sum_of_x2: 1.0, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x / self.sum_of_x2.sqrt(), + code: self.code.to_owned(), + } + } + + fn operator_add(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_sub(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_mul(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_and(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_or(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_xor(&self, _: Self) -> Self::Owned { + unimplemented!() + } +} diff --git a/crates/vector/src/scalar8.rs b/crates/vector/src/scalar8.rs deleted file mode 100644 index 87186819..00000000 --- a/crates/vector/src/scalar8.rs +++ /dev/null @@ -1,305 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -use crate::{VectorBorrowed, VectorOwned}; -use distance::Distance; -use std::ops::RangeBounds; - -#[derive(Debug, Clone)] -pub struct Scalar8Owned { - sum_of_x2: f32, - k: f32, - b: f32, - sum_of_code: f32, - code: Vec, -} - -impl Scalar8Owned { - #[inline(always)] - pub fn new(sum_of_x2: f32, k: f32, b: f32, sum_of_code: f32, code: Vec) -> Self { - Self::new_checked(sum_of_x2, k, b, sum_of_code, code).expect("invalid data") - } - - #[inline(always)] - pub fn new_checked( - sum_of_x2: f32, - k: f32, - b: f32, - sum_of_code: f32, - code: Vec, - ) -> Option { - if !(1..=65535).contains(&code.len()) { - return None; - } - #[allow(unsafe_code)] - Some(unsafe { Self::new_unchecked(sum_of_x2, k, b, sum_of_code, code) }) - } - - /// # Safety - /// - /// * `code.len()` must not be zero. - #[allow(unsafe_code)] - #[inline(always)] - pub unsafe fn new_unchecked( - sum_of_x2: f32, - k: f32, - b: f32, - sum_of_code: f32, - code: Vec, - ) -> Self { - Self { - sum_of_x2, - k, - b, - sum_of_code, - code, - } - } -} - -impl VectorOwned for Scalar8Owned { - type Borrowed<'a> = Scalar8Borrowed<'a>; - - #[inline(always)] - fn as_borrowed(&self) -> Scalar8Borrowed<'_> { - Scalar8Borrowed { - sum_of_x2: self.sum_of_x2, - k: self.k, - b: self.b, - sum_of_code: self.sum_of_code, - code: self.code.as_slice(), - } - } - - #[inline(always)] - fn zero(dims: u32) -> Self { - Self { - sum_of_x2: 0.0, - k: 0.0, - b: 0.0, - sum_of_code: 0.0, - code: vec![0; dims as usize], - } - } -} - -#[derive(Debug, Clone, Copy)] -pub struct Scalar8Borrowed<'a> { - sum_of_x2: f32, - k: f32, - b: f32, - sum_of_code: f32, - code: &'a [u8], -} - -impl<'a> Scalar8Borrowed<'a> { - #[inline(always)] - pub fn new(sum_of_x2: f32, k: f32, b: f32, sum_of_code: f32, code: &'a [u8]) -> Self { - Self::new_checked(sum_of_x2, k, b, sum_of_code, code).expect("invalid data") - } - - #[inline(always)] - pub fn new_checked( - sum_of_x2: f32, - k: f32, - b: f32, - sum_of_code: f32, - code: &'a [u8], - ) -> Option { - if !(1..=65535).contains(&code.len()) { - return None; - } - #[allow(unsafe_code)] - Some(unsafe { Self::new_unchecked(sum_of_x2, k, b, sum_of_code, code) }) - } - - /// # Safety - /// - /// * `code.len()` must not be zero. - #[allow(unsafe_code)] - #[inline(always)] - pub unsafe fn new_unchecked( - sum_of_x2: f32, - k: f32, - b: f32, - sum_of_code: f32, - code: &'a [u8], - ) -> Self { - Self { - sum_of_x2, - k, - b, - sum_of_code, - code, - } - } - - #[inline(always)] - pub fn sum_of_x2(&self) -> f32 { - self.sum_of_x2 - } - - #[inline(always)] - pub fn k(&self) -> f32 { - self.k - } - - #[inline(always)] - pub fn b(&self) -> f32 { - self.b - } - - #[inline(always)] - pub fn sum_of_code(&self) -> f32 { - self.sum_of_code - } - - #[inline(always)] - pub fn code(&self) -> &'a [u8] { - self.code - } -} - -impl VectorBorrowed for Scalar8Borrowed<'_> { - type Owned = Scalar8Owned; - - #[inline(always)] - fn dims(&self) -> u32 { - self.code.len() as u32 - } - - #[inline(always)] - fn own(&self) -> Scalar8Owned { - Scalar8Owned { - sum_of_x2: self.sum_of_x2, - k: self.k, - b: self.b, - sum_of_code: self.sum_of_code, - code: self.code.to_owned(), - } - } - - #[inline(always)] - fn norm(&self) -> f32 { - self.sum_of_x2.sqrt() - } - - #[inline(always)] - fn operator_dot(self, rhs: Self) -> Distance { - assert_eq!(self.code.len(), rhs.code.len()); - let xy = - self.k * rhs.k * simd::u8::reduce_sum_of_x_as_u32_y_as_u32(self.code, rhs.code) as f32 - + self.b * rhs.b * self.code.len() as f32 - + self.k * rhs.b * self.sum_of_code - + self.b * rhs.k * rhs.sum_of_code; - Distance::from(-xy) - } - - #[inline(always)] - fn operator_l2s(self, rhs: Self) -> Distance { - assert_eq!(self.code.len(), rhs.code.len()); - let xy = - self.k * rhs.k * simd::u8::reduce_sum_of_x_as_u32_y_as_u32(self.code, rhs.code) as f32 - + self.b * rhs.b * self.code.len() as f32 - + self.k * rhs.b * self.sum_of_code - + self.b * rhs.k * rhs.sum_of_code; - let x2 = self.sum_of_x2; - let y2 = rhs.sum_of_x2; - Distance::from(x2 + y2 - 2.0 * xy) - } - - #[inline(always)] - fn operator_cos(self, rhs: Self) -> Distance { - assert_eq!(self.code.len(), rhs.code.len()); - let xy = - self.k * rhs.k * simd::u8::reduce_sum_of_x_as_u32_y_as_u32(self.code, rhs.code) as f32 - + self.b * rhs.b * self.code.len() as f32 - + self.k * rhs.b * self.sum_of_code - + self.b * rhs.k * rhs.sum_of_code; - let x2 = self.sum_of_x2; - let y2 = rhs.sum_of_x2; - Distance::from(1.0 - xy / (x2 * y2).sqrt()) - } - - #[inline(always)] - fn operator_hamming(self, _: Self) -> Distance { - unimplemented!() - } - - #[inline(always)] - fn operator_jaccard(self, _: Self) -> Distance { - unimplemented!() - } - - #[inline(always)] - fn function_normalize(&self) -> Scalar8Owned { - let l = self.sum_of_x2.sqrt(); - Scalar8Owned { - sum_of_x2: 1.0, - k: self.k / l, - b: self.b / l, - sum_of_code: self.sum_of_code, - code: self.code.to_owned(), - } - } - - fn operator_add(&self, _: Self) -> Self::Owned { - unimplemented!() - } - - fn operator_sub(&self, _: Self) -> Self::Owned { - unimplemented!() - } - - fn operator_mul(&self, _: Self) -> Self::Owned { - unimplemented!() - } - - fn operator_and(&self, _: Self) -> Self::Owned { - unimplemented!() - } - - fn operator_or(&self, _: Self) -> Self::Owned { - unimplemented!() - } - - fn operator_xor(&self, _: Self) -> Self::Owned { - unimplemented!() - } - - #[inline(always)] - fn subvector(&self, bounds: impl RangeBounds) -> Option { - let start_bound = bounds.start_bound().map(|x| *x as usize); - let end_bound = bounds.end_bound().map(|x| *x as usize); - let code = self.code.get((start_bound, end_bound))?; - if code.is_empty() { - return None; - } - Self::Owned::new_checked( - { - // recover it as much as possible - let mut result = 0.0; - for &x in code { - let y = self.k * (x as f32) + self.b; - result += y * y; - } - result - }, - self.k, - self.b, - simd::u8::reduce_sum_of_x_as_u32(code) as f32, - code.to_owned(), - ) - } -} diff --git a/crates/vector/src/svect.rs b/crates/vector/src/svect.rs index 19d41104..1b9ebcee 100644 --- a/crates/vector/src/svect.rs +++ b/crates/vector/src/svect.rs @@ -15,7 +15,6 @@ use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; use simd::Floating; -use std::ops::{Bound, RangeBounds}; #[derive(Debug, Clone)] pub struct SVectOwned { @@ -94,11 +93,6 @@ impl VectorOwned for SVectOwned { values: &self.values, } } - - #[inline(always)] - fn zero(dims: u32) -> Self { - Self::new(dims, vec![], vec![]) - } } #[derive(Debug, Clone, Copy)] @@ -375,34 +369,6 @@ impl VectorBorrowed for SVectBorrowed<'_, S> { fn operator_xor(&self, _: Self) -> Self::Owned { unimplemented!() } - - #[inline(always)] - fn subvector(&self, bounds: impl RangeBounds) -> Option { - let start = match bounds.start_bound().cloned() { - Bound::Included(x) => x, - Bound::Excluded(u32::MAX) => return None, - Bound::Excluded(x) => x + 1, - Bound::Unbounded => 0, - }; - let end = match bounds.end_bound().cloned() { - Bound::Included(u32::MAX) => return None, - Bound::Included(x) => x + 1, - Bound::Excluded(x) => x, - Bound::Unbounded => self.dims, - }; - if start >= end || end > self.dims { - return None; - } - let dims = end - start; - let s = self.indexes.partition_point(|&x| x < start); - let e = self.indexes.partition_point(|&x| x < end); - let indexes = self.indexes[s..e] - .iter() - .map(|x| x - start) - .collect::>(); - let values = self.values[s..e].to_vec(); - Self::Owned::new_checked(dims, indexes, values) - } } impl PartialEq for SVectBorrowed<'_, S> { diff --git a/crates/vector/src/vect.rs b/crates/vector/src/vect.rs index 8f40cf2b..ac1c50bf 100644 --- a/crates/vector/src/vect.rs +++ b/crates/vector/src/vect.rs @@ -16,7 +16,6 @@ use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; use simd::Floating; use std::cmp::Ordering; -use std::ops::RangeBounds; #[derive(Debug, Clone)] #[repr(transparent)] @@ -69,11 +68,6 @@ impl VectorOwned for VectOwned { fn as_borrowed(&self) -> VectBorrowed<'_, S> { VectBorrowed(self.0.as_slice()) } - - #[inline(always)] - fn zero(dims: u32) -> Self { - Self::new(vec![S::zero(); dims as usize]) - } } #[derive(Debug, Clone, Copy)] @@ -187,17 +181,6 @@ impl VectorBorrowed for VectBorrowed<'_, S> { fn operator_xor(&self, _: Self) -> Self::Owned { unimplemented!() } - - #[inline(always)] - fn subvector(&self, bounds: impl RangeBounds) -> Option { - let start_bound = bounds.start_bound().map(|x| *x as usize); - let end_bound = bounds.end_bound().map(|x| *x as usize); - let slice = self.0.get((start_bound, end_bound))?; - if slice.is_empty() { - return None; - } - Self::Owned::new_checked(slice.to_vec()) - } } impl PartialEq for VectBorrowed<'_, S> { diff --git a/src/datatype/binary_scalar8.rs b/src/datatype/binary_rabitq8.rs similarity index 83% rename from src/datatype/binary_scalar8.rs rename to src/datatype/binary_rabitq8.rs index e1c35efe..cad31204 100644 --- a/src/datatype/binary_scalar8.rs +++ b/src/datatype/binary_rabitq8.rs @@ -12,21 +12,21 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::datatype::memory_scalar8::{Scalar8Input, Scalar8Output}; +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; use pgrx::datum::Internal; use pgrx::pg_sys::Oid; use vector::VectorBorrowed; -use vector::scalar8::Scalar8Borrowed; +use vector::rabitq8::Rabitq8Borrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_send(vector: Scalar8Input<'_>) -> Vec { +fn _vchord_rabitq8_send(vector: Rabitq8Input<'_>) -> Vec { let vector = vector.as_borrowed(); let mut stream = Vec::::new(); stream.extend(vector.dims().to_be_bytes()); stream.extend(vector.sum_of_x2().to_be_bytes()); - stream.extend(vector.k().to_be_bytes()); - stream.extend(vector.b().to_be_bytes()); + stream.extend(vector.norm_of_lattice().to_be_bytes()); stream.extend(vector.sum_of_code().to_be_bytes()); + stream.extend(vector.sum_of_abs_x().to_be_bytes()); for &c in vector.code() { stream.extend(c.to_be_bytes()); } @@ -34,7 +34,7 @@ fn _vchord_scalar8_send(vector: Scalar8Input<'_>) -> Vec { } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Scalar8Output { +fn _vchord_rabitq8_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Rabitq8Output { let _ = (oid, typmod); let buf = unsafe { internal.get_mut::().unwrap() }; @@ -50,19 +50,19 @@ fn _vchord_scalar8_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Scalar buf.cursor += 4; f32::from_be_bytes(raw) }; - let k = { + let norm_of_lattice = { assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; buf.cursor += 4; f32::from_be_bytes(raw) }; - let b = { + let sum_of_code = { assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; buf.cursor += 4; f32::from_be_bytes(raw) }; - let sum_of_code = { + let sum_of_abs_x = { assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; buf.cursor += 4; @@ -81,8 +81,10 @@ fn _vchord_scalar8_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Scalar result }; - if let Some(x) = Scalar8Borrowed::new_checked(sum_of_x2, k, b, sum_of_code, &code) { - Scalar8Output::new(x) + if let Some(x) = + Rabitq8Borrowed::new_checked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, &code) + { + Rabitq8Output::new(x) } else { pgrx::error!("detect data corruption"); } diff --git a/src/datatype/functions_rabitq8.rs b/src/datatype/functions_rabitq8.rs new file mode 100644 index 00000000..870bcf44 --- /dev/null +++ b/src/datatype/functions_rabitq8.rs @@ -0,0 +1,47 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::datatype::memory_halfvec::HalfvecInput; +use crate::datatype::memory_rabitq8::Rabitq8Output; +use crate::datatype::memory_vector::VectorInput; +use simd::{Floating, f16}; +use vector::rabitq8::Rabitq8Borrowed; + +#[pgrx::pg_extern(sql = "")] +fn _vchord_vector_quantize_to_rabitq8(vector: VectorInput) -> Rabitq8Output { + let mut vector = vector.as_borrowed().slice().to_vec(); + rabitq::rotate::rotate_inplace(&mut vector); + let code = rabitq::byte::ugly_code(&vector); + Rabitq8Output::new(Rabitq8Borrowed::new( + code.0.dis_u_2, + code.0.norm_of_lattice, + code.0.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + &code.1, + )) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_halfvec_quantize_to_rabitq8(vector: HalfvecInput) -> Rabitq8Output { + let mut vector = f16::vector_to_f32(vector.as_borrowed().slice()); + rabitq::rotate::rotate_inplace(&mut vector); + let code = rabitq::byte::ugly_code(&vector); + Rabitq8Output::new(Rabitq8Borrowed::new( + code.0.dis_u_2, + code.0.norm_of_lattice, + code.0.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + &code.1, + )) +} diff --git a/src/datatype/functions_scalar8.rs b/src/datatype/functions_scalar8.rs deleted file mode 100644 index 99990b15..00000000 --- a/src/datatype/functions_scalar8.rs +++ /dev/null @@ -1,39 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -use crate::datatype::memory_halfvec::HalfvecInput; -use crate::datatype::memory_scalar8::Scalar8Output; -use crate::datatype::memory_vector::VectorInput; -use simd::{Floating, f16}; -use vector::scalar8::Scalar8Borrowed; - -#[pgrx::pg_extern(sql = "")] -fn _vchord_vector_quantize_to_scalar8(vector: VectorInput) -> Scalar8Output { - let vector = vector.as_borrowed(); - let sum_of_x2 = f32::reduce_sum_of_x2(vector.slice()); - let (k, b, code) = - simd::quantize::quantize(f32::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); - let sum_of_code = simd::u8::reduce_sum_of_x_as_u32(&code) as f32; - Scalar8Output::new(Scalar8Borrowed::new(sum_of_x2, k, b, sum_of_code, &code)) -} - -#[pgrx::pg_extern(sql = "")] -fn _vchord_halfvec_quantize_to_scalar8(vector: HalfvecInput) -> Scalar8Output { - let vector = vector.as_borrowed(); - let sum_of_x2 = f16::reduce_sum_of_x2(vector.slice()); - let (k, b, code) = - simd::quantize::quantize(f16::vector_to_f32_borrowed(vector.slice()).as_ref(), 255.0); - let sum_of_code = simd::u8::reduce_sum_of_x_as_u32(&code) as f32; - Scalar8Output::new(Scalar8Borrowed::new(sum_of_x2, k, b, sum_of_code, &code)) -} diff --git a/src/datatype/memory_scalar8.rs b/src/datatype/memory_rabitq8.rs similarity index 68% rename from src/datatype/memory_scalar8.rs rename to src/datatype/memory_rabitq8.rs index f030b865..67970a59 100644 --- a/src/datatype/memory_scalar8.rs +++ b/src/datatype/memory_rabitq8.rs @@ -18,35 +18,35 @@ use pgrx::pgrx_sql_entity_graph::metadata::*; use std::marker::PhantomData; use std::ptr::NonNull; use vector::VectorBorrowed; -use vector::scalar8::Scalar8Borrowed; +use vector::rabitq8::Rabitq8Borrowed; #[repr(C)] -struct Scalar8Header { +struct Rabitq8Header { varlena: u32, dims: u16, unused: u16, sum_of_x2: f32, - k: f32, - b: f32, + norm_of_lattice: f32, sum_of_code: f32, + sum_of_abs_x: f32, elements: [u8; 0], } -impl Scalar8Header { +impl Rabitq8Header { fn size_of(len: usize) -> usize { if len > 65535 { panic!("vector is too large"); } - (size_of::() + size_of::() * len).next_multiple_of(8) + size_of::() + size_of::() * len } - unsafe fn as_borrowed<'a>(this: NonNull) -> Scalar8Borrowed<'a> { + unsafe fn as_borrowed<'a>(this: NonNull) -> Rabitq8Borrowed<'a> { unsafe { let this = this.as_ptr(); - Scalar8Borrowed::new( + Rabitq8Borrowed::new( (&raw const (*this).sum_of_x2).read(), - (&raw const (*this).k).read(), - (&raw const (*this).b).read(), + (&raw const (*this).norm_of_lattice).read(), (&raw const (*this).sum_of_code).read(), + (&raw const (*this).sum_of_abs_x).read(), std::slice::from_raw_parts( (&raw const (*this).elements).cast(), (&raw const (*this).dims).read() as usize, @@ -56,10 +56,10 @@ impl Scalar8Header { } } -pub struct Scalar8Input<'a>(NonNull, PhantomData<&'a ()>, bool); +pub struct Rabitq8Input<'a>(NonNull, PhantomData<&'a ()>, bool); -impl Scalar8Input<'_> { - unsafe fn from_ptr(p: NonNull) -> Self { +impl Rabitq8Input<'_> { + unsafe fn from_ptr(p: NonNull) -> Self { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() }; @@ -70,18 +70,18 @@ impl Scalar8Input<'_> { #[cfg(target_endian = "little")] let size = varlena as usize >> 2; let dims = q.byte_add(4).cast::().read(); - assert_eq!(Scalar8Header::size_of(dims as _), size); + assert_eq!(Rabitq8Header::size_of(dims as _), size); let unused = q.byte_add(6).cast::().read(); assert_eq!(unused, 0); } - Scalar8Input(q, PhantomData, p != q) + Rabitq8Input(q, PhantomData, p != q) } - pub fn as_borrowed(&self) -> Scalar8Borrowed<'_> { - unsafe { Scalar8Header::as_borrowed(self.0) } + pub fn as_borrowed(&self) -> Rabitq8Borrowed<'_> { + unsafe { Rabitq8Header::as_borrowed(self.0) } } } -impl Drop for Scalar8Input<'_> { +impl Drop for Rabitq8Input<'_> { fn drop(&mut self) { if self.2 { unsafe { @@ -91,10 +91,10 @@ impl Drop for Scalar8Input<'_> { } } -pub struct Scalar8Output(NonNull); +pub struct Rabitq8Output(NonNull); -impl Scalar8Output { - unsafe fn from_ptr(p: NonNull) -> Self { +impl Rabitq8Output { + unsafe fn from_ptr(p: NonNull) -> Self { let q = unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() }; @@ -105,18 +105,18 @@ impl Scalar8Output { #[cfg(target_endian = "little")] let size = varlena as usize >> 2; let dims = q.byte_add(4).cast::().read(); - assert_eq!(Scalar8Header::size_of(dims as _), size); + assert_eq!(Rabitq8Header::size_of(dims as _), size); let unused = q.byte_add(6).cast::().read(); assert_eq!(unused, 0); } Self(q) } - pub fn new(vector: Scalar8Borrowed<'_>) -> Self { + pub fn new(vector: Rabitq8Borrowed<'_>) -> Self { unsafe { let code = vector.code(); - let size = Scalar8Header::size_of(code.len()); + let size = Rabitq8Header::size_of(code.len()); - let ptr = pgrx::pg_sys::palloc0(size) as *mut Scalar8Header; + let ptr = pgrx::pg_sys::palloc0(size) as *mut Rabitq8Header; // SET_VARSIZE_4B #[cfg(target_endian = "big")] (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); @@ -125,9 +125,9 @@ impl Scalar8Output { (&raw mut (*ptr).dims).write(vector.dims() as _); (&raw mut (*ptr).unused).write(0); (&raw mut (*ptr).sum_of_x2).write(vector.sum_of_x2()); - (&raw mut (*ptr).k).write(vector.k()); - (&raw mut (*ptr).b).write(vector.b()); + (&raw mut (*ptr).norm_of_lattice).write(vector.norm_of_lattice()); (&raw mut (*ptr).sum_of_code).write(vector.sum_of_code()); + (&raw mut (*ptr).sum_of_abs_x).write(vector.sum_of_abs_x()); std::ptr::copy_nonoverlapping( code.as_ptr(), (&raw mut (*ptr).elements).cast(), @@ -136,17 +136,17 @@ impl Scalar8Output { Self(NonNull::new(ptr).unwrap()) } } - pub fn as_borrowed(&self) -> Scalar8Borrowed<'_> { - unsafe { Scalar8Header::as_borrowed(self.0) } + pub fn as_borrowed(&self) -> Rabitq8Borrowed<'_> { + unsafe { Rabitq8Header::as_borrowed(self.0) } } - fn into_raw(self) -> *mut Scalar8Header { + fn into_raw(self) -> *mut Rabitq8Header { let result = self.0.as_ptr(); std::mem::forget(self); result } } -impl Drop for Scalar8Output { +impl Drop for Rabitq8Output { fn drop(&mut self) { unsafe { pgrx::pg_sys::pfree(self.0.as_ptr().cast()); @@ -156,7 +156,7 @@ impl Drop for Scalar8Output { // FromDatum -impl FromDatum for Scalar8Input<'_> { +impl FromDatum for Rabitq8Input<'_> { unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { if is_null { None @@ -167,7 +167,7 @@ impl FromDatum for Scalar8Input<'_> { } } -impl FromDatum for Scalar8Output { +impl FromDatum for Rabitq8Output { unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { if is_null { None @@ -180,7 +180,7 @@ impl FromDatum for Scalar8Output { // IntoDatum -impl IntoDatum for Scalar8Output { +impl IntoDatum for Rabitq8Output { fn into_datum(self) -> Option { Some(Datum::from(self.into_raw())) } @@ -196,8 +196,24 @@ impl IntoDatum for Scalar8Output { // UnboxDatum -unsafe impl pgrx::datum::UnboxDatum for Scalar8Output { - type As<'src> = Scalar8Output; +unsafe impl<'a> pgrx::datum::UnboxDatum for Rabitq8Input<'a> { + type As<'src> + = Rabitq8Input<'src> + where + 'a: 'src; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +unsafe impl pgrx::datum::UnboxDatum for Rabitq8Output { + type As<'src> = Rabitq8Output; #[inline] unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> where @@ -211,27 +227,27 @@ unsafe impl pgrx::datum::UnboxDatum for Scalar8Output { // SqlTranslatable -unsafe impl SqlTranslatable for Scalar8Input<'_> { +unsafe impl SqlTranslatable for Rabitq8Input<'_> { fn argument_sql() -> Result { - Ok(SqlMapping::As(String::from("scalar8"))) + Ok(SqlMapping::As(String::from("rabitq8"))) } fn return_sql() -> Result { - Ok(Returns::One(SqlMapping::As(String::from("scalar8")))) + Ok(Returns::One(SqlMapping::As(String::from("rabitq8")))) } } -unsafe impl SqlTranslatable for Scalar8Output { +unsafe impl SqlTranslatable for Rabitq8Output { fn argument_sql() -> Result { - Ok(SqlMapping::As(String::from("scalar8"))) + Ok(SqlMapping::As(String::from("rabitq8"))) } fn return_sql() -> Result { - Ok(Returns::One(SqlMapping::As(String::from("scalar8")))) + Ok(Returns::One(SqlMapping::As(String::from("rabitq8")))) } } // ArgAbi -unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Scalar8Input<'fcx> { +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Rabitq8Input<'fcx> { unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { let index = arg.index(); unsafe { @@ -243,7 +259,7 @@ unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Scalar8Input<'fcx> { // BoxRet -unsafe impl pgrx::callconv::BoxRet for Scalar8Output { +unsafe impl pgrx::callconv::BoxRet for Rabitq8Output { unsafe fn box_into<'fcx>( self, fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 8da4ee16..2f4f6548 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -12,13 +12,14 @@ // // Copyright (c) 2025 TensorChord Inc. -mod binary_scalar8; -mod functions_scalar8; +mod binary_rabitq8; +mod functions_rabitq8; pub mod memory_halfvec; -pub mod memory_scalar8; +pub mod memory_rabitq8; pub mod memory_vector; mod operators_halfvec; -mod operators_scalar8; +mod operators_rabitq8; mod operators_vector; -mod text_scalar8; +mod text_rabitq8; pub mod typmod; +mod typmod_rabitq8; diff --git a/src/datatype/operators_scalar8.rs b/src/datatype/operators_rabitq8.rs similarity index 63% rename from src/datatype/operators_scalar8.rs rename to src/datatype/operators_rabitq8.rs index cf8ee5ef..505e06e1 100644 --- a/src/datatype/operators_scalar8.rs +++ b/src/datatype/operators_rabitq8.rs @@ -12,47 +12,48 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::datatype::memory_scalar8::{Scalar8Input, Scalar8Output}; +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; +use pgrx::datum::Array; use std::num::NonZero; use vector::VectorBorrowed; -use vector::scalar8::Scalar8Borrowed; +use vector::rabitq8::Rabitq8Borrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_operator_ip(lhs: Scalar8Input<'_>, rhs: Scalar8Input<'_>) -> f32 { +fn _vchord_rabitq8_operator_l2(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> f32 { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); if lhs.dims() != rhs.dims() { pgrx::error!("dimension is not matched"); } - Scalar8Borrowed::operator_dot(lhs, rhs).to_f32() + Rabitq8Borrowed::operator_l2s(lhs, rhs).to_f32().sqrt() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_operator_l2(lhs: Scalar8Input<'_>, rhs: Scalar8Input<'_>) -> f32 { +fn _vchord_rabitq8_operator_ip(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> f32 { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); if lhs.dims() != rhs.dims() { pgrx::error!("dimension is not matched"); } - Scalar8Borrowed::operator_l2s(lhs, rhs).to_f32().sqrt() + Rabitq8Borrowed::operator_dot(lhs, rhs).to_f32() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_operator_cosine(lhs: Scalar8Input<'_>, rhs: Scalar8Input<'_>) -> f32 { +fn _vchord_rabitq8_operator_cosine(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> f32 { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); if lhs.dims() != rhs.dims() { pgrx::error!("dimension is not matched"); } - Scalar8Borrowed::operator_cos(lhs, rhs).to_f32() + Rabitq8Borrowed::operator_cos(lhs, rhs).to_f32() } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_sphere_ip_in( - lhs: Scalar8Input<'_>, - rhs: pgrx::composite_type!("sphere_scalar8"), +fn _vchord_rabitq8_sphere_l2_in( + lhs: Rabitq8Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq8"), ) -> bool { - let center: Scalar8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + let center: Rabitq8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), @@ -67,16 +68,16 @@ fn _vchord_scalar8_sphere_ip_in( if lhs.dims() != center.dims() { pgrx::error!("dimension is not matched"); } - let d = Scalar8Borrowed::operator_dot(lhs, center).to_f32(); + let d = Rabitq8Borrowed::operator_l2s(lhs, center).to_f32().sqrt(); d < radius } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_sphere_l2_in( - lhs: Scalar8Input<'_>, - rhs: pgrx::composite_type!("sphere_scalar8"), +fn _vchord_rabitq8_sphere_ip_in( + lhs: Rabitq8Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq8"), ) -> bool { - let center: Scalar8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + let center: Rabitq8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), @@ -91,16 +92,16 @@ fn _vchord_scalar8_sphere_l2_in( if lhs.dims() != center.dims() { pgrx::error!("dimension is not matched"); } - let d = Scalar8Borrowed::operator_l2s(lhs, center).to_f32().sqrt(); + let d = Rabitq8Borrowed::operator_dot(lhs, center).to_f32(); d < radius } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_sphere_cosine_in( - lhs: Scalar8Input<'_>, - rhs: pgrx::composite_type!("sphere_scalar8"), +fn _vchord_rabitq8_sphere_cosine_in( + lhs: Rabitq8Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq8"), ) -> bool { - let center: Scalar8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + let center: Rabitq8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { Ok(Some(s)) => s, Ok(None) => pgrx::error!("Bad input: empty center at sphere"), Err(_) => unreachable!(), @@ -115,6 +116,24 @@ fn _vchord_scalar8_sphere_cosine_in( if lhs.dims() != center.dims() { pgrx::error!("dimension is not matched"); } - let d = Scalar8Borrowed::operator_cos(lhs, center).to_f32(); + let d = Rabitq8Borrowed::operator_cos(lhs, center).to_f32(); d < radius } + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq8_operator_maxsim( + lhs: Array<'_, Rabitq8Input<'_>>, + rhs: Array<'_, Rabitq8Input<'_>>, +) -> f32 { + let mut maxsim = 0.0f32; + for rhs in rhs.iter().flatten() { + let mut d = f32::INFINITY; + for lhs in lhs.iter().flatten() { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + d = d.min(Rabitq8Borrowed::operator_dot(lhs, rhs).to_f32()); + } + maxsim += d; + } + maxsim +} diff --git a/src/datatype/text_scalar8.rs b/src/datatype/text_rabitq8.rs similarity index 87% rename from src/datatype/text_scalar8.rs rename to src/datatype/text_rabitq8.rs index e123df51..e81b084b 100644 --- a/src/datatype/text_scalar8.rs +++ b/src/datatype/text_rabitq8.rs @@ -12,13 +12,13 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::datatype::memory_scalar8::{Scalar8Input, Scalar8Output}; +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; use pgrx::pg_sys::Oid; use std::ffi::{CStr, CString}; -use vector::scalar8::Scalar8Borrowed; +use vector::rabitq8::Rabitq8Borrowed; #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_in(input: &CStr, oid: Oid, typmod: i32) -> Scalar8Output { +fn _vchord_rabitq8_in(input: &CStr, oid: Oid, typmod: i32) -> Rabitq8Output { let _ = (oid, typmod); let mut input = input.to_bytes().iter(); let mut p0 = Vec::::new(); @@ -118,26 +118,28 @@ fn _vchord_scalar8_in(input: &CStr, oid: Oid, typmod: i32) -> Scalar8Output { pgrx::error!("vector must have at least 1 dimension"); } let sum_of_x2 = p0[0]; - let k = p0[1]; - let b = p0[2]; - let sum_of_code = p0[3]; + let norm_of_lattice = p0[1]; + let sum_of_code = p0[2]; + let sum_of_abs_x = p0[3]; let code = p1; - if let Some(x) = Scalar8Borrowed::new_checked(sum_of_x2, k, b, sum_of_code, &code) { - Scalar8Output::new(x) + if let Some(x) = + Rabitq8Borrowed::new_checked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, &code) + { + Rabitq8Output::new(x) } else { pgrx::error!("incorrect vector"); } } #[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_scalar8_out(vector: Scalar8Input<'_>) -> CString { +fn _vchord_rabitq8_out(vector: Rabitq8Input<'_>) -> CString { let vector = vector.as_borrowed(); let mut buffer = String::new(); buffer.push('('); buffer.push_str(format!("{}", vector.sum_of_x2()).as_str()); - buffer.push_str(format!(", {}", vector.k()).as_str()); - buffer.push_str(format!(", {}", vector.b()).as_str()); + buffer.push_str(format!(", {}", vector.norm_of_lattice()).as_str()); buffer.push_str(format!(", {}", vector.sum_of_code()).as_str()); + buffer.push_str(format!(", {}", vector.sum_of_abs_x()).as_str()); buffer.push(')'); buffer.push('['); if let Some(&x) = vector.code().first() { diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index 1fd60055..46fd7fe7 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -12,7 +12,6 @@ // // Copyright (c) 2025 TensorChord Inc. -use std::ffi::{CStr, CString}; use std::num::NonZero; #[derive(Debug, Clone, Copy)] @@ -22,7 +21,7 @@ pub enum Typmod { } impl Typmod { - pub fn parse_from_i32(x: i32) -> Option { + pub fn new(x: i32) -> Option { use Typmod::*; if x == -1 { Some(Any) @@ -32,20 +31,6 @@ impl Typmod { None } } - pub fn into_option_string(self) -> Option { - use Typmod::*; - match self { - Any => None, - Dims(x) => Some(x.get().to_string()), - } - } - pub fn into_i32(self) -> i32 { - use Typmod::*; - match self { - Any => -1, - Dims(x) => x.get() as i32, - } - } pub fn dims(self) -> Option> { use Typmod::*; match self { @@ -54,30 +39,3 @@ impl Typmod { } } } - -#[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_typmod_in_65535(list: pgrx::datum::Array<&CStr>) -> i32 { - if list.is_empty() { - -1 - } else if list.len() == 1 { - let s = list.get(0).unwrap().unwrap().to_str().unwrap(); - let d = s.parse::().ok(); - if let Some(d @ 1..=65535) = d { - let typmod = Typmod::Dims(NonZero::new(d).unwrap()); - typmod.into_i32() - } else { - pgrx::error!("Modifier of the type is invalid.") - } - } else { - pgrx::error!("Modifier of the type is invalid.") - } -} - -#[pgrx::pg_extern(immutable, strict, parallel_safe)] -fn _vchord_typmod_out(typmod: i32) -> CString { - let typmod = Typmod::parse_from_i32(typmod).unwrap(); - match typmod.into_option_string() { - Some(s) => CString::new(format!("({s})")).unwrap(), - None => CString::new("()").unwrap(), - } -} diff --git a/src/datatype/typmod_rabitq8.rs b/src/datatype/typmod_rabitq8.rs new file mode 100644 index 00000000..748888a1 --- /dev/null +++ b/src/datatype/typmod_rabitq8.rs @@ -0,0 +1,37 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use std::ffi::CStr; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_typmod_in(list: pgrx::datum::Array<&CStr>) -> i32 { + if list.is_empty() { + -1 + } else if list.len() == 1 { + let s = list.get(0).unwrap().unwrap().to_str().unwrap(); + if let Ok(d) = s.parse::() { + if d < 1 { + pgrx::error!("dimensions for type rabitq8 must be at least 1"); + } + if d > 65535 { + pgrx::error!("dimensions for type rabitq8 cannot exceed 65535"); + } + d + } else { + pgrx::error!("invalid type modifier") + } + } else { + pgrx::error!("invalid type modifier") + } +} diff --git a/src/index/opclass.rs b/src/index/opclass.rs index 0b182bf4..6c8ddd45 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -42,6 +42,21 @@ fn _vchordg_support_halfvec_ip_ops() -> String { "vchordg_halfvec_ip_ops".to_string() } +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq8_l2_ops() -> String { + "vchordg_rabitq8_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq8_cosine_ops() -> String { + "vchordg_rabitq8_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq8_ip_ops() -> String { + "vchordg_rabitq8_ip_ops".to_string() +} + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_vector_l2_ops() -> String { "vchordrq_vector_l2_ops".to_string() @@ -72,6 +87,21 @@ fn _vchordrq_support_halfvec_cosine_ops() -> String { "vchordrq_halfvec_cosine_ops".to_string() } +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq8_l2_ops() -> String { + "vchordrq_rabitq8_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq8_ip_ops() -> String { + "vchordrq_rabitq8_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq8_cosine_ops() -> String { + "vchordrq_rabitq8_cosine_ops".to_string() +} + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_vector_maxsim_ops() -> String { "vchordrq_vector_maxsim_ops".to_string() @@ -82,6 +112,11 @@ fn _vchordrq_support_halfvec_maxsim_ops() -> String { "vchordrq_halfvec_maxsim_ops".to_string() } +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq8_maxsim_ops() -> String { + "vchordrq_rabitq8_maxsim_ops".to_string() +} + pub struct Sphere { pub center: T, pub radius: f32, diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index 602a09a2..2faa02e6 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -131,8 +131,8 @@ pub unsafe extern "C-unwind" fn ambuild( if vector_options.d != DistanceKind::L2S && (vchordg_options.index.alpha != [1.0] && vchordg_options.index.alpha != [1.0, 1.2]) { - let errors = "alpha not equal to `1.0` are only applicable to l2 and cosine distance."; - pgrx::warning!("warning while validating options: {errors}"); + let warnings = "alpha not equal to `1.0` are only applicable to l2 and cosine distance."; + pgrx::warning!("warning while validating options: {warnings}"); } let index = unsafe { PostgresRelation::new(index_relation) }; let reporter = PostgresReporter { @@ -584,7 +584,7 @@ unsafe fn options( pgrx::error!("multicolumn index is not supported"); } // get dims - let typmod = Typmod::parse_from_i32(atts[0].atttypmod).unwrap(); + let typmod = Typmod::new(atts[0].atttypmod).unwrap(); let dims = if let Some(dims) = typmod.dims() { dims.get() } else { diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index f7ae5984..068f7aae 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -389,10 +389,13 @@ pub unsafe extern "C-unwind" fn amrescan( scanner.scanning = match opfamily { Opfamily::VectorL2 | Opfamily::VectorCosine + | Opfamily::VectorIp | Opfamily::HalfvecL2 | Opfamily::HalfvecCosine - | Opfamily::VectorIp - | Opfamily::HalfvecIp => { + | Opfamily::HalfvecIp + | Opfamily::Rabitq8L2 + | Opfamily::Rabitq8Cosine + | Opfamily::Rabitq8Ip => { let mut builder = DefaultBuilder::new(opfamily); for i in 0..(*scan).numberOfOrderBys { let data = (*scan).orderByData.add(i as usize); diff --git a/src/index/vchordg/dispatch.rs b/src/index/vchordg/dispatch.rs index 2c186a7c..7ddb1793 100644 --- a/src/index/vchordg/dispatch.rs +++ b/src/index/vchordg/dispatch.rs @@ -24,6 +24,7 @@ use std::num::NonZero; use vchordg::operator::Op; use vchordg::types::*; use vector::VectorOwned; +use vector::rabitq8::Rabitq8Owned; use vector::vect::{VectBorrowed, VectOwned}; pub fn prewarm(opfamily: Opfamily, index: &R) -> String @@ -44,6 +45,12 @@ where (VectorKind::Vecf16, DistanceKind::Dot) => { vchordg::prewarm::<_, Op, Dot>>(index) } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordg::prewarm::<_, Op>(index) + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordg::prewarm::<_, Op>(index) + } } } @@ -60,15 +67,21 @@ pub fn bulkdelete( (VectorKind::Vecf32, DistanceKind::L2S) => { vchordg::bulkdelete::<_, Op, L2S>>(index, &check, &callback); } - (VectorKind::Vecf16, DistanceKind::L2S) => { - vchordg::bulkdelete::<_, Op, L2S>>(index, &check, &callback); - } (VectorKind::Vecf32, DistanceKind::Dot) => { vchordg::bulkdelete::<_, Op, Dot>>(index, &check, &callback); } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + } (VectorKind::Vecf16, DistanceKind::Dot) => { vchordg::bulkdelete::<_, Op, Dot>>(index, &check, &callback); } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordg::bulkdelete::<_, Op>(index, &check, &callback); + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordg::bulkdelete::<_, Op>(index, &check, &callback); + } } } @@ -81,15 +94,21 @@ where (VectorKind::Vecf32, DistanceKind::L2S) => { vchordg::maintain::<_, Op, L2S>>(index, &check); } - (VectorKind::Vecf16, DistanceKind::L2S) => { - vchordg::maintain::<_, Op, L2S>>(index, &check); - } (VectorKind::Vecf32, DistanceKind::Dot) => { vchordg::maintain::<_, Op, Dot>>(index, &check); } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::maintain::<_, Op, L2S>>(index, &check); + } (VectorKind::Vecf16, DistanceKind::Dot) => { vchordg::maintain::<_, Op, Dot>>(index, &check); } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordg::maintain::<_, Op>(index, &check); + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordg::maintain::<_, Op>(index, &check); + } } } @@ -102,15 +121,21 @@ where (VectorKind::Vecf32, DistanceKind::L2S) => { vchordg::build::<_, Op, L2S>>(vector_options, vchordg_options, index) } - (VectorKind::Vecf16, DistanceKind::L2S) => { - vchordg::build::<_, Op, L2S>>(vector_options, vchordg_options, index) - } (VectorKind::Vecf32, DistanceKind::Dot) => { vchordg::build::<_, Op, Dot>>(vector_options, vchordg_options, index) } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::build::<_, Op, L2S>>(vector_options, vchordg_options, index) + } (VectorKind::Vecf16, DistanceKind::Dot) => { vchordg::build::<_, Op, Dot>>(vector_options, vchordg_options, index) } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordg::build::<_, Op>(vector_options, vchordg_options, index) + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordg::build::<_, Op>(vector_options, vchordg_options, index) + } } } @@ -135,10 +160,10 @@ where make_vector_plain_prefetcher, ) } - (OwnedVector::Vecf16(unprojected), DistanceKind::L2S) => { - assert!(opfamily.vector_kind() == VectorKind::Vecf16); + (OwnedVector::Vecf32(unprojected), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); let projected = RandomProject::project(unprojected.as_borrowed()); - vchordg::insert::<_, Op, L2S>>( + vchordg::insert::<_, Op, Dot>>( index, projected.as_borrowed(), payload, @@ -147,10 +172,10 @@ where make_vector_plain_prefetcher, ) } - (OwnedVector::Vecf32(unprojected), DistanceKind::Dot) => { - assert!(opfamily.vector_kind() == VectorKind::Vecf32); + (OwnedVector::Vecf16(unprojected), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); let projected = RandomProject::project(unprojected.as_borrowed()); - vchordg::insert::<_, Op, Dot>>( + vchordg::insert::<_, Op, L2S>>( index, projected.as_borrowed(), payload, @@ -171,6 +196,28 @@ where make_vector_plain_prefetcher, ) } + (OwnedVector::Rabitq8(unprojected), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq8); + vchordg::insert::<_, Op>( + index, + unprojected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Rabitq8(unprojected), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq8); + vchordg::insert::<_, Op>( + index, + unprojected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } } } diff --git a/src/index/vchordg/opclass.rs b/src/index/vchordg/opclass.rs index 298687b2..966c6bea 100644 --- a/src/index/vchordg/opclass.rs +++ b/src/index/vchordg/opclass.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use crate::index::opclass::Sphere; use distance::Distance; @@ -31,6 +32,9 @@ pub enum Opfamily { HalfvecL2, HalfvecCosine, HalfvecIp, + Rabitq8L2, + Rabitq8Cosine, + Rabitq8Ip, } impl Opfamily { @@ -49,6 +53,12 @@ impl Opfamily { (Self::HalfvecCosine, _) => unreachable!(), (Self::HalfvecIp, B::Vecf16(x)) => O::Vecf16(x.own()), (Self::HalfvecIp, _) => unreachable!(), + (Self::Rabitq8L2, B::Rabitq8(x)) => O::Rabitq8(x.own()), + (Self::Rabitq8L2, _) => unreachable!(), + (Self::Rabitq8Cosine, B::Rabitq8(x)) => O::Rabitq8(x.function_normalize()), + (Self::Rabitq8Cosine, _) => unreachable!(), + (Self::Rabitq8Ip, B::Rabitq8(x)) => O::Rabitq8(x.own()), + (Self::Rabitq8Ip, _) => unreachable!(), } } pub unsafe fn store(self, datum: Datum) -> Option> { @@ -64,6 +74,10 @@ impl Opfamily { let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] } + Self::Rabitq8L2 | Self::Rabitq8Cosine | Self::Rabitq8Ip => { + let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Rabitq8(vector.as_borrowed())), 0)] + } }; Some(store) } @@ -83,6 +97,10 @@ impl Opfamily { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Vecf16(vector.as_borrowed())) } + Self::Rabitq8L2 | Self::Rabitq8Cosine | Self::Rabitq8Ip => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) + } }; let radius = tuple.get_by_index::(attno_2).unwrap()?; Some(Sphere { center, radius }) @@ -100,27 +118,32 @@ impl Opfamily { let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; self.input(BorrowedVector::Vecf16(vector.as_borrowed())) } + Self::Rabitq8L2 | Self::Rabitq8Cosine | Self::Rabitq8Ip => { + let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) + } }; Some(vector) } pub fn output(self, x: Distance) -> f32 { match self { - Self::VectorCosine | Self::HalfvecCosine => x.to_f32() * 0.5, - Self::VectorL2 | Self::HalfvecL2 => x.to_f32().sqrt(), - Self::VectorIp | Self::HalfvecIp => x.to_f32(), + Self::VectorCosine | Self::HalfvecCosine | Self::Rabitq8Cosine => x.to_f32() * 0.5, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 => x.to_f32().sqrt(), + Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip => x.to_f32(), } } pub const fn distance_kind(self) -> DistanceKind { match self { - Self::VectorL2 | Self::HalfvecL2 => DistanceKind::L2S, - Self::VectorCosine | Self::HalfvecCosine => DistanceKind::L2S, - Self::VectorIp | Self::HalfvecIp => DistanceKind::Dot, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 => DistanceKind::L2S, + Self::VectorCosine | Self::HalfvecCosine | Self::Rabitq8Cosine => DistanceKind::L2S, + Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip => DistanceKind::Dot, } } pub const fn vector_kind(self) -> VectorKind { match self { Self::VectorL2 | Self::VectorCosine | Self::VectorIp => VectorKind::Vecf32, Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => VectorKind::Vecf16, + Self::Rabitq8L2 | Self::Rabitq8Cosine | Self::Rabitq8Ip => VectorKind::Rabitq8, } } } @@ -161,11 +184,14 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { let result = match result_string.as_str() { "vchordg_vector_l2_ops" => Opfamily::VectorL2, + "vchordg_vector_ip_ops" => Opfamily::VectorIp, "vchordg_vector_cosine_ops" => Opfamily::VectorCosine, "vchordg_halfvec_l2_ops" => Opfamily::HalfvecL2, - "vchordg_halfvec_cosine_ops" => Opfamily::HalfvecCosine, - "vchordg_vector_ip_ops" => Opfamily::VectorIp, "vchordg_halfvec_ip_ops" => Opfamily::HalfvecIp, + "vchordg_halfvec_cosine_ops" => Opfamily::HalfvecCosine, + "vchordg_rabitq8_l2_ops" => Opfamily::Rabitq8L2, + "vchordg_rabitq8_ip_ops" => Opfamily::Rabitq8Ip, + "vchordg_rabitq8_cosine_ops" => Opfamily::Rabitq8Cosine, _ => pgrx::error!("unknown operator class"), }; diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs index b0945bf4..820242ef 100644 --- a/src/index/vchordg/scanners/default.rs +++ b/src/index/vchordg/scanners/default.rs @@ -18,7 +18,7 @@ use crate::index::scanners::{Io, SearchBuilder}; use crate::index::vchordg::dispatch::*; use crate::index::vchordg::opclass::Opfamily; use crate::index::vchordg::scanners::SearchOptions; -use crate::recorder::{Recorder, halfvec_out, vector_out}; +use crate::recorder::{Recorder, text}; use distance::Distance; use index::accessor::{Dot, L2S}; use index::bump::Bump; @@ -29,6 +29,7 @@ use vchordg::operator::{self}; use vchordg::search; use vchordg::types::{DistanceKind, OwnedVector, VectorKind}; use vector::VectorOwned; +use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; pub struct DefaultBuilder { @@ -47,12 +48,15 @@ impl SearchBuilder for DefaultBuilder { fn new(opfamily: Opfamily) -> Self { assert!(matches!( opfamily, - Opfamily::HalfvecCosine - | Opfamily::HalfvecL2 - | Opfamily::VectorCosine + Opfamily::VectorCosine | Opfamily::VectorL2 | Opfamily::VectorIp + | Opfamily::HalfvecCosine + | Opfamily::HalfvecL2 | Opfamily::HalfvecIp + | Opfamily::Rabitq8Cosine + | Opfamily::Rabitq8L2 + | Opfamily::Rabitq8Ip )); Self { opfamily, @@ -219,9 +223,9 @@ impl SearchBuilder for DefaultBuilder { ), } } - (VectorKind::Vecf16, DistanceKind::L2S) => { - type Op = operator::Op, L2S>; - let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { + (VectorKind::Vecf32, DistanceKind::Dot) => { + type Op = operator::Op, Dot>; + let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() @@ -314,9 +318,9 @@ impl SearchBuilder for DefaultBuilder { ), } } - (VectorKind::Vecf32, DistanceKind::Dot) => { - type Op = operator::Op, Dot>; - let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { + (VectorKind::Vecf16, DistanceKind::L2S) => { + type Op = operator::Op, L2S>; + let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { VectBorrowed::new(bump.alloc_slice(vector.slice())) } else { unreachable!() @@ -504,6 +508,202 @@ impl SearchBuilder for DefaultBuilder { ), } } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + type Op = operator::Op; + let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { + let vector = vector.as_borrowed(); + Rabitq8Borrowed::new( + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + bump.alloc_slice(vector.code()), + ) + } else { + unreachable!() + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + type Op = operator::Op; + let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { + let vector = vector.as_borrowed(); + Rabitq8Borrowed::new( + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + bump.alloc_slice(vector.code()), + ) + } else { + unreachable!() + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } }; let iter = if let Some(threshold) = threshold { Box::new(iter.take_while(move |(distance, _)| distance.to_f32() < threshold)) @@ -518,10 +718,13 @@ impl SearchBuilder for DefaultBuilder { if recorder.is_enabled() { match &vector { OwnedVector::Vecf32(v) => { - recorder.send(&vector_out(v.as_borrowed())); + recorder.send(&text::vector_out(v.as_borrowed())); } OwnedVector::Vecf16(v) => { - recorder.send(&halfvec_out(v.as_borrowed())); + recorder.send(&text::halfvec_out(v.as_borrowed())); + } + OwnedVector::Rabitq8(v) => { + recorder.send(&text::rabitq8_out(v.as_borrowed())); } } } diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 0c016eb4..5bf29648 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -34,6 +34,7 @@ use std::num::NonZero; use std::ops::Deref; use vchordrq::types::*; use vchordrq::{InsertChooser, MaintainChooser}; +use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; #[derive(Debug, Clone, Copy)] @@ -216,6 +217,10 @@ pub unsafe extern "C-unwind" fn ambuild( if let Err(errors) = Validate::validate(&vchordrq_options) { pgrx::error!("error while validating options: {}", errors); } + if vector_options.v == VectorKind::Rabitq8 && vchordrq_options.index.residual_quantization { + let errors = "residual_quantization is not supported for rabitq8 type"; + pgrx::error!("error while validating options: {errors}"); + } let opfamily = unsafe { opfamily(index_relation) }; let reporter = PostgresReporter { _phantom: PhantomData, @@ -250,7 +255,7 @@ pub unsafe extern "C-unwind" fn ambuild( }; for structure in structures.iter_mut() { for centroid in structure.centroids.iter_mut() { - *centroid = rabitq::rotate::rotate(centroid); + rabitq::rotate::rotate_inplace(centroid); } } reporter.phase(BuildPhase::from_code(BuildPhaseCode::Build)); @@ -1226,7 +1231,7 @@ unsafe fn options( pgrx::error!("multicolumn index is not supported"); } // get dims - let typmod = Typmod::parse_from_i32(atts[0].atttypmod).unwrap(); + let typmod = Typmod::new(atts[0].atttypmod).unwrap(); let dims = if let Some(dims) = typmod.dims() { dims.get() } else { @@ -1322,6 +1327,7 @@ fn make_internal_build( let mut x = match vector { OwnedVector::Vecf32(x) => VectOwned::normalize(x), OwnedVector::Vecf16(x) => VectOwned::normalize(x), + OwnedVector::Rabitq8(x) => Rabitq8Owned::normalize(x), }; assert_eq!( vector_options.dims, @@ -1477,6 +1483,9 @@ fn make_internal_build( let x = match vector { OwnedVector::Vecf32(x) => VectOwned::normalize(x), OwnedVector::Vecf16(x) => VectOwned::normalize(x), + OwnedVector::Rabitq8(x) => { + Rabitq8Owned::normalize(x) + } }; assert_eq!( vector_options.dims, diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 69641f13..74a7b999 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -208,6 +208,9 @@ pub unsafe extern "C-unwind" fn amcostestimate( | Opfamily::VectorCosine | Opfamily::VectorIp | Opfamily::VectorL2 + | Opfamily::Rabitq8Cosine + | Opfamily::Rabitq8Ip + | Opfamily::Rabitq8L2 ) { *index_startup_cost = 0.0; *index_total_cost = 0.0; @@ -470,7 +473,10 @@ pub unsafe extern "C-unwind" fn amrescan( | Opfamily::VectorCosine | Opfamily::HalfvecL2 | Opfamily::HalfvecIp - | Opfamily::HalfvecCosine => { + | Opfamily::HalfvecCosine + | Opfamily::Rabitq8L2 + | Opfamily::Rabitq8Ip + | Opfamily::Rabitq8Cosine => { let mut builder = DefaultBuilder::new(opfamily); for i in 0..(*scan).numberOfOrderBys { let data = (*scan).orderByData.add(i as usize); @@ -490,7 +496,7 @@ pub unsafe extern "C-unwind" fn amrescan( builder.build(index, options, fetcher, bump, recorder) })) } - Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim => { + Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim | Opfamily::Rabitq8Maxsim => { let mut builder = MaxsimBuilder::new(opfamily); for i in 0..(*scan).numberOfOrderBys { let data = (*scan).orderByData.add(i as usize); diff --git a/src/index/vchordrq/build.rs b/src/index/vchordrq/build.rs index 007fc9cb..81370b63 100644 --- a/src/index/vchordrq/build.rs +++ b/src/index/vchordrq/build.rs @@ -13,8 +13,9 @@ // Copyright (c) 2025 TensorChord Inc. use simd::{Floating, f16}; -use vector::VectorOwned; +use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; +use vector::{VectorBorrowed, VectorOwned}; pub type Normalized = Vec; @@ -42,3 +43,28 @@ impl Normalize for VectOwned { Self::new(f16::vector_from_f32(&x)) } } + +impl Normalize for Rabitq8Owned { + fn normalize(vector: Self) -> Normalized { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dims() as _); + for c in vector.code().iter().copied() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + result + } + + fn denormalize(x: Normalized) -> Self { + let code = rabitq::byte::ugly_code(x.as_slice()); + Rabitq8Owned::new( + code.0.dis_u_2, + code.0.norm_of_lattice, + code.0.sum_of_code, + f32::reduce_sum_of_abs_x(&x), + code.1, + ) + } +} diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs index cdbb0608..1a5569bd 100644 --- a/src/index/vchordrq/dispatch.rs +++ b/src/index/vchordrq/dispatch.rs @@ -28,6 +28,7 @@ use vchordrq::operator::Op; use vchordrq::types::*; use vchordrq::{FastHeap, InsertChooser, MaintainChooser}; use vector::VectorOwned; +use vector::rabitq8::Rabitq8Owned; use vector::vect::{VectBorrowed, VectOwned}; pub fn prewarm(opfamily: Opfamily, index: &R, height: i32) -> String @@ -49,6 +50,12 @@ where (VectorKind::Vecf16, DistanceKind::Dot) => { vchordrq::prewarm::<_, Op, Dot>>(index, height, make_h0_plain_prefetcher) } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordrq::prewarm::<_, Op>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordrq::prewarm::<_, Op>(index, height, make_h0_plain_prefetcher) + } } } @@ -78,6 +85,14 @@ pub fn bulkdelete( vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + } } } @@ -124,6 +139,18 @@ pub fn maintain( check, ) } + (VectorKind::Rabitq8, DistanceKind::L2S) => vchordrq::maintain::<_, Op>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ), + (VectorKind::Rabitq8, DistanceKind::Dot) => vchordrq::maintain::<_, Op>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ), }; pgrx::info!( "maintain: number_of_formerly_allocated_pages = {}", @@ -173,6 +200,18 @@ pub fn build( index, map_structures(structures, Normalize::denormalize), ), + (VectorKind::Rabitq8, DistanceKind::L2S) => vchordrq::build::<_, Op>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + (VectorKind::Rabitq8, DistanceKind::Dot) => vchordrq::build::<_, Op>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), } } @@ -271,6 +310,44 @@ pub fn insert( skip_freespaces, ) } + (OwnedVector::Rabitq8(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq8); + let key = vchordrq::insert_vector::<_, Op>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op>( + index, + payload, + vector.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Rabitq8(vector), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq8); + let key = vchordrq::insert_vector::<_, Op>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op>( + index, + payload, + vector.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } } } diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs index ae02dc6a..13bbd4b9 100644 --- a/src/index/vchordrq/opclass.rs +++ b/src/index/vchordrq/opclass.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use crate::index::opclass::Sphere; use distance::Distance; @@ -31,8 +32,12 @@ pub enum Opfamily { HalfvecL2, HalfvecIp, HalfvecCosine, + Rabitq8L2, + Rabitq8Ip, + Rabitq8Cosine, VectorMaxsim, HalfvecMaxsim, + Rabitq8Maxsim, } impl Opfamily { @@ -47,6 +52,10 @@ impl Opfamily { (B::Vecf16(x), Self::HalfvecIp | Self::HalfvecMaxsim) => O::Vecf16(x.own()), (B::Vecf16(x), Self::HalfvecCosine) => O::Vecf16(x.function_normalize()), (B::Vecf16(_), _) => unreachable!(), + (B::Rabitq8(x), Self::Rabitq8L2) => O::Rabitq8(x.own()), + (B::Rabitq8(x), Self::Rabitq8Ip | Self::Rabitq8Maxsim) => O::Rabitq8(x.own()), + (B::Rabitq8(x), Self::Rabitq8Cosine) => O::Rabitq8(x.function_normalize()), + (B::Rabitq8(_), _) => unreachable!(), } } pub unsafe fn store(self, datum: Datum) -> Option> { @@ -62,6 +71,10 @@ impl Opfamily { let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine => { + let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Rabitq8(vector.as_borrowed())), 0)] + } Self::VectorMaxsim => { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; @@ -87,6 +100,19 @@ impl Opfamily { } result } + Self::Rabitq8Maxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())), + i as u16, + )); + } + result + } }; Some(store) } @@ -106,6 +132,10 @@ impl Opfamily { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Vecf16(vector.as_borrowed())) } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine | Self::Rabitq8Maxsim => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) + } }; let radius = tuple.get_by_index::(attno_2).unwrap()?; Some(Sphere { center, radius }) @@ -123,6 +153,10 @@ impl Opfamily { let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; self.input(BorrowedVector::Vecf16(vector.as_borrowed())) } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine | Self::Rabitq8Maxsim => { + let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) + } }; Some(vector) } @@ -150,27 +184,43 @@ impl Opfamily { } result } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine | Self::Rabitq8Maxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Rabitq8(vector.as_borrowed()))); + } + result + } }; Some(vectors) } pub fn output(self, x: Distance) -> f32 { match self { - Self::VectorCosine | Self::HalfvecCosine => x.to_f32() + 1.0f32, - Self::VectorL2 | Self::HalfvecL2 => x.to_f32().sqrt(), - Self::VectorIp | Self::HalfvecIp | Self::VectorMaxsim | Self::HalfvecMaxsim => { - x.to_f32() - } + Self::VectorCosine | Self::HalfvecCosine | Self::Rabitq8Cosine => x.to_f32() + 1.0f32, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 => x.to_f32().sqrt(), + Self::VectorIp + | Self::HalfvecIp + | Self::Rabitq8Ip + | Self::VectorMaxsim + | Self::HalfvecMaxsim + | Self::Rabitq8Maxsim => x.to_f32(), } } pub const fn distance_kind(self) -> DistanceKind { match self { - Self::VectorL2 | Self::HalfvecL2 => DistanceKind::L2S, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 => DistanceKind::L2S, Self::VectorIp | Self::HalfvecIp + | Self::Rabitq8Ip | Self::VectorCosine | Self::HalfvecCosine + | Self::Rabitq8Cosine | Self::VectorMaxsim - | Self::HalfvecMaxsim => DistanceKind::Dot, + | Self::HalfvecMaxsim + | Self::Rabitq8Maxsim => DistanceKind::Dot, } } pub const fn vector_kind(self) -> VectorKind { @@ -181,6 +231,9 @@ impl Opfamily { Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { VectorKind::Vecf16 } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine | Self::Rabitq8Maxsim => { + VectorKind::Rabitq8 + } } } } @@ -226,8 +279,12 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { "vchordrq_halfvec_l2_ops" => Opfamily::HalfvecL2, "vchordrq_halfvec_ip_ops" => Opfamily::HalfvecIp, "vchordrq_halfvec_cosine_ops" => Opfamily::HalfvecCosine, + "vchordrq_rabitq8_l2_ops" => Opfamily::Rabitq8L2, + "vchordrq_rabitq8_ip_ops" => Opfamily::Rabitq8Ip, + "vchordrq_rabitq8_cosine_ops" => Opfamily::Rabitq8Cosine, "vchordrq_vector_maxsim_ops" => Opfamily::VectorMaxsim, "vchordrq_halfvec_maxsim_ops" => Opfamily::HalfvecMaxsim, + "vchordrq_rabitq8_maxsim_ops" => Opfamily::Rabitq8Maxsim, _ => pgrx::error!("unknown operator class"), }; diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index 388396f2..4897d979 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -19,7 +19,7 @@ use crate::index::vchordrq::dispatch::*; use crate::index::vchordrq::filter::filter; use crate::index::vchordrq::opclass::Opfamily; use crate::index::vchordrq::scanners::SearchOptions; -use crate::recorder::{Recorder, halfvec_out, vector_out}; +use crate::recorder::{Recorder, text}; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; use index::accessor::{Dot, L2S}; @@ -32,6 +32,7 @@ use std::num::NonZero; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; use vchordrq::{RerankMethod, default_search, how, rerank_heap, rerank_index}; use vector::VectorOwned; +use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; pub struct DefaultBuilder { @@ -56,6 +57,9 @@ impl SearchBuilder for DefaultBuilder { | Opfamily::VectorCosine | Opfamily::VectorIp | Opfamily::VectorL2 + | Opfamily::Rabitq8Cosine + | Opfamily::Rabitq8Ip + | Opfamily::Rabitq8L2 )); Self { opfamily, @@ -675,6 +679,280 @@ impl SearchBuilder for DefaultBuilder { } } } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + type Op = vchordrq::operator::Op; + let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq8(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq8(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + type Op = vchordrq::operator::Op; + let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq8(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq8(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } }; let iter = if let Some(threshold) = threshold { Box::new(iter.take_while(move |(x, _)| *x < threshold)) @@ -689,10 +967,13 @@ impl SearchBuilder for DefaultBuilder { if recorder.is_enabled() { match &vector { OwnedVector::Vecf32(v) => { - recorder.send(&vector_out(v.as_borrowed())); + recorder.send(&text::vector_out(v.as_borrowed())); } OwnedVector::Vecf16(v) => { - recorder.send(&halfvec_out(v.as_borrowed())); + recorder.send(&text::halfvec_out(v.as_borrowed())); + } + OwnedVector::Rabitq8(v) => { + recorder.send(&text::rabitq8_out(v.as_borrowed())); } } } diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 934a2e94..8f83fa8b 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -34,6 +34,7 @@ use std::num::NonZero; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; use vchordrq::{RerankMethod, how, maxsim_search, rerank_index}; use vector::VectorOwned; +use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; pub struct MaxsimBuilder { @@ -51,7 +52,7 @@ impl SearchBuilder for MaxsimBuilder { fn new(opfamily: Opfamily) -> Self { assert!(matches!( opfamily, - Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim + Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim | Opfamily::Rabitq8Maxsim )); Self { opfamily, @@ -408,6 +409,146 @@ impl SearchBuilder for MaxsimBuilder { (accu_set, rough_set, estimation_by_threshold) })) } + VectorKind::Rabitq8 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq8(vector) = vector { + vector + } else { + unreachable!() + } + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); + } + (accu_set, rough_set, estimation_by_threshold) + })) + } }; let mut updates = Vec::new(); let mut estimations = Vec::new(); diff --git a/src/recorder/mod.rs b/src/recorder/mod.rs index 77a09d4b..29f89886 100644 --- a/src/recorder/mod.rs +++ b/src/recorder/mod.rs @@ -12,15 +12,15 @@ // // Copyright (c) 2025 TensorChord Inc. -pub use text::{halfvec_out, vector_out}; pub use types::{DefaultRecorder, Recorder}; pub use worker::dump; mod hook; -mod text; mod types; mod worker; +pub mod text; + pub fn init() { hook::init(); } diff --git a/src/recorder/text.rs b/src/recorder/text.rs index 321599af..1b1a2100 100644 --- a/src/recorder/text.rs +++ b/src/recorder/text.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use simd::f16; +use vector::rabitq8::Rabitq8Borrowed; use vector::vect::VectBorrowed; pub fn vector_out(vector: VectBorrowed<'_, f32>) -> String { @@ -38,3 +39,25 @@ pub fn halfvec_out(vector: VectBorrowed<'_, f16>) -> String { result.push(']'); result } + +pub fn rabitq8_out(vector: Rabitq8Borrowed<'_>) -> String { + let mut result = String::new(); + result.push('('); + result.push_str(&vector.sum_of_x2().to_string()); + result.push(','); + result.push_str(&vector.norm_of_lattice().to_string()); + result.push(','); + result.push_str(&vector.sum_of_code().to_string()); + result.push(','); + result.push_str(&vector.sum_of_abs_x().to_string()); + result.push(')'); + result.push('['); + for x in vector.code() { + if !result.ends_with('[') { + result.push(','); + } + result.push_str(&x.to_string()); + } + result.push(']'); + result +} diff --git a/src/sql/bootstrap.sql b/src/sql/bootstrap.sql index acc6a551..743fea00 100644 --- a/src/sql/bootstrap.sql +++ b/src/sql/bootstrap.sql @@ -1,6 +1,6 @@ -- List of shell types -CREATE TYPE scalar8; CREATE TYPE sphere_vector; CREATE TYPE sphere_halfvec; -CREATE TYPE sphere_scalar8; +CREATE TYPE rabitq8; +CREATE TYPE sphere_rabitq8; diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index a1aadcdc..a7e4d923 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -1,15 +1,12 @@ --- List of data types - -CREATE TYPE scalar8 ( - INPUT = _vchord_scalar8_in, - OUTPUT = _vchord_scalar8_out, - RECEIVE = _vchord_scalar8_recv, - SEND = _vchord_scalar8_send, - TYPMOD_IN = _vchord_typmod_in_65535, - TYPMOD_OUT = _vchord_typmod_out, - STORAGE = EXTERNAL, - INTERNALLENGTH = VARIABLE, - ALIGNMENT = double +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external ); CREATE TYPE sphere_vector AS ( @@ -22,31 +19,36 @@ CREATE TYPE sphere_halfvec AS ( radius REAL ); -CREATE TYPE sphere_scalar8 AS ( - center scalar8, +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, radius REAL ); +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + -- List of operators CREATE OPERATOR <-> ( - PROCEDURE = _vchord_scalar8_operator_l2, - LEFTARG = scalar8, - RIGHTARG = scalar8, + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, COMMUTATOR = <-> ); CREATE OPERATOR <#> ( - PROCEDURE = _vchord_scalar8_operator_ip, - LEFTARG = scalar8, - RIGHTARG = scalar8, + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, COMMUTATOR = <#> ); CREATE OPERATOR <=> ( - PROCEDURE = _vchord_scalar8_operator_cosine, - LEFTARG = scalar8, - RIGHTARG = scalar8, + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, COMMUTATOR = <=> ); @@ -65,9 +67,9 @@ CREATE OPERATOR <<->> ( ); CREATE OPERATOR <<->> ( - PROCEDURE = _vchord_scalar8_sphere_l2_in, - LEFTARG = scalar8, - RIGHTARG = sphere_scalar8, + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8, COMMUTATOR = <<->> ); @@ -86,9 +88,9 @@ CREATE OPERATOR <<#>> ( ); CREATE OPERATOR <<#>> ( - PROCEDURE = _vchord_scalar8_sphere_ip_in, - LEFTARG = scalar8, - RIGHTARG = sphere_scalar8, + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8, COMMUTATOR = <<#>> ); @@ -107,9 +109,9 @@ CREATE OPERATOR <<=>> ( ); CREATE OPERATOR <<=>> ( - PROCEDURE = _vchord_scalar8_sphere_cosine_in, - LEFTARG = scalar8, - RIGHTARG = sphere_scalar8, + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8, COMMUTATOR = <<=>> ); @@ -125,6 +127,12 @@ CREATE OPERATOR @# ( RIGHTARG = halfvec[] ); +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + -- List of functions CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector @@ -133,14 +141,14 @@ IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; -CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; -CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 -IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; -CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 -IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; @@ -319,14 +327,21 @@ CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; -- List of operator classes @@ -366,6 +381,24 @@ CREATE OPERATOR CLASS halfvec_cosine_ops OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + CREATE OPERATOR CLASS vector_maxsim_ops FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, @@ -376,6 +409,11 @@ CREATE OPERATOR CLASS halfvec_maxsim_ops OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + CREATE OPERATOR CLASS vector_l2_ops FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, @@ -412,6 +450,24 @@ CREATE OPERATOR CLASS halfvec_cosine_ops OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + -- List of views CREATE VIEW vchordrq_sampled_queries AS diff --git a/tests/general/distance.slt b/tests/general/distance.slt index e03f32f6..eb798772 100644 --- a/tests/general/distance.slt +++ b/tests/general/distance.slt @@ -27,33 +27,3 @@ query I SELECT round(('[1,2,3]'::halfvec <=> '[2,3,4]'::halfvec):: numeric, 3); ---- 0.007 - -query I -SELECT round((quantize_to_scalar8('[1,2,3]'::vector) <-> quantize_to_scalar8('[2,3,4]'::vector)):: numeric, 1); ----- -1.7 - -query I -SELECT round((quantize_to_scalar8('[1,2,3]'::vector) <#> quantize_to_scalar8('[2,3,4]'::vector)):: numeric, 1); ----- --20.0 - -query I -SELECT round((quantize_to_scalar8('[1,2,3]'::vector) <=> quantize_to_scalar8('[2,3,4]'::vector)):: numeric, 2); ----- -0.01 - -query I -SELECT round((quantize_to_scalar8('[1,2,3]'::halfvec) <-> quantize_to_scalar8('[2,3,4]'::halfvec)):: numeric, 1); ----- -1.7 - -query I -SELECT round((quantize_to_scalar8('[1,2,3]'::halfvec) <#> quantize_to_scalar8('[2,3,4]'::halfvec)):: numeric, 1); ----- --20.0 - -query I -SELECT round((quantize_to_scalar8('[1,2,3]'::halfvec) <=> quantize_to_scalar8('[2,3,4]'::halfvec)):: numeric, 2); ----- -0.01 diff --git a/tests/vchordg/index_vector.slt b/tests/vchordg/index_vector.slt new file mode 100644 index 00000000..d6e15ac3 --- /dev/null +++ b/tests/vchordg/index_vector.slt @@ -0,0 +1,255 @@ +statement ok +CREATE TABLE t (index serial primary key, val vector(64)); + +statement ok +INSERT INTO t (val) +SELECT + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordg (val vector_l2_ops); + +query I +SELECT index FROM t ORDER BY val <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg (val vector_ip_ops); + +query I +SELECT index FROM t ORDER BY val <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg (val vector_cosine_ops); + +query I +SELECT index FROM t ORDER BY val <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((val::halfvec(64)) halfvec_l2_ops); + +query I +SELECT index FROM t ORDER BY val <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((val::halfvec(64)) halfvec_ip_ops); + +query I +SELECT index FROM t ORDER BY val <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((val::halfvec(64)) halfvec_cosine_ops); + +query I +SELECT index FROM t ORDER BY val <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/index_multivector.slt b/tests/vchordrq/index_multivector.slt new file mode 100644 index 00000000..f1e331e7 --- /dev/null +++ b/tests/vchordrq/index_multivector.slt @@ -0,0 +1,104 @@ +statement ok +CREATE TABLE t (index serial primary key, val vector(64)[2]); + +statement ok +CREATE FUNCTION pg_temp.quantize_to_rabitq8(val vector[]) +RETURNS rabitq8[] AS $$ + SELECT array_agg(public.quantize_to_rabitq8(x)) + FROM unnest(val) AS x; +$$ LANGUAGE sql IMMUTABLE; + +statement ok +INSERT INTO t (val) +SELECT + ARRAY[ + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector), + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) + ] +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordrq (val vector_maxsim_ops); + +query I +SELECT index FROM t ORDER BY +val @# ARRAY[ + array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector, + array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector +] +LIMIT 10; +---- +1207 +919 +1639 +1821 +174 +79 +1076 +1125 +239 +194 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)[2]) halfvec_maxsim_ops); + +query I +SELECT index FROM t ORDER BY +val::halfvec(64)[] @# ARRAY[ + array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec, + array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec +] +LIMIT 10; +---- +1207 +919 +1639 +1821 +174 +79 +1076 +1125 +239 +194 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((pg_temp.quantize_to_rabitq8(val)::rabitq8(64)[2]) rabitq8_maxsim_ops); + +query I +SELECT index FROM t ORDER BY +pg_temp.quantize_to_rabitq8(val)::rabitq8(64)[] @# ARRAY[ + quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector), + quantize_to_rabitq8(array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) +] +LIMIT 10; +---- +1207 +919 +1821 +1639 +174 +79 +1076 +1125 +239 +194 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/index_vector.slt b/tests/vchordrq/index_vector.slt new file mode 100644 index 00000000..dd2ea412 --- /dev/null +++ b/tests/vchordrq/index_vector.slt @@ -0,0 +1,255 @@ +statement ok +CREATE TABLE t (index serial primary key, val vector(64)); + +statement ok +INSERT INTO t (val) +SELECT + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordrq (val vector_l2_ops); + +query I +SELECT index FROM t ORDER BY val <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq (val vector_ip_ops); + +query I +SELECT index FROM t ORDER BY val <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq (val vector_cosine_ops); + +query I +SELECT index FROM t ORDER BY val <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)) halfvec_l2_ops); + +query I +SELECT index FROM t ORDER BY val <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)) halfvec_ip_ops); + +query I +SELECT index FROM t ORDER BY val <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)) halfvec_cosine_ops); + +query I +SELECT index FROM t ORDER BY val <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; From 1c6dc09557c16787d2a0bedb9b7cbc383ad6368c Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 4 Dec 2025 17:04:50 +0800 Subject: [PATCH 254/324] feat: add `simd::halfbyte` (#387) Signed-off-by: usamoi --- crates/index/src/accessor.rs | 4 +- crates/rabitq/src/bit.rs | 4 +- crates/rabitq/src/byte.rs | 2 +- crates/simd/src/{u8.rs => byte.rs} | 128 +++++---- crates/simd/src/halfbyte.rs | 404 +++++++++++++++++++++++++++++ crates/simd/src/lib.rs | 3 +- 6 files changed, 474 insertions(+), 71 deletions(-) rename crates/simd/src/{u8.rs => byte.rs} (88%) create mode 100644 crates/simd/src/halfbyte.rs diff --git a/crates/index/src/accessor.rs b/crates/index/src/accessor.rs index 0d1a9eb1..cc1bdbbb 100644 --- a/crates/index/src/accessor.rs +++ b/crates/index/src/accessor.rs @@ -409,7 +409,7 @@ impl Accessor2 for DimensionDistanceAccessor for DimensionDistanceAccessor BinaryLut { let (k, b, qvector) = simd::quantize::quantize(vector, ((1 << BITS) - 1) as f32); let qvector_sum = if vector.len() <= (65535_usize / ((1 << BITS) - 1)) { - simd::u8::reduce_sum_of_x_as_u16(&qvector) as f32 + simd::byte::reduce_sum_of_x_as_u16(&qvector) as f32 } else { - simd::u8::reduce_sum_of_x_as_u32(&qvector) as f32 + simd::byte::reduce_sum_of_x_as_u32(&qvector) as f32 }; ( BinaryLutMetadata { diff --git a/crates/rabitq/src/byte.rs b/crates/rabitq/src/byte.rs index 9f145dbe..49b36028 100644 --- a/crates/rabitq/src/byte.rs +++ b/crates/rabitq/src/byte.rs @@ -44,7 +44,7 @@ pub mod binary { } pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { - simd::u8::reduce_sum_of_x_as_u32_y_as_u32(x, y) + simd::byte::reduce_sum_of_x_as_u32_y_as_u32(x, y) } pub fn half_process_dot( diff --git a/crates/simd/src/u8.rs b/crates/simd/src/byte.rs similarity index 88% rename from crates/simd/src/u8.rs rename to crates/simd/src/byte.rs index 2985413a..54081f96 100644 --- a/crates/simd/src/u8.rs +++ b/crates/simd/src/byte.rs @@ -22,8 +22,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let lo = _mm512_set1_epi16(0x00ff_u16 as i16); - let hi = _mm512_set1_epi16(0xff00_u16 as i16); + let lo = _mm512_set1_epi16(0x00ff_i16); let mut _0 = _mm512_setzero_si512(); let mut _1 = _mm512_setzero_si512(); let mut _2 = _mm512_setzero_si512(); @@ -31,38 +30,39 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { while n >= 64 { let x = unsafe { _mm512_loadu_epi8(a.cast()) }; let y = unsafe { _mm512_loadu_epi8(b.cast()) }; - let x_l = _mm512_and_si512(x, lo); - let x_h = _mm512_and_si512(_mm512_srli_epi16(_mm512_and_si512(x, hi), 8), lo); - let y_l = _mm512_and_si512(y, lo); - let y_h = _mm512_and_si512(_mm512_srli_epi16(_mm512_and_si512(y, hi), 8), lo); - let l = _mm512_mullo_epi16(x_l, y_l); - let h = _mm512_mullo_epi16(x_h, y_h); + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_srli_epi16(x, 8); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_srli_epi16(y, 8); + let z_0 = _mm512_mullo_epi16(x_0, y_0); + let z_1 = _mm512_mullo_epi16(x_1, y_1); a = unsafe { a.add(64) }; b = unsafe { b.add(64) }; n -= 64; - _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(l, 0))); - _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(l, 1))); - _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(h, 0))); - _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(h, 1))); + _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 0))); + _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 1))); + _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 0))); + _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 1))); } if n > 0 { let mask = _bzhi_u64(0xffffffffffffffff, n as u32); let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; - let x_l = _mm512_and_si512(x, lo); - let x_h = _mm512_and_si512(_mm512_srli_epi16(_mm512_and_si512(x, hi), 8), lo); - let y_l = _mm512_and_si512(y, lo); - let y_h = _mm512_and_si512(_mm512_srli_epi16(_mm512_and_si512(y, hi), 8), lo); - let l = _mm512_mullo_epi16(x_l, y_l); - let h = _mm512_mullo_epi16(x_h, y_h); - _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(l, 0))); - _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(l, 1))); - _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(h, 0))); - _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(h, 1))); - } - let _5 = _mm512_add_epi32(_0, _1); - let _6 = _mm512_add_epi32(_2, _3); - _mm512_reduce_add_epi32(_mm512_add_epi32(_5, _6)) as u32 + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_srli_epi16(x, 8); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_srli_epi16(y, 8); + let z_0 = _mm512_mullo_epi16(x_0, y_0); + let z_1 = _mm512_mullo_epi16(x_1, y_1); + _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 0))); + _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 1))); + _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 0))); + _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 1))); + } + let r_0 = _mm512_add_epi32(_0, _2); + let r_1 = _mm512_add_epi32(_1, _3); + let r_2 = _mm512_add_epi32(r_0, r_1); + _mm512_reduce_add_epi32(r_2) as u32 } #[cfg(all(target_arch = "x86_64", test, not(miri)))] @@ -101,8 +101,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let lo = _mm256_set1_epi16(0x00ff_u16 as i16); - let hi = _mm256_set1_epi16(0xff00_u16 as i16); + let lo = _mm256_set1_epi16(0x00ff_i16); let mut _0 = _mm256_setzero_si256(); let mut _1 = _mm256_setzero_si256(); let mut _2 = _mm256_setzero_si256(); @@ -110,23 +109,23 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { while n >= 32 { let x = unsafe { _mm256_loadu_si256(a.cast()) }; let y = unsafe { _mm256_loadu_si256(b.cast()) }; - let x_l = _mm256_and_si256(x, lo); - let x_h = _mm256_and_si256(_mm256_srli_epi16(_mm256_and_si256(x, hi), 8), lo); - let y_l = _mm256_and_si256(y, lo); - let y_h = _mm256_and_si256(_mm256_srli_epi16(_mm256_and_si256(y, hi), 8), lo); - let l = _mm256_mullo_epi16(x_l, y_l); - let h = _mm256_mullo_epi16(x_h, y_h); + let x_0 = _mm256_and_si256(x, lo); + let x_1 = _mm256_srli_epi16(x, 8); + let y_0 = _mm256_and_si256(y, lo); + let y_1 = _mm256_srli_epi16(y, 8); + let z_0 = _mm256_mullo_epi16(x_0, y_0); + let z_1 = _mm256_mullo_epi16(x_1, y_1); a = unsafe { a.add(32) }; b = unsafe { b.add(32) }; n -= 32; - _0 = _mm256_add_epi32(_0, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(l, 0))); - _1 = _mm256_add_epi32(_1, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(l, 1))); - _2 = _mm256_add_epi32(_2, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(h, 0))); - _3 = _mm256_add_epi32(_3, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(h, 1))); + _0 = _mm256_add_epi32(_0, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_0, 0))); + _1 = _mm256_add_epi32(_1, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_0, 1))); + _2 = _mm256_add_epi32(_2, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_1, 0))); + _3 = _mm256_add_epi32(_3, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_1, 1))); } let mut sum = emulate_mm256_reduce_add_epi32(_mm256_add_epi32( - _mm256_add_epi32(_0, _1), - _mm256_add_epi32(_2, _3), + _mm256_add_epi32(_0, _2), + _mm256_add_epi32(_1, _3), )) as u32; // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { @@ -176,8 +175,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let lo = _mm_set1_epi16(0x00ff_u16 as i16); - let hi = _mm_set1_epi16(0xff00_u16 as i16); + let lo = _mm_set1_epi16(0x00ff_i16); let mut _0 = _mm_setzero_si128(); let mut _1 = _mm_setzero_si128(); let mut _2 = _mm_setzero_si128(); @@ -185,19 +183,19 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { while n >= 16 { let x = unsafe { _mm_loadu_si128(a.cast()) }; let y = unsafe { _mm_loadu_si128(b.cast()) }; - let x_l = _mm_and_si128(x, lo); - let x_h = _mm_and_si128(_mm_srli_epi16(_mm_and_si128(x, hi), 8), lo); - let y_l = _mm_and_si128(y, lo); - let y_h = _mm_and_si128(_mm_srli_epi16(_mm_and_si128(y, hi), 8), lo); - let l = _mm_mullo_epi16(x_l, y_l); - let h = _mm_mullo_epi16(x_h, y_h); + let x_0 = _mm_and_si128(x, lo); + let x_1 = _mm_srli_epi16(x, 8); + let y_0 = _mm_and_si128(y, lo); + let y_1 = _mm_srli_epi16(y, 8); + let z_0 = _mm_mullo_epi16(x_0, y_0); + let z_1 = _mm_mullo_epi16(x_1, y_1); a = unsafe { a.add(16) }; b = unsafe { b.add(16) }; n -= 16; - _0 = _mm_add_epi32(_0, _mm_cvtepu16_epi32(l)); - _1 = _mm_add_epi32(_1, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(l, l))); - _2 = _mm_add_epi32(_2, _mm_cvtepu16_epi32(h)); - _3 = _mm_add_epi32(_3, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(h, h))); + _0 = _mm_add_epi32(_0, _mm_cvtepu16_epi32(z_0)); + _1 = _mm_add_epi32(_1, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_0, z_0))); + _2 = _mm_add_epi32(_2, _mm_cvtepu16_epi32(z_1)); + _3 = _mm_add_epi32(_3, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_1, z_1))); } let mut sum = emulate_mm_reduce_add_epi32(_mm_add_epi32( _mm_add_epi32(_0, _1), @@ -256,23 +254,23 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let mut _2 = vdupq_n_u32(0); let mut _3 = vdupq_n_u32(0); while n >= 16 { - let x = vreinterpretq_u16_u8(unsafe { vld1q_u8(a.cast()) }); - let y = vreinterpretq_u16_u8(unsafe { vld1q_u8(b.cast()) }); - let x_l = vandq_u16(x, lo); - let x_h = vshrq_n_u16(x, 8); - let y_l = vandq_u16(y, lo); - let y_h = vshrq_n_u16(y, 8); - let l = vmulq_u16(x_l, y_l); - let h = vmulq_u16(x_h, y_h); + let x = unsafe { vld1q_u16(a.cast()) }; + let y = unsafe { vld1q_u16(b.cast()) }; + let x_0 = vandq_u16(x, lo); + let x_1 = vshrq_n_u16(x, 8); + let y_0 = vandq_u16(y, lo); + let y_1 = vshrq_n_u16(y, 8); + let z_0 = vmulq_u16(x_0, y_0); + let z_1 = vmulq_u16(x_1, y_1); a = unsafe { a.add(16) }; b = unsafe { b.add(16) }; n -= 16; - _0 = vaddq_u32(_0, vmovl_u16(vget_low_u16(l))); - _1 = vaddq_u32(_1, vmovl_u16(vget_high_u16(l))); - _2 = vaddq_u32(_2, vmovl_u16(vget_low_u16(h))); - _3 = vaddq_u32(_3, vmovl_u16(vget_high_u16(h))); + _0 = vaddq_u32(_0, vmovl_u16(vget_low_u16(z_0))); + _1 = vaddq_u32(_1, vmovl_u16(vget_high_u16(z_0))); + _2 = vaddq_u32(_2, vmovl_u16(vget_low_u16(z_1))); + _3 = vaddq_u32(_3, vmovl_u16(vget_high_u16(z_1))); } - let mut sum = vaddvq_u32(vaddq_u32(vaddq_u32(_0, _1), vaddq_u32(_2, _3))); + let mut sum = vaddvq_u32(vaddq_u32(vaddq_u32(_0, _2), vaddq_u32(_1, _3))); // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { let x = unsafe { a.read() }; diff --git a/crates/simd/src/halfbyte.rs b/crates/simd/src/halfbyte.rs new file mode 100644 index 00000000..008c7fcd --- /dev/null +++ b/crates/simd/src/halfbyte.rs @@ -0,0 +1,404 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +mod reduce_sum_of_x_as_u32_y_as_u32 { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_as_u32_y_as_u32_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm512_set1_epi16(0x000f_i16); + let mut _0 = _mm512_setzero_si512(); + let mut _1 = _mm512_setzero_si512(); + let mut _2 = _mm512_setzero_si512(); + let mut _3 = _mm512_setzero_si512(); + let mut _4 = _mm512_setzero_si512(); + let mut _5 = _mm512_setzero_si512(); + let mut _6 = _mm512_setzero_si512(); + let mut _7 = _mm512_setzero_si512(); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + let y = unsafe { _mm512_loadu_epi8(b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); + let x_2 = _mm512_and_si512(_mm512_srli_epi16(x, 8), lo); + let x_3 = _mm512_srli_epi16(x, 12); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); + let y_2 = _mm512_and_si512(_mm512_srli_epi16(y, 8), lo); + let y_3 = _mm512_srli_epi16(y, 12); + let z_0 = _mm512_mullo_epi16(x_0, y_0); + let z_1 = _mm512_mullo_epi16(x_1, y_1); + let z_2 = _mm512_mullo_epi16(x_2, y_2); + let z_3 = _mm512_mullo_epi16(x_3, y_3); + a = unsafe { a.add(64) }; + b = unsafe { b.add(64) }; + n -= 64; + _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 0))); + _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 1))); + _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 0))); + _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 1))); + _4 = _mm512_add_epi32(_4, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_2, 0))); + _5 = _mm512_add_epi32(_5, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_2, 1))); + _6 = _mm512_add_epi32(_6, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_3, 0))); + _7 = _mm512_add_epi32(_7, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_3, 1))); + } + if n > 0 { + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); + let x_2 = _mm512_and_si512(_mm512_srli_epi16(x, 8), lo); + let x_3 = _mm512_srli_epi16(x, 12); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); + let y_2 = _mm512_and_si512(_mm512_srli_epi16(y, 8), lo); + let y_3 = _mm512_srli_epi16(y, 12); + let z_0 = _mm512_mullo_epi16(x_0, y_0); + let z_1 = _mm512_mullo_epi16(x_1, y_1); + let z_2 = _mm512_mullo_epi16(x_2, y_2); + let z_3 = _mm512_mullo_epi16(x_3, y_3); + _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 0))); + _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 1))); + _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 0))); + _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 1))); + _4 = _mm512_add_epi32(_4, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_2, 0))); + _5 = _mm512_add_epi32(_5, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_2, 1))); + _6 = _mm512_add_epi32(_6, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_3, 0))); + _7 = _mm512_add_epi32(_7, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_3, 1))); + } + let r_0 = _mm512_add_epi32(_0, _4); + let r_1 = _mm512_add_epi32(_1, _5); + let r_2 = _mm512_add_epi32(_2, _6); + let r_3 = _mm512_add_epi32(_3, _7); + let r_4 = _mm512_add_epi32(r_0, r_2); + let r_5 = _mm512_add_epi32(r_1, r_3); + let r_6 = _mm512_add_epi32(r_4, r_5); + _mm512_reduce_add_epi32(r_6) as u32 + } + + #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_as_u32_y_as_u32_v4_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v4(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_as_u32_y_as_u32_v3(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::emulate_mm256_reduce_add_epi32; + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm256_set1_epi16(0x000f_i16); + let mut _0 = _mm256_setzero_si256(); + let mut _1 = _mm256_setzero_si256(); + let mut _2 = _mm256_setzero_si256(); + let mut _3 = _mm256_setzero_si256(); + let mut _4 = _mm256_setzero_si256(); + let mut _5 = _mm256_setzero_si256(); + let mut _6 = _mm256_setzero_si256(); + let mut _7 = _mm256_setzero_si256(); + while n >= 32 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let x_0 = _mm256_and_si256(x, lo); + let x_1 = _mm256_and_si256(_mm256_srli_epi16(x, 4), lo); + let x_2 = _mm256_and_si256(_mm256_srli_epi16(x, 8), lo); + let x_3 = _mm256_srli_epi16(x, 12); + let y_0 = _mm256_and_si256(y, lo); + let y_1 = _mm256_and_si256(_mm256_srli_epi16(y, 4), lo); + let y_2 = _mm256_and_si256(_mm256_srli_epi16(y, 8), lo); + let y_3 = _mm256_srli_epi16(y, 12); + let z_0 = _mm256_mullo_epi16(x_0, y_0); + let z_1 = _mm256_mullo_epi16(x_1, y_1); + let z_2 = _mm256_mullo_epi16(x_2, y_2); + let z_3 = _mm256_mullo_epi16(x_3, y_3); + a = unsafe { a.add(32) }; + b = unsafe { b.add(32) }; + n -= 32; + _0 = _mm256_add_epi32(_0, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_0, 0))); + _1 = _mm256_add_epi32(_1, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_0, 1))); + _2 = _mm256_add_epi32(_2, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_1, 0))); + _3 = _mm256_add_epi32(_3, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_1, 1))); + _4 = _mm256_add_epi32(_4, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_2, 0))); + _5 = _mm256_add_epi32(_5, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_2, 1))); + _6 = _mm256_add_epi32(_6, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_3, 0))); + _7 = _mm256_add_epi32(_7, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_3, 1))); + } + let mut sum = emulate_mm256_reduce_add_epi32(_mm256_add_epi32( + _mm256_add_epi32(_mm256_add_epi32(_0, _4), _mm256_add_epi32(_1, _5)), + _mm256_add_epi32(_mm256_add_epi32(_2, _6), _mm256_add_epi32(_3, _7)), + )) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + sum += (x as u32 & 0xf) * (y as u32 & 0xf); + sum += (x as u32 >> 4) * (y as u32 >> 4); + } + sum + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_as_u32_y_as_u32_v3_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v3(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_x_as_u32_y_as_u32_v2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::emulate_mm_reduce_add_epi32; + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm_set1_epi16(0x000f_i16); + let mut _0 = _mm_setzero_si128(); + let mut _1 = _mm_setzero_si128(); + let mut _2 = _mm_setzero_si128(); + let mut _3 = _mm_setzero_si128(); + let mut _4 = _mm_setzero_si128(); + let mut _5 = _mm_setzero_si128(); + let mut _6 = _mm_setzero_si128(); + let mut _7 = _mm_setzero_si128(); + while n >= 16 { + let x = unsafe { _mm_loadu_si128(a.cast()) }; + let y = unsafe { _mm_loadu_si128(b.cast()) }; + let x_0 = _mm_and_si128(x, lo); + let x_1 = _mm_and_si128(_mm_srli_epi16(x, 4), lo); + let x_2 = _mm_and_si128(_mm_srli_epi16(x, 8), lo); + let x_3 = _mm_srli_epi16(x, 12); + let y_0 = _mm_and_si128(y, lo); + let y_1 = _mm_and_si128(_mm_srli_epi16(y, 4), lo); + let y_2 = _mm_and_si128(_mm_srli_epi16(y, 8), lo); + let y_3 = _mm_srli_epi16(y, 12); + let z_0 = _mm_mullo_epi16(x_0, y_0); + let z_1 = _mm_mullo_epi16(x_1, y_1); + let z_2 = _mm_mullo_epi16(x_2, y_2); + let z_3 = _mm_mullo_epi16(x_3, y_3); + a = unsafe { a.add(16) }; + b = unsafe { b.add(16) }; + n -= 16; + _0 = _mm_add_epi32(_0, _mm_cvtepu16_epi32(z_0)); + _1 = _mm_add_epi32(_1, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_0, z_0))); + _2 = _mm_add_epi32(_2, _mm_cvtepu16_epi32(z_1)); + _3 = _mm_add_epi32(_3, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_1, z_1))); + _4 = _mm_add_epi32(_4, _mm_cvtepu16_epi32(z_2)); + _5 = _mm_add_epi32(_5, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_2, z_2))); + _6 = _mm_add_epi32(_6, _mm_cvtepu16_epi32(z_3)); + _7 = _mm_add_epi32(_7, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_3, z_3))); + } + let mut sum = emulate_mm_reduce_add_epi32(_mm_add_epi32( + _mm_add_epi32(_mm_add_epi32(_0, _4), _mm_add_epi32(_1, _5)), + _mm_add_epi32(_mm_add_epi32(_2, _6), _mm_add_epi32(_3, _7)), + )) as u32; + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + sum += (x as u32 & 0xf) * (y as u32 & 0xf); + sum += (x as u32 >> 4) * (y as u32 >> 4); + } + sum + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_as_u32_y_as_u32_v2_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x_as_u32_y_as_u32_a2(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::aarch64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = vdupq_n_u16(0x000f_u16); + let mut _0 = vdupq_n_u32(0); + let mut _1 = vdupq_n_u32(0); + let mut _2 = vdupq_n_u32(0); + let mut _3 = vdupq_n_u32(0); + let mut _4 = vdupq_n_u32(0); + let mut _5 = vdupq_n_u32(0); + let mut _6 = vdupq_n_u32(0); + let mut _7 = vdupq_n_u32(0); + while n >= 16 { + let x = unsafe { vld1q_u16(a.cast()) }; + let y = unsafe { vld1q_u16(b.cast()) }; + let x_0 = vandq_u16(x, lo); + let x_1 = vandq_u16(vshrq_n_u16(x, 4), lo); + let x_2 = vandq_u16(vshrq_n_u16(x, 8), lo); + let x_3 = vshrq_n_u16(x, 12); + let y_0 = vandq_u16(y, lo); + let y_1 = vandq_u16(vshrq_n_u16(y, 4), lo); + let y_2 = vandq_u16(vshrq_n_u16(y, 8), lo); + let y_3 = vshrq_n_u16(y, 12); + let z_0 = vmulq_u16(x_0, y_0); + let z_1 = vmulq_u16(x_1, y_1); + let z_2 = vmulq_u16(x_2, y_2); + let z_3 = vmulq_u16(x_3, y_3); + a = unsafe { a.add(16) }; + b = unsafe { b.add(16) }; + n -= 16; + _0 = vaddq_u32(_0, vmovl_u16(vget_low_u16(z_0))); + _1 = vaddq_u32(_1, vmovl_u16(vget_high_u16(z_0))); + _2 = vaddq_u32(_2, vmovl_u16(vget_low_u16(z_1))); + _3 = vaddq_u32(_3, vmovl_u16(vget_high_u16(z_1))); + _4 = vaddq_u32(_4, vmovl_u16(vget_low_u16(z_2))); + _5 = vaddq_u32(_5, vmovl_u16(vget_high_u16(z_2))); + _6 = vaddq_u32(_6, vmovl_u16(vget_low_u16(z_3))); + _7 = vaddq_u32(_7, vmovl_u16(vget_high_u16(z_3))); + } + let mut sum = vaddvq_u32(vaddq_u32( + vaddq_u32(vaddq_u32(_0, _4), vaddq_u32(_1, _5)), + vaddq_u32(vaddq_u32(_2, _6), vaddq_u32(_3, _7)), + )); + // this hint is used to disable loop unrolling + while std::hint::black_box(n) > 0 { + let x = unsafe { a.read() }; + let y = unsafe { b.read() }; + a = unsafe { a.add(1) }; + b = unsafe { b.add(1) }; + n -= 1; + sum += (x as u32 & 0xf) * (y as u32 & 0xf); + sum += (x as u32 >> 4) * (y as u32 >> 4); + } + sum + } + + #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[test] + fn reduce_sum_of_x_as_u32_y_as_u32_a2_test() { + use rand::Rng; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { + assert_eq!(s.len(), t.len()); + let n = s.len(); + let mut result = 0; + for i in 0..n { + result += (s[i] as u32 & 0xf) * (t[i] as u32 & 0xf); + result += (s[i] as u32 >> 4) * (t[i] as u32 >> 4); + } + result + } +} + +#[inline(always)] +pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { + reduce_sum_of_x_as_u32_y_as_u32::reduce_sum_of_x_as_u32_y_as_u32(s, t) +} diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 6798a9b1..2becec62 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -29,11 +29,12 @@ mod floating_f16; mod floating_f32; pub mod bit; +pub mod byte; pub mod fast_scan; pub mod fht; +pub mod halfbyte; pub mod quantize; pub mod rotate; -pub mod u8; #[cfg(not(feature = "nightly_f16"))] pub use half::f16; From 86042f3a15f21466c359ee3e0e44636f637dffe7 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 9 Dec 2025 16:21:57 +0800 Subject: [PATCH 255/324] feat: rabitq4 (#388) Signed-off-by: usamoi --- crates/index/src/accessor.rs | 87 ++++- crates/k_means/src/index.rs | 3 +- crates/rabitq/src/bit.rs | 10 +- crates/rabitq/src/bits.rs | 18 +- crates/rabitq/src/byte.rs | 28 +- crates/rabitq/src/extended.rs | 38 +- crates/rabitq/src/halfbyte.rs | 92 +++++ crates/rabitq/src/lib.rs | 1 + crates/vchordg/src/build.rs | 2 +- crates/vchordg/src/insert.rs | 19 +- crates/vchordg/src/maintain.rs | 8 +- crates/vchordg/src/operator.rs | 238 +++++++++--- crates/vchordg/src/search.rs | 15 +- crates/vchordg/src/tuples.rs | 12 +- crates/vchordg/src/types.rs | 18 +- crates/vchordrq/src/build.rs | 8 +- crates/vchordrq/src/cost.rs | 6 +- crates/vchordrq/src/insert.rs | 29 +- crates/vchordrq/src/maintain.rs | 6 +- crates/vchordrq/src/operator.rs | 359 ++++++++++++++---- crates/vchordrq/src/rerank.rs | 10 +- crates/vchordrq/src/search.rs | 30 +- crates/vchordrq/src/tuples.rs | 12 +- crates/vchordrq/src/types.rs | 17 +- crates/vector/src/bvect.rs | 72 ++-- crates/vector/src/lib.rs | 3 +- crates/vector/src/rabitq4.rs | 384 ++++++++++++++++++++ crates/vector/src/rabitq8.rs | 139 ++++--- crates/vector/src/svect.rs | 70 ++-- crates/vector/src/vect.rs | 4 +- src/datatype/binary_rabitq4.rs | 96 +++++ src/datatype/binary_rabitq8.rs | 23 +- src/datatype/functions_rabitq4.rs | 56 +++ src/datatype/functions_rabitq8.rs | 33 +- src/datatype/memory_halfvec.rs | 14 +- src/datatype/memory_rabitq4.rs | 276 ++++++++++++++ src/datatype/memory_rabitq8.rs | 28 +- src/datatype/memory_vector.rs | 14 +- src/datatype/mod.rs | 6 + src/datatype/operators_halfvec.rs | 6 +- src/datatype/operators_rabitq4.rs | 139 +++++++ src/datatype/operators_rabitq8.rs | 12 +- src/datatype/operators_vector.rs | 6 +- src/datatype/text_rabitq4.rs | 160 ++++++++ src/datatype/text_rabitq8.rs | 19 +- src/datatype/typmod.rs | 4 +- src/datatype/typmod_rabitq4.rs | 37 ++ src/index/opclass.rs | 35 ++ src/index/vchordg/am/am_build.rs | 8 +- src/index/vchordg/am/mod.rs | 5 +- src/index/vchordg/dispatch.rs | 47 +++ src/index/vchordg/opclass.rs | 50 ++- src/index/vchordg/scanners/default.rs | 213 ++++++++++- src/index/vchordrq/am/am_build.rs | 53 +-- src/index/vchordrq/am/mod.rs | 18 +- src/index/vchordrq/build.rs | 50 ++- src/index/vchordrq/dispatch.rs | 77 ++++ src/index/vchordrq/opclass.rs | 73 +++- src/index/vchordrq/scanners/default.rs | 281 ++++++++++++++ src/index/vchordrq/scanners/maxsim.rs | 146 +++++++- src/recorder/text.rs | 25 +- src/sql/bootstrap.sql | 2 + src/sql/finalize.sql | 122 +++++++ tests/vchordg/index_vector.slt | 129 +------ tests/vchordg/index_vector_rabitq.slt | 216 +++++++++++ tests/vchordrq/index_multivector.slt | 33 +- tests/vchordrq/index_multivector_rabitq.slt | 89 +++++ tests/vchordrq/index_vector.slt | 129 +------ tests/vchordrq/index_vector_rabitq.slt | 237 ++++++++++++ 69 files changed, 3918 insertions(+), 787 deletions(-) create mode 100644 crates/rabitq/src/halfbyte.rs create mode 100644 crates/vector/src/rabitq4.rs create mode 100644 src/datatype/binary_rabitq4.rs create mode 100644 src/datatype/functions_rabitq4.rs create mode 100644 src/datatype/memory_rabitq4.rs create mode 100644 src/datatype/operators_rabitq4.rs create mode 100644 src/datatype/text_rabitq4.rs create mode 100644 src/datatype/typmod_rabitq4.rs create mode 100644 tests/vchordg/index_vector_rabitq.slt create mode 100644 tests/vchordrq/index_multivector_rabitq.slt create mode 100644 tests/vchordrq/index_vector_rabitq.slt diff --git a/crates/index/src/accessor.rs b/crates/index/src/accessor.rs index cc1bdbbb..d2046a3b 100644 --- a/crates/index/src/accessor.rs +++ b/crates/index/src/accessor.rs @@ -16,6 +16,7 @@ use distance::Distance; use rabitq::byte::CodeMetadata; use simd::{Floating, f16}; use std::marker::PhantomData; +use vector::rabitq4::Rabitq4Owned; use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; @@ -388,34 +389,97 @@ impl Accessor2 for DistanceAccessor, Dot> { } } +pub trait DefaultWithDimension { + fn default_with_dimension(dim: u32) -> Self; +} + +impl DefaultWithDimension for T { + fn default_with_dimension(_dim: u32) -> Self { + Self::default() + } +} + +#[derive(Debug)] +pub struct ByteDistanceAccessor(u32, u32, PhantomData V>, PhantomData D>); + +impl DefaultWithDimension for ByteDistanceAccessor { + #[inline(always)] + fn default_with_dimension(dim: u32) -> Self { + Self(dim, 0, PhantomData, PhantomData) + } +} + +impl Accessor2 for ByteDistanceAccessor { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[u8], input: &[u8]) { + self.1 += rabitq::byte::binary::accumulate(target, input); + } + + #[inline(always)] + fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { + Distance::from_f32( + rabitq::byte::binary::half_process_l2s( + self.0, + self.1, + CodeMetadata::from_array(std::array::from_fn(|i| target[i])), + CodeMetadata::from_array(std::array::from_fn(|i| input[i])), + ) + .0, + ) + } +} + +impl Accessor2 for ByteDistanceAccessor { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[u8], input: &[u8]) { + self.1 += rabitq::byte::binary::accumulate(target, input); + } + + #[inline(always)] + fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { + Distance::from_f32( + rabitq::byte::binary::half_process_dot( + self.0, + self.1, + CodeMetadata::from_array(std::array::from_fn(|i| target[i])), + CodeMetadata::from_array(std::array::from_fn(|i| input[i])), + ) + .0, + ) + } +} + #[derive(Debug)] -pub struct DimensionDistanceAccessor( +pub struct HalfbyteDistanceAccessor( u32, u32, PhantomData V>, PhantomData D>, ); -impl Default for DimensionDistanceAccessor { +impl DefaultWithDimension for HalfbyteDistanceAccessor { #[inline(always)] - fn default() -> Self { - Self(0, 0, PhantomData, PhantomData) + fn default_with_dimension(dim: u32) -> Self { + Self(dim, 0, PhantomData, PhantomData) } } -impl Accessor2 for DimensionDistanceAccessor { +impl Accessor2 for HalfbyteDistanceAccessor { type Output = Distance; #[inline(always)] fn push(&mut self, target: &[u8], input: &[u8]) { - self.0 += target.len() as u32; - self.1 += simd::byte::reduce_sum_of_x_as_u32_y_as_u32(target, input); + self.1 += rabitq::halfbyte::binary::accumulate(target, input); } #[inline(always)] fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { Distance::from_f32( - rabitq::byte::binary::half_process_l2( + rabitq::halfbyte::binary::half_process_l2s( self.0, self.1, CodeMetadata::from_array(std::array::from_fn(|i| target[i])), @@ -426,19 +490,18 @@ impl Accessor2 for DimensionDistanceAccessor for DimensionDistanceAccessor { +impl Accessor2 for HalfbyteDistanceAccessor { type Output = Distance; #[inline(always)] fn push(&mut self, target: &[u8], input: &[u8]) { - self.0 += target.len() as u32; - self.1 += simd::byte::reduce_sum_of_x_as_u32_y_as_u32(target, input); + self.1 += rabitq::halfbyte::binary::accumulate(target, input); } #[inline(always)] fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { Distance::from_f32( - rabitq::byte::binary::half_process_dot( + rabitq::halfbyte::binary::half_process_dot( self.0, self.1, CodeMetadata::from_array(std::array::from_fn(|i| target[i])), diff --git a/crates/k_means/src/index.rs b/crates/k_means/src/index.rs index 612e3365..2a6ddb9c 100644 --- a/crates/k_means/src/index.rs +++ b/crates/k_means/src/index.rs @@ -61,7 +61,8 @@ pub fn rabitq_index( if i % 32 == 0 { sum = rabitq::bit::block::accumulate(&blocks[i / 32], &lut.1); } - let (rough, err) = rabitq::bit::block::half_process_l2(sum[i % 32], metadata[i], lut.0); + let (rough, err) = + rabitq::bit::block::half_process_l2s(sum[i % 32], metadata[i], lut.0); let lowerbound = rough - err * 1.9; if lowerbound < result.0 { let dis = f32::reduce_sum_of_d2(sample, centroid); diff --git a/crates/rabitq/src/bit.rs b/crates/rabitq/src/bit.rs index 1c83b49f..bfaa0bd7 100644 --- a/crates/rabitq/src/bit.rs +++ b/crates/rabitq/src/bit.rs @@ -201,7 +201,7 @@ pub mod binary { } #[inline(always)] - pub fn half_process_l2( + pub fn half_process_l2s( sum: u32, CodeMetadata { dis_u_2, @@ -223,7 +223,7 @@ pub mod binary { } #[inline(always)] - pub fn half_process_l2_residual( + pub fn half_process_l2s_residual( sum: u32, CodeMetadata { dis_u_2, @@ -362,7 +362,7 @@ pub mod block { } #[inline(always)] - pub fn full_process_l2( + pub fn full_process_l2s( (dis_u_2, _, factor_ip, factor_err, t): BlockCode<'_>, lut: &BlockLut, ) -> [(f32, f32); 32] { @@ -402,7 +402,7 @@ pub mod block { } #[inline(always)] - pub fn half_process_l2( + pub fn half_process_l2s( sum: u32, CodeMetadata { dis_u_2, @@ -419,7 +419,7 @@ pub mod block { } #[inline(always)] - pub fn half_process_l2_residual( + pub fn half_process_l2s_residual( sum: u32, CodeMetadata { dis_u_2, diff --git a/crates/rabitq/src/bits.rs b/crates/rabitq/src/bits.rs index 66bb62f8..456f0f12 100644 --- a/crates/rabitq/src/bits.rs +++ b/crates/rabitq/src/bits.rs @@ -89,28 +89,28 @@ pub mod binary { pub fn half_process_dot( bits: Bits, - n: u32, - value: u32, + dim: u32, + sum: u32, code: CodeMetadata, lut: BinaryLutMetadata, ) -> (f32,) { let rough = match bits { - Bits::_1 => crate::extended::half_process_dot::<1, BITS>(n, value, code, lut), - Bits::_2 => crate::extended::half_process_dot::<2, BITS>(n, value, code, lut), + Bits::_1 => crate::extended::half_process_dot::<1, BITS>(dim, sum, code, lut), + Bits::_2 => crate::extended::half_process_dot::<2, BITS>(dim, sum, code, lut), }; (rough,) } - pub fn half_process_l2( + pub fn half_process_l2s( bits: Bits, - n: u32, - value: u32, + dim: u32, + sum: u32, code: CodeMetadata, lut: BinaryLutMetadata, ) -> (f32,) { let rough = match bits { - Bits::_1 => crate::extended::half_process_l2::<1, BITS>(n, value, code, lut), - Bits::_2 => crate::extended::half_process_l2::<2, BITS>(n, value, code, lut), + Bits::_1 => crate::extended::half_process_l2s::<1, BITS>(dim, sum, code, lut), + Bits::_2 => crate::extended::half_process_l2s::<2, BITS>(dim, sum, code, lut), }; (rough,) } diff --git a/crates/rabitq/src/byte.rs b/crates/rabitq/src/byte.rs index 49b36028..4af1e3a4 100644 --- a/crates/rabitq/src/byte.rs +++ b/crates/rabitq/src/byte.rs @@ -23,6 +23,10 @@ pub fn ugly_code(vector: &[f32]) -> Code { crate::extended::ugly_code::<8>(vector) } +pub fn pack_code(input: &[u8]) -> Vec { + input.to_vec() +} + pub mod binary { use crate::extended::CodeMetadata; @@ -35,12 +39,12 @@ pub mod binary { #[deprecated] pub fn preprocess(vector: &[f32]) -> BinaryLut { let (metadata, elements) = crate::extended::code::(vector); - (metadata, elements) + (metadata, super::pack_code(&elements)) } pub fn ugly_preprocess(vector: &[f32]) -> BinaryLut { let (metadata, elements) = crate::extended::ugly_code::(vector); - (metadata, elements) + (metadata, super::pack_code(&elements)) } pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { @@ -48,32 +52,32 @@ pub mod binary { } pub fn half_process_dot( - n: u32, - value: u32, + dim: u32, + sum: u32, code: CodeMetadata, lut: BinaryLutMetadata, ) -> (f32,) { - let rough = crate::extended::half_process_dot::<8, BITS>(n, value, code, lut); + let rough = crate::extended::half_process_dot::<8, BITS>(dim, sum, code, lut); (rough,) } - pub fn half_process_l2( - n: u32, - value: u32, + pub fn half_process_l2s( + dim: u32, + sum: u32, code: CodeMetadata, lut: BinaryLutMetadata, ) -> (f32,) { - let rough = crate::extended::half_process_l2::<8, BITS>(n, value, code, lut); + let rough = crate::extended::half_process_l2s::<8, BITS>(dim, sum, code, lut); (rough,) } pub fn half_process_cos( - n: u32, - value: u32, + dim: u32, + sum: u32, code: CodeMetadata, lut: BinaryLutMetadata, ) -> (f32,) { - let rough = crate::extended::half_process_cos::<8, BITS>(n, value, code, lut); + let rough = crate::extended::half_process_cos::<8, BITS>(dim, sum, code, lut); (rough,) } } diff --git a/crates/rabitq/src/extended.rs b/crates/rabitq/src/extended.rs index 53b04212..0e8c52a9 100644 --- a/crates/rabitq/src/extended.rs +++ b/crates/rabitq/src/extended.rs @@ -153,9 +153,9 @@ pub fn ugly_code(vector: &[f32]) -> Code { ) } -pub fn half_process_l2( - n: u32, - value: u32, +pub fn half_process_l2s( + dim: u32, + sum: u32, lhs: CodeMetadata, rhs: CodeMetadata, ) -> f32 { @@ -165,7 +165,7 @@ pub fn half_process_l2( let c_x = ((1 << X) - 1) as f32 * 0.5; let c_y = ((1 << Y) - 1) as f32 * 0.5; - let ip = value as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + n as f32 * c_x * c_y; + let ip = sum as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + dim as f32 * c_x * c_y; lhs.dis_u_2 + rhs.dis_u_2 - 2.0 * ip @@ -174,8 +174,8 @@ pub fn half_process_l2( } pub fn half_process_dot( - n: u32, - value: u32, + dim: u32, + sum: u32, lhs: CodeMetadata, rhs: CodeMetadata, ) -> f32 { @@ -185,13 +185,13 @@ pub fn half_process_dot( let c_x = ((1 << X) - 1) as f32 * 0.5; let c_y = ((1 << Y) - 1) as f32 * 0.5; - let ip = value as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + n as f32 * c_x * c_y; + let ip = sum as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + dim as f32 * c_x * c_y; -ip * (lhs.dis_u_2.sqrt() / lhs.norm_of_lattice) * (rhs.dis_u_2.sqrt() / rhs.norm_of_lattice) } pub fn half_process_cos( - n: u32, - value: u32, + dim: u32, + sum: u32, lhs: CodeMetadata, rhs: CodeMetadata, ) -> f32 { @@ -201,7 +201,7 @@ pub fn half_process_cos( let c_x = ((1 << X) - 1) as f32 * 0.5; let c_y = ((1 << Y) - 1) as f32 * 0.5; - let ip = value as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + n as f32 * c_x * c_y; + let ip = sum as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + dim as f32 * c_x * c_y; -ip / lhs.norm_of_lattice / rhs.norm_of_lattice } @@ -209,15 +209,15 @@ fn find_scale(o: &[f32]) -> f32 { assert!((1..=8).contains(&B)); let mask = (1_u32 << (B - 1)) - 1; - let dims = o.len(); + let dim = o.len(); - let mut code = Vec::::with_capacity(dims); + let mut code = Vec::::with_capacity(dim); let mut numerator = 0.0f64; let mut sqr_denominator = 0.0f64; let (mut y_m, mut x_m); - for i in 0..dims { + for i in 0..dim { code.push(0); let value = 0.5; numerator += value * o[i] as f64; @@ -230,7 +230,7 @@ fn find_scale(o: &[f32]) -> f32 { } let mut events = Vec::<(f64, usize)>::new(); - for i in 0..dims { + for i in 0..dim { for c in 1..=mask { let x = (c as f64) / o[i] as f64; events.push((x, i)); @@ -254,14 +254,14 @@ fn find_scale(o: &[f32]) -> f32 { fn ugly_find_scale(o: &[f32]) -> (f32, Vec) { assert!((1..=8).contains(&B)); - let dims = o.len(); + let dim = o.len(); - let mut code = Vec::::with_capacity(dims); + let mut code = Vec::::with_capacity(dim); let mut numerator_m = 0.0f64; let mut sqr_denominator_m = 0.0f64; let scale = (1 << (B - 1)) as f32 / f32::reduce_min_max_of_x(o).1; - for i in 0..dims { + for i in 0..dim { code.push((o[i] as f64 * scale as f64) as u8); let value = code[i] as f64 + 0.5; numerator_m += value * o[i] as f64; @@ -269,9 +269,9 @@ fn ugly_find_scale(o: &[f32]) -> (f32, Vec) { } let mut y_m = numerator_m / sqr_denominator_m.sqrt(); - let mut delta = vec![0_i32; dims]; + let mut delta = vec![0_i32; dim]; for _ in 0..8 { - for i in 0..dims { + for i in 0..dim { if code[i] < (1 << (B - 1)) - 1 { let numerator = numerator_m + o[i] as f64; let sqr_denominator = sqr_denominator_m + 2.0 * (code[i] as f64 + 1.0); diff --git a/crates/rabitq/src/halfbyte.rs b/crates/rabitq/src/halfbyte.rs new file mode 100644 index 00000000..bd45c986 --- /dev/null +++ b/crates/rabitq/src/halfbyte.rs @@ -0,0 +1,92 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +pub use crate::extended::{Code, CodeMetadata}; + +#[deprecated] +pub fn code(vector: &[f32]) -> Code { + crate::extended::code::<4>(vector) +} + +pub fn ugly_code(vector: &[f32]) -> Code { + crate::extended::ugly_code::<4>(vector) +} + +pub fn pack_code(input: &[u8]) -> Vec { + let f = |t: &[u8; 2]| t[0] | t[1] << 4; + let (arrays, remainder) = input.as_chunks::<2>(); + let mut buffer = [0u8; 2]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() +} + +pub mod binary { + use crate::extended::CodeMetadata; + + const BITS: usize = 4; + + pub type BinaryLutMetadata = CodeMetadata; + pub type BinaryLut = (BinaryLutMetadata, Vec); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); + + #[deprecated] + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::code::(vector); + (metadata, super::pack_code(&elements)) + } + + pub fn ugly_preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::ugly_code::(vector); + (metadata, super::pack_code(&elements)) + } + + pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { + simd::halfbyte::reduce_sum_of_x_as_u32_y_as_u32(x, y) + } + + pub fn half_process_dot( + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_dot::<4, BITS>(dim, sum, code, lut); + (rough,) + } + + pub fn half_process_l2s( + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_l2s::<4, BITS>(dim, sum, code, lut); + (rough,) + } + + pub fn half_process_cos( + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_cos::<4, BITS>(dim, sum, code, lut); + (rough,) + } +} diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index 1925af5b..0aae7682 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -16,5 +16,6 @@ pub mod bit; pub mod bits; pub mod byte; mod extended; +pub mod halfbyte; pub mod packing; pub mod rotate; diff --git a/crates/vchordg/src/build.rs b/crates/vchordg/src/build.rs index 2f13f734..beaa093b 100644 --- a/crates/vchordg/src/build.rs +++ b/crates/vchordg/src/build.rs @@ -52,7 +52,7 @@ pub fn build( assert_eq!(vector_guard.id(), 2); drop(vector_guard); let serialized = MetaTuple::serialize(&MetaTuple { - dims: vector_options.dims, + dim: vector_options.dim, bits: index_options.bits, m: index_options.m, alpha: index_options.alpha, diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index 4fd7b76a..5b72d923 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -21,7 +21,7 @@ use crate::types::DistanceKind; use crate::vectors::{by_prefetch, by_read, copy_all, copy_nothing, copy_outs, update}; use crate::visited::Visited; use always_equal::AlwaysEqual; -use index::accessor::LAccess; +use index::accessor::{DefaultWithDimension, LAccess}; use index::bump::Bump; use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; @@ -44,8 +44,8 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); - let dims = meta_tuple.dims(); - assert_eq!(dims, vector.dims(), "unmatched dimensions"); + let dim = meta_tuple.dim(); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); let start = meta_tuple.start(); let bits = Bits::try_from(meta_tuple.bits()).expect("data corruption"); let m = meta_tuple.m(); @@ -149,7 +149,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let pointers_s: &[_] = bump.alloc_slice(vertex_tuple.pointers()); let score_s = O::process( bits, - dims, + dim, (vertex_tuple.metadata(), vertex_tuple.elements()), &lut, ); @@ -159,7 +159,10 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( while let Some(((_, AlwaysEqual((pointers_u, u))), guards)) = candidates.pop() { let Ok((dis_u, outs_u, _, _)) = crate::vectors::read::( by_prefetch::(guards, pointers_u.iter().copied()), - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), copy_outs, ) else { // the link is broken @@ -188,7 +191,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let pointers_v: &[_] = bump.alloc_slice(vertex_tuple.pointers()); let score_v = O::process( bits, - dims, + dim, (vertex_tuple.metadata(), vertex_tuple.elements()), &lut, ); @@ -218,7 +221,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; let Ok((vector_u, _, _, _)) = crate::vectors::read::( by_read(index, pointers_u.iter().copied()), - CloneAccessor::::default(), + CloneAccessor::::default_with_dimension(dim), copy_nothing, ) else { // the link is broken @@ -271,7 +274,7 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; let Ok((vector_u, _, _, _)) = crate::vectors::read::( by_read(index, pointers_u.iter().copied()), - CloneAccessor::::default(), + CloneAccessor::::default_with_dimension(dim), copy_nothing, ) else { // the link is broken diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs index 1451fd5b..de25edbc 100644 --- a/crates/vchordg/src/maintain.rs +++ b/crates/vchordg/src/maintain.rs @@ -18,6 +18,7 @@ use crate::tuples::{MetaTuple, VertexTuple, WithReader}; use crate::types::DistanceKind; use crate::vectors::{by_read, copy_all, copy_nothing, copy_outs, update}; use always_equal::AlwaysEqual; +use index::accessor::DefaultWithDimension; use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; use std::cmp::Reverse; use vector::VectorOwned; @@ -29,6 +30,7 @@ where let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); let m = meta_tuple.m(); let alpha = meta_tuple.alpha().to_vec(); let start = meta_tuple.start(); @@ -58,7 +60,7 @@ where 'occ: loop { let Ok((vector_u, neighbours_u, _, version)) = crate::vectors::read::( by_read::(index, pointers_u.iter().copied()), - CloneAccessor::::default(), + CloneAccessor::::default_with_dimension(dim), copy_all, ) else { // the link is broken @@ -79,7 +81,7 @@ where drop(vertex_guard); let Ok((vector_v, outs_v, _, _)) = crate::vectors::read::( by_read::(index, pointers_v.iter().copied()), - CloneAccessor::::default(), + CloneAccessor::::default_with_dimension(dim), copy_outs, ) else { // the link is broken @@ -106,7 +108,7 @@ where drop(vertex_guard); let Ok((vector_v, _, _, _)) = crate::vectors::read::( by_read::(index, pointers_v.iter().copied()), - CloneAccessor::::default(), + CloneAccessor::::default_with_dimension(dim), copy_nothing, ) else { // the link is broken diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs index 4cf31f57..6eead52e 100644 --- a/crates/vchordg/src/operator.rs +++ b/crates/vchordg/src/operator.rs @@ -15,12 +15,14 @@ use crate::types::DistanceKind; use distance::Distance; use index::accessor::{ - Accessor1, Accessor2, DimensionDistanceAccessor, DistanceAccessor, Dot, L2S, + Accessor1, Accessor2, ByteDistanceAccessor, DefaultWithDimension, DistanceAccessor, Dot, + HalfbyteDistanceAccessor, L2S, }; use rabitq::bits::Bits; use simd::{Floating, f16}; use std::fmt::Debug; use std::marker::PhantomData; +use vector::rabitq4::Rabitq4Owned; use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; use vector::{VectorBorrowed, VectorOwned}; @@ -35,7 +37,7 @@ pub trait Vector: VectorOwned { vector: Self::Borrowed<'_>, m: usize, ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)); - fn pack(elements: Vec, metadata: Self::Metadata) -> Self; + fn pack(dim: u32, elements: Vec, metadata: Self::Metadata) -> Self; fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code; fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut; @@ -68,7 +70,7 @@ impl Vector for VectOwned { ) } - fn pack(elements: Vec, (): Self::Metadata) -> Self { + fn pack(_: u32, elements: Vec, (): Self::Metadata) -> Self { VectOwned::new(elements) } @@ -108,7 +110,7 @@ impl Vector for VectOwned { ) } - fn pack(elements: Vec, (): Self::Metadata) -> Self { + fn pack(_: u32, elements: Vec, (): Self::Metadata) -> Self { VectOwned::new(elements) } @@ -128,7 +130,7 @@ impl Vector for Rabitq8Owned { fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { ( - vector.code(), + vector.packed_code(), [ vector.sum_of_x2(), vector.norm_of_lattice(), @@ -148,28 +150,25 @@ impl Vector for Rabitq8Owned { vector.sum_of_code(), vector.sum_of_abs_x(), ]; - let slice = vector.code(); + let slice = vector.packed_code(); let tailing = (size_of::() * m) .next_multiple_of(crate::tuples::ALIGN); assert!(tailing <= 8000); - if slice.len() <= (8000 - tailing) / size_of::() { + if slice.len() <= 8000 - tailing { return (vec![], (slice, metadata)); } - let (l, r) = slice.split_at(slice.len() - (8000 - tailing) / size_of::()); - ( - l.chunks(8000 / size_of::()).collect::>(), - (r, metadata), - ) + let (l, r) = slice.split_at(slice.len() - (8000 - tailing)); + (l.chunks(8000).collect::>(), (r, metadata)) } - fn pack(code: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { - Rabitq8Owned::new(_0, _1, _2, _3, code) + fn pack(dim: u32, code: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq8Owned::new(dim, _0, _1, _2, _3, code) } fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code { let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); - let mut result = Vec::with_capacity(vector.dims() as _); - for c in vector.code().iter().copied() { + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { let base = -0.5 * ((1 << 8) - 1) as f32; result.push((base + c as f32) * scale); } @@ -178,8 +177,8 @@ impl Vector for Rabitq8Owned { fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut { let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); - let mut result = Vec::with_capacity(vector.dims() as _); - for c in vector.code().iter().copied() { + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { let base = -0.5 * ((1 << 8) - 1) as f32; result.push((base + c as f32) * scale); } @@ -187,12 +186,75 @@ impl Vector for Rabitq8Owned { } } +impl Vector for Rabitq4Owned { + type Metadata = [f32; 4]; + + type Element = u8; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + ( + vector.packed_code(), + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)) { + let metadata = [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ]; + let slice = vector.packed_code(); + let tailing = (size_of::() * m) + .next_multiple_of(crate::tuples::ALIGN); + assert!(tailing <= 8000); + if slice.len() <= 8000 - tailing { + return (vec![], (slice, metadata)); + } + let (l, r) = slice.split_at(slice.len() - (8000 - tailing)); + (l.chunks(8000).collect::>(), (r, metadata)) + } + + fn pack(dim: u32, code: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq4Owned::new(dim, _0, _1, _2, _3, code) + } + + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bits::code(bits, &result) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bits::binary::preprocess(&result) + } +} + pub trait Operator: 'static + Debug + Copy { const DISTANCE: DistanceKind; type Vector: Vector; - type DistanceAccessor: Default + type DistanceAccessor: DefaultWithDimension + Accessor2< ::Element, ::Element, @@ -203,7 +265,7 @@ pub trait Operator: 'static + Debug + Copy { fn process( bits: Bits, - n: u32, + dim: u32, code: ([f32; 3], &[u64]), lut: &rabitq::bits::binary::BinaryLut, ) -> Distance; @@ -233,16 +295,16 @@ impl Operator for Op, L2S> { fn process( bits: Bits, - n: u32, + dim: u32, code: ([f32; 3], &[u64]), lut: &rabitq::bits::binary::BinaryLut, ) -> Distance { use rabitq::bits::CodeMetadata; - let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); - let (distance,) = rabitq::bits::binary::half_process_l2( + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2s( bits, - n, - value, + dim, + sum, CodeMetadata::from_array(code.0), lut.0, ); @@ -266,16 +328,16 @@ impl Operator for Op, Dot> { fn process( bits: Bits, - n: u32, + dim: u32, code: ([f32; 3], &[u64]), lut: &rabitq::bits::binary::BinaryLut, ) -> Distance { use rabitq::bits::CodeMetadata; - let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); let (distance,) = rabitq::bits::binary::half_process_dot( bits, - n, - value, + dim, + sum, CodeMetadata::from_array(code.0), lut.0, ); @@ -299,16 +361,16 @@ impl Operator for Op, L2S> { fn process( bits: Bits, - n: u32, + dim: u32, code: ([f32; 3], &[u64]), lut: &rabitq::bits::binary::BinaryLut, ) -> Distance { use rabitq::bits::CodeMetadata; - let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); - let (distance,) = rabitq::bits::binary::half_process_l2( + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2s( bits, - n, - value, + dim, + sum, CodeMetadata::from_array(code.0), lut.0, ); @@ -332,16 +394,16 @@ impl Operator for Op, Dot> { fn process( bits: Bits, - n: u32, + dim: u32, code: ([f32; 3], &[u64]), lut: &rabitq::bits::binary::BinaryLut, ) -> Distance { use rabitq::bits::CodeMetadata; - let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); let (distance,) = rabitq::bits::binary::half_process_dot( bits, - n, - value, + dim, + sum, CodeMetadata::from_array(code.0), lut.0, ); @@ -361,20 +423,20 @@ impl Operator for Op { type Vector = Rabitq8Owned; - type DistanceAccessor = DimensionDistanceAccessor; + type DistanceAccessor = ByteDistanceAccessor; fn process( bits: Bits, - n: u32, + dim: u32, code: ([f32; 3], &[u64]), lut: &rabitq::bits::binary::BinaryLut, ) -> Distance { use rabitq::bits::CodeMetadata; - let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); - let (distance,) = rabitq::bits::binary::half_process_l2( + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2s( bits, - n, - value, + dim, + sum, CodeMetadata::from_array(code.0), lut.0, ); @@ -394,20 +456,86 @@ impl Operator for Op { type Vector = Rabitq8Owned; - type DistanceAccessor = DimensionDistanceAccessor; + type DistanceAccessor = ByteDistanceAccessor; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_dot( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_dot(rhs) + } +} + +impl Operator for Op { + const DISTANCE: DistanceKind = DistanceKind::L2S; + + type Vector = Rabitq4Owned; + + type DistanceAccessor = HalfbyteDistanceAccessor; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2s( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_l2s(rhs) + } +} + +impl Operator for Op { + const DISTANCE: DistanceKind = DistanceKind::Dot; + + type Vector = Rabitq4Owned; + + type DistanceAccessor = HalfbyteDistanceAccessor; fn process( bits: Bits, - n: u32, + dim: u32, code: ([f32; 3], &[u64]), lut: &rabitq::bits::binary::BinaryLut, ) -> Distance { use rabitq::bits::CodeMetadata; - let value = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); let (distance,) = rabitq::bits::binary::half_process_dot( bits, - n, - value, + dim, + sum, CodeMetadata::from_array(code.0), lut.0, ); @@ -423,11 +551,11 @@ impl Operator for Op { } #[derive(Debug, Clone)] -pub struct CloneAccessor(Vec); +pub struct CloneAccessor(u32, Vec); -impl Default for CloneAccessor { - fn default() -> Self { - Self(Vec::new()) +impl DefaultWithDimension for CloneAccessor { + fn default_with_dimension(dim: u32) -> Self { + Self(dim, Vec::new()) } } @@ -435,10 +563,10 @@ impl Accessor1 for CloneAccessor { type Output = V; fn push(&mut self, input: &[V::Element]) { - self.0.extend(input); + self.1.extend(input); } fn finish(self, metadata: V::Metadata) -> Self::Output { - V::pack(self.0, metadata) + V::pack(self.0, self.1, metadata) } } diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs index 5e260492..ee5ee353 100644 --- a/crates/vchordg/src/search.rs +++ b/crates/vchordg/src/search.rs @@ -21,7 +21,7 @@ use crate::vectors::{by_prefetch, copy_outs}; use crate::visited::Visited; use always_equal::AlwaysEqual; use distance::Distance; -use index::accessor::LAccess; +use index::accessor::{DefaultWithDimension, LAccess}; use index::bump::Bump; use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use index::relation::{Page, RelationRead}; @@ -46,10 +46,10 @@ where let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); - let dims = meta_tuple.dims(); + let dim = meta_tuple.dim(); let start = meta_tuple.start(); let bits = Bits::try_from(meta_tuple.bits()).expect("data corruption"); - assert_eq!(dims, vector.dims(), "unmatched dimensions"); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); let ef = ef_search; let beam = beam_search; drop(meta_guard); @@ -70,7 +70,7 @@ where let pointers_s: &[_] = bump.alloc_slice(vertex_tuple.pointers()); let score_s = O::process( bits, - dims, + dim, (vertex_tuple.metadata(), vertex_tuple.elements()), &lut, ); @@ -80,7 +80,10 @@ where while let Some(((_, AlwaysEqual(pointers_u)), guards)) = candidates.pop() { let Ok((dis_u, outs_u, payload_u, _)) = crate::vectors::read::( by_prefetch::(guards, pointers_u.iter().copied()), - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), copy_outs, ) else { // the link is broken @@ -109,7 +112,7 @@ where let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); let score_v = O::process( bits, - dims, + dim, (vertex_tuple.metadata(), vertex_tuple.elements()), &lut, ); diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 1555e4c9..a456f5bb 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -50,7 +50,7 @@ pub trait WithWriter: Tuple { #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct MetaTupleHeader { version: u64, - dims: u32, + dim: u32, bits: u8, _padding: [Padding; 7], m: u32, @@ -63,7 +63,7 @@ struct MetaTupleHeader { } pub struct MetaTuple { - pub dims: u32, + pub dim: u32, pub bits: u8, pub m: u32, pub alpha: Vec, @@ -79,7 +79,7 @@ impl Tuple for MetaTuple { let mut buffer = Vec::::new(); match self { MetaTuple { - dims, + dim, bits, m, alpha, @@ -101,7 +101,7 @@ impl Tuple for MetaTuple { buffer[size_of::()..][..size_of::()].copy_from_slice( MetaTupleHeader { version: VERSION, - dims: *dims, + dim: *dim, bits: *bits, m: *m, alpha_s, @@ -146,8 +146,8 @@ pub struct MetaTupleReader<'a> { } impl<'a> MetaTupleReader<'a> { - pub fn dims(self) -> u32 { - self.header.dims + pub fn dim(self) -> u32 { + self.header.dim } pub fn bits(self) -> u8 { self.header.bits diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index c1d94cf5..4fbb9a48 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -15,6 +15,7 @@ use serde::{Deserialize, Serialize}; use simd::f16; use validator::{Validate, ValidationError}; +use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; @@ -85,6 +86,7 @@ pub enum OwnedVector { Vecf32(VectOwned), Vecf16(VectOwned), Rabitq8(Rabitq8Owned), + Rabitq4(Rabitq4Owned), } #[derive(Debug, Clone, Copy)] @@ -92,6 +94,7 @@ pub enum BorrowedVector<'a> { Vecf32(VectBorrowed<'a, f32>), Vecf16(VectBorrowed<'a, f16>), Rabitq8(Rabitq8Borrowed<'a>), + Rabitq4(Rabitq4Borrowed<'a>), } #[repr(u8)] @@ -107,30 +110,21 @@ pub enum VectorKind { Vecf32, Vecf16, Rabitq8, -} - -impl VectorKind { - pub fn element_size(self) -> u32 { - match self { - VectorKind::Vecf32 => size_of::() as _, - VectorKind::Vecf16 => size_of::() as _, - VectorKind::Rabitq8 => size_of::() as _, - } - } + Rabitq4, } #[derive(Debug, Clone, Validate)] #[validate(schema(function = "Self::validate_self"))] pub struct VectorOptions { #[validate(range(min = 1))] - pub dims: u32, + pub dim: u32, pub v: VectorKind, pub d: DistanceKind, } impl VectorOptions { pub fn validate_self(&self) -> Result<(), ValidationError> { - match (self.v, self.d, self.dims) { + match (self.v, self.d, self.dim) { (_, _, 1..=30000) => Ok(()), _ => Err(ValidationError::new("invalid vector options")), } diff --git a/crates/vchordrq/src/build.rs b/crates/vchordrq/src/build.rs index f0905df0..9c00197c 100644 --- a/crates/vchordrq/src/build.rs +++ b/crates/vchordrq/src/build.rs @@ -29,7 +29,7 @@ pub fn build( ) where R::Page: Page, { - let dims = vector_options.dims; + let dim = vector_options.dim; let is_residual = vchordrq_options.residual_quantization; let mut meta = TapeWriter::<_, MetaTuple>::create(index, false); assert_eq!(meta.first(), 0); @@ -93,7 +93,7 @@ pub fn build( }); level.push(jump.first()); } else { - let mut tape = H1TapeWriter::create(index, O::Vector::count(dims as _), false); + let mut tape = H1TapeWriter::create(index, O::Vector::count(dim) as _, false); let centroid = structures[i].centroids[j].as_borrowed(); for child in structures[i].children[j].iter().copied() { let vector = structures[i - 1].centroids[child as usize].as_borrowed(); @@ -108,14 +108,14 @@ pub fn build( }); } let (mut tape, chunk) = tape.into_inner(); - H1TapeWriter::flush(&mut tape, O::Vector::count(dims as _), chunk); + H1TapeWriter::flush(&mut tape, O::Vector::count(dim) as _, chunk); level.push(tape.first()); } } pointer_of_firsts.push(level); } meta.push(MetaTuple { - dims, + dim, height_of_root: structures.len() as u32, is_residual, rerank_in_heap: vchordrq_options.rerank_in_table, diff --git a/crates/vchordrq/src/cost.rs b/crates/vchordrq/src/cost.rs index 5eb2dc6d..3dbb4365 100644 --- a/crates/vchordrq/src/cost.rs +++ b/crates/vchordrq/src/cost.rs @@ -16,7 +16,7 @@ use crate::tuples::{MetaTuple, WithReader}; use index::relation::{Page, RelationRead}; pub struct Cost { - pub dims: u32, + pub dim: u32, pub cells: Vec, } @@ -25,10 +25,10 @@ pub fn cost(index: &R) -> Cost { let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); - let dims = meta_tuple.dims(); + let dim = meta_tuple.dim(); let cells = meta_tuple.cells().to_vec(); drop(meta_guard); - Cost { dims, cells } + Cost { dim, cells } } diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index bd30f62f..32ae080a 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -12,13 +12,14 @@ // // Copyright (c) 2025 TensorChord Inc. +use crate::closure_lifetime_binder::id_0; use crate::linked_vec::LinkedVec; use crate::operator::*; use crate::tuples::*; use crate::{Opaque, centroids, tape, vectors}; use always_equal::AlwaysEqual; use distance::Distance; -use index::accessor::{FunctionalAccessor, LAccess}; +use index::accessor::{DefaultWithDimension, FunctionalAccessor, LAccess}; use index::bump::Bump; use index::fetch::BorrowedIter; use index::prefetcher::{Prefetcher, PrefetcherHeapFamily}; @@ -47,9 +48,9 @@ where let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); - let dims = meta_tuple.dims(); + let dim = meta_tuple.dim(); let rerank_in_heap = meta_tuple.rerank_in_heap(); - assert_eq!(dims, vector.dims(), "unmatched dimensions"); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); let vectors_first = { let l = meta_tuple.vectors_first(); let n = NonZero::new(l.len()).expect("data corruption"); @@ -80,11 +81,11 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); - let dims = meta_tuple.dims(); + let dim = meta_tuple.dim(); let is_residual = meta_tuple.is_residual(); let height_of_root = meta_tuple.height_of_root(); let freepages_first = meta_tuple.freepages_first(); - assert_eq!(dims, vector.dims(), "unmatched dimensions"); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); let epsilon = 1.9; type State = (Reverse, AlwaysEqual, AlwaysEqual); @@ -95,7 +96,10 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let distance = centroids::read::( prefetch.map(|id| index.read(id)), head, - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), ); let norm = meta_tuple.centroid_norm(); let first = meta_tuple.first(); @@ -141,7 +145,10 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( let distance = centroids::read::( prefetch, head, - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), ); cache.push((Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))); } @@ -169,7 +176,13 @@ pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( .iter() .map(|&id| index.read(id)), jump_tuple.centroid_head(), - FunctionalAccessor::new(Vec::new(), Vec::extend_from_slice, O::Vector::pack), + FunctionalAccessor::new( + (dim, Vec::new()), + id_0(|(_, elements): &mut (_, Vec<_>), slice| { + elements.extend_from_slice(slice) + }), + |(dim, elements), metadata| O::Vector::pack(dim, elements, metadata), + ), ) }), ); diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs index e9492679..24f4b255 100644 --- a/crates/vchordrq/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -47,7 +47,7 @@ where let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); - let dims = meta_tuple.dims(); + let dim = meta_tuple.dim(); let height_of_root = meta_tuple.height_of_root(); let freepages_first = meta_tuple.freepages_first(); @@ -118,7 +118,7 @@ where }) }); - let mut tape = FrozenTapeWriter::create(&hooked_index, O::Vector::count(dims as _), false); + let mut tape = FrozenTapeWriter::create(&hooked_index, O::Vector::count(dim) as _, false); let mut trace_directory = Vec::new(); let mut trace_forzen = Vec::new(); @@ -182,7 +182,7 @@ where let signs = elements .iter() .flat_map(|x| std::array::from_fn::<_, 64, _>(|i| *x & (1 << i) != 0)) - .take(dims as _) + .take(dim as _) .collect::>(); ( ( diff --git a/crates/vchordrq/src/operator.rs b/crates/vchordrq/src/operator.rs index 8b0688ce..54395aae 100644 --- a/crates/vchordrq/src/operator.rs +++ b/crates/vchordrq/src/operator.rs @@ -14,7 +14,8 @@ use distance::Distance; use index::accessor::{ - Accessor1, Accessor2, DimensionDistanceAccessor, DistanceAccessor, Dot, L2S, RAccess, + Accessor1, Accessor2, ByteDistanceAccessor, DefaultWithDimension, DistanceAccessor, Dot, + HalfbyteDistanceAccessor, L2S, RAccess, }; use rabitq::bit::CodeMetadata; use rabitq::bit::binary::BinaryLut; @@ -22,6 +23,7 @@ use rabitq::bit::block::{BlockLut, STEP}; use simd::{Floating, f16}; use std::fmt::Debug; use std::marker::PhantomData; +use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; use vector::{VectorBorrowed, VectorOwned}; @@ -72,7 +74,7 @@ impl Default for CloneAccessor { } } -impl Accessor1 for CloneAccessor { +impl Accessor1 for CloneAccessor { type Output = V; #[inline(always)] @@ -81,8 +83,8 @@ impl Accessor1 for CloneAccessor { } #[inline(always)] - fn finish(self, metadata: V::Metadata) -> Self::Output { - V::pack(self.0, metadata) + fn finish(self, (metadata, dim): (V::Metadata, u32)) -> Self::Output { + V::pack(dim, self.0, metadata) } } @@ -93,11 +95,11 @@ pub trait Vector: VectorOwned { fn split(vector: Self::Borrowed<'_>) -> (Vec<&[Self::Element]>, Self::Metadata); - fn count(n: usize) -> usize; + fn count(dim: u32) -> u32; fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); - fn pack(elements: Vec, metadata: Self::Metadata) -> Self; + fn pack(dim: u32, elements: Vec, metadata: Self::Metadata) -> Self; fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut; @@ -126,12 +128,12 @@ impl Vector for VectOwned { ) } - fn count(n: usize) -> usize { - match n { + fn count(dim: u32) -> u32 { + match dim { 0 => unreachable!(), 1..=960 => 1, 961..=1280 => 2, - 1281.. => n.div_ceil(1920), + 1281.. => dim.div_ceil(1920), } } @@ -139,7 +141,7 @@ impl Vector for VectOwned { (vector.slice(), ()) } - fn pack(elements: Vec, (): Self::Metadata) -> Self { + fn pack(_: u32, elements: Vec, (): Self::Metadata) -> Self { VectOwned::new(elements) } @@ -178,12 +180,12 @@ impl Vector for VectOwned { ) } - fn count(n: usize) -> usize { - match n { + fn count(dim: u32) -> u32 { + match dim { 0 => unreachable!(), 1..=1920 => 1, 1921..=2560 => 2, - 2561.. => n.div_ceil(3840), + 2561.. => dim.div_ceil(3840), } } @@ -191,7 +193,7 @@ impl Vector for VectOwned { (vector.slice(), ()) } - fn pack(elements: Vec, (): Self::Metadata) -> Self { + fn pack(_: u32, elements: Vec, (): Self::Metadata) -> Self { VectOwned::new(elements) } @@ -219,11 +221,11 @@ impl Vector for Rabitq8Owned { fn split(vector: Self::Borrowed<'_>) -> (Vec<&[u8]>, [f32; 4]) { ( - match vector.code().len() { + match vector.packed_code().len() { 0 => unreachable!(), - 1..=3840 => vec![vector.code()], - 3841..=5120 => vec![&vector.code()[..1280], &vector.code()[1280..]], - 5121.. => vector.code().chunks(7680).collect(), + 1..=3840 => vec![vector.packed_code()], + 3841..=5120 => vec![&vector.packed_code()[..1280], &vector.packed_code()[1280..]], + 5121.. => vector.packed_code().chunks(7680).collect(), }, [ vector.sum_of_x2(), @@ -234,18 +236,18 @@ impl Vector for Rabitq8Owned { ) } - fn count(n: usize) -> usize { - match n { + fn count(dim: u32) -> u32 { + match dim { 0 => unreachable!(), 1..=3840 => 1, 3841..=5120 => 2, - 5121.. => n.div_ceil(7680), + 5121.. => dim.div_ceil(7680), } } fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { ( - vector.code(), + vector.packed_code(), [ vector.sum_of_x2(), vector.norm_of_lattice(), @@ -255,14 +257,14 @@ impl Vector for Rabitq8Owned { ) } - fn pack(elements: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { - Rabitq8Owned::new(_0, _1, _2, _3, elements) + fn pack(dim: u32, elements: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq8Owned::new(dim, _0, _1, _2, _3, elements) } fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); - let mut result = Vec::with_capacity(vector.dims() as _); - for c in vector.code().iter().copied() { + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { let base = -0.5 * ((1 << 8) - 1) as f32; result.push((base + c as f32) * scale); } @@ -271,8 +273,8 @@ impl Vector for Rabitq8Owned { fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); - let mut result = Vec::with_capacity(vector.dims() as _); - for c in vector.code().iter().copied() { + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { let base = -0.5 * ((1 << 8) - 1) as f32; result.push((base + c as f32) * scale); } @@ -280,15 +282,15 @@ impl Vector for Rabitq8Owned { } fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code { - let n = vector.dims(); + let n = vector.dim(); let sum_of_abs_x = vector.sum_of_abs_x(); let sum_of_x_2 = vector.sum_of_x2(); ( CodeMetadata { dis_u_2: sum_of_x_2, factor_cnt: { - let cnt_pos = vector.code().iter().filter(|&&x| x >= 128).count(); - let cnt_neg = vector.code().iter().filter(|&&x| x <= 127).count(); + let cnt_pos = vector.unpacked_code().filter(|&x| x >= 128).count(); + let cnt_neg = vector.unpacked_code().filter(|&x| x <= 127).count(); cnt_pos as f32 - cnt_neg as f32 }, factor_ip: sum_of_x_2 / sum_of_abs_x, @@ -299,11 +301,112 @@ impl Vector for Rabitq8Owned { }, }, { - let vector = vector.code(); - let n = vector.len(); + let vector = vector.unpacked_code(); let mut signs = Vec::new(); - for i in 0..n { - signs.push(vector[i] >= 128); + for x in vector { + signs.push(x >= 128); + } + signs + }, + ) + } + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { + vector.sum_of_x2() + } +} + +impl Vector for Rabitq4Owned { + type Metadata = [f32; 4]; + + type Element = u8; + + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[u8]>, [f32; 4]) { + ( + match vector.packed_code().len() { + 0 => unreachable!(), + 1..=3840 => vec![vector.packed_code()], + 3841..=5120 => vec![&vector.packed_code()[..1280], &vector.packed_code()[1280..]], + 5121.. => vector.packed_code().chunks(7680).collect(), + }, + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn count(dim: u32) -> u32 { + match dim { + 0 => unreachable!(), + 1..=7680 => 1, + 7681..=10240 => 2, + 10241.. => dim.div_ceil(15360), + } + } + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + ( + vector.packed_code(), + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn pack(dim: u32, elements: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq4Owned::new(dim, _0, _1, _2, _3, elements) + } + + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bit::block::preprocess(&result) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bit::preprocess(&result) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code { + let n = vector.dim(); + let sum_of_abs_x = vector.sum_of_abs_x(); + let sum_of_x_2 = vector.sum_of_x2(); + ( + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector.unpacked_code().filter(|&x| x >= 8).count(); + let cnt_neg = vector.unpacked_code().filter(|&x| x <= 7).count(); + cnt_pos as f32 - cnt_neg as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + factor_err: { + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (n as f32).sqrt(); + dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (n as f32 - 1.0).sqrt() + }, + }, + { + let vector = vector.unpacked_code(); + let mut signs = Vec::new(); + for x in vector { + signs.push(x >= 8); } signs }, @@ -318,7 +421,7 @@ impl Vector for Rabitq8Owned { pub trait Operator: 'static + Debug + Copy { type Vector: Vector; - type DistanceAccessor: Default + type DistanceAccessor: DefaultWithDimension + Accessor2< ::Element, ::Element, @@ -374,9 +477,9 @@ impl Operator for Op, L2S> { (&lut.1, ()), BlockAccessor([0_u32; 32], move |value, code, delta| { if !is_residual { - rabitq::bit::block::half_process_l2(value, code, lut.0) + rabitq::bit::block::half_process_l2s(value, code, lut.0) } else { - rabitq::bit::block::half_process_l2_residual(value, code, lut.0, dis_f, delta) + rabitq::bit::block::half_process_l2s_residual(value, code, lut.0, dis_f, delta) } }), ) @@ -397,9 +500,9 @@ impl Operator for Op, L2S> { factor_err: metadata[3], }; if !is_residual { - rabitq::bit::binary::half_process_l2(value, code, lut.0) + rabitq::bit::binary::half_process_l2s(value, code, lut.0) } else { - rabitq::bit::binary::half_process_l2_residual(value, code, lut.0, dis_f, delta) + rabitq::bit::binary::half_process_l2s_residual(value, code, lut.0, dis_f, delta) } } } @@ -413,15 +516,15 @@ impl Operator for Op, L2S> { let code = Self::Vector::code(residual.as_borrowed()); let delta = { use std::iter::zip; - let dims = vector.dims(); + let dim = vector.dim(); let t = zip(&code.1, centroid.slice()) .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) .sum::() - / (dims as f32).sqrt(); + / (dim as f32).sqrt(); let sum_of_x_2 = code.0.dis_u_2; let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; let dis_u = sum_of_x_2.sqrt(); - let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dim as f32).sqrt(); 2.0 * dis_u * t / x_0 }; (code, delta) @@ -447,12 +550,12 @@ impl Operator for Op, Dot> { { RAccess::new( (&lut.1, ()), - BlockAccessor([0_u32; 32], move |value, code, delta| { + BlockAccessor([0_u32; 32], move |sum, code, delta| { if !is_residual { - rabitq::bit::block::half_process_dot(value, code, lut.0) + rabitq::bit::block::half_process_dot(sum, code, lut.0) } else { rabitq::bit::block::half_process_dot_residual( - value, code, lut.0, dis_f, delta, norm, + sum, code, lut.0, dis_f, delta, norm, ) } }), @@ -466,7 +569,7 @@ impl Operator for Op, Dot> { norm: f32, ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let sum = rabitq::bit::binary::accumulate(elements, &lut.1); let code = CodeMetadata { dis_u_2: metadata[0], factor_cnt: metadata[1], @@ -474,11 +577,9 @@ impl Operator for Op, Dot> { factor_err: metadata[3], }; if !is_residual { - rabitq::bit::binary::half_process_dot(value, code, lut.0) + rabitq::bit::binary::half_process_dot(sum, code, lut.0) } else { - rabitq::bit::binary::half_process_dot_residual( - value, code, lut.0, dis_f, delta, norm, - ) + rabitq::bit::binary::half_process_dot_residual(sum, code, lut.0, dis_f, delta, norm) } } } @@ -492,15 +593,15 @@ impl Operator for Op, Dot> { let code = Self::Vector::code(residual.as_borrowed()); let delta = { use std::iter::zip; - let dims = vector.dims(); + let dim = vector.dim(); let t = zip(&code.1, centroid.slice()) .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) .sum::() - / (dims as f32).sqrt(); + / (dim as f32).sqrt(); let sum_of_x_2 = code.0.dis_u_2; let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; let dis_u = sum_of_x_2.sqrt(); - let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dim as f32).sqrt(); dis_u * t / x_0 - f32::reduce_sum_of_xy(residual.slice(), centroid.slice()) }; (code, delta) @@ -528,9 +629,9 @@ impl Operator for Op, L2S> { (&lut.1, ()), BlockAccessor([0_u32; 32], move |value, code, delta| { if !is_residual { - rabitq::bit::block::half_process_l2(value, code, lut.0) + rabitq::bit::block::half_process_l2s(value, code, lut.0) } else { - rabitq::bit::block::half_process_l2_residual(value, code, lut.0, dis_f, delta) + rabitq::bit::block::half_process_l2s_residual(value, code, lut.0, dis_f, delta) } }), ) @@ -551,9 +652,9 @@ impl Operator for Op, L2S> { factor_err: metadata[3], }; if !is_residual { - rabitq::bit::binary::half_process_l2(value, code, lut.0) + rabitq::bit::binary::half_process_l2s(value, code, lut.0) } else { - rabitq::bit::binary::half_process_l2_residual(value, code, lut.0, dis_f, delta) + rabitq::bit::binary::half_process_l2s_residual(value, code, lut.0, dis_f, delta) } } } @@ -567,16 +668,16 @@ impl Operator for Op, L2S> { let code = Self::Vector::code(residual.as_borrowed()); let delta = { use std::iter::zip; - let dims = vector.dims(); + let dim = vector.dim(); let t = zip(&code.1, centroid.slice()) .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) .map(simd::F16::_to_f32) .sum::() - / (dims as f32).sqrt(); + / (dim as f32).sqrt(); let sum_of_x_2 = code.0.dis_u_2; let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; let dis_u = sum_of_x_2.sqrt(); - let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dim as f32).sqrt(); 2.0 * dis_u * t / x_0 }; (code, delta) @@ -602,12 +703,12 @@ impl Operator for Op, Dot> { { RAccess::new( (&lut.1, ()), - BlockAccessor([0_u32; 32], move |value, code, delta| { + BlockAccessor([0_u32; 32], move |sum, code, delta| { if !is_residual { - rabitq::bit::block::half_process_dot(value, code, lut.0) + rabitq::bit::block::half_process_dot(sum, code, lut.0) } else { rabitq::bit::block::half_process_dot_residual( - value, code, lut.0, dis_f, delta, norm, + sum, code, lut.0, dis_f, delta, norm, ) } }), @@ -621,7 +722,7 @@ impl Operator for Op, Dot> { norm: f32, ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { move |metadata: [f32; 4], elements: &[u64], delta: f32| { - let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let sum = rabitq::bit::binary::accumulate(elements, &lut.1); let code = CodeMetadata { dis_u_2: metadata[0], factor_cnt: metadata[1], @@ -629,11 +730,9 @@ impl Operator for Op, Dot> { factor_err: metadata[3], }; if !is_residual { - rabitq::bit::binary::half_process_dot(value, code, lut.0) + rabitq::bit::binary::half_process_dot(sum, code, lut.0) } else { - rabitq::bit::binary::half_process_dot_residual( - value, code, lut.0, dis_f, delta, norm, - ) + rabitq::bit::binary::half_process_dot_residual(sum, code, lut.0, dis_f, delta, norm) } } } @@ -647,16 +746,16 @@ impl Operator for Op, Dot> { let code = Self::Vector::code(residual.as_borrowed()); let delta = { use std::iter::zip; - let dims = vector.dims(); + let dim = vector.dim(); let t = zip(&code.1, centroid.slice()) .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) .map(simd::F16::_to_f32) .sum::() - / (dims as f32).sqrt(); + / (dim as f32).sqrt(); let sum_of_x_2 = code.0.dis_u_2; let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; let dis_u = sum_of_x_2.sqrt(); - let x_0 = sum_of_abs_x / dis_u / (dims as f32).sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dim as f32).sqrt(); dis_u * t / x_0 - f16::reduce_sum_of_xy(residual.slice(), centroid.slice()) }; (code, delta) @@ -671,7 +770,7 @@ impl Operator for Op, Dot> { impl Operator for Op { type Vector = Rabitq8Owned; - type DistanceAccessor = DimensionDistanceAccessor; + type DistanceAccessor = ByteDistanceAccessor; fn block_access( lut: &BlockLut, @@ -684,7 +783,7 @@ impl Operator for Op { RAccess::new( (&lut.1, ()), BlockAccessor([0_u32; 32], move |value, code, _delta| { - rabitq::bit::block::half_process_l2(value, code, lut.0) + rabitq::bit::block::half_process_l2s(value, code, lut.0) }), ) } @@ -704,7 +803,7 @@ impl Operator for Op { factor_ip: metadata[2], factor_err: metadata[3], }; - rabitq::bit::binary::half_process_l2(value, code, lut.0) + rabitq::bit::binary::half_process_l2s(value, code, lut.0) } } @@ -722,7 +821,58 @@ impl Operator for Op { impl Operator for Op { type Vector = Rabitq8Owned; - type DistanceAccessor = DimensionDistanceAccessor; + type DistanceAccessor = ByteDistanceAccessor; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + assert!(!is_residual); + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |sum, code, _delta| { + rabitq::bit::block::half_process_dot(sum, code, lut.0) + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + assert!(!is_residual); + move |metadata: [f32; 4], elements: &[u64], _delta: f32| { + let sum = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + rabitq::bit::binary::half_process_dot(sum, code, lut.0) + } + } + + fn build( + vector: Rabitq8Borrowed<'_>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if centroid.is_some() { + unimplemented!(); + } + (Self::Vector::code(vector), 0.0) + } +} + +impl Operator for Op { + type Vector = Rabitq4Owned; + + type DistanceAccessor = HalfbyteDistanceAccessor; fn block_access( lut: &BlockLut, @@ -735,7 +885,7 @@ impl Operator for Op { RAccess::new( (&lut.1, ()), BlockAccessor([0_u32; 32], move |value, code, _delta| { - rabitq::bit::block::half_process_dot(value, code, lut.0) + rabitq::bit::block::half_process_l2s(value, code, lut.0) }), ) } @@ -755,12 +905,63 @@ impl Operator for Op { factor_ip: metadata[2], factor_err: metadata[3], }; - rabitq::bit::binary::half_process_dot(value, code, lut.0) + rabitq::bit::binary::half_process_l2s(value, code, lut.0) } } fn build( - vector: Rabitq8Borrowed<'_>, + vector: Rabitq4Borrowed<'_>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if centroid.is_some() { + unimplemented!(); + } + (Self::Vector::code(vector), 0.0) + } +} + +impl Operator for Op { + type Vector = Rabitq4Owned; + + type DistanceAccessor = HalfbyteDistanceAccessor; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + assert!(!is_residual); + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |sum, code, _delta| { + rabitq::bit::block::half_process_dot(sum, code, lut.0) + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + assert!(!is_residual); + move |metadata: [f32; 4], elements: &[u64], _delta: f32| { + let sum = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + rabitq::bit::binary::half_process_dot(sum, code, lut.0) + } + } + + fn build( + vector: Rabitq4Borrowed<'_>, centroid: Option, ) -> (rabitq::bit::Code, f32) { if centroid.is_some() { diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs index 446630c7..53b1adb5 100644 --- a/crates/vchordrq/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -18,7 +18,7 @@ use crate::tuples::{MetaTuple, WithReader}; use crate::{RerankMethod, vectors}; use always_equal::AlwaysEqual; use distance::Distance; -use index::accessor::{Accessor2, LTryAccess}; +use index::accessor::{Accessor2, DefaultWithDimension, LTryAccess}; use index::fetch::BorrowedIter; use index::packed::PackedRefMut; use index::prefetcher::Prefetcher; @@ -27,7 +27,7 @@ use std::cmp::Reverse; use std::collections::BinaryHeap; use std::marker::PhantomData; use std::num::NonZero; -use vector::VectorOwned; +use vector::{VectorBorrowed, VectorOwned}; type Result = (Reverse, AlwaysEqual>); @@ -89,6 +89,7 @@ pub fn rerank_index< vector: O::Vector, prefetcher: P, ) -> Reranker, P::Guards, u16) -> Option, P, W> { + let dim = vector.as_borrowed().dim(); Reranker { prefetcher, cache: BinaryHeap::new(), @@ -99,7 +100,7 @@ pub fn rerank_index< payload, LTryAccess::new( O::Vector::unpack(vector.as_borrowed()), - O::DistanceAccessor::default(), + O::DistanceAccessor::default_with_dimension(dim), ), ) }), @@ -118,6 +119,7 @@ pub fn rerank_heap< prefetcher: P, mut fetch: impl FnMut(NonZero) -> Option + 'b, ) -> Reranker, P::Guards, u16) -> Option, P, W> { + let dim = vector.as_borrowed().dim(); Reranker { prefetcher, cache: BinaryHeap::new(), @@ -125,7 +127,7 @@ pub fn rerank_heap< let unpack = O::Vector::unpack(vector.as_borrowed()); let vector = fetch(payload)?; let vector = O::Vector::unpack(vector.as_borrowed()); - let mut accessor = O::DistanceAccessor::default(); + let mut accessor = O::DistanceAccessor::default_with_dimension(dim); accessor.push(unpack.0, vector.0); let distance = accessor.finish(unpack.1, vector.1); Some(distance) diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index 583569a2..4855c61f 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -20,7 +20,7 @@ use crate::tuples::*; use crate::{Opaque, centroids, tape}; use always_equal::AlwaysEqual; use distance::Distance; -use index::accessor::{FunctionalAccessor, LAccess}; +use index::accessor::{DefaultWithDimension, FunctionalAccessor, LAccess}; use index::bump::Bump; use index::fetch::BorrowedIter; use index::packed::{PackedRefMut4, PackedRefMut8}; @@ -51,11 +51,11 @@ where let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); - let dims = meta_tuple.dims(); + let dim = meta_tuple.dim(); let is_residual = meta_tuple.is_residual(); let height_of_root = meta_tuple.height_of_root(); let cells = meta_tuple.cells().to_vec(); - assert_eq!(dims, vector.dims(), "unmatched dimensions"); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); if height_of_root as usize != 1 + probes.len() { panic!( "usage: need {} probes, but {} probes provided", @@ -73,7 +73,10 @@ where let distance = centroids::read::( prefetch.map(|id| index.read(id)), head, - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), ); let norm = meta_tuple.centroid_norm(); let first = meta_tuple.first(); @@ -118,7 +121,10 @@ where let distance = centroids::read::( prefetch, head, - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), ); cache.push((Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))); } @@ -212,11 +218,11 @@ where let meta_guard = index.read(0); let meta_bytes = meta_guard.get(1).expect("data corruption"); let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); - let dims = meta_tuple.dims(); + let dim = meta_tuple.dim(); let is_residual = meta_tuple.is_residual(); let height_of_root = meta_tuple.height_of_root(); let cells = meta_tuple.cells().to_vec(); - assert_eq!(dims, vector.dims(), "unmatched dimensions"); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); if height_of_root as usize != 1 + probes.len() { panic!( "usage: need {} probes, but {} probes provided", @@ -234,7 +240,10 @@ where let distance = centroids::read::( prefetch.map(|id| index.read(id)), head, - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), ); let norm = meta_tuple.centroid_norm(); let first = meta_tuple.first(); @@ -279,7 +288,10 @@ where let distance = centroids::read::( prefetch, head, - LAccess::new(O::Vector::unpack(vector), O::DistanceAccessor::default()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), ); cache.push((Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))); } diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index 29a6ea7a..e02574b4 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -49,7 +49,7 @@ pub trait WithWriter: Tuple { #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] struct MetaTupleHeader { version: u64, - dims: u32, + dim: u32, height_of_root: u32, is_residual: Bool, rerank_in_heap: Bool, @@ -70,7 +70,7 @@ struct MetaTupleHeader { } pub struct MetaTuple { - pub dims: u32, + pub dim: u32, pub height_of_root: u32, pub is_residual: bool, pub rerank_in_heap: bool, @@ -90,7 +90,7 @@ impl Tuple for MetaTuple { let mut buffer = Vec::::new(); match self { MetaTuple { - dims, + dim, height_of_root, is_residual, rerank_in_heap, @@ -130,7 +130,7 @@ impl Tuple for MetaTuple { buffer[size_of::()..][..size_of::()].copy_from_slice( MetaTupleHeader { version: VERSION, - dims: *dims, + dim: *dim, height_of_root: *height_of_root, is_residual: (*is_residual).into(), rerank_in_heap: (*rerank_in_heap).into(), @@ -195,8 +195,8 @@ pub struct MetaTupleReader<'a> { } impl<'a> MetaTupleReader<'a> { - pub fn dims(self) -> u32 { - self.header.dims + pub fn dim(self) -> u32 { + self.header.dim } pub fn height_of_root(self) -> u32 { self.header.height_of_root diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index 23bd43c4..7097cb8d 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -15,6 +15,7 @@ use serde::{Deserialize, Serialize}; use simd::f16; use validator::{Validate, ValidationError}; +use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; @@ -57,6 +58,7 @@ pub enum OwnedVector { Vecf32(VectOwned), Vecf16(VectOwned), Rabitq8(Rabitq8Owned), + Rabitq4(Rabitq4Owned), } #[derive(Debug, Clone, Copy)] @@ -64,6 +66,7 @@ pub enum BorrowedVector<'a> { Vecf32(VectBorrowed<'a, f32>), Vecf16(VectBorrowed<'a, f16>), Rabitq8(Rabitq8Borrowed<'a>), + Rabitq4(Rabitq4Borrowed<'a>), } #[repr(u8)] @@ -79,14 +82,16 @@ pub enum VectorKind { Vecf32, Vecf16, Rabitq8, + Rabitq4, } impl VectorKind { - pub fn element_size(self) -> u32 { + pub fn number_of_bits_of_an_elements(self) -> u32 { match self { - VectorKind::Vecf32 => size_of::() as _, - VectorKind::Vecf16 => size_of::() as _, - VectorKind::Rabitq8 => size_of::() as _, + VectorKind::Vecf32 => 32, + VectorKind::Vecf16 => 16, + VectorKind::Rabitq8 => 8, + VectorKind::Rabitq4 => 8, } } } @@ -95,14 +100,14 @@ impl VectorKind { #[validate(schema(function = "Self::validate_self"))] pub struct VectorOptions { #[validate(range(min = 1))] - pub dims: u32, + pub dim: u32, pub v: VectorKind, pub d: DistanceKind, } impl VectorOptions { pub fn validate_self(&self) -> Result<(), ValidationError> { - match (self.v, self.d, self.dims) { + match (self.v, self.d, self.dim) { (_, _, 1..=60000) => Ok(()), _ => Err(ValidationError::new("invalid vector options")), } diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs index dd549d34..a1afca58 100644 --- a/crates/vector/src/bvect.rs +++ b/crates/vector/src/bvect.rs @@ -20,42 +20,42 @@ pub const BVECTOR_WIDTH: u32 = u64::BITS; // When using binary vector, please ensure that the padding bits are always zero. #[derive(Debug, Clone)] pub struct BVectOwned { - dims: u32, + dim: u32, data: Vec, } impl BVectOwned { #[inline(always)] - pub fn new(dims: u32, data: Vec) -> Self { - Self::new_checked(dims, data).expect("invalid data") + pub fn new(dim: u32, data: Vec) -> Self { + Self::new_checked(dim, data).expect("invalid data") } #[inline(always)] - pub fn new_checked(dims: u32, data: Vec) -> Option { - if !(1..=65535).contains(&dims) { + pub fn new_checked(dim: u32, data: Vec) -> Option { + if !(1..=65535).contains(&dim) { return None; } - if data.len() != dims.div_ceil(BVECTOR_WIDTH) as usize { + if data.len() != dim.div_ceil(BVECTOR_WIDTH) as usize { return None; } - if dims % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dims % BVECTOR_WIDTH) != 0 { + if dim % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dim % BVECTOR_WIDTH) != 0 { return None; } #[allow(unsafe_code)] unsafe { - Some(Self::new_unchecked(dims, data)) + Some(Self::new_unchecked(dim, data)) } } /// # Safety /// - /// * `dims` must be in `1..=65535`. + /// * `dim` must be in `1..=65535`. /// * `data` must be of the correct length. /// * The padding bits must be zero. #[allow(unsafe_code)] #[inline(always)] - pub unsafe fn new_unchecked(dims: u32, data: Vec) -> Self { - Self { dims, data } + pub unsafe fn new_unchecked(dim: u32, data: Vec) -> Self { + Self { dim, data } } } @@ -65,7 +65,7 @@ impl VectorOwned for BVectOwned { #[inline(always)] fn as_borrowed(&self) -> BVectBorrowed<'_> { BVectBorrowed { - dims: self.dims, + dim: self.dim, data: &self.data, } } @@ -73,42 +73,42 @@ impl VectorOwned for BVectOwned { #[derive(Debug, Clone, Copy)] pub struct BVectBorrowed<'a> { - dims: u32, + dim: u32, data: &'a [u64], } impl<'a> BVectBorrowed<'a> { #[inline(always)] - pub fn new(dims: u32, data: &'a [u64]) -> Self { - Self::new_checked(dims, data).expect("invalid data") + pub fn new(dim: u32, data: &'a [u64]) -> Self { + Self::new_checked(dim, data).expect("invalid data") } #[inline(always)] - pub fn new_checked(dims: u32, data: &'a [u64]) -> Option { - if !(1..=65535).contains(&dims) { + pub fn new_checked(dim: u32, data: &'a [u64]) -> Option { + if !(1..=65535).contains(&dim) { return None; } - if data.len() != dims.div_ceil(BVECTOR_WIDTH) as usize { + if data.len() != dim.div_ceil(BVECTOR_WIDTH) as usize { return None; } - if dims % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dims % BVECTOR_WIDTH) != 0 { + if dim % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dim % BVECTOR_WIDTH) != 0 { return None; } #[allow(unsafe_code)] unsafe { - Some(Self::new_unchecked(dims, data)) + Some(Self::new_unchecked(dim, data)) } } /// # Safety /// - /// * `dims` must be in `1..=65535`. + /// * `dim` must be in `1..=65535`. /// * `data` must be of the correct length. /// * The padding bits must be zero. #[allow(unsafe_code)] #[inline(always)] - pub unsafe fn new_unchecked(dims: u32, data: &'a [u64]) -> Self { - Self { dims, data } + pub unsafe fn new_unchecked(dim: u32, data: &'a [u64]) -> Self { + Self { dim, data } } #[inline(always)] @@ -118,7 +118,7 @@ impl<'a> BVectBorrowed<'a> { #[inline(always)] pub fn get(&self, index: u32) -> bool { - assert!(index < self.dims); + assert!(index < self.dim); self.data[(index / BVECTOR_WIDTH) as usize] & (1 << (index % BVECTOR_WIDTH)) != 0 } @@ -126,7 +126,7 @@ impl<'a> BVectBorrowed<'a> { pub fn iter(self) -> impl Iterator + 'a { let mut index = 0_u32; std::iter::from_fn(move || { - if index < self.dims { + if index < self.dim { let result = self.data[(index / BVECTOR_WIDTH) as usize] & (1 << (index % BVECTOR_WIDTH)) != 0; @@ -143,13 +143,13 @@ impl VectorBorrowed for BVectBorrowed<'_> { type Owned = BVectOwned; #[inline(always)] - fn dims(&self) -> u32 { - self.dims + fn dim(&self) -> u32 { + self.dim } fn own(&self) -> BVectOwned { BVectOwned { - dims: self.dims, + dim: self.dim, data: self.data.to_vec(), } } @@ -203,27 +203,27 @@ impl VectorBorrowed for BVectBorrowed<'_> { } fn operator_and(&self, rhs: Self) -> Self::Owned { - assert_eq!(self.dims, rhs.dims); + assert_eq!(self.dim, rhs.dim); let data = simd::bit::vector_and(self.data, rhs.data); - BVectOwned::new(self.dims, data) + BVectOwned::new(self.dim, data) } fn operator_or(&self, rhs: Self) -> Self::Owned { - assert_eq!(self.dims, rhs.dims); + assert_eq!(self.dim, rhs.dim); let data = simd::bit::vector_or(self.data, rhs.data); - BVectOwned::new(self.dims, data) + BVectOwned::new(self.dim, data) } fn operator_xor(&self, rhs: Self) -> Self::Owned { - assert_eq!(self.dims, rhs.dims); + assert_eq!(self.dim, rhs.dim); let data = simd::bit::vector_xor(self.data, rhs.data); - BVectOwned::new(self.dims, data) + BVectOwned::new(self.dim, data) } } impl PartialEq for BVectBorrowed<'_> { fn eq(&self, other: &Self) -> bool { - if self.dims != other.dims { + if self.dim != other.dim { return false; } for (&l, &r) in self.data.iter().zip(other.data.iter()) { @@ -240,7 +240,7 @@ impl PartialEq for BVectBorrowed<'_> { impl PartialOrd for BVectBorrowed<'_> { fn partial_cmp(&self, other: &Self) -> Option { use std::cmp::Ordering; - if self.dims != other.dims { + if self.dim != other.dim { return None; } for (&l, &r) in self.data.iter().zip(other.data.iter()) { diff --git a/crates/vector/src/lib.rs b/crates/vector/src/lib.rs index 67403990..1009043a 100644 --- a/crates/vector/src/lib.rs +++ b/crates/vector/src/lib.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. pub mod bvect; +pub mod rabitq4; pub mod rabitq8; pub mod svect; pub mod vect; @@ -28,7 +29,7 @@ pub trait VectorBorrowed: Copy { fn own(&self) -> Self::Owned; - fn dims(&self) -> u32; + fn dim(&self) -> u32; fn norm(&self) -> f32; diff --git a/crates/vector/src/rabitq4.rs b/crates/vector/src/rabitq4.rs new file mode 100644 index 00000000..5c24abfd --- /dev/null +++ b/crates/vector/src/rabitq4.rs @@ -0,0 +1,384 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; + +#[derive(Debug, Clone)] +pub struct Rabitq4Owned { + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, +} + +impl Rabitq4Owned { + #[inline(always)] + pub fn new( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, + ) -> Self { + Self::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + .expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, + ) -> Option { + if !(1..=65535).contains(&dim) { + return None; + } + if dim.div_ceil(2) as usize != packed_code.len() { + return None; + } + if dim % 2 == 1 && packed_code.last().copied().unwrap_or_default() >> 4 != 0 { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { + Self::new_unchecked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + }) + } + + /// # Safety + /// + /// * `dim` must not be zero. + /// * `dim` must be less than 65536. + /// * `dim` must be equal to 1/2 of `code.len()`, rounding to infinity. + /// * `packed_code` is filled with zero bits correctly. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, + ) -> Self { + Self { + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + } + } +} + +impl VectorOwned for Rabitq4Owned { + type Borrowed<'a> = Rabitq4Borrowed<'a>; + + #[inline(always)] + fn as_borrowed(&self) -> Rabitq4Borrowed<'_> { + Rabitq4Borrowed { + dim: self.dim, + sum_of_x2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x, + packed_code: self.packed_code.as_slice(), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct Rabitq4Borrowed<'a> { + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], +} + +impl<'a> Rabitq4Borrowed<'a> { + #[inline(always)] + pub fn new( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], + ) -> Self { + Self::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + .expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], + ) -> Option { + if !(1..=65535).contains(&packed_code.len()) { + return None; + } + if dim.div_ceil(2) as usize != packed_code.len() { + return None; + } + if dim % 2 == 1 && packed_code.last().copied().unwrap_or_default() >> 4 != 0 { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { + Self::new_unchecked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + }) + } + + /// # Safety + /// + /// * `dim` must not be zero. + /// * `dim` must be less than 65536. + /// * `dim` must be equal to 1/2 of `code.len()`, rounding to infinity. + /// * `packed_code` is filled with zero bits correctly. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], + ) -> Self { + Self { + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + } + } + + #[inline(always)] + pub fn sum_of_x2(&self) -> f32 { + self.sum_of_x2 + } + + #[inline(always)] + pub fn norm_of_lattice(&self) -> f32 { + self.norm_of_lattice + } + + #[inline(always)] + pub fn sum_of_code(&self) -> f32 { + self.sum_of_code + } + + #[inline(always)] + pub fn sum_of_abs_x(&self) -> f32 { + self.sum_of_abs_x + } + + #[inline(always)] + pub fn packed_code(&self) -> &'a [u8] { + self.packed_code + } + + #[inline(always)] + pub fn unpacked_code(&self) -> impl Iterator { + self.packed_code + .iter() + .flat_map(|x| [x & 0xf, x >> 4]) + .take(self.dim as _) + } +} + +impl VectorBorrowed for Rabitq4Borrowed<'_> { + type Owned = Rabitq4Owned; + + #[inline(always)] + fn dim(&self) -> u32 { + self.dim + } + + #[inline(always)] + fn own(&self) -> Rabitq4Owned { + Rabitq4Owned { + dim: self.dim, + sum_of_x2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x, + packed_code: self.packed_code.to_owned(), + } + } + + #[inline(always)] + fn norm(&self) -> f32 { + self.sum_of_x2.sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + let dim = self.dim(); + let sum = rabitq::halfbyte::binary::accumulate(self.packed_code, rhs.packed_code); + Distance::from_f32( + rabitq::halfbyte::binary::half_process_dot( + dim, + sum, + rabitq::halfbyte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::halfbyte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_l2s(self, rhs: Self) -> Distance { + let dim = self.dim(); + let sum = rabitq::halfbyte::binary::accumulate(self.packed_code, rhs.packed_code); + Distance::from_f32( + rabitq::halfbyte::binary::half_process_l2s( + dim, + sum, + rabitq::halfbyte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::halfbyte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_cos(self, rhs: Self) -> Distance { + let dim = self.dim(); + let sum = rabitq::halfbyte::binary::accumulate(self.packed_code, rhs.packed_code); + Distance::from_f32( + rabitq::halfbyte::binary::half_process_cos( + dim, + sum, + rabitq::halfbyte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::halfbyte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_hamming(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_jaccard(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn function_normalize(&self) -> Rabitq4Owned { + Rabitq4Owned { + dim: self.dim, + sum_of_x2: 1.0, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x / self.sum_of_x2.sqrt(), + packed_code: self.packed_code.to_owned(), + } + } + + fn operator_add(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_sub(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_mul(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_and(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_or(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_xor(&self, _: Self) -> Self::Owned { + unimplemented!() + } +} diff --git a/crates/vector/src/rabitq8.rs b/crates/vector/src/rabitq8.rs index 0cbd074f..738b5fa2 100644 --- a/crates/vector/src/rabitq8.rs +++ b/crates/vector/src/rabitq8.rs @@ -17,61 +17,85 @@ use distance::Distance; #[derive(Debug, Clone)] pub struct Rabitq8Owned { + dim: u32, sum_of_x2: f32, norm_of_lattice: f32, sum_of_code: f32, sum_of_abs_x: f32, - code: Vec, + packed_code: Vec, } impl Rabitq8Owned { #[inline(always)] pub fn new( + dim: u32, sum_of_x2: f32, norm_of_lattice: f32, sum_of_code: f32, sum_of_abs_x: f32, - code: Vec, + packed_code: Vec, ) -> Self { - Self::new_checked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, code) - .expect("invalid data") + Self::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + .expect("invalid data") } #[inline(always)] pub fn new_checked( + dim: u32, sum_of_x2: f32, norm_of_lattice: f32, sum_of_code: f32, sum_of_abs_x: f32, - code: Vec, + packed_code: Vec, ) -> Option { - if !(1..=65535).contains(&code.len()) { + if !(1..=65535).contains(&dim) { + return None; + } + if dim.div_ceil(1) as usize != packed_code.len() { return None; } #[allow(unsafe_code)] Some(unsafe { - Self::new_unchecked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, code) + Self::new_unchecked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) }) } /// # Safety /// - /// * `code.len()` must not be zero. + /// * `dim` must not be zero. + /// * `dim` must be less than 65536. + /// * `dim` must be equal to `code.len()`. #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked( + dim: u32, sum_of_x2: f32, norm_of_lattice: f32, sum_of_code: f32, sum_of_abs_x: f32, - code: Vec, + packed_code: Vec, ) -> Self { Self { + dim, sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, - code, + packed_code, } } } @@ -82,72 +106,97 @@ impl VectorOwned for Rabitq8Owned { #[inline(always)] fn as_borrowed(&self) -> Rabitq8Borrowed<'_> { Rabitq8Borrowed { + dim: self.dim, sum_of_x2: self.sum_of_x2, norm_of_lattice: self.norm_of_lattice, sum_of_code: self.sum_of_code, sum_of_abs_x: self.sum_of_abs_x, - code: self.code.as_slice(), + packed_code: self.packed_code.as_slice(), } } } #[derive(Debug, Clone, Copy)] pub struct Rabitq8Borrowed<'a> { + dim: u32, sum_of_x2: f32, norm_of_lattice: f32, sum_of_code: f32, sum_of_abs_x: f32, - code: &'a [u8], + packed_code: &'a [u8], } impl<'a> Rabitq8Borrowed<'a> { #[inline(always)] pub fn new( + dim: u32, sum_of_x2: f32, norm_of_lattice: f32, sum_of_code: f32, sum_of_abs_x: f32, - code: &'a [u8], + packed_code: &'a [u8], ) -> Self { - Self::new_checked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, code) - .expect("invalid data") + Self::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + .expect("invalid data") } #[inline(always)] pub fn new_checked( + dim: u32, sum_of_x2: f32, norm_of_lattice: f32, sum_of_code: f32, sum_of_abs_x: f32, - code: &'a [u8], + packed_code: &'a [u8], ) -> Option { - if !(1..=65535).contains(&code.len()) { + if !(1..=65535).contains(&dim) { + return None; + } + if dim.div_ceil(1) as usize != packed_code.len() { return None; } #[allow(unsafe_code)] Some(unsafe { - Self::new_unchecked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, code) + Self::new_unchecked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) }) } /// # Safety /// - /// * `code.len()` must not be zero. + /// * `dim` must not be zero. + /// * `dim` must be less than 65536. + /// * `dim` must be equal to `code.len()`. #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked( + dim: u32, sum_of_x2: f32, norm_of_lattice: f32, sum_of_code: f32, sum_of_abs_x: f32, - code: &'a [u8], + packed_code: &'a [u8], ) -> Self { Self { + dim, sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, - code, + packed_code, } } @@ -172,8 +221,13 @@ impl<'a> Rabitq8Borrowed<'a> { } #[inline(always)] - pub fn code(&self) -> &'a [u8] { - self.code + pub fn packed_code(&self) -> &'a [u8] { + self.packed_code + } + + #[inline(always)] + pub fn unpacked_code(&self) -> impl Iterator { + self.packed_code.iter().copied() } } @@ -181,18 +235,19 @@ impl VectorBorrowed for Rabitq8Borrowed<'_> { type Owned = Rabitq8Owned; #[inline(always)] - fn dims(&self) -> u32 { - self.code.len() as u32 + fn dim(&self) -> u32 { + self.dim } #[inline(always)] fn own(&self) -> Rabitq8Owned { Rabitq8Owned { + dim: self.dim, sum_of_x2: self.sum_of_x2, norm_of_lattice: self.norm_of_lattice, sum_of_code: self.sum_of_code, sum_of_abs_x: self.sum_of_abs_x, - code: self.code.to_owned(), + packed_code: self.packed_code.to_owned(), } } @@ -203,13 +258,12 @@ impl VectorBorrowed for Rabitq8Borrowed<'_> { #[inline(always)] fn operator_dot(self, rhs: Self) -> Distance { - assert_eq!(self.code.len(), rhs.code.len()); - let n = self.code.len() as u32; - let value = rabitq::byte::binary::accumulate(self.code, rhs.code); + let dim = self.dim(); + let sum = rabitq::byte::binary::accumulate(self.packed_code, rhs.packed_code); Distance::from_f32( rabitq::byte::binary::half_process_dot( - n, - value, + dim, + sum, rabitq::byte::CodeMetadata { dis_u_2: self.sum_of_x2, norm_of_lattice: self.norm_of_lattice, @@ -227,13 +281,12 @@ impl VectorBorrowed for Rabitq8Borrowed<'_> { #[inline(always)] fn operator_l2s(self, rhs: Self) -> Distance { - assert_eq!(self.code.len(), rhs.code.len()); - let n = self.code.len() as u32; - let value = rabitq::byte::binary::accumulate(self.code, rhs.code); + let dim = self.dim(); + let sum = rabitq::byte::binary::accumulate(self.packed_code, rhs.packed_code); Distance::from_f32( - rabitq::byte::binary::half_process_l2( - n, - value, + rabitq::byte::binary::half_process_l2s( + dim, + sum, rabitq::byte::CodeMetadata { dis_u_2: self.sum_of_x2, norm_of_lattice: self.norm_of_lattice, @@ -251,13 +304,12 @@ impl VectorBorrowed for Rabitq8Borrowed<'_> { #[inline(always)] fn operator_cos(self, rhs: Self) -> Distance { - assert_eq!(self.code.len(), rhs.code.len()); - let n = self.code.len() as u32; - let value = rabitq::byte::binary::accumulate(self.code, rhs.code); + let dim = self.dim(); + let sum = rabitq::byte::binary::accumulate(self.packed_code, rhs.packed_code); Distance::from_f32( rabitq::byte::binary::half_process_cos( - n, - value, + dim, + sum, rabitq::byte::CodeMetadata { dis_u_2: self.sum_of_x2, norm_of_lattice: self.norm_of_lattice, @@ -286,11 +338,12 @@ impl VectorBorrowed for Rabitq8Borrowed<'_> { #[inline(always)] fn function_normalize(&self) -> Rabitq8Owned { Rabitq8Owned { + dim: self.dim, sum_of_x2: 1.0, norm_of_lattice: self.norm_of_lattice, sum_of_code: self.sum_of_code, sum_of_abs_x: self.sum_of_abs_x / self.sum_of_x2.sqrt(), - code: self.code.to_owned(), + packed_code: self.packed_code.to_owned(), } } diff --git a/crates/vector/src/svect.rs b/crates/vector/src/svect.rs index 1b9ebcee..caa3f07b 100644 --- a/crates/vector/src/svect.rs +++ b/crates/vector/src/svect.rs @@ -18,20 +18,20 @@ use simd::Floating; #[derive(Debug, Clone)] pub struct SVectOwned { - dims: u32, + dim: u32, indexes: Vec, values: Vec, } impl SVectOwned { #[inline(always)] - pub fn new(dims: u32, indexes: Vec, values: Vec) -> Self { - Self::new_checked(dims, indexes, values).expect("invalid data") + pub fn new(dim: u32, indexes: Vec, values: Vec) -> Self { + Self::new_checked(dim, indexes, values).expect("invalid data") } #[inline(always)] - pub fn new_checked(dims: u32, indexes: Vec, values: Vec) -> Option { - if !(1..=1_048_575).contains(&dims) { + pub fn new_checked(dim: u32, indexes: Vec, values: Vec) -> Option { + if !(1..=1_048_575).contains(&dim) { return None; } if indexes.len() != values.len() { @@ -43,7 +43,7 @@ impl SVectOwned { return None; } } - if len != 0 && !(indexes[len - 1] < dims) { + if len != 0 && !(indexes[len - 1] < dim) { return None; } if S::reduce_or_of_is_zero_x(&values) { @@ -51,21 +51,21 @@ impl SVectOwned { } #[allow(unsafe_code)] unsafe { - Some(Self::new_unchecked(dims, indexes, values)) + Some(Self::new_unchecked(dim, indexes, values)) } } /// # Safety /// - /// * `dims` must be in `1..=1_048_575`. + /// * `dim` must be in `1..=1_048_575`. /// * `indexes.len()` must be equal to `values.len()`. - /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dims`. + /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dim`. /// * A floating number in `values` must not be positive zero or negative zero. #[allow(unsafe_code)] #[inline(always)] - pub unsafe fn new_unchecked(dims: u32, indexes: Vec, values: Vec) -> Self { + pub unsafe fn new_unchecked(dim: u32, indexes: Vec, values: Vec) -> Self { Self { - dims, + dim, indexes, values, } @@ -88,7 +88,7 @@ impl VectorOwned for SVectOwned { #[inline(always)] fn as_borrowed(&self) -> SVectBorrowed<'_, S> { SVectBorrowed { - dims: self.dims, + dim: self.dim, indexes: &self.indexes, values: &self.values, } @@ -97,20 +97,20 @@ impl VectorOwned for SVectOwned { #[derive(Debug, Clone, Copy)] pub struct SVectBorrowed<'a, S> { - dims: u32, + dim: u32, indexes: &'a [u32], values: &'a [S], } impl<'a, S: Floating> SVectBorrowed<'a, S> { #[inline(always)] - pub fn new(dims: u32, indexes: &'a [u32], values: &'a [S]) -> Self { - Self::new_checked(dims, indexes, values).expect("invalid data") + pub fn new(dim: u32, indexes: &'a [u32], values: &'a [S]) -> Self { + Self::new_checked(dim, indexes, values).expect("invalid data") } #[inline(always)] - pub fn new_checked(dims: u32, indexes: &'a [u32], values: &'a [S]) -> Option { - if !(1..=1_048_575).contains(&dims) { + pub fn new_checked(dim: u32, indexes: &'a [u32], values: &'a [S]) -> Option { + if !(1..=1_048_575).contains(&dim) { return None; } if indexes.len() != values.len() { @@ -122,7 +122,7 @@ impl<'a, S: Floating> SVectBorrowed<'a, S> { return None; } } - if len != 0 && !(indexes[len - 1] < dims) { + if len != 0 && !(indexes[len - 1] < dim) { return None; } for i in 0..len { @@ -132,21 +132,21 @@ impl<'a, S: Floating> SVectBorrowed<'a, S> { } #[allow(unsafe_code)] unsafe { - Some(Self::new_unchecked(dims, indexes, values)) + Some(Self::new_unchecked(dim, indexes, values)) } } /// # Safety /// - /// * `dims` must be in `1..=1_048_575`. + /// * `dim` must be in `1..=1_048_575`. /// * `indexes.len()` must be equal to `values.len()`. - /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dims`. + /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dim`. /// * A floating number in `values` must not be positive zero or negative zero. #[inline(always)] #[allow(unsafe_code)] - pub unsafe fn new_unchecked(dims: u32, indexes: &'a [u32], values: &'a [S]) -> Self { + pub unsafe fn new_unchecked(dim: u32, indexes: &'a [u32], values: &'a [S]) -> Self { Self { - dims, + dim, indexes, values, } @@ -177,14 +177,14 @@ impl VectorBorrowed for SVectBorrowed<'_, S> { type Owned = SVectOwned; #[inline(always)] - fn dims(&self) -> u32 { - self.dims + fn dim(&self) -> u32 { + self.dim } #[inline(always)] fn own(&self) -> SVectOwned { SVectOwned { - dims: self.dims, + dim: self.dim, indexes: self.indexes.to_vec(), values: self.values.to_vec(), } @@ -242,11 +242,11 @@ impl VectorBorrowed for SVectBorrowed<'_, S> { } indexes.truncate(j); values.truncate(j); - SVectOwned::new(self.dims, indexes, values) + SVectOwned::new(self.dim, indexes, values) } fn operator_add(&self, rhs: Self) -> Self::Owned { - assert_eq!(self.dims, rhs.dims); + assert_eq!(self.dim, rhs.dim); let size1 = self.len(); let size2 = rhs.len(); let mut pos1 = 0; @@ -280,11 +280,11 @@ impl VectorBorrowed for SVectBorrowed<'_, S> { } indexes.truncate(pos); values.truncate(pos); - SVectOwned::new(self.dims, indexes, values) + SVectOwned::new(self.dim, indexes, values) } fn operator_sub(&self, rhs: Self) -> Self::Owned { - assert_eq!(self.dims, rhs.dims); + assert_eq!(self.dim, rhs.dim); let size1 = self.len(); let size2 = rhs.len(); let mut pos1 = 0; @@ -318,11 +318,11 @@ impl VectorBorrowed for SVectBorrowed<'_, S> { } indexes.truncate(pos); values.truncate(pos); - SVectOwned::new(self.dims, indexes, values) + SVectOwned::new(self.dim, indexes, values) } fn operator_mul(&self, rhs: Self) -> Self::Owned { - assert_eq!(self.dims, rhs.dims); + assert_eq!(self.dim, rhs.dim); let size1 = self.len(); let size2 = rhs.len(); let mut pos1 = 0; @@ -355,7 +355,7 @@ impl VectorBorrowed for SVectBorrowed<'_, S> { } indexes.truncate(pos); values.truncate(pos); - SVectOwned::new(self.dims, indexes, values) + SVectOwned::new(self.dim, indexes, values) } fn operator_and(&self, _: Self) -> Self::Owned { @@ -373,7 +373,7 @@ impl VectorBorrowed for SVectBorrowed<'_, S> { impl PartialEq for SVectBorrowed<'_, S> { fn eq(&self, other: &Self) -> bool { - if self.dims != other.dims { + if self.dim != other.dim { return false; } if self.indexes.len() != other.indexes.len() { @@ -396,7 +396,7 @@ impl PartialEq for SVectBorrowed<'_, S> { impl PartialOrd for SVectBorrowed<'_, S> { fn partial_cmp(&self, other: &Self) -> Option { use std::cmp::Ordering; - if self.dims != other.dims { + if self.dim != other.dim { return None; } let mut lhs = self diff --git a/crates/vector/src/vect.rs b/crates/vector/src/vect.rs index ac1c50bf..b78f2a68 100644 --- a/crates/vector/src/vect.rs +++ b/crates/vector/src/vect.rs @@ -39,6 +39,7 @@ impl VectOwned { /// # Safety /// /// * `slice.len()` must not be zero. + /// * `slice.len()` must be less than 65536. #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked(slice: Vec) -> Self { @@ -92,6 +93,7 @@ impl<'a, S: Floating> VectBorrowed<'a, S> { /// # Safety /// /// * `slice.len()` must not be zero. + /// * `slice.len()` must be less than 65536. #[allow(unsafe_code)] #[inline(always)] pub unsafe fn new_unchecked(slice: &'a [S]) -> Self { @@ -108,7 +110,7 @@ impl VectorBorrowed for VectBorrowed<'_, S> { type Owned = VectOwned; #[inline(always)] - fn dims(&self) -> u32 { + fn dim(&self) -> u32 { self.0.len() as u32 } diff --git a/src/datatype/binary_rabitq4.rs b/src/datatype/binary_rabitq4.rs new file mode 100644 index 00000000..d3a9d56c --- /dev/null +++ b/src/datatype/binary_rabitq4.rs @@ -0,0 +1,96 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use pgrx::datum::Internal; +use pgrx::pg_sys::Oid; +use vector::VectorBorrowed; +use vector::rabitq4::Rabitq4Borrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_send(vector: Rabitq4Input<'_>) -> Vec { + let vector = vector.as_borrowed(); + let mut stream = Vec::::new(); + stream.extend(vector.dim().to_be_bytes()); + stream.extend(vector.sum_of_x2().to_be_bytes()); + stream.extend(vector.norm_of_lattice().to_be_bytes()); + stream.extend(vector.sum_of_code().to_be_bytes()); + stream.extend(vector.sum_of_abs_x().to_be_bytes()); + for &c in vector.packed_code() { + stream.extend(c.to_be_bytes()); + } + stream +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Rabitq4Output { + let _ = (oid, typmod); + let buf = unsafe { internal.get_mut::().unwrap() }; + + let dim = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + u32::from_be_bytes(raw) + }; + let sum_of_x2 = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let norm_of_lattice = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let sum_of_code = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let sum_of_abs_x = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let packed_code = { + let mut result = Vec::new(); + for _ in 0..dim.div_ceil(2) { + result.push({ + assert!(buf.cursor < i32::MAX - 1 && buf.cursor + 1 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 1]>().read() }; + buf.cursor += 1; + u8::from_be_bytes(raw) + }); + } + result + }; + + if let Some(x) = Rabitq4Borrowed::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + &packed_code, + ) { + Rabitq4Output::new(x) + } else { + pgrx::error!("detect data corruption"); + } +} diff --git a/src/datatype/binary_rabitq8.rs b/src/datatype/binary_rabitq8.rs index cad31204..2c1af908 100644 --- a/src/datatype/binary_rabitq8.rs +++ b/src/datatype/binary_rabitq8.rs @@ -22,12 +22,12 @@ use vector::rabitq8::Rabitq8Borrowed; fn _vchord_rabitq8_send(vector: Rabitq8Input<'_>) -> Vec { let vector = vector.as_borrowed(); let mut stream = Vec::::new(); - stream.extend(vector.dims().to_be_bytes()); + stream.extend(vector.dim().to_be_bytes()); stream.extend(vector.sum_of_x2().to_be_bytes()); stream.extend(vector.norm_of_lattice().to_be_bytes()); stream.extend(vector.sum_of_code().to_be_bytes()); stream.extend(vector.sum_of_abs_x().to_be_bytes()); - for &c in vector.code() { + for &c in vector.packed_code() { stream.extend(c.to_be_bytes()); } stream @@ -38,7 +38,7 @@ fn _vchord_rabitq8_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Rabitq let _ = (oid, typmod); let buf = unsafe { internal.get_mut::().unwrap() }; - let dims = { + let dim = { assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; buf.cursor += 4; @@ -68,9 +68,9 @@ fn _vchord_rabitq8_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Rabitq buf.cursor += 4; f32::from_be_bytes(raw) }; - let code = { - let mut result = Vec::with_capacity(dims as _); - for _ in 0..dims { + let packed_code = { + let mut result = Vec::new(); + for _ in 0..dim.div_ceil(1) { result.push({ assert!(buf.cursor < i32::MAX - 1 && buf.cursor + 1 <= buf.len); let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 1]>().read() }; @@ -81,9 +81,14 @@ fn _vchord_rabitq8_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Rabitq result }; - if let Some(x) = - Rabitq8Borrowed::new_checked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, &code) - { + if let Some(x) = Rabitq8Borrowed::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + &packed_code, + ) { Rabitq8Output::new(x) } else { pgrx::error!("detect data corruption"); diff --git a/src/datatype/functions_rabitq4.rs b/src/datatype/functions_rabitq4.rs new file mode 100644 index 00000000..031f28b6 --- /dev/null +++ b/src/datatype/functions_rabitq4.rs @@ -0,0 +1,56 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::datatype::memory_halfvec::HalfvecInput; +use crate::datatype::memory_rabitq4::Rabitq4Output; +use crate::datatype::memory_vector::VectorInput; +use simd::{Floating, f16}; +use vector::VectorBorrowed; +use vector::rabitq4::Rabitq4Borrowed; + +#[pgrx::pg_extern(sql = "")] +fn _vchord_vector_quantize_to_rabitq4(vector: VectorInput) -> Rabitq4Output { + let vector = vector.as_borrowed(); + let dim = vector.dim(); + let mut vector = vector.slice().to_vec(); + rabitq::rotate::rotate_inplace(&mut vector); + let (metadata, elements) = rabitq::halfbyte::ugly_code(&vector); + let elements = rabitq::halfbyte::pack_code(&elements); + Rabitq4Output::new(Rabitq4Borrowed::new( + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + &elements, + )) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_halfvec_quantize_to_rabitq4(vector: HalfvecInput) -> Rabitq4Output { + let vector = vector.as_borrowed(); + let dim = vector.dim(); + let mut vector = f16::vector_to_f32(vector.slice()); + rabitq::rotate::rotate_inplace(&mut vector); + let (metadata, elements) = rabitq::halfbyte::ugly_code(&vector); + let elements = rabitq::halfbyte::pack_code(&elements); + Rabitq4Output::new(Rabitq4Borrowed::new( + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + &elements, + )) +} diff --git a/src/datatype/functions_rabitq8.rs b/src/datatype/functions_rabitq8.rs index 870bcf44..1765216b 100644 --- a/src/datatype/functions_rabitq8.rs +++ b/src/datatype/functions_rabitq8.rs @@ -16,32 +16,41 @@ use crate::datatype::memory_halfvec::HalfvecInput; use crate::datatype::memory_rabitq8::Rabitq8Output; use crate::datatype::memory_vector::VectorInput; use simd::{Floating, f16}; +use vector::VectorBorrowed; use vector::rabitq8::Rabitq8Borrowed; #[pgrx::pg_extern(sql = "")] fn _vchord_vector_quantize_to_rabitq8(vector: VectorInput) -> Rabitq8Output { - let mut vector = vector.as_borrowed().slice().to_vec(); + let vector = vector.as_borrowed(); + let dim = vector.dim(); + let mut vector = vector.slice().to_vec(); rabitq::rotate::rotate_inplace(&mut vector); - let code = rabitq::byte::ugly_code(&vector); + let (metadata, elements) = rabitq::byte::ugly_code(&vector); + let elements = rabitq::byte::pack_code(&elements); Rabitq8Output::new(Rabitq8Borrowed::new( - code.0.dis_u_2, - code.0.norm_of_lattice, - code.0.sum_of_code, + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, f32::reduce_sum_of_abs_x(&vector), - &code.1, + &elements, )) } #[pgrx::pg_extern(sql = "")] fn _vchord_halfvec_quantize_to_rabitq8(vector: HalfvecInput) -> Rabitq8Output { - let mut vector = f16::vector_to_f32(vector.as_borrowed().slice()); + let vector = vector.as_borrowed(); + let dim = vector.dim(); + let mut vector = f16::vector_to_f32(vector.slice()); rabitq::rotate::rotate_inplace(&mut vector); - let code = rabitq::byte::ugly_code(&vector); + let (metadata, elements) = rabitq::byte::ugly_code(&vector); + let elements = rabitq::byte::pack_code(&elements); Rabitq8Output::new(Rabitq8Borrowed::new( - code.0.dis_u_2, - code.0.norm_of_lattice, - code.0.sum_of_code, + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, f32::reduce_sum_of_abs_x(&vector), - &code.1, + &elements, )) } diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs index 89d94538..95e5f6cd 100644 --- a/src/datatype/memory_halfvec.rs +++ b/src/datatype/memory_halfvec.rs @@ -24,7 +24,7 @@ use vector::vect::VectBorrowed; #[repr(C)] struct HalfvecHeader { varlena: u32, - dims: u16, + dim: u16, unused: u16, elements: [f16; 0], } @@ -41,7 +41,7 @@ impl HalfvecHeader { let this = this.as_ptr(); VectBorrowed::new(std::slice::from_raw_parts( (&raw const (*this).elements).cast(), - (&raw const (*this).dims).read() as usize, + (&raw const (*this).dim).read() as usize, )) } } @@ -60,8 +60,8 @@ impl HalfvecInput<'_> { let size = varlena as usize; #[cfg(target_endian = "little")] let size = varlena as usize >> 2; - let dims = q.byte_add(4).cast::().read(); - assert_eq!(HalfvecHeader::size_of(dims as _), size); + let dim = q.byte_add(4).cast::().read(); + assert_eq!(HalfvecHeader::size_of(dim as _), size); let unused = q.byte_add(6).cast::().read(); assert_eq!(unused, 0); } @@ -95,8 +95,8 @@ impl HalfvecOutput { let size = varlena as usize; #[cfg(target_endian = "little")] let size = varlena as usize >> 2; - let dims = q.byte_add(4).cast::().read(); - assert_eq!(HalfvecHeader::size_of(dims as _), size); + let dim = q.byte_add(4).cast::().read(); + assert_eq!(HalfvecHeader::size_of(dim as _), size); let unused = q.byte_add(6).cast::().read(); assert_eq!(unused, 0); } @@ -114,7 +114,7 @@ impl HalfvecOutput { (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); #[cfg(target_endian = "little")] (&raw mut (*ptr).varlena).write((size << 2) as u32); - (&raw mut (*ptr).dims).write(vector.dims() as _); + (&raw mut (*ptr).dim).write(vector.dim() as _); (&raw mut (*ptr).unused).write(0); std::ptr::copy_nonoverlapping( slice.as_ptr(), diff --git a/src/datatype/memory_rabitq4.rs b/src/datatype/memory_rabitq4.rs new file mode 100644 index 00000000..3d51643b --- /dev/null +++ b/src/datatype/memory_rabitq4.rs @@ -0,0 +1,276 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use pgrx::datum::{FromDatum, IntoDatum}; +use pgrx::pg_sys::{Datum, Oid}; +use pgrx::pgrx_sql_entity_graph::metadata::*; +use std::marker::PhantomData; +use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::rabitq4::Rabitq4Borrowed; + +#[repr(C)] +struct Rabitq4Header { + varlena: u32, + dim: u16, + unused: u16, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + elements: [u8; 0], +} + +impl Rabitq4Header { + fn size_of_by_dim(dim: usize) -> usize { + Self::size_of_by_len(dim.div_ceil(2)) + } + fn size_of_by_len(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + size_of::() + size_of::() * len + } + unsafe fn as_borrowed<'a>(this: NonNull) -> Rabitq4Borrowed<'a> { + unsafe { + let this = this.as_ptr(); + Rabitq4Borrowed::new( + (&raw const (*this).dim).read() as u32, + (&raw const (*this).sum_of_x2).read(), + (&raw const (*this).norm_of_lattice).read(), + (&raw const (*this).sum_of_code).read(), + (&raw const (*this).sum_of_abs_x).read(), + std::slice::from_raw_parts( + (&raw const (*this).elements).cast(), + (&raw const (*this).dim).read().div_ceil(2) as usize, + ), + ) + } + } +} + +pub struct Rabitq4Input<'a>(NonNull, PhantomData<&'a ()>, bool); + +impl Rabitq4Input<'_> { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(Rabitq4Header::size_of_by_dim(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + Rabitq4Input(q, PhantomData, p != q) + } + pub fn as_borrowed(&self) -> Rabitq4Borrowed<'_> { + unsafe { Rabitq4Header::as_borrowed(self.0) } + } +} + +impl Drop for Rabitq4Input<'_> { + fn drop(&mut self) { + if self.2 { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } + } +} + +pub struct Rabitq4Output(NonNull); + +impl Rabitq4Output { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(Rabitq4Header::size_of_by_dim(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + Self(q) + } + pub fn new(vector: Rabitq4Borrowed<'_>) -> Self { + unsafe { + let packed_code = vector.packed_code(); + let size = Rabitq4Header::size_of_by_len(packed_code.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut Rabitq4Header; + // SET_VARSIZE_4B + #[cfg(target_endian = "big")] + (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); + #[cfg(target_endian = "little")] + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dim).write(vector.dim() as _); + (&raw mut (*ptr).unused).write(0); + (&raw mut (*ptr).sum_of_x2).write(vector.sum_of_x2()); + (&raw mut (*ptr).norm_of_lattice).write(vector.norm_of_lattice()); + (&raw mut (*ptr).sum_of_code).write(vector.sum_of_code()); + (&raw mut (*ptr).sum_of_abs_x).write(vector.sum_of_abs_x()); + std::ptr::copy_nonoverlapping( + packed_code.as_ptr(), + (&raw mut (*ptr).elements).cast(), + packed_code.len(), + ); + Self(NonNull::new(ptr).unwrap()) + } + } + pub fn as_borrowed(&self) -> Rabitq4Borrowed<'_> { + unsafe { Rabitq4Header::as_borrowed(self.0) } + } + fn into_raw(self) -> *mut Rabitq4Header { + let result = self.0.as_ptr(); + std::mem::forget(self); + result + } +} + +impl Drop for Rabitq4Output { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } +} + +// FromDatum + +impl FromDatum for Rabitq4Input<'_> { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +impl FromDatum for Rabitq4Output { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +// IntoDatum + +impl IntoDatum for Rabitq4Output { + fn into_datum(self) -> Option { + Some(Datum::from(self.into_raw())) + } + + fn type_oid() -> Oid { + Oid::INVALID + } + + fn is_compatible_with(_: Oid) -> bool { + true + } +} + +// UnboxDatum + +unsafe impl<'a> pgrx::datum::UnboxDatum for Rabitq4Input<'a> { + type As<'src> + = Rabitq4Input<'src> + where + 'a: 'src; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +unsafe impl pgrx::datum::UnboxDatum for Rabitq4Output { + type As<'src> = Rabitq4Output; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +// SqlTranslatable + +unsafe impl SqlTranslatable for Rabitq4Input<'_> { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("rabitq4"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("rabitq4")))) + } +} + +unsafe impl SqlTranslatable for Rabitq4Output { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("rabitq4"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("rabitq4")))) + } +} + +// ArgAbi + +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Rabitq4Input<'fcx> { + unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { + let index = arg.index(); + unsafe { + arg.unbox_arg_using_from_datum() + .unwrap_or_else(|| panic!("argument {index} must not be null")) + } + } +} + +// BoxRet + +unsafe impl pgrx::callconv::BoxRet for Rabitq4Output { + unsafe fn box_into<'fcx>( + self, + fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, + ) -> pgrx::datum::Datum<'fcx> { + match self.into_datum() { + Some(datum) => unsafe { fcinfo.return_raw_datum(datum) }, + None => fcinfo.return_null(), + } + } +} diff --git a/src/datatype/memory_rabitq8.rs b/src/datatype/memory_rabitq8.rs index 67970a59..a6d04b13 100644 --- a/src/datatype/memory_rabitq8.rs +++ b/src/datatype/memory_rabitq8.rs @@ -23,7 +23,7 @@ use vector::rabitq8::Rabitq8Borrowed; #[repr(C)] struct Rabitq8Header { varlena: u32, - dims: u16, + dim: u16, unused: u16, sum_of_x2: f32, norm_of_lattice: f32, @@ -33,7 +33,10 @@ struct Rabitq8Header { } impl Rabitq8Header { - fn size_of(len: usize) -> usize { + fn size_of_by_dim(dim: usize) -> usize { + Self::size_of_by_len(dim.div_ceil(1)) + } + fn size_of_by_len(len: usize) -> usize { if len > 65535 { panic!("vector is too large"); } @@ -43,13 +46,14 @@ impl Rabitq8Header { unsafe { let this = this.as_ptr(); Rabitq8Borrowed::new( + (&raw const (*this).dim).read() as u32, (&raw const (*this).sum_of_x2).read(), (&raw const (*this).norm_of_lattice).read(), (&raw const (*this).sum_of_code).read(), (&raw const (*this).sum_of_abs_x).read(), std::slice::from_raw_parts( (&raw const (*this).elements).cast(), - (&raw const (*this).dims).read() as usize, + (&raw const (*this).dim).read().div_ceil(1) as usize, ), ) } @@ -69,8 +73,8 @@ impl Rabitq8Input<'_> { let size = varlena as usize; #[cfg(target_endian = "little")] let size = varlena as usize >> 2; - let dims = q.byte_add(4).cast::().read(); - assert_eq!(Rabitq8Header::size_of(dims as _), size); + let dim = q.byte_add(4).cast::().read(); + assert_eq!(Rabitq8Header::size_of_by_dim(dim as _), size); let unused = q.byte_add(6).cast::().read(); assert_eq!(unused, 0); } @@ -104,8 +108,8 @@ impl Rabitq8Output { let size = varlena as usize; #[cfg(target_endian = "little")] let size = varlena as usize >> 2; - let dims = q.byte_add(4).cast::().read(); - assert_eq!(Rabitq8Header::size_of(dims as _), size); + let dim = q.byte_add(4).cast::().read(); + assert_eq!(Rabitq8Header::size_of_by_dim(dim as _), size); let unused = q.byte_add(6).cast::().read(); assert_eq!(unused, 0); } @@ -113,8 +117,8 @@ impl Rabitq8Output { } pub fn new(vector: Rabitq8Borrowed<'_>) -> Self { unsafe { - let code = vector.code(); - let size = Rabitq8Header::size_of(code.len()); + let packed_code = vector.packed_code(); + let size = Rabitq8Header::size_of_by_len(packed_code.len()); let ptr = pgrx::pg_sys::palloc0(size) as *mut Rabitq8Header; // SET_VARSIZE_4B @@ -122,16 +126,16 @@ impl Rabitq8Output { (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); #[cfg(target_endian = "little")] (&raw mut (*ptr).varlena).write((size << 2) as u32); - (&raw mut (*ptr).dims).write(vector.dims() as _); + (&raw mut (*ptr).dim).write(vector.dim() as _); (&raw mut (*ptr).unused).write(0); (&raw mut (*ptr).sum_of_x2).write(vector.sum_of_x2()); (&raw mut (*ptr).norm_of_lattice).write(vector.norm_of_lattice()); (&raw mut (*ptr).sum_of_code).write(vector.sum_of_code()); (&raw mut (*ptr).sum_of_abs_x).write(vector.sum_of_abs_x()); std::ptr::copy_nonoverlapping( - code.as_ptr(), + packed_code.as_ptr(), (&raw mut (*ptr).elements).cast(), - code.len(), + packed_code.len(), ); Self(NonNull::new(ptr).unwrap()) } diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs index 1de17baa..bb5cdd92 100644 --- a/src/datatype/memory_vector.rs +++ b/src/datatype/memory_vector.rs @@ -23,7 +23,7 @@ use vector::vect::VectBorrowed; #[repr(C)] struct VectorHeader { varlena: u32, - dims: u16, + dim: u16, unused: u16, elements: [f32; 0], } @@ -40,7 +40,7 @@ impl VectorHeader { let this = this.as_ptr(); VectBorrowed::new(std::slice::from_raw_parts( (&raw const (*this).elements).cast(), - (&raw const (*this).dims).read() as usize, + (&raw const (*this).dim).read() as usize, )) } } @@ -59,8 +59,8 @@ impl VectorInput<'_> { let size = varlena as usize; #[cfg(target_endian = "little")] let size = varlena as usize >> 2; - let dims = q.byte_add(4).cast::().read(); - assert_eq!(VectorHeader::size_of(dims as _), size); + let dim = q.byte_add(4).cast::().read(); + assert_eq!(VectorHeader::size_of(dim as _), size); let unused = q.byte_add(6).cast::().read(); assert_eq!(unused, 0); } @@ -94,8 +94,8 @@ impl VectorOutput { let size = varlena as usize; #[cfg(target_endian = "little")] let size = varlena as usize >> 2; - let dims = q.byte_add(4).cast::().read(); - assert_eq!(VectorHeader::size_of(dims as _), size); + let dim = q.byte_add(4).cast::().read(); + assert_eq!(VectorHeader::size_of(dim as _), size); let unused = q.byte_add(6).cast::().read(); assert_eq!(unused, 0); } @@ -113,7 +113,7 @@ impl VectorOutput { (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); #[cfg(target_endian = "little")] (&raw mut (*ptr).varlena).write((size << 2) as u32); - (&raw mut (*ptr).dims).write(vector.dims() as _); + (&raw mut (*ptr).dim).write(vector.dim() as _); (&raw mut (*ptr).unused).write(0); std::ptr::copy_nonoverlapping( slice.as_ptr(), diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 2f4f6548..8fcb90fa 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -12,14 +12,20 @@ // // Copyright (c) 2025 TensorChord Inc. +mod binary_rabitq4; mod binary_rabitq8; +mod functions_rabitq4; mod functions_rabitq8; pub mod memory_halfvec; +pub mod memory_rabitq4; pub mod memory_rabitq8; pub mod memory_vector; mod operators_halfvec; +mod operators_rabitq4; mod operators_rabitq8; mod operators_vector; +mod text_rabitq4; mod text_rabitq8; pub mod typmod; +mod typmod_rabitq4; mod typmod_rabitq8; diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index 41ff8568..96bc5ce0 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -35,7 +35,7 @@ fn _vchord_halfvec_sphere_l2_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); - if lhs.dims() != center.dims() { + if lhs.dim() != center.dim() { pgrx::error!("dimension is not matched"); } let d = VectBorrowed::operator_l2s(lhs, center).to_f32().sqrt(); @@ -59,7 +59,7 @@ fn _vchord_halfvec_sphere_ip_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); - if lhs.dims() != center.dims() { + if lhs.dim() != center.dim() { pgrx::error!("dimension is not matched"); } let d = VectBorrowed::operator_dot(lhs, center).to_f32(); @@ -83,7 +83,7 @@ fn _vchord_halfvec_sphere_cosine_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); - if lhs.dims() != center.dims() { + if lhs.dim() != center.dim() { pgrx::error!("dimension is not matched"); } let d = VectBorrowed::operator_cos(lhs, center).to_f32(); diff --git a/src/datatype/operators_rabitq4.rs b/src/datatype/operators_rabitq4.rs new file mode 100644 index 00000000..3c74cab2 --- /dev/null +++ b/src/datatype/operators_rabitq4.rs @@ -0,0 +1,139 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use pgrx::datum::Array; +use std::num::NonZero; +use vector::VectorBorrowed; +use vector::rabitq4::Rabitq4Borrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_operator_l2(lhs: Rabitq4Input<'_>, rhs: Rabitq4Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + Rabitq4Borrowed::operator_l2s(lhs, rhs).to_f32().sqrt() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_operator_ip(lhs: Rabitq4Input<'_>, rhs: Rabitq4Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + Rabitq4Borrowed::operator_dot(lhs, rhs).to_f32() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_operator_cosine(lhs: Rabitq4Input<'_>, rhs: Rabitq4Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + Rabitq4Borrowed::operator_cos(lhs, rhs).to_f32() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_sphere_l2_in( + lhs: Rabitq4Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq4"), +) -> bool { + let center: Rabitq4Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = Rabitq4Borrowed::operator_l2s(lhs, center).to_f32().sqrt(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_sphere_ip_in( + lhs: Rabitq4Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq4"), +) -> bool { + let center: Rabitq4Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = Rabitq4Borrowed::operator_dot(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_sphere_cosine_in( + lhs: Rabitq4Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq4"), +) -> bool { + let center: Rabitq4Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = Rabitq4Borrowed::operator_cos(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq4_operator_maxsim( + lhs: Array<'_, Rabitq4Input<'_>>, + rhs: Array<'_, Rabitq4Input<'_>>, +) -> f32 { + let mut maxsim = 0.0f32; + for rhs in rhs.iter().flatten() { + let mut d = f32::INFINITY; + for lhs in lhs.iter().flatten() { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + d = d.min(Rabitq4Borrowed::operator_dot(lhs, rhs).to_f32()); + } + maxsim += d; + } + maxsim +} diff --git a/src/datatype/operators_rabitq8.rs b/src/datatype/operators_rabitq8.rs index 505e06e1..7eaf0226 100644 --- a/src/datatype/operators_rabitq8.rs +++ b/src/datatype/operators_rabitq8.rs @@ -22,7 +22,7 @@ use vector::rabitq8::Rabitq8Borrowed; fn _vchord_rabitq8_operator_l2(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> f32 { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); - if lhs.dims() != rhs.dims() { + if lhs.dim() != rhs.dim() { pgrx::error!("dimension is not matched"); } Rabitq8Borrowed::operator_l2s(lhs, rhs).to_f32().sqrt() @@ -32,7 +32,7 @@ fn _vchord_rabitq8_operator_l2(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> fn _vchord_rabitq8_operator_ip(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> f32 { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); - if lhs.dims() != rhs.dims() { + if lhs.dim() != rhs.dim() { pgrx::error!("dimension is not matched"); } Rabitq8Borrowed::operator_dot(lhs, rhs).to_f32() @@ -42,7 +42,7 @@ fn _vchord_rabitq8_operator_ip(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> fn _vchord_rabitq8_operator_cosine(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> f32 { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); - if lhs.dims() != rhs.dims() { + if lhs.dim() != rhs.dim() { pgrx::error!("dimension is not matched"); } Rabitq8Borrowed::operator_cos(lhs, rhs).to_f32() @@ -65,7 +65,7 @@ fn _vchord_rabitq8_sphere_l2_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); - if lhs.dims() != center.dims() { + if lhs.dim() != center.dim() { pgrx::error!("dimension is not matched"); } let d = Rabitq8Borrowed::operator_l2s(lhs, center).to_f32().sqrt(); @@ -89,7 +89,7 @@ fn _vchord_rabitq8_sphere_ip_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); - if lhs.dims() != center.dims() { + if lhs.dim() != center.dim() { pgrx::error!("dimension is not matched"); } let d = Rabitq8Borrowed::operator_dot(lhs, center).to_f32(); @@ -113,7 +113,7 @@ fn _vchord_rabitq8_sphere_cosine_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); - if lhs.dims() != center.dims() { + if lhs.dim() != center.dim() { pgrx::error!("dimension is not matched"); } let d = Rabitq8Borrowed::operator_cos(lhs, center).to_f32(); diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index c648bfa7..d7f8137d 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -35,7 +35,7 @@ fn _vchord_vector_sphere_l2_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); - if lhs.dims() != center.dims() { + if lhs.dim() != center.dim() { pgrx::error!("dimension is not matched"); } let d = VectBorrowed::operator_l2s(lhs, center).to_f32().sqrt(); @@ -59,7 +59,7 @@ fn _vchord_vector_sphere_ip_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); - if lhs.dims() != center.dims() { + if lhs.dim() != center.dim() { pgrx::error!("dimension is not matched"); } let d = VectBorrowed::operator_dot(lhs, center).to_f32(); @@ -83,7 +83,7 @@ fn _vchord_vector_sphere_cosine_in( }; let lhs = lhs.as_borrowed(); let center = center.as_borrowed(); - if lhs.dims() != center.dims() { + if lhs.dim() != center.dim() { pgrx::error!("dimension is not matched"); } let d = VectBorrowed::operator_cos(lhs, center).to_f32(); diff --git a/src/datatype/text_rabitq4.rs b/src/datatype/text_rabitq4.rs new file mode 100644 index 00000000..fa6af638 --- /dev/null +++ b/src/datatype/text_rabitq4.rs @@ -0,0 +1,160 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use pgrx::pg_sys::Oid; +use std::ffi::{CStr, CString}; +use vector::rabitq4::Rabitq4Borrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_in(input: &CStr, oid: Oid, typmod: i32) -> Rabitq4Output { + let _ = (oid, typmod); + let mut input = input.to_bytes().iter(); + let mut p0 = Vec::::new(); + let mut p1 = Vec::::new(); + { + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + match c { + b' ' => (), + b'(' => break, + _ => pgrx::error!("incorrect vector"), + } + } + } + { + let mut s = Option::::None; + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + s = match (s, c) { + (s, b' ') => s, + (None, c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + Some(String::from(c as char)) + } + (Some(s), c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + let mut x = s; + x.push(c as char); + Some(x) + } + (Some(s), b',') => { + p0.push(s.parse().expect("failed to parse number")); + None + } + (None, b',') => { + pgrx::error!("incorrect vector") + } + (Some(s), b')') => { + p0.push(s.parse().expect("failed to parse number")); + break; + } + (None, b')') => break, + _ => pgrx::error!("incorrect vector"), + }; + } + } + { + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + match c { + b' ' => (), + b'[' => break, + _ => pgrx::error!("incorrect vector"), + } + } + } + { + let mut s = Option::::None; + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + s = match (s, c) { + (s, b' ') => s, + (None, c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + Some(String::from(c as char)) + } + (Some(s), c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + let mut x = s; + x.push(c as char); + Some(x) + } + (Some(s), b',') => { + p1.push(s.parse().expect("failed to parse number")); + None + } + (None, b',') => { + pgrx::error!("incorrect vector") + } + (Some(s), b']') => { + p1.push(s.parse().expect("failed to parse number")); + break; + } + (None, b']') => break, + _ => pgrx::error!("incorrect vector"), + }; + } + } + if p0.len() != 4 { + pgrx::error!("incorrect vector"); + } + if p1.is_empty() { + pgrx::error!("vector must have at least 1 dimension"); + } + let sum_of_x2 = p0[0]; + let norm_of_lattice = p0[1]; + let sum_of_code = p0[2]; + let sum_of_abs_x = p0[3]; + let unpacked_code = p1; + let packed_code = rabitq::halfbyte::pack_code(&unpacked_code); + if let Some(x) = Rabitq4Borrowed::new_checked( + unpacked_code.len() as _, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + &packed_code, + ) { + Rabitq4Output::new(x) + } else { + pgrx::error!("incorrect vector"); + } +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_out(vector: Rabitq4Input<'_>) -> CString { + let vector = vector.as_borrowed(); + let mut buffer = String::new(); + buffer.push('('); + buffer.push_str(format!("{}", vector.sum_of_x2()).as_str()); + buffer.push_str(format!(", {}", vector.norm_of_lattice()).as_str()); + buffer.push_str(format!(", {}", vector.sum_of_code()).as_str()); + buffer.push_str(format!(", {}", vector.sum_of_abs_x()).as_str()); + buffer.push(')'); + buffer.push('['); + let mut unpacked_code = vector.unpacked_code(); + if let Some(x) = unpacked_code.next() { + buffer.push_str(format!("{x}").as_str()); + } + for x in unpacked_code { + buffer.push_str(format!(", {x}").as_str()); + } + buffer.push(']'); + CString::new(buffer).unwrap() +} diff --git a/src/datatype/text_rabitq8.rs b/src/datatype/text_rabitq8.rs index e81b084b..52df1efb 100644 --- a/src/datatype/text_rabitq8.rs +++ b/src/datatype/text_rabitq8.rs @@ -121,10 +121,16 @@ fn _vchord_rabitq8_in(input: &CStr, oid: Oid, typmod: i32) -> Rabitq8Output { let norm_of_lattice = p0[1]; let sum_of_code = p0[2]; let sum_of_abs_x = p0[3]; - let code = p1; - if let Some(x) = - Rabitq8Borrowed::new_checked(sum_of_x2, norm_of_lattice, sum_of_code, sum_of_abs_x, &code) - { + let unpacked_code = p1; + let packed_code = rabitq::byte::pack_code(&unpacked_code); + if let Some(x) = Rabitq8Borrowed::new_checked( + unpacked_code.len() as _, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + &packed_code, + ) { Rabitq8Output::new(x) } else { pgrx::error!("incorrect vector"); @@ -142,10 +148,11 @@ fn _vchord_rabitq8_out(vector: Rabitq8Input<'_>) -> CString { buffer.push_str(format!(", {}", vector.sum_of_abs_x()).as_str()); buffer.push(')'); buffer.push('['); - if let Some(&x) = vector.code().first() { + let mut unpacked_code = vector.unpacked_code(); + if let Some(x) = unpacked_code.next() { buffer.push_str(format!("{x}").as_str()); } - for &x in vector.code().iter().skip(1) { + for x in unpacked_code { buffer.push_str(format!(", {x}").as_str()); } buffer.push(']'); diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index 46fd7fe7..da6a6bf3 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -31,11 +31,11 @@ impl Typmod { None } } - pub fn dims(self) -> Option> { + pub fn dim(self) -> Option> { use Typmod::*; match self { Any => None, - Dims(dims) => Some(dims), + Dims(dim) => Some(dim), } } } diff --git a/src/datatype/typmod_rabitq4.rs b/src/datatype/typmod_rabitq4.rs new file mode 100644 index 00000000..59970207 --- /dev/null +++ b/src/datatype/typmod_rabitq4.rs @@ -0,0 +1,37 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +use std::ffi::CStr; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_typmod_in(list: pgrx::datum::Array<&CStr>) -> i32 { + if list.is_empty() { + -1 + } else if list.len() == 1 { + let s = list.get(0).unwrap().unwrap().to_str().unwrap(); + if let Ok(d) = s.parse::() { + if d < 1 { + pgrx::error!("dimensions for type rabitq4 must be at least 1"); + } + if d > 65535 { + pgrx::error!("dimensions for type rabitq4 cannot exceed 65535"); + } + d + } else { + pgrx::error!("invalid type modifier") + } + } else { + pgrx::error!("invalid type modifier") + } +} diff --git a/src/index/opclass.rs b/src/index/opclass.rs index 6c8ddd45..607c163f 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -57,6 +57,21 @@ fn _vchordg_support_rabitq8_ip_ops() -> String { "vchordg_rabitq8_ip_ops".to_string() } +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq4_l2_ops() -> String { + "vchordg_rabitq4_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq4_cosine_ops() -> String { + "vchordg_rabitq4_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq4_ip_ops() -> String { + "vchordg_rabitq4_ip_ops".to_string() +} + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_vector_l2_ops() -> String { "vchordrq_vector_l2_ops".to_string() @@ -102,6 +117,21 @@ fn _vchordrq_support_rabitq8_cosine_ops() -> String { "vchordrq_rabitq8_cosine_ops".to_string() } +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq4_l2_ops() -> String { + "vchordrq_rabitq4_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq4_ip_ops() -> String { + "vchordrq_rabitq4_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq4_cosine_ops() -> String { + "vchordrq_rabitq4_cosine_ops".to_string() +} + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordrq_support_vector_maxsim_ops() -> String { "vchordrq_vector_maxsim_ops".to_string() @@ -117,6 +147,11 @@ fn _vchordrq_support_rabitq8_maxsim_ops() -> String { "vchordrq_rabitq8_maxsim_ops".to_string() } +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq4_maxsim_ops() -> String { + "vchordrq_rabitq4_maxsim_ops".to_string() +} + pub struct Sphere { pub center: T, pub radius: f32, diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index 2faa02e6..78609109 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -583,10 +583,10 @@ unsafe fn options( if atts.len() != 1 { pgrx::error!("multicolumn index is not supported"); } - // get dims + // get dim let typmod = Typmod::new(atts[0].atttypmod).unwrap(); - let dims = if let Some(dims) = typmod.dims() { - dims.get() + let dim = if let Some(dim) = typmod.dim() { + dim.get() } else { pgrx::error!( "Dimensions type modifier of a vector column is needed for building the index." @@ -595,7 +595,7 @@ unsafe fn options( // get v, d let opfamily = unsafe { opfamily(index_relation) }; let vector = VectorOptions { - dims, + dim, v: opfamily.vector_kind(), d: opfamily.distance_kind(), }; diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index 068f7aae..6b2a385a 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -395,7 +395,10 @@ pub unsafe extern "C-unwind" fn amrescan( | Opfamily::HalfvecIp | Opfamily::Rabitq8L2 | Opfamily::Rabitq8Cosine - | Opfamily::Rabitq8Ip => { + | Opfamily::Rabitq8Ip + | Opfamily::Rabitq4L2 + | Opfamily::Rabitq4Cosine + | Opfamily::Rabitq4Ip => { let mut builder = DefaultBuilder::new(opfamily); for i in 0..(*scan).numberOfOrderBys { let data = (*scan).orderByData.add(i as usize); diff --git a/src/index/vchordg/dispatch.rs b/src/index/vchordg/dispatch.rs index 7ddb1793..548a77c6 100644 --- a/src/index/vchordg/dispatch.rs +++ b/src/index/vchordg/dispatch.rs @@ -24,6 +24,7 @@ use std::num::NonZero; use vchordg::operator::Op; use vchordg::types::*; use vector::VectorOwned; +use vector::rabitq4::Rabitq4Owned; use vector::rabitq8::Rabitq8Owned; use vector::vect::{VectBorrowed, VectOwned}; @@ -51,6 +52,12 @@ where (VectorKind::Rabitq8, DistanceKind::Dot) => { vchordg::prewarm::<_, Op>(index) } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordg::prewarm::<_, Op>(index) + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordg::prewarm::<_, Op>(index) + } } } @@ -82,6 +89,12 @@ pub fn bulkdelete( (VectorKind::Rabitq8, DistanceKind::Dot) => { vchordg::bulkdelete::<_, Op>(index, &check, &callback); } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordg::bulkdelete::<_, Op>(index, &check, &callback); + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordg::bulkdelete::<_, Op>(index, &check, &callback); + } } } @@ -109,6 +122,12 @@ where (VectorKind::Rabitq8, DistanceKind::Dot) => { vchordg::maintain::<_, Op>(index, &check); } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordg::maintain::<_, Op>(index, &check); + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordg::maintain::<_, Op>(index, &check); + } } } @@ -136,6 +155,12 @@ where (VectorKind::Rabitq8, DistanceKind::Dot) => { vchordg::build::<_, Op>(vector_options, vchordg_options, index) } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordg::build::<_, Op>(vector_options, vchordg_options, index) + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordg::build::<_, Op>(vector_options, vchordg_options, index) + } } } @@ -218,6 +243,28 @@ where make_vector_plain_prefetcher, ) } + (OwnedVector::Rabitq4(unprojected), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq4); + vchordg::insert::<_, Op>( + index, + unprojected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Rabitq4(unprojected), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq4); + vchordg::insert::<_, Op>( + index, + unprojected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } } } diff --git a/src/index/vchordg/opclass.rs b/src/index/vchordg/opclass.rs index 966c6bea..460a0fda 100644 --- a/src/index/vchordg/opclass.rs +++ b/src/index/vchordg/opclass.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use crate::index::opclass::Sphere; @@ -35,6 +36,9 @@ pub enum Opfamily { Rabitq8L2, Rabitq8Cosine, Rabitq8Ip, + Rabitq4L2, + Rabitq4Cosine, + Rabitq4Ip, } impl Opfamily { @@ -59,6 +63,12 @@ impl Opfamily { (Self::Rabitq8Cosine, _) => unreachable!(), (Self::Rabitq8Ip, B::Rabitq8(x)) => O::Rabitq8(x.own()), (Self::Rabitq8Ip, _) => unreachable!(), + (Self::Rabitq4L2, B::Rabitq4(x)) => O::Rabitq4(x.own()), + (Self::Rabitq4L2, _) => unreachable!(), + (Self::Rabitq4Cosine, B::Rabitq4(x)) => O::Rabitq4(x.function_normalize()), + (Self::Rabitq4Cosine, _) => unreachable!(), + (Self::Rabitq4Ip, B::Rabitq4(x)) => O::Rabitq4(x.own()), + (Self::Rabitq4Ip, _) => unreachable!(), } } pub unsafe fn store(self, datum: Datum) -> Option> { @@ -78,6 +88,10 @@ impl Opfamily { let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; vec![(self.input(BorrowedVector::Rabitq8(vector.as_borrowed())), 0)] } + Self::Rabitq4L2 | Self::Rabitq4Cosine | Self::Rabitq4Ip => { + let vector = unsafe { Rabitq4Input::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Rabitq4(vector.as_borrowed())), 0)] + } }; Some(store) } @@ -101,6 +115,10 @@ impl Opfamily { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) } + Self::Rabitq4L2 | Self::Rabitq4Cosine | Self::Rabitq4Ip => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())) + } }; let radius = tuple.get_by_index::(attno_2).unwrap()?; Some(Sphere { center, radius }) @@ -122,21 +140,37 @@ impl Opfamily { let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) } + Self::Rabitq4L2 | Self::Rabitq4Cosine | Self::Rabitq4Ip => { + let vector = unsafe { Rabitq4Input::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())) + } }; Some(vector) } pub fn output(self, x: Distance) -> f32 { match self { - Self::VectorCosine | Self::HalfvecCosine | Self::Rabitq8Cosine => x.to_f32() * 0.5, - Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 => x.to_f32().sqrt(), - Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip => x.to_f32(), + Self::VectorCosine + | Self::HalfvecCosine + | Self::Rabitq8Cosine + | Self::Rabitq4Cosine => x.to_f32() * 0.5, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 | Self::Rabitq4L2 => { + x.to_f32().sqrt() + } + Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip | Self::Rabitq4Ip => x.to_f32(), } } pub const fn distance_kind(self) -> DistanceKind { match self { - Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 => DistanceKind::L2S, - Self::VectorCosine | Self::HalfvecCosine | Self::Rabitq8Cosine => DistanceKind::L2S, - Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip => DistanceKind::Dot, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 | Self::Rabitq4L2 => { + DistanceKind::L2S + } + Self::VectorCosine + | Self::HalfvecCosine + | Self::Rabitq8Cosine + | Self::Rabitq4Cosine => DistanceKind::L2S, + Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip | Self::Rabitq4Ip => { + DistanceKind::Dot + } } } pub const fn vector_kind(self) -> VectorKind { @@ -144,6 +178,7 @@ impl Opfamily { Self::VectorL2 | Self::VectorCosine | Self::VectorIp => VectorKind::Vecf32, Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => VectorKind::Vecf16, Self::Rabitq8L2 | Self::Rabitq8Cosine | Self::Rabitq8Ip => VectorKind::Rabitq8, + Self::Rabitq4L2 | Self::Rabitq4Cosine | Self::Rabitq4Ip => VectorKind::Rabitq4, } } } @@ -192,6 +227,9 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { "vchordg_rabitq8_l2_ops" => Opfamily::Rabitq8L2, "vchordg_rabitq8_ip_ops" => Opfamily::Rabitq8Ip, "vchordg_rabitq8_cosine_ops" => Opfamily::Rabitq8Cosine, + "vchordg_rabitq4_l2_ops" => Opfamily::Rabitq4L2, + "vchordg_rabitq4_ip_ops" => Opfamily::Rabitq4Ip, + "vchordg_rabitq4_cosine_ops" => Opfamily::Rabitq4Cosine, _ => pgrx::error!("unknown operator class"), }; diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs index 820242ef..31582979 100644 --- a/src/index/vchordg/scanners/default.rs +++ b/src/index/vchordg/scanners/default.rs @@ -28,9 +28,10 @@ use std::num::NonZero; use vchordg::operator::{self}; use vchordg::search; use vchordg::types::{DistanceKind, OwnedVector, VectorKind}; -use vector::VectorOwned; +use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; +use vector::{VectorBorrowed, VectorOwned}; pub struct DefaultBuilder { opfamily: Opfamily, @@ -57,6 +58,9 @@ impl SearchBuilder for DefaultBuilder { | Opfamily::Rabitq8Cosine | Opfamily::Rabitq8L2 | Opfamily::Rabitq8Ip + | Opfamily::Rabitq4Cosine + | Opfamily::Rabitq4L2 + | Opfamily::Rabitq4Ip )); Self { opfamily, @@ -513,11 +517,12 @@ impl SearchBuilder for DefaultBuilder { let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { let vector = vector.as_borrowed(); Rabitq8Borrowed::new( + vector.dim(), vector.sum_of_x2(), vector.norm_of_lattice(), vector.sum_of_code(), vector.sum_of_abs_x(), - bump.alloc_slice(vector.code()), + bump.alloc_slice(vector.packed_code()), ) } else { unreachable!() @@ -611,11 +616,210 @@ impl SearchBuilder for DefaultBuilder { let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { let vector = vector.as_borrowed(); Rabitq8Borrowed::new( + vector.dim(), vector.sum_of_x2(), vector.norm_of_lattice(), vector.sum_of_code(), vector.sum_of_abs_x(), - bump.alloc_slice(vector.code()), + bump.alloc_slice(vector.packed_code()), + ) + } else { + unreachable!() + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + type Op = operator::Op; + let unprojected = if let OwnedVector::Rabitq4(vector) = vector.clone() { + let vector = vector.as_borrowed(); + Rabitq4Borrowed::new( + vector.dim(), + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + bump.alloc_slice(vector.packed_code()), + ) + } else { + unreachable!() + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + type Op = operator::Op; + let unprojected = if let OwnedVector::Rabitq4(vector) = vector.clone() { + let vector = vector.as_borrowed(); + Rabitq4Borrowed::new( + vector.dim(), + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + bump.alloc_slice(vector.packed_code()), ) } else { unreachable!() @@ -726,6 +930,9 @@ impl SearchBuilder for DefaultBuilder { OwnedVector::Rabitq8(v) => { recorder.send(&text::rabitq8_out(v.as_borrowed())); } + OwnedVector::Rabitq4(v) => { + recorder.send(&text::rabitq4_out(v.as_borrowed())); + } } } Box::new(iter.map(move |(distance, pointer)| { diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 5bf29648..4da27b65 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -34,6 +34,7 @@ use std::num::NonZero; use std::ops::Deref; use vchordrq::types::*; use vchordrq::{InsertChooser, MaintainChooser}; +use vector::rabitq4::Rabitq4Owned; use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; @@ -221,6 +222,10 @@ pub unsafe extern "C-unwind" fn ambuild( let errors = "residual_quantization is not supported for rabitq8 type"; pgrx::error!("error while validating options: {errors}"); } + if vector_options.v == VectorKind::Rabitq4 && vchordrq_options.index.residual_quantization { + let errors = "residual_quantization is not supported for rabitq4 type"; + pgrx::error!("error while validating options: {errors}"); + } let opfamily = unsafe { opfamily(index_relation) }; let reporter = PostgresReporter { _phantom: PhantomData, @@ -1230,10 +1235,10 @@ unsafe fn options( if atts.len() != 1 { pgrx::error!("multicolumn index is not supported"); } - // get dims + // get dim let typmod = Typmod::new(atts[0].atttypmod).unwrap(); - let dims = if let Some(dims) = typmod.dims() { - dims.get() + let dim = if let Some(dim) = typmod.dim() { + dim.get() } else { pgrx::error!( "Dimensions type modifier of a vector column is needed for building the index." @@ -1242,7 +1247,7 @@ unsafe fn options( // get v, d let opfamily = unsafe { opfamily(index_relation) }; let vector = VectorOptions { - dims, + dim, v: opfamily.vector_kind(), d: opfamily.distance_kind(), }; @@ -1266,7 +1271,7 @@ fn make_default_build( _default_build: VchordrqDefaultBuildOptions, ) -> Vec> { vec![Structure:: { - centroids: vec![vec![0.0f32; vector_options.dims as usize]], + centroids: vec![vec![0.0f32; vector_options.dim as usize]], children: vec![vec![]], }] } @@ -1281,15 +1286,15 @@ fn make_internal_build( use humansize::{BINARY, format_size}; use std::iter::once; let (reduction, sample_dim) = match internal_build.kmeans_dimension { - None => (None, vector_options.dims as usize), - Some(d) if d < vector_options.dims => (Some(d as usize), d as usize), + None => (None, vector_options.dim as usize), + Some(d) if d < vector_options.dim => (Some(d as usize), d as usize), Some(d) => { pgrx::warning!( "ignoring `kmeans_dimension = {}` because it is less than the vector dimension {}", d, - vector_options.dims + vector_options.dim ); - (None, vector_options.dims as usize) + (None, vector_options.dim as usize) } }; { @@ -1328,9 +1333,10 @@ fn make_internal_build( OwnedVector::Vecf32(x) => VectOwned::normalize(x), OwnedVector::Vecf16(x) => VectOwned::normalize(x), OwnedVector::Rabitq8(x) => Rabitq8Owned::normalize(x), + OwnedVector::Rabitq4(x) => Rabitq4Owned::normalize(x), }; assert_eq!( - vector_options.dims, + vector_options.dim, x.len() as u32, "invalid vector dimensions" ); @@ -1359,7 +1365,7 @@ fn make_internal_build( for w in internal_build.lists.iter().rev().copied().chain(once(1)) { let mut input = if let Some(structure) = result.last() { let mut input = - Square::with_capacity(vector_options.dims as _, structure.centroids.len()); + Square::with_capacity(vector_options.dim as _, structure.centroids.len()); for slice in structure.centroids.iter() { input.push_slice(slice); } @@ -1370,7 +1376,7 @@ fn make_internal_build( unreachable!() }; let num_points = input.len(); - let num_dims = input.d(); + let num_dim = input.d(); let num_lists = w as usize; let num_iterations = internal_build.kmeans_iterations as _; if result.is_empty() { @@ -1382,7 +1388,7 @@ fn make_internal_build( } if num_lists > 1 { pgrx::info!( - "clustering: starting, clustering {num_points} vectors of {num_dims} dimension into {num_lists} clusters, in {num_iterations} iterations" + "clustering: starting, clustering {num_points} vectors of {num_dim} dimension into {num_lists} clusters, in {num_iterations} iterations" ); } let mut f = 'f: { @@ -1390,7 +1396,7 @@ fn make_internal_build( if result.last().is_some() { break 'f k_means::lloyd_k_means( &pool, - num_dims, + num_dim, view, num_lists, [7; 32], @@ -1400,7 +1406,7 @@ fn make_internal_build( match internal_build.kmeans_algorithm { KMeansAlgorithm::Lloyd {} => k_means::lloyd_k_means( &pool, - num_dims, + num_dim, view, num_lists, [7; 32], @@ -1408,7 +1414,7 @@ fn make_internal_build( ), KMeansAlgorithm::Hierarchical {} => k_means::hierarchical_k_means( &pool, - num_dims, + num_dim, view, num_lists, [7; 32], @@ -1443,7 +1449,7 @@ fn make_internal_build( let is_spherical = internal_build.spherical_centroids; let num_threads = { let config = internal_build.build_threads as f64; - let ratio = sample_dim as f64 / vector_options.dims as f64; + let ratio = sample_dim as f64 / vector_options.dim as f64; (config * ratio).ceil().max(1.0).min(config) as usize }; let (tx, rx) = crossbeam_channel::bounded::>(1024); @@ -1452,7 +1458,7 @@ fn make_internal_build( let mut handles = Vec::with_capacity(num_threads); for _ in 0..num_threads { let rx = rx.clone(); - let mut sum = Square::from_zeros(vector_options.dims as _, num_lists); + let mut sum = Square::from_zeros(vector_options.dim as _, num_lists); let mut count = vec![0.0f32; num_lists]; handles.push(scope.spawn(move || { while let Ok(sample) = rx.recv() { @@ -1486,9 +1492,12 @@ fn make_internal_build( OwnedVector::Rabitq8(x) => { Rabitq8Owned::normalize(x) } + OwnedVector::Rabitq4(x) => { + Rabitq4Owned::normalize(x) + } }; assert_eq!( - vector_options.dims, + vector_options.dim, x.len() as u32, "invalid vector dimensions" ); @@ -1510,7 +1519,7 @@ fn make_internal_build( .map(|handle| handle.join().expect("failed to spawn threads")) .collect::>() }); - let mut sum = Square::from_zeros(vector_options.dims as _, num_lists); + let mut sum = Square::from_zeros(vector_options.dim as _, num_lists); let mut count = vec![0.0f32; num_lists]; for (sum_1, count_1) in list { for i in 0..num_lists { @@ -1624,7 +1633,7 @@ fn make_external_build( "external build: there are at least two lines have same id, id = {id}" ); } - if vector_options.dims != vector.as_borrowed().dims() { + if vector_options.dim != vector.as_borrowed().dim() { pgrx::error!("external build: incorrect dimension, id = {id}"); } vectors.insert(id, vector.as_borrowed().slice().to_vec()); @@ -1642,7 +1651,7 @@ fn make_external_build( result.push(Structure { centroids: vec![{ // compute the vector on root, without normalizing it - let mut sum = vec![0.0f32; vector_options.dims as _]; + let mut sum = vec![0.0f32; vector_options.dim as _]; for vector in vectors.values() { f32::vector_add_inplace(&mut sum, vector); } diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 74a7b999..8fa36323 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -211,6 +211,9 @@ pub unsafe extern "C-unwind" fn amcostestimate( | Opfamily::Rabitq8Cosine | Opfamily::Rabitq8Ip | Opfamily::Rabitq8L2 + | Opfamily::Rabitq4Cosine + | Opfamily::Rabitq4Ip + | Opfamily::Rabitq4L2 ) { *index_startup_cost = 0.0; *index_total_cost = 0.0; @@ -244,9 +247,10 @@ pub unsafe extern "C-unwind" fn amcostestimate( let page_count = { let mut pages = 0_f64; pages += 1.0; - pages += node_count * cost.dims as f64 / 60000.0; + pages += node_count * cost.dim as f64 / 60000.0; pages += probes.iter().sum::() as f64 * { - let x = opfamily.vector_kind().element_size() * cost.dims; + let x = (opfamily.vector_kind().number_of_bits_of_an_elements() * cost.dim) + .div_ceil(8); x.div_ceil(3840 * x.div_ceil(5120).min(2)) as f64 }; pages += cost.cells[0] as f64; @@ -476,7 +480,10 @@ pub unsafe extern "C-unwind" fn amrescan( | Opfamily::HalfvecCosine | Opfamily::Rabitq8L2 | Opfamily::Rabitq8Ip - | Opfamily::Rabitq8Cosine => { + | Opfamily::Rabitq8Cosine + | Opfamily::Rabitq4L2 + | Opfamily::Rabitq4Ip + | Opfamily::Rabitq4Cosine => { let mut builder = DefaultBuilder::new(opfamily); for i in 0..(*scan).numberOfOrderBys { let data = (*scan).orderByData.add(i as usize); @@ -496,7 +503,10 @@ pub unsafe extern "C-unwind" fn amrescan( builder.build(index, options, fetcher, bump, recorder) })) } - Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim | Opfamily::Rabitq8Maxsim => { + Opfamily::VectorMaxsim + | Opfamily::HalfvecMaxsim + | Opfamily::Rabitq8Maxsim + | Opfamily::Rabitq4Maxsim => { let mut builder = MaxsimBuilder::new(opfamily); for i in 0..(*scan).numberOfOrderBys { let data = (*scan).orderByData.add(i as usize); diff --git a/src/index/vchordrq/build.rs b/src/index/vchordrq/build.rs index 81370b63..4e91a340 100644 --- a/src/index/vchordrq/build.rs +++ b/src/index/vchordrq/build.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use simd::{Floating, f16}; +use vector::rabitq4::Rabitq4Owned; use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; use vector::{VectorBorrowed, VectorOwned}; @@ -48,8 +49,8 @@ impl Normalize for Rabitq8Owned { fn normalize(vector: Self) -> Normalized { let vector = vector.as_borrowed(); let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); - let mut result = Vec::with_capacity(vector.dims() as _); - for c in vector.code().iter().copied() { + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { let base = -0.5 * ((1 << 8) - 1) as f32; result.push((base + c as f32) * scale); } @@ -57,14 +58,45 @@ impl Normalize for Rabitq8Owned { result } - fn denormalize(x: Normalized) -> Self { - let code = rabitq::byte::ugly_code(x.as_slice()); + fn denormalize(vector: Normalized) -> Self { + let dim = vector.len() as u32; + let (metadata, elements) = rabitq::byte::ugly_code(&vector); + let elements = rabitq::byte::pack_code(&elements); Rabitq8Owned::new( - code.0.dis_u_2, - code.0.norm_of_lattice, - code.0.sum_of_code, - f32::reduce_sum_of_abs_x(&x), - code.1, + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + elements, + ) + } +} + +impl Normalize for Rabitq4Owned { + fn normalize(vector: Self) -> Normalized { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + result + } + + fn denormalize(vector: Normalized) -> Self { + let dim = vector.len() as u32; + let (metadata, elements) = rabitq::halfbyte::ugly_code(&vector); + let elements = rabitq::halfbyte::pack_code(&elements); + Rabitq4Owned::new( + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + elements, ) } } diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs index 1a5569bd..c4a80fd7 100644 --- a/src/index/vchordrq/dispatch.rs +++ b/src/index/vchordrq/dispatch.rs @@ -28,6 +28,7 @@ use vchordrq::operator::Op; use vchordrq::types::*; use vchordrq::{FastHeap, InsertChooser, MaintainChooser}; use vector::VectorOwned; +use vector::rabitq4::Rabitq4Owned; use vector::rabitq8::Rabitq8Owned; use vector::vect::{VectBorrowed, VectOwned}; @@ -56,6 +57,12 @@ where (VectorKind::Rabitq8, DistanceKind::Dot) => { vchordrq::prewarm::<_, Op>(index, height, make_h0_plain_prefetcher) } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordrq::prewarm::<_, Op>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordrq::prewarm::<_, Op>(index, height, make_h0_plain_prefetcher) + } } } @@ -93,6 +100,14 @@ pub fn bulkdelete( vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + } } } @@ -151,6 +166,18 @@ pub fn maintain( chooser, check, ), + (VectorKind::Rabitq4, DistanceKind::L2S) => vchordrq::maintain::<_, Op>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ), + (VectorKind::Rabitq4, DistanceKind::Dot) => vchordrq::maintain::<_, Op>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ), }; pgrx::info!( "maintain: number_of_formerly_allocated_pages = {}", @@ -212,6 +239,18 @@ pub fn build( index, map_structures(structures, Normalize::denormalize), ), + (VectorKind::Rabitq4, DistanceKind::L2S) => vchordrq::build::<_, Op>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + (VectorKind::Rabitq4, DistanceKind::Dot) => vchordrq::build::<_, Op>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), } } @@ -348,6 +387,44 @@ pub fn insert( skip_freespaces, ) } + (OwnedVector::Rabitq4(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq4); + let key = vchordrq::insert_vector::<_, Op>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op>( + index, + payload, + vector.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Rabitq4(vector), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq4); + let key = vchordrq::insert_vector::<_, Op>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op>( + index, + payload, + vector.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } } } diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs index 13bbd4b9..f3d1a4d5 100644 --- a/src/index/vchordrq/opclass.rs +++ b/src/index/vchordrq/opclass.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use crate::index::opclass::Sphere; @@ -35,9 +36,13 @@ pub enum Opfamily { Rabitq8L2, Rabitq8Ip, Rabitq8Cosine, + Rabitq4L2, + Rabitq4Ip, + Rabitq4Cosine, VectorMaxsim, HalfvecMaxsim, Rabitq8Maxsim, + Rabitq4Maxsim, } impl Opfamily { @@ -56,6 +61,10 @@ impl Opfamily { (B::Rabitq8(x), Self::Rabitq8Ip | Self::Rabitq8Maxsim) => O::Rabitq8(x.own()), (B::Rabitq8(x), Self::Rabitq8Cosine) => O::Rabitq8(x.function_normalize()), (B::Rabitq8(_), _) => unreachable!(), + (B::Rabitq4(x), Self::Rabitq4L2) => O::Rabitq4(x.own()), + (B::Rabitq4(x), Self::Rabitq4Ip | Self::Rabitq4Maxsim) => O::Rabitq4(x.own()), + (B::Rabitq4(x), Self::Rabitq4Cosine) => O::Rabitq4(x.function_normalize()), + (B::Rabitq4(_), _) => unreachable!(), } } pub unsafe fn store(self, datum: Datum) -> Option> { @@ -75,6 +84,10 @@ impl Opfamily { let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; vec![(self.input(BorrowedVector::Rabitq8(vector.as_borrowed())), 0)] } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine => { + let vector = unsafe { Rabitq4Input::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Rabitq4(vector.as_borrowed())), 0)] + } Self::VectorMaxsim => { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; @@ -113,6 +126,19 @@ impl Opfamily { } result } + Self::Rabitq4Maxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())), + i as u16, + )); + } + result + } }; Some(store) } @@ -136,6 +162,10 @@ impl Opfamily { let vector = tuple.get_by_index::(attno_1).unwrap()?; self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine | Self::Rabitq4Maxsim => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())) + } }; let radius = tuple.get_by_index::(attno_2).unwrap()?; Some(Sphere { center, radius }) @@ -157,6 +187,10 @@ impl Opfamily { let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine | Self::Rabitq4Maxsim => { + let vector = unsafe { Rabitq4Input::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())) + } }; Some(vector) } @@ -194,33 +228,55 @@ impl Opfamily { } result } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine | Self::Rabitq4Maxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Rabitq4(vector.as_borrowed()))); + } + result + } }; Some(vectors) } pub fn output(self, x: Distance) -> f32 { match self { - Self::VectorCosine | Self::HalfvecCosine | Self::Rabitq8Cosine => x.to_f32() + 1.0f32, - Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 => x.to_f32().sqrt(), + Self::VectorCosine + | Self::HalfvecCosine + | Self::Rabitq8Cosine + | Self::Rabitq4Cosine => x.to_f32() + 1.0f32, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 | Self::Rabitq4L2 => { + x.to_f32().sqrt() + } Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip + | Self::Rabitq4Ip | Self::VectorMaxsim | Self::HalfvecMaxsim - | Self::Rabitq8Maxsim => x.to_f32(), + | Self::Rabitq8Maxsim + | Self::Rabitq4Maxsim => x.to_f32(), } } pub const fn distance_kind(self) -> DistanceKind { match self { - Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 => DistanceKind::L2S, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 | Self::Rabitq4L2 => { + DistanceKind::L2S + } Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip + | Self::Rabitq4Ip | Self::VectorCosine | Self::HalfvecCosine | Self::Rabitq8Cosine + | Self::Rabitq4Cosine | Self::VectorMaxsim | Self::HalfvecMaxsim - | Self::Rabitq8Maxsim => DistanceKind::Dot, + | Self::Rabitq8Maxsim + | Self::Rabitq4Maxsim => DistanceKind::Dot, } } pub const fn vector_kind(self) -> VectorKind { @@ -234,6 +290,9 @@ impl Opfamily { Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine | Self::Rabitq8Maxsim => { VectorKind::Rabitq8 } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine | Self::Rabitq4Maxsim => { + VectorKind::Rabitq4 + } } } } @@ -282,9 +341,13 @@ pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { "vchordrq_rabitq8_l2_ops" => Opfamily::Rabitq8L2, "vchordrq_rabitq8_ip_ops" => Opfamily::Rabitq8Ip, "vchordrq_rabitq8_cosine_ops" => Opfamily::Rabitq8Cosine, + "vchordrq_rabitq4_l2_ops" => Opfamily::Rabitq4L2, + "vchordrq_rabitq4_ip_ops" => Opfamily::Rabitq4Ip, + "vchordrq_rabitq4_cosine_ops" => Opfamily::Rabitq4Cosine, "vchordrq_vector_maxsim_ops" => Opfamily::VectorMaxsim, "vchordrq_halfvec_maxsim_ops" => Opfamily::HalfvecMaxsim, "vchordrq_rabitq8_maxsim_ops" => Opfamily::Rabitq8Maxsim, + "vchordrq_rabitq4_maxsim_ops" => Opfamily::Rabitq4Maxsim, _ => pgrx::error!("unknown operator class"), }; diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index 4897d979..ba94953b 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -32,6 +32,7 @@ use std::num::NonZero; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; use vchordrq::{RerankMethod, default_search, how, rerank_heap, rerank_index}; use vector::VectorOwned; +use vector::rabitq4::Rabitq4Owned; use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; @@ -60,6 +61,9 @@ impl SearchBuilder for DefaultBuilder { | Opfamily::Rabitq8Cosine | Opfamily::Rabitq8Ip | Opfamily::Rabitq8L2 + | Opfamily::Rabitq4Cosine + | Opfamily::Rabitq4Ip + | Opfamily::Rabitq4L2 )); Self { opfamily, @@ -953,6 +957,280 @@ impl SearchBuilder for DefaultBuilder { } } } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + type Op = vchordrq::operator::Op; + let unprojected = if let OwnedVector::Rabitq4(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq4(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq4(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + type Op = vchordrq::operator::Op; + let unprojected = if let OwnedVector::Rabitq4(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq4(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq4(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } }; let iter = if let Some(threshold) = threshold { Box::new(iter.take_while(move |(x, _)| *x < threshold)) @@ -975,6 +1253,9 @@ impl SearchBuilder for DefaultBuilder { OwnedVector::Rabitq8(v) => { recorder.send(&text::rabitq8_out(v.as_borrowed())); } + OwnedVector::Rabitq4(v) => { + recorder.send(&text::rabitq4_out(v.as_borrowed())); + } } } Box::new(iter.map(move |(distance, pointer)| { diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 8f83fa8b..b82e302b 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -34,6 +34,7 @@ use std::num::NonZero; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; use vchordrq::{RerankMethod, how, maxsim_search, rerank_index}; use vector::VectorOwned; +use vector::rabitq4::Rabitq4Owned; use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; @@ -52,7 +53,10 @@ impl SearchBuilder for MaxsimBuilder { fn new(opfamily: Opfamily) -> Self { assert!(matches!( opfamily, - Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim | Opfamily::Rabitq8Maxsim + Opfamily::VectorMaxsim + | Opfamily::HalfvecMaxsim + | Opfamily::Rabitq8Maxsim + | Opfamily::Rabitq4Maxsim )); Self { opfamily, @@ -549,6 +553,146 @@ impl SearchBuilder for MaxsimBuilder { (accu_set, rough_set, estimation_by_threshold) })) } + VectorKind::Rabitq4 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq4(vector) = vector { + vector + } else { + unreachable!() + } + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = + rerank_index::(unprojected[i].clone(), prefetcher); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); + } + (accu_set, rough_set, estimation_by_threshold) + })) + } }; let mut updates = Vec::new(); let mut estimations = Vec::new(); diff --git a/src/recorder/text.rs b/src/recorder/text.rs index 1b1a2100..40a8dd8a 100644 --- a/src/recorder/text.rs +++ b/src/recorder/text.rs @@ -13,6 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use simd::f16; +use vector::rabitq4::Rabitq4Borrowed; use vector::rabitq8::Rabitq8Borrowed; use vector::vect::VectBorrowed; @@ -52,7 +53,29 @@ pub fn rabitq8_out(vector: Rabitq8Borrowed<'_>) -> String { result.push_str(&vector.sum_of_abs_x().to_string()); result.push(')'); result.push('['); - for x in vector.code() { + for x in vector.unpacked_code() { + if !result.ends_with('[') { + result.push(','); + } + result.push_str(&x.to_string()); + } + result.push(']'); + result +} + +pub fn rabitq4_out(vector: Rabitq4Borrowed<'_>) -> String { + let mut result = String::new(); + result.push('('); + result.push_str(&vector.sum_of_x2().to_string()); + result.push(','); + result.push_str(&vector.norm_of_lattice().to_string()); + result.push(','); + result.push_str(&vector.sum_of_code().to_string()); + result.push(','); + result.push_str(&vector.sum_of_abs_x().to_string()); + result.push(')'); + result.push('['); + for x in vector.unpacked_code() { if !result.ends_with('[') { result.push(','); } diff --git a/src/sql/bootstrap.sql b/src/sql/bootstrap.sql index 743fea00..c9f28449 100644 --- a/src/sql/bootstrap.sql +++ b/src/sql/bootstrap.sql @@ -4,3 +4,5 @@ CREATE TYPE sphere_vector; CREATE TYPE sphere_halfvec; CREATE TYPE rabitq8; CREATE TYPE sphere_rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq4; diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index a7e4d923..4e0847a6 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -9,6 +9,15 @@ CREATE TYPE rabitq8 ( STORAGE = external ); +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + CREATE TYPE sphere_vector AS ( center vector, radius REAL @@ -24,11 +33,19 @@ CREATE TYPE sphere_rabitq8 AS ( radius REAL ); +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + -- List of internal functions CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + -- List of operators CREATE OPERATOR <-> ( @@ -52,6 +69,27 @@ CREATE OPERATOR <=> ( COMMUTATOR = <=> ); +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + CREATE OPERATOR <<->> ( PROCEDURE = _vchord_vector_sphere_l2_in, LEFTARG = vector, @@ -73,6 +111,13 @@ CREATE OPERATOR <<->> ( COMMUTATOR = <<->> ); +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4, + COMMUTATOR = <<->> +); + CREATE OPERATOR <<#>> ( PROCEDURE = _vchord_vector_sphere_ip_in, LEFTARG = vector, @@ -94,6 +139,13 @@ CREATE OPERATOR <<#>> ( COMMUTATOR = <<#>> ); +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4, + COMMUTATOR = <<#>> +); + CREATE OPERATOR <<=>> ( PROCEDURE = _vchord_vector_sphere_cosine_in, LEFTARG = vector, @@ -115,6 +167,13 @@ CREATE OPERATOR <<=>> ( COMMUTATOR = <<=>> ); +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4, + COMMUTATOR = <<=>> +); + CREATE OPERATOR @# ( PROCEDURE = _vchord_vector_operator_maxsim, LEFTARG = vector[], @@ -133,6 +192,12 @@ CREATE OPERATOR @# ( RIGHTARG = rabitq8[] ); +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + -- List of functions CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector @@ -144,12 +209,21 @@ IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; @@ -330,9 +404,13 @@ CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; @@ -342,6 +420,9 @@ CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; -- List of operator classes @@ -399,6 +480,24 @@ CREATE OPERATOR CLASS rabitq8_cosine_ops OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + CREATE OPERATOR CLASS vector_maxsim_ops FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, @@ -414,6 +513,11 @@ CREATE OPERATOR CLASS rabitq8_maxsim_ops OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + CREATE OPERATOR CLASS vector_l2_ops FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, @@ -468,6 +572,24 @@ CREATE OPERATOR CLASS rabitq8_cosine_ops OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); + -- List of views CREATE VIEW vchordrq_sampled_queries AS diff --git a/tests/vchordg/index_vector.slt b/tests/vchordg/index_vector.slt index d6e15ac3..8e3a4382 100644 --- a/tests/vchordg/index_vector.slt +++ b/tests/vchordg/index_vector.slt @@ -1,3 +1,6 @@ +statement ok +SET enable_seqscan TO off; + statement ok CREATE TABLE t (index serial primary key, val vector(64)); @@ -75,7 +78,7 @@ statement ok CREATE INDEX ti ON t USING vchordg ((val::halfvec(64)) halfvec_l2_ops); query I -SELECT index FROM t ORDER BY val <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +SELECT index FROM t ORDER BY val::halfvec <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; ---- 1608 155 @@ -95,7 +98,7 @@ statement ok CREATE INDEX ti ON t USING vchordg ((val::halfvec(64)) halfvec_ip_ops); query I -SELECT index FROM t ORDER BY val <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +SELECT index FROM t ORDER BY val::halfvec <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; ---- 1608 155 @@ -115,130 +118,10 @@ statement ok CREATE INDEX ti ON t USING vchordg ((val::halfvec(64)) halfvec_cosine_ops); query I -SELECT index FROM t ORDER BY val <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; ----- -1608 -155 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_l2_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ----- -155 -1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_ip_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ----- -155 -1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_cosine_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +SELECT index FROM t ORDER BY val::halfvec <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; ---- -155 1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ----- 155 -1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ----- -155 -1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ----- -155 -1608 1643 174 818 diff --git a/tests/vchordg/index_vector_rabitq.slt b/tests/vchordg/index_vector_rabitq.slt new file mode 100644 index 00000000..f33dd8ea --- /dev/null +++ b/tests/vchordg/index_vector_rabitq.slt @@ -0,0 +1,216 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE t (index serial primary key, val vector(64)); + +statement ok +INSERT INTO t (val) +SELECT + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1639 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1639 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1639 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_l2_ops); + +query I +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 +) AS s(index); +---- +6162 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_ip_ops); + +query I +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 +) AS s(index); +---- +6162 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_cosine_ops); + +query I +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 +) AS s(index); +---- +6162 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_l2_ops); + +query I +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 +) AS s(index); +---- +7290 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_ip_ops); + +query I +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 +) AS s(index); +---- +7290 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_cosine_ops); + +query I +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 +) AS s(index); +---- +7290 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/index_multivector.slt b/tests/vchordrq/index_multivector.slt index f1e331e7..044835dc 100644 --- a/tests/vchordrq/index_multivector.slt +++ b/tests/vchordrq/index_multivector.slt @@ -1,12 +1,8 @@ statement ok -CREATE TABLE t (index serial primary key, val vector(64)[2]); +SET enable_seqscan TO off; statement ok -CREATE FUNCTION pg_temp.quantize_to_rabitq8(val vector[]) -RETURNS rabitq8[] AS $$ - SELECT array_agg(public.quantize_to_rabitq8(x)) - FROM unnest(val) AS x; -$$ LANGUAGE sql IMMUTABLE; +CREATE TABLE t (index serial primary key, val vector(64)[2]); statement ok INSERT INTO t (val) @@ -75,30 +71,5 @@ LIMIT 10; statement ok DROP INDEX ti; -statement ok -CREATE INDEX ti ON t USING vchordrq ((pg_temp.quantize_to_rabitq8(val)::rabitq8(64)[2]) rabitq8_maxsim_ops); - -query I -SELECT index FROM t ORDER BY -pg_temp.quantize_to_rabitq8(val)::rabitq8(64)[] @# ARRAY[ - quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector), - quantize_to_rabitq8(array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) -] -LIMIT 10; ----- -1207 -919 -1821 -1639 -174 -79 -1076 -1125 -239 -194 - -statement ok -DROP INDEX ti; - statement ok DROP TABLE t; diff --git a/tests/vchordrq/index_multivector_rabitq.slt b/tests/vchordrq/index_multivector_rabitq.slt new file mode 100644 index 00000000..7756da41 --- /dev/null +++ b/tests/vchordrq/index_multivector_rabitq.slt @@ -0,0 +1,89 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE pg_temp.t (index serial primary key, val vector(64)[2]); + +statement ok +INSERT INTO pg_temp.t (val) +SELECT + ARRAY[ + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector), + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) + ] +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE FUNCTION pg_temp.quantize_to_rabitq8(val vector[]) +RETURNS rabitq8[] AS $$ + SELECT array_agg(public.quantize_to_rabitq8(x)) + FROM unnest(val) AS x; +$$ LANGUAGE sql IMMUTABLE; + +statement ok +CREATE FUNCTION pg_temp.quantize_to_rabitq4(val vector[]) +RETURNS rabitq4[] AS $$ + SELECT array_agg(public.quantize_to_rabitq4(x)) + FROM unnest(val) AS x; +$$ LANGUAGE sql IMMUTABLE; + +statement ok +CREATE INDEX ti ON pg_temp.t USING vchordrq ((pg_temp.quantize_to_rabitq8(val)::rabitq8(64)[2]) rabitq8_maxsim_ops); + +query I +SELECT index FROM pg_temp.t ORDER BY +pg_temp.quantize_to_rabitq8(val)::rabitq8(64)[] @# ARRAY[ + quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector), + quantize_to_rabitq8(array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) +] +LIMIT 10; +---- +1207 +919 +1821 +1639 +174 +79 +1076 +1125 +239 +194 + +statement ok +DROP INDEX pg_temp.ti; + +statement ok +CREATE INDEX ti ON pg_temp.t USING vchordrq ((pg_temp.quantize_to_rabitq4(val)::rabitq4(64)[2]) rabitq4_maxsim_ops); + +query I +SELECT index FROM pg_temp.t ORDER BY +pg_temp.quantize_to_rabitq4(val)::rabitq4(64)[] @# ARRAY[ + quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector), + quantize_to_rabitq4(array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) +] +LIMIT 10; +---- +1207 +919 +1821 +1076 +1639 +174 +1125 +79 +1163 +537 + +statement ok +DROP INDEX pg_temp.ti; + +statement ok +DROP TABLE pg_temp.t; diff --git a/tests/vchordrq/index_vector.slt b/tests/vchordrq/index_vector.slt index dd2ea412..af69591f 100644 --- a/tests/vchordrq/index_vector.slt +++ b/tests/vchordrq/index_vector.slt @@ -1,3 +1,6 @@ +statement ok +SET enable_seqscan TO off; + statement ok CREATE TABLE t (index serial primary key, val vector(64)); @@ -75,7 +78,7 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)) halfvec_l2_ops); query I -SELECT index FROM t ORDER BY val <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +SELECT index FROM t ORDER BY val::halfvec <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; ---- 1608 155 @@ -95,7 +98,7 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)) halfvec_ip_ops); query I -SELECT index FROM t ORDER BY val <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +SELECT index FROM t ORDER BY val::halfvec <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; ---- 1608 155 @@ -115,130 +118,10 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)) halfvec_cosine_ops); query I -SELECT index FROM t ORDER BY val <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; ----- -1608 -155 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_l2_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ----- -155 -1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_ip_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ----- -155 -1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_cosine_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +SELECT index FROM t ORDER BY val::halfvec <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; ---- -155 1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ----- 155 -1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ----- -155 -1608 -1643 -174 -818 -1603 -1629 -60 -218 -1080 - -statement ok -DROP INDEX ti; - -statement ok -CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); - -query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ----- -155 -1608 1643 174 818 diff --git a/tests/vchordrq/index_vector_rabitq.slt b/tests/vchordrq/index_vector_rabitq.slt new file mode 100644 index 00000000..6233d53e --- /dev/null +++ b/tests/vchordrq/index_vector_rabitq.slt @@ -0,0 +1,237 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE t (index serial primary key, val vector(64)); + +statement ok +INSERT INTO t (val) +SELECT + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1639 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1639 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1608 +1643 +174 +818 +1603 +1629 +60 +218 +1639 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_l2_ops); + +query I +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 +) AS s(index); +---- +6162 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_ip_ops); + +query I +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 +) AS s(index); +---- +6162 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_cosine_ops); + +query I +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 +) AS s(index); +---- +6162 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +818 +1608 +155 +60 +174 +1639 +1080 +218 +1207 +331 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +818 +1608 +155 +60 +174 +1639 +1080 +218 +1207 +331 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +818 +1608 +155 +60 +174 +1639 +1080 +218 +1207 +331 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; From f48362ec8d5c99be1114864b2ff278b7ae31546c Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 10 Dec 2025 15:31:01 +0800 Subject: [PATCH 256/324] fix: aarch64 reduce_min_max_of_x nan handling (#389) Signed-off-by: usamoi --- crates/simd/cshim/aarch64.c | 8 +-- crates/simd/src/floating_f32.rs | 109 +++++++++++++++++--------------- 2 files changed, 61 insertions(+), 56 deletions(-) diff --git a/crates/simd/cshim/aarch64.c b/crates/simd/cshim/aarch64.c index 5b8307f3..27d76d17 100644 --- a/crates/simd/cshim/aarch64.c +++ b/crates/simd/cshim/aarch64.c @@ -217,11 +217,11 @@ fp32_reduce_min_max_of_x_a3_256(float *restrict this, size_t n, float *out_min, for (size_t i = 0; i < n; i += svcntw()) { svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); svfloat32_t x = svld1_f32(mask, this + i); - min = svmin_f32_m(mask, min, x); - max = svmax_f32_m(mask, max, x); + min = svminnm_f32_m(mask, min, x); + max = svmaxnm_f32_m(mask, max, x); } - *out_min = svminv_f32(svptrue_b32(), min); - *out_max = svmaxv_f32(svptrue_b32(), max); + *out_min = svminnmv_f32(svptrue_b32(), min); + *out_max = svmaxnmv_f32(svptrue_b32(), max); } #endif diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs index 5c291a57..1b52a1b3 100644 --- a/crates/simd/src/floating_f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -1006,9 +1006,10 @@ mod reduce_min_max_of_x { let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 200; - let x = (0..n) + let mut x = (0..n) .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); + (x[0], x[1]) = (f32::NAN, -f32::NAN); for z in 50..200 { let x = &x[..z]; let specialized = unsafe { reduce_min_max_of_x_v4(x) }; @@ -1060,9 +1061,10 @@ mod reduce_min_max_of_x { let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 200; - let x = (0..n) + let mut x = (0..n) .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); + (x[0], x[1]) = (f32::NAN, -f32::NAN); for z in 50..200 { let x = &x[..z]; let specialized = unsafe { reduce_min_max_of_x_v3(x) }; @@ -1114,9 +1116,10 @@ mod reduce_min_max_of_x { let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 200; - let x = (0..n) + let mut x = (0..n) .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); + (x[0], x[1]) = (f32::NAN, -f32::NAN); for z in 50..200 { let x = &x[..z]; let specialized = unsafe { reduce_min_max_of_x_v2(x) }; @@ -1127,6 +1130,51 @@ mod reduce_min_max_of_x { } } + #[inline] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] + fn reduce_min_max_of_x_a3_256(this: &[f32]) -> (f32, f32) { + let mut min = f32::INFINITY; + let mut max = -f32::INFINITY; + unsafe { + unsafe extern "C" { + unsafe fn fp32_reduce_min_max_of_x_a3_256( + this: *const f32, + n: usize, + out_min: &mut f32, + out_max: &mut f32, + ); + } + fp32_reduce_min_max_of_x_a3_256(this.as_ptr(), this.len(), &mut min, &mut max); + } + (min, max) + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] + #[test] + fn reduce_min_max_of_x_a3_256_test() { + use rand::Rng; + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + (x[0], x[1]) = (f32::NAN, -f32::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_a3_256(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] @@ -1140,11 +1188,11 @@ mod reduce_min_max_of_x { let x = unsafe { vld1q_f32(a) }; a = unsafe { a.add(4) }; n -= 4; - min = vminq_f32(x, min); - max = vmaxq_f32(x, max); + min = vminnmq_f32(x, min); + max = vmaxnmq_f32(x, max); } - let mut min = vminvq_f32(min); - let mut max = vmaxvq_f32(max); + let mut min = vminnmvq_f32(min); + let mut max = vmaxnmvq_f32(max); // this hint is used to disable loop unrolling while std::hint::black_box(n) > 0 { let x = unsafe { a.read() }; @@ -1167,9 +1215,10 @@ mod reduce_min_max_of_x { let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 200; - let x = (0..n) + let mut x = (0..n) .map(|_| rng.random_range(-1.0..=1.0)) .collect::>(); + (x[0], x[1]) = (f32::NAN, -f32::NAN); for z in 50..200 { let x = &x[..z]; let specialized = unsafe { reduce_min_max_of_x_a2(x) }; @@ -1180,50 +1229,6 @@ mod reduce_min_max_of_x { } } - #[inline] - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - #[crate::target_cpu(enable = "a3.256")] - fn reduce_min_max_of_x_a3_256(this: &[f32]) -> (f32, f32) { - let mut min = f32::INFINITY; - let mut max = -f32::INFINITY; - unsafe { - unsafe extern "C" { - unsafe fn fp32_reduce_min_max_of_x_a3_256( - this: *const f32, - n: usize, - out_min: &mut f32, - out_max: &mut f32, - ); - } - fp32_reduce_min_max_of_x_a3_256(this.as_ptr(), this.len(), &mut min, &mut max); - } - (min, max) - } - - #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] - #[test] - fn reduce_min_max_of_x_a3_256_test() { - use rand::Rng; - if !crate::is_cpu_detected!("a3.256") { - println!("test {} ... skipped (a3.256)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 200; - let x = (0..n) - .map(|_| rng.random_range(-1.0..=1.0)) - .collect::>(); - for z in 50..200 { - let x = &x[..z]; - let specialized = unsafe { reduce_min_max_of_x_a3_256(x) }; - let fallback = fallback(x); - assert_eq!(specialized.0, fallback.0,); - assert_eq!(specialized.1, fallback.1,); - } - } - } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; From e1611304e8e13b4a7e528fa99e17b79298e91f1e Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 15 Dec 2025 15:37:49 +0800 Subject: [PATCH 257/324] chore: don't use clang by default for simd c shim (#393) Some C code is compiled with system's C compiler by default, while SIMD C shim is compiled with Clang by default, which causes issues when using LTO for packaging. Therefore, compile all C code with system's C compiler. Signed-off-by: usamoi --- .github/workflows/check.yml | 12 ++++++++- crates/simd/build.rs | 50 ------------------------------------- 2 files changed, 11 insertions(+), 51 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 292fbfe4..706affbf 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -124,6 +124,11 @@ jobs: rustup default nightly rustup component add miri + - name: Install Clang + run: | + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 + echo CC=clang-18 >> $GITHUB_ENV + - name: Checkout uses: actions/checkout@v4 @@ -174,6 +179,11 @@ jobs: - name: Install Rust run: rustup update + - name: Install Clang + run: | + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 + echo CC=clang-18 >> $GITHUB_ENV + - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 @@ -238,7 +248,7 @@ jobs: - name: Install Clang run: | curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 - sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 + echo CC=clang-18 >> $GITHUB_ENV - name: Install PostgreSQL & pgvector run: | diff --git a/crates/simd/build.rs b/crates/simd/build.rs index bba723b8..40461381 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -14,69 +14,19 @@ use std::env::var; use std::error::Error; -use std::ffi::OsString; -use std::path::Path; - -fn compiler_version(cc: impl AsRef) -> Option { - let cc = cc.as_ref(); - if let Ok(r) = std::process::Command::new(cc).arg("-dumpversion").output() - && r.status.success() - && let Some(major) = r.stdout.split(|c| !c.is_ascii_digit()).next() - && let Ok(major) = std::str::from_utf8(major) - && let Ok(major) = major.parse::() - { - return Some(major); - } - None -} - -fn compiler(host: &str, target: &str, clang_version: u16, gcc_version: u16) -> Option { - let keys = [ - &format!("CC_{target}"), - &format!("CC_{}", target.replace("-", "_")), - "TARGET_CC", - "CC", - ]; - if keys.iter().any(|key| std::env::var_os(key).is_some()) { - return None; - } - if host == target { - if let Ok(cc) = which::which("clang") - && compiler_version(&cc) >= Some(clang_version) - { - return Some(cc.into()); - } - if let Ok(cc) = which::which("gcc") - && compiler_version(&cc) >= Some(gcc_version) - { - return Some(cc.into()); - } - } - None -} fn main() -> Result<(), Box> { println!("cargo::rerun-if-changed=cshim"); - let host = var("HOST")?; - let target = var("TARGET")?; let target_arch = var("CARGO_CFG_TARGET_ARCH")?; match target_arch.as_str() { "aarch64" => { let mut build = cc::Build::new(); - if let Some(compiler) = compiler(&host, &target, 16, 14) { - build.compiler(compiler); - } build.file("./cshim/aarch64.c"); build.opt_level(3); build.compile("simd_cshim"); } - "powerpc64" => {} - "s390x" => {} "x86_64" => { let mut build = cc::Build::new(); - if let Some(compiler) = compiler(&host, &target, 16, 12) { - build.compiler(compiler); - } build.file("./cshim/x86_64.c"); build.opt_level(3); build.compile("simd_cshim"); From 121f2d3d9743b57e312145060297754b268056e6 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 16 Dec 2025 16:26:47 +0800 Subject: [PATCH 258/324] fix: slips in dump-codegen (#394) Signed-off-by: usamoi --- devtools/dump-codegen.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devtools/dump-codegen.sh b/devtools/dump-codegen.sh index 5d3718fd..7ea0f795 100755 --- a/devtools/dump-codegen.sh +++ b/devtools/dump-codegen.sh @@ -29,7 +29,7 @@ Datum amhandler() { IndexAmRoutine *amroutine = makeNode(IndexAmRoutine); amroutine->amsupport = 1; amroutine->amcanorderbyop = true; - (Datum) amroutine; + return (Datum) amroutine; } EOF From 61fce4d7eef5a412a15c2acc8bc546337a6ecce3 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 24 Dec 2025 15:43:58 +0800 Subject: [PATCH 259/324] fix: ci (#396) Signed-off-by: usamoi --- .github/workflows/check.yml | 61 ++++++++++++++++++--------------- .github/workflows/release.yml | 2 +- crates/simd/src/floating_f32.rs | 3 -- crates/simd/src/lib.rs | 1 - 4 files changed, 35 insertions(+), 32 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 706affbf..79e6ea65 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -23,7 +23,7 @@ jobs: run: | curl -fsSL https://github.com/tamasfe/taplo/releases/download/0.10.0/taplo-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo - curl -fsSL https://github.com/EmbarkStudios/cargo-deny/releases/download/0.18.2/cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl/cargo-deny | install -m 755 /dev/stdin /usr/local/bin/cargo-deny + curl -fsSL https://github.com/EmbarkStudios/cargo-deny/releases/download/0.19.2/cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - cargo-deny-0.19.2-$(uname -m)-unknown-linux-musl/cargo-deny | install -m 755 /dev/stdin /usr/local/bin/cargo-deny - name: Install Rust run: rustup update @@ -134,8 +134,7 @@ jobs: - name: Cargo Test (Miri) env: - # https://github.com/rust-lang/rust/issues/149314 - RUSTFLAGS: "-Ctarget-cpu=sapphirerapids -Zcodegen-backend=llvm" + RUSTFLAGS: "-Ctarget-cpu=sapphirerapids" MIRIFLAGS: "-Zmiri-strict-provenance" run: | cargo miri test --locked --target x86_64-unknown-linux-gnu \ @@ -250,6 +249,9 @@ jobs: curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 echo CC=clang-18 >> $GITHUB_ENV + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 + - name: Install PostgreSQL & pgvector run: | sudo apt-get remove -y '^postgres.*' '^libpq.*' @@ -284,9 +286,6 @@ jobs: sudo apt-get install -y --no-install-recommends postgresql-${{ matrix.version }}-pgvector - - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 - - name: Checkout uses: actions/checkout@v4 @@ -360,6 +359,9 @@ jobs: - name: Install Clang run: echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 + - name: Install PostgreSQL & pgvector run: | brew install postgresql@${{ matrix.version }} @@ -377,11 +379,8 @@ jobs: mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install - make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config - sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config install - - - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' + sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' install - name: Checkout uses: actions/checkout@v4 @@ -455,6 +454,9 @@ jobs: - name: Install Rust run: rustup update + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 + - name: Install PostgreSQL & pgvector run: | 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } @@ -500,8 +502,8 @@ jobs: nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" install Pop-Location - - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + - name: Install Clang + run: Add-Content -Path $env:GITHUB_ENV -Value "CC=clang-cl.exe" - name: Checkout uses: actions/checkout@v4 @@ -586,7 +588,12 @@ jobs: echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV - name: Install Clang - run: sudo apk add --no-cache clang18-dev + run: | + sudo apk add --no-cache clang18 clang18-dev + echo CC=clang-18 >> $GITHUB_ENV + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 - name: Install PostgreSQL & pgvector run: | @@ -615,11 +622,8 @@ jobs: mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install - make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config - make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config install - - - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' install - name: Checkout uses: actions/checkout@v4 @@ -786,7 +790,12 @@ jobs: echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV - name: Install Clang - run: sudo apt-get install -y clang lld + run: | + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 + echo CC=clang-18 >> $GITHUB_ENV + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@v0.0.9 - name: Install PostgreSQL & pgvector run: | @@ -825,11 +834,8 @@ jobs: mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install - make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" - sudo make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC=clang CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install - - - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC='sccache clang' CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" + sudo make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC='sccache clang' CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install - name: Checkout uses: actions/checkout@v4 @@ -848,7 +854,6 @@ jobs: "-Dwarnings", ] [env] - CC_${{ matrix.rust_triple }} = "clang" CFLAGS_${{ matrix.rust_triple }} = "--target=${{ matrix.clang_triple }} --sysroot=/sysroot" BINDGEN_EXTRA_CLANG_ARGS_${{ matrix.rust_triple }} = "--sysroot=/sysroot" EOF @@ -942,7 +947,9 @@ jobs: echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV - name: Install Clang - run: sudo pacman -S --noconfirm clang + run: | + sudo pacman -S --noconfirm clang + echo CC=clang >> $GITHUB_ENV - name: Set up Sccache uses: mozilla-actions/sccache-action@v0.0.9 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 17c24812..afddd91d 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -60,7 +60,7 @@ jobs: - name: Install Clang run: | curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 - sudo update-alternatives --install /usr/bin/clang clang $(which clang-18) 255 + echo CC=clang-18 >> $GITHUB_ENV - name: Install PostgreSQL run: | diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs index 1b52a1b3..8e31dce1 100644 --- a/crates/simd/src/floating_f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -965,9 +965,6 @@ mod reduce_sum_of_x2 { } mod reduce_min_max_of_x { - // Semanctics of `f32::min` is different from `_mm256_min_ps`, - // which may lead to issues... - #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 2becec62..c50bb998 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -14,7 +14,6 @@ #![allow(unsafe_code)] #![cfg_attr(feature = "nightly_f16", feature(f16))] -#![cfg_attr(target_arch = "s390x", feature(stdarch_s390x_feature_detection))] #![cfg_attr(target_arch = "s390x", feature(s390x_target_feature))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] #![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc_feature_detection))] From 37405168b0131454ba9ff5548b90c6fc63759c71 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 6 Jan 2026 18:07:45 +0800 Subject: [PATCH 260/324] refactor: rewrite byte and halfbyte simd routines (#399) Removes SVE implementations that are not used for distance computation. Signed-off-by: usamoi --- Cargo.lock | 163 +++ crates/rabitq/src/bit.rs | 6 +- crates/rabitq/src/byte.rs | 2 +- crates/rabitq/src/halfbyte.rs | 2 +- crates/simd/Cargo.toml | 5 + crates/simd/benches/bench.rs | 115 +++ crates/simd/build.rs | 11 +- crates/simd/cshim/aarch64.c | 257 ----- crates/simd/cshim/aarch64_byte.c | 52 + crates/simd/cshim/aarch64_fp16.c | 148 +++ crates/simd/cshim/aarch64_halfbyte.c | 64 ++ crates/simd/cshim/aarch64_sve_fp16.c | 60 ++ crates/simd/cshim/aarch64_sve_fp32.c | 62 ++ crates/simd/cshim/{x86_64.c => x86_64_fp16.c} | 32 +- crates/simd/src/bit.rs | 272 +++-- crates/simd/src/byte.rs | 695 ++++++------- crates/simd/src/emulate.rs | 105 +- crates/simd/src/fast_scan.rs | 21 +- crates/simd/src/floating_f16.rs | 313 +++--- crates/simd/src/floating_f32.rs | 964 +++++++----------- crates/simd/src/halfbyte.rs | 425 ++++---- crates/simd/src/lib.rs | 4 +- crates/simd/src/quantize.rs | 118 +-- 23 files changed, 1953 insertions(+), 1943 deletions(-) create mode 100644 crates/simd/benches/bench.rs delete mode 100644 crates/simd/cshim/aarch64.c create mode 100644 crates/simd/cshim/aarch64_byte.c create mode 100644 crates/simd/cshim/aarch64_fp16.c create mode 100644 crates/simd/cshim/aarch64_halfbyte.c create mode 100644 crates/simd/cshim/aarch64_sve_fp16.c create mode 100644 crates/simd/cshim/aarch64_sve_fp32.c rename crates/simd/cshim/{x86_64.c => x86_64_fp16.c} (81%) diff --git a/Cargo.lock b/Cargo.lock index 525c9b2f..580b26c1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,10 +17,25 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + [[package]] name = "always_equal" version = "0.0.0" +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + [[package]] name = "annotate-snippets" version = "0.11.5" @@ -93,6 +108,12 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + [[package]] name = "bindgen" version = "0.71.1" @@ -146,6 +167,12 @@ dependencies = [ "toml", ] +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.2.44" @@ -181,6 +208,33 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half 2.7.1", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -265,6 +319,41 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "criterion" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d883447757bb0ee46f233e9dc22eb84d93a9508c9b868687b274fc431d886bf" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed943f81ea2faa8dcecbbfa50164acf95d555afec96a27871663b300e387b2e4" +dependencies = [ + "cast", + "itertools", +] + [[package]] name = "crossbeam-channel" version = "0.5.15" @@ -870,6 +959,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "object" version = "0.37.3" @@ -894,6 +992,12 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "owo-colors" version = "4.2.3" @@ -903,6 +1007,16 @@ dependencies = [ "supports-color", ] +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "paste" version = "1.0.15" @@ -1075,6 +1189,34 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "potential_utf" version = "0.1.4" @@ -1381,6 +1523,7 @@ name = "simd" version = "0.0.0" dependencies = [ "cc", + "criterion", "half 2.7.1", "rand", "seq-macro", @@ -1512,6 +1655,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "toml" version = "0.9.8" @@ -1803,6 +1956,16 @@ dependencies = [ "bitflags", ] +[[package]] +name = "web-sys" +version = "0.3.82" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a1f95c0d03a47f4ae1f7a64643a6bb97465d9b740f0fa8f90ea33915c99a9a1" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "which" version = "8.0.0" diff --git a/crates/rabitq/src/bit.rs b/crates/rabitq/src/bit.rs index bfaa0bd7..efeef07b 100644 --- a/crates/rabitq/src/bit.rs +++ b/crates/rabitq/src/bit.rs @@ -175,11 +175,7 @@ pub mod binary { pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BinaryLut { let (k, b, qvector) = simd::quantize::quantize(vector, ((1 << BITS) - 1) as f32); - let qvector_sum = if vector.len() <= (65535_usize / ((1 << BITS) - 1)) { - simd::byte::reduce_sum_of_x_as_u16(&qvector) as f32 - } else { - simd::byte::reduce_sum_of_x_as_u32(&qvector) as f32 - }; + let qvector_sum = simd::byte::reduce_sum_of_x(&qvector) as f32; ( BinaryLutMetadata { dis_v_2, diff --git a/crates/rabitq/src/byte.rs b/crates/rabitq/src/byte.rs index 4af1e3a4..01d9a71e 100644 --- a/crates/rabitq/src/byte.rs +++ b/crates/rabitq/src/byte.rs @@ -48,7 +48,7 @@ pub mod binary { } pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { - simd::byte::reduce_sum_of_x_as_u32_y_as_u32(x, y) + simd::byte::reduce_sum_of_xy(x, y) } pub fn half_process_dot( diff --git a/crates/rabitq/src/halfbyte.rs b/crates/rabitq/src/halfbyte.rs index bd45c986..f60c6420 100644 --- a/crates/rabitq/src/halfbyte.rs +++ b/crates/rabitq/src/halfbyte.rs @@ -57,7 +57,7 @@ pub mod binary { } pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { - simd::halfbyte::reduce_sum_of_x_as_u32_y_as_u32(x, y) + simd::halfbyte::reduce_sum_of_xy(x, y) } pub fn half_process_dot( diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 181661b1..5f9acc54 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -4,6 +4,10 @@ version.workspace = true edition.workspace = true publish = false +[[bench]] +name = "bench" +harness = false + [features] init = [] nightly_f16 = ["zerocopy/float-nightly"] @@ -16,6 +20,7 @@ seq-macro.workspace = true zerocopy.workspace = true [dev-dependencies] +criterion = "0.8.1" rand.workspace = true [build-dependencies] diff --git a/crates/simd/benches/bench.rs b/crates/simd/benches/bench.rs new file mode 100644 index 00000000..a8c55542 --- /dev/null +++ b/crates/simd/benches/bench.rs @@ -0,0 +1,115 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +#![allow(unsafe_code)] + +use criterion::{Criterion, criterion_group, criterion_main}; + +fn byte_reduce_sum_of_xy(c: &mut Criterion) { + use rand::Rng; + let mut rng = rand::rng(); + let x = (0..4095).map(|_| rng.random()).collect::>(); + let y = (0..4095).map(|_| rng.random()).collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512vnni") { + c.bench_function("byte::reduce_sum_of_xy::v4_avx512vnni", |b| { + b.iter(|| unsafe { + simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_v4_avx512vnni(&x, &y) + }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("byte::reduce_sum_of_xy::v4", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("byte::reduce_sum_of_xy::v3", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") { + c.bench_function("byte::reduce_sum_of_xy::v2", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_v2(&x, &y) }) + }); + } +} + +fn byte_reduce_sum_of_x(c: &mut Criterion) { + use rand::Rng; + let mut rng = rand::rng(); + let this = (0..4095).map(|_| rng.random()).collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("byte::reduce_sum_of_x::v4", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_x::reduce_sum_of_x_v4(&this) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("byte::reduce_sum_of_x::v3", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_x::reduce_sum_of_x_v3(&this) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") { + c.bench_function("byte::reduce_sum_of_x::v2", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_x::reduce_sum_of_x_v2(&this) }) + }); + } +} + +fn halfbyte_reduce_sum_of_xy(c: &mut Criterion) { + use rand::Rng; + let mut rng = rand::rng(); + let x = (0..2047).map(|_| rng.random()).collect::>(); + let y = (0..2047).map(|_| rng.random()).collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512vnni") { + c.bench_function("halfbyte::reduce_sum_of_xy::v4_avx512vnni", |b| { + b.iter(|| unsafe { + simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_v4_avx512vnni(&x, &y) + }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("halfbyte::reduce_sum_of_xy::v4", |b| { + b.iter(|| unsafe { simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("halfbyte::reduce_sum_of_xy::v3", |b| { + b.iter(|| unsafe { simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") { + c.bench_function("halfbyte::reduce_sum_of_xy::v2", |b| { + b.iter(|| unsafe { simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_v2(&x, &y) }) + }); + } +} + +criterion_group!( + benches, + byte_reduce_sum_of_xy, + byte_reduce_sum_of_x, + halfbyte_reduce_sum_of_xy +); +criterion_main!(benches); diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 40461381..3229041e 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -18,16 +18,23 @@ use std::error::Error; fn main() -> Result<(), Box> { println!("cargo::rerun-if-changed=cshim"); let target_arch = var("CARGO_CFG_TARGET_ARCH")?; + let target_endian = var("CARGO_CFG_TARGET_ENDIAN")?; match target_arch.as_str() { "aarch64" => { let mut build = cc::Build::new(); - build.file("./cshim/aarch64.c"); + build.file("./cshim/aarch64_fp16.c"); + build.file("./cshim/aarch64_byte.c"); + build.file("./cshim/aarch64_halfbyte.c"); + if target_endian == "little" { + build.file("./cshim/aarch64_sve_fp16.c"); + build.file("./cshim/aarch64_sve_fp32.c"); + } build.opt_level(3); build.compile("simd_cshim"); } "x86_64" => { let mut build = cc::Build::new(); - build.file("./cshim/x86_64.c"); + build.file("./cshim/x86_64_fp16.c"); build.opt_level(3); build.compile("simd_cshim"); } diff --git a/crates/simd/cshim/aarch64.c b/crates/simd/cshim/aarch64.c deleted file mode 100644 index 27d76d17..00000000 --- a/crates/simd/cshim/aarch64.c +++ /dev/null @@ -1,257 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -#if defined(__clang__) -#if !(__clang_major__ >= 16) -#error "Clang version must be at least 16." -#endif -#elif defined(__GNUC__) -#if !(__GNUC__ >= 14) -#error "GCC version must be at least 14." -#endif -#else -#error "This file requires Clang or GCC." -#endif - -#include -#if defined(__AARCH64EL__) -#include -#endif -#include -#include - -typedef __fp16 f16; -typedef float f32; - -__attribute__((target("+fp16"))) float -fp16_reduce_sum_of_xy_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { - float16x8_t xy_0 = vdupq_n_f16(0.0); - float16x8_t xy_1 = vdupq_n_f16(0.0); - float16x8_t xy_2 = vdupq_n_f16(0.0); - float16x8_t xy_3 = vdupq_n_f16(0.0); - while (n >= 32) { - float16x8_t x_0 = vld1q_f16(a + 0); - float16x8_t x_1 = vld1q_f16(a + 8); - float16x8_t x_2 = vld1q_f16(a + 16); - float16x8_t x_3 = vld1q_f16(a + 24); - float16x8_t y_0 = vld1q_f16(b + 0); - float16x8_t y_1 = vld1q_f16(b + 8); - float16x8_t y_2 = vld1q_f16(b + 16); - float16x8_t y_3 = vld1q_f16(b + 24); - a += 32; - b += 32; - n -= 32; - xy_0 = vfmaq_f16(xy_0, x_0, y_0); - xy_1 = vfmaq_f16(xy_1, x_1, y_1); - xy_2 = vfmaq_f16(xy_2, x_2, y_2); - xy_3 = vfmaq_f16(xy_3, x_3, y_3); - } - if (n > 0) { - f16 A[32] = {}; - f16 B[32] = {}; - for (size_t i = 0; i < n; i += 1) { - A[i] = a[i]; - B[i] = b[i]; - } - float16x8_t x_0 = vld1q_f16(A + 0); - float16x8_t x_1 = vld1q_f16(A + 8); - float16x8_t x_2 = vld1q_f16(A + 16); - float16x8_t x_3 = vld1q_f16(A + 24); - float16x8_t y_0 = vld1q_f16(B + 0); - float16x8_t y_1 = vld1q_f16(B + 8); - float16x8_t y_2 = vld1q_f16(B + 16); - float16x8_t y_3 = vld1q_f16(B + 24); - xy_0 = vfmaq_f16(xy_0, x_0, y_0); - xy_1 = vfmaq_f16(xy_1, x_1, y_1); - xy_2 = vfmaq_f16(xy_2, x_2, y_2); - xy_3 = vfmaq_f16(xy_3, x_3, y_3); - } - float16x8_t xy = vaddq_f16(vaddq_f16(xy_0, xy_1), vaddq_f16(xy_2, xy_3)); - float32x4_t lo = vcvt_f32_f16(vget_low_f16(xy)); - float32x4_t hi = vcvt_f32_f16(vget_high_f16(xy)); - return vaddvq_f32(lo) + vaddvq_f32(hi); -} - -#if defined(__AARCH64EL__) -__attribute__((target("+sve"))) float -fp16_reduce_sum_of_xy_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { - svfloat16_t xy = svdup_f16(0.0); - for (size_t i = 0; i < n; i += svcnth()) { - svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); - svfloat16_t x = svld1_f16(mask, a + i); - svfloat16_t y = svld1_f16(mask, b + i); - xy = svmla_f16_m(mask, xy, x, y); - } - return svaddv_f16(svptrue_b16(), xy); -} -#endif - -__attribute__((target("+fp16"))) float -fp16_reduce_sum_of_d2_a2_fp16(f16 *restrict a, f16 *restrict b, size_t n) { - float16x8_t d2_0 = vdupq_n_f16(0.0); - float16x8_t d2_1 = vdupq_n_f16(0.0); - float16x8_t d2_2 = vdupq_n_f16(0.0); - float16x8_t d2_3 = vdupq_n_f16(0.0); - while (n >= 32) { - float16x8_t x_0 = vld1q_f16(a + 0); - float16x8_t x_1 = vld1q_f16(a + 8); - float16x8_t x_2 = vld1q_f16(a + 16); - float16x8_t x_3 = vld1q_f16(a + 24); - float16x8_t y_0 = vld1q_f16(b + 0); - float16x8_t y_1 = vld1q_f16(b + 8); - float16x8_t y_2 = vld1q_f16(b + 16); - float16x8_t y_3 = vld1q_f16(b + 24); - a += 32; - b += 32; - n -= 32; - float16x8_t d_0 = vsubq_f16(x_0, y_0); - float16x8_t d_1 = vsubq_f16(x_1, y_1); - float16x8_t d_2 = vsubq_f16(x_2, y_2); - float16x8_t d_3 = vsubq_f16(x_3, y_3); - d2_0 = vfmaq_f16(d2_0, d_0, d_0); - d2_1 = vfmaq_f16(d2_1, d_1, d_1); - d2_2 = vfmaq_f16(d2_2, d_2, d_2); - d2_3 = vfmaq_f16(d2_3, d_3, d_3); - } - if (n > 0) { - f16 A[32] = {}; - f16 B[32] = {}; - for (size_t i = 0; i < n; i += 1) { - A[i] = a[i]; - B[i] = b[i]; - } - float16x8_t x_0 = vld1q_f16(A + 0); - float16x8_t x_1 = vld1q_f16(A + 8); - float16x8_t x_2 = vld1q_f16(A + 16); - float16x8_t x_3 = vld1q_f16(A + 24); - float16x8_t y_0 = vld1q_f16(B + 0); - float16x8_t y_1 = vld1q_f16(B + 8); - float16x8_t y_2 = vld1q_f16(B + 16); - float16x8_t y_3 = vld1q_f16(B + 24); - float16x8_t d_0 = vsubq_f16(x_0, y_0); - float16x8_t d_1 = vsubq_f16(x_1, y_1); - float16x8_t d_2 = vsubq_f16(x_2, y_2); - float16x8_t d_3 = vsubq_f16(x_3, y_3); - d2_0 = vfmaq_f16(d2_0, d_0, d_0); - d2_1 = vfmaq_f16(d2_1, d_1, d_1); - d2_2 = vfmaq_f16(d2_2, d_2, d_2); - d2_3 = vfmaq_f16(d2_3, d_3, d_3); - } - float16x8_t d2 = vaddq_f16(vaddq_f16(d2_0, d2_1), vaddq_f16(d2_2, d2_3)); - float32x4_t lo = vcvt_f32_f16(vget_low_f16(d2)); - float32x4_t hi = vcvt_f32_f16(vget_high_f16(d2)); - return vaddvq_f32(lo) + vaddvq_f32(hi); -} - -#if defined(__AARCH64EL__) -__attribute__((target("+sve"))) float -fp16_reduce_sum_of_d2_a3_512(f16 *restrict a, f16 *restrict b, size_t n) { - svfloat16_t d2 = svdup_f16(0.0); - for (size_t i = 0; i < n; i += svcnth()) { - svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); - svfloat16_t x = svld1_f16(mask, a + i); - svfloat16_t y = svld1_f16(mask, b + i); - svfloat16_t d = svsub_f16_z(mask, x, y); - d2 = svmla_f16_m(mask, d2, d, d); - } - return svaddv_f16(svptrue_b16(), d2); -} -#endif - -#if defined(__AARCH64EL__) -__attribute__((target("+sve"))) float -fp32_reduce_sum_of_x_a3_256(float *restrict this, size_t n) { - svfloat32_t sum = svdup_f32(0.0); - for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); - svfloat32_t x = svld1_f32(mask, this + i); - sum = svadd_f32_m(mask, sum, x); - } - return svaddv_f32(svptrue_b32(), sum); -} -#endif - -#if defined(__AARCH64EL__) -__attribute__((target("+sve"))) float -fp32_reduce_sum_of_abs_x_a3_256(float *restrict this, size_t n) { - svfloat32_t sum = svdup_f32(0.0); - for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); - svfloat32_t x = svld1_f32(mask, this + i); - sum = svadd_f32_m(mask, sum, svabs_f32_z(mask, x)); - } - return svaddv_f32(svptrue_b32(), sum); -} -#endif - -#if defined(__AARCH64EL__) -__attribute__((target("+sve"))) float -fp32_reduce_sum_of_x2_a3_256(float *restrict this, size_t n) { - svfloat32_t sum = svdup_f32(0.0); - for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); - svfloat32_t x = svld1_f32(mask, this + i); - sum = svmla_f32_m(mask, sum, x, x); - } - return svaddv_f32(svptrue_b32(), sum); -} -#endif - -#if defined(__AARCH64EL__) -__attribute__((target("+sve"))) void -fp32_reduce_min_max_of_x_a3_256(float *restrict this, size_t n, float *out_min, - float *out_max) { - svfloat32_t min = svdup_f32(1.0 / 0.0); - svfloat32_t max = svdup_f32(-1.0 / 0.0); - for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); - svfloat32_t x = svld1_f32(mask, this + i); - min = svminnm_f32_m(mask, min, x); - max = svmaxnm_f32_m(mask, max, x); - } - *out_min = svminnmv_f32(svptrue_b32(), min); - *out_max = svmaxnmv_f32(svptrue_b32(), max); -} -#endif - -#if defined(__AARCH64EL__) -__attribute__((target("+sve"))) float -fp32_reduce_sum_of_xy_a3_256(float *restrict lhs, float *restrict rhs, - size_t n) { - svfloat32_t sum = svdup_f32(0.0); - for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); - svfloat32_t x = svld1_f32(mask, lhs + i); - svfloat32_t y = svld1_f32(mask, rhs + i); - sum = svmla_f32_m(mask, sum, x, y); - } - return svaddv_f32(svptrue_b32(), sum); -} -#endif - -#if defined(__AARCH64EL__) -__attribute__((target("+sve"))) float -fp32_reduce_sum_of_d2_a3_256(float *restrict lhs, float *restrict rhs, - size_t n) { - svfloat32_t sum = svdup_f32(0.0); - for (size_t i = 0; i < n; i += svcntw()) { - svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); - svfloat32_t x = svld1_f32(mask, lhs + i); - svfloat32_t y = svld1_f32(mask, rhs + i); - svfloat32_t d = svsub_f32_z(mask, x, y); - sum = svmla_f32_m(mask, sum, d, d); - } - return svaddv_f32(svptrue_b32(), sum); -} -#endif diff --git a/crates/simd/cshim/aarch64_byte.c b/crates/simd/cshim/aarch64_byte.c new file mode 100644 index 00000000..ddd3818e --- /dev/null +++ b/crates/simd/cshim/aarch64_byte.c @@ -0,0 +1,52 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." +#endif + +#include +#include +#include + +__attribute__((target("+dotprod"))) uint32_t +byte_reduce_sum_of_xy_a2_dotprod(size_t n, uint8_t *restrict a, + uint8_t *restrict b) { + uint32x4_t sum = vdupq_n_u32(0); + while (n >= 16) { + uint8x16_t x = vld1q_u8(a); + uint8x16_t y = vld1q_u8(b); + sum = vdotq_u32(sum, x, y); + n -= 16, a += 16, b += 16; + } + if (n > 0) { + uint8_t _a[16] = {}, _b[16] = {}; + for (size_t i = 0; i < n; i += 1) { + _a[i] = a[i], _b[i] = b[i]; + } + a = _a, b = _b; + uint8x16_t x = vld1q_u8(a); + uint8x16_t y = vld1q_u8(b); + sum = vdotq_u32(sum, x, y); + } + return vaddvq_u32(sum); +} diff --git a/crates/simd/cshim/aarch64_fp16.c b/crates/simd/cshim/aarch64_fp16.c new file mode 100644 index 00000000..c72b14fb --- /dev/null +++ b/crates/simd/cshim/aarch64_fp16.c @@ -0,0 +1,148 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." +#endif + +#include +#include +#include + +typedef __fp16 f16; +typedef float f32; + +__attribute__((target("+fp16"))) float +fp16_reduce_sum_of_xy_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { + float16x8_t sum_0 = vdupq_n_f16(0.0); + float16x8_t sum_1 = vdupq_n_f16(0.0); + float16x8_t sum_2 = vdupq_n_f16(0.0); + float16x8_t sum_3 = vdupq_n_f16(0.0); + while (n >= 32) { + float16x8_t x_0 = vld1q_f16(a + 0); + float16x8_t x_1 = vld1q_f16(a + 8); + float16x8_t x_2 = vld1q_f16(a + 16); + float16x8_t x_3 = vld1q_f16(a + 24); + float16x8_t y_0 = vld1q_f16(b + 0); + float16x8_t y_1 = vld1q_f16(b + 8); + float16x8_t y_2 = vld1q_f16(b + 16); + float16x8_t y_3 = vld1q_f16(b + 24); + sum_0 = vfmaq_f16(sum_0, x_0, y_0); + sum_1 = vfmaq_f16(sum_1, x_1, y_1); + sum_2 = vfmaq_f16(sum_2, x_2, y_2); + sum_3 = vfmaq_f16(sum_3, x_3, y_3); + n -= 32, a += 32, b += 32; + } + if (n > 0) { + f16 _a[32] = {}, _b[32] = {}; + for (size_t i = 0; i < n; i += 1) { + _a[i] = a[i], _b[i] = b[i]; + } + a = _a, b = _b; + float16x8_t x_0 = vld1q_f16(_a + 0); + float16x8_t x_1 = vld1q_f16(_a + 8); + float16x8_t x_2 = vld1q_f16(_a + 16); + float16x8_t x_3 = vld1q_f16(_a + 24); + float16x8_t y_0 = vld1q_f16(_b + 0); + float16x8_t y_1 = vld1q_f16(_b + 8); + float16x8_t y_2 = vld1q_f16(_b + 16); + float16x8_t y_3 = vld1q_f16(_b + 24); + sum_0 = vfmaq_f16(sum_0, x_0, y_0); + sum_1 = vfmaq_f16(sum_1, x_1, y_1); + sum_2 = vfmaq_f16(sum_2, x_2, y_2); + sum_3 = vfmaq_f16(sum_3, x_3, y_3); + } + float32x4_t s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + float32x4_t s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + float32x4_t s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + float32x4_t s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + float32x4_t s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + float32x4_t s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + float32x4_t s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + float32x4_t s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + float32x4_t s = + vpaddq_f32(vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7))); + return vaddvq_f32(s); +} + +__attribute__((target("+fp16"))) float +fp16_reduce_sum_of_d2_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { + float16x8_t sum_0 = vdupq_n_f16(0.0); + float16x8_t sum_1 = vdupq_n_f16(0.0); + float16x8_t sum_2 = vdupq_n_f16(0.0); + float16x8_t sum_3 = vdupq_n_f16(0.0); + while (n >= 32) { + float16x8_t x_0 = vld1q_f16(a + 0); + float16x8_t x_1 = vld1q_f16(a + 8); + float16x8_t x_2 = vld1q_f16(a + 16); + float16x8_t x_3 = vld1q_f16(a + 24); + float16x8_t y_0 = vld1q_f16(b + 0); + float16x8_t y_1 = vld1q_f16(b + 8); + float16x8_t y_2 = vld1q_f16(b + 16); + float16x8_t y_3 = vld1q_f16(b + 24); + float16x8_t d_0 = vsubq_f16(x_0, y_0); + float16x8_t d_1 = vsubq_f16(x_1, y_1); + float16x8_t d_2 = vsubq_f16(x_2, y_2); + float16x8_t d_3 = vsubq_f16(x_3, y_3); + sum_0 = vfmaq_f16(sum_0, d_0, d_0); + sum_1 = vfmaq_f16(sum_1, d_1, d_1); + sum_2 = vfmaq_f16(sum_2, d_2, d_2); + sum_3 = vfmaq_f16(sum_3, d_3, d_3); + n -= 32, a += 32, b += 32; + } + if (n > 0) { + f16 _a[32] = {}, _b[32] = {}; + for (size_t i = 0; i < n; i += 1) { + _a[i] = a[i], _b[i] = b[i]; + } + a = _a, b = _b; + float16x8_t x_0 = vld1q_f16(a + 0); + float16x8_t x_1 = vld1q_f16(a + 8); + float16x8_t x_2 = vld1q_f16(a + 16); + float16x8_t x_3 = vld1q_f16(a + 24); + float16x8_t y_0 = vld1q_f16(b + 0); + float16x8_t y_1 = vld1q_f16(b + 8); + float16x8_t y_2 = vld1q_f16(b + 16); + float16x8_t y_3 = vld1q_f16(b + 24); + float16x8_t d_0 = vsubq_f16(x_0, y_0); + float16x8_t d_1 = vsubq_f16(x_1, y_1); + float16x8_t d_2 = vsubq_f16(x_2, y_2); + float16x8_t d_3 = vsubq_f16(x_3, y_3); + sum_0 = vfmaq_f16(sum_0, d_0, d_0); + sum_1 = vfmaq_f16(sum_1, d_1, d_1); + sum_2 = vfmaq_f16(sum_2, d_2, d_2); + sum_3 = vfmaq_f16(sum_3, d_3, d_3); + } + float32x4_t s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + float32x4_t s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + float32x4_t s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + float32x4_t s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + float32x4_t s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + float32x4_t s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + float32x4_t s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + float32x4_t s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + float32x4_t s = + vpaddq_f32(vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7))); + return vaddvq_f32(s); +} diff --git a/crates/simd/cshim/aarch64_halfbyte.c b/crates/simd/cshim/aarch64_halfbyte.c new file mode 100644 index 00000000..53026e0a --- /dev/null +++ b/crates/simd/cshim/aarch64_halfbyte.c @@ -0,0 +1,64 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." +#endif + +#include +#include +#include + +__attribute__((target("+dotprod"))) uint32_t +halfbyte_reduce_sum_of_xy_a2_dotprod(size_t n, uint8_t *restrict a, + uint8_t *restrict b) { + uint8x16_t lo = vdupq_n_u8(0x0f); + uint32x4_t _0 = vdupq_n_u32(0); + uint32x4_t _1 = vdupq_n_u32(0); + while (n >= 16) { + uint8x16_t x = vld1q_u8(a); + uint8x16_t y = vld1q_u8(b); + uint8x16_t x_0 = vandq_u8(x, lo); + uint8x16_t x_1 = vshrq_n_u8(x, 4); + uint8x16_t y_0 = vandq_u8(y, lo); + uint8x16_t y_1 = vshrq_n_u8(y, 4); + _0 = vdotq_u32(_0, x_0, y_0); + _1 = vdotq_u32(_1, x_1, y_1); + n -= 16, a += 16, b += 16; + } + if (n > 0) { + uint8_t _a[16] = {}, _b[16] = {}; + for (size_t i = 0; i < n; i += 1) { + _a[i] = a[i], _b[i] = b[i]; + } + a = _a, b = _b; + uint8x16_t x = vld1q_u8(a); + uint8x16_t y = vld1q_u8(b); + uint8x16_t x_0 = vandq_u8(x, lo); + uint8x16_t x_1 = vshrq_n_u8(x, 4); + uint8x16_t y_0 = vandq_u8(y, lo); + uint8x16_t y_1 = vshrq_n_u8(y, 4); + _0 = vdotq_u32(_0, x_0, y_0); + _1 = vdotq_u32(_1, x_1, y_1); + } + return vaddvq_u32(vaddq_u32(_0, _1)); +} diff --git a/crates/simd/cshim/aarch64_sve_fp16.c b/crates/simd/cshim/aarch64_sve_fp16.c new file mode 100644 index 00000000..77111827 --- /dev/null +++ b/crates/simd/cshim/aarch64_sve_fp16.c @@ -0,0 +1,60 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." +#endif + +#include +#if defined(__AARCH64EL__) +#include +#endif +#include +#include + +typedef __fp16 f16; +typedef float f32; + +__attribute__((target("+sve"))) float +fp16_reduce_sum_of_xy_a3_512(size_t n, f16 *restrict a, f16 *restrict b) { + svfloat16_t sum = svdup_f16(0.0); + for (size_t i = 0; i < n; i += svcnth()) { + svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); + svfloat16_t x = svld1_f16(mask, a + i); + svfloat16_t y = svld1_f16(mask, b + i); + sum = svmla_f16_m(mask, sum, x, y); + } + return svaddv_f16(svptrue_b16(), sum); +} + +__attribute__((target("+sve"))) float +fp16_reduce_sum_of_d2_a3_512(size_t n, f16 *restrict a, f16 *restrict b) { + svfloat16_t sum = svdup_f16(0.0); + for (size_t i = 0; i < n; i += svcnth()) { + svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); + svfloat16_t x = svld1_f16(mask, a + i); + svfloat16_t y = svld1_f16(mask, b + i); + svfloat16_t d = svsub_f16_z(mask, x, y); + sum = svmla_f16_m(mask, sum, d, d); + } + return svaddv_f16(svptrue_b16(), sum); +} diff --git a/crates/simd/cshim/aarch64_sve_fp32.c b/crates/simd/cshim/aarch64_sve_fp32.c new file mode 100644 index 00000000..15e3de6e --- /dev/null +++ b/crates/simd/cshim/aarch64_sve_fp32.c @@ -0,0 +1,62 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025 TensorChord Inc. + +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." +#endif + +#include +#if defined(__AARCH64EL__) +#include +#endif +#include +#include + +typedef __fp16 f16; +typedef float f32; + +__attribute__((target("+sve"))) float +fp32_reduce_sum_of_xy_a3_256(size_t n, float *restrict lhs, + float *restrict rhs) { + svfloat32_t sum = svdup_f32(0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); + svfloat32_t x = svld1_f32(mask, lhs + i); + svfloat32_t y = svld1_f32(mask, rhs + i); + sum = svmla_f32_m(mask, sum, x, y); + } + return svaddv_f32(svptrue_b32(), sum); +} + +__attribute__((target("+sve"))) float +fp32_reduce_sum_of_d2_a3_256(size_t n, float *restrict lhs, + float *restrict rhs) { + svfloat32_t sum = svdup_f32(0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); + svfloat32_t x = svld1_f32(mask, lhs + i); + svfloat32_t y = svld1_f32(mask, rhs + i); + svfloat32_t d = svsub_f32_z(mask, x, y); + sum = svmla_f32_m(mask, sum, d, d); + } + return svaddv_f32(svptrue_b32(), sum); +} diff --git a/crates/simd/cshim/x86_64.c b/crates/simd/cshim/x86_64_fp16.c similarity index 81% rename from crates/simd/cshim/x86_64.c rename to crates/simd/cshim/x86_64_fp16.c index 0f469dd6..f9f40ea1 100644 --- a/crates/simd/cshim/x86_64.c +++ b/crates/simd/cshim/x86_64_fp16.c @@ -60,46 +60,42 @@ typedef float f32; __attribute__((target("avx512bw,avx512cd,avx512dq,avx512vl,bmi,bmi2,lzcnt," "movbe,popcnt,avx512fp16"))) float -fp16_reduce_sum_of_xy_v4_avx512fp16(f16 *restrict a, f16 *restrict b, - size_t n) { - __m512h xy = _mm512_setzero_ph(); +fp16_reduce_sum_of_xy_v4_avx512fp16(size_t n, f16 *restrict a, + f16 *restrict b) { + __m512h sum = _mm512_setzero_ph(); while (n >= 32) { __m512h x = _mm512_loadu_ph(a); __m512h y = _mm512_loadu_ph(b); - a += 32; - b += 32; - n -= 32; - xy = _mm512_fmadd_ph(x, y, xy); + sum = _mm512_fmadd_ph(x, y, sum); + n -= 32, a += 32, b += 32; } if (n > 0) { unsigned int mask = _bzhi_u32(0xffffffff, n); __m512h x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a)); __m512h y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b)); - xy = _mm512_fmadd_ph(x, y, xy); + sum = _mm512_fmadd_ph(x, y, sum); } - return _mm512_reduce_add_ph(xy); + return _mm512_reduce_add_ph(sum); } __attribute__((target("avx512bw,avx512cd,avx512dq,avx512vl,bmi,bmi2,lzcnt," "movbe,popcnt,avx512fp16"))) float -fp16_reduce_sum_of_d2_v4_avx512fp16(f16 *restrict a, f16 *restrict b, - size_t n) { - __m512h d2 = _mm512_setzero_ph(); +fp16_reduce_sum_of_d2_v4_avx512fp16(size_t n, f16 *restrict a, + f16 *restrict b) { + __m512h sum = _mm512_setzero_ph(); while (n >= 32) { __m512h x = _mm512_loadu_ph(a); __m512h y = _mm512_loadu_ph(b); - a += 32; - b += 32; - n -= 32; __m512h d = _mm512_sub_ph(x, y); - d2 = _mm512_fmadd_ph(d, d, d2); + sum = _mm512_fmadd_ph(d, d, sum); + n -= 32, a += 32, b += 32; } if (n > 0) { unsigned int mask = _bzhi_u32(0xffffffff, n); __m512h x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a)); __m512h y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b)); __m512h d = _mm512_sub_ph(x, y); - d2 = _mm512_fmadd_ph(d, d, d2); + sum = _mm512_fmadd_ph(d, d, sum); } - return _mm512_reduce_add_ph(d2); + return _mm512_reduce_add_ph(sum); } diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index e372f860..24541a8a 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -24,7 +24,7 @@ mod reduce_sum_of_and { #[target_feature(enable = "avx512vpopcntdq")] fn reduce_sum_of_and_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut and = _mm512_setzero_si512(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); @@ -32,10 +32,8 @@ mod reduce_sum_of_and { while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; let y = unsafe { _mm512_loadu_si512(b.cast()) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -46,7 +44,7 @@ mod reduce_sum_of_and { _mm512_reduce_add_epi64(and) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_and_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { @@ -67,7 +65,7 @@ mod reduce_sum_of_and { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_and_v4(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); @@ -78,10 +76,6 @@ mod reduce_sum_of_and { while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; let y = unsafe { _mm512_loadu_si512(b.cast()) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; - // let and = _mm512_and_si512(x, y); let and_lo = _mm512_and_si512(and, mask_0); let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); @@ -90,12 +84,12 @@ mod reduce_sum_of_and { let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); sum_and = _mm512_add_epi64(sum_and, and_sad); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; - // let and = _mm512_and_si512(x, y); let and_lo = _mm512_and_si512(and, mask_0); let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); @@ -108,7 +102,7 @@ mod reduce_sum_of_and { _mm512_reduce_add_epi64(sum_and) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_and_v4_test() { if !crate::is_cpu_detected!("v4") { @@ -129,8 +123,8 @@ mod reduce_sum_of_and { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_and_v3(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm256_reduce_add_epi64; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); @@ -141,10 +135,6 @@ mod reduce_sum_of_and { while n >= 4 { let x = unsafe { _mm256_loadu_si256(a.cast()) }; let y = unsafe { _mm256_loadu_si256(b.cast()) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; - // let and = _mm256_and_si256(x, y); let and_lo = _mm256_and_si256(and, mask_0); let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); @@ -153,21 +143,26 @@ mod reduce_sum_of_and { let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); sum_and = _mm256_add_epi64(sum_and, and_sad); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut and = emulate_mm256_reduce_add_epi64(sum_and) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - and += (x & y).count_ones(); + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let and = _mm256_and_si256(x, y); + let and_lo = _mm256_and_si256(and, mask_0); + let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); + let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); + sum_and = _mm256_add_epi64(sum_and, and_sad); } - and + emulate_mm256_reduce_add_epi64(sum_and) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_and_v3_test() { if !crate::is_cpu_detected!("v3") { @@ -207,7 +202,7 @@ mod reduce_sum_of_or { #[target_feature(enable = "avx512vpopcntdq")] fn reduce_sum_of_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut or = _mm512_setzero_si512(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); @@ -215,10 +210,8 @@ mod reduce_sum_of_or { while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; let y = unsafe { _mm512_loadu_si512(b.cast()) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -229,7 +222,7 @@ mod reduce_sum_of_or { _mm512_reduce_add_epi64(or) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_or_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { @@ -250,7 +243,7 @@ mod reduce_sum_of_or { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_or_v4(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); @@ -261,10 +254,6 @@ mod reduce_sum_of_or { while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; let y = unsafe { _mm512_loadu_si512(b.cast()) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; - // let or = _mm512_or_si512(x, y); let or_lo = _mm512_and_si512(or, mask_0); let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); @@ -273,12 +262,12 @@ mod reduce_sum_of_or { let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); sum_or = _mm512_add_epi64(sum_or, or_sad); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; - // let or = _mm512_or_si512(x, y); let or_lo = _mm512_and_si512(or, mask_0); let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); @@ -291,7 +280,7 @@ mod reduce_sum_of_or { _mm512_reduce_add_epi64(sum_or) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_or_v4_test() { if !crate::is_cpu_detected!("v4") { @@ -312,8 +301,8 @@ mod reduce_sum_of_or { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_or_v3(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm256_reduce_add_epi64; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); @@ -324,10 +313,6 @@ mod reduce_sum_of_or { while n >= 4 { let x = unsafe { _mm256_loadu_si256(a.cast()) }; let y = unsafe { _mm256_loadu_si256(b.cast()) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; - // let or = _mm256_or_si256(x, y); let or_lo = _mm256_and_si256(or, mask_0); let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); @@ -336,21 +321,26 @@ mod reduce_sum_of_or { let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); sum_or = _mm256_add_epi64(sum_or, or_sad); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut or = emulate_mm256_reduce_add_epi64(sum_or) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - or += (x | y).count_ones(); + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let or = _mm256_or_si256(x, y); + let or_lo = _mm256_and_si256(or, mask_0); + let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); + let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); + sum_or = _mm256_add_epi64(sum_or, or_sad); } - or + emulate_mm256_reduce_add_epi64(sum_or) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_or_v3_test() { if !crate::is_cpu_detected!("v3") { @@ -390,7 +380,7 @@ mod reduce_sum_of_xor { #[target_feature(enable = "avx512vpopcntdq")] fn reduce_sum_of_xor_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut xor = _mm512_setzero_si512(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); @@ -398,10 +388,8 @@ mod reduce_sum_of_xor { while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; let y = unsafe { _mm512_loadu_si512(b.cast()) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; xor = _mm512_add_epi64(xor, _mm512_popcnt_epi64(_mm512_xor_si512(x, y))); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -412,7 +400,7 @@ mod reduce_sum_of_xor { _mm512_reduce_add_epi64(xor) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xor_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { @@ -433,7 +421,7 @@ mod reduce_sum_of_xor { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_xor_v4(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); @@ -444,10 +432,6 @@ mod reduce_sum_of_xor { while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; let y = unsafe { _mm512_loadu_si512(b.cast()) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; - // let xor = _mm512_xor_si512(x, y); let xor_lo = _mm512_and_si512(xor, mask_0); let xor_hi = _mm512_and_si512(_mm512_srli_epi16(xor, 4), mask_0); @@ -456,12 +440,12 @@ mod reduce_sum_of_xor { let xor_res = _mm512_add_epi8(xor_res_lo, xor_res_hi); let xor_sad = _mm512_sad_epu8(xor_res, _mm512_setzero_si512()); sum_xor = _mm512_add_epi64(sum_xor, xor_sad); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; - // let xor = _mm512_xor_si512(x, y); let xor_lo = _mm512_and_si512(xor, mask_0); let xor_hi = _mm512_and_si512(_mm512_srli_epi16(xor, 4), mask_0); @@ -474,7 +458,7 @@ mod reduce_sum_of_xor { _mm512_reduce_add_epi64(sum_xor) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xor_v4_test() { if !crate::is_cpu_detected!("v4") { @@ -495,8 +479,8 @@ mod reduce_sum_of_xor { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_xor_v3(lhs: &[u64], rhs: &[u64]) -> u32 { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm256_reduce_add_epi64; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); @@ -507,10 +491,6 @@ mod reduce_sum_of_xor { while n >= 4 { let x = unsafe { _mm256_loadu_si256(a.cast()) }; let y = unsafe { _mm256_loadu_si256(b.cast()) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; - // let xor = _mm256_xor_si256(x, y); let xor_lo = _mm256_and_si256(xor, mask_0); let xor_hi = _mm256_and_si256(_mm256_srli_epi16(xor, 4), mask_0); @@ -519,21 +499,26 @@ mod reduce_sum_of_xor { let xor_res = _mm256_add_epi8(xor_res_lo, xor_res_hi); let xor_sad = _mm256_sad_epu8(xor_res, _mm256_setzero_si256()); sum_xor = _mm256_add_epi64(sum_xor, xor_sad); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut xor = emulate_mm256_reduce_add_epi64(sum_xor) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - xor += (x ^ y).count_ones(); + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let xor = _mm256_xor_si256(x, y); + let xor_lo = _mm256_and_si256(xor, mask_0); + let xor_hi = _mm256_and_si256(_mm256_srli_epi16(xor, 4), mask_0); + let xor_res_lo = _mm256_shuffle_epi8(lut, xor_lo); + let xor_res_hi = _mm256_shuffle_epi8(lut, xor_hi); + let xor_res = _mm256_add_epi8(xor_res_lo, xor_res_hi); + let xor_sad = _mm256_sad_epu8(xor_res, _mm256_setzero_si256()); + sum_xor = _mm256_add_epi64(sum_xor, xor_sad); } - xor + emulate_mm256_reduce_add_epi64(sum_xor) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xor_v3_test() { if !crate::is_cpu_detected!("v3") { @@ -573,7 +558,7 @@ mod reduce_sum_of_and_or { #[target_feature(enable = "avx512vpopcntdq")] fn reduce_sum_of_and_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut and = _mm512_setzero_si512(); let mut or = _mm512_setzero_si512(); let mut a = lhs.as_ptr(); @@ -582,11 +567,9 @@ mod reduce_sum_of_and_or { while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; let y = unsafe { _mm512_loadu_si512(b.cast()) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -601,7 +584,7 @@ mod reduce_sum_of_and_or { ) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_and_or_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { @@ -622,7 +605,7 @@ mod reduce_sum_of_and_or { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_and_or_v4(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); @@ -634,10 +617,6 @@ mod reduce_sum_of_and_or { while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; let y = unsafe { _mm512_loadu_si512(b.cast()) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; - // let and = _mm512_and_si512(x, y); let and_lo = _mm512_and_si512(and, mask_0); let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); @@ -646,7 +625,6 @@ mod reduce_sum_of_and_or { let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); sum_and = _mm512_add_epi64(sum_and, and_sad); - // let or = _mm512_or_si512(x, y); let or_lo = _mm512_and_si512(or, mask_0); let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); @@ -655,12 +633,12 @@ mod reduce_sum_of_and_or { let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); sum_or = _mm512_add_epi64(sum_or, or_sad); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; - // let and = _mm512_and_si512(x, y); let and_lo = _mm512_and_si512(and, mask_0); let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); @@ -669,7 +647,6 @@ mod reduce_sum_of_and_or { let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); sum_and = _mm512_add_epi64(sum_and, and_sad); - // let or = _mm512_or_si512(x, y); let or_lo = _mm512_and_si512(or, mask_0); let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); @@ -685,7 +662,7 @@ mod reduce_sum_of_and_or { ) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_and_or_v4_test() { if !crate::is_cpu_detected!("v4") { @@ -706,8 +683,8 @@ mod reduce_sum_of_and_or { #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_and_or_v3(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { assert!(lhs.len() == rhs.len()); - use crate::emulate::emulate_mm256_reduce_add_epi64; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); @@ -719,10 +696,6 @@ mod reduce_sum_of_and_or { while n >= 4 { let x = unsafe { _mm256_loadu_si256(a.cast()) }; let y = unsafe { _mm256_loadu_si256(b.cast()) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; - // let and = _mm256_and_si256(x, y); let and_lo = _mm256_and_si256(and, mask_0); let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); @@ -731,7 +704,6 @@ mod reduce_sum_of_and_or { let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); sum_and = _mm256_add_epi64(sum_and, and_sad); - // let or = _mm256_or_si256(x, y); let or_lo = _mm256_and_si256(or, mask_0); let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); @@ -740,23 +712,37 @@ mod reduce_sum_of_and_or { let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); sum_or = _mm256_add_epi64(sum_or, or_sad); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut and = emulate_mm256_reduce_add_epi64(sum_and) as u32; - let mut or = emulate_mm256_reduce_add_epi64(sum_or) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - and += (x & y).count_ones(); - or += (x | y).count_ones(); + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let and = _mm256_and_si256(x, y); + let and_lo = _mm256_and_si256(and, mask_0); + let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); + let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); + sum_and = _mm256_add_epi64(sum_and, and_sad); + let or = _mm256_or_si256(x, y); + let or_lo = _mm256_and_si256(or, mask_0); + let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); + let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); + sum_or = _mm256_add_epi64(sum_or, or_sad); } - (and, or) + ( + emulate_mm256_reduce_add_epi64(sum_and) as u32, + emulate_mm256_reduce_add_epi64(sum_or) as u32, + ) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_and_or_v3_test() { if !crate::is_cpu_detected!("v3") { @@ -797,15 +783,14 @@ mod reduce_sum_of_x { #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vpopcntdq")] fn reduce_sum_of_x_v4_avx512vpopcntdq(this: &[u64]) -> u32 { - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut sum = _mm512_setzero_si512(); let mut a = this.as_ptr(); let mut n = this.len(); while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; - a = unsafe { a.add(8) }; - n -= 8; sum = _mm512_add_epi64(sum, _mm512_popcnt_epi64(x)); + (n, a) = unsafe { (n - 8, a.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -815,7 +800,7 @@ mod reduce_sum_of_x { _mm512_reduce_add_epi64(sum) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x_v4_avx512vpopcntdq_test() { if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { @@ -834,7 +819,7 @@ mod reduce_sum_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_x_v4(this: &[u64]) -> u32 { - use std::arch::x86_64::*; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; let mask_0 = _mm512_set1_epi8(0x0f); @@ -843,8 +828,6 @@ mod reduce_sum_of_x { let mut n = this.len(); while n >= 8 { let x = unsafe { _mm512_loadu_si512(a.cast()) }; - a = unsafe { a.add(8) }; - n -= 8; let lo = _mm512_and_si512(x, mask_0); let hi = _mm512_and_si512(_mm512_srli_epi16(x, 4), mask_0); let res_lo = _mm512_shuffle_epi8(lut, lo); @@ -852,6 +835,7 @@ mod reduce_sum_of_x { let res = _mm512_add_epi8(res_lo, res_hi); let sad = _mm512_sad_epu8(res, _mm512_setzero_si512()); sum = _mm512_add_epi64(sum, sad); + (n, a) = unsafe { (n - 8, a.add(8)) }; } if n > 0 { let mask = _bzhi_u32(0xff, n as u32) as u8; @@ -867,7 +851,7 @@ mod reduce_sum_of_x { _mm512_reduce_add_epi64(sum) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x_v4_test() { if !crate::is_cpu_detected!("v4") { @@ -886,8 +870,8 @@ mod reduce_sum_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_x_v3(this: &[u64]) -> u32 { - use crate::emulate::emulate_mm256_reduce_add_epi64; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; let mask_0 = _mm256_set1_epi8(0x0f); @@ -896,8 +880,6 @@ mod reduce_sum_of_x { let mut n = this.len(); while n >= 4 { let x = unsafe { _mm256_loadu_si256(a.cast()) }; - a = unsafe { a.add(4) }; - n -= 4; let lo = _mm256_and_si256(x, mask_0); let hi = _mm256_and_si256(_mm256_srli_epi16(x, 4), mask_0); let res_lo = _mm256_shuffle_epi8(lut, lo); @@ -905,20 +887,24 @@ mod reduce_sum_of_x { let res = _mm256_add_epi8(res_lo, res_hi); let sad = _mm256_sad_epu8(res, _mm256_setzero_si256()); sum = _mm256_add_epi64(sum, sad); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut sum = emulate_mm256_reduce_add_epi64(sum) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - let p = x.count_ones(); - sum += p; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let lo = _mm256_and_si256(x, mask_0); + let hi = _mm256_and_si256(_mm256_srli_epi16(x, 4), mask_0); + let res_lo = _mm256_shuffle_epi8(lut, lo); + let res_hi = _mm256_shuffle_epi8(lut, hi); + let res = _mm256_add_epi8(res_lo, res_hi); + let sad = _mm256_sad_epu8(res, _mm256_setzero_si256()); + sum = _mm256_add_epi64(sum, sad); } - sum + emulate_mm256_reduce_add_epi64(sum) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x_v3_test() { if !crate::is_cpu_detected!("v3") { diff --git a/crates/simd/src/byte.rs b/crates/simd/src/byte.rs index 54081f96..4513b983 100644 --- a/crates/simd/src/byte.rs +++ b/crates/simd/src/byte.rs @@ -12,11 +12,76 @@ // // Copyright (c) 2025 TensorChord Inc. -mod reduce_sum_of_x_as_u32_y_as_u32 { +pub mod reduce_sum_of_xy { + /// # Safety + /// + /// * Don't call it. Internal use. #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_x_as_u32_y_as_u32_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + #[target_feature(enable = "avx512vnni")] + pub fn reduce_sum_of_xy_v4_avx512vnni(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let sign = _mm512_set1_epi8(-128); + let i8_1 = _mm512_set1_epi8(1); + let i16_1 = _mm512_set1_epi16(1); + let mut _0 = _mm512_setzero_si512(); + let mut _1 = _mm512_setzero_si512(); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + let y = unsafe { _mm512_loadu_epi8(b.cast()) }; + _0 = _mm512_dpbusd_epi32(_0, x, _mm512_xor_si512(sign, y)); + _1 = _mm512_add_epi32(_1, _mm512_madd_epi16(_mm512_maddubs_epi16(x, i8_1), i16_1)); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n > 0 { + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; + _0 = _mm512_dpbusd_epi32(_0, x, _mm512_xor_si512(sign, y)); + _1 = _mm512_add_epi32(_1, _mm512_madd_epi16(_mm512_maddubs_epi16(x, i8_1), i16_1)); + } + _mm512_reduce_add_epi32(_mm512_add_epi32(_0, _mm512_slli_epi32(_1, 7))) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_avx512vnni_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vnni") { + println!("test {} ... skipped (v4:avx512vnni)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4_avx512vnni(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + /// # Safety + /// + /// * Don't call it. Internal use. + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + pub fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); @@ -25,8 +90,6 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let lo = _mm512_set1_epi16(0x00ff_i16); let mut _0 = _mm512_setzero_si512(); let mut _1 = _mm512_setzero_si512(); - let mut _2 = _mm512_setzero_si512(); - let mut _3 = _mm512_setzero_si512(); while n >= 64 { let x = unsafe { _mm512_loadu_epi8(a.cast()) }; let y = unsafe { _mm512_loadu_epi8(b.cast()) }; @@ -34,15 +97,11 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let x_1 = _mm512_srli_epi16(x, 8); let y_0 = _mm512_and_si512(y, lo); let y_1 = _mm512_srli_epi16(y, 8); - let z_0 = _mm512_mullo_epi16(x_0, y_0); - let z_1 = _mm512_mullo_epi16(x_1, y_1); - a = unsafe { a.add(64) }; - b = unsafe { b.add(64) }; - n -= 64; - _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 0))); - _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 1))); - _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 0))); - _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 1))); + let z_0 = _mm512_madd_epi16(x_0, y_0); + let z_1 = _mm512_madd_epi16(x_1, y_1); + _0 = _mm512_add_epi32(_0, z_0); + _1 = _mm512_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; } if n > 0 { let mask = _bzhi_u64(0xffffffffffffffff, n as u32); @@ -52,22 +111,18 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let x_1 = _mm512_srli_epi16(x, 8); let y_0 = _mm512_and_si512(y, lo); let y_1 = _mm512_srli_epi16(y, 8); - let z_0 = _mm512_mullo_epi16(x_0, y_0); - let z_1 = _mm512_mullo_epi16(x_1, y_1); - _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 0))); - _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 1))); - _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 0))); - _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 1))); - } - let r_0 = _mm512_add_epi32(_0, _2); - let r_1 = _mm512_add_epi32(_1, _3); - let r_2 = _mm512_add_epi32(r_0, r_1); - _mm512_reduce_add_epi32(r_2) as u32 + let z_0 = _mm512_madd_epi16(x_0, y_0); + let z_1 = _mm512_madd_epi16(x_1, y_1); + _0 = _mm512_add_epi32(_0, z_0); + _1 = _mm512_add_epi32(_1, z_1); + } + _mm512_reduce_add_epi32(_mm512_add_epi32(_0, _1)) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u32_y_as_u32_v4_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_test() { use rand::Rng; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -81,7 +136,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v4(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -91,11 +146,14 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } } + /// # Safety + /// + /// * Don't call it. Internal use. #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - fn reduce_sum_of_x_as_u32_y_as_u32_v3(lhs: &[u8], rhs: &[u8]) -> u32 { - use crate::emulate::emulate_mm256_reduce_add_epi32; + pub fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::{emulate_mm256_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); @@ -104,8 +162,6 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let lo = _mm256_set1_epi16(0x00ff_i16); let mut _0 = _mm256_setzero_si256(); let mut _1 = _mm256_setzero_si256(); - let mut _2 = _mm256_setzero_si256(); - let mut _3 = _mm256_setzero_si256(); while n >= 32 { let x = unsafe { _mm256_loadu_si256(a.cast()) }; let y = unsafe { _mm256_loadu_si256(b.cast()) }; @@ -113,35 +169,32 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let x_1 = _mm256_srli_epi16(x, 8); let y_0 = _mm256_and_si256(y, lo); let y_1 = _mm256_srli_epi16(y, 8); - let z_0 = _mm256_mullo_epi16(x_0, y_0); - let z_1 = _mm256_mullo_epi16(x_1, y_1); - a = unsafe { a.add(32) }; - b = unsafe { b.add(32) }; - n -= 32; - _0 = _mm256_add_epi32(_0, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_0, 0))); - _1 = _mm256_add_epi32(_1, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_0, 1))); - _2 = _mm256_add_epi32(_2, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_1, 0))); - _3 = _mm256_add_epi32(_3, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_1, 1))); - } - let mut sum = emulate_mm256_reduce_add_epi32(_mm256_add_epi32( - _mm256_add_epi32(_0, _2), - _mm256_add_epi32(_1, _3), - )) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - sum += x as u32 * y as u32; + let z_0 = _mm256_madd_epi16(x_0, y_0); + let z_1 = _mm256_madd_epi16(x_1, y_1); + _0 = _mm256_add_epi32(_0, z_0); + _1 = _mm256_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; } - sum + if n > 0 { + let (_a, _b) = unsafe { partial_load!(32, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let x_0 = _mm256_and_si256(x, lo); + let x_1 = _mm256_srli_epi16(x, 8); + let y_0 = _mm256_and_si256(y, lo); + let y_1 = _mm256_srli_epi16(y, 8); + let z_0 = _mm256_madd_epi16(x_0, y_0); + let z_1 = _mm256_madd_epi16(x_1, y_1); + _0 = _mm256_add_epi32(_0, z_0); + _1 = _mm256_add_epi32(_1, z_1); + } + emulate_mm256_reduce_add_epi32(_mm256_add_epi32(_0, _1)) as u32 } #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u32_y_as_u32_v3_test() { + fn reduce_sum_of_xy_v3_test() { use rand::Rng; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -155,7 +208,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v3(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -165,11 +218,14 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } } + /// # Safety + /// + /// * Don't call it. Internal use. #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - fn reduce_sum_of_x_as_u32_y_as_u32_v2(lhs: &[u8], rhs: &[u8]) -> u32 { - use crate::emulate::emulate_mm_reduce_add_epi32; + pub fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::{emulate_mm_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); @@ -178,8 +234,6 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let lo = _mm_set1_epi16(0x00ff_i16); let mut _0 = _mm_setzero_si128(); let mut _1 = _mm_setzero_si128(); - let mut _2 = _mm_setzero_si128(); - let mut _3 = _mm_setzero_si128(); while n >= 16 { let x = unsafe { _mm_loadu_si128(a.cast()) }; let y = unsafe { _mm_loadu_si128(b.cast()) }; @@ -187,35 +241,32 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let x_1 = _mm_srli_epi16(x, 8); let y_0 = _mm_and_si128(y, lo); let y_1 = _mm_srli_epi16(y, 8); - let z_0 = _mm_mullo_epi16(x_0, y_0); - let z_1 = _mm_mullo_epi16(x_1, y_1); - a = unsafe { a.add(16) }; - b = unsafe { b.add(16) }; - n -= 16; - _0 = _mm_add_epi32(_0, _mm_cvtepu16_epi32(z_0)); - _1 = _mm_add_epi32(_1, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_0, z_0))); - _2 = _mm_add_epi32(_2, _mm_cvtepu16_epi32(z_1)); - _3 = _mm_add_epi32(_3, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_1, z_1))); - } - let mut sum = emulate_mm_reduce_add_epi32(_mm_add_epi32( - _mm_add_epi32(_0, _1), - _mm_add_epi32(_2, _3), - )) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - sum += x as u32 * y as u32; + let z_0 = _mm_madd_epi16(x_0, y_0); + let z_1 = _mm_madd_epi16(x_1, y_1); + _0 = _mm_add_epi32(_0, z_0); + _1 = _mm_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; } - sum + if n > 0 { + let (_a, _b) = unsafe { partial_load!(16, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm_loadu_si128(a.cast()) }; + let y = unsafe { _mm_loadu_si128(b.cast()) }; + let x_0 = _mm_and_si128(x, lo); + let x_1 = _mm_srli_epi16(x, 8); + let y_0 = _mm_and_si128(y, lo); + let y_1 = _mm_srli_epi16(y, 8); + let z_0 = _mm_madd_epi16(x_0, y_0); + let z_1 = _mm_madd_epi16(x_1, y_1); + _0 = _mm_add_epi32(_0, z_0); + _1 = _mm_add_epi32(_1, z_1); + } + emulate_mm_reduce_add_epi32(_mm_add_epi32(_0, _1)) as u32 } #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u32_y_as_u32_v2_test() { + fn reduce_sum_of_xy_v2_test() { use rand::Rng; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); @@ -229,7 +280,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v2(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v2(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -242,53 +293,26 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] - fn reduce_sum_of_x_as_u32_y_as_u32_a2(lhs: &[u8], rhs: &[u8]) -> u32 { - use core::arch::aarch64::*; - assert_eq!(lhs.len(), rhs.len()); - let mut n = lhs.len(); - let mut a = lhs.as_ptr(); - let mut b = rhs.as_ptr(); - let lo = vdupq_n_u16(0x00ff_u16); - let mut _0 = vdupq_n_u32(0); - let mut _1 = vdupq_n_u32(0); - let mut _2 = vdupq_n_u32(0); - let mut _3 = vdupq_n_u32(0); - while n >= 16 { - let x = unsafe { vld1q_u16(a.cast()) }; - let y = unsafe { vld1q_u16(b.cast()) }; - let x_0 = vandq_u16(x, lo); - let x_1 = vshrq_n_u16(x, 8); - let y_0 = vandq_u16(y, lo); - let y_1 = vshrq_n_u16(y, 8); - let z_0 = vmulq_u16(x_0, y_0); - let z_1 = vmulq_u16(x_1, y_1); - a = unsafe { a.add(16) }; - b = unsafe { b.add(16) }; - n -= 16; - _0 = vaddq_u32(_0, vmovl_u16(vget_low_u16(z_0))); - _1 = vaddq_u32(_1, vmovl_u16(vget_high_u16(z_0))); - _2 = vaddq_u32(_2, vmovl_u16(vget_low_u16(z_1))); - _3 = vaddq_u32(_3, vmovl_u16(vget_high_u16(z_1))); - } - let mut sum = vaddvq_u32(vaddq_u32(vaddq_u32(_0, _2), vaddq_u32(_1, _3))); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - sum += x as u32 * y as u32; + #[target_feature(enable = "dotprod")] + fn reduce_sum_of_xy_a2_dotprod(lhs: &[u8], rhs: &[u8]) -> u32 { + unsafe extern "C" { + #[link_name = "byte_reduce_sum_of_xy_a2_dotprod"] + unsafe fn f(n: usize, a: *const u8, b: *const u8) -> u32; } - sum + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] - fn reduce_sum_of_x_as_u32_y_as_u32_a2_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_dotprod_test() { use rand::Rng; - if !crate::is_cpu_detected!("a2") { - println!("test {} ... skipped (a2)", module_path!()); + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("dotprod") { + println!("test {} ... skipped (a2:dotprod)", module_path!()); return; } let mut rng = rand::rng(); @@ -299,7 +323,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_a2(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_a2_dotprod(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -309,194 +333,44 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] - pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { - assert_eq!(s.len(), t.len()); - let n = s.len(); - let mut result = 0; - for i in 0..n { - result += (s[i] as u32) * (t[i] as u32); - } - result - } -} - -#[inline(always)] -pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { - reduce_sum_of_x_as_u32_y_as_u32::reduce_sum_of_x_as_u32_y_as_u32(s, t) -} - -mod reduce_sum_of_x_as_u16 { - #[inline] - #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_x_as_u16_v4(this: &[u8]) -> u16 { - use crate::emulate::emulate_mm512_reduce_add_epi16; - use std::arch::x86_64::*; - let us = _mm512_set1_epi16(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm512_setzero_si512(); - while n >= 32 { - let x = unsafe { _mm256_loadu_si256(a.cast()) }; - a = unsafe { a.add(32) }; - n -= 32; - sum = _mm512_add_epi16(_mm512_and_si512(us, _mm512_cvtepi8_epi16(x)), sum); - } - if n > 0 { - let mask = _bzhi_u32(0xffffffff, n as u32); - let x = unsafe { _mm256_maskz_loadu_epi8(mask, a.cast()) }; - sum = _mm512_add_epi16(_mm512_and_si512(us, _mm512_cvtepi8_epi16(x)), sum); - } - emulate_mm512_reduce_add_epi16(sum) as u16 - } - - #[cfg(all(target_arch = "x86_64", test, not(miri)))] - #[test] - fn reduce_sum_of_x_as_u16_v4_test() { - use rand::Rng; - if !crate::is_cpu_detected!("v4") { - println!("test {} ... skipped (v4)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); - for z in 3984..4016 { - let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u16_v4(this) }; - let fallback = fallback(this); - assert_eq!(specialized, fallback); - } - } - } - - #[inline] - #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v3")] - fn reduce_sum_of_x_as_u16_v3(this: &[u8]) -> u16 { - use crate::emulate::emulate_mm256_reduce_add_epi16; - use std::arch::x86_64::*; - let us = _mm256_set1_epi16(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm256_setzero_si256(); - while n >= 16 { - let x = unsafe { _mm_loadu_si128(a.cast()) }; - a = unsafe { a.add(16) }; - n -= 16; - sum = _mm256_add_epi16(_mm256_and_si256(us, _mm256_cvtepi8_epi16(x)), sum); - } - let mut sum = emulate_mm256_reduce_add_epi16(sum) as u16; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - sum += x as u16; - } - sum - } - - #[cfg(all(target_arch = "x86_64", test))] - #[test] - fn reduce_sum_of_x_as_u16_v3_test() { - use rand::Rng; - if !crate::is_cpu_detected!("v3") { - println!("test {} ... skipped (v3)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); - for z in 3984..4016 { - let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u16_v3(this) }; - let fallback = fallback(this); - assert_eq!(specialized, fallback); - } - } - } - - #[inline] - #[cfg(target_arch = "x86_64")] - #[crate::target_cpu(enable = "v2")] - fn reduce_sum_of_x_as_u16_v2(this: &[u8]) -> u16 { - use crate::emulate::emulate_mm_reduce_add_epi16; - use std::arch::x86_64::*; - let us = _mm_set1_epi16(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = _mm_setzero_si128(); - while n >= 8 { - let x = unsafe { _mm_loadu_si64(a.cast()) }; - a = unsafe { a.add(8) }; - n -= 8; - sum = _mm_add_epi16(_mm_and_si128(us, _mm_cvtepi8_epi16(x)), sum); - } - let mut sum = emulate_mm_reduce_add_epi16(sum) as u16; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - sum += x as u16; - } - sum - } - - #[cfg(all(target_arch = "x86_64", test))] - #[test] - fn reduce_sum_of_x_as_u16_v2_test() { - use rand::Rng; - if !crate::is_cpu_detected!("v2") { - println!("test {} ... skipped (v2)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); - for z in 3984..4016 { - let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u16_v2(this) }; - let fallback = fallback(this); - assert_eq!(specialized, fallback); - } - } - } - #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] - fn reduce_sum_of_x_as_u16_a2(this: &[u8]) -> u16 { - use std::arch::aarch64::*; - let us = vdupq_n_u16(255); - let mut n = this.len(); - let mut a = this.as_ptr(); - let mut sum = vdupq_n_u16(0); - while n >= 8 { - let x = unsafe { vld1_u8(a) }; - a = unsafe { a.add(8) }; - n -= 8; - sum = vaddq_u16(vandq_u16(us, vmovl_u8(x)), sum); - } - let mut sum = vaddvq_u16(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - sum += x as u16; + fn reduce_sum_of_xy_a2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut _0 = vdupq_n_u32(0); + let mut _1 = vdupq_n_u32(0); + while n >= 16 { + let x = unsafe { vld1q_u8(a.cast()) }; + let y = unsafe { vld1q_u8(b.cast()) }; + let lo = vmull_u8(vget_low_u8(x), vget_low_u8(y)); + let hi = vmull_u8(vget_high_u8(x), vget_high_u8(y)); + _0 = vaddq_u32(_0, vpaddlq_u16(lo)); + _1 = vaddq_u32(_1, vpaddlq_u16(hi)); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; } - sum + if n > 0 { + let (_a, _b) = unsafe { partial_load!(16, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { vld1q_u8(a.cast()) }; + let y = unsafe { vld1q_u8(b.cast()) }; + let lo = vmull_u8(vget_low_u8(x), vget_low_u8(y)); + let hi = vmull_u8(vget_high_u8(x), vget_high_u8(y)); + _0 = vaddq_u32(_0, vpaddlq_u16(lo)); + _1 = vaddq_u32(_1, vpaddlq_u16(hi)); + } + vaddvq_u32(vaddq_u32(_0, _1)) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] - fn reduce_sum_of_x_as_u16_a2_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_test() { use rand::Rng; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -505,59 +379,69 @@ mod reduce_sum_of_x_as_u16 { let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); for z in 3984..4016 { - let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u16_a2(this) }; - let fallback = fallback(this); - assert_eq!(specialized, fallback); + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); } } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] - pub fn reduce_sum_of_x_as_u16(this: &[u8]) -> u16 { - let n = this.len(); - let mut sum = 0; + #[crate::multiversion(@"v4:avx512vnni", @"v4", @"v3", @"v2", @"a2:dotprod", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + assert_eq!(s.len(), t.len()); + let n = s.len(); + let mut result = 0; for i in 0..n { - sum += this[i] as u16; + result += (s[i] as u32) * (t[i] as u32); } - sum + result } } #[inline(always)] -pub fn reduce_sum_of_x_as_u16(vector: &[u8]) -> u16 { - reduce_sum_of_x_as_u16::reduce_sum_of_x_as_u16(vector) +pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + reduce_sum_of_xy::reduce_sum_of_xy(s, t) } -mod reduce_sum_of_x_as_u32 { +pub mod reduce_sum_of_x { + /// # Safety + /// + /// * Don't call it. Internal use. #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_x_as_u32_v4(this: &[u8]) -> u32 { - use std::arch::x86_64::*; - let us = _mm512_set1_epi32(255); + pub fn reduce_sum_of_x_v4(this: &[u8]) -> u32 { + use core::arch::x86_64::*; + let i8_1 = _mm512_set1_epi8(1); + let i16_1 = _mm512_set1_epi16(1); let mut n = this.len(); let mut a = this.as_ptr(); let mut sum = _mm512_setzero_si512(); - while n >= 16 { - let x = unsafe { _mm_loadu_epi8(a.cast()) }; - a = unsafe { a.add(16) }; - n -= 16; - sum = _mm512_add_epi32(_mm512_and_si512(us, _mm512_cvtepi8_epi32(x)), sum); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + sum = _mm512_add_epi32(sum, _mm512_madd_epi16(_mm512_maddubs_epi16(x, i8_1), i16_1)); + (n, a) = unsafe { (n - 64, a.add(64)) }; } if n > 0 { - let mask = _bzhi_u32(0xffff, n as u32) as u16; - let x = unsafe { _mm_maskz_loadu_epi8(mask, a.cast()) }; - sum = _mm512_add_epi32(_mm512_and_si512(us, _mm512_cvtepi8_epi32(x)), sum); + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + sum = _mm512_add_epi32(sum, _mm512_madd_epi16(_mm512_maddubs_epi16(x, i8_1), i16_1)); } _mm512_reduce_add_epi32(sum) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u32_v4_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_v4_test() { use rand::Rng; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -566,46 +450,47 @@ mod reduce_sum_of_x_as_u32 { let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random()).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_v4(this) }; + let specialized = unsafe { reduce_sum_of_x_v4(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } } } + /// # Safety + /// + /// * Don't call it. Internal use. #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - fn reduce_sum_of_x_as_u32_v3(this: &[u8]) -> u32 { - use crate::emulate::emulate_mm256_reduce_add_epi32; - use std::arch::x86_64::*; - let us = _mm256_set1_epi32(255); + pub fn reduce_sum_of_x_v3(this: &[u8]) -> u32 { + use crate::emulate::{emulate_mm256_reduce_add_epi32, partial_load}; + use core::arch::x86_64::*; + let i8_1 = _mm256_set1_epi8(1); + let i16_1 = _mm256_set1_epi16(1); let mut n = this.len(); let mut a = this.as_ptr(); let mut sum = _mm256_setzero_si256(); - while n >= 8 { - let x = unsafe { _mm_loadl_epi64(a.cast()) }; - a = unsafe { a.add(8) }; - n -= 8; - sum = _mm256_add_epi32(_mm256_and_si256(us, _mm256_cvtepi8_epi32(x)), sum); - } - let mut sum = emulate_mm256_reduce_add_epi32(sum) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - sum += x as u32; + while n >= 32 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + sum = _mm256_add_epi32(sum, _mm256_madd_epi16(_mm256_maddubs_epi16(x, i8_1), i16_1)); + (n, a) = unsafe { (n - 32, a.add(32)) }; } - sum + if n > 0 { + let (_a,) = unsafe { partial_load!(32, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + sum = _mm256_add_epi32(sum, _mm256_madd_epi16(_mm256_maddubs_epi16(x, i8_1), i16_1)); + } + emulate_mm256_reduce_add_epi32(sum) as u32 } #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u32_v3_test() { + fn reduce_sum_of_x_v3_test() { use rand::Rng; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -614,46 +499,47 @@ mod reduce_sum_of_x_as_u32 { let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random()).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_v3(this) }; + let specialized = unsafe { reduce_sum_of_x_v3(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } } } + /// # Safety + /// + /// * Don't call it. Internal use. #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - fn reduce_sum_of_x_as_u32_v2(this: &[u8]) -> u32 { - use crate::emulate::emulate_mm_reduce_add_epi32; - use std::arch::x86_64::*; - let us = _mm_set1_epi32(255); + pub fn reduce_sum_of_x_v2(this: &[u8]) -> u32 { + use crate::emulate::{emulate_mm_reduce_add_epi32, partial_load}; + use core::arch::x86_64::*; + let i8_1 = _mm_set1_epi8(1); + let i16_1 = _mm_set1_epi16(1); let mut n = this.len(); let mut a = this.as_ptr(); let mut sum = _mm_setzero_si128(); - while n >= 4 { - let x = unsafe { _mm_cvtsi32_si128(a.cast::().read_unaligned()) }; - a = unsafe { a.add(4) }; - n -= 4; - sum = _mm_add_epi32(_mm_and_si128(us, _mm_cvtepi8_epi32(x)), sum); - } - let mut sum = emulate_mm_reduce_add_epi32(sum) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - sum += x as u32; + while n >= 16 { + let x = unsafe { _mm_loadu_si128(a.cast()) }; + sum = _mm_add_epi32(sum, _mm_madd_epi16(_mm_maddubs_epi16(x, i8_1), i16_1)); + (n, a) = unsafe { (n - 16, a.add(16)) }; } - sum + if n > 0 { + let (_a,) = unsafe { partial_load!(16, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_si128(a.cast()) }; + sum = _mm_add_epi32(sum, _mm_madd_epi16(_mm_maddubs_epi16(x, i8_1), i16_1)); + } + emulate_mm_reduce_add_epi32(sum) as u32 } #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u32_v2_test() { + fn reduce_sum_of_x_v2_test() { use rand::Rng; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); @@ -662,10 +548,10 @@ mod reduce_sum_of_x_as_u32 { let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random()).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_v2(this) }; + let specialized = unsafe { reduce_sum_of_x_v2(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } @@ -675,33 +561,30 @@ mod reduce_sum_of_x_as_u32 { #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] - fn reduce_sum_of_x_as_u32_a2(this: &[u8]) -> u32 { - use std::arch::aarch64::*; + fn reduce_sum_of_x_a2(this: &[u8]) -> u32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; let mut n = this.len(); let mut a = this.as_ptr(); - let mut sum_0 = vdupq_n_u32(0); - let mut sum_1 = vdupq_n_u32(0); - while n >= 8 { - let x = unsafe { vmovl_u8(vld1_u8(a.cast())) }; - a = unsafe { a.add(8) }; - n -= 8; - sum_0 = vaddq_u32(vmovl_u16(vget_low_u16(x)), sum_0); - sum_1 = vaddq_u32(vmovl_u16(vget_high_u16(x)), sum_1); - } - let mut sum = vaddvq_u32(vaddq_u32(sum_0, sum_1)); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - sum += x as u32; + let mut sum = vdupq_n_u32(0); + while n >= 16 { + let x = unsafe { vld1q_u8(a.cast()) }; + sum = vaddq_u32(sum, vpaddlq_u16(vpaddlq_u8(x))); + (n, a) = unsafe { (n - 16, a.add(16)) }; } - sum + if n > 0 { + let (_a,) = unsafe { partial_load!(16, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_u8(a.cast()) }; + sum = vaddq_u32(sum, vpaddlq_u16(vpaddlq_u8(x))); + } + vaddvq_u32(sum) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] - fn reduce_sum_of_x_as_u32_a2_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_a2_test() { use rand::Rng; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -710,10 +593,10 @@ mod reduce_sum_of_x_as_u32 { let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4016; - let this = (0..n).map(|_| rng.random_range(0..16)).collect::>(); + let this = (0..n).map(|_| rng.random()).collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_a2(this) }; + let specialized = unsafe { reduce_sum_of_x_a2(this) }; let fallback = fallback(this); assert_eq!(specialized, fallback); } @@ -721,7 +604,7 @@ mod reduce_sum_of_x_as_u32 { } #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] - pub fn reduce_sum_of_x_as_u32(this: &[u8]) -> u32 { + pub fn reduce_sum_of_x(this: &[u8]) -> u32 { let n = this.len(); let mut sum = 0; for i in 0..n { @@ -732,6 +615,6 @@ mod reduce_sum_of_x_as_u32 { } #[inline(always)] -pub fn reduce_sum_of_x_as_u32(vector: &[u8]) -> u32 { - reduce_sum_of_x_as_u32::reduce_sum_of_x_as_u32(vector) +pub fn reduce_sum_of_x(vector: &[u8]) -> u32 { + reduce_sum_of_x::reduce_sum_of_x(vector) } diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index 4f298d47..8d11495d 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -12,6 +12,33 @@ // // Copyright (c) 2025 TensorChord Inc. +macro_rules! partial_load { + (@internal_0) => { + ::core::mem::zeroed() + }; + (@internal_0 $x:expr) => { + $x + }; + ($N:literal, $n:ident, $($p:ident $(= $x:expr)?),+) => { + ( + $( + { + ::core::hint::assert_unchecked($n < $N); + let mut result = [$crate::emulate::partial_load!(@internal_0 $($x)?); $N]; + // LLVM loves `memcpy`, which is much slower. + // So we use an explicit loop here. + for i in 0..$n { + result[i] = $p.add(i).read(); + } + result + }, + )+ + ) + }; +} + +pub(crate) use partial_load; + // VP2INTERSECT emulation. // Díez-Cañas, G. (2021). Faster-Than-Native Alternatives for x86 VP2INTERSECT // Instructions. arXiv preprint arXiv:2112.06342. @@ -19,10 +46,10 @@ #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] pub fn emulate_mm512_2intersect_epi32( - a: std::arch::x86_64::__m512i, - b: std::arch::x86_64::__m512i, -) -> (std::arch::x86_64::__mmask16, std::arch::x86_64::__mmask16) { - use std::arch::x86_64::*; + a: core::arch::x86_64::__m512i, + b: core::arch::x86_64::__m512i, +) -> (core::arch::x86_64::__mmask16, core::arch::x86_64::__mmask16) { + use core::arch::x86_64::*; let a1 = _mm512_alignr_epi32(a, a, 4); let a2 = _mm512_alignr_epi32(a, a, 8); @@ -73,8 +100,8 @@ pub fn emulate_mm512_2intersect_epi32( #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] -pub fn emulate_mm256_reduce_add_ps(mut x: std::arch::x86_64::__m256) -> f32 { - use std::arch::x86_64::*; +pub fn emulate_mm256_reduce_add_ps(mut x: core::arch::x86_64::__m256) -> f32 { + use core::arch::x86_64::*; x = _mm256_add_ps(x, _mm256_permute2f128_ps(x, x, 1)); x = _mm256_hadd_ps(x, x); x = _mm256_hadd_ps(x, x); @@ -84,52 +111,18 @@ pub fn emulate_mm256_reduce_add_ps(mut x: std::arch::x86_64::__m256) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] -pub fn emulate_mm_reduce_add_ps(mut x: std::arch::x86_64::__m128) -> f32 { - use std::arch::x86_64::*; +pub fn emulate_mm_reduce_add_ps(mut x: core::arch::x86_64::__m128) -> f32 { + use core::arch::x86_64::*; x = _mm_hadd_ps(x, x); x = _mm_hadd_ps(x, x); _mm_cvtss_f32(x) } -#[inline] -#[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v4")] -pub fn emulate_mm512_reduce_add_epi16(x: std::arch::x86_64::__m512i) -> i16 { - use std::arch::x86_64::*; - i16::wrapping_add( - _mm256_reduce_add_epi16(_mm512_castsi512_si256(x)), - _mm256_reduce_add_epi16(_mm512_extracti32x8_epi32(x, 1)), - ) -} - -#[inline] -#[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v3")] -pub fn emulate_mm256_reduce_add_epi16(mut x: std::arch::x86_64::__m256i) -> i16 { - use std::arch::x86_64::*; - x = _mm256_add_epi16(x, _mm256_permute2f128_si256(x, x, 1)); - x = _mm256_hadd_epi16(x, x); - x = _mm256_hadd_epi16(x, x); - let x = _mm256_cvtsi256_si32(x); - i16::wrapping_add(x as i16, (x >> 16) as i16) -} - -#[inline] -#[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v2")] -pub fn emulate_mm_reduce_add_epi16(mut x: std::arch::x86_64::__m128i) -> i16 { - use std::arch::x86_64::*; - x = _mm_hadd_epi16(x, x); - x = _mm_hadd_epi16(x, x); - let x = _mm_cvtsi128_si32(x); - i16::wrapping_add(x as i16, (x >> 16) as i16) -} - #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] -pub fn emulate_mm256_reduce_add_epi32(mut x: std::arch::x86_64::__m256i) -> i32 { - use std::arch::x86_64::*; +pub fn emulate_mm256_reduce_add_epi32(mut x: core::arch::x86_64::__m256i) -> i32 { + use core::arch::x86_64::*; x = _mm256_add_epi32(x, _mm256_permute2f128_si256(x, x, 1)); x = _mm256_hadd_epi32(x, x); x = _mm256_hadd_epi32(x, x); @@ -139,8 +132,8 @@ pub fn emulate_mm256_reduce_add_epi32(mut x: std::arch::x86_64::__m256i) -> i32 #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] -pub fn emulate_mm_reduce_add_epi32(mut x: std::arch::x86_64::__m128i) -> i32 { - use std::arch::x86_64::*; +pub fn emulate_mm_reduce_add_epi32(mut x: core::arch::x86_64::__m128i) -> i32 { + use core::arch::x86_64::*; x = _mm_hadd_epi32(x, x); x = _mm_hadd_epi32(x, x); _mm_cvtsi128_si32(x) @@ -149,9 +142,9 @@ pub fn emulate_mm_reduce_add_epi32(mut x: std::arch::x86_64::__m128i) -> i32 { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] -pub fn emulate_mm256_reduce_min_ps(x: std::arch::x86_64::__m256) -> f32 { +pub fn emulate_mm256_reduce_min_ps(x: core::arch::x86_64::__m256) -> f32 { use crate::aligned::Aligned16; - use std::arch::x86_64::*; + use core::arch::x86_64::*; let lo = _mm256_castps256_ps128(x); let hi = _mm256_extractf128_ps(x, 1); let min = _mm_min_ps(lo, hi); @@ -165,9 +158,9 @@ pub fn emulate_mm256_reduce_min_ps(x: std::arch::x86_64::__m256) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] -pub fn emulate_mm_reduce_min_ps(x: std::arch::x86_64::__m128) -> f32 { +pub fn emulate_mm_reduce_min_ps(x: core::arch::x86_64::__m128) -> f32 { use crate::aligned::Aligned16; - use std::arch::x86_64::*; + use core::arch::x86_64::*; let min = x; let mut x = Aligned16([0.0f32; 4]); unsafe { @@ -179,9 +172,9 @@ pub fn emulate_mm_reduce_min_ps(x: std::arch::x86_64::__m128) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] -pub fn emulate_mm256_reduce_max_ps(x: std::arch::x86_64::__m256) -> f32 { +pub fn emulate_mm256_reduce_max_ps(x: core::arch::x86_64::__m256) -> f32 { use crate::aligned::Aligned16; - use std::arch::x86_64::*; + use core::arch::x86_64::*; let lo = _mm256_castps256_ps128(x); let hi = _mm256_extractf128_ps(x, 1); let max = _mm_max_ps(lo, hi); @@ -195,9 +188,9 @@ pub fn emulate_mm256_reduce_max_ps(x: std::arch::x86_64::__m256) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] -pub fn emulate_mm_reduce_max_ps(x: std::arch::x86_64::__m128) -> f32 { +pub fn emulate_mm_reduce_max_ps(x: core::arch::x86_64::__m128) -> f32 { use crate::aligned::Aligned16; - use std::arch::x86_64::*; + use core::arch::x86_64::*; let max = x; let mut x = Aligned16([0.0f32; 4]); unsafe { @@ -209,8 +202,8 @@ pub fn emulate_mm_reduce_max_ps(x: std::arch::x86_64::__m128) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] -pub fn emulate_mm256_reduce_add_epi64(mut x: std::arch::x86_64::__m256i) -> i64 { - use std::arch::x86_64::*; +pub fn emulate_mm256_reduce_add_epi64(mut x: core::arch::x86_64::__m256i) -> i64 { + use core::arch::x86_64::*; x = _mm256_add_epi64(x, _mm256_permute2f128_si256(x, x, 1)); _mm256_extract_epi64(x, 0) + _mm256_extract_epi64(x, 1) } diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index a0221584..cd6c7bcf 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -55,7 +55,7 @@ mod scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - use std::arch::x86_64::*; + use core::arch::x86_64::*; #[inline] #[crate::target_cpu(enable = "v4")] @@ -157,7 +157,7 @@ mod scan { result } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn scan_v4_test() { if !crate::is_cpu_detected!("v4") { @@ -187,7 +187,7 @@ mod scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - use std::arch::x86_64::*; + use core::arch::x86_64::*; #[inline] #[crate::target_cpu(enable = "v3")] @@ -260,7 +260,7 @@ mod scan { result } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn scan_v3_test() { if !crate::is_cpu_detected!("v3") { @@ -289,7 +289,7 @@ mod scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut accu_0 = _mm_setzero_si128(); let mut accu_1 = _mm_setzero_si128(); @@ -333,7 +333,7 @@ mod scan { result } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn scan_v2_test() { if !crate::is_cpu_detected!("v2") { @@ -362,7 +362,7 @@ mod scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - use std::arch::aarch64::*; + use core::arch::aarch64::*; let mut accu_0 = vdupq_n_u16(0); let mut accu_1 = vdupq_n_u16(0); @@ -405,8 +405,9 @@ mod scan { result } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] + #[cfg_attr(miri, ignore)] fn scan_a2_test() { if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -435,7 +436,7 @@ mod scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - use std::arch::s390x::*; + use core::arch::s390x::*; use std::mem::transmute; use {vector_unsigned_char as u8x16, vector_unsigned_short as u16x8}; @@ -518,7 +519,7 @@ mod scan { assert_eq!(code.len(), lut.len()); let n = code.len(); - use std::arch::powerpc64::*; + use core::arch::powerpc64::*; use std::mem::transmute; use {vector_unsigned_char as u8x16, vector_unsigned_short as u16x8}; diff --git a/crates/simd/src/floating_f16.rs b/crates/simd/src/floating_f16.rs index 136d6ca3..fb32c13c 100644 --- a/crates/simd/src/floating_f16.rs +++ b/crates/simd/src/floating_f16.rs @@ -244,21 +244,20 @@ mod reduce_sum_of_xy { #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512fp16")] pub fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { - assert!(lhs.len() == rhs.len()); - unsafe { - unsafe extern "C" { - unsafe fn fp16_reduce_sum_of_xy_v4_avx512fp16( - a: *const (), - b: *const (), - n: usize, - ) -> f32; - } - fp16_reduce_sum_of_xy_v4_avx512fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + unsafe extern "C" { + #[link_name = "fp16_reduce_sum_of_xy_v4_avx512fp16"] + unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v4_avx512fp16_test() { use rand::Rng; const EPSILON: f32 = 2.0; @@ -293,30 +292,29 @@ mod reduce_sum_of_xy { #[crate::target_cpu(enable = "v4")] pub fn reduce_sum_of_xy_v4(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut xy = _mm512_setzero_ps(); + let mut sum = _mm512_setzero_ps(); while n >= 16 { let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; let y = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())) }; - a = unsafe { a.add(16) }; - b = unsafe { b.add(16) }; - n -= 16; - xy = _mm512_fmadd_ps(x, y, xy); + sum = _mm512_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; } if n > 0 { let mask = _bzhi_u32(0xffff, n as u32) as u16; let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; let y = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())) }; - xy = _mm512_fmadd_ps(x, y, xy); + sum = _mm512_fmadd_ps(x, y, sum); } - _mm512_reduce_add_ps(xy) + _mm512_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v4_test() { use rand::Rng; const EPSILON: f32 = 2.0; @@ -346,36 +344,32 @@ mod reduce_sum_of_xy { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] pub fn reduce_sum_of_xy_v3(lhs: &[f16], rhs: &[f16]) -> f32 { - use crate::emulate::emulate_mm256_reduce_add_ps; + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut xy = _mm256_setzero_ps(); + let mut sum = _mm256_setzero_ps(); while n >= 8 { let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; - xy = _mm256_fmadd_ps(x, y, xy); + sum = _mm256_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } - let mut xy = emulate_mm256_reduce_add_ps(xy); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read()._to_f32() }; - let y = unsafe { b.read()._to_f32() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - xy += x * y; + if n > 0 { + let (_a, _b) = unsafe { partial_load!(8, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; + sum = _mm256_fmadd_ps(x, y, sum); } - xy + emulate_mm256_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v3_test() { use rand::Rng; const EPSILON: f32 = 2.0; @@ -406,30 +400,28 @@ mod reduce_sum_of_xy { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "fp16")] - pub fn reduce_sum_of_xy_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { - assert!(lhs.len() == rhs.len()); - unsafe { - unsafe extern "C" { - unsafe fn fp16_reduce_sum_of_xy_a2_fp16( - a: *const (), - b: *const (), - n: usize, - ) -> f32; - } - fp16_reduce_sum_of_xy_a2_fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.512")] + fn reduce_sum_of_xy_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { + unsafe extern "C" { + #[link_name = "fp16_reduce_sum_of_xy_a3_512"] + unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test))] #[test] - fn reduce_sum_of_xy_a2_fp16_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a3_512_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { - println!("test {} ... skipped (a2:fp16)", module_path!()); + if !crate::is_cpu_detected!("a3.512") { + println!("test {} ... skipped (a3.512)", module_path!()); return; } let mut rng = rand::rng(); @@ -444,7 +436,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_a2_fp16(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_a3_512(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -455,26 +447,29 @@ mod reduce_sum_of_xy { } #[inline] - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - #[crate::target_cpu(enable = "a3.512")] - pub fn reduce_sum_of_xy_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { - assert!(lhs.len() == rhs.len()); - unsafe { - unsafe extern "C" { - unsafe fn fp16_reduce_sum_of_xy_a3_512(a: *const (), b: *const (), n: usize) - -> f32; - } - fp16_reduce_sum_of_xy_a3_512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + pub fn reduce_sum_of_xy_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + unsafe extern "C" { + #[link_name = "fp16_reduce_sum_of_xy_a2_fp16"] + unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } } - #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] - fn reduce_sum_of_xy_a3_512_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_fp16_test() { use rand::Rng; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("a3.512") { - println!("test {} ... skipped (a3.512)", module_path!()); + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -489,7 +484,7 @@ mod reduce_sum_of_xy { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_a3_512(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_a2_fp16(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -503,11 +498,11 @@ mod reduce_sum_of_xy { pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); - let mut xy = 0.0f32; + let mut sum = 0.0f32; for i in 0..n { - xy += lhs[i]._to_f32() * rhs[i]._to_f32(); + sum += lhs[i]._to_f32() * rhs[i]._to_f32(); } - xy + sum } } @@ -519,21 +514,20 @@ mod reduce_sum_of_d2 { #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512fp16")] pub fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { - assert!(lhs.len() == rhs.len()); - unsafe { - unsafe extern "C" { - unsafe fn fp16_reduce_sum_of_d2_v4_avx512fp16( - a: *const (), - b: *const (), - n: usize, - ) -> f32; - } - fp16_reduce_sum_of_d2_v4_avx512fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + unsafe extern "C" { + #[link_name = "fp16_reduce_sum_of_d2_v4_avx512fp16"] + unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_v4_avx512fp16_test() { use rand::Rng; const EPSILON: f32 = 6.4; @@ -568,32 +562,31 @@ mod reduce_sum_of_d2 { #[crate::target_cpu(enable = "v4")] pub fn reduce_sum_of_d2_v4(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut d2 = _mm512_setzero_ps(); + let mut sum = _mm512_setzero_ps(); while n >= 16 { let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; let y = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())) }; - a = unsafe { a.add(16) }; - b = unsafe { b.add(16) }; - n -= 16; let d = _mm512_sub_ps(x, y); - d2 = _mm512_fmadd_ps(d, d, d2); + sum = _mm512_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; } if n > 0 { let mask = _bzhi_u32(0xffff, n as u32) as u16; let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; let y = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())) }; let d = _mm512_sub_ps(x, y); - d2 = _mm512_fmadd_ps(d, d, d2); + sum = _mm512_fmadd_ps(d, d, sum); } - _mm512_reduce_add_ps(d2) + _mm512_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_v4_test() { use rand::Rng; const EPSILON: f32 = 2.0; @@ -627,38 +620,35 @@ mod reduce_sum_of_d2 { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] pub fn reduce_sum_of_d2_v3(lhs: &[f16], rhs: &[f16]) -> f32 { - use crate::emulate::emulate_mm256_reduce_add_ps; + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut d2 = _mm256_setzero_ps(); + let mut sum = _mm256_setzero_ps(); while n >= 8 { let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; let d = _mm256_sub_ps(x, y); - d2 = _mm256_fmadd_ps(d, d, d2); + sum = _mm256_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } - let mut d2 = emulate_mm256_reduce_add_ps(d2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read()._to_f32() }; - let y = unsafe { b.read()._to_f32() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - let d = x - y; - d2 += d * d; + if n > 0 { + let (_a, _b) = unsafe { partial_load!(8, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; + let d = _mm256_sub_ps(x, y); + sum = _mm256_fmadd_ps(d, d, sum); } - d2 + emulate_mm256_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_v3_test() { use rand::Rng; const EPSILON: f32 = 2.0; @@ -689,30 +679,28 @@ mod reduce_sum_of_d2 { } #[inline] - #[cfg(target_arch = "aarch64")] - #[crate::target_cpu(enable = "a2")] - #[target_feature(enable = "fp16")] - pub fn reduce_sum_of_d2_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { - assert!(lhs.len() == rhs.len()); - unsafe { - unsafe extern "C" { - unsafe fn fp16_reduce_sum_of_d2_a2_fp16( - a: *const (), - b: *const (), - n: usize, - ) -> f32; - } - fp16_reduce_sum_of_d2_a2_fp16(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.512")] + fn reduce_sum_of_d2_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { + unsafe extern "C" { + #[link_name = "fp16_reduce_sum_of_d2_a3_512"] + unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", target_endian = "little", test))] #[test] - fn reduce_sum_of_d2_a2_fp16_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_a3_512_test() { use rand::Rng; const EPSILON: f32 = 6.4; - if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { - println!("test {} ... skipped (a2:fp16)", module_path!()); + if !crate::is_cpu_detected!("a3.512") { + println!("test {} ... skipped (a3.512)", module_path!()); return; } let mut rng = rand::rng(); @@ -727,7 +715,7 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_a2_fp16(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_d2_a3_512(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -738,26 +726,29 @@ mod reduce_sum_of_d2 { } #[inline] - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - #[crate::target_cpu(enable = "a3.512")] - pub fn reduce_sum_of_d2_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { - assert!(lhs.len() == rhs.len()); - unsafe { - unsafe extern "C" { - unsafe fn fp16_reduce_sum_of_d2_a3_512(a: *const (), b: *const (), n: usize) - -> f32; - } - fp16_reduce_sum_of_d2_a3_512(lhs.as_ptr().cast(), rhs.as_ptr().cast(), lhs.len()) + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + pub fn reduce_sum_of_d2_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + unsafe extern "C" { + #[link_name = "fp16_reduce_sum_of_d2_a2_fp16"] + unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr().cast(); + let b = rhs.as_ptr().cast(); + unsafe { f(n, a, b) } } - #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] - fn reduce_sum_of_d2_a3_512_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_a2_fp16_test() { use rand::Rng; const EPSILON: f32 = 6.4; - if !crate::is_cpu_detected!("a3.512") { - println!("test {} ... skipped (a3.512)", module_path!()); + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); return; } let mut rng = rand::rng(); @@ -772,7 +763,7 @@ mod reduce_sum_of_d2 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_a3_512(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_d2_a2_fp16(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( (specialized - fallback).abs() < EPSILON, @@ -786,12 +777,12 @@ mod reduce_sum_of_d2 { pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); - let mut d2 = 0.0_f32; + let mut sum = 0.0_f32; for i in 0..n { let d = lhs[i]._to_f32() - rhs[i]._to_f32(); - d2 += d * d; + sum += d * d; } - d2 + sum } } @@ -810,11 +801,11 @@ mod reduce_sum_of_xy_sparse { assert_eq!(ridx.len(), rval.len()); let (mut lp, ln) = (0, lidx.len()); let (mut rp, rn) = (0, ridx.len()); - let mut xy = 0.0f32; + let mut sum = 0.0f32; while lp < ln && rp < rn { match Ord::cmp(&lidx[lp], &ridx[rp]) { Ordering::Equal => { - xy += lval[lp]._to_f32() * rval[rp]._to_f32(); + sum += lval[lp]._to_f32() * rval[rp]._to_f32(); lp += 1; rp += 1; } @@ -826,7 +817,7 @@ mod reduce_sum_of_xy_sparse { } } } - xy + sum } } @@ -845,32 +836,32 @@ mod reduce_sum_of_d2_sparse { assert_eq!(ridx.len(), rval.len()); let (mut lp, ln) = (0, lidx.len()); let (mut rp, rn) = (0, ridx.len()); - let mut d2 = 0.0f32; + let mut sum = 0.0f32; while lp < ln && rp < rn { match Ord::cmp(&lidx[lp], &ridx[rp]) { Ordering::Equal => { let d = lval[lp]._to_f32() - rval[rp]._to_f32(); - d2 += d * d; + sum += d * d; lp += 1; rp += 1; } Ordering::Less => { - d2 += lval[lp]._to_f32() * lval[lp]._to_f32(); + sum += lval[lp]._to_f32() * lval[lp]._to_f32(); lp += 1; } Ordering::Greater => { - d2 += rval[rp]._to_f32() * rval[rp]._to_f32(); + sum += rval[rp]._to_f32() * rval[rp]._to_f32(); rp += 1; } } } for i in lp..ln { - d2 += lval[i]._to_f32() * lval[i]._to_f32(); + sum += lval[i]._to_f32() * lval[i]._to_f32(); } for i in rp..rn { - d2 += rval[i]._to_f32() * rval[i]._to_f32(); + sum += rval[i]._to_f32() * rval[i]._to_f32(); } - d2 + sum } } diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs index 8e31dce1..994f4e7d 100644 --- a/crates/simd/src/floating_f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -165,15 +165,14 @@ mod reduce_sum_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_x_v4(this: &[f32]) -> f32 { - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut sum = _mm512_setzero_ps(); while n >= 16 { let x = unsafe { _mm512_loadu_ps(a) }; - a = unsafe { a.add(16) }; - n -= 16; sum = _mm512_add_ps(x, sum); + (n, a) = unsafe { (n - 16, a.add(16)) }; } if n > 0 { let mask = _bzhi_u32(0xffff, n as u32) as u16; @@ -183,7 +182,7 @@ mod reduce_sum_of_x { _mm512_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x_v4_test() { use rand::Rng; @@ -214,32 +213,28 @@ mod reduce_sum_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_x_v3(this: &[f32]) -> f32 { - use crate::emulate::emulate_mm256_reduce_add_ps; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut sum = _mm256_setzero_ps(); while n >= 8 { let x = unsafe { _mm256_loadu_ps(a) }; - a = unsafe { a.add(8) }; - n -= 8; sum = _mm256_add_ps(x, sum); + (n, a) = unsafe { (n - 8, a.add(8)) }; } if n >= 4 { let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; - a = unsafe { a.add(4) }; - n -= 4; sum = _mm256_add_ps(x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut sum = emulate_mm256_reduce_add_ps(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - sum += x; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + sum = _mm256_add_ps(x, sum); } - sum + emulate_mm256_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -273,26 +268,23 @@ mod reduce_sum_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] fn reduce_sum_of_x_v2(this: &[f32]) -> f32 { - use crate::emulate::emulate_mm_reduce_add_ps; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut sum = _mm_setzero_ps(); while n >= 4 { let x = unsafe { _mm_loadu_ps(a) }; - a = unsafe { a.add(4) }; - n -= 4; sum = _mm_add_ps(x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut sum = emulate_mm_reduce_add_ps(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - sum += x; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + sum = _mm_add_ps(x, sum); } - sum + emulate_mm_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -326,29 +318,28 @@ mod reduce_sum_of_x { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_x_a2(this: &[f32]) -> f32 { - use std::arch::aarch64::*; + use crate::emulate::partial_load; + use core::arch::aarch64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut sum = vdupq_n_f32(0.0); while n >= 4 { let x = unsafe { vld1q_f32(a) }; - a = unsafe { a.add(4) }; - n -= 4; sum = vaddq_f32(x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut sum = vaddvq_f32(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - sum += x; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + sum = vaddq_f32(x, sum); } - sum + vaddvq_f32(sum) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_x_a2_test() { use rand::Rng; const EPSILON: f32 = 0.008; @@ -374,46 +365,7 @@ mod reduce_sum_of_x { } } - #[inline] - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - #[crate::target_cpu(enable = "a3.256")] - fn reduce_sum_of_x_a3_256(this: &[f32]) -> f32 { - unsafe { - unsafe extern "C" { - unsafe fn fp32_reduce_sum_of_x_a3_256(this: *const f32, n: usize) -> f32; - } - fp32_reduce_sum_of_x_a3_256(this.as_ptr(), this.len()) - } - } - - #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] - #[test] - fn reduce_sum_of_x_a3_256_test() { - use rand::Rng; - const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("a3.256") { - println!("test {} ... skipped (a3.256)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let this = (0..n) - .map(|_| rng.random_range(-1.0..=1.0)) - .collect::>(); - for z in 3984..4016 { - let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x_a3_256(this) }; - let fallback = fallback(this); - assert!( - (specialized - fallback).abs() < EPSILON, - "specialized = {specialized}, fallback = {fallback}." - ); - } - } - } - - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -429,16 +381,15 @@ mod reduce_sum_of_abs_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_abs_x_v4(this: &[f32]) -> f32 { - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut sum = _mm512_setzero_ps(); while n >= 16 { let x = unsafe { _mm512_loadu_ps(a) }; let abs_x = _mm512_abs_ps(x); - a = unsafe { a.add(16) }; - n -= 16; sum = _mm512_add_ps(abs_x, sum); + (n, a) = unsafe { (n - 16, a.add(16)) }; } if n > 0 { let mask = _bzhi_u32(0xffff, n as u32) as u16; @@ -449,7 +400,7 @@ mod reduce_sum_of_abs_x { _mm512_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_abs_x_v4_test() { use rand::Rng; @@ -480,8 +431,8 @@ mod reduce_sum_of_abs_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_abs_x_v3(this: &[f32]) -> f32 { - use crate::emulate::emulate_mm256_reduce_add_ps; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; let abs = _mm256_castsi256_ps(_mm256_srli_epi32(_mm256_set1_epi32(-1), 1)); let mut n = this.len(); let mut a = this.as_ptr(); @@ -489,27 +440,23 @@ mod reduce_sum_of_abs_x { while n >= 8 { let x = unsafe { _mm256_loadu_ps(a) }; let abs_x = _mm256_and_ps(abs, x); - a = unsafe { a.add(8) }; - n -= 8; sum = _mm256_add_ps(abs_x, sum); + (n, a) = unsafe { (n - 8, a.add(8)) }; } if n >= 4 { let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; let abs_x = _mm256_and_ps(abs, x); - a = unsafe { a.add(4) }; - n -= 4; sum = _mm256_add_ps(abs_x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut sum = emulate_mm256_reduce_add_ps(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let abs_x = x.abs(); - a = unsafe { a.add(1) }; - n -= 1; - sum += abs_x; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let abs_x = _mm256_and_ps(abs, x); + sum = _mm256_add_ps(abs_x, sum); } - sum + emulate_mm256_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -543,8 +490,8 @@ mod reduce_sum_of_abs_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] fn reduce_sum_of_abs_x_v2(this: &[f32]) -> f32 { - use crate::emulate::emulate_mm_reduce_add_ps; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; let abs = _mm_castsi128_ps(_mm_srli_epi32(_mm_set1_epi32(-1), 1)); let mut n = this.len(); let mut a = this.as_ptr(); @@ -552,20 +499,17 @@ mod reduce_sum_of_abs_x { while n >= 4 { let x = unsafe { _mm_loadu_ps(a) }; let abs_x = _mm_and_ps(abs, x); - a = unsafe { a.add(4) }; - n -= 4; sum = _mm_add_ps(abs_x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut sum = emulate_mm_reduce_add_ps(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let abs_x = x.abs(); - a = unsafe { a.add(1) }; - n -= 1; - sum += abs_x; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + let abs_x = _mm_and_ps(abs, x); + sum = _mm_add_ps(abs_x, sum); } - sum + emulate_mm_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -599,31 +543,30 @@ mod reduce_sum_of_abs_x { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_abs_x_a2(this: &[f32]) -> f32 { - use std::arch::aarch64::*; + use crate::emulate::partial_load; + use core::arch::aarch64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut sum = vdupq_n_f32(0.0); while n >= 4 { let x = unsafe { vld1q_f32(a) }; let abs_x = vabsq_f32(x); - a = unsafe { a.add(4) }; - n -= 4; sum = vaddq_f32(abs_x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut sum = vaddvq_f32(sum); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let abs_x = x.abs(); - a = unsafe { a.add(1) }; - n -= 1; - sum += abs_x; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + let abs_x = vabsq_f32(x); + sum = vaddq_f32(abs_x, sum); } - sum + vaddvq_f32(sum) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_abs_x_a2_test() { use rand::Rng; const EPSILON: f32 = 0.009; @@ -649,46 +592,7 @@ mod reduce_sum_of_abs_x { } } - #[inline] - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - #[crate::target_cpu(enable = "a3.256")] - fn reduce_sum_of_abs_x_a3_256(this: &[f32]) -> f32 { - unsafe { - unsafe extern "C" { - unsafe fn fp32_reduce_sum_of_abs_x_a3_256(this: *const f32, n: usize) -> f32; - } - fp32_reduce_sum_of_abs_x_a3_256(this.as_ptr(), this.len()) - } - } - - #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] - #[test] - fn reduce_sum_of_abs_x_a3_256_test() { - use rand::Rng; - const EPSILON: f32 = 0.009; - if !crate::is_cpu_detected!("a3.256") { - println!("test {} ... skipped (a3.256)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let this = (0..n) - .map(|_| rng.random_range(-1.0..=1.0)) - .collect::>(); - for z in 3984..4016 { - let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_abs_x_a3_256(this) }; - let fallback = fallback(this); - assert!( - (specialized - fallback).abs() < EPSILON, - "specialized = {specialized}, fallback = {fallback}." - ); - } - } - } - - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { let n = this.len(); let mut sum = 0.0f32; @@ -704,25 +608,24 @@ mod reduce_sum_of_x2 { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_x2_v4(this: &[f32]) -> f32 { - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); - let mut x2 = _mm512_setzero_ps(); + let mut sum = _mm512_setzero_ps(); while n >= 16 { let x = unsafe { _mm512_loadu_ps(a) }; - a = unsafe { a.add(16) }; - n -= 16; - x2 = _mm512_fmadd_ps(x, x, x2); + sum = _mm512_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 16, a.add(16)) }; } if n > 0 { let mask = _bzhi_u32(0xffff, n as u32) as u16; let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; - x2 = _mm512_fmadd_ps(x, x, x2); + sum = _mm512_fmadd_ps(x, x, sum); } - _mm512_reduce_add_ps(x2) + _mm512_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x2_v4_test() { use rand::Rng; @@ -753,32 +656,28 @@ mod reduce_sum_of_x2 { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_x2_v3(this: &[f32]) -> f32 { - use crate::emulate::emulate_mm256_reduce_add_ps; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); - let mut x2 = _mm256_setzero_ps(); + let mut sum = _mm256_setzero_ps(); while n >= 8 { let x = unsafe { _mm256_loadu_ps(a) }; - a = unsafe { a.add(8) }; - n -= 8; - x2 = _mm256_fmadd_ps(x, x, x2); + sum = _mm256_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 8, a.add(8)) }; } if n >= 4 { let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; - a = unsafe { a.add(4) }; - n -= 4; - x2 = _mm256_fmadd_ps(x, x, x2); + sum = _mm256_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut x2 = emulate_mm256_reduce_add_ps(x2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - x2 += x * x; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + sum = _mm256_fmadd_ps(x, x, sum); } - x2 + emulate_mm256_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -813,26 +712,23 @@ mod reduce_sum_of_x2 { #[crate::target_cpu(enable = "v2")] #[target_feature(enable = "fma")] fn reduce_sum_of_x2_v2_fma(this: &[f32]) -> f32 { - use crate::emulate::emulate_mm_reduce_add_ps; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); - let mut x2 = _mm_setzero_ps(); + let mut sum = _mm_setzero_ps(); while n >= 4 { let x = unsafe { _mm_loadu_ps(a) }; - a = unsafe { a.add(4) }; - n -= 4; - x2 = _mm_fmadd_ps(x, x, x2); + sum = _mm_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut x2 = emulate_mm_reduce_add_ps(x2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - x2 += x * x; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + sum = _mm_fmadd_ps(x, x, sum); } - x2 + emulate_mm_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -866,29 +762,28 @@ mod reduce_sum_of_x2 { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_x2_a2(this: &[f32]) -> f32 { - use std::arch::aarch64::*; + use crate::emulate::partial_load; + use core::arch::aarch64::*; let mut n = this.len(); let mut a = this.as_ptr(); - let mut x2 = vdupq_n_f32(0.0); + let mut sum = vdupq_n_f32(0.0); while n >= 4 { let x = unsafe { vld1q_f32(a) }; - a = unsafe { a.add(4) }; - n -= 4; - x2 = vfmaq_f32(x2, x, x); + sum = vfmaq_f32(sum, x, x); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut x2 = vaddvq_f32(x2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - x2 += x * x; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + sum = vfmaq_f32(sum, x, x); } - x2 + vaddvq_f32(sum) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_x2_a2_test() { use rand::Rng; const EPSILON: f32 = 0.008; @@ -914,53 +809,14 @@ mod reduce_sum_of_x2 { } } - #[inline] - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - #[crate::target_cpu(enable = "a3.256")] - fn reduce_sum_of_x2_a3_256(this: &[f32]) -> f32 { - unsafe { - unsafe extern "C" { - unsafe fn fp32_reduce_sum_of_x2_a3_256(this: *const f32, n: usize) -> f32; - } - fp32_reduce_sum_of_x2_a3_256(this.as_ptr(), this.len()) - } - } - - #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] - #[test] - fn reduce_sum_of_x2_a3_256_test() { - use rand::Rng; - const EPSILON: f32 = 0.008; - if !crate::is_cpu_detected!("a3.256") { - println!("test {} ... skipped (a3.256)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let this = (0..n) - .map(|_| rng.random_range(-1.0..=1.0)) - .collect::>(); - for z in 3984..4016 { - let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x2_a3_256(this) }; - let fallback = fallback(this); - assert!( - (specialized - fallback).abs() < EPSILON, - "specialized = {specialized}, fallback = {fallback}." - ); - } - } - } - - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_x2(this: &[f32]) -> f32 { let n = this.len(); - let mut x2 = 0.0f32; + let mut sum = 0.0f32; for i in 0..n { - x2 += this[i] * this[i]; + sum += this[i] * this[i]; } - x2 + sum } } @@ -969,17 +825,16 @@ mod reduce_min_max_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] fn reduce_min_max_of_x_v4(this: &[f32]) -> (f32, f32) { - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut min = _mm512_set1_ps(f32::INFINITY); let mut max = _mm512_set1_ps(f32::NEG_INFINITY); while n >= 16 { let x = unsafe { _mm512_loadu_ps(a) }; - a = unsafe { a.add(16) }; - n -= 16; min = _mm512_min_ps(x, min); max = _mm512_max_ps(x, max); + (n, a) = unsafe { (n - 16, a.add(16)) }; } if n > 0 { let mask = _bzhi_u32(0xffff, n as u32) as u16; @@ -992,8 +847,9 @@ mod reduce_min_max_of_x { (min, max) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_min_max_of_x_v4_test() { use rand::Rng; if !crate::is_cpu_detected!("v4") { @@ -1021,30 +877,37 @@ mod reduce_min_max_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn reduce_min_max_of_x_v3(this: &[f32]) -> (f32, f32) { - use crate::emulate::{emulate_mm256_reduce_max_ps, emulate_mm256_reduce_min_ps}; - use std::arch::x86_64::*; + use crate::emulate::{ + emulate_mm256_reduce_max_ps, emulate_mm256_reduce_min_ps, partial_load, + }; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut min = _mm256_set1_ps(f32::INFINITY); let mut max = _mm256_set1_ps(f32::NEG_INFINITY); while n >= 8 { let x = unsafe { _mm256_loadu_ps(a) }; - a = unsafe { a.add(8) }; - n -= 8; min = _mm256_min_ps(x, min); max = _mm256_max_ps(x, max); + (n, a) = unsafe { (n - 8, a.add(8)) }; } - let mut min = emulate_mm256_reduce_min_ps(min); - let mut max = emulate_mm256_reduce_max_ps(max); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - min = x.min(min); - max = x.max(max); + if n >= 4 { + let x = _mm256_setr_m128(unsafe { _mm_loadu_ps(a) }, _mm_set1_ps(f32::NAN)); + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - (min, max) + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a = f32::NAN) }; + (a,) = (_a.as_ptr(),); + let x = _mm256_setr_m128(unsafe { _mm_loadu_ps(a) }, _mm_set1_ps(f32::NAN)); + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + } + ( + emulate_mm256_reduce_min_ps(min), + emulate_mm256_reduce_max_ps(max), + ) } #[cfg(all(target_arch = "x86_64", test))] @@ -1076,30 +939,26 @@ mod reduce_min_max_of_x { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] fn reduce_min_max_of_x_v2(this: &[f32]) -> (f32, f32) { - use crate::emulate::{emulate_mm_reduce_max_ps, emulate_mm_reduce_min_ps}; - use std::arch::x86_64::*; + use crate::emulate::{emulate_mm_reduce_max_ps, emulate_mm_reduce_min_ps, partial_load}; + use core::arch::x86_64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut min = _mm_set1_ps(f32::INFINITY); let mut max = _mm_set1_ps(f32::NEG_INFINITY); while n >= 4 { let x = unsafe { _mm_loadu_ps(a) }; - a = unsafe { a.add(4) }; - n -= 4; min = _mm_min_ps(x, min); max = _mm_max_ps(x, max); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut min = emulate_mm_reduce_min_ps(min); - let mut max = emulate_mm_reduce_max_ps(max); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - min = x.min(min); - max = x.max(max); + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a = f32::NAN) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + min = _mm_min_ps(x, min); + max = _mm_max_ps(x, max); } - (min, max) + (emulate_mm_reduce_min_ps(min), emulate_mm_reduce_max_ps(max)) } #[cfg(all(target_arch = "x86_64", test))] @@ -1127,82 +986,35 @@ mod reduce_min_max_of_x { } } - #[inline] - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - #[crate::target_cpu(enable = "a3.256")] - fn reduce_min_max_of_x_a3_256(this: &[f32]) -> (f32, f32) { - let mut min = f32::INFINITY; - let mut max = -f32::INFINITY; - unsafe { - unsafe extern "C" { - unsafe fn fp32_reduce_min_max_of_x_a3_256( - this: *const f32, - n: usize, - out_min: &mut f32, - out_max: &mut f32, - ); - } - fp32_reduce_min_max_of_x_a3_256(this.as_ptr(), this.len(), &mut min, &mut max); - } - (min, max) - } - - #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] - #[test] - fn reduce_min_max_of_x_a3_256_test() { - use rand::Rng; - if !crate::is_cpu_detected!("a3.256") { - println!("test {} ... skipped (a3.256)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 200; - let mut x = (0..n) - .map(|_| rng.random_range(-1.0..=1.0)) - .collect::>(); - (x[0], x[1]) = (f32::NAN, -f32::NAN); - for z in 50..200 { - let x = &x[..z]; - let specialized = unsafe { reduce_min_max_of_x_a3_256(x) }; - let fallback = fallback(x); - assert_eq!(specialized.0, fallback.0,); - assert_eq!(specialized.1, fallback.1,); - } - } - } - #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_min_max_of_x_a2(this: &[f32]) -> (f32, f32) { - use std::arch::aarch64::*; + use crate::emulate::partial_load; + use core::arch::aarch64::*; let mut n = this.len(); let mut a = this.as_ptr(); let mut min = vdupq_n_f32(f32::INFINITY); let mut max = vdupq_n_f32(f32::NEG_INFINITY); while n >= 4 { let x = unsafe { vld1q_f32(a) }; - a = unsafe { a.add(4) }; - n -= 4; min = vminnmq_f32(x, min); max = vmaxnmq_f32(x, max); + (n, a) = unsafe { (n - 4, a.add(4)) }; } - let mut min = vminnmvq_f32(min); - let mut max = vmaxnmvq_f32(max); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - a = unsafe { a.add(1) }; - n -= 1; - min = x.min(min); - max = x.max(max); + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a = f32::NAN) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + min = vminnmq_f32(x, min); + max = vmaxnmq_f32(x, max); } - (min, max) + (vminnmvq_f32(min), vmaxnmvq_f32(max)) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_min_max_of_x_a2_test() { use rand::Rng; if !crate::is_cpu_detected!("a2") { @@ -1226,7 +1038,7 @@ mod reduce_min_max_of_x { } } - #[crate::multiversion(@"v4", @"v3", @"v2", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { let mut min = f32::INFINITY; let mut max = f32::NEG_INFINITY; @@ -1245,29 +1057,27 @@ mod reduce_sum_of_xy { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_xy_v4(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut xy = _mm512_setzero_ps(); + let mut sum = _mm512_setzero_ps(); while n >= 16 { let x = unsafe { _mm512_loadu_ps(a) }; let y = unsafe { _mm512_loadu_ps(b) }; - a = unsafe { a.add(16) }; - b = unsafe { b.add(16) }; - n -= 16; - xy = _mm512_fmadd_ps(x, y, xy); + sum = _mm512_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; } if n > 0 { let mask = _bzhi_u32(0xffff, n as u32) as u16; let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; let y = unsafe { _mm512_maskz_loadu_ps(mask, b) }; - xy = _mm512_fmadd_ps(x, y, xy); + sum = _mm512_fmadd_ps(x, y, sum); } - _mm512_reduce_add_ps(xy) + _mm512_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xy_v4_test() { use rand::Rng; @@ -1302,40 +1112,33 @@ mod reduce_sum_of_xy { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_xy_v3(lhs: &[f32], rhs: &[f32]) -> f32 { - use crate::emulate::emulate_mm256_reduce_add_ps; + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut xy = _mm256_setzero_ps(); + let mut sum = _mm256_setzero_ps(); while n >= 8 { let x = unsafe { _mm256_loadu_ps(a) }; let y = unsafe { _mm256_loadu_ps(b) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; - xy = _mm256_fmadd_ps(x, y, xy); + sum = _mm256_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n >= 4 { let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; - xy = _mm256_fmadd_ps(x, y, xy); + sum = _mm256_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut xy = emulate_mm256_reduce_add_ps(xy); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - xy += x * y; + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; + sum = _mm256_fmadd_ps(x, y, sum); } - xy + emulate_mm256_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -1374,32 +1177,27 @@ mod reduce_sum_of_xy { #[crate::target_cpu(enable = "v2")] #[target_feature(enable = "fma")] fn reduce_sum_of_xy_v2_fma(lhs: &[f32], rhs: &[f32]) -> f32 { - use crate::emulate::emulate_mm_reduce_add_ps; + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut xy = _mm_setzero_ps(); + let mut sum = _mm_setzero_ps(); while n >= 4 { let x = unsafe { _mm_loadu_ps(a) }; let y = unsafe { _mm_loadu_ps(b) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; - xy = _mm_fmadd_ps(x, y, xy); + sum = _mm_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut xy = emulate_mm_reduce_add_ps(xy); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - xy += x * y; + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm_loadu_ps(a) }; + let y = unsafe { _mm_loadu_ps(b) }; + sum = _mm_fmadd_ps(x, y, sum); } - xy + emulate_mm_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -1433,39 +1231,58 @@ mod reduce_sum_of_xy { } } + #[inline] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] + fn reduce_sum_of_xy_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { + unsafe extern "C" { + #[link_name = "fp32_reduce_sum_of_xy_a3_256"] + unsafe fn f(n: usize, a: *const f32, b: *const f32) -> f32; + } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a3_256_test() { + // + } + #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_xy_a2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); - use std::arch::aarch64::*; + use crate::emulate::partial_load; + use core::arch::aarch64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut xy = vdupq_n_f32(0.0); + let mut sum = vdupq_n_f32(0.0); while n >= 4 { let x = unsafe { vld1q_f32(a) }; let y = unsafe { vld1q_f32(b) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; - xy = vfmaq_f32(xy, x, y); + sum = vfmaq_f32(sum, x, y); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut xy = vaddvq_f32(xy); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - xy += x * y; + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { vld1q_f32(a) }; + let y = unsafe { vld1q_f32(b) }; + sum = vfmaq_f32(sum, x, y); } - xy + vaddvq_f32(sum) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a2_test() { use rand::Rng; const EPSILON: f32 = 0.004; @@ -1495,63 +1312,15 @@ mod reduce_sum_of_xy { } } - #[inline] - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - #[crate::target_cpu(enable = "a3.256")] - fn reduce_sum_of_xy_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { - assert!(lhs.len() == rhs.len()); - unsafe { - unsafe extern "C" { - unsafe fn fp32_reduce_sum_of_xy_a3_256( - a: *const f32, - b: *const f32, - n: usize, - ) -> f32; - } - fp32_reduce_sum_of_xy_a3_256(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) - } - } - - #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] - #[test] - fn reduce_sum_of_xy_a3_256_test() { - use rand::Rng; - const EPSILON: f32 = 0.004; - if !crate::is_cpu_detected!("a3.256") { - println!("test {} ... skipped (a3.256)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let lhs = (0..n) - .map(|_| rng.random_range(-1.0..=1.0)) - .collect::>(); - let rhs = (0..n) - .map(|_| rng.random_range(-1.0..=1.0)) - .collect::>(); - for z in 3984..4016 { - let lhs = &lhs[..z]; - let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_xy_a3_256(lhs, rhs) }; - let fallback = fallback(lhs, rhs); - assert!( - (specialized - fallback).abs() < EPSILON, - "specialized = {specialized}, fallback = {fallback}." - ); - } - } - } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_xy(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); - let mut xy = 0.0f32; + let mut sum = 0.0f32; for i in 0..n { - xy += lhs[i] * rhs[i]; + sum += lhs[i] * rhs[i]; } - xy + sum } } @@ -1561,31 +1330,29 @@ mod reduce_sum_of_d2 { #[crate::target_cpu(enable = "v4")] fn reduce_sum_of_d2_v4(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut d2 = _mm512_setzero_ps(); + let mut sum = _mm512_setzero_ps(); while n >= 16 { let x = unsafe { _mm512_loadu_ps(a) }; let y = unsafe { _mm512_loadu_ps(b) }; - a = unsafe { a.add(16) }; - b = unsafe { b.add(16) }; - n -= 16; let d = _mm512_sub_ps(x, y); - d2 = _mm512_fmadd_ps(d, d, d2); + sum = _mm512_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; } if n > 0 { let mask = _bzhi_u32(0xffff, n as u32) as u16; let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; let y = unsafe { _mm512_maskz_loadu_ps(mask, b) }; let d = _mm512_sub_ps(x, y); - d2 = _mm512_fmadd_ps(d, d, d2); + sum = _mm512_fmadd_ps(d, d, sum); } - _mm512_reduce_add_ps(d2) + _mm512_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_d2_v4_test() { use rand::Rng; @@ -1620,43 +1387,36 @@ mod reduce_sum_of_d2 { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn reduce_sum_of_d2_v3(lhs: &[f32], rhs: &[f32]) -> f32 { - use crate::emulate::emulate_mm256_reduce_add_ps; + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut d2 = _mm256_setzero_ps(); + let mut sum = _mm256_setzero_ps(); while n >= 8 { let x = unsafe { _mm256_loadu_ps(a) }; let y = unsafe { _mm256_loadu_ps(b) }; - a = unsafe { a.add(8) }; - b = unsafe { b.add(8) }; - n -= 8; let d = _mm256_sub_ps(x, y); - d2 = _mm256_fmadd_ps(d, d, d2); + sum = _mm256_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; } if n >= 4 { let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; let d = _mm256_sub_ps(x, y); - d2 = _mm256_fmadd_ps(d, d, d2); + sum = _mm256_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut d2 = emulate_mm256_reduce_add_ps(d2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - let d = x - y; - d2 += d * d; + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; + let d = _mm256_sub_ps(x, y); + sum = _mm256_fmadd_ps(d, d, sum); } - d2 + emulate_mm256_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -1695,34 +1455,29 @@ mod reduce_sum_of_d2 { #[crate::target_cpu(enable = "v2")] #[target_feature(enable = "fma")] fn reduce_sum_of_d2_v2_fma(lhs: &[f32], rhs: &[f32]) -> f32 { - use crate::emulate::emulate_mm_reduce_add_ps; + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; assert!(lhs.len() == rhs.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut d2 = _mm_setzero_ps(); + let mut sum = _mm_setzero_ps(); while n >= 4 { let x = unsafe { _mm_loadu_ps(a) }; let y = unsafe { _mm_loadu_ps(b) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; let d = _mm_sub_ps(x, y); - d2 = _mm_fmadd_ps(d, d, d2); + sum = _mm_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut d2 = emulate_mm_reduce_add_ps(d2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - let d = x - y; - d2 += d * d; + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm_loadu_ps(a) }; + let y = unsafe { _mm_loadu_ps(b) }; + let d = _mm_sub_ps(x, y); + sum = _mm_fmadd_ps(d, d, sum); } - d2 + emulate_mm_reduce_add_ps(sum) } #[cfg(all(target_arch = "x86_64", test))] @@ -1756,41 +1511,60 @@ mod reduce_sum_of_d2 { } } + #[inline] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] + fn reduce_sum_of_d2_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { + unsafe extern "C" { + #[link_name = "fp32_reduce_sum_of_d2_a3_256"] + unsafe fn f(n: usize, a: *const f32, b: *const f32) -> f32; + } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_a3_256_test() { + // + } + #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn reduce_sum_of_d2_a2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); - use std::arch::aarch64::*; + use crate::emulate::partial_load; + use core::arch::aarch64::*; let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let mut d2 = vdupq_n_f32(0.0); + let mut sum = vdupq_n_f32(0.0); while n >= 4 { let x = unsafe { vld1q_f32(a) }; let y = unsafe { vld1q_f32(b) }; - a = unsafe { a.add(4) }; - b = unsafe { b.add(4) }; - n -= 4; let d = vsubq_f32(x, y); - d2 = vfmaq_f32(d2, d, d); + sum = vfmaq_f32(sum, d, d); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; } - let mut d2 = vaddvq_f32(d2); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - let d = x - y; - d2 += d * d; + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { vld1q_f32(a) }; + let y = unsafe { vld1q_f32(b) }; + let d = vsubq_f32(x, y); + sum = vfmaq_f32(sum, d, d); } - d2 + vaddvq_f32(sum) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_a2_test() { use rand::Rng; const EPSILON: f32 = 0.02; @@ -1820,64 +1594,16 @@ mod reduce_sum_of_d2 { } } - #[inline] - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - #[crate::target_cpu(enable = "a3.256")] - fn reduce_sum_of_d2_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { - assert!(lhs.len() == rhs.len()); - unsafe { - unsafe extern "C" { - unsafe fn fp32_reduce_sum_of_d2_a3_256( - a: *const f32, - b: *const f32, - n: usize, - ) -> f32; - } - fp32_reduce_sum_of_d2_a3_256(lhs.as_ptr(), rhs.as_ptr(), lhs.len()) - } - } - - #[cfg(all(target_arch = "aarch64", target_endian = "little", test, not(miri)))] - #[test] - fn reduce_sum_of_d2_a3_256_test() { - use rand::Rng; - const EPSILON: f32 = 0.02; - if !crate::is_cpu_detected!("a3.256") { - println!("test {} ... skipped (a3.256)", module_path!()); - return; - } - let mut rng = rand::rng(); - for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { - let n = 4016; - let lhs = (0..n) - .map(|_| rng.random_range(-1.0..=1.0)) - .collect::>(); - let rhs = (0..n) - .map(|_| rng.random_range(-1.0..=1.0)) - .collect::>(); - for z in 3984..4016 { - let lhs = &lhs[..z]; - let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_d2_a3_256(lhs, rhs) }; - let fallback = fallback(lhs, rhs); - assert!( - (specialized - fallback).abs() < EPSILON, - "specialized = {specialized}, fallback = {fallback}." - ); - } - } - } - #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_d2(lhs: &[f32], rhs: &[f32]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); - let mut d2 = 0.0f32; + let mut sum = 0.0f32; for i in 0..n { let d = lhs[i] - rhs[i]; - d2 += d * d; + sum += d * d; } - d2 + sum } } @@ -1893,15 +1619,15 @@ mod reduce_sum_of_xy_sparse { let (mut rp, rn) = (0, ri.len()); let (li, lv) = (li.as_ptr(), lv.as_ptr()); let (ri, rv) = (ri.as_ptr(), rv.as_ptr()); - use std::arch::x86_64::*; - let mut xy = _mm512_setzero_ps(); + use core::arch::x86_64::*; + let mut sum = _mm512_setzero_ps(); while lp + 16 <= ln && rp + 16 <= rn { let lx = unsafe { _mm512_loadu_epi32(li.add(lp).cast()) }; let rx = unsafe { _mm512_loadu_epi32(ri.add(rp).cast()) }; let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))) }; let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))) }; - xy = _mm512_fmadd_ps(lv, rv, xy); + sum = _mm512_fmadd_ps(lv, rv, sum); let lt = unsafe { li.add(lp + 16 - 1).read() }; let rt = unsafe { ri.add(rp + 16 - 1).read() }; lp += (lt <= rt) as usize * 16; @@ -1919,17 +1645,18 @@ mod reduce_sum_of_xy_sparse { let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))) }; let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))) }; - xy = _mm512_fmadd_ps(lv, rv, xy); + sum = _mm512_fmadd_ps(lv, rv, sum); let lt = unsafe { li.add(lp + lw - 1).read() }; let rt = unsafe { ri.add(rp + rw - 1).read() }; lp += (lt <= rt) as usize * lw; rp += (lt >= rt) as usize * rw; } - _mm512_reduce_add_ps(xy) + _mm512_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_sparse_v4_test() { use rand::Rng; const EPSILON: f32 = 0.000001; @@ -1977,11 +1704,11 @@ mod reduce_sum_of_xy_sparse { assert_eq!(ridx.len(), rval.len()); let (mut lp, ln) = (0, lidx.len()); let (mut rp, rn) = (0, ridx.len()); - let mut xy = 0.0f32; + let mut sum = 0.0f32; while lp < ln && rp < rn { match Ord::cmp(&lidx[lp], &ridx[rp]) { Ordering::Equal => { - xy += lval[lp] * rval[rp]; + sum += lval[lp] * rval[rp]; lp += 1; rp += 1; } @@ -1993,7 +1720,7 @@ mod reduce_sum_of_xy_sparse { } } } - xy + sum } } @@ -2009,8 +1736,8 @@ mod reduce_sum_of_d2_sparse { let (mut rp, rn) = (0, ri.len()); let (li, lv) = (li.as_ptr(), lv.as_ptr()); let (ri, rv) = (ri.as_ptr(), rv.as_ptr()); - use std::arch::x86_64::*; - let mut d2 = _mm512_setzero_ps(); + use core::arch::x86_64::*; + let mut sum = _mm512_setzero_ps(); while lp + 16 <= ln && rp + 16 <= rn { let lx = unsafe { _mm512_loadu_epi32(li.add(lp).cast()) }; let rx = unsafe { _mm512_loadu_epi32(ri.add(rp).cast()) }; @@ -2018,9 +1745,9 @@ mod reduce_sum_of_d2_sparse { let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))) }; let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))) }; let d = _mm512_sub_ps(lv, rv); - d2 = _mm512_fmadd_ps(d, d, d2); - d2 = _mm512_sub_ps(d2, _mm512_mul_ps(lv, lv)); - d2 = _mm512_sub_ps(d2, _mm512_mul_ps(rv, rv)); + sum = _mm512_fmadd_ps(d, d, sum); + sum = _mm512_sub_ps(sum, _mm512_mul_ps(lv, lv)); + sum = _mm512_sub_ps(sum, _mm512_mul_ps(rv, rv)); let lt = unsafe { li.add(lp + 16 - 1).read() }; let rt = unsafe { ri.add(rp + 16 - 1).read() }; lp += (lt <= rt) as usize * 16; @@ -2039,9 +1766,9 @@ mod reduce_sum_of_d2_sparse { let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))) }; let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))) }; let d = _mm512_sub_ps(lv, rv); - d2 = _mm512_fmadd_ps(d, d, d2); - d2 = _mm512_sub_ps(d2, _mm512_mul_ps(lv, lv)); - d2 = _mm512_sub_ps(d2, _mm512_mul_ps(rv, rv)); + sum = _mm512_fmadd_ps(d, d, sum); + sum = _mm512_sub_ps(sum, _mm512_mul_ps(lv, lv)); + sum = _mm512_sub_ps(sum, _mm512_mul_ps(rv, rv)); let lt = unsafe { li.add(lp + lw - 1).read() }; let rt = unsafe { ri.add(rp + rw - 1).read() }; lp += (lt <= rt) as usize * lw; @@ -2051,35 +1778,36 @@ mod reduce_sum_of_d2_sparse { let mut lp = 0; while lp + 16 <= ln { let d = unsafe { _mm512_loadu_ps(lv.add(lp)) }; - d2 = _mm512_fmadd_ps(d, d, d2); + sum = _mm512_fmadd_ps(d, d, sum); lp += 16; } if lp < ln { let lw = ln - lp; let lm = _bzhi_u32(0xffff, lw as _) as u16; let d = unsafe { _mm512_maskz_loadu_ps(lm, lv.add(lp)) }; - d2 = _mm512_fmadd_ps(d, d, d2); + sum = _mm512_fmadd_ps(d, d, sum); } } { let mut rp = 0; while rp + 16 <= rn { let d = unsafe { _mm512_loadu_ps(rv.add(rp)) }; - d2 = _mm512_fmadd_ps(d, d, d2); + sum = _mm512_fmadd_ps(d, d, sum); rp += 16; } if rp < rn { let rw = rn - rp; let rm = _bzhi_u32(0xffff, rw as _) as u16; let d = unsafe { _mm512_maskz_loadu_ps(rm, rv.add(rp)) }; - d2 = _mm512_fmadd_ps(d, d, d2); + sum = _mm512_fmadd_ps(d, d, sum); } } - _mm512_reduce_add_ps(d2) + _mm512_reduce_add_ps(sum) } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_sparse_v4_test() { use rand::Rng; const EPSILON: f32 = 0.0004; @@ -2127,32 +1855,32 @@ mod reduce_sum_of_d2_sparse { assert_eq!(ridx.len(), rval.len()); let (mut lp, ln) = (0, lidx.len()); let (mut rp, rn) = (0, ridx.len()); - let mut d2 = 0.0f32; + let mut sum = 0.0f32; while lp < ln && rp < rn { match Ord::cmp(&lidx[lp], &ridx[rp]) { Ordering::Equal => { let d = lval[lp] - rval[rp]; - d2 += d * d; + sum += d * d; lp += 1; rp += 1; } Ordering::Less => { - d2 += lval[lp] * lval[lp]; + sum += lval[lp] * lval[lp]; lp += 1; } Ordering::Greater => { - d2 += rval[rp] * rval[rp]; + sum += rval[rp] * rval[rp]; rp += 1; } } } for i in lp..ln { - d2 += lval[i] * lval[i]; + sum += lval[i] * lval[i]; } for i in rp..rn { - d2 += rval[i] * rval[i]; + sum += rval[i] * rval[i]; } - d2 + sum } } diff --git a/crates/simd/src/halfbyte.rs b/crates/simd/src/halfbyte.rs index 008c7fcd..7831ea0a 100644 --- a/crates/simd/src/halfbyte.rs +++ b/crates/simd/src/halfbyte.rs @@ -12,51 +12,33 @@ // // Copyright (c) 2025 TensorChord Inc. -mod reduce_sum_of_x_as_u32_y_as_u32 { +pub mod reduce_sum_of_xy { + /// # Safety + /// + /// * Don't call it. Internal use. #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - fn reduce_sum_of_x_as_u32_y_as_u32_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + #[target_feature(enable = "avx512vnni")] + pub fn reduce_sum_of_xy_v4_avx512vnni(lhs: &[u8], rhs: &[u8]) -> u32 { use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let lo = _mm512_set1_epi16(0x000f_i16); + let lo = _mm512_set1_epi8(0x0f_i8); let mut _0 = _mm512_setzero_si512(); let mut _1 = _mm512_setzero_si512(); - let mut _2 = _mm512_setzero_si512(); - let mut _3 = _mm512_setzero_si512(); - let mut _4 = _mm512_setzero_si512(); - let mut _5 = _mm512_setzero_si512(); - let mut _6 = _mm512_setzero_si512(); - let mut _7 = _mm512_setzero_si512(); while n >= 64 { let x = unsafe { _mm512_loadu_epi8(a.cast()) }; let y = unsafe { _mm512_loadu_epi8(b.cast()) }; let x_0 = _mm512_and_si512(x, lo); let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); - let x_2 = _mm512_and_si512(_mm512_srli_epi16(x, 8), lo); - let x_3 = _mm512_srli_epi16(x, 12); let y_0 = _mm512_and_si512(y, lo); let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); - let y_2 = _mm512_and_si512(_mm512_srli_epi16(y, 8), lo); - let y_3 = _mm512_srli_epi16(y, 12); - let z_0 = _mm512_mullo_epi16(x_0, y_0); - let z_1 = _mm512_mullo_epi16(x_1, y_1); - let z_2 = _mm512_mullo_epi16(x_2, y_2); - let z_3 = _mm512_mullo_epi16(x_3, y_3); - a = unsafe { a.add(64) }; - b = unsafe { b.add(64) }; - n -= 64; - _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 0))); - _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 1))); - _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 0))); - _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 1))); - _4 = _mm512_add_epi32(_4, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_2, 0))); - _5 = _mm512_add_epi32(_5, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_2, 1))); - _6 = _mm512_add_epi32(_6, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_3, 0))); - _7 = _mm512_add_epi32(_7, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_3, 1))); + _0 = _mm512_dpbusd_epi32(_0, x_0, y_0); + _1 = _mm512_dpbusd_epi32(_1, x_1, y_1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; } if n > 0 { let mask = _bzhi_u64(0xffffffffffffffff, n as u32); @@ -64,38 +46,90 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; let x_0 = _mm512_and_si512(x, lo); let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); - let x_2 = _mm512_and_si512(_mm512_srli_epi16(x, 8), lo); - let x_3 = _mm512_srli_epi16(x, 12); let y_0 = _mm512_and_si512(y, lo); let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); - let y_2 = _mm512_and_si512(_mm512_srli_epi16(y, 8), lo); - let y_3 = _mm512_srli_epi16(y, 12); - let z_0 = _mm512_mullo_epi16(x_0, y_0); - let z_1 = _mm512_mullo_epi16(x_1, y_1); - let z_2 = _mm512_mullo_epi16(x_2, y_2); - let z_3 = _mm512_mullo_epi16(x_3, y_3); - _0 = _mm512_add_epi32(_0, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 0))); - _1 = _mm512_add_epi32(_1, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_0, 1))); - _2 = _mm512_add_epi32(_2, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 0))); - _3 = _mm512_add_epi32(_3, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_1, 1))); - _4 = _mm512_add_epi32(_4, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_2, 0))); - _5 = _mm512_add_epi32(_5, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_2, 1))); - _6 = _mm512_add_epi32(_6, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_3, 0))); - _7 = _mm512_add_epi32(_7, _mm512_cvtepu16_epi32(_mm512_extracti32x8_epi32(z_3, 1))); + _0 = _mm512_dpbusd_epi32(_0, x_0, y_0); + _1 = _mm512_dpbusd_epi32(_1, x_1, y_1); } - let r_0 = _mm512_add_epi32(_0, _4); - let r_1 = _mm512_add_epi32(_1, _5); - let r_2 = _mm512_add_epi32(_2, _6); - let r_3 = _mm512_add_epi32(_3, _7); - let r_4 = _mm512_add_epi32(r_0, r_2); - let r_5 = _mm512_add_epi32(r_1, r_3); - let r_6 = _mm512_add_epi32(r_4, r_5); - _mm512_reduce_add_epi32(r_6) as u32 + _mm512_reduce_add_epi32(_mm512_add_epi32(_0, _1)) as u32 } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u32_y_as_u32_v4_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_avx512vnni_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vnni") { + println!("test {} ... skipped (v4:avx512vnni)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4_avx512vnni(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + /// # Safety + /// + /// * Don't call it. Internal use. + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + pub fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let i16_1 = _mm512_set1_epi16(1); + let lo = _mm512_set1_epi8(0x0f_i8); + let mut _0 = _mm512_setzero_si512(); + let mut _1 = _mm512_setzero_si512(); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + let y = unsafe { _mm512_loadu_epi8(b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); + let z_0 = _mm512_madd_epi16(_mm512_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm512_madd_epi16(_mm512_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm512_add_epi32(_0, z_0); + _1 = _mm512_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n > 0 { + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); + let z_0 = _mm512_madd_epi16(_mm512_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm512_madd_epi16(_mm512_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm512_add_epi32(_0, z_0); + _1 = _mm512_add_epi32(_1, z_1); + } + _mm512_reduce_add_epi32(_mm512_add_epi32(_0, _1)) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_test() { use rand::Rng; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -109,7 +143,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v4(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -119,72 +153,56 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } } + /// # Safety + /// + /// * Don't call it. Internal use. #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - fn reduce_sum_of_x_as_u32_y_as_u32_v3(lhs: &[u8], rhs: &[u8]) -> u32 { - use crate::emulate::emulate_mm256_reduce_add_epi32; + pub fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::{emulate_mm256_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let lo = _mm256_set1_epi16(0x000f_i16); + let i16_1 = _mm256_set1_epi16(1); + let lo = _mm256_set1_epi8(0x0f_i8); let mut _0 = _mm256_setzero_si256(); let mut _1 = _mm256_setzero_si256(); - let mut _2 = _mm256_setzero_si256(); - let mut _3 = _mm256_setzero_si256(); - let mut _4 = _mm256_setzero_si256(); - let mut _5 = _mm256_setzero_si256(); - let mut _6 = _mm256_setzero_si256(); - let mut _7 = _mm256_setzero_si256(); while n >= 32 { let x = unsafe { _mm256_loadu_si256(a.cast()) }; let y = unsafe { _mm256_loadu_si256(b.cast()) }; let x_0 = _mm256_and_si256(x, lo); let x_1 = _mm256_and_si256(_mm256_srli_epi16(x, 4), lo); - let x_2 = _mm256_and_si256(_mm256_srli_epi16(x, 8), lo); - let x_3 = _mm256_srli_epi16(x, 12); let y_0 = _mm256_and_si256(y, lo); let y_1 = _mm256_and_si256(_mm256_srli_epi16(y, 4), lo); - let y_2 = _mm256_and_si256(_mm256_srli_epi16(y, 8), lo); - let y_3 = _mm256_srli_epi16(y, 12); - let z_0 = _mm256_mullo_epi16(x_0, y_0); - let z_1 = _mm256_mullo_epi16(x_1, y_1); - let z_2 = _mm256_mullo_epi16(x_2, y_2); - let z_3 = _mm256_mullo_epi16(x_3, y_3); - a = unsafe { a.add(32) }; - b = unsafe { b.add(32) }; - n -= 32; - _0 = _mm256_add_epi32(_0, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_0, 0))); - _1 = _mm256_add_epi32(_1, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_0, 1))); - _2 = _mm256_add_epi32(_2, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_1, 0))); - _3 = _mm256_add_epi32(_3, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_1, 1))); - _4 = _mm256_add_epi32(_4, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_2, 0))); - _5 = _mm256_add_epi32(_5, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_2, 1))); - _6 = _mm256_add_epi32(_6, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_3, 0))); - _7 = _mm256_add_epi32(_7, _mm256_cvtepu16_epi32(_mm256_extracti128_si256(z_3, 1))); + let z_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm256_madd_epi16(_mm256_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm256_add_epi32(_0, z_0); + _1 = _mm256_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; } - let mut sum = emulate_mm256_reduce_add_epi32(_mm256_add_epi32( - _mm256_add_epi32(_mm256_add_epi32(_0, _4), _mm256_add_epi32(_1, _5)), - _mm256_add_epi32(_mm256_add_epi32(_2, _6), _mm256_add_epi32(_3, _7)), - )) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - sum += (x as u32 & 0xf) * (y as u32 & 0xf); - sum += (x as u32 >> 4) * (y as u32 >> 4); + if n > 0 { + let (_a, _b) = unsafe { partial_load!(32, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let x_0 = _mm256_and_si256(x, lo); + let x_1 = _mm256_and_si256(_mm256_srli_epi16(x, 4), lo); + let y_0 = _mm256_and_si256(y, lo); + let y_1 = _mm256_and_si256(_mm256_srli_epi16(y, 4), lo); + let z_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm256_madd_epi16(_mm256_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm256_add_epi32(_0, z_0); + _1 = _mm256_add_epi32(_1, z_1); } - sum + emulate_mm256_reduce_add_epi32(_mm256_add_epi32(_0, _1)) as u32 } #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u32_y_as_u32_v3_test() { + fn reduce_sum_of_xy_v3_test() { use rand::Rng; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -198,7 +216,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v3(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -208,72 +226,56 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } } + /// # Safety + /// + /// * Don't call it. Internal use. #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - fn reduce_sum_of_x_as_u32_y_as_u32_v2(lhs: &[u8], rhs: &[u8]) -> u32 { - use crate::emulate::emulate_mm_reduce_add_epi32; + pub fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::{emulate_mm_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let lo = _mm_set1_epi16(0x000f_i16); + let i16_1 = _mm_set1_epi16(1); + let lo = _mm_set1_epi8(0x0f_i8); let mut _0 = _mm_setzero_si128(); let mut _1 = _mm_setzero_si128(); - let mut _2 = _mm_setzero_si128(); - let mut _3 = _mm_setzero_si128(); - let mut _4 = _mm_setzero_si128(); - let mut _5 = _mm_setzero_si128(); - let mut _6 = _mm_setzero_si128(); - let mut _7 = _mm_setzero_si128(); while n >= 16 { let x = unsafe { _mm_loadu_si128(a.cast()) }; let y = unsafe { _mm_loadu_si128(b.cast()) }; let x_0 = _mm_and_si128(x, lo); let x_1 = _mm_and_si128(_mm_srli_epi16(x, 4), lo); - let x_2 = _mm_and_si128(_mm_srli_epi16(x, 8), lo); - let x_3 = _mm_srli_epi16(x, 12); let y_0 = _mm_and_si128(y, lo); let y_1 = _mm_and_si128(_mm_srli_epi16(y, 4), lo); - let y_2 = _mm_and_si128(_mm_srli_epi16(y, 8), lo); - let y_3 = _mm_srli_epi16(y, 12); - let z_0 = _mm_mullo_epi16(x_0, y_0); - let z_1 = _mm_mullo_epi16(x_1, y_1); - let z_2 = _mm_mullo_epi16(x_2, y_2); - let z_3 = _mm_mullo_epi16(x_3, y_3); - a = unsafe { a.add(16) }; - b = unsafe { b.add(16) }; - n -= 16; - _0 = _mm_add_epi32(_0, _mm_cvtepu16_epi32(z_0)); - _1 = _mm_add_epi32(_1, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_0, z_0))); - _2 = _mm_add_epi32(_2, _mm_cvtepu16_epi32(z_1)); - _3 = _mm_add_epi32(_3, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_1, z_1))); - _4 = _mm_add_epi32(_4, _mm_cvtepu16_epi32(z_2)); - _5 = _mm_add_epi32(_5, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_2, z_2))); - _6 = _mm_add_epi32(_6, _mm_cvtepu16_epi32(z_3)); - _7 = _mm_add_epi32(_7, _mm_cvtepu16_epi32(_mm_unpackhi_epi64(z_3, z_3))); + let z_0 = _mm_madd_epi16(_mm_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm_madd_epi16(_mm_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm_add_epi32(_0, z_0); + _1 = _mm_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; } - let mut sum = emulate_mm_reduce_add_epi32(_mm_add_epi32( - _mm_add_epi32(_mm_add_epi32(_0, _4), _mm_add_epi32(_1, _5)), - _mm_add_epi32(_mm_add_epi32(_2, _6), _mm_add_epi32(_3, _7)), - )) as u32; - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - sum += (x as u32 & 0xf) * (y as u32 & 0xf); - sum += (x as u32 >> 4) * (y as u32 >> 4); + if n > 0 { + let (_a, _b) = unsafe { partial_load!(16, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm_loadu_si128(a.cast()) }; + let y = unsafe { _mm_loadu_si128(b.cast()) }; + let x_0 = _mm_and_si128(x, lo); + let x_1 = _mm_and_si128(_mm_srli_epi16(x, 4), lo); + let y_0 = _mm_and_si128(y, lo); + let y_1 = _mm_and_si128(_mm_srli_epi16(y, 4), lo); + let z_0 = _mm_madd_epi16(_mm_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm_madd_epi16(_mm_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm_add_epi32(_0, z_0); + _1 = _mm_add_epi32(_1, z_1); } - sum + emulate_mm_reduce_add_epi32(_mm_add_epi32(_0, _1)) as u32 } #[cfg(all(target_arch = "x86_64", test))] #[test] - fn reduce_sum_of_x_as_u32_y_as_u32_v2_test() { + fn reduce_sum_of_xy_v2_test() { use rand::Rng; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); @@ -287,7 +289,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_v2(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_v2(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -300,68 +302,97 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] - fn reduce_sum_of_x_as_u32_y_as_u32_a2(lhs: &[u8], rhs: &[u8]) -> u32 { + #[target_feature(enable = "dotprod")] + fn reduce_sum_of_xy_a2_dotprod(lhs: &[u8], rhs: &[u8]) -> u32 { + unsafe extern "C" { + #[link_name = "halfbyte_reduce_sum_of_xy_a2_dotprod"] + unsafe fn f(n: usize, a: *const u8, b: *const u8) -> u32; + } + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_dotprod_test() { + use rand::Rng; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("dotprod") { + println!("test {} ... skipped (a2:dotprod)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2_dotprod(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_xy_a2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::partial_load; use core::arch::aarch64::*; assert_eq!(lhs.len(), rhs.len()); + let lo = vdupq_n_u8(0x0f_u8); let mut n = lhs.len(); let mut a = lhs.as_ptr(); let mut b = rhs.as_ptr(); - let lo = vdupq_n_u16(0x000f_u16); let mut _0 = vdupq_n_u32(0); let mut _1 = vdupq_n_u32(0); - let mut _2 = vdupq_n_u32(0); - let mut _3 = vdupq_n_u32(0); - let mut _4 = vdupq_n_u32(0); - let mut _5 = vdupq_n_u32(0); - let mut _6 = vdupq_n_u32(0); - let mut _7 = vdupq_n_u32(0); while n >= 16 { - let x = unsafe { vld1q_u16(a.cast()) }; - let y = unsafe { vld1q_u16(b.cast()) }; - let x_0 = vandq_u16(x, lo); - let x_1 = vandq_u16(vshrq_n_u16(x, 4), lo); - let x_2 = vandq_u16(vshrq_n_u16(x, 8), lo); - let x_3 = vshrq_n_u16(x, 12); - let y_0 = vandq_u16(y, lo); - let y_1 = vandq_u16(vshrq_n_u16(y, 4), lo); - let y_2 = vandq_u16(vshrq_n_u16(y, 8), lo); - let y_3 = vshrq_n_u16(y, 12); - let z_0 = vmulq_u16(x_0, y_0); - let z_1 = vmulq_u16(x_1, y_1); - let z_2 = vmulq_u16(x_2, y_2); - let z_3 = vmulq_u16(x_3, y_3); - a = unsafe { a.add(16) }; - b = unsafe { b.add(16) }; - n -= 16; - _0 = vaddq_u32(_0, vmovl_u16(vget_low_u16(z_0))); - _1 = vaddq_u32(_1, vmovl_u16(vget_high_u16(z_0))); - _2 = vaddq_u32(_2, vmovl_u16(vget_low_u16(z_1))); - _3 = vaddq_u32(_3, vmovl_u16(vget_high_u16(z_1))); - _4 = vaddq_u32(_4, vmovl_u16(vget_low_u16(z_2))); - _5 = vaddq_u32(_5, vmovl_u16(vget_high_u16(z_2))); - _6 = vaddq_u32(_6, vmovl_u16(vget_low_u16(z_3))); - _7 = vaddq_u32(_7, vmovl_u16(vget_high_u16(z_3))); + let x = unsafe { vld1q_u8(a.cast()) }; + let y = unsafe { vld1q_u8(b.cast()) }; + let x_0 = vandq_u8(x, lo); + let x_1 = vshrq_n_u8(x, 4); + let y_0 = vandq_u8(y, lo); + let y_1 = vshrq_n_u8(y, 4); + let z_0 = vmulq_u8(x_0, y_0); + let z_1 = vmulq_u8(x_1, y_1); + let lo = vaddl_u8(vget_low_u8(z_0), vget_low_u8(z_1)); + let hi = vaddl_high_u8(z_0, z_1); + _0 = vaddq_u32(_0, vpaddlq_u16(lo)); + _1 = vaddq_u32(_1, vpaddlq_u16(hi)); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; } - let mut sum = vaddvq_u32(vaddq_u32( - vaddq_u32(vaddq_u32(_0, _4), vaddq_u32(_1, _5)), - vaddq_u32(vaddq_u32(_2, _6), vaddq_u32(_3, _7)), - )); - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let y = unsafe { b.read() }; - a = unsafe { a.add(1) }; - b = unsafe { b.add(1) }; - n -= 1; - sum += (x as u32 & 0xf) * (y as u32 & 0xf); - sum += (x as u32 >> 4) * (y as u32 >> 4); + if n > 0 { + let (_a, _b) = unsafe { partial_load!(16, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { vld1q_u8(a.cast()) }; + let y = unsafe { vld1q_u8(b.cast()) }; + let x_0 = vandq_u8(x, lo); + let x_1 = vshrq_n_u8(x, 4); + let y_0 = vandq_u8(y, lo); + let y_1 = vshrq_n_u8(y, 4); + let z_0 = vmulq_u8(x_0, y_0); + let z_1 = vmulq_u8(x_1, y_1); + let lo = vaddl_u8(vget_low_u8(z_0), vget_low_u8(z_1)); + let hi = vaddl_high_u8(z_0, z_1); + _0 = vaddq_u32(_0, vpaddlq_u16(lo)); + _1 = vaddq_u32(_1, vpaddlq_u16(hi)); } - sum + vaddvq_u32(vaddq_u32(_0, _1)) } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] - fn reduce_sum_of_x_as_u32_y_as_u32_a2_test() { + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_test() { use rand::Rng; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -375,7 +406,7 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { for z in 3984..4016 { let lhs = &lhs[..z]; let rhs = &rhs[..z]; - let specialized = unsafe { reduce_sum_of_x_as_u32_y_as_u32_a2(lhs, rhs) }; + let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; let fallback = fallback(lhs, rhs); assert!( specialized == fallback, @@ -385,8 +416,8 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } } - #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] - pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { + #[crate::multiversion(@"v4:avx512vnni", @"v4", @"v3", @"v2", @"a2:dotprod", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { assert_eq!(s.len(), t.len()); let n = s.len(); let mut result = 0; @@ -399,6 +430,6 @@ mod reduce_sum_of_x_as_u32_y_as_u32 { } #[inline(always)] -pub fn reduce_sum_of_x_as_u32_y_as_u32(s: &[u8], t: &[u8]) -> u32 { - reduce_sum_of_x_as_u32_y_as_u32::reduce_sum_of_x_as_u32_y_as_u32(s, t) +pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + reduce_sum_of_xy::reduce_sum_of_xy(s, t) } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index c50bb998..c967e548 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -109,7 +109,8 @@ pub trait Floating: fn vector_abs_inplace(this: &mut [Self]); } -mod internal { +#[doc(hidden)] +pub mod internal { #[cfg(target_arch = "x86_64")] simd_macros::define_is_cpu_detected!("x86_64"); @@ -211,7 +212,6 @@ mod internal { } #[cfg(target_arch = "aarch64")] - #[expect(dead_code)] pub fn is_a3_128_detected() -> bool { std::arch::is_aarch64_feature_detected!("sve") } diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index 90c88fcd..cb01126c 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -18,7 +18,7 @@ mod mul_add_round { #[crate::target_cpu(enable = "v4")] fn mul_add_round_v4(this: &[f32], k: f32, b: f32) -> Vec { let mut r = Vec::::with_capacity(this.len()); - use std::arch::x86_64::*; + use core::arch::x86_64::*; let lk = _mm512_set1_ps(k); let lb = _mm512_set1_ps(b); let mut n = this.len(); @@ -29,12 +29,8 @@ mod mul_add_round { let v = _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); let v = _mm512_cvtps_epi32(v); let vfl = _mm512_cvtepi32_epi8(v); - unsafe { - _mm_storeu_si128(p.cast(), vfl); - } - n -= 16; - a = unsafe { a.add(16) }; - p = unsafe { p.add(16) }; + unsafe { _mm_storeu_si128(p.cast(), vfl) }; + (n, a, p) = unsafe { (n - 16, a.add(16), p.add(16)) }; } if n > 0 { let mask = _bzhi_u32(0xffff, n as u32) as u16; @@ -42,9 +38,7 @@ mod mul_add_round { let v = _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); let v = _mm512_cvtps_epi32(v); let vfl = _mm512_cvtepi32_epi8(v); - unsafe { - _mm_mask_storeu_epi8(p.cast(), mask, vfl); - } + unsafe { _mm_mask_storeu_epi8(p.cast(), mask, vfl) }; } unsafe { r.set_len(this.len()); @@ -52,8 +46,9 @@ mod mul_add_round { r } - #[cfg(all(target_arch = "x86_64", test, not(miri)))] + #[cfg(all(target_arch = "x86_64", test))] #[test] + #[cfg_attr(miri, ignore)] fn mul_add_round_v4_test() { if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -77,8 +72,9 @@ mod mul_add_round { #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] fn mul_add_round_v3(this: &[f32], k: f32, b: f32) -> Vec { - let mut r = Vec::::with_capacity(this.len()); - use std::arch::x86_64::*; + let mut r = Vec::::with_capacity(this.len().next_multiple_of(8)); + use crate::emulate::partial_load; + use core::arch::x86_64::*; let cons = _mm256_setr_epi8( 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 @@ -98,23 +94,20 @@ mod mul_add_round { let vlo = _mm256_extract_epi32::<0>(vs) as u32; let vhi = _mm256_extract_epi32::<4>(vs) as u32; let vfl = vlo as u64 | ((vhi as u64) << 32); - unsafe { - p.cast::().write_unaligned(vfl); - } - n -= 8; - a = unsafe { a.add(8) }; - p = unsafe { p.add(8) }; + unsafe { p.cast::().write_unaligned(vfl) }; + (n, a, p) = unsafe { (n - 8, a.add(8), p.add(8)) }; } - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let v = x.mul_add(k, b).round_ties_even() as u8; - unsafe { - p.write(v); - } - n -= 1; - a = unsafe { a.add(1) }; - p = unsafe { p.add(1) }; + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_loadu_ps(a) }; + let v = _mm256_fmadd_ps(x, lk, lb); + let v = _mm256_cvtps_epi32(_mm256_round_ps(v, 0x00)); + let vs = _mm256_shuffle_epi8(v, cons); + let vlo = _mm256_extract_epi32::<0>(vs) as u32; + let vhi = _mm256_extract_epi32::<4>(vs) as u32; + let vfl = vlo as u64 | ((vhi as u64) << 32); + unsafe { p.cast::().write_unaligned(vfl) }; } unsafe { r.set_len(this.len()); @@ -148,8 +141,9 @@ mod mul_add_round { #[crate::target_cpu(enable = "v2")] #[target_feature(enable = "fma")] fn mul_add_round_v2_fma(this: &[f32], k: f32, b: f32) -> Vec { - let mut r = Vec::::with_capacity(this.len()); - use std::arch::x86_64::*; + let mut r = Vec::::with_capacity(this.len().next_multiple_of(4)); + use crate::emulate::partial_load; + use core::arch::x86_64::*; let cons = _mm_setr_epi8( 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 @@ -165,23 +159,18 @@ mod mul_add_round { let v = _mm_cvtps_epi32(_mm_round_ps(v, 0x00)); let vs = _mm_shuffle_epi8(v, cons); let vfl = _mm_extract_epi32::<0>(vs) as u32; - unsafe { - p.cast::().write_unaligned(vfl); - } - n -= 4; - a = unsafe { a.add(4) }; - p = unsafe { p.add(4) }; + unsafe { p.cast::().write_unaligned(vfl) }; + (n, a, p) = unsafe { (n - 4, a.add(4), p.add(4)) }; } - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let v = x.mul_add(k, b).round_ties_even() as u8; - unsafe { - p.write(v); - } - n -= 1; - a = unsafe { a.add(1) }; - p = unsafe { p.add(1) }; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + let v = _mm_fmadd_ps(x, lk, lb); + let v = _mm_cvtps_epi32(_mm_round_ps(v, 0x00)); + let vs = _mm_shuffle_epi8(v, cons); + let vfl = _mm_extract_epi32::<0>(vs) as u32; + unsafe { p.cast::().write_unaligned(vfl) }; } unsafe { r.set_len(this.len()); @@ -214,8 +203,9 @@ mod mul_add_round { #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] fn mul_add_round_a2(this: &[f32], k: f32, b: f32) -> Vec { - let mut r = Vec::::with_capacity(this.len()); - use std::arch::aarch64::*; + let mut r = Vec::::with_capacity(this.len().next_multiple_of(4)); + use crate::emulate::partial_load; + use core::arch::aarch64::*; #[cfg(target_endian = "little")] const CONS: [u8; 16] = [ 0, 4, 8, 12, 0xff, 0xff, 0xff, 0xff, // 0..8 @@ -238,23 +228,18 @@ mod mul_add_round { let v = vcvtnq_u32_f32(v); let vs = vqtbl1q_u8(vreinterpretq_u8_u32(v), cons); let vfl = vgetq_lane_u32::<0>(vreinterpretq_u32_u8(vs)); - unsafe { - p.cast::().write_unaligned(vfl); - } - n -= 4; - a = unsafe { a.add(4) }; - p = unsafe { p.add(4) }; + unsafe { p.cast::().write_unaligned(vfl) }; + (n, a, p) = unsafe { (n - 4, a.add(4), p.add(4)) }; } - // this hint is used to disable loop unrolling - while std::hint::black_box(n) > 0 { - let x = unsafe { a.read() }; - let v = x.mul_add(k, b).round_ties_even() as u8; - unsafe { - p.write(v); - } - n -= 1; - a = unsafe { a.add(1) }; - p = unsafe { p.add(1) }; + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + let v = vfmaq_f32(lb, x, lk); + let v = vcvtnq_u32_f32(v); + let vs = vqtbl1q_u8(vreinterpretq_u8_u32(v), cons); + let vfl = vgetq_lane_u32::<0>(vreinterpretq_u32_u8(vs)); + unsafe { p.cast::().write_unaligned(vfl) }; } unsafe { r.set_len(this.len()); @@ -262,8 +247,9 @@ mod mul_add_round { r } - #[cfg(all(target_arch = "aarch64", test, not(miri)))] + #[cfg(all(target_arch = "aarch64", test))] #[test] + #[cfg_attr(miri, ignore)] fn mul_add_round_a2_test() { if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); From cdb7156f55fd09be48f1603e26c5cf3978924b72 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 6 Jan 2026 20:22:59 +0800 Subject: [PATCH 261/324] fix: add sve tests (#400) Signed-off-by: usamoi --- crates/simd/src/byte.rs | 28 +++++------------ crates/simd/src/emulate.rs | 20 ++++++------ crates/simd/src/floating_f32.rs | 54 +++++++++++++++++++++++++++++++-- crates/simd/src/halfbyte.rs | 16 +++------- crates/simd/src/quantize.rs | 40 ++++++++++++++++-------- 5 files changed, 101 insertions(+), 57 deletions(-) diff --git a/crates/simd/src/byte.rs b/crates/simd/src/byte.rs index 4513b983..4b94ea93 100644 --- a/crates/simd/src/byte.rs +++ b/crates/simd/src/byte.rs @@ -13,9 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. pub mod reduce_sum_of_xy { - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -75,9 +73,7 @@ pub mod reduce_sum_of_xy { } } - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -146,9 +142,7 @@ pub mod reduce_sum_of_xy { } } - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] @@ -218,9 +212,7 @@ pub mod reduce_sum_of_xy { } } - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] @@ -412,9 +404,7 @@ pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { } pub mod reduce_sum_of_x { - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -460,9 +450,7 @@ pub mod reduce_sum_of_x { } } - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] @@ -509,9 +497,7 @@ pub mod reduce_sum_of_x { } } - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index 8d11495d..5434088f 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -44,7 +44,7 @@ pub(crate) use partial_load; // Instructions. arXiv preprint arXiv:2112.06342. #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v4")] +#[target_feature(enable = "avx512f")] pub fn emulate_mm512_2intersect_epi32( a: core::arch::x86_64::__m512i, b: core::arch::x86_64::__m512i, @@ -99,7 +99,7 @@ pub fn emulate_mm512_2intersect_epi32( #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v3")] +#[target_feature(enable = "avx")] pub fn emulate_mm256_reduce_add_ps(mut x: core::arch::x86_64::__m256) -> f32 { use core::arch::x86_64::*; x = _mm256_add_ps(x, _mm256_permute2f128_ps(x, x, 1)); @@ -110,7 +110,7 @@ pub fn emulate_mm256_reduce_add_ps(mut x: core::arch::x86_64::__m256) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v2")] +#[target_feature(enable = "ssse3")] pub fn emulate_mm_reduce_add_ps(mut x: core::arch::x86_64::__m128) -> f32 { use core::arch::x86_64::*; x = _mm_hadd_ps(x, x); @@ -120,7 +120,7 @@ pub fn emulate_mm_reduce_add_ps(mut x: core::arch::x86_64::__m128) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v3")] +#[target_feature(enable = "avx2")] pub fn emulate_mm256_reduce_add_epi32(mut x: core::arch::x86_64::__m256i) -> i32 { use core::arch::x86_64::*; x = _mm256_add_epi32(x, _mm256_permute2f128_si256(x, x, 1)); @@ -131,7 +131,7 @@ pub fn emulate_mm256_reduce_add_epi32(mut x: core::arch::x86_64::__m256i) -> i32 #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v2")] +#[target_feature(enable = "ssse3")] pub fn emulate_mm_reduce_add_epi32(mut x: core::arch::x86_64::__m128i) -> i32 { use core::arch::x86_64::*; x = _mm_hadd_epi32(x, x); @@ -141,7 +141,7 @@ pub fn emulate_mm_reduce_add_epi32(mut x: core::arch::x86_64::__m128i) -> i32 { #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v3")] +#[target_feature(enable = "avx")] pub fn emulate_mm256_reduce_min_ps(x: core::arch::x86_64::__m256) -> f32 { use crate::aligned::Aligned16; use core::arch::x86_64::*; @@ -157,7 +157,7 @@ pub fn emulate_mm256_reduce_min_ps(x: core::arch::x86_64::__m256) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v2")] +#[target_feature(enable = "sse")] pub fn emulate_mm_reduce_min_ps(x: core::arch::x86_64::__m128) -> f32 { use crate::aligned::Aligned16; use core::arch::x86_64::*; @@ -171,7 +171,7 @@ pub fn emulate_mm_reduce_min_ps(x: core::arch::x86_64::__m128) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v3")] +#[target_feature(enable = "avx")] pub fn emulate_mm256_reduce_max_ps(x: core::arch::x86_64::__m256) -> f32 { use crate::aligned::Aligned16; use core::arch::x86_64::*; @@ -187,7 +187,7 @@ pub fn emulate_mm256_reduce_max_ps(x: core::arch::x86_64::__m256) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v2")] +#[target_feature(enable = "sse")] pub fn emulate_mm_reduce_max_ps(x: core::arch::x86_64::__m128) -> f32 { use crate::aligned::Aligned16; use core::arch::x86_64::*; @@ -201,7 +201,7 @@ pub fn emulate_mm_reduce_max_ps(x: core::arch::x86_64::__m128) -> f32 { #[inline] #[cfg(target_arch = "x86_64")] -#[crate::target_cpu(enable = "v3")] +#[target_feature(enable = "avx2")] pub fn emulate_mm256_reduce_add_epi64(mut x: core::arch::x86_64::__m256i) -> i64 { use core::arch::x86_64::*; x = _mm256_add_epi64(x, _mm256_permute2f128_si256(x, x, 1)); diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs index 994f4e7d..4605914e 100644 --- a/crates/simd/src/floating_f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -1250,7 +1250,32 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a3_256_test() { - // + use rand::Rng; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a3_256(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } } #[inline] @@ -1530,7 +1555,32 @@ mod reduce_sum_of_d2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_a3_256_test() { - // + use rand::Rng; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_a3_256(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } } #[inline] diff --git a/crates/simd/src/halfbyte.rs b/crates/simd/src/halfbyte.rs index 7831ea0a..ad91a940 100644 --- a/crates/simd/src/halfbyte.rs +++ b/crates/simd/src/halfbyte.rs @@ -13,9 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. pub mod reduce_sum_of_xy { - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -81,9 +79,7 @@ pub mod reduce_sum_of_xy { } } - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -153,9 +149,7 @@ pub mod reduce_sum_of_xy { } } - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] @@ -226,9 +220,7 @@ pub mod reduce_sum_of_xy { } } - /// # Safety - /// - /// * Don't call it. Internal use. + #[doc(hidden)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index cb01126c..50c050b9 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -50,17 +50,21 @@ mod mul_add_round { #[test] #[cfg_attr(miri, ignore)] fn mul_add_round_v4_test() { + use rand::Rng; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; } + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4010; - let x = (0..n).map(|_| rand::random::<_>()).collect::>(); + let x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); for z in 3990..4010 { let x = &x[..z]; - let k = 20.0; - let b = 20.0; + let k = 127.0; + let b = 127.0; let specialized = unsafe { mul_add_round_v4(x, k, b) }; let fallback = fallback(x, k, b); assert_eq!(specialized, fallback); @@ -118,17 +122,21 @@ mod mul_add_round { #[cfg(all(target_arch = "x86_64", test))] #[test] fn mul_add_round_v3_test() { + use rand::Rng; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; } + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4010; - let x = (0..n).map(|_| rand::random::<_>()).collect::>(); + let x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); for z in 3990..4010 { let x = &x[..z]; - let k = 20.0; - let b = 20.0; + let k = 127.0; + let b = 127.0; let specialized = unsafe { mul_add_round_v3(x, k, b) }; let fallback = fallback(x, k, b); assert_eq!(specialized, fallback); @@ -181,17 +189,21 @@ mod mul_add_round { #[cfg(all(target_arch = "x86_64", test))] #[test] fn mul_add_round_v2_fma_test() { + use rand::Rng; if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { println!("test {} ... skipped (v2:fma)", module_path!()); return; } + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4010; - let x = (0..n).map(|_| rand::random::<_>()).collect::>(); + let x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); for z in 3990..4010 { let x = &x[..z]; - let k = 20.0; - let b = 20.0; + let k = 127.0; + let b = 127.0; let specialized = unsafe { mul_add_round_v2_fma(x, k, b) }; let fallback = fallback(x, k, b); assert_eq!(specialized, fallback); @@ -251,17 +263,21 @@ mod mul_add_round { #[test] #[cfg_attr(miri, ignore)] fn mul_add_round_a2_test() { + use rand::Rng; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); return; } + let mut rng = rand::rng(); for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { let n = 4010; - let x = (0..n).map(|_| rand::random::<_>()).collect::>(); + let x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); for z in 3990..4010 { let x = &x[..z]; - let k = 20.0; - let b = 20.0; + let k = 127.0; + let b = 127.0; let specialized = unsafe { mul_add_round_a2(x, k, b) }; let fallback = fallback(x, k, b); assert_eq!(specialized, fallback); From 7101aa7c48561e4a23fa435ad49d45f1c8235f70 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 7 Jan 2026 20:12:32 +0800 Subject: [PATCH 262/324] refactor: rewrite f16 simd routines for precision (#401) Signed-off-by: usamoi --- crates/simd/Cargo.toml | 2 + crates/simd/benches/bench.rs | 234 +++++++++++++- crates/simd/build.rs | 1 + crates/simd/cshim/aarch64_fp16.c | 154 +++++++-- crates/simd/cshim/aarch64_sve_fp16.c | 83 ++++- crates/simd/cshim/x86_64_fp16.c | 72 ++++- crates/simd/src/byte.rs | 37 ++- crates/simd/src/fast_scan.rs | 12 +- crates/simd/src/floating_f16.rs | 450 +++++++++++++++++++++++++-- crates/simd/src/floating_f32.rs | 12 + crates/simd/src/halfbyte.rs | 21 +- crates/simd/src/lib.rs | 14 + crates/simd/src/rotate.rs | 4 +- crates/simd_macros/src/lib.rs | 18 ++ 14 files changed, 997 insertions(+), 117 deletions(-) diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 5f9acc54..87357b92 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -7,9 +7,11 @@ publish = false [[bench]] name = "bench" harness = false +required-features = ["internal"] [features] init = [] +internal = [] nightly_f16 = ["zerocopy/float-nightly"] [dependencies] diff --git a/crates/simd/benches/bench.rs b/crates/simd/benches/bench.rs index a8c55542..28f3b5ba 100644 --- a/crates/simd/benches/bench.rs +++ b/crates/simd/benches/bench.rs @@ -16,11 +16,197 @@ use criterion::{Criterion, criterion_group, criterion_main}; +fn floating_f32_reduce_sum_of_xy(c: &mut Criterion) { + use rand::Rng; + let mut rng = rand::rng(); + let x = (0..4095) + .map(|_| rng.random_range(-1.0..=1.0f32)) + .collect::>(); + let y = (0..4095) + .map(|_| rng.random_range(-1.0..=1.0f32)) + .collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("floating_f32::reduce_sum_of_xy::v4", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("floating_f32::reduce_sum_of_xy::v3", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") && simd::is_feature_detected!("fma") { + c.bench_function("floating_f32::reduce_sum_of_xy::v2_fma", |b| { + b.iter(|| unsafe { + simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_v2_fma(&x, &y) + }) + }); + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + if simd::is_cpu_detected!("a3.256") { + c.bench_function("floating_f32::reduce_sum_of_xy::a3_256", |b| { + b.iter(|| unsafe { + simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_a3_256(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("floating_f32::reduce_sum_of_xy::a2", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_a2(&x, &y) }) + }); + } +} + +fn floating_f32_reduce_sum_of_d2(c: &mut Criterion) { + use rand::Rng; + let mut rng = rand::rng(); + let x = (0..4095) + .map(|_| rng.random_range(-1.0..=1.0f32)) + .collect::>(); + let y = (0..4095) + .map(|_| rng.random_range(-1.0..=1.0f32)) + .collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("floating_f32::reduce_sum_of_d2::v4", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("floating_f32::reduce_sum_of_d2::v3", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") && simd::is_feature_detected!("fma") { + c.bench_function("floating_f32::reduce_sum_of_d2::v2_fma", |b| { + b.iter(|| unsafe { + simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_v2_fma(&x, &y) + }) + }); + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + if simd::is_cpu_detected!("a3.256") { + c.bench_function("floating_f32::reduce_sum_of_d2::a3_256", |b| { + b.iter(|| unsafe { + simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_a3_256(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("floating_f32::reduce_sum_of_d2::a2", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_a2(&x, &y) }) + }); + } +} + +fn floating_f16_reduce_sum_of_xy(c: &mut Criterion) { + use rand::Rng; + use simd::{F16, f16}; + let mut rng = rand::rng(); + let x = (0..4095) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let y = (0..4095) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512fp16") { + c.bench_function("floating_f16::reduce_sum_of_xy::v4_avx512fp16", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_v4_avx512fp16(&x, &y) + }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("floating_f16::reduce_sum_of_xy::v4", |b| { + b.iter(|| unsafe { simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("floating_f16::reduce_sum_of_xy::v3", |b| { + b.iter(|| unsafe { simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a3.512") { + c.bench_function("floating_f16::reduce_sum_of_xy::a3_512", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_a3_512(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") && simd::is_feature_detected!("fp16") { + c.bench_function("floating_f16::reduce_sum_of_xy::a2_fp16", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_a2_fp16(&x, &y) + }) + }); + } +} + +fn floating_f16_reduce_sum_of_d2(c: &mut Criterion) { + use rand::Rng; + use simd::{F16, f16}; + let mut rng = rand::rng(); + let x = (0..4095) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let y = (0..4095) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512fp16") { + c.bench_function("floating_f16::reduce_sum_of_d2::v4_avx512fp16", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_v4_avx512fp16(&x, &y) + }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("floating_f16::reduce_sum_of_d2::v4", |b| { + b.iter(|| unsafe { simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("floating_f16::reduce_sum_of_d2::v3", |b| { + b.iter(|| unsafe { simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a3.512") { + c.bench_function("floating_f16::reduce_sum_of_d2::a3_512", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_a3_512(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") && simd::is_feature_detected!("fp16") { + c.bench_function("floating_f16::reduce_sum_of_d2::a2_fp16", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_a2_fp16(&x, &y) + }) + }); + } +} + fn byte_reduce_sum_of_xy(c: &mut Criterion) { use rand::Rng; let mut rng = rand::rng(); - let x = (0..4095).map(|_| rng.random()).collect::>(); - let y = (0..4095).map(|_| rng.random()).collect::>(); + let x = (0..4095).map(|_| rng.random::()).collect::>(); + let y = (0..4095).map(|_| rng.random::()).collect::>(); #[cfg(target_arch = "x86_64")] if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512vnni") { c.bench_function("byte::reduce_sum_of_xy::v4_avx512vnni", |b| { @@ -47,12 +233,24 @@ fn byte_reduce_sum_of_xy(c: &mut Criterion) { b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_v2(&x, &y) }) }); } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") && simd::is_feature_detected!("dotprod") { + c.bench_function("byte::reduce_sum_of_xy::a2_prod", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_a2_dotprod(&x, &y) }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("byte::reduce_sum_of_xy::a2", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_a2(&x, &y) }) + }); + } } fn byte_reduce_sum_of_x(c: &mut Criterion) { use rand::Rng; let mut rng = rand::rng(); - let this = (0..4095).map(|_| rng.random()).collect::>(); + let this = (0..4095).map(|_| rng.random::()).collect::>(); #[cfg(target_arch = "x86_64")] if simd::is_cpu_detected!("v4") { c.bench_function("byte::reduce_sum_of_x::v4", |b| { @@ -71,13 +269,19 @@ fn byte_reduce_sum_of_x(c: &mut Criterion) { b.iter(|| unsafe { simd::byte::reduce_sum_of_x::reduce_sum_of_x_v2(&this) }) }); } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("byte::reduce_sum_of_x::a2", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_x::reduce_sum_of_x_a2(&this) }) + }); + } } fn halfbyte_reduce_sum_of_xy(c: &mut Criterion) { use rand::Rng; let mut rng = rand::rng(); - let x = (0..2047).map(|_| rng.random()).collect::>(); - let y = (0..2047).map(|_| rng.random()).collect::>(); + let x = (0..2047).map(|_| rng.random::()).collect::>(); + let y = (0..2047).map(|_| rng.random::()).collect::>(); #[cfg(target_arch = "x86_64")] if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512vnni") { c.bench_function("halfbyte::reduce_sum_of_xy::v4_avx512vnni", |b| { @@ -104,12 +308,30 @@ fn halfbyte_reduce_sum_of_xy(c: &mut Criterion) { b.iter(|| unsafe { simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_v2(&x, &y) }) }); } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") && simd::is_feature_detected!("dotprod") { + c.bench_function("halfbyte::reduce_sum_of_xy::a2_prod", |b| { + b.iter(|| unsafe { + simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_a2_dotprod(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("halfbyte::reduce_sum_of_xy::a2", |b| { + b.iter(|| unsafe { simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_a2(&x, &y) }) + }); + } } criterion_group!( benches, + floating_f32_reduce_sum_of_xy, + floating_f32_reduce_sum_of_d2, + floating_f16_reduce_sum_of_xy, + floating_f16_reduce_sum_of_d2, byte_reduce_sum_of_xy, byte_reduce_sum_of_x, - halfbyte_reduce_sum_of_xy + halfbyte_reduce_sum_of_xy, ); criterion_main!(benches); diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 3229041e..63af110d 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -19,6 +19,7 @@ fn main() -> Result<(), Box> { println!("cargo::rerun-if-changed=cshim"); let target_arch = var("CARGO_CFG_TARGET_ARCH")?; let target_endian = var("CARGO_CFG_TARGET_ENDIAN")?; + #[allow(clippy::single_match)] match target_arch.as_str() { "aarch64" => { let mut build = cc::Build::new(); diff --git a/crates/simd/cshim/aarch64_fp16.c b/crates/simd/cshim/aarch64_fp16.c index c72b14fb..ad7ffb93 100644 --- a/crates/simd/cshim/aarch64_fp16.c +++ b/crates/simd/cshim/aarch64_fp16.c @@ -37,7 +37,38 @@ fp16_reduce_sum_of_xy_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { float16x8_t sum_1 = vdupq_n_f16(0.0); float16x8_t sum_2 = vdupq_n_f16(0.0); float16x8_t sum_3 = vdupq_n_f16(0.0); - while (n >= 32) { + float16x8_t sum_4 = vdupq_n_f16(0.0); + float16x8_t sum_5 = vdupq_n_f16(0.0); + float16x8_t sum_6 = vdupq_n_f16(0.0); + float16x8_t sum_7 = vdupq_n_f16(0.0); + while (n >= 64) { + float16x8_t x_0 = vld1q_f16(a + 0); + float16x8_t x_1 = vld1q_f16(a + 8); + float16x8_t x_2 = vld1q_f16(a + 16); + float16x8_t x_3 = vld1q_f16(a + 24); + float16x8_t x_4 = vld1q_f16(a + 32); + float16x8_t x_5 = vld1q_f16(a + 40); + float16x8_t x_6 = vld1q_f16(a + 48); + float16x8_t x_7 = vld1q_f16(a + 56); + float16x8_t y_0 = vld1q_f16(b + 0); + float16x8_t y_1 = vld1q_f16(b + 8); + float16x8_t y_2 = vld1q_f16(b + 16); + float16x8_t y_3 = vld1q_f16(b + 24); + float16x8_t y_4 = vld1q_f16(b + 32); + float16x8_t y_5 = vld1q_f16(b + 40); + float16x8_t y_6 = vld1q_f16(b + 48); + float16x8_t y_7 = vld1q_f16(b + 56); + sum_0 = vfmaq_f16(sum_0, x_0, y_0); + sum_1 = vfmaq_f16(sum_1, x_1, y_1); + sum_2 = vfmaq_f16(sum_2, x_2, y_2); + sum_3 = vfmaq_f16(sum_3, x_3, y_3); + sum_4 = vfmaq_f16(sum_4, x_4, y_4); + sum_5 = vfmaq_f16(sum_5, x_5, y_5); + sum_6 = vfmaq_f16(sum_6, x_6, y_6); + sum_7 = vfmaq_f16(sum_7, x_7, y_7); + n -= 64, a += 64, b += 64; + } + if (n >= 32) { float16x8_t x_0 = vld1q_f16(a + 0); float16x8_t x_1 = vld1q_f16(a + 8); float16x8_t x_2 = vld1q_f16(a + 16); @@ -52,24 +83,30 @@ fp16_reduce_sum_of_xy_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { sum_3 = vfmaq_f16(sum_3, x_3, y_3); n -= 32, a += 32, b += 32; } + if (n >= 16) { + float16x8_t x_4 = vld1q_f16(a + 0); + float16x8_t x_5 = vld1q_f16(a + 8); + float16x8_t y_4 = vld1q_f16(b + 0); + float16x8_t y_5 = vld1q_f16(b + 8); + sum_4 = vfmaq_f16(sum_4, x_4, y_4); + sum_5 = vfmaq_f16(sum_5, x_5, y_5); + n -= 16, a += 16, b += 16; + } + if (n >= 8) { + float16x8_t x_6 = vld1q_f16(a + 0); + float16x8_t y_6 = vld1q_f16(b + 0); + sum_6 = vfmaq_f16(sum_6, x_6, y_6); + n -= 8, a += 8, b += 8; + } if (n > 0) { - f16 _a[32] = {}, _b[32] = {}; + f16 _a[8] = {}, _b[8] = {}; for (size_t i = 0; i < n; i += 1) { _a[i] = a[i], _b[i] = b[i]; } a = _a, b = _b; - float16x8_t x_0 = vld1q_f16(_a + 0); - float16x8_t x_1 = vld1q_f16(_a + 8); - float16x8_t x_2 = vld1q_f16(_a + 16); - float16x8_t x_3 = vld1q_f16(_a + 24); - float16x8_t y_0 = vld1q_f16(_b + 0); - float16x8_t y_1 = vld1q_f16(_b + 8); - float16x8_t y_2 = vld1q_f16(_b + 16); - float16x8_t y_3 = vld1q_f16(_b + 24); - sum_0 = vfmaq_f16(sum_0, x_0, y_0); - sum_1 = vfmaq_f16(sum_1, x_1, y_1); - sum_2 = vfmaq_f16(sum_2, x_2, y_2); - sum_3 = vfmaq_f16(sum_3, x_3, y_3); + float16x8_t x_7 = vld1q_f16(a); + float16x8_t y_7 = vld1q_f16(b); + sum_7 = vfmaq_f16(sum_7, x_7, y_7); } float32x4_t s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); float32x4_t s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); @@ -79,9 +116,19 @@ fp16_reduce_sum_of_xy_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { float32x4_t s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); float32x4_t s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); float32x4_t s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); - float32x4_t s = + float32x4_t s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + float32x4_t s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + float32x4_t s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + float32x4_t s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + float32x4_t s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + float32x4_t s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + float32x4_t s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + float32x4_t s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + float32x4_t s = vpaddq_f32( vpaddq_f32(vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), - vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7))); + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7))), + vpaddq_f32(vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)))); return vaddvq_f32(s); } @@ -91,31 +138,46 @@ fp16_reduce_sum_of_d2_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { float16x8_t sum_1 = vdupq_n_f16(0.0); float16x8_t sum_2 = vdupq_n_f16(0.0); float16x8_t sum_3 = vdupq_n_f16(0.0); - while (n >= 32) { + float16x8_t sum_4 = vdupq_n_f16(0.0); + float16x8_t sum_5 = vdupq_n_f16(0.0); + float16x8_t sum_6 = vdupq_n_f16(0.0); + float16x8_t sum_7 = vdupq_n_f16(0.0); + while (n >= 64) { float16x8_t x_0 = vld1q_f16(a + 0); float16x8_t x_1 = vld1q_f16(a + 8); float16x8_t x_2 = vld1q_f16(a + 16); float16x8_t x_3 = vld1q_f16(a + 24); + float16x8_t x_4 = vld1q_f16(a + 32); + float16x8_t x_5 = vld1q_f16(a + 40); + float16x8_t x_6 = vld1q_f16(a + 48); + float16x8_t x_7 = vld1q_f16(a + 56); float16x8_t y_0 = vld1q_f16(b + 0); float16x8_t y_1 = vld1q_f16(b + 8); float16x8_t y_2 = vld1q_f16(b + 16); float16x8_t y_3 = vld1q_f16(b + 24); + float16x8_t y_4 = vld1q_f16(b + 32); + float16x8_t y_5 = vld1q_f16(b + 40); + float16x8_t y_6 = vld1q_f16(b + 48); + float16x8_t y_7 = vld1q_f16(b + 56); float16x8_t d_0 = vsubq_f16(x_0, y_0); float16x8_t d_1 = vsubq_f16(x_1, y_1); float16x8_t d_2 = vsubq_f16(x_2, y_2); float16x8_t d_3 = vsubq_f16(x_3, y_3); + float16x8_t d_4 = vsubq_f16(x_4, y_4); + float16x8_t d_5 = vsubq_f16(x_5, y_5); + float16x8_t d_6 = vsubq_f16(x_6, y_6); + float16x8_t d_7 = vsubq_f16(x_7, y_7); sum_0 = vfmaq_f16(sum_0, d_0, d_0); sum_1 = vfmaq_f16(sum_1, d_1, d_1); sum_2 = vfmaq_f16(sum_2, d_2, d_2); sum_3 = vfmaq_f16(sum_3, d_3, d_3); - n -= 32, a += 32, b += 32; + sum_4 = vfmaq_f16(sum_4, d_4, d_4); + sum_5 = vfmaq_f16(sum_5, d_5, d_5); + sum_6 = vfmaq_f16(sum_6, d_6, d_6); + sum_7 = vfmaq_f16(sum_7, d_7, d_7); + n -= 64, a += 64, b += 64; } - if (n > 0) { - f16 _a[32] = {}, _b[32] = {}; - for (size_t i = 0; i < n; i += 1) { - _a[i] = a[i], _b[i] = b[i]; - } - a = _a, b = _b; + if (n >= 32) { float16x8_t x_0 = vld1q_f16(a + 0); float16x8_t x_1 = vld1q_f16(a + 8); float16x8_t x_2 = vld1q_f16(a + 16); @@ -132,6 +194,36 @@ fp16_reduce_sum_of_d2_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { sum_1 = vfmaq_f16(sum_1, d_1, d_1); sum_2 = vfmaq_f16(sum_2, d_2, d_2); sum_3 = vfmaq_f16(sum_3, d_3, d_3); + n -= 32, a += 32, b += 32; + } + if (n >= 16) { + float16x8_t x_4 = vld1q_f16(a + 0); + float16x8_t x_5 = vld1q_f16(a + 8); + float16x8_t y_4 = vld1q_f16(b + 0); + float16x8_t y_5 = vld1q_f16(b + 8); + float16x8_t d_4 = vsubq_f16(x_4, y_4); + float16x8_t d_5 = vsubq_f16(x_5, y_5); + sum_4 = vfmaq_f16(sum_4, d_4, d_4); + sum_5 = vfmaq_f16(sum_5, d_5, d_5); + n -= 16, a += 16, b += 16; + } + if (n >= 8) { + float16x8_t x_6 = vld1q_f16(a + 0); + float16x8_t y_6 = vld1q_f16(b + 0); + float16x8_t d_6 = vsubq_f16(x_6, y_6); + sum_6 = vfmaq_f16(sum_6, d_6, d_6); + n -= 8, a += 8, b += 8; + } + if (n > 0) { + f16 _a[8] = {}, _b[8] = {}; + for (size_t i = 0; i < n; i += 1) { + _a[i] = a[i], _b[i] = b[i]; + } + a = _a, b = _b; + float16x8_t x_7 = vld1q_f16(a); + float16x8_t y_7 = vld1q_f16(b); + float16x8_t d_7 = vsubq_f16(x_7, y_7); + sum_7 = vfmaq_f16(sum_7, d_7, d_7); } float32x4_t s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); float32x4_t s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); @@ -141,8 +233,18 @@ fp16_reduce_sum_of_d2_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { float32x4_t s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); float32x4_t s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); float32x4_t s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); - float32x4_t s = + float32x4_t s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + float32x4_t s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + float32x4_t s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + float32x4_t s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + float32x4_t s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + float32x4_t s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + float32x4_t s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + float32x4_t s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + float32x4_t s = vpaddq_f32( vpaddq_f32(vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), - vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7))); + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7))), + vpaddq_f32(vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)))); return vaddvq_f32(s); } diff --git a/crates/simd/cshim/aarch64_sve_fp16.c b/crates/simd/cshim/aarch64_sve_fp16.c index 77111827..ba1023cd 100644 --- a/crates/simd/cshim/aarch64_sve_fp16.c +++ b/crates/simd/cshim/aarch64_sve_fp16.c @@ -36,25 +36,78 @@ typedef float f32; __attribute__((target("+sve"))) float fp16_reduce_sum_of_xy_a3_512(size_t n, f16 *restrict a, f16 *restrict b) { - svfloat16_t sum = svdup_f16(0.0); - for (size_t i = 0; i < n; i += svcnth()) { - svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); - svfloat16_t x = svld1_f16(mask, a + i); - svfloat16_t y = svld1_f16(mask, b + i); - sum = svmla_f16_m(mask, sum, x, y); + const size_t vl = svcnth(); + svfloat16_t _0 = svdup_f16(0.0); + svfloat16_t _1 = svdup_f16(0.0); + while (n >= vl * 2) { + svbool_t mask = svptrue_b16(); + svfloat16_t x_0 = svld1_f16(mask, a + 0 * vl); + svfloat16_t y_0 = svld1_f16(mask, b + 0 * vl); + svfloat16_t x_1 = svld1_f16(mask, a + 1 * vl); + svfloat16_t y_1 = svld1_f16(mask, b + 1 * vl); + _0 = svmla_f16_x(mask, _0, x_0, y_0); + _1 = svmla_f16_x(mask, _1, x_1, y_1); + n -= vl * 2, a += vl * 2, b += vl * 2; } - return svaddv_f16(svptrue_b16(), sum); + if (n >= vl) { + svbool_t mask = svptrue_b16(); + svfloat16_t x_0 = svld1_f16(mask, a + 0 * vl); + svfloat16_t y_0 = svld1_f16(mask, b + 0 * vl); + _0 = svmla_f16_x(mask, _0, x_0, y_0); + n -= vl * 1, a += vl * 1, b += vl * 1; + } + if (n > 0) { + svbool_t mask = svwhilelt_b16((int64_t)0, (int64_t)n); + svfloat16_t x_0 = svld1_f16(mask, a); + svfloat16_t y_0 = svld1_f16(mask, b); + _0 = svmla_f16_m(mask, _0, x_0, y_0); + } + svbool_t mask = svptrue_b32(); + svfloat32_t s_0 = svcvt_f32_f16_x(mask, _0); + svfloat32_t s_1 = svcvt_f32_f16_x(mask, svext_f16(_0, _0, 1)); + svfloat32_t s_2 = svcvt_f32_f16_x(mask, _1); + svfloat32_t s_3 = svcvt_f32_f16_x(mask, svext_f16(_1, _1, 1)); + return svaddv_f32(mask, svadd_f32_x(mask, svadd_f32_x(mask, s_0, s_2), + svadd_f32_x(mask, s_1, s_3))); } __attribute__((target("+sve"))) float fp16_reduce_sum_of_d2_a3_512(size_t n, f16 *restrict a, f16 *restrict b) { - svfloat16_t sum = svdup_f16(0.0); - for (size_t i = 0; i < n; i += svcnth()) { - svbool_t mask = svwhilelt_b16((int64_t)i, (int64_t)n); - svfloat16_t x = svld1_f16(mask, a + i); - svfloat16_t y = svld1_f16(mask, b + i); - svfloat16_t d = svsub_f16_z(mask, x, y); - sum = svmla_f16_m(mask, sum, d, d); + const size_t vl = svcnth(); + svfloat16_t _0 = svdup_f16(0.0); + svfloat16_t _1 = svdup_f16(0.0); + while (n >= vl * 2) { + svbool_t mask = svptrue_b16(); + svfloat16_t x_0 = svld1_f16(mask, a + 0 * vl); + svfloat16_t y_0 = svld1_f16(mask, b + 0 * vl); + svfloat16_t x_1 = svld1_f16(mask, a + 1 * vl); + svfloat16_t y_1 = svld1_f16(mask, b + 1 * vl); + svfloat16_t d_0 = svsub_f16_z(mask, x_0, y_0); + svfloat16_t d_1 = svsub_f16_z(mask, x_1, y_1); + _0 = svmla_f16_x(mask, _0, d_0, d_0); + _1 = svmla_f16_x(mask, _1, d_1, d_1); + n -= vl * 2, a += vl * 2, b += vl * 2; + } + if (n >= vl) { + svbool_t mask = svptrue_b16(); + svfloat16_t x_0 = svld1_f16(mask, a + 0 * vl); + svfloat16_t y_0 = svld1_f16(mask, b + 0 * vl); + svfloat16_t d_0 = svsub_f16_z(mask, x_0, y_0); + _0 = svmla_f16_x(mask, _0, d_0, d_0); + n -= vl * 1, a += vl * 1, b += vl * 1; + } + if (n > 0) { + svbool_t mask = svwhilelt_b16((int64_t)0, (int64_t)n); + svfloat16_t x_0 = svld1_f16(mask, a); + svfloat16_t y_0 = svld1_f16(mask, b); + svfloat16_t d_0 = svsub_f16_z(mask, x_0, y_0); + _0 = svmla_f16_m(mask, _0, d_0, d_0); } - return svaddv_f16(svptrue_b16(), sum); + svbool_t mask = svptrue_b32(); + svfloat32_t s_0 = svcvt_f32_f16_x(mask, _0); + svfloat32_t s_1 = svcvt_f32_f16_x(mask, svext_f16(_0, _0, 1)); + svfloat32_t s_2 = svcvt_f32_f16_x(mask, _1); + svfloat32_t s_3 = svcvt_f32_f16_x(mask, svext_f16(_1, _1, 1)); + return svaddv_f32(mask, svadd_f32_x(mask, svadd_f32_x(mask, s_0, s_2), + svadd_f32_x(mask, s_1, s_3))); } diff --git a/crates/simd/cshim/x86_64_fp16.c b/crates/simd/cshim/x86_64_fp16.c index f9f40ea1..79a204b5 100644 --- a/crates/simd/cshim/x86_64_fp16.c +++ b/crates/simd/cshim/x86_64_fp16.c @@ -62,40 +62,80 @@ __attribute__((target("avx512bw,avx512cd,avx512dq,avx512vl,bmi,bmi2,lzcnt," "movbe,popcnt,avx512fp16"))) float fp16_reduce_sum_of_xy_v4_avx512fp16(size_t n, f16 *restrict a, f16 *restrict b) { - __m512h sum = _mm512_setzero_ph(); + __m512h _0 = _mm512_setzero_ph(); + __m512h _1 = _mm512_setzero_ph(); + while (n >= 64) { + __m512h x_0 = _mm512_loadu_ph(a + 0); + __m512h x_1 = _mm512_loadu_ph(a + 32); + __m512h y_0 = _mm512_loadu_ph(b + 0); + __m512h y_1 = _mm512_loadu_ph(b + 32); + _0 = _mm512_fmadd_ph(x_0, y_0, _0); + _1 = _mm512_fmadd_ph(x_1, y_1, _1); + n -= 64, a += 64, b += 64; + } while (n >= 32) { - __m512h x = _mm512_loadu_ph(a); - __m512h y = _mm512_loadu_ph(b); - sum = _mm512_fmadd_ph(x, y, sum); + __m512h x_0 = _mm512_loadu_ph(a + 0); + __m512h y_0 = _mm512_loadu_ph(b + 0); + _0 = _mm512_fmadd_ph(x_0, y_0, _0); n -= 32, a += 32, b += 32; } if (n > 0) { unsigned int mask = _bzhi_u32(0xffffffff, n); __m512h x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a)); __m512h y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b)); - sum = _mm512_fmadd_ph(x, y, sum); + _1 = _mm512_fmadd_ph(x, y, _1); } - return _mm512_reduce_add_ph(sum); + __m512 s_0 = + _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + __m512 s_1 = + _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + __m512 s_2 = + _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + __m512 s_3 = + _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + return _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(s_0, s_2), _mm512_add_ps(s_1, s_3))); } __attribute__((target("avx512bw,avx512cd,avx512dq,avx512vl,bmi,bmi2,lzcnt," "movbe,popcnt,avx512fp16"))) float fp16_reduce_sum_of_d2_v4_avx512fp16(size_t n, f16 *restrict a, f16 *restrict b) { - __m512h sum = _mm512_setzero_ph(); + __m512h _0 = _mm512_setzero_ph(); + __m512h _1 = _mm512_setzero_ph(); + while (n >= 64) { + __m512h x_0 = _mm512_loadu_ph(a + 0); + __m512h x_1 = _mm512_loadu_ph(a + 32); + __m512h y_0 = _mm512_loadu_ph(b + 0); + __m512h y_1 = _mm512_loadu_ph(b + 32); + __m512h d_0 = _mm512_sub_ph(x_0, y_0); + __m512h d_1 = _mm512_sub_ph(x_1, y_1); + _0 = _mm512_fmadd_ph(d_0, d_0, _0); + _1 = _mm512_fmadd_ph(d_1, d_1, _1); + n -= 64, a += 64, b += 64; + } while (n >= 32) { - __m512h x = _mm512_loadu_ph(a); - __m512h y = _mm512_loadu_ph(b); - __m512h d = _mm512_sub_ph(x, y); - sum = _mm512_fmadd_ph(d, d, sum); + __m512h x_0 = _mm512_loadu_ph(a + 0); + __m512h y_0 = _mm512_loadu_ph(b + 0); + __m512h d_0 = _mm512_sub_ph(x_0, y_0); + _0 = _mm512_fmadd_ph(d_0, d_0, _0); n -= 32, a += 32, b += 32; } if (n > 0) { unsigned int mask = _bzhi_u32(0xffffffff, n); - __m512h x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a)); - __m512h y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b)); - __m512h d = _mm512_sub_ph(x, y); - sum = _mm512_fmadd_ph(d, d, sum); + __m512h x_1 = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a)); + __m512h y_1 = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b)); + __m512h d_1 = _mm512_sub_ph(x_1, y_1); + _1 = _mm512_fmadd_ph(d_1, d_1, _1); } - return _mm512_reduce_add_ph(sum); + __m512 s_0 = + _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + __m512 s_1 = + _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + __m512 s_2 = + _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + __m512 s_3 = + _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + return _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(s_0, s_2), _mm512_add_ps(s_1, s_3))); } diff --git a/crates/simd/src/byte.rs b/crates/simd/src/byte.rs index 4b94ea93..f553709e 100644 --- a/crates/simd/src/byte.rs +++ b/crates/simd/src/byte.rs @@ -12,13 +12,14 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod reduce_sum_of_xy { - #[doc(hidden)] +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_xy { + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vnni")] - pub fn reduce_sum_of_xy_v4_avx512vnni(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_xy_v4_avx512vnni(lhs: &[u8], rhs: &[u8]) -> u32 { use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); @@ -73,11 +74,11 @@ pub mod reduce_sum_of_xy { } } - #[doc(hidden)] + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - pub fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); @@ -142,11 +143,11 @@ pub mod reduce_sum_of_xy { } } - #[doc(hidden)] + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - pub fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { use crate::emulate::{emulate_mm256_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); @@ -212,11 +213,11 @@ pub mod reduce_sum_of_xy { } } - #[doc(hidden)] + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - pub fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { use crate::emulate::{emulate_mm_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); @@ -282,6 +283,7 @@ pub mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] @@ -325,6 +327,7 @@ pub mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] @@ -403,12 +406,13 @@ pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { reduce_sum_of_xy::reduce_sum_of_xy(s, t) } -pub mod reduce_sum_of_x { - #[doc(hidden)] +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_x { + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - pub fn reduce_sum_of_x_v4(this: &[u8]) -> u32 { + fn reduce_sum_of_x_v4(this: &[u8]) -> u32 { use core::arch::x86_64::*; let i8_1 = _mm512_set1_epi8(1); let i16_1 = _mm512_set1_epi16(1); @@ -450,11 +454,11 @@ pub mod reduce_sum_of_x { } } - #[doc(hidden)] + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - pub fn reduce_sum_of_x_v3(this: &[u8]) -> u32 { + fn reduce_sum_of_x_v3(this: &[u8]) -> u32 { use crate::emulate::{emulate_mm256_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; let i8_1 = _mm256_set1_epi8(1); @@ -497,11 +501,11 @@ pub mod reduce_sum_of_x { } } - #[doc(hidden)] + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - pub fn reduce_sum_of_x_v2(this: &[u8]) -> u32 { + fn reduce_sum_of_x_v2(this: &[u8]) -> u32 { use crate::emulate::{emulate_mm_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; let i8_1 = _mm_set1_epi8(1); @@ -544,6 +548,7 @@ pub mod reduce_sum_of_x { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index cd6c7bcf..deeccb68 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -484,8 +484,9 @@ mod scan { } } - #[cfg(all(target_arch = "s390x", test, not(miri)))] + #[cfg(all(target_arch = "s390x", test))] #[test] + #[cfg(miri, ignore)] fn scan_z13_test() { if !crate::is_cpu_detected!("z13") { println!("test {} ... skipped (z13)", module_path!()); @@ -605,8 +606,9 @@ mod scan { #[cfg(target_arch = "powerpc64")] scan_powerpc64!(scan_p9, "p9"); - #[cfg(all(target_arch = "powerpc64", test, not(miri)))] + #[cfg(all(target_arch = "powerpc64", test))] #[test] + #[cfg(miri, ignore)] fn scan_p9_test() { if !crate::is_cpu_detected!("p9") { println!("test {} ... skipped (p9)", module_path!()); @@ -630,8 +632,9 @@ mod scan { #[cfg(target_arch = "powerpc64")] scan_powerpc64!(scan_p8, "p8"); - #[cfg(all(target_arch = "powerpc64", test, not(miri)))] + #[cfg(all(target_arch = "powerpc64", test))] #[test] + #[cfg(miri, ignore)] fn scan_p8_test() { if !crate::is_cpu_detected!("p8") { println!("test {} ... skipped (p8)", module_path!()); @@ -655,8 +658,9 @@ mod scan { #[cfg(target_arch = "powerpc64")] scan_powerpc64!(scan_p7, "p7"); - #[cfg(all(target_arch = "powerpc64", test, not(miri)))] + #[cfg(all(target_arch = "powerpc64", test))] #[test] + #[cfg(miri, ignore)] fn scan_p7_test() { if !crate::is_cpu_detected!("p7") { println!("test {} ... skipped (p7)", module_path!()); diff --git a/crates/simd/src/floating_f16.rs b/crates/simd/src/floating_f16.rs index fb32c13c..e7b0d7ea 100644 --- a/crates/simd/src/floating_f16.rs +++ b/crates/simd/src/floating_f16.rs @@ -163,12 +163,107 @@ mod reduce_or_of_is_zero_x { } mod reduce_sum_of_x { - // FIXME: add manually-implemented SIMD version - use crate::{F16, f16}; + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + sum = _mm512_add_ps(sum, x); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + sum = _mm512_add_ps(sum, x); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_v4_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let specialized = unsafe { reduce_sum_of_x_v4(&this) }; + let fallback = fallback(&this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_v3(this: &[f16]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_add_ps(sum, x); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_add_ps(sum, x); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_v3_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_x(this: &[f16]) -> f32 { let n = this.len(); @@ -181,12 +276,108 @@ mod reduce_sum_of_x { } mod reduce_sum_of_abs_x { - // FIXME: add manually-implemented SIMD version - use crate::{F16, f16}; + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_abs_x_v4(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + sum = _mm512_add_ps(sum, _mm512_abs_ps(x)); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + sum = _mm512_add_ps(sum, _mm512_abs_ps(x)); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_v4_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let specialized = unsafe { reduce_sum_of_abs_x_v4(&this) }; + let fallback = fallback(&this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_abs_x_v3(this: &[f16]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let abs = _mm256_castsi256_ps(_mm256_srli_epi32(_mm256_set1_epi32(-1), 1)); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_add_ps(sum, _mm256_and_ps(abs, x)); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_add_ps(sum, _mm256_and_ps(abs, x)); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_v3_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { let n = this.len(); @@ -199,12 +390,107 @@ mod reduce_sum_of_abs_x { } mod reduce_sum_of_x2 { - // FIXME: add manually-implemented SIMD version - use crate::{F16, f16}; + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x2_v4(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + sum = _mm512_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + sum = _mm512_fmadd_ps(x, x, sum); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_v4_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let specialized = unsafe { reduce_sum_of_x2_v4(&this) }; + let fallback = fallback(&this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x2_v3(this: &[f16]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_fmadd_ps(x, x, sum); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_v3_test() { + use rand::Rng; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_x2(this: &[f16]) -> f32 { let n = this.len(); @@ -217,12 +503,118 @@ mod reduce_sum_of_x2 { } mod reduce_min_max_of_x { - // FIXME: add manually-implemented SIMD version - use crate::{F16, f16}; + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_min_max_of_x_v4(this: &[f16]) -> (f32, f32) { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm512_set1_ps(f32::INFINITY); + let mut max = _mm512_set1_ps(f32::NEG_INFINITY); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + min = _mm512_min_ps(x, min); + max = _mm512_max_ps(x, max); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + min = _mm512_mask_min_ps(min, mask, x, min); + max = _mm512_mask_max_ps(max, mask, x, max); + } + let min = _mm512_reduce_min_ps(min); + let max = _mm512_reduce_max_ps(max); + (min, max) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_min_max_of_x_v4_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v4(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0); + assert_eq!(specialized.1, fallback.1); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_min_max_of_x_v3(this: &[f16]) -> (f32, f32) { + use crate::emulate::{ + emulate_mm256_reduce_max_ps, emulate_mm256_reduce_min_ps, partial_load, + }; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm256_set1_ps(f32::INFINITY); + let mut max = _mm256_set1_ps(f32::NEG_INFINITY); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a = f16::NAN) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + } + ( + emulate_mm256_reduce_min_ps(min), + emulate_mm256_reduce_max_ps(max), + ) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_min_max_of_x_v3_test() { + use rand::Rng; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v3(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + #[crate::multiversion( - "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { let mut min = f32::INFINITY; @@ -236,14 +628,16 @@ mod reduce_min_max_of_x { } } +#[cfg_attr(feature = "internal", simd_macros::public)] mod reduce_sum_of_xy { use crate::{F16, f16}; + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512fp16")] - pub fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { unsafe extern "C" { #[link_name = "fp16_reduce_sum_of_xy_v4_avx512fp16"] unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; @@ -287,10 +681,11 @@ mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - pub fn reduce_sum_of_xy_v4(lhs: &[f16], rhs: &[f16]) -> f32 { + fn reduce_sum_of_xy_v4(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); use core::arch::x86_64::*; let mut n = lhs.len(); @@ -340,10 +735,11 @@ mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - pub fn reduce_sum_of_xy_v3(lhs: &[f16], rhs: &[f16]) -> f32 { + fn reduce_sum_of_xy_v3(lhs: &[f16], rhs: &[f16]) -> f32 { use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; use core::arch::x86_64::*; assert!(lhs.len() == rhs.len()); @@ -399,6 +795,7 @@ mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(all(target_arch = "aarch64", target_endian = "little"))] #[crate::target_cpu(enable = "a3.512")] @@ -446,11 +843,12 @@ mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "fp16")] - pub fn reduce_sum_of_xy_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + fn reduce_sum_of_xy_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { unsafe extern "C" { #[link_name = "fp16_reduce_sum_of_xy_a2_fp16"] unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; @@ -506,14 +904,16 @@ mod reduce_sum_of_xy { } } +#[cfg_attr(feature = "internal", simd_macros::public)] mod reduce_sum_of_d2 { use crate::{F16, f16}; + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512fp16")] - pub fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { unsafe extern "C" { #[link_name = "fp16_reduce_sum_of_d2_v4_avx512fp16"] unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; @@ -530,7 +930,7 @@ mod reduce_sum_of_d2 { #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_v4_avx512fp16_test() { use rand::Rng; - const EPSILON: f32 = 6.4; + const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { println!("test {} ... skipped (v4:avx512fp16)", module_path!()); return; @@ -557,10 +957,11 @@ mod reduce_sum_of_d2 { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - pub fn reduce_sum_of_d2_v4(lhs: &[f16], rhs: &[f16]) -> f32 { + fn reduce_sum_of_d2_v4(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); use core::arch::x86_64::*; let mut n = lhs.len(); @@ -616,10 +1017,11 @@ mod reduce_sum_of_d2 { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - pub fn reduce_sum_of_d2_v3(lhs: &[f16], rhs: &[f16]) -> f32 { + fn reduce_sum_of_d2_v3(lhs: &[f16], rhs: &[f16]) -> f32 { use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; assert!(lhs.len() == rhs.len()); use core::arch::x86_64::*; @@ -678,6 +1080,7 @@ mod reduce_sum_of_d2 { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(all(target_arch = "aarch64", target_endian = "little"))] #[crate::target_cpu(enable = "a3.512")] @@ -698,7 +1101,7 @@ mod reduce_sum_of_d2 { #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_a3_512_test() { use rand::Rng; - const EPSILON: f32 = 6.4; + const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("a3.512") { println!("test {} ... skipped (a3.512)", module_path!()); return; @@ -725,11 +1128,12 @@ mod reduce_sum_of_d2 { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "fp16")] - pub fn reduce_sum_of_d2_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + fn reduce_sum_of_d2_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { unsafe extern "C" { #[link_name = "fp16_reduce_sum_of_d2_a2_fp16"] unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; @@ -746,7 +1150,7 @@ mod reduce_sum_of_d2 { #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_a2_fp16_test() { use rand::Rng; - const EPSILON: f32 = 6.4; + const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { println!("test {} ... skipped (a2:fp16)", module_path!()); return; diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs index 4605914e..f6a09c74 100644 --- a/crates/simd/src/floating_f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -1051,7 +1051,9 @@ mod reduce_min_max_of_x { } } +#[cfg_attr(feature = "internal", simd_macros::public)] mod reduce_sum_of_xy { + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -1108,6 +1110,7 @@ mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] @@ -1172,6 +1175,7 @@ mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] @@ -1231,6 +1235,7 @@ mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(all(target_arch = "aarch64", target_endian = "little"))] #[crate::target_cpu(enable = "a3.256")] @@ -1278,6 +1283,7 @@ mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] @@ -1349,7 +1355,9 @@ mod reduce_sum_of_xy { } } +#[cfg_attr(feature = "internal", simd_macros::public)] mod reduce_sum_of_d2 { + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -1408,6 +1416,7 @@ mod reduce_sum_of_d2 { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] @@ -1475,6 +1484,7 @@ mod reduce_sum_of_d2 { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] @@ -1536,6 +1546,7 @@ mod reduce_sum_of_d2 { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(all(target_arch = "aarch64", target_endian = "little"))] #[crate::target_cpu(enable = "a3.256")] @@ -1583,6 +1594,7 @@ mod reduce_sum_of_d2 { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] diff --git a/crates/simd/src/halfbyte.rs b/crates/simd/src/halfbyte.rs index ad91a940..f1e80273 100644 --- a/crates/simd/src/halfbyte.rs +++ b/crates/simd/src/halfbyte.rs @@ -12,13 +12,14 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod reduce_sum_of_xy { - #[doc(hidden)] +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_xy { + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512vnni")] - pub fn reduce_sum_of_xy_v4_avx512vnni(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_xy_v4_avx512vnni(lhs: &[u8], rhs: &[u8]) -> u32 { use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); @@ -79,11 +80,11 @@ pub mod reduce_sum_of_xy { } } - #[doc(hidden)] + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] - pub fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); let mut n = lhs.len(); @@ -149,11 +150,11 @@ pub mod reduce_sum_of_xy { } } - #[doc(hidden)] + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v3")] - pub fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { use crate::emulate::{emulate_mm256_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); @@ -220,11 +221,11 @@ pub mod reduce_sum_of_xy { } } - #[doc(hidden)] + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v2")] - pub fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { + fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { use crate::emulate::{emulate_mm_reduce_add_epi32, partial_load}; use core::arch::x86_64::*; assert_eq!(lhs.len(), rhs.len()); @@ -291,6 +292,7 @@ pub mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] @@ -334,6 +336,7 @@ pub mod reduce_sum_of_xy { } } + #[cfg_attr(feature = "internal", simd_macros::public)] #[inline] #[cfg(target_arch = "aarch64")] #[crate::target_cpu(enable = "a2")] diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index c967e548..c9ed2a75 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -24,9 +24,19 @@ mod aligned; mod emulate; + +#[cfg(not(feature = "internal"))] mod floating_f16; + +#[cfg(feature = "internal")] +pub mod floating_f16; + +#[cfg(not(feature = "internal"))] mod floating_f32; +#[cfg(feature = "internal")] +pub mod floating_f32; + pub mod bit; pub mod byte; pub mod fast_scan; @@ -53,10 +63,12 @@ pub trait F16: Sized { impl F16 for f16 { const _ZERO: Self = f16::ZERO; + #[inline(always)] fn _from_f32(x: f32) -> Self { f16::from_f32(x) } + #[inline(always)] fn _to_f32(self) -> f32 { self.to_f32() } @@ -66,10 +78,12 @@ impl F16 for f16 { impl F16 for f16 { const _ZERO: Self = 0.0; + #[inline(always)] fn _from_f32(x: f32) -> Self { x as _ } + #[inline(always)] fn _to_f32(self) -> f32 { self as _ } diff --git a/crates/simd/src/rotate.rs b/crates/simd/src/rotate.rs index 15916229..1a8304c9 100644 --- a/crates/simd/src/rotate.rs +++ b/crates/simd/src/rotate.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod givens { +mod givens { #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] @@ -30,7 +30,7 @@ pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { givens::givens(lhs, rhs) } -pub mod flip { +mod flip { #[crate::multiversion( "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs index 094a5066..9152fd6a 100644 --- a/crates/simd_macros/src/lib.rs +++ b/crates/simd_macros/src/lib.rs @@ -259,3 +259,21 @@ pub fn define_is_cpu_detected(input: proc_macro::TokenStream) -> proc_macro::Tok } .into() } + +#[proc_macro_attribute] +pub fn public( + _attr: proc_macro::TokenStream, + item: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + use quote::ToTokens; + + let mut input: syn::Item = syn::parse_macro_input!(item); + + if let syn::Item::Fn(syn::ItemFn { ref mut vis, .. }) + | syn::Item::Mod(syn::ItemMod { ref mut vis, .. }) = input + { + *vis = syn::Visibility::Public(Default::default()); + } + + input.to_token_stream().into() +} From 1b20c583dae293ce0cf66bf27cbbad7205afd654 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 8 Jan 2026 14:22:52 +0800 Subject: [PATCH 263/324] fix: lto (#402) https://github.com/rust-lang/rust/issues/51009 Signed-off-by: usamoi --- crates/xtask/src/main.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index 59f1d11a..006be754 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -205,8 +205,11 @@ fn build( target: &str, ) -> Result> { let mut command = Command::new("cargo"); + command.args(["rustc"]); + if !matches!(profile, "dev" | "test") { + command.args(["--crate-type", "cdylib"]); + } command - .args(["build"]) .args(["-p", "vchord", "--lib"]) .args(["--profile", profile]) .args(["--target", target]) From 0854c036d0bcee7a0009813f28a46cd28ab2dbb0 Mon Sep 17 00:00:00 2001 From: Rene Leonhardt <65483435+reneleonhardt@users.noreply.github.com> Date: Sat, 10 Jan 2026 09:24:49 +0100 Subject: [PATCH 264/324] chore: update dependencies (#403) * removes PostgreSQL 13 support and CI * configures dependabot for github actions and cargo * updates `cargo-deny` and `sqllogictest` in CI --------- Signed-off-by: Rene Leonhardt <65483435+reneleonhardt@users.noreply.github.com> --- .github/dependabot.yml | 23 ++++++ .github/workflows/check.yml | 87 ++++++++++------------- .github/workflows/cla.yml | 2 +- .github/workflows/release.yml | 8 +-- Cargo.lock | 53 ++++++++++---- Cargo.toml | 3 +- META.json.in | 4 +- crates/xtask/Cargo.toml | 2 +- crates/xtask/src/main.rs | 2 +- src/index/gucs.rs | 12 ++-- src/index/sample.rs | 1 - src/index/scanners.rs | 2 +- src/index/storage.rs | 10 +-- src/index/storage/buffered.rs | 2 +- src/index/vchordg/am/am_build.rs | 8 +-- src/index/vchordg/am/mod.rs | 30 +------- src/index/vchordrq/am/am_build.rs | 8 +-- src/index/vchordrq/am/am_vacuumcleanup.rs | 8 +-- src/index/vchordrq/am/mod.rs | 22 +----- src/lib.rs | 2 +- tests/vchordrq/pushdown_plan.slt | 14 ---- 21 files changed, 130 insertions(+), 173 deletions(-) create mode 100644 .github/dependabot.yml diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 00000000..0de32ee5 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,23 @@ +version: 2 +updates: + - package-ecosystem: cargo + directory: / + schedule: + interval: weekly + # https://blog.yossarian.net/2025/11/21/We-should-all-be-using-dependency-cooldowns +# cooldown: # applies only to non-security updates +# semver-patch: 7 # wait 7 days before applying patch updates +# semver-minor: 14 +# semver-major: 28 + # dependabot is forbidden to suggest patch / minor / major updates and let CI test for regressions automatically + groups: + cargo: + applies-to: version-updates + exclude-patterns: ["*"] + + - package-ecosystem: github-actions + directory: / + schedule: + interval: weekly + cooldown: + default-days: 7 diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 79e6ea65..9c3a6b4e 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -23,13 +23,13 @@ jobs: run: | curl -fsSL https://github.com/tamasfe/taplo/releases/download/0.10.0/taplo-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo - curl -fsSL https://github.com/EmbarkStudios/cargo-deny/releases/download/0.19.2/cargo-deny-0.18.2-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - cargo-deny-0.19.2-$(uname -m)-unknown-linux-musl/cargo-deny | install -m 755 /dev/stdin /usr/local/bin/cargo-deny + curl -fsSL https://github.com/EmbarkStudios/cargo-deny/releases/download/0.18.9/cargo-deny-0.18.9-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - cargo-deny-0.18.9-$(uname -m)-unknown-linux-musl/cargo-deny | install -m 755 /dev/stdin /usr/local/bin/cargo-deny - name: Install Rust run: rustup update - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Typos uses: crate-ci/typos@master @@ -38,7 +38,7 @@ jobs: run: taplo fmt --check - name: Ruff - uses: astral-sh/ruff-action@v1 + uses: astral-sh/ruff-action@57714a7c8a2e59f32539362ba31877a1957dded1 # v3.5.1 - name: Rustfmt run: cargo fmt --check -- --config-path /dev/null --config imports_granularity=Module @@ -130,7 +130,7 @@ jobs: echo CC=clang-18 >> $GITHUB_ENV - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Cargo Test (Miri) env: @@ -166,9 +166,9 @@ jobs: sudo apt-get update if [ "$(uname -m)" == "x86_64" ]; then - wget https://downloadmirror.intel.com/843185/sde-external-9.48.0-2024-11-25-lin.tar.xz -O /tmp/sde-external.tar.xz + wget https://downloadmirror.intel.com/859732/sde-external-9.58.0-2025-06-16-lin.tar.xz -O /tmp/sde-external.tar.xz sudo tar -xf /tmp/sde-external.tar.xz -C /opt - sudo mv /opt/sde-external-9.48.0-2024-11-25-lin /opt/sde + sudo mv /opt/sde-external-9.58.0-2025-06-16-lin /opt/sde fi if [ "$(uname -m)" == "aarch64" ]; then @@ -184,10 +184,10 @@ jobs: echo CC=clang-18 >> $GITHUB_ENV - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Cargo Test run: cargo test --locked --workspace --exclude vchord --no-fail-fast -- --no-capture @@ -222,7 +222,7 @@ jobs: strategy: matrix: - version: ["13", "14", "15", "16", "17", "18"] + version: ["14", "15", "16", "17", "18"] arch: ["x86_64", "aarch64"] runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} @@ -239,7 +239,7 @@ jobs: run: | sudo apt-get update - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - name: Install Rust run: rustup update @@ -250,7 +250,7 @@ jobs: echo CC=clang-18 >> $GITHUB_ENV - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - name: Install PostgreSQL & pgvector run: | @@ -287,7 +287,7 @@ jobs: sudo apt-get install -y --no-install-recommends postgresql-${{ matrix.version }}-pgvector - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Clippy run: make PROFILE=dev clippy @@ -299,7 +299,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu path: ./build @@ -333,7 +333,7 @@ jobs: strategy: matrix: - version: ["13", "14", "15", "16", "17", "18"] + version: ["14", "15", "16", "17", "18"] arch: ["aarch64", "x86_64"] runs-on: ${{ matrix.arch == 'aarch64' && 'macos-15' || 'macos-15-intel' }} @@ -351,7 +351,7 @@ jobs: run: | brew update - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest - name: Install Rust run: rustup update @@ -360,7 +360,7 @@ jobs: run: echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - name: Install PostgreSQL & pgvector run: | @@ -383,7 +383,7 @@ jobs: sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' install - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Clippy run: make PROFILE=dev clippy @@ -395,7 +395,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin path: ./build @@ -431,7 +431,7 @@ jobs: strategy: matrix: - version: ["13", "14", "15", "16", "17", "18"] + version: ["14", "15", "16", "17", "18"] arch: ["x86_64"] runs-on: "windows-2022" @@ -447,7 +447,7 @@ jobs: steps: - name: Set up Environment run: | - Invoke-WebRequest -Uri https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-${{ matrix.arch }}-pc-windows-msvc.zip -OutFile "$env:TEMP\sqllogictest-install.zip" + Invoke-WebRequest -Uri https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-${{ matrix.arch }}-pc-windows-msvc.zip -OutFile "$env:TEMP\sqllogictest-install.zip" Expand-Archive -Path "$env:TEMP\sqllogictest-install.zip" -DestinationPath "D:\sqllogictest-install" -Force Add-Content -Path $env:GITHUB_PATH -Value "D:\sqllogictest-install" @@ -455,29 +455,26 @@ jobs: run: rustup update - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - name: Install PostgreSQL & pgvector run: | 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } - if ( "${{ matrix.version }}" -eq "13" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-13.22-1-windows-x64-binaries.zip" - } if ( "${{ matrix.version }}" -eq "14" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-14.19-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-14.20-2-windows-x64-binaries.zip" } if ( "${{ matrix.version }}" -eq "15" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-15.14-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-15.15-2-windows-x64-binaries.zip" } if ( "${{ matrix.version }}" -eq "16" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-16.10-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-16.11-2-windows-x64-binaries.zip" } if ( "${{ matrix.version }}" -eq "17" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-17.6-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-17.7-2-windows-x64-binaries.zip" } if ( "${{ matrix.version }}" -eq "18" ) { - $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-18.0-1-windows-x64-binaries.zip" + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-18.1-2-windows-x64-binaries.zip" } Invoke-WebRequest -Uri $postgresql_url -OutFile "$env:TEMP\postgresql-install.zip" Expand-Archive -Path "$env:TEMP\postgresql-install.zip" -DestinationPath "D:\postgresql-install" -Force @@ -506,7 +503,7 @@ jobs: run: Add-Content -Path $env:GITHUB_ENV -Value "CC=clang-cl.exe" - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Clippy run: make PROFILE=dev clippy @@ -518,7 +515,7 @@ jobs: run: make install - name: Upload Artifacts - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc path: ./build @@ -579,7 +576,7 @@ jobs: mkdir /opt/bin ln -s /usr/bin/node /opt/bin/node - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - name: Install Rust run: | @@ -593,7 +590,7 @@ jobs: echo CC=clang-18 >> $GITHUB_ENV - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - name: Install PostgreSQL & pgvector run: | @@ -626,7 +623,7 @@ jobs: make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' install - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Patch run: | @@ -648,7 +645,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-musl path: ./build @@ -695,12 +692,6 @@ jobs: rust_triple: "s390x-unknown-linux-gnu" gcc_version: "14" debian_version: "trixie" - - version: "13" - platform: "ppc64le" - clang_triple: "powerpc64le-linux-gnu" - rust_triple: "powerpc64le-unknown-linux-gnu" - gcc_version: "12" - debian_version: "bookworm" - version: "14" platform: "ppc64le" clang_triple: "powerpc64le-linux-gnu" @@ -780,7 +771,7 @@ jobs: echo QEMU_LD_PREFIX=$QEMU_LD_PREFIX >> $GITHUB_ENV echo QEMU_CPU=$QEMU_CPU >> $GITHUB_ENV - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest - name: Install Rust run: | @@ -795,7 +786,7 @@ jobs: echo CC=clang-18 >> $GITHUB_ENV - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - name: Install PostgreSQL & pgvector run: | @@ -838,7 +829,7 @@ jobs: sudo make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC='sccache clang' CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Patch run: | @@ -871,7 +862,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.clang_triple }} path: ./build @@ -938,7 +929,7 @@ jobs: sudo pacman -S --noconfirm liburing numactl python valgrind - curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.28.4/sqllogictest-bin-v0.28.4-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | sudo install -m 755 /dev/stdin /usr/local/bin/sqllogictest + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | sudo install -m 755 /dev/stdin /usr/local/bin/sqllogictest - name: Install Rust run: | @@ -952,7 +943,7 @@ jobs: echo CC=clang >> $GITHUB_ENV - name: Set up Sccache - uses: mozilla-actions/sccache-action@v0.0.9 + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - name: Install PostgreSQL & pgvector run: | @@ -1035,7 +1026,7 @@ jobs: popd - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Cargo Test run: cargo test --locked --workspace --exclude vchord --no-fail-fast -- --no-capture diff --git a/.github/workflows/cla.yml b/.github/workflows/cla.yml index 7bc964da..97e090b8 100644 --- a/.github/workflows/cla.yml +++ b/.github/workflows/cla.yml @@ -18,7 +18,7 @@ jobs: steps: - name: "CLA Assistant" if: (github.event.comment.body == 'recheck' || github.event.comment.body == 'I have read the CLA Document and I hereby sign the CLA') || github.event_name == 'pull_request_target' - uses: contributor-assistant/github-action@v2.6.1 + uses: contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08 # v2.6.1 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # the below token should have repo scope and must be manually added by you in the repository's secret diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index afddd91d..dc89b198 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -21,7 +21,7 @@ jobs: steps: - name: Semver id: semver - uses: actions/github-script@v7 + uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0 with: script: | const tag = "${{ github.event.inputs.tag }}" || "${{ github.event.release.tag_name }}"; @@ -39,7 +39,7 @@ jobs: needs: ["semver"] strategy: matrix: - version: ["13", "14", "15", "16", "17", "18"] + version: ["14", "15", "16", "17", "18"] arch: ["x86_64", "aarch64"] runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} @@ -74,7 +74,7 @@ jobs: pg_config - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Build run: | @@ -122,7 +122,7 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - name: Upload env: diff --git a/Cargo.lock b/Cargo.lock index 580b26c1..c966fe45 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -580,6 +580,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -636,22 +642,25 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ - "foldhash", + "foldhash 0.1.5", ] [[package]] name = "hashbrown" -version = "0.16.0" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "foldhash 0.2.0", +] [[package]] name = "hashlink" -version = "0.10.0" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7382cf6263419f2d8df38c55d7da83da5c18aef87fc7a7fc1fb1e344edfe14c1" +checksum = "ea0b22561a9c04a7cb1a302c013e0259cd3b4bb619f145b32f72b8b4bcbed230" dependencies = [ - "hashbrown 0.15.5", + "hashbrown 0.16.1", ] [[package]] @@ -805,7 +814,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6717a8d2a5a929a1a2eb43a12812498ed141a0bcfb7e8f7844fbdbe4303bba9f" dependencies = [ "equivalent", - "hashbrown 0.16.0", + "hashbrown 0.16.1", ] [[package]] @@ -891,9 +900,9 @@ dependencies = [ [[package]] name = "libsqlite3-sys" -version = "0.35.0" +version = "0.36.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "133c182a6a2c87864fe97778797e46c7e999672690dc9fa3ee8e241aa4a9c13f" +checksum = "95b4103cffefa72eb8428cb6b47d6627161e51c2739fc5e3b734584157bc642a" dependencies = [ "cc", "pkg-config", @@ -970,9 +979,9 @@ dependencies = [ [[package]] name = "object" -version = "0.37.3" +version = "0.38.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" +checksum = "271638cd5fa9cca89c4c304675ca658efc4e64a66c716b7cfe1afb4b9611dbbc" dependencies = [ "flate2", "memchr", @@ -1377,9 +1386,9 @@ checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" [[package]] name = "rusqlite" -version = "0.37.0" +version = "0.38.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "165ca6e57b20e1351573e3729b958bc62f0e48025386970b6e4d29e7a7e71f3f" +checksum = "f1c93dd1c9683b438c392c492109cb702b8090b2bfc8fed6f6e4eb4523f17af3" dependencies = [ "bitflags", "fallible-iterator", @@ -1387,6 +1396,7 @@ dependencies = [ "hashlink", "libsqlite3-sys", "smallvec", + "sqlite-wasm-rs", ] [[package]] @@ -1557,6 +1567,19 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +[[package]] +name = "sqlite-wasm-rs" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05e98301bf8b0540c7de45ecd760539b9c62f5772aed172f08efba597c11cd5d" +dependencies = [ + "cc", + "hashbrown 0.16.1", + "js-sys", + "thiserror", + "wasm-bindgen", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -1949,9 +1972,9 @@ dependencies = [ [[package]] name = "wasmparser" -version = "0.236.1" +version = "0.243.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9b1e81f3eb254cf7404a82cee6926a4a3ccc5aad80cc3d43608a070c67aa1d7" +checksum = "f6d8db401b0528ec316dfbe579e6ab4152d61739cfe076706d2009127970159d" dependencies = [ "bitflags", ] diff --git a/Cargo.toml b/Cargo.toml index 1e952a64..26b539bc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,7 +14,6 @@ path = "./src/bin/pgrx_embed.rs" [features] default = ["simd/init"] -pg13 = ["pgrx/pg13", "pgrx-catalog/pg13"] pg14 = ["pgrx/pg14", "pgrx-catalog/pg14"] pg15 = ["pgrx/pg15", "pgrx-catalog/pg15"] pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] @@ -44,7 +43,7 @@ pgrx = { version = "=0.16.1", default-features = false, features = ["cshim"] } pgrx-catalog = "0.3.1" rand.workspace = true rayon.workspace = true -rusqlite = { version = "0.37.0", features = ["bundled"] } +rusqlite = { version = "0.38.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true toml = "0.9.8" diff --git a/META.json.in b/META.json.in index a823cf15..ef79c686 100644 --- a/META.json.in +++ b/META.json.in @@ -11,7 +11,7 @@ "prereqs": { "runtime": { "requires": { - "PostgreSQL": "13.0.0" + "PostgreSQL": "14.0.0" } } }, @@ -46,4 +46,4 @@ "nearest neighbor search", "approximate nearest neighbors" ] -} \ No newline at end of file +} diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index 3e03f42f..0dd5fa75 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -6,7 +6,7 @@ publish = false [dependencies] clap = { version = "4.5.51", features = ["derive", "env"] } -object = { version = "0.37.3", features = ["read", "wasm"] } +object = { version = "0.38.1", features = ["read", "wasm"] } serde.workspace = true serde_json = "1.0.145" shlex = "1.3.0" diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index 006be754..f3102340 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -100,7 +100,7 @@ impl RustcCfg { } fn ext_suffix(&self, fork: &str) -> Result<&'static str, Box> { if self.is_macos { - Ok(if matches!(fork, "pg13" | "pg14" | "pg15") { + Ok(if matches!(fork, "pg14" | "pg15") { ".so" } else { ".dylib" diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 5b74bc2b..f29ede82 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -42,14 +42,14 @@ static VCHORDG_BEAM_SEARCH: GucSetting = GucSetting::::new(1); static VCHORDG_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); static VCHORDG_IO_SEARCH: GucSetting = GucSetting::::new( - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] PostgresIo::PrefetchBuffer, #[cfg(any(feature = "pg17", feature = "pg18"))] PostgresIo::ReadStream, ); static VCHORDG_IO_RERANK: GucSetting = GucSetting::::new( - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] PostgresIo::PrefetchBuffer, #[cfg(any(feature = "pg17", feature = "pg18"))] PostgresIo::ReadStream, @@ -70,14 +70,14 @@ static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); static VCHORDRQ_PREFILTER: GucSetting = GucSetting::::new(false); static VCHORDRQ_IO_SEARCH: GucSetting = GucSetting::::new( - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] PostgresIo::PrefetchBuffer, #[cfg(any(feature = "pg17", feature = "pg18"))] PostgresIo::ReadStream, ); static VCHORDRQ_IO_RERANK: GucSetting = GucSetting::::new( - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] PostgresIo::PrefetchBuffer, #[cfg(any(feature = "pg17", feature = "pg18"))] PostgresIo::ReadStream, @@ -193,7 +193,7 @@ pub fn init() { GucFlags::default(), ); unsafe { - #[cfg(any(feature = "pg13", feature = "pg14"))] + #[cfg(feature = "pg14")] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrq".as_ptr()); @@ -253,7 +253,7 @@ pub fn init() { GucFlags::default(), ); unsafe { - #[cfg(any(feature = "pg13", feature = "pg14"))] + #[cfg(feature = "pg14")] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordg".as_ptr()); #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordg".as_ptr()); diff --git a/src/index/sample.rs b/src/index/sample.rs index bbdbc3f5..2d8c6d3e 100644 --- a/src/index/sample.rs +++ b/src/index/sample.rs @@ -264,7 +264,6 @@ static TSM: AssertSync = AssertSync(sys::TsmRoutine { #[allow(non_camel_case_types)] #[allow(non_snake_case)] mod sys { - #[cfg(not(feature = "pg13"))] #[cfg(not(feature = "pg14"))] #[cfg(not(feature = "pg15"))] #[cfg(not(feature = "pg16"))] diff --git a/src/index/scanners.rs b/src/index/scanners.rs index e0d11c7a..eaf671a0 100644 --- a/src/index/scanners.rs +++ b/src/index/scanners.rs @@ -23,7 +23,7 @@ pub enum Io { Plain, Simple, #[cfg_attr( - any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"), + any(feature = "pg14", feature = "pg15", feature = "pg16"), expect(dead_code) )] Stream, diff --git a/src/index/storage.rs b/src/index/storage.rs index d6896e97..e91039b6 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -349,7 +349,7 @@ impl RelationWrite for PostgresRelation { GENERIC_XLOG_FULL_IMAGE, GenericXLogRegisterBuffer, GenericXLogStart, }; let buf; - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15"))] + #[cfg(any(feature = "pg14", feature = "pg15"))] { use pgrx::pg_sys::{ BUFFER_LOCK_EXCLUSIVE, ExclusiveLock, ForkNumber, LockBuffer, @@ -515,7 +515,7 @@ where { type Item = PostgresBufferReadGuard; - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] fn next(&mut self) -> Option { unimplemented!() } @@ -575,12 +575,12 @@ where type Inner = Chain, Flatten>>; - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] fn next(&mut self) -> Option<(I::Item, Self::Guards)> { panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); } - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] fn next_if bool>( &mut self, _predicate: P, @@ -641,7 +641,7 @@ impl RelationReadStreamTypes for PostgresRelation { } impl RelationReadStream for PostgresRelation { - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] fn read_stream<'b, I: Iterator>(&'b self, _iter: I, _hints: Hints) -> Self::ReadStream<'b, I> where I::Item: Fetch<'b>, diff --git a/src/index/storage/buffered.rs b/src/index/storage/buffered.rs index 151987a3..6c438179 100644 --- a/src/index/storage/buffered.rs +++ b/src/index/storage/buffered.rs @@ -76,7 +76,7 @@ impl RelationWrite for BufferedPostgresRelation { fn search(&self, freespace: usize) -> Option> { self.postgres.search(freespace) } - #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15"))] + #[cfg(any(feature = "pg14", feature = "pg15"))] fn extend(&self, opaque: O, tracking_freespace: bool) -> Self::WriteGuard<'_> { self.postgres.extend(opaque, tracking_freespace) } diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index 78609109..edcef0ba 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -560,13 +560,7 @@ unsafe fn options( index_relation: pgrx::pg_sys::Relation, ) -> (VectorOptions, VchordgIndexingOptions) { let att = unsafe { &mut *(*index_relation).rd_att }; - #[cfg(any( - feature = "pg13", - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17" - ))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] let atts = unsafe { att.attrs.as_slice(att.natts as _) }; #[cfg(feature = "pg18")] let atts = unsafe { diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index 6b2a385a..06c47e9c 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -182,20 +182,6 @@ pub unsafe extern "C-unwind" fn amcostestimate( } } -#[cfg(feature = "pg13")] -#[pgrx::pg_guard] -pub unsafe extern "C-unwind" fn aminsert( - index_relation: pgrx::pg_sys::Relation, - values: *mut Datum, - is_null: *mut bool, - heap_tid: pgrx::pg_sys::ItemPointer, - heap_relation: pgrx::pg_sys::Relation, - _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, - _index_info: *mut pgrx::pg_sys::IndexInfo, -) -> bool { - unsafe { aminsertinner(index_relation, heap_relation, values, is_null, heap_tid) } -} - #[cfg(any( feature = "pg14", feature = "pg15", @@ -255,13 +241,7 @@ pub unsafe extern "C-unwind" fn ambulkdelete( let opfamily = unsafe { opfamily((*info).index) }; let index = unsafe { PostgresRelation::new((*info).index) }; let check = || unsafe { - #[cfg(any( - feature = "pg13", - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17" - ))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] pgrx::pg_sys::vacuum_delay_point(); #[cfg(feature = "pg18")] pgrx::pg_sys::vacuum_delay_point(false); @@ -293,13 +273,7 @@ pub unsafe extern "C-unwind" fn amvacuumcleanup( let opfamily = unsafe { opfamily((*info).index) }; let index = unsafe { PostgresRelation::new((*info).index) }; let check = || unsafe { - #[cfg(any( - feature = "pg13", - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17" - ))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] pgrx::pg_sys::vacuum_delay_point(); #[cfg(feature = "pg18")] pgrx::pg_sys::vacuum_delay_point(false); diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 4da27b65..b3ae1eb6 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -1212,13 +1212,7 @@ unsafe fn options( index_relation: pgrx::pg_sys::Relation, ) -> (VectorOptions, VchordrqIndexingOptions) { let att = unsafe { &mut *(*index_relation).rd_att }; - #[cfg(any( - feature = "pg13", - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17" - ))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] let atts = unsafe { att.attrs.as_slice(att.natts as _) }; #[cfg(feature = "pg18")] let atts = unsafe { diff --git a/src/index/vchordrq/am/am_vacuumcleanup.rs b/src/index/vchordrq/am/am_vacuumcleanup.rs index c7e4f835..193bedc0 100644 --- a/src/index/vchordrq/am/am_vacuumcleanup.rs +++ b/src/index/vchordrq/am/am_vacuumcleanup.rs @@ -49,13 +49,7 @@ unsafe fn sequential_vacuumcleanup( let opfamily = unsafe { opfamily(index_relation) }; let index = unsafe { PostgresRelation::new(index_relation) }; let check = || unsafe { - #[cfg(any( - feature = "pg13", - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17" - ))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] pgrx::pg_sys::vacuum_delay_point(); #[cfg(feature = "pg18")] pgrx::pg_sys::vacuum_delay_point(false); diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 8fa36323..125d1874 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -273,20 +273,6 @@ pub unsafe extern "C-unwind" fn amcostestimate( } } -#[cfg(feature = "pg13")] -#[pgrx::pg_guard] -pub unsafe extern "C-unwind" fn aminsert( - index_relation: pgrx::pg_sys::Relation, - values: *mut Datum, - is_null: *mut bool, - heap_tid: pgrx::pg_sys::ItemPointer, - _heap_relation: pgrx::pg_sys::Relation, - _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, - _index_info: *mut pgrx::pg_sys::IndexInfo, -) -> bool { - unsafe { aminsertinner(index_relation, values, is_null, heap_tid) } -} - #[cfg(any( feature = "pg14", feature = "pg15", @@ -363,13 +349,7 @@ pub unsafe extern "C-unwind" fn ambulkdelete( let opfamily = unsafe { opfamily((*info).index) }; let index = unsafe { PostgresRelation::new((*info).index) }; let check = || unsafe { - #[cfg(any( - feature = "pg13", - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17" - ))] + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] pgrx::pg_sys::vacuum_delay_point(); #[cfg(feature = "pg18")] pgrx::pg_sys::vacuum_delay_point(false); diff --git a/src/lib.rs b/src/lib.rs index 594d7824..6a18e50d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -55,7 +55,7 @@ unsafe extern "C-unwind" fn _pg_init() { index::init(); recorder::init(); unsafe { - #[cfg(any(feature = "pg13", feature = "pg14"))] + #[cfg(feature = "pg14")] pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] pgrx::pg_sys::MarkGUCPrefixReserved(c"vchord".as_ptr()); diff --git a/tests/vchordrq/pushdown_plan.slt b/tests/vchordrq/pushdown_plan.slt index b586edf2..62327229 100644 --- a/tests/vchordrq/pushdown_plan.slt +++ b/tests/vchordrq/pushdown_plan.slt @@ -29,14 +29,6 @@ SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) ORDER BY val0 # 1 vector key + 0 order_by key + original style -onlyif pg13 -query I -EXPLAIN (COSTS FALSE, TIMING FALSE) -SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; ----- - Seq Scan on t - Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) - onlyif pg14 query I EXPLAIN (COSTS FALSE, TIMING FALSE) @@ -179,12 +171,6 @@ ORDER BY val0 <-> '[0, 0, 0]'; # 0 vector key + 0 order_by key(variable) -onlyif pg13 -query I -EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; ----- - Seq Scan on t - onlyif pg14 query I EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; From 97732d081dda6ae6b312d6995c93fba9cd192c18 Mon Sep 17 00:00:00 2001 From: usamoi Date: Sat, 10 Jan 2026 16:50:43 +0800 Subject: [PATCH 265/324] fix: remove dependabot for cargo (#409) It looks like Dependabot also opens PRs for non-security updates, which is really annoying, so disable it. Signed-off-by: usamoi --- .github/dependabot.yml | 15 --- Cargo.lock | 240 ++++++++++++++-------------------------- Cargo.toml | 6 +- crates/simd/Cargo.toml | 2 +- crates/xtask/Cargo.toml | 6 +- 5 files changed, 90 insertions(+), 179 deletions(-) diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 0de32ee5..9cbb9b7c 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -1,20 +1,5 @@ version: 2 updates: - - package-ecosystem: cargo - directory: / - schedule: - interval: weekly - # https://blog.yossarian.net/2025/11/21/We-should-all-be-using-dependency-cooldowns -# cooldown: # applies only to non-security updates -# semver-patch: 7 # wait 7 days before applying patch updates -# semver-minor: 14 -# semver-major: 28 - # dependabot is forbidden to suggest patch / minor / major updates and let CI test for regressions automatically - groups: - cargo: - applies-to: version-updates - exclude-patterns: ["*"] - - package-ecosystem: github-actions directory: / schedule: diff --git a/Cargo.lock b/Cargo.lock index c966fe45..facac879 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -78,22 +78,22 @@ dependencies = [ [[package]] name = "anstyle-query" -version = "1.1.4" +version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e231f6134f61b71076a3eab506c379d4f36122f2af15a9ff04415ea4c3339e2" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys", ] [[package]] name = "anstyle-wincon" -version = "3.0.10" +version = "3.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e0633414522a32ffaac8ac6cc8f748e090c5717661fddeea04219e2344f5f2a" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys", ] [[package]] @@ -153,9 +153,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.19.0" +version = "3.19.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43" +checksum = "5dd9dc738b7a8311c7ade152424974d8115f2cdad61e8dab8dac9f2362298510" [[package]] name = "cargo_toml" @@ -175,9 +175,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.2.44" +version = "1.2.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37521ac7aabe3d13122dc382493e20c9416f299d2ccd5b3a5340a2570cdeb0f3" +checksum = "cd4932aefd12402b36c60956a4fe0035421f544799057659ff86f923657aada3" dependencies = [ "find-msvc-tools", "shlex", @@ -248,9 +248,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.51" +version = "4.5.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c26d721170e0295f191a69bd9a1f93efcdb0aff38684b61ab5750468972e5f5" +checksum = "c6e6ff9dcd79cff5cd969a17a545d79e84ab086e444102a591e288a8aa3ce394" dependencies = [ "clap_builder", "clap_derive", @@ -258,9 +258,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.51" +version = "4.5.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75835f0c7bf681bfd05abe44e965760fea999a5286c6eb2d59883634fd02011a" +checksum = "fa42cf4d2b7a41bc8f663a7cab4031ebafa1bf3875705bfaf8466dc60ab52c00" dependencies = [ "anstream", "anstyle", @@ -507,7 +507,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys", ] [[package]] @@ -548,9 +548,9 @@ dependencies = [ [[package]] name = "find-msvc-tools" -version = "0.1.4" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52051878f80a721bb68ebfbc930e07b65ba72f2da88968ea5c06fd6ca3d3a127" +checksum = "f449e6c6c08c865631d4890cfacf252b3d396c9bcc83adb6623cdb02a8336c41" [[package]] name = "fixedbitset" @@ -726,9 +726,9 @@ checksum = "7aedcccd01fc5fe81e6b489c15b247b8b0690feb23304303a9e560f37efc560a" [[package]] name = "icu_properties" -version = "2.1.1" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e93fcd3157766c0c8da2f8cff6ce651a31f0810eaa1c51ec363ef790bbb5fb99" +checksum = "020bfc02fe870ec3a66d93e677ccca0562506e5872c650f893269e08615d74ec" dependencies = [ "icu_collections", "icu_locale_core", @@ -740,9 +740,9 @@ dependencies = [ [[package]] name = "icu_properties_data" -version = "2.1.1" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02845b3647bb045f1100ecd6480ff52f34c35f82d9880e029d329c21d1054899" +checksum = "616c294cf8d725c6afcd8f55abc17c56464ef6211f9ed59cccffe534129c77af" [[package]] name = "icu_provider" @@ -809,9 +809,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.12.0" +version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6717a8d2a5a929a1a2eb43a12812498ed141a0bcfb7e8f7844fbdbe4303bba9f" +checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" dependencies = [ "equivalent", "hashbrown 0.16.1", @@ -840,15 +840,15 @@ dependencies = [ [[package]] name = "itoa" -version = "1.0.15" +version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" +checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" [[package]] name = "js-sys" -version = "0.3.82" +version = "0.3.83" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b011eec8cc36da2aab2d5cff675ec18454fad408585853910a202391cf9f8e65" +checksum = "464a3709c7f55f1f721e5389aa6ea4e3bc6aba669353300af094b29ffbdde1d8" dependencies = [ "once_cell", "wasm-bindgen", @@ -868,9 +868,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.177" +version = "0.2.180" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976" +checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc" [[package]] name = "libloading" @@ -1268,18 +1268,18 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.103" +version = "1.0.105" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ee95bc4ef87b8d5ba32e8b7714ccc834865276eab0aed5c9958d00ec45f49e8" +checksum = "535d180e0ecab6268a3e718bb9fd44db66bbbc256257165fc699dadf70d16fe7" dependencies = [ "unicode-ident", ] [[package]] name = "quote" -version = "1.0.42" +version = "1.0.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a338cc41d27e6cc6dce6cefc13a0729dfbb81c262b1f519331575dd80ef3067f" +checksum = "dc74d9a594b72ae6656596548f56f667211f8a97b3d4c3d467150794690dc40a" dependencies = [ "proc-macro2", ] @@ -1407,15 +1407,15 @@ checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" [[package]] name = "rustix" -version = "1.1.2" +version = "1.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd15f8a2c5551a84d56efdc1cd049089e409ac19a3072d5037a17fd70719ff3e" +checksum = "146c9e247ccc180c1f61615433868c99f3de3ae256a30a43b49f67c2d9171f34" dependencies = [ "bitflags", "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys", ] [[package]] @@ -1433,12 +1433,6 @@ dependencies = [ "twox-hash", ] -[[package]] -name = "ryu" -version = "1.0.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" - [[package]] name = "same-file" version = "1.0.6" @@ -1502,22 +1496,22 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.145" +version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "402a6f66d8c709116cf22f558eab210f5a50187f702eb4d7e5ef38d9a7f1c79c" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ "itoa", "memchr", - "ryu", "serde", "serde_core", + "zmij", ] [[package]] name = "serde_spanned" -version = "1.0.3" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e24345aa0fe688594e73770a5f6d1b216508b4f93484c0026d521acd30134392" +checksum = "f8bbf91e5a4d6315eee45e704372590b30e260ee83af6639d64557f51b067776" dependencies = [ "serde_core", ] @@ -1544,9 +1538,9 @@ dependencies = [ [[package]] name = "simd-adler32" -version = "0.3.7" +version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" +checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" [[package]] name = "simd_macros" @@ -1603,9 +1597,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.109" +version = "2.0.114" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f17c7e013e88258aa9543dcbe81aca68a667a9ac37cd69c9fbc07858bfe0e2f" +checksum = "d4d107df263a3013ef9b1879b0df87d706ff80f65a86ea879bd9c31f9b307c2a" dependencies = [ "proc-macro2", "quote", @@ -1637,15 +1631,15 @@ checksum = "591ef38edfb78ca4771ee32cf494cb8771944bee237a9b91fc9c1424ac4b777b" [[package]] name = "tempfile" -version = "3.23.0" +version = "3.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d31c77bdf42a745371d260a26ca7163f1e0924b64afa0b688e61b5a9fa02f16" +checksum = "655da9c7eb6305c55742045d5a8d2037996d61d8de95806335c7c86ce0f82e9c" dependencies = [ "fastrand", "getrandom", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys", ] [[package]] @@ -1690,9 +1684,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.9.8" +version = "0.9.11+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0dc8b1fb61449e27716ec0e1bdf0f6b8f3e8f6b05391e8497b8b6d7804ea6d8" +checksum = "f3afc9a848309fe1aaffaed6e1546a7a14de1f935dc9d89d32afd9a44bab7c46" dependencies = [ "indexmap", "serde_core", @@ -1705,27 +1699,27 @@ dependencies = [ [[package]] name = "toml_datetime" -version = "0.7.3" +version = "0.7.5+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2cdb639ebbc97961c51720f858597f7f24c4fc295327923af55b74c3c724533" +checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" dependencies = [ "serde_core", ] [[package]] name = "toml_parser" -version = "1.0.4" +version = "1.0.6+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0cbe268d35bdb4bb5a56a2de88d0ad0eb70af5384a99d648cd4b3d04039800e" +checksum = "a3198b4b0a8e11f09dd03e133c0280504d0801269e9afa46362ffde1cbeebf44" dependencies = [ "winnow", ] [[package]] name = "toml_writer" -version = "1.0.4" +version = "1.0.6+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df8b2b54733674ad286d16267dcfc7a71ed5c776e4ac7aa3c3e2561f7c637bf2" +checksum = "ab16f14aed21ee8bfd8ec22513f7287cd4a91aa92e44edfe2c17ddd004e92607" [[package]] name = "twox-hash" @@ -1759,9 +1753,9 @@ checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" [[package]] name = "url" -version = "2.5.7" +version = "2.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08bc136a29a3d1758e07a9cca267be308aeebf5cfd5a10f3f67ab2097683ef5b" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" dependencies = [ "form_urlencoded", "idna", @@ -1783,9 +1777,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.18.1" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f87b8aa10b915a06587d0dec516c282ff295b475d94abf425d62b57710070a2" +checksum = "e2e054861b4bd027cd373e18e8d8d8e6548085000e41290d95ce0c373a654b4a" dependencies = [ "getrandom", "js-sys", @@ -1927,9 +1921,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.105" +version = "0.2.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da95793dfc411fbbd93f5be7715b0578ec61fe87cb1a42b12eb625caa5c5ea60" +checksum = "0d759f433fa64a2d763d1340820e46e111a7a5ab75f993d1852d70b03dbb80fd" dependencies = [ "cfg-if", "once_cell", @@ -1940,9 +1934,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.105" +version = "0.2.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04264334509e04a7bf8690f2384ef5265f05143a4bff3889ab7a3269adab59c2" +checksum = "48cb0d2638f8baedbc542ed444afc0644a29166f1595371af4fecf8ce1e7eeb3" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1950,9 +1944,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.105" +version = "0.2.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "420bc339d9f322e562942d52e115d57e950d12d88983a14c79b86859ee6c7ebc" +checksum = "cefb59d5cd5f92d9dcf80e4683949f15ca4b511f4ac0a6e14d4e1ac60c6ecd40" dependencies = [ "bumpalo", "proc-macro2", @@ -1963,9 +1957,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.105" +version = "0.2.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76f218a38c84bcb33c25ec7059b07847d465ce0e0a76b995e134a45adcb6af76" +checksum = "cbc538057e648b67f72a982e708d485b2efa771e1ac05fec311f9f63e5800db4" dependencies = [ "unicode-ident", ] @@ -1981,9 +1975,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.82" +version = "0.3.83" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a1f95c0d03a47f4ae1f7a64643a6bb97465d9b740f0fa8f90ea33915c99a9a1" +checksum = "9b32828d774c412041098d182a8b38b16ea816958e07cf40eec2bc080ae137ac" dependencies = [ "js-sys", "wasm-bindgen", @@ -2022,7 +2016,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys", ] [[package]] @@ -2037,15 +2031,6 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets", -] - [[package]] name = "windows-sys" version = "0.61.2" @@ -2055,76 +2040,11 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-targets" -version = "0.53.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" -dependencies = [ - "windows-link", - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_gnullvm", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", -] - -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" - -[[package]] -name = "windows_aarch64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" - -[[package]] -name = "windows_i686_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" - -[[package]] -name = "windows_i686_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" - -[[package]] -name = "windows_i686_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" - -[[package]] -name = "windows_x86_64_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" - -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" - -[[package]] -name = "windows_x86_64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" - [[package]] name = "winnow" -version = "0.7.13" +version = "0.7.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21a0236b59786fed61e2a80582dd500fe61f18b5dca67a4a067d0bc9039339cf" +checksum = "5a5364e9d77fcdeeaa6062ced926ee3381faa2ee02d3eb83a5c27a8825540829" [[package]] name = "winsafe" @@ -2200,18 +2120,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.27" +version = "0.8.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0894878a5fa3edfd6da3f88c4805f4c8558e2b996227a3d864f47fe11e38282c" +checksum = "668f5168d10b9ee831de31933dc111a459c97ec93225beb307aed970d1372dfd" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.27" +version = "0.8.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88d2b8d9c68ad2b9e4340d7832716a4d21a22a1154777ad56ea55c51a9cf3831" +checksum = "2c7962b26b0a8685668b671ee4b54d007a67d4eaf05fda79ac0ecf41e32270f1" dependencies = [ "proc-macro2", "quote", @@ -2271,3 +2191,9 @@ dependencies = [ "quote", "syn", ] + +[[package]] +name = "zmij" +version = "1.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fc5a66a20078bf1251bde995aa2fdcc4b800c70b5d92dd2c62abc5c60f679f8" diff --git a/Cargo.toml b/Cargo.toml index 26b539bc..d6d7f0cf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,7 +46,7 @@ rayon.workspace = true rusqlite = { version = "0.38.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true -toml = "0.9.8" +toml = "0.9.11+spec-1.1.0" validator.workspace = true wyhash.workspace = true zerocopy.workspace = true @@ -66,7 +66,7 @@ version = "0.0.0" edition = "2024" [workspace.dependencies] -bumpalo = "3.19.0" +bumpalo = "3.19.1" dary_heap = "0.3.8" paste = "1.0.15" rand = "0.9.2" @@ -76,7 +76,7 @@ seq-macro = "0.3.6" serde = { version = "1.0.228", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } wyhash = "0.6.0" -zerocopy = { version = "0.8.27", features = ["derive"] } +zerocopy = { version = "0.8.33", features = ["derive"] } [workspace.lints] # complexity diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 87357b92..8fae83aa 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -26,7 +26,7 @@ criterion = "0.8.1" rand.workspace = true [build-dependencies] -cc = "1.2.44" +cc = "1.2.52" which = "8.0.0" [lints] diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index 0dd5fa75..6b922f9c 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -5,13 +5,13 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.51", features = ["derive", "env"] } +clap = { version = "4.5.54", features = ["derive", "env"] } object = { version = "0.38.1", features = ["read", "wasm"] } serde.workspace = true -serde_json = "1.0.145" +serde_json = "1.0.149" shlex = "1.3.0" target-triple = "1.0.0" -tempfile = "3.23.0" +tempfile = "3.24.0" [lints] workspace = true From a50209014d8914cfd39081a7c7c96f6be507d70d Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 23 Jan 2026 15:12:15 +0800 Subject: [PATCH 266/324] chore: remove `vec_deque_pop_front_if` (#412) `vec_deque_pop_if` is stablized since Rust 1.93. Signed-off-by: usamoi --- Cargo.lock | 89 +++++++++++++++------------ Cargo.toml | 2 +- crates/index/src/prefetcher.rs | 17 +---- crates/simd/Cargo.toml | 2 +- src/index/storage.rs | 17 +---- src/index/vchordrq/scanners/maxsim.rs | 4 +- 6 files changed, 55 insertions(+), 76 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index facac879..618138d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -175,9 +175,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.2.52" +version = "1.2.53" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd4932aefd12402b36c60956a4fe0035421f544799057659ff86f923657aada3" +checksum = "755d2fce177175ffca841e9a06afdb2c4ab0f593d53b4dee48147dfaade85932" dependencies = [ "find-msvc-tools", "shlex", @@ -282,9 +282,9 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.7.6" +version = "0.7.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1d728cc89cf3aee9ff92b05e62b19ee65a02b5702cff7d5a377e32c6ae29d8d" +checksum = "c3e64b0cc0439b12df2fa678eae89a1c56a529fd067a9115f7827f1fffd22b32" [[package]] name = "codepage" @@ -548,9 +548,9 @@ dependencies = [ [[package]] name = "find-msvc-tools" -version = "0.1.7" +version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f449e6c6c08c865631d4890cfacf252b3d396c9bcc83adb6623cdb02a8336c41" +checksum = "8591b0bcc8a98a64310a2fae1bb3e9b8564dd10e381e6e28010fde8e8e8568db" [[package]] name = "fixedbitset" @@ -560,9 +560,9 @@ checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" [[package]] name = "flate2" -version = "1.1.5" +version = "1.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfe33edd8e85a12a67454e37f8c75e730830d83e313556ab9ebf9ee7fbeb3bfb" +checksum = "b375d6465b98090a5f25b1c7703f3859783755aa9a80433b36e0379a3ec2f369" dependencies = [ "crc32fast", "miniz_oxide", @@ -846,9 +846,9 @@ checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" [[package]] name = "js-sys" -version = "0.3.83" +version = "0.3.85" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "464a3709c7f55f1f721e5389aa6ea4e3bc6aba669353300af094b29ffbdde1d8" +checksum = "8c942ebf8e95485ca0d52d97da7c5a2c387d0e7f0ba4c35e93bfcaee045955b3" dependencies = [ "once_cell", "wasm-bindgen", @@ -1268,18 +1268,18 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.105" +version = "1.0.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "535d180e0ecab6268a3e718bb9fd44db66bbbc256257165fc699dadf70d16fe7" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" dependencies = [ "unicode-ident", ] [[package]] name = "quote" -version = "1.0.43" +version = "1.0.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc74d9a594b72ae6656596548f56f667211f8a97b3d4c3d467150794690dc40a" +checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4" dependencies = [ "proc-macro2", ] @@ -1328,9 +1328,9 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.9.3" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" dependencies = [ "getrandom", ] @@ -1384,6 +1384,16 @@ version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" +[[package]] +name = "rsqlite-vfs" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8a1f2315036ef6b1fbacd1972e8ee7688030b0a2121edfc2a6550febd41574d" +dependencies = [ + "hashbrown 0.16.1", + "thiserror", +] + [[package]] name = "rusqlite" version = "0.38.0" @@ -1563,14 +1573,13 @@ checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" [[package]] name = "sqlite-wasm-rs" -version = "0.5.1" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05e98301bf8b0540c7de45ecd760539b9c62f5772aed172f08efba597c11cd5d" +checksum = "2f4206ed3a67690b9c29b77d728f6acc3ce78f16bf846d83c94f76400320181b" dependencies = [ "cc", - "hashbrown 0.16.1", "js-sys", - "thiserror", + "rsqlite-vfs", "wasm-bindgen", ] @@ -1644,18 +1653,18 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.17" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f63587ca0f12b72a0600bcba1d40081f830876000bb46dd2337a3051618f4fc8" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.17" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ff15c8ecd7de3849db632e14d18d2571fa09dfc5ed93479bc4485c7a517c913" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", @@ -1912,18 +1921,18 @@ dependencies = [ [[package]] name = "wasip2" -version = "1.0.1+wasi-0.2.4" +version = "1.0.2+wasi-0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" +checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5" dependencies = [ "wit-bindgen", ] [[package]] name = "wasm-bindgen" -version = "0.2.106" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d759f433fa64a2d763d1340820e46e111a7a5ab75f993d1852d70b03dbb80fd" +checksum = "64024a30ec1e37399cf85a7ffefebdb72205ca1c972291c51512360d90bd8566" dependencies = [ "cfg-if", "once_cell", @@ -1934,9 +1943,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.106" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48cb0d2638f8baedbc542ed444afc0644a29166f1595371af4fecf8ce1e7eeb3" +checksum = "008b239d9c740232e71bd39e8ef6429d27097518b6b30bdf9086833bd5b6d608" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1944,9 +1953,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.106" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cefb59d5cd5f92d9dcf80e4683949f15ca4b511f4ac0a6e14d4e1ac60c6ecd40" +checksum = "5256bae2d58f54820e6490f9839c49780dff84c65aeab9e772f15d5f0e913a55" dependencies = [ "bumpalo", "proc-macro2", @@ -1957,9 +1966,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.106" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbc538057e648b67f72a982e708d485b2efa771e1ac05fec311f9f63e5800db4" +checksum = "1f01b580c9ac74c8d8f0c0e4afb04eeef2acf145458e52c03845ee9cd23e3d12" dependencies = [ "unicode-ident", ] @@ -1975,9 +1984,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.83" +version = "0.3.85" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b32828d774c412041098d182a8b38b16ea816958e07cf40eec2bc080ae137ac" +checksum = "312e32e551d92129218ea9a2452120f4aabc03529ef03e4d0d82fb2780608598" dependencies = [ "js-sys", "wasm-bindgen", @@ -2054,9 +2063,9 @@ checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" [[package]] name = "wit-bindgen" -version = "0.46.0" +version = "0.51.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" [[package]] name = "writeable" @@ -2194,6 +2203,6 @@ dependencies = [ [[package]] name = "zmij" -version = "1.0.12" +version = "1.0.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2fc5a66a20078bf1251bde995aa2fdcc4b800c70b5d92dd2c62abc5c60f679f8" +checksum = "dfcd145825aace48cff44a8844de64bf75feec3080e0aa5cdbde72961ae51a65" diff --git a/Cargo.toml b/Cargo.toml index d6d7f0cf..d1ee1f92 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,7 +46,7 @@ rayon.workspace = true rusqlite = { version = "0.38.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true -toml = "0.9.11+spec-1.1.0" +toml = "0.9.11" validator.workspace = true wyhash.workspace = true zerocopy.workspace = true diff --git a/crates/index/src/prefetcher.rs b/crates/index/src/prefetcher.rs index cd53a263..d0bd662a 100644 --- a/crates/index/src/prefetcher.rs +++ b/crates/index/src/prefetcher.rs @@ -291,7 +291,7 @@ where } self.window.push_back(e); } - let e = vec_deque_pop_front_if(&mut self.window, predicate)?; + let e = self.window.pop_front_if(move |x| predicate(x))?; let list = e.fetch(); Some(( e, @@ -417,18 +417,3 @@ pub trait PrefetcherHeapFamily<'b, R> { fn is_not_plain(&self) -> bool; } - -// Emulate unstable library feature `vec_deque_pop_if`. -// See https://github.com/rust-lang/rust/issues/135889. - -fn vec_deque_pop_front_if( - this: &mut VecDeque, - predicate: impl FnOnce(&T) -> bool, -) -> Option { - let first = this.front()?; - if predicate(first) { - this.pop_front() - } else { - None - } -} diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 8fae83aa..7b8a5c6c 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -26,7 +26,7 @@ criterion = "0.8.1" rand.workspace = true [build-dependencies] -cc = "1.2.52" +cc = "1.2.53" which = "8.0.0" [lints] diff --git a/src/index/storage.rs b/src/index/storage.rs index e91039b6..bad6b812 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -477,7 +477,7 @@ where } self.window.push_back(e); } - vec_deque_pop_front_if(&mut self.window, predicate) + self.window.pop_front_if(move |x| predicate(x)) } #[allow(dead_code)] pub fn pop_item(&mut self) -> Option { @@ -708,21 +708,6 @@ fn lp_flags(x: pgrx::pg_sys::ItemIdData) -> u32 { (x >> 15) & 0b11 } -// Emulate unstable library feature `vec_deque_pop_if`. -// See https://github.com/rust-lang/rust/issues/135889. - -fn vec_deque_pop_front_if( - this: &mut VecDeque, - predicate: impl FnOnce(&T) -> bool, -) -> Option { - let first = this.front()?; - if predicate(first) { - this.pop_front() - } else { - None - } -} - // Emulate unstable library feature `box_vec_non_null`. // See https://github.com/rust-lang/rust/issues/130364. diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index b82e302b..460d137a 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -753,7 +753,7 @@ impl SearchBuilder for MaxsimBuilder { // Emulate unstable library feature `binary_heap_into_iter_sorted`. // See https://github.com/rust-lang/rust/issues/59278. -pub trait IntoIterSortedPolyfill { +trait IntoIterSortedPolyfill { fn into_iter_sorted_polyfill(self) -> IntoIterSorted; } @@ -764,7 +764,7 @@ impl IntoIterSortedPolyfill for BinaryHeap { } #[derive(Clone, Debug)] -pub struct IntoIterSorted { +struct IntoIterSorted { inner: BinaryHeap, } From bad7db05be38c8625346d3e59d859e445fa6ed9a Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 29 Jan 2026 16:49:48 +0800 Subject: [PATCH 267/324] build(deps): bump actions/checkout from 6.0.1 to 6.0.2 (#413) --- .github/workflows/check.yml | 18 +++++++++--------- .github/workflows/release.yml | 4 ++-- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 9c3a6b4e..b2ba2467 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -29,7 +29,7 @@ jobs: run: rustup update - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Typos uses: crate-ci/typos@master @@ -130,7 +130,7 @@ jobs: echo CC=clang-18 >> $GITHUB_ENV - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Cargo Test (Miri) env: @@ -187,7 +187,7 @@ jobs: uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Cargo Test run: cargo test --locked --workspace --exclude vchord --no-fail-fast -- --no-capture @@ -287,7 +287,7 @@ jobs: sudo apt-get install -y --no-install-recommends postgresql-${{ matrix.version }}-pgvector - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Clippy run: make PROFILE=dev clippy @@ -383,7 +383,7 @@ jobs: sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' install - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Clippy run: make PROFILE=dev clippy @@ -503,7 +503,7 @@ jobs: run: Add-Content -Path $env:GITHUB_ENV -Value "CC=clang-cl.exe" - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Clippy run: make PROFILE=dev clippy @@ -623,7 +623,7 @@ jobs: make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' install - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Patch run: | @@ -829,7 +829,7 @@ jobs: sudo make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC='sccache clang' CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Patch run: | @@ -1026,7 +1026,7 @@ jobs: popd - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Cargo Test run: cargo test --locked --workspace --exclude vchord --no-fail-fast -- --no-capture diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index dc89b198..684988a7 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -74,7 +74,7 @@ jobs: pg_config - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Build run: | @@ -122,7 +122,7 @@ jobs: steps: - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Upload env: From cf07d5c5b1f14b434cec9f2e3605bb1dc851d782 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 3 Feb 2026 17:16:14 +0800 Subject: [PATCH 268/324] readme: sync with docs (#416) Signed-off-by: usamoi --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index fa7f0c54..bdacb991 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,7 @@ For more usage, please read: - [Prefilter](https://docs.vectorchord.ai/vectorchord/usage/prefilter.html) - [Prefetch](https://docs.vectorchord.ai/vectorchord/usage/prefetch.html) - [Rerank in Table](https://docs.vectorchord.ai/vectorchord/usage/rerank-in-table.html) +- [Partitioning Tuning](https://docs.vectorchord.ai/vectorchord/usage/partitioning-tuning.html) - [External Build](https://docs.vectorchord.ai/vectorchord/usage/external-index-precomputation.html) ## License From 2d70595486e7699414d32d251106864c1c3ec621 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 6 Feb 2026 17:07:44 +0800 Subject: [PATCH 269/324] feat: specify search parameters in index options (#418) ```sql CREATE INDEX i ON t USING vchordrq (embedding vector_l2_ops) with (options = $$ build.internal.lists = [1024] $$, probes = '32'); -- vchordrq.probes is not set; `vchordrq.probes = 32` according to the index option SELECT id FROM t ORDER BY t <-> '[1,1,1]'; SET vchordrq.probes TO '64'; -- vchordrq.probes is set; `vchordrq.probes = 64` according to the GUC SELECT id FROM t ORDER BY t <-> '[1,1,1]'; SET vchordrq.probes TO default; -- vchordrq.probes is not set again; `vchordrq.probes = 32` according to the index option SELECT id FROM t ORDER BY t <-> '[1,1,1]'; ALTER INDEX i SET (probes = '64'); -- vchordrq.probes is not set; `vchordrq.probes = 64` according to the index option SELECT id FROM t ORDER BY t <-> '[1,1,1]'; ``` #392 Signed-off-by: usamoi --- .github/workflows/check.yml | 3 + crates/simd/src/emulate.rs | 2 + crates/simd/src/fast_scan.rs | 120 +++++---- crates/simd/src/lib.rs | 18 +- src/index/gucs.rs | 339 ++++++++++++++++++++++--- src/index/vchordg/am/am_build.rs | 9 +- src/index/vchordg/am/mod.rs | 60 ++++- src/index/vchordrq/am/am_build.rs | 9 +- src/index/vchordrq/am/mod.rs | 147 +++++++++-- src/index/vchordrq/types.rs | 1 + tests/vchordg/index_vector_rabitq.slt | 21 +- tests/vchordrq/index_vector_rabitq.slt | 66 ++--- tests/vchordrq/options.slt | 46 ++++ 13 files changed, 656 insertions(+), 185 deletions(-) create mode 100644 tests/vchordrq/options.slt diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index b2ba2467..f8cc8edd 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -361,6 +361,8 @@ jobs: - name: Set up Sccache uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 + with: + version: "v0.12.0" - name: Install PostgreSQL & pgvector run: | @@ -559,6 +561,7 @@ jobs: volumes: - /opt:/opt:rw,rshared - /opt:/__e/node20:ro,rshared + - /opt:/__e/node24:ro,rshared env: SCCACHE_GHA_ENABLED: "true" diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index 5434088f..b84e5e8f 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -12,6 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. +#[allow(unused_macros)] macro_rules! partial_load { (@internal_0) => { ::core::mem::zeroed() @@ -37,6 +38,7 @@ macro_rules! partial_load { }; } +#[allow(unused_imports)] pub(crate) use partial_load; // VP2INTERSECT emulation. diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index deeccb68..22de0496 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -165,13 +165,15 @@ mod scan { return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); for n in 90..110 { - let code = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); - let lut = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); + let code = &code[..n]; + let lut = &lut[..n]; unsafe { assert_eq!(scan_v4(&code, &lut), fallback(&code, &lut)); } @@ -268,13 +270,15 @@ mod scan { return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); for n in 90..110 { - let code = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); - let lut = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); + let code = &code[..n]; + let lut = &lut[..n]; unsafe { assert_eq!(scan_v3(&code, &lut), fallback(&code, &lut)); } @@ -341,13 +345,15 @@ mod scan { return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); for n in 90..110 { - let code = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); - let lut = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); + let code = &code[..n]; + let lut = &lut[..n]; unsafe { assert_eq!(scan_v2(&code, &lut), fallback(&code, &lut)); } @@ -414,13 +420,15 @@ mod scan { return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); for n in 90..110 { - let code = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); - let lut = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); + let code = &code[..n]; + let lut = &lut[..n]; unsafe { assert_eq!(scan_a2(&code, &lut), fallback(&code, &lut)); } @@ -486,20 +494,22 @@ mod scan { #[cfg(all(target_arch = "s390x", test))] #[test] - #[cfg(miri, ignore)] + #[cfg_attr(miri, ignore)] fn scan_z13_test() { if !crate::is_cpu_detected!("z13") { println!("test {} ... skipped (z13)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); for n in 90..110 { - let code = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); - let lut = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); + let code = &code[..n]; + let lut = &lut[..n]; unsafe { assert_eq!(scan_z13(&code, &lut), fallback(&code, &lut)); } @@ -608,20 +618,22 @@ mod scan { #[cfg(all(target_arch = "powerpc64", test))] #[test] - #[cfg(miri, ignore)] + #[cfg_attr(miri, ignore)] fn scan_p9_test() { if !crate::is_cpu_detected!("p9") { println!("test {} ... skipped (p9)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); for n in 90..110 { - let code = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); - let lut = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); + let code = &code[..n]; + let lut = &lut[..n]; unsafe { assert_eq!(scan_p9(&code, &lut), fallback(&code, &lut)); } @@ -634,20 +646,22 @@ mod scan { #[cfg(all(target_arch = "powerpc64", test))] #[test] - #[cfg(miri, ignore)] + #[cfg_attr(miri, ignore)] fn scan_p8_test() { if !crate::is_cpu_detected!("p8") { println!("test {} ... skipped (p8)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); for n in 90..110 { - let code = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); - let lut = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); + let code = &code[..n]; + let lut = &lut[..n]; unsafe { assert_eq!(scan_p8(&code, &lut), fallback(&code, &lut)); } @@ -660,20 +674,22 @@ mod scan { #[cfg(all(target_arch = "powerpc64", test))] #[test] - #[cfg(miri, ignore)] + #[cfg_attr(miri, ignore)] fn scan_p7_test() { if !crate::is_cpu_detected!("p7") { println!("test {} ... skipped (p7)", module_path!()); return; } for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); for n in 90..110 { - let code = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); - let lut = (0..n) - .map(|_| std::array::from_fn(|_| rand::random())) - .collect::>(); + let code = &code[..n]; + let lut = &lut[..n]; unsafe { assert_eq!(scan_p7(&code, &lut), fallback(&code, &lut)); } diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index c9ed2a75..4ca9340e 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -65,12 +65,26 @@ impl F16 for f16 { #[inline(always)] fn _from_f32(x: f32) -> Self { - f16::from_f32(x) + #[cfg(not(miri))] + { + f16::from_f32(x) + } + #[cfg(miri)] + { + f16::from_f32_const(x) + } } #[inline(always)] fn _to_f32(self) -> f32 { - self.to_f32() + #[cfg(not(miri))] + { + self.to_f32() + } + #[cfg(miri)] + { + self.to_f32_const() + } } } diff --git a/src/index/gucs.rs b/src/index/gucs.rs index f29ede82..ee3cc726 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -14,7 +14,7 @@ use crate::index::scanners::Io; use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting, PostgresGucEnum}; -use std::ffi::CString; +use std::ffi::{CStr, CString}; #[derive(Debug, Clone, Copy, PostgresGucEnum)] pub enum PostgresIo { @@ -37,6 +37,8 @@ static VCHORDG_ENABLE_SCAN: GucSetting = GucSetting::::new(true); static VCHORDG_EF_SEARCH: GucSetting = GucSetting::::new(64); +static mut VCHORDG_EF_SEARCH_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); + static VCHORDG_BEAM_SEARCH: GucSetting = GucSetting::::new(1); static VCHORDG_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); @@ -59,14 +61,22 @@ static VCHORDRQ_ENABLE_SCAN: GucSetting = GucSetting::::new(true); static VCHORDRQ_PROBES: GucSetting> = GucSetting::>::new(Some(c"")); +static mut VCHORDRQ_PROBES_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); + static VCHORDRQ_EPSILON: GucSetting = GucSetting::::new(1.9); +static mut VCHORDRQ_EPSILON_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); + static VCHORDRQ_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); static VCHORDRQ_MAXSIM_REFINE: GucSetting = GucSetting::::new(0); +static mut VCHORDRQ_MAXSIM_REFINE_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); + static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); +static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); + static VCHORDRQ_PREFILTER: GucSetting = GucSetting::::new(false); static VCHORDRQ_IO_SEARCH: GucSetting = GucSetting::::new( @@ -258,14 +268,85 @@ pub fn init() { #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordg".as_ptr()); } + assert!(crate::is_main()); + let targets = vec![ + (c"vchordg.ef_search", &raw mut VCHORDG_EF_SEARCH_CONFIG), + (c"vchordrq.epsilon", &raw mut VCHORDRQ_EPSILON_CONFIG), + ( + c"vchordrq.maxsim_refine", + &raw mut VCHORDRQ_MAXSIM_REFINE_CONFIG, + ), + ( + c"vchordrq.maxsim_threshold", + &raw mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG, + ), + (c"vchordrq.probes", &raw mut VCHORDRQ_PROBES_CONFIG), + ]; + #[cfg(any(feature = "pg14", feature = "pg15"))] + unsafe { + let len = pgrx::pg_sys::GetNumConfigOptions() as usize; + let arr = sys::get_guc_variables(); + let mut sources = (0..len).map(|i| arr.add(i).read()); + debug_assert!(targets.is_sorted_by(|(a, _), (b, _)| guc_name_compare(a, b).is_le())); + for (name, ptr) in targets { + *ptr = loop { + if let Some(source) = sources.next() { + if !(*source).name.is_null() && CStr::from_ptr((*source).name) == name { + break source; + } else { + continue; + } + } else { + pgrx::error!("failed to find GUC {name:?}"); + } + }; + assert!(check(*ptr, name), "failed to find GUC {name:?}"); + } + } + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + unsafe { + use pgrx::pg_sys::PGERROR; + for (name, ptr) in targets { + *ptr = sys::find_option(name.as_ptr(), false, false, PGERROR as _); + assert!(check(*ptr, name), "failed to find GUC {name:?}"); + } + } +} + +unsafe fn check(p: *mut sys::config_generic, name: &CStr) -> bool { + if p.is_null() { + return false; + } + if unsafe { (*p).flags } & pgrx::pg_sys::GUC_CUSTOM_PLACEHOLDER as core::ffi::c_int != 0 { + return false; + } + if unsafe { (*p).name }.is_null() { + return false; + } + if unsafe { CStr::from_ptr((*p).name) != name } { + return false; + } + true } pub fn vchordg_enable_scan() -> bool { VCHORDG_ENABLE_SCAN.get() } -pub fn vchordg_ef_search() -> u32 { - VCHORDG_EF_SEARCH.get() as u32 +pub fn vchordg_ef_search(index: pgrx::pg_sys::Relation) -> u32 { + fn parse(x: i32) -> u32 { + x as u32 + } + assert!(crate::is_main()); + const DEFAULT: i32 = 64; + if unsafe { (*VCHORDG_EF_SEARCH_CONFIG).source } != pgrx::pg_sys::GucSource::PGC_S_DEFAULT { + let value = VCHORDG_EF_SEARCH.get(); + parse(value) + } else { + use crate::index::vchordg::am::Reloption; + let value = unsafe { Reloption::ef_search((*index).rd_options as _, DEFAULT) }; + parse(value) + } } pub fn vchordg_beam_search() -> u32 { @@ -299,36 +380,55 @@ pub fn vchordrq_enable_scan() -> bool { VCHORDRQ_ENABLE_SCAN.get() } -pub fn vchordrq_probes() -> Vec { - match VCHORDRQ_PROBES.get() { - None => Vec::new(), - Some(probes) => { - let mut result = Vec::new(); - let mut current = None; - for &c in probes.to_bytes() { - match c { - b' ' => continue, - b',' => result.push(current.take().expect("empty probes")), - b'0'..=b'9' => { - if let Some(x) = current.as_mut() { - *x = *x * 10 + (c - b'0') as u32; - } else { - current = Some((c - b'0') as u32); - } +pub unsafe fn vchordrq_probes(index: pgrx::pg_sys::Relation) -> Vec { + fn parse(value: &CStr) -> Vec { + let mut result = Vec::new(); + let mut current = None; + for &c in value.to_bytes() { + match c { + b' ' => continue, + b',' => result.push(current.take().expect("empty probes")), + b'0'..=b'9' => { + if let Some(x) = current.as_mut() { + *x = *x * 10 + (c - b'0') as u32; + } else { + current = Some((c - b'0') as u32); } - c => pgrx::error!("unknown character in probes: ASCII = {c}"), } + c => pgrx::error!("unknown character in probes: ASCII = {c}"), } - if let Some(current) = current { - result.push(current); - } - result } + if let Some(current) = current { + result.push(current); + } + result + } + assert!(crate::is_main()); + const DEFAULT: &CStr = c""; + if unsafe { (*VCHORDRQ_PROBES_CONFIG).source } != pgrx::pg_sys::GucSource::PGC_S_DEFAULT { + let value = VCHORDRQ_PROBES.get(); + parse(value.as_deref().unwrap_or(DEFAULT)) + } else { + use crate::index::vchordrq::am::Reloption; + let value = unsafe { Reloption::probes((*index).rd_options as _, DEFAULT) }; + parse(value) } } -pub fn vchordrq_epsilon() -> f32 { - VCHORDRQ_EPSILON.get() as f32 +pub unsafe fn vchordrq_epsilon(index: pgrx::pg_sys::Relation) -> f32 { + fn parse(x: f64) -> f32 { + x as f32 + } + assert!(crate::is_main()); + const DEFAULT: f64 = 1.9; + if unsafe { (*VCHORDRQ_EPSILON_CONFIG).source } != pgrx::pg_sys::GucSource::PGC_S_DEFAULT { + let value = VCHORDRQ_EPSILON.get(); + parse(value) + } else { + use crate::index::vchordrq::am::Reloption; + let value = unsafe { Reloption::epsilon((*index).rd_options as _, DEFAULT) }; + parse(value) + } } pub fn vchordrq_max_scan_tuples() -> Option { @@ -336,12 +436,39 @@ pub fn vchordrq_max_scan_tuples() -> Option { if x < 0 { None } else { Some(x as u32) } } -pub fn vchordrq_maxsim_refine() -> u32 { - VCHORDRQ_MAXSIM_REFINE.get() as u32 +pub fn vchordrq_maxsim_refine(index: pgrx::pg_sys::Relation) -> u32 { + fn parse(x: i32) -> u32 { + x as u32 + } + assert!(crate::is_main()); + const DEFAULT: i32 = 0; + if unsafe { (*VCHORDRQ_MAXSIM_REFINE_CONFIG).source } != pgrx::pg_sys::GucSource::PGC_S_DEFAULT + { + let value = VCHORDRQ_MAXSIM_REFINE.get(); + parse(value) + } else { + use crate::index::vchordrq::am::Reloption; + let value = unsafe { Reloption::maxsim_refine((*index).rd_options as _, DEFAULT) }; + parse(value) + } } -pub fn vchordrq_maxsim_threshold() -> u32 { - VCHORDRQ_MAXSIM_THRESHOLD.get() as u32 +pub fn vchordrq_maxsim_threshold(index: pgrx::pg_sys::Relation) -> u32 { + fn parse(x: i32) -> u32 { + x as u32 + } + assert!(crate::is_main()); + const DEFAULT: i32 = 0; + if unsafe { (*VCHORDRQ_MAXSIM_THRESHOLD_CONFIG).source } + != pgrx::pg_sys::GucSource::PGC_S_DEFAULT + { + let value = VCHORDRQ_MAXSIM_THRESHOLD.get(); + parse(value) + } else { + use crate::index::vchordrq::am::Reloption; + let value = unsafe { Reloption::maxsim_threshold((*index).rd_options as _, DEFAULT) }; + parse(value) + } } pub fn vchordrq_prefilter() -> bool { @@ -377,3 +504,155 @@ pub fn vchordrq_query_sampling_max_records() -> u32 { pub fn vchordrq_query_sampling_rate() -> f64 { VCHORDRQ_QUERY_SAMPLING_RATE.get() } + +#[allow(non_camel_case_types)] +#[allow(non_snake_case)] +mod sys { + #[cfg(not(feature = "pg14"))] + #[cfg(not(feature = "pg15"))] + #[cfg(not(feature = "pg16"))] + #[cfg(not(feature = "pg17"))] + #[cfg(not(feature = "pg18"))] + compile_error!("bindings are not checked"); + + pub mod config_type { + pub type Type = ::core::ffi::c_uint; + } + + #[repr(C)] + #[derive(Copy, Clone)] + pub union config_var_val { + pub boolval: bool, + pub intval: ::core::ffi::c_int, + pub realval: f64, + pub stringval: *mut ::core::ffi::c_char, + pub enumval: ::core::ffi::c_int, + } + + #[repr(C)] + #[derive(Copy, Clone)] + pub struct config_var_value { + pub val: config_var_val, + pub extra: *mut ::core::ffi::c_void, + } + + pub mod config_group { + pub type Type = ::core::ffi::c_uint; + } + + pub mod GucStackState { + pub type Type = ::core::ffi::c_uint; + } + + #[repr(C)] + #[derive(Copy, Clone)] + pub struct guc_stack { + pub prev: *mut guc_stack, + pub nest_level: ::core::ffi::c_int, + pub state: GucStackState::Type, + pub source: pgrx::pg_sys::GucSource::Type, + pub scontext: pgrx::pg_sys::GucContext::Type, + pub masked_scontext: pgrx::pg_sys::GucContext::Type, + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] + pub srole: pgrx::pg_sys::Oid, + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] + pub masked_srole: pgrx::pg_sys::Oid, + pub prior: config_var_value, + pub masked: config_var_value, + } + + pub type GucStack = guc_stack; + + #[repr(C)] + #[derive(Debug, Copy, Clone)] + pub struct config_generic { + pub name: *const ::core::ffi::c_char, + pub context: pgrx::pg_sys::GucContext::Type, + pub group: config_group::Type, + pub short_desc: *const ::core::ffi::c_char, + pub long_desc: *const ::core::ffi::c_char, + pub flags: ::core::ffi::c_int, + pub vartype: config_type::Type, + pub status: ::core::ffi::c_int, + pub source: pgrx::pg_sys::GucSource::Type, + pub reset_source: pgrx::pg_sys::GucSource::Type, + pub scontext: pgrx::pg_sys::GucContext::Type, + pub reset_scontext: pgrx::pg_sys::GucContext::Type, + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] + pub srole: pgrx::pg_sys::Oid, + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] + pub reset_srole: pgrx::pg_sys::Oid, + pub stack: *mut GucStack, + pub extra: *mut ::core::ffi::c_void, + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + pub nondef_link: pgrx::pg_sys::dlist_node, + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + pub stack_link: pgrx::pg_sys::slist_node, + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + pub report_link: pgrx::pg_sys::slist_node, + pub last_reported: *mut ::core::ffi::c_char, + pub sourcefile: *mut ::core::ffi::c_char, + pub sourceline: ::core::ffi::c_int, + } + + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + #[inline] + pub unsafe fn find_option( + name: *const ::core::ffi::c_char, + create_placeholders: bool, + skip_errors: bool, + elevel: ::core::ffi::c_int, + ) -> *mut config_generic { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + #[cfg_attr(target_os = "windows", link(name = "postgres"))] + unsafe extern "C-unwind" { + #[link_name = "find_option"] + unsafe fn f( + name: *const ::core::ffi::c_char, + create_placeholders: bool, + skip_errors: bool, + elevel: ::core::ffi::c_int, + ) -> *mut config_generic; + } + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + unsafe { + pg_guard_ffi_boundary(move || f(name, create_placeholders, skip_errors, elevel)) + } + } + + #[cfg(any(feature = "pg14", feature = "pg15"))] + #[inline] + pub unsafe fn get_guc_variables() -> *mut *mut config_generic { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + #[cfg_attr(target_os = "windows", link(name = "postgres"))] + unsafe extern "C-unwind" { + #[link_name = "get_guc_variables"] + unsafe fn f() -> *mut *mut config_generic; + } + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + unsafe { + pg_guard_ffi_boundary(move || f()) + } + } +} + +#[allow(dead_code)] +fn guc_name_compare(a: &CStr, b: &CStr) -> std::cmp::Ordering { + let (a, b) = (a.to_bytes_with_nul(), b.to_bytes_with_nul()); + let mut i = 0; + while a[i] != 0 && b[i] != 0 { + let a = a[i].to_ascii_lowercase(); + let b = b[i].to_ascii_lowercase(); + if a != b { + return Ord::cmp(&a, &b); + } + i += 1; + } + if b[i] != 0 { + std::cmp::Ordering::Less + } else if a[i] != 0 { + std::cmp::Ordering::Greater + } else { + std::cmp::Ordering::Equal + } +} diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index edcef0ba..8fc2d549 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -594,18 +594,15 @@ unsafe fn options( d: opfamily.distance_kind(), }; // get indexing, segment, optimizing - let rabitq = 'rabitq: { + let indexing_options = { let reloption = unsafe { (*index_relation).rd_options as *const Reloption }; - if reloption.is_null() || unsafe { (*reloption).options == 0 } { - break 'rabitq Default::default(); - } - let s = unsafe { Reloption::options(reloption) }.to_string_lossy(); + let s = unsafe { Reloption::options(reloption, c"") }.to_string_lossy(); match toml::from_str::(&s) { Ok(p) => p, Err(e) => pgrx::error!("failed to parse options: {}", e), } }; - (vector, rabitq) + (vector, indexing_options) } mod vchordg_cached { diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index 06c47e9c..68fa97bf 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -31,29 +31,52 @@ use std::ptr::NonNull; use std::sync::OnceLock; #[repr(C)] -struct Reloption { +pub struct Reloption { vl_len_: i32, - pub options: i32, + options: i32, + ef_search: i32, } impl Reloption { - unsafe fn options<'a>(this: *const Self) -> &'a CStr { + unsafe fn options<'a>(this: *const Self, default: &'static CStr) -> &'a CStr { unsafe { - let ptr = this - .cast::() - .add((&raw const (*this).options).read() as _); + if this.is_null() { + return default; + } + let count = (&raw const (*this).options).read(); + if count == 0 { + return default; + } + let ptr = this.cast::().add(count as _); CStr::from_ptr(ptr.cast()) } } + pub unsafe fn ef_search(this: *const Self, default: i32) -> i32 { + unsafe { + if this.is_null() { + return default; + } + (*this).ef_search + } + } } -const TABLE: &[pgrx::pg_sys::relopt_parse_elt] = &[pgrx::pg_sys::relopt_parse_elt { - optname: c"options".as_ptr(), - opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, - offset: std::mem::offset_of!(Reloption, options) as i32, - #[cfg(feature = "pg18")] - isset_offset: 0, -}]; +const TABLE: &[pgrx::pg_sys::relopt_parse_elt] = &[ + pgrx::pg_sys::relopt_parse_elt { + optname: c"options".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, + offset: std::mem::offset_of!(Reloption, options) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"ef_search".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_INT, + offset: std::mem::offset_of!(Reloption, ef_search) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, +]; static RELOPT_KIND: OnceLock = OnceLock::new(); @@ -70,6 +93,15 @@ pub fn init() { None, pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, ); + pgrx::pg_sys::add_int_reloption( + kind as _, + c"ef_search".as_ptr(), + c"Search parameter `vchordg.ef_search`".as_ptr(), + 64, + 1, + 65535, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); } kind }); @@ -329,7 +361,7 @@ pub unsafe extern "C-unwind" fn amrescan( let opfamily = opfamily((*scan).indexRelation); let index = PostgresRelation::new((*scan).indexRelation); let options = SearchOptions { - ef_search: gucs::vchordg_ef_search(), + ef_search: gucs::vchordg_ef_search((*scan).indexRelation), beam_search: gucs::vchordg_beam_search(), max_scan_tuples: gucs::vchordg_max_scan_tuples(), io_search: gucs::vchordg_io_search(), diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index b3ae1eb6..4674895d 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -1246,18 +1246,15 @@ unsafe fn options( d: opfamily.distance_kind(), }; // get indexing, segment, optimizing - let rabitq = 'rabitq: { + let indexing_options = { let reloption = unsafe { (*index_relation).rd_options as *const Reloption }; - if reloption.is_null() || unsafe { (*reloption).options == 0 } { - break 'rabitq Default::default(); - } - let s = unsafe { Reloption::options(reloption) }.to_string_lossy(); + let s = unsafe { Reloption::options(reloption, c"") }.to_string_lossy(); match toml::from_str::(&s) { Ok(p) => p, Err(e) => pgrx::error!("failed to parse options: {}", e), } }; - (vector, rabitq) + (vector, indexing_options) } fn make_default_build( diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 125d1874..32ed017a 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -34,29 +34,105 @@ use std::sync::OnceLock; use vchordrq::InsertChooser; #[repr(C)] -struct Reloption { +pub struct Reloption { vl_len_: i32, - pub options: i32, + options: i32, + probes: i32, + epsilon: f64, + maxsim_refine: i32, + maxsim_threshold: i32, } impl Reloption { - unsafe fn options<'a>(this: *const Self) -> &'a CStr { + unsafe fn options<'a>(this: *const Self, default: &'static CStr) -> &'a CStr { unsafe { - let ptr = this - .cast::() - .add((&raw const (*this).options).read() as _); + if this.is_null() { + return default; + } + let count = (&raw const (*this).options).read(); + if count == 0 { + return default; + } + let ptr = this.cast::().add(count as _); + CStr::from_ptr(ptr.cast()) + } + } + pub unsafe fn probes<'a>(this: *const Self, default: &'static CStr) -> &'a CStr { + unsafe { + if this.is_null() { + return default; + } + let count = (&raw const (*this).probes).read(); + if count == 0 { + return default; + } + let ptr = this.cast::().add(count as _); CStr::from_ptr(ptr.cast()) } } + pub unsafe fn epsilon(this: *const Self, default: f64) -> f64 { + unsafe { + if this.is_null() { + return default; + } + (*this).epsilon + } + } + pub unsafe fn maxsim_refine(this: *const Self, default: i32) -> i32 { + unsafe { + if this.is_null() { + return default; + } + (*this).maxsim_refine + } + } + pub unsafe fn maxsim_threshold(this: *const Self, default: i32) -> i32 { + unsafe { + if this.is_null() { + return default; + } + (*this).maxsim_threshold + } + } } -const TABLE: &[pgrx::pg_sys::relopt_parse_elt] = &[pgrx::pg_sys::relopt_parse_elt { - optname: c"options".as_ptr(), - opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, - offset: std::mem::offset_of!(Reloption, options) as i32, - #[cfg(feature = "pg18")] - isset_offset: 0, -}]; +const TABLE: &[pgrx::pg_sys::relopt_parse_elt] = &[ + pgrx::pg_sys::relopt_parse_elt { + optname: c"options".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, + offset: std::mem::offset_of!(Reloption, options) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"probes".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, + offset: std::mem::offset_of!(Reloption, probes) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"epsilon".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_REAL, + offset: std::mem::offset_of!(Reloption, epsilon) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"maxsim_refine".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_INT, + offset: std::mem::offset_of!(Reloption, maxsim_refine) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"maxsim_threshold".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_INT, + offset: std::mem::offset_of!(Reloption, maxsim_threshold) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, +]; static RELOPT_KIND: OnceLock = OnceLock::new(); @@ -73,6 +149,41 @@ pub fn init() { None, pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, ); + pgrx::pg_sys::add_string_reloption( + kind as _, + c"probes".as_ptr(), + c"Search parameter `vchordrq.probes`".as_ptr(), + c"".as_ptr(), + None, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + pgrx::pg_sys::add_real_reloption( + kind as _, + c"epsilon".as_ptr(), + c"Search parameter `vchordrq.epsilon`".as_ptr(), + 1.9, + 0.0, + 4.0, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + pgrx::pg_sys::add_int_reloption( + kind as _, + c"maxsim_refine".as_ptr(), + c"Search parameter `vchordrq.maxsim_refine`".as_ptr(), + 0, + 0, + i32::MAX, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + pgrx::pg_sys::add_int_reloption( + kind as _, + c"maxsim_threshold".as_ptr(), + c"Search parameter `vchordrq.maxsim_threshold`".as_ptr(), + 0, + 0, + i32::MAX, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); } kind }); @@ -223,7 +334,7 @@ pub unsafe extern "C-unwind" fn amcostestimate( return; } let index = PostgresRelation::::new(relation.raw()); - let probes = gucs::vchordrq_probes(); + let probes = gucs::vchordrq_probes(relation.raw()); let cost = vchordrq::cost(&index); if cost.cells.len() != 1 + probes.len() { panic!( @@ -414,11 +525,11 @@ pub unsafe extern "C-unwind" fn amrescan( let opfamily = opfamily((*scan).indexRelation); let index = PostgresRelation::new((*scan).indexRelation); let options = SearchOptions { - epsilon: gucs::vchordrq_epsilon(), - probes: gucs::vchordrq_probes(), + epsilon: gucs::vchordrq_epsilon((*scan).indexRelation), + probes: gucs::vchordrq_probes((*scan).indexRelation), max_scan_tuples: gucs::vchordrq_max_scan_tuples(), - maxsim_refine: gucs::vchordrq_maxsim_refine(), - maxsim_threshold: gucs::vchordrq_maxsim_threshold(), + maxsim_refine: gucs::vchordrq_maxsim_refine((*scan).indexRelation), + maxsim_threshold: gucs::vchordrq_maxsim_threshold((*scan).indexRelation), io_search: gucs::vchordrq_io_search(), io_rerank: gucs::vchordrq_io_rerank(), prefilter: gucs::vchordrq_prefilter(), diff --git a/src/index/vchordrq/types.rs b/src/index/vchordrq/types.rs index f16f5bec..d9f773b2 100644 --- a/src/index/vchordrq/types.rs +++ b/src/index/vchordrq/types.rs @@ -177,6 +177,7 @@ pub struct VchordrqIndexingOptions { #[serde(flatten)] #[validate(nested)] pub index: VchordrqIndexOptions, + #[serde(default)] #[validate(nested)] pub build: VchordrqBuildOptions, } diff --git a/tests/vchordg/index_vector_rabitq.slt b/tests/vchordg/index_vector_rabitq.slt index f33dd8ea..50334194 100644 --- a/tests/vchordg/index_vector_rabitq.slt +++ b/tests/vchordg/index_vector_rabitq.slt @@ -78,7 +78,7 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; ---- 155 1608 @@ -89,7 +89,6 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_r 1629 60 218 -1639 statement ok DROP INDEX ti; @@ -98,7 +97,7 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; ---- 155 1608 @@ -109,7 +108,6 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_r 1629 60 218 -1639 statement ok DROP INDEX ti; @@ -118,7 +116,7 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; ---- 155 1608 @@ -129,7 +127,6 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_r 1629 60 218 -1639 statement ok DROP INDEX ti; @@ -138,11 +135,11 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_l2_ops); query I -SELECT SUM(index) FROM ( +SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 ) AS s(index); ---- -6162 +1 statement ok DROP INDEX ti; @@ -151,11 +148,11 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_ip_ops); query I -SELECT SUM(index) FROM ( +SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 ) AS s(index); ---- -6162 +1 statement ok DROP INDEX ti; @@ -164,11 +161,11 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_cosine_ops); query I -SELECT SUM(index) FROM ( +SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 ) AS s(index); ---- -6162 +1 statement ok DROP INDEX ti; diff --git a/tests/vchordrq/index_vector_rabitq.slt b/tests/vchordrq/index_vector_rabitq.slt index 6233d53e..29841b14 100644 --- a/tests/vchordrq/index_vector_rabitq.slt +++ b/tests/vchordrq/index_vector_rabitq.slt @@ -78,7 +78,7 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; ---- 155 1608 @@ -89,7 +89,6 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_r 1629 60 218 -1639 statement ok DROP INDEX ti; @@ -98,7 +97,7 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; ---- 155 1608 @@ -109,7 +108,6 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_r 1629 60 218 -1639 statement ok DROP INDEX ti; @@ -118,7 +116,7 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; ---- 155 1608 @@ -129,7 +127,6 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_r 1629 60 218 -1639 statement ok DROP INDEX ti; @@ -138,11 +135,11 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_l2_ops); query I -SELECT SUM(index) FROM ( +SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 ) AS s(index); ---- -6162 +1 statement ok DROP INDEX ti; @@ -151,11 +148,11 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_ip_ops); query I -SELECT SUM(index) FROM ( +SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 ) AS s(index); ---- -6162 +1 statement ok DROP INDEX ti; @@ -164,11 +161,11 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_cosine_ops); query I -SELECT SUM(index) FROM ( +SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 ) AS s(index); ---- -6162 +1 statement ok DROP INDEX ti; @@ -177,18 +174,11 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_l2_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 +) AS s(index); ---- -818 -1608 -155 -60 -174 -1639 -1080 -218 -1207 -331 +7290 statement ok DROP INDEX ti; @@ -197,18 +187,11 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_ip_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 +) AS s(index); ---- -818 -1608 -155 -60 -174 -1639 -1080 -218 -1207 -331 +7290 statement ok DROP INDEX ti; @@ -217,18 +200,11 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_cosine_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +SELECT SUM(index) FROM ( +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 +) AS s(index); ---- -818 -1608 -155 -60 -174 -1639 -1080 -218 -1207 -331 +7290 statement ok DROP INDEX ti; diff --git a/tests/vchordrq/options.slt b/tests/vchordrq/options.slt new file mode 100644 index 00000000..972917c1 --- /dev/null +++ b/tests/vchordrq/options.slt @@ -0,0 +1,46 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) VALUES ('[1,1,1]'), ('[2,2,2]'); + +statement ok +CREATE INDEX i ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +build.internal.lists = [2] +$$, probes = '2'); + +query T +SELECT val FROM t ORDER BY val <-> '[0,0,0]' LIMIT 10; +---- +[1,1,1] +[2,2,2] + +statement ok +SET vchordrq.probes TO '1'; + +query T +SELECT val FROM t ORDER BY val <-> '[0,0,0]' LIMIT 10; +---- +[1,1,1] + +statement ok +ALTER INDEX i SET (probes = '0'); + +query T +SELECT val FROM t ORDER BY val <-> '[0,0,0]' LIMIT 10; +---- +[1,1,1] + +statement ok +SET vchordrq.probes TO DEFAULT; + +query T +SELECT val FROM t ORDER BY val <-> '[0,0,0]' LIMIT 10; +---- + +statement ok +DROP TABLE t; From 74992ab17a4bc24b8be441d6367fea00b4d1e369 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 10 Feb 2026 16:09:05 +0800 Subject: [PATCH 270/324] fix: validate table name in external build (#423) Signed-off-by: usamoi --- src/index/vchordrq/types.rs | 34 ++++++++++++++++++++ tests/vchordrq/external_build_sql_inject.slt | 11 +++++++ 2 files changed, 45 insertions(+) create mode 100644 tests/vchordrq/external_build_sql_inject.slt diff --git a/src/index/vchordrq/types.rs b/src/index/vchordrq/types.rs index d9f773b2..f2645c07 100644 --- a/src/index/vchordrq/types.rs +++ b/src/index/vchordrq/types.rs @@ -109,9 +109,43 @@ impl Default for VchordrqInternalBuildOptions { #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] pub struct VchordrqExternalBuildOptions { + #[validate(custom(function = VchordrqExternalBuildOptions::validate_table))] pub table: String, } +impl VchordrqExternalBuildOptions { + fn validate_table(table: &str) -> Result<(), ValidationError> { + let (schema_name, table_name) = if let Some((left, right)) = table.split_once(".") { + (Some(left), right) + } else { + (None, table) + }; + fn check(s: &str) -> bool { + if s.is_empty() { + return false; + } + if !matches!(s.as_bytes()[0], b'A'..=b'Z' | b'a'..=b'z' | b'_') { + return false; + } + for c in s.as_bytes().iter().copied() { + if !matches!(c, b'0'..=b'9' | b'A'..=b'Z' | b'a'..=b'z' | b'_' | b'$') { + return false; + } + } + true + } + if let Some(schema_name) = schema_name { + if !check(schema_name) { + return Err(ValidationError::new("table name is not well-formed")); + } + } + if !check(table_name) { + return Err(ValidationError::new("table name is not well-formed")); + } + Ok(()) + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] #[serde(rename_all = "snake_case")] diff --git a/tests/vchordrq/external_build_sql_inject.slt b/tests/vchordrq/external_build_sql_inject.slt new file mode 100644 index 00000000..6129ede9 --- /dev/null +++ b/tests/vchordrq/external_build_sql_inject.slt @@ -0,0 +1,11 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement error table name is not well-formed +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +build.external.table = "s; SELECT s" +$$); + +statement ok +DROP TABLE t; From 597d7cfe30666b763d160852050f221044d69e5c Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 10 Feb 2026 16:48:51 +0800 Subject: [PATCH 271/324] chore: remove `./scripts` (#420) Alternative to https://github.com/tensorchord/VectorChord/pull/419 See also https://github.com/tensorchord/pgvecto.rs-docs/pull/164 Signed-off-by: usamoi --- scripts/README.md | 58 ---------- scripts/bench.py | 267 ---------------------------------------------- scripts/dump.py | 81 -------------- scripts/index.py | 247 ------------------------------------------ scripts/train.py | 218 ------------------------------------- 5 files changed, 871 deletions(-) delete mode 100644 scripts/README.md delete mode 100644 scripts/bench.py delete mode 100644 scripts/dump.py delete mode 100644 scripts/index.py delete mode 100644 scripts/train.py diff --git a/scripts/README.md b/scripts/README.md deleted file mode 100644 index 41115ab3..00000000 --- a/scripts/README.md +++ /dev/null @@ -1,58 +0,0 @@ -## Run External Index Precomputation Toolkit - -1. Install requirements - -```shell -# PYTHON = 3.11 -# When using CPU to train k-means clustering -conda install conda-forge::pgvector-python numpy pytorch::faiss-cpu conda-forge::psycopg h5py tqdm -# or -pip install pgvector numpy faiss-cpu psycopg h5py tqdm - -# When using GPU to train k-means clustering -conda install conda-forge::pgvector-python numpy pytorch::faiss-gpu conda-forge::psycopg h5py tqdm -``` - -1. Prepare dataset in `hdf5` format - - - If you already have your vectors stored in `PostgreSQL` using `pgvector`, you can export them to a local file by: - ```shell - python scripts/dump.py -n [table name] -c [column name] -d [dim] -o export.hdf5 --url postgresql://USERNAME:PASSWORD@localhost:5432/DBNAME - ``` - - - If you don't have any data, but would like to give it a try, you can choose one of these datasets: - - ```shell - wget http://ann-benchmarks.com/sift-128-euclidean.hdf5 # num=1M dim=128 metric=l2 - wget http://ann-benchmarks.com/gist-960-euclidean.hdf5 # num=1M dim=960 metric=l2 - wget https://myscale-datasets.s3.ap-southeast-1.amazonaws.com/laion-5m-test-ip.hdf5 # num=5M dim=768 metric=dot - wget https://myscale-datasets.s3.ap-southeast-1.amazonaws.com/laion-20m-test-ip.hdf5 # num=20M dim=768 metric=dot - wget https://myscale-datasets.s3.ap-southeast-1.amazonaws.com/laion-100m-test-ip.hdf5 # num=100M dim=768 metric=dot - ``` - -2. Preform clustering of centroids from vectors - - ```shell - # For small dataset size from 1M to 5M - python scripts/train.py -i [dataset file(export.hdf5)] -o [centroid filename(centroid.npy)] --lists [lists] -m [metric(l2/cos/dot)] - # For large datasets size, 5M to 100M in size, use GPU and mmap chunks - python scripts/train.py -i [dataset file(export.hdf5)] -o [centroid filename(centroid.npy)] --lists [lists] -m [metric(l2/cos/dot)] -g --mmap - ``` - - `lists` is the number of centroids for clustering, and a typical value for large datasets(>5M) could range from: - - $$ - 4*\sqrt{len(vectors)} \le lists \le 16*\sqrt{len(vectors)} - $$ - -3. To insert vectors and centroids into the database, and then create an index - - ```shell - python scripts/index.py -n [table name] -i [dataset file(export.hdf5)] -c [centroid filename(centroid.npy)] -m [metric(l2/cos/dot)] -d [dim] --url postgresql://postgres:123@localhost:5432/postgres - ``` - -4. Let's start our tour to check the benchmark result of VectorChord - - ```shell - python scripts/bench.py -n [table name] -i [dataset file(export.hdf5)] -m [metric(l2/cos/dot)] --nprob 100 --epsilon 1.0 --url postgresql://postgres:123@localhost:5432/postgres - ``` diff --git a/scripts/bench.py b/scripts/bench.py deleted file mode 100644 index 0517b608..00000000 --- a/scripts/bench.py +++ /dev/null @@ -1,267 +0,0 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025 TensorChord Inc. - -import time -import argparse -from pathlib import Path -from tqdm import tqdm -import multiprocessing as mp -import numpy as np - -import psycopg -import h5py -from pgvector.psycopg import register_vector - - -def build_arg_parse(): - parser = argparse.ArgumentParser() - parser.add_argument( - "-m", - "--metric", - help="Metric to pick, in l2 or cos", - choices=["l2", "cos", "dot"], - default="l2", - ) - parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) - parser.add_argument("-i", "--input", help="input filepath", required=True) - parser.add_argument( - "--url", - help="url, like `postgresql://postgres:123@localhost:5432/postgres`", - required=True, - ) - parser.add_argument( - "-t", "--top", help="Dimension", type=int, choices=[10, 100], default=10 - ) - parser.add_argument( - "--nprob", help="argument probes for query", default=100, type=int - ) - parser.add_argument( - "--epsilon", help="argument epsilon for query", type=float, default=1.0 - ) - parser.add_argument( - "--processes", help="Number of parallel processes to use", type=int, default=1 - ) - return parser - - -def create_connection(url, nprob, epsilon): - keepalive_kwargs = { - "keepalives": 1, - "keepalives_idle": 30, - "keepalives_interval": 5, - "keepalives_count": 5, - } - conn = psycopg.connect( - conninfo=url, - dbname="postgres", - autocommit=True, - **keepalive_kwargs, - ) - try: - conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") - except Exception: - pass - # Tuning - conn.execute("SET jit=false") - conn.execute("SET effective_io_concurrency=200") - - conn.execute(f"SET vchordrq.probes={nprob}") - conn.execute(f"SET vchordrq.epsilon={epsilon}") - try: - conn.execute(f"SELECT vchordrq_prewarm('{args.name}_embedding_idx'::regclass)") - except Exception: - pass - register_vector(conn) - return conn - - -def calculate_coverage(time_intervals): - if not time_intervals: - return 0 - sorted_intervals = sorted(time_intervals, key=lambda x: x[0]) - merged = [] - current_start, current_end = sorted_intervals[0] - for interval in sorted_intervals[1:]: - next_start, next_end = interval - if next_start <= current_end: - current_end = max(current_end, next_end) - else: - merged.append((current_start, current_end)) - current_start, current_end = next_start, next_end - merged.append((current_start, current_end)) - total_length = 0 - for start, end in merged: - total_length += end - start - return total_length - - -def process_batch(args): - """Process a batch of queries in a single process""" - batch_queries, batch_answers, k, metric_ops, url, name, nprob, epsilon = args - - # Create a new connection for this process - conn = create_connection(url, nprob, epsilon) - - hits = 0 - results = [] - - for query, ground_truth in zip(batch_queries, batch_answers): - start = time.perf_counter() - result = conn.execute( - f"SELECT id FROM {name} ORDER BY embedding {metric_ops} %s LIMIT {k}", - (query,), - ).fetchall() - end = time.perf_counter() - - result_ids = set([p[0] for p in result[:k]]) - ground_truth_ids = set(ground_truth[:k].tolist()) - hit = len(result_ids & ground_truth_ids) - hits += hit - - results.append((hit, (start, end))) - - conn.close() - return results - - -def calculate_metrics(all_results, k, m, num_processes=1): - """Calculate recall, QPS, and latency percentiles from results""" - hits, latencies = zip(*all_results) - - if isinstance(latencies[0], list | tuple): - # parallel_bench - total_time = calculate_coverage(latencies) - latencies = [(end - start) for start, end in latencies] - else: - # sequential_bench - total_time = sum(latencies) - - total_hits = sum(hits) - - recall = total_hits / (k * m * num_processes) - qps = (m * num_processes) / total_time - - # Calculate latency percentiles (in milliseconds) - latencies_ms = np.array(latencies) * 1000 - p50 = np.percentile(latencies_ms, 50) - p99 = np.percentile(latencies_ms, 99) - - return recall, qps, p50, p99 - - -def parallel_bench( - name, test, answer, metric_ops, num_processes, url, top, nprob, epsilon -): - """Run benchmark in parallel using multiple processes""" - m = test.shape[0] - batches = [] - - for _ in range(num_processes): - batch = ( - test, - answer, - top, - metric_ops, - url, - name, - nprob, - epsilon, - ) - batches.append(batch) - - # Create process pool and execute batches - with mp.Pool(processes=num_processes) as pool: - batch_results = list(pool.map(process_batch, batches)) - - # Flatten results from all batches - all_results = [result for batch in batch_results for result in batch] - - # Calculate metrics - recall, qps, p50, p99 = calculate_metrics(all_results, top, m, num_processes) - - print(f"Top: {top}") - print(f" Recall: {recall:.4f}") - print(f" QPS: {qps:.2f}") - print(f" P50 latency: {p50:.2f}ms") - print(f" P99 latency: {p99:.2f}ms") - - -def sequential_bench(name, test, answer, metric_ops, conn, top): - """Original sequential benchmark implementation with latency tracking""" - m = test.shape[0] - results = [] - pbar = tqdm(enumerate(test), total=m) - for i, query in pbar: - start = time.perf_counter() - result = conn.execute( - f"SELECT id FROM {name} ORDER BY embedding {metric_ops} %s LIMIT {top}", - (query,), - ).fetchall() - end = time.perf_counter() - - query_time = end - start - hit = len(set([p[0] for p in result[:top]]) & set(answer[i][:top].tolist())) - results.append((hit, query_time)) - - # Update progress bar with running metrics - curr_results = results[: i + 1] - curr_recall, curr_qps, curr_p50, _ = calculate_metrics(curr_results, top, i + 1) - pbar.set_description( - f"recall: {curr_recall:.4f} QPS: {curr_qps:.2f} P50: {curr_p50:.2f}ms" - ) - - # Calculate final metrics - recall, qps, p50, p99 = calculate_metrics(results, top, m) - - print(f"Top: {top}") - print(f" Recall: {recall:.4f}") - print(f" QPS: {qps:.2f}") - print(f" P50 latency: {p50:.2f}ms") - print(f" P99 latency: {p99:.2f}ms") - - -if __name__ == "__main__": - parser = build_arg_parse() - args = parser.parse_args() - print(args) - - dataset = h5py.File(Path(args.input), "r") - test = dataset["test"][:] - answer = dataset["neighbors"][:] - - if args.metric == "l2": - metric_ops = "<->" - elif args.metric == "cos": - metric_ops = "<=>" - elif args.metric == "dot": - metric_ops = "<#>" - else: - raise ValueError - - if args.processes > 1: - parallel_bench( - args.name, - test, - answer, - metric_ops, - args.processes, - args.url, - args.top, - args.nprob, - args.epsilon, - ) - else: - conn = create_connection(args.url, args.nprob, args.epsilon) - sequential_bench(args.name, test, answer, metric_ops, conn, args.top) \ No newline at end of file diff --git a/scripts/dump.py b/scripts/dump.py deleted file mode 100644 index 68814a99..00000000 --- a/scripts/dump.py +++ /dev/null @@ -1,81 +0,0 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025 TensorChord Inc. - -import argparse - -import h5py -import numpy as np -import psycopg -from tqdm import tqdm -from pgvector.psycopg import register_vector - - -def build_arg_parse(): - parser = argparse.ArgumentParser(description="Dump embeddings to a local file") - parser.add_argument( - "--url", - help="url, like `postgresql://USERNAME:PASSWORD@localhost:5432/DBNAME`", - default="postgresql://postgres:postgres@localhost:5432/postgres" - ) - parser.add_argument("-n", "--name", help="Table name", required=True) - parser.add_argument("-c", "--column", help="Column name", default="embedding") - parser.add_argument("-d", "--dim", help="Dimension", type=int, required=True) - parser.add_argument("-o", "--output", help="Output filepath", required=True) - return parser - - -def create_connection(url): - keepalive_kwargs = { - "keepalives": 1, - "keepalives_idle": 30, - "keepalives_interval": 5, - "keepalives_count": 5, - } - conn = psycopg.connect( - conninfo=url, - autocommit=True, - **keepalive_kwargs, - ) - register_vector(conn) - return conn - - -def extract_vectors(conn, name, column): - with conn.execute(f"SELECT {column} FROM {name}") as cursor: - for row in cursor: - yield row[0] - - -def write_to_h5(filepath, vecs, dim): - with h5py.File(filepath, "w") as file: - dataset = file.create_dataset( - "train", (0, dim), maxshape=(None, dim), chunks=True, dtype=np.float32 - ) - current_size = 0 - for vec in tqdm(vecs): - if dataset.shape[0] == current_size: - dataset.resize(current_size + 64, axis=0) - dataset[current_size] = vec - current_size += 1 - dataset.resize(current_size, axis=0) - dataset.flush() - - -if __name__ == "__main__": - parser = build_arg_parse() - args = parser.parse_args() - print(args) - - conn = create_connection(args.url) - write_to_h5(args.output, extract_vectors(conn, args.name, args.column), args.dim) diff --git a/scripts/index.py b/scripts/index.py deleted file mode 100644 index 39aa4710..00000000 --- a/scripts/index.py +++ /dev/null @@ -1,247 +0,0 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025 TensorChord Inc. - -import asyncio -import math -from time import perf_counter -import argparse -from pathlib import Path -import multiprocessing - -import psycopg -import h5py -from pgvector.psycopg import register_vector_async -import numpy as np -from tqdm import tqdm - -KEEPALIVE_KWARGS = { - "keepalives": 1, - "keepalives_idle": 30, - "keepalives_interval": 5, - "keepalives_count": 5, -} -CHUNKS = 10 - - -def build_arg_parse(): - parser = argparse.ArgumentParser(description="Build index with K-means centroids") - parser.add_argument( - "-m", - "--metric", - help="Distance metric", - default="l2", - choices=["l2", "cos", "dot"], - ) - parser.add_argument("-n", "--name", help="Dataset name, like: sift", required=True) - parser.add_argument("-i", "--input", help="Input filepath", required=False) - parser.add_argument( - "--url", help="url, like `postgresql://postgres:123@localhost:5432/postgres`", required=True - ) - parser.add_argument("-d", "--dim", help="Dimension", type=int, required=True) - # Remember to set `max_worker_processes` at server start - parser.add_argument( - "-w", - "--workers", - help="Workers to build index", - type=int, - required=False, - default=max(multiprocessing.cpu_count() - 1, 1), - ) - parser.add_argument( - "--chunks", - help="chunks for in-memory mode. If OOM, increase it", - type=int, - default=CHUNKS, - ) - # External build - parser.add_argument( - "-c", "--centroids", help="K-means centroids file", required=False - ) - # Internal build - parser.add_argument("--lists", help="Number of centroids", type=int, required=False) - # Do not build index - parser.add_argument("--noindex", help="Do not build index", action='store_true', required=False) - - return parser - - -def get_ivf_ops_config(metric, workers, k=None, name=None): - assert name is not None or k is not None - external_centroids_cfg = """ - [build.external] - table = 'public.{name}_centroids' - """ - if metric == "l2": - metric_ops = "vector_l2_ops" - config = "residual_quantization = true" - internal_centroids_cfg = f""" - [build.internal] - lists = [{k}] - build_threads = {workers} - spherical_centroids = false - """ - elif metric == "cos": - metric_ops = "vector_cosine_ops" - config = "residual_quantization = false" - internal_centroids_cfg = f""" - [build.internal] - lists = [{k}] - build_threads = {workers} - spherical_centroids = true - """ - elif metric == "dot": - metric_ops = "vector_ip_ops" - config = "residual_quantization = false" - internal_centroids_cfg = f""" - [build.internal] - lists = [{k}] - build_threads = {workers} - spherical_centroids = true - """ - else: - raise ValueError - - build_config = ( - external_centroids_cfg.format(name=name) if name else internal_centroids_cfg - ) - return metric_ops, "\n".join([config, build_config]) - - -async def create_connection(url): - conn = await psycopg.AsyncConnection.connect( - conninfo=url, - autocommit=True, - **KEEPALIVE_KWARGS, - ) - await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") - await register_vector_async(conn) - return conn - - -async def add_centroids(conn, name, centroids): - n, dim = centroids.shape - root = np.mean(centroids, axis=0) - await conn.execute(f"DROP TABLE IF EXISTS public.{name}_centroids") - await conn.execute( - f"CREATE TABLE public.{name}_centroids (id integer, parent integer, vector vector({dim}))" - ) - async with conn.cursor().copy( - f"COPY public.{name}_centroids (id, parent, vector) FROM STDIN WITH (FORMAT BINARY)" - ) as copy: - copy.set_types(["integer", "integer", "vector"]) - await copy.write_row((0, None, root)) - for i, centroid in tqdm(enumerate(centroids), desc="Adding centroids", total=n): - await copy.write_row((i + 1, 0, centroid)) - while conn.pgconn.flush() == 1: - await asyncio.sleep(0) - - -async def add_embeddings(conn, name, dim, train, chunks, workers): - await conn.execute(f"DROP TABLE IF EXISTS {name}") - await conn.execute(f"CREATE TABLE {name} (id integer, embedding vector({dim}))") - await conn.execute(f"ALTER TABLE {name} SET (parallel_workers = {workers})") - - n, dim = train.shape - chunk_size = math.ceil(n / chunks) - pbar = tqdm(desc="Adding embeddings", total=n) - for i in range(chunks): - chunk_start = i * chunk_size - chunk_len = min(chunk_size, n - i * chunk_size) - data = train[chunk_start : chunk_start + chunk_len] - - async with conn.cursor().copy( - f"COPY {name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)" - ) as copy: - copy.set_types(["integer", "vector"]) - - for i, vec in enumerate(data): - await copy.write_row((chunk_start + i, vec)) - while conn.pgconn.flush() == 1: - await asyncio.sleep(0) - pbar.update(chunk_len) - pbar.close() - - -async def build_index( - conn, name, workers, metric_ops, ivf_config, finish: asyncio.Event -): - await conn.execute("CREATE EXTENSION IF NOT EXISTS vchord") - start_time = perf_counter() - await conn.execute(f"SET max_parallel_maintenance_workers TO {workers}") - await conn.execute(f"SET max_parallel_workers TO {workers}") - await conn.execute( - f"CREATE INDEX {name}_embedding_idx ON {name} USING vchordrq (embedding {metric_ops}) WITH (options = $${ivf_config}$$)" - ) - print(f"Index build time: {perf_counter() - start_time:.2f}s") - finish.set() - - -async def monitor_index_build(conn, finish: asyncio.Event): - async with conn.cursor() as acur: - blocks_total = None - while blocks_total is None: - await asyncio.sleep(1) - await acur.execute("SELECT blocks_total FROM pg_stat_progress_create_index") - blocks_total = await acur.fetchone() - total = 0 if blocks_total is None else blocks_total[0] - pbar = tqdm(smoothing=0.0, total=total, desc="Building index") - while True: - if finish.is_set(): - pbar.update(pbar.total - pbar.n) - return - await acur.execute("SELECT blocks_done FROM pg_stat_progress_create_index") - blocks_done = await acur.fetchone() - done = 0 if blocks_done is None else blocks_done[0] - pbar.update(done - pbar.n) - await asyncio.sleep(1) - pbar.close() - - -async def main(dataset): - conn = await create_connection(args.url) - if args.centroids: - centroids = np.load(args.centroids, allow_pickle=False) - await add_centroids(conn, args.name, centroids) - - if args.input: - dataset = h5py.File(Path(args.input), "r") - await add_embeddings(conn, args.name, args.dim, dataset["train"], args.chunks, args.workers) - - if not args.noindex: - metric_ops, ivf_config = get_ivf_ops_config( - args.metric, args.workers, args.lists, args.name if args.centroids else None - ) - index_finish = asyncio.Event() - # Need a separate connection for monitor process - monitor_conn = await create_connection(args.url) - monitor_task = monitor_index_build( - monitor_conn, - index_finish, - ) - index_task = build_index( - conn, - args.name, - args.workers, - metric_ops, - ivf_config, - index_finish, - ) - await asyncio.gather(index_task, monitor_task) - - -if __name__ == "__main__": - parser = build_arg_parse() - args = parser.parse_args() - print(args) - asyncio.run(main(args)) diff --git a/scripts/train.py b/scripts/train.py deleted file mode 100644 index 54f86c28..00000000 --- a/scripts/train.py +++ /dev/null @@ -1,218 +0,0 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025 TensorChord Inc. - -import itertools -from multiprocessing import Pool, cpu_count -from time import perf_counter -import argparse -from pathlib import Path -from sys import version_info -from tqdm import tqdm -from numpy import linalg as LA - -if version_info >= (3, 12): - raise RuntimeError("h5py doesn't support 3.12") - -import h5py -from faiss import Kmeans -import numpy as np - -DEFAULT_LISTS = 4096 -N_ITER = 10 -CHUNKS = 10 -SEED = 42 -MAX_POINTS_PER_CLUSTER = 256 - - -def build_arg_parse(): - parser = argparse.ArgumentParser(description="Train K-means centroids") - parser.add_argument("-i", "--input", help="input filepath", required=True) - parser.add_argument("-o", "--output", help="output filepath", required=True) - parser.add_argument( - "--lists", - "--lists-1", - help="Number of centroids", - type=int, - required=False, - default=DEFAULT_LISTS, - ) - parser.add_argument("--lists-2", type=int, help="lower layer lists (if enabled)") - parser.add_argument( - "--niter", help="number of iterations", type=int, default=N_ITER - ) - parser.add_argument("-m", "--metric", choices=["l2", "cos", "dot"], default="l2") - parser.add_argument( - "-g", "--gpu", help="enable GPU for KMeans", action="store_true" - ) - parser.add_argument( - "--mmap", - help="not load by iter, instead use numpy chunk mmap, faster for large dataset", - action="store_true", - ) - parser.add_argument( - "--chunks", - help="chunks for in-memory mode. If OOM, increase it", - type=int, - default=CHUNKS, - ) - return parser - - -def reservoir_sampling(iterator, k: int): - """Reservoir sampling from an iterator.""" - res = [] - for _ in tqdm(range(k), desc="Collect train subset"): - try: - res.append(next(iterator)) - except StopIteration: - return np.vstack(res) - for i, vec in tqdm(enumerate(iterator, k + 1), total=k, desc="Random Pick"): - j = np.random.randint(0, i) - if j < k: - res[j] = vec - return np.vstack(res) - - -def _slice_chunk(args: tuple[int, str, np.ndarray]): - k, file_path, chunk, start_idx = args - dataset = h5py.File(Path(file_path), "r") - data = dataset["train"] - start, end = min(chunk), max(chunk) - indexes = [c - start for c in chunk] - source = data[start : end + 1] - select = source[indexes] - delta, dim = select.shape - - output = np.memmap("index.mmap", dtype=np.float32, mode="r+", shape=(k, dim)) - output[start_idx : start_idx + delta, :] = select - output.flush() - - -def reservoir_sampling_np(data, file_path, k: int, chunks: int): - """Reservoir sampling in memory by numpy.""" - index = np.random.permutation(len(data))[:k] - indices = np.sort(index) - num_processes = cpu_count() - # Split indices into chunks for parallel processing - index_chunks = np.array_split(indices, chunks) - _, dim = data.shape - np.memmap("index.mmap", dtype=np.float32, mode="w+", shape=(k, dim)) - # Create arguments for each chunk - start_idx_acu = [0] - start_idx_acu.extend( - list(itertools.accumulate([len(c) for c in index_chunks[:-1]])) - ) - chunk_args = [ - (k, file_path, chunk, start_idx_acu[i]) for i, chunk in enumerate(index_chunks) - ] - # Process chunks in parallel - with Pool(processes=num_processes) as pool: - list(pool.map(_slice_chunk, chunk_args)) - - -def filter_by_label(iter, labels, target): - for i, vec in enumerate(iter): - if labels[i] == target: - yield vec - - -def kmeans_cluster( - data, - file_path, - k, - child_k, - niter, - metric, - gpu=False, - mmap=False, - chunks=CHUNKS, -): - n, dim = data.shape - if n > MAX_POINTS_PER_CLUSTER * k and not mmap: - train = reservoir_sampling(iter(data), MAX_POINTS_PER_CLUSTER * args.lists) - elif n > MAX_POINTS_PER_CLUSTER * k and mmap: - reservoir_sampling_np( - data, file_path, MAX_POINTS_PER_CLUSTER * args.lists, chunks - ) - train = np.array( - np.memmap( - "index.mmap", - dtype=np.float32, - mode="r", - shape=(MAX_POINTS_PER_CLUSTER * k, dim), - ) - ) - else: - train = data[:] - if metric == "cos": - train = train / LA.norm(train, axis=1, keepdims=True) - kmeans = Kmeans( - dim, k, gpu=gpu, verbose=True, niter=niter, seed=SEED, spherical=metric != "l2" - ) - kmeans.train(train) - if not child_k: - return kmeans.centroids - - # train the lower layer k-means - labels = np.zeros(n, dtype=np.uint32) - for i, vec in tqdm(enumerate(data), desc="Assigning labels"): - _, label = kmeans.assign(vec.reshape((1, -1))) - labels[i] = label[0] - - centroids = [] - total_k = k * child_k - for i in tqdm(range(k), desc="training k-means for child layers"): - samples = np.sum(labels == i) / n * total_k * MAX_POINTS_PER_CLUSTER - child_train = reservoir_sampling( - filter_by_label(iter(data), labels, i), samples - ) - child_kmeans = Kmeans( - dim, - child_k, - gpu=gpu, - verbose=True, - niter=niter, - seed=SEED, - spherical=metric != "l2", - ) - child_kmeans.train(child_train) - centroids.append(child_kmeans.centroids) - return np.vstack(centroids) - - -if __name__ == "__main__": - parser = build_arg_parse() - args = parser.parse_args() - print(args) - - dataset = h5py.File(Path(args.input), "r") - n, dim = dataset["train"].shape - - start_time = perf_counter() - centroids = kmeans_cluster( - dataset["train"], - args.input, - args.lists, - args.lists_2, - args.niter, - args.metric, - args.gpu, - args.mmap, - args.chunks, - ) - print( - f"K-means (k=({args.lists}, {args.lists_2})): {perf_counter() - start_time:.2f}s" - ) - - np.save(Path(args.output), centroids, allow_pickle=False) From 6f667ba24cd5a3395c8cc0969e68aac3f106fdc3 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 10 Feb 2026 18:46:21 +0800 Subject: [PATCH 272/324] chore: upload 1.1.0 schema scripts (#421) Signed-off-by: usamoi --- .github/workflows/release.yml | 2 +- crates/vchordg/src/tuples.rs | 2 +- crates/vchordrq/src/tuples.rs | 2 +- sql/install/vchord--1.1.0.sql | 1291 ++++++++++++++++++++++++++ sql/upgrade/vchord--1.0.0--1.1.0.sql | 405 ++++++++ src/lib.rs | 2 + src/sql/finalize.sql | 36 +- src/upgrade.rs | 14 +- 8 files changed, 1726 insertions(+), 28 deletions(-) create mode 100644 sql/install/vchord--1.1.0.sql create mode 100644 sql/upgrade/vchord--1.0.0--1.1.0.sql diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 684988a7..775192a2 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -110,7 +110,7 @@ jobs: Maintainer: Tensorchord Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM Homepage: https://vectorchord.ai/ - License: AGPL-3.0-only or Elastic-2.0" \ + License: AGPL-3.0-only OR Elastic-2.0" \ > ~/debian/DEBIAN/control (cd ~/debian && find usr -type f -print0 | xargs -0 md5sum) > ~/debian/DEBIAN/md5sums dpkg-deb --root-owner-group -Zxz --build ~/debian ~/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index a456f5bb..0f5da7e7 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -21,7 +21,7 @@ use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordg\0"); -const VERSION: u64 = 1000; +const VERSION: u64 = 1001; #[inline(always)] fn tag(source: &[u8]) -> Tag { diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index e02574b4..a66bacfd 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -20,7 +20,7 @@ use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); -const VERSION: u64 = 1000; +const VERSION: u64 = 1001; #[inline(always)] fn tag(source: &[u8]) -> Tag { diff --git a/sql/install/vchord--1.1.0.sql b/sql/install/vchord--1.1.0.sql new file mode 100644 index 00000000..89241be2 --- /dev/null +++ b/sql/install/vchord--1.1.0.sql @@ -0,0 +1,1291 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:47 +-- bootstrap +-- List of shell types + +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE rabitq8; +CREATE TYPE sphere_rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq4; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ + "rhs" halfvec[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:40 +-- vchord::datatype::functions_rabitq4::_vchord_halfvec_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:40 +-- vchord::datatype::functions_rabitq8::_vchord_halfvec_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:20 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_in +CREATE FUNCTION "_vchord_rabitq4_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:41 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_cosine +CREATE FUNCTION "_vchord_rabitq4_operator_cosine"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:31 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_ip +CREATE FUNCTION "_vchord_rabitq4_operator_ip"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:21 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_l2 +CREATE FUNCTION "_vchord_rabitq4_operator_l2"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:123 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:140 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_out +CREATE FUNCTION "_vchord_rabitq4_out"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:36 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_recv +CREATE FUNCTION "_vchord_rabitq4_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:21 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_send +CREATE FUNCTION "_vchord_rabitq4_send"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:99 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq4_sphere_cosine_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:75 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq4_sphere_ip_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:51 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq4_sphere_l2_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq4.rs:17 +-- vchord::datatype::typmod_rabitq4::_vchord_rabitq4_typmod_in +CREATE FUNCTION "_vchord_rabitq4_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:20 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_in +CREATE FUNCTION "_vchord_rabitq8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:41 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_cosine +CREATE FUNCTION "_vchord_rabitq8_operator_cosine"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:31 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_ip +CREATE FUNCTION "_vchord_rabitq8_operator_ip"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:21 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_l2 +CREATE FUNCTION "_vchord_rabitq8_operator_l2"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:123 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:140 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_out +CREATE FUNCTION "_vchord_rabitq8_out"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:36 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_recv +CREATE FUNCTION "_vchord_rabitq8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:21 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_send +CREATE FUNCTION "_vchord_rabitq8_send"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:99 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq8_sphere_cosine_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:75 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq8_sphere_ip_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:51 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq8_sphere_l2_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq8.rs:17 +-- vchord::datatype::typmod_rabitq8::_vchord_rabitq8_typmod_in +CREATE FUNCTION "_vchord_rabitq8_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ + "rhs" vector[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:22 +-- vchord::datatype::functions_rabitq4::_vchord_vector_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:22 +-- vchord::datatype::functions_rabitq8::_vchord_vector_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:110 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordg_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordg_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordg_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordg_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordg_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordg_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:192 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:100 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:95 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:90 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:140 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:130 +-- vchord::index::opclass::_vchordrq_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:125 +-- vchord::index::opclass::_vchordrq_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:120 +-- vchord::index::opclass::_vchordrq_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:150 +-- vchord::index::opclass::_vchordrq_support_rabitq4_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:115 +-- vchord::index::opclass::_vchordrq_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:110 +-- vchord::index::opclass::_vchordrq_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:105 +-- vchord::index::opclass::_vchordrq_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:145 +-- vchord::index::opclass::_vchordrq_support_rabitq8_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:85 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:135 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:48 +-- finalize +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external +); + +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, + radius REAL +); + +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/upgrade/vchord--1.0.0--1.1.0.sql b/sql/upgrade/vchord--1.0.0--1.1.0.sql new file mode 100644 index 00000000..555c3e6a --- /dev/null +++ b/sql/upgrade/vchord--1.0.0--1.1.0.sql @@ -0,0 +1,405 @@ +-- List of shell types + +CREATE TYPE rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq8; +CREATE TYPE sphere_rabitq4; + +-- List of internal changes + +DROP OPERATOR <<=>>(sphere_scalar8, scalar8); +DROP OPERATOR <<#>>(sphere_scalar8, scalar8); +DROP OPERATOR <<->>(sphere_scalar8, scalar8); +DROP FUNCTION quantize_to_scalar8(halfvec); +DROP FUNCTION quantize_to_scalar8(vector); +DROP FUNCTION sphere(scalar8, real); +DROP OPERATOR <<=>>(scalar8, sphere_scalar8); +DROP OPERATOR <<#>>(scalar8, sphere_scalar8); +DROP OPERATOR <<->>(scalar8, sphere_scalar8); +DROP OPERATOR <=>(scalar8, scalar8); +DROP OPERATOR <#>(scalar8, scalar8); +DROP OPERATOR <->(scalar8, scalar8); +DROP FUNCTION _vchord_scalar8_operator_cosine; +DROP FUNCTION _vchord_scalar8_operator_ip; +DROP FUNCTION _vchord_scalar8_operator_l2; +DROP FUNCTION _vchord_scalar8_sphere_cosine_in; +DROP FUNCTION _vchord_scalar8_sphere_ip_in; +DROP FUNCTION _vchord_scalar8_sphere_l2_in; +DROP TYPE sphere_scalar8; +ALTER TYPE scalar8 SET (TYPMOD_IN = NONE); +DROP FUNCTION _vchord_typmod_in_65535; +ALTER TYPE scalar8 SET (TYPMOD_OUT = NONE); +DROP FUNCTION _vchord_typmod_out; +ALTER TYPE scalar8 SET (RECEIVE = NONE); +DROP FUNCTION _vchord_scalar8_recv; +ALTER TYPE scalar8 SET (SEND = NONE); +DROP FUNCTION _vchord_scalar8_send; +DROP TYPE scalar8 CASCADE; +-- DROP FUNCTION _vchord_scalar8_in; +-- DROP FUNCTION _vchord_scalar8_out; + +CREATE FUNCTION _vchord_rabitq4_in(input cstring, oid oid, typmod INT) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_cosine(lhs rabitq4, rhs rabitq4) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_cosine_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_ip(lhs rabitq4, rhs rabitq4) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_ip_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_l2(lhs rabitq4, rhs rabitq4) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_l2_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_out(vector rabitq4) RETURNS cstring +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_out_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_recv(internal internal, oid oid, typmod INT) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_recv_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_send(vector rabitq4) RETURNS bytea +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_send_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_sphere_cosine_in(lhs rabitq4, rhs sphere_rabitq4) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_cosine_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_sphere_ip_in(lhs rabitq4, rhs sphere_rabitq4) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_ip_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_sphere_l2_in(lhs rabitq4, rhs sphere_rabitq4) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_l2_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_typmod_in(list cstring[]) RETURNS INT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_typmod_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_in(input cstring, oid oid, typmod INT) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_operator_cosine(lhs rabitq8, rhs rabitq8) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_cosine_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_operator_ip(lhs rabitq8, rhs rabitq8) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_ip_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_operator_l2(lhs rabitq8, rhs rabitq8) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_l2_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_out(vector rabitq8) RETURNS cstring +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_out_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_recv(internal internal, oid oid, typmod INT) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_recv_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_send(vector rabitq8) RETURNS bytea +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_send_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_sphere_cosine_in(lhs rabitq8, rhs sphere_rabitq8) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_cosine_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_sphere_ip_in(lhs rabitq8, rhs sphere_rabitq8) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_ip_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_sphere_l2_in(lhs rabitq8, rhs sphere_rabitq8) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_l2_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_typmod_in(list cstring[]) RETURNS INT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_typmod_in_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq4_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq4_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_ip_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq4_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_l2_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq8_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq8_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_ip_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq8_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_l2_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq4_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq4_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_ip_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq4_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_l2_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq4_maxsim_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_maxsim_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq8_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq8_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_ip_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq8_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_l2_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq8_maxsim_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_maxsim_ops_wrapper'; + +DROP OPERATOR <<=>>(sphere_halfvec, halfvec); +DROP OPERATOR <<#>>(sphere_halfvec, halfvec); +DROP OPERATOR <<->>(sphere_halfvec, halfvec); +DROP OPERATOR <<=>>(sphere_vector, vector); +DROP OPERATOR <<#>>(sphere_vector, vector); +DROP OPERATOR <<->>(sphere_vector, vector); + +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external +); + +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, + radius REAL +); + +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + +-- List of functions + +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + +-- List of operator families + +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); diff --git a/src/lib.rs b/src/lib.rs index 6a18e50d..f9e8c20c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -42,6 +42,8 @@ pgrx::pg_module_magic!( const { STR } } ); +const _: &str = include_str!("./sql/bootstrap.sql"); +const _: &str = include_str!("./sql/finalize.sql"); pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap); pgrx::extension_sql_file!("./sql/finalize.sql", finalize); diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 4e0847a6..c344e102 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -93,85 +93,73 @@ CREATE OPERATOR <=> ( CREATE OPERATOR <<->> ( PROCEDURE = _vchord_vector_sphere_l2_in, LEFTARG = vector, - RIGHTARG = sphere_vector, - COMMUTATOR = <<->> + RIGHTARG = sphere_vector ); CREATE OPERATOR <<->> ( PROCEDURE = _vchord_halfvec_sphere_l2_in, LEFTARG = halfvec, - RIGHTARG = sphere_halfvec, - COMMUTATOR = <<->> + RIGHTARG = sphere_halfvec ); CREATE OPERATOR <<->> ( PROCEDURE = _vchord_rabitq8_sphere_l2_in, LEFTARG = rabitq8, - RIGHTARG = sphere_rabitq8, - COMMUTATOR = <<->> + RIGHTARG = sphere_rabitq8 ); CREATE OPERATOR <<->> ( PROCEDURE = _vchord_rabitq4_sphere_l2_in, LEFTARG = rabitq4, - RIGHTARG = sphere_rabitq4, - COMMUTATOR = <<->> + RIGHTARG = sphere_rabitq4 ); CREATE OPERATOR <<#>> ( PROCEDURE = _vchord_vector_sphere_ip_in, LEFTARG = vector, - RIGHTARG = sphere_vector, - COMMUTATOR = <<#>> + RIGHTARG = sphere_vector ); CREATE OPERATOR <<#>> ( PROCEDURE = _vchord_halfvec_sphere_ip_in, LEFTARG = halfvec, - RIGHTARG = sphere_halfvec, - COMMUTATOR = <<#>> + RIGHTARG = sphere_halfvec ); CREATE OPERATOR <<#>> ( PROCEDURE = _vchord_rabitq8_sphere_ip_in, LEFTARG = rabitq8, - RIGHTARG = sphere_rabitq8, - COMMUTATOR = <<#>> + RIGHTARG = sphere_rabitq8 ); CREATE OPERATOR <<#>> ( PROCEDURE = _vchord_rabitq4_sphere_ip_in, LEFTARG = rabitq4, - RIGHTARG = sphere_rabitq4, - COMMUTATOR = <<#>> + RIGHTARG = sphere_rabitq4 ); CREATE OPERATOR <<=>> ( PROCEDURE = _vchord_vector_sphere_cosine_in, LEFTARG = vector, - RIGHTARG = sphere_vector, - COMMUTATOR = <<=>> + RIGHTARG = sphere_vector ); CREATE OPERATOR <<=>> ( PROCEDURE = _vchord_halfvec_sphere_cosine_in, LEFTARG = halfvec, - RIGHTARG = sphere_halfvec, - COMMUTATOR = <<=>> + RIGHTARG = sphere_halfvec ); CREATE OPERATOR <<=>> ( PROCEDURE = _vchord_rabitq8_sphere_cosine_in, LEFTARG = rabitq8, - RIGHTARG = sphere_rabitq8, - COMMUTATOR = <<=>> + RIGHTARG = sphere_rabitq8 ); CREATE OPERATOR <<=>> ( PROCEDURE = _vchord_rabitq4_sphere_cosine_in, LEFTARG = rabitq4, - RIGHTARG = sphere_rabitq4, - COMMUTATOR = <<=>> + RIGHTARG = sphere_rabitq4 ); CREATE OPERATOR @# ( diff --git a/src/upgrade.rs b/src/upgrade.rs index 15f81f67..1dc273ed 100644 --- a/src/upgrade.rs +++ b/src/upgrade.rs @@ -18,7 +18,6 @@ // * https://www.postgresql.org/message-id/CACX+KaPOzzRHEt4w_=iqKbTpMKjyrUGVng1C749yP3r6dprtcg@mail.gmail.com // * https://github.com/tensorchord/pgvecto.rs/issues/397 -#[allow(unused_macros)] macro_rules! symbol { ($t:ident) => { paste::paste! { @@ -42,3 +41,16 @@ macro_rules! symbol { } }; } + +symbol!(_vchord_vector_quantize_to_scalar8_wrapper); +symbol!(_vchord_halfvec_quantize_to_scalar8_wrapper); +symbol!(_vchord_scalar8_operator_cosine_wrapper); +symbol!(_vchord_scalar8_operator_ip_wrapper); +symbol!(_vchord_scalar8_operator_l2_wrapper); +symbol!(_vchord_scalar8_sphere_cosine_in_wrapper); +symbol!(_vchord_scalar8_sphere_ip_in_wrapper); +symbol!(_vchord_scalar8_sphere_l2_in_wrapper); +symbol!(_vchord_typmod_in_65535_wrapper); +symbol!(_vchord_typmod_out_wrapper); +symbol!(_vchord_scalar8_recv_wrapper); +symbol!(_vchord_scalar8_send_wrapper); From 0aeb82298e9321ee755ca9e191ff44d7a6c8d0be Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 11 Feb 2026 15:11:04 +0800 Subject: [PATCH 273/324] readme: sync with docs (#422) https://github.com/tensorchord/pgvecto.rs-docs/pull/165 Signed-off-by: usamoi --- README.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index bdacb991..8174d9c6 100644 --- a/README.md +++ b/README.md @@ -67,8 +67,9 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg18-v1.0.0 + -d ghcr.io/tensorchord/vchord-postgres:pg18-v1.1.0 ``` + > [!NOTE] > In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). @@ -110,9 +111,11 @@ For more usage, please read: - [Indexing](https://docs.vectorchord.ai/vectorchord/usage/indexing.html) - [Multi-Vector Retrieval](https://docs.vectorchord.ai/vectorchord/usage/indexing-with-maxsim-operators.html) - [Graph Index](https://docs.vectorchord.ai/vectorchord/usage/graph-index.html) +- [Quantization Types](https://docs.vectorchord.ai/vectorchord/usage/quantization-types.html) - [Similarity Filter](https://docs.vectorchord.ai/vectorchord/usage/range-query.html) - [PostgreSQL Tuning](https://docs.vectorchord.ai/vectorchord/usage/performance-tuning.html) - [Monitoring](https://docs.vectorchord.ai/vectorchord/usage/monitoring.html) +- [Fallback Parameters](https://docs.vectorchord.ai/vectorchord/usage/fallback-parameters.html) - [Measure Recall](https://docs.vectorchord.ai/vectorchord/usage/measure-recall.html) - [Prewarm](https://docs.vectorchord.ai/vectorchord/usage/prewarm.html) - [Prefilter](https://docs.vectorchord.ai/vectorchord/usage/prefilter.html) From 66470e0244c048808f05bfa57df501301607592b Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 12 Feb 2026 15:53:58 +0800 Subject: [PATCH 274/324] ci: remove ruff action (#425) Signed-off-by: usamoi --- .github/workflows/check.yml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index f8cc8edd..894bcf97 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -37,9 +37,6 @@ jobs: - name: Taplo run: taplo fmt --check - - name: Ruff - uses: astral-sh/ruff-action@57714a7c8a2e59f32539362ba31877a1957dded1 # v3.5.1 - - name: Rustfmt run: cargo fmt --check -- --config-path /dev/null --config imports_granularity=Module From 0e0230f9401cc1be49f8f71cd91e1008abc05ce4 Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 12 Feb 2026 16:55:53 +0800 Subject: [PATCH 275/324] chore: update dependencies (#426) Signed-off-by: usamoi --- Cargo.lock | 397 +++++++++++++++++++++++------ Cargo.toml | 12 +- crates/feistel/src/lib.rs | 4 +- crates/k_means/src/flat.rs | 2 +- crates/k_means/src/hierarchical.rs | 2 +- crates/k_means/src/index.rs | 2 +- crates/k_means/src/quick.rs | 2 +- crates/k_means/src/rabitq.rs | 2 +- crates/rabitq/build.rs | 2 +- crates/simd/Cargo.toml | 4 +- crates/simd/benches/bench.rs | 14 +- crates/simd/src/byte.rs | 20 +- crates/simd/src/fht.rs | 2 +- crates/simd/src/floating_f16.rs | 36 +-- crates/simd/src/floating_f32.rs | 60 ++--- crates/simd/src/halfbyte.rs | 12 +- crates/simd/src/quantize.rs | 8 +- crates/xtask/Cargo.toml | 4 +- src/index/gucs.rs | 148 +---------- src/index/sample.rs | 77 +----- src/index/vchordrq/am/mod.rs | 6 +- src/recorder/types.rs | 2 +- 22 files changed, 423 insertions(+), 395 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 618138d4..ebbd5472 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -98,9 +98,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.100" +version = "1.0.101" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" +checksum = "5f0e0fee31ef5ed1ba1316088939cea399010ed7731dba877ed44aeb407a75ea" [[package]] name = "arrayvec" @@ -116,9 +116,9 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "bindgen" -version = "0.71.1" +version = "0.72.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" dependencies = [ "annotate-snippets", "bitflags", @@ -164,7 +164,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "374b7c592d9c00c1f4972ea58390ac6b18cbb6ab79011f3bdc90a0b82ca06b77" dependencies = [ "serde", - "toml", + "toml 0.9.12+spec-1.1.0", ] [[package]] @@ -175,9 +175,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.2.53" +version = "1.2.55" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "755d2fce177175ffca841e9a06afdb2c4ab0f593d53b4dee48147dfaade85932" +checksum = "47b26a0954ae34af09b50f0de26458fa95369a0d478d8236d3f93082b219bd29" dependencies = [ "find-msvc-tools", "shlex", @@ -208,6 +208,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core 0.10.0", +] + [[package]] name = "ciborium" version = "0.2.2" @@ -248,9 +259,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.54" +version = "4.5.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6e6ff9dcd79cff5cd969a17a545d79e84ab086e444102a591e288a8aa3ce394" +checksum = "63be97961acde393029492ce0be7a1af7e323e6bae9511ebfac33751be5e6806" dependencies = [ "clap_builder", "clap_derive", @@ -258,9 +269,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.54" +version = "4.5.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa42cf4d2b7a41bc8f663a7cab4031ebafa1bf3875705bfaf8466dc60ab52c00" +checksum = "7f13174bda5dfd69d7e947827e5af4b0f2f94a4a3ee92912fba07a66150f21e2" dependencies = [ "anstream", "anstyle", @@ -270,9 +281,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.5.49" +version = "4.5.55" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2a0b5487afeab2deb2ff4e03a807ad1a03ac532ff5a2cee5d86884440c7f7671" +checksum = "a92793da1a46a5f2a02a6f4c46c6496b28c43638adea8306fcb0caa1634f24e5" dependencies = [ "heck", "proc-macro2", @@ -282,9 +293,9 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.7.7" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3e64b0cc0439b12df2fa678eae89a1c56a529fd067a9115f7827f1fffd22b32" +checksum = "3a822ea5bc7590f9d40f1ba12c0dc3c2760f3482c6984db1573ad11031420831" [[package]] name = "codepage" @@ -310,6 +321,15 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -321,9 +341,9 @@ dependencies = [ [[package]] name = "criterion" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d883447757bb0ee46f233e9dc22eb84d93a9508c9b868687b274fc431d886bf" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" dependencies = [ "alloca", "anes", @@ -346,9 +366,9 @@ dependencies = [ [[package]] name = "criterion-plot" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed943f81ea2faa8dcecbbfa50164acf95d555afec96a27871663b300e387b2e4" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" dependencies = [ "cast", "itertools", @@ -548,9 +568,9 @@ dependencies = [ [[package]] name = "find-msvc-tools" -version = "0.1.8" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8591b0bcc8a98a64310a2fae1bb3e9b8564dd10e381e6e28010fde8e8e8568db" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" [[package]] name = "fixedbitset" @@ -560,9 +580,9 @@ checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" [[package]] name = "flate2" -version = "1.1.8" +version = "1.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b375d6465b98090a5f25b1c7703f3859783755aa9a80433b36e0379a3ec2f369" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", "miniz_oxide", @@ -613,6 +633,20 @@ dependencies = [ "wasip2", ] +[[package]] +name = "getrandom" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139ef39800118c7683f2fd3c98c1b23c09ae076556b435f8e9064ae108aaeeec" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "rand_core 0.10.0", + "wasip2", + "wasip3", +] + [[package]] name = "glob" version = "0.3.3" @@ -759,6 +793,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + [[package]] name = "ident_case" version = "1.0.1" @@ -815,6 +855,8 @@ checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" dependencies = [ "equivalent", "hashbrown 0.16.1", + "serde", + "serde_core", ] [[package]] @@ -866,11 +908,17 @@ dependencies = [ "simd", ] +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + [[package]] name = "libc" -version = "0.2.180" +version = "0.2.181" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc" +checksum = "459427e2af2b9c839b132acb702a1c654d95e10f8c326bfc2ad11310e458b1c5" [[package]] name = "libloading" @@ -884,9 +932,9 @@ dependencies = [ [[package]] name = "libm" -version = "0.2.15" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libmimalloc-sys" @@ -921,11 +969,17 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + [[package]] name = "memchr" -version = "2.7.6" +version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] name = "mimalloc" @@ -986,7 +1040,7 @@ dependencies = [ "flate2", "memchr", "ruzstd", - "wasmparser", + "wasmparser 0.243.0", ] [[package]] @@ -1062,9 +1116,9 @@ dependencies = [ [[package]] name = "pgrx" -version = "0.16.1" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdcfb88f7fa9ba42b4ea9d1f85a1d968bbb407cc30308b35f73bdfe6c966f64b" +checksum = "dd42eed31a68e3e6d7cec284cdaacd8ffa9b2b3ebb3f1df9ec2fd3558834c6e7" dependencies = [ "bitflags", "bitvec", @@ -1083,9 +1137,9 @@ dependencies = [ [[package]] name = "pgrx-bindgen" -version = "0.16.1" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00e35193b7e71e2f612d336cecd00db0f049f4cc609f2b1c9a34755b5ec559d7" +checksum = "d30fa3688ab1cb6429ca46725eb32d5d9f186806c0f0aa6419f86e418b8a81ab" dependencies = [ "bindgen", "cc", @@ -1102,9 +1156,9 @@ dependencies = [ [[package]] name = "pgrx-catalog" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3342c1e98f08dc57ceb8a44835f362ae3de4422bada1fbbccfcd74a4b55dc978" +checksum = "d7b8320f441bd6aa7460a397d1a16dfb915979a1f665c85d9a7bc10d57e38cb8" dependencies = [ "paste", "pgrx", @@ -1112,9 +1166,9 @@ dependencies = [ [[package]] name = "pgrx-macros" -version = "0.16.1" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dab542dd4041773874f90cd8e3448195749548dc3fb1daf501e7e11ebfb1dd22" +checksum = "1ca7cc537daa4cc004f12ad1d31756b506358aa900251a82c3280e719da88fa4" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", @@ -1124,9 +1178,9 @@ dependencies = [ [[package]] name = "pgrx-pg-config" -version = "0.16.1" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eff9b29df94c3f9fcb0cde220f92eea6975ed05962784a98fb557754ad663501" +checksum = "1cd8670ad408093a83615fc68071929d7677b40021a475649a626319c464f0a3" dependencies = [ "cargo_toml", "codepage", @@ -1137,16 +1191,16 @@ dependencies = [ "serde", "serde_json", "thiserror", - "toml", + "toml 0.9.12+spec-1.1.0", "url", "winapi", ] [[package]] name = "pgrx-pg-sys" -version = "0.16.1" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "934f2536953ccb6722bef2cfdfb1f8d6d3cd4a4f2c508d56ec85b649c5680c2b" +checksum = "d9f81faae87d533b573156261b17521c0c7f66f5ca5dab2b9e6ddbf5a63dbd14" dependencies = [ "cee-scape", "libc", @@ -1158,9 +1212,9 @@ dependencies = [ [[package]] name = "pgrx-sql-entity-graph" -version = "0.16.1" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07a767cb9faa612f1ba7f13718136d4006950d6f253b414ef487a03e85c47a94" +checksum = "4800fa91a481a83d3c1aaffd1324fcd2e587df0ed342c71d6a66d0899027e505" dependencies = [ "convert_case", "eyre", @@ -1244,6 +1298,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + [[package]] name = "proc-macro-error-attr2" version = "2.0.0" @@ -1308,22 +1372,23 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "rand" -version = "0.9.2" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" dependencies = [ - "rand_chacha", - "rand_core", + "chacha20", + "getrandom 0.4.1", + "rand_core 0.10.0", ] [[package]] name = "rand_chacha" -version = "0.9.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +checksum = "3e6af7f3e25ded52c41df4e0b1af2d047e45896c2f3281792ed68a1c243daedb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.10.0", ] [[package]] @@ -1331,9 +1396,12 @@ name = "rand_core" version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" -dependencies = [ - "getrandom", -] + +[[package]] +name = "rand_core" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" [[package]] name = "rayon" @@ -1357,9 +1425,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.12.2" +version = "1.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843bc0191f75f3e22651ae5f1e72939ab2f72a4bc30fa80a066bd66edefc24d4" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" dependencies = [ "aho-corasick", "memchr", @@ -1369,9 +1437,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.13" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5276caf25ac86c8d810222b3dbb938e512c55c6831a10f3e6ed1c93b84041f1c" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" dependencies = [ "aho-corasick", "memchr", @@ -1380,9 +1448,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.8" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" +checksum = "a96887878f22d7bad8a3b6dc5b7440e0ada9a245242924394987b21cf2210a4c" [[package]] name = "rsqlite-vfs" @@ -1458,6 +1526,12 @@ version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" +[[package]] +name = "semver" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" + [[package]] name = "seq-macro" version = "0.3.6" @@ -1606,9 +1680,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.114" +version = "2.0.115" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4d107df263a3013ef9b1879b0df87d706ff80f65a86ea879bd9c31f9b307c2a" +checksum = "6e614ed320ac28113fa64972c4262d5dbc89deacdfd00c34a3e4cea073243c12" dependencies = [ "proc-macro2", "quote", @@ -1640,12 +1714,12 @@ checksum = "591ef38edfb78ca4771ee32cf494cb8771944bee237a9b91fc9c1424ac4b777b" [[package]] name = "tempfile" -version = "3.24.0" +version = "3.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "655da9c7eb6305c55742045d5a8d2037996d61d8de95806335c7c86ce0f82e9c" +checksum = "0136791f7c95b1f6dd99f9cc786b91bb81c3800b639b3478e561ddb7be95e5f1" dependencies = [ "fastrand", - "getrandom", + "getrandom 0.4.1", "once_cell", "rustix", "windows-sys", @@ -1693,14 +1767,29 @@ dependencies = [ [[package]] name = "toml" -version = "0.9.11+spec-1.1.0" +version = "0.9.12+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime 0.7.5+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow", +] + +[[package]] +name = "toml" +version = "1.0.0+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3afc9a848309fe1aaffaed6e1546a7a14de1f935dc9d89d32afd9a44bab7c46" +checksum = "d1d7e18e3dd1d31e0ee5e863a8091ffec2fcc271636586042452b656a22c8ee1" dependencies = [ "indexmap", "serde_core", "serde_spanned", - "toml_datetime", + "toml_datetime 1.0.0+spec-1.1.0", "toml_parser", "toml_writer", "winnow", @@ -1715,11 +1804,20 @@ dependencies = [ "serde_core", ] +[[package]] +name = "toml_datetime" +version = "1.0.0+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32c2555c699578a4f59f0cc68e5116c8d7cabbd45e1409b989d4be085b53f13e" +dependencies = [ + "serde_core", +] + [[package]] name = "toml_parser" -version = "1.0.6+spec-1.1.0" +version = "1.0.7+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3198b4b0a8e11f09dd03e133c0280504d0801269e9afa46362ffde1cbeebf44" +checksum = "247eaa3197818b831697600aadf81514e577e0cba5eab10f7e064e78ae154df1" dependencies = [ "winnow", ] @@ -1744,9 +1842,9 @@ checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] name = "unicode-ident" -version = "1.0.22" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" +checksum = "537dd038a89878be9b64dd4bd1b260315c1bb94f4d784956b81e27a088d9a09e" [[package]] name = "unicode-segmentation" @@ -1760,6 +1858,12 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + [[package]] name = "url" version = "2.5.8" @@ -1786,11 +1890,11 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.19.0" +version = "1.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2e054861b4bd027cd373e18e8d8d8e6548085000e41290d95ce0c373a654b4a" +checksum = "ee48d38b119b0cd71fe4141b30f5ba9c7c5d9f4e7a3a8b4a674e4b6ef789976f" dependencies = [ - "getrandom", + "getrandom 0.3.4", "js-sys", "wasm-bindgen", ] @@ -1851,7 +1955,7 @@ dependencies = [ "serde", "simd", "small_iter", - "toml", + "toml 1.0.0+spec-1.1.0", "validator", "vchordg", "vchordrq", @@ -1928,6 +2032,15 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "wasm-bindgen" version = "0.2.108" @@ -1973,6 +2086,28 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser 0.244.0", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap", + "wasm-encoder", + "wasmparser 0.244.0", +] + [[package]] name = "wasmparser" version = "0.243.0" @@ -1982,6 +2117,18 @@ dependencies = [ "bitflags", ] +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags", + "hashbrown 0.15.5", + "indexmap", + "semver", +] + [[package]] name = "web-sys" version = "0.3.85" @@ -2066,6 +2213,88 @@ name = "wit-bindgen" version = "0.51.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] + +[[package]] +name = "wit-bindgen-core" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck", + "indexmap", + "prettyplease", + "syn", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn", + "wit-bindgen-core", + "wit-bindgen-rust", +] + +[[package]] +name = "wit-component" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser 0.244.0", + "wit-parser", +] + +[[package]] +name = "wit-parser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser 0.244.0", +] [[package]] name = "writeable" @@ -2079,7 +2308,7 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ca4d373340c479fd1e779f7a763acee85da3e423b1a9a9acccf97babcc92edbb" dependencies = [ - "rand_core", + "rand_core 0.9.5", ] [[package]] @@ -2129,18 +2358,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.33" +version = "0.8.39" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "668f5168d10b9ee831de31933dc111a459c97ec93225beb307aed970d1372dfd" +checksum = "db6d35d663eadb6c932438e763b262fe1a70987f9ae936e60158176d710cae4a" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.33" +version = "0.8.39" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c7962b26b0a8685668b671ee4b54d007a67d4eaf05fda79ac0ecf41e32270f1" +checksum = "4122cd3169e94605190e77839c9a40d40ed048d305bfdc146e7df40ab0f3e517" dependencies = [ "proc-macro2", "quote", @@ -2203,6 +2432,6 @@ dependencies = [ [[package]] name = "zmij" -version = "1.0.16" +version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfcd145825aace48cff44a8844de64bf75feec3080e0aa5cdbde72961ae51a65" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/Cargo.toml b/Cargo.toml index d1ee1f92..743081fe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,14 +39,14 @@ crossbeam-channel = "0.5.15" dary_heap.workspace = true humansize = "2.1.3" paste.workspace = true -pgrx = { version = "=0.16.1", default-features = false, features = ["cshim"] } -pgrx-catalog = "0.3.1" +pgrx = { version = "=0.17.0", default-features = false, features = ["cshim"] } +pgrx-catalog = "0.3.2" rand.workspace = true rayon.workspace = true rusqlite = { version = "0.38.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true -toml = "0.9.11" +toml = "1.0.0" validator.workspace = true wyhash.workspace = true zerocopy.workspace = true @@ -69,14 +69,14 @@ edition = "2024" bumpalo = "3.19.1" dary_heap = "0.3.8" paste = "1.0.15" -rand = "0.9.2" -rand_chacha = "0.9.0" +rand = "0.10.0" +rand_chacha = "0.10.0" rayon = "1.11.0" seq-macro = "0.3.6" serde = { version = "1.0.228", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } wyhash = "0.6.0" -zerocopy = { version = "0.8.33", features = ["derive"] } +zerocopy = { version = "0.8.39", features = ["derive"] } [workspace.lints] # complexity diff --git a/crates/feistel/src/lib.rs b/crates/feistel/src/lib.rs index 54b16380..1098e7b3 100644 --- a/crates/feistel/src/lib.rs +++ b/crates/feistel/src/lib.rs @@ -102,8 +102,8 @@ fn is_a_permutation() { fn sample() { let n = if cfg!(not(miri)) { 6370_u32 } else { 637_u32 }; let width = (n.ilog2() + 1).next_multiple_of(2); - let key_0 = rand::Rng::random(&mut rand::rng()); - let key_1 = rand::Rng::random(&mut rand::rng()); + let key_0 = rand::RngExt::random(&mut rand::rng()); + let key_1 = rand::RngExt::random(&mut rand::rng()); let secret = move |round: u32, x: u32| { let buffer = [round.to_le_bytes(), x.to_le_bytes(), key_0, key_1]; wyhash::wyhash(buffer.as_flattened(), 0) as u32 diff --git a/crates/k_means/src/flat.rs b/crates/k_means/src/flat.rs index f1770616..0d032b6c 100644 --- a/crates/k_means/src/flat.rs +++ b/crates/k_means/src/flat.rs @@ -16,7 +16,7 @@ use crate::index::{flat_index as prefect_index, flat_index as index}; use crate::square::{Square, SquareMut}; use crate::{KMeans, This}; use rand::rngs::StdRng; -use rand::{Rng, SeedableRng}; +use rand::{RngExt, SeedableRng}; use rayon::prelude::*; struct Flat<'a> { diff --git a/crates/k_means/src/hierarchical.rs b/crates/k_means/src/hierarchical.rs index 180e0f08..05d9bbf4 100644 --- a/crates/k_means/src/hierarchical.rs +++ b/crates/k_means/src/hierarchical.rs @@ -266,7 +266,7 @@ fn test_modified_webster_method() { #[test] fn test_partition() { - fn gen_random_alloc(rows: usize, groups: usize, rng: &mut impl Rng) -> Vec> { + fn gen_random_alloc(rows: usize, groups: usize, rng: &mut impl RngExt) -> Vec> { let mut idx: Vec = (0..rows).collect(); idx.shuffle(rng); let mut alloc = Vec::with_capacity(groups); diff --git a/crates/k_means/src/index.rs b/crates/k_means/src/index.rs index 2a6ddb9c..0bb59100 100644 --- a/crates/k_means/src/index.rs +++ b/crates/k_means/src/index.rs @@ -14,7 +14,7 @@ use crate::This; use crate::square::{Square, SquareMut}; -use rand::Rng; +use rand::RngExt; use rayon::prelude::*; use simd::Floating; diff --git a/crates/k_means/src/quick.rs b/crates/k_means/src/quick.rs index d07db425..72fd8fca 100644 --- a/crates/k_means/src/quick.rs +++ b/crates/k_means/src/quick.rs @@ -16,7 +16,7 @@ use crate::KMeans; use crate::index::flat_index as prefect_index; use crate::square::{Square, SquareMut}; use rand::rngs::StdRng; -use rand::{Rng, SeedableRng}; +use rand::{RngExt, SeedableRng}; use rayon::iter::{IntoParallelIterator, ParallelIterator}; struct Quick { diff --git a/crates/k_means/src/rabitq.rs b/crates/k_means/src/rabitq.rs index a7e567e8..a40ff34b 100644 --- a/crates/k_means/src/rabitq.rs +++ b/crates/k_means/src/rabitq.rs @@ -16,7 +16,7 @@ use crate::index::{flat_index as prefect_index, rabitq_index as index}; use crate::square::{Square, SquareMut}; use crate::{KMeans, This}; use rand::rngs::StdRng; -use rand::{Rng, SeedableRng}; +use rand::{RngExt, SeedableRng}; use rayon::prelude::*; struct RaBitQ<'a> { diff --git a/crates/rabitq/build.rs b/crates/rabitq/build.rs index 54e24743..18e81d0a 100644 --- a/crates/rabitq/build.rs +++ b/crates/rabitq/build.rs @@ -17,7 +17,7 @@ use std::error::Error; use std::path::PathBuf; fn main() -> Result<(), Box> { - use rand::{Rng, SeedableRng}; + use rand::{RngExt, SeedableRng}; use rand_chacha::ChaCha12Rng; let mut rng = ChaCha12Rng::from_seed([7; 32]); let mut bits = (0..262144).map(|_| rng.random::()).collect::>(); diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 7b8a5c6c..6309ab1f 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -22,11 +22,11 @@ seq-macro.workspace = true zerocopy.workspace = true [dev-dependencies] -criterion = "0.8.1" +criterion = "0.8.2" rand.workspace = true [build-dependencies] -cc = "1.2.53" +cc = "1.2.55" which = "8.0.0" [lints] diff --git a/crates/simd/benches/bench.rs b/crates/simd/benches/bench.rs index 28f3b5ba..200d9919 100644 --- a/crates/simd/benches/bench.rs +++ b/crates/simd/benches/bench.rs @@ -17,7 +17,7 @@ use criterion::{Criterion, criterion_group, criterion_main}; fn floating_f32_reduce_sum_of_xy(c: &mut Criterion) { - use rand::Rng; + use rand::RngExt; let mut rng = rand::rng(); let x = (0..4095) .map(|_| rng.random_range(-1.0..=1.0f32)) @@ -62,7 +62,7 @@ fn floating_f32_reduce_sum_of_xy(c: &mut Criterion) { } fn floating_f32_reduce_sum_of_d2(c: &mut Criterion) { - use rand::Rng; + use rand::RngExt; let mut rng = rand::rng(); let x = (0..4095) .map(|_| rng.random_range(-1.0..=1.0f32)) @@ -107,7 +107,7 @@ fn floating_f32_reduce_sum_of_d2(c: &mut Criterion) { } fn floating_f16_reduce_sum_of_xy(c: &mut Criterion) { - use rand::Rng; + use rand::RngExt; use simd::{F16, f16}; let mut rng = rand::rng(); let x = (0..4095) @@ -155,7 +155,7 @@ fn floating_f16_reduce_sum_of_xy(c: &mut Criterion) { } fn floating_f16_reduce_sum_of_d2(c: &mut Criterion) { - use rand::Rng; + use rand::RngExt; use simd::{F16, f16}; let mut rng = rand::rng(); let x = (0..4095) @@ -203,7 +203,7 @@ fn floating_f16_reduce_sum_of_d2(c: &mut Criterion) { } fn byte_reduce_sum_of_xy(c: &mut Criterion) { - use rand::Rng; + use rand::RngExt; let mut rng = rand::rng(); let x = (0..4095).map(|_| rng.random::()).collect::>(); let y = (0..4095).map(|_| rng.random::()).collect::>(); @@ -248,7 +248,7 @@ fn byte_reduce_sum_of_xy(c: &mut Criterion) { } fn byte_reduce_sum_of_x(c: &mut Criterion) { - use rand::Rng; + use rand::RngExt; let mut rng = rand::rng(); let this = (0..4095).map(|_| rng.random::()).collect::>(); #[cfg(target_arch = "x86_64")] @@ -278,7 +278,7 @@ fn byte_reduce_sum_of_x(c: &mut Criterion) { } fn halfbyte_reduce_sum_of_xy(c: &mut Criterion) { - use rand::Rng; + use rand::RngExt; let mut rng = rand::rng(); let x = (0..2047).map(|_| rng.random::()).collect::>(); let y = (0..2047).map(|_| rng.random::()).collect::>(); diff --git a/crates/simd/src/byte.rs b/crates/simd/src/byte.rs index f553709e..444f9ca6 100644 --- a/crates/simd/src/byte.rs +++ b/crates/simd/src/byte.rs @@ -51,7 +51,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v4_avx512vnni_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vnni") { println!("test {} ... skipped (v4:avx512vnni)", module_path!()); return; @@ -120,7 +120,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v4_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; @@ -190,7 +190,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xy_v3_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; @@ -260,7 +260,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xy_v2_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); return; @@ -304,7 +304,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a2_dotprod_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("dotprod") { println!("test {} ... skipped (a2:dotprod)", module_path!()); return; @@ -366,7 +366,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a2_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); return; @@ -436,7 +436,7 @@ mod reduce_sum_of_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_x_v4_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; @@ -483,7 +483,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x_v3_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; @@ -530,7 +530,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x_v2_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); return; @@ -576,7 +576,7 @@ mod reduce_sum_of_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_x_a2_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); return; diff --git a/crates/simd/src/fht.rs b/crates/simd/src/fht.rs index 1378ff1a..d99a671b 100644 --- a/crates/simd/src/fht.rs +++ b/crates/simd/src/fht.rs @@ -151,7 +151,7 @@ mod tests { #[test] fn fht() { - use rand::Rng; + use rand::RngExt; use std::iter::zip; const EPSILON: f32 = 1e-6; let mut rng = rand::rng(); diff --git a/crates/simd/src/floating_f16.rs b/crates/simd/src/floating_f16.rs index e7b0d7ea..8a9f7131 100644 --- a/crates/simd/src/floating_f16.rs +++ b/crates/simd/src/floating_f16.rs @@ -190,7 +190,7 @@ mod reduce_sum_of_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_x_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -238,7 +238,7 @@ mod reduce_sum_of_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_x_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -303,7 +303,7 @@ mod reduce_sum_of_abs_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_abs_x_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -352,7 +352,7 @@ mod reduce_sum_of_abs_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_abs_x_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -417,7 +417,7 @@ mod reduce_sum_of_x2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_x2_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -465,7 +465,7 @@ mod reduce_sum_of_x2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_x2_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -535,7 +535,7 @@ mod reduce_min_max_of_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_min_max_of_x_v4_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; @@ -591,7 +591,7 @@ mod reduce_min_max_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_min_max_of_x_v3_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; @@ -653,7 +653,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v4_avx512fp16_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { println!("test {} ... skipped (v4:avx512fp16)", module_path!()); @@ -711,7 +711,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -767,7 +767,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -815,7 +815,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a3_512_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("a3.512") { println!("test {} ... skipped (a3.512)", module_path!()); @@ -864,7 +864,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a2_fp16_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { println!("test {} ... skipped (a2:fp16)", module_path!()); @@ -929,7 +929,7 @@ mod reduce_sum_of_d2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_v4_avx512fp16_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { println!("test {} ... skipped (v4:avx512fp16)", module_path!()); @@ -989,7 +989,7 @@ mod reduce_sum_of_d2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -1052,7 +1052,7 @@ mod reduce_sum_of_d2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -1100,7 +1100,7 @@ mod reduce_sum_of_d2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_a3_512_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("a3.512") { println!("test {} ... skipped (a3.512)", module_path!()); @@ -1149,7 +1149,7 @@ mod reduce_sum_of_d2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_a2_fp16_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 2.0; if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { println!("test {} ... skipped (a2:fp16)", module_path!()); diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs index f6a09c74..6bcb476f 100644 --- a/crates/simd/src/floating_f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -185,7 +185,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -240,7 +240,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -290,7 +290,7 @@ mod reduce_sum_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x_v2_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); @@ -341,7 +341,7 @@ mod reduce_sum_of_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_x_a2_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -403,7 +403,7 @@ mod reduce_sum_of_abs_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_abs_x_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.009; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -462,7 +462,7 @@ mod reduce_sum_of_abs_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_abs_x_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.009; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -515,7 +515,7 @@ mod reduce_sum_of_abs_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_abs_x_v2_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.009; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); @@ -568,7 +568,7 @@ mod reduce_sum_of_abs_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_abs_x_a2_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.009; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -628,7 +628,7 @@ mod reduce_sum_of_x2 { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x2_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -683,7 +683,7 @@ mod reduce_sum_of_x2 { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x2_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -734,7 +734,7 @@ mod reduce_sum_of_x2 { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_x2_v2_fma_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { println!("test {} ... skipped (v2:fma)", module_path!()); @@ -785,7 +785,7 @@ mod reduce_sum_of_x2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_x2_a2_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.008; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -851,7 +851,7 @@ mod reduce_min_max_of_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_min_max_of_x_v4_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; @@ -913,7 +913,7 @@ mod reduce_min_max_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_min_max_of_x_v3_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; @@ -964,7 +964,7 @@ mod reduce_min_max_of_x { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_min_max_of_x_v2_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); return; @@ -1016,7 +1016,7 @@ mod reduce_min_max_of_x { #[test] #[cfg_attr(miri, ignore)] fn reduce_min_max_of_x_a2_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); return; @@ -1082,7 +1082,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xy_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.004; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -1147,7 +1147,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xy_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.004; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -1207,7 +1207,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xy_v2_fma_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.004; if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { println!("test {} ... skipped (v2:fma)", module_path!()); @@ -1255,7 +1255,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a3_256_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.004; if !crate::is_cpu_detected!("a3.256") { println!("test {} ... skipped (a3.256)", module_path!()); @@ -1315,7 +1315,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a2_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.004; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -1388,7 +1388,7 @@ mod reduce_sum_of_d2 { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_d2_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.02; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -1456,7 +1456,7 @@ mod reduce_sum_of_d2 { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_d2_v3_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.02; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); @@ -1518,7 +1518,7 @@ mod reduce_sum_of_d2 { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_d2_v2_fma_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.02; if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { println!("test {} ... skipped (v2:fma)", module_path!()); @@ -1566,7 +1566,7 @@ mod reduce_sum_of_d2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_a3_256_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.02; if !crate::is_cpu_detected!("a3.256") { println!("test {} ... skipped (a3.256)", module_path!()); @@ -1628,7 +1628,7 @@ mod reduce_sum_of_d2 { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_a2_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.02; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); @@ -1720,7 +1720,7 @@ mod reduce_sum_of_xy_sparse { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_sparse_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.000001; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -1728,7 +1728,7 @@ mod reduce_sum_of_xy_sparse { } let mut rng = rand::rng(); pub fn sample_u32_sorted( - rng: &mut (impl Rng + ?Sized), + rng: &mut (impl RngExt + ?Sized), length: u32, amount: u32, ) -> Vec { @@ -1871,7 +1871,7 @@ mod reduce_sum_of_d2_sparse { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_d2_sparse_v4_test() { - use rand::Rng; + use rand::RngExt; const EPSILON: f32 = 0.0004; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); @@ -1879,7 +1879,7 @@ mod reduce_sum_of_d2_sparse { } let mut rng = rand::rng(); pub fn sample_u32_sorted( - rng: &mut (impl Rng + ?Sized), + rng: &mut (impl RngExt + ?Sized), length: u32, amount: u32, ) -> Vec { diff --git a/crates/simd/src/halfbyte.rs b/crates/simd/src/halfbyte.rs index f1e80273..03a9ae49 100644 --- a/crates/simd/src/halfbyte.rs +++ b/crates/simd/src/halfbyte.rs @@ -57,7 +57,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v4_avx512vnni_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vnni") { println!("test {} ... skipped (v4:avx512vnni)", module_path!()); return; @@ -127,7 +127,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_v4_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; @@ -198,7 +198,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xy_v3_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; @@ -269,7 +269,7 @@ mod reduce_sum_of_xy { #[cfg(all(target_arch = "x86_64", test))] #[test] fn reduce_sum_of_xy_v2_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v2") { println!("test {} ... skipped (v2)", module_path!()); return; @@ -313,7 +313,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a2_dotprod_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("dotprod") { println!("test {} ... skipped (a2:dotprod)", module_path!()); return; @@ -388,7 +388,7 @@ mod reduce_sum_of_xy { #[test] #[cfg_attr(miri, ignore)] fn reduce_sum_of_xy_a2_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); return; diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index 50c050b9..98d54b81 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -50,7 +50,7 @@ mod mul_add_round { #[test] #[cfg_attr(miri, ignore)] fn mul_add_round_v4_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v4") { println!("test {} ... skipped (v4)", module_path!()); return; @@ -122,7 +122,7 @@ mod mul_add_round { #[cfg(all(target_arch = "x86_64", test))] #[test] fn mul_add_round_v3_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v3") { println!("test {} ... skipped (v3)", module_path!()); return; @@ -189,7 +189,7 @@ mod mul_add_round { #[cfg(all(target_arch = "x86_64", test))] #[test] fn mul_add_round_v2_fma_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { println!("test {} ... skipped (v2:fma)", module_path!()); return; @@ -263,7 +263,7 @@ mod mul_add_round { #[test] #[cfg_attr(miri, ignore)] fn mul_add_round_a2_test() { - use rand::Rng; + use rand::RngExt; if !crate::is_cpu_detected!("a2") { println!("test {} ... skipped (a2)", module_path!()); return; diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index 6b922f9c..8b092d7e 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -5,13 +5,13 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.54", features = ["derive", "env"] } +clap = { version = "4.5.58", features = ["derive", "env"] } object = { version = "0.38.1", features = ["read", "wasm"] } serde.workspace = true serde_json = "1.0.149" shlex = "1.3.0" target-triple = "1.0.0" -tempfile = "3.24.0" +tempfile = "3.25.0" [lints] workspace = true diff --git a/src/index/gucs.rs b/src/index/gucs.rs index ee3cc726..039909a7 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -37,7 +37,7 @@ static VCHORDG_ENABLE_SCAN: GucSetting = GucSetting::::new(true); static VCHORDG_EF_SEARCH: GucSetting = GucSetting::::new(64); -static mut VCHORDG_EF_SEARCH_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); +static mut VCHORDG_EF_SEARCH_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); static VCHORDG_BEAM_SEARCH: GucSetting = GucSetting::::new(1); @@ -61,21 +61,22 @@ static VCHORDRQ_ENABLE_SCAN: GucSetting = GucSetting::::new(true); static VCHORDRQ_PROBES: GucSetting> = GucSetting::>::new(Some(c"")); -static mut VCHORDRQ_PROBES_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); +static mut VCHORDRQ_PROBES_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); static VCHORDRQ_EPSILON: GucSetting = GucSetting::::new(1.9); -static mut VCHORDRQ_EPSILON_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); +static mut VCHORDRQ_EPSILON_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); static VCHORDRQ_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); static VCHORDRQ_MAXSIM_REFINE: GucSetting = GucSetting::::new(0); -static mut VCHORDRQ_MAXSIM_REFINE_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); +static mut VCHORDRQ_MAXSIM_REFINE_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); -static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut sys::config_generic = core::ptr::null_mut(); +static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut pgrx::pg_sys::config_generic = + core::ptr::null_mut(); static VCHORDRQ_PREFILTER: GucSetting = GucSetting::::new(false); @@ -285,7 +286,7 @@ pub fn init() { #[cfg(any(feature = "pg14", feature = "pg15"))] unsafe { let len = pgrx::pg_sys::GetNumConfigOptions() as usize; - let arr = sys::get_guc_variables(); + let arr = pgrx::pg_sys::get_guc_variables(); let mut sources = (0..len).map(|i| arr.add(i).read()); debug_assert!(targets.is_sorted_by(|(a, _), (b, _)| guc_name_compare(a, b).is_le())); for (name, ptr) in targets { @@ -307,13 +308,13 @@ pub fn init() { unsafe { use pgrx::pg_sys::PGERROR; for (name, ptr) in targets { - *ptr = sys::find_option(name.as_ptr(), false, false, PGERROR as _); + *ptr = pgrx::pg_sys::find_option(name.as_ptr(), false, false, PGERROR as _); assert!(check(*ptr, name), "failed to find GUC {name:?}"); } } } -unsafe fn check(p: *mut sys::config_generic, name: &CStr) -> bool { +unsafe fn check(p: *mut pgrx::pg_sys::config_generic, name: &CStr) -> bool { if p.is_null() { return false; } @@ -505,137 +506,6 @@ pub fn vchordrq_query_sampling_rate() -> f64 { VCHORDRQ_QUERY_SAMPLING_RATE.get() } -#[allow(non_camel_case_types)] -#[allow(non_snake_case)] -mod sys { - #[cfg(not(feature = "pg14"))] - #[cfg(not(feature = "pg15"))] - #[cfg(not(feature = "pg16"))] - #[cfg(not(feature = "pg17"))] - #[cfg(not(feature = "pg18"))] - compile_error!("bindings are not checked"); - - pub mod config_type { - pub type Type = ::core::ffi::c_uint; - } - - #[repr(C)] - #[derive(Copy, Clone)] - pub union config_var_val { - pub boolval: bool, - pub intval: ::core::ffi::c_int, - pub realval: f64, - pub stringval: *mut ::core::ffi::c_char, - pub enumval: ::core::ffi::c_int, - } - - #[repr(C)] - #[derive(Copy, Clone)] - pub struct config_var_value { - pub val: config_var_val, - pub extra: *mut ::core::ffi::c_void, - } - - pub mod config_group { - pub type Type = ::core::ffi::c_uint; - } - - pub mod GucStackState { - pub type Type = ::core::ffi::c_uint; - } - - #[repr(C)] - #[derive(Copy, Clone)] - pub struct guc_stack { - pub prev: *mut guc_stack, - pub nest_level: ::core::ffi::c_int, - pub state: GucStackState::Type, - pub source: pgrx::pg_sys::GucSource::Type, - pub scontext: pgrx::pg_sys::GucContext::Type, - pub masked_scontext: pgrx::pg_sys::GucContext::Type, - #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] - pub srole: pgrx::pg_sys::Oid, - #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] - pub masked_srole: pgrx::pg_sys::Oid, - pub prior: config_var_value, - pub masked: config_var_value, - } - - pub type GucStack = guc_stack; - - #[repr(C)] - #[derive(Debug, Copy, Clone)] - pub struct config_generic { - pub name: *const ::core::ffi::c_char, - pub context: pgrx::pg_sys::GucContext::Type, - pub group: config_group::Type, - pub short_desc: *const ::core::ffi::c_char, - pub long_desc: *const ::core::ffi::c_char, - pub flags: ::core::ffi::c_int, - pub vartype: config_type::Type, - pub status: ::core::ffi::c_int, - pub source: pgrx::pg_sys::GucSource::Type, - pub reset_source: pgrx::pg_sys::GucSource::Type, - pub scontext: pgrx::pg_sys::GucContext::Type, - pub reset_scontext: pgrx::pg_sys::GucContext::Type, - #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] - pub srole: pgrx::pg_sys::Oid, - #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] - pub reset_srole: pgrx::pg_sys::Oid, - pub stack: *mut GucStack, - pub extra: *mut ::core::ffi::c_void, - #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] - pub nondef_link: pgrx::pg_sys::dlist_node, - #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] - pub stack_link: pgrx::pg_sys::slist_node, - #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] - pub report_link: pgrx::pg_sys::slist_node, - pub last_reported: *mut ::core::ffi::c_char, - pub sourcefile: *mut ::core::ffi::c_char, - pub sourceline: ::core::ffi::c_int, - } - - #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] - #[inline] - pub unsafe fn find_option( - name: *const ::core::ffi::c_char, - create_placeholders: bool, - skip_errors: bool, - elevel: ::core::ffi::c_int, - ) -> *mut config_generic { - use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; - #[cfg_attr(target_os = "windows", link(name = "postgres"))] - unsafe extern "C-unwind" { - #[link_name = "find_option"] - unsafe fn f( - name: *const ::core::ffi::c_char, - create_placeholders: bool, - skip_errors: bool, - elevel: ::core::ffi::c_int, - ) -> *mut config_generic; - } - #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] - unsafe { - pg_guard_ffi_boundary(move || f(name, create_placeholders, skip_errors, elevel)) - } - } - - #[cfg(any(feature = "pg14", feature = "pg15"))] - #[inline] - pub unsafe fn get_guc_variables() -> *mut *mut config_generic { - use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; - #[cfg_attr(target_os = "windows", link(name = "postgres"))] - unsafe extern "C-unwind" { - #[link_name = "get_guc_variables"] - unsafe fn f() -> *mut *mut config_generic; - } - #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] - unsafe { - pg_guard_ffi_boundary(move || f()) - } - } -} - #[allow(dead_code)] fn guc_name_compare(a: &CStr, b: &CStr) -> std::cmp::Ordering { let (a, b) = (a.to_bytes_with_nul(), b.to_bytes_with_nul()); diff --git a/src/index/sample.rs b/src/index/sample.rs index 2d8c6d3e..5bb54735 100644 --- a/src/index/sample.rs +++ b/src/index/sample.rs @@ -203,8 +203,8 @@ impl Tuple for HeapTuple<'_> { fn sample(n: u32) -> Box> { let width = (n.ilog2() + 1).next_multiple_of(2); - let key_0 = rand::Rng::random(&mut rand::rng()); - let key_1 = rand::Rng::random(&mut rand::rng()); + let key_0 = rand::RngExt::random(&mut rand::rng()); + let key_1 = rand::RngExt::random(&mut rand::rng()); let secret = move |round: u32, x: u32| { let buffer = [round.to_le_bytes(), x.to_le_bytes(), key_0, key_1]; wyhash::wyhash(buffer.as_flattened(), 0) as u32 @@ -254,80 +254,9 @@ struct AssertSync(T); unsafe impl Sync for AssertSync {} -static TSM: AssertSync = AssertSync(sys::TsmRoutine { +static TSM: AssertSync = AssertSync(pgrx::pg_sys::TsmRoutine { type_: pgrx::pg_sys::NodeTag::T_TsmRoutine, NextSampleBlock: Some(feistel_rows_nextsampleblock), NextSampleTuple: Some(feistel_rows_nextsampletuple), ..unsafe { core::mem::zeroed() } }); - -#[allow(non_camel_case_types)] -#[allow(non_snake_case)] -mod sys { - #[cfg(not(feature = "pg14"))] - #[cfg(not(feature = "pg15"))] - #[cfg(not(feature = "pg16"))] - #[cfg(not(feature = "pg17"))] - #[cfg(not(feature = "pg18"))] - compile_error!("bindings are not checked"); - - use core::ffi::c_int; - use pgrx::pg_sys::{ - BlockNumber, Datum, List, NodeTag, OffsetNumber, PlannerInfo, RelOptInfo, SampleScanState, - }; - - pub type SampleScanGetSampleSize_function = Option< - unsafe extern "C-unwind" fn( - root: *mut PlannerInfo, - baserel: *mut RelOptInfo, - paramexprs: *mut List, - pages: *mut BlockNumber, - tuples: *mut f64, - ), - >; - - pub type InitSampleScan_function = - Option; - - pub type BeginSampleScan_function = Option< - unsafe extern "C-unwind" fn( - node: *mut SampleScanState, - params: *mut Datum, - nparams: c_int, - seed: u32, - ), - >; - - pub type NextSampleBlock_function = Option< - unsafe extern "C-unwind" fn( - node: *mut SampleScanState, - nblocks: BlockNumber, - ) -> BlockNumber, - >; - - pub type NextSampleTuple_function = Option< - unsafe extern "C-unwind" fn( - node: *mut SampleScanState, - blockno: BlockNumber, - maxoffset: OffsetNumber, - ) -> OffsetNumber, - >; - - pub type EndSampleScan_function = - Option; - - #[repr(C)] - #[derive(Debug, Copy, Clone)] - pub struct TsmRoutine { - pub type_: NodeTag, - pub parameterTypes: *mut List, - pub repeatable_across_queries: bool, - pub repeatable_across_scans: bool, - pub SampleScanGetSampleSize: SampleScanGetSampleSize_function, - pub InitSampleScan: InitSampleScan_function, - pub BeginSampleScan: BeginSampleScan_function, - pub NextSampleBlock: NextSampleBlock_function, - pub NextSampleTuple: NextSampleTuple_function, - pub EndSampleScan: EndSampleScan_function, - } -} diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 32ed017a..d69d1328 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -24,7 +24,7 @@ use crate::index::vchordrq::scanners::*; use crate::recorder::DefaultRecorder; use pgrx::datum::Internal; use pgrx::pg_sys::Datum; -use rand::RngCore; +use rand::RngExt; use std::cell::LazyCell; use std::ffi::CStr; use std::num::NonZero; @@ -412,9 +412,9 @@ unsafe fn aminsertinner( ctid: pgrx::pg_sys::ItemPointer, ) -> bool { struct RngChooser(T); - impl InsertChooser for RngChooser { + impl InsertChooser for RngChooser { fn choose(&mut self, n: NonZero) -> usize { - rand::Rng::random_range(&mut self.0, 0..n.get()) + RngExt::random_range(&mut self.0, 0..n.get()) } } diff --git a/src/recorder/types.rs b/src/recorder/types.rs index 39ec9a97..f32076e1 100644 --- a/src/recorder/types.rs +++ b/src/recorder/types.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use crate::recorder::worker::push; -use rand::Rng; +use rand::RngExt; use std::cell::RefMut; pub trait Recorder { From 0e5d3c0b84cbafc72f33588d25866d900c525ce3 Mon Sep 17 00:00:00 2001 From: cutecutecat Date: Sat, 14 Feb 2026 17:34:02 +0800 Subject: [PATCH 276/324] docs: new readme (#417) Preview: https://github.com/cutecutecat/VectorChord/tree/new-readme --------- Signed-off-by: cutecutecat Co-authored-by: Jinjing Zhou --- README.md | 38 +++++++++++++++----------------------- 1 file changed, 15 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index 8174d9c6..f71df9f9 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@

VectorChord

-

Effortlessly host 100 million 768-dimensional vectors (250GB+) on an AWS i4i.xlarge instance ($250/month), featuring 4 vCPUs and 32GB of RAM with VectorChord.

+

Ready for the Billion-Scale Era. Host 100M vectors on a single i4i.xlarge ($247/mo) and (scale seamlessly to 1B+)[https://blog.vectorchord.ai/scaling-vector-search-to-1-billion-on-postgresql].

@@ -16,19 +16,16 @@ [![][github-downloads-shield]][github-downloads-link] [![][discord-shield]][discord-link] [![][X-shield]][X-link] -[![][cloud-shield]][cloud-link] [![][deepwiki-shield]][deepwiki-link] [![][license-1-shield]][license-1-link] [![][license-2-shield]][license-2-link]
+VectorChord (vchord) is a PostgreSQL extension engineered for scalable, high-performance, and cost-effective vector search. -> [!NOTE] -> VectorChord serves as the successor to [pgvecto.rs](https://github.com/tensorchord/pgvecto.rs) [![][previous-docker-pulls-shield]][previous-docker-pulls-link] with better stability and performance. If you are interested in this new solution, you may find the [migration guide](https://docs.vectorchord.ai/vectorchord/admin/migration.html) helpful. - -VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search. +To efficiently store vectors while preserving search quality, VectorChord applies RaBitQ[^1] compression together with autonomous reranking. With VectorChord, you can store 400,000 vectors for just $1, enabling significant savings: 6x more vectors compared to Pinecone's optimized storage and 26x more than pgvector/pgvecto.rs for the same price. -With VectorChord, you can store 400,000 vectors for just $1, enabling significant savings: 6x more vectors compared to Pinecone's optimized storage and 26x more than pgvector/pgvecto.rs for the same price[^1]. +[^1]: Gao, Jianyang, and Cheng Long. "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search." Proceedings of the ACM on Management of Data 2.3 (2024): 1-27. ![][image-compare] @@ -36,27 +33,22 @@ With VectorChord, you can store 400,000 vectors for just $1, enabling significan VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: -**⚡ Enhanced Performance**: Delivering optimized operations with up to 5x faster queries, 16x higher insert throughput, and 16x quicker[^1] index building compared to pgvector's HNSW implementation. - -[^1]: Based on [MyScale Benchmark](https://myscale.github.io/benchmark/) with 768-dimensional vectors and 95% recall. Please check out our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. +**💰 Affordable Vector Search**: Host 100M × 768-dimensional vectors → AWS i4i.xlarge ($247/month)[^2], host 1B × 96-dimensional vectors → i7ie.6xlarge ($2246/month)[^3], helping you keep infrastructure costs down while maintaining competitive search quality. -**💰 Affordable Vector Search**: Query 100M 768-dimensional vectors using just 32GB of memory, achieving 35ms P50 latency with top10 recall@95%, helping you keep infrastructure costs down while maintaining high search quality. +[^2]: Please check out our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. +[^3]: Please check out our [blog post](https://blog.vectorchord.ai/scaling-vector-search-to-1-billion-on-postgresql) for more details. -**🔌 Seamless Integration**: Fully compatible with pgvector data types and syntax while providing optimal defaults out of the box - no manual parameter tuning needed. Just drop in VectorChord for enhanced performance. +**⚡ Accelerated Index Build**: Index 100 million vectors in just 20 minutes. Powered by hierarchical K-means and highly optimized disk operations, VectorChord eliminates the bottleneck of vector indexing on a single machine with limited hardware resources. -**🔧 Accelerated Index Build**: Leverage IVF to build indexes externally (e.g., on GPU) for faster KMeans clustering, combined with RaBitQ[^2] compression to efficiently store vectors while maintaining search quality through autonomous reranking. +[^4]: Please check out our [blog post](https://blog.vectorchord.ai/how-we-made-100m-vector-indexing-in-20-minutes-possible-on-postgresql#heading-hierarchical-k-means) for more technique details and [document](https://docs.vectorchord.ai/vectorchord/usage/partitioning-tuning.html#hierarchical-k-means) for usages. -[^2]: Gao, Jianyang, and Cheng Long. "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search." Proceedings of the ACM on Management of Data 2.3 (2024): 1-27. +**📈 Smoothly Scale Up**: Scale with confidence as your data grows. Through dimensionality reduction and sampling[^5], VectorChord effectively controls memory growth, enabling 1B-vector indexes to be built on machines with 128GB of memory in practice. -**📏 Long Vector Support**: Store and search vectors up to 60,000[^3] dimensions, enabling the use of the best high-dimensional models like text-embedding-3-large with ease. +[^5]: Please check out our [blog post](https://blog.vectorchord.ai/how-we-made-100m-vector-indexing-in-20-minutes-possible-on-postgresql#heading-dimensionality-reduction) for more technique details and [document](https://docs.vectorchord.ai/vectorchord/usage/partitioning-tuning.html#reduce-sampling-factor) for usages. -[^3]: There is a [limitation](https://github.com/pgvector/pgvector#vector-type) at pgvector of 16,000 dimensions now. If you really have a large dimension(`16,000 with any questions or requests regarding the licenses. +[cost-estimation]: https://github.com/user-attachments/assets/168fe550-6465-4eee-a224-8c848c301e3d [image-compare]: https://github.com/user-attachments/assets/2d985f1e-7093-4c3a-8bf3-9f0b92c0e7e7 [license-1-link]: https://github.com/tensorchord/VectorChord#license [license-1-shield]: https://img.shields.io/badge/License-AGPLv3-green?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCIgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiBmaWxsPSIjZmZmZmZmIj48cGF0aCBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik0xMi43NSAyLjc1YS43NS43NSAwIDAwLTEuNSAwVjQuNUg5LjI3NmExLjc1IDEuNzUgMCAwMC0uOTg1LjMwM0w2LjU5NiA1Ljk1N0EuMjUuMjUgMCAwMTYuNDU1IDZIMi4zNTNhLjc1Ljc1IDAgMTAwIDEuNUgzLjkzTC41NjMgMTUuMThhLjc2Mi43NjIgMCAwMC4yMS44OGMuMDguMDY0LjE2MS4xMjUuMzA5LjIyMS4xODYuMTIxLjQ1Mi4yNzguNzkyLjQzMy42OC4zMTEgMS42NjIuNjIgMi44NzYuNjJhNi45MTkgNi45MTkgMCAwMDIuODc2LS42MmMuMzQtLjE1NS42MDYtLjMxMi43OTItLjQzMy4xNS0uMDk3LjIzLS4xNTguMzEtLjIyM2EuNzUuNzUgMCAwMC4yMDktLjg3OEw1LjU2OSA3LjVoLjg4NmMuMzUxIDAgLjY5NC0uMTA2Ljk4NC0uMzAzbDEuNjk2LTEuMTU0QS4yNS4yNSAwIDAxOS4yNzUgNmgxLjk3NXYxNC41SDYuNzYzYS43NS43NSAwIDAwMCAxLjVoMTAuNDc0YS43NS43NSAwIDAwMC0xLjVIMTIuNzVWNmgxLjk3NGMuMDUgMCAuMS4wMTUuMTQuMDQzbDEuNjk3IDEuMTU0Yy4yOS4xOTcuNjMzLjMwMy45ODQuMzAzaC44ODZsLTMuMzY4IDcuNjhhLjc1Ljc1IDAgMDAuMjMuODk2Yy4wMTIuMDA5IDAgMCAuMDAyIDBhMy4xNTQgMy4xNTQgMCAwMC4zMS4yMDZjLjE4NS4xMTIuNDUuMjU2Ljc5LjRhNy4zNDMgNy4zNDMgMCAwMDIuODU1LjU2OCA3LjM0MyA3LjM0MyAwIDAwMi44NTYtLjU2OWMuMzM4LS4xNDMuNjA0LS4yODcuNzktLjM5OWEzLjUgMy41IDAgMDAuMzEtLjIwNi43NS43NSAwIDAwLjIzLS44OTZMMjAuMDcgNy41aDEuNTc4YS43NS43NSAwIDAwMC0xLjVoLTQuMTAyYS4yNS4yNSAwIDAxLS4xNC0uMDQzbC0xLjY5Ny0xLjE1NGExLjc1IDEuNzUgMCAwMC0uOTg0LS4zMDNIMTIuNzVWMi43NXpNMi4xOTMgMTUuMTk4YTUuNDE4IDUuNDE4IDAgMDAyLjU1Ny42MzUgNS40MTggNS40MTggMCAwMDIuNTU3LS42MzVMNC43NSA5LjM2OGwtMi41NTcgNS44M3ptMTQuNTEtLjAyNGMuMDgyLjA0LjE3NC4wODMuMjc1LjEyNi41My4yMjMgMS4zMDUuNDUgMi4yNzIuNDVhNS44NDYgNS44NDYgMCAwMDIuNTQ3LS41NzZMMTkuMjUgOS4zNjdsLTIuNTQ3IDUuODA3eiI+PC9wYXRoPjwvc3ZnPgo= @@ -157,8 +151,6 @@ You may choose either license based on your needs. We welcome any commercial col [discord-shield]: https://img.shields.io/discord/974584200327991326?&logoColor=white&color=5865F2&style=flat&logo=discord&cacheSeconds=60 [X-link]: https://x.com/TensorChord [X-shield]: https://img.shields.io/badge/follow-%40tensorchord-1DA1F2?logo=x&style=flat&logoColor=white&color=1da1f2 -[cloud-link]: https://cloud.vectorchord.ai/ -[cloud-shield]: https://img.shields.io/badge/VectorChord_Cloud-Try_For_Free-F2B263.svg?labelColor=DAFDBA&logo=data:image/svg%2bxml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSIxMzMiIGhlaWdodD0iMTMyIiBmaWxsPSJub25lIj48cGF0aCBmaWxsPSIjRTZEQjNEIiBkPSJNNDguNCAzNy41YzAtMSAwLTEuNS0uMi0xLjhhMSAxIDAgMCAwLS44LS40Yy0uMyAwLS43LjMtMS42IDFMMjcuNiA1MC4xYy0xLjIuOC0xLjcgMS4zLTIuMiAxLjhhNSA1IDAgMCAwLS44IDEuNmMtLjIuNy0uMiAxLjQtLjIgMi45djM3LjNjMCAxLjIgMCAxLjguMyAyIC4yLjMuNS41LjguNC40IDAgLjgtLjQgMS42LTEuM2wxOS0xOC42IDEuNS0xLjhjLjMtLjQuNS0xIC42LTEuNC4yLS42LjItMS4yLjItMi40VjM3LjVaTTM1LjIgMTA1LjNjLS44LjgtMS4yIDEuMy0xLjIgMS42IDAgLjQgMCAuNy4zLjkuMy4yLjkuMiAyIC4yaDM3YzEuMyAwIDIgMCAyLjUtLjJhNSA1IDAgMCAwIDEuNS0uNmMuNi0uNCAxLS45IDEuOS0xLjhMOTYuNiA4NmMuNy0uOSAxLjEtMS4zIDEuMS0xLjZhMSAxIDAgMCAwLS4zLS44Yy0uMy0uMy0uOS0uMy0yLS4zaC0zNWMtMS4yIDAtMS44IDAtMi40LjJhNSA1IDAgMCAwLTEuNC42Yy0uNS4zLTEgLjctMS44IDEuNmwtMTkuNiAxOS42Wk05Ni4zIDcwLjFjMSAwIDEuNCAwIDEuNy0uMi40LS4xLjYtLjQuOC0uN2wuMS0xLjdWMzUuM2wtLjEtMS44Yy0uMi0uMy0uNC0uNS0uOC0uNy0uMy0uMi0uOC0uMi0xLjctLjJINjQuMWMtMSAwLTEuNCAwLTEuNy4yLS40LjItLjYuNC0uOC43bC0uMSAxLjh2MzIuMmwuMSAxLjdjLjIuMy40LjYuOC43LjMuMi44LjIgMS43LjJoMzIuMloiLz48cGF0aCBmaWxsPSIjMTAxNTA5IiBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik00My4yIDIxLjVjLTEuMyAwLTIgMC0yLjMtLjMtLjMtLjMtLjUtLjYtLjUtMSAwLS41LjQtMSAxLjEtMi4xTDUzLjEgMi4zYy42LS44LjktMS4yIDEuMi0xLjNoMWMuNC4xLjcuNSAxLjIgMS4zbDExLjYgMTUuOGMuOCAxIDEuMiAxLjYgMS4yIDIgMCAuNS0uMi44LS41IDEtLjQuNC0xIC40LTIuNC40SDU5Yy0uMy4yLS41LjQtLjYuNy0uMi4zLS4yLjYtLjIgMS40VjY4YzAgMS45IDAgMi44LjQgMy41LjMuNi44IDEuMSAxLjQgMS41LjcuMyAxLjcuMyAzLjUuM2g0NGwxLjMtLjFjLjMtLjEuNS0uMy42LS42bC4xLTEuNFY2NWMwLTEuMyAwLTIgLjMtMi4zLjItLjMuNi0uNSAxLS41czEgLjMgMiAxTDEzMCA3NWMuOC42IDEuMy45IDEuNCAxLjJ2MWMtLjEuNC0uNi43LTEuNCAxLjNsLTE3LjMgMTEuN2MtMSAuOC0xLjYgMS4xLTIgMS4xLS40IDAtLjgtLjItMS0uNS0uMy0uNC0uMy0xLS4zLTIuM1Y4MmwtLjEtMS40Yy0uMS0uMi0uMy0uNC0uNi0uNS0uMy0uMi0uNi0uMi0xLjQtLjJINjAuNWMtMS42IDAtMi40IDAtMy4xLjItLjcuMi0xLjMuNC0yIC44bC0yLjMgMi0yOC42IDI4LjNjLS41LjUtLjguOC0uOSAxLjF2LjhjLjEuMy40LjYuOSAxLjFsMy43IDMuN2MuOCAxIDEuMyAxLjMgMS4zIDEuOCAwIC40IDAgLjctLjMgMS0uMy4zLS45LjUtMiAuOGwtMjAgNC43Yy0xLjEuMi0xLjcuMy0yIC4yLS40LS4xLS43LS40LS44LS44LS4xLS4zIDAtMSAuMy0ybDUtMTkuOWMuMy0xLjEuNC0xLjcuOC0yIC4zLS4zLjctLjQgMS0uMy41IDAgLjkuNSAxLjggMS40bDMuNiAzLjcgMSAuOWguOWMuMy0uMS42LS40IDEtLjlsMjguNi0yOC4yYzEuMi0xLjEgMS43LTEuNyAyLjEtMi4zLjQtLjYuNy0xLjMuOC0yIC4yLS43LjItMS41LjItMy4yVjIzLjZsLS4xLTEuNGMtLjEtLjMtLjMtLjUtLjYtLjZsLTEuNC0uMWgtNi4yWiIgY2xpcC1ydWxlPSJldmVub2RkIi8+PC9zdmc+ [deepwiki-link]: https://deepwiki.com/tensorchord/VectorChord [deepwiki-shield]: https://deepwiki.com/badge.svg [blog-link]: https://blog.vectorchord.ai/ From c7deb14f4c6852b4f217a60b7bc1f2c55f9b652f Mon Sep 17 00:00:00 2001 From: usamoi Date: Sat, 14 Feb 2026 18:37:29 +0800 Subject: [PATCH 277/324] readme: fix formatting (#429) Signed-off-by: usamoi --- README.md | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index f71df9f9..76d32a34 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,11 @@
-

VectorChord

-

Ready for the Billion-Scale Era. Host 100M vectors on a single i4i.xlarge ($247/mo) and (scale seamlessly to 1B+)[https://blog.vectorchord.ai/scaling-vector-search-to-1-billion-on-postgresql].

-
-
+# VectorChord +**Ready for the Billion-Scale Era. Host 100M vectors on a single i4i.xlarge ($247/mo) and [scale seamlessly to 1B+](https://blog.vectorchord.ai/scaling-vector-search-to-1-billion-on-postgresql).** [Official Site][official-site-link] · [Blog][blog-link] · [Docs][docs-link] · [Feedback][github-issues-link] · [Contact Us][email-link] - [![][github-release-shield]][github-release-link] [![][docker-release-shield]][docker-release-link] [![][docker-pulls-shield]][docker-pulls-link] @@ -19,6 +16,7 @@ [![][deepwiki-shield]][deepwiki-link] [![][license-1-shield]][license-1-link] [![][license-2-shield]][license-2-link] +
VectorChord (vchord) is a PostgreSQL extension engineered for scalable, high-performance, and cost-effective vector search. From a105cbcb45e2b01bf906c03663aef0294d63bb49 Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 24 Feb 2026 16:53:16 +0800 Subject: [PATCH 278/324] feat: dequantize (#431) closes #427 Signed-off-by: usamoi --- src/datatype/functions_rabitq4.rs | 34 ++++++++++++++++++++++++++++--- src/datatype/functions_rabitq8.rs | 34 ++++++++++++++++++++++++++++--- src/datatype/memory_halfvec.rs | 1 - src/datatype/memory_vector.rs | 1 - src/sql/finalize.sql | 12 +++++++++++ tests/general/dequantize.slt | 19 +++++++++++++++++ 6 files changed, 93 insertions(+), 8 deletions(-) create mode 100644 tests/general/dequantize.slt diff --git a/src/datatype/functions_rabitq4.rs b/src/datatype/functions_rabitq4.rs index 031f28b6..63b46556 100644 --- a/src/datatype/functions_rabitq4.rs +++ b/src/datatype/functions_rabitq4.rs @@ -12,12 +12,13 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::datatype::memory_halfvec::HalfvecInput; -use crate::datatype::memory_rabitq4::Rabitq4Output; -use crate::datatype::memory_vector::VectorInput; +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use simd::{Floating, f16}; use vector::VectorBorrowed; use vector::rabitq4::Rabitq4Borrowed; +use vector::vect::VectBorrowed; #[pgrx::pg_extern(sql = "")] fn _vchord_vector_quantize_to_rabitq4(vector: VectorInput) -> Rabitq4Output { @@ -54,3 +55,30 @@ fn _vchord_halfvec_quantize_to_rabitq4(vector: HalfvecInput) -> Rabitq4Output { &elements, )) } + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq4_dequantize_to_vector(vector: Rabitq4Input) -> VectorOutput { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + VectorOutput::new(VectBorrowed::new(&result)) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq4_dequantize_to_halfvec(vector: Rabitq4Input) -> HalfvecOutput { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + let result = f16::vector_from_f32(&result); + HalfvecOutput::new(VectBorrowed::new(&result)) +} diff --git a/src/datatype/functions_rabitq8.rs b/src/datatype/functions_rabitq8.rs index 1765216b..929e9c89 100644 --- a/src/datatype/functions_rabitq8.rs +++ b/src/datatype/functions_rabitq8.rs @@ -12,12 +12,13 @@ // // Copyright (c) 2025 TensorChord Inc. -use crate::datatype::memory_halfvec::HalfvecInput; -use crate::datatype::memory_rabitq8::Rabitq8Output; -use crate::datatype::memory_vector::VectorInput; +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use simd::{Floating, f16}; use vector::VectorBorrowed; use vector::rabitq8::Rabitq8Borrowed; +use vector::vect::VectBorrowed; #[pgrx::pg_extern(sql = "")] fn _vchord_vector_quantize_to_rabitq8(vector: VectorInput) -> Rabitq8Output { @@ -54,3 +55,30 @@ fn _vchord_halfvec_quantize_to_rabitq8(vector: HalfvecInput) -> Rabitq8Output { &elements, )) } + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq8_dequantize_to_vector(vector: Rabitq8Input) -> VectorOutput { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + VectorOutput::new(VectBorrowed::new(&result)) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq8_dequantize_to_halfvec(vector: Rabitq8Input) -> HalfvecOutput { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + let result = f16::vector_from_f32(&result); + HalfvecOutput::new(VectBorrowed::new(&result)) +} diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs index 95e5f6cd..e0b1dd81 100644 --- a/src/datatype/memory_halfvec.rs +++ b/src/datatype/memory_halfvec.rs @@ -102,7 +102,6 @@ impl HalfvecOutput { } Self(q) } - #[expect(dead_code)] pub fn new(vector: VectBorrowed<'_, f16>) -> Self { unsafe { let slice = vector.slice(); diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs index bb5cdd92..d982fdd4 100644 --- a/src/datatype/memory_vector.rs +++ b/src/datatype/memory_vector.rs @@ -101,7 +101,6 @@ impl VectorOutput { } Self(q) } - #[expect(dead_code)] pub fn new(vector: VectBorrowed<'_, f32>) -> Self { unsafe { let slice = vector.slice(); diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index c344e102..39c0493a 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -206,12 +206,24 @@ IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_ CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; +CREATE FUNCTION dequantize_to_vector(rabitq8) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq8) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_halfvec_wrapper'; + CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; +CREATE FUNCTION dequantize_to_vector(rabitq4) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq4) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_halfvec_wrapper'; + CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; diff --git a/tests/general/dequantize.slt b/tests/general/dequantize.slt new file mode 100644 index 00000000..adf30c39 --- /dev/null +++ b/tests/general/dequantize.slt @@ -0,0 +1,19 @@ +query I +SELECT dequantize_to_vector(quantize_to_rabitq8('[1,2,3,4,5,6,7,8]'::vector)) <-> '[1,2,3,4,5,6,7,8]'::vector < 0.07; +---- +t + +query I +SELECT dequantize_to_halfvec(quantize_to_rabitq8('[1,2,3,4,5,6,7,8]'::halfvec)) <-> '[1,2,3,4,5,6,7,8]'::halfvec < 0.07; +---- +t + +query I +SELECT dequantize_to_vector(quantize_to_rabitq4('[1,2,3,4,5,6,7,8]'::vector)) <-> '[1,2,3,4,5,6,7,8]'::vector < 1.00; +---- +t + +query I +SELECT dequantize_to_halfvec(quantize_to_rabitq4('[1,2,3,4,5,6,7,8]'::halfvec)) <-> '[1,2,3,4,5,6,7,8]'::halfvec < 1.00; +---- +t From 097ec1953a94d33e002caf6ba160702a43b27c0d Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 27 Feb 2026 20:05:30 +0800 Subject: [PATCH 279/324] perf: disable SAQ for rabitq8 (#434) Signed-off-by: usamoi --- crates/rabitq/src/extended.rs | 52 +------- tests/vchordg/index_vector_rabitq.slt | 125 +++++++++++++------- tests/vchordrq/index_multivector_rabitq.slt | 8 +- tests/vchordrq/index_vector_rabitq.slt | 113 ++++++++++++------ 4 files changed, 171 insertions(+), 127 deletions(-) diff --git a/crates/rabitq/src/extended.rs b/crates/rabitq/src/extended.rs index 0e8c52a9..1c1759ce 100644 --- a/crates/rabitq/src/extended.rs +++ b/crates/rabitq/src/extended.rs @@ -112,7 +112,7 @@ pub fn ugly_code(vector: &[f32]) -> Code { f32::vector_mul_scalar_inplace(&mut vector, 1.0 / dis_u_2.sqrt()); vector }; - let (scale, delta) = { + let scale = { let mut o = normalized_vector.clone(); f32::vector_abs_inplace(&mut o); ugly_find_scale::(&o) @@ -120,7 +120,6 @@ pub fn ugly_code(vector: &[f32]) -> Code { let mut code = Vec::with_capacity(n as _); for i in 0..n { let v = scale * normalized_vector[i]; - let v = v + v.signum() * delta[i] as f32; let c = v.floor().clamp(min as f32, max as f32) as i32; code.push((c + (1 << (BITS - 1))) as _); } @@ -251,55 +250,10 @@ fn find_scale(o: &[f32]) -> f32 { x_m as f32 + f32::EPSILON } -fn ugly_find_scale(o: &[f32]) -> (f32, Vec) { +fn ugly_find_scale(o: &[f32]) -> f32 { assert!((1..=8).contains(&B)); - let dim = o.len(); - - let mut code = Vec::::with_capacity(dim); - let mut numerator_m = 0.0f64; - let mut sqr_denominator_m = 0.0f64; - - let scale = (1 << (B - 1)) as f32 / f32::reduce_min_max_of_x(o).1; - for i in 0..dim { - code.push((o[i] as f64 * scale as f64) as u8); - let value = code[i] as f64 + 0.5; - numerator_m += value * o[i] as f64; - sqr_denominator_m += value * value; - } - let mut y_m = numerator_m / sqr_denominator_m.sqrt(); - - let mut delta = vec![0_i32; dim]; - for _ in 0..8 { - for i in 0..dim { - if code[i] < (1 << (B - 1)) - 1 { - let numerator = numerator_m + o[i] as f64; - let sqr_denominator = sqr_denominator_m + 2.0 * (code[i] as f64 + 1.0); - let y = numerator / sqr_denominator.sqrt(); - if y > y_m { - y_m = y; - numerator_m = numerator; - sqr_denominator_m = sqr_denominator; - code[i] += 1; - delta[i] += 1; - } - } - if code[i] > 0 { - let numerator = numerator_m - o[i] as f64; - let sqr_denominator = sqr_denominator_m - 2.0 * code[i] as f64; - let y = numerator / sqr_denominator.sqrt(); - if y > y_m { - y_m = y; - numerator_m = numerator; - sqr_denominator_m = sqr_denominator; - code[i] -= 1; - delta[i] -= 1; - } - } - } - } - - (scale, delta) + (1 << (B - 1)) as f32 / f32::reduce_min_max_of_x(o).1 } pub fn pack_code(input: &[u8]) -> [Vec; BITS] { diff --git a/tests/vchordg/index_vector_rabitq.slt b/tests/vchordg/index_vector_rabitq.slt index 50334194..5310be9a 100644 --- a/tests/vchordg/index_vector_rabitq.slt +++ b/tests/vchordg/index_vector_rabitq.slt @@ -20,9 +20,8 @@ CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabi query I SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -155 1608 -1643 +155 174 818 1603 @@ -30,6 +29,7 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(ar 60 218 1080 +1568 statement ok DROP INDEX ti; @@ -40,9 +40,8 @@ CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabi query I SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -155 1608 -1643 +155 174 818 1603 @@ -50,6 +49,7 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(ar 60 218 1080 +1639 statement ok DROP INDEX ti; @@ -60,9 +60,8 @@ CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabi query I SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -155 1608 -1643 +155 174 818 1603 @@ -70,6 +69,7 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(ar 60 218 1080 +1568 statement ok DROP INDEX ti; @@ -78,17 +78,18 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -155 1608 -1643 +155 174 -818 1603 +818 1629 60 218 +1080 +79 statement ok DROP INDEX ti; @@ -97,17 +98,18 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -155 1608 -1643 +155 174 818 1603 1629 60 218 +1080 +79 statement ok DROP INDEX ti; @@ -116,17 +118,18 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -155 1608 -1643 +155 174 818 1603 1629 60 218 +1080 +79 statement ok DROP INDEX ti; @@ -135,11 +138,18 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_l2_ops); query I -SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -1 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; @@ -148,11 +158,18 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_ip_ops); query I -SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -1 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; @@ -161,11 +178,18 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_cosine_ops); query I -SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -1 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; @@ -174,11 +198,18 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_l2_ops); query I -SELECT SUM(index) FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -7290 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1839 +1568 statement ok DROP INDEX ti; @@ -187,11 +218,18 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_ip_ops); query I -SELECT SUM(index) FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -7290 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; @@ -200,11 +238,18 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_cosine_ops); query I -SELECT SUM(index) FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -7290 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; diff --git a/tests/vchordrq/index_multivector_rabitq.slt b/tests/vchordrq/index_multivector_rabitq.slt index 7756da41..0b7d824b 100644 --- a/tests/vchordrq/index_multivector_rabitq.slt +++ b/tests/vchordrq/index_multivector_rabitq.slt @@ -74,13 +74,13 @@ LIMIT 10; 1207 919 1821 -1076 1639 +537 174 -1125 79 -1163 -537 +194 +1920 +239 statement ok DROP INDEX pg_temp.ti; diff --git a/tests/vchordrq/index_vector_rabitq.slt b/tests/vchordrq/index_vector_rabitq.slt index 29841b14..33e7bc16 100644 --- a/tests/vchordrq/index_vector_rabitq.slt +++ b/tests/vchordrq/index_vector_rabitq.slt @@ -20,8 +20,8 @@ CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rab query I SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -155 1608 +155 1643 174 818 @@ -40,8 +40,8 @@ CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rab query I SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -155 1608 +155 1643 174 818 @@ -60,8 +60,8 @@ CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rab query I SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -155 1608 +155 1643 174 818 @@ -78,17 +78,18 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -155 1608 +155 1643 174 -818 1603 +818 1629 60 218 +1080 statement ok DROP INDEX ti; @@ -97,10 +98,10 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -155 1608 +155 1643 174 818 @@ -108,6 +109,7 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_r 1629 60 218 +1080 statement ok DROP INDEX ti; @@ -116,10 +118,10 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -155 1608 +155 1643 174 818 @@ -127,6 +129,7 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_r 1629 60 218 +1080 statement ok DROP INDEX ti; @@ -135,11 +138,18 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_l2_ops); query I -SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -1 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; @@ -148,11 +158,18 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_ip_ops); query I -SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -1 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; @@ -161,11 +178,18 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_cosine_ops); query I -SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; ---- -1 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; @@ -174,11 +198,18 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_l2_ops); query I -SELECT SUM(index) FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -7290 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1839 +1568 statement ok DROP INDEX ti; @@ -187,11 +218,18 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_ip_ops); query I -SELECT SUM(index) FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -7290 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; @@ -200,11 +238,18 @@ statement ok CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_cosine_ops); query I -SELECT SUM(index) FROM ( -SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10 -) AS s(index); +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; ---- -7290 +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 statement ok DROP INDEX ti; From 7a5aee581961d16d6d66bfa0ab38d65ef9c42e8f Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 27 Feb 2026 20:05:52 +0800 Subject: [PATCH 280/324] chore: upload 1.1.1 schema scripts (#433) Signed-off-by: usamoi --- sql/install/vchord--1.1.1.sql | 1323 ++++++++++++++++++++++++++ sql/upgrade/vchord--1.1.0--1.1.1.sql | 13 + 2 files changed, 1336 insertions(+) create mode 100644 sql/install/vchord--1.1.1.sql create mode 100644 sql/upgrade/vchord--1.1.0--1.1.1.sql diff --git a/sql/install/vchord--1.1.1.sql b/sql/install/vchord--1.1.1.sql new file mode 100644 index 00000000..c5ecd748 --- /dev/null +++ b/sql/install/vchord--1.1.1.sql @@ -0,0 +1,1323 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:47 +-- bootstrap +-- List of shell types + +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE rabitq8; +CREATE TYPE sphere_rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq4; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ + "rhs" halfvec[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:41 +-- vchord::datatype::functions_rabitq4::_vchord_halfvec_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:41 +-- vchord::datatype::functions_rabitq8::_vchord_halfvec_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:72 +-- vchord::datatype::functions_rabitq4::_vchord_rabitq4_dequantize_to_halfvec +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:59 +-- vchord::datatype::functions_rabitq4::_vchord_rabitq4_dequantize_to_vector +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:20 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_in +CREATE FUNCTION "_vchord_rabitq4_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:41 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_cosine +CREATE FUNCTION "_vchord_rabitq4_operator_cosine"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:31 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_ip +CREATE FUNCTION "_vchord_rabitq4_operator_ip"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:21 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_l2 +CREATE FUNCTION "_vchord_rabitq4_operator_l2"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:123 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:140 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_out +CREATE FUNCTION "_vchord_rabitq4_out"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:36 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_recv +CREATE FUNCTION "_vchord_rabitq4_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:21 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_send +CREATE FUNCTION "_vchord_rabitq4_send"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:99 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq4_sphere_cosine_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:75 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq4_sphere_ip_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:51 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq4_sphere_l2_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq4.rs:17 +-- vchord::datatype::typmod_rabitq4::_vchord_rabitq4_typmod_in +CREATE FUNCTION "_vchord_rabitq4_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:72 +-- vchord::datatype::functions_rabitq8::_vchord_rabitq8_dequantize_to_halfvec +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:59 +-- vchord::datatype::functions_rabitq8::_vchord_rabitq8_dequantize_to_vector +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:20 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_in +CREATE FUNCTION "_vchord_rabitq8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:41 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_cosine +CREATE FUNCTION "_vchord_rabitq8_operator_cosine"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:31 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_ip +CREATE FUNCTION "_vchord_rabitq8_operator_ip"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:21 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_l2 +CREATE FUNCTION "_vchord_rabitq8_operator_l2"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:123 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:140 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_out +CREATE FUNCTION "_vchord_rabitq8_out"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:36 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_recv +CREATE FUNCTION "_vchord_rabitq8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:21 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_send +CREATE FUNCTION "_vchord_rabitq8_send"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:99 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq8_sphere_cosine_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:75 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq8_sphere_ip_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:51 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq8_sphere_l2_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq8.rs:17 +-- vchord::datatype::typmod_rabitq8::_vchord_rabitq8_typmod_in +CREATE FUNCTION "_vchord_rabitq8_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ + "rhs" vector[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:23 +-- vchord::datatype::functions_rabitq4::_vchord_vector_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:23 +-- vchord::datatype::functions_rabitq8::_vchord_vector_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:110 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordg_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordg_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordg_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordg_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordg_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordg_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:192 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:100 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:95 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:90 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:140 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:130 +-- vchord::index::opclass::_vchordrq_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:125 +-- vchord::index::opclass::_vchordrq_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:120 +-- vchord::index::opclass::_vchordrq_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:150 +-- vchord::index::opclass::_vchordrq_support_rabitq4_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:115 +-- vchord::index::opclass::_vchordrq_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:110 +-- vchord::index::opclass::_vchordrq_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:105 +-- vchord::index::opclass::_vchordrq_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:145 +-- vchord::index::opclass::_vchordrq_support_rabitq8_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:85 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:135 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:48 +-- finalize +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external +); + +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, + radius REAL +); + +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq8) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq8) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq4) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq4) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/upgrade/vchord--1.1.0--1.1.1.sql b/sql/upgrade/vchord--1.1.0--1.1.1.sql new file mode 100644 index 00000000..6490df1e --- /dev/null +++ b/sql/upgrade/vchord--1.1.0--1.1.1.sql @@ -0,0 +1,13 @@ +-- List of functions + +CREATE FUNCTION dequantize_to_vector(rabitq8) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq8) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq4) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq4) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_halfvec_wrapper'; From 922ae4713441ecb584e6c460af3ec247cc0dfa2a Mon Sep 17 00:00:00 2001 From: usamoi Date: Sat, 28 Feb 2026 10:49:04 +0800 Subject: [PATCH 281/324] fix: tests (#437) Signed-off-by: usamoi --- .github/workflows/check.yml | 3 ++- tests/vchordg/index_vector_rabitq.slt | 6 ++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 894bcf97..ca525f07 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -1009,7 +1009,8 @@ jobs: sudo -iu postgres createuser -s -r $(whoami) sudo -iu postgres createdb -O $(whoami) $(whoami) if [ "${{ matrix.version }}" -ge "18" ]; then - sudo -iu postgres psql -c 'ALTER SYSTEM SET io_method = io_uring' + # valgrind works not well + sudo -iu postgres psql -c 'ALTER SYSTEM SET io_method = worker' fi sudo -iu postgres psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' diff --git a/tests/vchordg/index_vector_rabitq.slt b/tests/vchordg/index_vector_rabitq.slt index 5310be9a..8837d9fd 100644 --- a/tests/vchordg/index_vector_rabitq.slt +++ b/tests/vchordg/index_vector_rabitq.slt @@ -18,7 +18,7 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_l2_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 9; ---- 1608 155 @@ -29,7 +29,6 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(ar 60 218 1080 -1568 statement ok DROP INDEX ti; @@ -58,7 +57,7 @@ statement ok CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_cosine_ops); query I -SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 9; ---- 1608 155 @@ -69,7 +68,6 @@ SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(ar 60 218 1080 -1568 statement ok DROP INDEX ti; From 80843b2218369f54a943a6f9d3bfaabce19144cd Mon Sep 17 00:00:00 2001 From: usamoi Date: Sat, 28 Feb 2026 12:10:55 +0800 Subject: [PATCH 282/324] readme: sync with docs (#435) https://github.com/tensorchord/pgvecto.rs-docs/pull/168 Signed-off-by: usamoi --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 76d32a34..e305937c 100644 --- a/README.md +++ b/README.md @@ -57,7 +57,7 @@ docker run \ --name vectorchord-demo \ -e POSTGRES_PASSWORD=mysecretpassword \ -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg18-v1.1.0 + -d ghcr.io/tensorchord/vchord-postgres:pg18-v1.1.1 ``` > [!NOTE] From 0e456b633f8e31308f8d88c13fbad7bf3eaa7cd5 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 5 Mar 2026 15:50:54 +0800 Subject: [PATCH 283/324] build(deps): bump actions/upload-artifact from 6.0.0 to 7.0.0 (#438) --- .github/workflows/check.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index ca525f07..b9858fc2 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -296,7 +296,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 + uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu path: ./build @@ -394,7 +394,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 + uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin path: ./build @@ -514,7 +514,7 @@ jobs: run: make install - name: Upload Artifacts - uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 + uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc path: ./build @@ -645,7 +645,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 + uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-musl path: ./build @@ -862,7 +862,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 + uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.clang_triple }} path: ./build From a26561da25dd343f030b6b6fd9fef29fc05b1577 Mon Sep 17 00:00:00 2001 From: usamoi Date: Fri, 6 Mar 2026 08:59:08 +0800 Subject: [PATCH 284/324] feat: more f16 simd routines (#398) Signed-off-by: usamoi --- crates/simd/build.rs | 7 - crates/simd/cshim/aarch64_fp16.c | 250 ------ crates/simd/cshim/x86_64_fp16.c | 141 --- crates/simd/src/emulate.rs | 18 + crates/simd/src/floating_f16.rs | 1381 +++++++++++++++++++++++++++++- 5 files changed, 1357 insertions(+), 440 deletions(-) delete mode 100644 crates/simd/cshim/aarch64_fp16.c delete mode 100644 crates/simd/cshim/x86_64_fp16.c diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 63af110d..338f9e17 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -23,7 +23,6 @@ fn main() -> Result<(), Box> { match target_arch.as_str() { "aarch64" => { let mut build = cc::Build::new(); - build.file("./cshim/aarch64_fp16.c"); build.file("./cshim/aarch64_byte.c"); build.file("./cshim/aarch64_halfbyte.c"); if target_endian == "little" { @@ -33,12 +32,6 @@ fn main() -> Result<(), Box> { build.opt_level(3); build.compile("simd_cshim"); } - "x86_64" => { - let mut build = cc::Build::new(); - build.file("./cshim/x86_64_fp16.c"); - build.opt_level(3); - build.compile("simd_cshim"); - } _ => {} } Ok(()) diff --git a/crates/simd/cshim/aarch64_fp16.c b/crates/simd/cshim/aarch64_fp16.c deleted file mode 100644 index ad7ffb93..00000000 --- a/crates/simd/cshim/aarch64_fp16.c +++ /dev/null @@ -1,250 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -#if defined(__clang__) -#if !(__clang_major__ >= 16) -#error "Clang version must be at least 16." -#endif -#elif defined(__GNUC__) -#if !(__GNUC__ >= 14) -#error "GCC version must be at least 14." -#endif -#else -#error "This file requires Clang or GCC." -#endif - -#include -#include -#include - -typedef __fp16 f16; -typedef float f32; - -__attribute__((target("+fp16"))) float -fp16_reduce_sum_of_xy_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { - float16x8_t sum_0 = vdupq_n_f16(0.0); - float16x8_t sum_1 = vdupq_n_f16(0.0); - float16x8_t sum_2 = vdupq_n_f16(0.0); - float16x8_t sum_3 = vdupq_n_f16(0.0); - float16x8_t sum_4 = vdupq_n_f16(0.0); - float16x8_t sum_5 = vdupq_n_f16(0.0); - float16x8_t sum_6 = vdupq_n_f16(0.0); - float16x8_t sum_7 = vdupq_n_f16(0.0); - while (n >= 64) { - float16x8_t x_0 = vld1q_f16(a + 0); - float16x8_t x_1 = vld1q_f16(a + 8); - float16x8_t x_2 = vld1q_f16(a + 16); - float16x8_t x_3 = vld1q_f16(a + 24); - float16x8_t x_4 = vld1q_f16(a + 32); - float16x8_t x_5 = vld1q_f16(a + 40); - float16x8_t x_6 = vld1q_f16(a + 48); - float16x8_t x_7 = vld1q_f16(a + 56); - float16x8_t y_0 = vld1q_f16(b + 0); - float16x8_t y_1 = vld1q_f16(b + 8); - float16x8_t y_2 = vld1q_f16(b + 16); - float16x8_t y_3 = vld1q_f16(b + 24); - float16x8_t y_4 = vld1q_f16(b + 32); - float16x8_t y_5 = vld1q_f16(b + 40); - float16x8_t y_6 = vld1q_f16(b + 48); - float16x8_t y_7 = vld1q_f16(b + 56); - sum_0 = vfmaq_f16(sum_0, x_0, y_0); - sum_1 = vfmaq_f16(sum_1, x_1, y_1); - sum_2 = vfmaq_f16(sum_2, x_2, y_2); - sum_3 = vfmaq_f16(sum_3, x_3, y_3); - sum_4 = vfmaq_f16(sum_4, x_4, y_4); - sum_5 = vfmaq_f16(sum_5, x_5, y_5); - sum_6 = vfmaq_f16(sum_6, x_6, y_6); - sum_7 = vfmaq_f16(sum_7, x_7, y_7); - n -= 64, a += 64, b += 64; - } - if (n >= 32) { - float16x8_t x_0 = vld1q_f16(a + 0); - float16x8_t x_1 = vld1q_f16(a + 8); - float16x8_t x_2 = vld1q_f16(a + 16); - float16x8_t x_3 = vld1q_f16(a + 24); - float16x8_t y_0 = vld1q_f16(b + 0); - float16x8_t y_1 = vld1q_f16(b + 8); - float16x8_t y_2 = vld1q_f16(b + 16); - float16x8_t y_3 = vld1q_f16(b + 24); - sum_0 = vfmaq_f16(sum_0, x_0, y_0); - sum_1 = vfmaq_f16(sum_1, x_1, y_1); - sum_2 = vfmaq_f16(sum_2, x_2, y_2); - sum_3 = vfmaq_f16(sum_3, x_3, y_3); - n -= 32, a += 32, b += 32; - } - if (n >= 16) { - float16x8_t x_4 = vld1q_f16(a + 0); - float16x8_t x_5 = vld1q_f16(a + 8); - float16x8_t y_4 = vld1q_f16(b + 0); - float16x8_t y_5 = vld1q_f16(b + 8); - sum_4 = vfmaq_f16(sum_4, x_4, y_4); - sum_5 = vfmaq_f16(sum_5, x_5, y_5); - n -= 16, a += 16, b += 16; - } - if (n >= 8) { - float16x8_t x_6 = vld1q_f16(a + 0); - float16x8_t y_6 = vld1q_f16(b + 0); - sum_6 = vfmaq_f16(sum_6, x_6, y_6); - n -= 8, a += 8, b += 8; - } - if (n > 0) { - f16 _a[8] = {}, _b[8] = {}; - for (size_t i = 0; i < n; i += 1) { - _a[i] = a[i], _b[i] = b[i]; - } - a = _a, b = _b; - float16x8_t x_7 = vld1q_f16(a); - float16x8_t y_7 = vld1q_f16(b); - sum_7 = vfmaq_f16(sum_7, x_7, y_7); - } - float32x4_t s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); - float32x4_t s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); - float32x4_t s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); - float32x4_t s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); - float32x4_t s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); - float32x4_t s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); - float32x4_t s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); - float32x4_t s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); - float32x4_t s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); - float32x4_t s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); - float32x4_t s_a = vcvt_f32_f16(vget_low_f16(sum_5)); - float32x4_t s_b = vcvt_f32_f16(vget_high_f16(sum_5)); - float32x4_t s_c = vcvt_f32_f16(vget_low_f16(sum_6)); - float32x4_t s_d = vcvt_f32_f16(vget_high_f16(sum_6)); - float32x4_t s_e = vcvt_f32_f16(vget_low_f16(sum_7)); - float32x4_t s_f = vcvt_f32_f16(vget_high_f16(sum_7)); - float32x4_t s = vpaddq_f32( - vpaddq_f32(vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), - vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7))), - vpaddq_f32(vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), - vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)))); - return vaddvq_f32(s); -} - -__attribute__((target("+fp16"))) float -fp16_reduce_sum_of_d2_a2_fp16(size_t n, f16 *restrict a, f16 *restrict b) { - float16x8_t sum_0 = vdupq_n_f16(0.0); - float16x8_t sum_1 = vdupq_n_f16(0.0); - float16x8_t sum_2 = vdupq_n_f16(0.0); - float16x8_t sum_3 = vdupq_n_f16(0.0); - float16x8_t sum_4 = vdupq_n_f16(0.0); - float16x8_t sum_5 = vdupq_n_f16(0.0); - float16x8_t sum_6 = vdupq_n_f16(0.0); - float16x8_t sum_7 = vdupq_n_f16(0.0); - while (n >= 64) { - float16x8_t x_0 = vld1q_f16(a + 0); - float16x8_t x_1 = vld1q_f16(a + 8); - float16x8_t x_2 = vld1q_f16(a + 16); - float16x8_t x_3 = vld1q_f16(a + 24); - float16x8_t x_4 = vld1q_f16(a + 32); - float16x8_t x_5 = vld1q_f16(a + 40); - float16x8_t x_6 = vld1q_f16(a + 48); - float16x8_t x_7 = vld1q_f16(a + 56); - float16x8_t y_0 = vld1q_f16(b + 0); - float16x8_t y_1 = vld1q_f16(b + 8); - float16x8_t y_2 = vld1q_f16(b + 16); - float16x8_t y_3 = vld1q_f16(b + 24); - float16x8_t y_4 = vld1q_f16(b + 32); - float16x8_t y_5 = vld1q_f16(b + 40); - float16x8_t y_6 = vld1q_f16(b + 48); - float16x8_t y_7 = vld1q_f16(b + 56); - float16x8_t d_0 = vsubq_f16(x_0, y_0); - float16x8_t d_1 = vsubq_f16(x_1, y_1); - float16x8_t d_2 = vsubq_f16(x_2, y_2); - float16x8_t d_3 = vsubq_f16(x_3, y_3); - float16x8_t d_4 = vsubq_f16(x_4, y_4); - float16x8_t d_5 = vsubq_f16(x_5, y_5); - float16x8_t d_6 = vsubq_f16(x_6, y_6); - float16x8_t d_7 = vsubq_f16(x_7, y_7); - sum_0 = vfmaq_f16(sum_0, d_0, d_0); - sum_1 = vfmaq_f16(sum_1, d_1, d_1); - sum_2 = vfmaq_f16(sum_2, d_2, d_2); - sum_3 = vfmaq_f16(sum_3, d_3, d_3); - sum_4 = vfmaq_f16(sum_4, d_4, d_4); - sum_5 = vfmaq_f16(sum_5, d_5, d_5); - sum_6 = vfmaq_f16(sum_6, d_6, d_6); - sum_7 = vfmaq_f16(sum_7, d_7, d_7); - n -= 64, a += 64, b += 64; - } - if (n >= 32) { - float16x8_t x_0 = vld1q_f16(a + 0); - float16x8_t x_1 = vld1q_f16(a + 8); - float16x8_t x_2 = vld1q_f16(a + 16); - float16x8_t x_3 = vld1q_f16(a + 24); - float16x8_t y_0 = vld1q_f16(b + 0); - float16x8_t y_1 = vld1q_f16(b + 8); - float16x8_t y_2 = vld1q_f16(b + 16); - float16x8_t y_3 = vld1q_f16(b + 24); - float16x8_t d_0 = vsubq_f16(x_0, y_0); - float16x8_t d_1 = vsubq_f16(x_1, y_1); - float16x8_t d_2 = vsubq_f16(x_2, y_2); - float16x8_t d_3 = vsubq_f16(x_3, y_3); - sum_0 = vfmaq_f16(sum_0, d_0, d_0); - sum_1 = vfmaq_f16(sum_1, d_1, d_1); - sum_2 = vfmaq_f16(sum_2, d_2, d_2); - sum_3 = vfmaq_f16(sum_3, d_3, d_3); - n -= 32, a += 32, b += 32; - } - if (n >= 16) { - float16x8_t x_4 = vld1q_f16(a + 0); - float16x8_t x_5 = vld1q_f16(a + 8); - float16x8_t y_4 = vld1q_f16(b + 0); - float16x8_t y_5 = vld1q_f16(b + 8); - float16x8_t d_4 = vsubq_f16(x_4, y_4); - float16x8_t d_5 = vsubq_f16(x_5, y_5); - sum_4 = vfmaq_f16(sum_4, d_4, d_4); - sum_5 = vfmaq_f16(sum_5, d_5, d_5); - n -= 16, a += 16, b += 16; - } - if (n >= 8) { - float16x8_t x_6 = vld1q_f16(a + 0); - float16x8_t y_6 = vld1q_f16(b + 0); - float16x8_t d_6 = vsubq_f16(x_6, y_6); - sum_6 = vfmaq_f16(sum_6, d_6, d_6); - n -= 8, a += 8, b += 8; - } - if (n > 0) { - f16 _a[8] = {}, _b[8] = {}; - for (size_t i = 0; i < n; i += 1) { - _a[i] = a[i], _b[i] = b[i]; - } - a = _a, b = _b; - float16x8_t x_7 = vld1q_f16(a); - float16x8_t y_7 = vld1q_f16(b); - float16x8_t d_7 = vsubq_f16(x_7, y_7); - sum_7 = vfmaq_f16(sum_7, d_7, d_7); - } - float32x4_t s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); - float32x4_t s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); - float32x4_t s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); - float32x4_t s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); - float32x4_t s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); - float32x4_t s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); - float32x4_t s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); - float32x4_t s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); - float32x4_t s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); - float32x4_t s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); - float32x4_t s_a = vcvt_f32_f16(vget_low_f16(sum_5)); - float32x4_t s_b = vcvt_f32_f16(vget_high_f16(sum_5)); - float32x4_t s_c = vcvt_f32_f16(vget_low_f16(sum_6)); - float32x4_t s_d = vcvt_f32_f16(vget_high_f16(sum_6)); - float32x4_t s_e = vcvt_f32_f16(vget_low_f16(sum_7)); - float32x4_t s_f = vcvt_f32_f16(vget_high_f16(sum_7)); - float32x4_t s = vpaddq_f32( - vpaddq_f32(vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), - vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7))), - vpaddq_f32(vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), - vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)))); - return vaddvq_f32(s); -} diff --git a/crates/simd/cshim/x86_64_fp16.c b/crates/simd/cshim/x86_64_fp16.c deleted file mode 100644 index 79a204b5..00000000 --- a/crates/simd/cshim/x86_64_fp16.c +++ /dev/null @@ -1,141 +0,0 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025 TensorChord Inc. - -#if defined(__clang__) -#if !(__clang_major__ >= 16) -#error "Clang version must be at least 16." -#endif -#elif defined(__GNUC__) -#if !(__GNUC__ >= 12) -#error "GCC version must be at least 12." -#endif -#else -#error "This file requires Clang or GCC." -#endif - -#include -#include -#include - -#if defined(__clang__) && defined(_MSC_VER) && (__clang_major__ <= 19) -// https://github.com/llvm/llvm-project/issues/53520 -// clang-format off -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -// clang-format on -#endif - -typedef _Float16 f16; -typedef float f32; - -__attribute__((target("avx512bw,avx512cd,avx512dq,avx512vl,bmi,bmi2,lzcnt," - "movbe,popcnt,avx512fp16"))) float -fp16_reduce_sum_of_xy_v4_avx512fp16(size_t n, f16 *restrict a, - f16 *restrict b) { - __m512h _0 = _mm512_setzero_ph(); - __m512h _1 = _mm512_setzero_ph(); - while (n >= 64) { - __m512h x_0 = _mm512_loadu_ph(a + 0); - __m512h x_1 = _mm512_loadu_ph(a + 32); - __m512h y_0 = _mm512_loadu_ph(b + 0); - __m512h y_1 = _mm512_loadu_ph(b + 32); - _0 = _mm512_fmadd_ph(x_0, y_0, _0); - _1 = _mm512_fmadd_ph(x_1, y_1, _1); - n -= 64, a += 64, b += 64; - } - while (n >= 32) { - __m512h x_0 = _mm512_loadu_ph(a + 0); - __m512h y_0 = _mm512_loadu_ph(b + 0); - _0 = _mm512_fmadd_ph(x_0, y_0, _0); - n -= 32, a += 32, b += 32; - } - if (n > 0) { - unsigned int mask = _bzhi_u32(0xffffffff, n); - __m512h x = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a)); - __m512h y = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b)); - _1 = _mm512_fmadd_ph(x, y, _1); - } - __m512 s_0 = - _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); - __m512 s_1 = - _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); - __m512 s_2 = - _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); - __m512 s_3 = - _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); - return _mm512_reduce_add_ps( - _mm512_add_ps(_mm512_add_ps(s_0, s_2), _mm512_add_ps(s_1, s_3))); -} - -__attribute__((target("avx512bw,avx512cd,avx512dq,avx512vl,bmi,bmi2,lzcnt," - "movbe,popcnt,avx512fp16"))) float -fp16_reduce_sum_of_d2_v4_avx512fp16(size_t n, f16 *restrict a, - f16 *restrict b) { - __m512h _0 = _mm512_setzero_ph(); - __m512h _1 = _mm512_setzero_ph(); - while (n >= 64) { - __m512h x_0 = _mm512_loadu_ph(a + 0); - __m512h x_1 = _mm512_loadu_ph(a + 32); - __m512h y_0 = _mm512_loadu_ph(b + 0); - __m512h y_1 = _mm512_loadu_ph(b + 32); - __m512h d_0 = _mm512_sub_ph(x_0, y_0); - __m512h d_1 = _mm512_sub_ph(x_1, y_1); - _0 = _mm512_fmadd_ph(d_0, d_0, _0); - _1 = _mm512_fmadd_ph(d_1, d_1, _1); - n -= 64, a += 64, b += 64; - } - while (n >= 32) { - __m512h x_0 = _mm512_loadu_ph(a + 0); - __m512h y_0 = _mm512_loadu_ph(b + 0); - __m512h d_0 = _mm512_sub_ph(x_0, y_0); - _0 = _mm512_fmadd_ph(d_0, d_0, _0); - n -= 32, a += 32, b += 32; - } - if (n > 0) { - unsigned int mask = _bzhi_u32(0xffffffff, n); - __m512h x_1 = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a)); - __m512h y_1 = _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b)); - __m512h d_1 = _mm512_sub_ph(x_1, y_1); - _1 = _mm512_fmadd_ph(d_1, d_1, _1); - } - __m512 s_0 = - _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); - __m512 s_1 = - _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); - __m512 s_2 = - _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); - __m512 s_3 = - _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); - return _mm512_reduce_add_ps( - _mm512_add_ps(_mm512_add_ps(s_0, s_2), _mm512_add_ps(s_1, s_3))); -} diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index b84e5e8f..e710ab6a 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -209,3 +209,21 @@ pub fn emulate_mm256_reduce_add_epi64(mut x: core::arch::x86_64::__m256i) -> i64 x = _mm256_add_epi64(x, _mm256_permute2f128_si256(x, x, 1)); _mm256_extract_epi64(x, 0) + _mm256_extract_epi64(x, 1) } + +#[inline] +#[cfg(target_arch = "aarch64")] +#[target_feature(enable = "neon")] +pub fn emulate_vreinterpret_f16_u16( + x: core::arch::aarch64::uint16x4_t, +) -> core::arch::aarch64::float16x4_t { + unsafe { core::mem::transmute(x) } +} + +#[inline] +#[cfg(target_arch = "aarch64")] +#[target_feature(enable = "neon")] +pub fn emulate_vreinterpretq_f16_u16( + x: core::arch::aarch64::uint16x8_t, +) -> core::arch::aarch64::float16x8_t { + unsafe { core::mem::transmute(x) } +} diff --git a/crates/simd/src/floating_f16.rs b/crates/simd/src/floating_f16.rs index 8a9f7131..addc17dc 100644 --- a/crates/simd/src/floating_f16.rs +++ b/crates/simd/src/floating_f16.rs @@ -165,6 +165,71 @@ mod reduce_or_of_is_zero_x { mod reduce_sum_of_x { use crate::{F16, f16}; + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_sum_of_x_v4_avx512fp16(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + _0 = _mm512_add_ph(_0, x_0); + _1 = _mm512_add_ph(_1, x_1); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + _0 = _mm512_add_ph(_0, x_0); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + _1 = _mm512_add_ph(_1, x_1); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_v4_avx512fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v4_avx512fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -262,8 +327,187 @@ mod reduce_sum_of_x { } } + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_sum_of_x_a2_fp16(this: &[f16]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + sum_0 = vaddq_f16(sum_0, x_0); + sum_1 = vaddq_f16(sum_1, x_1); + sum_2 = vaddq_f16(sum_2, x_2); + sum_3 = vaddq_f16(sum_3, x_3); + sum_4 = vaddq_f16(sum_4, x_4); + sum_5 = vaddq_f16(sum_5, x_5); + sum_6 = vaddq_f16(sum_6, x_6); + sum_7 = vaddq_f16(sum_7, x_7); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + sum_0 = vaddq_f16(sum_0, x_0); + sum_1 = vaddq_f16(sum_1, x_1); + sum_2 = vaddq_f16(sum_2, x_2); + sum_3 = vaddq_f16(sum_3, x_3); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + sum_4 = vaddq_f16(sum_4, x_4); + sum_5 = vaddq_f16(sum_5, x_5); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + sum_6 = vaddq_f16(sum_6, x_6); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + sum_7 = vaddq_f16(sum_7, x_7); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_a2_fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_a2_fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x_a2(this: &[f16]) -> f32 { + use crate::emulate::emulate_vreinterpret_f16_u16; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vaddq_f32(sum, x); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let mut _a = [f16::_ZERO; 4]; + let mut _b = [f16::_ZERO; 4]; + unsafe { + std::ptr::copy_nonoverlapping(a.cast(), _a.as_mut_ptr(), n); + } + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vaddq_f32(sum, x); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_a2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + #[crate::multiversion( - @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + @"v4:avx512fp16", @"v4", @"v3", "v2", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_x(this: &[f16]) -> f32 { let n = this.len(); @@ -278,6 +522,71 @@ mod reduce_sum_of_x { mod reduce_sum_of_abs_x { use crate::{F16, f16}; + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_sum_of_abs_x_v4_avx512fp16(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + _0 = _mm512_add_ph(_0, _mm512_abs_ph(x_0)); + _1 = _mm512_add_ph(_1, _mm512_abs_ph(x_1)); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + while n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + _0 = _mm512_add_ph(_0, _mm512_abs_ph(x_0)); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + _1 = _mm512_add_ph(_1, _mm512_abs_ph(x_1)); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_v4_avx512fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v4_avx512fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -376,8 +685,187 @@ mod reduce_sum_of_abs_x { } } + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_sum_of_abs_x_a2_fp16(this: &[f16]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + sum_0 = vaddq_f16(sum_0, vabsq_f16(x_0)); + sum_1 = vaddq_f16(sum_1, vabsq_f16(x_1)); + sum_2 = vaddq_f16(sum_2, vabsq_f16(x_2)); + sum_3 = vaddq_f16(sum_3, vabsq_f16(x_3)); + sum_4 = vaddq_f16(sum_4, vabsq_f16(x_4)); + sum_5 = vaddq_f16(sum_5, vabsq_f16(x_5)); + sum_6 = vaddq_f16(sum_6, vabsq_f16(x_6)); + sum_7 = vaddq_f16(sum_7, vabsq_f16(x_7)); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + sum_0 = vaddq_f16(sum_0, vabsq_f16(x_0)); + sum_1 = vaddq_f16(sum_1, vabsq_f16(x_1)); + sum_2 = vaddq_f16(sum_2, vabsq_f16(x_2)); + sum_3 = vaddq_f16(sum_3, vabsq_f16(x_3)); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + sum_4 = vaddq_f16(sum_4, vabsq_f16(x_4)); + sum_5 = vaddq_f16(sum_5, vabsq_f16(x_5)); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + sum_6 = vaddq_f16(sum_6, vabsq_f16(x_6)); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + sum_7 = vaddq_f16(sum_7, vabsq_f16(x_7)); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_a2_fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_a2_fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_abs_x_a2(this: &[f16]) -> f32 { + use crate::emulate::emulate_vreinterpret_f16_u16; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vaddq_f32(sum, vabsq_f32(x)); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let mut _a = [f16::_ZERO; 4]; + let mut _b = [f16::_ZERO; 4]; + unsafe { + std::ptr::copy_nonoverlapping(a.cast(), _a.as_mut_ptr(), n); + } + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vaddq_f32(sum, vabsq_f32(x)); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_a2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + #[crate::multiversion( - @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + @"v4:avx512fp16", @"v4", @"v3", "v2", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { let n = this.len(); @@ -392,6 +880,71 @@ mod reduce_sum_of_abs_x { mod reduce_sum_of_x2 { use crate::{F16, f16}; + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_sum_of_x2_v4_avx512fp16(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + _0 = _mm512_fmadd_ph(x_0, x_0, _0); + _1 = _mm512_fmadd_ph(x_1, x_1, _1); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + _0 = _mm512_fmadd_ph(x_0, x_0, _0); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + _1 = _mm512_fmadd_ph(x_1, x_1, _1); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_v4_avx512fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v4_avx512fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -464,11 +1017,186 @@ mod reduce_sum_of_x2 { #[cfg(all(target_arch = "x86_64", test))] #[test] #[cfg_attr(miri, ignore)] - fn reduce_sum_of_x2_v3_test() { + fn reduce_sum_of_x2_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_sum_of_x2_a2_fp16(this: &[f16]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + sum_0 = vfmaq_f16(sum_0, x_0, x_0); + sum_1 = vfmaq_f16(sum_1, x_1, x_1); + sum_2 = vfmaq_f16(sum_2, x_2, x_2); + sum_3 = vfmaq_f16(sum_3, x_3, x_3); + sum_4 = vfmaq_f16(sum_4, x_4, x_4); + sum_5 = vfmaq_f16(sum_5, x_5, x_5); + sum_6 = vfmaq_f16(sum_6, x_6, x_6); + sum_7 = vfmaq_f16(sum_7, x_7, x_7); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + sum_0 = vfmaq_f16(sum_0, x_0, x_0); + sum_1 = vfmaq_f16(sum_1, x_1, x_1); + sum_2 = vfmaq_f16(sum_2, x_2, x_2); + sum_3 = vfmaq_f16(sum_3, x_3, x_3); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + sum_4 = vfmaq_f16(sum_4, x_4, x_4); + sum_5 = vfmaq_f16(sum_5, x_5, x_5); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + sum_6 = vfmaq_f16(sum_6, x_6, x_6); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + sum_7 = vfmaq_f16(sum_7, x_7, x_7); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_a2_fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_a2_fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x2_a2(this: &[f16]) -> f32 { + use crate::emulate::{emulate_vreinterpret_f16_u16, partial_load}; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vfmaq_f32(sum, x, x); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vfmaq_f32(sum, x, x); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_a2_test() { use rand::RngExt; const EPSILON: f32 = 2.0; - if !crate::is_cpu_detected!("v3") { - println!("test {} ... skipped (v3)", module_path!()); + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); return; } let mut rng = rand::rng(); @@ -479,7 +1207,7 @@ mod reduce_sum_of_x2 { .collect::>(); for z in 3984..4016 { let this = &this[..z]; - let specialized = unsafe { reduce_sum_of_x2_v3(this) }; + let specialized = unsafe { reduce_sum_of_x2_a2(this) }; let fallback = fallback(this); assert!( (specialized - fallback).abs() < EPSILON, @@ -490,7 +1218,7 @@ mod reduce_sum_of_x2 { } #[crate::multiversion( - @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + @"v4:avx512fp16", @"v4", @"v3", "v2", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_sum_of_x2(this: &[f16]) -> f32 { let n = this.len(); @@ -505,6 +1233,67 @@ mod reduce_sum_of_x2 { mod reduce_min_max_of_x { use crate::{F16, f16}; + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_min_max_of_x_v4_avx512fp16(this: &[f16]) -> (f32, f32) { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm512_cvtepi16_ph(_mm512_set1_epi16(f16::INFINITY.to_bits() as _)); + let mut max = _mm512_cvtepi16_ph(_mm512_set1_epi16(f16::NEG_INFINITY.to_bits() as _)); + while n >= 32 { + let x = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.cast())) }; + min = _mm512_min_ph(x, min); + max = _mm512_max_ph(x, max); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + min = _mm512_mask_min_ph(min, mask, x, min); + max = _mm512_mask_max_ph(max, mask, x, max); + } + let min = { + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(min), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(min), 1)); + _mm512_reduce_min_ps(_mm512_min_ps(s_0, s_1)) + }; + let max = { + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(max), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(max), 1)); + _mm512_reduce_max_ps(_mm512_max_ps(s_0, s_1)) + }; + (min, max) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_min_max_of_x_v4_avx512fp16_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v4_avx512fp16(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0); + assert_eq!(specialized.1, fallback.1); + } + } + } + #[inline] #[cfg(target_arch = "x86_64")] #[crate::target_cpu(enable = "v4")] @@ -613,8 +1402,122 @@ mod reduce_min_max_of_x { } } + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_min_max_of_x_a2_fp16(this: &[f16]) -> (f32, f32) { + use crate::emulate::{emulate_vreinterpretq_f16_u16, partial_load}; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = emulate_vreinterpretq_f16_u16(vdupq_n_u16(0x7C00u16)); + let mut max = emulate_vreinterpretq_f16_u16(vdupq_n_u16(0xFC00u16)); + while n >= 8 { + let x = emulate_vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + min = vminnmq_f16(x, min); + max = vmaxnmq_f16(x, max); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a = f16::NAN) }; + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + min = vminnmq_f16(x, min); + max = vmaxnmq_f16(x, max); + } + ( + vminnmvq_f32(vcvt_f32_f16(vminnm_f16( + vget_low_f16(min), + vget_high_f16(min), + ))), + vmaxnmvq_f32(vcvt_f32_f16(vmaxnm_f16( + vget_low_f16(max), + vget_high_f16(max), + ))), + ) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + fn reduce_min_max_of_x_a2_fp16_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_a2_fp16(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_min_max_of_x_a2(this: &[f16]) -> (f32, f32) { + use crate::emulate::{emulate_vreinterpret_f16_u16, partial_load}; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = vdupq_n_f32(f32::INFINITY); + let mut max = vdupq_n_f32(f32::NEG_INFINITY); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + min = vminnmq_f32(x, min); + max = vmaxnmq_f32(x, max); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a = f16::NAN) }; + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + min = vminnmq_f32(x, min); + max = vmaxnmq_f32(x, max); + } + (vminnmvq_f32(min), vmaxnmvq_f32(max)) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + fn reduce_min_max_of_x_a2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_a2(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + #[crate::multiversion( - @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + @"v4:avx512fp16", @"v4", @"v3", "v2", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" )] pub fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { let mut min = f32::INFINITY; @@ -638,15 +1541,42 @@ mod reduce_sum_of_xy { #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512fp16")] fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { - unsafe extern "C" { - #[link_name = "fp16_reduce_sum_of_xy_v4_avx512fp16"] - unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; - } assert!(lhs.len() == rhs.len()); - let n = lhs.len(); - let a = lhs.as_ptr(); - let b = rhs.as_ptr(); - unsafe { f(n, a, b) } + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let y_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + let y_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(32).cast())) }; + _0 = _mm512_fmadd_ph(x_0, y_0, _0); + _1 = _mm512_fmadd_ph(x_1, y_1, _1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let y_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(0).cast())) }; + _0 = _mm512_fmadd_ph(x_0, y_0, _0); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + let y_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())) }; + _1 = _mm512_fmadd_ph(x_1, y_1, _1); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) } #[cfg(all(target_arch = "x86_64", test))] @@ -849,15 +1779,111 @@ mod reduce_sum_of_xy { #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "fp16")] fn reduce_sum_of_xy_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { - unsafe extern "C" { - #[link_name = "fp16_reduce_sum_of_xy_a2_fp16"] - unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; - } + use crate::emulate::partial_load; + use core::arch::aarch64::*; assert!(lhs.len() == rhs.len()); - let n = lhs.len(); - let a = lhs.as_ptr(); - let b = rhs.as_ptr(); - unsafe { f(n, a, b) } + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + let y_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let y_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(16).cast()) }); + let y_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(24).cast()) }); + let y_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(32).cast()) }); + let y_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(40).cast()) }); + let y_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(48).cast()) }); + let y_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(56).cast()) }); + sum_0 = vfmaq_f16(sum_0, x_0, y_0); + sum_1 = vfmaq_f16(sum_1, x_1, y_1); + sum_2 = vfmaq_f16(sum_2, x_2, y_2); + sum_3 = vfmaq_f16(sum_3, x_3, y_3); + sum_4 = vfmaq_f16(sum_4, x_4, y_4); + sum_5 = vfmaq_f16(sum_5, x_5, y_5); + sum_6 = vfmaq_f16(sum_6, x_6, y_6); + sum_7 = vfmaq_f16(sum_7, x_7, y_7); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let y_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let y_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(16).cast()) }); + let y_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(24).cast()) }); + sum_0 = vfmaq_f16(sum_0, x_0, y_0); + sum_1 = vfmaq_f16(sum_1, x_1, y_1); + sum_2 = vfmaq_f16(sum_2, x_2, y_2); + sum_3 = vfmaq_f16(sum_3, x_3, y_3); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let y_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + sum_4 = vfmaq_f16(sum_4, x_4, y_4); + sum_5 = vfmaq_f16(sum_5, x_5, y_5); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let y_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + sum_6 = vfmaq_f16(sum_6, x_6, y_6); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(8, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + let y_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.cast()) }); + sum_7 = vfmaq_f16(sum_7, x_7, y_7); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) } #[cfg(all(target_arch = "aarch64", test))] @@ -892,7 +1918,70 @@ mod reduce_sum_of_xy { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_xy_a2(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::{emulate_vreinterpret_f16_u16, partial_load}; + use core::arch::aarch64::*; + assert!(lhs.len() == rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + let y = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(b.cast()) }); + let y = vcvt_f32_f16(y); + sum = vfmaq_f32(sum, x, y); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + let y = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(b.cast()) }); + let y = vcvt_f32_f16(y); + sum = vfmaq_f32(sum, x, y); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); @@ -914,15 +2003,46 @@ mod reduce_sum_of_d2 { #[crate::target_cpu(enable = "v4")] #[target_feature(enable = "avx512fp16")] fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { - unsafe extern "C" { - #[link_name = "fp16_reduce_sum_of_d2_v4_avx512fp16"] - unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; - } assert!(lhs.len() == rhs.len()); - let n = lhs.len(); - let a = lhs.as_ptr(); - let b = rhs.as_ptr(); - unsafe { f(n, a, b) } + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let y_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + let y_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(32).cast())) }; + let d_0 = _mm512_sub_ph(x_0, y_0); + let d_1 = _mm512_sub_ph(x_1, y_1); + _0 = _mm512_fmadd_ph(d_0, d_0, _0); + _1 = _mm512_fmadd_ph(d_1, d_1, _1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let y_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(0).cast())) }; + let d_0 = _mm512_sub_ph(x_0, y_0); + _0 = _mm512_fmadd_ph(d_0, d_0, _0); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + let y_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())) }; + let d_1 = _mm512_sub_ph(x_1, y_1); + _1 = _mm512_fmadd_ph(d_1, d_1, _1); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) } #[cfg(all(target_arch = "x86_64", test))] @@ -1134,15 +2254,127 @@ mod reduce_sum_of_d2 { #[crate::target_cpu(enable = "a2")] #[target_feature(enable = "fp16")] fn reduce_sum_of_d2_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { - unsafe extern "C" { - #[link_name = "fp16_reduce_sum_of_d2_a2_fp16"] - unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; - } + use crate::emulate::partial_load; + use core::arch::aarch64::*; assert!(lhs.len() == rhs.len()); - let n = lhs.len(); - let a = lhs.as_ptr().cast(); - let b = rhs.as_ptr().cast(); - unsafe { f(n, a, b) } + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + let y_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let y_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(16).cast()) }); + let y_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(24).cast()) }); + let y_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(32).cast()) }); + let y_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(40).cast()) }); + let y_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(48).cast()) }); + let y_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(56).cast()) }); + let d_0 = vsubq_f16(x_0, y_0); + let d_1 = vsubq_f16(x_1, y_1); + let d_2 = vsubq_f16(x_2, y_2); + let d_3 = vsubq_f16(x_3, y_3); + let d_4 = vsubq_f16(x_4, y_4); + let d_5 = vsubq_f16(x_5, y_5); + let d_6 = vsubq_f16(x_6, y_6); + let d_7 = vsubq_f16(x_7, y_7); + sum_0 = vfmaq_f16(sum_0, d_0, d_0); + sum_1 = vfmaq_f16(sum_1, d_1, d_1); + sum_2 = vfmaq_f16(sum_2, d_2, d_2); + sum_3 = vfmaq_f16(sum_3, d_3, d_3); + sum_4 = vfmaq_f16(sum_4, d_4, d_4); + sum_5 = vfmaq_f16(sum_5, d_5, d_5); + sum_6 = vfmaq_f16(sum_6, d_6, d_6); + sum_7 = vfmaq_f16(sum_7, d_7, d_7); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let y_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let y_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(16).cast()) }); + let y_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(24).cast()) }); + let d_0 = vsubq_f16(x_0, y_0); + let d_1 = vsubq_f16(x_1, y_1); + let d_2 = vsubq_f16(x_2, y_2); + let d_3 = vsubq_f16(x_3, y_3); + sum_0 = vfmaq_f16(sum_0, d_0, d_0); + sum_1 = vfmaq_f16(sum_1, d_1, d_1); + sum_2 = vfmaq_f16(sum_2, d_2, d_2); + sum_3 = vfmaq_f16(sum_3, d_3, d_3); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let y_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let d_4 = vsubq_f16(x_4, y_4); + let d_5 = vsubq_f16(x_5, y_5); + sum_4 = vfmaq_f16(sum_4, d_4, d_4); + sum_5 = vfmaq_f16(sum_5, d_5, d_5); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let y_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let d_6 = vsubq_f16(x_6, y_6); + sum_6 = vfmaq_f16(sum_6, d_6, d_6); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(8, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + let y_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.cast()) }); + let d_7 = vsubq_f16(x_7, y_7); + sum_7 = vfmaq_f16(sum_7, d_7, d_7); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) } #[cfg(all(target_arch = "aarch64", test))] @@ -1177,7 +2409,72 @@ mod reduce_sum_of_d2 { } } - #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_d2_a2(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::{emulate_vreinterpret_f16_u16, partial_load}; + use core::arch::aarch64::*; + assert!(lhs.len() == rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + let y = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(b.cast()) }); + let y = vcvt_f32_f16(y); + let d = vsubq_f32(x, y); + sum = vfmaq_f32(sum, d, d); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + let y = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(b.cast()) }); + let y = vcvt_f32_f16(y); + let d = vsubq_f32(x, y); + sum = vfmaq_f32(sum, d, d); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { assert!(lhs.len() == rhs.len()); let n = lhs.len(); From 2934c9441b93b7c802327a73c1d6bf40d1e74bdf Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 11 Mar 2026 18:43:35 +0800 Subject: [PATCH 285/324] ci: empty OPTFLAGS if sccache is used (#441) https://github.com/tensorchord/VectorChord/pull/436#issuecomment-3976244820 Signed-off-by: usamoi --- .github/workflows/check.yml | 12 +++--- crates/index/src/tuples.rs | 2 +- crates/vchordg/src/lib.rs | 2 +- crates/vchordg/src/tuples.rs | 18 ++++---- crates/vchordrq/src/lib.rs | 2 +- crates/vchordrq/src/tuples.rs | 68 ++++++++++++++++--------------- src/index/vchordg/am/am_build.rs | 2 +- src/index/vchordrq/am/am_build.rs | 4 +- 8 files changed, 57 insertions(+), 53 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index b9858fc2..d92607cb 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -378,8 +378,8 @@ jobs: mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install - make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' - sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' install + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' OPTFLAGS="" + sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' OPTFLAGS="" install - name: Checkout uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -619,8 +619,8 @@ jobs: mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install - make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' - make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' install + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' OPTFLAGS="" + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' OPTFLAGS="" install - name: Checkout uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -1022,8 +1022,8 @@ jobs: mkdir ~/pgvector-install curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install pushd ~/pgvector-install/pgvector-0.8.1 - make -j$(nproc) - sudo make install + make OPTFLAGS="" -j$(nproc) + sudo make OPTFLAGS="" install popd - name: Checkout diff --git a/crates/index/src/tuples.rs b/crates/index/src/tuples.rs index 8483db05..76dd51fc 100644 --- a/crates/index/src/tuples.rs +++ b/crates/index/src/tuples.rs @@ -137,7 +137,7 @@ impl<'a> MutChecker<'a> { #[test] fn aliasing_test() { #[repr(C, align(8))] - #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct ExampleHeader { elements_s: u16, elements_e: u16, diff --git a/crates/vchordg/src/lib.rs b/crates/vchordg/src/lib.rs index fbedf22d..f84de74a 100644 --- a/crates/vchordg/src/lib.rs +++ b/crates/vchordg/src/lib.rs @@ -38,7 +38,7 @@ pub use search::search; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] -#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, Copy, FromBytes, IntoBytes, Immutable, KnownLayout)] pub struct Opaque { pub next: u32, pub link: u32, diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 0f5da7e7..7f26d314 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -47,7 +47,7 @@ pub trait WithWriter: Tuple { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct MetaTupleHeader { version: u64, dim: u32, @@ -206,7 +206,7 @@ impl<'a> MetaTupleWriter<'a> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VertexTupleHeader { metadata: [f32; 3], elements_s: u16, @@ -217,7 +217,7 @@ struct VertexTupleHeader { _padding_0: [Padding; 4], } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub struct VertexTuple { pub metadata: [f32; 3], pub elements: Vec, @@ -292,7 +292,7 @@ impl WithWriter for VertexTuple { } } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub struct VertexTupleReader<'a> { header: &'a VertexTupleHeader, elements: &'a [u64], @@ -339,7 +339,7 @@ impl VertexTupleWriter<'_> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader0 { payload: Option>, elements_s: u16, @@ -352,7 +352,7 @@ struct VectorTupleHeader0 { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader1 { payload: Option>, elements_s: u16, @@ -360,7 +360,7 @@ struct VectorTupleHeader1 { index: u32, } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub enum VectorTuple { _0 { payload: Option>, @@ -528,7 +528,7 @@ impl<'a, V: Vector> Clone for VectorTupleReader<'a, V> { } } -#[derive(Debug, PartialEq)] +#[derive(Debug)] pub struct VectorTupleReader0<'a, V: Vector> { header: &'a VectorTupleHeader0, elements: &'a [V::Element], @@ -562,7 +562,7 @@ impl<'a, V: Vector> VectorTupleReader0<'a, V> { } } -#[derive(Debug, PartialEq)] +#[derive(Debug)] pub struct VectorTupleReader1<'a, V: Vector> { header: &'a VectorTupleHeader1, elements: &'a [V::Element], diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index e2938dab..726f29cc 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -50,7 +50,7 @@ pub use search::{default_search, maxsim_search}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] -#[derive(Debug, Clone, Copy, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, Copy, FromBytes, IntoBytes, Immutable, KnownLayout)] pub struct Opaque { pub next: u32, pub skip: u32, diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index a66bacfd..6f1072c0 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -46,7 +46,7 @@ pub trait WithWriter: Tuple { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct MetaTupleHeader { version: u64, dim: u32, @@ -234,13 +234,13 @@ impl<'a> MetaTupleReader<'a> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FreepagesTupleHeader { first: u32, _padding_0: [Padding; 4], } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub struct FreepagesTuple {} impl Tuple for FreepagesTuple { @@ -275,7 +275,7 @@ impl FreepagesTupleWriter<'_> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct CentroidTupleHeader0 { metadata_s: u16, elements_s: u16, @@ -284,7 +284,7 @@ struct CentroidTupleHeader0 { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct CentroidTupleHeader1 { head: u16, elements_s: u16, @@ -292,7 +292,7 @@ struct CentroidTupleHeader1 { _padding_0: [Padding; 2], } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub enum CentroidTuple { _0 { metadata: V::Metadata, @@ -431,7 +431,7 @@ impl<'a, V: Vector> CentroidTupleReader<'a, V> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader0 { payload: Option>, metadata_s: u16, @@ -441,7 +441,7 @@ struct VectorTupleHeader0 { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VectorTupleHeader1 { payload: Option>, head: u16, @@ -450,7 +450,7 @@ struct VectorTupleHeader1 { _padding_0: [Padding; 2], } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub enum VectorTuple { _0 { payload: Option>, @@ -606,7 +606,7 @@ impl<'a, V: Vector> VectorTupleReader<'a, V> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct DirectoryTupleHeader0 { elements_s: u16, elements_e: u16, @@ -614,14 +614,14 @@ struct DirectoryTupleHeader0 { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct DirectoryTupleHeader1 { elements_s: u16, elements_e: u16, _padding_0: [Padding; 4], } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub enum DirectoryTuple { _0 { elements: Vec }, _1 { elements: Vec }, @@ -722,20 +722,22 @@ impl WithReader for DirectoryTuple { } } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub enum DirectoryTupleReader<'a> { _0(DirectoryTupleReader0<'a>), _1(DirectoryTupleReader1<'a>), } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub struct DirectoryTupleReader0<'a> { + #[allow(dead_code)] header: &'a DirectoryTupleHeader0, elements: &'a [u32], } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub struct DirectoryTupleReader1<'a> { + #[allow(dead_code)] header: &'a DirectoryTupleHeader1, elements: &'a [u32], } @@ -753,7 +755,7 @@ impl<'a> DirectoryTupleReader1<'a> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct H1TupleHeader0 { metadata: [[f32; 32]; 4], delta: [f32; 32], @@ -769,7 +771,7 @@ struct H1TupleHeader0 { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct H1TupleHeader1 { elements_s: u16, elements_e: u16, @@ -777,7 +779,7 @@ struct H1TupleHeader1 { } #[allow(clippy::large_enum_variant)] -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub enum H1Tuple { _0 { metadata: [[f32; 32]; 4], @@ -921,21 +923,22 @@ impl WithReader for H1Tuple { } } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub enum H1TupleReader<'a> { _0(H1TupleReader0<'a>), _1(H1TupleReader1<'a>), } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub struct H1TupleReader0<'a> { header: &'a H1TupleHeader0, prefetch: &'a [[u32; 32]], elements: &'a [[u8; 16]], } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub struct H1TupleReader1<'a> { + #[allow(dead_code)] header: &'a H1TupleHeader1, elements: &'a [[u8; 16]], } @@ -974,7 +977,7 @@ impl<'a> H1TupleReader1<'a> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct JumpTupleHeader { centroid_prefetch_s: u16, centroid_prefetch_e: u16, @@ -986,7 +989,7 @@ struct JumpTupleHeader { tuples: u64, } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub struct JumpTuple { pub centroid_prefetch: Vec, pub centroid_head: u16, @@ -1096,7 +1099,7 @@ impl JumpTupleWriter<'_> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FrozenTupleHeader0 { metadata: [[f32; 32]; 4], delta: [f32; 32], @@ -1110,7 +1113,7 @@ struct FrozenTupleHeader0 { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FrozenTupleHeader1 { elements_s: u16, elements_e: u16, @@ -1118,7 +1121,7 @@ struct FrozenTupleHeader1 { } #[allow(clippy::large_enum_variant)] -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub enum FrozenTuple { _0 { metadata: [[f32; 32]; 4], @@ -1278,13 +1281,13 @@ impl WithWriter for FrozenTuple { } } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub enum FrozenTupleReader<'a> { _0(FrozenTupleReader0<'a>), _1(FrozenTupleReader1<'a>), } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub struct FrozenTupleReader0<'a> { header: &'a FrozenTupleHeader0, prefetch: &'a [[u32; 32]], @@ -1312,8 +1315,9 @@ impl<'a> FrozenTupleReader0<'a> { } } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub struct FrozenTupleReader1<'a> { + #[allow(dead_code)] header: &'a FrozenTupleHeader1, elements: &'a [[u8; 16]], } @@ -1353,7 +1357,7 @@ impl FrozenTupleWriter0<'_> { } #[repr(C, align(8))] -#[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct AppendableTupleHeader { metadata: [f32; 4], delta: f32, @@ -1367,7 +1371,7 @@ struct AppendableTupleHeader { payload: Option>, } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub struct AppendableTuple { pub metadata: [f32; 4], pub delta: f32, @@ -1441,7 +1445,7 @@ impl WithWriter for AppendableTuple { } } -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy)] pub struct AppendableTupleReader<'a> { header: &'a AppendableTupleHeader, prefetch: &'a [u32], diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index 8fc2d549..d1c3c465 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -612,7 +612,7 @@ mod vchordg_cached { use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] - #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VchordgCachedHeader0 {} pub enum VchordgCached { diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 4674895d..a33434ca 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -476,11 +476,11 @@ mod vchordrq_cached { use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] - #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VchordrqCachedHeader0 {} #[repr(C, align(8))] - #[derive(Debug, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] + #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct VchordrqCachedHeader1 { mapping_s: usize, mapping_e: usize, From cd0b1032dcba2db04353a81f656f506e38dc133f Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 16 Mar 2026 15:40:40 +0800 Subject: [PATCH 286/324] chore: move `index::accessor` to another crate (#442) Signed-off-by: usamoi --- Cargo.lock | 14 +++++++++++++- Cargo.toml | 1 + crates/index/Cargo.toml | 1 - crates/index/src/lib.rs | 2 +- crates/index_accessor/Cargo.toml | 17 +++++++++++++++++ crates/vchordg/Cargo.toml | 1 + crates/vchordg/src/insert.rs | 2 +- crates/vchordg/src/maintain.rs | 2 +- crates/vchordg/src/operator.rs | 2 +- crates/vchordg/src/search.rs | 2 +- crates/vchordg/src/vectors.rs | 2 +- crates/vchordrq/Cargo.toml | 1 + crates/vchordrq/src/bulkdelete.rs | 2 +- crates/vchordrq/src/cache.rs | 2 +- crates/vchordrq/src/centroids.rs | 2 +- crates/vchordrq/src/insert.rs | 2 +- crates/vchordrq/src/maintain.rs | 2 +- crates/vchordrq/src/operator.rs | 2 +- crates/vchordrq/src/prewarm.rs | 2 +- crates/vchordrq/src/rerank.rs | 2 +- crates/vchordrq/src/search.rs | 2 +- crates/vchordrq/src/tape.rs | 2 +- crates/vchordrq/src/vectors.rs | 2 +- src/index/vchordg/am/am_build.rs | 2 +- src/index/vchordg/dispatch.rs | 2 +- src/index/vchordg/scanners/default.rs | 2 +- src/index/vchordrq/am/am_build.rs | 2 +- src/index/vchordrq/am/mod.rs | 11 +---------- src/index/vchordrq/dispatch.rs | 2 +- src/index/vchordrq/scanners/default.rs | 2 +- src/index/vchordrq/scanners/maxsim.rs | 2 +- 31 files changed, 58 insertions(+), 36 deletions(-) create mode 100644 crates/index_accessor/Cargo.toml diff --git a/Cargo.lock b/Cargo.lock index ebbd5472..19a975a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -841,12 +841,21 @@ dependencies = [ "dary_heap", "distance", "rabitq", - "simd", "small_iter", "vector", "zerocopy", ] +[[package]] +name = "index_accessor" +version = "0.0.0" +dependencies = [ + "distance", + "rabitq", + "simd", + "vector", +] + [[package]] name = "indexmap" version = "2.13.0" @@ -1942,6 +1951,7 @@ dependencies = [ "feistel", "humansize", "index", + "index_accessor", "k_means", "mimalloc", "paste", @@ -1971,6 +1981,7 @@ dependencies = [ "always_equal", "distance", "index", + "index_accessor", "min-max-heap", "rabitq", "rand", @@ -1988,6 +1999,7 @@ dependencies = [ "always_equal", "distance", "index", + "index_accessor", "pin-project", "rabitq", "rand", diff --git a/Cargo.toml b/Cargo.toml index 743081fe..99b9aebc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,6 +25,7 @@ always_equal = { path = "./crates/always_equal" } distance = { path = "./crates/distance" } feistel = { path = "./crates/feistel" } index = { path = "./crates/index" } +index_accessor = { path = "./crates/index_accessor" } k_means = { path = "./crates/k_means" } rabitq = { path = "./crates/rabitq" } simd = { path = "./crates/simd" } diff --git a/crates/index/Cargo.toml b/crates/index/Cargo.toml index deb0d497..7ade4785 100644 --- a/crates/index/Cargo.toml +++ b/crates/index/Cargo.toml @@ -8,7 +8,6 @@ publish = false always_equal = { path = "../always_equal" } distance = { path = "../distance" } rabitq = { path = "../rabitq" } -simd = { path = "../simd" } small_iter = { path = "../small_iter" } vector = { path = "../vector" } diff --git a/crates/index/src/lib.rs b/crates/index/src/lib.rs index 1ba2ed4e..6ec70b6b 100644 --- a/crates/index/src/lib.rs +++ b/crates/index/src/lib.rs @@ -12,7 +12,7 @@ // // Copyright (c) 2025 TensorChord Inc. -pub mod accessor; +// pub mod accessor; pub mod bump; pub mod fetch; pub mod packed; diff --git a/crates/index_accessor/Cargo.toml b/crates/index_accessor/Cargo.toml new file mode 100644 index 00000000..e23da0f3 --- /dev/null +++ b/crates/index_accessor/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "index_accessor" +version.workspace = true +edition.workspace = true +publish = false + +[lib] +path = "../index/src/accessor.rs" + +[dependencies] +distance = { path = "../distance" } +rabitq = { path = "../rabitq" } +simd = { path = "../simd" } +vector = { path = "../vector" } + +[lints] +workspace = true diff --git a/crates/vchordg/Cargo.toml b/crates/vchordg/Cargo.toml index 7b5a1658..c1772834 100644 --- a/crates/vchordg/Cargo.toml +++ b/crates/vchordg/Cargo.toml @@ -8,6 +8,7 @@ publish = false always_equal = { path = "../always_equal" } distance = { path = "../distance" } index = { path = "../index" } +index_accessor = { path = "../index_accessor" } rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index 5b72d923..a133ae6d 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -21,10 +21,10 @@ use crate::types::DistanceKind; use crate::vectors::{by_prefetch, by_read, copy_all, copy_nothing, copy_outs, update}; use crate::visited::Visited; use always_equal::AlwaysEqual; -use index::accessor::{DefaultWithDimension, LAccess}; use index::bump::Bump; use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; +use index_accessor::{DefaultWithDimension, LAccess}; use rabitq::bits::Bits; use std::cmp::Reverse; use std::collections::VecDeque; diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs index de25edbc..6d43daad 100644 --- a/crates/vchordg/src/maintain.rs +++ b/crates/vchordg/src/maintain.rs @@ -18,8 +18,8 @@ use crate::tuples::{MetaTuple, VertexTuple, WithReader}; use crate::types::DistanceKind; use crate::vectors::{by_read, copy_all, copy_nothing, copy_outs, update}; use always_equal::AlwaysEqual; -use index::accessor::DefaultWithDimension; use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; +use index_accessor::DefaultWithDimension; use std::cmp::Reverse; use vector::VectorOwned; diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs index 6eead52e..bc279567 100644 --- a/crates/vchordg/src/operator.rs +++ b/crates/vchordg/src/operator.rs @@ -14,7 +14,7 @@ use crate::types::DistanceKind; use distance::Distance; -use index::accessor::{ +use index_accessor::{ Accessor1, Accessor2, ByteDistanceAccessor, DefaultWithDimension, DistanceAccessor, Dot, HalfbyteDistanceAccessor, L2S, }; diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs index ee5ee353..33a378e9 100644 --- a/crates/vchordg/src/search.rs +++ b/crates/vchordg/src/search.rs @@ -21,10 +21,10 @@ use crate::vectors::{by_prefetch, copy_outs}; use crate::visited::Visited; use always_equal::AlwaysEqual; use distance::Distance; -use index::accessor::{DefaultWithDimension, LAccess}; use index::bump::Bump; use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use index::relation::{Page, RelationRead}; +use index_accessor::{DefaultWithDimension, LAccess}; use rabitq::bits::Bits; use std::cmp::Reverse; use std::collections::VecDeque; diff --git a/crates/vchordg/src/vectors.rs b/crates/vchordg/src/vectors.rs index 95b83e25..272ba48f 100644 --- a/crates/vchordg/src/vectors.rs +++ b/crates/vchordg/src/vectors.rs @@ -15,8 +15,8 @@ use crate::operator::{Operator, Vector}; use crate::tuples::*; use distance::Distance; -use index::accessor::Accessor1; use index::relation::{Page, RelationRead, RelationWrite}; +use index_accessor::Accessor1; use std::collections::VecDeque; use std::num::{NonZero, Wrapping}; diff --git a/crates/vchordrq/Cargo.toml b/crates/vchordrq/Cargo.toml index 89c4c360..ebac0e59 100644 --- a/crates/vchordrq/Cargo.toml +++ b/crates/vchordrq/Cargo.toml @@ -8,6 +8,7 @@ publish = false always_equal = { path = "../always_equal" } distance = { path = "../distance" } index = { path = "../index" } +index_accessor = { path = "../index_accessor" } rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } diff --git a/crates/vchordrq/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs index 6bdbe530..d5e64acd 100644 --- a/crates/vchordrq/src/bulkdelete.rs +++ b/crates/vchordrq/src/bulkdelete.rs @@ -17,8 +17,8 @@ use crate::operator::Operator; use crate::tape::by_next; use crate::tuples::*; use crate::{Opaque, tape}; -use index::accessor::FunctionalAccessor; use index::relation::{Page, RelationRead, RelationWrite}; +use index_accessor::FunctionalAccessor; use std::num::NonZero; pub fn bulkdelete( diff --git a/crates/vchordrq/src/cache.rs b/crates/vchordrq/src/cache.rs index 89c9178e..65c28347 100644 --- a/crates/vchordrq/src/cache.rs +++ b/crates/vchordrq/src/cache.rs @@ -15,8 +15,8 @@ use crate::closure_lifetime_binder::{id_0, id_1}; use crate::tuples::{MetaTuple, WithReader}; use crate::{Opaque, tape}; -use index::accessor::FunctionalAccessor; use index::relation::{Page, PageGuard, RelationRead}; +use index_accessor::FunctionalAccessor; pub fn cache(index: &R, level: i32) -> Vec where diff --git a/crates/vchordrq/src/centroids.rs b/crates/vchordrq/src/centroids.rs index 44221a86..521987d5 100644 --- a/crates/vchordrq/src/centroids.rs +++ b/crates/vchordrq/src/centroids.rs @@ -14,8 +14,8 @@ use crate::operator::*; use crate::tuples::*; -use index::accessor::Accessor1; use index::relation::{Page, RelationRead}; +use index_accessor::Accessor1; pub fn read< 'a, diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index 32ae080a..bf256b08 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -19,11 +19,11 @@ use crate::tuples::*; use crate::{Opaque, centroids, tape, vectors}; use always_equal::AlwaysEqual; use distance::Distance; -use index::accessor::{DefaultWithDimension, FunctionalAccessor, LAccess}; use index::bump::Bump; use index::fetch::BorrowedIter; use index::prefetcher::{Prefetcher, PrefetcherHeapFamily}; use index::relation::{Page, RelationRead, RelationWrite}; +use index_accessor::{DefaultWithDimension, FunctionalAccessor, LAccess}; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZero; diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs index 24f4b255..291fa2e5 100644 --- a/crates/vchordrq/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -17,11 +17,11 @@ use crate::operator::{Operator, Vector}; use crate::tape_writer::{DirectoryTapeWriter, FrozenTapeWriter}; use crate::tuples::*; use crate::{Branch, Opaque, freepages, tape}; -use index::accessor::FunctionalAccessor; use index::prefetcher::PrefetcherSequenceFamily; use index::relation::{ Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, }; +use index_accessor::FunctionalAccessor; use rabitq::packing::unpack; use std::cell::RefCell; diff --git a/crates/vchordrq/src/operator.rs b/crates/vchordrq/src/operator.rs index 54395aae..3e1016d4 100644 --- a/crates/vchordrq/src/operator.rs +++ b/crates/vchordrq/src/operator.rs @@ -13,7 +13,7 @@ // Copyright (c) 2025 TensorChord Inc. use distance::Distance; -use index::accessor::{ +use index_accessor::{ Accessor1, Accessor2, ByteDistanceAccessor, DefaultWithDimension, DistanceAccessor, Dot, HalfbyteDistanceAccessor, L2S, RAccess, }; diff --git a/crates/vchordrq/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs index d56314e4..bfc8cb73 100644 --- a/crates/vchordrq/src/prewarm.rs +++ b/crates/vchordrq/src/prewarm.rs @@ -17,9 +17,9 @@ use crate::operator::Operator; use crate::tape::{by_directory, by_next}; use crate::tuples::*; use crate::{Opaque, centroids, tape}; -use index::accessor::FunctionalAccessor; use index::prefetcher::PrefetcherSequenceFamily; use index::relation::{Page, RelationRead}; +use index_accessor::FunctionalAccessor; use std::fmt::Write; pub fn prewarm<'b, R: RelationRead, O: Operator>( diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs index 53b1adb5..40905842 100644 --- a/crates/vchordrq/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -18,11 +18,11 @@ use crate::tuples::{MetaTuple, WithReader}; use crate::{RerankMethod, vectors}; use always_equal::AlwaysEqual; use distance::Distance; -use index::accessor::{Accessor2, DefaultWithDimension, LTryAccess}; use index::fetch::BorrowedIter; use index::packed::PackedRefMut; use index::prefetcher::Prefetcher; use index::relation::{Page, RelationRead}; +use index_accessor::{Accessor2, DefaultWithDimension, LTryAccess}; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::marker::PhantomData; diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index 4855c61f..9bd0409c 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -20,12 +20,12 @@ use crate::tuples::*; use crate::{Opaque, centroids, tape}; use always_equal::AlwaysEqual; use distance::Distance; -use index::accessor::{DefaultWithDimension, FunctionalAccessor, LAccess}; use index::bump::Bump; use index::fetch::BorrowedIter; use index::packed::{PackedRefMut4, PackedRefMut8}; use index::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; use index::relation::{Page, RelationRead}; +use index_accessor::{DefaultWithDimension, FunctionalAccessor, LAccess}; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZero; diff --git a/crates/vchordrq/src/tape.rs b/crates/vchordrq/src/tape.rs index aa135dd2..16502785 100644 --- a/crates/vchordrq/src/tape.rs +++ b/crates/vchordrq/src/tape.rs @@ -14,9 +14,9 @@ use crate::tuples::*; use crate::{Opaque, freepages}; -use index::accessor::Accessor1; use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; +use index_accessor::Accessor1; use std::marker::PhantomData; use std::num::NonZero; diff --git a/crates/vchordrq/src/vectors.rs b/crates/vchordrq/src/vectors.rs index 0767425c..d48150c0 100644 --- a/crates/vchordrq/src/vectors.rs +++ b/crates/vchordrq/src/vectors.rs @@ -15,8 +15,8 @@ use crate::operator::*; use crate::tuples::*; use crate::{Opaque, tape}; -use index::accessor::TryAccessor1; use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; +use index_accessor::TryAccessor1; use std::num::NonZero; use vector::VectorOwned; diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index d1c3c465..0349acc5 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -593,7 +593,7 @@ unsafe fn options( v: opfamily.vector_kind(), d: opfamily.distance_kind(), }; - // get indexing, segment, optimizing + // get indexing options let indexing_options = { let reloption = unsafe { (*index_relation).rd_options as *const Reloption }; let s = unsafe { Reloption::options(reloption, c"") }.to_string_lossy(); diff --git a/src/index/vchordg/dispatch.rs b/src/index/vchordg/dispatch.rs index 548a77c6..6f1dbff8 100644 --- a/src/index/vchordg/dispatch.rs +++ b/src/index/vchordg/dispatch.rs @@ -13,12 +13,12 @@ // Copyright (c) 2025 TensorChord Inc. use crate::index::vchordg::opclass::Opfamily; -use index::accessor::{Dot, L2S}; use index::fetch::Fetch; use index::prefetcher::*; use index::relation::{ Hints, Page, RelationPrefetch, RelationRead, RelationReadStream, RelationWrite, }; +use index_accessor::{Dot, L2S}; use simd::f16; use std::num::NonZero; use vchordg::operator::Op; diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs index 31582979..a5bebc42 100644 --- a/src/index/vchordg/scanners/default.rs +++ b/src/index/vchordg/scanners/default.rs @@ -20,9 +20,9 @@ use crate::index::vchordg::opclass::Opfamily; use crate::index::vchordg::scanners::SearchOptions; use crate::recorder::{Recorder, text}; use distance::Distance; -use index::accessor::{Dot, L2S}; use index::bump::Bump; use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadStream}; +use index_accessor::{Dot, L2S}; use simd::f16; use std::num::NonZero; use vchordg::operator::{self}; diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index a33434ca..1c504537 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -1245,7 +1245,7 @@ unsafe fn options( v: opfamily.vector_kind(), d: opfamily.distance_kind(), }; - // get indexing, segment, optimizing + // get indexing options let indexing_options = { let reloption = unsafe { (*index_relation).rd_options as *const Reloption }; let s = unsafe { Reloption::options(reloption, c"") }.to_string_lossy(); diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index d69d1328..e75f95c7 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -401,15 +401,6 @@ pub unsafe extern "C-unwind" fn aminsert( _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, _index_unchanged: bool, _index_info: *mut pgrx::pg_sys::IndexInfo, -) -> bool { - unsafe { aminsertinner(index_relation, values, is_null, heap_tid) } -} - -unsafe fn aminsertinner( - index_relation: pgrx::pg_sys::Relation, - values: *mut Datum, - is_null: *mut bool, - ctid: pgrx::pg_sys::ItemPointer, ) -> bool { struct RngChooser(T); impl InsertChooser for RngChooser { @@ -421,7 +412,7 @@ unsafe fn aminsertinner( let opfamily = unsafe { opfamily(index_relation) }; let index = unsafe { PostgresRelation::new(index_relation) }; let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; - let ctid = unsafe { ctid.read() }; + let ctid = unsafe { heap_tid.read() }; if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { for (vector, extra) in store { let key = ctid_to_key(ctid); diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs index c4a80fd7..550c1b61 100644 --- a/src/index/vchordrq/dispatch.rs +++ b/src/index/vchordrq/dispatch.rs @@ -14,13 +14,13 @@ use crate::index::vchordrq::build::{Normalize, Normalized}; use crate::index::vchordrq::opclass::Opfamily; -use index::accessor::{Dot, L2S}; use index::bump::Bump; use index::fetch::Fetch; use index::prefetcher::*; use index::relation::{ Hints, Page, RelationPrefetch, RelationRead, RelationReadStream, RelationWrite, }; +use index_accessor::{Dot, L2S}; use simd::f16; use std::collections::BinaryHeap; use std::num::NonZero; diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index ba94953b..9c6719ea 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -22,11 +22,11 @@ use crate::index::vchordrq::scanners::SearchOptions; use crate::recorder::{Recorder, text}; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; -use index::accessor::{Dot, L2S}; use index::bump::Bump; use index::packed::PackedRefMut4; use index::prefetcher::*; use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadStream}; +use index_accessor::{Dot, L2S}; use simd::f16; use std::num::NonZero; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 460d137a..c63e91b0 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -22,11 +22,11 @@ use crate::recorder::Recorder; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; use distance::Distance; -use index::accessor::Dot; use index::bump::Bump; use index::packed::PackedRefMut8; use index::prefetcher::*; use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadStream}; +use index_accessor::Dot; use simd::f16; use std::cmp::Reverse; use std::collections::BinaryHeap; From 04d2db8a0055a514dd68cbb76a554bb3a2767c68 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 16 Mar 2026 16:08:23 +0800 Subject: [PATCH 287/324] chore: remove unused deps and update deps (#443) Signed-off-by: usamoi --- Cargo.lock | 208 +++++++++++++++---------------------- Cargo.toml | 9 +- crates/index/Cargo.toml | 3 - crates/rabitq/Cargo.toml | 2 - crates/simd/Cargo.toml | 3 +- crates/vchordrq/Cargo.toml | 2 +- crates/xtask/Cargo.toml | 4 +- 7 files changed, 92 insertions(+), 139 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 19a975a9..903b939b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -48,9 +48,9 @@ dependencies = [ [[package]] name = "anstream" -version = "0.6.21" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43d5b281e737544384e969a5ccad3f1cdd24b48086a0fc1b2a5262a26b8f4f4a" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" dependencies = [ "anstyle", "anstyle-parse", @@ -63,15 +63,15 @@ dependencies = [ [[package]] name = "anstyle" -version = "1.0.13" +version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5192cca8006f1fd4f7237516f40fa183bb07f8fbdfedaa0036de5ea9b0b45e78" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" [[package]] name = "anstyle-parse" -version = "0.2.7" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e7644824f0aa2c7b9384579234ef10eb7efb6a0deb83f9630a49594dd9c15c2" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" dependencies = [ "utf8parse", ] @@ -98,9 +98,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.101" +version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f0e0fee31ef5ed1ba1316088939cea399010ed7731dba877ed44aeb407a75ea" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" [[package]] name = "arrayvec" @@ -135,9 +135,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.10.0" +version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" +checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" [[package]] name = "bitvec" @@ -153,9 +153,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.19.1" +version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dd9dc738b7a8311c7ade152424974d8115f2cdad61e8dab8dac9f2362298510" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" [[package]] name = "cargo_toml" @@ -175,9 +175,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.2.55" +version = "1.2.57" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47b26a0954ae34af09b50f0de26458fa95369a0d478d8236d3f93082b219bd29" +checksum = "7a0dd1ca384932ff3641c8718a02769f1698e7563dc6974ffd03346116310423" dependencies = [ "find-msvc-tools", "shlex", @@ -259,9 +259,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.58" +version = "4.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63be97961acde393029492ce0be7a1af7e323e6bae9511ebfac33751be5e6806" +checksum = "b193af5b67834b676abd72466a96c1024e6a6ad978a1f484bd90b85c94041351" dependencies = [ "clap_builder", "clap_derive", @@ -269,9 +269,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.58" +version = "4.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f13174bda5dfd69d7e947827e5af4b0f2f94a4a3ee92912fba07a66150f21e2" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" dependencies = [ "anstream", "anstyle", @@ -281,9 +281,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.5.55" +version = "4.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a92793da1a46a5f2a02a6f4c46c6496b28c43638adea8306fcb0caa1634f24e5" +checksum = "1110bd8a634a1ab8cb04345d8d878267d57c3cf1b38d91b71af6686408bbca6a" dependencies = [ "heck", "proc-macro2", @@ -293,9 +293,9 @@ dependencies = [ [[package]] name = "clap_lex" -version = "1.0.0" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a822ea5bc7590f9d40f1ba12c0dc3c2760f3482c6984db1573ad11031420831" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" [[package]] name = "codepage" @@ -308,9 +308,9 @@ dependencies = [ [[package]] name = "colorchoice" -version = "1.0.4" +version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" [[package]] name = "convert_case" @@ -508,12 +508,6 @@ dependencies = [ "syn", ] -[[package]] -name = "env_home" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe" - [[package]] name = "equivalent" version = "1.0.2" @@ -623,21 +617,9 @@ checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" [[package]] name = "getrandom" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" -dependencies = [ - "cfg-if", - "libc", - "r-efi", - "wasip2", -] - -[[package]] -name = "getrandom" -version = "0.4.1" +version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "139ef39800118c7683f2fd3c98c1b23c09ae076556b435f8e9064ae108aaeeec" +checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" dependencies = [ "cfg-if", "libc", @@ -839,10 +821,7 @@ dependencies = [ "always_equal", "bumpalo", "dary_heap", - "distance", - "rabitq", "small_iter", - "vector", "zerocopy", ] @@ -897,9 +876,9 @@ checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" [[package]] name = "js-sys" -version = "0.3.85" +version = "0.3.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c942ebf8e95485ca0d52d97da7c5a2c387d0e7f0ba4c35e93bfcaee045955b3" +checksum = "b49715b7073f385ba4bc528e5747d02e66cb39c6146efb66b781f131f0fb399c" dependencies = [ "once_cell", "wasm-bindgen", @@ -925,9 +904,9 @@ checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" [[package]] name = "libc" -version = "0.2.181" +version = "0.2.183" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "459427e2af2b9c839b132acb702a1c654d95e10f8c326bfc2ad11310e458b1c5" +checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" [[package]] name = "libloading" @@ -957,9 +936,9 @@ dependencies = [ [[package]] name = "libsqlite3-sys" -version = "0.36.0" +version = "0.37.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95b4103cffefa72eb8428cb6b47d6627161e51c2739fc5e3b734584157bc642a" +checksum = "b1f111c8c41e7c61a49cd34e44c7619462967221a6443b0ec299e0ac30cfb9b1" dependencies = [ "cc", "pkg-config", @@ -968,9 +947,9 @@ dependencies = [ [[package]] name = "linux-raw-sys" -version = "0.11.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" [[package]] name = "litemap" @@ -1054,9 +1033,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.21.3" +version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "once_cell_polyfill" @@ -1072,9 +1051,9 @@ checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" [[package]] name = "owo-colors" -version = "4.2.3" +version = "4.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c6901729fa79e91a0913333229e9ca5dc725089d1c363b2f4b4760709dc4a52" +checksum = "d211803b9b6b570f68772237e415a029d5a50c65d382910b879fb19d3271f94d" dependencies = [ "supports-color", ] @@ -1237,18 +1216,18 @@ dependencies = [ [[package]] name = "pin-project" -version = "1.1.10" +version = "1.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677f1add503faace112b9f1373e43e9e054bfdd22ff1a63c1bc485eaec6a6a8a" +checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517" dependencies = [ "pin-project-internal", ] [[package]] name = "pin-project-internal" -version = "1.1.10" +version = "1.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e918e4ff8c4549eb882f14b3a4bc8c8bc93de829416eacf579f1207a8fbf861" +checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6" dependencies = [ "proc-macro2", "quote", @@ -1350,18 +1329,18 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.44" +version = "1.0.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" dependencies = [ "proc-macro2", ] [[package]] name = "r-efi" -version = "5.3.0" +version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "rabitq" @@ -1386,7 +1365,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" dependencies = [ "chacha20", - "getrandom 0.4.1", + "getrandom", "rand_core 0.10.0", ] @@ -1457,9 +1436,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.9" +version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a96887878f22d7bad8a3b6dc5b7440e0ada9a245242924394987b21cf2210a4c" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" [[package]] name = "rsqlite-vfs" @@ -1473,9 +1452,9 @@ dependencies = [ [[package]] name = "rusqlite" -version = "0.38.0" +version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1c93dd1c9683b438c392c492109cb702b8090b2bfc8fed6f6e4eb4523f17af3" +checksum = "a0d2b0146dd9661bf67bb107c0bb2a55064d556eeb3fc314151b957f313bcd4e" dependencies = [ "bitflags", "fallible-iterator", @@ -1494,9 +1473,9 @@ checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" [[package]] name = "rustix" -version = "1.1.3" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "146c9e247ccc180c1f61615433868c99f3de3ae256a30a43b49f67c2d9171f34" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ "bitflags", "errno", @@ -1625,7 +1604,6 @@ dependencies = [ "rand", "seq-macro", "simd_macros", - "which", "zerocopy", ] @@ -1689,9 +1667,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.115" +version = "2.0.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e614ed320ac28113fa64972c4262d5dbc89deacdfd00c34a3e4cea073243c12" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" dependencies = [ "proc-macro2", "quote", @@ -1723,12 +1701,12 @@ checksum = "591ef38edfb78ca4771ee32cf494cb8771944bee237a9b91fc9c1424ac4b777b" [[package]] name = "tempfile" -version = "3.25.0" +version = "3.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0136791f7c95b1f6dd99f9cc786b91bb81c3800b639b3478e561ddb7be95e5f1" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.1", + "getrandom", "once_cell", "rustix", "windows-sys", @@ -1791,9 +1769,9 @@ dependencies = [ [[package]] name = "toml" -version = "1.0.0+spec-1.1.0" +version = "1.0.6+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1d7e18e3dd1d31e0ee5e863a8091ffec2fcc271636586042452b656a22c8ee1" +checksum = "399b1124a3c9e16766831c6bba21e50192572cdd98706ea114f9502509686ffc" dependencies = [ "indexmap", "serde_core", @@ -1824,9 +1802,9 @@ dependencies = [ [[package]] name = "toml_parser" -version = "1.0.7+spec-1.1.0" +version = "1.0.9+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "247eaa3197818b831697600aadf81514e577e0cba5eab10f7e064e78ae154df1" +checksum = "702d4415e08923e7e1ef96cd5727c0dfed80b4d2fa25db9647fe5eb6f7c5a4c4" dependencies = [ "winnow", ] @@ -1851,9 +1829,9 @@ checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] name = "unicode-ident" -version = "1.0.23" +version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "537dd038a89878be9b64dd4bd1b260315c1bb94f4d784956b81e27a088d9a09e" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" [[package]] name = "unicode-segmentation" @@ -1899,11 +1877,11 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.20.0" +version = "1.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee48d38b119b0cd71fe4141b30f5ba9c7c5d9f4e7a3a8b4a674e4b6ef789976f" +checksum = "a68d3c8f01c0cfa54a75291d83601161799e4a89a39e0929f4b0354d88757a37" dependencies = [ - "getrandom 0.3.4", + "getrandom", "js-sys", "wasm-bindgen", ] @@ -1964,8 +1942,7 @@ dependencies = [ "seq-macro", "serde", "simd", - "small_iter", - "toml 1.0.0+spec-1.1.0", + "toml 1.0.6+spec-1.1.0", "validator", "vchordg", "vchordrq", @@ -2055,9 +2032,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.108" +version = "0.2.114" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64024a30ec1e37399cf85a7ffefebdb72205ca1c972291c51512360d90bd8566" +checksum = "6532f9a5c1ece3798cb1c2cfdba640b9b3ba884f5db45973a6f442510a87d38e" dependencies = [ "cfg-if", "once_cell", @@ -2068,9 +2045,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.108" +version = "0.2.114" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "008b239d9c740232e71bd39e8ef6429d27097518b6b30bdf9086833bd5b6d608" +checksum = "18a2d50fcf105fb33bb15f00e7a77b772945a2ee45dcf454961fd843e74c18e6" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -2078,9 +2055,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.108" +version = "0.2.114" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5256bae2d58f54820e6490f9839c49780dff84c65aeab9e772f15d5f0e913a55" +checksum = "03ce4caeaac547cdf713d280eda22a730824dd11e6b8c3ca9e42247b25c631e3" dependencies = [ "bumpalo", "proc-macro2", @@ -2091,9 +2068,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.108" +version = "0.2.114" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f01b580c9ac74c8d8f0c0e4afb04eeef2acf145458e52c03845ee9cd23e3d12" +checksum = "75a326b8c223ee17883a4251907455a2431acc2791c98c26279376490c378c16" dependencies = [ "unicode-ident", ] @@ -2143,25 +2120,14 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.85" +version = "0.3.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "312e32e551d92129218ea9a2452120f4aabc03529ef03e4d0d82fb2780608598" +checksum = "854ba17bb104abfb26ba36da9729addc7ce7f06f5c0f90f3c391f8461cca21f9" dependencies = [ "js-sys", "wasm-bindgen", ] -[[package]] -name = "which" -version = "8.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3fabb953106c3c8eea8306e4393700d7657561cb43122571b172bbfb7c7ba1d" -dependencies = [ - "env_home", - "rustix", - "winsafe", -] - [[package]] name = "winapi" version = "0.3.9" @@ -2210,15 +2176,9 @@ dependencies = [ [[package]] name = "winnow" -version = "0.7.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a5364e9d77fcdeeaa6062ced926ee3381faa2ee02d3eb83a5c27a8825540829" - -[[package]] -name = "winsafe" -version = "0.0.19" +version = "0.7.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" [[package]] name = "wit-bindgen" @@ -2370,18 +2330,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.39" +version = "0.8.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db6d35d663eadb6c932438e763b262fe1a70987f9ae936e60158176d710cae4a" +checksum = "f2578b716f8a7a858b7f02d5bd870c14bf4ddbbcf3a4c05414ba6503640505e3" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.39" +version = "0.8.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4122cd3169e94605190e77839c9a40d40ed048d305bfdc146e7df40ab0f3e517" +checksum = "7e6cc098ea4d3bd6246687de65af3f920c430e236bee1e3bf2e441463f08a02f" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index 99b9aebc..e7dcd96f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,7 +29,6 @@ index_accessor = { path = "./crates/index_accessor" } k_means = { path = "./crates/k_means" } rabitq = { path = "./crates/rabitq" } simd = { path = "./crates/simd" } -small_iter = { path = "./crates/small_iter" } vchordg = { path = "./crates/vchordg" } vchordrq = { path = "./crates/vchordrq" } vector = { path = "./crates/vector" } @@ -44,10 +43,10 @@ pgrx = { version = "=0.17.0", default-features = false, features = ["cshim"] } pgrx-catalog = "0.3.2" rand.workspace = true rayon.workspace = true -rusqlite = { version = "0.38.0", features = ["bundled"] } +rusqlite = { version = "0.39.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true -toml = "1.0.0" +toml = "1.0.6" validator.workspace = true wyhash.workspace = true zerocopy.workspace = true @@ -67,7 +66,7 @@ version = "0.0.0" edition = "2024" [workspace.dependencies] -bumpalo = "3.19.1" +bumpalo = "3.20.2" dary_heap = "0.3.8" paste = "1.0.15" rand = "0.10.0" @@ -77,7 +76,7 @@ seq-macro = "0.3.6" serde = { version = "1.0.228", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } wyhash = "0.6.0" -zerocopy = { version = "0.8.39", features = ["derive"] } +zerocopy = { version = "0.8.42", features = ["derive"] } [workspace.lints] # complexity diff --git a/crates/index/Cargo.toml b/crates/index/Cargo.toml index 7ade4785..6d21f44c 100644 --- a/crates/index/Cargo.toml +++ b/crates/index/Cargo.toml @@ -6,10 +6,7 @@ publish = false [dependencies] always_equal = { path = "../always_equal" } -distance = { path = "../distance" } -rabitq = { path = "../rabitq" } small_iter = { path = "../small_iter" } -vector = { path = "../vector" } bumpalo.workspace = true dary_heap.workspace = true diff --git a/crates/rabitq/Cargo.toml b/crates/rabitq/Cargo.toml index dc6e6c91..72e65fc6 100644 --- a/crates/rabitq/Cargo.toml +++ b/crates/rabitq/Cargo.toml @@ -7,8 +7,6 @@ publish = false [dependencies] simd = { path = "../simd" } -rand.workspace = true -rand_chacha.workspace = true zerocopy.workspace = true [build-dependencies] diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 6309ab1f..e6d87c71 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -26,8 +26,7 @@ criterion = "0.8.2" rand.workspace = true [build-dependencies] -cc = "1.2.55" -which = "8.0.0" +cc = "1.2.57" [lints] workspace = true diff --git a/crates/vchordrq/Cargo.toml b/crates/vchordrq/Cargo.toml index ebac0e59..531118a4 100644 --- a/crates/vchordrq/Cargo.toml +++ b/crates/vchordrq/Cargo.toml @@ -13,7 +13,7 @@ rabitq = { path = "../rabitq" } simd = { path = "../simd" } vector = { path = "../vector" } -pin-project = "1.1.10" +pin-project = "1.1.11" serde.workspace = true validator.workspace = true zerocopy.workspace = true diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index 8b092d7e..83f9522d 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -5,13 +5,13 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.5.58", features = ["derive", "env"] } +clap = { version = "4.6.0", features = ["derive", "env"] } object = { version = "0.38.1", features = ["read", "wasm"] } serde.workspace = true serde_json = "1.0.149" shlex = "1.3.0" target-triple = "1.0.0" -tempfile = "3.25.0" +tempfile = "3.27.0" [lints] workspace = true From 729d33c86daab45cc9b8aa78d6497c1fafe386ca Mon Sep 17 00:00:00 2001 From: usamoi Date: Thu, 19 Mar 2026 17:53:35 +0800 Subject: [PATCH 288/324] chore: update license header and bundle copyright in the debian package (#445) Signed-off-by: usamoi --- .github/workflows/check.yml | 2 +- .github/workflows/release.yml | 24 ++++++++++++++++--- build.rs | 2 +- crates/always_equal/src/lib.rs | 2 +- crates/distance/src/lib.rs | 2 +- crates/feistel/src/lib.rs | 2 +- crates/index/src/accessor.rs | 2 +- crates/index/src/bump.rs | 2 +- crates/index/src/fetch.rs | 2 +- crates/index/src/lib.rs | 2 +- crates/index/src/packed.rs | 2 +- crates/index/src/prefetcher.rs | 2 +- crates/index/src/relation.rs | 2 +- crates/index/src/tuples.rs | 2 +- crates/k_means/src/flat.rs | 2 +- crates/k_means/src/hierarchical.rs | 2 +- crates/k_means/src/index.rs | 2 +- crates/k_means/src/lib.rs | 2 +- crates/k_means/src/quick.rs | 2 +- crates/k_means/src/rabitq.rs | 2 +- crates/k_means/src/square.rs | 2 +- crates/rabitq/build.rs | 2 +- crates/rabitq/src/bit.rs | 2 +- crates/rabitq/src/bits.rs | 2 +- crates/rabitq/src/byte.rs | 2 +- crates/rabitq/src/extended.rs | 2 +- crates/rabitq/src/halfbyte.rs | 2 +- crates/rabitq/src/lib.rs | 2 +- crates/rabitq/src/packing.rs | 2 +- crates/rabitq/src/rotate.rs | 2 +- crates/simd/benches/bench.rs | 2 +- crates/simd/build.rs | 2 +- crates/simd/cshim/aarch64_byte.c | 2 +- crates/simd/cshim/aarch64_halfbyte.c | 2 +- crates/simd/cshim/aarch64_sve_fp16.c | 2 +- crates/simd/cshim/aarch64_sve_fp32.c | 2 +- crates/simd/src/aligned.rs | 2 +- crates/simd/src/bit.rs | 2 +- crates/simd/src/byte.rs | 2 +- crates/simd/src/emulate.rs | 2 +- crates/simd/src/fast_scan.rs | 2 +- crates/simd/src/fht.rs | 2 +- crates/simd/src/floating_f16.rs | 2 +- crates/simd/src/floating_f32.rs | 2 +- crates/simd/src/halfbyte.rs | 2 +- crates/simd/src/lib.rs | 2 +- crates/simd/src/quantize.rs | 2 +- crates/simd/src/rotate.rs | 2 +- crates/simd_macros/src/lib.rs | 2 +- crates/simd_macros/src/target.rs | 2 +- crates/small_iter/src/borrowed.rs | 2 +- crates/small_iter/src/lib.rs | 2 +- crates/small_iter/src/stack.rs | 2 +- crates/vchordg/src/build.rs | 2 +- crates/vchordg/src/bulkdelete.rs | 2 +- crates/vchordg/src/candidates.rs | 2 +- crates/vchordg/src/insert.rs | 2 +- crates/vchordg/src/lib.rs | 2 +- crates/vchordg/src/maintain.rs | 2 +- crates/vchordg/src/operator.rs | 2 +- crates/vchordg/src/prewarm.rs | 2 +- crates/vchordg/src/prune.rs | 2 +- crates/vchordg/src/results.rs | 2 +- crates/vchordg/src/search.rs | 2 +- crates/vchordg/src/tuples.rs | 2 +- crates/vchordg/src/types.rs | 2 +- crates/vchordg/src/vectors.rs | 2 +- crates/vchordg/src/visited.rs | 2 +- crates/vchordrq/src/build.rs | 2 +- crates/vchordrq/src/bulkdelete.rs | 2 +- crates/vchordrq/src/cache.rs | 2 +- crates/vchordrq/src/centroids.rs | 2 +- .../vchordrq/src/closure_lifetime_binder.rs | 2 +- crates/vchordrq/src/consume.rs | 2 +- crates/vchordrq/src/cost.rs | 2 +- crates/vchordrq/src/fast_heap.rs | 2 +- crates/vchordrq/src/freepages.rs | 2 +- crates/vchordrq/src/insert.rs | 2 +- crates/vchordrq/src/lib.rs | 2 +- crates/vchordrq/src/linked_vec.rs | 2 +- crates/vchordrq/src/maintain.rs | 2 +- crates/vchordrq/src/operator.rs | 2 +- crates/vchordrq/src/prewarm.rs | 2 +- crates/vchordrq/src/rerank.rs | 2 +- crates/vchordrq/src/search.rs | 2 +- crates/vchordrq/src/tape.rs | 2 +- crates/vchordrq/src/tape_writer.rs | 2 +- crates/vchordrq/src/tuples.rs | 2 +- crates/vchordrq/src/types.rs | 2 +- crates/vchordrq/src/vectors.rs | 2 +- crates/vector/src/bvect.rs | 2 +- crates/vector/src/lib.rs | 2 +- crates/vector/src/rabitq4.rs | 2 +- crates/vector/src/rabitq8.rs | 2 +- crates/vector/src/svect.rs | 2 +- crates/vector/src/vect.rs | 2 +- crates/xtask/src/main.rs | 2 +- src/bin/pgrx_embed.rs | 2 +- src/datatype/binary_rabitq4.rs | 2 +- src/datatype/binary_rabitq8.rs | 2 +- src/datatype/functions_rabitq4.rs | 2 +- src/datatype/functions_rabitq8.rs | 2 +- src/datatype/memory_halfvec.rs | 2 +- src/datatype/memory_rabitq4.rs | 2 +- src/datatype/memory_rabitq8.rs | 2 +- src/datatype/memory_vector.rs | 2 +- src/datatype/mod.rs | 2 +- src/datatype/operators_halfvec.rs | 2 +- src/datatype/operators_rabitq4.rs | 2 +- src/datatype/operators_rabitq8.rs | 2 +- src/datatype/operators_vector.rs | 2 +- src/datatype/text_rabitq4.rs | 2 +- src/datatype/text_rabitq8.rs | 2 +- src/datatype/typmod.rs | 2 +- src/datatype/typmod_rabitq4.rs | 2 +- src/datatype/typmod_rabitq8.rs | 2 +- src/index/fetcher.rs | 2 +- src/index/functions.rs | 2 +- src/index/gucs.rs | 2 +- src/index/hook.rs | 2 +- src/index/mod.rs | 2 +- src/index/opclass.rs | 2 +- src/index/sample.rs | 2 +- src/index/scanners.rs | 2 +- src/index/storage.rs | 2 +- src/index/storage/buffered.rs | 2 +- src/index/traverse.rs | 2 +- src/index/vchordg/am/am_build.rs | 2 +- src/index/vchordg/am/mod.rs | 2 +- src/index/vchordg/dispatch.rs | 2 +- src/index/vchordg/mod.rs | 2 +- src/index/vchordg/opclass.rs | 2 +- src/index/vchordg/scanners/default.rs | 2 +- src/index/vchordg/scanners/mod.rs | 2 +- src/index/vchordg/types.rs | 2 +- src/index/vchordrq/am/am_build.rs | 2 +- src/index/vchordrq/am/am_vacuumcleanup.rs | 2 +- src/index/vchordrq/am/mod.rs | 2 +- src/index/vchordrq/build.rs | 2 +- src/index/vchordrq/dispatch.rs | 2 +- src/index/vchordrq/filter.rs | 2 +- src/index/vchordrq/mod.rs | 2 +- src/index/vchordrq/opclass.rs | 2 +- src/index/vchordrq/scanners/default.rs | 2 +- src/index/vchordrq/scanners/maxsim.rs | 2 +- src/index/vchordrq/scanners/mod.rs | 2 +- src/index/vchordrq/types.rs | 2 +- src/lib.rs | 2 +- src/recorder/hook.rs | 2 +- src/recorder/mod.rs | 2 +- src/recorder/text.rs | 2 +- src/recorder/types.rs | 2 +- src/recorder/worker.rs | 2 +- src/upgrade.rs | 2 +- 154 files changed, 174 insertions(+), 156 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index d92607cb..52fbd9e7 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -58,7 +58,7 @@ jobs: regarding the licenses, please contact us at: vectorchord-inquiry@tensorchord.ai - Copyright (c) 2025 TensorChord Inc. + Copyright (c) 2025-2026 TensorChord Inc. EOF ) COUNT=$(echo "$HEADER" | wc -l) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 775192a2..214776ba 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -98,6 +98,22 @@ jobs: VERSION: ${{ matrix.version }} run: | make DESTDIR="~/debian" install + mkdir -p ~/debian/usr/share/doc/postgresql-${VERSION}-vchord/ + echo "Format: https://www.debian.org/doc/packaging-manuals/copyright-format/1.0/ + Upstream-Name: VectorChord + Source: https://github.com/tensorchord/VectorChord + + Files: * + Copyright: Copyright (c) 2025-2026 TensorChord Inc. + License: AGPL-3.0-only OR Elastic-2.0 + + License: AGPL-3.0-only + $(cat ./licenses/LICENSE.AGPLv3 | sed 's/^$/./' | sed 's/^/ /') + + License: Elastic-2.0 + $(cat ./licenses/LICENSE.ELv2 | sed 's/^$/./' | sed 's/^/ /') + " \ + >> ~/debian/usr/share/doc/postgresql-${VERSION}-vchord/copyright ARCH=$(uname -m) PLATFORM=$(dpkg --print-architecture) mkdir -p ~/debian/DEBIAN @@ -108,9 +124,11 @@ jobs: Priority: optional Architecture: ${PLATFORM} Maintainer: Tensorchord - Description: Vector database plugin for Postgres, written in Rust, specifically designed for LLM - Homepage: https://vectorchord.ai/ - License: AGPL-3.0-only OR Elastic-2.0" \ + Description: Scalable, fast, and disk-friendly vector search in Postgres + VectorChord is a PostgreSQL extension designed for scalable, + high-performance, and disk-efficient vector similarity search. It + supports L2 distance, inner product, cosine distance and maxsim. + Homepage: https://vectorchord.ai/" \ > ~/debian/DEBIAN/control (cd ~/debian && find usr -type f -print0 | xargs -0 md5sum) > ~/debian/DEBIAN/md5sums dpkg-deb --root-owner-group -Zxz --build ~/debian ~/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb diff --git a/build.rs b/build.rs index bc506af2..642a488a 100644 --- a/build.rs +++ b/build.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use std::collections::HashMap; use std::env::{var, var_os}; diff --git a/crates/always_equal/src/lib.rs b/crates/always_equal/src/lib.rs index ff4b51fd..f95d470a 100644 --- a/crates/always_equal/src/lib.rs +++ b/crates/always_equal/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use std::cmp::Ordering; use std::hash::Hash; diff --git a/crates/distance/src/lib.rs b/crates/distance/src/lib.rs index 557f84e6..cdaca6d8 100644 --- a/crates/distance/src/lib.rs +++ b/crates/distance/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; diff --git a/crates/feistel/src/lib.rs b/crates/feistel/src/lib.rs index 1098e7b3..b7987f54 100644 --- a/crates/feistel/src/lib.rs +++ b/crates/feistel/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub fn feistel(width: u32, x: I, round: u32, secret: impl Fn(u32, I) -> I) -> I where diff --git a/crates/index/src/accessor.rs b/crates/index/src/accessor.rs index d2046a3b..e10a41d2 100644 --- a/crates/index/src/accessor.rs +++ b/crates/index/src/accessor.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use distance::Distance; use rabitq::byte::CodeMetadata; diff --git a/crates/index/src/bump.rs b/crates/index/src/bump.rs index e80e242e..8416a686 100644 --- a/crates/index/src/bump.rs +++ b/crates/index/src/bump.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub trait Bump: 'static { #[allow(clippy::mut_from_ref)] diff --git a/crates/index/src/fetch.rs b/crates/index/src/fetch.rs index 685c90dc..a8ec1137 100644 --- a/crates/index/src/fetch.rs +++ b/crates/index/src/fetch.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::packed::PackedRefMut; use always_equal::AlwaysEqual; diff --git a/crates/index/src/lib.rs b/crates/index/src/lib.rs index 6ec70b6b..7513f10c 100644 --- a/crates/index/src/lib.rs +++ b/crates/index/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. // pub mod accessor; pub mod bump; diff --git a/crates/index/src/packed.rs b/crates/index/src/packed.rs index 65d93db7..f8ee67c9 100644 --- a/crates/index/src/packed.rs +++ b/crates/index/src/packed.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub trait PackedRefMut { type T; diff --git a/crates/index/src/prefetcher.rs b/crates/index/src/prefetcher.rs index d0bd662a..626d4a64 100644 --- a/crates/index/src/prefetcher.rs +++ b/crates/index/src/prefetcher.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::fetch::Fetch; use crate::relation::{ diff --git a/crates/index/src/relation.rs b/crates/index/src/relation.rs index e6023542..d6c4e832 100644 --- a/crates/index/src/relation.rs +++ b/crates/index/src/relation.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::fetch::Fetch; use std::ops::{Deref, DerefMut}; diff --git a/crates/index/src/tuples.rs b/crates/index/src/tuples.rs index 76dd51fc..0934ce7e 100644 --- a/crates/index/src/tuples.rs +++ b/crates/index/src/tuples.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use std::marker::PhantomData; use std::num::NonZero; diff --git a/crates/k_means/src/flat.rs b/crates/k_means/src/flat.rs index 0d032b6c..3ebf88ec 100644 --- a/crates/k_means/src/flat.rs +++ b/crates/k_means/src/flat.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::{flat_index as prefect_index, flat_index as index}; use crate::square::{Square, SquareMut}; diff --git a/crates/k_means/src/hierarchical.rs b/crates/k_means/src/hierarchical.rs index 05d9bbf4..7531db64 100644 --- a/crates/k_means/src/hierarchical.rs +++ b/crates/k_means/src/hierarchical.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::square::{Square, SquareMut}; use crate::{KMeans, This}; diff --git a/crates/k_means/src/index.rs b/crates/k_means/src/index.rs index 0bb59100..6da6524d 100644 --- a/crates/k_means/src/index.rs +++ b/crates/k_means/src/index.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::This; use crate::square::{Square, SquareMut}; diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs index 0c8a5f3d..510a62e7 100644 --- a/crates/k_means/src/lib.rs +++ b/crates/k_means/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod flat; mod hierarchical; diff --git a/crates/k_means/src/quick.rs b/crates/k_means/src/quick.rs index 72fd8fca..ac69da6a 100644 --- a/crates/k_means/src/quick.rs +++ b/crates/k_means/src/quick.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::KMeans; use crate::index::flat_index as prefect_index; diff --git a/crates/k_means/src/rabitq.rs b/crates/k_means/src/rabitq.rs index a40ff34b..61999979 100644 --- a/crates/k_means/src/rabitq.rs +++ b/crates/k_means/src/rabitq.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::{flat_index as prefect_index, rabitq_index as index}; use crate::square::{Square, SquareMut}; diff --git a/crates/k_means/src/square.rs b/crates/k_means/src/square.rs index 3fa52008..769f1fde 100644 --- a/crates/k_means/src/square.rs +++ b/crates/k_means/src/square.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #[derive(Debug, Clone)] pub struct Square { diff --git a/crates/rabitq/build.rs b/crates/rabitq/build.rs index 18e81d0a..007f32ef 100644 --- a/crates/rabitq/build.rs +++ b/crates/rabitq/build.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use std::env::var; use std::error::Error; diff --git a/crates/rabitq/src/bit.rs b/crates/rabitq/src/bit.rs index efeef07b..900d4528 100644 --- a/crates/rabitq/src/bit.rs +++ b/crates/rabitq/src/bit.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use binary::BinaryLut; use block::BlockLut; diff --git a/crates/rabitq/src/bits.rs b/crates/rabitq/src/bits.rs index 456f0f12..a1853916 100644 --- a/crates/rabitq/src/bits.rs +++ b/crates/rabitq/src/bits.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub use crate::extended::{Code, CodeMetadata}; diff --git a/crates/rabitq/src/byte.rs b/crates/rabitq/src/byte.rs index 01d9a71e..71681c71 100644 --- a/crates/rabitq/src/byte.rs +++ b/crates/rabitq/src/byte.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub use crate::extended::{Code, CodeMetadata}; diff --git a/crates/rabitq/src/extended.rs b/crates/rabitq/src/extended.rs index 1c1759ce..789211f0 100644 --- a/crates/rabitq/src/extended.rs +++ b/crates/rabitq/src/extended.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use simd::Floating; diff --git a/crates/rabitq/src/halfbyte.rs b/crates/rabitq/src/halfbyte.rs index f60c6420..725b7d4f 100644 --- a/crates/rabitq/src/halfbyte.rs +++ b/crates/rabitq/src/halfbyte.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub use crate::extended::{Code, CodeMetadata}; diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs index 0aae7682..cc412b7f 100644 --- a/crates/rabitq/src/lib.rs +++ b/crates/rabitq/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub mod bit; pub mod bits; diff --git a/crates/rabitq/src/packing.rs b/crates/rabitq/src/packing.rs index 79c8e2a4..ce21115d 100644 --- a/crates/rabitq/src/packing.rs +++ b/crates/rabitq/src/packing.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub fn pack(x: [&[u8]; 32]) -> Vec<[u8; 16]> { let n = { diff --git a/crates/rabitq/src/rotate.rs b/crates/rabitq/src/rotate.rs index fb54a96b..93806163 100644 --- a/crates/rabitq/src/rotate.rs +++ b/crates/rabitq/src/rotate.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. const BITS: &[u8; 262144] = include_bytes!(concat!(env!("OUT_DIR"), "/bits")); diff --git a/crates/simd/benches/bench.rs b/crates/simd/benches/bench.rs index 200d9919..db9ba0df 100644 --- a/crates/simd/benches/bench.rs +++ b/crates/simd/benches/bench.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #![allow(unsafe_code)] diff --git a/crates/simd/build.rs b/crates/simd/build.rs index 338f9e17..57a73c01 100644 --- a/crates/simd/build.rs +++ b/crates/simd/build.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use std::env::var; use std::error::Error; diff --git a/crates/simd/cshim/aarch64_byte.c b/crates/simd/cshim/aarch64_byte.c index ddd3818e..a9519510 100644 --- a/crates/simd/cshim/aarch64_byte.c +++ b/crates/simd/cshim/aarch64_byte.c @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #if defined(__clang__) #if !(__clang_major__ >= 16) diff --git a/crates/simd/cshim/aarch64_halfbyte.c b/crates/simd/cshim/aarch64_halfbyte.c index 53026e0a..3b0faa07 100644 --- a/crates/simd/cshim/aarch64_halfbyte.c +++ b/crates/simd/cshim/aarch64_halfbyte.c @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #if defined(__clang__) #if !(__clang_major__ >= 16) diff --git a/crates/simd/cshim/aarch64_sve_fp16.c b/crates/simd/cshim/aarch64_sve_fp16.c index ba1023cd..064ba4be 100644 --- a/crates/simd/cshim/aarch64_sve_fp16.c +++ b/crates/simd/cshim/aarch64_sve_fp16.c @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #if defined(__clang__) #if !(__clang_major__ >= 16) diff --git a/crates/simd/cshim/aarch64_sve_fp32.c b/crates/simd/cshim/aarch64_sve_fp32.c index 15e3de6e..4ceadeaf 100644 --- a/crates/simd/cshim/aarch64_sve_fp32.c +++ b/crates/simd/cshim/aarch64_sve_fp32.c @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #if defined(__clang__) #if !(__clang_major__ >= 16) diff --git a/crates/simd/src/aligned.rs b/crates/simd/src/aligned.rs index e8891bfa..b06b1b4c 100644 --- a/crates/simd/src/aligned.rs +++ b/crates/simd/src/aligned.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #[allow(dead_code)] #[derive(Debug, Clone, Copy)] diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs index 24541a8a..d0160b25 100644 --- a/crates/simd/src/bit.rs +++ b/crates/simd/src/bit.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #[inline(always)] pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { diff --git a/crates/simd/src/byte.rs b/crates/simd/src/byte.rs index 444f9ca6..0c7c4bed 100644 --- a/crates/simd/src/byte.rs +++ b/crates/simd/src/byte.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #[cfg_attr(feature = "internal", simd_macros::public)] mod reduce_sum_of_xy { diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs index e710ab6a..1647ee2d 100644 --- a/crates/simd/src/emulate.rs +++ b/crates/simd/src/emulate.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #[allow(unused_macros)] macro_rules! partial_load { diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index 22de0496..06502460 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. /* diff --git a/crates/simd/src/fht.rs b/crates/simd/src/fht.rs index d99a671b..2fec563e 100644 --- a/crates/simd/src/fht.rs +++ b/crates/simd/src/fht.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #[inline(always)] fn basic_1(x: &mut [f32]) { diff --git a/crates/simd/src/floating_f16.rs b/crates/simd/src/floating_f16.rs index addc17dc..1707955d 100644 --- a/crates/simd/src/floating_f16.rs +++ b/crates/simd/src/floating_f16.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::{F16, Floating, f16}; diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs index 6bcb476f..b298c107 100644 --- a/crates/simd/src/floating_f32.rs +++ b/crates/simd/src/floating_f32.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::Floating; diff --git a/crates/simd/src/halfbyte.rs b/crates/simd/src/halfbyte.rs index 03a9ae49..e1aa0163 100644 --- a/crates/simd/src/halfbyte.rs +++ b/crates/simd/src/halfbyte.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #[cfg_attr(feature = "internal", simd_macros::public)] mod reduce_sum_of_xy { diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 4ca9340e..601da07d 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #![allow(unsafe_code)] #![cfg_attr(feature = "nightly_f16", feature(f16))] diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs index 98d54b81..40df3b39 100644 --- a/crates/simd/src/quantize.rs +++ b/crates/simd/src/quantize.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod mul_add_round { #[inline] diff --git a/crates/simd/src/rotate.rs b/crates/simd/src/rotate.rs index 1a8304c9..5ed0de39 100644 --- a/crates/simd/src/rotate.rs +++ b/crates/simd/src/rotate.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod givens { #[crate::multiversion( diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs index 9152fd6a..ed6f436b 100644 --- a/crates/simd_macros/src/lib.rs +++ b/crates/simd_macros/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod target; diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs index 5903f72a..0f3c60c7 100644 --- a/crates/simd_macros/src/target.rs +++ b/crates/simd_macros/src/target.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub struct TargetCpu { pub target_cpu: &'static str, diff --git a/crates/small_iter/src/borrowed.rs b/crates/small_iter/src/borrowed.rs index 3e76300e..ff90cc8d 100644 --- a/crates/small_iter/src/borrowed.rs +++ b/crates/small_iter/src/borrowed.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::stack::StackIter; use std::marker::PhantomData; diff --git a/crates/small_iter/src/lib.rs b/crates/small_iter/src/lib.rs index 886062e5..60b4aa34 100644 --- a/crates/small_iter/src/lib.rs +++ b/crates/small_iter/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub mod borrowed; pub mod stack; diff --git a/crates/small_iter/src/stack.rs b/crates/small_iter/src/stack.rs index 6e78ecd1..1bab4a4a 100644 --- a/crates/small_iter/src/stack.rs +++ b/crates/small_iter/src/stack.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use std::mem::MaybeUninit; diff --git a/crates/vchordg/src/build.rs b/crates/vchordg/src/build.rs index beaa093b..54b69d2e 100644 --- a/crates/vchordg/src/build.rs +++ b/crates/vchordg/src/build.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::Opaque; use crate::operator::Operator; diff --git a/crates/vchordg/src/bulkdelete.rs b/crates/vchordg/src/bulkdelete.rs index 3cbfffbb..2e619837 100644 --- a/crates/vchordg/src/bulkdelete.rs +++ b/crates/vchordg/src/bulkdelete.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::Opaque; use crate::operator::Operator; diff --git a/crates/vchordg/src/candidates.rs b/crates/vchordg/src/candidates.rs index 2a30b5a9..2c32019f 100644 --- a/crates/vchordg/src/candidates.rs +++ b/crates/vchordg/src/candidates.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use index::fetch::Fetch; use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs index a133ae6d..ff11ee9f 100644 --- a/crates/vchordg/src/insert.rs +++ b/crates/vchordg/src/insert.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::Opaque; use crate::candidates::Candidates; diff --git a/crates/vchordg/src/lib.rs b/crates/vchordg/src/lib.rs index f84de74a..a16936c0 100644 --- a/crates/vchordg/src/lib.rs +++ b/crates/vchordg/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod build; mod bulkdelete; diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs index 6d43daad..db3303cc 100644 --- a/crates/vchordg/src/maintain.rs +++ b/crates/vchordg/src/maintain.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::Opaque; use crate::operator::{CloneAccessor, Operator}; diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs index bc279567..d5961c7d 100644 --- a/crates/vchordg/src/operator.rs +++ b/crates/vchordg/src/operator.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::types::DistanceKind; use distance::Distance; diff --git a/crates/vchordg/src/prewarm.rs b/crates/vchordg/src/prewarm.rs index 6c0cda0b..9ccc076f 100644 --- a/crates/vchordg/src/prewarm.rs +++ b/crates/vchordg/src/prewarm.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::Opaque; use crate::operator::Operator; diff --git a/crates/vchordg/src/prune.rs b/crates/vchordg/src/prune.rs index b3f1b552..138ebc1f 100644 --- a/crates/vchordg/src/prune.rs +++ b/crates/vchordg/src/prune.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use always_equal::AlwaysEqual; use distance::Distance; diff --git a/crates/vchordg/src/results.rs b/crates/vchordg/src/results.rs index c04b6fa1..8a31396a 100644 --- a/crates/vchordg/src/results.rs +++ b/crates/vchordg/src/results.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use always_equal::AlwaysEqual; use distance::Distance; diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs index 33a378e9..23a6230d 100644 --- a/crates/vchordg/src/search.rs +++ b/crates/vchordg/src/search.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::Opaque; use crate::candidates::Candidates; diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs index 7f26d314..42efddf8 100644 --- a/crates/vchordg/src/tuples.rs +++ b/crates/vchordg/src/tuples.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::operator::Vector; use distance::Distance; diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index 4fbb9a48..62ec3cc4 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use serde::{Deserialize, Serialize}; use simd::f16; diff --git a/crates/vchordg/src/vectors.rs b/crates/vchordg/src/vectors.rs index 272ba48f..e2933fba 100644 --- a/crates/vchordg/src/vectors.rs +++ b/crates/vchordg/src/vectors.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::operator::{Operator, Vector}; use crate::tuples::*; diff --git a/crates/vchordg/src/visited.rs b/crates/vchordg/src/visited.rs index 8f0cd169..8599dff2 100644 --- a/crates/vchordg/src/visited.rs +++ b/crates/vchordg/src/visited.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::Id; use std::collections::HashSet; diff --git a/crates/vchordrq/src/build.rs b/crates/vchordrq/src/build.rs index 9c00197c..a7371353 100644 --- a/crates/vchordrq/src/build.rs +++ b/crates/vchordrq/src/build.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::operator::{Operator, Vector}; use crate::tape::TapeWriter; diff --git a/crates/vchordrq/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs index d5e64acd..ae54ec54 100644 --- a/crates/vchordrq/src/bulkdelete.rs +++ b/crates/vchordrq/src/bulkdelete.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1}; use crate::operator::Operator; diff --git a/crates/vchordrq/src/cache.rs b/crates/vchordrq/src/cache.rs index 65c28347..d81ca98f 100644 --- a/crates/vchordrq/src/cache.rs +++ b/crates/vchordrq/src/cache.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1}; use crate::tuples::{MetaTuple, WithReader}; diff --git a/crates/vchordrq/src/centroids.rs b/crates/vchordrq/src/centroids.rs index 521987d5..d2024c9e 100644 --- a/crates/vchordrq/src/centroids.rs +++ b/crates/vchordrq/src/centroids.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::operator::*; use crate::tuples::*; diff --git a/crates/vchordrq/src/closure_lifetime_binder.rs b/crates/vchordrq/src/closure_lifetime_binder.rs index f247e9e2..86237624 100644 --- a/crates/vchordrq/src/closure_lifetime_binder.rs +++ b/crates/vchordrq/src/closure_lifetime_binder.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. // Use stable language features as an alternative to `closure_lifetime_binder`. // See https://github.com/rust-lang/rust/issues/97362. diff --git a/crates/vchordrq/src/consume.rs b/crates/vchordrq/src/consume.rs index 97749161..e68d9c0d 100644 --- a/crates/vchordrq/src/consume.rs +++ b/crates/vchordrq/src/consume.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::tuples::*; use crate::{Opaque, freepages}; diff --git a/crates/vchordrq/src/cost.rs b/crates/vchordrq/src/cost.rs index 3dbb4365..60579401 100644 --- a/crates/vchordrq/src/cost.rs +++ b/crates/vchordrq/src/cost.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::tuples::{MetaTuple, WithReader}; use index::relation::{Page, RelationRead}; diff --git a/crates/vchordrq/src/fast_heap.rs b/crates/vchordrq/src/fast_heap.rs index d4c61711..c36f91a9 100644 --- a/crates/vchordrq/src/fast_heap.rs +++ b/crates/vchordrq/src/fast_heap.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use index::prefetcher::Sequence; use std::collections::BinaryHeap; diff --git a/crates/vchordrq/src/freepages.rs b/crates/vchordrq/src/freepages.rs index 7657f6e0..9850f452 100644 --- a/crates/vchordrq/src/freepages.rs +++ b/crates/vchordrq/src/freepages.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::Opaque; use crate::tuples::*; diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs index bf256b08..3d5308e8 100644 --- a/crates/vchordrq/src/insert.rs +++ b/crates/vchordrq/src/insert.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::closure_lifetime_binder::id_0; use crate::linked_vec::LinkedVec; diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index 726f29cc..d0efef05 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod build; mod bulkdelete; diff --git a/crates/vchordrq/src/linked_vec.rs b/crates/vchordrq/src/linked_vec.rs index 650a962c..a42f23fc 100644 --- a/crates/vchordrq/src/linked_vec.rs +++ b/crates/vchordrq/src/linked_vec.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub struct LinkedVec { inner: Vec>, diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs index 291fa2e5..3d79ad7c 100644 --- a/crates/vchordrq/src/maintain.rs +++ b/crates/vchordrq/src/maintain.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1, id_2, id_3}; use crate::operator::{Operator, Vector}; diff --git a/crates/vchordrq/src/operator.rs b/crates/vchordrq/src/operator.rs index 3e1016d4..d3a893d2 100644 --- a/crates/vchordrq/src/operator.rs +++ b/crates/vchordrq/src/operator.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use distance::Distance; use index_accessor::{ diff --git a/crates/vchordrq/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs index bfc8cb73..02cb9b03 100644 --- a/crates/vchordrq/src/prewarm.rs +++ b/crates/vchordrq/src/prewarm.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1, id_2}; use crate::operator::Operator; diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs index 40905842..08f089ff 100644 --- a/crates/vchordrq/src/rerank.rs +++ b/crates/vchordrq/src/rerank.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::closure_lifetime_binder::id_4; use crate::operator::*; diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs index 9bd0409c..5e5a9e02 100644 --- a/crates/vchordrq/src/search.rs +++ b/crates/vchordrq/src/search.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::closure_lifetime_binder::{id_0, id_1, id_2}; use crate::linked_vec::LinkedVec; diff --git a/crates/vchordrq/src/tape.rs b/crates/vchordrq/src/tape.rs index 16502785..fb0d8878 100644 --- a/crates/vchordrq/src/tape.rs +++ b/crates/vchordrq/src/tape.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::tuples::*; use crate::{Opaque, freepages}; diff --git a/crates/vchordrq/src/tape_writer.rs b/crates/vchordrq/src/tape_writer.rs index 258ead51..565d796f 100644 --- a/crates/vchordrq/src/tape_writer.rs +++ b/crates/vchordrq/src/tape_writer.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::tape::TapeWriter; use crate::tuples::*; diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index 6f1072c0..8438c1ae 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::operator::Vector; use index::tuples::{Bool, MutChecker, Padding, RefChecker}; diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index 7097cb8d..c41ba501 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use serde::{Deserialize, Serialize}; use simd::f16; diff --git a/crates/vchordrq/src/vectors.rs b/crates/vchordrq/src/vectors.rs index d48150c0..3531efa2 100644 --- a/crates/vchordrq/src/vectors.rs +++ b/crates/vchordrq/src/vectors.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::operator::*; use crate::tuples::*; diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs index a1afca58..716d26ea 100644 --- a/crates/vector/src/bvect.rs +++ b/crates/vector/src/bvect.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; diff --git a/crates/vector/src/lib.rs b/crates/vector/src/lib.rs index 1009043a..07301250 100644 --- a/crates/vector/src/lib.rs +++ b/crates/vector/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub mod bvect; pub mod rabitq4; diff --git a/crates/vector/src/rabitq4.rs b/crates/vector/src/rabitq4.rs index 5c24abfd..4010c786 100644 --- a/crates/vector/src/rabitq4.rs +++ b/crates/vector/src/rabitq4.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; diff --git a/crates/vector/src/rabitq8.rs b/crates/vector/src/rabitq8.rs index 738b5fa2..d02a40c5 100644 --- a/crates/vector/src/rabitq8.rs +++ b/crates/vector/src/rabitq8.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; diff --git a/crates/vector/src/svect.rs b/crates/vector/src/svect.rs index caa3f07b..f04afe86 100644 --- a/crates/vector/src/svect.rs +++ b/crates/vector/src/svect.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; diff --git a/crates/vector/src/vect.rs b/crates/vector/src/vect.rs index b78f2a68..84a658f3 100644 --- a/crates/vector/src/vect.rs +++ b/crates/vector/src/vect.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::{VectorBorrowed, VectorOwned}; use distance::Distance; diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index f3102340..b2b5c65d 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use clap::{Args, Parser, Subcommand}; use object::{Object, ObjectSymbol}; diff --git a/src/bin/pgrx_embed.rs b/src/bin/pgrx_embed.rs index 6e0b06da..3307dd2a 100644 --- a/src/bin/pgrx_embed.rs +++ b/src/bin/pgrx_embed.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #![allow(unsafe_code)] diff --git a/src/datatype/binary_rabitq4.rs b/src/datatype/binary_rabitq4.rs index d3a9d56c..b623c62d 100644 --- a/src/datatype/binary_rabitq4.rs +++ b/src/datatype/binary_rabitq4.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; use pgrx::datum::Internal; diff --git a/src/datatype/binary_rabitq8.rs b/src/datatype/binary_rabitq8.rs index 2c1af908..023fed84 100644 --- a/src/datatype/binary_rabitq8.rs +++ b/src/datatype/binary_rabitq8.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; use pgrx::datum::Internal; diff --git a/src/datatype/functions_rabitq4.rs b/src/datatype/functions_rabitq4.rs index 63b46556..eb40e559 100644 --- a/src/datatype/functions_rabitq4.rs +++ b/src/datatype/functions_rabitq4.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; diff --git a/src/datatype/functions_rabitq8.rs b/src/datatype/functions_rabitq8.rs index 929e9c89..b5c56cad 100644 --- a/src/datatype/functions_rabitq8.rs +++ b/src/datatype/functions_rabitq8.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs index e0b1dd81..ca185bdd 100644 --- a/src/datatype/memory_halfvec.rs +++ b/src/datatype/memory_halfvec.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use pgrx::datum::{FromDatum, IntoDatum}; use pgrx::pg_sys::{Datum, Oid}; diff --git a/src/datatype/memory_rabitq4.rs b/src/datatype/memory_rabitq4.rs index 3d51643b..2f1e9433 100644 --- a/src/datatype/memory_rabitq4.rs +++ b/src/datatype/memory_rabitq4.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use pgrx::datum::{FromDatum, IntoDatum}; use pgrx::pg_sys::{Datum, Oid}; diff --git a/src/datatype/memory_rabitq8.rs b/src/datatype/memory_rabitq8.rs index a6d04b13..b5ef8a4d 100644 --- a/src/datatype/memory_rabitq8.rs +++ b/src/datatype/memory_rabitq8.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use pgrx::datum::{FromDatum, IntoDatum}; use pgrx::pg_sys::{Datum, Oid}; diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs index d982fdd4..6cd6017c 100644 --- a/src/datatype/memory_vector.rs +++ b/src/datatype/memory_vector.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use pgrx::datum::{FromDatum, IntoDatum}; use pgrx::pg_sys::{Datum, Oid}; diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 8fcb90fa..3725f050 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod binary_rabitq4; mod binary_rabitq8; diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index 96bc5ce0..0ae307fe 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; use pgrx::datum::Array; diff --git a/src/datatype/operators_rabitq4.rs b/src/datatype/operators_rabitq4.rs index 3c74cab2..bf0ff920 100644 --- a/src/datatype/operators_rabitq4.rs +++ b/src/datatype/operators_rabitq4.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; use pgrx::datum::Array; diff --git a/src/datatype/operators_rabitq8.rs b/src/datatype/operators_rabitq8.rs index 7eaf0226..d8d3cb72 100644 --- a/src/datatype/operators_rabitq8.rs +++ b/src/datatype/operators_rabitq8.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; use pgrx::datum::Array; diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index d7f8137d..76b3bf8c 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_vector::{VectorInput, VectorOutput}; use pgrx::datum::Array; diff --git a/src/datatype/text_rabitq4.rs b/src/datatype/text_rabitq4.rs index fa6af638..a239eb1c 100644 --- a/src/datatype/text_rabitq4.rs +++ b/src/datatype/text_rabitq4.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; use pgrx::pg_sys::Oid; diff --git a/src/datatype/text_rabitq8.rs b/src/datatype/text_rabitq8.rs index 52df1efb..f21e6586 100644 --- a/src/datatype/text_rabitq8.rs +++ b/src/datatype/text_rabitq8.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; use pgrx::pg_sys::Oid; diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index da6a6bf3..6ab9b44c 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use std::num::NonZero; diff --git a/src/datatype/typmod_rabitq4.rs b/src/datatype/typmod_rabitq4.rs index 59970207..32ac0f93 100644 --- a/src/datatype/typmod_rabitq4.rs +++ b/src/datatype/typmod_rabitq4.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use std::ffi::CStr; diff --git a/src/datatype/typmod_rabitq8.rs b/src/datatype/typmod_rabitq8.rs index 748888a1..881bb35d 100644 --- a/src/datatype/typmod_rabitq8.rs +++ b/src/datatype/typmod_rabitq8.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use std::ffi::CStr; diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs index bba55206..704edfc8 100644 --- a/src/index/fetcher.rs +++ b/src/index/fetcher.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use pgrx::pg_sys::{BlockIdData, Datum, ItemPointerData}; use std::cell::LazyCell; diff --git a/src/index/functions.rs b/src/index/functions.rs index 922e46db..d21d08b3 100644 --- a/src/index/functions.rs +++ b/src/index/functions.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::storage::PostgresRelation; use crate::recorder::dump; diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 039909a7..8c0e9fd3 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::scanners::Io; use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting, PostgresGucEnum}; diff --git a/src/index/hook.rs b/src/index/hook.rs index bfada43a..7a35f0b0 100644 --- a/src/index/hook.rs +++ b/src/index/hook.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #[pgrx::pg_guard] unsafe extern "C-unwind" fn rewrite_plan_state( diff --git a/src/index/mod.rs b/src/index/mod.rs index f4cd19d4..154e7489 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod fetcher; mod functions; diff --git a/src/index/opclass.rs b/src/index/opclass.rs index 607c163f..32ee8b20 100644 --- a/src/index/opclass.rs +++ b/src/index/opclass.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vchordg_support_vector_l2_ops() -> String { diff --git a/src/index/sample.rs b/src/index/sample.rs index 5bb54735..3bd25493 100644 --- a/src/index/sample.rs +++ b/src/index/sample.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use pgrx::pg_sys::{Datum, ItemPointerData}; use std::ptr::NonNull; diff --git a/src/index/scanners.rs b/src/index/scanners.rs index eaf671a0..dd48848b 100644 --- a/src/index/scanners.rs +++ b/src/index/scanners.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::fetcher::Fetcher; use crate::recorder::Recorder; diff --git a/src/index/storage.rs b/src/index/storage.rs index bad6b812..946df2c2 100644 --- a/src/index/storage.rs +++ b/src/index/storage.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub mod buffered; diff --git a/src/index/storage/buffered.rs b/src/index/storage/buffered.rs index 6c438179..3748dd52 100644 --- a/src/index/storage/buffered.rs +++ b/src/index/storage/buffered.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::storage::{ PostgresBufferReadGuard, PostgresBufferWriteGuard, PostgresPage, PostgresRelation, diff --git a/src/index/traverse.rs b/src/index/traverse.rs index b30ff39b..c8d35699 100644 --- a/src/index/traverse.rs +++ b/src/index/traverse.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use pgrx::pg_sys::{Datum, ItemPointerData}; diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index 0349acc5..819d79be 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::typmod::Typmod; use crate::index::storage::PostgresRelation; diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index 68fa97bf..5ed8c9f3 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod am_build; diff --git a/src/index/vchordg/dispatch.rs b/src/index/vchordg/dispatch.rs index 6f1dbff8..02115da3 100644 --- a/src/index/vchordg/dispatch.rs +++ b/src/index/vchordg/dispatch.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::vchordg::opclass::Opfamily; use index::fetch::Fetch; diff --git a/src/index/vchordg/mod.rs b/src/index/vchordg/mod.rs index 157200da..b01977c7 100644 --- a/src/index/vchordg/mod.rs +++ b/src/index/vchordg/mod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub mod am; pub mod dispatch; diff --git a/src/index/vchordg/opclass.rs b/src/index/vchordg/opclass.rs index 460a0fda..7bbbaaa4 100644 --- a/src/index/vchordg/opclass.rs +++ b/src/index/vchordg/opclass.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs index a5bebc42..22fd5b33 100644 --- a/src/index/vchordg/scanners/default.rs +++ b/src/index/vchordg/scanners/default.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::fetcher::{Fetcher, pointer_to_kv}; use crate::index::opclass::Sphere; diff --git a/src/index/vchordg/scanners/mod.rs b/src/index/vchordg/scanners/mod.rs index a9545ccc..9809809d 100644 --- a/src/index/vchordg/scanners/mod.rs +++ b/src/index/vchordg/scanners/mod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod default; diff --git a/src/index/vchordg/types.rs b/src/index/vchordg/types.rs index 84f88fb2..6be44ea3 100644 --- a/src/index/vchordg/types.rs +++ b/src/index/vchordg/types.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use serde::{Deserialize, Serialize}; use validator::Validate; diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 1c504537..5803c772 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::typmod::Typmod; use crate::index::fetcher::*; diff --git a/src/index/vchordrq/am/am_vacuumcleanup.rs b/src/index/vchordrq/am/am_vacuumcleanup.rs index 193bedc0..12e19075 100644 --- a/src/index/vchordrq/am/am_vacuumcleanup.rs +++ b/src/index/vchordrq/am/am_vacuumcleanup.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::vchordrq::am::PostgresRelation; use crate::index::vchordrq::opclass::opfamily; diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index e75f95c7..b41482d0 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod am_build; mod am_vacuumcleanup; diff --git a/src/index/vchordrq/build.rs b/src/index/vchordrq/build.rs index 4e91a340..bbdc55ae 100644 --- a/src/index/vchordrq/build.rs +++ b/src/index/vchordrq/build.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use simd::{Floating, f16}; use vector::rabitq4::Rabitq4Owned; diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs index 550c1b61..ec4b5b51 100644 --- a/src/index/vchordrq/dispatch.rs +++ b/src/index/vchordrq/dispatch.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::vchordrq::build::{Normalize, Normalized}; use crate::index::vchordrq::opclass::Opfamily; diff --git a/src/index/vchordrq/filter.rs b/src/index/vchordrq/filter.rs index f63c7729..89470f5b 100644 --- a/src/index/vchordrq/filter.rs +++ b/src/index/vchordrq/filter.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use index::prefetcher::Sequence; diff --git a/src/index/vchordrq/mod.rs b/src/index/vchordrq/mod.rs index 400542cf..38d75856 100644 --- a/src/index/vchordrq/mod.rs +++ b/src/index/vchordrq/mod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub mod am; mod build; diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs index f3d1a4d5..4ceaa46b 100644 --- a/src/index/vchordrq/opclass.rs +++ b/src/index/vchordrq/opclass.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs index 9c6719ea..d91b4273 100644 --- a/src/index/vchordrq/scanners/default.rs +++ b/src/index/vchordrq/scanners/default.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::fetcher::*; use crate::index::opclass::Sphere; diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index c63e91b0..1c51c97f 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::index::fetcher::*; use crate::index::scanners::{Io, SearchBuilder}; diff --git a/src/index/vchordrq/scanners/mod.rs b/src/index/vchordrq/scanners/mod.rs index a1c7e8f1..b345da37 100644 --- a/src/index/vchordrq/scanners/mod.rs +++ b/src/index/vchordrq/scanners/mod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. mod default; mod maxsim; diff --git a/src/index/vchordrq/types.rs b/src/index/vchordrq/types.rs index f2645c07..cff94858 100644 --- a/src/index/vchordrq/types.rs +++ b/src/index/vchordrq/types.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use serde::{Deserialize, Serialize}; use validator::{Validate, ValidationError, ValidationErrors}; diff --git a/src/lib.rs b/src/lib.rs index f9e8c20c..93fcc081 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. #![allow(unsafe_code)] #![deny(ffi_unwind_calls)] diff --git a/src/recorder/hook.rs b/src/recorder/hook.rs index 8dadee99..c0e5e786 100644 --- a/src/recorder/hook.rs +++ b/src/recorder/hook.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::recorder::worker::{delete_database, delete_index}; diff --git a/src/recorder/mod.rs b/src/recorder/mod.rs index 29f89886..1488a4ac 100644 --- a/src/recorder/mod.rs +++ b/src/recorder/mod.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. pub use types::{DefaultRecorder, Recorder}; pub use worker::dump; diff --git a/src/recorder/text.rs b/src/recorder/text.rs index 40a8dd8a..19ed4546 100644 --- a/src/recorder/text.rs +++ b/src/recorder/text.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use simd::f16; use vector::rabitq4::Rabitq4Borrowed; diff --git a/src/recorder/types.rs b/src/recorder/types.rs index f32076e1..ab40ab02 100644 --- a/src/recorder/types.rs +++ b/src/recorder/types.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::recorder::worker::push; use rand::RngExt; diff --git a/src/recorder/worker.rs b/src/recorder/worker.rs index cbef6b11..fd942ae0 100644 --- a/src/recorder/worker.rs +++ b/src/recorder/worker.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. use crate::recorder::types::PgRefCell; use std::cell::RefMut; diff --git a/src/upgrade.rs b/src/upgrade.rs index 1dc273ed..119c6802 100644 --- a/src/upgrade.rs +++ b/src/upgrade.rs @@ -10,7 +10,7 @@ // regarding the licenses, please contact us at: // vectorchord-inquiry@tensorchord.ai // -// Copyright (c) 2025 TensorChord Inc. +// Copyright (c) 2025-2026 TensorChord Inc. // Referenced symbols must exist in the dynamic library when dropping functions. // So we should never remove symbols used by schema, otherwise there will be errors in upgrade. From 5bc138fab437105a37039b00e60a38e61277e60d Mon Sep 17 00:00:00 2001 From: Narek Galstyan Date: Mon, 23 Mar 2026 20:24:13 -0400 Subject: [PATCH 289/324] fix: WHPG compatibility (#444) Signed-off-by: Narek Galstyan Co-authored-by: Artjoms Iskovs --- src/index/vchordrq/am/am_build.rs | 7 ++++--- src/index/vchordrq/dispatch.rs | 6 +++--- src/sql/finalize.sql | 8 ++++---- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 5803c772..647d2478 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -517,9 +517,10 @@ mod vchordrq_cached { } let pages_s = buffer.len(); buffer.extend(pages.iter().flat_map(|x| unsafe { - std::mem::transmute::<&PostgresPage, &[u8; 8192]>( - x.as_ref(), - ) + std::mem::transmute::< + &PostgresPage, + &[u8; pgrx::pg_sys::BLCKSZ as usize], + >(x.as_ref()) })); let pages_e = buffer.len(); while buffer.len() % ALIGN != 0 { diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs index ec4b5b51..01d70a87 100644 --- a/src/index/vchordrq/dispatch.rs +++ b/src/index/vchordrq/dispatch.rs @@ -179,15 +179,15 @@ pub fn maintain( check, ), }; - pgrx::info!( + pgrx::debug1!( "maintain: number_of_formerly_allocated_pages = {}", maintain.number_of_formerly_allocated_pages ); - pgrx::info!( + pgrx::debug1!( "maintain: number_of_freshly_allocated_pages = {}", maintain.number_of_freshly_allocated_pages ); - pgrx::info!( + pgrx::debug1!( "maintain: number_of_freed_pages = {}", maintain.number_of_freed_pages ); diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 39c0493a..a4f8f38a 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -189,16 +189,16 @@ CREATE OPERATOR @# ( -- List of functions CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector -IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_vector'; CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec -IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_halfvec'; CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 -IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq8'; CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 -IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq4'; CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; From af55342470903b3949c6d4b73a02d2630c6f1a6e Mon Sep 17 00:00:00 2001 From: usamoi Date: Tue, 24 Mar 2026 16:36:04 +0800 Subject: [PATCH 290/324] fix: don't create sqlite database until stats are written (#448) Signed-off-by: usamoi --- src/recorder/worker.rs | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/recorder/worker.rs b/src/recorder/worker.rs index fd942ae0..b56be6d8 100644 --- a/src/recorder/worker.rs +++ b/src/recorder/worker.rs @@ -17,14 +17,14 @@ use std::cell::RefMut; use std::fs; use std::path::Path; -// Safety: The directory name must start with "pgsql_tmp" to be excluded by pg_basebackup +// The directory name must start with "pgsql_tmp" to be excluded by pg_basebackup const RECORDER_DIR: &str = "pgsql_tmp_vchord_sampling"; const RECORDER_VERSION: u32 = 1; static CONNECTION: PgRefCell> = PgRefCell::>::new(None); -fn get<'a>() -> Option> { +fn get<'a>(create: bool) -> Option> { if unsafe { !pgrx::pg_sys::IsBackendPid(pgrx::pg_sys::MyProcPid) } { return None; } @@ -36,10 +36,16 @@ fn get<'a>() -> Option> { if connection.is_none() && let Err(err) = || -> rusqlite::Result<()> { if !Path::new(RECORDER_DIR).exists() { - let _ = fs::create_dir_all(RECORDER_DIR); + if create { + let _ = fs::create_dir_all(RECORDER_DIR); + } } let p = format!("{RECORDER_DIR}/database_{database_oid}.sqlite"); - let mut conn = rusqlite::Connection::open(&p)?; + let mut flags = rusqlite::OpenFlags::default(); + if !create { + flags.remove(rusqlite::OpenFlags::SQLITE_OPEN_CREATE); + } + let mut conn = rusqlite::Connection::open_with_flags(&p, flags)?; conn.pragma_update(Some("main"), "journal_mode", "WAL")?; conn.pragma_update(Some("main"), "synchronous", "NORMAL")?; let tx = conn.transaction()?; @@ -72,7 +78,7 @@ fn get<'a>() -> Option> { } pub fn push(index: u32, sample: &str, max_records: u32) { - let mut connection = match get() { + let mut connection = match get(true) { Some(c) => c, None => return, }; @@ -105,7 +111,7 @@ pub fn push(index: u32, sample: &str, max_records: u32) { } pub fn delete_index(index: u32) { - let connection = match get() { + let connection = match get(false) { Some(c) => c, None => return, }; @@ -122,7 +128,7 @@ pub fn delete_database(database_oid: u32) { } pub fn dump(index: u32) -> Vec { - let connection = match get() { + let connection = match get(false) { Some(c) => c, None => return Vec::new(), }; From 5d55cd2e9af96062e82cb26d4e82f7b92e19ef26 Mon Sep 17 00:00:00 2001 From: usamoi Date: Wed, 8 Apr 2026 18:03:14 +0800 Subject: [PATCH 291/324] chore: update dependencies (#449) Signed-off-by: usamoi --- Cargo.lock | 201 +++++++++++++++-------------- Cargo.toml | 4 +- crates/simd/Cargo.toml | 2 +- crates/vchordrq/src/tape_writer.rs | 22 ++-- crates/xtask/Cargo.toml | 2 +- src/index/vchordg/am/mod.rs | 21 +-- src/index/vchordrq/am/mod.rs | 7 - 7 files changed, 123 insertions(+), 136 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 903b939b..b59c2982 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -175,9 +175,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.2.57" +version = "1.2.59" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a0dd1ca384932ff3641c8718a02769f1698e7563dc6974ffd03346116310423" +checksum = "b7a4d3ec6524d28a329fc53654bbadc9bdd7b0431f5d65f1a56ffb28a1ee5283" dependencies = [ "find-msvc-tools", "shlex", @@ -548,9 +548,9 @@ checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" [[package]] name = "fastrand" -version = "2.3.0" +version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" [[package]] name = "feistel" @@ -696,12 +696,13 @@ dependencies = [ [[package]] name = "icu_collections" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c6b649701667bbe825c3b7e6388cb521c23d88644678e83c0c4d0a621a34b43" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" dependencies = [ "displaydoc", "potential_utf", + "utf8_iter", "yoke", "zerofrom", "zerovec", @@ -709,9 +710,9 @@ dependencies = [ [[package]] name = "icu_locale_core" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "edba7861004dd3714265b4db54a3c390e880ab658fec5f7db895fae2046b5bb6" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" dependencies = [ "displaydoc", "litemap", @@ -722,9 +723,9 @@ dependencies = [ [[package]] name = "icu_normalizer" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f6c8828b67bf8908d82127b2054ea1b4427ff0230ee9141c54251934ab1b599" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" dependencies = [ "icu_collections", "icu_normalizer_data", @@ -736,15 +737,15 @@ dependencies = [ [[package]] name = "icu_normalizer_data" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7aedcccd01fc5fe81e6b489c15b247b8b0690feb23304303a9e560f37efc560a" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" [[package]] name = "icu_properties" -version = "2.1.2" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "020bfc02fe870ec3a66d93e677ccca0562506e5872c650f893269e08615d74ec" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" dependencies = [ "icu_collections", "icu_locale_core", @@ -756,15 +757,15 @@ dependencies = [ [[package]] name = "icu_properties_data" -version = "2.1.2" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "616c294cf8d725c6afcd8f55abc17c56464ef6211f9ed59cccffe534129c77af" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" [[package]] name = "icu_provider" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85962cf0ce02e1e0a629cc34e7ca3e373ce20dda4c4d7294bbd0bf1fdb59e614" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" dependencies = [ "displaydoc", "icu_locale_core", @@ -837,9 +838,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.13.0" +version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" +checksum = "45a8a2b9cb3e0b0c1803dbb0758ffac5de2f425b23c28f518faabd9d805342ff" dependencies = [ "equivalent", "hashbrown 0.16.1", @@ -870,15 +871,15 @@ dependencies = [ [[package]] name = "itoa" -version = "1.0.17" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "js-sys" -version = "0.3.91" +version = "0.3.94" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b49715b7073f385ba4bc528e5747d02e66cb39c6146efb66b781f131f0fb399c" +checksum = "2e04e2ef80ce82e13552136fabeef8a5ed1f985a96805761cbb9a2c34e7664d9" dependencies = [ "once_cell", "wasm-bindgen", @@ -904,9 +905,9 @@ checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" [[package]] name = "libc" -version = "0.2.183" +version = "0.2.184" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" +checksum = "48f5d2a454e16a5ea0f4ced81bd44e4cfc7bd3a507b61887c99fd3538b28e4af" [[package]] name = "libloading" @@ -953,9 +954,9 @@ checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" [[package]] name = "litemap" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "log" @@ -1021,14 +1022,14 @@ dependencies = [ [[package]] name = "object" -version = "0.38.1" +version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "271638cd5fa9cca89c4c304675ca658efc4e64a66c716b7cfe1afb4b9611dbbc" +checksum = "63944c133d03f44e75866bbd160b95af0ec3f6a13d936d69d31c81078cbc5baf" dependencies = [ "flate2", "memchr", "ruzstd", - "wasmparser 0.243.0", + "wasmparser 0.245.1", ] [[package]] @@ -1270,9 +1271,9 @@ dependencies = [ [[package]] name = "potential_utf" -version = "0.1.4" +version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b73949432f5e2a09657003c25bca5e19a0e9c84f8058ca374f49e0ebe605af77" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" dependencies = [ "zerovec", ] @@ -1467,9 +1468,9 @@ dependencies = [ [[package]] name = "rustc-hash" -version = "2.1.1" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" [[package]] name = "rustix" @@ -1516,9 +1517,9 @@ checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" [[package]] name = "semver" -version = "1.0.27" +version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" [[package]] name = "seq-macro" @@ -1581,9 +1582,9 @@ dependencies = [ [[package]] name = "serde_spanned" -version = "1.0.4" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8bbf91e5a4d6315eee45e704372590b30e260ee83af6639d64557f51b067776" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" dependencies = [ "serde_core", ] @@ -1609,9 +1610,9 @@ dependencies = [ [[package]] name = "simd-adler32" -version = "0.3.8" +version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" [[package]] name = "simd_macros" @@ -1734,9 +1735,9 @@ dependencies = [ [[package]] name = "tinystr" -version = "0.8.2" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42d3e9c45c09de15d06dd8acf5f4e0e399e85927b7f00711024eb7ae10fa4869" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" dependencies = [ "displaydoc", "zerovec", @@ -1764,22 +1765,22 @@ dependencies = [ "toml_datetime 0.7.5+spec-1.1.0", "toml_parser", "toml_writer", - "winnow", + "winnow 0.7.15", ] [[package]] name = "toml" -version = "1.0.6+spec-1.1.0" +version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "399b1124a3c9e16766831c6bba21e50192572cdd98706ea114f9502509686ffc" +checksum = "81f3d15e84cbcd896376e6730314d59fb5a87f31e4b038454184435cd57defee" dependencies = [ "indexmap", "serde_core", "serde_spanned", - "toml_datetime 1.0.0+spec-1.1.0", + "toml_datetime 1.1.1+spec-1.1.0", "toml_parser", "toml_writer", - "winnow", + "winnow 1.0.1", ] [[package]] @@ -1793,27 +1794,27 @@ dependencies = [ [[package]] name = "toml_datetime" -version = "1.0.0+spec-1.1.0" +version = "1.1.1+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32c2555c699578a4f59f0cc68e5116c8d7cabbd45e1409b989d4be085b53f13e" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" dependencies = [ "serde_core", ] [[package]] name = "toml_parser" -version = "1.0.9+spec-1.1.0" +version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "702d4415e08923e7e1ef96cd5727c0dfed80b4d2fa25db9647fe5eb6f7c5a4c4" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" dependencies = [ - "winnow", + "winnow 1.0.1", ] [[package]] name = "toml_writer" -version = "1.0.6+spec-1.1.0" +version = "1.1.1+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab16f14aed21ee8bfd8ec22513f7287cd4a91aa92e44edfe2c17ddd004e92607" +checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" [[package]] name = "twox-hash" @@ -1835,9 +1836,9 @@ checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" [[package]] name = "unicode-segmentation" -version = "1.12.0" +version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" +checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" [[package]] name = "unicode-width" @@ -1877,9 +1878,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.22.0" +version = "1.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a68d3c8f01c0cfa54a75291d83601161799e4a89a39e0929f4b0354d88757a37" +checksum = "5ac8b6f42ead25368cf5b098aeb3dc8a1a2c05a3eee8a9a1a68c640edbfc79d9" dependencies = [ "getrandom", "js-sys", @@ -1942,7 +1943,7 @@ dependencies = [ "seq-macro", "serde", "simd", - "toml 1.0.6+spec-1.1.0", + "toml 1.1.2+spec-1.1.0", "validator", "vchordg", "vchordrq", @@ -2032,9 +2033,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.114" +version = "0.2.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6532f9a5c1ece3798cb1c2cfdba640b9b3ba884f5db45973a6f442510a87d38e" +checksum = "0551fc1bb415591e3372d0bc4780db7e587d84e2a7e79da121051c5c4b89d0b0" dependencies = [ "cfg-if", "once_cell", @@ -2045,9 +2046,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.114" +version = "0.2.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "18a2d50fcf105fb33bb15f00e7a77b772945a2ee45dcf454961fd843e74c18e6" +checksum = "7fbdf9a35adf44786aecd5ff89b4563a90325f9da0923236f6104e603c7e86be" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -2055,9 +2056,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.114" +version = "0.2.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03ce4caeaac547cdf713d280eda22a730824dd11e6b8c3ca9e42247b25c631e3" +checksum = "dca9693ef2bab6d4e6707234500350d8dad079eb508dca05530c85dc3a529ff2" dependencies = [ "bumpalo", "proc-macro2", @@ -2068,9 +2069,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.114" +version = "0.2.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75a326b8c223ee17883a4251907455a2431acc2791c98c26279376490c378c16" +checksum = "39129a682a6d2d841b6c429d0c51e5cb0ed1a03829d8b3d1e69a011e62cb3d3b" dependencies = [ "unicode-ident", ] @@ -2099,30 +2100,30 @@ dependencies = [ [[package]] name = "wasmparser" -version = "0.243.0" +version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6d8db401b0528ec316dfbe579e6ab4152d61739cfe076706d2009127970159d" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ "bitflags", + "hashbrown 0.15.5", + "indexmap", + "semver", ] [[package]] name = "wasmparser" -version = "0.244.0" +version = "0.245.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +checksum = "4f08c9adee0428b7bddf3890fc27e015ac4b761cc608c822667102b8bfd6995e" dependencies = [ "bitflags", - "hashbrown 0.15.5", - "indexmap", - "semver", ] [[package]] name = "web-sys" -version = "0.3.91" +version = "0.3.94" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "854ba17bb104abfb26ba36da9729addc7ce7f06f5c0f90f3c391f8461cca21f9" +checksum = "cd70027e39b12f0849461e08ffc50b9cd7688d942c1c8e3c7b22273236b4dd0a" dependencies = [ "js-sys", "wasm-bindgen", @@ -2180,6 +2181,12 @@ version = "0.7.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" +[[package]] +name = "winnow" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09dac053f1cd375980747450bfc7250c264eaae0583872e845c0c7cd578872b5" + [[package]] name = "wit-bindgen" version = "0.51.0" @@ -2270,9 +2277,9 @@ dependencies = [ [[package]] name = "writeable" -version = "0.6.2" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" [[package]] name = "wyhash" @@ -2307,9 +2314,9 @@ dependencies = [ [[package]] name = "yoke" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72d6e5c6afb84d73944e5cedb052c4680d5657337201555f9f2a16b7406d4954" +checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca" dependencies = [ "stable_deref_trait", "yoke-derive", @@ -2318,9 +2325,9 @@ dependencies = [ [[package]] name = "yoke-derive" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", @@ -2330,18 +2337,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.42" +version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2578b716f8a7a858b7f02d5bd870c14bf4ddbbcf3a4c05414ba6503640505e3" +checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.42" +version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e6cc098ea4d3bd6246687de65af3f920c430e236bee1e3bf2e441463f08a02f" +checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4" dependencies = [ "proc-macro2", "quote", @@ -2350,18 +2357,18 @@ dependencies = [ [[package]] name = "zerofrom" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50cc42e0333e05660c3587f3bf9d0478688e15d870fab3346451ce7f8c9fbea5" +checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df" dependencies = [ "zerofrom-derive", ] [[package]] name = "zerofrom-derive" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", @@ -2371,9 +2378,9 @@ dependencies = [ [[package]] name = "zerotrie" -version = "0.2.3" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2a59c17a5562d507e4b54960e8569ebee33bee890c70aa3fe7b97e85a9fd7851" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" dependencies = [ "displaydoc", "yoke", @@ -2382,9 +2389,9 @@ dependencies = [ [[package]] name = "zerovec" -version = "0.11.5" +version = "0.11.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c28719294829477f525be0186d13efa9a3c602f7ec202ca9e353d310fb9a002" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" dependencies = [ "yoke", "zerofrom", @@ -2393,9 +2400,9 @@ dependencies = [ [[package]] name = "zerovec-derive" -version = "0.11.2" +version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index e7dcd96f..264e546a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,7 +46,7 @@ rayon.workspace = true rusqlite = { version = "0.39.0", features = ["bundled"] } seq-macro.workspace = true serde.workspace = true -toml = "1.0.6" +toml = "1.1.2" validator.workspace = true wyhash.workspace = true zerocopy.workspace = true @@ -76,7 +76,7 @@ seq-macro = "0.3.6" serde = { version = "1.0.228", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } wyhash = "0.6.0" -zerocopy = { version = "0.8.42", features = ["derive"] } +zerocopy = { version = "0.8.48", features = ["derive"] } [workspace.lints] # complexity diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index e6d87c71..65c5d560 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -26,7 +26,7 @@ criterion = "0.8.2" rand.workspace = true [build-dependencies] -cc = "1.2.57" +cc = "1.2.59" [lints] workspace = true diff --git a/crates/vchordrq/src/tape_writer.rs b/crates/vchordrq/src/tape_writer.rs index 565d796f..3bb71a20 100644 --- a/crates/vchordrq/src/tape_writer.rs +++ b/crates/vchordrq/src/tape_writer.rs @@ -37,11 +37,13 @@ where } } pub fn push(&mut self, branch: &[u32]) { - let mut remain = branch.to_vec(); + let mut remain = branch; loop { let freespace = self.tape.freespace(); if DirectoryTuple::estimate_size_0(remain.len()) <= freespace as usize { - self.tape.tape_put(DirectoryTuple::_0 { elements: remain }); + self.tape.tape_put(DirectoryTuple::_0 { + elements: remain.to_vec(), + }); break; } if let Some(w) = DirectoryTuple::fit_1(freespace) { @@ -49,7 +51,7 @@ where self.tape.tape_put(DirectoryTuple::_1 { elements: left.to_vec(), }); - remain = right.to_vec(); + remain = right; } else { self.tape.tape_move(); } @@ -84,8 +86,9 @@ where pub fn push(&mut self, branch: Branch) { self.branches.push(branch); if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { - let mut remain = + let elements = padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.code.1))); + let mut remain = elements.as_slice(); loop { let freespace = self.tape.freespace(); if H1Tuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { @@ -102,7 +105,7 @@ where head: chunk.each_ref().map(|x| x.head), first: chunk.each_ref().map(|x| x.extra), len: chunk.len() as _, - elements: remain, + elements: remain.to_vec(), }); break; } @@ -111,7 +114,7 @@ where self.tape.tape_put(H1Tuple::_1 { elements: left.to_vec(), }); - remain = right.to_vec(); + remain = right; } else { self.tape.tape_move(); } @@ -184,8 +187,9 @@ where pub fn push(&mut self, branch: Branch>) { self.branches.push(branch); if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { - let mut remain = + let elements = padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.code.1))); + let mut remain = elements.as_slice(); loop { let freespace = self.tape.freespace(); if FrozenTuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { @@ -200,7 +204,7 @@ where prefetch: fix_good(chunk.each_ref().map(|x| x.prefetch.as_slice())), head: chunk.each_ref().map(|x| x.head), payload: chunk.each_ref().map(|x| Some(x.extra)), - elements: remain, + elements: remain.to_vec(), }); break; } @@ -209,7 +213,7 @@ where self.tape.tape_put(FrozenTuple::_1 { elements: left.to_vec(), }); - remain = right.to_vec(); + remain = right; } else { self.tape.tape_move(); } diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index 83f9522d..1645c81f 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -6,7 +6,7 @@ publish = false [dependencies] clap = { version = "4.6.0", features = ["derive", "env"] } -object = { version = "0.38.1", features = ["read", "wasm"] } +object = { version = "0.39.0", features = ["read", "wasm"] } serde.workspace = true serde_json = "1.0.149" shlex = "1.3.0" diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs index 5ed8c9f3..39707100 100644 --- a/src/index/vchordg/am/mod.rs +++ b/src/index/vchordg/am/mod.rs @@ -214,38 +214,21 @@ pub unsafe extern "C-unwind" fn amcostestimate( } } -#[cfg(any( - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17", - feature = "pg18" -))] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn aminsert( index_relation: pgrx::pg_sys::Relation, values: *mut Datum, is_null: *mut bool, heap_tid: pgrx::pg_sys::ItemPointer, - heap_relation: pgrx::pg_sys::Relation, + _heap_relation: pgrx::pg_sys::Relation, _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, _index_unchanged: bool, _index_info: *mut pgrx::pg_sys::IndexInfo, -) -> bool { - unsafe { aminsertinner(index_relation, heap_relation, values, is_null, heap_tid) } -} - -unsafe fn aminsertinner( - index_relation: pgrx::pg_sys::Relation, - _heap_relation: pgrx::pg_sys::Relation, - values: *mut Datum, - is_null: *mut bool, - ctid: pgrx::pg_sys::ItemPointer, ) -> bool { let opfamily = unsafe { opfamily(index_relation) }; let index = unsafe { PostgresRelation::new(index_relation) }; let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; - let ctid = unsafe { ctid.read() }; + let ctid = unsafe { heap_tid.read() }; if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { for (vector, extra) in store { let key = ctid_to_key(ctid); diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index b41482d0..5040b587 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -384,13 +384,6 @@ pub unsafe extern "C-unwind" fn amcostestimate( } } -#[cfg(any( - feature = "pg14", - feature = "pg15", - feature = "pg16", - feature = "pg17", - feature = "pg18" -))] #[pgrx::pg_guard] pub unsafe extern "C-unwind" fn aminsert( index_relation: pgrx::pg_sys::Relation, From 4e1b234e440afd0bcd3c5c1f8838bb19fd220106 Mon Sep 17 00:00:00 2001 From: usamoi Date: Mon, 20 Apr 2026 21:29:59 +0800 Subject: [PATCH 292/324] chore: update dependencies (#452) Signed-off-by: usamoi --- .github/workflows/release.yml | 2 +- Cargo.lock | 115 +++++++++++++++++++--------------- Cargo.toml | 9 +-- crates/simd/Cargo.toml | 2 +- crates/simd/src/fast_scan.rs | 6 +- crates/xtask/Cargo.toml | 2 +- src/index/vchordg/opclass.rs | 3 +- src/index/vchordrq/opclass.rs | 3 +- 8 files changed, 79 insertions(+), 63 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 214776ba..2c426c2a 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -21,7 +21,7 @@ jobs: steps: - name: Semver id: semver - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | const tag = "${{ github.event.inputs.tag }}" || "${{ github.event.release.tag_name }}"; diff --git a/Cargo.lock b/Cargo.lock index b59c2982..f324f59e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -135,9 +135,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.11.0" +version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" +checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" [[package]] name = "bitvec" @@ -175,9 +175,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.2.59" +version = "1.2.60" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7a4d3ec6524d28a329fc53654bbadc9bdd7b0431f5d65f1a56ffb28a1ee5283" +checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20" dependencies = [ "find-msvc-tools", "shlex", @@ -216,7 +216,7 @@ checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" dependencies = [ "cfg-if", "cpufeatures", - "rand_core 0.10.0", + "rand_core 0.10.1", ] [[package]] @@ -259,9 +259,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.0" +version = "4.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b193af5b67834b676abd72466a96c1024e6a6ad978a1f484bd90b85c94041351" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" dependencies = [ "clap_builder", "clap_derive", @@ -281,9 +281,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.6.0" +version = "4.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1110bd8a634a1ab8cb04345d8d878267d57c3cf1b38d91b71af6686408bbca6a" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" dependencies = [ "heck", "proc-macro2", @@ -451,9 +451,9 @@ dependencies = [ [[package]] name = "dary_heap" -version = "0.3.8" +version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06d2e3287df1c007e74221c49ca10a95d557349e54b3a75dc2fb14712c751f04" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" [[package]] name = "displaydoc" @@ -624,7 +624,7 @@ dependencies = [ "cfg-if", "libc", "r-efi", - "rand_core 0.10.0", + "rand_core 0.10.1", "wasip2", "wasip3", ] @@ -670,6 +670,12 @@ dependencies = [ "foldhash 0.2.0", ] +[[package]] +name = "hashbrown" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" + [[package]] name = "hashlink" version = "0.11.0" @@ -838,12 +844,12 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.13.1" +version = "2.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45a8a2b9cb3e0b0c1803dbb0758ffac5de2f425b23c28f518faabd9d805342ff" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown 0.16.1", + "hashbrown 0.17.0", "serde", "serde_core", ] @@ -877,9 +883,9 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "js-sys" -version = "0.3.94" +version = "0.3.95" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e04e2ef80ce82e13552136fabeef8a5ed1f985a96805761cbb9a2c34e7664d9" +checksum = "2964e92d1d9dc3364cae4d718d93f227e3abb088e747d92e0395bfdedf1c12ca" dependencies = [ "once_cell", "wasm-bindgen", @@ -905,9 +911,9 @@ checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" [[package]] name = "libc" -version = "0.2.184" +version = "0.2.185" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48f5d2a454e16a5ea0f4ced81bd44e4cfc7bd3a507b61887c99fd3538b28e4af" +checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f" [[package]] name = "libloading" @@ -927,12 +933,11 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libmimalloc-sys" -version = "0.1.44" +version = "0.1.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "667f4fec20f29dfc6bc7357c582d91796c169ad7e2fce709468aefeb2c099870" +checksum = "bc89deee4af0429081d2a518c0431ae068222a5a262a3bc6ff4d8535ec2e02fe" dependencies = [ "cc", - "libc", ] [[package]] @@ -972,9 +977,9 @@ checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] name = "mimalloc" -version = "0.1.48" +version = "0.1.49" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1ee66a4b64c74f4ef288bcbb9192ad9c3feaad75193129ac8509af543894fd8" +checksum = "aca3c01a711f395b4257b81674c0e90e8dd1f1e62c4b7db45f684cc7a4fcb18a" dependencies = [ "libmimalloc-sys", ] @@ -1237,9 +1242,9 @@ dependencies = [ [[package]] name = "pkg-config" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" [[package]] name = "plotters" @@ -1361,13 +1366,13 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "rand" -version = "0.10.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" +checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" dependencies = [ "chacha20", "getrandom", - "rand_core 0.10.0", + "rand_core 0.10.1", ] [[package]] @@ -1377,7 +1382,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3e6af7f3e25ded52c41df4e0b1af2d047e45896c2f3281792ed68a1c243daedb" dependencies = [ "ppv-lite86", - "rand_core 0.10.0", + "rand_core 0.10.1", ] [[package]] @@ -1388,15 +1393,15 @@ checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" [[package]] name = "rand_core" -version = "0.10.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" [[package]] name = "rayon" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" dependencies = [ "either", "rayon-core", @@ -1635,9 +1640,9 @@ checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" [[package]] name = "sqlite-wasm-rs" -version = "0.5.2" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f4206ed3a67690b9c29b77d728f6acc3ce78f16bf846d83c94f76400320181b" +checksum = "1b2c760607300407ddeaee518acf28c795661b7108c75421303dbefb237d3a36" dependencies = [ "cc", "js-sys", @@ -1878,9 +1883,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.23.0" +version = "1.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ac8b6f42ead25368cf5b098aeb3dc8a1a2c05a3eee8a9a1a68c640edbfc79d9" +checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" dependencies = [ "getrandom", "js-sys", @@ -2015,11 +2020,11 @@ dependencies = [ [[package]] name = "wasip2" -version = "1.0.2+wasi-0.2.9" +version = "1.0.3+wasi-0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5" +checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" dependencies = [ - "wit-bindgen", + "wit-bindgen 0.57.1", ] [[package]] @@ -2028,14 +2033,14 @@ version = "0.4.0+wasi-0.3.0-rc-2026-01-06" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" dependencies = [ - "wit-bindgen", + "wit-bindgen 0.51.0", ] [[package]] name = "wasm-bindgen" -version = "0.2.117" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0551fc1bb415591e3372d0bc4780db7e587d84e2a7e79da121051c5c4b89d0b0" +checksum = "0bf938a0bacb0469e83c1e148908bd7d5a6010354cf4fb73279b7447422e3a89" dependencies = [ "cfg-if", "once_cell", @@ -2046,9 +2051,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.117" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fbdf9a35adf44786aecd5ff89b4563a90325f9da0923236f6104e603c7e86be" +checksum = "eeff24f84126c0ec2db7a449f0c2ec963c6a49efe0698c4242929da037ca28ed" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -2056,9 +2061,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.117" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dca9693ef2bab6d4e6707234500350d8dad079eb508dca05530c85dc3a529ff2" +checksum = "9d08065faf983b2b80a79fd87d8254c409281cf7de75fc4b773019824196c904" dependencies = [ "bumpalo", "proc-macro2", @@ -2069,9 +2074,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.117" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39129a682a6d2d841b6c429d0c51e5cb0ed1a03829d8b3d1e69a011e62cb3d3b" +checksum = "5fd04d9e306f1907bd13c6361b5c6bfc7b3b3c095ed3f8a9246390f8dbdee129" dependencies = [ "unicode-ident", ] @@ -2121,9 +2126,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.94" +version = "0.3.95" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd70027e39b12f0849461e08ffc50b9cd7688d942c1c8e3c7b22273236b4dd0a" +checksum = "4f2dfbb17949fa2088e5d39408c48368947b86f7834484e87b73de55bc14d97d" dependencies = [ "js-sys", "wasm-bindgen", @@ -2196,6 +2201,12 @@ dependencies = [ "wit-bindgen-rust-macro", ] +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + [[package]] name = "wit-bindgen-core" version = "0.51.0" diff --git a/Cargo.toml b/Cargo.toml index 264e546a..5868e4d6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,7 @@ name = "vchord" version.workspace = true edition.workspace = true publish = false +autotests = false [lib] name = "vchord" @@ -52,7 +53,7 @@ wyhash.workspace = true zerocopy.workspace = true [target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos")))'.dependencies] -mimalloc = { version = "0.1.48", features = ["local_dynamic_tls"] } +mimalloc = { version = "0.1.49", features = ["local_dynamic_tls"] } [lints] workspace = true @@ -67,11 +68,11 @@ edition = "2024" [workspace.dependencies] bumpalo = "3.20.2" -dary_heap = "0.3.8" +dary_heap = "0.3.9" paste = "1.0.15" -rand = "0.10.0" +rand = "0.10.1" rand_chacha = "0.10.0" -rayon = "1.11.0" +rayon = "1.12.0" seq-macro = "0.3.6" serde = { version = "1.0.228", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] } diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml index 65c5d560..1b6faf59 100644 --- a/crates/simd/Cargo.toml +++ b/crates/simd/Cargo.toml @@ -26,7 +26,7 @@ criterion = "0.8.2" rand.workspace = true [build-dependencies] -cc = "1.2.59" +cc = "1.2.60" [lints] workspace = true diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs index 06502460..a02dfc33 100644 --- a/crates/simd/src/fast_scan.rs +++ b/crates/simd/src/fast_scan.rs @@ -446,7 +446,8 @@ mod scan { use core::arch::s390x::*; use std::mem::transmute; - use {vector_unsigned_char as u8x16, vector_unsigned_short as u16x8}; + use vector_unsigned_char as u8x16; + use vector_unsigned_short as u16x8; let _0001_u16x8 = vec_splat_u16::<0x0001>(); let _00ff_u16x8 = vec_splat_u16::<0x00ff>(); @@ -532,7 +533,8 @@ mod scan { use core::arch::powerpc64::*; use std::mem::transmute; - use {vector_unsigned_char as u8x16, vector_unsigned_short as u16x8}; + use vector_unsigned_char as u8x16; + use vector_unsigned_short as u16x8; #[cfg(target_endian = "big")] let revb = transmute::<[u8; 16], u8x16>([ diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index 1645c81f..b9d11560 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -5,7 +5,7 @@ edition.workspace = true publish = false [dependencies] -clap = { version = "4.6.0", features = ["derive", "env"] } +clap = { version = "4.6.1", features = ["derive", "env"] } object = { version = "0.39.0", features = ["read", "wasm"] } serde.workspace = true serde_json = "1.0.149" diff --git a/src/index/vchordg/opclass.rs b/src/index/vchordg/opclass.rs index 7bbbaaa4..2132ae22 100644 --- a/src/index/vchordg/opclass.rs +++ b/src/index/vchordg/opclass.rs @@ -43,7 +43,8 @@ pub enum Opfamily { impl Opfamily { fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { - use {BorrowedVector as B, OwnedVector as O}; + use BorrowedVector as B; + use OwnedVector as O; match (self, vector) { (Self::VectorL2, B::Vecf32(x)) => O::Vecf32(x.own()), (Self::VectorL2, _) => unreachable!(), diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs index 4ceaa46b..7c51f5db 100644 --- a/src/index/vchordrq/opclass.rs +++ b/src/index/vchordrq/opclass.rs @@ -47,7 +47,8 @@ pub enum Opfamily { impl Opfamily { fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { - use {BorrowedVector as B, OwnedVector as O}; + use BorrowedVector as B; + use OwnedVector as O; match (vector, self) { (B::Vecf32(x), Self::VectorL2) => O::Vecf32(x.own()), (B::Vecf32(x), Self::VectorIp | Self::VectorMaxsim) => O::Vecf32(x.own()), From b5ac9124154d84094a8409f7e2032b7756458dd9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 23 Apr 2026 20:26:25 +0800 Subject: [PATCH 293/324] build(deps): bump actions/upload-artifact from 7.0.0 to 7.0.1 (#454) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps [actions/upload-artifact](https://github.com/actions/upload-artifact) from 7.0.0 to 7.0.1.
Release notes

Sourced from actions/upload-artifact's releases.

v7.0.1

What's Changed

Full Changelog: https://github.com/actions/upload-artifact/compare/v7...v7.0.1

Commits
  • 043fb46 Merge pull request #797 from actions/yacaovsnc/update-dependency
  • 634250c Include changes in typespec/ts-http-runtime 0.3.5
  • e454baa Readme: bump all the example versions to v7 (#796)
  • 74fad66 Update the readme with direct upload details (#795)
  • See full diff in compare view

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=actions/upload-artifact&package-manager=github_actions&previous-version=7.0.0&new-version=7.0.1)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/check.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 52fbd9e7..df084841 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -296,7 +296,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu path: ./build @@ -394,7 +394,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin path: ./build @@ -514,7 +514,7 @@ jobs: run: make install - name: Upload Artifacts - uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc path: ./build @@ -645,7 +645,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-musl path: ./build @@ -862,7 +862,7 @@ jobs: run: sudo make PG_CONFIG=$PG_CONFIG install - name: Upload Artifacts - uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.clang_triple }} path: ./build From c2dea34ea657894f655ab9c196328de1d3760d28 Mon Sep 17 00:00:00 2001 From: spenc-r Date: Thu, 21 May 2026 17:25:30 +0000 Subject: [PATCH 294/324] feat(vchordrq): estimate filtered candidate selectivity Use the planner's filtered row estimate as filter selectivity for deciding how many IVFFlat candidates are needed to satisfy LIMIT after heap-side quals. Report indexSelectivity from the expected retrieved candidate count instead of the final survivor fraction, so cost_index() prices heap fetch and qpqual evaluation against candidate rows rather than every table row or only the rows that pass filters. Clamp the candidate-processing term by node_count, and tighten the cost estimator sqllogictest coverage around btree alternatives, candidate-cost floors, and extreme selectivity. --- src/index/vchordrq/am/mod.rs | 58 ++-- tests/vchordrq/cost_estimator.slt | 426 ++++++++++++++++++++++++++++++ 2 files changed, 467 insertions(+), 17 deletions(-) create mode 100644 tests/vchordrq/cost_estimator.slt diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 5040b587..f995c600 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -292,20 +292,31 @@ pub unsafe extern "C-unwind" fn amcostestimate( *index_pages = 1.0; return; } - let selectivity = { - use pgrx::pg_sys::{ - JoinType, add_predicate_to_index_quals, clauselist_selectivity, - get_quals_from_indexclauses, + // Vchordrq indexes only the vector column, so ordinary filters on + // other columns are heap-side quals rather than index quals. We use + // the planner's filtered row estimate to decide how many distance + // candidates we likely need to produce enough survivors for LIMIT. + // + // Keep this separate from the value returned as `indexSelectivity`: + // `cost_index()` interprets that output as the fraction of parent + // table rows the index scan retrieves, not the fraction surviving all + // filters. With LIMIT, the retrieved fraction is better represented + // by the candidate count computed below. + let (total_rows, filter_selectivity) = { + let baserel = (*index_opt_info).rel; + let total_rows = (*baserel).tuples; + let param_info = (*path).path.param_info; + let filtered_rows = if !param_info.is_null() { + (*param_info).ppi_rows + } else { + (*baserel).rows }; - let index_quals = get_quals_from_indexclauses((*path).indexclauses); - let selectivity_quals = add_predicate_to_index_quals(index_opt_info, index_quals); - clauselist_selectivity( - root, - selectivity_quals, - (*(*index_opt_info).rel).relid as _, - JoinType::JOIN_INNER, - std::ptr::null_mut(), - ) + let filter_selectivity = if total_rows > 0.0 { + (filtered_rows / total_rows).clamp(1e-9, 1.0) + } else { + 1.0 + }; + (total_rows, filter_selectivity) }; // index exists if !(*index_opt_info).hypothetical { @@ -367,18 +378,31 @@ pub unsafe extern "C-unwind" fn amcostestimate( pages += cost.cells[0] as f64; pages }; - let next_count = - f64::max(1.0, (*root).limit_tuples) * f64::min(1000.0, 1.0 / selectivity); + // `next_count` represents candidates we expect to process to + // surface `limit_tuples` survivors after filter rejection. Clamp + // by `node_count` so the estimate cannot exceed the candidates + // the IVF visits at the configured probe count. + let next_count = if (*root).limit_tuples > 0.0 { + ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)) + .min(node_count) + } else { + node_count + }; + let scan_selectivity = if total_rows > 0.0 { + (next_count / total_rows).clamp(1e-9, 1.0) + } else { + 1.0 + }; *index_startup_cost = 0.001 * node_count; *index_total_cost = 0.001 * node_count + next_count; - *index_selectivity = selectivity; + *index_selectivity = scan_selectivity; *index_correlation = 0.0; *index_pages = page_count; return; } *index_startup_cost = 0.0; *index_total_cost = 0.0; - *index_selectivity = selectivity; + *index_selectivity = filter_selectivity; *index_correlation = 0.0; *index_pages = 1.0; } diff --git a/tests/vchordrq/cost_estimator.slt b/tests/vchordrq/cost_estimator.slt new file mode 100644 index 00000000..0f4ad266 --- /dev/null +++ b/tests/vchordrq/cost_estimator.slt @@ -0,0 +1,426 @@ +# Tests for amcostestimate in src/index/vchordrq/am/mod.rs. +# +# The cost estimator must: +# 1. Use the baserel's planner-computed row estimate to derive filter +# selectivity, then report index selectivity from the candidate count +# that the scan is expected to retrieve. +# 2. Cap the candidate-processing term (`next_count`) by the IVF candidate +# budget (`node_count`), so absurdly small selectivity cannot inflate +# the reported cost past the work the index actually does. +# 3. Degrade safely when planner stats are missing or extreme. +# +# --------------------------------------------------------------------------- +# Setup +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = off; + +statement ok +SET max_parallel_workers_per_gather = 0; + +statement ok +SET random_page_cost = 1.1; + +statement ok +CREATE TABLE cost_test ( + id int PRIMARY KEY, + category int NOT NULL, + v vector(3) NOT NULL +); + +# 10 000 rows, 100 evenly-sized categories. category=k matches 1% of rows. +# Vectors are deterministic per row so the test is reproducible. +statement ok +INSERT INTO cost_test +SELECT i, + (i % 100), + ARRAY[ + (i % 97) / 97.0, + (i % 89) / 89.0, + (i % 83) / 83.0 + ]::real[]::vector +FROM generate_series(1, 10000) i; + +statement ok +CREATE INDEX cost_test_v ON cost_test USING vchordrq (v vector_l2_ops); + +statement ok +CREATE INDEX cost_test_cat ON cost_test (category); + +# A predicate with an explicitly high procost that PG can't optimize away. +# Three things matter here: +# * `LANGUAGE plpgsql` — prevents SQL inlining (an inlined `IS NOT NULL` +# check on a PRIMARY KEY column gets folded to `true` and dropped from +# the plan, defeating the test). +# * `COST 10000` — sets `pg_proc.procost` so cost_index() charges +# 10000 * cpu_operator_cost ≈ 25 cost-units per tuple of filter eval. +# This mimics PostGIS ST_DWithin, JSONB containment, expensive regex, +# etc. — the real-world shapes where the bug bites. +# * Body is `mod($1,7) <> 99` so the result is always true but PG can't +# prove that across the function boundary, so the filter survives +# constant-folding. +statement ok +CREATE OR REPLACE FUNCTION slow_true(int) RETURNS boolean + LANGUAGE plpgsql IMMUTABLE PARALLEL SAFE COST 10000 AS +$$ BEGIN RETURN mod($1, 7) <> 99; END $$; + +statement ok +ANALYZE cost_test; + +# Helper: returns the index name used at the top of the plan, or NULL if +# no Index Scan / Bitmap Index Scan appears. Robust across PG versions. +statement ok +CREATE OR REPLACE FUNCTION plan_top_index(query text) RETURNS text + LANGUAGE plpgsql AS $$ +DECLARE + line text; + m text[]; +BEGIN + FOR line IN EXECUTE 'EXPLAIN (FORMAT TEXT, COSTS OFF) ' || query LOOP + -- 'Index Scan using on ...' or 'Index Only Scan using on ...' + m := regexp_match(line, 'Index (?:Only )?Scan using (\S+) on '); + IF m IS NOT NULL THEN RETURN m[1]; END IF; + -- 'Bitmap Index Scan on ' + m := regexp_match(line, 'Bitmap Index Scan on (\S+)'); + IF m IS NOT NULL THEN RETURN m[1]; END IF; + END LOOP; + RETURN NULL; +END $$; + +# Helper: returns the top plan node's total cost. +statement ok +CREATE OR REPLACE FUNCTION top_total_cost(query text) RETURNS double precision + LANGUAGE plpgsql AS $$ +DECLARE + rec text; + m text[]; +BEGIN + FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP + m := regexp_match(rec, '\.\.([0-9]+(?:\.[0-9]+)?) rows='); + IF m IS NOT NULL THEN RETURN m[1]::double precision; END IF; + END LOOP; + RETURN NULL; +END $$; + +# Helper: returns total cost for the first plan node matching `pat`. +statement ok +CREATE OR REPLACE FUNCTION plan_node_total_cost(query text, pat text) RETURNS double precision + LANGUAGE plpgsql AS $$ +DECLARE + rec text; + m text[]; +BEGIN + FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP + IF rec LIKE pat THEN + m := regexp_match(rec, '\.\.([0-9]+(?:\.[0-9]+)?) rows='); + IF m IS NOT NULL THEN RETURN m[1]::double precision; END IF; + END IF; + END LOOP; + RETURN NULL; +END $$; + +# --------------------------------------------------------------------------- +# Case 1: pure ORDER BY, no filter. +# Vchord must win. (Passes pre- and post-fix; sanity check.) +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +# --------------------------------------------------------------------------- +# Case 2: tight filter (selectivity 1%) with an alternative index. +# Btree must still beat vchord — the IVF candidate budget is too large +# to be competitive at this selectivity. (Passes pre- and post-fix.) +# This guards against over-correction: the fix must not make vchord win +# on tight filters, which was the original motivation for PR #234. +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_cat + +statement ok +DROP INDEX cost_test_cat; + +# --------------------------------------------------------------------------- +# Case 3: selective heap filter combined with an expensive predicate and no +# supporting non-vector index. Pre-fix: index_selectivity=1.0 charges +# filter eval against every table row, so seq scan + sort can look cheaper. +# Post-fix: vchord estimates the candidates needed for the LIMIT and wins. +# The node-cost floor guards against under-pricing the scan as if it fetched +# only the final surviving rows. +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = on; + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10', + '%Index Scan using cost_test_v on cost_test%' +) > 10000; +---- +t + +statement ok +SET enable_seqscan = off; + +# --------------------------------------------------------------------------- +# Case 4: filter on a column with no supporting non-vector index, combined +# with the expensive predicate. The only alternatives are seq scan (disabled) +# or vchord. Vchord must be chosen. Pre-fix this still picks vchord because +# seq is disabled, but the *reported* cost must drop enough that LIMIT-aware +# planners above this node (joins, gather, partitions) don't overestimate. +# We assert plan shape + that the cost is bounded below 1e9 (disable_cost). +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +# disable_cost in PG 14-18 is 1.0e10; a healthy estimate is well below that. +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) < 1e9; +---- +t + +# --------------------------------------------------------------------------- +# Case 5: no LIMIT clause. PlannerInfo.limit_tuples is -1, so the +# estimator should fall back to the IVF candidate budget instead of doing +# LIMIT/selectivity math. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test ORDER BY v <-> ''[0,0,0]''::vector' +) IS NOT NULL; +---- +t + +# --------------------------------------------------------------------------- +# Case 6: LIMIT larger than the number of expected survivors. +# The `next_count.min(node_count)` clamp must hold — the AM cannot claim +# to process more candidates than the IVF visits. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE category = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000' +) IS NOT NULL; +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category + 0 = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category + 0 = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000', + '%Index Scan using cost_test_v on cost_test%' +) < 100000; +---- +t + +# --------------------------------------------------------------------------- +# Case 7: near-zero filter selectivity. The fix clamps to 1e-9 instead of +# 0 so 1/selectivity does not produce +Inf. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS NOT NULL; +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10', + '%Index Scan using cost_test_v on cost_test%' +) < 100000; +---- +t + +# --------------------------------------------------------------------------- +# Case 8: empty table. baserel->tuples is 0, fix falls back to +# selectivity = 1.0. Plan and result must not crash. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_empty (id int, v vector(3)); + +statement ok +CREATE INDEX ON cost_test_empty USING vchordrq (v vector_l2_ops); + +statement ok +ANALYZE cost_test_empty; + +query I +SELECT count(*) FROM ( + SELECT id FROM cost_test_empty ORDER BY v <-> '[0,0,0]'::vector LIMIT 10 +) t; +---- +0 + +# --------------------------------------------------------------------------- +# Case 9: never-ANALYZE'd table. baserel->tuples may be negative or zero +# in the planner's view. The fallback in the fix returns 1.0, matching the +# pre-fix behavior exactly so we don't regress on cold tables. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_cold (id int, v vector(3)); + +statement ok +INSERT INTO cost_test_cold +SELECT i, ARRAY[(i%7)/7.0, (i%11)/11.0, (i%13)/13.0]::real[]::vector +FROM generate_series(1, 100) i; + +statement ok +CREATE INDEX ON cost_test_cold USING vchordrq (v vector_l2_ops); + +# Intentionally no ANALYZE. + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test_cold ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS NOT NULL; +---- +t + +# --------------------------------------------------------------------------- +# Case 10: partial vchord index. The filter-selectivity estimate reflects +# baserestrictinfo, including clauses that overlap the partial-index +# predicate. We verify that this builds, plans, and runs. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_partial (id int, kind int, v vector(3)); + +statement ok +INSERT INTO cost_test_partial +SELECT i, (i % 3), ARRAY[(i%7)/7.0, (i%11)/11.0, (i%13)/13.0]::real[]::vector +FROM generate_series(1, 2000) i; + +statement ok +CREATE INDEX cost_test_partial_v ON cost_test_partial USING vchordrq (v vector_l2_ops) + WHERE kind = 0; + +statement ok +ANALYZE cost_test_partial; + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test_partial WHERE kind = 0 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 5' +); +---- +cost_test_partial_v + +# --------------------------------------------------------------------------- +# Case 11: vchordrq.enable_scan = off must still disable the path +# regardless of selectivity. Re-enable seqscan inside this case so there is +# a non-disabled alternative for the planner to fall back to — otherwise on +# PG18 every path looks "disabled" and the planner picks vchord by default. +# Use IS DISTINCT FROM so a seq-scan top (plan_top_index returns NULL) +# still compares as "not vchord." +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = on; + +statement ok +SET vchordrq.enable_scan = off; + +query B +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS DISTINCT FROM 'cost_test_v'; +---- +t + +statement ok +RESET vchordrq.enable_scan; + +statement ok +RESET enable_seqscan; + +# --------------------------------------------------------------------------- +# Cleanup +# --------------------------------------------------------------------------- + +statement ok +DROP TABLE cost_test; + +statement ok +DROP TABLE cost_test_empty; + +statement ok +DROP TABLE cost_test_cold; + +statement ok +DROP TABLE cost_test_partial; + +statement ok +DROP FUNCTION slow_true(int); + +statement ok +DROP FUNCTION plan_top_index(text); + +statement ok +DROP FUNCTION plan_node_total_cost(text, text); + +statement ok +DROP FUNCTION top_total_cost(text); + +statement ok +RESET enable_seqscan; + +statement ok +RESET max_parallel_workers_per_gather; + +statement ok +RESET random_page_cost; From ddcd0190fb4af5065508366b597dfe96f734d17b Mon Sep 17 00:00:00 2001 From: Amey Pawar <138877912+ameyypawar@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:57:45 +0530 Subject: [PATCH 295/324] fix(vchordg): avoid panic when validating an empty alpha option `validate_alpha` indexes `alpha[0]` after the sorted/range checks, both of which pass on an empty slice (`[].is_sorted()` and `[].iter().all(..)` are true). So an empty `alpha` (e.g. `WITH (options='[index] alpha = []')`) panics with an index-out-of-bounds instead of returning a ValidationError. Use `alpha.first()` so the empty case yields the existing "`alpha` should contain `1.0`" error. Behavior is unchanged for every non-empty input. Add unit tests covering the empty slice plus valid/invalid cases. --- crates/vchordg/src/types.rs | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index 62ec3cc4..5a5a6169 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -56,7 +56,7 @@ impl VchordgIndexOptions { if !alpha.iter().all(|x| (1.0..2.0).contains(x)) { return Err(ValidationError::new("alpha is too large or too small")); } - if alpha[0] != 1.0 { + if alpha.first() != Some(&1.0) { return Err(ValidationError::new("`alpha` should contain `1.0`")); } Ok(()) @@ -144,3 +144,25 @@ impl Structure { self.children.is_empty() } } + +#[cfg(test)] +mod tests { + use super::VchordgIndexOptions; + + #[test] + fn validate_alpha_handles_empty_without_panicking() { + let empty: &[f32] = &[]; + assert!(VchordgIndexOptions::validate_alpha(empty).is_err()); + } + + #[test] + fn validate_alpha_accepts_default() { + assert!(VchordgIndexOptions::validate_alpha(&[1.0, 1.2]).is_ok()); + } + + #[test] + fn validate_alpha_rejects_unsorted_or_out_of_range() { + assert!(VchordgIndexOptions::validate_alpha(&[1.2, 1.0]).is_err()); + assert!(VchordgIndexOptions::validate_alpha(&[2.0]).is_err()); + } +} From e39b10eb985c530201a44674028dfa961739a497 Mon Sep 17 00:00:00 2001 From: Amey Pawar <138877912+ameyypawar@users.noreply.github.com> Date: Thu, 9 Jul 2026 07:52:57 +0530 Subject: [PATCH 296/324] fix(gucs): report an error instead of panicking on an empty `probes` entry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `vchordrq.probes` parser panics via `.expect("empty probes")` when an entry is empty (e.g. `SET vchordrq.probes = ',5'` or `'1,,2'`), even though the sibling match arm already reports malformed input cleanly with `pgrx::error!`. Handle the empty entry the same way — a proper error rather than a panic. Behavior is unchanged for all valid input. --- src/index/gucs.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 8c0e9fd3..f2724593 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -388,7 +388,10 @@ pub unsafe fn vchordrq_probes(index: pgrx::pg_sys::Relation) -> Vec { for &c in value.to_bytes() { match c { b' ' => continue, - b',' => result.push(current.take().expect("empty probes")), + b',' => match current.take() { + Some(value) => result.push(value), + None => pgrx::error!("empty entry in probes"), + }, b'0'..=b'9' => { if let Some(x) = current.as_mut() { *x = *x * 10 + (c - b'0') as u32; From 286df294ff8cc5a6ba1d8ec6eb58cf8b42aa015f Mon Sep 17 00:00:00 2001 From: spenc-r Date: Thu, 21 May 2026 17:25:30 +0000 Subject: [PATCH 297/324] feat(vchordrq): estimate filtered candidate selectivity Use the planner's filtered row estimate as filter selectivity for deciding how many IVFFlat candidates are needed to satisfy LIMIT after heap-side quals. Report indexSelectivity from the expected retrieved candidate count instead of the final survivor fraction, so cost_index() prices heap fetch and qpqual evaluation against candidate rows rather than every table row or only the rows that pass filters. Clamp the candidate-processing term by node_count, and tighten the cost estimator sqllogictest coverage around btree alternatives, candidate-cost floors, and extreme selectivity. --- src/index/vchordrq/am/mod.rs | 58 ++-- tests/vchordrq/cost_estimator.slt | 426 ++++++++++++++++++++++++++++++ 2 files changed, 467 insertions(+), 17 deletions(-) create mode 100644 tests/vchordrq/cost_estimator.slt diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 5040b587..f995c600 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -292,20 +292,31 @@ pub unsafe extern "C-unwind" fn amcostestimate( *index_pages = 1.0; return; } - let selectivity = { - use pgrx::pg_sys::{ - JoinType, add_predicate_to_index_quals, clauselist_selectivity, - get_quals_from_indexclauses, + // Vchordrq indexes only the vector column, so ordinary filters on + // other columns are heap-side quals rather than index quals. We use + // the planner's filtered row estimate to decide how many distance + // candidates we likely need to produce enough survivors for LIMIT. + // + // Keep this separate from the value returned as `indexSelectivity`: + // `cost_index()` interprets that output as the fraction of parent + // table rows the index scan retrieves, not the fraction surviving all + // filters. With LIMIT, the retrieved fraction is better represented + // by the candidate count computed below. + let (total_rows, filter_selectivity) = { + let baserel = (*index_opt_info).rel; + let total_rows = (*baserel).tuples; + let param_info = (*path).path.param_info; + let filtered_rows = if !param_info.is_null() { + (*param_info).ppi_rows + } else { + (*baserel).rows }; - let index_quals = get_quals_from_indexclauses((*path).indexclauses); - let selectivity_quals = add_predicate_to_index_quals(index_opt_info, index_quals); - clauselist_selectivity( - root, - selectivity_quals, - (*(*index_opt_info).rel).relid as _, - JoinType::JOIN_INNER, - std::ptr::null_mut(), - ) + let filter_selectivity = if total_rows > 0.0 { + (filtered_rows / total_rows).clamp(1e-9, 1.0) + } else { + 1.0 + }; + (total_rows, filter_selectivity) }; // index exists if !(*index_opt_info).hypothetical { @@ -367,18 +378,31 @@ pub unsafe extern "C-unwind" fn amcostestimate( pages += cost.cells[0] as f64; pages }; - let next_count = - f64::max(1.0, (*root).limit_tuples) * f64::min(1000.0, 1.0 / selectivity); + // `next_count` represents candidates we expect to process to + // surface `limit_tuples` survivors after filter rejection. Clamp + // by `node_count` so the estimate cannot exceed the candidates + // the IVF visits at the configured probe count. + let next_count = if (*root).limit_tuples > 0.0 { + ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)) + .min(node_count) + } else { + node_count + }; + let scan_selectivity = if total_rows > 0.0 { + (next_count / total_rows).clamp(1e-9, 1.0) + } else { + 1.0 + }; *index_startup_cost = 0.001 * node_count; *index_total_cost = 0.001 * node_count + next_count; - *index_selectivity = selectivity; + *index_selectivity = scan_selectivity; *index_correlation = 0.0; *index_pages = page_count; return; } *index_startup_cost = 0.0; *index_total_cost = 0.0; - *index_selectivity = selectivity; + *index_selectivity = filter_selectivity; *index_correlation = 0.0; *index_pages = 1.0; } diff --git a/tests/vchordrq/cost_estimator.slt b/tests/vchordrq/cost_estimator.slt new file mode 100644 index 00000000..0f4ad266 --- /dev/null +++ b/tests/vchordrq/cost_estimator.slt @@ -0,0 +1,426 @@ +# Tests for amcostestimate in src/index/vchordrq/am/mod.rs. +# +# The cost estimator must: +# 1. Use the baserel's planner-computed row estimate to derive filter +# selectivity, then report index selectivity from the candidate count +# that the scan is expected to retrieve. +# 2. Cap the candidate-processing term (`next_count`) by the IVF candidate +# budget (`node_count`), so absurdly small selectivity cannot inflate +# the reported cost past the work the index actually does. +# 3. Degrade safely when planner stats are missing or extreme. +# +# --------------------------------------------------------------------------- +# Setup +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = off; + +statement ok +SET max_parallel_workers_per_gather = 0; + +statement ok +SET random_page_cost = 1.1; + +statement ok +CREATE TABLE cost_test ( + id int PRIMARY KEY, + category int NOT NULL, + v vector(3) NOT NULL +); + +# 10 000 rows, 100 evenly-sized categories. category=k matches 1% of rows. +# Vectors are deterministic per row so the test is reproducible. +statement ok +INSERT INTO cost_test +SELECT i, + (i % 100), + ARRAY[ + (i % 97) / 97.0, + (i % 89) / 89.0, + (i % 83) / 83.0 + ]::real[]::vector +FROM generate_series(1, 10000) i; + +statement ok +CREATE INDEX cost_test_v ON cost_test USING vchordrq (v vector_l2_ops); + +statement ok +CREATE INDEX cost_test_cat ON cost_test (category); + +# A predicate with an explicitly high procost that PG can't optimize away. +# Three things matter here: +# * `LANGUAGE plpgsql` — prevents SQL inlining (an inlined `IS NOT NULL` +# check on a PRIMARY KEY column gets folded to `true` and dropped from +# the plan, defeating the test). +# * `COST 10000` — sets `pg_proc.procost` so cost_index() charges +# 10000 * cpu_operator_cost ≈ 25 cost-units per tuple of filter eval. +# This mimics PostGIS ST_DWithin, JSONB containment, expensive regex, +# etc. — the real-world shapes where the bug bites. +# * Body is `mod($1,7) <> 99` so the result is always true but PG can't +# prove that across the function boundary, so the filter survives +# constant-folding. +statement ok +CREATE OR REPLACE FUNCTION slow_true(int) RETURNS boolean + LANGUAGE plpgsql IMMUTABLE PARALLEL SAFE COST 10000 AS +$$ BEGIN RETURN mod($1, 7) <> 99; END $$; + +statement ok +ANALYZE cost_test; + +# Helper: returns the index name used at the top of the plan, or NULL if +# no Index Scan / Bitmap Index Scan appears. Robust across PG versions. +statement ok +CREATE OR REPLACE FUNCTION plan_top_index(query text) RETURNS text + LANGUAGE plpgsql AS $$ +DECLARE + line text; + m text[]; +BEGIN + FOR line IN EXECUTE 'EXPLAIN (FORMAT TEXT, COSTS OFF) ' || query LOOP + -- 'Index Scan using on ...' or 'Index Only Scan using on ...' + m := regexp_match(line, 'Index (?:Only )?Scan using (\S+) on '); + IF m IS NOT NULL THEN RETURN m[1]; END IF; + -- 'Bitmap Index Scan on ' + m := regexp_match(line, 'Bitmap Index Scan on (\S+)'); + IF m IS NOT NULL THEN RETURN m[1]; END IF; + END LOOP; + RETURN NULL; +END $$; + +# Helper: returns the top plan node's total cost. +statement ok +CREATE OR REPLACE FUNCTION top_total_cost(query text) RETURNS double precision + LANGUAGE plpgsql AS $$ +DECLARE + rec text; + m text[]; +BEGIN + FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP + m := regexp_match(rec, '\.\.([0-9]+(?:\.[0-9]+)?) rows='); + IF m IS NOT NULL THEN RETURN m[1]::double precision; END IF; + END LOOP; + RETURN NULL; +END $$; + +# Helper: returns total cost for the first plan node matching `pat`. +statement ok +CREATE OR REPLACE FUNCTION plan_node_total_cost(query text, pat text) RETURNS double precision + LANGUAGE plpgsql AS $$ +DECLARE + rec text; + m text[]; +BEGIN + FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP + IF rec LIKE pat THEN + m := regexp_match(rec, '\.\.([0-9]+(?:\.[0-9]+)?) rows='); + IF m IS NOT NULL THEN RETURN m[1]::double precision; END IF; + END IF; + END LOOP; + RETURN NULL; +END $$; + +# --------------------------------------------------------------------------- +# Case 1: pure ORDER BY, no filter. +# Vchord must win. (Passes pre- and post-fix; sanity check.) +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +# --------------------------------------------------------------------------- +# Case 2: tight filter (selectivity 1%) with an alternative index. +# Btree must still beat vchord — the IVF candidate budget is too large +# to be competitive at this selectivity. (Passes pre- and post-fix.) +# This guards against over-correction: the fix must not make vchord win +# on tight filters, which was the original motivation for PR #234. +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_cat + +statement ok +DROP INDEX cost_test_cat; + +# --------------------------------------------------------------------------- +# Case 3: selective heap filter combined with an expensive predicate and no +# supporting non-vector index. Pre-fix: index_selectivity=1.0 charges +# filter eval against every table row, so seq scan + sort can look cheaper. +# Post-fix: vchord estimates the candidates needed for the LIMIT and wins. +# The node-cost floor guards against under-pricing the scan as if it fetched +# only the final surviving rows. +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = on; + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10', + '%Index Scan using cost_test_v on cost_test%' +) > 10000; +---- +t + +statement ok +SET enable_seqscan = off; + +# --------------------------------------------------------------------------- +# Case 4: filter on a column with no supporting non-vector index, combined +# with the expensive predicate. The only alternatives are seq scan (disabled) +# or vchord. Vchord must be chosen. Pre-fix this still picks vchord because +# seq is disabled, but the *reported* cost must drop enough that LIMIT-aware +# planners above this node (joins, gather, partitions) don't overestimate. +# We assert plan shape + that the cost is bounded below 1e9 (disable_cost). +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +# disable_cost in PG 14-18 is 1.0e10; a healthy estimate is well below that. +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) < 1e9; +---- +t + +# --------------------------------------------------------------------------- +# Case 5: no LIMIT clause. PlannerInfo.limit_tuples is -1, so the +# estimator should fall back to the IVF candidate budget instead of doing +# LIMIT/selectivity math. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test ORDER BY v <-> ''[0,0,0]''::vector' +) IS NOT NULL; +---- +t + +# --------------------------------------------------------------------------- +# Case 6: LIMIT larger than the number of expected survivors. +# The `next_count.min(node_count)` clamp must hold — the AM cannot claim +# to process more candidates than the IVF visits. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE category = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000' +) IS NOT NULL; +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category + 0 = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category + 0 = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000', + '%Index Scan using cost_test_v on cost_test%' +) < 100000; +---- +t + +# --------------------------------------------------------------------------- +# Case 7: near-zero filter selectivity. The fix clamps to 1e-9 instead of +# 0 so 1/selectivity does not produce +Inf. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS NOT NULL; +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10', + '%Index Scan using cost_test_v on cost_test%' +) < 100000; +---- +t + +# --------------------------------------------------------------------------- +# Case 8: empty table. baserel->tuples is 0, fix falls back to +# selectivity = 1.0. Plan and result must not crash. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_empty (id int, v vector(3)); + +statement ok +CREATE INDEX ON cost_test_empty USING vchordrq (v vector_l2_ops); + +statement ok +ANALYZE cost_test_empty; + +query I +SELECT count(*) FROM ( + SELECT id FROM cost_test_empty ORDER BY v <-> '[0,0,0]'::vector LIMIT 10 +) t; +---- +0 + +# --------------------------------------------------------------------------- +# Case 9: never-ANALYZE'd table. baserel->tuples may be negative or zero +# in the planner's view. The fallback in the fix returns 1.0, matching the +# pre-fix behavior exactly so we don't regress on cold tables. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_cold (id int, v vector(3)); + +statement ok +INSERT INTO cost_test_cold +SELECT i, ARRAY[(i%7)/7.0, (i%11)/11.0, (i%13)/13.0]::real[]::vector +FROM generate_series(1, 100) i; + +statement ok +CREATE INDEX ON cost_test_cold USING vchordrq (v vector_l2_ops); + +# Intentionally no ANALYZE. + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test_cold ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS NOT NULL; +---- +t + +# --------------------------------------------------------------------------- +# Case 10: partial vchord index. The filter-selectivity estimate reflects +# baserestrictinfo, including clauses that overlap the partial-index +# predicate. We verify that this builds, plans, and runs. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_partial (id int, kind int, v vector(3)); + +statement ok +INSERT INTO cost_test_partial +SELECT i, (i % 3), ARRAY[(i%7)/7.0, (i%11)/11.0, (i%13)/13.0]::real[]::vector +FROM generate_series(1, 2000) i; + +statement ok +CREATE INDEX cost_test_partial_v ON cost_test_partial USING vchordrq (v vector_l2_ops) + WHERE kind = 0; + +statement ok +ANALYZE cost_test_partial; + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test_partial WHERE kind = 0 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 5' +); +---- +cost_test_partial_v + +# --------------------------------------------------------------------------- +# Case 11: vchordrq.enable_scan = off must still disable the path +# regardless of selectivity. Re-enable seqscan inside this case so there is +# a non-disabled alternative for the planner to fall back to — otherwise on +# PG18 every path looks "disabled" and the planner picks vchord by default. +# Use IS DISTINCT FROM so a seq-scan top (plan_top_index returns NULL) +# still compares as "not vchord." +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = on; + +statement ok +SET vchordrq.enable_scan = off; + +query B +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS DISTINCT FROM 'cost_test_v'; +---- +t + +statement ok +RESET vchordrq.enable_scan; + +statement ok +RESET enable_seqscan; + +# --------------------------------------------------------------------------- +# Cleanup +# --------------------------------------------------------------------------- + +statement ok +DROP TABLE cost_test; + +statement ok +DROP TABLE cost_test_empty; + +statement ok +DROP TABLE cost_test_cold; + +statement ok +DROP TABLE cost_test_partial; + +statement ok +DROP FUNCTION slow_true(int); + +statement ok +DROP FUNCTION plan_top_index(text); + +statement ok +DROP FUNCTION plan_node_total_cost(text, text); + +statement ok +DROP FUNCTION top_total_cost(text); + +statement ok +RESET enable_seqscan; + +statement ok +RESET max_parallel_workers_per_gather; + +statement ok +RESET random_page_cost; From 4dc1bbedbe85369205fb4b05bd6d797798f4214b Mon Sep 17 00:00:00 2001 From: Amey Pawar <138877912+ameyypawar@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:57:45 +0530 Subject: [PATCH 298/324] fix(vchordg): avoid panic when validating an empty alpha option `validate_alpha` indexes `alpha[0]` after the sorted/range checks, both of which pass on an empty slice (`[].is_sorted()` and `[].iter().all(..)` are true). So an empty `alpha` (e.g. `WITH (options='[index] alpha = []')`) panics with an index-out-of-bounds instead of returning a ValidationError. Use `alpha.first()` so the empty case yields the existing "`alpha` should contain `1.0`" error. Behavior is unchanged for every non-empty input. Add unit tests covering the empty slice plus valid/invalid cases. --- crates/vchordg/src/types.rs | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index 62ec3cc4..5a5a6169 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -56,7 +56,7 @@ impl VchordgIndexOptions { if !alpha.iter().all(|x| (1.0..2.0).contains(x)) { return Err(ValidationError::new("alpha is too large or too small")); } - if alpha[0] != 1.0 { + if alpha.first() != Some(&1.0) { return Err(ValidationError::new("`alpha` should contain `1.0`")); } Ok(()) @@ -144,3 +144,25 @@ impl Structure { self.children.is_empty() } } + +#[cfg(test)] +mod tests { + use super::VchordgIndexOptions; + + #[test] + fn validate_alpha_handles_empty_without_panicking() { + let empty: &[f32] = &[]; + assert!(VchordgIndexOptions::validate_alpha(empty).is_err()); + } + + #[test] + fn validate_alpha_accepts_default() { + assert!(VchordgIndexOptions::validate_alpha(&[1.0, 1.2]).is_ok()); + } + + #[test] + fn validate_alpha_rejects_unsorted_or_out_of_range() { + assert!(VchordgIndexOptions::validate_alpha(&[1.2, 1.0]).is_err()); + assert!(VchordgIndexOptions::validate_alpha(&[2.0]).is_err()); + } +} From 3545e75b71aa0c1e03a830633d04d59154fe8de5 Mon Sep 17 00:00:00 2001 From: Amey Pawar <138877912+ameyypawar@users.noreply.github.com> Date: Thu, 9 Jul 2026 07:52:57 +0530 Subject: [PATCH 299/324] fix(gucs): report an error instead of panicking on an empty `probes` entry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `vchordrq.probes` parser panics via `.expect("empty probes")` when an entry is empty (e.g. `SET vchordrq.probes = ',5'` or `'1,,2'`), even though the sibling match arm already reports malformed input cleanly with `pgrx::error!`. Handle the empty entry the same way — a proper error rather than a panic. Behavior is unchanged for all valid input. --- src/index/gucs.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 8c0e9fd3..f2724593 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -388,7 +388,10 @@ pub unsafe fn vchordrq_probes(index: pgrx::pg_sys::Relation) -> Vec { for &c in value.to_bytes() { match c { b' ' => continue, - b',' => result.push(current.take().expect("empty probes")), + b',' => match current.take() { + Some(value) => result.push(value), + None => pgrx::error!("empty entry in probes"), + }, b'0'..=b'9' => { if let Some(x) = current.as_mut() { *x = *x * 10 + (c - b'0') as u32; From d8f7c8b0c94593c7fa068def3a7e8330aca34df8 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Mon, 13 Jul 2026 18:40:47 +0800 Subject: [PATCH 300/324] feat(vchordrq): add exact TileMaxSim reranking --- Cargo.lock | 1 + Cargo.toml | 3 + crates/vchordrq/src/build.rs | 1 + crates/vchordrq/src/bulkdelete.rs | 28 +- crates/vchordrq/src/cost.rs | 8 +- crates/vchordrq/src/lib.rs | 6 + crates/vchordrq/src/maxsim_cost.rs | 206 +++ crates/vchordrq/src/statistics.rs | 30 + crates/vchordrq/src/tuples.rs | 151 +- crates/vchordrq/src/types.rs | 34 + devtools/test_tilemaxsim_reference_sidecar.py | 383 +++++ devtools/tilemaxsim_reference_sidecar.py | 780 +++++++++ services/Dockerfile.tilemaxsim | 12 + services/benchmark_tilemaxsim_cuda.py | 152 ++ services/build_tilemaxsim_tensor_cache.py | 235 +++ services/test_tilemaxsim_cuda_sidecar.py | 313 ++++ services/tilemaxsim_cuda_sidecar.py | 802 +++++++++ src/datatype/mod.rs | 14 + src/datatype/operators_halfvec.rs | 8 + src/datatype/operators_rabitq4.rs | 8 + src/datatype/operators_rabitq8.rs | 8 + src/datatype/operators_vector.rs | 8 + src/index/fetcher.rs | 108 +- src/index/gucs.rs | 139 ++ src/index/vchordrq/am/am_build.rs | 31 +- src/index/vchordrq/am/mod.rs | 89 +- src/index/vchordrq/dispatch.rs | 27 +- src/index/vchordrq/opclass.rs | 8 + src/index/vchordrq/scanners/maxsim.rs | 1345 ++++++++------- .../vchordrq/scanners/maxsim/candidate.rs | 192 +++ .../vchordrq/scanners/maxsim/external.rs | 551 ++++++ src/index/vchordrq/scanners/maxsim/gpu.rs | 1495 +++++++++++++++++ src/index/vchordrq/scanners/maxsim/rerank.rs | 315 ++++ src/index/vchordrq/scanners/maxsim/search.rs | 578 +++++++ src/index/vchordrq/scanners/mod.rs | 8 + src/sql/finalize.sql | 821 +++++++++ tests/vchordrq/cost_estimator.slt | 171 ++ tests/vchordrq/maxsim_correctness.slt | 221 +++ tests/vchordrq/maxsim_source_registry.slt | 341 ++++ 39 files changed, 8933 insertions(+), 698 deletions(-) create mode 100644 crates/vchordrq/src/maxsim_cost.rs create mode 100644 crates/vchordrq/src/statistics.rs create mode 100644 devtools/test_tilemaxsim_reference_sidecar.py create mode 100644 devtools/tilemaxsim_reference_sidecar.py create mode 100644 services/Dockerfile.tilemaxsim create mode 100644 services/benchmark_tilemaxsim_cuda.py create mode 100644 services/build_tilemaxsim_tensor_cache.py create mode 100644 services/test_tilemaxsim_cuda_sidecar.py create mode 100644 services/tilemaxsim_cuda_sidecar.py create mode 100644 src/index/vchordrq/scanners/maxsim/candidate.rs create mode 100644 src/index/vchordrq/scanners/maxsim/external.rs create mode 100644 src/index/vchordrq/scanners/maxsim/gpu.rs create mode 100644 src/index/vchordrq/scanners/maxsim/rerank.rs create mode 100644 src/index/vchordrq/scanners/maxsim/search.rs create mode 100644 tests/vchordrq/maxsim_correctness.slt create mode 100644 tests/vchordrq/maxsim_source_registry.slt diff --git a/Cargo.lock b/Cargo.lock index f324f59e..80722f2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1937,6 +1937,7 @@ dependencies = [ "index", "index_accessor", "k_means", + "libc", "mimalloc", "paste", "pgrx", diff --git a/Cargo.toml b/Cargo.toml index 5868e4d6..32654e1f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,6 +55,9 @@ zerocopy.workspace = true [target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos")))'.dependencies] mimalloc = { version = "0.1.49", features = ["local_dynamic_tls"] } +[target.'cfg(unix)'.dependencies] +libc = "0.2.182" + [lints] workspace = true diff --git a/crates/vchordrq/src/build.rs b/crates/vchordrq/src/build.rs index a7371353..d622fbc6 100644 --- a/crates/vchordrq/src/build.rs +++ b/crates/vchordrq/src/build.rs @@ -119,6 +119,7 @@ pub fn build( height_of_root: structures.len() as u32, is_residual, rerank_in_heap: vchordrq_options.rerank_in_table, + indexed_vectors: Some(0), centroids_first: centroids.first(), vectors_first: vectors, centroid_prefetch: pointer_of_centroids diff --git a/crates/vchordrq/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs index ae54ec54..0aae88fe 100644 --- a/crates/vchordrq/src/bulkdelete.rs +++ b/crates/vchordrq/src/bulkdelete.rs @@ -25,7 +25,8 @@ pub fn bulkdelete( index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, -) where +) -> u64 +where R::Page: Page, { let meta_guard = index.read(0); @@ -38,6 +39,8 @@ pub fn bulkdelete( drop(meta_guard); + let mut live = 0_u64; + let step = |state: State| { let mut results = Vec::new(); for first in state { @@ -64,19 +67,21 @@ pub fn bulkdelete( while current != u32::MAX { check(); let read = index.read(current); - let flag = 'flag: { + let (flag, page_live) = 'scan: { + let mut page_live = 0_u64; for i in 1..=read.len() { let bytes = read.get(i).expect("data corruption"); let tuple = FrozenTuple::deserialize_ref(bytes); if let FrozenTupleReader::_0(tuple) = tuple { for p in tuple.payload().iter() { if Some(true) == p.map(&callback) { - break 'flag true; + break 'scan (true, 0); } + page_live += u64::from(p.is_some()); } } } - false + (false, page_live) }; if flag { drop(read); @@ -89,9 +94,12 @@ pub fn bulkdelete( if Some(true) == p.map(&callback) { *p = None; } + live += u64::from(p.is_some()); } } } + } else { + live += page_live; } current = directory.next().unwrap_or(u32::MAX); } @@ -101,16 +109,18 @@ pub fn bulkdelete( while current != u32::MAX { check(); let read = index.read(current); - let flag = 'flag: { + let (flag, page_live) = 'scan: { + let mut page_live = 0_u64; for i in 1..=read.len() { let bytes = read.get(i).expect("data corruption"); let tuple = AppendableTuple::deserialize_ref(bytes); let p = tuple.payload(); if Some(true) == p.map(&callback) { - break 'flag true; + break 'scan (true, 0); } + page_live += u64::from(p.is_some()); } - false + (false, page_live) }; if flag { drop(read); @@ -122,14 +132,18 @@ pub fn bulkdelete( if Some(true) == p.map(&callback) { *p = None; } + live += u64::from(p.is_some()); } current = write.get_opaque().next; } else { + live += page_live; current = read.get_opaque().next; } } } } + + live } pub fn bulkdelete_vectors( diff --git a/crates/vchordrq/src/cost.rs b/crates/vchordrq/src/cost.rs index 60579401..c8e40252 100644 --- a/crates/vchordrq/src/cost.rs +++ b/crates/vchordrq/src/cost.rs @@ -18,6 +18,7 @@ use index::relation::{Page, RelationRead}; pub struct Cost { pub dim: u32, pub cells: Vec, + pub indexed_vectors: Option, } #[must_use] @@ -27,8 +28,13 @@ pub fn cost(index: &R) -> Cost { let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dim = meta_tuple.dim(); let cells = meta_tuple.cells().to_vec(); + let indexed_vectors = meta_tuple.indexed_vectors(); drop(meta_guard); - Cost { dim, cells } + Cost { + dim, + cells, + indexed_vectors, + } } diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index d0efef05..b60c2e3b 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -24,9 +24,11 @@ mod freepages; mod insert; mod linked_vec; mod maintain; +mod maxsim_cost; mod prewarm; mod rerank; mod search; +mod statistics; mod tape; mod tape_writer; mod tuples; @@ -43,9 +45,13 @@ pub use cost::cost; pub use fast_heap::FastHeap; pub use insert::{InsertChooser, insert, insert_vector}; pub use maintain::{MaintainChooser, maintain}; +pub use maxsim_cost::{ + MaxsimCostBackend, MaxsimCostEstimate, MaxsimCostInput, estimate_maxsim_cost, +}; pub use prewarm::prewarm; pub use rerank::{how, rerank_heap, rerank_index}; pub use search::{default_search, maxsim_search}; +pub use statistics::set_indexed_vectors; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; diff --git a/crates/vchordrq/src/maxsim_cost.rs b/crates/vchordrq/src/maxsim_cost.rs new file mode 100644 index 00000000..ca843d10 --- /dev/null +++ b/crates/vchordrq/src/maxsim_cost.rs @@ -0,0 +1,206 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the terms of the ELv2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[derive(Clone, Copy, Debug)] +pub enum MaxsimCostBackend { + CoarseOnly, + CpuExact, + Gpu, + Auto, +} + +#[derive(Clone, Copy, Debug)] +pub struct MaxsimCostInput { + pub heap_rows: f64, + pub index_tokens: f64, + pub token_nodes_per_query: f64, + pub base_index_pages: f64, + pub dimension: u32, + pub element_bits: u32, + pub query_tokens: u32, + pub limit_tuples: Option, + pub filter_selectivity: f64, + pub candidate_limit: Option, + pub backend: MaxsimCostBackend, +} + +#[derive(Clone, Copy, Debug)] +pub struct MaxsimCostEstimate { + pub startup_cost: f64, + pub total_cost: f64, + pub selectivity: f64, + pub index_pages: f64, +} + +pub fn estimate_maxsim_cost(input: MaxsimCostInput) -> MaxsimCostEstimate { + let heap_rows = input.heap_rows.max(1.0); + let index_tokens = input.index_tokens.max(heap_rows); + let query_tokens = f64::from(input.query_tokens.max(1)); + let average_document_tokens = (index_tokens / heap_rows).clamp(1.0, 65_536.0); + let token_visits = input.token_nodes_per_query.max(1.0) * query_tokens; + + // We do not yet persist page-level candidate statistics. Until then, use a + // conservative occupancy estimate: token visits are sampled from the token + // index, and a page is a candidate when at least one of its average tokens + // is visited. This is intentionally bounded by the heap row count. + let token_visit_fraction = (token_visits / index_tokens).clamp(0.0, 1.0); + let candidate_probability = 1.0 - (1.0 - token_visit_fraction).powf(average_document_tokens); + let generated_pages = (heap_rows * candidate_probability).clamp(1.0, heap_rows); + + let filtered_limit = input + .limit_tuples + .map(|limit| limit.max(1.0) / input.filter_selectivity.clamp(1e-9, 1.0)); + let exact_candidate_count = input + .candidate_limit + .map_or(generated_pages, |limit| { + f64::from(limit).min(generated_pages) + }) + .clamp(1.0, heap_rows); + let returned_pages = match input.backend { + MaxsimCostBackend::CoarseOnly => filtered_limit + .unwrap_or(generated_pages) + .min(generated_pages), + MaxsimCostBackend::CpuExact | MaxsimCostBackend::Gpu | MaxsimCostBackend::Auto => { + exact_candidate_count + } + }; + + // Candidate generation and aggregation are eager in the current scanner, + // so their full work belongs to startup cost even when SQL has a small + // LIMIT. The constants are deliberately conservative placeholders until + // committed corpus benchmarks replace them with fitted values. + let search_cost = 0.001 * token_visits; + let aggregation_cost = 0.01 * token_visits + 0.05 * generated_pages; + let exact_components = exact_candidate_count + * average_document_tokens + * query_tokens + * f64::from(input.dimension.max(1)); + let tensor_bytes = exact_candidate_count + * average_document_tokens + * f64::from(input.dimension.max(1)) + * f64::from(input.element_bits.max(1)) + / 8.0; + let cpu_exact_cost = exact_candidate_count + exact_components * 1e-6; + let gpu_exact_cost = 5.0 + tensor_bytes * 1e-7 + exact_components * 5e-8; + let backend_cost = match input.backend { + MaxsimCostBackend::CoarseOnly => 0.0, + MaxsimCostBackend::CpuExact => cpu_exact_cost, + MaxsimCostBackend::Gpu => gpu_exact_cost, + // Price a small but nonzero fallback risk. Runtime still performs a + // complete CPU rerank on every GPU failure. + MaxsimCostBackend::Auto => gpu_exact_cost + 0.05 * cpu_exact_cost, + }; + let startup_cost = search_cost + aggregation_cost + backend_cost; + let total_cost = startup_cost + returned_pages; + let selectivity = (returned_pages / heap_rows).clamp(1e-9, 1.0); + let index_pages = + input.base_index_pages.max(1.0) * (1.0 + 0.25 * (query_tokens - 1.0).max(0.0)); + + MaxsimCostEstimate { + startup_cost, + total_cost, + selectivity, + index_pages, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn input(backend: MaxsimCostBackend) -> MaxsimCostInput { + MaxsimCostInput { + heap_rows: 34_054.0, + index_tokens: 25_438_338.0, + token_nodes_per_query: 10_000.0, + base_index_pages: 20_000.0, + dimension: 320, + element_bits: 16, + query_tokens: 32, + limit_tuples: Some(20.0), + filter_selectivity: 1.0, + candidate_limit: Some(256), + backend, + } + } + + #[test] + fn maxsim_is_never_zero_cost() { + for backend in [ + MaxsimCostBackend::CoarseOnly, + MaxsimCostBackend::CpuExact, + MaxsimCostBackend::Gpu, + MaxsimCostBackend::Auto, + ] { + let estimate = estimate_maxsim_cost(input(backend)); + assert!(estimate.startup_cost > 0.0); + assert!(estimate.total_cost >= estimate.startup_cost); + assert!((1e-9..=1.0).contains(&estimate.selectivity)); + assert!(estimate.index_pages >= 1.0); + } + } + + #[test] + fn query_token_count_increases_eager_work() { + let one = estimate_maxsim_cost(MaxsimCostInput { + query_tokens: 1, + ..input(MaxsimCostBackend::CpuExact) + }); + let many = estimate_maxsim_cost(MaxsimCostInput { + query_tokens: 64, + ..input(MaxsimCostBackend::CpuExact) + }); + assert!(many.startup_cost > one.startup_cost); + assert!(many.index_pages > one.index_pages); + } + + #[test] + fn exact_candidate_limit_bounds_rows_and_cost() { + let small = estimate_maxsim_cost(MaxsimCostInput { + candidate_limit: Some(128), + ..input(MaxsimCostBackend::CpuExact) + }); + let large = estimate_maxsim_cost(MaxsimCostInput { + candidate_limit: Some(2048), + ..input(MaxsimCostBackend::CpuExact) + }); + assert!(small.selectivity < large.selectivity); + assert!(small.startup_cost < large.startup_cost); + } + + #[test] + fn auto_prices_more_than_gpu_for_fallback_risk() { + let gpu = estimate_maxsim_cost(input(MaxsimCostBackend::Gpu)); + let auto = estimate_maxsim_cost(input(MaxsimCostBackend::Auto)); + assert!(auto.startup_cost > gpu.startup_cost); + } + + #[test] + fn missing_stats_remain_finite() { + let estimate = estimate_maxsim_cost(MaxsimCostInput { + heap_rows: -1.0, + index_tokens: 0.0, + token_nodes_per_query: 0.0, + base_index_pages: 0.0, + filter_selectivity: 0.0, + limit_tuples: None, + candidate_limit: None, + ..input(MaxsimCostBackend::CoarseOnly) + }); + assert!(estimate.startup_cost.is_finite()); + assert!(estimate.total_cost.is_finite()); + assert!(estimate.selectivity.is_finite()); + assert!(estimate.index_pages.is_finite()); + } +} diff --git a/crates/vchordrq/src/statistics.rs b/crates/vchordrq/src/statistics.rs new file mode 100644 index 00000000..f651bec5 --- /dev/null +++ b/crates/vchordrq/src/statistics.rs @@ -0,0 +1,30 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the terms of the ELv2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::tuples::{MetaTuple, WithWriter}; +use index::relation::{Page, RelationWrite}; + +/// Store the number of live vector nodes observed by the latest complete +/// build or vacuum pass. +/// +/// This is deliberately refreshed in bulk instead of on every insert. A +/// per-insert update would serialize all writers on the metapage, which is a +/// poor tradeoff for a planner statistic. Like PostgreSQL's relation +/// statistics, the value may be stale between maintenance passes. +pub fn set_indexed_vectors(index: &R, indexed_vectors: u64) { + let mut meta_guard = index.write(0, false); + let meta_bytes = meta_guard.get_mut(1).expect("data corruption"); + let mut meta_tuple = MetaTuple::deserialize_mut(meta_bytes); + meta_tuple.set_indexed_vectors(indexed_vectors); +} diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index 8438c1ae..85dee786 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -21,6 +21,8 @@ pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); const VERSION: u64 = 1001; +const STATISTICS_VERSION: u16 = 1; +const MAX_INDEXED_VECTORS: u64 = (1_u64 << 48) - 1; #[inline(always)] fn tag(source: &[u8]) -> Tag { @@ -55,12 +57,13 @@ struct MetaTupleHeader { rerank_in_heap: Bool, cells_s: u16, cells_e: u16, - _padding_0: [Padding; 2], + statistics_version: u16, centroids_first: u32, vectors_first_s: u16, vectors_first_e: u16, freepages_first: u32, - _padding_1: [Padding; 6], + indexed_vectors_low: u32, + indexed_vectors_high: u16, // tree centroid_prefetch_s: u16, centroid_prefetch_e: u16, @@ -69,11 +72,24 @@ struct MetaTupleHeader { first: u32, } +// Statistics deliberately replace the old 2-byte and 6-byte padding regions. +// Keep these assertions in non-test builds: changing any offset would require +// an index format version bump and REINDEX instead of the compatibility path. +const _: () = { + assert!(size_of::() == 56); + assert!(std::mem::offset_of!(MetaTupleHeader, statistics_version) == 22); + assert!(std::mem::offset_of!(MetaTupleHeader, centroids_first) == 24); + assert!(std::mem::offset_of!(MetaTupleHeader, indexed_vectors_low) == 36); + assert!(std::mem::offset_of!(MetaTupleHeader, indexed_vectors_high) == 40); + assert!(std::mem::offset_of!(MetaTupleHeader, centroid_prefetch_s) == 42); +}; + pub struct MetaTuple { pub dim: u32, pub height_of_root: u32, pub is_residual: bool, pub rerank_in_heap: bool, + pub indexed_vectors: Option, pub cells: Vec, pub centroids_first: u32, pub vectors_first: Vec, @@ -94,6 +110,7 @@ impl Tuple for MetaTuple { height_of_root, is_residual, rerank_in_heap, + indexed_vectors, cells, centroids_first, vectors_first, @@ -103,6 +120,12 @@ impl Tuple for MetaTuple { centroid_norm, first, } => { + if let Some(indexed_vectors) = indexed_vectors { + assert!( + *indexed_vectors <= MAX_INDEXED_VECTORS, + "indexed vector count exceeds the on-disk 48-bit limit" + ); + } buffer.extend((MAGIC as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); // cells @@ -136,17 +159,20 @@ impl Tuple for MetaTuple { rerank_in_heap: (*rerank_in_heap).into(), cells_s, cells_e, + statistics_version: indexed_vectors + .map(|_| STATISTICS_VERSION) + .unwrap_or(0), centroids_first: *centroids_first, vectors_first_s, vectors_first_e, freepages_first: *freepages_first, + indexed_vectors_low: indexed_vectors.unwrap_or(0) as u32, + indexed_vectors_high: (indexed_vectors.unwrap_or(0) >> 32) as u16, centroid_prefetch_s, centroid_prefetch_e, centroid_head: *centroid_head, centroid_norm: *centroid_norm, first: *first, - _padding_0: Default::default(), - _padding_1: Default::default(), } .as_bytes(), ); @@ -186,6 +212,28 @@ impl WithReader for MetaTuple { } } +impl WithWriter for MetaTuple { + type Writer<'a> = MetaTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> MetaTupleWriter<'_> { + let tag = tag(source); + match tag { + MAGIC => { + let mut checker = MutChecker::new(source); + let header: &mut MetaTupleHeader = checker.prefix(size_of::()); + if VERSION != header.version { + panic!( + "deserialization: bad version number; {}", + "after upgrading VectorChord, please use REINDEX to rebuild the index." + ); + } + MetaTupleWriter { header } + } + _ => panic!("deserialization: bad magic number"), + } + } +} + #[derive(Debug, Clone, Copy)] pub struct MetaTupleReader<'a> { header: &'a MetaTupleHeader, @@ -207,6 +255,16 @@ impl<'a> MetaTupleReader<'a> { pub fn rerank_in_heap(self) -> bool { self.header.rerank_in_heap.into() } + pub fn indexed_vectors(self) -> Option { + match self.header.statistics_version { + 0 => None, + STATISTICS_VERSION => Some( + u64::from(self.header.indexed_vectors_low) + | (u64::from(self.header.indexed_vectors_high) << 32), + ), + _ => panic!("deserialization: unsupported statistics version"), + } + } pub fn cells(self) -> &'a [u32] { self.cells } @@ -233,6 +291,23 @@ impl<'a> MetaTupleReader<'a> { } } +#[derive(Debug)] +pub struct MetaTupleWriter<'a> { + header: &'a mut MetaTupleHeader, +} + +impl MetaTupleWriter<'_> { + pub fn set_indexed_vectors(&mut self, indexed_vectors: u64) { + assert!( + indexed_vectors <= MAX_INDEXED_VECTORS, + "indexed vector count exceeds the on-disk 48-bit limit" + ); + self.header.statistics_version = STATISTICS_VERSION; + self.header.indexed_vectors_low = indexed_vectors as u32; + self.header.indexed_vectors_high = (indexed_vectors >> 32) as u16; + } +} + #[repr(C, align(8))] #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FreepagesTupleHeader { @@ -1485,3 +1560,71 @@ impl AppendableTupleWriter<'_> { &mut self.header.payload } } + +#[cfg(test)] +mod tests { + use super::*; + + fn meta(indexed_vectors: Option) -> MetaTuple { + MetaTuple { + dim: 3, + height_of_root: 1, + is_residual: false, + rerank_in_heap: false, + indexed_vectors, + cells: vec![4], + centroids_first: 1, + vectors_first: vec![2], + freepages_first: 3, + centroid_prefetch: vec![4], + centroid_head: 0, + centroid_norm: 1.0, + first: 5, + } + } + + #[test] + fn meta_statistics_reuse_padding_without_moving_fields() { + assert_eq!(size_of::(), 56); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, statistics_version), + 22 + ); + assert_eq!(std::mem::offset_of!(MetaTupleHeader, centroids_first), 24); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, indexed_vectors_low), + 36 + ); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, indexed_vectors_high), + 40 + ); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, centroid_prefetch_s), + 42 + ); + } + + #[test] + fn meta_statistics_roundtrip_full_48_bit_range() { + for expected in [0, 1, u32::MAX as u64 + 1, MAX_INDEXED_VECTORS] { + let bytes = meta(Some(expected)).serialize(); + assert_eq!( + MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + Some(expected) + ); + } + } + + #[test] + fn old_meta_padding_reads_as_missing_statistics_and_can_be_upgraded() { + let mut bytes = meta(None).serialize(); + assert_eq!(MetaTuple::deserialize_ref(&bytes).indexed_vectors(), None); + + MetaTuple::deserialize_mut(&mut bytes).set_indexed_vectors(42); + assert_eq!( + MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + Some(42) + ); + } +} diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index c41ba501..d9b85cbc 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -12,12 +12,14 @@ // // Copyright (c) 2025-2026 TensorChord Inc. +use distance::Distance; use serde::{Deserialize, Serialize}; use simd::f16; use validator::{Validate, ValidationError}; use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; +use vector::{VectorBorrowed, VectorOwned}; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] @@ -61,6 +63,38 @@ pub enum OwnedVector { Rabitq4(Rabitq4Owned), } +impl OwnedVector { + pub fn dim(&self) -> u32 { + match self { + Self::Vecf32(vector) => vector.as_borrowed().dim(), + Self::Vecf16(vector) => vector.as_borrowed().dim(), + Self::Rabitq8(vector) => vector.as_borrowed().dim(), + Self::Rabitq4(vector) => vector.as_borrowed().dim(), + } + } + + pub fn operator_dot(&self, rhs: &Self) -> Option { + if self.dim() != rhs.dim() { + return None; + } + match (self, rhs) { + (Self::Vecf32(lhs), Self::Vecf32(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Vecf16(lhs), Self::Vecf16(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Rabitq8(lhs), Self::Rabitq8(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Rabitq4(lhs), Self::Rabitq4(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + _ => None, + } + } +} + #[derive(Debug, Clone, Copy)] pub enum BorrowedVector<'a> { Vecf32(VectBorrowed<'a, f32>), diff --git a/devtools/test_tilemaxsim_reference_sidecar.py b/devtools/test_tilemaxsim_reference_sidecar.py new file mode 100644 index 00000000..ee5ce666 --- /dev/null +++ b/devtools/test_tilemaxsim_reference_sidecar.py @@ -0,0 +1,383 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the terms of the ELv2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +from __future__ import annotations + +import hashlib +import os +import socket +import stat +import struct +import tempfile +import threading +import time +import unittest +from pathlib import Path + +try: + from . import tilemaxsim_reference_sidecar as sidecar +except ImportError: + import tilemaxsim_reference_sidecar as sidecar + + +def request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + candidates: list[tuple[int, list[list[float]]]], +) -> bytes: + dimension = len(query[0]) + code = "f" if dtype == sidecar.DTYPE_F32 else "e" + body = bytearray( + sidecar.REQUEST_FIXED.pack( + dimension, + len(query), + len(candidates), + dtype, + sidecar.SCORING_SUM_QUERY_MAX_DOCUMENT_DOT, + 0, + ) + ) + body.extend(struct.pack(f"<{len(query) * dimension}{code}", *sum(query, []))) + for candidate_id, tensor in candidates: + body.extend(sidecar.CANDIDATE_FIXED.pack(candidate_id, len(tensor))) + body.extend(struct.pack(f"<{len(tensor) * dimension}{code}", *sum(tensor, []))) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.VERSION, + sidecar.REQUEST_KIND, + request_id, + len(body), + ) + + body + ) + + +def external_request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + model_contract_id: str, + candidates: list[tuple[int, str, list[list[float]]]], +) -> tuple[bytes, dict[str, bytes]]: + dimension = len(query[0]) + code = "f" if dtype == sidecar.DTYPE_F32 else "e" + contract = model_contract_id.encode() + body = bytearray( + sidecar.EXTERNAL_REQUEST_FIXED.pack( + dimension, + len(query), + len(candidates), + dtype, + sidecar.SCORING_SUM_QUERY_MAX_DOCUMENT_DOT, + 0, + len(contract), + ) + ) + body.extend(contract) + body.extend(struct.pack(f"<{len(query) * dimension}{code}", *sum(query, []))) + objects = {} + for candidate_id, tensor_ref, tensor in candidates: + payload = struct.pack(f"<{len(tensor) * dimension}{code}", *sum(tensor, [])) + objects[tensor_ref] = payload + reference = tensor_ref.encode() + checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}".encode() + body.extend( + sidecar.EXTERNAL_CANDIDATE_FIXED.pack( + candidate_id, len(tensor), len(reference), len(checksum) + ) + ) + body.extend(reference) + body.extend(checksum) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.EXTERNAL_VERSION, + sidecar.REQUEST_KIND, + request_id, + len(body), + ) + + body, + objects, + ) + + +def decode_response(frame: bytes) -> tuple[int, int, list[tuple[int, float]] | str]: + magic, version, kind, request_id, body_len = sidecar.HEADER.unpack_from(frame) + assert magic == sidecar.MAGIC + assert version in (sidecar.VERSION, sidecar.EXTERNAL_VERSION) + assert kind == sidecar.RESPONSE_KIND + assert len(frame) == sidecar.HEADER.size + body_len + status, count_or_length = sidecar.RESPONSE_FIXED.unpack_from( + frame, sidecar.HEADER.size + ) + offset = sidecar.HEADER.size + sidecar.RESPONSE_FIXED.size + if status: + return request_id, status, frame[offset : offset + count_or_length].decode() + results = [] + for _ in range(count_or_length): + results.append(sidecar.RESULT.unpack_from(frame, offset)) + offset += sidecar.RESULT.size + assert offset == len(frame) + return request_id, status, results + + +class ReferenceSidecarTest(unittest.TestCase): + def test_f32_exact_scores_and_opaque_ids(self) -> None: + frame = request_frame( + 41, + sidecar.DTYPE_F32, + [[1.0, 0.0], [0.0, 1.0]], + [ + (17, [[1.0, 0.0], [0.0, 1.0]]), + (3, [[0.5, 0.5]]), + ], + ) + request_id, status, results = decode_response(sidecar.process_frame(frame)) + + self.assertEqual(request_id, 41) + self.assertEqual(status, 0) + self.assertEqual(results, [(17, 2.0), (3, 1.0)]) + + def test_f16_exact_scores(self) -> None: + frame = request_frame( + 42, + sidecar.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + [(0, [[0.75, 0.0], [0.0, 0.5]])], + ) + _, status, results = decode_response(sidecar.process_frame(frame)) + + self.assertEqual(status, 0) + self.assertEqual(results, [(0, 1.25)]) + + def test_shared_raw_decoder_preserves_inline_payloads(self) -> None: + frame = request_frame( + 420, + sidecar.DTYPE_F16, + [[1.0, 0.0]], + [(11, [[0.5, 0.25]])], + ) + request = sidecar.parse_request_frame(frame) + self.assertIsInstance(request, sidecar.InlineTensorRequest) + assert isinstance(request, sidecar.InlineTensorRequest) + self.assertEqual(request.request_id, 420) + self.assertEqual(request.query_rows, 1) + self.assertEqual(request.dimension, 2) + self.assertEqual([item.candidate_id for item in request.candidates], [11]) + self.assertEqual(len(request.query_payload), 4) + self.assertEqual(len(request.candidates[0].payload), 4) + + def test_duplicate_id_and_non_finite_input_fail_closed(self) -> None: + duplicate = request_frame( + 43, + sidecar.DTYPE_F32, + [[1.0]], + [(7, [[1.0]]), (7, [[2.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(duplicate)) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("duplicate candidate ID", message) + + non_finite = request_frame( + 44, + sidecar.DTYPE_F32, + [[float("nan")]], + [(0, [[1.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(non_finite)) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("non-finite", message) + + def test_truncated_and_oversized_frames_fail_closed(self) -> None: + valid = request_frame(45, sidecar.DTYPE_F32, [[1.0]], [(0, [[1.0]])]) + _, status, message = decode_response(sidecar.process_frame(valid[:-1])) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("length mismatch", message) + + _, status, message = decode_response( + sidecar.process_frame(valid, sidecar.Limits(max_request_bytes=32)) + ) + self.assertEqual(status, sidecar.STATUS_RESOURCE_LIMIT) + self.assertIn("byte limit", message) + + def test_header_reserved_trailing_and_token_limit_fail_closed(self) -> None: + valid = request_frame(47, sidecar.DTYPE_F32, [[1.0]], [(9, [[1.0]])]) + invalid_frames = [] + for offset, value in ((0, 0), (4, 3), (6, 2)): + invalid = bytearray(valid) + invalid[offset] = value + invalid_frames.append(bytes(invalid)) + + reserved = bytearray(valid) + reserved[sidecar.HEADER.size + 14] = 1 + invalid_frames.append(bytes(reserved)) + + trailing = bytearray(valid) + trailing.extend(b"x") + struct.pack_into(" None: + frame, objects = external_request_frame( + 48, + sidecar.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "colqwen@immutable-revision", + [ + (77, "object://fixture/page-1", [[1.0, 0.0], [0.0, 1.0]]), + (4, "object://fixture/page-2", [[0.5, 0.5]]), + ], + ) + seen = [] + + def resolver(request: sidecar.ExternalTensorRequest) -> bytes: + seen.append(request) + return objects[request.tensor_ref] + + request_id, status, results = decode_response( + sidecar.process_frame(frame, resolver=resolver) + ) + + self.assertEqual(request_id, 48) + self.assertEqual(status, 0) + self.assertEqual(results, [(77, 2.0), (4, 1.0)]) + self.assertEqual( + {request.model_contract_id for request in seen}, + {"colqwen@immutable-revision"}, + ) + self.assertEqual({request.tensor_ref for request in seen}, set(objects)) + + parsed = sidecar.parse_request_frame(frame) + self.assertIsInstance(parsed, sidecar.ParsedExternalTensorRequest) + assert isinstance(parsed, sidecar.ParsedExternalTensorRequest) + self.assertEqual(parsed.model_contract_id, "colqwen@immutable-revision") + self.assertEqual( + [candidate.candidate_id for candidate in parsed.candidates], [77, 4] + ) + + def test_external_v2_fails_closed_without_resolver_or_on_checksum_mismatch( + self, + ) -> None: + frame, objects = external_request_frame( + 49, + sidecar.DTYPE_F32, + [[1.0]], + "contract@1", + [(0, "object://immutable/page", [[2.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(frame)) + self.assertEqual(status, sidecar.STATUS_COMPUTE_ERROR) + self.assertIn("resolver is not configured", message) + + def corrupt_resolver(request: sidecar.ExternalTensorRequest) -> bytes: + return objects[request.tensor_ref] + b"x" + + _, status, message = decode_response( + sidecar.process_frame(frame, resolver=corrupt_resolver) + ) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("byte length", message) + + def test_external_v2_validates_complete_control_frame_before_resolution( + self, + ) -> None: + frame, objects = external_request_frame( + 50, + sidecar.DTYPE_F32, + [[1.0, 0.0]], + "contract@1", + [(0, "object://immutable/page", [[1.0, 0.0]])], + ) + invalid = bytearray(frame) + contract_length = len("contract@1") + candidate_offset = ( + sidecar.HEADER.size + + sidecar.EXTERNAL_REQUEST_FIXED.size + + contract_length + + 8 + ) + reference_offset = candidate_offset + sidecar.EXTERNAL_CANDIDATE_FIXED.size + invalid[reference_offset] = 0 + called = False + + def resolver(request: sidecar.ExternalTensorRequest) -> bytes: + nonlocal called + called = True + return objects[request.tensor_ref] + + _, status, message = decode_response( + sidecar.process_frame(bytes(invalid), resolver=resolver) + ) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("control characters", message) + self.assertFalse(called) + + def test_unix_socket_end_to_end(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "tilemaxsim.sock" + thread = threading.Thread( + target=sidecar.serve, + args=(path, sidecar.Limits()), + kwargs={"once": True}, + daemon=True, + ) + thread.start() + for _ in range(100): + if path.exists() and stat.S_IMODE(path.stat().st_mode) == 0o600: + break + time.sleep(0.01) + else: + self.fail("sidecar socket was not created with mode 0600") + self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600) + + frame = request_frame( + 46, + sidecar.DTYPE_F32, + [[1.0, 0.0], [0.0, 1.0]], + [(5, [[1.0, 0.0], [0.0, 1.0]])], + ) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(path)) + connection.sendall(frame) + header = connection.recv(sidecar.HEADER.size) + while len(header) < sidecar.HEADER.size: + header += connection.recv(sidecar.HEADER.size - len(header)) + body_len = sidecar.HEADER.unpack(header)[4] + body = b"" + while len(body) < body_len: + body += connection.recv(body_len - len(body)) + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + + request_id, status, results = decode_response(header + body) + self.assertEqual(request_id, 46) + self.assertEqual(status, 0) + self.assertEqual(results, [(5, 2.0)]) + self.assertFalse(path.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/devtools/tilemaxsim_reference_sidecar.py b/devtools/tilemaxsim_reference_sidecar.py new file mode 100644 index 00000000..8d9f7ee1 --- /dev/null +++ b/devtools/tilemaxsim_reference_sidecar.py @@ -0,0 +1,780 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the terms of the ELv2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""CPU reference implementation of the VectorChord TileMaxSim IPC sidecar. + +This executable is intentionally simple and single-threaded. It is a protocol +oracle and end-to-end development aid, not a production or GPU implementation. +The CLI serves inline v1 requests. External-descriptor v2 requests require an +explicit resolver injected by a test or embedding application and fail closed +when no resolver is configured. +""" + +from __future__ import annotations + +import argparse +import hashlib +import hmac +import math +import os +import signal +import socket +import stat +import struct +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterable + +MAGIC = b"VCTM" +VERSION = 1 +EXTERNAL_VERSION = 2 +REQUEST_KIND = 1 +RESPONSE_KIND = 2 +SCORING_SUM_QUERY_MAX_DOCUMENT_DOT = 1 +DTYPE_F32 = 1 +DTYPE_F16 = 2 + +HEADER = struct.Struct("<4sHHQQ") +REQUEST_FIXED = struct.Struct(" None: + super().__init__(message) + self.status = status + + +@dataclass(frozen=True) +class Limits: + max_request_bytes: int = 64 * 1024 * 1024 + max_batch_tokens: int = 1_000_000 + max_tensor_bytes: int = 1024 * 1024 * 1024 + max_candidates: int = 65_536 + + +@dataclass(frozen=True) +class ExternalTensorRequest: + model_contract_id: str + tensor_ref: str + rows: int + dimension: int + dtype: int + checksum: str + + +@dataclass(frozen=True) +class InlineTensorCandidate: + candidate_id: int + rows: int + payload: bytes + + +@dataclass(frozen=True) +class InlineTensorRequest: + request_id: int + dimension: int + query_rows: int + dtype: int + query_payload: bytes + candidates: tuple[InlineTensorCandidate, ...] + + +@dataclass(frozen=True) +class ExternalTensorCandidate: + candidate_id: int + descriptor: ExternalTensorRequest + + +@dataclass(frozen=True) +class ParsedExternalTensorRequest: + request_id: int + dimension: int + query_rows: int + dtype: int + model_contract_id: str + query_payload: bytes + candidates: tuple[ExternalTensorCandidate, ...] + + +ParsedRequest = InlineTensorRequest | ParsedExternalTensorRequest + + +# A reference resolver returns the exact row-major scalar bytes whose SHA-256 +# digest is stored in the descriptor. Production resolvers may adapt richer +# immutable object formats before returning this canonical tensor payload. +ExternalTensorResolver = Callable[[ExternalTensorRequest], bytes] + + +class Reader: + def __init__(self, payload: bytes) -> None: + self.payload = payload + self.offset = 0 + + def take(self, count: int) -> bytes: + end = self.offset + count + if count < 0 or end > len(self.payload): + raise SidecarError(STATUS_INVALID_REQUEST, "truncated request") + chunk = self.payload[self.offset : end] + self.offset = end + return chunk + + def unpack(self, layout: struct.Struct) -> tuple: + return layout.unpack(self.take(layout.size)) + + def finish(self) -> None: + if self.offset != len(self.payload): + raise SidecarError(STATUS_INVALID_REQUEST, "trailing request bytes") + + +def checked_elements(rows: int, dimension: int) -> int: + if rows <= 0 or dimension <= 0: + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor rows and dimension must be positive" + ) + elements = rows * dimension + if elements > (1 << 63) - 1: + raise SidecarError(STATUS_RESOURCE_LIMIT, "tensor shape is too large") + return elements + + +def dtype_size(dtype: int) -> int: + if dtype == DTYPE_F32: + return 4 + if dtype == DTYPE_F16: + return 2 + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported tensor dtype") + + +def checked_tensor_bytes(rows: int, dimension: int, dtype: int) -> int: + return checked_elements(rows, dimension) * dtype_size(dtype) + + +def validate_finite_tensor_payload( + payload: bytes, rows: int, dimension: int, dtype: int +) -> None: + expected = checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor byte length does not match its shape" + ) + code = "f" if dtype == DTYPE_F32 else "e" + try: + values = struct.iter_unpack(f"<{code}", payload) + if any(not math.isfinite(value[0]) for value in values): + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor contains non-finite value" + ) + except struct.error as error: + raise SidecarError(STATUS_INVALID_REQUEST, str(error)) from error + + +def read_text(reader: Reader, length: int, maximum: int, field: str) -> str: + if length <= 0 or length > maximum: + raise SidecarError(STATUS_RESOURCE_LIMIT, f"invalid {field} length") + try: + value = reader.take(length).decode("utf-8") + except UnicodeDecodeError as error: + raise SidecarError(STATUS_INVALID_REQUEST, f"{field} is not UTF-8") from error + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise SidecarError( + STATUS_INVALID_REQUEST, f"{field} contains control characters" + ) + return value + + +def read_tensor( + reader: Reader, rows: int, dimension: int, dtype: int +) -> list[tuple[float, ...]]: + elements = checked_elements(rows, dimension) + if dtype == DTYPE_F32: + code = "f" + element_size = 4 + elif dtype == DTYPE_F16: + code = "e" + element_size = 2 + else: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported tensor dtype") + raw = reader.take(elements * element_size) + try: + values = struct.unpack(f"<{elements}{code}", raw) + except struct.error as error: + raise SidecarError(STATUS_INVALID_REQUEST, str(error)) from error + if not all(math.isfinite(value) for value in values): + raise SidecarError(STATUS_INVALID_REQUEST, "tensor contains non-finite value") + return [ + tuple(values[offset : offset + dimension]) + for offset in range(0, elements, dimension) + ] + + +def decode_resolved_tensor( + payload: bytes, request: ExternalTensorRequest +) -> list[tuple[float, ...]]: + expected_bytes = checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + if len(payload) != expected_bytes: + raise SidecarError( + STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match descriptor", + ) + expected_checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}" + if not hmac.compare_digest(expected_checksum, request.checksum): + raise SidecarError(STATUS_INVALID_REQUEST, "resolved tensor checksum mismatch") + reader = Reader(payload) + tensor = read_tensor(reader, request.rows, request.dimension, request.dtype) + reader.finish() + return tensor + + +def tilemaxsim( + query: list[tuple[float, ...]], document: list[tuple[float, ...]] +) -> float: + if not query or not document: + raise SidecarError(STATUS_INVALID_REQUEST, "tensor must not be empty") + score = 0.0 + try: + for query_vector in query: + best = -math.inf + for document_vector in document: + dot = math.fsum( + left * right + for left, right in zip(query_vector, document_vector, strict=True) + ) + best = max(best, dot) + score += best + except (OverflowError, ValueError) as error: + raise SidecarError(STATUS_COMPUTE_ERROR, str(error)) from error + if not math.isfinite(score): + raise SidecarError(STATUS_COMPUTE_ERROR, "TileMaxSim result is non-finite") + return score + + +def success_response( + request_id: int, + results: Iterable[tuple[int, float]], + version: int = VERSION, +) -> bytes: + results = list(results) + body = bytearray(RESPONSE_FIXED.pack(0, len(results))) + for candidate_id, similarity in results: + body.extend(RESULT.pack(candidate_id, similarity)) + return HEADER.pack(MAGIC, version, RESPONSE_KIND, request_id, len(body)) + body + + +def error_response( + request_id: int, + status_code: int, + message: str, + version: int = VERSION, +) -> bytes: + encoded = message.encode("utf-8", errors="replace")[:MAX_ERROR_BYTES] + body = ERROR_FIXED.pack(status_code or STATUS_COMPUTE_ERROR, len(encoded)) + encoded + return HEADER.pack(MAGIC, version, RESPONSE_KIND, request_id, len(body)) + body + + +def validate_request_fixed( + dimension: int, + query_rows: int, + candidate_count: int, + dtype: int, + scoring: int, + reserved: int, + limits: Limits, +) -> None: + if dimension == 0 or dimension > 60_000: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid tensor dimension") + if query_rows == 0: + raise SidecarError(STATUS_INVALID_REQUEST, "query tensor is empty") + if candidate_count > limits.max_candidates: + raise SidecarError(STATUS_RESOURCE_LIMIT, "too many candidates") + dtype_size(dtype) + if scoring != SCORING_SUM_QUERY_MAX_DOCUMENT_DOT: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported scoring function") + if reserved != 0: + raise SidecarError(STATUS_INVALID_REQUEST, "reserved field must be zero") + if query_rows > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if checked_tensor_bytes(query_rows, dimension, dtype) > limits.max_tensor_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit") + + +def process_inline_request(reader: Reader, limits: Limits) -> list[tuple[int, float]]: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + ) = reader.unpack(REQUEST_FIXED) + validate_request_fixed( + dimension, query_rows, candidate_count, dtype, scoring, reserved, limits + ) + + query = read_tensor(reader, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = checked_tensor_bytes(query_rows, dimension, dtype) + candidates: list[tuple[int, list[tuple[float, ...]]]] = [] + candidate_ids: set[int] = set() + for _ in range(candidate_count): + candidate_id, rows = reader.unpack(CANDIDATE_FIXED) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + candidates.append((candidate_id, read_tensor(reader, rows, dimension, dtype))) + reader.finish() + return [ + (candidate_id, tilemaxsim(query, document)) + for candidate_id, document in candidates + ] + + +def process_external_request( + reader: Reader, + limits: Limits, + resolver: ExternalTensorResolver | None, +) -> list[tuple[int, float]]: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + validate_request_fixed( + dimension, query_rows, candidate_count, dtype, scoring, reserved, limits + ) + model_contract_id = read_text(reader, contract_length, 512, "model contract") + query = read_tensor(reader, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = checked_tensor_bytes(query_rows, dimension, dtype) + descriptors: list[tuple[int, ExternalTensorRequest]] = [] + candidate_ids: set[int] = set() + for _ in range(candidate_count): + candidate_id, rows, reference_length, checksum_length = reader.unpack( + EXTERNAL_CANDIDATE_FIXED + ) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + tensor_ref = read_text(reader, reference_length, 4096, "tensor reference") + checksum = read_text(reader, checksum_length, 512, "tensor checksum") + digest = checksum.removeprefix("sha256:") + if ( + not checksum.startswith("sha256:") + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise SidecarError( + STATUS_INVALID_REQUEST, + "tensor checksum must be a lowercase sha256 digest", + ) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + descriptors.append( + ( + candidate_id, + ExternalTensorRequest( + model_contract_id=model_contract_id, + tensor_ref=tensor_ref, + rows=rows, + dimension=dimension, + dtype=dtype, + checksum=checksum, + ), + ) + ) + reader.finish() + + # Validate the complete control frame before performing any external I/O. + if resolver is None: + raise SidecarError( + STATUS_COMPUTE_ERROR, "external tensor resolver is not configured" + ) + results = [] + for candidate_id, descriptor in descriptors: + payload = resolver(descriptor) + if not isinstance(payload, bytes): + raise SidecarError( + STATUS_COMPUTE_ERROR, + "external tensor resolver returned a non-bytes value", + ) + document = decode_resolved_tensor(payload, descriptor) + results.append((candidate_id, tilemaxsim(query, document))) + return results + + +def parse_request_frame( + frame: bytes, + limits: Limits = Limits(), + *, + validate_finite: bool = True, +) -> ParsedRequest: + """Decode and validate a request without resolving or computing tensors. + + The production CUDA sidecar uses this function so the protocol oracle and + deployable executor share one strict wire decoder. All control data and the + inline query are validated before a v2 resolver performs external I/O. + """ + + if len(frame) < HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "truncated frame header") + magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) + if magic != MAGIC: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") + if version not in (VERSION, EXTERNAL_VERSION): + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") + if kind != REQUEST_KIND: + raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") + if body_len != len(frame) - HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "request length mismatch") + if len(frame) > limits.max_request_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds byte limit") + + reader = Reader(frame[HEADER.size :]) + if version == VERSION: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + ) = reader.unpack(REQUEST_FIXED) + validate_request_fixed( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + limits, + ) + query_payload = reader.take(checked_tensor_bytes(query_rows, dimension, dtype)) + if validate_finite: + validate_finite_tensor_payload(query_payload, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = len(query_payload) + candidate_ids: set[int] = set() + candidates = [] + for _ in range(candidate_count): + candidate_id, rows = reader.unpack(CANDIDATE_FIXED) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + payload_bytes = checked_tensor_bytes(rows, dimension, dtype) + total_tokens += rows + total_tensor_bytes += payload_bytes + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + payload = reader.take(payload_bytes) + if validate_finite: + validate_finite_tensor_payload(payload, rows, dimension, dtype) + candidates.append(InlineTensorCandidate(candidate_id, rows, payload)) + reader.finish() + return InlineTensorRequest( + request_id, + dimension, + query_rows, + dtype, + query_payload, + tuple(candidates), + ) + + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + validate_request_fixed( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + limits, + ) + model_contract_id = read_text(reader, contract_length, 512, "model contract") + query_payload = reader.take(checked_tensor_bytes(query_rows, dimension, dtype)) + if validate_finite: + validate_finite_tensor_payload(query_payload, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = len(query_payload) + candidate_ids = set() + candidates = [] + for _ in range(candidate_count): + candidate_id, rows, reference_length, checksum_length = reader.unpack( + EXTERNAL_CANDIDATE_FIXED + ) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + tensor_ref = read_text(reader, reference_length, 4096, "tensor reference") + checksum = read_text(reader, checksum_length, 512, "tensor checksum") + digest = checksum.removeprefix("sha256:") + if ( + not checksum.startswith("sha256:") + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise SidecarError( + STATUS_INVALID_REQUEST, + "tensor checksum must be a lowercase sha256 digest", + ) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + candidates.append( + ExternalTensorCandidate( + candidate_id, + ExternalTensorRequest( + model_contract_id=model_contract_id, + tensor_ref=tensor_ref, + rows=rows, + dimension=dimension, + dtype=dtype, + checksum=checksum, + ), + ) + ) + reader.finish() + return ParsedExternalTensorRequest( + request_id, + dimension, + query_rows, + dtype, + model_contract_id, + query_payload, + tuple(candidates), + ) + + +def process_frame( + frame: bytes, + limits: Limits = Limits(), + resolver: ExternalTensorResolver | None = None, +) -> bytes: + request_id = 0 + response_version = VERSION + try: + if len(frame) < HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "truncated frame header") + magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) + if version in (VERSION, EXTERNAL_VERSION): + response_version = version + if magic != MAGIC: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") + if version not in (VERSION, EXTERNAL_VERSION): + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") + if kind != REQUEST_KIND: + raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") + if body_len != len(frame) - HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "request length mismatch") + if len(frame) > limits.max_request_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds byte limit") + + reader = Reader(frame[HEADER.size :]) + if version == VERSION: + results = process_inline_request(reader, limits) + else: + results = process_external_request(reader, limits, resolver) + return success_response(request_id, results, response_version) + except SidecarError as error: + return error_response(request_id, error.status, str(error), response_version) + except Exception as error: # Keep protocol failures inside the response boundary. + return error_response( + request_id, STATUS_COMPUTE_ERROR, str(error), response_version + ) + + +def receive_exact(connection: socket.socket, count: int) -> bytes: + chunks = bytearray() + while len(chunks) < count: + chunk = connection.recv(count - len(chunks)) + if not chunk: + raise SidecarError( + STATUS_INVALID_REQUEST, "connection closed during request" + ) + chunks.extend(chunk) + return bytes(chunks) + + +def handle_connection( + connection: socket.socket, + limits: Limits, + resolver: ExternalTensorResolver | None = None, +) -> None: + request_id = 0 + response_version = VERSION + try: + header = receive_exact(connection, HEADER.size) + _, version, _, request_id, body_len = HEADER.unpack(header) + if version in (VERSION, EXTERNAL_VERSION): + response_version = version + if body_len > limits.max_request_bytes - HEADER.size: + response = error_response( + request_id, + STATUS_RESOURCE_LIMIT, + "request exceeds byte limit", + response_version, + ) + else: + body = receive_exact(connection, body_len) + response = process_frame(header + body, limits, resolver) + except SidecarError as error: + response = error_response( + request_id, error.status, str(error), response_version + ) + except Exception as error: + response = error_response( + request_id, STATUS_COMPUTE_ERROR, str(error), response_version + ) + connection.sendall(response) + + +def remove_stale_socket(path: Path) -> None: + try: + mode = path.lstat().st_mode + except FileNotFoundError: + return + if not stat.S_ISSOCK(mode): + raise RuntimeError(f"refusing to replace non-socket path: {path}") + path.unlink() + + +def serve( + socket_path: Path, + limits: Limits, + socket_mode: int = 0o600, + once: bool = False, + stop: threading.Event | None = None, + resolver: ExternalTensorResolver | None = None, +) -> None: + stop = stop or threading.Event() + remove_stale_socket(socket_path) + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + listener.bind(os.fspath(socket_path)) + os.chmod(socket_path, socket_mode) + listener.listen(16) + listener.settimeout(0.25) + bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino + while not stop.is_set(): + try: + connection, _ = listener.accept() + except TimeoutError: + continue + with connection: + handle_connection(connection, limits, resolver) + if once: + break + finally: + listener.close() + try: + current = socket_path.lstat() + if (current.st_dev, current.st_ino) == bound_identity: + socket_path.unlink() + except (FileNotFoundError, UnboundLocalError): + pass + + +def parse_mode(value: str) -> int: + mode = int(value, 8) + if mode < 0 or mode > 0o777: + raise argparse.ArgumentTypeError("socket mode must be between 000 and 777") + return mode + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=parse_mode, default=0o600) + parser.add_argument("--max-request-bytes", type=int, default=64 * 1024 * 1024) + parser.add_argument("--max-batch-tokens", type=int, default=1_000_000) + parser.add_argument("--max-tensor-bytes", type=int, default=1024 * 1024 * 1024) + parser.add_argument("--max-candidates", type=int, default=65_536) + parser.add_argument("--once", action="store_true") + args = parser.parse_args() + limits = Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + if ( + min( + limits.max_request_bytes, + limits.max_batch_tokens, + limits.max_tensor_bytes, + limits.max_candidates, + ) + <= 0 + ): + parser.error("all limits must be positive") + + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + serve(args.socket, limits, args.socket_mode, args.once, stop) + + +if __name__ == "__main__": + main() diff --git a/services/Dockerfile.tilemaxsim b/services/Dockerfile.tilemaxsim new file mode 100644 index 00000000..9f58b88d --- /dev/null +++ b/services/Dockerfile.tilemaxsim @@ -0,0 +1,12 @@ +ARG BASE_IMAGE +FROM ${BASE_IMAGE} + +WORKDIR /opt/vectorchord + +COPY devtools/tilemaxsim_reference_sidecar.py devtools/tilemaxsim_reference_sidecar.py +COPY services/tilemaxsim_cuda_sidecar.py services/tilemaxsim_cuda_sidecar.py + +ENV PYTHONPATH=/opt/vectorchord \ + PYTHONUNBUFFERED=1 + +ENTRYPOINT ["python3", "-m", "services.tilemaxsim_cuda_sidecar"] diff --git a/services/benchmark_tilemaxsim_cuda.py b/services/benchmark_tilemaxsim_cuda.py new file mode 100644 index 00000000..698e4b84 --- /dev/null +++ b/services/benchmark_tilemaxsim_cuda.py @@ -0,0 +1,152 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Reproducible synthetic load probe for the CUDA TileMaxSim executor.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import time + +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import TorchTileMaxsimEngine, positive_int + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + index = max(0, math.ceil(fraction * len(ordered)) - 1) + return ordered[index] + + +def canonical_payload(tensor: torch.Tensor, dtype: int) -> bytes: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + return tensor.to(dtype=scalar_dtype).contiguous().numpy().tobytes() + + +def normalized_tensor( + shape: tuple[int, ...], generator: torch.Generator +) -> torch.Tensor: + tensor = torch.randn(shape, dtype=torch.float32, generator=generator) + return torch.nn.functional.normalize(tensor, dim=-1) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--dtype", choices=("f16", "f32"), default="f16") + parser.add_argument("--dimension", type=positive_int, default=320) + parser.add_argument("--query-rows", type=positive_int, default=32) + parser.add_argument("--document-rows", type=positive_int, default=747) + parser.add_argument("--candidates", type=positive_int, default=128) + parser.add_argument("--warmup", type=positive_int, default=3) + parser.add_argument("--iterations", type=positive_int, default=10) + parser.add_argument("--seed", type=int, default=20260713) + parser.add_argument( + "--max-device-bytes", type=positive_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument("--allow-tf32", action="store_true") + args = parser.parse_args() + + dtype = protocol.DTYPE_F32 if args.dtype == "f32" else protocol.DTYPE_F16 + generator = torch.Generator(device="cpu").manual_seed(args.seed) + query = normalized_tensor((args.query_rows, args.dimension), generator) + document_tensors = normalized_tensor( + (args.candidates, args.document_rows, args.dimension), generator + ) + query_payload = canonical_payload(query, dtype) + documents = [ + ( + candidate_id, + args.document_rows, + canonical_payload(document_tensors[candidate_id], dtype), + ) + for candidate_id in range(args.candidates) + ] + del document_tensors + + engine = TorchTileMaxsimEngine( + args.device, args.max_device_bytes, args.allow_tf32, 1 + ) + if engine.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(engine.device) + + total_samples: list[float] = [] + queue_samples: list[float] = [] + compute_samples: list[float] = [] + score_checksum = 0.0 + for iteration in range(args.warmup + args.iterations): + started = time.perf_counter() + results, queue_ms, compute_ms = engine.score( + query_payload, + args.query_rows, + args.dimension, + dtype, + documents, + time.monotonic() + 300, + lambda: False, + ) + total_ms = (time.perf_counter() - started) * 1000.0 + if iteration >= args.warmup: + total_samples.append(total_ms) + queue_samples.append(queue_ms) + compute_samples.append(compute_ms) + score_checksum = math.fsum(score for _, score in results) + + output = { + "benchmark": "tilemaxsim_cuda_synthetic_v1", + "device": str(engine.device), + "device_name": ( + torch.cuda.get_device_name(engine.device) + if engine.device.type == "cuda" + else "cpu" + ), + "torch_version": torch.__version__, + "dtype": args.dtype, + "dimension": args.dimension, + "query_rows": args.query_rows, + "document_rows": args.document_rows, + "candidates": args.candidates, + "candidate_tokens": args.candidates * args.document_rows, + "seed": args.seed, + "warmup": args.warmup, + "iterations": args.iterations, + "allow_tf32": args.allow_tf32, + "max_device_bytes": args.max_device_bytes, + "latency_ms": { + "mean": round(statistics.fmean(total_samples), 3), + "p50": round(percentile(total_samples, 0.50), 3), + "p95": round(percentile(total_samples, 0.95), 3), + "p99": round(percentile(total_samples, 0.99), 3), + "queue_mean": round(statistics.fmean(queue_samples), 3), + "compute_mean": round(statistics.fmean(compute_samples), 3), + }, + "score_checksum": score_checksum, + } + if engine.device.type == "cuda": + output["cuda_peak_allocated_bytes"] = torch.cuda.max_memory_allocated( + engine.device + ) + output["cuda_peak_reserved_bytes"] = torch.cuda.max_memory_reserved( + engine.device + ) + print(json.dumps(output, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/build_tilemaxsim_tensor_cache.py b/services/build_tilemaxsim_tensor_cache.py new file mode 100644 index 00000000..a48a1728 --- /dev/null +++ b/services/build_tilemaxsim_tensor_cache.py @@ -0,0 +1,235 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Publish NPY page tensors into the sidecar's immutable SHA-256 cache.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +import tempfile +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np + +from services.tilemaxsim_cuda_sidecar import positive_int + + +def canonical_tensor(path: Path, expected_rows: int, expected_dim: int) -> np.ndarray: + tensor = np.load(path, mmap_mode="r", allow_pickle=False) + if tensor.ndim != 2 or tensor.shape != (expected_rows, expected_dim): + raise ValueError( + f"{path}: expected shape {(expected_rows, expected_dim)}, got {tensor.shape}" + ) + if tensor.dtype == np.dtype("float16"): + little_dtype = np.dtype(" str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def write_payload(path: Path, payload: memoryview, digest: str, fsync: bool) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + if path.stat().st_size != len(payload): + raise ValueError(f"existing cache payload has wrong size: {path}") + if file_digest(path) != digest: + raise ValueError(f"existing cache payload has wrong checksum: {path}") + return + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".tmp", dir=path.parent + ) + try: + with os.fdopen(descriptor, "wb", closefd=True) as stream: + stream.write(payload) + stream.flush() + if fsync: + os.fsync(stream.fileno()) + try: + os.link(temporary_name, path) + except FileExistsError: + if path.stat().st_size != len(payload): + raise ValueError(f"concurrent cache payload has wrong size: {path}") + if file_digest(path) != digest: + raise ValueError(f"concurrent cache payload has wrong checksum: {path}") + if fsync: + directory_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + + +def process_record( + record: dict[str, object], + source_root: Path, + cache_root: Path, + fsync: bool, + dry_run: bool, +) -> dict[str, object]: + page_key = record.get("page_key") + relative = record.get("embedding_file") + rows = record.get("n_tokens") + dimension = record.get("dim") + if not isinstance(page_key, str) or not page_key: + raise ValueError("manifest record has no page_key") + if not isinstance(relative, str) or not relative: + raise ValueError(f"manifest page {page_key} has no embedding_file") + if not isinstance(rows, int) or rows <= 0: + raise ValueError(f"manifest page {page_key} has invalid n_tokens") + if not isinstance(dimension, int) or dimension <= 0: + raise ValueError(f"manifest page {page_key} has invalid dim") + source = (source_root / relative).resolve(strict=True) + try: + source.relative_to(source_root) + except ValueError as error: + raise ValueError( + f"embedding path escapes the source root: {relative}" + ) from error + tensor = canonical_tensor(source, rows, dimension) + payload = memoryview(tensor).cast("B") + digest = hashlib.sha256(payload).hexdigest() + destination = cache_root / digest[:2] / f"{digest}.bin" + if not dry_run: + write_payload(destination, payload, digest, fsync) + dtype_name = "float16" if tensor.dtype == np.dtype("float16") else "float32" + return { + "page_key": page_key, + "tensor_ref": f"sha256://{digest}", + "tensor_rows": rows, + "tensor_dim": dimension, + "tensor_dtype": dtype_name, + "tensor_checksum": f"sha256:{digest}", + "canonical_bytes": len(payload), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--cache-root", required=True, type=Path) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--workers", type=positive_int, default=4) + parser.add_argument("--no-fsync", action="store_true") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + + source_root = args.manifest.resolve(strict=True).parent + cache_root = args.cache_root.resolve() + if not cache_root.is_absolute(): + parser.error("--cache-root must be absolute") + records = [] + with args.manifest.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError( + f"invalid JSON at manifest line {line_number}" + ) from error + if not isinstance(record, dict): + raise ValueError(f"manifest line {line_number} is not an object") + records.append(record) + if not records: + raise ValueError("manifest is empty") + + cache_root.mkdir(parents=True, exist_ok=True) + descriptors = [] + with ThreadPoolExecutor(max_workers=args.workers) as workers: + results = workers.map( + lambda record: process_record( + record, + source_root, + cache_root, + not args.no_fsync, + args.dry_run, + ), + records, + ) + for completed, item in enumerate(results, 1): + descriptors.append(item) + if completed % 1000 == 0 or completed == len(records): + print( + json.dumps( + {"event": "tensor_cache_progress", "completed": completed}, + separators=(",", ":"), + ), + file=sys.stderr, + flush=True, + ) + + output_parent = args.descriptor_manifest.parent + output_parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{args.descriptor_manifest.name}.", + suffix=".tmp", + dir=output_parent, + text=True, + ) + try: + with os.fdopen(descriptor, "w", encoding="utf-8", closefd=True) as stream: + for item in descriptors: + stream.write(json.dumps(item, separators=(",", ":"), sort_keys=True)) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_name, args.descriptor_manifest) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + + total_bytes = sum(int(item["canonical_bytes"]) for item in descriptors) + print( + json.dumps( + { + "pages": len(descriptors), + "canonical_bytes": total_bytes, + "cache_root": os.fspath(cache_root), + "descriptor_manifest": os.fspath(args.descriptor_manifest), + "dry_run": args.dry_run, + }, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/services/test_tilemaxsim_cuda_sidecar.py b/services/test_tilemaxsim_cuda_sidecar.py new file mode 100644 index 00000000..093245e3 --- /dev/null +++ b/services/test_tilemaxsim_cuda_sidecar.py @@ -0,0 +1,313 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +from __future__ import annotations + +import hashlib +import os +import socket +import stat +import struct +import tempfile +import threading +import time +import unittest +from pathlib import Path + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from devtools.test_tilemaxsim_reference_sidecar import ( + decode_response, + external_request_frame, + request_frame, +) +from services import tilemaxsim_cuda_sidecar as cuda_sidecar +from services.build_tilemaxsim_tensor_cache import process_record + + +class CapturingMetrics(cuda_sidecar.JsonMetrics): + def __init__(self) -> None: + super().__init__() + self.events: list[dict[str, object]] = [] + + def emit(self, fields: dict[str, object]) -> None: + with self.lock: + self.events.append(fields.copy()) + + +def write_content_addressed(root: Path, payload: bytes) -> tuple[str, str]: + digest = hashlib.sha256(payload).hexdigest() + directory = root / digest[:2] + directory.mkdir(parents=True, exist_ok=True) + (directory / f"{digest}.bin").write_bytes(payload) + return f"sha256://{digest}", f"sha256:{digest}" + + +class CudaSidecarTest(unittest.TestCase): + def test_cache_builder_publishes_resolver_compatible_payload(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_root = root / "source" + cache_root = root / "cache" + source_root.mkdir() + tensor = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + cuda_sidecar.validate_finite_payload( + struct.pack("<2e", 1.0, 0.0), 1, 2, protocol.DTYPE_F16 + ) + with self.assertRaisesRegex(protocol.SidecarError, "non-finite"): + cuda_sidecar.validate_finite_payload( + struct.pack("<2f", 1.0, float("nan")), + 1, + 2, + protocol.DTYPE_F32, + ) + + def test_content_addressed_resolver_validates_and_caches(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 1024) + try: + request = protocol.ExternalTensorRequest( + "model@1", + tensor_ref, + 2, + 2, + protocol.DTYPE_F16, + checksum, + ) + first = resolver.resolve(request) + second = resolver.resolve(request) + self.assertEqual(first.payload, payload) + self.assertFalse(first.cache_hit) + self.assertTrue(second.cache_hit) + + bad = protocol.ExternalTensorRequest( + "model@1", + tensor_ref, + 2, + 2, + protocol.DTYPE_F16, + "sha256:" + "0" * 64, + ) + with self.assertRaisesRegex(protocol.SidecarError, "disagree"): + resolver.resolve(bad) + finally: + resolver.close() + + def test_content_addressed_resolver_rejects_symlink(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) / "root" + root.mkdir() + payload = struct.pack(" None: + query = [[1.0, 0.0], [0.0, 1.0]] + candidates = [ + (17, [[1.0, 0.0], [0.0, 1.0]]), + (3, [[0.5, 0.5], [0.25, 0.25]]), + ] + frame = request_frame(41, protocol.DTYPE_F32, query, candidates) + parsed = protocol.parse_request_frame(frame) + self.assertIsInstance(parsed, protocol.InlineTensorRequest) + assert isinstance(parsed, protocol.InlineTensorRequest) + documents = [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in parsed.candidates + ] + # 64 bytes fits one candidate but not both, exercising internal + # all-or-nothing device chunking. + engine = cuda_sidecar.TorchTileMaxsimEngine("cpu", 64, False, 1) + results, _, _ = engine.score( + parsed.query_payload, + parsed.query_rows, + parsed.dimension, + parsed.dtype, + documents, + time.monotonic() + 2, + lambda: False, + ) + _, status, oracle = decode_response(protocol.process_frame(frame)) + self.assertEqual(status, 0) + self.assertEqual(results, oracle) + + def test_compute_capacity_wait_uses_overall_deadline(self) -> None: + engine = cuda_sidecar.TorchTileMaxsimEngine("cpu", 1024, False, 1) + self.assertTrue(engine.compute_slots.acquire(blocking=False)) + try: + started = time.monotonic() + with self.assertRaisesRegex(protocol.SidecarError, "CUDA capacity"): + engine.score( + struct.pack(" None: + query = [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25]] + candidates = [ + (7, [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25]]), + (2, [[0.5, 0.5, 0.0], [-0.5, 0.25, 1.0]]), + ] + frame = request_frame(52, protocol.DTYPE_F16, query, candidates) + parsed = protocol.parse_request_frame(frame) + assert isinstance(parsed, protocol.InlineTensorRequest) + engine = cuda_sidecar.TorchTileMaxsimEngine("cuda:0", 1024 * 1024, False, 1) + results, _, _ = engine.score( + parsed.query_payload, + parsed.query_rows, + parsed.dimension, + parsed.dtype, + [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in parsed.candidates + ], + time.monotonic() + 5, + lambda: False, + ) + _, status, oracle = decode_response(protocol.process_frame(frame)) + self.assertEqual(status, 0) + assert isinstance(oracle, list) + self.assertEqual([item[0] for item in results], [item[0] for item in oracle]) + for (_, actual), (_, expected) in zip(results, oracle, strict=True): + self.assertAlmostEqual(actual, expected, places=5) + + def test_v2_unix_socket_end_to_end(self) -> None: + with tempfile.TemporaryDirectory() as directory: + directory_path = Path(directory) + root = directory_path / "tensors" + root.mkdir() + tensor = [[1.0, 0.0], [0.0, 1.0]] + payload = struct.pack("<4e", *sum(tensor, [])) + tensor_ref, _ = write_content_addressed(root, payload) + frame, _ = external_request_frame( + 61, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, tensor)], + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 1024) + metrics = CapturingMetrics() + service = cuda_sidecar.TileMaxsimService( + protocol.Limits(), + resolver, + cuda_sidecar.TorchTileMaxsimEngine("cpu", 1024 * 1024, False, 1), + 2000, + metrics, + ) + socket_path = directory_path / "tilemaxsim.sock" + stop = threading.Event() + thread = threading.Thread( + target=cuda_sidecar.serve, + args=(socket_path, 0o600, 4, 2, service, stop), + kwargs={"once": True}, + daemon=True, + ) + thread.start() + for _ in range(100): + if socket_path.exists(): + break + time.sleep(0.01) + else: + self.fail("CUDA sidecar socket was not created") + self.assertEqual(stat.S_IMODE(socket_path.stat().st_mode), 0o600) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_len = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_len) + thread.join(timeout=3) + resolver.close() + self.assertFalse(thread.is_alive()) + _, status, results = decode_response(response) + self.assertEqual(status, 0) + self.assertEqual(results, [(9, 2.0)]) + request_events = [ + event + for event in metrics.events + if event.get("event") == "tilemaxsim_request" + ] + self.assertEqual(len(request_events), 1) + self.assertEqual(request_events[0]["source"], "content_addressed") + self.assertEqual(request_events[0]["status"], "ok") + + +if __name__ == "__main__": + unittest.main() diff --git a/services/tilemaxsim_cuda_sidecar.py b/services/tilemaxsim_cuda_sidecar.py new file mode 100644 index 00000000..ce744119 --- /dev/null +++ b/services/tilemaxsim_cuda_sidecar.py @@ -0,0 +1,802 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Bounded CUDA executor for VectorChord TileMaxSim IPC v1 and v2. + +Version 1 consumes inline tensors. Version 2 resolves canonical tensor payloads +from an operations-configured, per-model content-addressed cache. Applications may +populate that cache from any object store; object-store credentials and routing +never enter PostgreSQL or this protocol. +""" + +from __future__ import annotations + +import argparse +import hashlib +import hmac +import json +import math +import os +import select +import signal +import socket +import stat +import struct +import sys +import threading +import time +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterable + +import numpy as np +import torch + +if __package__ in (None, ""): + sys.path.insert(0, os.fspath(Path(__file__).resolve().parents[1])) + +from devtools import tilemaxsim_reference_sidecar as protocol + + +def validate_finite_payload( + payload: bytes, rows: int, dimension: int, dtype: int +) -> None: + expected = protocol.checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor byte length does not match its shape", + ) + scalar_dtype = " None: + self.maximum_bytes = maximum_bytes + self.current_bytes = 0 + self.entries: OrderedDict[tuple[object, ...], bytes] = OrderedDict() + self.lock = threading.Lock() + + def get(self, key: tuple[object, ...]) -> bytes | None: + if self.maximum_bytes == 0: + return None + with self.lock: + payload = self.entries.get(key) + if payload is not None: + self.entries.move_to_end(key) + return payload + + def put(self, key: tuple[object, ...], payload: bytes) -> None: + if self.maximum_bytes == 0 or len(payload) > self.maximum_bytes: + return + with self.lock: + previous = self.entries.pop(key, None) + if previous is not None: + self.current_bytes -= len(previous) + self.entries[key] = payload + self.current_bytes += len(payload) + while self.current_bytes > self.maximum_bytes: + _, evicted = self.entries.popitem(last=False) + self.current_bytes -= len(evicted) + + +class ContentAddressedResolver: + """Resolve ``sha256://`` inside an allowlisted model cache root. + + A payload with digest ``abcdef...`` is stored as + ``/ab/abcdef....bin``. Directory and file symlinks are + rejected with ``openat(O_NOFOLLOW)``. The digest in the reference, the + registered checksum, the exact byte length, and the file content must all + agree before a payload is returned. + """ + + def __init__(self, roots: dict[str, Path], cache_bytes: int) -> None: + self.root_fds: dict[str, int] = {} + try: + for contract, path in roots.items(): + self.root_fds[contract] = os.open( + os.fspath(path), + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + ) + except Exception: + self.close() + raise + self.cache = PayloadCache(cache_bytes) + + def close(self) -> None: + for descriptor in self.root_fds.values(): + os.close(descriptor) + self.root_fds.clear() + + @staticmethod + def _digest(request: protocol.ExternalTensorRequest) -> str: + prefix = "sha256://" + if not request.tensor_ref.startswith(prefix): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "unsupported tensor reference; expected sha256://", + ) + digest = request.tensor_ref[len(prefix) :] + if len(digest) != 64 or any( + character not in "0123456789abcdef" for character in digest + ): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "invalid content-addressed tensor reference", + ) + if not hmac.compare_digest(request.checksum, f"sha256:{digest}"): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor reference and checksum disagree", + ) + return digest + + @staticmethod + def _read_exact_file(root_fd: int, digest: str, expected_bytes: int) -> bytes: + directory_fd = -1 + payload_fd = -1 + try: + directory_fd = os.open( + digest[:2], + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + payload_fd = os.open( + f"{digest}.bin", + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=directory_fd, + ) + metadata = os.fstat(payload_fd) + if not stat.S_ISREG(metadata.st_mode): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor is not a regular file", + ) + if metadata.st_size != expected_bytes: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match descriptor", + ) + chunks = bytearray() + while len(chunks) < expected_bytes: + chunk = os.read( + payload_fd, min(1024 * 1024, expected_bytes - len(chunks)) + ) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor file ended early", + ) + chunks.extend(chunk) + if os.read(payload_fd, 1): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor file grew during read", + ) + return bytes(chunks) + except FileNotFoundError as error: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "content-addressed tensor is missing" + ) from error + except OSError as error: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + f"content-addressed tensor read failed: {error.strerror}", + ) from error + finally: + if payload_fd >= 0: + os.close(payload_fd) + if directory_fd >= 0: + os.close(directory_fd) + + def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: + root_fd = self.root_fds.get(request.model_contract_id) + if root_fd is None: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "model contract has no configured tensor cache root", + ) + digest = self._digest(request) + expected_bytes = protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + key = ( + request.model_contract_id, + digest, + request.rows, + request.dimension, + request.dtype, + ) + cached = self.cache.get(key) + if cached is not None: + return ResolvedPayload(cached, True) + payload = self._read_exact_file(root_fd, digest, expected_bytes) + actual = hashlib.sha256(payload).hexdigest() + if not hmac.compare_digest(actual, digest): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, "resolved tensor checksum mismatch" + ) + validate_finite_payload(payload, request.rows, request.dimension, request.dtype) + self.cache.put(key, payload) + return ResolvedPayload(payload, False) + + +class TorchTileMaxsimEngine: + def __init__( + self, + device_name: str, + max_device_bytes: int, + allow_tf32: bool, + max_cuda_inflight: int, + ) -> None: + self.device = torch.device(device_name) + if self.device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested but torch.cuda.is_available() is false" + ) + if self.device.type not in ("cuda", "cpu"): + raise RuntimeError("device must be CUDA or CPU") + self.max_device_bytes = max_device_bytes + self.compute_slots = threading.BoundedSemaphore(max_cuda_inflight) + if self.device.type == "cuda": + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.backends.cudnn.allow_tf32 = allow_tf32 + with torch.inference_mode(): + left = torch.zeros((1, 1), dtype=torch.float32, device=self.device) + _ = left @ left + torch.cuda.synchronize(self.device) + + @staticmethod + def _cpu_tensor( + payload: bytes, rows: int, dimension: int, dtype: int + ) -> torch.Tensor: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + # bytearray gives torch a writable, owned buffer; clone detaches the + # resulting tensor before that temporary buffer leaves scope. + tensor = torch.frombuffer(bytearray(payload), dtype=scalar_dtype).reshape( + rows, dimension + ) + if scalar_dtype == torch.float32: + return tensor.clone() + return tensor.to(dtype=torch.float32) + + def _groups( + self, + query_rows: int, + dimension: int, + documents: list[tuple[int, int, bytes]], + ) -> Iterable[list[tuple[int, int, bytes]]]: + query_bytes = query_rows * dimension * 4 + group: list[tuple[int, int, bytes]] = [] + group_rows = 0 + for document in documents: + rows = document[1] + next_rows = group_rows + rows + # Device residency includes the f32 query, f32 documents, and the + # q-by-total-document-token similarity matrix. + required = ( + query_bytes + next_rows * dimension * 4 + query_rows * next_rows * 4 + ) + if required > self.max_device_bytes and group: + yield group + group = [] + group_rows = 0 + next_rows = rows + required = query_bytes + rows * dimension * 4 + query_rows * rows * 4 + if required > self.max_device_bytes: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one candidate exceeds the CUDA device-byte limit", + ) + group.append(document) + group_rows = next_rows + if group: + yield group + + def score( + self, + query_payload: bytes, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, int, bytes]], + deadline: float, + cancelled: Callable[[], bool], + ) -> tuple[list[tuple[int, float]], float, float]: + if not documents: + return [], 0.0, 0.0 + query_cpu = self._cpu_tensor(query_payload, query_rows, dimension, dtype) + results: list[tuple[int, float]] = [] + queue_started = time.monotonic() + remaining = deadline - time.monotonic() + if remaining <= 0 or not self.compute_slots.acquire(timeout=remaining): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired while waiting for CUDA capacity", + ) + queue_ms = (time.monotonic() - queue_started) * 1000.0 + compute_started = time.monotonic() + try: + with torch.inference_mode(): + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + query_device = query_cpu.to(self.device) + for group in self._groups(query_rows, dimension, documents): + if cancelled(): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request peer disconnected" + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + cpu_documents = [ + self._cpu_tensor(payload, rows, dimension, dtype) + for _, rows, payload in group + ] + document_device = torch.cat(cpu_documents).to(self.device) + similarities = query_device @ document_device.transpose(0, 1) + scores = [] + offset = 0 + for _, rows, _ in group: + scores.append( + similarities[:, offset : offset + rows] + .amax(dim=1) + .sum(dtype=torch.float32) + ) + offset += rows + host_scores = torch.stack(scores).to(device="cpu").tolist() + for (candidate_id, _, _), score in zip( + group, host_scores, strict=True + ): + if not math.isfinite(score): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "TileMaxSim result is non-finite", + ) + results.append((candidate_id, score)) + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + finally: + self.compute_slots.release() + return ( + results, + queue_ms, + (time.monotonic() - compute_started) * 1000.0, + ) + + +class JsonMetrics: + def __init__(self) -> None: + self.lock = threading.Lock() + + def emit(self, fields: dict[str, object]) -> None: + with self.lock: + print(json.dumps(fields, separators=(",", ":"), sort_keys=True), flush=True) + + +class TileMaxsimService: + def __init__( + self, + limits: protocol.Limits, + resolver: ContentAddressedResolver, + engine: TorchTileMaxsimEngine, + request_timeout_ms: int, + metrics: JsonMetrics, + ) -> None: + self.limits = limits + self.resolver = resolver + self.engine = engine + self.request_timeout_seconds = request_timeout_ms / 1000.0 + self.metrics = metrics + + @staticmethod + def _peer_disconnected(connection: socket.socket) -> bool: + poller = select.poll() + poller.register(connection, select.POLLHUP | select.POLLERR | select.POLLNVAL) + return bool(poller.poll(0)) + + @staticmethod + def _receive_exact_until( + connection: socket.socket, count: int, deadline: float + ) -> bytes: + chunks = bytearray() + while len(chunks) < count: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("request deadline expired during socket read") + connection.settimeout(remaining) + chunk = connection.recv(count - len(chunks)) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "connection closed during request", + ) + chunks.extend(chunk) + return bytes(chunks) + + def process_frame( + self, + frame: bytes, + connection: socket.socket, + deadline: float, + peer_credentials: tuple[int, int, int] | None, + ) -> bytes: + request_id = 0 + version = protocol.VERSION + started = time.monotonic() + metrics: dict[str, object] = {"event": "tilemaxsim_request"} + if peer_credentials is not None: + metrics["peer_pid"], metrics["peer_uid"], metrics["peer_gid"] = ( + peer_credentials + ) + try: + if len(frame) >= protocol.HEADER.size: + _, wire_version, _, request_id, _ = protocol.HEADER.unpack_from(frame) + if wire_version in (protocol.VERSION, protocol.EXTERNAL_VERSION): + version = wire_version + request = protocol.parse_request_frame( + frame, self.limits, validate_finite=False + ) + metrics.update( + request_id=request.request_id, + protocol_version=version, + query_rows=request.query_rows, + dimension=request.dimension, + candidate_count=len(request.candidates), + ) + validate_finite_payload( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + ) + resolve_started = time.monotonic() + cache_hits = 0 + documents: list[tuple[int, int, bytes]] = [] + if isinstance(request, protocol.InlineTensorRequest): + for candidate in request.candidates: + validate_finite_payload( + candidate.payload, + candidate.rows, + request.dimension, + request.dtype, + ) + documents = [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in request.candidates + ] + metrics["source"] = "inline" + else: + metrics["source"] = "content_addressed" + for candidate in request.candidates: + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during tensor resolution", + ) + resolved = self.resolver.resolve(candidate.descriptor) + cache_hits += int(resolved.cache_hit) + documents.append( + ( + candidate.candidate_id, + candidate.descriptor.rows, + resolved.payload, + ) + ) + metrics["cache_hits"] = cache_hits + metrics["resolve_ms"] = round( + (time.monotonic() - resolve_started) * 1000.0, 3 + ) + metrics["document_tokens"] = sum(rows for _, rows, _ in documents) + results, queue_ms, compute_ms = self.engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + metrics["queue_ms"] = round(queue_ms, 3) + metrics["compute_ms"] = round(compute_ms, 3) + metrics["status"] = "ok" + return protocol.success_response(request.request_id, results, version) + except protocol.SidecarError as error: + metrics.update(status="error", error_class=type(error).__name__) + return protocol.error_response( + request_id, error.status, str(error), version + ) + except torch.OutOfMemoryError: + if self.engine.device.type == "cuda": + torch.cuda.empty_cache() + metrics.update(status="error", error_class="CudaOutOfMemory") + return protocol.error_response( + request_id, + protocol.STATUS_RESOURCE_LIMIT, + "CUDA out of memory", + version, + ) + except Exception as error: + metrics.update(status="error", error_class=type(error).__name__) + return protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + f"TileMaxSim compute failed: {error}", + version, + ) + finally: + metrics["total_ms"] = round((time.monotonic() - started) * 1000.0, 3) + self.metrics.emit(metrics) + + def handle(self, connection: socket.socket) -> None: + deadline = time.monotonic() + self.request_timeout_seconds + request_id = 0 + version = protocol.VERSION + peer_credentials = None + if hasattr(socket, "SO_PEERCRED"): + try: + raw_credentials = connection.getsockopt( + socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i") + ) + peer_credentials = struct.unpack("3i", raw_credentials) + except OSError: + pass + try: + header = self._receive_exact_until( + connection, protocol.HEADER.size, deadline + ) + _, wire_version, _, request_id, body_len = protocol.HEADER.unpack(header) + if wire_version in (protocol.VERSION, protocol.EXTERNAL_VERSION): + version = wire_version + if body_len > self.limits.max_request_bytes - protocol.HEADER.size: + response = protocol.error_response( + request_id, + protocol.STATUS_RESOURCE_LIMIT, + "request exceeds byte limit", + version, + ) + else: + body = self._receive_exact_until(connection, body_len, deadline) + response = self.process_frame( + header + body, connection, deadline, peer_credentials + ) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("request deadline expired before socket write") + connection.settimeout(remaining) + connection.sendall(response) + except (BrokenPipeError, ConnectionResetError): + return + except (TimeoutError, socket.timeout): + try: + connection.sendall( + protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during socket I/O", + version, + ) + ) + except OSError: + pass + except Exception as error: + try: + connection.sendall( + protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + str(error), + version, + ) + ) + except OSError: + pass + + +def serve( + socket_path: Path, + socket_mode: int, + backlog: int, + max_inflight: int, + service: TileMaxsimService, + stop: threading.Event, + once: bool = False, +) -> None: + protocol.remove_stale_socket(socket_path) + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + slots = threading.BoundedSemaphore(max_inflight) + workers = ThreadPoolExecutor( + max_workers=max_inflight, thread_name_prefix="tilemaxsim" + ) + active = 0 + active_lock = threading.Lock() + + def handle(connection: socket.socket) -> None: + nonlocal active + try: + with connection: + service.handle(connection) + finally: + with active_lock: + active -= 1 + slots.release() + + try: + listener.bind(os.fspath(socket_path)) + os.chmod(socket_path, socket_mode) + listener.listen(backlog) + listener.settimeout(0.25) + bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino + service.metrics.emit( + { + "event": "tilemaxsim_ready", + "device": str(service.engine.device), + "max_inflight": max_inflight, + "socket": os.fspath(socket_path), + } + ) + accepted = 0 + while not stop.is_set(): + if not slots.acquire(timeout=0.25): + continue + try: + connection, _ = listener.accept() + except TimeoutError: + slots.release() + continue + with active_lock: + active += 1 + current_active = active + service.metrics.emit( + {"event": "tilemaxsim_accept", "inflight": current_active} + ) + workers.submit(handle, connection) + accepted += 1 + if once and accepted == 1: + break + finally: + listener.close() + workers.shutdown(wait=True, cancel_futures=False) + try: + current = socket_path.lstat() + if (current.st_dev, current.st_ino) == bound_identity: + socket_path.unlink() + except (FileNotFoundError, UnboundLocalError): + pass + + +def parse_mode(value: str) -> int: + mode = int(value, 8) + if mode < 0 or mode > 0o777: + raise argparse.ArgumentTypeError("socket mode must be between 000 and 777") + return mode + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def nonnegative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("value must be nonnegative") + return parsed + + +def contract_roots( + values: list[str], parser: argparse.ArgumentParser +) -> dict[str, Path]: + roots = {} + for value in values: + if "=" not in value: + parser.error("--contract-root must be MODEL_CONTRACT_ID=/absolute/path") + contract, raw_path = value.split("=", 1) + path = Path(raw_path) + if not contract or not path.is_absolute(): + parser.error("--contract-root must contain a nonempty ID and absolute path") + if contract in roots: + parser.error(f"duplicate --contract-root for {contract!r}") + roots[contract] = path + return roots + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=parse_mode, default=0o600) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--contract-root", action="append", default=[]) + parser.add_argument( + "--max-request-bytes", type=positive_int, default=64 * 1024 * 1024 + ) + parser.add_argument("--max-batch-tokens", type=positive_int, default=1_000_000) + parser.add_argument( + "--max-tensor-bytes", type=positive_int, default=1024 * 1024 * 1024 + ) + parser.add_argument("--max-candidates", type=positive_int, default=65_536) + parser.add_argument( + "--max-device-bytes", type=positive_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument( + "--cache-bytes", type=nonnegative_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument("--request-timeout-ms", type=positive_int, default=2000) + parser.add_argument("--max-inflight", type=positive_int, default=8) + parser.add_argument("--max-cuda-inflight", type=positive_int, default=1) + parser.add_argument("--backlog", type=positive_int, default=64) + parser.add_argument("--allow-tf32", action="store_true") + parser.add_argument("--once", action="store_true") + args = parser.parse_args() + + roots = contract_roots(args.contract_root, parser) + limits = protocol.Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + resolver = ContentAddressedResolver(roots, args.cache_bytes) + metrics = JsonMetrics() + try: + engine = TorchTileMaxsimEngine( + args.device, + args.max_device_bytes, + args.allow_tf32, + args.max_cuda_inflight, + ) + service = TileMaxsimService( + limits, resolver, engine, args.request_timeout_ms, metrics + ) + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + serve( + args.socket, + args.socket_mode, + args.backlog, + args.max_inflight, + service, + stop, + args.once, + ) + finally: + resolver.close() + + +if __name__ == "__main__": + main() diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 3725f050..ac550103 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -29,3 +29,17 @@ mod text_rabitq8; pub mod typmod; mod typmod_rabitq4; mod typmod_rabitq8; + +pub(crate) const MAX_MAXSIM_VECTORS: usize = u16::MAX as usize + 1; + +pub(crate) fn validate_maxsim_array_len(len: usize) { + if len == 0 { + pgrx::error!("MaxSim arrays must contain at least one vector"); + } + if len > MAX_MAXSIM_VECTORS { + pgrx::error!( + "MaxSim arrays cannot contain more than {} vectors", + MAX_MAXSIM_VECTORS + ); + } +} diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index 0ae307fe..f20ebc8d 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -95,12 +95,20 @@ fn _vchord_halfvec_operator_maxsim( lhs: Array<'_, HalfvecInput<'_>>, rhs: Array<'_, HalfvecInput<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/datatype/operators_rabitq4.rs b/src/datatype/operators_rabitq4.rs index bf0ff920..d1980fe9 100644 --- a/src/datatype/operators_rabitq4.rs +++ b/src/datatype/operators_rabitq4.rs @@ -125,12 +125,20 @@ fn _vchord_rabitq4_operator_maxsim( lhs: Array<'_, Rabitq4Input<'_>>, rhs: Array<'_, Rabitq4Input<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(Rabitq4Borrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/datatype/operators_rabitq8.rs b/src/datatype/operators_rabitq8.rs index d8d3cb72..7cac9aac 100644 --- a/src/datatype/operators_rabitq8.rs +++ b/src/datatype/operators_rabitq8.rs @@ -125,12 +125,20 @@ fn _vchord_rabitq8_operator_maxsim( lhs: Array<'_, Rabitq8Input<'_>>, rhs: Array<'_, Rabitq8Input<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(Rabitq8Borrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index 76b3bf8c..c125a18b 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -95,12 +95,20 @@ fn _vchord_vector_operator_maxsim( lhs: Array<'_, VectorInput<'_>>, rhs: Array<'_, VectorInput<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs index 704edfc8..fa836316 100644 --- a/src/index/fetcher.rs +++ b/src/index/fetcher.rs @@ -24,6 +24,21 @@ pub trait FilterableTuple: Tuple { pub trait Tuple { fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); + + /// Read one user attribute from the already fetched heap slot. + /// + /// Callers that may expose the value outside PostgreSQL must call + /// [`FilterableTuple::filter`] first. Attribute numbers are PostgreSQL + /// one-based `attnum` values, not zero-based Rust indexes. + #[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] + fn attribute(&mut self, attnum: i16) -> Option; +} + +#[derive(Clone, Copy)] +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +pub struct TupleAttribute { + pub datum: Datum, + pub is_null: bool, } pub trait Fetcher { @@ -52,6 +67,7 @@ pub struct HeapFetcher { heap_relation: pgrx::pg_sys::Relation, snapshot: pgrx::pg_sys::Snapshot, heapfetch: *mut pgrx::pg_sys::IndexFetchTableData, + owns_heapfetch: bool, slot: *mut pgrx::pg_sys::TupleTableSlot, values: [Datum; 32], is_nulls: [bool; 32], @@ -77,6 +93,7 @@ impl HeapFetcher { heap_relation, snapshot, heapfetch, + owns_heapfetch: false, slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), values: [Datum::null(); 32], is_nulls: [true; 32], @@ -84,6 +101,44 @@ impl HeapFetcher { } } } + + /// Create a heap fetch state that is not owned by an `IndexScanDesc`. + /// + /// This is used by the restricted external MaxSim executor to resolve the + /// root TIDs stored in the index through HOT chains before SQL-visible + /// descriptor projection. The table AM owns the rules for doing that; + /// looking up `ctid` directly does not follow a HOT chain. + pub unsafe fn new_standalone( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + ) -> Self { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + unsafe { + let table_am = (*heap_relation).rd_tableam; + if table_am.is_null() { + panic!("unknown heap access method"); + } + let index_fetch_begin = (*table_am) + .index_fetch_begin + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + let heapfetch = pg_guard_ffi_boundary(|| index_fetch_begin(heap_relation)); + if heapfetch.is_null() { + panic!("heap access method returned a null index fetch state"); + } + let mut fetcher = Self::new( + index_relation, + heap_relation, + snapshot, + heapfetch, + std::ptr::null_mut(), + ); + fetcher.owns_heapfetch = true; + fetcher + } + } } impl Drop for HeapFetcher { @@ -93,6 +148,16 @@ impl Drop for HeapFetcher { // free common resources pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); pgrx::pg_sys::FreeExecutorState(self.estate); + if self.owns_heapfetch { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + let table_am = (*self.heap_relation).rd_tableam; + let index_fetch_end = (*table_am) + .index_fetch_end + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| index_fetch_end(self.heapfetch)); + } } } } @@ -147,7 +212,14 @@ impl Fetcher for HeapFetcher { false }; if found { - Some(HeapTuple { this: self }) + // The heap table AM rewrites the requested root TID to the + // snapshot-visible HOT-chain member. The slot itself keeps + // the rewritten index TID as well, but carrying it explicitly + // avoids depending on slot representation details. + Some(HeapTuple { + this: self, + current_ctid: ctid, + }) } else { None } @@ -157,6 +229,16 @@ impl Fetcher for HeapFetcher { pub struct HeapTuple<'a> { this: &'a mut HeapFetcher, + current_ctid: ItemPointerData, +} + +impl HeapTuple<'_> { + /// Return the physical TID of the tuple version materialized in the slot. + /// This may differ from the root TID supplied to `Fetcher::fetch` after a + /// HOT update. + pub fn ctid(&self) -> ItemPointerData { + self.current_ctid + } } impl Tuple for HeapTuple<'_> { @@ -175,6 +257,30 @@ impl Tuple for HeapTuple<'_> { (&this.values, &this.is_nulls) } } + + fn attribute(&mut self, attnum: i16) -> Option { + unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + let slot = self.this.slot; + let tuple_descriptor = (*slot).tts_tupleDescriptor; + if attnum <= 0 + || tuple_descriptor.is_null() + || i32::from(attnum) > (*tuple_descriptor).natts + { + return None; + } + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + pgrx::pg_sys::slot_getsomeattrs_int(slot, i32::from(attnum)); + }); + let offset = usize::try_from(attnum - 1).ok()?; + Some(TupleAttribute { + datum: *(*slot).tts_values.add(offset), + is_null: *(*slot).tts_isnull.add(offset), + }) + } + } } impl FilterableTuple for HeapTuple<'_> { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index f2724593..26cde468 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -27,6 +27,18 @@ pub enum PostgresIo { ReadStream, } +#[derive(Debug, Clone, Copy, PostgresGucEnum)] +pub enum PostgresMaxsimBackend { + #[name = c"coarse_only"] + CoarseOnly, + #[name = c"cpu_exact"] + CpuExact, + #[name = c"gpu"] + Gpu, + #[name = c"auto"] + Auto, +} + static VCHORDRQ_QUERY_SAMPLING_ENABLE: GucSetting = GucSetting::::new(false); static VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS: GucSetting = GucSetting::::new(0); @@ -78,6 +90,24 @@ static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); +static VCHORDRQ_MAXSIM_CANDIDATE_LIMIT: GucSetting = GucSetting::::new(-1); + +static VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS: GucSetting = GucSetting::::new(32); + +static VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS: GucSetting = GucSetting::::new(256); + +static VCHORDRQ_MAXSIM_BACKEND: GucSetting = + GucSetting::::new(PostgresMaxsimBackend::CoarseOnly); + +static VCHORDRQ_MAXSIM_GPU_ENDPOINT: GucSetting> = + GucSetting::>::new(Some(c"")); + +static VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS: GucSetting = GucSetting::::new(2000); + +static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS: GucSetting = GucSetting::::new(1_000_000); + +static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES: GucSetting = GucSetting::::new(1_073_741_824); + static VCHORDRQ_PREFILTER: GucSetting = GucSetting::::new(false); static VCHORDRQ_IO_SEARCH: GucSetting = GucSetting::::new( @@ -151,6 +181,82 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_candidate_limit", + c"Maximum number of page candidates produced by MaxSim aggregation.", + c"Use -1 for no page-candidate limit.", + &VCHORDRQ_MAXSIM_CANDIDATE_LIMIT, + -1, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_planner_query_tokens", + c"Expected MaxSim query-token count used by the planner.", + c"Set this to the measured deployment average until expression statistics are available.", + &VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS, + 1, + 65_536, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_planner_document_tokens", + c"Fallback MaxSim document-token count used by the planner.", + c"Used for indexes that predate the native indexed-vector statistic; set it to the measured deployment average.", + &VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS, + 1, + 65_536, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordrq.maxsim_backend", + c"Page-level MaxSim rerank backend.", + c"GPU and auto modes require the native sidecar integration.", + &VCHORDRQ_MAXSIM_BACKEND, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_string_guc( + c"vchordrq.maxsim_gpu_endpoint", + c"Unix-socket endpoint for the TileMaxSim sidecar.", + c"An empty endpoint disables GPU transport.", + &VCHORDRQ_MAXSIM_GPU_ENDPOINT, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_timeout_ms", + c"Overall TileMaxSim sidecar deadline in milliseconds.", + c"The deadline covers connection, request write, and response read.", + &VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS, + 1, + 600_000, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_max_batch_tokens", + c"Maximum query plus document tokens in one TileMaxSim request.", + c"Requests exceeding the limit fail before connecting to the sidecar.", + &VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS, + 1, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_max_batch_bytes", + c"Maximum encoded TileMaxSim request size in bytes.", + c"Requests exceeding the limit fail before connecting to the sidecar.", + &VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES, + 1024, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); GucRegistry::define_bool_guc( c"vchordrq.prefilter", c"`prefilter` argument of vchordrq.", @@ -475,6 +581,39 @@ pub fn vchordrq_maxsim_threshold(index: pgrx::pg_sys::Relation) -> u32 { } } +pub fn vchordrq_maxsim_candidate_limit() -> Option { + let value = VCHORDRQ_MAXSIM_CANDIDATE_LIMIT.get(); + if value < 0 { None } else { Some(value as u32) } +} + +pub fn vchordrq_maxsim_planner_query_tokens() -> u32 { + VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_planner_document_tokens() -> u32 { + VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_backend() -> PostgresMaxsimBackend { + VCHORDRQ_MAXSIM_BACKEND.get() +} + +pub fn vchordrq_maxsim_gpu_endpoint() -> Option { + VCHORDRQ_MAXSIM_GPU_ENDPOINT.get() +} + +pub fn vchordrq_maxsim_gpu_timeout_ms() -> u32 { + VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS.get() as u32 +} + +pub fn vchordrq_maxsim_gpu_max_batch_tokens() -> u32 { + VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_gpu_max_batch_bytes() -> u32 { + VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES.get() as u32 +} + pub fn vchordrq_prefilter() -> bool { VCHORDRQ_PREFILTER.get() } diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 647d2478..95b4bc45 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -349,9 +349,10 @@ pub unsafe extern "C-unwind" fn ambuild( reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); order }, - |indtuples| { + |indtuples, indexed_vectors| { reporter.tuples_done(indtuples); reporter.tuples_total(indtuples); + vchordrq::set_indexed_vectors(&index, indexed_vectors); // enter the barrier let shared = leader.vchordrqshared; pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); @@ -429,9 +430,10 @@ pub unsafe extern "C-unwind" fn ambuild( || { reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); }, - |indtuples| { + |indtuples, indexed_vectors| { reporter.tuples_done(indtuples); reporter.tuples_total(indtuples); + vchordrq::set_indexed_vectors(&index, indexed_vectors); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); }, || {}, @@ -460,6 +462,7 @@ struct VchordrqShared { barrier_enter_0: i32, nparticipants: u32, indtuples: u64, + indexed_vectors: u64, barrier_leave_0: bool, barrier_enter_1: i32, barrier_leave_1: bool, @@ -721,6 +724,7 @@ impl VchordrqLeader { barrier_leave_2: false, mutex: std::mem::zeroed(), indtuples: 0, + indexed_vectors: 0, }); pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_enter_0); pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_leave_0); @@ -875,7 +879,7 @@ pub unsafe extern "C-unwind" fn vchordrq_parallel_build_main( pgrx::pg_sys::ConditionVariableCancelSleep(); order }, - |_| { + |_, _| { // enter the barrier let shared = vchordrqshared; pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); @@ -941,7 +945,7 @@ unsafe fn parallel_build( vchordrqcached: *const u8, mut callback: impl FnMut(u64), sync_0: impl FnOnce() -> u32, - sync_1: impl FnOnce(u64), + sync_1: impl FnOnce(u64, u64), sync_2: impl FnOnce(), ) { use vchordrq_cached::VchordrqCachedReader; @@ -989,6 +993,7 @@ unsafe fn parallel_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + let indexed_vectors = store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1010,6 +1015,7 @@ unsafe fn parallel_build( { pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); (*vchordrqshared).indtuples += 1; + (*vchordrqshared).indexed_vectors += indexed_vectors; indtuples = (*vchordrqshared).indtuples; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } @@ -1029,6 +1035,7 @@ unsafe fn parallel_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + let indexed_vectors = store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1050,6 +1057,7 @@ unsafe fn parallel_build( { pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); (*vchordrqshared).indtuples += 1; + (*vchordrqshared).indexed_vectors += indexed_vectors; indtuples = (*vchordrqshared).indtuples; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } @@ -1065,7 +1073,13 @@ unsafe fn parallel_build( drop(index); - sync_1(unsafe { (*vchordrqshared).indtuples }); + let (indtuples, indexed_vectors) = unsafe { + ( + (*vchordrqshared).indtuples, + (*vchordrqshared).indexed_vectors, + ) + }; + sync_1(indtuples, indexed_vectors); let index = unsafe { PostgresRelation::new(index_relation) }; @@ -1085,7 +1099,7 @@ unsafe fn sequential_build( vchordrqcached: &[u8], mut callback: impl FnMut(u64), sync_0: impl FnOnce(), - sync_1: impl FnOnce(u64), + sync_1: impl FnOnce(u64, u64), sync_2: impl FnOnce(), ) { use vchordrq_cached::VchordrqCachedReader; @@ -1125,6 +1139,7 @@ unsafe fn sequential_build( let index = unsafe { BufferedPostgresRelation::new(index_relation) }; let mut indtuples = 0; + let mut indexed_vectors = 0_u64; match cached { VchordrqCachedReader::_0(_) => { traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { @@ -1134,6 +1149,7 @@ unsafe fn sequential_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + indexed_vectors += store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1166,6 +1182,7 @@ unsafe fn sequential_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + indexed_vectors += store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1194,7 +1211,7 @@ unsafe fn sequential_build( drop(index); - sync_1(indtuples); + sync_1(indtuples, indexed_vectors); let index = unsafe { PostgresRelation::new(index_relation) }; diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index f995c600..976692b6 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -322,28 +322,13 @@ pub unsafe extern "C-unwind" fn amcostestimate( if !(*index_opt_info).hypothetical { let relation = Index::open((*index_opt_info).indexoid, pgrx::pg_sys::NoLock as _); let opfamily = opfamily(relation.raw()); - if !matches!( + let is_maxsim = matches!( opfamily, - Opfamily::HalfvecCosine - | Opfamily::HalfvecIp - | Opfamily::HalfvecL2 - | Opfamily::VectorCosine - | Opfamily::VectorIp - | Opfamily::VectorL2 - | Opfamily::Rabitq8Cosine - | Opfamily::Rabitq8Ip - | Opfamily::Rabitq8L2 - | Opfamily::Rabitq4Cosine - | Opfamily::Rabitq4Ip - | Opfamily::Rabitq4L2 - ) { - *index_startup_cost = 0.0; - *index_total_cost = 0.0; - *index_selectivity = 1.0; - *index_correlation = 0.0; - *index_pages = 1.0; - return; - } + Opfamily::VectorMaxsim + | Opfamily::HalfvecMaxsim + | Opfamily::Rabitq8Maxsim + | Opfamily::Rabitq4Maxsim + ); let index = PostgresRelation::::new(relation.raw()); let probes = gucs::vchordrq_probes(relation.raw()); let cost = vchordrq::cost(&index); @@ -354,15 +339,25 @@ pub unsafe extern "C-unwind" fn amcostestimate( probes.len() ); } + let estimated_index_vectors = if is_maxsim { + cost.indexed_vectors.map_or_else( + || { + (*index_opt_info).tuples.max(total_rows).max(1.0) + * f64::from(gucs::vchordrq_maxsim_planner_document_tokens()) + }, + |indexed_vectors| indexed_vectors as f64, + ) + } else { + (*index_opt_info).tuples.max(0.0) + }; let node_count = { - let tuples = (*index_opt_info).tuples as u32; let mut count = 0.0; - let r = cost.cells.iter().copied().rev(); - let numerator = std::iter::once(1).chain(probes.clone()); + let r = cost.cells.iter().copied().rev().map(f64::from); + let numerator = std::iter::once(1.0).chain(probes.iter().copied().map(f64::from)); let denumerator = r.clone(); - let scale = r.skip(1).chain(std::iter::once(tuples)); + let scale = r.skip(1).chain(std::iter::once(estimated_index_vectors)); for (scale, (numerator, denumerator)) in scale.zip(numerator.zip(denumerator)) { - count += (scale as f64) * 1.0f64.min((numerator as f64) / (denumerator as f64)); + count += scale * 1.0f64.min(numerator / denumerator); } count }; @@ -378,13 +373,41 @@ pub unsafe extern "C-unwind" fn amcostestimate( pages += cost.cells[0] as f64; pages }; + if is_maxsim { + let backend = match gucs::vchordrq_maxsim_backend() { + gucs::PostgresMaxsimBackend::CoarseOnly => { + vchordrq::MaxsimCostBackend::CoarseOnly + } + gucs::PostgresMaxsimBackend::CpuExact => vchordrq::MaxsimCostBackend::CpuExact, + gucs::PostgresMaxsimBackend::Gpu => vchordrq::MaxsimCostBackend::Gpu, + gucs::PostgresMaxsimBackend::Auto => vchordrq::MaxsimCostBackend::Auto, + }; + let estimate = vchordrq::estimate_maxsim_cost(vchordrq::MaxsimCostInput { + heap_rows: total_rows, + index_tokens: estimated_index_vectors, + token_nodes_per_query: node_count, + base_index_pages: page_count, + dimension: cost.dim, + element_bits: opfamily.vector_kind().number_of_bits_of_an_elements(), + query_tokens: gucs::vchordrq_maxsim_planner_query_tokens(), + limit_tuples: ((*root).limit_tuples > 0.0).then_some((*root).limit_tuples), + filter_selectivity, + candidate_limit: gucs::vchordrq_maxsim_candidate_limit(), + backend, + }); + *index_startup_cost = estimate.startup_cost; + *index_total_cost = estimate.total_cost; + *index_selectivity = estimate.selectivity; + *index_correlation = 0.0; + *index_pages = estimate.index_pages; + return; + } // `next_count` represents candidates we expect to process to // surface `limit_tuples` survivors after filter rejection. Clamp // by `node_count` so the estimate cannot exceed the candidates // the IVF visits at the configured probe count. let next_count = if (*root).limit_tuples > 0.0 { - ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)) - .min(node_count) + ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)).min(node_count) } else { node_count }; @@ -482,7 +505,9 @@ pub unsafe extern "C-unwind" fn ambulkdelete( pg_guard_ffi_boundary(|| callback(&mut ctid, callback_state)) } }; - crate::index::vchordrq::dispatch::bulkdelete(opfamily, &index, check, callback); + let indexed_vectors = + crate::index::vchordrq::dispatch::bulkdelete(opfamily, &index, check, callback); + vchordrq::set_indexed_vectors(&index, indexed_vectors); stats } @@ -538,6 +563,12 @@ pub unsafe extern "C-unwind" fn amrescan( max_scan_tuples: gucs::vchordrq_max_scan_tuples(), maxsim_refine: gucs::vchordrq_maxsim_refine((*scan).indexRelation), maxsim_threshold: gucs::vchordrq_maxsim_threshold((*scan).indexRelation), + maxsim_candidate_limit: gucs::vchordrq_maxsim_candidate_limit(), + maxsim_backend: gucs::vchordrq_maxsim_backend(), + maxsim_gpu_endpoint: gucs::vchordrq_maxsim_gpu_endpoint(), + maxsim_gpu_timeout_ms: gucs::vchordrq_maxsim_gpu_timeout_ms(), + maxsim_gpu_max_batch_tokens: gucs::vchordrq_maxsim_gpu_max_batch_tokens(), + maxsim_gpu_max_batch_bytes: gucs::vchordrq_maxsim_gpu_max_batch_bytes(), io_search: gucs::vchordrq_io_search(), io_rerank: gucs::vchordrq_io_rerank(), prefilter: gucs::vchordrq_prefilter(), diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs index 01d70a87..36baee63 100644 --- a/src/index/vchordrq/dispatch.rs +++ b/src/index/vchordrq/dispatch.rs @@ -71,42 +71,51 @@ pub fn bulkdelete( index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, -) where +) -> u64 +where R: RelationRead + RelationWrite, R::Page: Page, { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + live } (VectorKind::Vecf32, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + live } (VectorKind::Vecf16, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + live } (VectorKind::Vecf16, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + live } (VectorKind::Rabitq8, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } (VectorKind::Rabitq8, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } (VectorKind::Rabitq4, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } (VectorKind::Rabitq4, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } } } diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs index 7c51f5db..075b1075 100644 --- a/src/index/vchordrq/opclass.rs +++ b/src/index/vchordrq/opclass.rs @@ -92,6 +92,7 @@ impl Opfamily { Self::VectorMaxsim => { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -105,6 +106,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -118,6 +120,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -131,6 +134,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -203,6 +207,7 @@ impl Opfamily { Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Vecf32(vector.as_borrowed()))); @@ -213,6 +218,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Vecf16(vector.as_borrowed()))); @@ -223,6 +229,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Rabitq8(vector.as_borrowed()))); @@ -233,6 +240,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Rabitq4(vector.as_borrowed()))); diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 1c51c97f..238dbe2e 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -12,7 +12,17 @@ // // Copyright (c) 2025-2026 TensorChord Inc. +mod candidate; +mod external; +mod gpu; +mod rerank; +mod search; + +use self::candidate::{LegacyPageCandidateGenerator, PageCandidateGenerator}; +use self::gpu::{GpuTileMaxsimBackend, UnixSocketTransport, report_gpu_fallback}; +use self::rerank::{CpuExactMaxsimBackend, ExactMaxsimBackend, HeapArrayTensorSource}; use crate::index::fetcher::*; +use crate::index::gucs::PostgresMaxsimBackend; use crate::index::scanners::{Io, SearchBuilder}; use crate::index::vchordrq::dispatch::*; use crate::index::vchordrq::filter::filter; @@ -21,7 +31,6 @@ use crate::index::vchordrq::scanners::SearchOptions; use crate::recorder::Recorder; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; -use distance::Distance; use index::bump::Bump; use index::packed::PackedRefMut8; use index::prefetcher::*; @@ -29,8 +38,8 @@ use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadS use index_accessor::Dot; use simd::f16; use std::cmp::Reverse; -use std::collections::BinaryHeap; use std::num::NonZero; +use std::time::Duration; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; use vchordrq::{RerankMethod, how, maxsim_search, rerank_index}; use vector::VectorOwned; @@ -99,10 +108,33 @@ impl SearchBuilder for MaxsimBuilder { } let maxsim_refine = options.maxsim_refine; let maxsim_threshold = options.maxsim_threshold; + let maxsim_candidate_limit = options + .maxsim_candidate_limit + .map_or(usize::MAX, |value| value as usize); + let has_maxsim_candidate_limit = options.maxsim_candidate_limit.is_some(); + let maxsim_backend = options.maxsim_backend; + let maxsim_gpu_endpoint = options + .maxsim_gpu_endpoint + .as_deref() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let maxsim_gpu_timeout = Duration::from_millis(options.maxsim_gpu_timeout_ms as u64); + let maxsim_gpu_max_batch_tokens = options.maxsim_gpu_max_batch_tokens as usize; + let maxsim_gpu_max_batch_bytes = options.maxsim_gpu_max_batch_bytes as usize; + if !matches!(maxsim_backend, PostgresMaxsimBackend::CoarseOnly) + && (!has_maxsim_candidate_limit || maxsim_candidate_limit == 0) + { + pgrx::error!("exact MaxSim requires a positive vchordrq.maxsim_candidate_limit"); + } let opfamily = self.opfamily; let Some(vectors) = vectors else { return Box::new(std::iter::empty()) as Box>; }; + let expected_dim = vchordrq::cost(index).dim; + if vectors.iter().any(|vector| vector.dim() != expected_dim) { + pgrx::error!("dimension is not matched"); + } + let rerank_query = vectors.clone(); let method = how(index); if !matches!(method, RerankMethod::Index) { pgrx::error!("maxsim search with rerank_in_table is not supported"); @@ -124,669 +156,702 @@ impl SearchBuilder for MaxsimBuilder { _, AlwaysEqual, _, _)>>, )| (rough, payload); - let iter: Box> = match opfamily.vector_kind() { - VectorKind::Vecf32 => { - type Op = vchordrq::operator::Op, Dot>; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Vecf32(vector) = vector { - vector - } else { - unreachable!() - } - }) - .collect::>(); - let projected = unprojected - .iter() - .map(|vector| RandomProject::project(vector.as_borrowed())) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + let mut token_searches = { + let fetcher = &mut fetcher; + let iter: Box> = match opfamily.vector_kind() { + VectorKind::Vecf32 => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } - } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) - } - VectorKind::Vecf16 => { - type Op = vchordrq::operator::Op, Dot>; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Vecf16(vector) = vector { - vector } else { - unreachable!() + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - }) - .collect::>(); - let projected = unprojected - .iter() - .map(|vector| RandomProject::project(vector.as_borrowed())) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Vecf16 => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } - } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) - } - VectorKind::Rabitq8 => { - type Op = vchordrq::operator::Op; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Rabitq8(vector) = vector { - vector } else { - unreachable!() + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - }) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Rabitq8 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq8(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } - } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) - } - VectorKind::Rabitq4 => { - type Op = vchordrq::operator::Op; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Rabitq4(vector) = vector { - vector } else { - unreachable!() + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - }) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Rabitq4 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq4(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) + (accu_set, rough_set, estimation_by_threshold) + })) + } + }; + iter + }; + let mut candidates = + LegacyPageCandidateGenerator.generate(n, &mut token_searches, maxsim_candidate_limit); + drop(token_searches); + let iter: Box> = match maxsim_backend { + PostgresMaxsimBackend::CoarseOnly => Box::new(candidates.map(|candidate| { + let distance = candidate.approximate_distance.to_f32(); + let recheck = false; + (distance, candidate.heap_key, recheck) + })), + PostgresMaxsimBackend::CpuExact => { + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let exact = + CpuExactMaxsimBackend.rerank(&rerank_query, &mut candidates, &mut source); + let exact = exact.unwrap_or_else(|error| pgrx::error!("{error}")); + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) })) } - }; - let mut updates = Vec::new(); - let mut estimations = Vec::new(); - for (query_id, (accu_set, rough_set, estimation_by_threshold)) in iter.enumerate() { - updates.reserve(accu_set.len() + rough_set.len()); - let is_empty = accu_set.is_empty() && rough_set.is_empty(); - let mut estimation_by_scope = Distance::NEG_INFINITY; - for (distance, payload) in accu_set { - estimation_by_scope = std::cmp::max(estimation_by_scope, distance); - let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + PostgresMaxsimBackend::Gpu => { + let candidates = candidates.collect::>(); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let transport = UnixSocketTransport::new(maxsim_gpu_endpoint); + let mut backend = GpuTileMaxsimBackend::new( + transport, + maxsim_gpu_timeout, + maxsim_gpu_max_batch_tokens, + maxsim_gpu_max_batch_bytes, + ); + let exact = backend.rerank(&rerank_query, &mut candidate_iter, &mut source); + let exact = exact.unwrap_or_else(|error| pgrx::error!("{error}")); + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) + })) } - for (distance, payload) in rough_set { - let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + PostgresMaxsimBackend::Auto => { + let candidates = candidates.collect::>(); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let transport = UnixSocketTransport::new(maxsim_gpu_endpoint); + let mut backend = GpuTileMaxsimBackend::new( + transport, + maxsim_gpu_timeout, + maxsim_gpu_max_batch_tokens, + maxsim_gpu_max_batch_bytes, + ); + let exact = backend.rerank(&rerank_query, &mut candidate_iter, &mut source); + let exact = match exact { + Ok(exact) => exact, + Err(error) => { + report_gpu_fallback(&error); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + CpuExactMaxsimBackend + .rerank(&rerank_query, &mut candidate_iter, &mut source) + .unwrap_or_else(|error| pgrx::error!("{error}")) + } + }; + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) + })) } - estimations.push(if !is_empty { - std::cmp::max(estimation_by_scope, estimation_by_threshold) - } else { - Distance::ZERO - }); - } - updates.sort_unstable_by_key(|&(key, ..)| key); - let iter = updates - .chunk_by(|(kl, ..), (kr, ..)| kl == kr) - .map(|chunk| { - let key = chunk[0].0; - let mut value = vec![None; n]; - for &(_, query_id, distance) in chunk { - let this = value[query_id].get_or_insert(Distance::INFINITY); - *this = std::cmp::min(*this, distance); - } - let mut maxsim = 0.0f32; - for (query_id, distance) in value.into_iter().enumerate() { - let d = distance.unwrap_or(estimations[query_id]); - maxsim += Distance::to_f32(d); - } - (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) - }) - .collect::>() - .into_iter_sorted_polyfill() - .map(|(Reverse(distance), AlwaysEqual(key))| { - let distance = distance.to_f32(); - let recheck = false; - (distance, key, recheck) - }); - let iter: Box> = Box::new(iter); - let iter = if let Some(max_scan_tuples) = options.max_scan_tuples { - Box::new(iter.take(max_scan_tuples as _)) - } else { - iter }; #[allow(clippy::let_and_return)] iter } } -// Emulate unstable library feature `binary_heap_into_iter_sorted`. -// See https://github.com/rust-lang/rust/issues/59278. - -trait IntoIterSortedPolyfill { - fn into_iter_sorted_polyfill(self) -> IntoIterSorted; -} - -impl IntoIterSortedPolyfill for BinaryHeap { - fn into_iter_sorted_polyfill(self) -> IntoIterSorted { - IntoIterSorted { inner: self } - } -} - -#[derive(Clone, Debug)] -struct IntoIterSorted { - inner: BinaryHeap, -} - -impl Iterator for IntoIterSorted { - type Item = T; - - #[inline] - fn next(&mut self) -> Option { - self.inner.pop() - } - - #[inline] - fn size_hint(&self) -> (usize, Option) { - let exact = self.inner.len(); - (exact, Some(exact)) - } -} - -impl ExactSizeIterator for IntoIterSorted {} - -impl std::iter::FusedIterator for IntoIterSorted {} - #[inline(always)] pub fn id_0(f: F) -> F where diff --git a/src/index/vchordrq/scanners/maxsim/candidate.rs b/src/index/vchordrq/scanners/maxsim/candidate.rs new file mode 100644 index 00000000..9aa3e22f --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/candidate.rs @@ -0,0 +1,192 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::fetcher::pointer_to_kv; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::num::NonZero; + +pub(super) type HeapKey = [u16; 3]; +pub(super) type TokenCandidate = (Distance, NonZero); +pub(super) type TokenSearchResult = (Vec, Vec, Distance); + +#[derive(Clone, Copy, Debug)] +pub(super) struct PageCandidate { + pub approximate_distance: Distance, + pub heap_key: HeapKey, +} + +pub(super) trait PageCandidateGenerator { + type Candidates: Iterator; + + fn generate( + &mut self, + query_count: usize, + token_searches: &mut dyn Iterator, + candidate_limit: usize, + ) -> Self::Candidates; +} + +#[derive(Default)] +pub(super) struct LegacyPageCandidateGenerator; + +impl PageCandidateGenerator for LegacyPageCandidateGenerator { + type Candidates = LegacyPageCandidates; + + fn generate( + &mut self, + query_count: usize, + token_searches: &mut dyn Iterator, + candidate_limit: usize, + ) -> Self::Candidates { + let mut updates = Vec::new(); + let mut estimations = Vec::with_capacity(query_count); + for (query_id, (accu_set, rough_set, estimation_by_threshold)) in token_searches.enumerate() + { + updates.reserve(accu_set.len() + rough_set.len()); + let is_empty = accu_set.is_empty() && rough_set.is_empty(); + let mut estimation_by_scope = Distance::NEG_INFINITY; + for (distance, payload) in accu_set { + estimation_by_scope = std::cmp::max(estimation_by_scope, distance); + let (key, _) = pointer_to_kv(payload); + updates.push((key, query_id, distance)); + } + for (distance, payload) in rough_set { + let (key, _) = pointer_to_kv(payload); + updates.push((key, query_id, distance)); + } + estimations.push(if !is_empty { + std::cmp::max(estimation_by_scope, estimation_by_threshold) + } else { + Distance::ZERO + }); + } + debug_assert_eq!(estimations.len(), query_count); + updates.sort_unstable_by_key(|&(key, ..)| key); + let inner = updates + .chunk_by(|(kl, ..), (kr, ..)| kl == kr) + .map(|chunk| { + let key = chunk[0].0; + let mut value = vec![None; query_count]; + for &(_, query_id, distance) in chunk { + let this = value[query_id].get_or_insert(Distance::INFINITY); + *this = std::cmp::min(*this, distance); + } + let mut maxsim = 0.0f32; + for (query_id, distance) in value.into_iter().enumerate() { + let distance = distance.unwrap_or(estimations[query_id]); + maxsim += distance.to_f32(); + } + (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) + }) + .collect::>(); + LegacyPageCandidates { + inner, + remaining: candidate_limit, + } + } +} + +pub(super) struct LegacyPageCandidates { + inner: BinaryHeap<(Reverse, AlwaysEqual)>, + remaining: usize, +} + +impl Iterator for LegacyPageCandidates { + type Item = PageCandidate; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + return None; + } + let (Reverse(approximate_distance), AlwaysEqual(heap_key)) = self.inner.pop()?; + self.remaining -= 1; + Some(PageCandidate { + approximate_distance, + heap_key, + }) + } + + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len().min(self.remaining); + (exact, Some(exact)) + } +} + +impl ExactSizeIterator for LegacyPageCandidates {} +impl std::iter::FusedIterator for LegacyPageCandidates {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::index::fetcher::kv_to_pointer; + + fn distance(value: f32) -> Distance { + Distance::from_f32(value) + } + + #[test] + fn aggregates_tokens_by_page_and_preserves_distance_order() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let mut token_searches = vec![ + ( + vec![ + (distance(-1.0), kv_to_pointer((page_1, 0))), + (distance(-0.5), kv_to_pointer((page_2, 0))), + ], + vec![(distance(-0.4), kv_to_pointer((page_2, 1)))], + distance(-0.75), + ), + ( + vec![(distance(-2.0), kv_to_pointer((page_1, 1)))], + vec![], + distance(-1.0), + ), + ] + .into_iter(); + let candidates = LegacyPageCandidateGenerator + .generate(2, &mut token_searches, usize::MAX) + .collect::>(); + + assert_eq!(candidates.len(), 2); + assert_eq!(candidates[0].heap_key, page_1); + assert_eq!(candidates[0].approximate_distance.to_f32(), -3.0); + assert_eq!(candidates[1].heap_key, page_2); + assert_eq!(candidates[1].approximate_distance.to_f32(), -1.5); + } + + #[test] + fn candidate_limit_is_applied_after_global_ordering() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let mut token_searches = vec![( + vec![ + (distance(-1.0), kv_to_pointer((page_1, 0))), + (distance(-2.0), kv_to_pointer((page_2, 0))), + ], + vec![], + Distance::ZERO, + )] + .into_iter(); + let candidates = LegacyPageCandidateGenerator + .generate(1, &mut token_searches, 1) + .collect::>(); + + assert_eq!(candidates.len(), 1); + assert_eq!(candidates[0].heap_key, page_2); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/external.rs b/src/index/vchordrq/scanners/maxsim/external.rs new file mode 100644 index 00000000..0c8d1eb0 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/external.rs @@ -0,0 +1,551 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): This software is also available under the ELv2, +// which has specific restrictions. +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::PageCandidate; +use super::rerank::RerankError; +use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; +use pgrx::datum::{FromDatum, IntoDatum}; +use std::ffi::CString; + +const MAX_MODEL_CONTRACT_BYTES: usize = 512; +const MAX_TENSOR_REF_BYTES: usize = 4096; +const MAX_CHECKSUM_BYTES: usize = 512; +const MAX_TENSOR_ROWS: u32 = 65_536; +const MAX_TENSOR_DIMENSION: u32 = 60_000; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ExternalTensorDtype { + F32, + F16, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ExternalTensorStorage { + SameHeap, + DescriptorRelation, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) struct ExternalTensorColumns { + pub model_contract: i16, + pub public_id: i16, + pub tensor_ref: i16, + pub tensor_rows: i16, + pub tensor_dimension: i16, + pub tensor_dtype: i16, + pub tensor_checksum: i16, +} + +impl ExternalTensorColumns { + pub(super) fn validate(self) -> Result { + let columns = [ + self.model_contract, + self.public_id, + self.tensor_ref, + self.tensor_rows, + self.tensor_dimension, + self.tensor_dtype, + self.tensor_checksum, + ]; + if columns.iter().any(|column| *column <= 0) { + return Err(RerankError::InvalidDescriptor( + "registered attribute number is invalid", + )); + } + for (index, column) in columns.iter().enumerate() { + if columns[..index].contains(column) { + return Err(RerankError::InvalidDescriptor( + "registered descriptor columns are not distinct", + )); + } + } + Ok(self) + } +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorDescriptor { + pub candidate: PageCandidate, + /// Stable application identifier. It stays inside PostgreSQL and is never + /// encoded into the sidecar request. + #[allow( + dead_code, + reason = "returned by the score API through a separate visible-row map" + )] + pub public_id: i64, + pub tensor_ref: String, + pub rows: u32, + pub dimension: u32, + pub dtype: ExternalTensorDtype, + pub checksum: String, +} + +pub(super) trait CandidateTensorDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError>; +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +pub(super) struct HeapTensorRefSource<'a, F> { + fetcher: &'a mut F, + model_contract_id: String, + columns: ExternalTensorColumns, +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorSourceBinding { + pub index_oid: pgrx::pg_sys::Oid, + pub heap_oid: pgrx::pg_sys::Oid, + pub descriptor_oid: Option, + pub storage: ExternalTensorStorage, + pub model_contract_id: String, + #[allow( + dead_code, + reason = "physical attnums are consumed by the optional Phase 3C heap source" + )] + pub columns: Option, + pub column_names: ExternalTensorColumnNames, +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorColumnNames { + pub model_contract: String, + pub public_id: String, + pub descriptor_public_id: Option, + pub tensor_ref: String, + pub tensor_rows: String, + pub tensor_dimension: String, + pub tensor_dtype: String, + pub tensor_checksum: String, +} + +/// Resolve the privilege-aware SQL registry boundary. Same-heap bindings also +/// resolve the physical attribute numbers consumed by [`HeapTensorRefSource`]; +/// independent descriptor relations are projected later through SPI. +/// +/// The SECURITY DEFINER SQL function performs ownership/SELECT checks and +/// revalidates the live index, heap relation, opclass, column types, and NOT +/// NULL constraints. This Rust layer deliberately calls that function rather +/// than reading the private registry table with extension privileges. +pub(super) fn resolve_external_tensor_source( + index_oid: pgrx::pg_sys::Oid, +) -> Result { + pgrx::spi::Spi::connect(|client| { + let schema_rows = client + .select( + "SELECT n.nspname::text AS schema_name + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'", + Some(1), + &[], + ) + .map_err(registry_error)?; + if schema_rows.is_empty() { + return Err(RerankError::Registry( + "vchord extension schema is unavailable".into(), + )); + } + let schema_name = schema_rows + .first() + .get_by_name::("schema_name") + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry("vchord extension schema is NULL".into()))?; + let resolver = pgrx::spi::quote_qualified_identifier( + schema_name, + "vchordrq_maxsim_source_info".to_string(), + ); + let query = format!( + "SELECT registered_index::oid AS index_oid, + heap_relation::oid AS heap_oid, + descriptor_relation::oid AS descriptor_oid, + model_contract_id, + source_storage, + model_contract_column::text AS model_contract_column, + public_id_column::text AS public_id_column, + descriptor_public_id_column::text AS descriptor_public_id_column, + tensor_ref_column::text AS tensor_ref_column, + tensor_rows_column::text AS tensor_rows_column, + tensor_dim_column::text AS tensor_dim_column, + tensor_dtype_column::text AS tensor_dtype_column, + tensor_checksum_column::text AS tensor_checksum_column + FROM {resolver}($1::regclass)" + ); + let prepared = client + .prepare(query.as_str(), &pgrx::oids_of![pgrx::pg_sys::Oid]) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[index_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered MaxSim tensor source resolution returned no row".into(), + )); + } + let row = rows.first(); + let resolved_index_oid = required_column::(&row, "index_oid")?; + let heap_oid = required_column::(&row, "heap_oid")?; + let descriptor_oid = optional_column::(&row, "descriptor_oid")?; + let model_contract_id = required_column::(&row, "model_contract_id")?; + let storage = required_column::(&row, "source_storage")?; + if resolved_index_oid != index_oid { + return Err(RerankError::Registry( + "registered MaxSim tensor source resolved a different index".into(), + )); + } + let storage = match storage.as_str() { + "external_ref" if descriptor_oid.is_none() => ExternalTensorStorage::SameHeap, + "external_relation" if descriptor_oid.is_some() => { + ExternalTensorStorage::DescriptorRelation + } + "heap_array" => { + return Err(RerankError::Registry( + "registered MaxSim tensor source is not external".into(), + )); + } + _ => { + return Err(RerankError::Registry( + "registered external MaxSim tensor source is inconsistent".into(), + )); + } + }; + + let column_names = ExternalTensorColumnNames { + model_contract: required_column::(&row, "model_contract_column")?, + public_id: required_column::(&row, "public_id_column")?, + descriptor_public_id: optional_column::(&row, "descriptor_public_id_column")?, + tensor_ref: required_column::(&row, "tensor_ref_column")?, + tensor_rows: required_column::(&row, "tensor_rows_column")?, + tensor_dimension: required_column::(&row, "tensor_dim_column")?, + tensor_dtype: required_column::(&row, "tensor_dtype_column")?, + tensor_checksum: required_column::(&row, "tensor_checksum_column")?, + }; + let columns = match storage { + ExternalTensorStorage::SameHeap => Some( + ExternalTensorColumns { + model_contract: resolve_attnum(heap_oid, &column_names.model_contract)?, + public_id: resolve_attnum(heap_oid, &column_names.public_id)?, + tensor_ref: resolve_attnum(heap_oid, &column_names.tensor_ref)?, + tensor_rows: resolve_attnum(heap_oid, &column_names.tensor_rows)?, + tensor_dimension: resolve_attnum(heap_oid, &column_names.tensor_dimension)?, + tensor_dtype: resolve_attnum(heap_oid, &column_names.tensor_dtype)?, + tensor_checksum: resolve_attnum(heap_oid, &column_names.tensor_checksum)?, + } + .validate()?, + ), + ExternalTensorStorage::DescriptorRelation => { + if column_names.descriptor_public_id.is_none() { + return Err(RerankError::Registry( + "registered descriptor relation has no public ID column".into(), + )); + } + None + } + }; + validate_model_contract(&model_contract_id)?; + Ok(ExternalTensorSourceBinding { + index_oid, + heap_oid, + descriptor_oid, + storage, + model_contract_id, + columns, + column_names, + }) + }) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry(format!("registry resolver returned NULL {name}"))) +} + +fn optional_column( + row: &pgrx::spi::SpiTupleTable<'_>, + name: &str, +) -> Result, RerankError> +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name).map_err(registry_error) +} + +fn resolve_attnum(heap_oid: pgrx::pg_sys::Oid, column_name: &str) -> Result { + let column_name = CString::new(column_name) + .map_err(|_| RerankError::Registry("registry column contains a NUL byte".into()))?; + let attnum = unsafe { pgrx::pg_sys::get_attnum(heap_oid, column_name.as_ptr()) }; + if attnum <= 0 { + return Err(RerankError::Registry( + "registered descriptor column disappeared during resolution".into(), + )); + } + Ok(attnum) +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +impl<'a, F> HeapTensorRefSource<'a, F> { + pub(super) fn new( + fetcher: &'a mut F, + model_contract_id: String, + columns: ExternalTensorColumns, + ) -> Result { + validate_model_contract(&model_contract_id)?; + Ok(Self { + fetcher, + model_contract_id, + columns: columns.validate()?, + }) + } +} + +impl CandidateTensorDescriptorSource for HeapTensorRefSource<'_, F> { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + let Some(mut tuple) = self.fetcher.fetch(candidate.heap_key) else { + return Ok(None); + }; + + // This is the data-exfiltration boundary: no descriptor value may be + // read before PostgreSQL has accepted the active base-relation qual. + if !tuple.filter() { + return Ok(None); + } + + let model_contract = read_attribute::(&mut tuple, self.columns.model_contract)?; + if model_contract != self.model_contract_id { + return Err(RerankError::ModelContractMismatch); + } + let public_id = read_attribute::(&mut tuple, self.columns.public_id)?; + let tensor_ref = read_attribute::(&mut tuple, self.columns.tensor_ref)?; + let rows = read_attribute::(&mut tuple, self.columns.tensor_rows)?; + let dimension = read_attribute::(&mut tuple, self.columns.tensor_dimension)?; + let dtype = read_attribute::(&mut tuple, self.columns.tensor_dtype)?; + let checksum = read_attribute::(&mut tuple, self.columns.tensor_checksum)?; + + validate_descriptor( + candidate, public_id, tensor_ref, rows, dimension, dtype, checksum, + ) + .map(Some) + } +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +fn read_attribute(tuple: &mut impl Tuple, attnum: i16) -> Result { + let attribute = tuple + .attribute(attnum) + .ok_or(RerankError::InvalidDescriptor( + "registered attribute is unavailable", + ))?; + unsafe { T::from_datum(attribute.datum, attribute.is_null) }.ok_or( + RerankError::InvalidDescriptor("registered descriptor value is NULL or malformed"), + ) +} + +fn validate_model_contract(value: &str) -> Result<(), RerankError> { + if value.is_empty() || value.len() > MAX_MODEL_CONTRACT_BYTES || contains_control(value) { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn validate_descriptor( + candidate: PageCandidate, + public_id: i64, + tensor_ref: String, + rows: i32, + dimension: i32, + dtype: String, + checksum: String, +) -> Result { + if tensor_ref.is_empty() + || tensor_ref.len() > MAX_TENSOR_REF_BYTES + || contains_control(&tensor_ref) + { + return Err(RerankError::InvalidDescriptor( + "tensor reference is empty, oversized, or contains control characters", + )); + } + let rows = u32::try_from(rows) + .ok() + .filter(|rows| (1..=MAX_TENSOR_ROWS).contains(rows)) + .ok_or(RerankError::InvalidDescriptor( + "tensor row count is invalid", + ))?; + let dimension = u32::try_from(dimension) + .ok() + .filter(|dimension| (1..=MAX_TENSOR_DIMENSION).contains(dimension)) + .ok_or(RerankError::InvalidDescriptor( + "tensor dimension is invalid", + ))?; + let dtype = match dtype.as_str() { + "float32" => ExternalTensorDtype::F32, + "float16" => ExternalTensorDtype::F16, + _ => { + return Err(RerankError::InvalidDescriptor( + "tensor dtype must be float16 or float32", + )); + } + }; + if checksum.len() > MAX_CHECKSUM_BYTES || !is_sha256_checksum(&checksum) { + return Err(RerankError::InvalidDescriptor( + "tensor checksum must be a lowercase sha256 digest", + )); + } + Ok(ExternalTensorDescriptor { + candidate, + public_id, + tensor_ref, + rows, + dimension, + dtype, + checksum, + }) +} + +fn contains_control(value: &str) -> bool { + value.chars().any(char::is_control) +} + +fn is_sha256_checksum(value: &str) -> bool { + value.strip_prefix("sha256:").is_some_and(|digest| { + digest.len() == 64 + && digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::index::fetcher::TupleAttribute; + use distance::Distance; + + fn candidate() -> PageCandidate { + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + } + } + + fn columns() -> ExternalTensorColumns { + ExternalTensorColumns { + model_contract: 1, + public_id: 2, + tensor_ref: 3, + tensor_rows: 4, + tensor_dimension: 5, + tensor_dtype: 6, + tensor_checksum: 7, + } + } + + struct RejectingFetcher; + struct RejectingTuple; + + impl Tuple for RejectingTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + panic!("external descriptor sources do not build index expressions") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + panic!("a rejected tuple must not expose descriptor attributes") + } + } + + impl FilterableTuple for RejectingTuple { + fn filter(&mut self) -> bool { + false + } + } + + impl Fetcher for RejectingFetcher { + type Tuple<'a> = RejectingTuple; + + fn fetch(&mut self, _key: [u16; 3]) -> Option> { + Some(RejectingTuple) + } + } + + #[test] + fn source_filters_before_reading_descriptor_attributes() { + let mut fetcher = RejectingFetcher; + let mut source = + HeapTensorRefSource::new(&mut fetcher, "contract@1".into(), columns()).unwrap(); + assert!(source.fetch(candidate()).unwrap().is_none()); + } + + #[test] + fn descriptor_validation_accepts_bounded_immutable_metadata() { + let descriptor = validate_descriptor( + candidate(), + 42, + "s3://immutable-bucket/page-42.tensor".into(), + 747, + 320, + "float16".into(), + format!("sha256:{}", "a".repeat(64)), + ) + .unwrap(); + assert_eq!(descriptor.public_id, 42); + assert_eq!(descriptor.rows, 747); + assert_eq!(descriptor.dimension, 320); + assert_eq!(descriptor.dtype, ExternalTensorDtype::F16); + } + + #[test] + fn descriptor_validation_rejects_ambiguous_or_unbounded_metadata() { + assert!(matches!( + columns_with_duplicate().validate(), + Err(RerankError::InvalidDescriptor(_)) + )); + for (rows, dimension, dtype, checksum) in [ + (0, 320, "float16", format!("sha256:{}", "a".repeat(64))), + (1, 0, "float16", format!("sha256:{}", "a".repeat(64))), + (1, 320, "bf16", format!("sha256:{}", "a".repeat(64))), + (1, 320, "float16", "sha256:short".into()), + ] { + assert!(matches!( + validate_descriptor( + candidate(), + 1, + "s3://immutable/tensor".into(), + rows, + dimension, + dtype.into(), + checksum, + ), + Err(RerankError::InvalidDescriptor(_)) + )); + } + } + + fn columns_with_duplicate() -> ExternalTensorColumns { + ExternalTensorColumns { + public_id: 1, + ..columns() + } + } +} diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs new file mode 100644 index 00000000..af8c5455 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -0,0 +1,1495 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorDtype, +}; +use super::rerank::{CandidateTensorSource, ExactMaxsimBackend, RerankError, RerankResults}; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::mem::size_of_val; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; +use vchordrq::types::OwnedVector; + +const MAGIC: &[u8; 4] = b"VCTM"; +const VERSION: u16 = 1; +const EXTERNAL_VERSION: u16 = 2; +const REQUEST_KIND: u16 = 1; +const RESPONSE_KIND: u16 = 2; +const HEADER_LEN: usize = 24; +const MAX_REMOTE_ERROR_BYTES: usize = 64 * 1024; + +static NEXT_REQUEST_ID: AtomicU64 = AtomicU64::new(1); +static LAST_FALLBACK_WARNING_SECONDS: AtomicU64 = AtomicU64::new(0); + +pub(super) fn report_gpu_fallback(error: &RerankError) { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |duration| duration.as_secs()); + let should_warn = LAST_FALLBACK_WARNING_SECONDS + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |last| { + (last == 0 || now.saturating_sub(last) >= 60).then_some(now) + }) + .is_ok(); + if should_warn { + pgrx::warning!("GPU MaxSim failed; using cpu_exact: {error}"); + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum TensorDtype { + F32 = 1, + F16 = 2, +} + +pub(super) trait TileMaxsimTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + max_response_bytes: usize, + ) -> Result, RerankError>; +} + +pub(super) struct GpuTileMaxsimBackend { + transport: T, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, +} + +impl GpuTileMaxsimBackend { + pub fn new( + transport: T, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, + ) -> Self { + Self { + transport, + timeout, + max_batch_tokens, + max_batch_bytes, + } + } +} + +impl ExactMaxsimBackend for GpuTileMaxsimBackend { + type Results = RerankResults; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let encoded = encode_request( + request_id, + query, + candidates, + source, + self.max_batch_tokens, + self.max_batch_bytes, + )?; + if encoded.heap_keys.is_empty() { + return Ok(RerankResults { + inner: BinaryHeap::new(), + }); + } + let max_response_bytes = HEADER_LEN + .checked_add(8) + .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) + .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) + .ok_or(RerankError::RequestTooLarge)?; + let response = + self.transport + .round_trip(&encoded.frame, self.timeout, max_response_bytes)?; + decode_response(&response, request_id, &encoded.heap_keys) + } +} + +/// GPU backend for Phase 3B external tensor descriptors. +/// +/// This codec is intentionally separate from [`ExactMaxsimBackend`]: an +/// external full tensor may have different logical values from the indexed +/// sketch, so it cannot be substituted into ordinary `@#` execution. +pub(super) struct GpuExternalTileMaxsimBackend { + transport: T, + model_contract_id: String, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, +} + +impl GpuExternalTileMaxsimBackend { + pub(super) fn new( + transport: T, + model_contract_id: String, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, + ) -> Self { + Self { + transport, + model_contract_id, + timeout, + max_batch_tokens, + max_batch_bytes, + } + } +} + +impl GpuExternalTileMaxsimBackend { + pub(super) fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let encoded = encode_external_request( + request_id, + &self.model_contract_id, + query, + candidates, + source, + self.max_batch_tokens, + self.max_batch_bytes, + )?; + if encoded.heap_keys.is_empty() { + return Ok(RerankResults { + inner: BinaryHeap::new(), + }); + } + let max_response_bytes = HEADER_LEN + .checked_add(8) + .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) + .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) + .ok_or(RerankError::RequestTooLarge)?; + let response = + self.transport + .round_trip(&encoded.frame, self.timeout, max_response_bytes)?; + decode_response_for_version(&response, EXTERNAL_VERSION, request_id, &encoded.heap_keys) + } +} + +struct EncodedRequest { + frame: Vec, + heap_keys: Vec, +} + +fn encode_external_request( + request_id: u64, + model_contract_id: &str, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, +) -> Result { + const MAX_MODEL_CONTRACT_BYTES: usize = 512; + const MAX_TENSOR_REF_BYTES: usize = 4096; + const MAX_CHECKSUM_BYTES: usize = 512; + + if model_contract_id.is_empty() + || model_contract_id.len() > MAX_MODEL_CONTRACT_BYTES + || model_contract_id.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + let (dtype, dimension) = tensor_metadata(query)?; + let external_dtype = match dtype { + TensorDtype::F32 => ExternalTensorDtype::F32, + TensorDtype::F16 => ExternalTensorDtype::F16, + }; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let mut total_tokens = query.len(); + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + let mut declared_tensor_bytes = tensor_bytes(query_rows, dimension, dtype)?; + if declared_tensor_bytes > max_batch_bytes { + return Err(RerankError::RequestTooLarge); + } + + let mut writer = BoundedWriter::new(max_batch_bytes); + writer.zeros(HEADER_LEN)?; + writer.u32(dimension)?; + writer.u32(query_rows)?; + let candidate_count_offset = writer.len(); + writer.u32(0)?; + writer.u8(dtype as u8)?; + writer.u8(1)?; // sum_query_max_document_dot + writer.u16(0)?; + writer + .u32(u32::try_from(model_contract_id.len()).map_err(|_| RerankError::RequestTooLarge)?)?; + writer.bytes(model_contract_id.as_bytes())?; + encode_tensor_values(&mut writer, query, dtype)?; + + let mut heap_keys = Vec::new(); + for candidate in candidates { + let Some(descriptor) = source.fetch(candidate)? else { + continue; + }; + validate_external_for_request( + &descriptor, + dimension, + external_dtype, + MAX_TENSOR_REF_BYTES, + MAX_CHECKSUM_BYTES, + )?; + total_tokens = total_tokens + .checked_add(descriptor.rows as usize) + .ok_or(RerankError::RequestTooLarge)?; + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + declared_tensor_bytes = declared_tensor_bytes + .checked_add(tensor_bytes(descriptor.rows, dimension, dtype)?) + .ok_or(RerankError::RequestTooLarge)?; + if declared_tensor_bytes > max_batch_bytes { + return Err(RerankError::RequestTooLarge); + } + + let candidate_id = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.u32(candidate_id)?; + writer.u32(descriptor.rows)?; + writer.u32( + u32::try_from(descriptor.tensor_ref.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.u32( + u32::try_from(descriptor.checksum.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.bytes(descriptor.tensor_ref.as_bytes())?; + writer.bytes(descriptor.checksum.as_bytes())?; + heap_keys.push(descriptor.candidate.heap_key); + } + + let candidate_count = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_u32(candidate_count_offset, candidate_count); + let body_len = writer + .len() + .checked_sub(HEADER_LEN) + .ok_or_else(|| RerankError::Protocol("invalid request length".into()))?; + writer.patch_bytes(0, MAGIC); + writer.patch_u16(4, EXTERNAL_VERSION); + writer.patch_u16(6, REQUEST_KIND); + writer.patch_u64(8, request_id); + writer.patch_u64( + 16, + u64::try_from(body_len).map_err(|_| RerankError::RequestTooLarge)?, + ); + Ok(EncodedRequest { + frame: writer.finish(), + heap_keys, + }) +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: TensorDtype) -> Result { + let scalar_bytes = match dtype { + TensorDtype::F32 => 4usize, + TensorDtype::F16 => 2usize, + }; + usize::try_from(rows) + .ok() + .and_then(|rows| rows.checked_mul(dimension as usize)) + .and_then(|elements| elements.checked_mul(scalar_bytes)) + .ok_or(RerankError::RequestTooLarge) +} + +fn validate_external_for_request( + descriptor: &ExternalTensorDescriptor, + dimension: u32, + dtype: ExternalTensorDtype, + max_tensor_ref_bytes: usize, + max_checksum_bytes: usize, +) -> Result<(), RerankError> { + if descriptor.rows == 0 + || descriptor.rows > 65_536 + || dimension > 60_000 + || descriptor.dimension != dimension + || descriptor.dtype != dtype + { + return Err(RerankError::TensorMismatch); + } + if descriptor.tensor_ref.is_empty() + || descriptor.tensor_ref.len() > max_tensor_ref_bytes + || descriptor.tensor_ref.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "tensor reference is empty, oversized, or contains control characters", + )); + } + if descriptor.checksum.len() > max_checksum_bytes + || !descriptor + .checksum + .strip_prefix("sha256:") + .is_some_and(|digest| { + digest.len() == 64 + && digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + }) + { + return Err(RerankError::InvalidDescriptor( + "tensor checksum must be a lowercase sha256 digest", + )); + } + Ok(()) +} + +fn encode_request( + request_id: u64, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, +) -> Result { + let (dtype, dimension) = tensor_metadata(query)?; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let mut total_tokens = query.len(); + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + + let mut writer = BoundedWriter::new(max_batch_bytes); + writer.zeros(HEADER_LEN)?; + writer.u32(dimension)?; + writer.u32(query_rows)?; + let candidate_count_offset = writer.len(); + writer.u32(0)?; + writer.u8(dtype as u8)?; + writer.u8(1)?; // sum_query_max_document_dot + writer.u16(0)?; + encode_tensor_values(&mut writer, query, dtype)?; + + let mut heap_keys = Vec::new(); + for candidate in candidates { + let Some(tensor) = source.fetch(candidate)? else { + continue; + }; + let (candidate_dtype, candidate_dimension) = tensor_metadata(&tensor.vectors)?; + if candidate_dtype != dtype || candidate_dimension != dimension { + return Err(RerankError::TensorMismatch); + } + total_tokens = total_tokens + .checked_add(tensor.vectors.len()) + .ok_or(RerankError::RequestTooLarge)?; + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + let candidate_id = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + let rows = u32::try_from(tensor.vectors.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.u32(candidate_id)?; + writer.u32(rows)?; + encode_tensor_values(&mut writer, &tensor.vectors, dtype)?; + heap_keys.push(tensor.candidate.heap_key); + } + let candidate_count = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_u32(candidate_count_offset, candidate_count); + let body_len = writer + .len() + .checked_sub(HEADER_LEN) + .ok_or_else(|| RerankError::Protocol("invalid request length".into()))?; + let body_len = u64::try_from(body_len).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_bytes(0, MAGIC); + writer.patch_u16(4, VERSION); + writer.patch_u16(6, REQUEST_KIND); + writer.patch_u64(8, request_id); + writer.patch_u64(16, body_len); + Ok(EncodedRequest { + frame: writer.finish(), + heap_keys, + }) +} + +fn tensor_metadata(vectors: &[OwnedVector]) -> Result<(TensorDtype, u32), RerankError> { + let Some(first) = vectors.first() else { + return Err(RerankError::TensorMismatch); + }; + let dtype = match first { + OwnedVector::Vecf32(_) => TensorDtype::F32, + OwnedVector::Vecf16(_) => TensorDtype::F16, + OwnedVector::Rabitq8(_) | OwnedVector::Rabitq4(_) => { + return Err(RerankError::UnsupportedTensorKind); + } + }; + let dimension = first.dim(); + for vector in vectors { + let this_dtype = match vector { + OwnedVector::Vecf32(_) => TensorDtype::F32, + OwnedVector::Vecf16(_) => TensorDtype::F16, + OwnedVector::Rabitq8(_) | OwnedVector::Rabitq4(_) => { + return Err(RerankError::UnsupportedTensorKind); + } + }; + if this_dtype != dtype || vector.dim() != dimension { + return Err(RerankError::TensorMismatch); + } + } + Ok((dtype, dimension)) +} + +fn encode_tensor_values( + writer: &mut BoundedWriter, + vectors: &[OwnedVector], + dtype: TensorDtype, +) -> Result<(), RerankError> { + for vector in vectors { + match (dtype, vector) { + (TensorDtype::F32, OwnedVector::Vecf32(vector)) => { + for value in vector.slice() { + writer.bytes(&value.to_le_bytes())?; + } + } + (TensorDtype::F16, OwnedVector::Vecf16(vector)) => { + for value in vector.slice() { + writer.u16(value.to_bits())?; + } + } + _ => return Err(RerankError::TensorMismatch), + } + } + Ok(()) +} + +fn decode_response( + frame: &[u8], + request_id: u64, + heap_keys: &[HeapKey], +) -> Result { + decode_response_for_version(frame, VERSION, request_id, heap_keys) +} + +fn decode_response_for_version( + frame: &[u8], + expected_version: u16, + request_id: u64, + heap_keys: &[HeapKey], +) -> Result { + let mut cursor = Cursor::new(frame); + if cursor.bytes(4)? != MAGIC { + return Err(RerankError::Protocol("invalid magic".into())); + } + if cursor.u16()? != expected_version { + return Err(RerankError::Protocol("unsupported version".into())); + } + if cursor.u16()? != RESPONSE_KIND { + return Err(RerankError::Protocol("unexpected message kind".into())); + } + if cursor.u64()? != request_id { + return Err(RerankError::Protocol("request ID mismatch".into())); + } + let body_len = usize::try_from(cursor.u64()?) + .map_err(|_| RerankError::Protocol("response body is too large".into()))?; + if body_len != frame.len().saturating_sub(HEADER_LEN) { + return Err(RerankError::Protocol("response length mismatch".into())); + } + let status = cursor.u32()?; + if status != 0 { + let length = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("remote error is too large".into()))?; + if length > MAX_REMOTE_ERROR_BYTES { + return Err(RerankError::Protocol("remote error is too large".into())); + } + let message = std::str::from_utf8(cursor.bytes(length)?) + .map_err(|_| RerankError::Protocol("remote error is not UTF-8".into()))?; + cursor.finish()?; + return Err(RerankError::Remote(message.into())); + } + let result_count = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("result count is too large".into()))?; + if result_count != heap_keys.len() { + return Err(RerankError::Protocol("partial result set".into())); + } + let mut seen = vec![false; heap_keys.len()]; + let mut results = BinaryHeap::new(); + for _ in 0..result_count { + let candidate_id = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("candidate ID is too large".into()))?; + let Some(heap_key) = heap_keys.get(candidate_id).copied() else { + return Err(RerankError::Protocol("unknown candidate ID".into())); + }; + if std::mem::replace(&mut seen[candidate_id], true) { + return Err(RerankError::Protocol("duplicate candidate ID".into())); + } + let similarity = f32::from_bits(cursor.u32()?); + if !similarity.is_finite() { + return Err(RerankError::Protocol("non-finite similarity".into())); + } + let distance = Distance::from_f32(-similarity); + results.push((Reverse(distance), Reverse(heap_key))); + } + cursor.finish()?; + if seen.iter().any(|seen| !seen) { + return Err(RerankError::Protocol("partial result set".into())); + } + Ok(RerankResults { inner: results }) +} + +struct BoundedWriter { + bytes: Vec, + limit: usize, +} + +impl BoundedWriter { + fn new(limit: usize) -> Self { + Self { + bytes: Vec::new(), + limit, + } + } + + fn len(&self) -> usize { + self.bytes.len() + } + + fn ensure(&self, additional: usize) -> Result<(), RerankError> { + let size = self + .bytes + .len() + .checked_add(additional) + .ok_or(RerankError::RequestTooLarge)?; + if size > self.limit { + return Err(RerankError::RequestTooLarge); + } + Ok(()) + } + + fn zeros(&mut self, count: usize) -> Result<(), RerankError> { + self.ensure(count)?; + self.bytes.resize(self.bytes.len() + count, 0); + Ok(()) + } + + fn bytes(&mut self, bytes: &[u8]) -> Result<(), RerankError> { + self.ensure(bytes.len())?; + self.bytes.extend_from_slice(bytes); + Ok(()) + } + + fn u8(&mut self, value: u8) -> Result<(), RerankError> { + self.bytes(&[value]) + } + + fn u16(&mut self, value: u16) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + + fn u32(&mut self, value: u32) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + + fn patch_bytes(&mut self, offset: usize, bytes: &[u8]) { + self.bytes[offset..offset + bytes.len()].copy_from_slice(bytes); + } + + fn patch_u16(&mut self, offset: usize, value: u16) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn patch_u32(&mut self, offset: usize, value: u32) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn patch_u64(&mut self, offset: usize, value: u64) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn finish(self) -> Vec { + self.bytes + } +} + +struct Cursor<'a> { + bytes: &'a [u8], + offset: usize, +} + +impl<'a> Cursor<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, offset: 0 } + } + + fn bytes(&mut self, count: usize) -> Result<&'a [u8], RerankError> { + let end = self + .offset + .checked_add(count) + .ok_or_else(|| RerankError::Protocol("message offset overflow".into()))?; + let bytes = self + .bytes + .get(self.offset..end) + .ok_or_else(|| RerankError::Protocol("truncated message".into()))?; + self.offset = end; + Ok(bytes) + } + + fn u16(&mut self) -> Result { + Ok(u16::from_le_bytes(self.bytes(2)?.try_into().unwrap())) + } + + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.bytes(4)?.try_into().unwrap())) + } + + fn u64(&mut self) -> Result { + Ok(u64::from_le_bytes(self.bytes(8)?.try_into().unwrap())) + } + + fn finish(self) -> Result<(), RerankError> { + if self.offset != self.bytes.len() { + return Err(RerankError::Protocol("trailing response bytes".into())); + } + Ok(()) + } +} + +pub(super) struct UnixSocketTransport { + endpoint: String, +} + +impl UnixSocketTransport { + pub fn new(endpoint: String) -> Self { + Self { endpoint } + } +} + +#[cfg(unix)] +impl TileMaxsimTransport for UnixSocketTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + max_response_bytes: usize, + ) -> Result, RerankError> { + use std::time::Instant; + + if self.endpoint.is_empty() { + return Err(RerankError::Transport("endpoint is empty".into())); + } + let deadline = Instant::now() + .checked_add(timeout) + .ok_or_else(|| RerankError::Transport("invalid timeout".into()))?; + let mut stream = connect_interruptible(&self.endpoint, deadline)?; + let poll = remaining_until(deadline)?.min(Duration::from_millis(50)); + stream + .set_read_timeout(Some(poll)) + .map_err(|error| RerankError::Transport(error.to_string()))?; + stream + .set_write_timeout(Some(poll)) + .map_err(|error| RerankError::Transport(error.to_string()))?; + + write_interruptible(&mut stream, request, deadline)?; + let mut header = [0u8; HEADER_LEN]; + read_interruptible(&mut stream, &mut header, deadline)?; + let body_len = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) + .map_err(|_| RerankError::Protocol("response body is too large".into()))?; + let response_len = HEADER_LEN + .checked_add(body_len) + .ok_or_else(|| RerankError::Protocol("response length overflow".into()))?; + if response_len > max_response_bytes { + return Err(RerankError::Protocol( + "response exceeds configured limit".into(), + )); + } + let mut response = Vec::with_capacity(response_len); + response.extend_from_slice(&header); + response.resize(response_len, 0); + read_interruptible(&mut stream, &mut response[HEADER_LEN..], deadline)?; + Ok(response) + } +} + +#[cfg(unix)] +fn connect_interruptible( + endpoint: &str, + deadline: std::time::Instant, +) -> Result { + use std::os::fd::{AsRawFd, FromRawFd, OwnedFd}; + + let (address, address_len) = unix_socket_address(endpoint)?; + let raw_fd = unsafe { libc::socket(libc::AF_UNIX, libc::SOCK_STREAM, 0) }; + if raw_fd < 0 { + return Err(last_transport_error()); + } + let fd = unsafe { OwnedFd::from_raw_fd(raw_fd) }; + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFD, + libc::F_SETFD, + libc::FD_CLOEXEC, + true, + )?; + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFL, + libc::F_SETFL, + libc::O_NONBLOCK, + true, + )?; + + let connected = unsafe { + libc::connect( + fd.as_raw_fd(), + (&raw const address).cast::(), + address_len, + ) + } == 0; + if !connected { + let error = std::io::Error::last_os_error(); + let raw_error = error.raw_os_error(); + if raw_error != Some(libc::EINPROGRESS) + && raw_error != Some(libc::EAGAIN) + && raw_error != Some(libc::EWOULDBLOCK) + { + return Err(RerankError::Transport(error.to_string())); + } + wait_for_connect(fd.as_raw_fd(), deadline)?; + } + + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFL, + libc::F_SETFL, + libc::O_NONBLOCK, + false, + )?; + Ok(std::os::unix::net::UnixStream::from(fd)) +} + +#[cfg(unix)] +fn unix_socket_address( + endpoint: &str, +) -> Result<(libc::sockaddr_un, libc::socklen_t), RerankError> { + let path = endpoint.as_bytes(); + let mut address = unsafe { std::mem::zeroed::() }; + if path.contains(&0) { + return Err(RerankError::Transport( + "endpoint contains a NUL byte".into(), + )); + } + if path.len() >= address.sun_path.len() { + return Err(RerankError::Transport("endpoint path is too long".into())); + } + address.sun_family = libc::AF_UNIX as libc::sa_family_t; + unsafe { + std::ptr::copy_nonoverlapping( + path.as_ptr().cast::(), + address.sun_path.as_mut_ptr(), + path.len(), + ); + } + let length = std::mem::offset_of!(libc::sockaddr_un, sun_path) + .checked_add(path.len()) + .and_then(|length| length.checked_add(1)) + .and_then(|length| libc::socklen_t::try_from(length).ok()) + .ok_or_else(|| RerankError::Transport("endpoint path is too long".into()))?; + #[cfg(any( + target_os = "aix", + target_os = "dragonfly", + target_os = "freebsd", + target_os = "haiku", + target_os = "hurd", + target_os = "ios", + target_os = "macos", + target_os = "netbsd", + target_os = "openbsd", + target_os = "tvos", + target_os = "visionos", + target_os = "watchos" + ))] + { + address.sun_len = u8::try_from(length) + .map_err(|_| RerankError::Transport("endpoint path is too long".into()))?; + } + Ok((address, length)) +} + +#[cfg(unix)] +fn update_fd_flag( + fd: std::os::fd::RawFd, + get_command: libc::c_int, + set_command: libc::c_int, + flag: libc::c_int, + enabled: bool, +) -> Result<(), RerankError> { + let current = unsafe { libc::fcntl(fd, get_command) }; + if current < 0 { + return Err(last_transport_error()); + } + let updated = if enabled { + current | flag + } else { + current & !flag + }; + if unsafe { libc::fcntl(fd, set_command, updated) } < 0 { + return Err(last_transport_error()); + } + Ok(()) +} + +#[cfg(unix)] +fn wait_for_connect( + fd: std::os::fd::RawFd, + deadline: std::time::Instant, +) -> Result<(), RerankError> { + loop { + pgrx::check_for_interrupts!(); + let remaining = remaining_until(deadline)?; + let timeout_ms = remaining.min(Duration::from_millis(50)).as_millis().max(1) as libc::c_int; + let mut poll_fd = libc::pollfd { + fd, + events: libc::POLLOUT, + revents: 0, + }; + let result = unsafe { libc::poll(&mut poll_fd, 1, timeout_ms) }; + if result == 0 { + continue; + } + if result < 0 { + let error = std::io::Error::last_os_error(); + if error.kind() == std::io::ErrorKind::Interrupted { + continue; + } + return Err(RerankError::Transport(error.to_string())); + } + let mut socket_error = 0; + let mut socket_error_len = size_of_val(&socket_error) as libc::socklen_t; + if unsafe { + libc::getsockopt( + fd, + libc::SOL_SOCKET, + libc::SO_ERROR, + (&raw mut socket_error).cast(), + &raw mut socket_error_len, + ) + } < 0 + { + return Err(last_transport_error()); + } + if socket_error != 0 { + return Err(RerankError::Transport( + std::io::Error::from_raw_os_error(socket_error).to_string(), + )); + } + return Ok(()); + } +} + +#[cfg(unix)] +fn remaining_until(deadline: std::time::Instant) -> Result { + deadline + .checked_duration_since(std::time::Instant::now()) + .filter(|remaining| !remaining.is_zero()) + .ok_or_else(|| RerankError::Transport("request timed out".into())) +} + +#[cfg(unix)] +fn last_transport_error() -> RerankError { + RerankError::Transport(std::io::Error::last_os_error().to_string()) +} + +#[cfg(unix)] +fn write_interruptible( + stream: &mut std::os::unix::net::UnixStream, + mut bytes: &[u8], + deadline: std::time::Instant, +) -> Result<(), RerankError> { + use std::io::Write; + + while !bytes.is_empty() { + match stream.write(bytes) { + Ok(0) => return Err(RerankError::Transport("connection closed".into())), + Ok(count) => bytes = &bytes[count..], + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::Interrupted + | std::io::ErrorKind::WouldBlock + | std::io::ErrorKind::TimedOut + ) => {} + Err(error) => return Err(RerankError::Transport(error.to_string())), + } + pgrx::check_for_interrupts!(); + if std::time::Instant::now() >= deadline { + return Err(RerankError::Transport("request timed out".into())); + } + } + Ok(()) +} + +#[cfg(unix)] +fn read_interruptible( + stream: &mut std::os::unix::net::UnixStream, + mut bytes: &mut [u8], + deadline: std::time::Instant, +) -> Result<(), RerankError> { + use std::io::Read; + + while !bytes.is_empty() { + match stream.read(bytes) { + Ok(0) => return Err(RerankError::Transport("connection closed".into())), + Ok(count) => bytes = &mut bytes[count..], + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::Interrupted + | std::io::ErrorKind::WouldBlock + | std::io::ErrorKind::TimedOut + ) => {} + Err(error) => return Err(RerankError::Transport(error.to_string())), + } + pgrx::check_for_interrupts!(); + if std::time::Instant::now() >= deadline { + return Err(RerankError::Transport("request timed out".into())); + } + } + Ok(()) +} + +#[cfg(not(unix))] +impl TileMaxsimTransport for UnixSocketTransport { + fn round_trip( + &mut self, + _request: &[u8], + _timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + if self.endpoint.is_empty() { + return Err(RerankError::Transport("endpoint is empty".into())); + } + Err(RerankError::Transport( + "Unix sockets are not supported on this platform".into(), + )) + } +} + +#[cfg(test)] +mod tests { + use super::super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorDtype, + }; + use super::super::rerank::CandidateTensor; + use super::*; + use std::collections::BTreeMap; + use vector::vect::VectOwned; + + struct MockTensorSource(BTreeMap>); + + impl CandidateTensorSource for MockTensorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some(CandidateTensor { + candidate, + vectors: self + .0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + })) + } + } + + struct MockDescriptorSource(BTreeMap); + + impl CandidateTensorDescriptorSource for MockDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some( + self.0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + )) + } + } + + struct MockTransport { + similarities: Vec<(u32, f32)>, + } + + impl TileMaxsimTransport for MockTransport { + fn round_trip( + &mut self, + request: &[u8], + _timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + let request_id = u64::from_le_bytes(request[8..16].try_into().unwrap()); + let version = u16::from_le_bytes(request[4..6].try_into().unwrap()); + Ok(success_response_with_version( + version, + request_id, + &self.similarities, + )) + } + } + + fn vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf32(VectOwned::new(values.to_vec())) + } + + fn half_vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf16(VectOwned::new( + values.iter().copied().map(simd::f16::from_f32).collect(), + )) + } + + fn success_response(request_id: u64, similarities: &[(u32, f32)]) -> Vec { + success_response_with_version(VERSION, request_id, similarities) + } + + fn success_response_with_version( + version: u16, + request_id: u64, + similarities: &[(u32, f32)], + ) -> Vec { + let body_len = 8 + similarities.len() * 8; + let mut response = Vec::with_capacity(HEADER_LEN + body_len); + response.extend_from_slice(MAGIC); + response.extend_from_slice(&version.to_le_bytes()); + response.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + response.extend_from_slice(&request_id.to_le_bytes()); + response.extend_from_slice(&(body_len as u64).to_le_bytes()); + response.extend_from_slice(&0u32.to_le_bytes()); + response.extend_from_slice(&(similarities.len() as u32).to_le_bytes()); + for (candidate_id, similarity) in similarities { + response.extend_from_slice(&candidate_id.to_le_bytes()); + response.extend_from_slice(&similarity.to_bits().to_le_bytes()); + } + response + } + + fn external_descriptor( + candidate: PageCandidate, + public_id: i64, + tensor_ref: &str, + rows: u32, + dimension: u32, + dtype: ExternalTensorDtype, + ) -> ExternalTensorDescriptor { + ExternalTensorDescriptor { + candidate, + public_id, + tensor_ref: tensor_ref.into(), + rows, + dimension, + dtype, + checksum: format!("sha256:{}", "a".repeat(64)), + } + } + + fn error_response(request_id: u64, message: &str) -> Vec { + let body_len = 8 + message.len(); + let mut response = Vec::with_capacity(HEADER_LEN + body_len); + response.extend_from_slice(MAGIC); + response.extend_from_slice(&VERSION.to_le_bytes()); + response.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + response.extend_from_slice(&request_id.to_le_bytes()); + response.extend_from_slice(&(body_len as u64).to_le_bytes()); + response.extend_from_slice(&1u32.to_le_bytes()); + response.extend_from_slice(&(message.len() as u32).to_le_bytes()); + response.extend_from_slice(message.as_bytes()); + response + } + + #[test] + fn gpu_backend_maps_positive_similarity_to_ascending_distance() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page_1, + }, + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page_2, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0])]), + (page_2, vec![vector(&[0.5, 0.0])]), + ])); + let transport = MockTransport { + similarities: vec![(0, 1.0), (1, 2.0)], + }; + let results = GpuTileMaxsimBackend::new(transport, Duration::from_secs(1), 100, 4096) + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results[0].heap_key, page_2); + assert_eq!(results[0].distance.to_f32(), -2.0); + assert_eq!(results[1].heap_key, page_1); + assert_eq!(results[1].distance.to_f32(), -1.0); + } + + #[test] + fn response_rejects_partial_and_duplicate_results() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let partial = success_response(7, &[(0, 1.0)]); + assert!(matches!( + decode_response(&partial, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let duplicate = success_response(7, &[(0, 1.0), (0, 2.0)]); + assert!(matches!( + decode_response(&duplicate, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + + #[test] + fn response_ids_may_arrive_out_of_order() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let response = success_response(7, &[(1, 0.5), (0, 1.0)]); + let results = decode_response(&response, 7, &keys) + .unwrap() + .collect::>(); + + assert_eq!(results[0].heap_key, keys[0]); + assert_eq!(results[0].distance.to_f32(), -1.0); + assert_eq!(results[1].heap_key, keys[1]); + assert_eq!(results[1].distance.to_f32(), -0.5); + } + + #[test] + fn response_rejects_unknown_non_finite_and_trailing_results() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let unknown = success_response(7, &[(0, 1.0), (2, 2.0)]); + assert!(matches!( + decode_response(&unknown, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let non_finite = success_response(7, &[(0, f32::NAN), (1, 2.0)]); + assert!(matches!( + decode_response(&non_finite, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let mut trailing = success_response(7, &[(0, 1.0), (1, 2.0)]); + trailing.push(0); + assert!(matches!( + decode_response(&trailing, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + + #[test] + fn response_rejects_invalid_header_fields() { + let keys = [[0, 0, 1]]; + let valid = success_response(7, &[(0, 1.0)]); + + for (offset, replacement) in [(0, 0u8), (4, 2), (6, 1), (8, 8), (16, 0)] { + let mut invalid = valid.clone(); + invalid[offset] = replacement; + assert!(matches!( + decode_response(&invalid, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + } + + #[test] + fn response_surfaces_remote_error() { + let response = error_response(7, "CUDA queue is unavailable"); + assert!(matches!( + decode_response(&response, 7, &[]), + Err(RerankError::Remote(message)) if message == "CUDA queue is unavailable" + )); + } + + #[test] + fn request_frame_is_versioned_and_length_prefixed() { + let page = [0, 0, 1]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([(page, vec![vector(&[0.5, 0.0])])])); + let encoded = encode_request(9, &query, &mut candidates, &mut source, 100, 4096).unwrap(); + + assert_eq!(&encoded.frame[0..4], MAGIC); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + VERSION + ); + assert_eq!( + u16::from_le_bytes(encoded.frame[6..8].try_into().unwrap()), + REQUEST_KIND + ); + assert_eq!( + u64::from_le_bytes(encoded.frame[8..16].try_into().unwrap()), + 9 + ); + assert_eq!( + u64::from_le_bytes(encoded.frame[16..24].try_into().unwrap()) as usize, + encoded.frame.len() - HEADER_LEN + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[32..36].try_into().unwrap()), + 1 + ); + assert_eq!(encoded.heap_keys, vec![page]); + } + + #[test] + fn external_request_encodes_contract_and_opaque_descriptor_ids() { + let page = [0, 0, 7]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, -0.5])]; + let tensor_ref = "s3://immutable/page-9001.tensor"; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor(candidate, 9001, tensor_ref, 2, 2, ExternalTensorDtype::F16), + )])); + let contract = "colqwen@immutable-revision"; + let encoded = encode_external_request( + 19, + contract, + &query, + &mut candidates, + &mut source, + 100, + 4096, + ) + .unwrap(); + + assert_eq!(&encoded.frame[0..4], MAGIC); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + EXTERNAL_VERSION + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[32..36].try_into().unwrap()), + 1 + ); + let contract_len = u32::from_le_bytes(encoded.frame[40..44].try_into().unwrap()) as usize; + assert_eq!(&encoded.frame[44..44 + contract_len], contract.as_bytes()); + let candidate_offset = 44 + contract_len + 4; // one 2-D f16 query row + assert_eq!( + u32::from_le_bytes( + encoded.frame[candidate_offset..candidate_offset + 4] + .try_into() + .unwrap() + ), + 0 + ); + assert_eq!( + u32::from_le_bytes( + encoded.frame[candidate_offset + 4..candidate_offset + 8] + .try_into() + .unwrap() + ), + 2 + ); + let reference_len = u32::from_le_bytes( + encoded.frame[candidate_offset + 8..candidate_offset + 12] + .try_into() + .unwrap(), + ) as usize; + let reference_offset = candidate_offset + 16; + assert_eq!( + &encoded.frame[reference_offset..reference_offset + reference_len], + tensor_ref.as_bytes() + ); + assert_eq!(encoded.heap_keys, vec![page]); + } + + #[test] + fn external_backend_maps_scores_without_exposing_public_ids() { + let page = [0, 0, 8]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + i64::MAX, + "object://immutable/tensor", + 4, + 2, + ExternalTensorDtype::F32, + ), + )])); + let transport = MockTransport { + similarities: vec![(0, 3.5)], + }; + let results = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_secs(1), + 100, + 4096, + ) + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].heap_key, page); + assert_eq!(results[0].distance.to_f32(), -3.5); + } + + #[test] + fn external_request_rejects_shape_and_declared_payload_overflow() { + let page = [0, 0, 9]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, 0.0])]; + let make_source = |dimension, rows| { + MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + 9, + "object://immutable/tensor", + rows, + dimension, + ExternalTensorDtype::F16, + ), + )])) + }; + + let mut candidates = vec![candidate].into_iter(); + assert!(matches!( + encode_external_request( + 1, + "contract@1", + &query, + &mut candidates, + &mut make_source(3, 1), + 100, + 4096, + ), + Err(RerankError::TensorMismatch) + )); + + let mut candidates = vec![candidate].into_iter(); + assert!(matches!( + encode_external_request( + 1, + "contract@1", + &query, + &mut candidates, + &mut make_source(2, 100), + 1000, + 64, + ), + Err(RerankError::RequestTooLarge) + )); + } + + #[test] + fn request_frame_encodes_f16_tensor_bits() { + let page = [0, 0, 1]; + let query = vec![half_vector(&[1.0, -0.5])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = + MockTensorSource(BTreeMap::from([(page, vec![half_vector(&[0.25, 2.0])])])); + let encoded = encode_request(9, &query, &mut candidates, &mut source, 100, 4096).unwrap(); + + assert_eq!(encoded.frame[36], TensorDtype::F16 as u8); + assert_eq!( + u16::from_le_bytes(encoded.frame[40..42].try_into().unwrap()), + simd::f16::from_f32(1.0).to_bits() + ); + assert_eq!( + u16::from_le_bytes(encoded.frame[42..44].try_into().unwrap()), + simd::f16::from_f32(-0.5).to_bits() + ); + } + + #[test] + fn request_limits_are_enforced_before_transport() { + let page = [0, 0, 1]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([(page, vec![vector(&[1.0, 0.0])])])); + assert!(matches!( + encode_request(1, &query, &mut candidates, &mut source, 1, 4096), + Err(RerankError::RequestTooLarge) + )); + } + + #[cfg(unix)] + #[test] + fn unix_socket_address_is_length_bounded_and_nul_terminated() { + let (address, length) = unix_socket_address("/tmp/vectorchord.sock").unwrap(); + let path_offset = std::mem::offset_of!(libc::sockaddr_un, sun_path); + + assert_eq!(address.sun_family, libc::AF_UNIX as libc::sa_family_t); + assert_eq!( + length as usize, + path_offset + "/tmp/vectorchord.sock".len() + 1 + ); + assert_eq!( + &address.sun_path[.."/tmp/vectorchord.sock".len()], + "/tmp/vectorchord.sock" + .as_bytes() + .iter() + .map(|byte| *byte as libc::c_char) + .collect::>() + ); + assert_eq!(address.sun_path["/tmp/vectorchord.sock".len()], 0); + + let too_long = "x".repeat(address.sun_path.len()); + assert!(matches!( + unix_socket_address(&too_long), + Err(RerankError::Transport(message)) if message == "endpoint path is too long" + )); + assert!(matches!( + unix_socket_address("invalid\0path"), + Err(RerankError::Transport(message)) if message == "endpoint contains a NUL byte" + )); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/rerank.rs b/src/index/vchordrq/scanners/maxsim/rerank.rs new file mode 100644 index 00000000..151ae4ea --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/rerank.rs @@ -0,0 +1,315 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::{HeapKey, PageCandidate}; +use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; +use crate::index::vchordrq::opclass::Opfamily; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use vchordrq::types::OwnedVector; + +pub(super) struct CandidateTensor { + pub candidate: PageCandidate, + pub vectors: Vec, +} + +pub(super) trait CandidateTensorSource { + fn fetch(&mut self, candidate: PageCandidate) -> Result, RerankError>; +} + +pub(super) struct HeapArrayTensorSource<'a, F> { + fetcher: &'a mut F, + opfamily: Opfamily, +} + +impl<'a, F> HeapArrayTensorSource<'a, F> { + pub fn new(fetcher: &'a mut F, opfamily: Opfamily) -> Self { + Self { fetcher, opfamily } + } +} + +impl CandidateTensorSource for HeapArrayTensorSource<'_, F> { + fn fetch(&mut self, candidate: PageCandidate) -> Result, RerankError> { + let Some(mut tuple) = self.fetcher.fetch(candidate.heap_key) else { + return Ok(None); + }; + // Exact sources are the last boundary before a tensor may leave the + // executor process. Re-evaluate the active base-relation scan qual + // here even if token reranking already prefiltered some hits. This + // keeps same-relation quals ahead of CPU/GPU tensor access and + // also covers configurations with token-level refine disabled. + if !tuple.filter() { + return Ok(None); + } + let (values, is_nulls) = tuple.build(); + if is_nulls[0] { + return Err(RerankError::TensorMismatch); + } + let vectors = + unsafe { self.opfamily.input_vectors(values[0]) }.ok_or(RerankError::TensorMismatch)?; + Ok(Some(CandidateTensor { candidate, vectors })) + } +} + +pub(super) trait ExactMaxsimBackend { + type Results: Iterator; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result; +} + +#[derive(Debug)] +pub(super) enum RerankError { + TensorMismatch, + ModelContractMismatch, + InvalidDescriptor(&'static str), + Registry(String), + Configuration(&'static str), + UnsupportedTensorKind, + RequestTooLarge, + Transport(String), + Protocol(String), + Remote(String), +} + +impl std::fmt::Display for RerankError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::TensorMismatch => write!(f, "MaxSim tensor kind or dimension is not matched"), + Self::ModelContractMismatch => write!(f, "MaxSim model contract is not matched"), + Self::InvalidDescriptor(message) => { + write!(f, "invalid external MaxSim tensor descriptor: {message}") + } + Self::Registry(message) => { + write!(f, "MaxSim tensor-source registry error: {message}") + } + Self::Configuration(message) => write!(f, "MaxSim configuration error: {message}"), + Self::UnsupportedTensorKind => { + write!(f, "GPU MaxSim supports only vector and halfvec tensors") + } + Self::RequestTooLarge => write!(f, "GPU MaxSim request exceeds configured limits"), + Self::Transport(message) => write!(f, "GPU MaxSim transport error: {message}"), + Self::Protocol(message) => write!(f, "GPU MaxSim protocol error: {message}"), + Self::Remote(message) => write!(f, "GPU MaxSim sidecar error: {message}"), + } + } +} + +impl std::error::Error for RerankError {} + +#[derive(Default)] +pub(super) struct CpuExactMaxsimBackend; + +impl ExactMaxsimBackend for CpuExactMaxsimBackend { + type Results = RerankResults; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let mut results = BinaryHeap::new(); + for candidate in candidates { + let Some(tensor) = source.fetch(candidate)? else { + continue; + }; + let Some(distance) = exact_maxsim_distance(query, &tensor.vectors) else { + return Err(RerankError::TensorMismatch); + }; + results.push((Reverse(distance), Reverse(tensor.candidate.heap_key))); + } + Ok(RerankResults { inner: results }) + } +} + +fn exact_maxsim_distance(query: &[OwnedVector], document: &[OwnedVector]) -> Option { + if query.is_empty() || document.is_empty() { + return None; + } + let mut maxsim = 0.0f32; + for query_vector in query { + let mut best = Distance::INFINITY; + for document_vector in document { + let distance = document_vector.operator_dot(query_vector)?; + best = std::cmp::min(best, distance); + } + maxsim += best.to_f32(); + } + Some(Distance::from_f32(maxsim)) +} + +#[derive(Clone, Copy, Debug)] +pub(super) struct RerankedPage { + pub distance: Distance, + pub heap_key: HeapKey, +} + +pub(super) struct RerankResults { + pub(super) inner: BinaryHeap<(Reverse, Reverse)>, +} + +impl Iterator for RerankResults { + type Item = RerankedPage; + + fn next(&mut self) -> Option { + let (Reverse(distance), Reverse(heap_key)) = self.inner.pop()?; + Some(RerankedPage { distance, heap_key }) + } + + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len(); + (exact, Some(exact)) + } +} + +impl ExactSizeIterator for RerankResults {} +impl std::iter::FusedIterator for RerankResults {} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeMap; + use vector::vect::VectOwned; + + struct MockTensorSource(BTreeMap>); + + impl CandidateTensorSource for MockTensorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some(CandidateTensor { + candidate, + vectors: self + .0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + })) + } + } + + fn vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf32(VectOwned::new(values.to_vec())) + } + + struct RejectingFetcher; + + struct RejectingTuple; + + impl Tuple for RejectingTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + panic!("a rejected tuple must not materialize its tensor") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + panic!("a rejected tuple must not expose heap attributes") + } + } + + impl FilterableTuple for RejectingTuple { + fn filter(&mut self) -> bool { + false + } + } + + impl Fetcher for RejectingFetcher { + type Tuple<'a> = RejectingTuple; + + fn fetch(&mut self, _key: HeapKey) -> Option> { + Some(RejectingTuple) + } + } + + #[test] + fn heap_source_applies_scan_qual_before_materializing_tensor() { + let mut fetcher = RejectingFetcher; + let mut source = HeapArrayTensorSource::new(&mut fetcher, Opfamily::VectorMaxsim); + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + }; + + assert!(source.fetch(candidate).unwrap().is_none()); + } + + #[test] + fn cpu_backend_reorders_candidates_by_exact_page_maxsim() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::from_f32(-2.0), + heap_key: page_2, + }, + PageCandidate { + approximate_distance: Distance::from_f32(-1.0), + heap_key: page_1, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]), + (page_2, vec![vector(&[0.5, 0.5])]), + ])); + let results = CpuExactMaxsimBackend + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].heap_key, page_1); + assert_eq!(results[0].distance.to_f32(), -2.0); + assert_eq!(results[1].heap_key, page_2); + assert_eq!(results[1].distance.to_f32(), -1.0); + } + + #[test] + fn exact_ties_are_ordered_by_heap_key() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::from_f32(0.0), + heap_key: page_2, + }, + PageCandidate { + approximate_distance: Distance::from_f32(0.0), + heap_key: page_1, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0])]), + (page_2, vec![vector(&[1.0, 0.0])]), + ])); + + let results = CpuExactMaxsimBackend + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!( + results.iter().map(|page| page.heap_key).collect::>(), + vec![page_1, page_2] + ); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/search.rs b/src/index/vchordrq/scanners/maxsim/search.rs new file mode 100644 index 00000000..915fb87f --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/search.rs @@ -0,0 +1,578 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::MaxsimBuilder; +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorSourceBinding, + ExternalTensorStorage, resolve_external_tensor_source, validate_descriptor, +}; +use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; +use super::rerank::RerankError; +use crate::index::fetcher::{ + Fetcher, FilterableTuple, HeapFetcher, Tuple, TupleAttribute, ctid_to_key, +}; +use crate::index::gucs::{self, PostgresMaxsimBackend}; +use crate::index::scanners::SearchBuilder; +use crate::index::storage::PostgresRelation; +use crate::index::vchordrq::opclass::{Opfamily, opfamily}; +use crate::index::vchordrq::scanners::SearchOptions; +use crate::recorder::DefaultRecorder; +use distance::Distance; +use pgrx::datum::{DatumWithOid, FromDatum}; +use pgrx::iter::TableIterator; +use pgrx::{AnyArray, IntoDatum, name}; +use std::collections::{BTreeMap, BTreeSet}; +use std::time::Duration; + +const MAX_EXPLICIT_CANDIDATES: i32 = 65_536; + +/// Restricted Phase 3B search surface. Candidate generation reads only the +/// named index. Descriptor projection happens later through SPI, under the +/// caller's normal SELECT privileges and active MVCC snapshot. +#[pgrx::pg_extern(sql = "")] +fn _vchordrq_maxsim_search_external( + index_oid: pgrx::pg_sys::Oid, + query: AnyArray, + candidate_limit: i32, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = execute_external_search(index_oid, query, candidate_limit, top_k) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +fn execute_external_search( + index_oid: pgrx::pg_sys::Oid, + query: AnyArray, + candidate_limit: i32, + top_k: i32, +) -> Result, RerankError> { + validate_search_limits(candidate_limit, top_k)?; + if !matches!(gucs::vchordrq_maxsim_backend(), PostgresMaxsimBackend::Gpu) { + return Err(RerankError::Configuration( + "external MaxSim search currently requires vchordrq.maxsim_backend = 'gpu'", + )); + } + + let binding = resolve_external_tensor_source(index_oid)?; + let index_lock = RelationLock::open(index_oid, pgrx::pg_sys::AccessShareLock as _)?; + let heap_lock = RelationLock::open(binding.heap_oid, pgrx::pg_sys::AccessShareLock as _)?; + let descriptor_lock = binding + .descriptor_oid + .map(|oid| RelationLock::open(oid, pgrx::pg_sys::AccessShareLock as _)) + .transpose()?; + if binding.index_oid != index_lock.oid() || binding.heap_oid != heap_lock.oid() { + return Err(RerankError::Registry( + "registered MaxSim tensor source changed during execution".into(), + )); + } + if descriptor_lock.as_ref().map(RelationLock::oid) != binding.descriptor_oid { + return Err(RerankError::Registry( + "registered descriptor relation changed during execution".into(), + )); + } + + let opfamily = unsafe { opfamily(index_lock.raw()) }; + if !matches!(opfamily, Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim) { + return Err(RerankError::UnsupportedTensorKind); + } + let indexed_type = unsafe { pgrx::pg_sys::get_atttype(index_oid, 1) }; + if indexed_type != query.oid() { + return Err(RerankError::TensorMismatch); + } + let query_vectors = + unsafe { opfamily.input_vectors(query.datum()) }.ok_or(RerankError::TensorMismatch)?; + + // The registry resolution happens before the index read, and this + // privilege-only SELECT is planned/executed before any candidate CTID is + // generated. It fails early when the caller cannot project the registered + // descriptor columns. The same query shape is used for the actual fetch. + preflight_descriptor_access(&binding)?; + + let candidates = generate_candidates( + index_lock.raw(), + opfamily, + query.datum(), + candidate_limit as u32, + )?; + let resolved = + resolve_visible_candidates(index_lock.raw(), heap_lock.raw(), candidates.into_iter())?; + let (mut source, public_ids) = load_visible_descriptors(&binding, &resolved)?; + let visible_candidates = resolved + .into_iter() + .filter_map(|resolved| { + public_ids + .contains_key(&resolved.candidate.heap_key) + .then_some(resolved.candidate) + }) + .collect::>(); + let endpoint = gucs::vchordrq_maxsim_gpu_endpoint() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let transport = UnixSocketTransport::new(endpoint); + let mut backend = GpuExternalTileMaxsimBackend::new( + transport, + binding.model_contract_id, + Duration::from_millis(gucs::vchordrq_maxsim_gpu_timeout_ms() as u64), + gucs::vchordrq_maxsim_gpu_max_batch_tokens() as usize, + gucs::vchordrq_maxsim_gpu_max_batch_bytes() as usize, + ); + let mut candidate_iter = visible_candidates.into_iter(); + let exact = backend.rerank(&query_vectors, &mut candidate_iter, &mut source)?; + let mut rows = exact + .map(|result| { + let public_id = public_ids.get(&result.heap_key).copied().ok_or_else(|| { + RerankError::Protocol("sidecar result has no visible public ID".into()) + })?; + Ok((result.distance, public_id)) + }) + .collect::, RerankError>>()?; + rows.sort_unstable_by(|(left_distance, left_id), (right_distance, right_id)| { + left_distance + .cmp(right_distance) + .then_with(|| left_id.cmp(right_id)) + }); + rows.truncate(top_k as usize); + Ok(rows + .into_iter() + .map(|(distance, public_id)| (public_id, -distance.to_f32())) + .collect::>() + .into_iter()) +} + +#[derive(Clone, Copy)] +struct ResolvedCandidate { + candidate: PageCandidate, + current_ctid: pgrx::pg_sys::ItemPointerData, +} + +fn resolve_visible_candidates( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + candidates: impl Iterator, +) -> Result, RerankError> { + let snapshot = unsafe { pgrx::pg_sys::GetActiveSnapshot() }; + if snapshot.is_null() { + return Err(RerankError::Configuration( + "external MaxSim search requires an active MVCC snapshot", + )); + } + if unsafe { (*snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC { + return Err(RerankError::Configuration( + "external MaxSim search requires an MVCC snapshot", + )); + } + + let mut fetcher = + unsafe { HeapFetcher::new_standalone(index_relation, heap_relation, snapshot) }; + let mut resolved = Vec::new(); + for candidate in candidates { + pgrx::check_for_interrupts!(); + if let Some(tuple) = fetcher.fetch(candidate.heap_key) { + resolved.push(ResolvedCandidate { + candidate, + current_ctid: tuple.ctid(), + }); + } + } + Ok(resolved) +} + +fn validate_search_limits(candidate_limit: i32, top_k: i32) -> Result<(), RerankError> { + if !(1..=MAX_EXPLICIT_CANDIDATES).contains(&candidate_limit) { + return Err(RerankError::Configuration( + "candidate_limit must be between 1 and 65536", + )); + } + if top_k <= 0 || top_k > candidate_limit { + return Err(RerankError::Configuration( + "top_k must be positive and no greater than candidate_limit", + )); + } + Ok(()) +} + +fn generate_candidates( + index_relation: pgrx::pg_sys::Relation, + opfamily: Opfamily, + query: pgrx::pg_sys::Datum, + candidate_limit: u32, +) -> Result, RerankError> { + let index = unsafe { PostgresRelation::::new(index_relation) }; + let mut builder = MaxsimBuilder::new(opfamily); + unsafe { builder.add(3, Some(query)) }; + let options = SearchOptions { + epsilon: unsafe { gucs::vchordrq_epsilon(index_relation) }, + probes: unsafe { gucs::vchordrq_probes(index_relation) }, + max_scan_tuples: None, + maxsim_refine: gucs::vchordrq_maxsim_refine(index_relation), + maxsim_threshold: gucs::vchordrq_maxsim_threshold(index_relation), + maxsim_candidate_limit: Some(candidate_limit), + maxsim_backend: PostgresMaxsimBackend::CoarseOnly, + maxsim_gpu_endpoint: None, + maxsim_gpu_timeout_ms: 1, + maxsim_gpu_max_batch_tokens: 1, + maxsim_gpu_max_batch_bytes: 1, + io_search: gucs::vchordrq_io_search(), + io_rerank: gucs::vchordrq_io_rerank(), + // General same-relation quals would require the optional Phase 3C + // CustomScan. The restricted function applies PostgreSQL row + // visibility during the following SPI descriptor fetch. + prefilter: false, + }; + let bump = bumpalo::Bump::new(); + let recorder = DefaultRecorder { + enable: false, + rate: None, + max_records: 0, + index: unsafe { (*index_relation).rd_id.to_u32() }, + }; + let candidates = builder + .build(&index, options, NoHeapFetcher, &bump, recorder) + .map(|(distance, heap_key, _)| PageCandidate { + approximate_distance: Distance::from_f32(distance), + heap_key, + }) + .collect(); + Ok(candidates) +} + +fn preflight_descriptor_access(binding: &ExternalTensorSourceBinding) -> Result<(), RerankError> { + let query = descriptor_query(binding, true)?; + pgrx::spi::Spi::connect(|client| { + let privilege = client + .prepare( + "SELECT pg_catalog.has_table_privilege($1, 'SELECT') AS allowed", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + for relation_oid in std::iter::once(binding.heap_oid).chain(binding.descriptor_oid) { + let allowed = client + .select(&privilege, Some(1), &[relation_oid.into()]) + .map_err(registry_error)? + .first() + .get_by_name::("allowed") + .map_err(registry_error)? + .unwrap_or(false); + if !allowed { + return Err(RerankError::Registry( + "table-level SELECT privilege is required for every external MaxSim relation" + .into(), + )); + } + } + client + .select(query.as_str(), Some(1), &[]) + .map(|_| ()) + .map_err(registry_error) + }) +} + +fn load_visible_descriptors( + binding: &ExternalTensorSourceBinding, + candidates: &[ResolvedCandidate], +) -> Result<(MaterializedDescriptorSource, BTreeMap), RerankError> { + if candidates.is_empty() { + return Ok((MaterializedDescriptorSource::default(), BTreeMap::new())); + } + let mut candidates_by_key = BTreeMap::new(); + for resolved in candidates { + if candidates_by_key + .insert(ctid_to_key(resolved.current_ctid), resolved.candidate) + .is_some() + { + return Err(RerankError::Registry( + "multiple index candidates resolved to the same visible CTID".into(), + )); + } + } + let ctids = candidates + .iter() + .map(|resolved| resolved.current_ctid) + .collect::>(); + let query = descriptor_query(binding, false)?; + let (descriptors, public_ids) = pgrx::spi::Spi::connect(|client| { + let tid_array_oid = unsafe { pgrx::pg_sys::get_array_type(pgrx::pg_sys::TIDOID) }; + let prepared = client + .prepare( + query.as_str(), + &[ + pgrx::pg_sys::PgOid::from(tid_array_oid), + pgrx::pg_sys::PgOid::from(pgrx::pg_sys::TEXTOID), + ], + ) + .map_err(registry_error)?; + let args: [DatumWithOid<'_>; 2] = [ctids.into(), binding.model_contract_id.clone().into()]; + let rows = client + .select(&prepared, Some(candidates.len() as _), &args) + .map_err(registry_error)?; + let mut descriptors = BTreeMap::new(); + let mut public_ids = BTreeMap::new(); + let mut unique_public_ids = BTreeSet::new(); + for row in rows { + let ctid = required_heap_column::(&row, "heap_tid")?; + let heap_key = ctid_to_key(ctid); + let candidate = candidates_by_key.get(&heap_key).copied().ok_or_else(|| { + RerankError::Registry("descriptor query returned an unknown CTID".into()) + })?; + let public_id = required_heap_column::(&row, "public_id")?; + if !unique_public_ids.insert(public_id) { + return Err(RerankError::InvalidDescriptor( + "public IDs are not unique in the visible candidate batch", + )); + } + let descriptor = validate_descriptor( + candidate, + public_id, + required_heap_column::(&row, "tensor_ref")?, + required_heap_column::(&row, "tensor_rows")?, + required_heap_column::(&row, "tensor_dimension")?, + required_heap_column::(&row, "tensor_dtype")?, + required_heap_column::(&row, "tensor_checksum")?, + )?; + if descriptors.insert(candidate.heap_key, descriptor).is_some() { + return Err(RerankError::Registry( + "descriptor query returned a duplicate CTID".into(), + )); + } + public_ids.insert(candidate.heap_key, public_id); + } + Ok((descriptors, public_ids)) + })?; + Ok((MaterializedDescriptorSource(descriptors), public_ids)) +} + +fn descriptor_query( + binding: &ExternalTensorSourceBinding, + preflight: bool, +) -> Result { + let heap_relation = relation_name(binding.heap_oid)?; + let names = &binding.column_names; + let model_contract = pgrx::spi::quote_identifier(&names.model_contract); + let public_id = pgrx::spi::quote_identifier(&names.public_id); + let tensor_ref = pgrx::spi::quote_identifier(&names.tensor_ref); + let tensor_rows = pgrx::spi::quote_identifier(&names.tensor_rows); + let tensor_dimension = pgrx::spi::quote_identifier(&names.tensor_dimension); + let tensor_dtype = pgrx::spi::quote_identifier(&names.tensor_dtype); + let tensor_checksum = pgrx::spi::quote_identifier(&names.tensor_checksum); + let predicate = if preflight { + "false".to_string() + } else { + format!("h.ctid = ANY($1) AND h.{model_contract} = $2") + }; + match binding.storage { + ExternalTensorStorage::SameHeap => Ok(format!( + "SELECT h.ctid AS heap_tid, + h.{model_contract} AS model_contract, + h.{public_id} AS public_id, + h.{tensor_ref} AS tensor_ref, + h.{tensor_rows} AS tensor_rows, + h.{tensor_dimension} AS tensor_dimension, + h.{tensor_dtype} AS tensor_dtype, + h.{tensor_checksum} AS tensor_checksum + FROM ONLY {heap_relation} AS h + WHERE {predicate}" + )), + ExternalTensorStorage::DescriptorRelation => { + let descriptor_oid = binding.descriptor_oid.ok_or_else(|| { + RerankError::Registry("registered descriptor relation is missing".into()) + })?; + let descriptor_relation = relation_name(descriptor_oid)?; + let descriptor_public_id = pgrx::spi::quote_identifier( + names.descriptor_public_id.as_deref().ok_or_else(|| { + RerankError::Registry( + "registered descriptor public ID column is missing".into(), + ) + })?, + ); + Ok(format!( + "SELECT h.ctid AS heap_tid, + h.{model_contract} AS model_contract, + h.{public_id} AS public_id, + d.{tensor_ref} AS tensor_ref, + d.{tensor_rows} AS tensor_rows, + d.{tensor_dimension} AS tensor_dimension, + d.{tensor_dtype} AS tensor_dtype, + d.{tensor_checksum} AS tensor_checksum + FROM ONLY {heap_relation} AS h + LEFT JOIN ONLY {descriptor_relation} AS d + ON d.{descriptor_public_id} = h.{public_id} + WHERE {predicate}" + )) + } + } +} + +fn relation_name(relation_oid: pgrx::pg_sys::Oid) -> Result { + pgrx::spi::Spi::connect(|client| { + let prepared = client + .prepare( + "SELECT n.nspname::text AS schema_name, c.relname::text AS relation_name + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE c.oid = $1", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[relation_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered relation disappeared".into(), + )); + } + let row = rows.first(); + Ok(pgrx::spi::quote_qualified_identifier( + required_column::(&row, "schema_name")?, + required_column::(&row, "relation_name")?, + )) + }) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn required_heap_column( + row: &pgrx::spi::SpiHeapTupleData<'_>, + name: &str, +) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +#[derive(Default)] +struct MaterializedDescriptorSource(BTreeMap); + +impl CandidateTensorDescriptorSource for MaterializedDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(self.0.remove(&candidate.heap_key)) + } +} + +struct RelationLock { + raw: pgrx::pg_sys::Relation, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl RelationLock { + fn open(oid: pgrx::pg_sys::Oid, lockmode: pgrx::pg_sys::LOCKMODE) -> Result { + let raw = unsafe { pgrx::pg_sys::relation_open(oid, lockmode) }; + if raw.is_null() { + return Err(RerankError::Registry("relation open returned NULL".into())); + } + Ok(Self { raw, lockmode }) + } + + fn raw(&self) -> pgrx::pg_sys::Relation { + self.raw + } + + fn oid(&self) -> pgrx::pg_sys::Oid { + unsafe { (*self.raw).rd_id } + } +} + +impl Drop for RelationLock { + fn drop(&mut self) { + unsafe { pgrx::pg_sys::relation_close(self.raw, self.lockmode) }; + } +} + +struct NoHeapFetcher; +struct NoHeapTuple; + +impl Tuple for NoHeapTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + unreachable!("restricted external candidate generation must not read the heap") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +impl FilterableTuple for NoHeapTuple { + fn filter(&mut self) -> bool { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +impl Fetcher for NoHeapFetcher { + type Tuple<'a> = NoHeapTuple; + + fn fetch(&mut self, _key: HeapKey) -> Option> { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn explicit_search_limits_are_bounded() { + assert!(validate_search_limits(256, 10).is_ok()); + for (candidates, top_k) in [(0, 1), (65_537, 1), (1, 0), (10, 11)] { + assert!(matches!( + validate_search_limits(candidates, top_k), + Err(RerankError::Configuration(_)) + )); + } + } + + #[test] + fn materialized_source_only_returns_visible_descriptors_once() { + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + }; + let descriptor = validate_descriptor( + candidate, + 42, + "s3://immutable/tensor".into(), + 1, + 2, + "float16".into(), + format!("sha256:{}", "a".repeat(64)), + ) + .unwrap(); + let mut source = + MaterializedDescriptorSource(BTreeMap::from([(candidate.heap_key, descriptor)])); + assert_eq!(source.fetch(candidate).unwrap().unwrap().public_id, 42); + assert!(source.fetch(candidate).unwrap().is_none()); + } +} diff --git a/src/index/vchordrq/scanners/mod.rs b/src/index/vchordrq/scanners/mod.rs index b345da37..c5e392bc 100644 --- a/src/index/vchordrq/scanners/mod.rs +++ b/src/index/vchordrq/scanners/mod.rs @@ -15,7 +15,9 @@ mod default; mod maxsim; +use crate::index::gucs::PostgresMaxsimBackend; use crate::index::scanners::Io; +use std::ffi::CString; pub use default::DefaultBuilder; pub use maxsim::MaxsimBuilder; @@ -27,6 +29,12 @@ pub struct SearchOptions { pub max_scan_tuples: Option, pub maxsim_refine: u32, pub maxsim_threshold: u32, + pub maxsim_candidate_limit: Option, + pub maxsim_backend: PostgresMaxsimBackend, + pub maxsim_gpu_endpoint: Option, + pub maxsim_gpu_timeout_ms: u32, + pub maxsim_gpu_max_batch_tokens: u32, + pub maxsim_gpu_max_batch_bytes: u32, pub io_search: Io, pub io_rerank: Io, pub prefilter: bool, diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index a4f8f38a..a7e7af4e 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -610,3 +610,824 @@ FROM WHERE am.amname = 'vchordrq' ) AS index_oids CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; + +-- Phase 3B tensor-source bindings + +CREATE TABLE _vchordrq_maxsim_sources ( + index_oid oid PRIMARY KEY, + heap_oid oid NOT NULL, + model_contract_id text NOT NULL + CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['heap_array', 'external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint, + tensor_rows_attnum smallint, + tensor_dim_attnum smallint, + tensor_dtype_attnum smallint, + tensor_checksum_attnum smallint, + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'heap_array'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NULL + AND tensor_rows_attnum IS NULL + AND tensor_dim_attnum IS NULL + AND tensor_dtype_attnum IS NULL + AND tensor_checksum_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_maxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_maxsim_source( + index_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name DEFAULT NULL, + tensor_rows_column name DEFAULT NULL, + tensor_dim_column name DEFAULT NULL, + tensor_dtype_column name DEFAULT NULL, + tensor_checksum_column name DEFAULT NULL, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + heap_oid oid; + index_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + descriptor_id_is_unique boolean; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be heap_array, external_ref, or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + IF caller_oid IS NULL THEN + RAISE EXCEPTION 'could not resolve caller role'; + END IF; + + SELECT x.indrelid, i.relowner + INTO heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF heap_oid IS NULL THEN + RAISE EXCEPTION 'relation % is not a valid single-key vchordrq MaxSim index', + index_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may register its MaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a MaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO descriptor_id_is_unique; + IF NOT descriptor_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION '% sources must not specify a descriptor relation', normalized_storage; + END IF; + tensor_relation_oid := heap_oid; + END IF; + + IF normalized_storage = 'heap_array' THEN + IF tensor_ref_column IS NOT NULL + OR tensor_rows_column IS NOT NULL + OR tensor_dim_column IS NOT NULL + OR tensor_dtype_column IS NOT NULL + OR tensor_checksum_column IS NOT NULL THEN + RAISE EXCEPTION 'heap_array sources must not specify external tensor columns'; + END IF; + ELSE + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor sources require ref, rows, dim, dtype, and checksum columns'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', + tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', + tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', + tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', + tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'tensor source columns must be distinct'; + END IF; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_maxsim_sources ( + index_oid, heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) + ON CONFLICT (index_oid) DO UPDATE SET + heap_oid = EXCLUDED.heap_oid, + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + index_relation::oid, + heap_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_maxsim_source(index_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + index_owner oid; + removed_count bigint; +BEGIN + IF index_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO index_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = index_relation::oid AND c.relkind = 'i'; + IF index_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may unregister its MaxSim tensor source'; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources WHERE index_oid = $1', + ext_schema + ) USING index_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_maxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_maxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_source_info(index_relation regclass) +RETURNS TABLE( + registered_index regclass, + heap_relation regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + bound_heap_oid oid; + live_heap_oid oid; + index_owner oid; + bound_model_contract_id text; + bound_storage text; + bound_descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + expected_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + + EXECUTE pg_catalog.format( + 'SELECT heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_maxsim_sources + WHERE index_oid = $1', + ext_schema + ) INTO + bound_heap_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + bound_descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING index_relation::oid; + IF bound_heap_oid IS NULL THEN + RAISE EXCEPTION 'MaxSim tensor source is not registered for index %', + index_relation; + END IF; + + SELECT x.indrelid, i.relowner + INTO live_heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF live_heap_oid IS NULL OR live_heap_oid <> bound_heap_oid THEN + RAISE EXCEPTION 'registered MaxSim tensor source is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, live_heap_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered MaxSim tensor source'; + END IF; + + IF bound_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid storage'; + END IF; + + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, 'bigint'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = bound_heap_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap columns'; + END IF; + + IF bound_storage = 'heap_array' THEN + IF bound_descriptor_oid IS NOT NULL + OR descriptor_public_id_attnum IS NOT NULL + OR tensor_ref_attnum IS NOT NULL + OR tensor_rows_attnum IS NOT NULL + OR tensor_dim_attnum IS NOT NULL + OR tensor_dtype_attnum IS NOT NULL + OR tensor_checksum_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap_array binding'; + END IF; + tensor_relation_oid := NULL; + ELSIF bound_storage = 'external_ref' THEN + IF bound_descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_ref binding'; + END IF; + tensor_relation_oid := bound_heap_oid; + ELSE + IF bound_descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_relation binding'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = bound_descriptor_oid + AND c.relkind IN ('r', 'm'); + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor relation is stale or invalid'; + END IF; + IF NOT pg_catalog.has_table_privilege(caller_oid, bound_descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'SELECT privilege on the registered descriptor relation is required'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid = 'bigint'::regtype + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID column is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = bound_descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := bound_descriptor_oid; + END IF; + + expected_columns := CASE WHEN bound_storage = 'heap_array' THEN 0 ELSE 5 END; + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.attnum IS NOT NULL; + IF valid_columns <> expected_columns THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid descriptor columns'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + index_relation, + bound_heap_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + bound_descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_search( + index_relation regclass, + query anyarray, + candidate_limit integer, + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_maxsim_search_external_wrapper'; + +COMMENT ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) +IS 'Restricted Phase 3B external-tensor MaxSim search; returns exact similarity under caller MVCC, SELECT privileges, and PostgreSQL row visibility.'; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_maxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + registry regclass; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RETURN; + END IF; + registry := pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_maxsim_sources', ext_schema) + ); + IF registry IS NULL THEN + RETURN; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE d.objid = s.index_oid + OR ( + d.objid = s.heap_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.model_contract_attnum + OR d.objsubid = s.public_id_attnum + OR ( + s.storage = ''external_ref'' + AND d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.descriptor_public_id_attnum + OR d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); diff --git a/tests/vchordrq/cost_estimator.slt b/tests/vchordrq/cost_estimator.slt index 0f4ad266..87cf85e7 100644 --- a/tests/vchordrq/cost_estimator.slt +++ b/tests/vchordrq/cost_estimator.slt @@ -388,6 +388,174 @@ RESET vchordrq.enable_scan; statement ok RESET enable_seqscan; +# --------------------------------------------------------------------------- +# Case 12: MaxSim must have a nonzero, query-token-aware, backend-aware cost. +# The old special branch returned total_cost=0 and selectivity=1 for every +# MaxSim path, hiding all token expansion and exact rerank work. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_maxsim ( + id int PRIMARY KEY, + v vector(3)[] NOT NULL +); + +statement ok +INSERT INTO cost_test_maxsim +SELECT i, ARRAY[ + '[1,0,0]'::vector, + '[0,1,0]'::vector, + '[0,0,1]'::vector, + '[0.5,0.5,0]'::vector +] +FROM generate_series(1, 1000) i; + +statement ok +CREATE INDEX cost_test_maxsim_v +ON cost_test_maxsim +USING vchordrq (v vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [4] +$$); + +statement ok +ANALYZE cost_test_maxsim; + +statement ok +SET enable_seqscan = off; + +statement ok +SET vchordrq.probes = '4'; + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4; + +statement ok +CREATE TEMP TABLE maxsim_cost_observations ( + name text PRIMARY KEY, + value double precision NOT NULL +); + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 1; + +# A newly built index has a native indexed-vector count. The document-token +# GUC is only a compatibility fallback for pre-statistics indexes, so changing +# it must not change this index's cost. +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'native_document_4', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4096; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'native_document_4096', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'native_document_4') + = + (SELECT value FROM maxsim_cost_observations WHERE name = 'native_document_4096'); +---- +t + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'query_1', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 64; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'query_64', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'query_64') + > + (SELECT value FROM maxsim_cost_observations WHERE name = 'query_1'); +---- +t + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 32; + +statement ok +SET vchordrq.maxsim_candidate_limit = 16; + +statement ok +SET vchordrq.maxsim_backend = 'cpu_exact'; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'cpu_16', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_candidate_limit = 256; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'cpu_256', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'cpu_256') + > + (SELECT value FROM maxsim_cost_observations WHERE name = 'cpu_16'); +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); +---- +cost_test_maxsim_v + +statement ok +RESET vchordrq.maxsim_backend; + +statement ok +RESET vchordrq.maxsim_candidate_limit; + +statement ok +RESET vchordrq.maxsim_planner_query_tokens; + +statement ok +RESET vchordrq.maxsim_planner_document_tokens; + +statement ok +RESET vchordrq.probes; + # --------------------------------------------------------------------------- # Cleanup # --------------------------------------------------------------------------- @@ -404,6 +572,9 @@ DROP TABLE cost_test_cold; statement ok DROP TABLE cost_test_partial; +statement ok +DROP TABLE cost_test_maxsim; + statement ok DROP FUNCTION slow_true(int); diff --git a/tests/vchordrq/maxsim_correctness.slt b/tests/vchordrq/maxsim_correctness.slt new file mode 100644 index 00000000..a5d9652f --- /dev/null +++ b/tests/vchordrq/maxsim_correctness.slt @@ -0,0 +1,221 @@ +# Deterministic MaxSim semantics and input-boundary coverage. + +statement ok +SET enable_seqscan TO off; + +# @# is a distance: negative late-interaction similarity, ordered ascending. +query I +SELECT round(( + ARRAY['[1,0]'::vector, '[0,1]'::vector] + @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +)::numeric, 3); +---- +-2.000 + +# MaxSim is asymmetric: the right-hand array contains query vectors. +query I +SELECT round(( + ARRAY['[1,0]'::vector, '[0,1]'::vector] + @# ARRAY['[1,1]'::vector] +)::numeric, 3); +---- +-1.000 + +query I +SELECT round(( + ARRAY['[1,1]'::vector] + @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +)::numeric, 3); +---- +-2.000 + +# halfvec uses the same sign and orientation contract. +query I +SELECT round(( + ARRAY['[1,0]'::halfvec, '[0,1]'::halfvec] + @# ARRAY['[1,0]'::halfvec, '[0,1]'::halfvec] +)::numeric, 3); +---- +-2.000 + +statement error MaxSim arrays must contain at least one vector +SELECT ARRAY[]::vector[] @# ARRAY['[1,0]'::vector]; + +statement error MaxSim arrays must contain at least one vector +SELECT ARRAY['[1,0]'::vector] @# ARRAY[]::vector[]; + +statement error MaxSim arrays must not contain NULL vectors +SELECT ARRAY['[1,0]'::vector, NULL::vector] + @# ARRAY['[1,0]'::vector]; + +statement error dimension is not matched +SELECT ARRAY['[1,0]'::vector] @# ARRAY['[1,0,0]'::vector]; + +statement error MaxSim arrays cannot contain more than 65536 vectors +SELECT ARRAY['[1]'::vector] + @# ARRAY(SELECT '[1]'::vector FROM generate_series(1, 65537)); + +# Indexing an empty document array must fail instead of silently omitting it. +statement ok +CREATE TABLE maxsim_invalid_document (val vector(2)[]); + +statement ok +INSERT INTO maxsim_invalid_document VALUES (ARRAY[]::vector[]); + +statement error MaxSim arrays must contain at least one vector +CREATE INDEX ON maxsim_invalid_document +USING vchordrq (val vector_maxsim_ops); + +statement ok +TRUNCATE maxsim_invalid_document; + +statement ok +INSERT INTO maxsim_invalid_document +SELECT ARRAY( + SELECT '[1,0]'::vector + FROM generate_series(1, 65537) +); + +statement error MaxSim arrays cannot contain more than 65536 vectors +CREATE INDEX ON maxsim_invalid_document +USING vchordrq (val vector_maxsim_ops); + +statement ok +DROP TABLE maxsim_invalid_document; + +# Lock down current ranking with a non-flat index before the Phase 3 refactor. +statement ok +CREATE TABLE maxsim_deterministic ( + id integer primary key, + val vector(2)[] not null +); + +statement ok +INSERT INTO maxsim_deterministic VALUES + (1, ARRAY['[1,0]'::vector, '[0,1]'::vector]), + (2, ARRAY['[0.8,0]'::vector, '[0,0.8]'::vector]), + (3, ARRAY['[0.5,0.5]'::vector]), + (4, ARRAY['[-1,0]'::vector, '[0,-1]'::vector]); + +statement ok +CREATE INDEX maxsim_deterministic_idx +ON maxsim_deterministic +USING vchordrq (val vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [2] +$$); + +statement ok +SET vchordrq.probes = '2'; + +statement ok +SET vchordrq.maxsim_refine = 100; + +statement ok +SET vchordrq.maxsim_threshold = 0; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +statement ok +SET vchordrq.maxsim_candidate_limit = 2; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 + +statement ok +SET vchordrq.maxsim_candidate_limit = -1; + +statement ok +SET vchordrq.maxsim_backend = 'cpu_exact'; + +statement error exact MaxSim requires a positive vchordrq.maxsim_candidate_limit +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 1; + +statement ok +SET vchordrq.maxsim_candidate_limit = 4; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +query I +SELECT round((val @# ARRAY['[1,0]'::vector, '[0,1]'::vector])::numeric, 3) +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +-2.000 +-1.600 +-1.000 +0.000 + +statement ok +SET vchordrq.maxsim_backend = 'gpu'; + +statement error GPU MaxSim transport error: endpoint is empty +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 1; + +statement ok +SET vchordrq.maxsim_backend = 'auto'; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +statement ok +SET vchordrq.maxsim_candidate_limit = -1; + +# An empty query must also fail on the index scan path. +statement error MaxSim arrays must contain at least one vector +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY[]::vector[] +LIMIT 1; + +statement error dimension is not matched +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0,0]'::vector] +LIMIT 1; + +statement ok +DROP TABLE maxsim_deterministic; diff --git a/tests/vchordrq/maxsim_source_registry.slt b/tests/vchordrq/maxsim_source_registry.slt new file mode 100644 index 00000000..c408a469 --- /dev/null +++ b/tests/vchordrq/maxsim_source_registry.slt @@ -0,0 +1,341 @@ +# Phase 3B tensor-source registration must bind to a real MaxSim index by +# relation/attribute OID, reject incompatible descriptors, and fail closed when +# a bound index/table column is dropped. + +statement ok +CREATE TABLE maxsim_source_test ( + id bigint PRIMARY KEY, + model_contract text NOT NULL, + embedding vector(2)[] NOT NULL, + single_embedding vector(2) NOT NULL, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL, + application_note text +); + +statement ok +INSERT INTO maxsim_source_test VALUES ( + 1, + 'colqwen@test', + ARRAY['[1,0]'::vector, '[0,1]'::vector], + '[1,0]'::vector, + 'tensor://1', + 2, + 2, + 'float32', + 'sha256:test' +); + +statement ok +CREATE INDEX maxsim_source_test_idx +ON maxsim_source_test +USING vchordrq (embedding vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [1] +$$); + +statement ok +CREATE INDEX maxsim_source_wrong_idx +ON maxsim_source_test +USING vchordrq (single_embedding vector_l2_ops) +WITH (options = $$ +[build.internal] +lists = [1] +$$); + +statement error not a valid single-key vchordrq MaxSim index +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_wrong_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +statement error tensor rows column tensor_dtype must be a NOT NULL integer column +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_ref', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_dtype', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => ' colqwen@test ', + storage => 'EXTERNAL_REF', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +query TT +SELECT s.storage, s.model_contract_id +FROM _vchordrq_maxsim_sources AS s +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +external_ref colqwen@test + +query TTT +SELECT source_storage, tensor_ref_column::text, public_id_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +external_ref tensor_ref id + +# Production external tensors may live in a compact descriptor relation keyed +# by the stable application public ID, avoiding a rewrite of the indexed heap. +statement ok +CREATE TABLE maxsim_descriptor_no_unique ( + public_id bigint NOT NULL, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement error must have a non-partial single-key unique index +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_relation', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'maxsim_descriptor_no_unique'::regclass, + descriptor_public_id_column => 'public_id' +); + +statement ok +CREATE TABLE maxsim_descriptor_test ( + public_id bigint PRIMARY KEY, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement ok +INSERT INTO maxsim_descriptor_test VALUES ( + 1, 'tensor://1', 2, 2, 'float32', 'sha256:test' +); + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_relation', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'maxsim_descriptor_test'::regclass, + descriptor_public_id_column => 'public_id' +); + +query TTTT +SELECT source_storage, descriptor_relation::text, + descriptor_public_id_column::text, tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +external_relation maxsim_descriptor_test public_id tensor_ref + +statement ok +ALTER TABLE maxsim_descriptor_test RENAME tensor_ref TO tensor_location; + +query T +SELECT tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +tensor_location + +statement ok +ALTER TABLE maxsim_descriptor_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +0 + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_ref', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +# Unbound application columns are outside the registry contract. Dropping one +# must not invalidate an otherwise live source binding. +statement ok +ALTER TABLE maxsim_source_test DROP COLUMN application_note; + +query T +SELECT model_contract_id +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +colqwen@test + +# The explicit score surface validates bounded work and exact query/index type +# before any sidecar access. +statement error candidate_limit must be between 1 and 65536 +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::vector], + 0, + 1 +); + +statement error top_k must be positive and no greater than candidate_limit +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::vector], + 1, + 2 +); + +statement ok +SET vchordrq.maxsim_backend = 'gpu'; + +statement error MaxSim tensor kind or dimension is not matched +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::halfvec], + 1, + 1 +); + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +query T +SELECT a.attname +FROM _vchordrq_maxsim_sources AS s +JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = s.heap_oid AND a.attnum = s.tensor_ref_attnum +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +tensor_ref + +# Attribute numbers survive renames, so bindings do not depend on SQL text. +statement ok +ALTER TABLE maxsim_source_test RENAME tensor_ref TO tensor_location; + +query T +SELECT a.attname +FROM _vchordrq_maxsim_sources AS s +JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = s.heap_oid AND a.attnum = s.tensor_ref_attnum +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +tensor_location + +query T +SELECT tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +tensor_location + +# Type-altering DDL does not drop an attribute, so runtime resolution must +# revalidate the complete descriptor and fail closed. +statement ok +ALTER TABLE maxsim_source_test +ALTER COLUMN tensor_dtype TYPE varchar(16); + +statement error registered MaxSim tensor source has invalid descriptor columns +SELECT * FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); + +statement ok +ALTER TABLE maxsim_source_test +ALTER COLUMN tensor_dtype TYPE text; + +# Dropping a bound descriptor column invalidates the complete binding. +statement ok +ALTER TABLE maxsim_source_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +0 + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +query B +SELECT tensor_ref_attnum IS NULL +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +t + +query B +SELECT vchordrq_unregister_maxsim_source('maxsim_source_test_idx'::regclass); +---- +t + +query B +SELECT vchordrq_unregister_maxsim_source('maxsim_source_test_idx'::regclass); +---- +f + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +statement ok +DROP INDEX maxsim_source_test_idx; + +query I +SELECT count(*) FROM _vchordrq_maxsim_sources; +---- +0 + +statement ok +DROP TABLE maxsim_source_test; + +statement ok +DROP TABLE maxsim_descriptor_test, maxsim_descriptor_no_unique; From b24dc2931ef49f6553f880895e9dfb195a134ccf Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:08:14 +0800 Subject: [PATCH 301/324] docs: describe public TileMaxSim fork --- README.md | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/README.md b/README.md index e305937c..edeb4ea2 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,41 @@
+# VectorChord TileMaxSim + +**An open-source VectorChord fork focused on exact multi-vector retrieval in PostgreSQL.** + +
+ +This fork extends VectorChord's `vchordrq` index with a bounded candidate stage +and exact late-interaction TileMaxSim reranking. It is intended for applications +that store one array of token vectors per document and need PostgreSQL-native +multi-vector search. + +## What this fork adds + +- Exact TileMaxSim reranking on CPU, plus an optional CUDA sidecar backend. +- A configurable upper bound on reranking candidates to keep work predictable. +- External tensor-source registration for deployments that keep full-precision + token tensors outside the indexed PostgreSQL value. +- PostgreSQL-aware permission, MVCC, row-visibility, cancellation, and timeout + handling for the external-tensor search path. +- Planner statistics and cost estimation for multi-vector queries. +- Deterministic correctness, registry, sidecar-protocol, and planner-cost tests. + +The implementation is currently under active development. Its SQL interfaces +and deployment packaging may change before a stable release. This repository +contains only the public implementation and public-facing project information; +private planning and application documentation are intentionally excluded. + +This work is based on +[supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord). The +original VectorChord README is retained below for upstream installation, +licensing, and project information. + +--- + +
+ # VectorChord **Ready for the Billion-Scale Era. Host 100M vectors on a single i4i.xlarge ($247/mo) and [scale seamlessly to 1B+](https://blog.vectorchord.ai/scaling-vector-search-to-1-billion-on-postgresql).** From d8024901087c97ff2db00b0419b0609294d46cb6 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Tue, 14 Jul 2026 10:57:05 +0800 Subject: [PATCH 302/324] feat(vchordrq): add managed GPU-resident TileMaxSim --- services/Dockerfile.tilemaxsim | 2 + services/test_tilemaxsim_cuda_sidecar.py | 277 +++++++ services/test_tilemaxsim_gpu_cache.py | 266 +++++++ services/tilemaxsim_cuda_sidecar.py | 565 ++++++++++++- services/tilemaxsim_gpu_cache.py | 430 ++++++++++ services/tilemaxsim_triton.py | 128 +++ src/index/gucs.rs | 14 + src/index/vchordrq/scanners/maxsim.rs | 74 +- .../vchordrq/scanners/maxsim/candidate.rs | 114 ++- src/index/vchordrq/scanners/maxsim/exact.rs | 500 ++++++++++++ .../vchordrq/scanners/maxsim/external.rs | 122 +++ src/index/vchordrq/scanners/maxsim/profile.rs | 181 +++++ src/index/vchordrq/scanners/maxsim/search.rs | 47 +- src/sql/finalize.sql | 752 ++++++++++++++++++ tests/vchordrq/tilemaxsim_source_registry.slt | 91 +++ 15 files changed, 3490 insertions(+), 73 deletions(-) create mode 100644 services/test_tilemaxsim_gpu_cache.py create mode 100644 services/tilemaxsim_gpu_cache.py create mode 100644 services/tilemaxsim_triton.py create mode 100644 src/index/vchordrq/scanners/maxsim/exact.rs create mode 100644 src/index/vchordrq/scanners/maxsim/profile.rs create mode 100644 tests/vchordrq/tilemaxsim_source_registry.slt diff --git a/services/Dockerfile.tilemaxsim b/services/Dockerfile.tilemaxsim index 9f58b88d..3ccb6732 100644 --- a/services/Dockerfile.tilemaxsim +++ b/services/Dockerfile.tilemaxsim @@ -5,6 +5,8 @@ WORKDIR /opt/vectorchord COPY devtools/tilemaxsim_reference_sidecar.py devtools/tilemaxsim_reference_sidecar.py COPY services/tilemaxsim_cuda_sidecar.py services/tilemaxsim_cuda_sidecar.py +COPY services/tilemaxsim_gpu_cache.py services/tilemaxsim_gpu_cache.py +COPY services/tilemaxsim_triton.py services/tilemaxsim_triton.py ENV PYTHONPATH=/opt/vectorchord \ PYTHONUNBUFFERED=1 diff --git a/services/test_tilemaxsim_cuda_sidecar.py b/services/test_tilemaxsim_cuda_sidecar.py index 093245e3..e015174d 100644 --- a/services/test_tilemaxsim_cuda_sidecar.py +++ b/services/test_tilemaxsim_cuda_sidecar.py @@ -15,10 +15,13 @@ from __future__ import annotations import hashlib +import json import os import socket import stat import struct +import subprocess +import sys import tempfile import threading import time @@ -57,6 +60,26 @@ def write_content_addressed(root: Path, payload: bytes) -> tuple[str, str]: class CudaSidecarTest(unittest.TestCase): + def test_cli_without_explicit_gpu_memory_keeps_tilemaxsim_disabled(self) -> None: + with tempfile.TemporaryDirectory() as directory: + socket_path = Path(directory) / "disabled.sock" + completed = subprocess.run( + [ + sys.executable, + "-m", + "services.tilemaxsim_cuda_sidecar", + "--socket", + os.fspath(socket_path), + ], + cwd=Path(__file__).resolve().parents[1], + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(completed.returncode, 2) + self.assertIn("TileMaxSim is disabled", completed.stderr) + self.assertFalse(socket_path.exists()) + def test_cache_builder_publishes_resolver_compatible_payload(self) -> None: with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -140,6 +163,14 @@ def test_content_addressed_resolver_validates_and_caches(self) -> None: finally: resolver.close() + def test_host_payload_cache_evicts_to_its_byte_budget(self) -> None: + cache = cuda_sidecar.PayloadCache(6) + cache.put(("first",), b"1234") + cache.put(("second",), b"5678") + self.assertIsNone(cache.get(("first",))) + self.assertEqual(cache.get(("second",)), b"5678") + self.assertEqual(cache.current_bytes, 4) + def test_content_addressed_resolver_rejects_symlink(self) -> None: with tempfile.TemporaryDirectory() as directory: root = Path(directory) / "root" @@ -246,6 +277,252 @@ def test_cuda_f16_matches_cpu_protocol_oracle(self) -> None: for (_, actual), (_, expected) in zip(results, oracle, strict=True): self.assertAlmostEqual(actual, expected, places=5) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_cli_gb_allocation_serves_tilemaxsim(self) -> None: + with tempfile.TemporaryDirectory() as directory: + directory_path = Path(directory) + root = directory_path / "tensors" + root.mkdir() + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, _ = write_content_addressed(root, payload) + frame, _ = external_request_frame( + 70, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, [[1.0, 0.0], [0.0, 1.0]])], + ) + device = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + socket_path = directory_path / "resident.sock" + process = subprocess.Popen( + [ + sys.executable, + "-m", + "services.tilemaxsim_cuda_sidecar", + "--socket", + os.fspath(socket_path), + "--gpu-memory-gb", + f"{device}=0.05", + "--gpu-workspace-gb", + "0.02", + "--host-cache-gb", + "0.01", + "--contract-root", + f"model@1={root}", + "--once", + ], + cwd=Path(__file__).resolve().parents[1], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + for _ in range(1000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + self.assertIsNone(process.poll()) + self.assertTrue(socket_path.exists()) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_len = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_len) + output, _ = process.communicate(timeout=10) + self.assertEqual(process.returncode, 0, output) + self.assertEqual(decode_response(response)[1:], (0, [(9, 2.0)])) + self.assertIn('"event":"tilemaxsim_ready"', output) + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=5) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_v2_gpu_resident_hit_does_not_resolve_payload_again(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + digest = checksum.removeprefix("sha256:") + frame, _ = external_request_frame( + 71, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, [[1.0, 0.0], [0.0, 1.0]])], + ) + device = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + pool = cuda_sidecar.GpuResourcePool( + [cuda_sidecar.GpuArenaSpec(f"cuda:{device}", 32 * 1024 * 1024)], + 16 * 1024 * 1024, + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 0) + try: + cache = cuda_sidecar.GpuTensorCache(pool, allow_eviction=False) + metrics = CapturingMetrics() + stream_engine = cuda_sidecar.TorchTileMaxsimEngine( + f"cuda:{device}", 16 * 1024 * 1024, False, 1 + ) + resident_engine = cuda_sidecar.ResidentTorchTileMaxsimEngine( + pool, 16 * 1024 * 1024, False, 1 + ) + service = cuda_sidecar.TileMaxsimService( + protocol.Limits(), + resolver, + stream_engine, + 2000, + metrics, + cache, + resident_engine, + pin_gpu_entries=True, + ) + client, server = socket.socketpair() + try: + first = service.process_frame( + frame, server, time.monotonic() + 2, None + ) + (root / digest[:2] / f"{digest}.bin").unlink() + second = service.process_frame( + frame, server, time.monotonic() + 2, None + ) + finally: + client.close() + server.close() + self.assertEqual(decode_response(first)[1:], (0, [(9, 2.0)])) + self.assertEqual(decode_response(second)[1:], (0, [(9, 2.0)])) + requests = [ + event + for event in metrics.events + if event.get("event") == "tilemaxsim_request" + ] + self.assertEqual(requests[0]["gpu_cache_misses"], 1) + self.assertEqual(requests[1]["gpu_cache_hits"], 1) + finally: + resolver.close() + pool.close() + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_lru_gpu_cache_streams_request_larger_than_its_arena(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + rows, dimension = 480, 320 + first_tensor = np.zeros((rows, dimension), dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + manifest = root / "descriptors.jsonl" + manifest.write_text( + json.dumps( + { + "page_key": "page-1", + "tensor_ref": tensor_ref, + "tensor_rows": 2, + "tensor_dim": 2, + "tensor_dtype": "float16", + "tensor_checksum": checksum, + "canonical_bytes": len(payload), + } + ) + + "\n", + encoding="utf-8", + ) + device = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + pool = cuda_sidecar.GpuResourcePool( + [cuda_sidecar.GpuArenaSpec(f"cuda:{device}", 32 * 1024 * 1024)], + 16 * 1024 * 1024, + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 0) + try: + cache = cuda_sidecar.GpuTensorCache(pool, allow_eviction=False) + metrics = CapturingMetrics() + cuda_sidecar.prewarm_resident_cache( + [("model@1", manifest)], resolver, cache, metrics + ) + status = cache.status() + self.assertEqual(status["entries"], 1) + self.assertEqual(status["pinned_entries"], 1) + self.assertEqual( + metrics.events[-1]["event"], "tilemaxsim_prewarm_complete" + ) + finally: + resolver.close() + pool.close() + def test_v2_unix_socket_end_to_end(self) -> None: with tempfile.TemporaryDirectory() as directory: directory_path = Path(directory) diff --git a/services/test_tilemaxsim_gpu_cache.py b/services/test_tilemaxsim_gpu_cache.py new file mode 100644 index 00000000..505d13b7 --- /dev/null +++ b/services/test_tilemaxsim_gpu_cache.py @@ -0,0 +1,266 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +from __future__ import annotations + +import time +import unittest + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import ResidentTorchTileMaxsimEngine +from services.tilemaxsim_gpu_cache import ( + FreeExtentAllocator, + GpuArenaSpec, + GpuResourcePool, + GpuTensorCache, + parse_gpu_memory_gb, + parse_memory_gb, +) + + +def available_device() -> str: + index = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + return f"cuda:{index}" + + +class GpuCacheUnitTest(unittest.TestCase): + def test_public_memory_configuration_uses_gb(self) -> None: + self.assertEqual(parse_memory_gb("20"), 20 * 1024**3) + self.assertEqual(parse_memory_gb("0.5"), 512 * 1024**2) + self.assertEqual( + parse_gpu_memory_gb("2=12"), + GpuArenaSpec("cuda:2", 12 * 1024**3), + ) + self.assertEqual( + parse_gpu_memory_gb("cuda:2=12.5"), + GpuArenaSpec("cuda:2", int(12.5 * 1024**3)), + ) + with self.assertRaisesRegex(ValueError, "GPU=GB"): + parse_gpu_memory_gb("cuda:0") + with self.assertRaisesRegex(ValueError, "byte suffixes"): + parse_gpu_memory_gb("0=20GiB") + + def test_extent_allocator_coalesces_released_ranges(self) -> None: + allocator = FreeExtentAllocator(4096, alignment=256) + first = allocator.allocate(300) + second = allocator.allocate(700) + self.assertEqual(first, (0, 512)) + self.assertEqual(second, (512, 768)) + assert first is not None and second is not None + allocator.release(*first) + allocator.release(*second) + self.assertEqual(allocator.extents, [(0, 4096)]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_pool_rejects_budget_larger_than_currently_free_memory(self) -> None: + device = available_device() + free_bytes, _ = torch.cuda.mem_get_info(torch.device(device)) + with self.assertRaisesRegex(RuntimeError, "cannot acquire"): + GpuResourcePool( + [GpuArenaSpec(device, free_bytes + 1024 * 1024)], + 1024 * 1024, + ) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_gpu_cache_evicts_only_released_entries(self) -> None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 16 * 1024 * 1024)], 8 * 1024 * 1024 + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + rows, dimension = 8192, 320 + payload = np.ones((rows, dimension), dtype=" None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 32 * 1024 * 1024)], 16 * 1024 * 1024 + ) + try: + cache = GpuTensorCache(pool, allow_eviction=False) + query = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 64 * 1024 * 1024)], 32 * 1024 * 1024 + ) + try: + generator = np.random.default_rng(7) + query = generator.standard_normal((44, 320)).astype("= 2, + "two CUDA devices are unavailable", + ) + def test_resident_engine_scores_shards_on_multiple_gpus(self) -> None: + pool = GpuResourcePool( + [ + GpuArenaSpec("cuda:0", 32 * 1024 * 1024), + GpuArenaSpec("cuda:1", 32 * 1024 * 1024), + ], + 16 * 1024 * 1024, + ) + try: + cache = GpuTensorCache(pool, allow_eviction=False) + identity = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" bytes: if directory_fd >= 0: os.close(directory_fd) - def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: + def key(self, request: protocol.ExternalTensorRequest) -> tuple[object, ...]: root_fd = self.root_fds.get(request.model_contract_id) if root_fd is None: raise protocol.SidecarError( @@ -222,9 +238,6 @@ def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: "model contract has no configured tensor cache root", ) digest = self._digest(request) - expected_bytes = protocol.checked_tensor_bytes( - request.rows, request.dimension, request.dtype - ) key = ( request.model_contract_id, digest, @@ -232,6 +245,15 @@ def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: request.dimension, request.dtype, ) + return key + + def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: + key = self.key(request) + digest = str(key[1]) + root_fd = self.root_fds[request.model_contract_id] + expected_bytes = protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) cached = self.cache.get(key) if cached is not None: return ResolvedPayload(cached, True) @@ -393,6 +415,190 @@ def score( ) +class ResidentTorchTileMaxsimEngine: + """Score tensors already owned by one or more process GPU arenas.""" + + def __init__( + self, + pool: GpuResourcePool, + max_workspace_bytes: int, + allow_tf32: bool, + max_cuda_inflight: int, + ) -> None: + self.pool = pool + self.device = pool.primary_device + self.max_workspace_bytes = max_workspace_bytes + self.compute_slots = threading.BoundedSemaphore(max_cuda_inflight) + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.backends.cudnn.allow_tf32 = allow_tf32 + with torch.inference_mode(): + for arena in pool.arenas: + left = torch.zeros((1, 1), dtype=torch.float32, device=arena.device) + _ = left @ left + for arena in pool.arenas: + torch.cuda.synchronize(arena.device) + + @staticmethod + def _cpu_tensor( + payload: bytes, rows: int, dimension: int, dtype: int + ) -> torch.Tensor: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + return torch.frombuffer(bytearray(payload), dtype=scalar_dtype).reshape( + rows, dimension + ) + + def _groups( + self, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, GpuTensorHandle]], + ) -> Iterable[list[tuple[int, GpuTensorHandle]]]: + scalar_bytes = 4 if dtype == protocol.DTYPE_F32 else 2 + query_bytes = query_rows * dimension * scalar_bytes + group: list[tuple[int, GpuTensorHandle]] = [] + group_rows = 0 + for document in documents: + rows = document[1].rows + next_rows = group_rows + rows + # The resident document remains inside the arena. torch.cat makes + # one device-local contiguous scoring view; the other temporaries + # are the query and q-by-document-token similarity matrix. + required = ( + query_bytes + + next_rows * dimension * scalar_bytes + + query_rows * next_rows * scalar_bytes + ) + if required > self.max_workspace_bytes and group: + yield group + group = [] + group_rows = 0 + next_rows = rows + required = ( + query_bytes + + rows * dimension * scalar_bytes + + query_rows * rows * scalar_bytes + ) + if required > self.max_workspace_bytes: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one resident candidate exceeds the configured GPU workspace", + ) + group.append(document) + group_rows = next_rows + if group: + yield group + + def score( + self, + query_payload: bytes, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, GpuTensorHandle]], + deadline: float, + cancelled: Callable[[], bool], + ) -> tuple[list[tuple[int, float]], float, float]: + if not documents: + return [], 0.0, 0.0 + query_cpu = self._cpu_tensor(query_payload, query_rows, dimension, dtype) + by_device: dict[str, list[tuple[int, GpuTensorHandle]]] = {} + for document in documents: + handle = document[1] + if handle.dimension != dimension or handle.dtype != dtype: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resident tensor contract disagrees with the query", + ) + by_device.setdefault(str(handle.device), []).append(document) + + queue_started = time.monotonic() + remaining = deadline - time.monotonic() + if remaining <= 0 or not self.compute_slots.acquire(timeout=remaining): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired while waiting for CUDA capacity", + ) + queue_ms = (time.monotonic() - queue_started) * 1000.0 + compute_started = time.monotonic() + pending: list[tuple[list[tuple[int, GpuTensorHandle]], torch.Tensor]] = [] + try: + with torch.inference_mode(): + for arena in self.pool.arenas: + device_documents = by_device.get(str(arena.device), []) + if not device_documents: + continue + if cancelled(): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request peer disconnected", + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + query_device = query_cpu.to(arena.device) + if dtype == protocol.DTYPE_F16 and ragged_tilemaxsim_fp16: + offsets = torch.tensor( + [ + handle.offset_bytes // 2 + for _, handle in device_documents + ], + dtype=torch.int64, + device=arena.device, + ) + rows = torch.tensor( + [handle.rows for _, handle in device_documents], + dtype=torch.int32, + device=arena.device, + ) + assert arena.storage is not None + device_scores = ragged_tilemaxsim_fp16( + query_device, + arena.storage.view(torch.float16), + offsets, + rows, + max(handle.rows for _, handle in device_documents), + ) + pending.append((device_documents, device_scores)) + continue + for group in self._groups( + query_rows, dimension, dtype, device_documents + ): + document_device = torch.cat( + [handle.tensor() for _, handle in group] + ) + similarities = query_device @ document_device.transpose(0, 1) + scores = [] + offset = 0 + for _, handle in group: + scores.append( + similarities[:, offset : offset + handle.rows] + .amax(dim=1) + .sum(dtype=torch.float32) + ) + offset += handle.rows + pending.append((group, torch.stack(scores))) + + results: list[tuple[int, float]] = [] + for group, device_scores in pending: + host_scores = device_scores.to(device="cpu", dtype=torch.float32) + for (candidate_id, _), score in zip( + group, host_scores.tolist(), strict=True + ): + if not math.isfinite(score): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "TileMaxSim result is non-finite", + ) + results.append((candidate_id, score)) + for arena in self.pool.arenas: + torch.cuda.synchronize(arena.device) + finally: + self.compute_slots.release() + return results, queue_ms, (time.monotonic() - compute_started) * 1000.0 + + class JsonMetrics: def __init__(self) -> None: self.lock = threading.Lock() @@ -410,12 +616,22 @@ def __init__( engine: TorchTileMaxsimEngine, request_timeout_ms: int, metrics: JsonMetrics, + gpu_cache: GpuTensorCache | None = None, + resident_engine: ResidentTorchTileMaxsimEngine | None = None, + pin_gpu_entries: bool = False, ) -> None: self.limits = limits self.resolver = resolver self.engine = engine self.request_timeout_seconds = request_timeout_ms / 1000.0 self.metrics = metrics + self.gpu_cache = gpu_cache + self.resident_engine = resident_engine + self.pin_gpu_entries = pin_gpu_entries + if (gpu_cache is None) != (resident_engine is None): + raise ValueError( + "GPU cache and resident engine must be configured together" + ) @staticmethod def _peer_disconnected(connection: socket.socket) -> bool: @@ -453,6 +669,7 @@ def process_frame( version = protocol.VERSION started = time.monotonic() metrics: dict[str, object] = {"event": "tilemaxsim_request"} + resident_documents: list[tuple[int, GpuTensorHandle]] = [] if peer_credentials is not None: metrics["peer_pid"], metrics["peer_uid"], metrics["peer_gid"] = ( peer_credentials @@ -480,7 +697,42 @@ def process_frame( ) resolve_started = time.monotonic() cache_hits = 0 + gpu_cache_hits = 0 + gpu_cache_misses = 0 + gpu_chunks = 0 + resident_results: list[tuple[int, float]] = [] + resident_queue_ms = 0.0 + resident_compute_ms = 0.0 + document_tokens = 0 documents: list[tuple[int, int, bytes]] = [] + + def flush_resident_documents() -> None: + nonlocal gpu_chunks, resident_queue_ms, resident_compute_ms + if not resident_documents: + return + assert self.resident_engine is not None + try: + batch_results, batch_queue_ms, batch_compute_ms = ( + self.resident_engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + resident_documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + ) + resident_results.extend(batch_results) + resident_queue_ms += batch_queue_ms + resident_compute_ms += batch_compute_ms + gpu_chunks += 1 + finally: + assert self.gpu_cache is not None + for _, resident_handle in resident_documents: + self.gpu_cache.release(resident_handle) + resident_documents.clear() + if isinstance(request, protocol.InlineTensorRequest): for candidate in request.candidates: validate_finite_payload( @@ -493,38 +745,99 @@ def process_frame( (candidate.candidate_id, candidate.rows, candidate.payload) for candidate in request.candidates ] + document_tokens = sum( + candidate.rows for candidate in request.candidates + ) metrics["source"] = "inline" else: metrics["source"] = "content_addressed" for candidate in request.candidates: + document_tokens += candidate.descriptor.rows if time.monotonic() >= deadline: raise protocol.SidecarError( protocol.STATUS_COMPUTE_ERROR, "request deadline expired during tensor resolution", ) - resolved = self.resolver.resolve(candidate.descriptor) - cache_hits += int(resolved.cache_hit) - documents.append( - ( - candidate.candidate_id, - candidate.descriptor.rows, - resolved.payload, + if self.gpu_cache is None: + resolved = self.resolver.resolve(candidate.descriptor) + cache_hits += int(resolved.cache_hit) + documents.append( + ( + candidate.candidate_id, + candidate.descriptor.rows, + resolved.payload, + ) ) - ) + else: + key = self.resolver.key(candidate.descriptor) + loaded_payload: bytes | None = None + + def load_payload() -> bytes: + nonlocal cache_hits, loaded_payload + if loaded_payload is None: + resolved = self.resolver.resolve(candidate.descriptor) + cache_hits += int(resolved.cache_hit) + loaded_payload = resolved.payload + return loaded_payload + + while True: + try: + handle, gpu_hit = self.gpu_cache.acquire( + key, + candidate.descriptor.rows, + candidate.descriptor.dimension, + candidate.descriptor.dtype, + load_payload, + pin=self.pin_gpu_entries, + ) + break + except protocol.SidecarError as error: + if ( + error.status != protocol.STATUS_RESOURCE_LIMIT + or not resident_documents + ): + raise + # A request may be larger than the configured + # GPU cache. Score and release the current + # working set, then admit the remaining + # candidates through the same bounded arenas. + flush_resident_documents() + gpu_cache_hits += int(gpu_hit) + gpu_cache_misses += int(not gpu_hit) + resident_documents.append((candidate.candidate_id, handle)) + if self.gpu_cache is not None: + flush_resident_documents() metrics["cache_hits"] = cache_hits + metrics["host_cache_hits"] = cache_hits + metrics["gpu_cache_hits"] = gpu_cache_hits + metrics["gpu_cache_misses"] = gpu_cache_misses + metrics["gpu_chunks"] = gpu_chunks metrics["resolve_ms"] = round( - (time.monotonic() - resolve_started) * 1000.0, 3 - ) - metrics["document_tokens"] = sum(rows for _, rows, _ in documents) - results, queue_ms, compute_ms = self.engine.score( - request.query_payload, - request.query_rows, - request.dimension, - request.dtype, - documents, - deadline, - lambda: self._peer_disconnected(connection), + max( + 0.0, + (time.monotonic() - resolve_started) * 1000.0 + - resident_queue_ms + - resident_compute_ms, + ), + 3, ) + metrics["document_tokens"] = document_tokens + if self.gpu_cache is not None and isinstance( + request, protocol.ParsedExternalTensorRequest + ): + results = resident_results + queue_ms = resident_queue_ms + compute_ms = resident_compute_ms + else: + results, queue_ms, compute_ms = self.engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + documents, + deadline, + lambda: self._peer_disconnected(connection), + ) metrics["queue_ms"] = round(queue_ms, 3) metrics["compute_ms"] = round(compute_ms, 3) metrics["status"] = "ok" @@ -535,8 +848,6 @@ def process_frame( request_id, error.status, str(error), version ) except torch.OutOfMemoryError: - if self.engine.device.type == "cuda": - torch.cuda.empty_cache() metrics.update(status="error", error_class="CudaOutOfMemory") return protocol.error_response( request_id, @@ -553,6 +864,9 @@ def process_frame( version, ) finally: + if self.gpu_cache is not None: + for _, handle in resident_documents: + self.gpu_cache.release(handle) metrics["total_ms"] = round((time.monotonic() - started) * 1000.0, 3) self.metrics.emit(metrics) @@ -655,14 +969,15 @@ def handle(connection: socket.socket) -> None: listener.listen(backlog) listener.settimeout(0.25) bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino - service.metrics.emit( - { - "event": "tilemaxsim_ready", - "device": str(service.engine.device), - "max_inflight": max_inflight, - "socket": os.fspath(socket_path), - } - ) + ready: dict[str, object] = { + "event": "tilemaxsim_ready", + "device": str(service.engine.device), + "max_inflight": max_inflight, + "socket": os.fspath(socket_path), + } + if service.gpu_cache is not None: + ready["gpu_cache"] = service.gpu_cache.status() + service.metrics.emit(ready) accepted = 0 while not stop.is_set(): if not slots.acquire(timeout=0.25): @@ -714,6 +1029,20 @@ def nonnegative_int(value: str) -> int: return parsed +def memory_gb(value: str) -> int: + try: + return parse_memory_gb(value) + except ValueError as error: + raise argparse.ArgumentTypeError(str(error)) from error + + +def gpu_memory_gb(value: str) -> GpuArenaSpec: + try: + return parse_gpu_memory_gb(value) + except (ValueError, RuntimeError) as error: + raise argparse.ArgumentTypeError(str(error)) from error + + def contract_roots( values: list[str], parser: argparse.ArgumentParser ) -> dict[str, Path]: @@ -731,12 +1060,133 @@ def contract_roots( return roots +def contract_manifests( + values: list[str], parser: argparse.ArgumentParser +) -> list[tuple[str, Path]]: + manifests = [] + for value in values: + if "=" not in value: + parser.error("--resident-manifest must be MODEL_CONTRACT_ID=/absolute/path") + contract, raw_path = value.split("=", 1) + path = Path(raw_path) + if not contract or not path.is_absolute(): + parser.error( + "--resident-manifest must contain a nonempty ID and absolute path" + ) + manifests.append((contract, path)) + return manifests + + +def prewarm_resident_cache( + manifests: list[tuple[str, Path]], + resolver: ContentAddressedResolver, + gpu_cache: GpuTensorCache, + metrics: JsonMetrics, +) -> None: + completed = 0 + loaded_bytes = 0 + started = time.monotonic() + for contract, path in manifests: + with path.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path}:{line_number}: invalid JSON") from error + if not isinstance(record, dict): + raise ValueError(f"{path}:{line_number}: record must be an object") + dtype_name = record.get("tensor_dtype") + if dtype_name == "float16": + dtype = protocol.DTYPE_F16 + elif dtype_name == "float32": + dtype = protocol.DTYPE_F32 + else: + raise ValueError(f"{path}:{line_number}: unsupported tensor_dtype") + tensor_ref = record.get("tensor_ref") + checksum = record.get("tensor_checksum") + rows = record.get("tensor_rows") + dimension = record.get("tensor_dim") + if not isinstance(tensor_ref, str) or not tensor_ref: + raise ValueError(f"{path}:{line_number}: invalid tensor_ref") + if not isinstance(checksum, str) or not checksum: + raise ValueError(f"{path}:{line_number}: invalid tensor_checksum") + if not isinstance(rows, int) or rows <= 0: + raise ValueError(f"{path}:{line_number}: invalid tensor_rows") + if not isinstance(dimension, int) or dimension <= 0: + raise ValueError(f"{path}:{line_number}: invalid tensor_dim") + request = protocol.ExternalTensorRequest( + contract, tensor_ref, rows, dimension, dtype, checksum + ) + expected_bytes = protocol.checked_tensor_bytes(rows, dimension, dtype) + declared_bytes = record.get("canonical_bytes") + if declared_bytes is not None and declared_bytes != expected_bytes: + raise ValueError( + f"{path}:{line_number}: canonical_bytes disagrees with shape" + ) + key = resolver.key(request) + handle, _ = gpu_cache.acquire( + key, + rows, + dimension, + dtype, + lambda request=request: resolver.resolve(request).payload, + pin=True, + ) + gpu_cache.release(handle) + completed += 1 + loaded_bytes += expected_bytes + if completed % 1000 == 0: + metrics.emit( + { + "event": "tilemaxsim_prewarm_progress", + "entries": completed, + "logical_bytes": loaded_bytes, + } + ) + if completed == 0: + raise ValueError("resident manifests contain no tensor descriptors") + metrics.emit( + { + "event": "tilemaxsim_prewarm_complete", + "entries": completed, + "logical_bytes": loaded_bytes, + "elapsed_ms": round((time.monotonic() - started) * 1000.0, 3), + "gpu_cache": gpu_cache.status(), + } + ) + + def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--socket", required=True, type=Path) parser.add_argument("--socket-mode", type=parse_mode, default=0o600) - parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--gpu-cache-mode", + choices=("lru", "resident"), + default="lru", + help="use evictable GPU arenas or pin a full descriptor manifest", + ) + parser.add_argument( + "--gpu-memory-gb", + action="append", + type=gpu_memory_gb, + default=[], + metavar="GPU=GB", + help="repeatable strict allocation, for example 1=20; enables TileMaxSim", + ) + parser.add_argument( + "--gpu-workspace-gb", + type=memory_gb, + default=2 * 1024**3, + help="per-GPU portion of the configured GB reserved for scoring temporaries", + ) parser.add_argument("--contract-root", action="append", default=[]) + parser.add_argument( + "--resident-manifest", + action="append", + default=[], + metavar="MODEL_CONTRACT_ID=/ABSOLUTE/PATH", + ) parser.add_argument( "--max-request-bytes", type=positive_int, default=64 * 1024 * 1024 ) @@ -746,10 +1196,10 @@ def main() -> None: ) parser.add_argument("--max-candidates", type=positive_int, default=65_536) parser.add_argument( - "--max-device-bytes", type=positive_int, default=8 * 1024 * 1024 * 1024 - ) - parser.add_argument( - "--cache-bytes", type=nonnegative_int, default=8 * 1024 * 1024 * 1024 + "--host-cache-gb", + type=memory_gb, + default=8 * 1024**3, + help="decoded host-memory tensor cache size in GB", ) parser.add_argument("--request-timeout-ms", type=positive_int, default=2000) parser.add_argument("--max-inflight", type=positive_int, default=8) @@ -760,23 +1210,52 @@ def main() -> None: args = parser.parse_args() roots = contract_roots(args.contract_root, parser) + manifests = contract_manifests(args.resident_manifest, parser) + if not args.gpu_memory_gb: + parser.error( + "TileMaxSim is disabled until at least one --gpu-memory-gb GPU=GB is configured" + ) + if args.gpu_cache_mode == "resident" and not manifests: + parser.error( + "--gpu-cache-mode resident requires at least one --resident-manifest" + ) + if args.gpu_cache_mode == "lru" and manifests: + parser.error("--resident-manifest is valid only in resident mode") limits = protocol.Limits( max_request_bytes=args.max_request_bytes, max_batch_tokens=args.max_batch_tokens, max_tensor_bytes=args.max_tensor_bytes, max_candidates=args.max_candidates, ) - resolver = ContentAddressedResolver(roots, args.cache_bytes) + resolver = ContentAddressedResolver(roots, args.host_cache_gb) metrics = JsonMetrics() + pool: GpuResourcePool | None = None try: + pool = GpuResourcePool(args.gpu_memory_gb, args.gpu_workspace_gb) engine = TorchTileMaxsimEngine( - args.device, - args.max_device_bytes, + str(pool.primary_device), + args.gpu_workspace_gb, + args.allow_tf32, + args.max_cuda_inflight, + ) + gpu_cache = GpuTensorCache(pool, allow_eviction=args.gpu_cache_mode == "lru") + resident_engine = ResidentTorchTileMaxsimEngine( + pool, + args.gpu_workspace_gb, args.allow_tf32, args.max_cuda_inflight, ) + if args.gpu_cache_mode == "resident": + prewarm_resident_cache(manifests, resolver, gpu_cache, metrics) service = TileMaxsimService( - limits, resolver, engine, args.request_timeout_ms, metrics + limits, + resolver, + engine, + args.request_timeout_ms, + metrics, + gpu_cache, + resident_engine, + pin_gpu_entries=args.gpu_cache_mode == "resident", ) stop = threading.Event() @@ -795,6 +1274,8 @@ def request_stop(_signum: int, _frame: object) -> None: args.once, ) finally: + if pool is not None: + pool.close() resolver.close() diff --git a/services/tilemaxsim_gpu_cache.py b/services/tilemaxsim_gpu_cache.py new file mode 100644 index 00000000..43d66d97 --- /dev/null +++ b/services/tilemaxsim_gpu_cache.py @@ -0,0 +1,430 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Process-owned GPU arenas and a bounded tensor cache for TileMaxSim.""" + +from __future__ import annotations + +import re +import threading +from collections import OrderedDict +from dataclasses import dataclass +from decimal import Decimal, InvalidOperation +from typing import Callable + +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol + + +_GPU_MEMORY_GB = re.compile(r"^(?:cuda:)?([0-9]+)=([0-9]+(?:\.[0-9]+)?)$") +_MEMORY_GB = re.compile(r"^[0-9]+(?:\.[0-9]+)?$") +GIB = 1024**3 + + +@dataclass(frozen=True) +class GpuArenaSpec: + device: str + total_bytes: int + + +def parse_gpu_memory_gb(value: str) -> GpuArenaSpec: + """Parse the public ``GPU=GB`` configuration into an internal byte budget.""" + + match = _GPU_MEMORY_GB.fullmatch(value.strip()) + if match is None: + raise ValueError( + "GPU memory must be GPU=GB, for example 1=20; byte suffixes are not accepted" + ) + raw_index, raw_gb = match.groups() + try: + total_bytes = int(Decimal(raw_gb) * GIB) + except InvalidOperation as error: + raise ValueError("GPU memory GB value is invalid") from error + if total_bytes <= 0: + raise ValueError("GPU memory GB value must be positive") + return GpuArenaSpec(f"cuda:{int(raw_index)}", total_bytes) + + +def parse_memory_gb(value: str) -> int: + if _MEMORY_GB.fullmatch(value.strip()) is None: + raise ValueError("memory size must be a positive number of GB") + try: + total_bytes = int(Decimal(value.strip()) * GIB) + except InvalidOperation as error: + raise ValueError("memory size must be a positive number of GB") from error + if total_bytes <= 0: + raise ValueError("memory size must be a positive number of GB") + return total_bytes + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +class FreeExtentAllocator: + """Best-fit allocator for contiguous slices of a preallocated byte arena.""" + + def __init__(self, capacity: int, alignment: int = 256) -> None: + if capacity <= 0: + raise ValueError("arena capacity must be positive") + self.capacity = capacity + self.alignment = alignment + self.extents: list[tuple[int, int]] = [(0, capacity)] + + @property + def free_bytes(self) -> int: + return sum(length for _, length in self.extents) + + @property + def largest_free_extent(self) -> int: + return max((length for _, length in self.extents), default=0) + + def allocation_bytes(self, payload_bytes: int) -> int: + return _align_up(payload_bytes, self.alignment) + + def allocate(self, payload_bytes: int) -> tuple[int, int] | None: + required = self.allocation_bytes(payload_bytes) + choices = [ + (length, index, start) + for index, (start, length) in enumerate(self.extents) + if length >= required + ] + if not choices: + return None + length, index, start = min(choices) + if length == required: + self.extents.pop(index) + else: + self.extents[index] = (start + required, length - required) + return start, required + + def release(self, start: int, length: int) -> None: + if start < 0 or length <= 0 or start + length > self.capacity: + raise ValueError("released extent is outside the arena") + self.extents.append((start, length)) + self.extents.sort() + merged: list[tuple[int, int]] = [] + for extent_start, extent_length in self.extents: + if merged and merged[-1][0] + merged[-1][1] == extent_start: + previous_start, previous_length = merged[-1] + merged[-1] = (previous_start, previous_length + extent_length) + else: + merged.append((extent_start, extent_length)) + self.extents = merged + + +class GpuArena: + """A CUDA byte buffer acquired atomically during process startup.""" + + def __init__(self, spec: GpuArenaSpec, workspace_bytes: int) -> None: + self.device = torch.device(spec.device) + if not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested but torch.cuda.is_available() is false" + ) + if self.device.index is None or self.device.index >= torch.cuda.device_count(): + raise RuntimeError(f"configured CUDA device is unavailable: {spec.device}") + if workspace_bytes <= 0 or workspace_bytes >= spec.total_bytes: + raise RuntimeError( + f"{spec.device} allocation must exceed its TileMaxSim workspace" + ) + self.total_bytes = spec.total_bytes + self.workspace_bytes = workspace_bytes + self.capacity = (spec.total_bytes - workspace_bytes) // 256 * 256 + self.reserved_workspace_bytes = spec.total_bytes - self.capacity + if self.capacity <= 0: + raise RuntimeError(f"{spec.device} has no aligned tensor-cache capacity") + self.storage: torch.Tensor | None = None + self.allocator = FreeExtentAllocator(self.capacity) + + with torch.cuda.device(self.device): + free_bytes, device_bytes = torch.cuda.mem_get_info(self.device) + self.device_bytes = device_bytes + if free_bytes < spec.total_bytes: + raise RuntimeError( + f"cannot acquire {spec.total_bytes} bytes on {spec.device}: " + f"only {free_bytes} bytes are free" + ) + try: + self.storage = torch.empty( + self.capacity, dtype=torch.uint8, device=self.device + ) + # Reserve the remaining configured budget in PyTorch's CUDA + # allocator. Releasing this temporary tensor leaves the block + # in the process-owned caching allocator for TileMaxSim + # workspaces instead of returning it to another process. + workspace = torch.empty( + self.reserved_workspace_bytes, + dtype=torch.uint8, + device=self.device, + ) + torch.cuda.synchronize(self.device) + del workspace + except Exception: + self.storage = None + torch.cuda.empty_cache() + raise + + def tensor( + self, offset: int, payload_bytes: int, rows: int, dimension: int, dtype: int + ) -> torch.Tensor: + if self.storage is None: + raise RuntimeError("GPU arena is closed") + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + return ( + self.storage.narrow(0, offset, payload_bytes) + .view(scalar_dtype) + .reshape(rows, dimension) + ) + + def copy_from_host(self, offset: int, payload: bytes) -> None: + if self.storage is None: + raise RuntimeError("GPU arena is closed") + source = torch.frombuffer(bytearray(payload), dtype=torch.uint8) + self.storage.narrow(0, offset, len(payload)).copy_(source) + + def status(self) -> dict[str, object]: + return { + "device": str(self.device), + "allocated_gb": round(self.total_bytes / GIB, 3), + "tensor_capacity_gb": round(self.capacity / GIB, 3), + "workspace_gb": round(self.workspace_bytes / GIB, 3), + "allocated_bytes": self.total_bytes, + "tensor_capacity_bytes": self.capacity, + "workspace_bytes": self.workspace_bytes, + "tensor_free_bytes": self.allocator.free_bytes, + "largest_free_extent_bytes": self.allocator.largest_free_extent, + } + + def close(self) -> None: + if self.storage is None: + return + with torch.cuda.device(self.device): + self.storage = None + torch.cuda.empty_cache() + + +class GpuResourcePool: + """Own all configured CUDA allocations or fail without a partial pool.""" + + def __init__(self, specs: list[GpuArenaSpec], workspace_bytes: int) -> None: + if not specs: + raise RuntimeError("at least one GPU allocation is required") + devices = [spec.device for spec in specs] + if len(devices) != len(set(devices)): + raise RuntimeError("each CUDA device may be configured only once") + self.arenas: list[GpuArena] = [] + try: + for spec in specs: + self.arenas.append(GpuArena(spec, workspace_bytes)) + except Exception: + self.close() + raise + + @property + def primary_device(self) -> torch.device: + return self.arenas[0].device + + def status(self) -> list[dict[str, object]]: + return [arena.status() for arena in self.arenas] + + def close(self) -> None: + for arena in self.arenas: + arena.close() + self.arenas.clear() + + +@dataclass +class _GpuCacheEntry: + key: tuple[object, ...] + arena: GpuArena + offset: int + allocated_bytes: int + payload_bytes: int + rows: int + dimension: int + dtype: int + references: int = 0 + pinned: bool = False + + +@dataclass(frozen=True) +class GpuTensorHandle: + entry: _GpuCacheEntry + + @property + def arena(self) -> GpuArena: + return self.entry.arena + + @property + def device(self) -> torch.device: + return self.entry.arena.device + + @property + def rows(self) -> int: + return self.entry.rows + + @property + def dimension(self) -> int: + return self.entry.dimension + + @property + def dtype(self) -> int: + return self.entry.dtype + + @property + def payload_bytes(self) -> int: + return self.entry.payload_bytes + + @property + def offset_bytes(self) -> int: + return self.entry.offset + + def tensor(self) -> torch.Tensor: + return self.entry.arena.tensor( + self.entry.offset, + self.entry.payload_bytes, + self.entry.rows, + self.entry.dimension, + self.entry.dtype, + ) + + +class GpuTensorCache: + """Thread-safe LRU directory over one or more process-owned GPU arenas.""" + + def __init__(self, pool: GpuResourcePool, allow_eviction: bool) -> None: + self.pool = pool + self.allow_eviction = allow_eviction + self.entries: OrderedDict[tuple[object, ...], _GpuCacheEntry] = OrderedDict() + self.lock = threading.Lock() + self.hits = 0 + self.misses = 0 + self.evictions = 0 + self.loaded_bytes = 0 + + def _find_arena(self, payload_bytes: int) -> GpuArena | None: + candidates = [ + arena + for arena in self.pool.arenas + if arena.allocator.largest_free_extent + >= arena.allocator.allocation_bytes(payload_bytes) + ] + if not candidates: + return None + return max(candidates, key=lambda arena: arena.allocator.free_bytes) + + def _evict_one(self) -> bool: + for key, entry in self.entries.items(): + if entry.references or entry.pinned: + continue + self.entries.pop(key) + entry.arena.allocator.release(entry.offset, entry.allocated_bytes) + self.evictions += 1 + return True + return False + + def _allocate(self, payload_bytes: int) -> tuple[GpuArena, int, int]: + arena = self._find_arena(payload_bytes) + while arena is None and self.allow_eviction and self._evict_one(): + arena = self._find_arena(payload_bytes) + if arena is None: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "configured GPU tensor arenas have insufficient contiguous capacity", + ) + allocated = arena.allocator.allocate(payload_bytes) + assert allocated is not None + return arena, allocated[0], allocated[1] + + def acquire( + self, + key: tuple[object, ...], + rows: int, + dimension: int, + dtype: int, + loader: Callable[[], bytes], + *, + pin: bool = False, + ) -> tuple[GpuTensorHandle, bool]: + with self.lock: + cached = self.entries.get(key) + if cached is not None: + cached.references += 1 + cached.pinned = cached.pinned or pin + self.entries.move_to_end(key) + self.hits += 1 + return GpuTensorHandle(cached), True + + payload = loader() + expected = protocol.checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match its shape", + ) + + with self.lock: + cached = self.entries.get(key) + if cached is not None: + cached.references += 1 + cached.pinned = cached.pinned or pin + self.entries.move_to_end(key) + self.hits += 1 + return GpuTensorHandle(cached), True + arena, offset, allocated_bytes = self._allocate(len(payload)) + try: + arena.copy_from_host(offset, payload) + except Exception: + arena.allocator.release(offset, allocated_bytes) + raise + entry = _GpuCacheEntry( + key, + arena, + offset, + allocated_bytes, + len(payload), + rows, + dimension, + dtype, + references=1, + pinned=pin, + ) + self.entries[key] = entry + self.misses += 1 + self.loaded_bytes += len(payload) + return GpuTensorHandle(entry), False + + def release(self, handle: GpuTensorHandle) -> None: + with self.lock: + entry = self.entries.get(handle.entry.key) + if entry is not handle.entry or entry.references <= 0: + raise RuntimeError("GPU tensor handle was released more than once") + entry.references -= 1 + + def status(self) -> dict[str, object]: + with self.lock: + return { + "entries": len(self.entries), + "pinned_entries": sum(entry.pinned for entry in self.entries.values()), + "active_references": sum( + entry.references for entry in self.entries.values() + ), + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "loaded_bytes": self.loaded_bytes, + "arenas": self.pool.status(), + } diff --git a/services/tilemaxsim_triton.py b/services/tilemaxsim_triton.py new file mode 100644 index 00000000..5571ffad --- /dev/null +++ b/services/tilemaxsim_triton.py @@ -0,0 +1,128 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Fused ragged FP16 TileMaxSim over a process-owned GPU tensor arena.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _ragged_tilemaxsim_fp16_kernel( + query, + documents, + document_offsets, + document_rows, + scores, + query_rows, + dimension: tl.constexpr, + max_document_rows: tl.constexpr, + block_query: tl.constexpr, + block_document: tl.constexpr, + block_dimension: tl.constexpr, +): + document_index = tl.program_id(0).to(tl.int64) + query_block_index = tl.program_id(1).to(tl.int64) + query_indices = (query_block_index * block_query + tl.arange(0, block_query)).to( + tl.int64 + ) + query_mask = query_indices < query_rows + document_base = tl.load(document_offsets + document_index).to(tl.int64) + valid_document_rows = tl.load(document_rows + document_index) + running_max = tl.full([block_query], value=float("-inf"), dtype=tl.float32) + + for document_start in range(0, max_document_rows, block_document): + document_indices = (document_start + tl.arange(0, block_document)).to(tl.int64) + document_mask = document_indices < valid_document_rows + similarities = tl.zeros([block_query, block_document], dtype=tl.float32) + for dimension_start in range(0, dimension, block_dimension): + dimension_indices = (dimension_start + tl.arange(0, block_dimension)).to( + tl.int64 + ) + dimension_mask = dimension_indices < dimension + query_pointers = ( + query + query_indices[:, None] * dimension + dimension_indices[None, :] + ) + query_tile = tl.load( + query_pointers, + mask=query_mask[:, None] & dimension_mask[None, :], + other=0.0, + ) + document_pointers = ( + documents + + document_base + + document_indices[:, None] * dimension + + dimension_indices[None, :] + ) + document_tile = tl.load( + document_pointers, + mask=document_mask[:, None] & dimension_mask[None, :], + other=0.0, + ) + similarities += tl.dot(query_tile, tl.trans(document_tile)) + similarities = tl.where(document_mask[None, :], similarities, float("-inf")) + running_max = tl.maximum(running_max, tl.max(similarities, axis=1)) + + running_max = tl.where(query_mask, running_max, 0.0) + tl.atomic_add(scores + document_index, tl.sum(running_max, axis=0)) + + +def ragged_tilemaxsim_fp16( + query: torch.Tensor, + document_arena: torch.Tensor, + document_offsets: torch.Tensor, + document_rows: torch.Tensor, + maximum_document_rows: int, +) -> torch.Tensor: + if query.device.type != "cuda" or document_arena.device != query.device: + raise ValueError("query and document arena must be on the same CUDA device") + if query.dtype != torch.float16 or document_arena.dtype != torch.float16: + raise ValueError("ragged TileMaxSim currently requires FP16 tensors") + if query.ndim != 2 or document_arena.ndim != 1: + raise ValueError("invalid query or document arena shape") + if document_offsets.dtype != torch.int64 or document_rows.dtype != torch.int32: + raise ValueError("invalid ragged TileMaxSim metadata dtype") + if document_offsets.shape != document_rows.shape: + raise ValueError("document offsets and rows must have the same shape") + count = document_offsets.numel() + if count == 0: + return torch.empty(0, dtype=torch.float32, device=query.device) + if maximum_document_rows <= 0: + raise ValueError("maximum document rows must be positive") + block_query = 32 + block_document = 32 + block_dimension = 128 + padded_document_rows = ( + (maximum_document_rows + block_document - 1) // block_document * block_document + ) + query_blocks = (query.shape[0] + block_query - 1) // block_query + scores = torch.zeros(count, dtype=torch.float32, device=query.device) + with torch.cuda.device(query.device): + _ragged_tilemaxsim_fp16_kernel[(count, query_blocks)]( + query, + document_arena, + document_offsets, + document_rows, + scores, + query.shape[0], + dimension=query.shape[1], + max_document_rows=padded_document_rows, + block_query=block_query, + block_document=block_document, + block_dimension=block_dimension, + ) + return scores diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 26cde468..52ad3fef 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -92,6 +92,8 @@ static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut pgrx::pg_sys::config_generic = static VCHORDRQ_MAXSIM_CANDIDATE_LIMIT: GucSetting = GucSetting::::new(-1); +static VCHORDRQ_MAXSIM_PROFILE: GucSetting = GucSetting::::new(false); + static VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS: GucSetting = GucSetting::::new(32); static VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS: GucSetting = GucSetting::::new(256); @@ -191,6 +193,14 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_bool_guc( + c"vchordrq.maxsim_profile", + c"Emit one client NOTICE with MaxSim phase timings and counters.", + c"Disabled by default; emitted profiles contain no tensor references or application IDs.", + &VCHORDRQ_MAXSIM_PROFILE, + GucContext::Userset, + GucFlags::default(), + ); GucRegistry::define_int_guc( c"vchordrq.maxsim_planner_query_tokens", c"Expected MaxSim query-token count used by the planner.", @@ -586,6 +596,10 @@ pub fn vchordrq_maxsim_candidate_limit() -> Option { if value < 0 { None } else { Some(value as u32) } } +pub fn vchordrq_maxsim_profile() -> bool { + VCHORDRQ_MAXSIM_PROFILE.get() +} + pub fn vchordrq_maxsim_planner_query_tokens() -> u32 { VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS.get() as u32 } diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 238dbe2e..4f76c0db 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -13,12 +13,14 @@ // Copyright (c) 2025-2026 TensorChord Inc. mod candidate; +mod exact; mod external; mod gpu; +mod profile; mod rerank; mod search; -use self::candidate::{LegacyPageCandidateGenerator, PageCandidateGenerator}; +use self::candidate::{DensePageCandidateGenerator, PageCandidateGenerator}; use self::gpu::{GpuTileMaxsimBackend, UnixSocketTransport, report_gpu_fallback}; use self::rerank::{CpuExactMaxsimBackend, ExactMaxsimBackend, HeapArrayTensorSource}; use crate::index::fetcher::*; @@ -176,6 +178,7 @@ impl SearchBuilder for MaxsimBuilder { .map(|vector| RandomProject::project(vector.as_borrowed())) .collect::>(); Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); let (results, estimation_by_threshold) = match options.io_search { Io::Plain => maxsim_search::<_, Op>( index, @@ -208,6 +211,15 @@ impl SearchBuilder for MaxsimBuilder { make_h0_stream_prefetcher.clone(), ), }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { let sequence = Heap::from(results); @@ -312,6 +324,13 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); (accu_set, rough_set, estimation_by_threshold) })) } @@ -332,6 +351,7 @@ impl SearchBuilder for MaxsimBuilder { .map(|vector| RandomProject::project(vector.as_borrowed())) .collect::>(); Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); let (results, estimation_by_threshold) = match options.io_search { Io::Plain => maxsim_search::<_, Op>( index, @@ -364,6 +384,15 @@ impl SearchBuilder for MaxsimBuilder { make_h0_stream_prefetcher.clone(), ), }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { let sequence = Heap::from(results); @@ -468,6 +497,13 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); (accu_set, rough_set, estimation_by_threshold) })) } @@ -484,6 +520,7 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); let (results, estimation_by_threshold) = match options.io_search { Io::Plain => maxsim_search::<_, Op>( index, @@ -516,6 +553,15 @@ impl SearchBuilder for MaxsimBuilder { make_h0_stream_prefetcher.clone(), ), }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { let sequence = Heap::from(results); @@ -620,6 +666,13 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); (accu_set, rough_set, estimation_by_threshold) })) } @@ -636,6 +689,7 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); let (results, estimation_by_threshold) = match options.io_search { Io::Plain => maxsim_search::<_, Op>( index, @@ -668,6 +722,15 @@ impl SearchBuilder for MaxsimBuilder { make_h0_stream_prefetcher.clone(), ), }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { let sequence = Heap::from(results); @@ -772,6 +835,13 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); (accu_set, rough_set, estimation_by_threshold) })) } @@ -779,7 +849,7 @@ impl SearchBuilder for MaxsimBuilder { iter }; let mut candidates = - LegacyPageCandidateGenerator.generate(n, &mut token_searches, maxsim_candidate_limit); + DensePageCandidateGenerator.generate(n, &mut token_searches, maxsim_candidate_limit); drop(token_searches); let iter: Box> = match maxsim_backend { PostgresMaxsimBackend::CoarseOnly => Box::new(candidates.map(|candidate| { diff --git a/src/index/vchordrq/scanners/maxsim/candidate.rs b/src/index/vchordrq/scanners/maxsim/candidate.rs index 9aa3e22f..af9b45b0 100644 --- a/src/index/vchordrq/scanners/maxsim/candidate.rs +++ b/src/index/vchordrq/scanners/maxsim/candidate.rs @@ -12,11 +12,13 @@ // // Copyright (c) 2025-2026 TensorChord Inc. +use super::profile; use crate::index::fetcher::pointer_to_kv; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; -use std::collections::BinaryHeap; +use std::collections::hash_map::Entry; +use std::collections::{BinaryHeap, HashMap}; use std::num::NonZero; pub(super) type HeapKey = [u16; 3]; @@ -41,10 +43,10 @@ pub(super) trait PageCandidateGenerator { } #[derive(Default)] -pub(super) struct LegacyPageCandidateGenerator; +pub(super) struct DensePageCandidateGenerator; -impl PageCandidateGenerator for LegacyPageCandidateGenerator { - type Candidates = LegacyPageCandidates; +impl PageCandidateGenerator for DensePageCandidateGenerator { + type Candidates = PageCandidates; fn generate( &mut self, @@ -52,60 +54,120 @@ impl PageCandidateGenerator for LegacyPageCandidateGenerator { token_searches: &mut dyn Iterator, candidate_limit: usize, ) -> Self::Candidates { - let mut updates = Vec::new(); + let mut page_lookup = HashMap::new(); + let mut page_keys = Vec::new(); + let mut best_by_page_query = Vec::new(); let mut estimations = Vec::with_capacity(query_count); + let mut hit_updates = 0u64; + let mut page_token_updates = 0u64; for (query_id, (accu_set, rough_set, estimation_by_threshold)) in token_searches.enumerate() { - updates.reserve(accu_set.len() + rough_set.len()); + debug_assert!(query_id < query_count); + let hit_collect_timer = profile::ProfileTimer::start(); + hit_updates += (accu_set.len() + rough_set.len()) as u64; let is_empty = accu_set.is_empty() && rough_set.is_empty(); let mut estimation_by_scope = Distance::NEG_INFINITY; for (distance, payload) in accu_set { estimation_by_scope = std::cmp::max(estimation_by_scope, distance); let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + page_token_updates += u64::from(record_page_token_distance( + &mut page_lookup, + &mut page_keys, + &mut best_by_page_query, + query_count, + query_id, + key, + distance, + )); } for (distance, payload) in rough_set { let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + page_token_updates += u64::from(record_page_token_distance( + &mut page_lookup, + &mut page_keys, + &mut best_by_page_query, + query_count, + query_id, + key, + distance, + )); } estimations.push(if !is_empty { std::cmp::max(estimation_by_scope, estimation_by_threshold) } else { Distance::ZERO }); + let hit_collect_elapsed = hit_collect_timer.elapsed(); + profile::update(|profile| { + profile.hit_collect_us += profile::duration_us(hit_collect_elapsed); + }); } debug_assert_eq!(estimations.len(), query_count); - updates.sort_unstable_by_key(|&(key, ..)| key); - let inner = updates - .chunk_by(|(kl, ..), (kr, ..)| kl == kr) - .map(|chunk| { - let key = chunk[0].0; - let mut value = vec![None; query_count]; - for &(_, query_id, distance) in chunk { - let this = value[query_id].get_or_insert(Distance::INFINITY); - *this = std::cmp::min(*this, distance); - } + profile::update(|profile| { + profile.hit_updates += hit_updates; + profile.page_token_updates += page_token_updates; + }); + let page_aggregate_timer = profile::ProfileTimer::start(); + let mut page_order = (0..page_keys.len()).collect::>(); + page_order.sort_unstable_by_key(|&page_index| page_keys[page_index]); + let inner = page_order + .into_iter() + .map(|page_index| { + let key = page_keys[page_index]; + let start = page_index * query_count; + let values = &best_by_page_query[start..start + query_count]; let mut maxsim = 0.0f32; - for (query_id, distance) in value.into_iter().enumerate() { + for (query_id, distance) in values.iter().copied().enumerate() { let distance = distance.unwrap_or(estimations[query_id]); maxsim += distance.to_f32(); } (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) }) .collect::>(); - LegacyPageCandidates { + let page_aggregate_elapsed = page_aggregate_timer.elapsed(); + profile::update(|profile| { + profile.page_aggregate_us += profile::duration_us(page_aggregate_elapsed); + profile.aggregated_pages += page_keys.len() as u64; + }); + PageCandidates { inner, remaining: candidate_limit, } } } -pub(super) struct LegacyPageCandidates { +fn record_page_token_distance( + page_lookup: &mut HashMap, + page_keys: &mut Vec, + best_by_page_query: &mut Vec>, + query_count: usize, + query_id: usize, + key: HeapKey, + distance: Distance, +) -> bool { + let page_index = match page_lookup.entry(key) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + let page_index = page_keys.len(); + entry.insert(page_index); + page_keys.push(key); + best_by_page_query.resize(best_by_page_query.len() + query_count, None); + page_index + } + }; + let slot = &mut best_by_page_query[page_index * query_count + query_id]; + let first_for_page_token = slot.is_none(); + let best = slot.get_or_insert(Distance::INFINITY); + *best = std::cmp::min(*best, distance); + first_for_page_token +} + +pub(super) struct PageCandidates { inner: BinaryHeap<(Reverse, AlwaysEqual)>, remaining: usize, } -impl Iterator for LegacyPageCandidates { +impl Iterator for PageCandidates { type Item = PageCandidate; fn next(&mut self) -> Option { @@ -126,8 +188,8 @@ impl Iterator for LegacyPageCandidates { } } -impl ExactSizeIterator for LegacyPageCandidates {} -impl std::iter::FusedIterator for LegacyPageCandidates {} +impl ExactSizeIterator for PageCandidates {} +impl std::iter::FusedIterator for PageCandidates {} #[cfg(test)] mod tests { @@ -158,7 +220,7 @@ mod tests { ), ] .into_iter(); - let candidates = LegacyPageCandidateGenerator + let candidates = DensePageCandidateGenerator .generate(2, &mut token_searches, usize::MAX) .collect::>(); @@ -182,7 +244,7 @@ mod tests { Distance::ZERO, )] .into_iter(); - let candidates = LegacyPageCandidateGenerator + let candidates = DensePageCandidateGenerator .generate(1, &mut token_searches, 1) .collect::>(); diff --git a/src/index/vchordrq/scanners/maxsim/exact.rs b/src/index/vchordrq/scanners/maxsim/exact.rs new file mode 100644 index 00000000..020e8c5e --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/exact.rs @@ -0,0 +1,500 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the ELv2, which has specific restrictions. +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorStorage, + TileMaxsimSourceBinding, resolve_tilemaxsim_source, validate_descriptor, +}; +use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; +use super::profile; +use super::rerank::RerankError; +use crate::datatype::memory_halfvec::HalfvecInput; +use crate::datatype::memory_vector::VectorInput; +use crate::index::gucs::{self, PostgresMaxsimBackend}; +use distance::Distance; +use pgrx::datum::{Array, DatumWithOid, FromDatum, IntoDatum}; +use pgrx::iter::TableIterator; +use pgrx::{name, pg_extern}; +use std::collections::{BTreeMap, BTreeSet}; +use std::time::Duration; +use vchordrq::types::OwnedVector; +use vector::VectorBorrowed; + +const MAX_EXPLICIT_CANDIDATES: usize = 65_536; + +/// Exact TileMaxSim for a caller-supplied candidate set of external tensors. +/// Candidate selection, tenancy, ACLs, graph traversal, and clustering are +/// intentionally outside this function. +#[pg_extern(sql = "")] +fn _vchordrq_tilemaxsim_rerank_vector( + source_oid: pgrx::pg_sys::Oid, + query: Array<'_, VectorInput<'_>>, + candidate_ids: Array<'_, i64>, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = collect_vector_query(query) + .and_then(|query| { + collect_candidate_ids(candidate_ids) + .and_then(|candidate_ids| execute_rerank(source_oid, query, candidate_ids, top_k)) + }) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +/// Half-precision overload of exact candidate-only TileMaxSim. +#[pg_extern(sql = "")] +fn _vchordrq_tilemaxsim_rerank_halfvec( + source_oid: pgrx::pg_sys::Oid, + query: Array<'_, HalfvecInput<'_>>, + candidate_ids: Array<'_, i64>, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = collect_halfvec_query(query) + .and_then(|query| { + collect_candidate_ids(candidate_ids) + .and_then(|candidate_ids| execute_rerank(source_oid, query, candidate_ids, top_k)) + }) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +fn collect_vector_query( + query: Array<'_, VectorInput<'_>>, +) -> Result, RerankError> { + let mut vectors = Vec::with_capacity(query.len()); + for vector in query.iter() { + let vector = vector.ok_or(RerankError::TensorMismatch)?; + vectors.push(OwnedVector::Vecf32(vector.as_borrowed().own())); + } + validate_query(&vectors)?; + Ok(vectors) +} + +fn collect_halfvec_query( + query: Array<'_, HalfvecInput<'_>>, +) -> Result, RerankError> { + let mut vectors = Vec::with_capacity(query.len()); + for vector in query.iter() { + let vector = vector.ok_or(RerankError::TensorMismatch)?; + vectors.push(OwnedVector::Vecf16(vector.as_borrowed().own())); + } + validate_query(&vectors)?; + Ok(vectors) +} + +fn validate_query(query: &[OwnedVector]) -> Result<(), RerankError> { + if query.is_empty() { + return Err(RerankError::Configuration( + "query tensor must contain at least one token", + )); + } + Ok(()) +} + +fn collect_candidate_ids(candidate_ids: Array<'_, i64>) -> Result, RerankError> { + let mut result = Vec::with_capacity(candidate_ids.len()); + let mut unique = BTreeSet::new(); + for public_id in candidate_ids.iter() { + let public_id = public_id.ok_or(RerankError::Configuration( + "candidate_ids must not contain NULL", + ))?; + if !unique.insert(public_id) { + return Err(RerankError::Configuration( + "candidate_ids must not contain duplicates", + )); + } + result.push(public_id); + } + if result.is_empty() || result.len() > MAX_EXPLICIT_CANDIDATES { + return Err(RerankError::Configuration( + "candidate_ids must contain between 1 and 65536 IDs", + )); + } + Ok(result) +} + +fn execute_rerank( + source_oid: pgrx::pg_sys::Oid, + query: Vec, + candidate_ids: Vec, + top_k: i32, +) -> Result, RerankError> { + validate_top_k(candidate_ids.len(), top_k)?; + if !matches!(gucs::vchordrq_maxsim_backend(), PostgresMaxsimBackend::Gpu) { + return Err(RerankError::Configuration( + "external TileMaxSim rerank requires vchordrq.maxsim_backend = 'gpu'", + )); + } + require_mvcc_snapshot()?; + + let profile_guard = profile::ProfileGuard::start(gucs::vchordrq_maxsim_profile()); + let total_timer = profile::ProfileTimer::start(); + profile::update(|profile| { + profile.query_tokens = query.len() as u64; + profile.generated_candidates = candidate_ids.len() as u64; + }); + + let preflight_timer = profile::ProfileTimer::start(); + let binding = resolve_tilemaxsim_source(source_oid)?; + let source_lock = RelationLock::open(source_oid, pgrx::pg_sys::AccessShareLock as _)?; + let descriptor_lock = binding + .descriptor_oid + .map(|oid| RelationLock::open(oid, pgrx::pg_sys::AccessShareLock as _)) + .transpose()?; + if source_lock.oid() != binding.source_oid + || descriptor_lock.as_ref().map(RelationLock::oid) != binding.descriptor_oid + { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source changed during execution".into(), + )); + } + preflight_descriptor_access(&binding)?; + profile::update(|profile| { + profile.preflight_us += profile::duration_us(preflight_timer.elapsed()); + }); + + let descriptor_timer = profile::ProfileTimer::start(); + let (mut source, public_ids) = load_visible_descriptors(&binding, &candidate_ids)?; + let mut candidates = public_ids + .keys() + .copied() + .map(|heap_key| PageCandidate { + approximate_distance: Distance::ZERO, + heap_key, + }) + .collect::>() + .into_iter(); + profile::update(|profile| { + profile.descriptor_us += profile::duration_us(descriptor_timer.elapsed()); + profile.visible_candidates = public_ids.len() as u64; + profile.descriptors = public_ids.len() as u64; + }); + + let endpoint = gucs::vchordrq_maxsim_gpu_endpoint() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let mut backend = GpuExternalTileMaxsimBackend::new( + UnixSocketTransport::new(endpoint), + binding.model_contract_id, + Duration::from_millis(gucs::vchordrq_maxsim_gpu_timeout_ms() as u64), + gucs::vchordrq_maxsim_gpu_max_batch_tokens() as usize, + gucs::vchordrq_maxsim_gpu_max_batch_bytes() as usize, + ); + let sidecar_timer = profile::ProfileTimer::start(); + let exact = backend.rerank(&query, &mut candidates, &mut source)?; + let mut rows = exact + .map(|result| { + let public_id = public_ids.get(&result.heap_key).copied().ok_or_else(|| { + RerankError::Protocol("sidecar result has no visible public ID".into()) + })?; + Ok((result.distance, public_id)) + }) + .collect::, RerankError>>()?; + profile::update(|profile| { + profile.sidecar_us += profile::duration_us(sidecar_timer.elapsed()); + }); + + let result_timer = profile::ProfileTimer::start(); + rows.sort_unstable_by(|(left_distance, left_id), (right_distance, right_id)| { + left_distance + .cmp(right_distance) + .then_with(|| left_id.cmp(right_id)) + }); + rows.truncate(top_k as usize); + let output = rows + .into_iter() + .map(|(distance, public_id)| (public_id, -distance.to_f32())) + .collect::>(); + profile::update(|profile| { + profile.result_finalize_us += profile::duration_us(result_timer.elapsed()); + profile.returned_rows = output.len() as u64; + }); + profile_guard.finish(total_timer.elapsed()); + Ok(output.into_iter()) +} + +fn validate_top_k(candidate_count: usize, top_k: i32) -> Result<(), RerankError> { + if top_k <= 0 || top_k as usize > candidate_count { + return Err(RerankError::Configuration( + "top_k must be positive and no greater than candidate_ids length", + )); + } + Ok(()) +} + +fn require_mvcc_snapshot() -> Result<(), RerankError> { + let snapshot = unsafe { pgrx::pg_sys::GetActiveSnapshot() }; + if snapshot.is_null() + || unsafe { (*snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + { + return Err(RerankError::Configuration( + "external TileMaxSim rerank requires an active MVCC snapshot", + )); + } + Ok(()) +} + +fn preflight_descriptor_access(binding: &TileMaxsimSourceBinding) -> Result<(), RerankError> { + let query = descriptor_query(binding, true)?; + pgrx::spi::Spi::connect(|client| { + let privilege = client + .prepare( + "SELECT pg_catalog.has_table_privilege($1, 'SELECT') AS allowed", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + for relation_oid in std::iter::once(binding.source_oid).chain(binding.descriptor_oid) { + let allowed = client + .select(&privilege, Some(1), &[relation_oid.into()]) + .map_err(registry_error)? + .first() + .get_by_name::("allowed") + .map_err(registry_error)? + .unwrap_or(false); + if !allowed { + return Err(RerankError::Registry( + "table-level SELECT privilege is required for every TileMaxSim relation".into(), + )); + } + } + client + .select(query.as_str(), Some(1), &[]) + .map(|_| ()) + .map_err(registry_error) + }) +} + +fn load_visible_descriptors( + binding: &TileMaxsimSourceBinding, + candidate_ids: &[i64], +) -> Result<(MaterializedDescriptorSource, BTreeMap), RerankError> { + let mut candidates_by_id = BTreeMap::new(); + for (ordinal, public_id) in candidate_ids.iter().copied().enumerate() { + candidates_by_id.insert(public_id, candidate_for_ordinal(ordinal)?); + } + let query = descriptor_query(binding, false)?; + let (descriptors, public_ids) = pgrx::spi::Spi::connect(|client| { + let int8_array_oid = unsafe { pgrx::pg_sys::get_array_type(pgrx::pg_sys::INT8OID) }; + let prepared = client + .prepare( + query.as_str(), + &[ + pgrx::pg_sys::PgOid::from(int8_array_oid), + pgrx::pg_sys::PgOid::from(pgrx::pg_sys::TEXTOID), + ], + ) + .map_err(registry_error)?; + let args: [DatumWithOid<'_>; 2] = [ + candidate_ids.to_vec().into(), + binding.model_contract_id.clone().into(), + ]; + let rows = client + .select(&prepared, Some(candidate_ids.len() as _), &args) + .map_err(registry_error)?; + let mut descriptors = BTreeMap::new(); + let mut public_ids = BTreeMap::new(); + for row in rows { + let public_id = required_heap_column::(&row, "public_id")?; + let candidate = candidates_by_id.get(&public_id).copied().ok_or_else(|| { + RerankError::Registry("descriptor query returned an unknown public ID".into()) + })?; + let descriptor = validate_descriptor( + candidate, + public_id, + required_heap_column::(&row, "tensor_ref")?, + required_heap_column::(&row, "tensor_rows")?, + required_heap_column::(&row, "tensor_dimension")?, + required_heap_column::(&row, "tensor_dtype")?, + required_heap_column::(&row, "tensor_checksum")?, + )?; + if descriptors.insert(candidate.heap_key, descriptor).is_some() + || public_ids.insert(candidate.heap_key, public_id).is_some() + { + return Err(RerankError::Registry( + "descriptor query returned a duplicate public ID".into(), + )); + } + } + Ok((descriptors, public_ids)) + })?; + Ok((MaterializedDescriptorSource(descriptors), public_ids)) +} + +fn candidate_for_ordinal(ordinal: usize) -> Result { + let ordinal = u32::try_from(ordinal).map_err(|_| RerankError::RequestTooLarge)?; + Ok(PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, (ordinal >> 16) as u16, ordinal as u16], + }) +} + +fn descriptor_query( + binding: &TileMaxsimSourceBinding, + preflight: bool, +) -> Result { + let source_relation = relation_name(binding.source_oid)?; + let names = &binding.column_names; + let model_contract = pgrx::spi::quote_identifier(&names.model_contract); + let public_id = pgrx::spi::quote_identifier(&names.public_id); + let tensor_ref = pgrx::spi::quote_identifier(&names.tensor_ref); + let tensor_rows = pgrx::spi::quote_identifier(&names.tensor_rows); + let tensor_dimension = pgrx::spi::quote_identifier(&names.tensor_dimension); + let tensor_dtype = pgrx::spi::quote_identifier(&names.tensor_dtype); + let tensor_checksum = pgrx::spi::quote_identifier(&names.tensor_checksum); + let predicate = if preflight { + "false".to_string() + } else { + format!("h.{public_id}::bigint = ANY($1) AND h.{model_contract} = $2") + }; + match binding.storage { + ExternalTensorStorage::SameHeap => Ok(format!( + "SELECT h.{public_id}::bigint AS public_id, + h.{tensor_ref} AS tensor_ref, + h.{tensor_rows} AS tensor_rows, + h.{tensor_dimension} AS tensor_dimension, + h.{tensor_dtype} AS tensor_dtype, + h.{tensor_checksum} AS tensor_checksum + FROM ONLY {source_relation} AS h + WHERE {predicate}" + )), + ExternalTensorStorage::DescriptorRelation => { + let descriptor_oid = binding.descriptor_oid.ok_or_else(|| { + RerankError::Registry("registered descriptor relation is missing".into()) + })?; + let descriptor_relation = relation_name(descriptor_oid)?; + let descriptor_public_id = pgrx::spi::quote_identifier( + names.descriptor_public_id.as_deref().ok_or_else(|| { + RerankError::Registry( + "registered descriptor public ID column is missing".into(), + ) + })?, + ); + Ok(format!( + "SELECT h.{public_id}::bigint AS public_id, + d.{tensor_ref} AS tensor_ref, + d.{tensor_rows} AS tensor_rows, + d.{tensor_dimension} AS tensor_dimension, + d.{tensor_dtype} AS tensor_dtype, + d.{tensor_checksum} AS tensor_checksum + FROM ONLY {source_relation} AS h + JOIN ONLY {descriptor_relation} AS d + ON d.{descriptor_public_id}::bigint = h.{public_id}::bigint + WHERE {predicate}" + )) + } + } +} + +fn relation_name(relation_oid: pgrx::pg_sys::Oid) -> Result { + pgrx::spi::Spi::connect(|client| { + let prepared = client + .prepare( + "SELECT n.nspname::text AS schema_name, c.relname::text AS relation_name + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE c.oid = $1", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[relation_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered relation disappeared".into(), + )); + } + let row = rows.first(); + Ok(pgrx::spi::quote_qualified_identifier( + required_column::(&row, "schema_name")?, + required_column::(&row, "relation_name")?, + )) + }) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry(format!("query returned NULL {name}"))) +} + +fn required_heap_column( + row: &pgrx::spi::SpiHeapTupleData<'_>, + name: &str, +) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +#[derive(Default)] +struct MaterializedDescriptorSource(BTreeMap); + +impl CandidateTensorDescriptorSource for MaterializedDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(self.0.remove(&candidate.heap_key)) + } +} + +struct RelationLock { + raw: pgrx::pg_sys::Relation, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl RelationLock { + fn open(oid: pgrx::pg_sys::Oid, lockmode: pgrx::pg_sys::LOCKMODE) -> Result { + let raw = unsafe { pgrx::pg_sys::relation_open(oid, lockmode) }; + if raw.is_null() { + return Err(RerankError::Registry("relation open returned NULL".into())); + } + Ok(Self { raw, lockmode }) + } + + fn oid(&self) -> pgrx::pg_sys::Oid { + unsafe { (*self.raw).rd_id } + } +} + +impl Drop for RelationLock { + fn drop(&mut self) { + unsafe { pgrx::pg_sys::relation_close(self.raw, self.lockmode) }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn explicit_candidates_are_bounded_and_unique() { + assert!(validate_top_k(256, 10).is_ok()); + assert!(validate_top_k(10, 0).is_err()); + assert!(validate_top_k(10, 11).is_err()); + let first = candidate_for_ordinal(0).unwrap().heap_key; + let last = candidate_for_ordinal(65_535).unwrap().heap_key; + assert_ne!(first, last); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/external.rs b/src/index/vchordrq/scanners/maxsim/external.rs index 0c8d1eb0..8b726da5 100644 --- a/src/index/vchordrq/scanners/maxsim/external.rs +++ b/src/index/vchordrq/scanners/maxsim/external.rs @@ -116,6 +116,20 @@ pub(super) struct ExternalTensorSourceBinding { pub column_names: ExternalTensorColumnNames, } +/// External tensor metadata bound directly to an application source relation. +/// +/// Unlike [`ExternalTensorSourceBinding`], this binding has no ANN index: the +/// caller supplies an already-authorized candidate ID set and VectorChord only +/// performs exact TileMaxSim over those candidates. +#[derive(Clone, Debug)] +pub(super) struct TileMaxsimSourceBinding { + pub source_oid: pgrx::pg_sys::Oid, + pub descriptor_oid: Option, + pub storage: ExternalTensorStorage, + pub model_contract_id: String, + pub column_names: ExternalTensorColumnNames, +} + #[derive(Clone, Debug)] pub(super) struct ExternalTensorColumnNames { pub model_contract: String, @@ -264,6 +278,114 @@ pub(super) fn resolve_external_tensor_source( }) } +/// Resolve a source-relation binding for exact candidate-only TileMaxSim. +/// +/// The SECURITY DEFINER SQL resolver revalidates relation ownership, caller +/// SELECT privileges, column types, uniqueness, and descriptor metadata. This +/// Rust layer deliberately sees only the validated projection. +pub(super) fn resolve_tilemaxsim_source( + source_oid: pgrx::pg_sys::Oid, +) -> Result { + pgrx::spi::Spi::connect(|client| { + let schema_rows = client + .select( + "SELECT n.nspname::text AS schema_name + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'", + Some(1), + &[], + ) + .map_err(registry_error)?; + if schema_rows.is_empty() { + return Err(RerankError::Registry( + "vchord extension schema is unavailable".into(), + )); + } + let schema_name = schema_rows + .first() + .get_by_name::("schema_name") + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry("vchord extension schema is NULL".into()))?; + let resolver = pgrx::spi::quote_qualified_identifier( + schema_name, + "vchordrq_tilemaxsim_source_info".to_string(), + ); + let query = format!( + "SELECT registered_source::oid AS source_oid, + descriptor_relation::oid AS descriptor_oid, + model_contract_id, + source_storage, + model_contract_column::text AS model_contract_column, + public_id_column::text AS public_id_column, + descriptor_public_id_column::text AS descriptor_public_id_column, + tensor_ref_column::text AS tensor_ref_column, + tensor_rows_column::text AS tensor_rows_column, + tensor_dim_column::text AS tensor_dim_column, + tensor_dtype_column::text AS tensor_dtype_column, + tensor_checksum_column::text AS tensor_checksum_column + FROM {resolver}($1::regclass)" + ); + let prepared = client + .prepare(query.as_str(), &pgrx::oids_of![pgrx::pg_sys::Oid]) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[source_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source resolution returned no row".into(), + )); + } + let row = rows.first(); + let resolved_source_oid = required_column::(&row, "source_oid")?; + let descriptor_oid = optional_column::(&row, "descriptor_oid")?; + let model_contract_id = required_column::(&row, "model_contract_id")?; + let storage = required_column::(&row, "source_storage")?; + if resolved_source_oid != source_oid { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source resolved a different relation".into(), + )); + } + let storage = match storage.as_str() { + "external_ref" if descriptor_oid.is_none() => ExternalTensorStorage::SameHeap, + "external_relation" if descriptor_oid.is_some() => { + ExternalTensorStorage::DescriptorRelation + } + _ => { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source is inconsistent".into(), + )); + } + }; + let column_names = ExternalTensorColumnNames { + model_contract: required_column::(&row, "model_contract_column")?, + public_id: required_column::(&row, "public_id_column")?, + descriptor_public_id: optional_column::(&row, "descriptor_public_id_column")?, + tensor_ref: required_column::(&row, "tensor_ref_column")?, + tensor_rows: required_column::(&row, "tensor_rows_column")?, + tensor_dimension: required_column::(&row, "tensor_dim_column")?, + tensor_dtype: required_column::(&row, "tensor_dtype_column")?, + tensor_checksum: required_column::(&row, "tensor_checksum_column")?, + }; + if matches!(storage, ExternalTensorStorage::DescriptorRelation) + && column_names.descriptor_public_id.is_none() + { + return Err(RerankError::Registry( + "registered descriptor relation has no public ID column".into(), + )); + } + validate_model_contract(&model_contract_id)?; + Ok(TileMaxsimSourceBinding { + source_oid, + descriptor_oid, + storage, + model_contract_id, + column_names, + }) + }) +} + fn registry_error(error: impl std::fmt::Display) -> RerankError { RerankError::Registry(error.to_string()) } diff --git a/src/index/vchordrq/scanners/maxsim/profile.rs b/src/index/vchordrq/scanners/maxsim/profile.rs new file mode 100644 index 00000000..e8a6dffe --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/profile.rs @@ -0,0 +1,181 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use std::cell::RefCell; +use std::time::{Duration, Instant}; + +#[derive(Default)] +pub(super) struct MaxsimProfile { + pub query_tokens: u64, + pub token_search_calls: u64, + pub token_search_us: u64, + pub token_search_results: u64, + pub token_refine_calls: u64, + pub token_refine_us: u64, + pub accurate_token_hits: u64, + pub rough_token_hits: u64, + pub hit_collect_us: u64, + pub hit_updates: u64, + pub page_token_updates: u64, + pub hit_sort_us: u64, + pub page_aggregate_us: u64, + pub aggregated_pages: u64, + pub preflight_us: u64, + pub candidate_generation_us: u64, + pub generated_candidates: u64, + pub visibility_us: u64, + pub visible_candidates: u64, + pub descriptor_us: u64, + pub descriptors: u64, + pub sidecar_us: u64, + pub result_finalize_us: u64, + pub returned_rows: u64, + pub total_us: u64, +} + +impl MaxsimProfile { + fn to_json(&self) -> String { + format!( + concat!( + "{{", + "\"schema_version\":1,", + "\"query_tokens\":{},", + "\"token_search_calls\":{},", + "\"token_search_us\":{},", + "\"token_search_results\":{},", + "\"token_refine_calls\":{},", + "\"token_refine_us\":{},", + "\"accurate_token_hits\":{},", + "\"rough_token_hits\":{},", + "\"hit_collect_us\":{},", + "\"hit_updates\":{},", + "\"page_token_updates\":{},", + "\"hit_sort_us\":{},", + "\"page_aggregate_us\":{},", + "\"aggregated_pages\":{},", + "\"preflight_us\":{},", + "\"candidate_generation_us\":{},", + "\"generated_candidates\":{},", + "\"visibility_us\":{},", + "\"visible_candidates\":{},", + "\"descriptor_us\":{},", + "\"descriptors\":{},", + "\"sidecar_us\":{},", + "\"result_finalize_us\":{},", + "\"returned_rows\":{},", + "\"total_us\":{}", + "}}" + ), + self.query_tokens, + self.token_search_calls, + self.token_search_us, + self.token_search_results, + self.token_refine_calls, + self.token_refine_us, + self.accurate_token_hits, + self.rough_token_hits, + self.hit_collect_us, + self.hit_updates, + self.page_token_updates, + self.hit_sort_us, + self.page_aggregate_us, + self.aggregated_pages, + self.preflight_us, + self.candidate_generation_us, + self.generated_candidates, + self.visibility_us, + self.visible_candidates, + self.descriptor_us, + self.descriptors, + self.sidecar_us, + self.result_finalize_us, + self.returned_rows, + self.total_us, + ) + } +} + +thread_local! { + static ACTIVE_PROFILE: RefCell> = const { RefCell::new(None) }; +} + +pub(super) struct ProfileGuard { + enabled: bool, + finished: bool, +} + +impl ProfileGuard { + pub fn start(enabled: bool) -> Self { + if enabled { + ACTIVE_PROFILE.with(|slot| { + *slot.borrow_mut() = Some(MaxsimProfile::default()); + }); + } + Self { + enabled, + finished: false, + } + } + + pub fn finish(mut self, total: Duration) { + if self.enabled { + let profile = ACTIVE_PROFILE.with(|slot| { + let mut profile = slot.borrow_mut().take(); + if let Some(profile) = profile.as_mut() { + profile.total_us = duration_us(total); + } + profile + }); + if let Some(profile) = profile { + pgrx::notice!("vchordrq_maxsim_profile {}", profile.to_json()); + } + } + self.finished = true; + } +} + +impl Drop for ProfileGuard { + fn drop(&mut self) { + if self.enabled && !self.finished { + ACTIVE_PROFILE.with(|slot| { + slot.borrow_mut().take(); + }); + } + } +} + +pub(super) struct ProfileTimer(Option); + +impl ProfileTimer { + pub fn start() -> Self { + let enabled = ACTIVE_PROFILE.with(|slot| slot.borrow().is_some()); + Self(enabled.then(Instant::now)) + } + + pub fn elapsed(self) -> Duration { + self.0.map_or(Duration::ZERO, |started| started.elapsed()) + } +} + +pub(super) fn update(f: impl FnOnce(&mut MaxsimProfile)) { + ACTIVE_PROFILE.with(|slot| { + if let Some(profile) = slot.borrow_mut().as_mut() { + f(profile); + } + }); +} + +pub(super) fn duration_us(duration: Duration) -> u64 { + duration.as_micros().min(u128::from(u64::MAX)) as u64 +} diff --git a/src/index/vchordrq/scanners/maxsim/search.rs b/src/index/vchordrq/scanners/maxsim/search.rs index 915fb87f..41ea6e8f 100644 --- a/src/index/vchordrq/scanners/maxsim/search.rs +++ b/src/index/vchordrq/scanners/maxsim/search.rs @@ -19,6 +19,7 @@ use super::external::{ ExternalTensorStorage, resolve_external_tensor_source, validate_descriptor, }; use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; +use super::profile; use super::rerank::RerankError; use crate::index::fetcher::{ Fetcher, FilterableTuple, HeapFetcher, Tuple, TupleAttribute, ctid_to_key, @@ -59,6 +60,8 @@ fn execute_external_search( candidate_limit: i32, top_k: i32, ) -> Result, RerankError> { + let profile_guard = profile::ProfileGuard::start(gucs::vchordrq_maxsim_profile()); + let total_timer = profile::ProfileTimer::start(); validate_search_limits(candidate_limit, top_k)?; if !matches!(gucs::vchordrq_maxsim_backend(), PostgresMaxsimBackend::Gpu) { return Err(RerankError::Configuration( @@ -66,6 +69,7 @@ fn execute_external_search( )); } + let preflight_timer = profile::ProfileTimer::start(); let binding = resolve_external_tensor_source(index_oid)?; let index_lock = RelationLock::open(index_oid, pgrx::pg_sys::AccessShareLock as _)?; let heap_lock = RelationLock::open(binding.heap_oid, pgrx::pg_sys::AccessShareLock as _)?; @@ -94,21 +98,41 @@ fn execute_external_search( } let query_vectors = unsafe { opfamily.input_vectors(query.datum()) }.ok_or(RerankError::TensorMismatch)?; + profile::update(|profile| { + profile.query_tokens = query_vectors.len() as u64; + }); // The registry resolution happens before the index read, and this // privilege-only SELECT is planned/executed before any candidate CTID is // generated. It fails early when the caller cannot project the registered // descriptor columns. The same query shape is used for the actual fetch. preflight_descriptor_access(&binding)?; + let preflight_elapsed = preflight_timer.elapsed(); + profile::update(|profile| { + profile.preflight_us += profile::duration_us(preflight_elapsed); + }); + let candidate_timer = profile::ProfileTimer::start(); let candidates = generate_candidates( index_lock.raw(), opfamily, query.datum(), candidate_limit as u32, )?; + let candidate_elapsed = candidate_timer.elapsed(); + profile::update(|profile| { + profile.candidate_generation_us += profile::duration_us(candidate_elapsed); + profile.generated_candidates = candidates.len() as u64; + }); + let visibility_timer = profile::ProfileTimer::start(); let resolved = resolve_visible_candidates(index_lock.raw(), heap_lock.raw(), candidates.into_iter())?; + let visibility_elapsed = visibility_timer.elapsed(); + profile::update(|profile| { + profile.visibility_us += profile::duration_us(visibility_elapsed); + profile.visible_candidates = resolved.len() as u64; + }); + let descriptor_timer = profile::ProfileTimer::start(); let (mut source, public_ids) = load_visible_descriptors(&binding, &resolved)?; let visible_candidates = resolved .into_iter() @@ -118,6 +142,11 @@ fn execute_external_search( .then_some(resolved.candidate) }) .collect::>(); + let descriptor_elapsed = descriptor_timer.elapsed(); + profile::update(|profile| { + profile.descriptor_us += profile::duration_us(descriptor_elapsed); + profile.descriptors = public_ids.len() as u64; + }); let endpoint = gucs::vchordrq_maxsim_gpu_endpoint() .map(|endpoint| endpoint.to_string_lossy().into_owned()) .unwrap_or_default(); @@ -130,6 +159,7 @@ fn execute_external_search( gucs::vchordrq_maxsim_gpu_max_batch_bytes() as usize, ); let mut candidate_iter = visible_candidates.into_iter(); + let sidecar_timer = profile::ProfileTimer::start(); let exact = backend.rerank(&query_vectors, &mut candidate_iter, &mut source)?; let mut rows = exact .map(|result| { @@ -139,17 +169,28 @@ fn execute_external_search( Ok((result.distance, public_id)) }) .collect::, RerankError>>()?; + let sidecar_elapsed = sidecar_timer.elapsed(); + profile::update(|profile| { + profile.sidecar_us += profile::duration_us(sidecar_elapsed); + }); + let result_finalize_timer = profile::ProfileTimer::start(); rows.sort_unstable_by(|(left_distance, left_id), (right_distance, right_id)| { left_distance .cmp(right_distance) .then_with(|| left_id.cmp(right_id)) }); rows.truncate(top_k as usize); - Ok(rows + let output = rows .into_iter() .map(|(distance, public_id)| (public_id, -distance.to_f32())) - .collect::>() - .into_iter()) + .collect::>(); + let result_finalize_elapsed = result_finalize_timer.elapsed(); + profile::update(|profile| { + profile.result_finalize_us += profile::duration_us(result_finalize_elapsed); + profile.returned_rows = output.len() as u64; + }); + profile_guard.finish(total_timer.elapsed()); + Ok(output.into_iter()) } #[derive(Clone, Copy)] diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index a7e7af4e..08aa1fcb 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -1431,3 +1431,755 @@ REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop ON sql_drop EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); + +-- Candidate-only exact TileMaxSim. This source registry is deliberately +-- relation-keyed rather than index-keyed: the caller owns candidate recall, +-- tenancy, ACL, graph, and clustering decisions. VectorChord only loads the +-- visible candidate tensors and computes exact TileMaxSim. + +CREATE TABLE _vchordrq_tilemaxsim_sources ( + source_oid oid PRIMARY KEY, + model_contract_id text NOT NULL CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint NOT NULL CHECK ( + tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_rows_attnum smallint NOT NULL CHECK ( + tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dim_attnum smallint NOT NULL CHECK ( + tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dtype_attnum smallint NOT NULL CHECK ( + tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_checksum_attnum smallint NOT NULL CHECK ( + tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_tilemaxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_tilemaxsim_source( + source_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + public_id_is_unique boolean; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor descriptor columns must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be external_ref or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.oid, c.relowner + INTO source_oid, source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid + AND c.relkind IN ('r', 'm'); + IF source_oid IS NULL THEN + RAISE EXCEPTION 'source_relation % must be a table or materialized view', source_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may register a TileMaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL integer or bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'public ID column % must have a non-partial single-key unique index', + public_id_column; + END IF; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a TileMaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL integer or bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION 'external_ref sources must not specify a descriptor relation'; + END IF; + tensor_relation_oid := source_oid; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'TileMaxSim tensor source columns must be distinct'; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_tilemaxsim_sources ( + source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) + ON CONFLICT (source_oid) DO UPDATE SET + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + source_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_tilemaxsim_source(source_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_owner oid; + removed_count bigint; +BEGIN + IF source_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid AND c.relkind IN ('r', 'm'); + IF ext_schema IS NULL + OR source_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may unregister its TileMaxSim source'; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources WHERE source_oid = $1', + ext_schema + ) USING source_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_source_info(source_relation regclass) +RETURNS TABLE( + registered_source regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + bound_model_contract_id text; + bound_storage text; + descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + EXECUTE pg_catalog.format( + 'SELECT source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_tilemaxsim_sources + WHERE source_oid = $1', + ext_schema + ) INTO + source_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING source_relation::oid; + IF source_oid IS NULL THEN + RAISE EXCEPTION 'TileMaxSim tensor source is not registered for relation %', source_relation; + END IF; + + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_oid AND c.relkind IN ('r', 'm'); + IF source_owner IS NULL THEN + RAISE EXCEPTION 'registered TileMaxSim source relation is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, source_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered TileMaxSim source'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, NULL::oid) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = source_oid + AND a.attnum = expected.attnum + AND (expected.atttypid IS NULL OR a.atttypid = expected.atttypid) + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.atttypid IS NOT NULL + OR a.atttypid IN ('integer'::regtype, 'bigint'::regtype); + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered TileMaxSim source columns are invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim source public ID is no longer unique'; + END IF; + + IF bound_storage = 'external_ref' THEN + IF descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered external_ref TileMaxSim binding is invalid'; + END IF; + tensor_relation_oid := source_oid; + ELSIF bound_storage = 'external_relation' THEN + IF descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered external_relation TileMaxSim binding is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_oid AND c.relkind IN ('r', 'm'); + IF NOT FOUND OR NOT pg_catalog.has_table_privilege(caller_oid, descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor relation is unavailable'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid IN ('integer'::regtype, 'bigint'::regtype) + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + RAISE EXCEPTION 'registered TileMaxSim storage is invalid'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 5 THEN + RAISE EXCEPTION 'registered TileMaxSim tensor descriptor columns are invalid'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + source_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query vector[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_vector_wrapper'; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query halfvec[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_halfvec_wrapper'; + +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) +IS 'Exact TileMaxSim over caller-supplied candidate IDs. Candidate recall, tenancy, ACL, graph, and clustering remain the caller responsibility.'; +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) +IS 'Exact TileMaxSim over caller-supplied candidate IDs. Candidate recall, tenancy, ACL, graph, and clustering remain the caller responsibility.'; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) FROM PUBLIC; +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) TO PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_tilemaxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL OR pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_tilemaxsim_sources', ext_schema) + ) IS NULL THEN + RETURN; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE ( + d.objid = s.source_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.model_contract_attnum, + s.public_id_attnum, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_ref_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_rows_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dim_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dtype_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_checksum_attnum END + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.descriptor_public_id_attnum, + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_tilemaxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_tilemaxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_tilemaxsim_source_sql_drop(); diff --git a/tests/vchordrq/tilemaxsim_source_registry.slt b/tests/vchordrq/tilemaxsim_source_registry.slt new file mode 100644 index 00000000..387a0d24 --- /dev/null +++ b/tests/vchordrq/tilemaxsim_source_registry.slt @@ -0,0 +1,91 @@ +# Candidate-only TileMaxSim binds external tensor descriptors directly to an +# application relation. It must not require or inspect a VectorChord ANN index. + +statement ok +CREATE TABLE tilemaxsim_source_test ( + id integer PRIMARY KEY, + model_contract text NOT NULL, + embedding vector(2) +); + +statement ok +CREATE TABLE tilemaxsim_descriptor_test ( + source_id integer PRIMARY KEY, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement ok +INSERT INTO tilemaxsim_source_test VALUES + (1, 'colqwen@test', '[1,0]'), + (2, 'colqwen@test', '[0,1]'); + +statement ok +INSERT INTO tilemaxsim_descriptor_test VALUES + (1, 'tensor://1', 2, 2, 'float32', 'sha256:test1'), + (2, 'tensor://2', 2, 2, 'float32', 'sha256:test2'); + +statement ok +SELECT vchordrq_register_tilemaxsim_source( + source_relation => 'tilemaxsim_source_test'::regclass, + model_contract_id => ' colqwen@test ', + storage => 'EXTERNAL_RELATION', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'tilemaxsim_descriptor_test'::regclass, + descriptor_public_id_column => 'source_id' +); + +query TTTT +SELECT registered_source::text, source_storage, + descriptor_relation::text, public_id_column::text +FROM vchordrq_tilemaxsim_source_info('tilemaxsim_source_test'::regclass); +---- +tilemaxsim_source_test external_relation tilemaxsim_descriptor_test id + +# Candidate validation occurs before external tensor or sidecar access. +statement error candidate_ids must not contain duplicates +SELECT * FROM vchordrq_tilemaxsim_rerank( + 'tilemaxsim_source_test'::regclass, + ARRAY['[1,0]'::vector], + ARRAY[1,1]::bigint[], + 1 +); + +statement error top_k must be positive and no greater than candidate_ids length +SELECT * FROM vchordrq_tilemaxsim_rerank( + 'tilemaxsim_source_test'::regclass, + ARRAY['[1,0]'::halfvec], + ARRAY[1,2]::bigint[], + 3 +); + +# Attribute OIDs survive renames and the resolver returns the live name. +statement ok +ALTER TABLE tilemaxsim_descriptor_test RENAME tensor_ref TO tensor_location; + +query T +SELECT tensor_ref_column::text +FROM vchordrq_tilemaxsim_source_info('tilemaxsim_source_test'::regclass); +---- +tensor_location + +# Dropping a bound descriptor column removes the binding through sql_drop. +statement ok +ALTER TABLE tilemaxsim_descriptor_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) FROM _vchordrq_tilemaxsim_sources; +---- +0 + +statement ok +DROP TABLE tilemaxsim_source_test, tilemaxsim_descriptor_test; From a1ed3cf65f865dee16b5d4296ce0346c7e103433 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Tue, 14 Jul 2026 12:31:08 +0800 Subject: [PATCH 303/324] feat(tilemaxsim): add sharded native cache daemon --- .gitignore | 3 + services/Dockerfile.tilemaxsim | 1 + services/Dockerfile.tilemaxsimd | 22 + services/benchmark_tilemaxsim_ablation.py | 633 +++++++++++++++++ services/build_tilemaxsim_tensor_cache.py | 154 ++++- services/test_tilemaxsim_gpu_cache.py | 87 +++ services/test_tilemaxsim_rust_daemon.py | 254 +++++++ services/test_tilemaxsim_shard.py | 181 +++++ services/tilemaxsim_cuda_sidecar.py | 640 +++++++++++++++--- services/tilemaxsim_gpu_cache.py | 566 ++++++++++++++-- services/tilemaxsim_shard.py | 354 ++++++++++ services/tilemaxsimd/Cargo.lock | 416 ++++++++++++ services/tilemaxsimd/Cargo.toml | 19 + services/tilemaxsimd/build.rs | 12 + .../tilemaxsimd/native/tilemaxsim_cuda.cu | 291 ++++++++ services/tilemaxsimd/src/cache.rs | 399 +++++++++++ services/tilemaxsimd/src/engine.rs | 434 ++++++++++++ services/tilemaxsimd/src/gpu.rs | 163 +++++ services/tilemaxsimd/src/main.rs | 321 +++++++++ services/tilemaxsimd/src/protocol.rs | 238 +++++++ services/tilemaxsimd/src/shard.rs | 428 ++++++++++++ 21 files changed, 5450 insertions(+), 166 deletions(-) create mode 100644 services/Dockerfile.tilemaxsimd create mode 100644 services/benchmark_tilemaxsim_ablation.py create mode 100644 services/test_tilemaxsim_rust_daemon.py create mode 100644 services/test_tilemaxsim_shard.py create mode 100644 services/tilemaxsim_shard.py create mode 100644 services/tilemaxsimd/Cargo.lock create mode 100644 services/tilemaxsimd/Cargo.toml create mode 100644 services/tilemaxsimd/build.rs create mode 100644 services/tilemaxsimd/native/tilemaxsim_cuda.cu create mode 100644 services/tilemaxsimd/src/cache.rs create mode 100644 services/tilemaxsimd/src/engine.rs create mode 100644 services/tilemaxsimd/src/gpu.rs create mode 100644 services/tilemaxsimd/src/main.rs create mode 100644 services/tilemaxsimd/src/protocol.rs create mode 100644 services/tilemaxsimd/src/shard.rs diff --git a/.gitignore b/.gitignore index 85d1619b..ab926f7e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,7 @@ /target +/services/tilemaxsimd/target +**/__pycache__/ +.ruff_cache/ **/*.rs.bk .vscode .ignore diff --git a/services/Dockerfile.tilemaxsim b/services/Dockerfile.tilemaxsim index 3ccb6732..2bf87252 100644 --- a/services/Dockerfile.tilemaxsim +++ b/services/Dockerfile.tilemaxsim @@ -6,6 +6,7 @@ WORKDIR /opt/vectorchord COPY devtools/tilemaxsim_reference_sidecar.py devtools/tilemaxsim_reference_sidecar.py COPY services/tilemaxsim_cuda_sidecar.py services/tilemaxsim_cuda_sidecar.py COPY services/tilemaxsim_gpu_cache.py services/tilemaxsim_gpu_cache.py +COPY services/tilemaxsim_shard.py services/tilemaxsim_shard.py COPY services/tilemaxsim_triton.py services/tilemaxsim_triton.py ENV PYTHONPATH=/opt/vectorchord \ diff --git a/services/Dockerfile.tilemaxsimd b/services/Dockerfile.tilemaxsimd new file mode 100644 index 00000000..d6dcb2f2 --- /dev/null +++ b/services/Dockerfile.tilemaxsimd @@ -0,0 +1,22 @@ +ARG CUDA_DEVEL_IMAGE=nvidia/cuda:12.6.3-devel-ubuntu24.04 +ARG CUDA_RUNTIME_IMAGE=nvidia/cuda:12.6.3-runtime-ubuntu24.04 + +FROM ${CUDA_DEVEL_IMAGE} AS build + +RUN apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + build-essential ca-certificates curl \ + && rm -rf /var/lib/apt/lists/* +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --profile minimal --default-toolchain 1.96.0 + +ENV PATH=/root/.cargo/bin:${PATH} +WORKDIR /build/tilemaxsimd +COPY services/tilemaxsimd/ . +RUN cargo build --release --locked + +FROM ${CUDA_RUNTIME_IMAGE} + +COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimd /usr/local/bin/tilemaxsimd + +ENTRYPOINT ["/usr/local/bin/tilemaxsimd"] diff --git a/services/benchmark_tilemaxsim_ablation.py b/services/benchmark_tilemaxsim_ablation.py new file mode 100644 index 00000000..5a6bbf18 --- /dev/null +++ b/services/benchmark_tilemaxsim_ablation.py @@ -0,0 +1,633 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Run storage, cache, H2D, and native-daemon TileMaxSim ablations.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import random +import select +import socket +import statistics +import subprocess +import sys +import tempfile +import time +from collections import OrderedDict +from pathlib import Path + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from devtools.test_tilemaxsim_reference_sidecar import decode_response +from services.tilemaxsim_cuda_sidecar import ContentAddressedResolver, PayloadCache +from services.tilemaxsim_gpu_cache import ( + FreeExtentAllocator, + GpuArenaSpec, + GpuResourcePool, + GpuTensorCache, + GpuTensorLoad, +) + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + return ordered[min(len(ordered) - 1, math.ceil(len(ordered) * fraction) - 1)] + + +def load_records(path: Path) -> list[dict[str, object]]: + with path.open(encoding="utf-8") as stream: + return [json.loads(line) for line in stream if line.strip()] + + +def request(record: dict[str, object], contract: str) -> protocol.ExternalTensorRequest: + dtype = protocol.DTYPE_F16 if record["tensor_dtype"] == "float16" else protocol.DTYPE_F32 + return protocol.ExternalTensorRequest( + contract, + str(record["tensor_ref"]), + int(record["tensor_rows"]), + int(record["tensor_dim"]), + dtype, + str(record["tensor_checksum"]), + ) + + +def evict_paths(paths: list[Path]) -> None: + if not hasattr(os, "posix_fadvise"): + return + for path in paths: + descriptor = os.open(path, os.O_RDONLY | os.O_CLOEXEC) + try: + os.posix_fadvise(descriptor, 0, 0, os.POSIX_FADV_DONTNEED) + finally: + os.close(descriptor) + + +def storage_ablation( + selected: list[dict[str, object]], contract: str, legacy_root: Path, shard_root: Path +) -> dict[str, object]: + requests = [request(record, contract) for record in selected] + legacy_paths = [ + legacy_root + / str(record["tensor_checksum"])[7:9] + / f"{str(record['tensor_checksum'])[7:]}.bin" + for record in selected + ] + shard_paths = sorted((shard_root / "shards").glob("*.vts")) + + evict_paths(legacy_paths) + resolver = ContentAddressedResolver({contract: legacy_root}, 0) + started = time.perf_counter() + try: + sequential = [resolver.resolve(item) for item in requests] + finally: + resolver.close() + sequential_ms = (time.perf_counter() - started) * 1000 + + evict_paths(legacy_paths) + resolver = ContentAddressedResolver({contract: legacy_root}, 0) + started = time.perf_counter() + try: + legacy_batch = resolver.resolve_many(requests) + finally: + resolver.close() + legacy_batch_ms = (time.perf_counter() - started) * 1000 + + evict_paths(shard_paths) + resolver = ContentAddressedResolver({contract: shard_root}, 0) + started = time.perf_counter() + try: + shard_batch = resolver.resolve_many(requests) + shard_status = resolver.status() + finally: + resolver.close() + shard_batch_ms = (time.perf_counter() - started) * 1000 + expected = [item.payload for item in sequential] + if expected != [item.payload for item in legacy_batch] or expected != [ + item.payload for item in shard_batch + ]: + raise RuntimeError("storage ablation payloads disagree") + return { + "candidates": len(selected), + "logical_bytes": sum(len(item) for item in expected), + "legacy_sequential_ms": sequential_ms, + "legacy_batch_ms": legacy_batch_ms, + "shard_batch_ms": shard_batch_ms, + "shard_batch_read_calls": shard_status["batch_read_calls"], + "shard_batch_read_bytes": shard_status["batch_read_bytes"], + } + + +def h2d_ablation( + selected: list[dict[str, object]], contract: str, shard_root: Path, device: int +) -> dict[str, object]: + requests = [request(record, contract) for record in selected] + resolver = ContentAddressedResolver({contract: shard_root}, 0) + try: + payloads = resolver.resolve_many(requests) + keys = [resolver.key(item) for item in requests] + finally: + resolver.close() + total_bytes = 768 * 1024**2 + workspace_bytes = 256 * 1024**2 + + pool = GpuResourcePool([GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + started = time.perf_counter() + handles = [] + for key, item, resolved in zip(keys, requests, payloads, strict=True): + handle, _ = cache.acquire( + key, + item.rows, + item.dimension, + item.dtype, + lambda payload=resolved.payload: payload, + ) + handles.append(handle) + torch.cuda.synchronize(device) + sequential_ms = (time.perf_counter() - started) * 1000 + sequential_status = pool.status()[0] + for handle in handles: + cache.release(handle) + finally: + pool.close() + + pool = GpuResourcePool([GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + loads = [ + GpuTensorLoad( + key, item.rows, item.dimension, item.dtype, resolved.payload + ) + for key, item, resolved in zip(keys, requests, payloads, strict=True) + ] + started = time.perf_counter() + batch = cache.acquire_many(loads) + torch.cuda.synchronize(device) + batch_ms = (time.perf_counter() - started) * 1000 + batch_status = pool.status()[0] + if batch.bypassed or batch.deferred: + raise RuntimeError("H2D ablation cache is undersized") + for handle in batch.handles: + assert handle is not None + cache.release(handle) + finally: + pool.close() + return { + "candidates": len(selected), + "logical_bytes": sum(len(item.payload) for item in payloads), + "per_tensor_h2d_ms": sequential_ms, + "batch_h2d_ms": batch_ms, + "per_tensor_copy_batches": sequential_status["h2d_batches"], + "batch_copy_batches": batch_status["h2d_batches"], + "batch_copy_calls": batch_status["h2d_copy_calls"], + } + + +def allocator_ablation(records: list[dict[str, object]]) -> dict[str, object]: + sizes = [int(record["canonical_bytes"]) for record in records] + rng = random.Random(991) + capacity = 256 * 1024**2 + extent = FreeExtentAllocator(capacity) + extent_live: list[tuple[int, int]] = [] + extent_failures = 0 + extent_fragmentation_failures = 0 + for _ in range(20_000): + if extent_live and rng.random() < 0.48: + extent.release(*extent_live.pop(rng.randrange(len(extent_live)))) + else: + size = rng.choice(sizes) + required = extent.allocation_bytes(size) + allocated = extent.allocate(size) + if allocated is None: + extent_failures += 1 + extent_fragmentation_failures += int(extent.free_bytes >= required) + else: + extent_live.append(allocated) + from services.tilemaxsim_gpu_cache import FixedBlockAllocator + + rng = random.Random(991) + slab = FixedBlockAllocator(capacity) + slab_live: list[tuple[int, ...]] = [] + slab_failures = 0 + slab_fragmentation_failures = 0 + requested = 0 + allocated_bytes = 0 + for _ in range(20_000): + if slab_live and rng.random() < 0.48: + slab.release(slab_live.pop(rng.randrange(len(slab_live)))) + else: + size = rng.choice(sizes) + required = slab.allocation_bytes(size) + allocated = slab.allocate(size) + if allocated is None: + slab_failures += 1 + slab_fragmentation_failures += int(slab.free_bytes >= required) + else: + slab_live.append(allocated) + requested += size + allocated_bytes += len(allocated) * slab.block_bytes + return { + "operations": 20_000, + "capacity_bytes": capacity, + "extent_allocation_failures": extent_failures, + "extent_fragmentation_failures": extent_fragmentation_failures, + "slab_allocation_failures": slab_failures, + "slab_fragmentation_failures": slab_fragmentation_failures, + "slab_internal_waste_ratio": (allocated_bytes - requested) / allocated_bytes, + } + + +class LegacyLru: + def __init__(self, maximum_bytes: int) -> None: + self.maximum_bytes = maximum_bytes + self.bytes = 0 + self.entries: OrderedDict[str, bytes] = OrderedDict() + + def access(self, key: str, size: int) -> bool: + if key in self.entries: + self.entries.move_to_end(key) + return True + payload = bytes(size) + if size <= self.maximum_bytes: + self.entries[key] = payload + self.bytes += size + while self.bytes > self.maximum_bytes: + _, evicted = self.entries.popitem(last=False) + self.bytes -= len(evicted) + return False + + +def policy_ablation(records: list[dict[str, object]]) -> dict[str, object]: + rng = random.Random(77) + universe = records[:5000] + sizes = {str(record["tensor_ref"]): int(record["canonical_bytes"]) for record in universe} + hot = list(sizes)[:100] + cold = list(sizes)[100:] + # Warm every hot object before injecting scans so TinyLFU admission is + # measured directly rather than starting from an empty cache. + trace = hot * 5 + for cycle in range(30): + trace.extend(rng.choices(hot, weights=range(len(hot), 0, -1), k=1000)) + start = cycle * 100 % len(cold) + trace.extend(cold[start : start + 300]) + # The budget holds exactly the hot set but not the scan tail. + budget = sum(sizes[key] for key in hot) + lru = LegacyLru(budget) + gdsf = PayloadCache(budget) + lru_hits = 0 + gdsf_hits = 0 + for key in trace: + size = sizes[key] + lru_hits += int(lru.access(key, size)) + cached = gdsf.get((key,)) + gdsf_hits += int(cached is not None) + if cached is None: + gdsf.put((key,), bytes(size)) + return { + "accesses": len(trace), + "budget_bytes": budget, + "lru_hit_ratio": lru_hits / len(trace), + "tinylfu_gdsf_hit_ratio": gdsf_hits / len(trace), + "tinylfu_gdsf_status": gdsf.status(), + } + + +def encode_frame( + records: list[dict[str, object]], contract: str, query: np.ndarray, request_id: int +) -> bytes: + dtype = protocol.DTYPE_F16 if query.dtype == np.dtype(" tuple[float, list[tuple[int, float]]]: + started = time.perf_counter() + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_bytes = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_bytes) + elapsed = (time.perf_counter() - started) * 1000 + _, status, parsed = decode_response(response) + if status != 0 or not isinstance(parsed, list): + raise RuntimeError(f"daemon request failed: {parsed}") + if len(parsed) != len(protocol.parse_request_frame(frame).candidates): + raise RuntimeError("daemon returned an incomplete result") + return elapsed, parsed + + +def daemon_ablation( + records: list[dict[str, object]], + contract: str, + shard_root: Path, + rust_binary: Path, + device: int, + repeats: int, +) -> dict[str, object]: + rng = np.random.default_rng(31) + query = rng.standard_normal((44, int(records[0]["tensor_dim"]))).astype(" dict[str, object]: + shard_paths = sorted((shard_root / "shards").glob("*.vts")) + results = {} + with tempfile.TemporaryDirectory() as directory: + commands = { + "python_triton": [ + sys.executable, + "-m", + "services.tilemaxsim_cuda_sidecar", + "--socket", + f"{directory}/python-resident.sock", + "--gpu-memory-gb", + f"{device}=20", + "--gpu-workspace-gb", + "2", + "--host-cache-gb", + "0.1", + "--contract-root", + f"{contract}={shard_root}", + "--gpu-cache-mode", + "resident", + "--resident-manifest", + f"{contract}={descriptor_manifest}", + "--prewarm-batch-size", + "256", + "--request-timeout-ms", + "20000", + ], + "rust_cuda": [ + os.fspath(rust_binary), + "--socket", + f"{directory}/rust-resident.sock", + "--gpu-memory-gb", + f"{device}=20", + "--gpu-workspace-gb", + "2", + "--host-cache-gb", + "0.1", + "--contract-root", + f"{contract}={shard_root}", + "--gpu-cache-mode", + "resident", + "--resident-manifest", + f"{contract}={descriptor_manifest}", + "--prewarm-batch-size", + "256", + ], + } + for name, command in commands.items(): + os.sync() + evict_paths(shard_paths) + started = time.perf_counter() + process = subprocess.Popen( + command, + cwd=Path(__file__).resolve().parents[1], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + lines = [] + prewarm_event = None + ready_event = None + deadline = time.monotonic() + 300 + try: + assert process.stdout is not None + while time.monotonic() < deadline: + readable, _, _ = select.select([process.stdout], [], [], 1.0) + if readable: + line = process.stdout.readline() + if line: + lines.append(line) + if line.startswith("{"): + event = json.loads(line) + if event.get("event") in ( + "tilemaxsim_prewarm_complete", + "tilemaxsim_rust_prewarm_complete", + ): + prewarm_event = event + if event.get("event") in ( + "tilemaxsim_ready", + "tilemaxsim_rust_ready", + ): + ready_event = event + break + if process.poll() is not None: + break + if ready_event is None: + remainder, _ = process.communicate(timeout=5) + raise RuntimeError( + f"{name} resident prewarm failed: {''.join(lines)}{remainder}" + ) + results[name] = { + "process_to_ready_ms": (time.perf_counter() - started) * 1000, + "prewarm_reported_ms": prewarm_event.get("elapsed_ms") + if prewarm_event + else None, + "cache": ready_event.get("gpu_cache", ready_event.get("cache")), + } + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=15) + process.communicate(timeout=5) + return results + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--legacy-root", required=True, type=Path) + parser.add_argument("--shard-root", required=True, type=Path) + parser.add_argument("--rust-binary", required=True, type=Path) + parser.add_argument("--contract", default="benchmark@1") + parser.add_argument("--device", required=True, type=int) + parser.add_argument("--sample-size", type=int, default=100) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--full-prewarm", action="store_true") + args = parser.parse_args() + records = load_records(args.descriptor_manifest) + selected = random.Random(20260714).sample(records, args.sample_size) + report = { + "corpus": { + "records": len(records), + "logical_bytes": sum(int(record["canonical_bytes"]) for record in records), + "sample_size": len(selected), + }, + "storage": storage_ablation(selected, args.contract, args.legacy_root, args.shard_root), + "h2d": h2d_ablation(selected, args.contract, args.shard_root, args.device), + "allocator": allocator_ablation(records), + "policy": policy_ablation(records), + "daemon": daemon_ablation( + selected, + args.contract, + args.shard_root, + args.rust_binary, + args.device, + args.repeats, + ), + } + if args.full_prewarm: + report["full_resident_prewarm"] = prewarm_ablation( + args.descriptor_manifest, + args.contract, + args.shard_root, + args.rust_binary, + args.device, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as stream: + json.dump(report, stream, indent=2, sort_keys=True) + stream.write("\n") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/build_tilemaxsim_tensor_cache.py b/services/build_tilemaxsim_tensor_cache.py index a48a1728..e6adf7ad 100644 --- a/services/build_tilemaxsim_tensor_cache.py +++ b/services/build_tilemaxsim_tensor_cache.py @@ -22,12 +22,15 @@ import os import sys import tempfile +from collections import deque from concurrent.futures import ThreadPoolExecutor +from decimal import Decimal, InvalidOperation from pathlib import Path import numpy as np from services.tilemaxsim_cuda_sidecar import positive_int +from services.tilemaxsim_shard import DEFAULT_SHARD_BYTES, ImmutableShardWriter def canonical_tensor(path: Path, expected_rows: int, expected_dim: int) -> np.ndarray: @@ -96,13 +99,10 @@ def write_payload(path: Path, payload: memoryview, digest: str, fsync: bool) -> pass -def process_record( +def prepare_record( record: dict[str, object], source_root: Path, - cache_root: Path, - fsync: bool, - dry_run: bool, -) -> dict[str, object]: +) -> tuple[dict[str, object], memoryview]: page_key = record.get("page_key") relative = record.get("embedding_file") rows = record.get("n_tokens") @@ -125,9 +125,6 @@ def process_record( tensor = canonical_tensor(source, rows, dimension) payload = memoryview(tensor).cast("B") digest = hashlib.sha256(payload).hexdigest() - destination = cache_root / digest[:2] / f"{digest}.bin" - if not dry_run: - write_payload(destination, payload, digest, fsync) dtype_name = "float16" if tensor.dtype == np.dtype("float16") else "float32" return { "page_key": page_key, @@ -137,7 +134,57 @@ def process_record( "tensor_dtype": dtype_name, "tensor_checksum": f"sha256:{digest}", "canonical_bytes": len(payload), - } + }, payload + + +def process_record( + record: dict[str, object], + source_root: Path, + cache_root: Path, + fsync: bool, + dry_run: bool, +) -> dict[str, object]: + """Publish one legacy per-tensor file. + + Kept for backwards compatibility and migration tests. New cache builds use + immutable shards by default. + """ + + descriptor, payload = prepare_record(record, source_root) + digest = str(descriptor["tensor_checksum"]).removeprefix("sha256:") + destination = cache_root / digest[:2] / f"{digest}.bin" + if not dry_run: + write_payload(destination, payload, digest, fsync) + return descriptor + + +def shard_size_gb(value: str) -> int: + try: + parsed = int(Decimal(value) * 1024**3) + except (InvalidOperation, ValueError) as error: + raise argparse.ArgumentTypeError("shard size must be a positive number of GB") from error + if parsed <= 0: + raise argparse.ArgumentTypeError("shard size must be a positive number of GB") + return parsed + + +def bounded_parallel_map(executor, function, items, maximum_pending: int): + """Preserve input order without retaining a corpus-sized future backlog.""" + + iterator = iter(items) + pending = deque() + for _ in range(maximum_pending): + try: + pending.append(executor.submit(function, next(iterator))) + except StopIteration: + break + while pending: + future = pending.popleft() + yield future.result() + try: + pending.append(executor.submit(function, next(iterator))) + except StopIteration: + pass def main() -> None: @@ -146,6 +193,18 @@ def main() -> None: parser.add_argument("--cache-root", required=True, type=Path) parser.add_argument("--descriptor-manifest", required=True, type=Path) parser.add_argument("--workers", type=positive_int, default=4) + parser.add_argument( + "--storage-format", + choices=("shards", "files"), + default="shards", + help="publish immutable shards (default) or legacy per-tensor files", + ) + parser.add_argument( + "--shard-size-gb", + type=shard_size_gb, + default=DEFAULT_SHARD_BYTES, + help="target immutable shard size in GB", + ) parser.add_argument("--no-fsync", action="store_true") parser.add_argument("--dry-run", action="store_true") args = parser.parse_args() @@ -171,28 +230,63 @@ def main() -> None: cache_root.mkdir(parents=True, exist_ok=True) descriptors = [] - with ThreadPoolExecutor(max_workers=args.workers) as workers: - results = workers.map( - lambda record: process_record( - record, - source_root, - cache_root, - not args.no_fsync, - args.dry_run, - ), - records, + shard_writer = None + if args.storage_format == "shards" and not args.dry_run: + shard_writer = ImmutableShardWriter( + cache_root, + target_bytes=args.shard_size_gb, + fsync=not args.no_fsync, ) - for completed, item in enumerate(results, 1): - descriptors.append(item) - if completed % 1000 == 0 or completed == len(records): - print( - json.dumps( - {"event": "tensor_cache_progress", "completed": completed}, - separators=(",", ":"), - ), - file=sys.stderr, - flush=True, + try: + with ThreadPoolExecutor(max_workers=args.workers) as workers: + if args.storage_format == "files": + results = ( + (item, None) + for item in bounded_parallel_map( + workers, + lambda record: process_record( + record, + source_root, + cache_root, + not args.no_fsync, + args.dry_run, + ), + records, + args.workers * 2, + ) ) + else: + results = bounded_parallel_map( + workers, + lambda record: prepare_record(record, source_root), + records, + args.workers * 2, + ) + for completed, (item, payload) in enumerate(results, 1): + descriptors.append(item) + if shard_writer is not None: + assert payload is not None + digest = str(item["tensor_checksum"]).removeprefix("sha256:") + shard_writer.add( + digest, + payload, + int(item["tensor_rows"]), + int(item["tensor_dim"]), + str(item["tensor_dtype"]), + ) + if completed % 1000 == 0 or completed == len(records): + print( + json.dumps( + {"event": "tensor_cache_progress", "completed": completed}, + separators=(",", ":"), + ), + file=sys.stderr, + flush=True, + ) + shard_index = shard_writer.finish() if shard_writer is not None else None + finally: + if shard_writer is not None: + shard_writer.close() output_parent = args.descriptor_manifest.parent output_parent.mkdir(parents=True, exist_ok=True) @@ -224,6 +318,8 @@ def main() -> None: "canonical_bytes": total_bytes, "cache_root": os.fspath(cache_root), "descriptor_manifest": os.fspath(args.descriptor_manifest), + "storage_format": args.storage_format, + "shard_index": os.fspath(shard_index) if shard_index else None, "dry_run": args.dry_run, }, sort_keys=True, diff --git a/services/test_tilemaxsim_gpu_cache.py b/services/test_tilemaxsim_gpu_cache.py index 505d13b7..ae228ac2 100644 --- a/services/test_tilemaxsim_gpu_cache.py +++ b/services/test_tilemaxsim_gpu_cache.py @@ -23,10 +23,12 @@ from devtools import tilemaxsim_reference_sidecar as protocol from services.tilemaxsim_cuda_sidecar import ResidentTorchTileMaxsimEngine from services.tilemaxsim_gpu_cache import ( + FixedBlockAllocator, FreeExtentAllocator, GpuArenaSpec, GpuResourcePool, GpuTensorCache, + GpuTensorLoad, parse_gpu_memory_gb, parse_memory_gb, ) @@ -68,6 +70,21 @@ def test_extent_allocator_coalesces_released_ranges(self) -> None: allocator.release(*second) self.assertEqual(allocator.extents, [(0, 4096)]) + def test_fixed_block_buddy_allocator_coalesces_slabs(self) -> None: + allocator = FixedBlockAllocator(8 * 256, block_bytes=256) + first = allocator.allocate(300) + second = allocator.allocate(300) + third = allocator.allocate(700) + self.assertEqual(first, (0, 1)) + self.assertEqual(second, (2, 3)) + self.assertEqual(third, (4, 5, 6, 7)) + assert first is not None and second is not None and third is not None + allocator.release(second) + allocator.release(first) + self.assertEqual(allocator.largest_free_extent, 4 * 256) + allocator.release(third) + self.assertEqual(allocator.largest_free_extent, 8 * 256) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") def test_pool_rejects_budget_larger_than_currently_free_memory(self) -> None: device = available_device() @@ -162,6 +179,76 @@ def test_resident_engine_scores_gpu_cache_handles(self) -> None: finally: pool.close() + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_batch_admission_uses_one_h2d_batch(self) -> None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 4 * 1024 * 1024)], 2 * 1024 * 1024 + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + first = np.ones((128, 320), dtype=" None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 768 * 1024)], 512 * 1024 + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + tensor = np.ones((128, 320), dtype=" None: device = available_device() diff --git a/services/test_tilemaxsim_rust_daemon.py b/services/test_tilemaxsim_rust_daemon.py new file mode 100644 index 00000000..630ff651 --- /dev/null +++ b/services/test_tilemaxsim_rust_daemon.py @@ -0,0 +1,254 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2025-2026 TensorChord Inc. + +from __future__ import annotations + +import hashlib +import json +import os +import socket +import subprocess +import tempfile +import time +import unittest +from pathlib import Path + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from devtools.test_tilemaxsim_reference_sidecar import ( + decode_response, + external_request_frame, +) +from services.tilemaxsim_shard import ImmutableShardWriter + + +class RustDaemonTest(unittest.TestCase): + def run_daemon( + self, + devices: list[int], + documents: list[np.ndarray] | None = None, + query: np.ndarray | None = None, + gpu_memory_gb: str = "0.05", + workspace_gb: str = "0.02", + resident: bool = False, + ) -> tuple[str, list[tuple[int, float]]]: + binary = ( + Path(__file__).parent + / "tilemaxsimd" + / "target" + / "release" + / "tilemaxsimd" + ) + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shard_root = root / "cache" + shard_root.mkdir() + if documents is None: + documents = [ + np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + self.run_daemon([device]) + + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.device_count() >= 2, + "two CUDA devices are unavailable", + ) + def test_multi_gpu_scheduler_uploads_and_scores_on_each_device(self) -> None: + output, _ = self.run_daemon([0, 1]) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + request_event = next( + event for event in events if event.get("event") == "tilemaxsim_rust_request" + ) + devices = request_event["cache"]["devices"] + self.assertEqual(len(devices), 2) + self.assertEqual([device["h2d_batches"] for device in devices], [1, 1]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_one_request_larger_than_gpu_cache_is_scored_in_chunks(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + rows, dimension = 480, 320 + first = np.zeros((rows, dimension), dtype=" None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + output, _ = self.run_daemon([device], resident=True) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + prewarm = next( + event + for event in events + if event.get("event") == "tilemaxsim_rust_prewarm_complete" + ) + self.assertEqual(prewarm["entries"], 2) + self.assertEqual(prewarm["cache"]["devices"][0]["gpu_pinned_entries"], 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/services/test_tilemaxsim_shard.py b/services/test_tilemaxsim_shard.py new file mode 100644 index 00000000..e5684375 --- /dev/null +++ b/services/test_tilemaxsim_shard.py @@ -0,0 +1,181 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2025-2026 TensorChord Inc. + +from __future__ import annotations + +import hashlib +import json +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import ContentAddressedResolver +from services.tilemaxsim_shard import ImmutableShardWriter, load_index + + +class ImmutableShardTest(unittest.TestCase): + def test_builder_defaults_to_shards_and_publishes_atomically(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source" + source.mkdir() + manifest = source / "pages.jsonl" + records = [] + for index in range(3): + tensor = np.full((2, 4), index + 1, dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + tensors = [ + np.arange(32, dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + tensor = np.eye(4, dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = np.ones((2, 4), dtype=" None: self.maximum_bytes = maximum_bytes self.current_bytes = 0 - self.entries: OrderedDict[tuple[object, ...], bytes] = OrderedDict() + self.entries: OrderedDict[ + tuple[object, ...], tuple[bytes, float] + ] = OrderedDict() self.lock = threading.Lock() + self.sketch = TinyLfuSketch() + self.inflation = 0.0 + self.hits = 0 + self.misses = 0 + self.evictions = 0 + self.admission_rejections = 0 def get(self, key: tuple[object, ...]) -> bytes | None: if self.maximum_bytes == 0: return None with self.lock: - payload = self.entries.get(key) - if payload is not None: + frequency = self.sketch.increment(key) + entry = self.entries.get(key) + if entry is not None: + payload, _ = entry + priority = self.inflation + frequency / len(payload) + self.entries[key] = (payload, priority) self.entries.move_to_end(key) - return payload + self.hits += 1 + return payload + self.misses += 1 + return None def put(self, key: tuple[object, ...], payload: bytes) -> None: if self.maximum_bytes == 0 or len(payload) > self.maximum_bytes: @@ -113,12 +140,39 @@ def put(self, key: tuple[object, ...], payload: bytes) -> None: with self.lock: previous = self.entries.pop(key, None) if previous is not None: - self.current_bytes -= len(previous) - self.entries[key] = payload + self.current_bytes -= len(previous[0]) + frequency = max(1, self.sketch.estimate(key)) + candidate_priority = self.inflation + frequency / len(payload) + while self.current_bytes + len(payload) > self.maximum_bytes: + victim_key, (victim_payload, victim_priority) = min( + self.entries.items(), key=lambda item: item[1][1] + ) + if candidate_priority < victim_priority: + if previous is not None: + self.entries[key] = previous + self.current_bytes += len(previous[0]) + self.admission_rejections += 1 + return + self.entries.pop(victim_key) + self.current_bytes -= len(victim_payload) + self.inflation = max(self.inflation, victim_priority) + candidate_priority = self.inflation + frequency / len(payload) + self.evictions += 1 + self.entries[key] = (payload, candidate_priority) self.current_bytes += len(payload) - while self.current_bytes > self.maximum_bytes: - _, evicted = self.entries.popitem(last=False) - self.current_bytes -= len(evicted) + + def status(self) -> dict[str, object]: + with self.lock: + return { + "entries": len(self.entries), + "bytes": self.current_bytes, + "maximum_bytes": self.maximum_bytes, + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "admission_rejections": self.admission_rejections, + "policy": "tinylfu-gdsf", + } class ContentAddressedResolver: @@ -131,20 +185,85 @@ class ContentAddressedResolver: agree before a payload is returned. """ - def __init__(self, roots: dict[str, Path], cache_bytes: int) -> None: + def __init__( + self, + roots: dict[str, Path], + cache_bytes: int, + verify_full_shards: bool = False, + ) -> None: self.root_fds: dict[str, int] = {} + self.shard_roots: dict[str, _OpenShardRoot] = {} + self.batch_read_calls = 0 + self.batch_read_bytes = 0 + self.batch_lock = threading.Lock() + self.verify_full_shards = verify_full_shards try: for contract, path in roots.items(): - self.root_fds[contract] = os.open( + root_fd = os.open( os.fspath(path), os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, ) + self.root_fds[contract] = root_fd + self._open_shard_root(contract, root_fd) except Exception: self.close() raise self.cache = PayloadCache(cache_bytes) + def _open_shard_root(self, contract: str, root_fd: int) -> None: + try: + index_fd = os.open( + INDEX_NAME, + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + except FileNotFoundError: + return + try: + with os.fdopen(index_fd, "r", encoding="utf-8", closefd=True) as stream: + index = parse_index(json.load(stream)) + except Exception: + raise + shard_directory_fd = -1 + shard_fds: dict[str, int] = {} + try: + shard_directory_fd = os.open( + "shards", + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + for relative, shard in index.shards.items(): + name = relative.removeprefix("shards/") + descriptor = os.open( + name, + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=shard_directory_fd, + ) + metadata = os.fstat(descriptor) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_size != shard.size: + os.close(descriptor) + raise ValueError(f"immutable shard metadata disagrees: {relative}") + shard_fds[relative] = descriptor + self.shard_roots[contract] = _OpenShardRoot( + index, + shard_directory_fd, + shard_fds, + {name: threading.Lock() for name in shard_fds}, + set(), + ) + except Exception: + for descriptor in shard_fds.values(): + os.close(descriptor) + if shard_directory_fd >= 0: + os.close(shard_directory_fd) + raise + def close(self) -> None: + for root in self.shard_roots.values(): + for descriptor in root.shard_fds.values(): + os.close(descriptor) + os.close(root.directory_fd) + self.shard_roots.clear() for descriptor in self.root_fds.values(): os.close(descriptor) self.root_fds.clear() @@ -247,25 +366,206 @@ def key(self, request: protocol.ExternalTensorRequest) -> tuple[object, ...]: ) return key - def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: - key = self.key(request) - digest = str(key[1]) - root_fd = self.root_fds[request.model_contract_id] - expected_bytes = protocol.checked_tensor_bytes( - request.rows, request.dimension, request.dtype + @staticmethod + def _validate_shard_entry( + request: protocol.ExternalTensorRequest, digest: str, entry: ShardEntry + ) -> None: + dtype_name = ( + "float32" if request.dtype == protocol.DTYPE_F32 else "float16" ) - cached = self.cache.get(key) - if cached is not None: - return ResolvedPayload(cached, True) - payload = self._read_exact_file(root_fd, digest, expected_bytes) - actual = hashlib.sha256(payload).hexdigest() - if not hmac.compare_digest(actual, digest): + if ( + entry.digest != digest + or entry.rows != request.rows + or entry.dimension != request.dimension + or entry.dtype != dtype_name + or entry.length + != protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + ): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "shard entry disagrees with the tensor descriptor", + ) + + @staticmethod + def _verify_shard(root: _OpenShardRoot, name: str) -> None: + if name in root.verified: + return + lock = root.verification_locks[name] + with lock: + if name in root.verified: + return + shard = root.index.shards[name] + expected = shard.checksum.removeprefix("sha256:") + digest = hashlib.sha256() + offset = 0 + descriptor = root.shard_fds[name] + while offset < shard.size: + chunk = os.pread(descriptor, min(8 * 1024 * 1024, shard.size - offset), offset) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "immutable tensor shard ended early", + ) + digest.update(chunk) + offset += len(chunk) + if not hmac.compare_digest(digest.hexdigest(), expected): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "immutable tensor shard checksum mismatch", + ) + root.verified.add(name) + + @staticmethod + def _coalesced_ranges( + entries: list[tuple[tuple[object, ...], ShardEntry]], + maximum_gap: int = 64 * 1024, + maximum_span: int = 8 * 1024 * 1024, + ) -> list[tuple[int, int, list[tuple[tuple[object, ...], ShardEntry]]]]: + ranges = [] + current: list[tuple[tuple[object, ...], ShardEntry]] = [] + start = 0 + end = 0 + for item in sorted(entries, key=lambda value: value[1].offset): + entry = item[1] + entry_end = entry.offset + entry.length + if current and ( + entry.offset - end > maximum_gap or entry_end - start > maximum_span + ): + ranges.append((start, end, current)) + current = [] + if not current: + start = entry.offset + end = entry_end + else: + end = max(end, entry_end) + current.append(item) + if current: + ranges.append((start, end, current)) + return ranges + + def _read_shard_range( + self, + root: _OpenShardRoot, + shard_name: str, + start: int, + end: int, + entries: list[tuple[tuple[object, ...], ShardEntry]], + ) -> dict[tuple[object, ...], bytes]: + payload = os.pread(root.shard_fds[shard_name], end - start, start) + if len(payload) != end - start: raise protocol.SidecarError( - protocol.STATUS_INVALID_REQUEST, "resolved tensor checksum mismatch" + protocol.STATUS_COMPUTE_ERROR, "immutable tensor shard ended early" + ) + with self.batch_lock: + self.batch_read_calls += 1 + self.batch_read_bytes += len(payload) + result = {} + for key, entry in entries: + tensor = payload[entry.offset - start : entry.offset - start + entry.length] + actual = hashlib.sha256(tensor).hexdigest() + if not hmac.compare_digest(actual, entry.digest): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved shard tensor checksum mismatch", + ) + dtype = protocol.DTYPE_F32 if entry.dtype == "float32" else protocol.DTYPE_F16 + validate_finite_payload(tensor, entry.rows, entry.dimension, dtype) + result[key] = tensor + return result + + def resolve_many( + self, requests: list[protocol.ExternalTensorRequest] + ) -> list[ResolvedPayload]: + if not requests: + return [] + keys = [self.key(request) for request in requests] + payloads: dict[tuple[object, ...], bytes] = {} + hits: dict[tuple[object, ...], bool] = {} + missing: dict[tuple[object, ...], protocol.ExternalTensorRequest] = {} + for key, request in zip(keys, requests, strict=True): + cached = self.cache.get(key) + if cached is not None: + payloads[key] = cached + hits[key] = True + elif key not in missing: + missing[key] = request + hits[key] = False + + shard_groups: dict[tuple[str, str], list[tuple[tuple[object, ...], ShardEntry]]] = {} + legacy: list[tuple[tuple[object, ...], protocol.ExternalTensorRequest]] = [] + for key, request in missing.items(): + contract = request.model_contract_id + shard_root = self.shard_roots.get(contract) + if shard_root is None: + legacy.append((key, request)) + continue + digest = str(key[1]) + entry = shard_root.index.entries.get(digest) + if entry is None: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "content-addressed tensor is missing from the shard index", + ) + self._validate_shard_entry(request, digest, entry) + shard_groups.setdefault((contract, entry.shard), []).append((key, entry)) + + jobs: list[tuple[_OpenShardRoot, str, int, int, list[tuple[tuple[object, ...], ShardEntry]]]] = [] + for (contract, shard_name), entries in shard_groups.items(): + root = self.shard_roots[contract] + if self.verify_full_shards: + self._verify_shard(root, shard_name) + for start, end, grouped in self._coalesced_ranges(entries): + jobs.append((root, shard_name, start, end, grouped)) + if jobs: + with ThreadPoolExecutor(max_workers=min(8, len(jobs))) as workers: + for resolved in workers.map(lambda job: self._read_shard_range(*job), jobs): + payloads.update(resolved) + + def read_legacy( + item: tuple[tuple[object, ...], protocol.ExternalTensorRequest] + ) -> tuple[tuple[object, ...], bytes]: + key, request = item + digest = str(key[1]) + expected_bytes = protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype ) - validate_finite_payload(payload, request.rows, request.dimension, request.dtype) - self.cache.put(key, payload) - return ResolvedPayload(payload, False) + payload = self._read_exact_file( + self.root_fds[request.model_contract_id], digest, expected_bytes + ) + actual = hashlib.sha256(payload).hexdigest() + if not hmac.compare_digest(actual, digest): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor checksum mismatch", + ) + validate_finite_payload( + payload, request.rows, request.dimension, request.dtype + ) + return key, payload + + if legacy: + with ThreadPoolExecutor(max_workers=min(8, len(legacy))) as workers: + for key, payload in workers.map(read_legacy, legacy): + payloads[key] = payload + for key in missing: + self.cache.put(key, payloads[key]) + return [ResolvedPayload(payloads[key], hits[key]) for key in keys] + + def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: + return self.resolve_many([request])[0] + + def status(self) -> dict[str, object]: + with self.batch_lock: + return { + "shard_contracts": len(self.shard_roots), + "verified_shards": sum( + len(root.verified) for root in self.shard_roots.values() + ), + "batch_read_calls": self.batch_read_calls, + "batch_read_bytes": self.batch_read_bytes, + } class TorchTileMaxsimEngine: @@ -751,15 +1051,28 @@ def flush_resident_documents() -> None: metrics["source"] = "inline" else: metrics["source"] = "content_addressed" - for candidate in request.candidates: - document_tokens += candidate.descriptor.rows - if time.monotonic() >= deadline: - raise protocol.SidecarError( - protocol.STATUS_COMPUTE_ERROR, - "request deadline expired during tensor resolution", - ) - if self.gpu_cache is None: - resolved = self.resolver.resolve(candidate.descriptor) + document_tokens = sum( + candidate.descriptor.rows for candidate in request.candidates + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during tensor resolution", + ) + if self.gpu_cache is None: + descriptors = [ + candidate.descriptor for candidate in request.candidates + ] + if hasattr(self.resolver, "resolve_many"): + resolved_payloads = self.resolver.resolve_many(descriptors) + else: + resolved_payloads = [ + self.resolver.resolve(descriptor) + for descriptor in descriptors + ] + for candidate, resolved in zip( + request.candidates, resolved_payloads, strict=True + ): cache_hits += int(resolved.cache_hit) documents.append( ( @@ -768,45 +1081,109 @@ def flush_resident_documents() -> None: resolved.payload, ) ) + else: + keys = [ + self.resolver.key(candidate.descriptor) + for candidate in request.candidates + ] + probed, miss_indices = self.gpu_cache.probe_many(keys) + for candidate, handle in zip( + request.candidates, probed, strict=True + ): + if handle is not None: + resident_documents.append((candidate.candidate_id, handle)) + gpu_cache_hits = len(request.candidates) - len(miss_indices) + gpu_cache_misses = len(miss_indices) + missing_candidates = [ + request.candidates[index] for index in miss_indices + ] + missing_descriptors = [ + candidate.descriptor for candidate in missing_candidates + ] + if hasattr(self.resolver, "resolve_many"): + resolved_payloads = self.resolver.resolve_many( + missing_descriptors + ) else: - key = self.resolver.key(candidate.descriptor) - loaded_payload: bytes | None = None - - def load_payload() -> bytes: - nonlocal cache_hits, loaded_payload - if loaded_payload is None: - resolved = self.resolver.resolve(candidate.descriptor) - cache_hits += int(resolved.cache_hit) - loaded_payload = resolved.payload - return loaded_payload - - while True: - try: - handle, gpu_hit = self.gpu_cache.acquire( - key, - candidate.descriptor.rows, - candidate.descriptor.dimension, - candidate.descriptor.dtype, - load_payload, - pin=self.pin_gpu_entries, + resolved_payloads = [ + self.resolver.resolve(descriptor) + for descriptor in missing_descriptors + ] + cache_hits += sum(item.cache_hit for item in resolved_payloads) + pending = [ + ( + candidate, + GpuTensorLoad( + key, + candidate.descriptor.rows, + candidate.descriptor.dimension, + candidate.descriptor.dtype, + resolved.payload, + self.pin_gpu_entries, + ), + False, + ) + for candidate, key, resolved in zip( + missing_candidates, + (keys[index] for index in miss_indices), + resolved_payloads, + strict=True, + ) + ] + admission_rejections = 0 + while pending: + batch = self.gpu_cache.acquire_many( + [item[1] for item in pending], + enforce_admission=( + not self.pin_gpu_entries + and not any(item[2] for item in pending) + ), + record_access=False, + count_stats=False, + ) + for (candidate, load, _force), handle in zip( + pending, batch.handles, strict=True + ): + if handle is not None: + resident_documents.append( + (candidate.candidate_id, handle) + ) + for index in batch.bypassed: + candidate, load, _force = pending[index] + stream_bytes = ( + request.query_rows * request.dimension * 4 + + load.rows * request.dimension * 4 + + request.query_rows * load.rows * 4 + ) + if stream_bytes <= self.engine.max_device_bytes: + documents.append( + (candidate.candidate_id, load.rows, load.payload) ) - break - except protocol.SidecarError as error: - if ( - error.status != protocol.STATUS_RESOURCE_LIMIT - or not resident_documents - ): - raise - # A request may be larger than the configured - # GPU cache. Score and release the current - # working set, then admit the remaining - # candidates through the same bounded arenas. - flush_resident_documents() - gpu_cache_hits += int(gpu_hit) - gpu_cache_misses += int(not gpu_hit) - resident_documents.append((candidate.candidate_id, handle)) - if self.gpu_cache is not None: - flush_resident_documents() + admission_rejections += len(batch.bypassed) + deferred = [pending[index] for index in batch.deferred] + forced = [ + (pending[index][0], pending[index][1], True) + for index in batch.bypassed + if ( + request.query_rows * request.dimension * 4 + + pending[index][1].rows * request.dimension * 4 + + request.query_rows * pending[index][1].rows * 4 + > self.engine.max_device_bytes + ) + ] + made_progress = len(deferred) < len(pending) + if resident_documents: + flush_resident_documents() + made_progress = True + if deferred and not made_progress: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "GPU block cache cannot make progress with its configured capacity", + ) + pending = forced + deferred + if resident_documents: + flush_resident_documents() + metrics["gpu_admission_rejections"] = admission_rejections metrics["cache_hits"] = cache_hits metrics["host_cache_hits"] = cache_hits metrics["gpu_cache_hits"] = gpu_cache_hits @@ -828,6 +1205,22 @@ def load_payload() -> bytes: results = resident_results queue_ms = resident_queue_ms compute_ms = resident_compute_ms + if documents: + stream_results, stream_queue_ms, stream_compute_ms = ( + self.engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + ) + results.extend(stream_results) + queue_ms += stream_queue_ms + compute_ms += stream_compute_ms + metrics["streamed_candidates"] = len(documents) else: results, queue_ms, compute_ms = self.engine.score( request.query_payload, @@ -1082,10 +1475,67 @@ def prewarm_resident_cache( resolver: ContentAddressedResolver, gpu_cache: GpuTensorCache, metrics: JsonMetrics, + batch_size: int = 256, ) -> None: completed = 0 loaded_bytes = 0 started = time.monotonic() + pending: list[tuple[protocol.ExternalTensorRequest, int]] = [] + + def flush() -> None: + nonlocal completed, loaded_bytes + if not pending: + return + keys = [resolver.key(request) for request, _ in pending] + probed, miss_indices = gpu_cache.probe_many(keys) + acquired = [handle for handle in probed if handle is not None] + try: + missing = [pending[index][0] for index in miss_indices] + resolved = resolver.resolve_many(missing) + loads = [ + GpuTensorLoad( + keys[index], + request.rows, + request.dimension, + request.dtype, + payload.payload, + True, + ) + for index, request, payload in zip( + miss_indices, missing, resolved, strict=True + ) + ] + batch = gpu_cache.acquire_many( + loads, + enforce_admission=False, + record_access=False, + count_stats=False, + ) + if batch.bypassed or batch.deferred or any( + handle is None for handle in batch.handles + ): + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "resident manifest exceeds the configured GPU block arenas", + ) + acquired.extend( + handle for handle in batch.handles if handle is not None + ) + finally: + for handle in acquired: + gpu_cache.release(handle) + completed += len(pending) + loaded_bytes += sum(size for _, size in pending) + if completed % 1000 < len(pending): + metrics.emit( + { + "event": "tilemaxsim_prewarm_progress", + "entries": completed, + "logical_bytes": loaded_bytes, + } + ) + pending.clear() + for contract, path in manifests: with path.open(encoding="utf-8") as stream: for line_number, line in enumerate(stream, 1): @@ -1123,26 +1573,10 @@ def prewarm_resident_cache( raise ValueError( f"{path}:{line_number}: canonical_bytes disagrees with shape" ) - key = resolver.key(request) - handle, _ = gpu_cache.acquire( - key, - rows, - dimension, - dtype, - lambda request=request: resolver.resolve(request).payload, - pin=True, - ) - gpu_cache.release(handle) - completed += 1 - loaded_bytes += expected_bytes - if completed % 1000 == 0: - metrics.emit( - { - "event": "tilemaxsim_prewarm_progress", - "entries": completed, - "logical_bytes": loaded_bytes, - } - ) + pending.append((request, expected_bytes)) + if len(pending) >= batch_size: + flush() + flush() if completed == 0: raise ValueError("resident manifests contain no tensor descriptors") metrics.emit( @@ -1204,8 +1638,14 @@ def main() -> None: parser.add_argument("--request-timeout-ms", type=positive_int, default=2000) parser.add_argument("--max-inflight", type=positive_int, default=8) parser.add_argument("--max-cuda-inflight", type=positive_int, default=1) + parser.add_argument("--prewarm-batch-size", type=positive_int, default=256) parser.add_argument("--backlog", type=positive_int, default=64) parser.add_argument("--allow-tf32", action="store_true") + parser.add_argument( + "--verify-full-shards", + action="store_true", + help="verify complete immutable shard digests lazily in addition to every tensor digest", + ) parser.add_argument("--once", action="store_true") args = parser.parse_args() @@ -1227,7 +1667,9 @@ def main() -> None: max_tensor_bytes=args.max_tensor_bytes, max_candidates=args.max_candidates, ) - resolver = ContentAddressedResolver(roots, args.host_cache_gb) + resolver = ContentAddressedResolver( + roots, args.host_cache_gb, args.verify_full_shards + ) metrics = JsonMetrics() pool: GpuResourcePool | None = None try: @@ -1246,7 +1688,13 @@ def main() -> None: args.max_cuda_inflight, ) if args.gpu_cache_mode == "resident": - prewarm_resident_cache(manifests, resolver, gpu_cache, metrics) + prewarm_resident_cache( + manifests, + resolver, + gpu_cache, + metrics, + args.prewarm_batch_size, + ) service = TileMaxsimService( limits, resolver, diff --git a/services/tilemaxsim_gpu_cache.py b/services/tilemaxsim_gpu_cache.py index 43d66d97..65a90271 100644 --- a/services/tilemaxsim_gpu_cache.py +++ b/services/tilemaxsim_gpu_cache.py @@ -21,8 +21,11 @@ from collections import OrderedDict from dataclasses import dataclass from decimal import Decimal, InvalidOperation +from hashlib import blake2b +from math import ceil from typing import Callable +import numpy as np import torch from devtools import tilemaxsim_reference_sidecar as protocol @@ -31,6 +34,8 @@ _GPU_MEMORY_GB = re.compile(r"^(?:cuda:)?([0-9]+)=([0-9]+(?:\.[0-9]+)?)$") _MEMORY_GB = re.compile(r"^[0-9]+(?:\.[0-9]+)?$") GIB = 1024**3 +DEFAULT_BLOCK_BYTES = 256 * 1024 +DEFAULT_STAGING_BYTES = 64 * 1024 * 1024 @dataclass(frozen=True) @@ -125,10 +130,159 @@ def release(self, start: int, length: int) -> None: self.extents = merged +class FixedBlockAllocator: + """Fixed-base-block buddy/slab allocator used by resident tensors. + + A tensor gets one contiguous power-of-two slab. This preserves the dense + memory layout required by the Tensor Core kernel while bounding internal + fragmentation and guaranteeing that released buddies coalesce. + """ + + def __init__(self, capacity: int, block_bytes: int = DEFAULT_BLOCK_BYTES) -> None: + if capacity <= 0 or block_bytes <= 0: + raise ValueError("arena capacity and block size must be positive") + if block_bytes % 256: + raise ValueError("GPU block size must be 256-byte aligned") + self.block_bytes = block_bytes + self.block_count = capacity // block_bytes + if self.block_count == 0: + raise ValueError("arena must contain at least one GPU block") + self.capacity = self.block_count * block_bytes + self._free: dict[int, set[tuple[int, int, int]]] = {} + self._allocated: dict[int, tuple[int, int, int]] = {} + start = 0 + remaining = self.block_count + while remaining: + order = remaining.bit_length() - 1 + size = 1 << order + self._free.setdefault(order, set()).add((start, start, order)) + start += size + remaining -= size + + @property + def free_blocks(self) -> int: + return sum(len(items) * (1 << order) for order, items in self._free.items()) + + @property + def free_bytes(self) -> int: + return self.free_blocks * self.block_bytes + + @property + def largest_free_extent(self) -> int: + return max( + (1 << order) * self.block_bytes + for order, items in self._free.items() + if items + ) if self.free_blocks else 0 + + def blocks_for(self, payload_bytes: int) -> int: + if payload_bytes <= 0: + raise ValueError("payload size must be positive") + raw = ceil(payload_bytes / self.block_bytes) + return 1 << (raw - 1).bit_length() + + def allocation_bytes(self, payload_bytes: int) -> int: + return self.blocks_for(payload_bytes) * self.block_bytes + + def allocate(self, payload_bytes: int) -> tuple[int, ...] | None: + required = self.blocks_for(payload_bytes) + order = required.bit_length() - 1 + available_order = next( + ( + candidate + for candidate in range(order, self.block_count.bit_length()) + if self._free.get(candidate) + ), + None, + ) + if available_order is None: + return None + start, root_start, root_order = self._free[available_order].pop() + while available_order > order: + available_order -= 1 + buddy = start + (1 << available_order) + self._free.setdefault(available_order, set()).add( + (buddy, root_start, root_order) + ) + blocks = tuple(range(start, start + required)) + self._allocated[start] = (order, root_start, root_order) + return blocks + + def release(self, blocks: tuple[int, ...]) -> None: + if ( + not blocks + or len(set(blocks)) != len(blocks) + or blocks != tuple(range(blocks[0], blocks[0] + len(blocks))) + ): + raise ValueError("released GPU blocks are invalid") + allocation = self._allocated.pop(blocks[0], None) + if allocation is None: + raise ValueError("GPU block was released more than once") + order, root_start, root_order = allocation + if len(blocks) != 1 << order: + raise ValueError("released GPU slab has the wrong size") + start = blocks[0] + while order < root_order: + buddy = root_start + (((start - root_start) ^ (1 << order))) + item = (buddy, root_start, root_order) + free = self._free.setdefault(order, set()) + if item not in free: + break + free.remove(item) + start = min(start, buddy) + order += 1 + self._free.setdefault(order, set()).add((start, root_start, root_order)) + + +class TinyLfuSketch: + """A small aging count-min sketch for cache admission and GDSF frequency.""" + + def __init__(self, width: int = 4096, depth: int = 4) -> None: + if width <= 0 or depth <= 0: + raise ValueError("TinyLFU dimensions must be positive") + self.width = width + self.depth = depth + self.tables = [[0] * width for _ in range(depth)] + self.samples = 0 + self.reset_at = width * 10 + + def _indices(self, key: tuple[object, ...]) -> tuple[int, ...]: + digest = blake2b(repr(key).encode("utf-8"), digest_size=16).digest() + return tuple( + int.from_bytes(digest[row * 4 : row * 4 + 4], "little") % self.width + for row in range(self.depth) + ) + + def increment(self, key: tuple[object, ...]) -> int: + indices = self._indices(key) + for row, index in enumerate(indices): + if self.tables[row][index] < 65535: + self.tables[row][index] += 1 + self.samples += 1 + estimate = min(self.tables[row][index] for row, index in enumerate(indices)) + if self.samples >= self.reset_at: + for table in self.tables: + for index, value in enumerate(table): + table[index] = value // 2 + self.samples //= 2 + return max(1, estimate) + + def estimate(self, key: tuple[object, ...]) -> int: + return min( + self.tables[row][index] + for row, index in enumerate(self._indices(key)) + ) + + class GpuArena: """A CUDA byte buffer acquired atomically during process startup.""" - def __init__(self, spec: GpuArenaSpec, workspace_bytes: int) -> None: + def __init__( + self, + spec: GpuArenaSpec, + workspace_bytes: int, + block_bytes: int = DEFAULT_BLOCK_BYTES, + ) -> None: self.device = torch.device(spec.device) if not torch.cuda.is_available(): raise RuntimeError( @@ -142,12 +296,18 @@ def __init__(self, spec: GpuArenaSpec, workspace_bytes: int) -> None: ) self.total_bytes = spec.total_bytes self.workspace_bytes = workspace_bytes - self.capacity = (spec.total_bytes - workspace_bytes) // 256 * 256 + raw_capacity = spec.total_bytes - workspace_bytes + self.allocator = FixedBlockAllocator(raw_capacity, block_bytes) + self.capacity = self.allocator.capacity self.reserved_workspace_bytes = spec.total_bytes - self.capacity if self.capacity <= 0: raise RuntimeError(f"{spec.device} has no aligned tensor-cache capacity") self.storage: torch.Tensor | None = None - self.allocator = FreeExtentAllocator(self.capacity) + self.host_staging: torch.Tensor | None = None + self.copy_stream: torch.cuda.Stream | None = None + self.h2d_batches = 0 + self.h2d_copy_calls = 0 + self.h2d_bytes = 0 with torch.cuda.device(self.device): free_bytes, device_bytes = torch.cuda.mem_get_info(self.device) @@ -172,28 +332,109 @@ def __init__(self, spec: GpuArenaSpec, workspace_bytes: int) -> None: ) torch.cuda.synchronize(self.device) del workspace + staging_bytes = min(self.capacity, DEFAULT_STAGING_BYTES) + staging_bytes = max( + self.allocator.block_bytes, + staging_bytes // self.allocator.block_bytes * self.allocator.block_bytes, + ) + self.host_staging = torch.empty( + staging_bytes, dtype=torch.uint8, pin_memory=True + ) + # Fault and pin every staging page during startup so the first + # cache-miss batch does not pay a request-path NUMA/page cost. + self.host_staging.zero_() + self.copy_stream = torch.cuda.Stream(device=self.device) except Exception: self.storage = None + self.host_staging = None + self.copy_stream = None torch.cuda.empty_cache() raise def tensor( - self, offset: int, payload_bytes: int, rows: int, dimension: int, dtype: int + self, + blocks: tuple[int, ...], + payload_bytes: int, + rows: int, + dimension: int, + dtype: int, ) -> torch.Tensor: - if self.storage is None: + if self.storage is None or self.host_staging is None or self.copy_stream is None: raise RuntimeError("GPU arena is closed") scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 - return ( - self.storage.narrow(0, offset, payload_bytes) - .view(scalar_dtype) - .reshape(rows, dimension) - ) - - def copy_from_host(self, offset: int, payload: bytes) -> None: + block_bytes = self.allocator.block_bytes + if all(right == left + 1 for left, right in zip(blocks, blocks[1:])): + raw = self.storage.narrow(0, blocks[0] * block_bytes, payload_bytes) + else: + raw = torch.cat( + [ + self.storage.narrow(0, block * block_bytes, block_bytes) + for block in blocks + ] + ).narrow(0, 0, payload_bytes) + return raw.view(scalar_dtype).reshape(rows, dimension) + + def copy_many_from_host( + self, items: list[tuple[tuple[int, ...], bytes]] + ) -> None: if self.storage is None: raise RuntimeError("GPU arena is closed") - source = torch.frombuffer(bytearray(payload), dtype=torch.uint8) - self.storage.narrow(0, offset, len(payload)).copy_(source) + if not items: + return + block_bytes = self.allocator.block_bytes + flattened: list[tuple[int, bytes, int, int]] = [] + for blocks, payload in items: + for ordinal, block in enumerate(blocks): + start = ordinal * block_bytes + length = min(block_bytes, max(0, len(payload) - start)) + flattened.append((block, payload, start, length)) + flattened.sort(key=lambda item: item[0]) + staging = self.host_staging + staging_array = staging.numpy() + staging_blocks = staging.numel() // block_bytes + copy_calls = 0 + with torch.cuda.device(self.device): + stream = self.copy_stream + for chunk_start in range(0, len(flattened), staging_blocks): + chunk = flattened[chunk_start : chunk_start + staging_blocks] + for index, (_, payload, source_start, source_length) in enumerate( + chunk + ): + if not source_length: + continue + start = index * block_bytes + staging_array[start : start + source_length] = np.frombuffer( + payload, + dtype=np.uint8, + count=source_length, + offset=source_start, + ) + with torch.cuda.stream(stream): + run_start = 0 + for index in range(1, len(chunk) + 1): + if ( + index < len(chunk) + and chunk[index][0] == chunk[index - 1][0] + 1 + ): + continue + first_block = chunk[run_start][0] + count = index - run_start + length = count * block_bytes + self.storage.narrow( + 0, first_block * block_bytes, length + ).copy_( + staging.narrow(0, run_start * block_bytes, length), + non_blocking=True, + ) + copy_calls += 1 + run_start = index + stream.synchronize() + self.h2d_batches += 1 + self.h2d_copy_calls += copy_calls + self.h2d_bytes += sum(len(payload) for _, payload in items) + + def copy_from_host(self, blocks: tuple[int, ...], payload: bytes) -> None: + self.copy_many_from_host([(blocks, payload)]) def status(self) -> dict[str, object]: return { @@ -206,6 +447,15 @@ def status(self) -> dict[str, object]: "workspace_bytes": self.workspace_bytes, "tensor_free_bytes": self.allocator.free_bytes, "largest_free_extent_bytes": self.allocator.largest_free_extent, + "block_bytes": self.allocator.block_bytes, + "block_count": self.allocator.block_count, + "free_blocks": self.allocator.free_blocks, + "h2d_batches": self.h2d_batches, + "h2d_copy_calls": self.h2d_copy_calls, + "h2d_bytes": self.h2d_bytes, + "host_staging_bytes": self.host_staging.numel() + if self.host_staging is not None + else 0, } def close(self) -> None: @@ -213,13 +463,20 @@ def close(self) -> None: return with torch.cuda.device(self.device): self.storage = None + self.host_staging = None + self.copy_stream = None torch.cuda.empty_cache() class GpuResourcePool: """Own all configured CUDA allocations or fail without a partial pool.""" - def __init__(self, specs: list[GpuArenaSpec], workspace_bytes: int) -> None: + def __init__( + self, + specs: list[GpuArenaSpec], + workspace_bytes: int, + block_bytes: int = DEFAULT_BLOCK_BYTES, + ) -> None: if not specs: raise RuntimeError("at least one GPU allocation is required") devices = [spec.device for spec in specs] @@ -228,7 +485,7 @@ def __init__(self, specs: list[GpuArenaSpec], workspace_bytes: int) -> None: self.arenas: list[GpuArena] = [] try: for spec in specs: - self.arenas.append(GpuArena(spec, workspace_bytes)) + self.arenas.append(GpuArena(spec, workspace_bytes, block_bytes)) except Exception: self.close() raise @@ -250,7 +507,7 @@ def close(self) -> None: class _GpuCacheEntry: key: tuple[object, ...] arena: GpuArena - offset: int + blocks: tuple[int, ...] allocated_bytes: int payload_bytes: int rows: int @@ -258,6 +515,7 @@ class _GpuCacheEntry: dtype: int references: int = 0 pinned: bool = False + priority: float = 0.0 @dataclass(frozen=True) @@ -290,11 +548,19 @@ def payload_bytes(self) -> int: @property def offset_bytes(self) -> int: - return self.entry.offset + return self.entry.blocks[0] * self.entry.arena.allocator.block_bytes + + @property + def block_ids(self) -> tuple[int, ...]: + return self.entry.blocks + + @property + def block_bytes(self) -> int: + return self.entry.arena.allocator.block_bytes def tensor(self) -> torch.Tensor: return self.entry.arena.tensor( - self.entry.offset, + self.entry.blocks, self.entry.payload_bytes, self.entry.rows, self.entry.dimension, @@ -302,8 +568,28 @@ def tensor(self) -> torch.Tensor: ) +@dataclass(frozen=True) +class GpuTensorLoad: + key: tuple[object, ...] + rows: int + dimension: int + dtype: int + payload: bytes + pin: bool = False + + +@dataclass(frozen=True) +class GpuAcquireBatch: + handles: tuple[GpuTensorHandle | None, ...] + bypassed: tuple[int, ...] + deferred: tuple[int, ...] + hits: int + misses: int + admitted: int + + class GpuTensorCache: - """Thread-safe LRU directory over one or more process-owned GPU arenas.""" + """Fixed-block GPU cache with TinyLFU admission and GDSF eviction.""" def __init__(self, pool: GpuResourcePool, allow_eviction: bool) -> None: self.pool = pool @@ -314,6 +600,9 @@ def __init__(self, pool: GpuResourcePool, allow_eviction: bool) -> None: self.misses = 0 self.evictions = 0 self.loaded_bytes = 0 + self.admission_rejections = 0 + self.inflation = 0.0 + self.sketch = TinyLfuSketch() def _find_arena(self, payload_bytes: int) -> GpuArena | None: candidates = [ @@ -326,28 +615,75 @@ def _find_arena(self, payload_bytes: int) -> GpuArena | None: return None return max(candidates, key=lambda arena: arena.allocator.free_bytes) - def _evict_one(self) -> bool: - for key, entry in self.entries.items(): - if entry.references or entry.pinned: - continue - self.entries.pop(key) - entry.arena.allocator.release(entry.offset, entry.allocated_bytes) - self.evictions += 1 - return True - return False - - def _allocate(self, payload_bytes: int) -> tuple[GpuArena, int, int]: + def _victim(self, arena: GpuArena | None = None) -> _GpuCacheEntry | None: + candidates = [ + entry + for entry in self.entries.values() + if not entry.references + and not entry.pinned + and (arena is None or entry.arena is arena) + ] + return min(candidates, key=lambda entry: entry.priority, default=None) + + def _evict(self, entry: _GpuCacheEntry) -> None: + current = self.entries.pop(entry.key, None) + if current is not entry: + raise RuntimeError("GPU eviction directory is inconsistent") + entry.arena.allocator.release(entry.blocks) + self.inflation = max(self.inflation, entry.priority) + self.evictions += 1 + + @staticmethod + def _entry_cost(entry: _GpuCacheEntry) -> int: + return len(entry.blocks) + + def _priority(self, key: tuple[object, ...], blocks: int) -> float: + return self.inflation + self.sketch.estimate(key) / max(1, blocks) + + def _allocate( + self, + key: tuple[object, ...], + payload_bytes: int, + enforce_admission: bool, + ) -> tuple[GpuArena, tuple[int, ...], int] | None: arena = self._find_arena(payload_bytes) - while arena is None and self.allow_eviction and self._evict_one(): - arena = self._find_arena(payload_bytes) + capable = [ + candidate + for candidate in self.pool.arenas + if candidate.allocator.capacity + >= candidate.allocator.allocation_bytes(payload_bytes) + ] + if not capable: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one tensor exceeds every configured GPU block arena", + ) + if arena is None and self.allow_eviction: + for candidate in sorted( + capable, key=lambda item: item.allocator.free_bytes, reverse=True + ): + required_blocks = candidate.allocator.blocks_for(payload_bytes) + required_bytes = candidate.allocator.allocation_bytes(payload_bytes) + while candidate.allocator.largest_free_extent < required_bytes: + victim = self._victim(candidate) + if victim is None: + break + candidate_priority = self._priority(key, required_blocks) + if enforce_admission and candidate_priority <= victim.priority: + self.admission_rejections += 1 + return None + self._evict(victim) + if candidate.allocator.largest_free_extent >= required_bytes: + arena = candidate + break if arena is None: raise protocol.SidecarError( protocol.STATUS_RESOURCE_LIMIT, - "configured GPU tensor arenas have insufficient contiguous capacity", + "configured GPU tensor arenas have insufficient free blocks", ) - allocated = arena.allocator.allocate(payload_bytes) - assert allocated is not None - return arena, allocated[0], allocated[1] + blocks = arena.allocator.allocate(payload_bytes) + assert blocks is not None + return arena, blocks, len(blocks) * arena.allocator.block_bytes def acquire( self, @@ -360,10 +696,12 @@ def acquire( pin: bool = False, ) -> tuple[GpuTensorHandle, bool]: with self.lock: + frequency = self.sketch.increment(key) cached = self.entries.get(key) if cached is not None: cached.references += 1 cached.pinned = cached.pinned or pin + cached.priority = self.inflation + frequency / self._entry_cost(cached) self.entries.move_to_end(key) self.hits += 1 return GpuTensorHandle(cached), True @@ -379,21 +717,25 @@ def acquire( with self.lock: cached = self.entries.get(key) if cached is not None: + frequency = self.sketch.increment(key) cached.references += 1 cached.pinned = cached.pinned or pin + cached.priority = self.inflation + frequency / self._entry_cost(cached) self.entries.move_to_end(key) self.hits += 1 return GpuTensorHandle(cached), True - arena, offset, allocated_bytes = self._allocate(len(payload)) + allocated = self._allocate(key, len(payload), enforce_admission=False) + assert allocated is not None + arena, blocks, allocated_bytes = allocated try: - arena.copy_from_host(offset, payload) + arena.copy_from_host(blocks, payload) except Exception: - arena.allocator.release(offset, allocated_bytes) + arena.allocator.release(blocks) raise entry = _GpuCacheEntry( key, arena, - offset, + blocks, allocated_bytes, len(payload), rows, @@ -401,12 +743,151 @@ def acquire( dtype, references=1, pinned=pin, + priority=self._priority(key, len(blocks)), ) self.entries[key] = entry self.misses += 1 self.loaded_bytes += len(payload) return GpuTensorHandle(entry), False + def acquire_many( + self, + loads: list[GpuTensorLoad], + *, + enforce_admission: bool = True, + record_access: bool = True, + count_stats: bool = True, + ) -> GpuAcquireBatch: + """Acquire a request working set and upload all new slabs in batches. + + ``deferred`` items could not be allocated while earlier handles in the + same request are referenced. The caller scores/releases the returned + handles and retries those items. ``bypassed`` items lost TinyLFU + admission and should be scored through the bounded streaming engine. + """ + + handles: list[GpuTensorHandle | None] = [None] * len(loads) + bypassed: list[int] = [] + deferred: list[int] = [] + new_entries: list[tuple[int, _GpuCacheEntry, bytes]] = [] + hits = 0 + misses = 0 + admitted = 0 + with self.lock: + try: + for index, load in enumerate(loads): + expected = protocol.checked_tensor_bytes( + load.rows, load.dimension, load.dtype + ) + if len(load.payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match its shape", + ) + frequency = ( + self.sketch.increment(load.key) + if record_access + else max(1, self.sketch.estimate(load.key)) + ) + cached = self.entries.get(load.key) + if cached is not None: + cached.references += 1 + cached.pinned = cached.pinned or load.pin + cached.priority = ( + self.inflation + frequency / self._entry_cost(cached) + ) + self.entries.move_to_end(load.key) + handles[index] = GpuTensorHandle(cached) + hits += 1 + continue + misses += 1 + try: + allocation = self._allocate( + load.key, + len(load.payload), + enforce_admission and not load.pin, + ) + except protocol.SidecarError as error: + if "one tensor exceeds" in str(error): + bypassed.append(index) + continue + deferred.append(index) + continue + if allocation is None: + bypassed.append(index) + continue + arena, blocks, allocated_bytes = allocation + entry = _GpuCacheEntry( + load.key, + arena, + blocks, + allocated_bytes, + len(load.payload), + load.rows, + load.dimension, + load.dtype, + references=1, + pinned=load.pin, + priority=self._priority(load.key, len(blocks)), + ) + self.entries[load.key] = entry + handles[index] = GpuTensorHandle(entry) + new_entries.append((index, entry, load.payload)) + admitted += 1 + + by_arena: dict[int, tuple[GpuArena, list[tuple[tuple[int, ...], bytes]]]] = {} + for _, entry, payload in new_entries: + bucket = by_arena.setdefault(id(entry.arena), (entry.arena, [])) + bucket[1].append((entry.blocks, payload)) + for arena, items in by_arena.values(): + arena.copy_many_from_host(items) + if count_stats: + self.hits += hits + self.misses += misses + self.loaded_bytes += sum(len(payload) for _, _, payload in new_entries) + except Exception: + new_ids = {id(entry) for _, entry, _ in new_entries} + for handle in handles: + if handle is None: + continue + entry = handle.entry + if id(entry) in new_ids: + if self.entries.pop(entry.key, None) is entry: + entry.arena.allocator.release(entry.blocks) + else: + entry.references -= 1 + raise + return GpuAcquireBatch( + tuple(handles), + tuple(bypassed), + tuple(deferred), + hits, + misses, + admitted, + ) + + def probe_many( + self, keys: list[tuple[object, ...]] + ) -> tuple[tuple[GpuTensorHandle | None, ...], tuple[int, ...]]: + """Acquire GPU hits without resolving the corresponding host payloads.""" + + handles: list[GpuTensorHandle | None] = [None] * len(keys) + misses = [] + with self.lock: + for index, key in enumerate(keys): + frequency = self.sketch.increment(key) + entry = self.entries.get(key) + if entry is None: + misses.append(index) + self.misses += 1 + continue + entry.references += 1 + entry.priority = self.inflation + frequency / self._entry_cost(entry) + self.entries.move_to_end(key) + handles[index] = GpuTensorHandle(entry) + self.hits += 1 + return tuple(handles), tuple(misses) + def release(self, handle: GpuTensorHandle) -> None: with self.lock: entry = self.entries.get(handle.entry.key) @@ -425,6 +906,9 @@ def status(self) -> dict[str, object]: "hits": self.hits, "misses": self.misses, "evictions": self.evictions, + "admission_rejections": self.admission_rejections, + "policy": "tinylfu-gdsf", + "gdsf_inflation": self.inflation, "loaded_bytes": self.loaded_bytes, "arenas": self.pool.status(), } diff --git a/services/tilemaxsim_shard.py b/services/tilemaxsim_shard.py new file mode 100644 index 00000000..f6307feb --- /dev/null +++ b/services/tilemaxsim_shard.py @@ -0,0 +1,354 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Immutable, content-addressed TileMaxSim tensor shard format. + +The data files contain canonical tensor bytes and alignment padding only. A +generation index maps tensor SHA-256 digests to byte ranges. Data files are +published under their own SHA-256, so a writer never mutates a visible shard. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO, Iterable + + +FORMAT = "vectorchord.tilemaxsim.shards" +VERSION = 1 +INDEX_NAME = "tilemaxsim-shards-v1.json" +SHARD_DIRECTORY = "shards" +DEFAULT_ALIGNMENT = 4096 +DEFAULT_SHARD_BYTES = 2 * 1024**3 + + +def align_up(value: int, alignment: int) -> int: + if alignment <= 0 or alignment & (alignment - 1): + raise ValueError("alignment must be a positive power of two") + return (value + alignment - 1) // alignment * alignment + + +@dataclass(frozen=True) +class ShardEntry: + digest: str + shard: str + offset: int + length: int + rows: int + dimension: int + dtype: str + + +@dataclass(frozen=True) +class ShardFile: + name: str + size: int + checksum: str + + +@dataclass(frozen=True) +class ShardIndex: + alignment: int + shards: dict[str, ShardFile] + entries: dict[str, ShardEntry] + + +def _digest_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(8 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +class ImmutableShardWriter: + """Deterministically publish bounded immutable shards and one atomic index.""" + + def __init__( + self, + root: Path, + target_bytes: int = DEFAULT_SHARD_BYTES, + alignment: int = DEFAULT_ALIGNMENT, + fsync: bool = True, + ) -> None: + if target_bytes <= 0: + raise ValueError("target shard bytes must be positive") + if target_bytes < alignment: + raise ValueError("target shard bytes must be at least one alignment unit") + align_up(0, alignment) + self.root = root + self.shard_root = root / SHARD_DIRECTORY + self.target_bytes = target_bytes + self.alignment = alignment + self.fsync = fsync + self.entries: dict[str, ShardEntry] = {} + self.shards: list[ShardFile] = [] + self._stream: BinaryIO | None = None + self._temporary: Path | None = None + self._hasher: hashlib._Hash | None = None + self._offset = 0 + self._pending: list[tuple[str, int, int, int, int, str]] = [] + self.root.mkdir(parents=True, exist_ok=True) + self.shard_root.mkdir(parents=True, exist_ok=True) + + def _open(self) -> None: + descriptor, name = tempfile.mkstemp( + prefix=".tilemaxsim-shard-", suffix=".tmp", dir=self.shard_root + ) + self._stream = os.fdopen(descriptor, "wb", closefd=True) + self._temporary = Path(name) + self._hasher = hashlib.sha256() + self._offset = 0 + self._pending = [] + + def _write(self, payload: bytes | memoryview) -> None: + assert self._stream is not None and self._hasher is not None + self._stream.write(payload) + self._hasher.update(payload) + self._offset += len(payload) + + def add( + self, + digest: str, + payload: bytes | memoryview, + rows: int, + dimension: int, + dtype: str, + ) -> None: + if len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest): + raise ValueError("invalid tensor SHA-256 digest") + if digest in self.entries or any(item[0] == digest for item in self._pending): + return + if not payload: + raise ValueError("tensor payload must not be empty") + if rows <= 0 or dimension <= 0 or dtype not in ("float16", "float32"): + raise ValueError("invalid tensor metadata") + padded = align_up(len(payload), self.alignment) + if self._stream is not None and self._offset and self._offset + padded > self.target_bytes: + self._finish_shard() + if self._stream is None: + self._open() + offset = self._offset + self._write(payload) + padding = padded - len(payload) + if padding: + self._write(bytes(min(padding, self.alignment))) + remaining = padding - min(padding, self.alignment) + while remaining: + chunk = min(remaining, self.alignment) + self._write(bytes(chunk)) + remaining -= chunk + self._pending.append((digest, offset, len(payload), rows, dimension, dtype)) + + def _finish_shard(self) -> None: + if self._stream is None: + return + assert self._temporary is not None and self._hasher is not None + stream = self._stream + stream.flush() + if self.fsync: + os.fsync(stream.fileno()) + stream.close() + digest = self._hasher.hexdigest() + name = f"sha256-{digest}.vts" + destination = self.shard_root / name + size = self._offset + if destination.exists(): + if destination.stat().st_size != size or _digest_file(destination) != digest: + raise ValueError(f"existing immutable shard is corrupt: {destination}") + self._temporary.unlink() + else: + os.replace(self._temporary, destination) + if self.fsync: + descriptor = os.open(self.shard_root, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + relative = f"{SHARD_DIRECTORY}/{name}" + shard = ShardFile(relative, size, f"sha256:{digest}") + self.shards.append(shard) + for tensor_digest, offset, length, rows, dimension, dtype in self._pending: + self.entries[tensor_digest] = ShardEntry( + tensor_digest, + relative, + offset, + length, + rows, + dimension, + dtype, + ) + self._stream = None + self._temporary = None + self._hasher = None + self._offset = 0 + self._pending = [] + + def finish(self) -> Path: + self._finish_shard() + if not self.entries: + raise ValueError("cannot publish an empty shard generation") + document = { + "format": FORMAT, + "version": VERSION, + "alignment": self.alignment, + "shards": [ + {"name": shard.name, "bytes": shard.size, "checksum": shard.checksum} + for shard in self.shards + ], + "entries": [ + { + "digest": entry.digest, + "shard": entry.shard, + "offset": entry.offset, + "length": entry.length, + "rows": entry.rows, + "dimension": entry.dimension, + "dtype": entry.dtype, + } + for entry in self.entries.values() + ], + } + descriptor, temporary = tempfile.mkstemp( + prefix=f".{INDEX_NAME}.", suffix=".tmp", dir=self.root, text=True + ) + try: + with os.fdopen(descriptor, "w", encoding="utf-8", closefd=True) as stream: + json.dump(document, stream, separators=(",", ":"), sort_keys=True) + stream.write("\n") + stream.flush() + if self.fsync: + os.fsync(stream.fileno()) + destination = self.root / INDEX_NAME + os.replace(temporary, destination) + if self.fsync: + directory = os.open(self.root, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory) + finally: + os.close(directory) + return destination + finally: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + + def close(self) -> None: + if self._stream is not None: + self._stream.close() + if self._temporary is not None: + try: + self._temporary.unlink() + except FileNotFoundError: + pass + self._stream = None + self._temporary = None + + +def parse_index(document: object) -> ShardIndex: + if not isinstance(document, dict): + raise ValueError("shard index must be a JSON object") + if document.get("format") != FORMAT or document.get("version") != VERSION: + raise ValueError("unsupported TileMaxSim shard index") + alignment = document.get("alignment") + if not isinstance(alignment, int): + raise ValueError("shard index has invalid alignment") + align_up(0, alignment) + raw_shards = document.get("shards") + raw_entries = document.get("entries") + if not isinstance(raw_shards, list) or not isinstance(raw_entries, list): + raise ValueError("shard index has invalid arrays") + shards: dict[str, ShardFile] = {} + for raw in raw_shards: + if not isinstance(raw, dict): + raise ValueError("shard index contains an invalid shard") + name, size, checksum = raw.get("name"), raw.get("bytes"), raw.get("checksum") + if ( + not isinstance(name, str) + or not name.startswith(f"{SHARD_DIRECTORY}/sha256-") + or "/" in name[len(SHARD_DIRECTORY) + 1 :] + or not isinstance(size, int) + or size <= 0 + or not isinstance(checksum, str) + or not checksum.startswith("sha256:") + or name != f"{SHARD_DIRECTORY}/sha256-{checksum.removeprefix('sha256:')}.vts" + or name in shards + ): + raise ValueError("shard index contains invalid shard metadata") + shards[name] = ShardFile(name, size, checksum) + entries: dict[str, ShardEntry] = {} + intervals: dict[str, list[tuple[int, int]]] = {name: [] for name in shards} + for raw in raw_entries: + if not isinstance(raw, dict): + raise ValueError("shard index contains an invalid tensor entry") + digest = raw.get("digest") + shard = raw.get("shard") + offset = raw.get("offset") + length = raw.get("length") + rows = raw.get("rows") + dimension = raw.get("dimension") + dtype = raw.get("dtype") + if ( + not isinstance(digest, str) + or len(digest) != 64 + or any(c not in "0123456789abcdef" for c in digest) + or not isinstance(shard, str) + or shard not in shards + or not isinstance(offset, int) + or offset < 0 + or offset % alignment + or not isinstance(length, int) + or length <= 0 + or not isinstance(rows, int) + or rows <= 0 + or not isinstance(dimension, int) + or dimension <= 0 + or dtype not in ("float16", "float32") + or digest in entries + ): + raise ValueError("shard index contains invalid tensor metadata") + scalar_bytes = 2 if dtype == "float16" else 4 + if length != rows * dimension * scalar_bytes: + raise ValueError("shard tensor length disagrees with its shape") + end = offset + length + if end > shards[shard].size: + raise ValueError("shard tensor range is outside its data file") + intervals[shard].append((offset, align_up(end, alignment))) + entries[digest] = ShardEntry( + digest, shard, offset, length, rows, dimension, dtype + ) + if not entries: + raise ValueError("shard index contains no tensor entries") + for shard, ranges in intervals.items(): + previous_end = 0 + for start, end in sorted(ranges): + if start < previous_end: + raise ValueError(f"overlapping tensor ranges in {shard}") + previous_end = end + return ShardIndex(alignment, shards, entries) + + +def load_index(path: Path) -> ShardIndex: + with path.open(encoding="utf-8") as stream: + return parse_index(json.load(stream)) + + +def iter_index_entries(index: ShardIndex) -> Iterable[ShardEntry]: + return index.entries.values() diff --git a/services/tilemaxsimd/Cargo.lock b/services/tilemaxsimd/Cargo.lock new file mode 100644 index 00000000..d55f046d --- /dev/null +++ b/services/tilemaxsimd/Cargo.lock @@ -0,0 +1,416 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cc" +version = "1.2.60" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom", + "libc", +] + +[[package]] +name = "libc" +version = "0.2.185" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f" + +[[package]] +name = "memchr" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tilemaxsimd" +version = "0.1.0" +dependencies = [ + "anyhow", + "cc", + "clap", + "hex", + "libc", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "typenum" +version = "1.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40ce102ab67701b8526c123c1bab5cbe42d7040ccfd0f64af1a385808d2f43de" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasip2" +version = "1.0.3+wasi-0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/services/tilemaxsimd/Cargo.toml b/services/tilemaxsimd/Cargo.toml new file mode 100644 index 00000000..30c5d82a --- /dev/null +++ b/services/tilemaxsimd/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "tilemaxsimd" +version = "0.1.0" +edition = "2024" +publish = false + +[dependencies] +anyhow = "1.0" +clap = { version = "4.5", features = ["derive"] } +hex = "0.4" +libc = "0.2" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +sha2 = "0.10" + +[build-dependencies] +cc = { version = "1.2", features = ["parallel"] } + +[workspace] diff --git a/services/tilemaxsimd/build.rs b/services/tilemaxsimd/build.rs new file mode 100644 index 00000000..15cbfae6 --- /dev/null +++ b/services/tilemaxsimd/build.rs @@ -0,0 +1,12 @@ +fn main() { + println!("cargo:rerun-if-changed=native/tilemaxsim_cuda.cu"); + cc::Build::new() + .cuda(true) + .flag("-O3") + .flag("--use_fast_math") + .flag("-lineinfo") + .file("native/tilemaxsim_cuda.cu") + .compile("tilemaxsim_cuda"); + println!("cargo:rustc-link-lib=cudart"); + println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64"); +} diff --git a/services/tilemaxsimd/native/tilemaxsim_cuda.cu b/services/tilemaxsimd/native/tilemaxsim_cuda.cu new file mode 100644 index 00000000..ff214e6e --- /dev/null +++ b/services/tilemaxsimd/native/tilemaxsim_cuda.cu @@ -0,0 +1,291 @@ +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +struct VctmGpu { + int device; + unsigned char *allocation; + size_t total_bytes; + size_t tensor_bytes; + size_t workspace_bytes; + unsigned char *host_staging; + size_t host_staging_bytes; + cudaStream_t upload_stream; + cudaStream_t compute_stream; +}; + +static int fail(char *error, size_t capacity, const char *message) { + if (error != nullptr && capacity != 0) { + std::snprintf(error, capacity, "%s", message); + } + return 1; +} + +static int cuda_fail(char *error, size_t capacity, const char *operation, + cudaError_t status) { + if (error != nullptr && capacity != 0) { + std::snprintf(error, capacity, "%s: %s", operation, + cudaGetErrorString(status)); + } + return 1; +} + +extern "C" int vctm_gpu_create(int device, size_t total_bytes, + size_t workspace_bytes, VctmGpu **output, + char *error, size_t error_capacity) { + if (output == nullptr || total_bytes == 0 || workspace_bytes == 0 || + workspace_bytes >= total_bytes) { + return fail(error, error_capacity, "invalid GPU arena configuration"); + } + cudaError_t status = cudaSetDevice(device); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaSetDevice", status); + } + size_t free_bytes = 0; + size_t device_bytes = 0; + status = cudaMemGetInfo(&free_bytes, &device_bytes); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaMemGetInfo", status); + } + if (free_bytes < total_bytes) { + return fail(error, error_capacity, + "configured GPU memory is not currently available"); + } + auto *gpu = new VctmGpu{}; + gpu->device = device; + gpu->total_bytes = total_bytes; + gpu->tensor_bytes = ((total_bytes - workspace_bytes) / 256) * 256; + gpu->workspace_bytes = total_bytes - gpu->tensor_bytes; + gpu->host_staging_bytes = + std::min(gpu->tensor_bytes, static_cast(64) * 1024 * 1024); + status = cudaMalloc(reinterpret_cast(&gpu->allocation), total_bytes); + if (status != cudaSuccess) { + delete gpu; + return cuda_fail(error, error_capacity, "cudaMalloc", status); + } + if ((status = cudaStreamCreateWithFlags(&gpu->upload_stream, + cudaStreamNonBlocking)) != cudaSuccess || + (status = cudaStreamCreateWithFlags(&gpu->compute_stream, + cudaStreamNonBlocking)) != cudaSuccess) { + if (gpu->upload_stream != nullptr) cudaStreamDestroy(gpu->upload_stream); + cudaFree(gpu->allocation); + delete gpu; + return cuda_fail(error, error_capacity, "cudaStreamCreate", status); + } + status = cudaHostAlloc(reinterpret_cast(&gpu->host_staging), + gpu->host_staging_bytes, cudaHostAllocPortable); + if (status != cudaSuccess) { + cudaStreamDestroy(gpu->upload_stream); + cudaStreamDestroy(gpu->compute_stream); + cudaFree(gpu->allocation); + delete gpu; + return cuda_fail(error, error_capacity, "cudaHostAlloc", status); + } + std::memset(gpu->host_staging, 0, gpu->host_staging_bytes); + *output = gpu; + return 0; +} + +extern "C" void vctm_gpu_destroy(VctmGpu *gpu) { + if (gpu == nullptr) return; + cudaSetDevice(gpu->device); + cudaStreamDestroy(gpu->upload_stream); + cudaStreamDestroy(gpu->compute_stream); + cudaFreeHost(gpu->host_staging); + cudaFree(gpu->allocation); + delete gpu; +} + +extern "C" size_t vctm_gpu_tensor_bytes(const VctmGpu *gpu) { + return gpu == nullptr ? 0 : gpu->tensor_bytes; +} + +extern "C" int vctm_gpu_upload_batch( + VctmGpu *gpu, const uint64_t *offsets, + const unsigned char *const *payloads, const size_t *lengths, size_t count, + char *error, size_t error_capacity) { + if (gpu == nullptr || offsets == nullptr || payloads == nullptr || + lengths == nullptr) { + return fail(error, error_capacity, "invalid upload batch"); + } + cudaError_t status = cudaSetDevice(gpu->device); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaSetDevice", status); + } + for (size_t i = 0; i < count; ++i) { + if (offsets[i] > gpu->tensor_bytes || + lengths[i] > gpu->tensor_bytes - offsets[i]) { + return fail(error, error_capacity, "upload is outside the tensor arena"); + } + } + size_t item = 0; + size_t item_offset = 0; + while (item < count) { + size_t staging_offset = 0; + while (item < count && staging_offset < gpu->host_staging_bytes) { + const size_t remaining = lengths[item] - item_offset; + const size_t chunk = + std::min(remaining, gpu->host_staging_bytes - staging_offset); + std::memcpy(gpu->host_staging + staging_offset, + payloads[item] + item_offset, chunk); + status = cudaMemcpyAsync(gpu->allocation + offsets[item] + item_offset, + gpu->host_staging + staging_offset, chunk, + cudaMemcpyHostToDevice, gpu->upload_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaMemcpyAsync(H2D)", status); + } + staging_offset += chunk; + item_offset += chunk; + if (item_offset == lengths[item]) { + item += 1; + item_offset = 0; + } + } + status = cudaStreamSynchronize(gpu->upload_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaStreamSynchronize(upload)", + status); + } + } + return 0; +} + +template +__device__ float scalar_to_float(Scalar value); + +template <> +__device__ float scalar_to_float(half value) { + return __half2float(value); +} + +template <> +__device__ float scalar_to_float(float value) { + return value; +} + +template +__global__ void tilemaxsim_kernel(const Scalar *query, uint32_t query_rows, + uint32_t dimension, + const unsigned char *documents, + const uint64_t *document_offsets, + const uint32_t *document_rows, float *scores) { + const uint32_t candidate = blockIdx.x; + const uint32_t query_row = blockIdx.y; + const uint32_t lane = threadIdx.x & 31; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t warps = blockDim.x >> 5; + const auto *document = reinterpret_cast( + documents + document_offsets[candidate]); + const Scalar *query_vector = query + static_cast(query_row) * dimension; + float best = -CUDART_INF_F; + for (uint32_t row = warp; row < document_rows[candidate]; row += warps) { + const Scalar *document_vector = + document + static_cast(row) * dimension; + float dot = 0.0f; + for (uint32_t index = lane; index < dimension; index += 32) { + dot = fmaf(scalar_to_float(query_vector[index]), + scalar_to_float(document_vector[index]), dot); + } + for (int delta = 16; delta != 0; delta >>= 1) { + dot += __shfl_down_sync(0xffffffff, dot, delta); + } + if (lane == 0) best = fmaxf(best, dot); + } + __shared__ float warp_best[8]; + if (lane == 0) warp_best[warp] = best; + __syncthreads(); + if (threadIdx.x == 0) { + float maximum = -CUDART_INF_F; + for (uint32_t index = 0; index < warps; ++index) { + maximum = fmaxf(maximum, warp_best[index]); + } + atomicAdd(scores + candidate, maximum); + } +} + +static size_t aligned(size_t value, size_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +extern "C" int vctm_gpu_score( + VctmGpu *gpu, const unsigned char *query, size_t query_bytes, + uint32_t query_rows, uint32_t dimension, uint8_t dtype, + const uint64_t *document_offsets, const uint32_t *document_rows, + size_t count, float *output, char *error, size_t error_capacity) { + if (gpu == nullptr || query == nullptr || document_offsets == nullptr || + document_rows == nullptr || output == nullptr || query_rows == 0 || + dimension == 0 || count == 0) { + return fail(error, error_capacity, "invalid TileMaxSim score request"); + } + cudaError_t status = cudaSetDevice(gpu->device); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaSetDevice", status); + } + unsigned char *workspace = gpu->allocation + gpu->tensor_bytes; + size_t cursor = 0; + const size_t query_offset = cursor; + cursor = aligned(cursor + query_bytes, 256); + const size_t offsets_offset = cursor; + cursor = aligned(cursor + count * sizeof(uint64_t), 256); + const size_t rows_offset = cursor; + cursor = aligned(cursor + count * sizeof(uint32_t), 256); + const size_t scores_offset = cursor; + cursor = aligned(cursor + count * sizeof(float), 256); + if (cursor > gpu->workspace_bytes) { + return fail(error, error_capacity, + "TileMaxSim request exceeds the configured GPU workspace"); + } + status = cudaMemcpyAsync(workspace + query_offset, query, query_bytes, + cudaMemcpyHostToDevice, gpu->compute_stream); + if (status == cudaSuccess) + status = cudaMemcpyAsync(workspace + offsets_offset, document_offsets, + count * sizeof(uint64_t), cudaMemcpyHostToDevice, + gpu->compute_stream); + if (status == cudaSuccess) + status = cudaMemcpyAsync(workspace + rows_offset, document_rows, + count * sizeof(uint32_t), cudaMemcpyHostToDevice, + gpu->compute_stream); + if (status == cudaSuccess) + status = cudaMemsetAsync(workspace + scores_offset, 0, + count * sizeof(float), gpu->compute_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "CUDA workspace initialization", + status); + } + dim3 grid(static_cast(count), query_rows); + dim3 block(256); + if (dtype == 2) { + tilemaxsim_kernel<<compute_stream>>>( + reinterpret_cast(workspace + query_offset), query_rows, + dimension, gpu->allocation, + reinterpret_cast(workspace + offsets_offset), + reinterpret_cast(workspace + rows_offset), + reinterpret_cast(workspace + scores_offset)); + } else if (dtype == 1) { + tilemaxsim_kernel<<compute_stream>>>( + reinterpret_cast(workspace + query_offset), query_rows, + dimension, gpu->allocation, + reinterpret_cast(workspace + offsets_offset), + reinterpret_cast(workspace + rows_offset), + reinterpret_cast(workspace + scores_offset)); + } else { + return fail(error, error_capacity, "unsupported tensor dtype"); + } + status = cudaGetLastError(); + if (status == cudaSuccess) + status = cudaMemcpyAsync(output, workspace + scores_offset, + count * sizeof(float), cudaMemcpyDeviceToHost, + gpu->compute_stream); + if (status == cudaSuccess) status = cudaStreamSynchronize(gpu->compute_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "TileMaxSim CUDA execution", status); + } + return 0; +} diff --git a/services/tilemaxsimd/src/cache.rs b/services/tilemaxsimd/src/cache.rs new file mode 100644 index 00000000..06ae3fa7 --- /dev/null +++ b/services/tilemaxsimd/src/cache.rs @@ -0,0 +1,399 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::hash::{Hash, Hasher}; + +#[derive(Debug)] +pub struct BuddyAllocator { + block_bytes: usize, + block_count: usize, + free: BTreeMap>, + allocated: HashMap, +} + +impl BuddyAllocator { + pub fn new(capacity: usize, block_bytes: usize) -> Result { + if capacity == 0 || block_bytes == 0 || !block_bytes.is_multiple_of(256) { + return Err("invalid fixed-block arena"); + } + let block_count = capacity / block_bytes; + if block_count == 0 { + return Err("fixed-block arena is too small"); + } + let mut free = BTreeMap::>::new(); + let mut start = 0; + let mut remaining = block_count; + while remaining != 0 { + let order = usize::BITS - 1 - remaining.leading_zeros(); + let size = 1_usize << order; + free.entry(order).or_default().insert((start, start, order)); + start += size; + remaining -= size; + } + Ok(Self { + block_bytes, + block_count, + free, + allocated: HashMap::new(), + }) + } + + pub fn capacity(&self) -> usize { + self.block_count * self.block_bytes + } + + pub fn block_bytes(&self) -> usize { + self.block_bytes + } + + pub fn allocation_bytes(&self, payload_bytes: usize) -> Option { + if payload_bytes == 0 { + return None; + } + let raw = payload_bytes.div_ceil(self.block_bytes); + raw.checked_next_power_of_two()? + .checked_mul(self.block_bytes) + } + + pub fn largest_free(&self) -> usize { + self.free + .iter() + .rev() + .find(|(_, entries)| !entries.is_empty()) + .map_or(0, |(order, _)| (1_usize << order) * self.block_bytes) + } + + pub fn allocate(&mut self, payload_bytes: usize) -> Option<(usize, usize)> { + let allocation_bytes = self.allocation_bytes(payload_bytes)?; + let blocks = allocation_bytes / self.block_bytes; + let order = blocks.trailing_zeros(); + let available_order = self + .free + .range(order..) + .find(|(_, entries)| !entries.is_empty()) + .map(|(candidate, _)| *candidate)?; + let item = self.free.get_mut(&available_order)?.pop_first()?; + let (start, root_start, root_order) = item; + let mut current_order = available_order; + while current_order > order { + current_order -= 1; + let buddy = start + (1_usize << current_order); + self.free + .entry(current_order) + .or_default() + .insert((buddy, root_start, root_order)); + } + self.allocated + .insert(start, (order, root_start, root_order)); + Some((start * self.block_bytes, allocation_bytes)) + } + + pub fn release(&mut self, offset: usize) -> Result<(), &'static str> { + if !offset.is_multiple_of(self.block_bytes) { + return Err("unaligned fixed-block release"); + } + let mut start = offset / self.block_bytes; + let (mut order, root_start, root_order) = self + .allocated + .remove(&start) + .ok_or("fixed-block slab was released twice")?; + while order < root_order { + let buddy = root_start + ((start - root_start) ^ (1_usize << order)); + let buddy_item = (buddy, root_start, root_order); + let entries = self.free.entry(order).or_default(); + if !entries.remove(&buddy_item) { + break; + } + start = start.min(buddy); + order += 1; + } + self.free + .entry(order) + .or_default() + .insert((start, root_start, root_order)); + Ok(()) + } +} + +#[derive(Debug)] +pub struct TinyLfu { + width: usize, + tables: Vec>, + samples: usize, +} + +impl TinyLfu { + pub fn new(width: usize) -> Self { + Self { + width, + tables: vec![vec![0; width]; 4], + samples: 0, + } + } + + fn indices(&self, key: &T) -> [usize; 4] { + std::array::from_fn(|row| { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + row.hash(&mut hasher); + key.hash(&mut hasher); + hasher.finish() as usize % self.width + }) + } + + pub fn increment(&mut self, key: &T) -> u16 { + let indices = self.indices(key); + for (row, index) in indices.into_iter().enumerate() { + self.tables[row][index] = self.tables[row][index].saturating_add(1); + } + self.samples += 1; + let estimate = self.estimate(key).max(1); + if self.samples >= self.width * 10 { + for table in &mut self.tables { + for value in table { + *value /= 2; + } + } + self.samples /= 2; + } + estimate + } + + pub fn estimate(&self, key: &T) -> u16 { + self.indices(key) + .into_iter() + .enumerate() + .map(|(row, index)| self.tables[row][index]) + .min() + .unwrap_or(0) + } +} + +#[derive(Clone, Debug)] +pub struct CacheEntry { + pub offset: u64, + pub allocated_bytes: usize, + pub payload_bytes: usize, + pub rows: u32, + pub dimension: u32, + pub dtype: u8, + pub references: usize, + pub pinned: bool, + priority: f64, +} + +#[derive(Debug, PartialEq)] +pub enum Admission { + Admitted { offset: u64, allocated_bytes: usize }, + Rejected, + Deferred, +} + +pub struct GpuCache { + allocator: BuddyAllocator, + entries: HashMap, + sketch: TinyLfu, + inflation: f64, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, +} + +impl GpuCache { + pub fn new(capacity: usize, block_bytes: usize) -> Result { + Ok(Self { + allocator: BuddyAllocator::new(capacity, block_bytes)?, + entries: HashMap::new(), + sketch: TinyLfu::new(4096), + inflation: 0.0, + hits: 0, + misses: 0, + evictions: 0, + admission_rejections: 0, + }) + } + + pub fn capacity(&self) -> usize { + self.allocator.capacity() + } + + pub fn block_bytes(&self) -> usize { + self.allocator.block_bytes() + } + + pub fn entry_count(&self) -> usize { + self.entries.len() + } + + pub fn pinned_entries(&self) -> usize { + self.entries.values().filter(|entry| entry.pinned).count() + } + + pub fn get(&mut self, key: &str) -> Option { + let frequency = self.sketch.increment(&key); + let Some(entry) = self.entries.get_mut(key) else { + self.misses += 1; + return None; + }; + entry.references += 1; + entry.priority = self.inflation + f64::from(frequency) / entry.allocated_bytes as f64; + self.hits += 1; + Some(entry.clone()) + } + + pub fn acquire_existing(&mut self, key: &str) -> Option { + let entry = self.entries.get_mut(key)?; + entry.references += 1; + Some(entry.clone()) + } + + pub fn record_access_miss(&mut self, key: &str) { + self.sketch.increment(&key); + self.misses += 1; + } + + pub fn contains(&self, key: &str) -> bool { + self.entries.contains_key(key) + } + + fn victim(&self) -> Option<(&String, &CacheEntry)> { + self.entries + .iter() + .filter(|(_, entry)| entry.references == 0 && !entry.pinned) + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + } + + fn evict_one(&mut self) -> bool { + let Some(key) = self.victim().map(|(key, _)| key.clone()) else { + return false; + }; + let entry = self.entries.remove(&key).expect("victim disappeared"); + self.allocator + .release(entry.offset as usize) + .expect("cache slab release failed"); + self.inflation = self.inflation.max(entry.priority); + self.evictions += 1; + true + } + + #[allow(clippy::too_many_arguments)] + pub fn admit( + &mut self, + key: String, + payload_bytes: usize, + rows: u32, + dimension: u32, + dtype: u8, + pinned: bool, + force: bool, + ) -> Admission { + let Some(allocation_bytes) = self.allocator.allocation_bytes(payload_bytes) else { + return Admission::Deferred; + }; + if allocation_bytes > self.allocator.capacity() { + return Admission::Deferred; + } + while self.allocator.largest_free() < allocation_bytes { + let Some((_, victim)) = self.victim() else { + return Admission::Deferred; + }; + let candidate = self.inflation + + f64::from(self.sketch.estimate(&key).max(1)) / allocation_bytes as f64; + if !force && !pinned && candidate <= victim.priority { + self.admission_rejections += 1; + return Admission::Rejected; + } + if !self.evict_one() { + return Admission::Deferred; + } + } + let Some((offset, allocated_bytes)) = self.allocator.allocate(payload_bytes) else { + return Admission::Deferred; + }; + let priority = + self.inflation + f64::from(self.sketch.estimate(&key).max(1)) / allocated_bytes as f64; + self.entries.insert( + key, + CacheEntry { + offset: offset as u64, + allocated_bytes, + payload_bytes, + rows, + dimension, + dtype, + references: 1, + pinned, + priority, + }, + ); + Admission::Admitted { + offset: offset as u64, + allocated_bytes, + } + } + + pub fn release(&mut self, key: &str) -> Result<(), &'static str> { + let entry = self + .entries + .get_mut(key) + .ok_or("GPU cache entry disappeared")?; + if entry.references == 0 { + return Err("GPU cache entry was released twice"); + } + entry.references -= 1; + Ok(()) + } + + #[cfg(test)] + pub fn entry(&self, key: &str) -> Option<&CacheEntry> { + self.entries.get(key) + } + + pub fn remove(&mut self, key: &str) -> Result<(), &'static str> { + let entry = self + .entries + .remove(key) + .ok_or("GPU cache entry disappeared")?; + self.allocator.release(entry.offset as usize) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn buddy_coalesces_fixed_slabs() { + let mut allocator = BuddyAllocator::new(8 * 256, 256).unwrap(); + let first = allocator.allocate(300).unwrap(); + let second = allocator.allocate(300).unwrap(); + let third = allocator.allocate(700).unwrap(); + assert_eq!(first, (0, 512)); + assert_eq!(second, (512, 512)); + assert_eq!(third, (1024, 1024)); + allocator.release(second.0).unwrap(); + allocator.release(first.0).unwrap(); + assert_eq!(allocator.largest_free(), 1024); + allocator.release(third.0).unwrap(); + assert_eq!(allocator.largest_free(), 2048); + } + + #[test] + fn tinylfu_protects_hot_entry() { + let mut cache = GpuCache::new(256, 256).unwrap(); + cache.record_access_miss("hot"); + assert!(matches!( + cache.admit("hot".into(), 64, 1, 32, 2, false, false), + Admission::Admitted { .. } + )); + cache.release("hot").unwrap(); + for _ in 0..3 { + cache.get("hot").unwrap(); + cache.release("hot").unwrap(); + } + cache.record_access_miss("cold"); + assert_eq!( + cache.admit("cold".into(), 64, 1, 32, 2, false, false), + Admission::Rejected + ); + assert!(cache.entry("hot").is_some()); + } +} diff --git a/services/tilemaxsimd/src/engine.rs b/services/tilemaxsimd/src/engine.rs new file mode 100644 index 00000000..fd5c1adf --- /dev/null +++ b/services/tilemaxsimd/src/engine.rs @@ -0,0 +1,434 @@ +use crate::cache::{Admission, GpuCache}; +use crate::gpu::Gpu; +use crate::protocol::{Descriptor, Request}; +use crate::shard::{ShardStore, cache_key}; +use anyhow::{Result, anyhow, bail}; + +struct MissingTensor { + candidate_index: usize, + descriptor: Descriptor, + key: String, + payload: Vec, +} + +struct ResidentTensor { + candidate_index: usize, + device: usize, + key: String, + offset: u64, + rows: u32, + transient: bool, +} + +struct DeviceState { + gpu: Gpu, + cache: GpuCache, + h2d_batches: u64, + h2d_bytes: u64, +} + +pub struct Engine { + devices: Vec, + store: ShardStore, + next_device: usize, +} + +impl Engine { + pub fn new(gpus: Vec, block_bytes: usize, store: ShardStore) -> Result { + if gpus.is_empty() { + bail!("at least one GPU is required"); + } + let devices = gpus + .into_iter() + .map(|gpu| { + let cache = GpuCache::new(gpu.tensor_bytes(), block_bytes) + .map_err(|message| anyhow!(message))?; + Ok(DeviceState { + gpu, + cache, + h2d_batches: 0, + h2d_bytes: 0, + }) + }) + .collect::>>()?; + Ok(Self { + devices, + store, + next_device: 0, + }) + } + + pub fn prewarm(&mut self, descriptors: &[Descriptor], batch_size: usize) -> Result<()> { + if batch_size == 0 { + bail!("resident prewarm batch size must be positive"); + } + for batch in descriptors.chunks(batch_size) { + let payloads = self.store.resolve_many(batch)?; + let mut uploads = (0..self.devices.len()) + .map(|_| Vec::<(u64, &[u8])>::new()) + .collect::>(); + let mut acquired = Vec::<(usize, String)>::new(); + for (descriptor, payload) in batch.iter().zip(&payloads) { + let key = cache_key(descriptor); + if let Some((device, _)) = + self.devices + .iter_mut() + .enumerate() + .find_map(|(index, device)| { + device + .cache + .acquire_existing(&key) + .map(|entry| (index, entry)) + }) + { + acquired.push((device, key)); + continue; + } + self.devices[self.next_device] + .cache + .record_access_miss(&key); + let mut admission = None; + for step in 0..self.devices.len() { + let device = (self.next_device + step) % self.devices.len(); + if let Admission::Admitted { offset, .. } = self.devices[device].cache.admit( + key.clone(), + payload.len(), + descriptor.rows, + descriptor.dimension, + descriptor.dtype, + true, + true, + ) { + admission = Some((device, offset)); + self.next_device = (device + 1) % self.devices.len(); + break; + } + } + let Some((device, offset)) = admission else { + bail!("resident manifest exceeds the configured Rust GPU block caches"); + }; + uploads[device].push((offset, payload.as_slice())); + acquired.push((device, key)); + } + for (device, items) in uploads.iter().enumerate() { + if items.is_empty() { + continue; + } + self.devices[device].gpu.upload_batch(items)?; + self.devices[device].h2d_batches += 1; + self.devices[device].h2d_bytes += items + .iter() + .map(|(_, payload)| payload.len() as u64) + .sum::(); + } + for (device, key) in acquired { + self.devices[device] + .cache + .release(&key) + .map_err(|message| anyhow!(message))?; + } + } + Ok(()) + } + + pub fn score(&mut self, request: &Request) -> Result> { + if request.candidates.is_empty() { + return Ok(Vec::new()); + } + let mut scores = vec![None; request.candidates.len()]; + let mut hit_chunks = (0..self.devices.len()) + .map(|_| Vec::::new()) + .collect::>(); + let mut missing_descriptors = Vec::new(); + let mut missing_indices = Vec::new(); + for (index, descriptor) in request.candidates.iter().enumerate() { + let key = cache_key(descriptor); + let hit_device = self + .devices + .iter() + .position(|device| device.cache.contains(&key)); + if let Some(device_index) = hit_device { + let entry = self.devices[device_index] + .cache + .get(&key) + .expect("cache hit disappeared"); + validate_entry(descriptor, &entry)?; + hit_chunks[device_index].push(ResidentTensor { + candidate_index: index, + device: device_index, + key, + offset: entry.offset, + rows: entry.rows, + transient: false, + }); + } else { + // Record one request-level miss on the device that will get the + // first admission opportunity. Other devices are not polluted. + self.devices[self.next_device] + .cache + .record_access_miss(&key); + missing_indices.push(index); + missing_descriptors.push(descriptor.clone()); + } + } + + let empty_uploads = (0..self.devices.len()) + .map(|_| Vec::<(u64, &[u8])>::new()) + .collect::>(); + self.execute_devices(request, &hit_chunks, &empty_uploads, &mut scores)?; + for chunk in &hit_chunks { + self.release_chunk(chunk)?; + } + + let payloads = self.store.resolve_many(&missing_descriptors)?; + let mut pending = missing_indices + .into_iter() + .zip(missing_descriptors) + .zip(payloads) + .map(|((candidate_index, descriptor), payload)| MissingTensor { + candidate_index, + key: cache_key(&descriptor), + descriptor, + payload, + }) + .collect::>(); + + while !pending.is_empty() { + let mut chunks = (0..self.devices.len()) + .map(|_| Vec::::new()) + .collect::>(); + let mut uploads = (0..self.devices.len()) + .map(|_| Vec::<(u64, &[u8])>::new()) + .collect::>(); + let mut consumed = 0; + for tensor in &pending { + if let Some((device_index, entry)) = + self.devices + .iter_mut() + .enumerate() + .find_map(|(index, device)| { + device + .cache + .acquire_existing(&tensor.key) + .map(|entry| (index, entry)) + }) + { + chunks[device_index].push(ResidentTensor { + candidate_index: tensor.candidate_index, + device: device_index, + key: tensor.key.clone(), + offset: entry.offset, + rows: entry.rows, + transient: false, + }); + consumed += 1; + continue; + } + + let mut admitted = None; + for step in 0..self.devices.len() { + let device_index = (self.next_device + step) % self.devices.len(); + let admission = self.devices[device_index].cache.admit( + tensor.key.clone(), + tensor.payload.len(), + tensor.descriptor.rows, + tensor.descriptor.dimension, + tensor.descriptor.dtype, + false, + false, + ); + if let Admission::Admitted { offset, .. } = admission { + admitted = Some((device_index, offset, false)); + self.next_device = (device_index + 1) % self.devices.len(); + break; + } + } + + if admitted.is_none() && chunks.iter().all(Vec::is_empty) { + // TinyLFU rejected the cold item on every device, but the + // request must still be computed. Use one transient slab + // and remove it after the chunk completes. + for step in 0..self.devices.len() { + let device_index = (self.next_device + step) % self.devices.len(); + if let Admission::Admitted { offset, .. } = + self.devices[device_index].cache.admit( + tensor.key.clone(), + tensor.payload.len(), + tensor.descriptor.rows, + tensor.descriptor.dimension, + tensor.descriptor.dtype, + false, + true, + ) + { + admitted = Some((device_index, offset, true)); + self.next_device = (device_index + 1) % self.devices.len(); + break; + } + } + } + + let Some((device_index, offset, transient)) = admitted else { + if chunks.iter().any(|chunk| !chunk.is_empty()) { + break; + } + bail!("one tensor cannot be scheduled in any configured Rust GPU block cache"); + }; + uploads[device_index].push((offset, tensor.payload.as_slice())); + chunks[device_index].push(ResidentTensor { + candidate_index: tensor.candidate_index, + device: device_index, + key: tensor.key.clone(), + offset, + rows: tensor.descriptor.rows, + transient, + }); + consumed += 1; + } + + if consumed == 0 { + bail!("Rust multi-GPU scheduler made no progress"); + } + self.execute_devices(request, &chunks, &uploads, &mut scores)?; + for chunk in &chunks { + self.release_chunk(chunk)?; + } + pending.drain(..consumed); + } + + request + .candidates + .iter() + .enumerate() + .map(|(index, descriptor)| { + scores[index] + .map(|score| (descriptor.candidate_id, score)) + .ok_or_else(|| anyhow!("missing native TileMaxSim result")) + }) + .collect() + } + + fn execute_devices( + &mut self, + request: &Request, + chunks: &[Vec], + uploads: &[Vec<(u64, &[u8])>], + scores: &mut [Option], + ) -> Result<()> { + let completed = std::thread::scope(|scope| -> Result>> { + let mut workers = Vec::new(); + for ((device, chunk), upload) in self.devices.iter_mut().zip(chunks).zip(uploads) { + if chunk.is_empty() { + continue; + } + workers.push(scope.spawn(move || -> Result> { + if !upload.is_empty() { + device.gpu.upload_batch(upload)?; + device.h2d_batches += 1; + device.h2d_bytes += upload + .iter() + .map(|(_, payload)| payload.len() as u64) + .sum::(); + } + let offsets = chunk.iter().map(|item| item.offset).collect::>(); + let rows = chunk.iter().map(|item| item.rows).collect::>(); + let computed = device.gpu.score( + &request.query, + request.query_rows, + request.dimension, + request.dtype, + &offsets, + &rows, + )?; + Ok(chunk + .iter() + .zip(computed) + .map(|(tensor, score)| (tensor.candidate_index, score)) + .collect()) + })); + } + let mut completed = Vec::new(); + for worker in workers { + completed.push( + worker + .join() + .map_err(|_| anyhow!("GPU worker panicked"))??, + ); + } + Ok(completed) + })?; + for device_scores in completed { + for (candidate_index, score) in device_scores { + scores[candidate_index] = Some(score); + } + } + Ok(()) + } + + fn release_chunk(&mut self, chunk: &[ResidentTensor]) -> Result<()> { + let mut transient = std::collections::HashSet::new(); + for tensor in chunk { + self.devices[tensor.device] + .cache + .release(&tensor.key) + .map_err(|message| anyhow!(message))?; + if tensor.transient { + transient.insert((tensor.device, tensor.key.clone())); + } + } + for (device, key) in transient { + self.devices[device] + .cache + .remove(&key) + .map_err(|message| anyhow!(message))?; + } + Ok(()) + } + + pub fn status_json(&self) -> serde_json::Value { + let (host_hits, host_misses, host_evictions, host_rejections) = self.store.host_status(); + let devices = self + .devices + .iter() + .enumerate() + .map(|(index, device)| { + serde_json::json!({ + "index": index, + "gpu_tensor_bytes": device.cache.capacity(), + "gpu_block_bytes": device.cache.block_bytes(), + "gpu_entries": device.cache.entry_count(), + "gpu_pinned_entries": device.cache.pinned_entries(), + "gpu_hits": device.cache.hits, + "gpu_misses": device.cache.misses, + "gpu_evictions": device.cache.evictions, + "gpu_admission_rejections": device.cache.admission_rejections, + "h2d_batches": device.h2d_batches, + "h2d_bytes": device.h2d_bytes, + }) + }) + .collect::>(); + serde_json::json!({ + "devices": devices, + "host_hits": host_hits, + "host_misses": host_misses, + "host_evictions": host_evictions, + "host_admission_rejections": host_rejections, + "batch_read_calls": self.store.batch_read_calls, + "batch_read_bytes": self.store.batch_read_bytes, + }) + } +} + +fn validate_entry(descriptor: &Descriptor, entry: &crate::cache::CacheEntry) -> Result<()> { + let scalar_bytes = if descriptor.dtype == 1 { 4 } else { 2 }; + let expected_bytes = descriptor.rows as usize * descriptor.dimension as usize * scalar_bytes; + if entry.rows != descriptor.rows + || entry.dimension != descriptor.dimension + || entry.dtype != descriptor.dtype + || entry.payload_bytes != expected_bytes + || entry.allocated_bytes < expected_bytes + { + bail!("GPU cache metadata disagrees with the tensor descriptor"); + } + Ok(()) +} diff --git a/services/tilemaxsimd/src/gpu.rs b/services/tilemaxsimd/src/gpu.rs new file mode 100644 index 00000000..2c6c2df0 --- /dev/null +++ b/services/tilemaxsimd/src/gpu.rs @@ -0,0 +1,163 @@ +use anyhow::{Result, anyhow, bail}; +use std::ffi::{CStr, c_char, c_int, c_uchar, c_void}; +use std::ptr::NonNull; + +#[repr(C)] +struct NativeGpu(c_void); + +unsafe extern "C" { + fn vctm_gpu_create( + device: c_int, + total_bytes: usize, + workspace_bytes: usize, + output: *mut *mut NativeGpu, + error: *mut c_char, + error_capacity: usize, + ) -> c_int; + fn vctm_gpu_destroy(gpu: *mut NativeGpu); + fn vctm_gpu_tensor_bytes(gpu: *const NativeGpu) -> usize; + fn vctm_gpu_upload_batch( + gpu: *mut NativeGpu, + offsets: *const u64, + payloads: *const *const c_uchar, + lengths: *const usize, + count: usize, + error: *mut c_char, + error_capacity: usize, + ) -> c_int; + fn vctm_gpu_score( + gpu: *mut NativeGpu, + query: *const c_uchar, + query_bytes: usize, + query_rows: u32, + dimension: u32, + dtype: u8, + document_offsets: *const u64, + document_rows: *const u32, + count: usize, + output: *mut f32, + error: *mut c_char, + error_capacity: usize, + ) -> c_int; +} + +pub struct Gpu { + native: NonNull, + tensor_bytes: usize, +} + +// SAFETY: `Gpu` uniquely owns the native handle. It may move to a scoped +// worker, but no method exposes the pointer and all calls require `&mut self`. +unsafe impl Send for Gpu {} + +impl Gpu { + pub fn create(device: i32, total_bytes: usize, workspace_bytes: usize) -> Result { + let mut native = std::ptr::null_mut(); + let mut error = [0_i8; 512]; + // SAFETY: the C API writes one opaque pointer and a bounded error string. + let status = unsafe { + vctm_gpu_create( + device, + total_bytes, + workspace_bytes, + &mut native, + error.as_mut_ptr(), + error.len(), + ) + }; + if status != 0 { + bail!(native_error(&error)); + } + let native = NonNull::new(native).ok_or_else(|| anyhow!("CUDA returned a null arena"))?; + // SAFETY: `native` is live until Drop. + let tensor_bytes = unsafe { vctm_gpu_tensor_bytes(native.as_ptr()) }; + Ok(Self { + native, + tensor_bytes, + }) + } + + pub fn tensor_bytes(&self) -> usize { + self.tensor_bytes + } + + pub fn upload_batch(&mut self, items: &[(u64, &[u8])]) -> Result<()> { + if items.is_empty() { + return Ok(()); + } + let offsets = items.iter().map(|item| item.0).collect::>(); + let payloads = items.iter().map(|item| item.1.as_ptr()).collect::>(); + let lengths = items.iter().map(|item| item.1.len()).collect::>(); + let mut error = [0_i8; 512]; + // SAFETY: all slices remain alive for the synchronous native batch call. + let status = unsafe { + vctm_gpu_upload_batch( + self.native.as_ptr(), + offsets.as_ptr(), + payloads.as_ptr(), + lengths.as_ptr(), + items.len(), + error.as_mut_ptr(), + error.len(), + ) + }; + if status != 0 { + bail!(native_error(&error)); + } + Ok(()) + } + + pub fn score( + &mut self, + query: &[u8], + query_rows: u32, + dimension: u32, + dtype: u8, + document_offsets: &[u64], + document_rows: &[u32], + ) -> Result> { + if document_offsets.len() != document_rows.len() || document_offsets.is_empty() { + bail!("invalid native document metadata"); + } + let mut output = vec![0.0_f32; document_offsets.len()]; + let mut error = [0_i8; 512]; + // SAFETY: the native call is synchronous and receives valid slice pointers. + let status = unsafe { + vctm_gpu_score( + self.native.as_ptr(), + query.as_ptr(), + query.len(), + query_rows, + dimension, + dtype, + document_offsets.as_ptr(), + document_rows.as_ptr(), + document_offsets.len(), + output.as_mut_ptr(), + error.as_mut_ptr(), + error.len(), + ) + }; + if status != 0 { + bail!(native_error(&error)); + } + if output.iter().any(|score| !score.is_finite()) { + bail!("native TileMaxSim returned a non-finite score"); + } + Ok(output) + } +} + +impl Drop for Gpu { + fn drop(&mut self) { + // SAFETY: this is the unique owned native pointer. + unsafe { vctm_gpu_destroy(self.native.as_ptr()) }; + } +} + +fn native_error(buffer: &[c_char]) -> String { + // SAFETY: the native helper always NUL-terminates a nonempty error buffer. + unsafe { CStr::from_ptr(buffer.as_ptr()) } + .to_string_lossy() + .into_owned() +} diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs new file mode 100644 index 00000000..847dc497 --- /dev/null +++ b/services/tilemaxsimd/src/main.rs @@ -0,0 +1,321 @@ +mod cache; +mod engine; +mod gpu; +mod protocol; +mod shard; + +use anyhow::{Context, Result, anyhow, bail}; +use clap::Parser; +use engine::Engine; +use gpu::Gpu; +use protocol::HEADER_BYTES; +use serde::Deserialize; +use shard::ShardStore; +use std::fs; +use std::io::{BufRead, Read, Write}; +use std::os::unix::fs::{FileTypeExt, PermissionsExt}; +use std::os::unix::net::{UnixListener, UnixStream}; +use std::path::PathBuf; +use std::time::Instant; + +const GIB: usize = 1024 * 1024 * 1024; + +#[derive(Parser)] +#[command(about = "Native Rust/CUDA TileMaxSim shard, cache, and scheduling daemon")] +struct Args { + #[arg(long)] + socket: PathBuf, + #[arg(long, required = true, value_parser = parse_gpu_memory)] + gpu_memory_gb: Vec, + #[arg(long, default_value = "2", value_parser = parse_gb)] + gpu_workspace_gb: usize, + #[arg(long, default_value = "8", value_parser = parse_gb)] + host_cache_gb: usize, + #[arg(long = "contract-root", required = true, value_parser = parse_contract_root)] + contract_roots: Vec<(String, PathBuf)>, + #[arg(long, default_value_t = 256)] + gpu_block_kib: usize, + #[arg(long, default_value_t = 64 * 1024 * 1024)] + max_request_bytes: usize, + #[arg(long, default_value = "600", value_parser = parse_mode)] + socket_mode: u32, + #[arg(long)] + once: bool, + #[arg(long)] + verify_full_shards: bool, + #[arg(long, default_value = "lru", value_parser = ["lru", "resident"])] + gpu_cache_mode: String, + #[arg(long = "resident-manifest", value_parser = parse_contract_root)] + resident_manifests: Vec<(String, PathBuf)>, + #[arg(long, default_value_t = 256)] + prewarm_batch_size: usize, +} + +#[derive(Clone)] +struct GpuMemory { + device: i32, + bytes: usize, +} + +fn parse_gpu_memory(value: &str) -> Result { + let (device, gb) = value + .strip_prefix("cuda:") + .unwrap_or(value) + .split_once('=') + .ok_or_else(|| "GPU memory must be GPU=GB, for example 1=20".to_owned())?; + let device = device + .parse::() + .map_err(|_| "GPU index must be a nonnegative integer".to_owned())?; + if device < 0 { + return Err("GPU index must be nonnegative".to_owned()); + } + let bytes = parse_gb(gb)?; + Ok(GpuMemory { device, bytes }) +} + +fn parse_gb(value: &str) -> Result { + if value.is_empty() + || !value + .bytes() + .all(|byte| byte.is_ascii_digit() || byte == b'.') + { + return Err("memory size must be a positive number of GB".to_owned()); + } + let gb = value + .parse::() + .map_err(|_| "memory size must be a positive number of GB".to_owned())?; + if !gb.is_finite() || gb <= 0.0 || gb * GIB as f64 > usize::MAX as f64 { + return Err("memory size must be a positive number of GB".to_owned()); + } + Ok((gb * GIB as f64) as usize) +} + +fn parse_contract_root(value: &str) -> Result<(String, PathBuf), String> { + let (contract, path) = value + .split_once('=') + .ok_or_else(|| "contract root must be MODEL_CONTRACT_ID=/absolute/path".to_owned())?; + let path = PathBuf::from(path); + if contract.is_empty() || !path.is_absolute() { + return Err("contract root must contain a nonempty ID and absolute path".to_owned()); + } + Ok((contract.to_owned(), path)) +} + +fn parse_mode(value: &str) -> Result { + let mode = u32::from_str_radix(value, 8).map_err(|_| "invalid octal socket mode".to_owned())?; + if mode > 0o777 { + return Err("socket mode must be between 000 and 777".to_owned()); + } + Ok(mode) +} + +fn main() -> Result<()> { + let args = Args::parse(); + let mut seen_devices = std::collections::HashSet::new(); + for specification in &args.gpu_memory_gb { + if args.gpu_workspace_gb >= specification.bytes { + bail!("every configured GPU allocation must exceed its workspace"); + } + if !seen_devices.insert(specification.device) { + bail!("each CUDA device may be configured only once"); + } + } + let block_bytes = args + .gpu_block_kib + .checked_mul(1024) + .ok_or_else(|| anyhow!("GPU block size overflow"))?; + if block_bytes == 0 || block_bytes % 256 != 0 { + bail!("GPU block size must be positive and 256-byte aligned"); + } + let store = ShardStore::open( + &args.contract_roots, + args.host_cache_gb, + args.verify_full_shards, + )?; + let gpus = args + .gpu_memory_gb + .iter() + .map(|specification| { + Gpu::create( + specification.device, + specification.bytes, + args.gpu_workspace_gb, + ) + }) + .collect::>>()?; + let mut engine = Engine::new(gpus, block_bytes, store)?; + if args.gpu_cache_mode == "resident" && args.resident_manifests.is_empty() { + bail!("resident GPU cache mode requires at least one resident manifest"); + } + if args.gpu_cache_mode == "lru" && !args.resident_manifests.is_empty() { + bail!("resident manifests are valid only in resident GPU cache mode"); + } + if args.gpu_cache_mode == "resident" { + let prewarm_started = Instant::now(); + let descriptors = load_resident_manifests(&args.resident_manifests)?; + engine.prewarm(&descriptors, args.prewarm_batch_size)?; + println!( + "{}", + serde_json::json!({ + "event": "tilemaxsim_rust_prewarm_complete", + "entries": descriptors.len(), + "elapsed_ms": prewarm_started.elapsed().as_secs_f64() * 1000.0, + "cache": engine.status_json(), + }) + ); + } + remove_stale_socket(&args.socket)?; + let listener = UnixListener::bind(&args.socket) + .with_context(|| format!("cannot bind {}", args.socket.display()))?; + fs::set_permissions(&args.socket, fs::Permissions::from_mode(args.socket_mode))?; + println!( + "{}", + serde_json::json!({ + "event": "tilemaxsim_rust_ready", + "socket": args.socket, + "devices": args.gpu_memory_gb.iter().map(|item| serde_json::json!({ + "device": item.device, + "allocated_bytes": item.bytes, + })).collect::>(), + "workspace_bytes": args.gpu_workspace_gb, + "cache": engine.status_json(), + }) + ); + let mut accepted = 0_usize; + for connection in listener.incoming() { + let mut connection = connection?; + let started = Instant::now(); + let response = match read_request(&mut connection, args.max_request_bytes) { + Ok(frame) => { + let request_id = header_request_id(&frame); + match protocol::parse(&frame).and_then(|request| { + let request_id = request.request_id; + engine + .score(&request) + .map(|results| protocol::success(request_id, &results)) + }) { + Ok(response) => response, + Err(error) => protocol::failure(request_id, 3, &format!("{error:#}")), + } + } + Err(error) => protocol::failure(0, 1, &format!("{error:#}")), + }; + connection.write_all(&response)?; + println!( + "{}", + serde_json::json!({ + "event": "tilemaxsim_rust_request", + "elapsed_ms": started.elapsed().as_secs_f64() * 1000.0, + "cache": engine.status_json(), + }) + ); + accepted += 1; + if args.once && accepted == 1 { + break; + } + } + drop(listener); + match fs::remove_file(&args.socket) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error.into()), + } + Ok(()) +} + +fn read_request(connection: &mut UnixStream, maximum: usize) -> Result> { + let mut header = [0_u8; HEADER_BYTES]; + connection.read_exact(&mut header)?; + let body_bytes = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) + .context("request body does not fit this host")?; + let total = HEADER_BYTES + .checked_add(body_bytes) + .ok_or_else(|| anyhow!("request length overflow"))?; + if total > maximum { + bail!("request exceeds byte limit"); + } + let mut frame = Vec::with_capacity(total); + frame.extend_from_slice(&header); + frame.resize(total, 0); + connection.read_exact(&mut frame[HEADER_BYTES..])?; + Ok(frame) +} + +fn header_request_id(frame: &[u8]) -> u64 { + if frame.len() < HEADER_BYTES { + 0 + } else { + u64::from_le_bytes(frame[8..16].try_into().unwrap()) + } +} + +fn remove_stale_socket(path: &PathBuf) -> Result<()> { + let Ok(metadata) = fs::symlink_metadata(path) else { + return Ok(()); + }; + if !metadata.file_type().is_socket() { + bail!("refusing to remove non-socket path {}", path.display()); + } + fs::remove_file(path)?; + Ok(()) +} + +#[derive(Deserialize)] +struct ResidentRecord { + tensor_ref: String, + tensor_rows: u32, + tensor_dim: u32, + tensor_dtype: String, + tensor_checksum: String, + canonical_bytes: Option, +} + +fn load_resident_manifests(values: &[(String, PathBuf)]) -> Result> { + let mut descriptors = Vec::new(); + for (contract, path) in values { + let file = fs::File::open(path) + .with_context(|| format!("cannot open resident manifest {}", path.display()))?; + for (line_number, line) in std::io::BufReader::new(file).lines().enumerate() { + let line = line?; + let record: ResidentRecord = serde_json::from_str(&line).with_context(|| { + format!( + "invalid resident manifest {}:{}", + path.display(), + line_number + 1 + ) + })?; + let digest = record + .tensor_ref + .strip_prefix("sha256://") + .ok_or_else(|| anyhow!("resident tensor reference is not SHA-256"))?; + if record.tensor_checksum != format!("sha256:{digest}") { + bail!("resident tensor reference and checksum disagree"); + } + let (dtype, scalar_bytes) = match record.tensor_dtype.as_str() { + "float16" => (2, 2), + "float32" => (1, 4), + _ => bail!("resident manifest has an unsupported tensor dtype"), + }; + let expected_bytes = + record.tensor_rows as usize * record.tensor_dim as usize * scalar_bytes; + if record + .canonical_bytes + .is_some_and(|bytes| bytes != expected_bytes) + { + bail!("resident canonical byte length disagrees with its shape"); + } + descriptors.push(protocol::Descriptor { + candidate_id: descriptors.len() as u32, + contract: contract.clone(), + digest: digest.to_owned(), + rows: record.tensor_rows, + dimension: record.tensor_dim, + dtype, + }); + } + } + if descriptors.is_empty() { + bail!("resident manifests contain no tensor descriptors"); + } + Ok(descriptors) +} diff --git a/services/tilemaxsimd/src/protocol.rs b/services/tilemaxsimd/src/protocol.rs new file mode 100644 index 00000000..d02cae01 --- /dev/null +++ b/services/tilemaxsimd/src/protocol.rs @@ -0,0 +1,238 @@ +use anyhow::{Result, anyhow, bail}; +use std::collections::HashSet; + +pub const HEADER_BYTES: usize = 24; +pub const VERSION_EXTERNAL: u16 = 2; +const MAGIC: &[u8; 4] = b"VCTM"; +const REQUEST_KIND: u16 = 1; +const RESPONSE_KIND: u16 = 2; + +#[derive(Clone, Debug)] +pub struct Descriptor { + pub candidate_id: u32, + pub contract: String, + pub digest: String, + pub rows: u32, + pub dimension: u32, + pub dtype: u8, +} + +#[derive(Debug)] +pub struct Request { + pub request_id: u64, + pub query_rows: u32, + pub dimension: u32, + pub dtype: u8, + pub query: Vec, + pub candidates: Vec, +} + +struct Reader<'a> { + bytes: &'a [u8], + offset: usize, +} + +impl<'a> Reader<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, offset: 0 } + } + + fn take(&mut self, count: usize) -> Result<&'a [u8]> { + let end = self + .offset + .checked_add(count) + .ok_or_else(|| anyhow!("request overflow"))?; + if end > self.bytes.len() { + bail!("truncated request"); + } + let result = &self.bytes[self.offset..end]; + self.offset = end; + Ok(result) + } + + fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + + fn u16(&mut self) -> Result { + Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap())) + } + + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + + fn text(&mut self, count: usize, maximum: usize, name: &str) -> Result { + if count == 0 || count > maximum { + bail!("invalid {name} length"); + } + let value = std::str::from_utf8(self.take(count)?)?; + if value.chars().any(|character| character.is_control()) { + bail!("{name} contains control characters"); + } + Ok(value.to_owned()) + } + + fn finish(self) -> Result<()> { + if self.offset != self.bytes.len() { + bail!("trailing request bytes"); + } + Ok(()) + } +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: u8) -> Result { + let scalar = match dtype { + 1 => 4, + 2 => 2, + _ => bail!("unsupported tensor dtype"), + }; + if rows == 0 || dimension == 0 || dimension > 60_000 { + bail!("invalid tensor shape"); + } + (rows as usize) + .checked_mul(dimension as usize) + .and_then(|value| value.checked_mul(scalar)) + .ok_or_else(|| anyhow!("tensor shape is too large")) +} + +fn validate_finite(payload: &[u8], dtype: u8) -> Result<()> { + if dtype == 1 { + for scalar in payload.chunks_exact(4) { + if !f32::from_bits(u32::from_le_bytes(scalar.try_into().unwrap())).is_finite() { + bail!("query contains a non-finite value"); + } + } + } else { + for scalar in payload.chunks_exact(2) { + let bits = u16::from_le_bytes(scalar.try_into().unwrap()); + if bits & 0x7c00 == 0x7c00 { + bail!("query contains a non-finite value"); + } + } + } + Ok(()) +} + +pub fn parse(frame: &[u8]) -> Result { + if frame.len() < HEADER_BYTES { + bail!("truncated request header"); + } + if &frame[..4] != MAGIC { + bail!("invalid protocol magic"); + } + let version = u16::from_le_bytes(frame[4..6].try_into().unwrap()); + let kind = u16::from_le_bytes(frame[6..8].try_into().unwrap()); + let request_id = u64::from_le_bytes(frame[8..16].try_into().unwrap()); + let body_bytes = u64::from_le_bytes(frame[16..24].try_into().unwrap()); + if version != VERSION_EXTERNAL || kind != REQUEST_KIND { + bail!("Rust daemon requires TileMaxSim external protocol v2"); + } + if usize::try_from(body_bytes).ok() != Some(frame.len() - HEADER_BYTES) { + bail!("request body length mismatch"); + } + let mut reader = Reader::new(&frame[HEADER_BYTES..]); + let dimension = reader.u32()?; + let query_rows = reader.u32()?; + let candidate_count = reader.u32()?; + let dtype = reader.u8()?; + let scoring = reader.u8()?; + let reserved = reader.u16()?; + let contract_bytes = reader.u32()? as usize; + if scoring != 1 || reserved != 0 { + bail!("unsupported scoring function or reserved bits"); + } + if candidate_count > 65_536 { + bail!("too many candidates"); + } + let contract = reader.text(contract_bytes, 512, "model contract")?; + let query_bytes = tensor_bytes(query_rows, dimension, dtype)?; + let query = reader.take(query_bytes)?.to_vec(); + validate_finite(&query, dtype)?; + let mut total_tokens = query_rows as usize; + let mut total_bytes = query_bytes; + let mut candidate_ids = HashSet::new(); + let mut candidates = Vec::with_capacity(candidate_count as usize); + for _ in 0..candidate_count { + let candidate_id = reader.u32()?; + let rows = reader.u32()?; + let reference_bytes = reader.u32()? as usize; + let checksum_bytes = reader.u32()? as usize; + if !candidate_ids.insert(candidate_id) { + bail!("duplicate candidate ID"); + } + let tensor_ref = reader.text(reference_bytes, 4096, "tensor reference")?; + let checksum = reader.text(checksum_bytes, 512, "tensor checksum")?; + let digest = tensor_ref + .strip_prefix("sha256://") + .ok_or_else(|| anyhow!("unsupported tensor reference"))?; + if digest.len() != 64 + || !digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + || checksum != format!("sha256:{digest}") + { + bail!("invalid content-addressed tensor descriptor"); + } + let bytes = tensor_bytes(rows, dimension, dtype)?; + total_tokens = total_tokens + .checked_add(rows as usize) + .ok_or_else(|| anyhow!("token overflow"))?; + total_bytes = total_bytes + .checked_add(bytes) + .ok_or_else(|| anyhow!("byte overflow"))?; + if total_tokens > 1_000_000 || total_bytes > 1024 * 1024 * 1024 { + bail!("request exceeds tensor limits"); + } + candidates.push(Descriptor { + candidate_id, + contract: contract.clone(), + digest: digest.to_owned(), + rows, + dimension, + dtype, + }); + } + reader.finish()?; + Ok(Request { + request_id, + query_rows, + dimension, + dtype, + query, + candidates, + }) +} + +pub fn success(request_id: u64, results: &[(u32, f32)]) -> Vec { + let body_bytes = 8 + results.len() * 8; + let mut frame = Vec::with_capacity(HEADER_BYTES + body_bytes); + frame.extend_from_slice(MAGIC); + frame.extend_from_slice(&VERSION_EXTERNAL.to_le_bytes()); + frame.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + frame.extend_from_slice(&request_id.to_le_bytes()); + frame.extend_from_slice(&(body_bytes as u64).to_le_bytes()); + frame.extend_from_slice(&0_u32.to_le_bytes()); + frame.extend_from_slice(&(results.len() as u32).to_le_bytes()); + for (candidate_id, score) in results { + frame.extend_from_slice(&candidate_id.to_le_bytes()); + frame.extend_from_slice(&score.to_le_bytes()); + } + frame +} + +pub fn failure(request_id: u64, status: u32, message: &str) -> Vec { + let message = message.as_bytes(); + let message = &message[..message.len().min(64 * 1024)]; + let body_bytes = 8 + message.len(); + let mut frame = Vec::with_capacity(HEADER_BYTES + body_bytes); + frame.extend_from_slice(MAGIC); + frame.extend_from_slice(&VERSION_EXTERNAL.to_le_bytes()); + frame.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + frame.extend_from_slice(&request_id.to_le_bytes()); + frame.extend_from_slice(&(body_bytes as u64).to_le_bytes()); + frame.extend_from_slice(&status.max(1).to_le_bytes()); + frame.extend_from_slice(&(message.len() as u32).to_le_bytes()); + frame.extend_from_slice(message); + frame +} diff --git a/services/tilemaxsimd/src/shard.rs b/services/tilemaxsimd/src/shard.rs new file mode 100644 index 00000000..1a60e3cb --- /dev/null +++ b/services/tilemaxsimd/src/shard.rs @@ -0,0 +1,428 @@ +use crate::cache::TinyLfu; +use crate::protocol::Descriptor; +use anyhow::{Context, Result, anyhow, bail}; +use serde::Deserialize; +use sha2::{Digest, Sha256}; +use std::collections::HashMap; +use std::fs::{File, OpenOptions}; +use std::io::Read; +use std::os::unix::fs::{FileExt, OpenOptionsExt}; +use std::path::{Path, PathBuf}; + +const INDEX_NAME: &str = "tilemaxsim-shards-v1.json"; + +#[derive(Deserialize)] +struct RawIndex { + format: String, + version: u32, + alignment: usize, + shards: Vec, + entries: Vec, +} + +#[derive(Deserialize)] +struct RawShard { + name: String, + #[serde(rename = "bytes")] + size: usize, + checksum: String, +} + +#[derive(Deserialize)] +struct RawEntry { + digest: String, + shard: String, + offset: u64, + length: usize, + rows: u32, + dimension: u32, + dtype: String, +} + +struct ShardFile { + file: File, + size: usize, + checksum: String, + verified: bool, +} + +#[derive(Clone)] +struct Entry { + shard: String, + offset: u64, + length: usize, + rows: u32, + dimension: u32, + dtype: u8, +} + +struct ContractStore { + shards: HashMap, + entries: HashMap, +} + +struct HostEntry { + payload: Vec, + priority: f64, +} + +struct HostCache { + maximum_bytes: usize, + current_bytes: usize, + entries: HashMap, + sketch: TinyLfu, + inflation: f64, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, +} + +impl HostCache { + fn new(maximum_bytes: usize) -> Self { + Self { + maximum_bytes, + current_bytes: 0, + entries: HashMap::new(), + sketch: TinyLfu::new(4096), + inflation: 0.0, + hits: 0, + misses: 0, + evictions: 0, + admission_rejections: 0, + } + } + + fn get(&mut self, key: &str) -> Option> { + let frequency = self.sketch.increment(&key); + let entry = self.entries.get_mut(key); + if let Some(entry) = entry { + entry.priority = self.inflation + f64::from(frequency) / entry.payload.len() as f64; + self.hits += 1; + Some(entry.payload.clone()) + } else { + self.misses += 1; + None + } + } + + fn put(&mut self, key: String, payload: Vec) { + if self.maximum_bytes == 0 || payload.len() > self.maximum_bytes { + return; + } + if let Some(previous) = self.entries.remove(&key) { + self.current_bytes -= previous.payload.len(); + } + let frequency = self.sketch.estimate(&key).max(1); + let mut priority = self.inflation + f64::from(frequency) / payload.len() as f64; + while self.current_bytes + payload.len() > self.maximum_bytes { + let Some((victim_key, victim_priority)) = self + .entries + .iter() + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + .map(|(key, entry)| (key.clone(), entry.priority)) + else { + return; + }; + if priority < victim_priority { + self.admission_rejections += 1; + return; + } + let victim = self + .entries + .remove(&victim_key) + .expect("host victim disappeared"); + self.current_bytes -= victim.payload.len(); + self.inflation = self.inflation.max(victim.priority); + priority = self.inflation + f64::from(frequency) / payload.len() as f64; + self.evictions += 1; + } + self.current_bytes += payload.len(); + self.entries.insert(key, HostEntry { payload, priority }); + } +} + +pub struct ShardStore { + contracts: HashMap, + host_cache: HostCache, + pub batch_read_calls: u64, + pub batch_read_bytes: u64, + verify_full_shards: bool, +} + +impl ShardStore { + pub fn open( + roots: &[(String, PathBuf)], + host_cache_bytes: usize, + verify_full_shards: bool, + ) -> Result { + let mut contracts = HashMap::new(); + for (contract, root) in roots { + if contract.is_empty() || contracts.contains_key(contract) { + bail!("invalid or duplicate model contract {contract:?}"); + } + contracts.insert(contract.clone(), open_contract(root)?); + } + if contracts.is_empty() { + bail!("at least one immutable shard contract is required"); + } + Ok(Self { + contracts, + host_cache: HostCache::new(host_cache_bytes), + batch_read_calls: 0, + batch_read_bytes: 0, + verify_full_shards, + }) + } + + pub fn resolve_many(&mut self, descriptors: &[Descriptor]) -> Result>> { + let mut output = vec![None; descriptors.len()]; + let mut groups = HashMap::<(String, String), Vec<(usize, Entry)>>::new(); + for (index, descriptor) in descriptors.iter().enumerate() { + let key = cache_key(descriptor); + if let Some(payload) = self.host_cache.get(&key) { + output[index] = Some(payload); + continue; + } + let contract = self + .contracts + .get(&descriptor.contract) + .ok_or_else(|| anyhow!("model contract has no immutable shard root"))?; + let entry = contract + .entries + .get(&descriptor.digest) + .ok_or_else(|| anyhow!("tensor is missing from the immutable shard index"))? + .clone(); + if entry.rows != descriptor.rows + || entry.dimension != descriptor.dimension + || entry.dtype != descriptor.dtype + || entry.length + != tensor_bytes(descriptor.rows, descriptor.dimension, descriptor.dtype)? + { + bail!("tensor descriptor disagrees with the immutable shard index"); + } + groups + .entry((descriptor.contract.clone(), entry.shard.clone())) + .or_default() + .push((index, entry)); + } + + for ((contract_name, shard_name), mut entries) in groups { + let contract = self.contracts.get_mut(&contract_name).unwrap(); + let shard = contract.shards.get_mut(&shard_name).unwrap(); + if self.verify_full_shards { + verify_shard(shard)?; + } + entries.sort_by_key(|(_, entry)| entry.offset); + let mut ranges = Vec::<(u64, u64, Vec<(usize, Entry)>)>::new(); + let mut cursor = 0; + while cursor < entries.len() { + let start = entries[cursor].1.offset; + let mut end = start + entries[cursor].1.length as u64; + let mut limit = cursor + 1; + while limit < entries.len() { + let candidate = &entries[limit].1; + let candidate_end = candidate.offset + candidate.length as u64; + if candidate.offset.saturating_sub(end) > 64 * 1024 + || candidate_end.saturating_sub(start) > 8 * 1024 * 1024 + { + break; + } + end = end.max(candidate_end); + limit += 1; + } + ranges.push((start, end, entries[cursor..limit].to_vec())); + cursor = limit; + } + let worker_count = ranges.len().min(8); + let ranges_per_worker = ranges.len().div_ceil(worker_count); + let resolved = std::thread::scope(|scope| -> Result> { + let mut workers = Vec::new(); + for ranges in ranges.chunks(ranges_per_worker) { + let file = shard.file.try_clone()?; + workers.push(scope.spawn(move || -> Result)>> { + let mut resolved = Vec::new(); + for (start, end, entries) in ranges { + let mut range = vec![0_u8; (end - start) as usize]; + file.read_exact_at(&mut range, *start)?; + for (output_index, entry) in entries { + let local = (entry.offset - start) as usize; + let payload = range[local..local + entry.length].to_vec(); + let actual = hex::encode(Sha256::digest(&payload)); + if actual != descriptors[*output_index].digest { + bail!("tensor checksum mismatch inside immutable shard"); + } + resolved.push((*output_index, payload)); + } + } + Ok(resolved) + })); + } + let mut resolved = Vec::new(); + for worker in workers { + resolved.extend( + worker + .join() + .map_err(|_| anyhow!("shard reader worker panicked"))??, + ); + } + Ok(resolved) + })?; + self.batch_read_calls += ranges.len() as u64; + self.batch_read_bytes += ranges + .iter() + .map(|(start, end, _)| end - start) + .sum::(); + for (output_index, payload) in resolved { + output[output_index] = Some(payload.clone()); + self.host_cache + .put(cache_key(&descriptors[output_index]), payload); + } + } + output + .into_iter() + .map(|payload| payload.ok_or_else(|| anyhow!("tensor resolution produced no payload"))) + .collect() + } + + pub fn host_status(&self) -> (u64, u64, u64, u64) { + ( + self.host_cache.hits, + self.host_cache.misses, + self.host_cache.evictions, + self.host_cache.admission_rejections, + ) + } +} + +fn open_contract(root: &Path) -> Result { + let index_path = root.join(INDEX_NAME); + let mut index_file = nofollow(&index_path)?; + let mut document = String::new(); + index_file.read_to_string(&mut document)?; + let raw: RawIndex = serde_json::from_str(&document).context("invalid immutable shard index")?; + if raw.format != "vectorchord.tilemaxsim.shards" || raw.version != 1 { + bail!("unsupported immutable shard index"); + } + if raw.alignment == 0 || !raw.alignment.is_power_of_two() { + bail!("invalid immutable shard alignment"); + } + let mut shards = HashMap::new(); + for shard in raw.shards { + let digest = shard + .checksum + .strip_prefix("sha256:") + .ok_or_else(|| anyhow!("invalid immutable shard checksum"))?; + if shard.name != format!("shards/sha256-{digest}.vts") + || !digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + { + bail!("invalid immutable shard name"); + } + let path = root.join(&shard.name); + let file = nofollow(&path)?; + if file.metadata()?.len() != shard.size as u64 { + bail!("immutable shard size mismatch"); + } + if shards + .insert( + shard.name, + ShardFile { + file, + size: shard.size, + checksum: digest.to_owned(), + verified: false, + }, + ) + .is_some() + { + bail!("duplicate immutable shard"); + } + } + let mut entries = HashMap::new(); + for entry in raw.entries { + let dtype = match entry.dtype.as_str() { + "float32" => 1, + "float16" => 2, + _ => bail!("invalid shard tensor dtype"), + }; + let shard = shards + .get(&entry.shard) + .ok_or_else(|| anyhow!("shard tensor references an unknown file"))?; + if !(entry.offset as usize).is_multiple_of(raw.alignment) + || entry.length != tensor_bytes(entry.rows, entry.dimension, dtype)? + || entry.offset + entry.length as u64 > shard.size as u64 + { + bail!("invalid shard tensor range"); + } + if entries + .insert( + entry.digest, + Entry { + shard: entry.shard, + offset: entry.offset, + length: entry.length, + rows: entry.rows, + dimension: entry.dimension, + dtype, + }, + ) + .is_some() + { + bail!("duplicate tensor digest in shard index"); + } + } + Ok(ContractStore { shards, entries }) +} + +fn nofollow(path: &Path) -> Result { + OpenOptions::new() + .read(true) + .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(path) + .with_context(|| format!("cannot open immutable shard path {}", path.display())) +} + +fn verify_shard(shard: &mut ShardFile) -> Result<()> { + if shard.verified { + return Ok(()); + } + let mut digest = Sha256::new(); + let mut offset = 0_u64; + let mut buffer = vec![0_u8; 8 * 1024 * 1024]; + while offset < shard.size as u64 { + let count = buffer.len().min(shard.size - offset as usize); + shard.file.read_exact_at(&mut buffer[..count], offset)?; + digest.update(&buffer[..count]); + offset += count as u64; + } + if hex::encode(digest.finalize()) != shard.checksum { + bail!("immutable shard checksum mismatch"); + } + shard.verified = true; + Ok(()) +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: u8) -> Result { + let scalar = match dtype { + 1 => 4, + 2 => 2, + _ => bail!("unsupported tensor dtype"), + }; + (rows as usize) + .checked_mul(dimension as usize) + .and_then(|value| value.checked_mul(scalar)) + .ok_or_else(|| anyhow!("tensor shape overflow")) +} + +pub fn cache_key(descriptor: &Descriptor) -> String { + format!( + "{}:{}:{}:{}:{}", + descriptor.contract, + descriptor.digest, + descriptor.rows, + descriptor.dimension, + descriptor.dtype + ) +} From ab3550923690e269df2b476a4b808dd1a8b8e8f0 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Tue, 14 Jul 2026 12:48:20 +0800 Subject: [PATCH 304/324] docs: publish TileMaxSim ablation results --- README.md | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/README.md b/README.md index edeb4ea2..86d6aabf 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,45 @@ multi-vector search. - Planner statistics and cost estimation for multi-vector queries. - Deterministic correctness, registry, sidecar-protocol, and planner-cost tests. +## Performance ablation + +We measured each cache-path optimization independently on the same development +machine. The corpus contained 34,054 tensor descriptors (34,027 unique tensors, +16.28 GB of logical FP16 tensor data). The request-level tests sampled 100 +candidates containing 47.81 MB of tensor data. Absolute latency depends on the +storage and GPU, so the same-run comparisons are more useful than the raw +numbers. + +| Optimization | Baseline | Optimized | Result | +| --- | ---: | ---: | ---: | +| Immutable shards and batched reads | 333.38 ms, sequential files | 56.52 ms | 5.90x faster | +| Shards versus batched legacy files | 87.56 ms | 56.52 ms | 1.55x faster | +| Batched host-to-device transfer | 43.07 ms, 100 transfers | 22.50 ms, one transfer | 1.91x faster | +| TinyLFU/GDSF admission | 69.98% LRU hit rate | 76.18% hit rate | +6.20 percentage points | +| Rust/CUDA cold request | 874.63 ms, Python/Triton | 122.19 ms | 7.16x faster | +| Rust/CUDA warm request p50 | 14.46 ms, Python/Triton | 1.77 ms | 8.17x faster | +| Rust/CUDA warm request p95 | 21.57 ms, Python/Triton | 1.83 ms | 11.79x faster | + +A full resident-cache run assigned 20 GiB to one GPU: 18 GiB for tensors and +2 GiB for the TileMaxSim workspace. All 34,027 unique tensors were pinned before +the service became ready. Process-to-ready time was 26.32 seconds for the +Python/Triton sidecar and 17.40 seconds for the Rust/CUDA daemon, a 33.9% +reduction. This is a one-time prewarm cost; resident warm requests do not read +the tensors from disk. + +The fixed-block buddy/slab allocator removed external-fragmentation failures in +a deterministic 20,000-operation allocation trace (220 with variable extents, +zero with slabs). It traded that unpredictability for 8.82% internal rounding +waste; total capacity failures were 220 for variable extents and 236 for slabs. +The allocator therefore improves predictability rather than claiming to create +additional memory. + +Rust/CUDA and Python/Triton produced identical top-10 results. The maximum +absolute score difference was 5.25e-6 and the mean difference was 3.45e-6. +The benchmark driver is +[`services/benchmark_tilemaxsim_ablation.py`](services/benchmark_tilemaxsim_ablation.py); +run it with `--help` for the corpus, cache-root, device, and output arguments. + The implementation is currently under active development. Its SQL interfaces and deployment packaging may change before a stable release. This repository contains only the public implementation and public-facing project information; From d25c93cdf660a2f158e3b20cd60cf87b6053d724 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Tue, 14 Jul 2026 13:57:10 +0800 Subject: [PATCH 305/324] perf(tilemaxsim): reduce GPU arena fragmentation --- README.md | 32 ++- services/benchmark_tilemaxsim_ablation.py | 242 ++++++++++++++----- services/test_tilemaxsim_gpu_cache.py | 39 +++- services/test_tilemaxsim_rust_daemon.py | 2 +- services/tilemaxsim_cuda_sidecar.py | 60 +++-- services/tilemaxsim_gpu_cache.py | 273 +++++++++++++--------- services/tilemaxsimd/src/cache.rs | 239 ++++++++++++++----- services/tilemaxsimd/src/engine.rs | 7 + services/tilemaxsimd/src/main.rs | 2 +- 9 files changed, 633 insertions(+), 263 deletions(-) diff --git a/README.md b/README.md index 86d6aabf..1d560cb0 100644 --- a/README.md +++ b/README.md @@ -35,25 +35,33 @@ numbers. | --- | ---: | ---: | ---: | | Immutable shards and batched reads | 333.38 ms, sequential files | 56.52 ms | 5.90x faster | | Shards versus batched legacy files | 87.56 ms | 56.52 ms | 1.55x faster | -| Batched host-to-device transfer | 43.07 ms, 100 transfers | 22.50 ms, one transfer | 1.91x faster | +| Batched host-to-device transfer | 37.06 ms, 100 transfers | 14.19 ms, one transfer | 2.61x faster | | TinyLFU/GDSF admission | 69.98% LRU hit rate | 76.18% hit rate | +6.20 percentage points | -| Rust/CUDA cold request | 874.63 ms, Python/Triton | 122.19 ms | 7.16x faster | -| Rust/CUDA warm request p50 | 14.46 ms, Python/Triton | 1.77 ms | 8.17x faster | -| Rust/CUDA warm request p95 | 21.57 ms, Python/Triton | 1.83 ms | 11.79x faster | +| Rust/CUDA cold request | 855.34 ms, Python/Triton | 93.86 ms | 9.11x faster | +| Rust/CUDA warm request p50 | 14.26 ms, Python/Triton | 2.09 ms | 6.82x faster | +| Rust/CUDA warm request p95 | 14.84 ms, Python/Triton | 2.20 ms | 6.75x faster | A full resident-cache run assigned 20 GiB to one GPU: 18 GiB for tensors and 2 GiB for the TileMaxSim workspace. All 34,027 unique tensors were pinned before -the service became ready. Process-to-ready time was 26.32 seconds for the -Python/Triton sidecar and 17.40 seconds for the Rust/CUDA daemon, a 33.9% +the service became ready. Process-to-ready time was 23.07 seconds for the +Python/Triton sidecar and 14.86 seconds for the Rust/CUDA daemon, a 35.6% reduction. This is a one-time prewarm cost; resident warm requests do not read the tensors from disk. -The fixed-block buddy/slab allocator removed external-fragmentation failures in -a deterministic 20,000-operation allocation trace (220 with variable extents, -zero with slabs). It traded that unpredictability for 8.82% internal rounding -waste; total capacity failures were 220 for variable extents and 236 for slabs. -The allocator therefore improves predictability rather than claiming to create -additional memory. +The GPU cache now suballocates exact contiguous page runs from one CUDA arena, +using best-fit size buckets and address-ordered coalescing. On the full corpus, +the 32 KiB default reduced allocated tensor space from 17.840 GB with the former +256 KiB power-of-two buddy allocator to 16.725 GB. It recovered 1.115 GB of GPU +space and reduced internal rounding waste from 8.82% to 2.74%. The default can +be overridden with `--gpu-block-kib`. + +In a deterministic 20,000-event churn trace, the former buddy allocator had 815 +failed cache-allocation attempts, while the segregated page-run allocator had +637; neither recorded an external-fragmentation failure on this workload. Exact +byte extents had 585 failures, all caused by external fragmentation. Page-run +metadata processing added about 1.4 microseconds per event versus the buddy +baseline. These are cache-admission attempts, not failed search requests: the +runtime evicts unpinned entries or streams oversized working sets in chunks. Rust/CUDA and Python/Triton produced identical top-10 results. The maximum absolute score difference was 5.25e-6 and the mean difference was 3.45e-6. diff --git a/services/benchmark_tilemaxsim_ablation.py b/services/benchmark_tilemaxsim_ablation.py index 5a6bbf18..cf3d29b0 100644 --- a/services/benchmark_tilemaxsim_ablation.py +++ b/services/benchmark_tilemaxsim_ablation.py @@ -34,6 +34,7 @@ from devtools.test_tilemaxsim_reference_sidecar import decode_response from services.tilemaxsim_cuda_sidecar import ContentAddressedResolver, PayloadCache from services.tilemaxsim_gpu_cache import ( + FixedBlockAllocator, FreeExtentAllocator, GpuArenaSpec, GpuResourcePool, @@ -42,6 +43,73 @@ ) +class _LegacyBuddyAllocator: + """Former power-of-two allocator retained only as an ablation baseline.""" + + def __init__(self, capacity: int, block_bytes: int = 256 * 1024) -> None: + self.block_bytes = block_bytes + self.block_count = capacity // block_bytes + self.capacity = self.block_count * block_bytes + self._free: dict[int, set[tuple[int, int, int]]] = {} + self._allocated: dict[int, tuple[int, int, int]] = {} + start = 0 + remaining = self.block_count + while remaining: + order = remaining.bit_length() - 1 + size = 1 << order + self._free.setdefault(order, set()).add((start, start, order)) + start += size + remaining -= size + + @property + def free_bytes(self) -> int: + return sum( + len(items) * (1 << order) * self.block_bytes + for order, items in self._free.items() + ) + + def allocation_bytes(self, payload_bytes: int) -> int: + raw = math.ceil(payload_bytes / self.block_bytes) + return (1 << (raw - 1).bit_length()) * self.block_bytes + + def allocate(self, payload_bytes: int) -> tuple[int, ...] | None: + required = self.allocation_bytes(payload_bytes) // self.block_bytes + order = required.bit_length() - 1 + available_order = next( + ( + candidate + for candidate in range(order, self.block_count.bit_length()) + if self._free.get(candidate) + ), + None, + ) + if available_order is None: + return None + start, root_start, root_order = self._free[available_order].pop() + while available_order > order: + available_order -= 1 + buddy = start + (1 << available_order) + self._free.setdefault(available_order, set()).add( + (buddy, root_start, root_order) + ) + self._allocated[start] = (order, root_start, root_order) + return tuple(range(start, start + required)) + + def release(self, blocks: tuple[int, ...]) -> None: + start = blocks[0] + order, root_start, root_order = self._allocated.pop(start) + while order < root_order: + buddy = root_start + ((start - root_start) ^ (1 << order)) + item = (buddy, root_start, root_order) + free = self._free.setdefault(order, set()) + if item not in free: + break + free.remove(item) + start = min(start, buddy) + order += 1 + self._free.setdefault(order, set()).add((start, root_start, root_order)) + + def percentile(samples: list[float], fraction: float) -> float: ordered = sorted(samples) return ordered[min(len(ordered) - 1, math.ceil(len(ordered) * fraction) - 1)] @@ -53,7 +121,11 @@ def load_records(path: Path) -> list[dict[str, object]]: def request(record: dict[str, object], contract: str) -> protocol.ExternalTensorRequest: - dtype = protocol.DTYPE_F16 if record["tensor_dtype"] == "float16" else protocol.DTYPE_F32 + dtype = ( + protocol.DTYPE_F16 + if record["tensor_dtype"] == "float16" + else protocol.DTYPE_F32 + ) return protocol.ExternalTensorRequest( contract, str(record["tensor_ref"]), @@ -76,7 +148,10 @@ def evict_paths(paths: list[Path]) -> None: def storage_ablation( - selected: list[dict[str, object]], contract: str, legacy_root: Path, shard_root: Path + selected: list[dict[str, object]], + contract: str, + legacy_root: Path, + shard_root: Path, ) -> dict[str, object]: requests = [request(record, contract) for record in selected] legacy_paths = [ @@ -143,7 +218,9 @@ def h2d_ablation( total_bytes = 768 * 1024**2 workspace_bytes = 256 * 1024**2 - pool = GpuResourcePool([GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes) + pool = GpuResourcePool( + [GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes + ) try: cache = GpuTensorCache(pool, allow_eviction=True) started = time.perf_counter() @@ -165,13 +242,13 @@ def h2d_ablation( finally: pool.close() - pool = GpuResourcePool([GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes) + pool = GpuResourcePool( + [GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes + ) try: cache = GpuTensorCache(pool, allow_eviction=True) loads = [ - GpuTensorLoad( - key, item.rows, item.dimension, item.dtype, resolved.payload - ) + GpuTensorLoad(key, item.rows, item.dimension, item.dtype, resolved.payload) for key, item, resolved in zip(keys, requests, payloads, strict=True) ] started = time.perf_counter() @@ -201,53 +278,82 @@ def allocator_ablation(records: list[dict[str, object]]) -> dict[str, object]: sizes = [int(record["canonical_bytes"]) for record in records] rng = random.Random(991) capacity = 256 * 1024**2 - extent = FreeExtentAllocator(capacity) - extent_live: list[tuple[int, int]] = [] - extent_failures = 0 - extent_fragmentation_failures = 0 + events: list[tuple[str, int, int]] = [] + abstract_live: list[int] = [] + next_identifier = 0 for _ in range(20_000): - if extent_live and rng.random() < 0.48: - extent.release(*extent_live.pop(rng.randrange(len(extent_live)))) + if abstract_live and rng.random() < 0.48: + index = rng.randrange(len(abstract_live)) + identifier = abstract_live.pop(index) + events.append(("release", identifier, 0)) else: size = rng.choice(sizes) - required = extent.allocation_bytes(size) - allocated = extent.allocate(size) - if allocated is None: - extent_failures += 1 - extent_fragmentation_failures += int(extent.free_bytes >= required) + identifier = next_identifier + next_identifier += 1 + abstract_live.append(identifier) + events.append(("allocate", identifier, size)) + + def run_trace( + allocator: FreeExtentAllocator | FixedBlockAllocator | _LegacyBuddyAllocator, + ) -> dict[str, object]: + live: dict[int, tuple[int, ...] | tuple[int, int]] = {} + failures = 0 + fragmentation_failures = 0 + requested_bytes = 0 + allocated_bytes = 0 + started = time.perf_counter() + for operation, identifier, size in events: + if operation == "release": + allocation = live.pop(identifier, None) + if allocation is None: + continue + if isinstance(allocator, FreeExtentAllocator): + allocator.release(*allocation) + else: + allocator.release(allocation) + continue + required = allocator.allocation_bytes(size) + allocation = allocator.allocate(size) + if allocation is None: + failures += 1 + fragmentation_failures += int(allocator.free_bytes >= required) else: - extent_live.append(allocated) - from services.tilemaxsim_gpu_cache import FixedBlockAllocator + live[identifier] = allocation + requested_bytes += size + allocated_bytes += ( + allocation[1] + if isinstance(allocator, FreeExtentAllocator) + else len(allocation) * allocator.block_bytes + ) + return { + "allocation_failures": failures, + "fragmentation_failures": fragmentation_failures, + "internal_waste_ratio": ( + (allocated_bytes - requested_bytes) / allocated_bytes + ), + "trace_ms": (time.perf_counter() - started) * 1000, + } - rng = random.Random(991) - slab = FixedBlockAllocator(capacity) - slab_live: list[tuple[int, ...]] = [] - slab_failures = 0 - slab_fragmentation_failures = 0 - requested = 0 - allocated_bytes = 0 - for _ in range(20_000): - if slab_live and rng.random() < 0.48: - slab.release(slab_live.pop(rng.randrange(len(slab_live)))) - else: - size = rng.choice(sizes) - required = slab.allocation_bytes(size) - allocated = slab.allocate(size) - if allocated is None: - slab_failures += 1 - slab_fragmentation_failures += int(slab.free_bytes >= required) - else: - slab_live.append(allocated) - requested += size - allocated_bytes += len(allocated) * slab.block_bytes + extent = FreeExtentAllocator(capacity) + legacy = _LegacyBuddyAllocator(capacity) + page_runs = FixedBlockAllocator(capacity) + legacy_full_bytes = sum(legacy.allocation_bytes(size) for size in sizes) + page_run_full_bytes = sum(page_runs.allocation_bytes(size) for size in sizes) return { "operations": 20_000, "capacity_bytes": capacity, - "extent_allocation_failures": extent_failures, - "extent_fragmentation_failures": extent_fragmentation_failures, - "slab_allocation_failures": slab_failures, - "slab_fragmentation_failures": slab_fragmentation_failures, - "slab_internal_waste_ratio": (allocated_bytes - requested) / allocated_bytes, + "exact_byte_extents": run_trace(extent), + "legacy_power_of_two_buddy": { + "block_bytes": legacy.block_bytes, + "full_corpus_allocated_bytes": legacy_full_bytes, + **run_trace(legacy), + }, + "segregated_page_runs": { + "block_bytes": page_runs.block_bytes, + "full_corpus_allocated_bytes": page_run_full_bytes, + "full_corpus_space_saved_bytes": legacy_full_bytes - page_run_full_bytes, + **run_trace(page_runs), + }, } @@ -274,7 +380,9 @@ def access(self, key: str, size: int) -> bool: def policy_ablation(records: list[dict[str, object]]) -> dict[str, object]: rng = random.Random(77) universe = records[:5000] - sizes = {str(record["tensor_ref"]): int(record["canonical_bytes"]) for record in universe} + sizes = { + str(record["tensor_ref"]): int(record["canonical_bytes"]) for record in universe + } hot = list(sizes)[:100] cold = list(sizes)[100:] # Warm every hot object before injecting scans so TinyLFU admission is @@ -330,9 +438,16 @@ def encode_frame( ) body.extend(tensor_ref) body.extend(checksum) - return protocol.HEADER.pack( - protocol.MAGIC, protocol.EXTERNAL_VERSION, protocol.REQUEST_KIND, request_id, len(body) - ) + body + return ( + protocol.HEADER.pack( + protocol.MAGIC, + protocol.EXTERNAL_VERSION, + protocol.REQUEST_KIND, + request_id, + len(body), + ) + + body + ) def request_round_trip( @@ -361,10 +476,13 @@ def daemon_ablation( rust_binary: Path, device: int, repeats: int, + gpu_block_kib: int = 32, ) -> dict[str, object]: rng = np.random.default_rng(31) query = rng.standard_normal((44, int(records[0]["tensor_dim"]))).astype(" dict[str, object]: shard_paths = sorted((shard_root / "shards").glob("*.vts")) results = {} @@ -488,6 +611,8 @@ def prewarm_ablation( f"{device}=20", "--gpu-workspace-gb", "2", + "--gpu-block-kib", + str(gpu_block_kib), "--host-cache-gb", "0.1", "--contract-root", @@ -509,6 +634,8 @@ def prewarm_ablation( f"{device}=20", "--gpu-workspace-gb", "2", + "--gpu-block-kib", + str(gpu_block_kib), "--host-cache-gb", "0.1", "--contract-root", @@ -590,6 +717,7 @@ def main() -> None: parser.add_argument("--device", required=True, type=int) parser.add_argument("--sample-size", type=int, default=100) parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--gpu-block-kib", type=int, default=32) parser.add_argument("--output", required=True, type=Path) parser.add_argument("--full-prewarm", action="store_true") args = parser.parse_args() @@ -601,7 +729,9 @@ def main() -> None: "logical_bytes": sum(int(record["canonical_bytes"]) for record in records), "sample_size": len(selected), }, - "storage": storage_ablation(selected, args.contract, args.legacy_root, args.shard_root), + "storage": storage_ablation( + selected, args.contract, args.legacy_root, args.shard_root + ), "h2d": h2d_ablation(selected, args.contract, args.shard_root, args.device), "allocator": allocator_ablation(records), "policy": policy_ablation(records), @@ -612,6 +742,7 @@ def main() -> None: args.rust_binary, args.device, args.repeats, + args.gpu_block_kib, ), } if args.full_prewarm: @@ -621,6 +752,7 @@ def main() -> None: args.shard_root, args.rust_binary, args.device, + args.gpu_block_kib, ) args.output.parent.mkdir(parents=True, exist_ok=True) with args.output.open("w", encoding="utf-8") as stream: diff --git a/services/test_tilemaxsim_gpu_cache.py b/services/test_tilemaxsim_gpu_cache.py index ae228ac2..57a4f71a 100644 --- a/services/test_tilemaxsim_gpu_cache.py +++ b/services/test_tilemaxsim_gpu_cache.py @@ -70,14 +70,15 @@ def test_extent_allocator_coalesces_released_ranges(self) -> None: allocator.release(*second) self.assertEqual(allocator.extents, [(0, 4096)]) - def test_fixed_block_buddy_allocator_coalesces_slabs(self) -> None: + def test_fixed_block_allocator_uses_exact_runs_and_coalesces(self) -> None: allocator = FixedBlockAllocator(8 * 256, block_bytes=256) first = allocator.allocate(300) second = allocator.allocate(300) third = allocator.allocate(700) self.assertEqual(first, (0, 1)) self.assertEqual(second, (2, 3)) - self.assertEqual(third, (4, 5, 6, 7)) + self.assertEqual(third, (4, 5, 6)) + self.assertEqual(allocator.free_bytes, 256) assert first is not None and second is not None and third is not None allocator.release(second) allocator.release(first) @@ -85,6 +86,20 @@ def test_fixed_block_buddy_allocator_coalesces_slabs(self) -> None: allocator.release(third) self.assertEqual(allocator.largest_free_extent, 8 * 256) + def test_fixed_block_allocator_reuses_best_fit_run(self) -> None: + allocator = FixedBlockAllocator(12 * 256, block_bytes=256) + first = allocator.allocate(2 * 256) + separator = allocator.allocate(256) + second = allocator.allocate(3 * 256) + tail = allocator.allocate(6 * 256) + assert first is not None and separator is not None + assert second is not None and tail is not None + allocator.release(first) + allocator.release(second) + reused = allocator.allocate(3 * 256) + self.assertEqual(reused, second) + self.assertEqual(allocator.largest_free_extent, 2 * 256) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") def test_pool_rejects_budget_larger_than_currently_free_memory(self) -> None: device = available_device() @@ -182,9 +197,7 @@ def test_resident_engine_scores_gpu_cache_handles(self) -> None: @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") def test_batch_admission_uses_one_h2d_batch(self) -> None: device = available_device() - pool = GpuResourcePool( - [GpuArenaSpec(device, 4 * 1024 * 1024)], 2 * 1024 * 1024 - ) + pool = GpuResourcePool([GpuArenaSpec(device, 4 * 1024 * 1024)], 2 * 1024 * 1024) try: cache = GpuTensorCache(pool, allow_eviction=True) first = np.ones((128, 320), dtype=" None: batch = cache.acquire_many( [ GpuTensorLoad( - ("model", "first"), 128, 320, protocol.DTYPE_F16, first.tobytes() + ("model", "first"), + 128, + 320, + protocol.DTYPE_F16, + first.tobytes(), ), GpuTensorLoad( - ("model", "second"), 128, 320, protocol.DTYPE_F16, second.tobytes() + ("model", "second"), + 128, + 320, + protocol.DTYPE_F16, + second.tobytes(), ), ] ) @@ -215,7 +236,9 @@ def test_batch_admission_uses_one_h2d_batch(self) -> None: def test_tinylfu_rejects_one_off_tensor_instead_of_polluting_hot_slab(self) -> None: device = available_device() pool = GpuResourcePool( - [GpuArenaSpec(device, 768 * 1024)], 512 * 1024 + [GpuArenaSpec(device, 768 * 1024)], + 512 * 1024, + block_bytes=256 * 1024, ) try: cache = GpuTensorCache(pool, allow_eviction=True) diff --git a/services/test_tilemaxsim_rust_daemon.py b/services/test_tilemaxsim_rust_daemon.py index 630ff651..0999e41f 100644 --- a/services/test_tilemaxsim_rust_daemon.py +++ b/services/test_tilemaxsim_rust_daemon.py @@ -222,7 +222,7 @@ def test_one_request_larger_than_gpu_cache_is_scored_in_chunks(self) -> None: [device], [first, second], query, - gpu_memory_gb="0.0012", + gpu_memory_gb="0.0010", workspace_gb="0.0005", ) events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] diff --git a/services/tilemaxsim_cuda_sidecar.py b/services/tilemaxsim_cuda_sidecar.py index 98f7d0b4..fc256fcd 100644 --- a/services/tilemaxsim_cuda_sidecar.py +++ b/services/tilemaxsim_cuda_sidecar.py @@ -107,9 +107,9 @@ class PayloadCache: def __init__(self, maximum_bytes: int) -> None: self.maximum_bytes = maximum_bytes self.current_bytes = 0 - self.entries: OrderedDict[ - tuple[object, ...], tuple[bytes, float] - ] = OrderedDict() + self.entries: OrderedDict[tuple[object, ...], tuple[bytes, float]] = ( + OrderedDict() + ) self.lock = threading.Lock() self.sketch = TinyLfuSketch() self.inflation = 0.0 @@ -370,9 +370,7 @@ def key(self, request: protocol.ExternalTensorRequest) -> tuple[object, ...]: def _validate_shard_entry( request: protocol.ExternalTensorRequest, digest: str, entry: ShardEntry ) -> None: - dtype_name = ( - "float32" if request.dtype == protocol.DTYPE_F32 else "float16" - ) + dtype_name = "float32" if request.dtype == protocol.DTYPE_F32 else "float16" if ( entry.digest != digest or entry.rows != request.rows @@ -402,7 +400,9 @@ def _verify_shard(root: _OpenShardRoot, name: str) -> None: offset = 0 descriptor = root.shard_fds[name] while offset < shard.size: - chunk = os.pread(descriptor, min(8 * 1024 * 1024, shard.size - offset), offset) + chunk = os.pread( + descriptor, min(8 * 1024 * 1024, shard.size - offset), offset + ) if not chunk: raise protocol.SidecarError( protocol.STATUS_COMPUTE_ERROR, @@ -470,7 +470,9 @@ def _read_shard_range( protocol.STATUS_INVALID_REQUEST, "resolved shard tensor checksum mismatch", ) - dtype = protocol.DTYPE_F32 if entry.dtype == "float32" else protocol.DTYPE_F16 + dtype = ( + protocol.DTYPE_F32 if entry.dtype == "float32" else protocol.DTYPE_F16 + ) validate_finite_payload(tensor, entry.rows, entry.dimension, dtype) result[key] = tensor return result @@ -493,7 +495,9 @@ def resolve_many( missing[key] = request hits[key] = False - shard_groups: dict[tuple[str, str], list[tuple[tuple[object, ...], ShardEntry]]] = {} + shard_groups: dict[ + tuple[str, str], list[tuple[tuple[object, ...], ShardEntry]] + ] = {} legacy: list[tuple[tuple[object, ...], protocol.ExternalTensorRequest]] = [] for key, request in missing.items(): contract = request.model_contract_id @@ -511,7 +515,15 @@ def resolve_many( self._validate_shard_entry(request, digest, entry) shard_groups.setdefault((contract, entry.shard), []).append((key, entry)) - jobs: list[tuple[_OpenShardRoot, str, int, int, list[tuple[tuple[object, ...], ShardEntry]]]] = [] + jobs: list[ + tuple[ + _OpenShardRoot, + str, + int, + int, + list[tuple[tuple[object, ...], ShardEntry]], + ] + ] = [] for (contract, shard_name), entries in shard_groups.items(): root = self.shard_roots[contract] if self.verify_full_shards: @@ -520,11 +532,13 @@ def resolve_many( jobs.append((root, shard_name, start, end, grouped)) if jobs: with ThreadPoolExecutor(max_workers=min(8, len(jobs))) as workers: - for resolved in workers.map(lambda job: self._read_shard_range(*job), jobs): + for resolved in workers.map( + lambda job: self._read_shard_range(*job), jobs + ): payloads.update(resolved) def read_legacy( - item: tuple[tuple[object, ...], protocol.ExternalTensorRequest] + item: tuple[tuple[object, ...], protocol.ExternalTensorRequest], ) -> tuple[tuple[object, ...], bytes]: key, request = item digest = str(key[1]) @@ -1511,16 +1525,16 @@ def flush() -> None: record_access=False, count_stats=False, ) - if batch.bypassed or batch.deferred or any( - handle is None for handle in batch.handles + if ( + batch.bypassed + or batch.deferred + or any(handle is None for handle in batch.handles) ): raise protocol.SidecarError( protocol.STATUS_RESOURCE_LIMIT, "resident manifest exceeds the configured GPU block arenas", ) - acquired.extend( - handle for handle in batch.handles if handle is not None - ) + acquired.extend(handle for handle in batch.handles if handle is not None) finally: for handle in acquired: gpu_cache.release(handle) @@ -1614,6 +1628,12 @@ def main() -> None: default=2 * 1024**3, help="per-GPU portion of the configured GB reserved for scoring temporaries", ) + parser.add_argument( + "--gpu-block-kib", + type=positive_int, + default=32, + help="base page size inside each preallocated GPU tensor arena", + ) parser.add_argument("--contract-root", action="append", default=[]) parser.add_argument( "--resident-manifest", @@ -1673,7 +1693,11 @@ def main() -> None: metrics = JsonMetrics() pool: GpuResourcePool | None = None try: - pool = GpuResourcePool(args.gpu_memory_gb, args.gpu_workspace_gb) + pool = GpuResourcePool( + args.gpu_memory_gb, + args.gpu_workspace_gb, + args.gpu_block_kib * 1024, + ) engine = TorchTileMaxsimEngine( str(pool.primary_device), args.gpu_workspace_gb, diff --git a/services/tilemaxsim_gpu_cache.py b/services/tilemaxsim_gpu_cache.py index 65a90271..78f507be 100644 --- a/services/tilemaxsim_gpu_cache.py +++ b/services/tilemaxsim_gpu_cache.py @@ -18,6 +18,7 @@ import re import threading +from bisect import bisect_left, insort from collections import OrderedDict from dataclasses import dataclass from decimal import Decimal, InvalidOperation @@ -34,7 +35,11 @@ _GPU_MEMORY_GB = re.compile(r"^(?:cuda:)?([0-9]+)=([0-9]+(?:\.[0-9]+)?)$") _MEMORY_GB = re.compile(r"^[0-9]+(?:\.[0-9]+)?$") GIB = 1024**3 -DEFAULT_BLOCK_BYTES = 256 * 1024 +# A 32 KiB page keeps the allocator metadata small while avoiding the 8.82% +# rounding loss observed with the former 256 KiB default on 472--483 KiB +# document tensors. The native daemon exposes the same default via +# ``--gpu-block-kib``. +DEFAULT_BLOCK_BYTES = 32 * 1024 DEFAULT_STAGING_BYTES = 64 * 1024 * 1024 @@ -131,11 +136,12 @@ def release(self, start: int, length: int) -> None: class FixedBlockAllocator: - """Fixed-base-block buddy/slab allocator used by resident tensors. + """Best-fit page-run allocator over one process-owned GPU arena. - A tensor gets one contiguous power-of-two slab. This preserves the dense - memory layout required by the Tensor Core kernel while bounding internal - fragmentation and guaranteeing that released buddies coalesce. + Tensors keep the dense layout required by the TileMaxSim kernel, but their + runs are rounded only to the base page instead of the next power of two. + Free runs are indexed by both address and exact size. Allocation therefore + finds the smallest suitable run, while release coalesces adjacent runs. """ def __init__(self, capacity: int, block_bytes: int = DEFAULT_BLOCK_BYTES) -> None: @@ -148,20 +154,50 @@ def __init__(self, capacity: int, block_bytes: int = DEFAULT_BLOCK_BYTES) -> Non if self.block_count == 0: raise ValueError("arena must contain at least one GPU block") self.capacity = self.block_count * block_bytes - self._free: dict[int, set[tuple[int, int, int]]] = {} - self._allocated: dict[int, tuple[int, int, int]] = {} - start = 0 - remaining = self.block_count - while remaining: - order = remaining.bit_length() - 1 - size = 1 << order - self._free.setdefault(order, set()).add((start, start, order)) - start += size - remaining -= size + self._free_by_start: dict[int, int] = {} + self._free_by_size: dict[int, set[int]] = {} + self._starts: list[int] = [] + self._sizes: list[int] = [] + self._allocated: dict[int, int] = {} + self._free_blocks = 0 + self._add_free_run(0, self.block_count) + + def _add_free_run(self, start: int, blocks: int) -> None: + if blocks <= 0 or start < 0 or start + blocks > self.block_count: + raise ValueError("free GPU run is outside the arena") + if start in self._free_by_start: + raise ValueError("duplicate free GPU run") + self._free_by_start[start] = blocks + insort(self._starts, start) + bucket = self._free_by_size.get(blocks) + if bucket is None: + bucket = set() + self._free_by_size[blocks] = bucket + insort(self._sizes, blocks) + bucket.add(start) + self._free_blocks += blocks + + def _remove_free_run(self, start: int, blocks: int) -> None: + if self._free_by_start.get(start) != blocks: + raise ValueError("free GPU run directory is inconsistent") + del self._free_by_start[start] + start_index = bisect_left(self._starts, start) + if start_index == len(self._starts) or self._starts[start_index] != start: + raise ValueError("free GPU address index is inconsistent") + self._starts.pop(start_index) + bucket = self._free_by_size[blocks] + bucket.remove(start) + if not bucket: + del self._free_by_size[blocks] + size_index = bisect_left(self._sizes, blocks) + if size_index == len(self._sizes) or self._sizes[size_index] != blocks: + raise ValueError("free GPU size index is inconsistent") + self._sizes.pop(size_index) + self._free_blocks -= blocks @property def free_blocks(self) -> int: - return sum(len(items) * (1 << order) for order, items in self._free.items()) + return self._free_blocks @property def free_bytes(self) -> int: @@ -169,44 +205,28 @@ def free_bytes(self) -> int: @property def largest_free_extent(self) -> int: - return max( - (1 << order) * self.block_bytes - for order, items in self._free.items() - if items - ) if self.free_blocks else 0 + return self._sizes[-1] * self.block_bytes if self._sizes else 0 def blocks_for(self, payload_bytes: int) -> int: if payload_bytes <= 0: raise ValueError("payload size must be positive") - raw = ceil(payload_bytes / self.block_bytes) - return 1 << (raw - 1).bit_length() + return ceil(payload_bytes / self.block_bytes) def allocation_bytes(self, payload_bytes: int) -> int: return self.blocks_for(payload_bytes) * self.block_bytes def allocate(self, payload_bytes: int) -> tuple[int, ...] | None: required = self.blocks_for(payload_bytes) - order = required.bit_length() - 1 - available_order = next( - ( - candidate - for candidate in range(order, self.block_count.bit_length()) - if self._free.get(candidate) - ), - None, - ) - if available_order is None: + size_index = bisect_left(self._sizes, required) + if size_index == len(self._sizes): return None - start, root_start, root_order = self._free[available_order].pop() - while available_order > order: - available_order -= 1 - buddy = start + (1 << available_order) - self._free.setdefault(available_order, set()).add( - (buddy, root_start, root_order) - ) - blocks = tuple(range(start, start + required)) - self._allocated[start] = (order, root_start, root_order) - return blocks + available = self._sizes[size_index] + start = min(self._free_by_size[available]) + self._remove_free_run(start, available) + if available > required: + self._add_free_run(start + required, available - required) + self._allocated[start] = required + return tuple(range(start, start + required)) def release(self, blocks: tuple[int, ...]) -> None: if ( @@ -215,23 +235,27 @@ def release(self, blocks: tuple[int, ...]) -> None: or blocks != tuple(range(blocks[0], blocks[0] + len(blocks))) ): raise ValueError("released GPU blocks are invalid") - allocation = self._allocated.pop(blocks[0], None) - if allocation is None: + allocated_blocks = self._allocated.pop(blocks[0], None) + if allocated_blocks is None: raise ValueError("GPU block was released more than once") - order, root_start, root_order = allocation - if len(blocks) != 1 << order: - raise ValueError("released GPU slab has the wrong size") + if len(blocks) != allocated_blocks: + raise ValueError("released GPU page run has the wrong size") start = blocks[0] - while order < root_order: - buddy = root_start + (((start - root_start) ^ (1 << order))) - item = (buddy, root_start, root_order) - free = self._free.setdefault(order, set()) - if item not in free: - break - free.remove(item) - start = min(start, buddy) - order += 1 - self._free.setdefault(order, set()).add((start, root_start, root_order)) + length = allocated_blocks + insertion = bisect_left(self._starts, start) + if insertion: + previous_start = self._starts[insertion - 1] + previous_length = self._free_by_start[previous_start] + if previous_start + previous_length == start: + self._remove_free_run(previous_start, previous_length) + start = previous_start + length += previous_length + next_start = start + length + next_length = self._free_by_start.get(next_start) + if next_length is not None: + self._remove_free_run(next_start, next_length) + length += next_length + self._add_free_run(start, length) class TinyLfuSketch: @@ -269,8 +293,7 @@ def increment(self, key: tuple[object, ...]) -> int: def estimate(self, key: tuple[object, ...]) -> int: return min( - self.tables[row][index] - for row, index in enumerate(self._indices(key)) + self.tables[row][index] for row, index in enumerate(self._indices(key)) ) @@ -335,7 +358,9 @@ def __init__( staging_bytes = min(self.capacity, DEFAULT_STAGING_BYTES) staging_bytes = max( self.allocator.block_bytes, - staging_bytes // self.allocator.block_bytes * self.allocator.block_bytes, + staging_bytes + // self.allocator.block_bytes + * self.allocator.block_bytes, ) self.host_staging = torch.empty( staging_bytes, dtype=torch.uint8, pin_memory=True @@ -359,7 +384,11 @@ def tensor( dimension: int, dtype: int, ) -> torch.Tensor: - if self.storage is None or self.host_staging is None or self.copy_stream is None: + if ( + self.storage is None + or self.host_staging is None + or self.copy_stream is None + ): raise RuntimeError("GPU arena is closed") scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 block_bytes = self.allocator.block_bytes @@ -374,60 +403,82 @@ def tensor( ).narrow(0, 0, payload_bytes) return raw.view(scalar_dtype).reshape(rows, dimension) - def copy_many_from_host( - self, items: list[tuple[tuple[int, ...], bytes]] - ) -> None: + def copy_many_from_host(self, items: list[tuple[tuple[int, ...], bytes]]) -> None: if self.storage is None: raise RuntimeError("GPU arena is closed") if not items: return block_bytes = self.allocator.block_bytes - flattened: list[tuple[int, bytes, int, int]] = [] + runs: list[tuple[int, int, bytes]] = [] for blocks, payload in items: - for ordinal, block in enumerate(blocks): - start = ordinal * block_bytes - length = min(block_bytes, max(0, len(payload) - start)) - flattened.append((block, payload, start, length)) - flattened.sort(key=lambda item: item[0]) + if ( + not blocks + or blocks != tuple(range(blocks[0], blocks[0] + len(blocks))) + or len(payload) > len(blocks) * block_bytes + ): + raise RuntimeError("GPU upload requires one valid contiguous page run") + runs.append((blocks[0], len(blocks), payload)) + runs.sort(key=lambda item: item[0]) staging = self.host_staging staging_array = staging.numpy() - staging_blocks = staging.numel() // block_bytes + staging_bytes = staging.numel() copy_calls = 0 with torch.cuda.device(self.device): stream = self.copy_stream - for chunk_start in range(0, len(flattened), staging_blocks): - chunk = flattened[chunk_start : chunk_start + staging_blocks] - for index, (_, payload, source_start, source_length) in enumerate( - chunk - ): - if not source_length: - continue - start = index * block_bytes - staging_array[start : start + source_length] = np.frombuffer( - payload, - dtype=np.uint8, - count=source_length, - offset=source_start, - ) - with torch.cuda.stream(stream): - run_start = 0 - for index in range(1, len(chunk) + 1): - if ( - index < len(chunk) - and chunk[index][0] == chunk[index - 1][0] + 1 - ): - continue - first_block = chunk[run_start][0] - count = index - run_start - length = count * block_bytes - self.storage.narrow( - 0, first_block * block_bytes, length - ).copy_( - staging.narrow(0, run_start * block_bytes, length), - non_blocking=True, + run_index = 0 + while run_index < len(runs): + first_block, run_blocks, payload = runs[run_index] + allocation_bytes = run_blocks * block_bytes + if allocation_bytes > staging_bytes: + source_offset = 0 + while source_offset < len(payload): + length = min(staging_bytes, len(payload) - source_offset) + staging_array[:length] = np.frombuffer( + payload, + dtype=np.uint8, + count=length, + offset=source_offset, ) + with torch.cuda.stream(stream): + self.storage.narrow( + 0, + first_block * block_bytes + source_offset, + length, + ).copy_(staging.narrow(0, 0, length), non_blocking=True) + stream.synchronize() copy_calls += 1 - run_start = index + source_offset += length + run_index += 1 + continue + + batch_first_block = first_block + expected_block = first_block + staging_offset = 0 + while run_index < len(runs): + start_block, count, current_payload = runs[run_index] + current_allocation = count * block_bytes + if ( + start_block != expected_block + or staging_offset + current_allocation > staging_bytes + ): + break + payload_end = staging_offset + len(current_payload) + staging_array[staging_offset:payload_end] = np.frombuffer( + current_payload, dtype=np.uint8 + ) + allocation_end = staging_offset + current_allocation + staging_array[payload_end:allocation_end] = 0 + staging_offset = allocation_end + expected_block += count + run_index += 1 + with torch.cuda.stream(stream): + self.storage.narrow( + 0, batch_first_block * block_bytes, staging_offset + ).copy_( + staging.narrow(0, 0, staging_offset), + non_blocking=True, + ) + copy_calls += 1 stream.synchronize() self.h2d_batches += 1 self.h2d_copy_calls += copy_calls @@ -793,8 +844,8 @@ def acquire_many( if cached is not None: cached.references += 1 cached.pinned = cached.pinned or load.pin - cached.priority = ( - self.inflation + frequency / self._entry_cost(cached) + cached.priority = self.inflation + frequency / self._entry_cost( + cached ) self.entries.move_to_end(load.key) handles[index] = GpuTensorHandle(cached) @@ -835,7 +886,9 @@ def acquire_many( new_entries.append((index, entry, load.payload)) admitted += 1 - by_arena: dict[int, tuple[GpuArena, list[tuple[tuple[int, ...], bytes]]]] = {} + by_arena: dict[ + int, tuple[GpuArena, list[tuple[tuple[int, ...], bytes]]] + ] = {} for _, entry, payload in new_entries: bucket = by_arena.setdefault(id(entry.arena), (entry.arena, [])) bucket[1].append((entry.blocks, payload)) @@ -897,7 +950,12 @@ def release(self, handle: GpuTensorHandle) -> None: def status(self) -> dict[str, object]: with self.lock: + allocated_bytes = sum( + entry.allocated_bytes for entry in self.entries.values() + ) + payload_bytes = sum(entry.payload_bytes for entry in self.entries.values()) return { + "allocator": "segregated-page-runs", "entries": len(self.entries), "pinned_entries": sum(entry.pinned for entry in self.entries.values()), "active_references": sum( @@ -910,5 +968,8 @@ def status(self) -> dict[str, object]: "policy": "tinylfu-gdsf", "gdsf_inflation": self.inflation, "loaded_bytes": self.loaded_bytes, + "allocated_bytes": allocated_bytes, + "payload_bytes": payload_bytes, + "internal_waste_bytes": allocated_bytes - payload_bytes, "arenas": self.pool.status(), } diff --git a/services/tilemaxsimd/src/cache.rs b/services/tilemaxsimd/src/cache.rs index 06ae3fa7..bcfcf57a 100644 --- a/services/tilemaxsimd/src/cache.rs +++ b/services/tilemaxsimd/src/cache.rs @@ -2,14 +2,16 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::hash::{Hash, Hasher}; #[derive(Debug)] -pub struct BuddyAllocator { +pub struct PageRunAllocator { block_bytes: usize, block_count: usize, - free: BTreeMap>, - allocated: HashMap, + free_by_start: BTreeMap, + free_by_size: BTreeMap>, + allocated: HashMap, + free_blocks: usize, } -impl BuddyAllocator { +impl PageRunAllocator { pub fn new(capacity: usize, block_bytes: usize) -> Result { if capacity == 0 || block_bytes == 0 || !block_bytes.is_multiple_of(256) { return Err("invalid fixed-block arena"); @@ -18,22 +20,59 @@ impl BuddyAllocator { if block_count == 0 { return Err("fixed-block arena is too small"); } - let mut free = BTreeMap::>::new(); - let mut start = 0; - let mut remaining = block_count; - while remaining != 0 { - let order = usize::BITS - 1 - remaining.leading_zeros(); - let size = 1_usize << order; - free.entry(order).or_default().insert((start, start, order)); - start += size; - remaining -= size; - } - Ok(Self { + let mut allocator = Self { block_bytes, block_count, - free, + free_by_start: BTreeMap::new(), + free_by_size: BTreeMap::new(), allocated: HashMap::new(), - }) + free_blocks: 0, + }; + allocator.add_free_run(0, block_count)?; + Ok(allocator) + } + + fn add_free_run(&mut self, start: usize, blocks: usize) -> Result<(), &'static str> { + if blocks == 0 + || start + .checked_add(blocks) + .is_none_or(|end| end > self.block_count) + { + return Err("free GPU page run is outside the arena"); + } + if self.free_by_start.insert(start, blocks).is_some() { + return Err("duplicate free GPU page run"); + } + self.free_by_size.entry(blocks).or_default().insert(start); + self.free_blocks = self + .free_blocks + .checked_add(blocks) + .ok_or("free GPU page count overflow")?; + Ok(()) + } + + fn remove_free_run(&mut self, start: usize, blocks: usize) -> Result<(), &'static str> { + if self.free_by_start.remove(&start) != Some(blocks) { + return Err("free GPU address index is inconsistent"); + } + let remove_size = { + let starts = self + .free_by_size + .get_mut(&blocks) + .ok_or("free GPU size index is inconsistent")?; + if !starts.remove(&start) { + return Err("free GPU size index is inconsistent"); + } + starts.is_empty() + }; + if remove_size { + self.free_by_size.remove(&blocks); + } + self.free_blocks = self + .free_blocks + .checked_sub(blocks) + .ok_or("free GPU page count underflow")?; + Ok(()) } pub fn capacity(&self) -> usize { @@ -48,41 +87,28 @@ impl BuddyAllocator { if payload_bytes == 0 { return None; } - let raw = payload_bytes.div_ceil(self.block_bytes); - raw.checked_next_power_of_two()? + payload_bytes + .div_ceil(self.block_bytes) .checked_mul(self.block_bytes) } pub fn largest_free(&self) -> usize { - self.free - .iter() - .rev() - .find(|(_, entries)| !entries.is_empty()) - .map_or(0, |(order, _)| (1_usize << order) * self.block_bytes) + self.free_by_size + .last_key_value() + .map_or(0, |(blocks, _)| blocks * self.block_bytes) } pub fn allocate(&mut self, payload_bytes: usize) -> Option<(usize, usize)> { let allocation_bytes = self.allocation_bytes(payload_bytes)?; - let blocks = allocation_bytes / self.block_bytes; - let order = blocks.trailing_zeros(); - let available_order = self - .free - .range(order..) - .find(|(_, entries)| !entries.is_empty()) - .map(|(candidate, _)| *candidate)?; - let item = self.free.get_mut(&available_order)?.pop_first()?; - let (start, root_start, root_order) = item; - let mut current_order = available_order; - while current_order > order { - current_order -= 1; - let buddy = start + (1_usize << current_order); - self.free - .entry(current_order) - .or_default() - .insert((buddy, root_start, root_order)); + let required = allocation_bytes / self.block_bytes; + let (&available, starts) = self.free_by_size.range(required..).next()?; + let start = *starts.first()?; + self.remove_free_run(start, available).ok()?; + if available > required { + self.add_free_run(start + required, available - required) + .ok()?; } - self.allocated - .insert(start, (order, root_start, root_order)); + self.allocated.insert(start, required); Some((start * self.block_bytes, allocation_bytes)) } @@ -91,24 +117,61 @@ impl BuddyAllocator { return Err("unaligned fixed-block release"); } let mut start = offset / self.block_bytes; - let (mut order, root_start, root_order) = self + let mut blocks = self .allocated .remove(&start) - .ok_or("fixed-block slab was released twice")?; - while order < root_order { - let buddy = root_start + ((start - root_start) ^ (1_usize << order)); - let buddy_item = (buddy, root_start, root_order); - let entries = self.free.entry(order).or_default(); - if !entries.remove(&buddy_item) { - break; + .ok_or("fixed-block page run was released twice")?; + let previous = self + .free_by_start + .range(..start) + .next_back() + .map(|(&previous_start, &previous_blocks)| (previous_start, previous_blocks)); + if let Some((previous_start, previous_blocks)) = previous + && previous_start + previous_blocks == start + { + self.remove_free_run(previous_start, previous_blocks)?; + start = previous_start; + blocks += previous_blocks; + } + let next_start = start + blocks; + if let Some(&next_blocks) = self.free_by_start.get(&next_start) { + self.remove_free_run(next_start, next_blocks)?; + blocks += next_blocks; + } + self.add_free_run(start, blocks) + } + + pub fn free_bytes(&self) -> usize { + self.free_blocks * self.block_bytes + } + + pub fn allocated_bytes(&self) -> usize { + self.allocated + .values() + .map(|blocks| blocks * self.block_bytes) + .sum() + } + + #[cfg(test)] + pub fn validate(&self) -> Result<(), &'static str> { + let address_blocks = self.free_by_start.values().sum::(); + let size_blocks = self + .free_by_size + .iter() + .map(|(blocks, starts)| blocks * starts.len()) + .sum::(); + if address_blocks != self.free_blocks || size_blocks != self.free_blocks { + return Err("free GPU page accounting is inconsistent"); + } + for (&start, &blocks) in &self.free_by_start { + if !self + .free_by_size + .get(&blocks) + .is_some_and(|starts| starts.contains(&start)) + { + return Err("free GPU page indexes disagree"); } - start = start.min(buddy); - order += 1; } - self.free - .entry(order) - .or_default() - .insert((start, root_start, root_order)); Ok(()) } } @@ -187,7 +250,7 @@ pub enum Admission { } pub struct GpuCache { - allocator: BuddyAllocator, + allocator: PageRunAllocator, entries: HashMap, sketch: TinyLfu, inflation: f64, @@ -200,7 +263,7 @@ pub struct GpuCache { impl GpuCache { pub fn new(capacity: usize, block_bytes: usize) -> Result { Ok(Self { - allocator: BuddyAllocator::new(capacity, block_bytes)?, + allocator: PageRunAllocator::new(capacity, block_bytes)?, entries: HashMap::new(), sketch: TinyLfu::new(4096), inflation: 0.0, @@ -227,6 +290,22 @@ impl GpuCache { self.entries.values().filter(|entry| entry.pinned).count() } + pub fn free_bytes(&self) -> usize { + self.allocator.free_bytes() + } + + pub fn largest_free_extent(&self) -> usize { + self.allocator.largest_free() + } + + pub fn allocated_bytes(&self) -> usize { + self.allocator.allocated_bytes() + } + + pub fn payload_bytes(&self) -> usize { + self.entries.values().map(|entry| entry.payload_bytes).sum() + } + pub fn get(&mut self, key: &str) -> Option { let frequency = self.sketch.increment(&key); let Some(entry) = self.entries.get_mut(key) else { @@ -360,20 +439,56 @@ impl GpuCache { mod tests { use super::*; + fn next_test_random(state: &mut u64) -> u64 { + // Deterministic xorshift64 sequence for allocator churn tests. This is + // test data, not a pointer, address, secret, or production RNG. + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + *state + } + #[test] - fn buddy_coalesces_fixed_slabs() { - let mut allocator = BuddyAllocator::new(8 * 256, 256).unwrap(); + fn page_runs_use_exact_blocks_and_coalesce() { + let mut allocator = PageRunAllocator::new(8 * 256, 256).unwrap(); let first = allocator.allocate(300).unwrap(); let second = allocator.allocate(300).unwrap(); let third = allocator.allocate(700).unwrap(); assert_eq!(first, (0, 512)); assert_eq!(second, (512, 512)); - assert_eq!(third, (1024, 1024)); + assert_eq!(third, (1024, 768)); + assert_eq!(allocator.free_bytes(), 256); allocator.release(second.0).unwrap(); allocator.release(first.0).unwrap(); assert_eq!(allocator.largest_free(), 1024); allocator.release(third.0).unwrap(); assert_eq!(allocator.largest_free(), 2048); + allocator.validate().unwrap(); + } + + #[test] + fn page_run_indexes_stay_consistent_under_churn() { + let mut allocator = PageRunAllocator::new(64 * 256, 256).unwrap(); + let mut live = Vec::new(); + let mut state = 0x5eed_u64; + for _ in 0..20_000 { + let random = next_test_random(&mut state); + if !live.is_empty() && random & 3 == 0 { + let index = random as usize % live.len(); + let (offset, _) = live.swap_remove(index); + allocator.release(offset).unwrap(); + } else { + let payload = ((random >> 16) as usize % (12 * 256)) + 1; + if let Some(allocation) = allocator.allocate(payload) { + live.push(allocation); + } + } + allocator.validate().unwrap(); + assert_eq!( + allocator.free_bytes() + allocator.allocated_bytes(), + allocator.capacity() + ); + } } #[test] diff --git a/services/tilemaxsimd/src/engine.rs b/services/tilemaxsimd/src/engine.rs index fd5c1adf..7668d67d 100644 --- a/services/tilemaxsimd/src/engine.rs +++ b/services/tilemaxsimd/src/engine.rs @@ -394,8 +394,15 @@ impl Engine { .map(|(index, device)| { serde_json::json!({ "index": index, + "gpu_allocator": "segregated-page-runs", "gpu_tensor_bytes": device.cache.capacity(), "gpu_block_bytes": device.cache.block_bytes(), + "gpu_free_bytes": device.cache.free_bytes(), + "gpu_largest_free_extent_bytes": device.cache.largest_free_extent(), + "gpu_allocated_bytes": device.cache.allocated_bytes(), + "gpu_payload_bytes": device.cache.payload_bytes(), + "gpu_internal_waste_bytes": device.cache.allocated_bytes() + - device.cache.payload_bytes(), "gpu_entries": device.cache.entry_count(), "gpu_pinned_entries": device.cache.pinned_entries(), "gpu_hits": device.cache.hits, diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs index 847dc497..b940a08b 100644 --- a/services/tilemaxsimd/src/main.rs +++ b/services/tilemaxsimd/src/main.rs @@ -33,7 +33,7 @@ struct Args { host_cache_gb: usize, #[arg(long = "contract-root", required = true, value_parser = parse_contract_root)] contract_roots: Vec<(String, PathBuf)>, - #[arg(long, default_value_t = 256)] + #[arg(long, default_value_t = 32)] gpu_block_kib: usize, #[arg(long, default_value_t = 64 * 1024 * 1024)] max_request_bytes: usize, From f7af4b262236d2f0cb33d57635200f4858a7f990 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:29:22 +0800 Subject: [PATCH 306/324] feat(tilemaxsim): harden native GPU service for production Add PostgreSQL-style daemon/control CLIs, bounded request scheduling, lifecycle and health controls, deployment templates, and CI coverage. Correct fork-specific copyright and release metadata to HuXinjing while preserving upstream notices on inherited source. --- .dockerignore | 10 + .github/workflows/check.yml | 143 +++- .github/workflows/release.yml | 5 +- META.json.in | 11 +- README.md | 8 + crates/vchordrq/src/maxsim_cost.rs | 14 +- crates/vchordrq/src/statistics.rs | 14 +- devtools/test_tilemaxsim_reference_sidecar.py | 14 +- devtools/tilemaxsim_reference_sidecar.py | 14 +- services/Dockerfile.tilemaxsimd | 20 +- services/benchmark_tilemaxsim_ablation.py | 10 +- services/benchmark_tilemaxsim_cuda.py | 14 +- services/build_tilemaxsim_tensor_cache.py | 14 +- .../packaging/tilemaxsimd-kubernetes.yaml | 122 +++ services/packaging/tilemaxsimd.conf.example | 3 + services/packaging/vchord-tilemaxsimd.service | 36 + services/test_tilemaxsim_cuda_sidecar.py | 14 +- services/test_tilemaxsim_gpu_cache.py | 14 +- services/test_tilemaxsim_rust_daemon.py | 538 +++++++++++- services/test_tilemaxsim_shard.py | 10 +- services/tilemaxsim_cuda_sidecar.py | 14 +- services/tilemaxsim_gpu_cache.py | 14 +- services/tilemaxsim_shard.py | 14 +- services/tilemaxsim_triton.py | 14 +- services/tilemaxsimd/Cargo.toml | 8 + services/tilemaxsimd/README.md | 205 +++++ services/tilemaxsimd/build.rs | 3 + .../tilemaxsimd/native/tilemaxsim_cuda.cu | 2 + .../src/bin/vchord-tilemaxsimctl.rs | 338 ++++++++ services/tilemaxsimd/src/cache.rs | 2 + services/tilemaxsimd/src/engine.rs | 34 +- services/tilemaxsimd/src/gpu.rs | 2 + services/tilemaxsimd/src/lifecycle.rs | 188 +++++ services/tilemaxsimd/src/main.rs | 541 +++++++++--- services/tilemaxsimd/src/protocol.rs | 60 ++ services/tilemaxsimd/src/server.rs | 796 ++++++++++++++++++ services/tilemaxsimd/src/shard.rs | 2 + .../vchordrq/scanners/maxsim/candidate.rs | 14 +- src/index/vchordrq/scanners/maxsim/exact.rs | 10 +- .../vchordrq/scanners/maxsim/external.rs | 10 +- src/index/vchordrq/scanners/maxsim/gpu.rs | 14 +- src/index/vchordrq/scanners/maxsim/profile.rs | 14 +- src/index/vchordrq/scanners/maxsim/rerank.rs | 14 +- src/index/vchordrq/scanners/maxsim/search.rs | 19 +- 44 files changed, 2925 insertions(+), 435 deletions(-) create mode 100644 .dockerignore create mode 100644 services/packaging/tilemaxsimd-kubernetes.yaml create mode 100644 services/packaging/tilemaxsimd.conf.example create mode 100644 services/packaging/vchord-tilemaxsimd.service create mode 100644 services/tilemaxsimd/README.md create mode 100644 services/tilemaxsimd/src/bin/vchord-tilemaxsimctl.rs create mode 100644 services/tilemaxsimd/src/lifecycle.rs create mode 100644 services/tilemaxsimd/src/server.rs diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..cbb66300 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,10 @@ +.git +**/target +build +dist +**/__pycache__ +**/*.pyc +.pytest_cache +.ruff_cache +.mypy_cache +.venv diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index df084841..6c28ef71 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -38,10 +38,14 @@ jobs: run: taplo fmt --check - name: Rustfmt - run: cargo fmt --check -- --config-path /dev/null --config imports_granularity=Module + run: | + cargo fmt --check -- --config-path /dev/null --config imports_granularity=Module + cargo fmt --check --manifest-path services/tilemaxsimd/Cargo.toml - name: Deny - run: cargo deny check + run: | + cargo deny check + cargo deny --manifest-path services/tilemaxsimd/Cargo.toml check - name: License Header run: | @@ -66,15 +70,47 @@ jobs: RS_HEADER=$(echo "$HEADER" | awk '{ if ($0 ~ /^$/) print "//"; else print "// " $0 }') C_HEADER=$(echo "$HEADER" | awk '{ if ($0 ~ /^$/) print "//"; else print "// " $0 }') PY_HEADER=$(echo "$HEADER" | awk '{ if ($0 ~ /^$/) print "#"; else print "# " $0 }') + FORK_RS_HEADER="// Copyright (c) 2026 HuXinjing" + FORK_PY_HEADER="# Copyright (c) 2026 HuXinjing" - RS_FILES=$(find . -type f -name "*.rs" -not -path "./target/*") - C_FILES=$(find . -type f -name "*.c" -not -path "./target/*") - PY_FILES=$(find . -type f -name "*.py" -not -path "./target/*") + RS_FILES=$(find . -type f -name "*.rs" -not -path "*/target/*") + C_FILES=$(find . -type f -name "*.c" -not -path "*/target/*") + CU_FILES=$(find . -type f -name "*.cu" -not -path "*/target/*") + PY_FILES=$(find . -type f -name "*.py" -not -path "*/target/*") FLAG="0" + is_fork_file() { + case "$1" in + ./crates/vchordrq/src/maxsim_cost.rs | \ + ./crates/vchordrq/src/statistics.rs | \ + ./devtools/*tilemaxsim* | \ + ./services/* | \ + ./services/tilemaxsimd/* | \ + ./src/index/vchordrq/scanners/maxsim/candidate.rs | \ + ./src/index/vchordrq/scanners/maxsim/exact.rs | \ + ./src/index/vchordrq/scanners/maxsim/external.rs | \ + ./src/index/vchordrq/scanners/maxsim/gpu.rs | \ + ./src/index/vchordrq/scanners/maxsim/profile.rs | \ + ./src/index/vchordrq/scanners/maxsim/rerank.rs | \ + ./src/index/vchordrq/scanners/maxsim/search.rs) + return 0 + ;; + *) + return 1 + ;; + esac + } + for p in $RS_FILES; do - if ! head -n "$COUNT" "$p" | cmp -s - <(echo "$RS_HEADER"); then + if is_fork_file "$p"; then + EXPECTED="$FORK_RS_HEADER" + HEADER_LINES=1 + else + EXPECTED="$RS_HEADER" + HEADER_LINES="$COUNT" + fi + if ! head -n "$HEADER_LINES" "$p" | cmp -s - <(echo "$EXPECTED"); then echo "license header mismatch in file $p" FLAG="1" fi @@ -87,8 +123,22 @@ jobs: fi done + for p in $CU_FILES; do + if ! head -n 1 "$p" | cmp -s - <(echo "$FORK_RS_HEADER"); then + echo "fork copyright header mismatch in file $p" + FLAG="1" + fi + done + for p in $PY_FILES; do - if ! head -n "$COUNT" "$p" | cmp -s - <(echo "$PY_HEADER"); then + if is_fork_file "$p"; then + EXPECTED="$FORK_PY_HEADER" + HEADER_LINES=1 + else + EXPECTED="$PY_HEADER" + HEADER_LINES="$COUNT" + fi + if ! head -n "$HEADER_LINES" "$p" | cmp -s - <(echo "$EXPECTED"); then echo "license header mismatch in file $p" FLAG="1" fi @@ -102,6 +152,85 @@ jobs: run: | ! grep -P '\t' -r ./sql + tilemaxsimd: + runs-on: ubuntu-24.04 + container: + image: nvidia/cuda:12.6.3-devel-ubuntu24.04 + + env: + CARGO_TERM_COLOR: always + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + apt-get update + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + build-essential ca-certificates curl git + rm -rf /var/lib/apt/lists/* + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --profile minimal --default-toolchain stable + echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Format + working-directory: services/tilemaxsimd + run: cargo fmt --all -- --check + + - name: Unit Test + working-directory: services/tilemaxsimd + run: cargo test --locked + + - name: Clippy + working-directory: services/tilemaxsimd + run: cargo clippy --locked --all-targets -- -D warnings + + - name: Release Build and CLI Smoke Test + working-directory: services/tilemaxsimd + run: | + cargo build --release --locked + target/release/vchord-tilemaxsimd --help + target/release/vchord-tilemaxsimd --version + target/release/vchord-tilemaxsimctl --help + target/release/vchord-tilemaxsimctl status --help + + tilemaxsimd_container: + runs-on: ubuntu-24.04 + + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Build Non-root Runtime Image + run: docker build -f services/Dockerfile.tilemaxsimd -t vchord-tilemaxsimd:ci . + + - name: Inspect Image and CLIs + run: | + test "$(docker inspect --format '{{.Config.User}}' vchord-tilemaxsimd:ci)" = "65532:65532" + docker run --rm vchord-tilemaxsimd:ci --help + docker run --rm vchord-tilemaxsimd:ci --version + docker run --rm --entrypoint vchord-tilemaxsimctl vchord-tilemaxsimd:ci --help + + tilemaxsimd_gpu: + if: ${{ vars.VECTORCHORD_GPU_CI == 'true' }} + runs-on: [self-hosted, linux, x64, tilemaxsim-gpu] + + env: + VECTORCHORD_TILEMAXSIM_SOAK_SECONDS: "3600" + + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Release Build + working-directory: services/tilemaxsimd + run: cargo build --release --locked + + - name: GPU Integration and Fault Tests + run: python -m unittest services.test_tilemaxsim_rust_daemon + miri: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +miri')) || diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 2c426c2a..af7503f1 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -101,10 +101,11 @@ jobs: mkdir -p ~/debian/usr/share/doc/postgresql-${VERSION}-vchord/ echo "Format: https://www.debian.org/doc/packaging-manuals/copyright-format/1.0/ Upstream-Name: VectorChord - Source: https://github.com/tensorchord/VectorChord + Source: https://github.com/HuXinjing/VectorChord Files: * - Copyright: Copyright (c) 2025-2026 TensorChord Inc. + Copyright: 2025-2026 TensorChord Inc. + 2026 HuXinjing License: AGPL-3.0-only OR Elastic-2.0 License: AGPL-3.0-only diff --git a/META.json.in b/META.json.in index ef79c686..32c70b38 100644 --- a/META.json.in +++ b/META.json.in @@ -4,6 +4,7 @@ "description": "VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search, it supports L2 distance, inner product, cosine distance and maxsim", "version": "@DISTVERSION@", "maintainer": [ + "HuXinjing", "TensorChord support@tensorchord.ai", "usamoi usamoi@tensorchord.ai" ], @@ -24,17 +25,17 @@ } }, "resources": { - "homepage": "https://github.com/tensorchord/VectorChord", + "homepage": "https://github.com/HuXinjing/VectorChord", "bugtracker": { - "web": "https://github.com/tensorchord/VectorChord/issues" + "web": "https://github.com/HuXinjing/VectorChord/issues" }, "repository": { - "url": "https://github.com/tensorchord/VectorChord.git", - "web": "https://github.com/tensorchord/VectorChord", + "url": "https://github.com/HuXinjing/VectorChord.git", + "web": "https://github.com/HuXinjing/VectorChord", "type": "git" } }, - "generated_by": "TensorChord", + "generated_by": "HuXinjing/VectorChord fork", "meta-spec": { "version": "1.0.0", "url": "http://pgxn.org/meta/spec.txt" diff --git a/README.md b/README.md index 1d560cb0..f92e0580 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,12 @@ The benchmark driver is [`services/benchmark_tilemaxsim_ablation.py`](services/benchmark_tilemaxsim_ablation.py); run it with `--help` for the corpus, cache-root, device, and output arguments. +The native daemon's build, CLI, security, cache sizing, systemd, container, +Kubernetes, upgrade, and backup procedures are documented in +[`services/tilemaxsimd/README.md`](services/tilemaxsimd/README.md). Both +`vchord-tilemaxsimd --help` and `vchord-tilemaxsimctl --help` are complete, +versioned command references suitable for installation alongside PostgreSQL. + The implementation is currently under active development. Its SQL interfaces and deployment packaging may change before a stable release. This repository contains only the public implementation and public-facing project information; @@ -76,6 +82,8 @@ private planning and application documentation are intentionally excluded. This work is based on [supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord). The +fork-specific TileMaxSim implementation is Copyright (c) 2026 HuXinjing; +upstream notices remain only on inherited source. The original VectorChord README is retained below for upstream installation, licensing, and project information. diff --git a/crates/vchordrq/src/maxsim_cost.rs b/crates/vchordrq/src/maxsim_cost.rs index ca843d10..f373b8f3 100644 --- a/crates/vchordrq/src/maxsim_cost.rs +++ b/crates/vchordrq/src/maxsim_cost.rs @@ -1,16 +1,4 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the terms of the ELv2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 HuXinjing #[derive(Clone, Copy, Debug)] pub enum MaxsimCostBackend { diff --git a/crates/vchordrq/src/statistics.rs b/crates/vchordrq/src/statistics.rs index f651bec5..4066022f 100644 --- a/crates/vchordrq/src/statistics.rs +++ b/crates/vchordrq/src/statistics.rs @@ -1,16 +1,4 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the terms of the ELv2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 HuXinjing use crate::tuples::{MetaTuple, WithWriter}; use index::relation::{Page, RelationWrite}; diff --git a/devtools/test_tilemaxsim_reference_sidecar.py b/devtools/test_tilemaxsim_reference_sidecar.py index ee5ce666..a0842aae 100644 --- a/devtools/test_tilemaxsim_reference_sidecar.py +++ b/devtools/test_tilemaxsim_reference_sidecar.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the terms of the ELv2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing from __future__ import annotations diff --git a/devtools/tilemaxsim_reference_sidecar.py b/devtools/tilemaxsim_reference_sidecar.py index 8d9f7ee1..ede4dd6b 100644 --- a/devtools/tilemaxsim_reference_sidecar.py +++ b/devtools/tilemaxsim_reference_sidecar.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the terms of the ELv2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing """CPU reference implementation of the VectorChord TileMaxSim IPC sidecar. diff --git a/services/Dockerfile.tilemaxsimd b/services/Dockerfile.tilemaxsimd index d6dcb2f2..e93c7773 100644 --- a/services/Dockerfile.tilemaxsimd +++ b/services/Dockerfile.tilemaxsimd @@ -17,6 +17,22 @@ RUN cargo build --release --locked FROM ${CUDA_RUNTIME_IMAGE} -COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimd /usr/local/bin/tilemaxsimd +ARG TILEMAXSIM_UID=65532 +ARG TILEMAXSIM_GID=65532 -ENTRYPOINT ["/usr/local/bin/tilemaxsimd"] +RUN groupadd --gid ${TILEMAXSIM_GID} vectorchord \ + && useradd --uid ${TILEMAXSIM_UID} --gid ${TILEMAXSIM_GID} \ + --no-create-home --home-dir /nonexistent --shell /usr/sbin/nologin vectorchord \ + && install --directory --owner ${TILEMAXSIM_UID} --group ${TILEMAXSIM_GID} \ + --mode 0750 /run/vectorchord + +COPY --from=build --chown=${TILEMAXSIM_UID}:${TILEMAXSIM_GID} /build/tilemaxsimd/target/release/vchord-tilemaxsimd /usr/local/bin/vchord-tilemaxsimd +COPY --from=build --chown=${TILEMAXSIM_UID}:${TILEMAXSIM_GID} /build/tilemaxsimd/target/release/vchord-tilemaxsimctl /usr/local/bin/vchord-tilemaxsimctl + +LABEL org.opencontainers.image.source="https://github.com/HuXinjing/VectorChord" + +USER ${TILEMAXSIM_UID}:${TILEMAXSIM_GID} +WORKDIR /var/lib/vectorchord +STOPSIGNAL SIGTERM + +ENTRYPOINT ["/usr/local/bin/vchord-tilemaxsimd"] diff --git a/services/benchmark_tilemaxsim_ablation.py b/services/benchmark_tilemaxsim_ablation.py index cf3d29b0..48018546 100644 --- a/services/benchmark_tilemaxsim_ablation.py +++ b/services/benchmark_tilemaxsim_ablation.py @@ -1,12 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing """Run storage, cache, H2D, and native-daemon TileMaxSim ablations.""" diff --git a/services/benchmark_tilemaxsim_cuda.py b/services/benchmark_tilemaxsim_cuda.py index 698e4b84..92fdcc48 100644 --- a/services/benchmark_tilemaxsim_cuda.py +++ b/services/benchmark_tilemaxsim_cuda.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing """Reproducible synthetic load probe for the CUDA TileMaxSim executor.""" diff --git a/services/build_tilemaxsim_tensor_cache.py b/services/build_tilemaxsim_tensor_cache.py index e6adf7ad..d5508a4a 100644 --- a/services/build_tilemaxsim_tensor_cache.py +++ b/services/build_tilemaxsim_tensor_cache.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing """Publish NPY page tensors into the sidecar's immutable SHA-256 cache.""" diff --git a/services/packaging/tilemaxsimd-kubernetes.yaml b/services/packaging/tilemaxsimd-kubernetes.yaml new file mode 100644 index 00000000..465671fb --- /dev/null +++ b/services/packaging/tilemaxsimd-kubernetes.yaml @@ -0,0 +1,122 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: vectorchord-postgres +spec: + replicas: 1 + selector: + matchLabels: + app: vectorchord-postgres + template: + metadata: + labels: + app: vectorchord-postgres + spec: + terminationGracePeriodSeconds: 35 + securityContext: + fsGroup: 65532 + seccompProfile: + type: RuntimeDefault + containers: + - name: postgres + image: example/vectorchord-postgres:replace-me + args: + - postgres + - -c + - shared_preload_libraries=vchord + - -c + - vchordrq.maxsim_backend=gpu + - -c + - vchordrq.maxsim_gpu_endpoint=/run/vectorchord/tilemaxsim.sock + securityContext: + runAsUser: 999 + runAsGroup: 999 + allowPrivilegeEscalation: false + volumeMounts: + - name: runtime + mountPath: /run/vectorchord + - name: tensors + mountPath: /var/lib/vectorchord/tensors + readOnly: true + - name: tilemaxsimd + image: example/vectorchord-tilemaxsimd:replace-me + args: + - --socket + - /run/vectorchord/tilemaxsim.sock + - --socket-mode + - "660" + - --gpu-memory-gb + - 0=20 + - --gpu-workspace-gb + - "2" + - --host-cache-gb + - "8" + - --contract-root + - model@1=/var/lib/vectorchord/tensors/model-v1 + - --allow-peer-uid + - "999" + - --max-batch-tokens + - "250000" + - --max-batch-mb + - "256" + - --ready-file + - /run/vectorchord/tilemaxsim.ready + - --pid-file + - /run/vectorchord/tilemaxsimd.pid + resources: + limits: + nvidia.com/gpu: 1 + startupProbe: + exec: + command: + - vchord-tilemaxsimctl + - status + - --socket + - /run/vectorchord/tilemaxsim.sock + - --ready-file + - /run/vectorchord/tilemaxsim.ready + - --quiet + failureThreshold: 120 + periodSeconds: 5 + readinessProbe: + exec: + command: + - vchord-tilemaxsimctl + - status + - --socket + - /run/vectorchord/tilemaxsim.sock + - --ready-file + - /run/vectorchord/tilemaxsim.ready + - --quiet + periodSeconds: 5 + livenessProbe: + exec: + command: + - vchord-tilemaxsimctl + - status + - --socket + - /run/vectorchord/tilemaxsim.sock + - --quiet + periodSeconds: 10 + timeoutSeconds: 2 + securityContext: + runAsNonRoot: true + runAsUser: 65532 + runAsGroup: 65532 + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: + - ALL + volumeMounts: + - name: runtime + mountPath: /run/vectorchord + - name: tensors + mountPath: /var/lib/vectorchord/tensors + readOnly: true + volumes: + - name: runtime + emptyDir: {} + - name: tensors + persistentVolumeClaim: + claimName: tilemaxsim-tensors diff --git a/services/packaging/tilemaxsimd.conf.example b/services/packaging/tilemaxsimd.conf.example new file mode 100644 index 00000000..2fdb0e72 --- /dev/null +++ b/services/packaging/tilemaxsimd.conf.example @@ -0,0 +1,3 @@ +# Replace the GPU size, model contract, tensor path, and PostgreSQL UID. +# Values named GB are GiB (1024^3 bytes). Keep this assignment on one line. +TILEMAXSIMD_ARGS="--socket /run/vectorchord/tilemaxsim.sock --socket-mode 660 --gpu-memory-gb 0=20 --gpu-workspace-gb 2 --host-cache-gb 8 --contract-root model@1=/var/lib/vectorchord/tensors/model-v1 --max-inflight 8 --backlog 64 --max-queued-requests 64 --max-batch-tokens 250000 --max-batch-mb 256 --request-timeout-ms 2000 --allow-peer-uid 999 --ready-file /run/vectorchord/tilemaxsim.ready --pid-file /run/vectorchord/tilemaxsimd.pid --shutdown-grace-ms 30000" diff --git a/services/packaging/vchord-tilemaxsimd.service b/services/packaging/vchord-tilemaxsimd.service new file mode 100644 index 00000000..b2ca37ab --- /dev/null +++ b/services/packaging/vchord-tilemaxsimd.service @@ -0,0 +1,36 @@ +[Unit] +Description=VectorChord TileMaxSim GPU service +After=local-fs.target +Before=postgresql.service + +[Service] +Type=simple +User=vectorchord +Group=vectorchord +EnvironmentFile=/etc/vectorchord/tilemaxsimd.conf +RuntimeDirectory=vectorchord +RuntimeDirectoryMode=0750 +UMask=0077 +ExecStart=/usr/local/bin/vchord-tilemaxsimd $TILEMAXSIMD_ARGS +ExecStartPost=/usr/local/bin/vchord-tilemaxsimctl status --socket /run/vectorchord/tilemaxsim.sock --ready-file /run/vectorchord/tilemaxsim.ready --timeout-ms 5000 +Restart=on-failure +RestartSec=2s +TimeoutStartSec=infinity +TimeoutStopSec=35s +KillSignal=SIGTERM +LimitMEMLOCK=infinity +NoNewPrivileges=yes +CapabilityBoundingSet= +AmbientCapabilities= +PrivateTmp=yes +ProtectControlGroups=yes +ProtectHome=yes +ProtectKernelModules=yes +ProtectKernelTunables=yes +ProtectSystem=strict +RestrictAddressFamilies=AF_UNIX +RestrictSUIDSGID=yes +SystemCallArchitectures=native + +[Install] +WantedBy=multi-user.target diff --git a/services/test_tilemaxsim_cuda_sidecar.py b/services/test_tilemaxsim_cuda_sidecar.py index e015174d..aa4f0067 100644 --- a/services/test_tilemaxsim_cuda_sidecar.py +++ b/services/test_tilemaxsim_cuda_sidecar.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing from __future__ import annotations diff --git a/services/test_tilemaxsim_gpu_cache.py b/services/test_tilemaxsim_gpu_cache.py index 57a4f71a..379a7b33 100644 --- a/services/test_tilemaxsim_gpu_cache.py +++ b/services/test_tilemaxsim_gpu_cache.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing from __future__ import annotations diff --git a/services/test_tilemaxsim_rust_daemon.py b/services/test_tilemaxsim_rust_daemon.py index 0999e41f..c674eddb 100644 --- a/services/test_tilemaxsim_rust_daemon.py +++ b/services/test_tilemaxsim_rust_daemon.py @@ -1,12 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing from __future__ import annotations @@ -16,8 +8,11 @@ import socket import subprocess import tempfile +import threading import time import unittest +from collections import deque +from concurrent.futures import ThreadPoolExecutor from pathlib import Path import numpy as np @@ -30,8 +25,128 @@ ) from services.tilemaxsim_shard import ImmutableShardWriter +SOAK_SECONDS = float(os.environ.get("VECTORCHORD_TILEMAXSIM_SOAK_SECONDS", "0")) + class RustDaemonTest(unittest.TestCase): + binary = ( + Path(__file__).parent + / "tilemaxsimd" + / "target" + / "release" + / "vchord-tilemaxsimd" + ) + control_binary = binary.with_name("vchord-tilemaxsimctl") + + def require_binary(self) -> Path: + if not self.binary.exists(): + self.skipTest("release vchord-tilemaxsimd binary has not been built") + return self.binary + + @staticmethod + def receive_response(connection: socket.socket) -> bytes: + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_bytes = protocol.HEADER.unpack(header)[4] + return header + protocol.receive_exact(connection, body_bytes) + + def wait_until_ready( + self, process: subprocess.Popen[str], socket_path: Path + ) -> None: + for _ in range(1000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + self.assertIsNone(process.poll()) + self.assertTrue(socket_path.exists()) + + @staticmethod + def write_fixture(root: Path) -> tuple[Path, bytes]: + shard_root = root / "cache" + shard_root.mkdir() + documents = [ + np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" tuple[subprocess.Popen[str], Path, bytes]: + binary = self.require_binary() + shard_root, frame = self.write_fixture(root) + socket_path = root / "tilemaxsimd.sock" + command = [ + os.fspath(binary), + "--socket", + os.fspath(socket_path), + "--gpu-memory-gb", + f"{device}=0.05", + "--gpu-workspace-gb", + "0.02", + "--host-cache-gb", + "0.01", + "--contract-root", + f"model@1={shard_root}", + *extra, + ] + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + self.wait_until_ready(process, socket_path) + return process, socket_path, frame + + def request(self, socket_path: Path, frame: bytes) -> tuple[int, int, list]: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.settimeout(5) + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + return decode_response(self.receive_response(connection)) + + @staticmethod + def free_device() -> int: + return max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + def run_daemon( self, devices: list[int], @@ -41,15 +156,7 @@ def run_daemon( workspace_gb: str = "0.02", resident: bool = False, ) -> tuple[str, list[tuple[int, float]]]: - binary = ( - Path(__file__).parent - / "tilemaxsimd" - / "target" - / "release" - / "tilemaxsimd" - ) - if not binary.exists(): - self.skipTest("release tilemaxsimd binary has not been built") + binary = self.require_binary() with tempfile.TemporaryDirectory() as directory: root = Path(directory) shard_root = root / "cache" @@ -147,18 +254,11 @@ def run_daemon( text=True, ) try: - for _ in range(1000): - if socket_path.exists() or process.poll() is not None: - break - time.sleep(0.01) - self.assertIsNone(process.poll()) - self.assertTrue(socket_path.exists()) + self.wait_until_ready(process, socket_path) with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: connection.connect(os.fspath(socket_path)) connection.sendall(frame) - header = protocol.receive_exact(connection, protocol.HEADER.size) - body_bytes = protocol.HEADER.unpack(header)[4] - response = header + protocol.receive_exact(connection, body_bytes) + response = self.receive_response(connection) output, _ = process.communicate(timeout=10) self.assertEqual(process.returncode, 0, output) request_id, status, results = decode_response(response) @@ -184,11 +284,7 @@ def run_daemon( @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") def test_external_v2_shard_round_trip_matches_protocol_oracle(self) -> None: - device = max( - range(torch.cuda.device_count()), - key=lambda index: torch.cuda.mem_get_info(index)[0], - ) - self.run_daemon([device]) + self.run_daemon([self.free_device()]) @unittest.skipUnless( torch.cuda.is_available() and torch.cuda.device_count() >= 2, @@ -206,10 +302,7 @@ def test_multi_gpu_scheduler_uploads_and_scores_on_each_device(self) -> None: @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") def test_one_request_larger_than_gpu_cache_is_scored_in_chunks(self) -> None: - device = max( - range(torch.cuda.device_count()), - key=lambda index: torch.cuda.mem_get_info(index)[0], - ) + device = self.free_device() rows, dimension = 480, 320 first = np.zeros((rows, dimension), dtype=" None: @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") def test_resident_manifest_is_pinned_before_socket_ready(self) -> None: - device = max( - range(torch.cuda.device_count()), - key=lambda index: torch.cuda.mem_get_info(index)[0], - ) + device = self.free_device() output, _ = self.run_daemon([device], resident=True) events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] prewarm = next( @@ -249,6 +339,374 @@ def test_resident_manifest_is_pinned_before_socket_ready(self) -> None: self.assertEqual(prewarm["entries"], 2) self.assertEqual(prewarm["cache"]["devices"][0]["gpu_pinned_entries"], 2) + def test_cli_help_version_and_invalid_values(self) -> None: + binary = self.require_binary() + help_result = subprocess.run( + [os.fspath(binary), "--help"], + check=False, + capture_output=True, + text=True, + ) + self.assertEqual(help_result.returncode, 0, help_result.stderr) + for expected in ( + "vchord-tilemaxsimd", + "--gpu-memory-gb ", + "--request-timeout-ms ", + "--allow-peer-uid ", + "EXAMPLES:", + ): + self.assertIn(expected, help_result.stdout) + version_result = subprocess.run( + [os.fspath(binary), "--version"], + check=False, + capture_output=True, + text=True, + ) + self.assertEqual(version_result.returncode, 0, version_result.stderr) + self.assertRegex(version_result.stdout, r"^vchord-tilemaxsimd \d+\.\d+\.\d+\n$") + invalid = subprocess.run( + [ + os.fspath(binary), + "--socket", + "/tmp/vchord-invalid.sock", + "--gpu-memory-gb", + "0=1", + "--contract-root", + "model@1=/tmp", + "--gpu-block-kib", + "3", + ], + check=False, + capture_output=True, + text=True, + ) + self.assertEqual(invalid.returncode, 2) + self.assertIn("power of two", invalid.stderr) + self.assertTrue(self.control_binary.exists()) + control_help = subprocess.run( + [os.fspath(self.control_binary), "--help"], + check=False, + capture_output=True, + text=True, + ) + self.assertEqual(control_help.returncode, 0, control_help.stderr) + self.assertIn("status", control_help.stdout) + self.assertIn("STATUS EXIT CODES:", control_help.stdout) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_half_packet_does_not_block_another_client(self) -> None: + with tempfile.TemporaryDirectory() as directory: + process, socket_path, frame = self.start_runtime_daemon( + Path(directory), + self.free_device(), + "--max-inflight", + "2", + "--request-timeout-ms", + "1500", + ) + slow = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + slow.connect(os.fspath(socket_path)) + slow.sendall(frame[:4]) + time.sleep(0.1) + started = time.monotonic() + _, status, _ = self.request(socket_path, frame) + elapsed = time.monotonic() - started + self.assertEqual(status, 0) + self.assertLess(elapsed, 1.0) + finally: + slow.close() + process.terminate() + output, _ = process.communicate(timeout=5) + self.assertEqual(process.returncode, 0, output) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_concurrent_clients_receive_correct_results(self) -> None: + with tempfile.TemporaryDirectory() as directory: + process, socket_path, frame = self.start_runtime_daemon( + Path(directory), + self.free_device(), + "--max-inflight", + "8", + "--request-timeout-ms", + "3000", + ) + try: + with ThreadPoolExecutor(max_workers=8) as executor: + responses = list( + executor.map(lambda _: self.request(socket_path, frame), range(8)) + ) + self.assertEqual([response[1] for response in responses], [0] * 8) + self.assertTrue( + all( + [candidate for candidate, _ in response[2]] == [11, 12] + for response in responses + ) + ) + finally: + process.terminate() + output, _ = process.communicate(timeout=5) + events = [ + json.loads(line) for line in output.splitlines() if line.startswith("{") + ] + completed = [ + event + for event in events + if event.get("event") == "tilemaxsim_rust_request" + and event.get("status") == "ok" + ] + self.assertEqual(len(completed), 8, output) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_full_connection_queue_returns_resource_limit(self) -> None: + with tempfile.TemporaryDirectory() as directory: + process, socket_path, frame = self.start_runtime_daemon( + Path(directory), + self.free_device(), + "--max-inflight", + "1", + "--backlog", + "1", + "--request-timeout-ms", + "1500", + ) + first = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + second = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + third = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + first.connect(os.fspath(socket_path)) + first.sendall(frame[:4]) + time.sleep(0.1) + second.connect(os.fspath(socket_path)) + second.sendall(frame[:4]) + time.sleep(0.1) + third.settimeout(2) + third.connect(os.fspath(socket_path)) + _, status, _ = decode_response(self.receive_response(third)) + self.assertEqual(status, 2) + finally: + first.close() + second.close() + third.close() + process.terminate() + output, _ = process.communicate(timeout=5) + self.assertEqual(process.returncode, 0, output) + self.assertIn('"status":"connection_backlog_full"', output) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_tensor_batch_limit_rejects_before_gpu_queue(self) -> None: + with tempfile.TemporaryDirectory() as directory: + process, socket_path, frame = self.start_runtime_daemon( + Path(directory), + self.free_device(), + "--max-batch-tokens", + "1", + ) + try: + _, status, _ = self.request(socket_path, frame) + self.assertEqual(status, 2) + finally: + process.terminate() + output, _ = process.communicate(timeout=5) + self.assertEqual(process.returncode, 0, output) + self.assertIn('"status":"batch_limit"', output) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_sigterm_removes_socket_ready_and_pid_files(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + ready_path = root / "ready.json" + pid_path = root / "tilemaxsimd.pid" + process, socket_path, frame = self.start_runtime_daemon( + root, + self.free_device(), + "--request-timeout-ms", + "300", + "--shutdown-grace-ms", + "1000", + "--ready-file", + os.fspath(ready_path), + "--pid-file", + os.fspath(pid_path), + ) + partial: socket.socket | None = None + output = "" + try: + for _ in range(200): + if ready_path.exists() and pid_path.exists(): + break + self.assertIsNone(process.poll()) + time.sleep(0.01) + self.assertTrue(ready_path.exists()) + self.assertEqual(int(pid_path.read_text().strip()), process.pid) + ready = json.loads(ready_path.read_text()) + self.assertEqual(ready["pid"], process.pid) + self.assertEqual(ready["schema_version"], 1) + status = subprocess.run( + [ + os.fspath(self.control_binary), + "status", + "--socket", + os.fspath(socket_path), + "--ready-file", + os.fspath(ready_path), + ], + check=False, + capture_output=True, + text=True, + ) + self.assertEqual(status.returncode, 0, status.stderr) + self.assertIn("accepting TileMaxSim connections", status.stdout) + partial = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + partial.connect(os.fspath(socket_path)) + partial.sendall(frame[:4]) + started = time.monotonic() + process.terminate() + output, _ = process.communicate(timeout=3) + elapsed = time.monotonic() - started + finally: + if partial is not None: + partial.close() + if process.poll() is None: + process.terminate() + process.wait(timeout=5) + self.assertEqual(process.returncode, 0, output) + self.assertLess(elapsed, 2) + self.assertFalse(socket_path.exists()) + self.assertFalse(ready_path.exists()) + self.assertFalse(pid_path.exists()) + self.assertIn('"event":"tilemaxsim_rust_stopped"', output) + self.assertIn('"drained":true', output) + stopped_status = subprocess.run( + [ + os.fspath(self.control_binary), + "status", + "--socket", + os.fspath(socket_path), + "--quiet", + ], + check=False, + ) + self.assertEqual(stopped_status.returncode, 1) + + @unittest.skipUnless( + torch.cuda.is_available() and SOAK_SECONDS > 0, + "set VECTORCHORD_TILEMAXSIM_SOAK_SECONDS on a GPU release runner", + ) + def test_configured_concurrency_and_disconnect_soak(self) -> None: + def snapshot(pid: int) -> tuple[int, int, int]: + output = subprocess.check_output( + [ + "nvidia-smi", + "--query-compute-apps=pid,used_memory", + "--format=csv,noheader,nounits", + ], + text=True, + ) + gpu_mib = next( + int(fields[1]) + for fields in ( + [value.strip() for value in line.split(",")] + for line in output.splitlines() + ) + if int(fields[0]) == pid + ) + status = Path(f"/proc/{pid}/status").read_text().splitlines() + rss_kib = int( + next(line for line in status if line.startswith("VmRSS:")).split()[1] + ) + file_descriptors = len(list(Path(f"/proc/{pid}/fd").iterdir())) + return gpu_mib, rss_kib, file_descriptors + + with tempfile.TemporaryDirectory() as directory: + process, socket_path, frame = self.start_runtime_daemon( + Path(directory), + self.free_device(), + "--max-inflight", + "32", + "--backlog", + "256", + "--max-queued-requests", + "256", + "--request-timeout-ms", + "5000", + "--shutdown-grace-ms", + "5000", + ) + recent_output: deque[str] = deque(maxlen=100) + + def drain_output() -> None: + assert process.stdout is not None + recent_output.extend(process.stdout) + + reader = threading.Thread(target=drain_output, daemon=True) + reader.start() + errors: list[str] = [] + try: + for _ in range(100): + self.assertEqual(self.request(socket_path, frame)[1], 0) + before = snapshot(process.pid) + deadline = time.monotonic() + SOAK_SECONDS + + def valid_worker() -> int: + completed = 0 + while time.monotonic() < deadline: + try: + status = self.request(socket_path, frame)[1] + if status == 0: + completed += 1 + else: + errors.append(f"valid request returned status {status}") + except Exception as error: # captured for the main test thread + errors.append(f"valid request failed: {error!r}") + return completed + + def disconnect_worker() -> int: + completed = 0 + while time.monotonic() < deadline: + try: + with socket.socket( + socket.AF_UNIX, socket.SOCK_STREAM + ) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame[:4]) + completed += 1 + time.sleep(0.005) + except Exception as error: # captured for the main test thread + errors.append(f"disconnect churn failed: {error!r}") + return completed + + with ThreadPoolExecutor(max_workers=10) as executor: + futures = [executor.submit(valid_worker) for _ in range(8)] + futures += [executor.submit(disconnect_worker) for _ in range(2)] + counts = [future.result() for future in futures] + self.assertFalse(errors, errors[:10]) + self.assertGreater(sum(counts[:8]), 0) + self.assertGreater(sum(counts[8:]), 0) + status = subprocess.run( + [ + os.fspath(self.control_binary), + "status", + "--socket", + os.fspath(socket_path), + "--quiet", + ], + check=False, + ) + self.assertEqual(status.returncode, 0) + after = snapshot(process.pid) + self.assertLessEqual(after[0] - before[0], 2) + self.assertLessEqual(after[1] - before[1], 64 * 1024) + self.assertLessEqual(after[2] - before[2], 2) + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=10) + reader.join(timeout=2) + if process.stdout is not None: + process.stdout.close() + self.assertEqual(process.returncode, 0, "".join(recent_output)) + if __name__ == "__main__": unittest.main() diff --git a/services/test_tilemaxsim_shard.py b/services/test_tilemaxsim_shard.py index e5684375..dfe66ec4 100644 --- a/services/test_tilemaxsim_shard.py +++ b/services/test_tilemaxsim_shard.py @@ -1,12 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing from __future__ import annotations diff --git a/services/tilemaxsim_cuda_sidecar.py b/services/tilemaxsim_cuda_sidecar.py index fc256fcd..a18d0a77 100644 --- a/services/tilemaxsim_cuda_sidecar.py +++ b/services/tilemaxsim_cuda_sidecar.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing """Bounded CUDA executor for VectorChord TileMaxSim IPC v1 and v2. diff --git a/services/tilemaxsim_gpu_cache.py b/services/tilemaxsim_gpu_cache.py index 78f507be..5f45934f 100644 --- a/services/tilemaxsim_gpu_cache.py +++ b/services/tilemaxsim_gpu_cache.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing """Process-owned GPU arenas and a bounded tensor cache for TileMaxSim.""" diff --git a/services/tilemaxsim_shard.py b/services/tilemaxsim_shard.py index f6307feb..8b06b761 100644 --- a/services/tilemaxsim_shard.py +++ b/services/tilemaxsim_shard.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing """Immutable, content-addressed TileMaxSim tensor shard format. diff --git a/services/tilemaxsim_triton.py b/services/tilemaxsim_triton.py index 5571ffad..c72ef4eb 100644 --- a/services/tilemaxsim_triton.py +++ b/services/tilemaxsim_triton.py @@ -1,16 +1,4 @@ -# This software is licensed under a dual license model: -# -# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -# distribute this software under the terms of the AGPLv3. -# -# Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the Elastic License v2, which has specific restrictions. -# -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 HuXinjing """Fused ragged FP16 TileMaxSim over a process-owned GPU tensor arena.""" diff --git a/services/tilemaxsimd/Cargo.toml b/services/tilemaxsimd/Cargo.toml index 30c5d82a..be5e3bed 100644 --- a/services/tilemaxsimd/Cargo.toml +++ b/services/tilemaxsimd/Cargo.toml @@ -4,6 +4,14 @@ version = "0.1.0" edition = "2024" publish = false +[[bin]] +name = "vchord-tilemaxsimd" +path = "src/main.rs" + +[[bin]] +name = "vchord-tilemaxsimctl" +path = "src/bin/vchord-tilemaxsimctl.rs" + [dependencies] anyhow = "1.0" clap = { version = "4.5", features = ["derive"] } diff --git a/services/tilemaxsimd/README.md b/services/tilemaxsimd/README.md new file mode 100644 index 00000000..c3f9e99a --- /dev/null +++ b/services/tilemaxsimd/README.md @@ -0,0 +1,205 @@ +# Native TileMaxSim service + +`vchord-tilemaxsimd` is the optional Rust/CUDA execution service for exact +TileMaxSim reranking. It owns explicitly assigned GPU memory, a decoded host +cache, immutable tensor-shard readers, cache admission, and the scoring +scheduler. Base pgvector and VectorChord queries do not start this process or +reserve GPU memory. If `vchordrq.maxsim_gpu_endpoint` is not configured, there +is no TileMaxSim daemon requirement. + +The service is Linux-only and requires an NVIDIA driver plus a CUDA toolkit at +build time. The supplied image builds against CUDA 12.6. The runtime driver +must satisfy NVIDIA's compatibility requirements for that CUDA version. + +## Build and command reference + +Build both installed commands from the repository root: + +```bash +cargo build --release --locked \ + --manifest-path services/tilemaxsimd/Cargo.toml +install -m 0755 services/tilemaxsimd/target/release/vchord-tilemaxsimd \ + /usr/local/bin/vchord-tilemaxsimd +install -m 0755 services/tilemaxsimd/target/release/vchord-tilemaxsimctl \ + /usr/local/bin/vchord-tilemaxsimctl +``` + +The build host needs `nvcc`, a C++ compiler, and Rust. The CLI is the canonical +reference: + +```bash +vchord-tilemaxsimd --help +vchord-tilemaxsimd --version +vchord-tilemaxsimctl --help +vchord-tilemaxsimctl status --help +``` + +All memory arguments exposed to operators use GB, interpreted as GiB +(`1024^3` bytes). Raw byte counts are intentionally not accepted for GPU or +host cache sizing. + +## Minimal start + +```bash +install -d -m 0750 /run/vectorchord +vchord-tilemaxsimd \ + --socket /run/vectorchord/tilemaxsim.sock \ + --gpu-memory-gb 0=20 \ + --gpu-workspace-gb 2 \ + --host-cache-gb 8 \ + --contract-root model@1=/var/lib/vectorchord/tensors/model-v1 \ + --ready-file /run/vectorchord/tilemaxsim.ready \ + --pid-file /run/vectorchord/tilemaxsimd.pid +``` + +Each `GPU=GB` value is a strict, process-owned allocation. Every configured +arena is allocated before the socket and readiness file are created. A missing +GPU, duplicate device, insufficient free memory, resident-manifest overflow, +or shard validation failure makes startup fail nonzero. There is no silent CPU +fallback and no partial ready state. + +The GPU allocation contains both tensor pages and the per-device workspace. +For example, `--gpu-memory-gb 0=20 --gpu-workspace-gb 2` leaves 18 GiB for +cached tensors. `--gpu-block-kib 32` is the default arena page size. Use `lru` +for bounded admission and eviction, or use `resident` with one or more +`--resident-manifest MODEL=PATH` arguments to pin the full declared working set +before readiness. + +## PostgreSQL connection + +The relevant PostgreSQL settings are ordinary VectorChord GUCs: + +```sql +ALTER SYSTEM SET vchordrq.maxsim_backend = 'gpu'; +ALTER SYSTEM SET vchordrq.maxsim_gpu_endpoint = + '/run/vectorchord/tilemaxsim.sock'; +ALTER SYSTEM SET vchordrq.maxsim_gpu_timeout_ms = 2000; +ALTER SYSTEM SET vchordrq.maxsim_gpu_max_batch_tokens = 250000; +ALTER SYSTEM SET vchordrq.maxsim_gpu_max_batch_bytes = 268435456; +SELECT pg_reload_conf(); +``` + +Keep PostgreSQL's timeout no greater than the daemon's +`--request-timeout-ms`, or use the same value for both. PostgreSQL also limits +encoded batch tokens and bytes before connecting. Candidate generation, +clustering, graph construction, tenant policy, and application routing remain +outside this daemon. + +By default only clients with the daemon's effective Unix UID are accepted. +When PostgreSQL runs under another account, add its numeric UID explicitly: + +```bash +--allow-peer-uid "$(id -u postgres)" +``` + +The socket directory and `--socket-mode` must also permit PostgreSQL to +connect. `SO_PEERCRED` is checked on every connection; an allowed GID refers to +the client's effective primary GID, not a supplementary group. Do not expose +the socket through a network proxy. + +## Health, shutdown, and observability + +```bash +vchord-tilemaxsimctl status \ + --socket /run/vectorchord/tilemaxsim.sock \ + --ready-file /run/vectorchord/tilemaxsim.ready +``` + +The status command sends a valid zero-candidate v2 request. It therefore checks +the listener, an I/O worker, the bounded request queue, and the GPU scheduler, +while reading no tensor and launching no scoring kernel. Exit status is `0` +when ready, `1` when unavailable, and `2` for invalid CLI usage. + +SIGINT and SIGTERM first remove the ready file and close the listening socket, +then drain accepted work for `--shutdown-grace-ms`. PID, ready, and socket paths +are guarded by device/inode identity so shutdown does not delete a replacement +file. A second instance refuses to unlink an active socket. A core worker panic +or an expired drain deadline produces a nonzero process exit. + +Standard output contains one JSON object per lifecycle or request event. Event +names are `tilemaxsim_rust_ready`, `tilemaxsim_rust_prewarm_complete`, +`tilemaxsim_rust_request`, and `tilemaxsim_rust_stopped`; each stable event has +`schema_version: 1`. Request events include queue, compute, total latency, peer +credentials, candidate count, client presence, and cache counters. Send stdout +and stderr to the service manager or a structured log collector. + +CUDA kernels cannot be safely preempted in the middle of a launch. Deadlines +and disconnect cancellation are enforced before and between scoring chunks; +one already-running kernel completes before its request is released. + +The GPU integration suite includes a configurable soak gate. Release runners +set `VECTORCHORD_TILEMAXSIM_SOAK_SECONDS=3600` (or longer) before running +`python -m unittest services.test_tilemaxsim_rust_daemon`. The soak mixes valid +concurrent requests with incomplete-client disconnects and fails on result, +health, GPU-memory, RSS, FD-recovery, or shutdown regressions. It is skipped +when the variable is absent so ordinary development runs remain short. + +## Bounded load behavior + +`--max-inflight` bounds clients being read or awaiting a response. +`--backlog` bounds both the kernel listen backlog and the accepted-connection +queue. `--max-queued-requests` bounds parsed work waiting for the single +request scheduler. Full queues return a resource-limit response instead of +growing memory without limit. `--max-request-mb` bounds each encoded frame; +`--max-batch-tokens` and `--max-batch-mb` independently bound the decoded +query-plus-candidate tensor working set before it enters the GPU queue. + +Tune PostgreSQL candidate limits and batch limits before increasing daemon +queues. More queue depth absorbs bursts but does not add GPU throughput and +increases tail latency. Multiple configured GPUs execute the chunks of one +request concurrently; requests themselves enter the scheduler in arrival +order. + +## systemd + +Create a dedicated unprivileged `vectorchord` user, make tensor roots readable +but not writable by that user, and install the public templates: + +```bash +install -D -m 0644 services/packaging/vchord-tilemaxsimd.service \ + /etc/systemd/system/vchord-tilemaxsimd.service +install -D -m 0640 services/packaging/tilemaxsimd.conf.example \ + /etc/vectorchord/tilemaxsimd.conf +systemctl daemon-reload +systemctl enable --now vchord-tilemaxsimd +systemctl status vchord-tilemaxsimd +``` + +Edit the example first. In particular, replace the PostgreSQL UID, GPU size, +contract ID, and tensor path. If PostgreSQL starts immediately after the unit, +keep the `Before=postgresql.service` ordering and make PostgreSQL require this +unit when GPU TileMaxSim is mandatory. + +## Container and Kubernetes + +The runtime image uses UID/GID `65532`, drops root, and handles SIGTERM. Build +it with: + +```bash +docker build -f services/Dockerfile.tilemaxsimd -t vectorchord-tilemaxsimd . +``` + +Mount the Unix-socket directory read-write, mount immutable tensor shards +read-only, and pass one or more GPUs through the NVIDIA container runtime. +Visible devices are renumbered inside many containers, so a container given +one host GPU usually configures it as `0=GB`. + +The Kubernetes example at +[`../packaging/tilemaxsimd-kubernetes.yaml`](../packaging/tilemaxsimd-kubernetes.yaml) +runs the daemon as a PostgreSQL sidecar with a shared `emptyDir`, non-root +security settings, an NVIDIA GPU limit, and startup/readiness/liveness probes. +Replace both example images, the PVC, sizes, contract ID, and PostgreSQL UID. + +## Upgrade, restore, and cache recovery + +For a planned daemon restart, stop or redirect GPU TileMaxSim queries, send +SIGTERM, wait for a clean exit, replace the binaries or image, and require a +successful status probe before restoring traffic. Do not run two daemon +versions against the same socket path. + +Back up PostgreSQL data together with the immutable shard set and the descriptor +manifests referenced by its tensor-source registry. Preserve content hashes, +relative shard metadata, file modes, and model contract IDs. GPU contents, the +decoded host cache, PID files, ready files, and sockets are disposable and must +not be backed up. After restore, validate shard access and start the daemon; LRU +caches refill on demand and resident caches rebuild completely before ready. diff --git a/services/tilemaxsimd/build.rs b/services/tilemaxsimd/build.rs index 15cbfae6..99f4a946 100644 --- a/services/tilemaxsimd/build.rs +++ b/services/tilemaxsimd/build.rs @@ -1,7 +1,10 @@ +// Copyright (c) 2026 HuXinjing + fn main() { println!("cargo:rerun-if-changed=native/tilemaxsim_cuda.cu"); cc::Build::new() .cuda(true) + .debug(false) .flag("-O3") .flag("--use_fast_math") .flag("-lineinfo") diff --git a/services/tilemaxsimd/native/tilemaxsim_cuda.cu b/services/tilemaxsimd/native/tilemaxsim_cuda.cu index ff214e6e..05dfd49a 100644 --- a/services/tilemaxsimd/native/tilemaxsim_cuda.cu +++ b/services/tilemaxsimd/native/tilemaxsim_cuda.cu @@ -1,3 +1,5 @@ +// Copyright (c) 2026 HuXinjing + #include #include #include diff --git a/services/tilemaxsimd/src/bin/vchord-tilemaxsimctl.rs b/services/tilemaxsimd/src/bin/vchord-tilemaxsimctl.rs new file mode 100644 index 00000000..b6b0d934 --- /dev/null +++ b/services/tilemaxsimd/src/bin/vchord-tilemaxsimctl.rs @@ -0,0 +1,338 @@ +// Copyright (c) 2026 HuXinjing + +use clap::{Args, Parser, Subcommand}; +use serde_json::Value; +use std::fs; +use std::io::{Read, Write}; +use std::os::fd::{AsRawFd, FromRawFd, OwnedFd}; +use std::os::unix::ffi::OsStrExt; +use std::os::unix::fs::FileTypeExt; +use std::os::unix::net::UnixStream; +use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant}; + +const AFTER_HELP: &str = r#"EXAMPLES: + Check only that the service socket accepts connections: + vchord-tilemaxsimctl status --socket /run/vectorchord/tilemaxsim.sock + + Also require a coherent daemon readiness file: + vchord-tilemaxsimctl status --socket /run/vectorchord/tilemaxsim.sock \ + --ready-file /run/vectorchord/tilemaxsim.ready + +STATUS EXIT CODES: + 0 The readiness file is valid (when requested) and the socket accepts. + 1 The daemon is not ready or cannot be reached. + 2 Command-line usage is invalid."#; + +#[derive(Debug, Parser)] +#[command( + name = "vchord-tilemaxsimctl", + version, + about = "Inspect a VectorChord TileMaxSim daemon", + long_about = "Perform local operational checks against vchord-tilemaxsimd.\n\ +The status command submits an empty v2 request through the I/O workers and GPU\n\ +scheduler. It reads no tensor data and launches no scoring kernel.", + after_long_help = AFTER_HELP, + subcommand_required = true, + arg_required_else_help = true +)] +struct Cli { + #[command(subcommand)] + command: Command, +} + +#[derive(Debug, Subcommand)] +enum Command { + /// Report whether the daemon is accepting local connections. + Status(StatusArgs), +} + +#[derive(Debug, Args)] +struct StatusArgs { + /// Unix-domain socket configured on vchord-tilemaxsimd. + #[arg(long, value_name = "PATH")] + socket: PathBuf, + + /// Also validate this daemon-created readiness JSON file. + #[arg(long, value_name = "PATH")] + ready_file: Option, + + /// Overall deadline for connecting and completing the empty probe. + #[arg( + long, + value_name = "MILLISECONDS", + default_value_t = 1_000, + value_parser = parse_positive_u32 + )] + timeout_ms: u32, + + /// Suppress normal status output; the exit code remains authoritative. + #[arg(short, long)] + quiet: bool, +} + +fn main() { + let Cli { command } = Cli::parse(); + let (socket, quiet, result) = match command { + Command::Status(args) => { + let result = status(&args); + (args.socket, args.quiet, result) + } + }; + match result { + Ok(()) => { + if !quiet { + println!("{} - accepting TileMaxSim connections", socket.display()); + } + } + Err(error) => { + if !quiet { + eprintln!("{} - no response: {error}", socket.display()); + } + std::process::exit(1); + } + } +} + +fn status(args: &StatusArgs) -> Result<(), String> { + let metadata = fs::symlink_metadata(&args.socket) + .map_err(|error| format!("cannot inspect socket: {error}"))?; + if !metadata.file_type().is_socket() { + return Err("configured path is not a Unix-domain socket".to_owned()); + } + if let Some(path) = &args.ready_file { + validate_ready_file(path, &args.socket)?; + } + let deadline = Instant::now() + Duration::from_millis(args.timeout_ms.into()); + let mut connection = connect_with_timeout(&args.socket, remaining(deadline)?)?; + probe_scheduler(&mut connection, deadline) +} + +fn parse_positive_u32(value: &str) -> Result { + let value = value + .parse::() + .map_err(|_| "value must be a positive integer".to_owned())?; + if value == 0 { + Err("value must be a positive integer".to_owned()) + } else { + Ok(value) + } +} + +fn validate_ready_file(path: &Path, socket: &Path) -> Result<(), String> { + let contents = + fs::read(path).map_err(|error| format!("cannot read readiness file: {error}"))?; + let value: Value = serde_json::from_slice(&contents) + .map_err(|error| format!("invalid readiness JSON: {error}"))?; + if value.get("schema_version").and_then(Value::as_u64) != Some(1) { + return Err("unsupported readiness schema".to_owned()); + } + if value.get("socket").and_then(Value::as_str) != socket.to_str() { + return Err("readiness file names a different socket".to_owned()); + } + let pid = value + .get("pid") + .and_then(Value::as_u64) + .and_then(|pid| libc::pid_t::try_from(pid).ok()) + .filter(|pid| *pid > 0) + .ok_or_else(|| "readiness file contains an invalid PID".to_owned())?; + // SAFETY: kill(pid, 0) only checks whether a process exists and is visible. + if unsafe { libc::kill(pid, 0) } != 0 { + let error = std::io::Error::last_os_error(); + if error.raw_os_error() != Some(libc::EPERM) { + return Err(format!("readiness PID is not running: {error}")); + } + } + Ok(()) +} + +fn connect_with_timeout(path: &Path, timeout: Duration) -> Result { + let bytes = path.as_os_str().as_bytes(); + if bytes.contains(&0) { + return Err("socket path contains a NUL byte".to_owned()); + } + let mut address = unsafe { std::mem::zeroed::() }; + if bytes.len() >= address.sun_path.len() { + return Err("socket path is too long".to_owned()); + } + address.sun_family = libc::AF_UNIX as libc::sa_family_t; + for (target, source) in address.sun_path.iter_mut().zip(bytes.iter().copied()) { + *target = source as libc::c_char; + } + // SAFETY: socket has no pointer arguments; successful ownership is moved to OwnedFd. + let raw = unsafe { + libc::socket( + libc::AF_UNIX, + libc::SOCK_STREAM | libc::SOCK_NONBLOCK | libc::SOCK_CLOEXEC, + 0, + ) + }; + if raw < 0 { + return Err(format!( + "cannot create health-check socket: {}", + std::io::Error::last_os_error() + )); + } + // SAFETY: `raw` is a new, uniquely owned file descriptor. + let descriptor = unsafe { OwnedFd::from_raw_fd(raw) }; + let address_length = + (std::mem::offset_of!(libc::sockaddr_un, sun_path) + bytes.len() + 1) as libc::socklen_t; + // SAFETY: the address is initialized for AF_UNIX and its exact length is supplied. + if unsafe { + libc::connect( + descriptor.as_raw_fd(), + std::ptr::addr_of!(address).cast(), + address_length, + ) + } == 0 + { + return Ok(UnixStream::from(descriptor)); + } + let error = std::io::Error::last_os_error(); + if !matches!(error.raw_os_error(), Some(libc::EINPROGRESS | libc::EAGAIN)) { + return Err(format!("cannot connect: {error}")); + } + let timeout_ms = i32::try_from(timeout.as_millis()).unwrap_or(i32::MAX); + let mut poll = libc::pollfd { + fd: descriptor.as_raw_fd(), + events: libc::POLLOUT, + revents: 0, + }; + // SAFETY: exactly one valid pollfd is supplied for the bounded wait. + let poll_result = unsafe { libc::poll(std::ptr::addr_of_mut!(poll), 1, timeout_ms) }; + if poll_result == 0 { + return Err("connection timed out".to_owned()); + } + if poll_result < 0 { + return Err(format!( + "connection poll failed: {}", + std::io::Error::last_os_error() + )); + } + let mut socket_error = 0; + let mut error_length = std::mem::size_of_val(&socket_error) as libc::socklen_t; + // SAFETY: both output pointers refer to initialized writable values of the required size. + if unsafe { + libc::getsockopt( + descriptor.as_raw_fd(), + libc::SOL_SOCKET, + libc::SO_ERROR, + std::ptr::addr_of_mut!(socket_error).cast(), + &mut error_length, + ) + } != 0 + { + return Err(format!( + "cannot inspect connection: {}", + std::io::Error::last_os_error() + )); + } + if socket_error != 0 { + return Err(format!( + "cannot connect: {}", + std::io::Error::from_raw_os_error(socket_error) + )); + } + Ok(UnixStream::from(descriptor)) +} + +fn probe_scheduler(connection: &mut UnixStream, deadline: Instant) -> Result<(), String> { + const REQUEST_ID: u64 = 0x5643_4845_414c_5448; + let contract = b"health"; + let body_bytes = 4 + 4 + 4 + 1 + 1 + 2 + 4 + contract.len() + 2; + let mut request = Vec::with_capacity(24 + body_bytes); + request.extend_from_slice(b"VCTM"); + request.extend_from_slice(&2_u16.to_le_bytes()); + request.extend_from_slice(&1_u16.to_le_bytes()); + request.extend_from_slice(&REQUEST_ID.to_le_bytes()); + request.extend_from_slice(&(body_bytes as u64).to_le_bytes()); + request.extend_from_slice(&1_u32.to_le_bytes()); + request.extend_from_slice(&1_u32.to_le_bytes()); + request.extend_from_slice(&0_u32.to_le_bytes()); + request.push(2); + request.push(1); + request.extend_from_slice(&0_u16.to_le_bytes()); + request.extend_from_slice(&(contract.len() as u32).to_le_bytes()); + request.extend_from_slice(contract); + request.extend_from_slice(&0_u16.to_le_bytes()); + + connection + .set_nonblocking(false) + .map_err(|error| format!("cannot configure probe socket: {error}"))?; + set_deadlines(connection, deadline)?; + connection + .write_all(&request) + .map_err(|error| format!("cannot write scheduler probe: {error}"))?; + let mut header = [0_u8; 24]; + set_deadlines(connection, deadline)?; + connection + .read_exact(&mut header) + .map_err(|error| format!("cannot read scheduler probe: {error}"))?; + if &header[..4] != b"VCTM" + || u16::from_le_bytes(header[4..6].try_into().unwrap()) != 2 + || u16::from_le_bytes(header[6..8].try_into().unwrap()) != 2 + || u64::from_le_bytes(header[8..16].try_into().unwrap()) != REQUEST_ID + { + return Err("daemon returned an invalid probe response header".to_owned()); + } + let body_length = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) + .map_err(|_| "daemon probe response is too large".to_owned())?; + if body_length > 64 * 1024 { + return Err("daemon probe response is too large".to_owned()); + } + let mut body = vec![0_u8; body_length]; + set_deadlines(connection, deadline)?; + connection + .read_exact(&mut body) + .map_err(|error| format!("cannot read scheduler probe body: {error}"))?; + if body.len() != 8 + || u32::from_le_bytes(body[..4].try_into().unwrap()) != 0 + || u32::from_le_bytes(body[4..8].try_into().unwrap()) != 0 + { + return Err("daemon scheduler rejected the empty probe".to_owned()); + } + Ok(()) +} + +fn set_deadlines(connection: &UnixStream, deadline: Instant) -> Result<(), String> { + let timeout = remaining(deadline)?; + connection + .set_read_timeout(Some(timeout)) + .and_then(|()| connection.set_write_timeout(Some(timeout))) + .map_err(|error| format!("cannot configure probe deadline: {error}")) +} + +fn remaining(deadline: Instant) -> Result { + let timeout = deadline.saturating_duration_since(Instant::now()); + if timeout.is_zero() { + Err("probe deadline expired".to_owned()) + } else { + Ok(timeout) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use clap::CommandFactory; + + #[test] + fn command_definition_and_exit_codes_are_documented() { + Cli::command().debug_assert(); + let help = Cli::command().render_long_help().to_string(); + assert!(help.contains("vchord-tilemaxsimctl")); + assert!(help.contains("STATUS EXIT CODES:")); + assert!(help.contains("status --socket")); + assert!( + Cli::try_parse_from([ + "vchord-tilemaxsimctl", + "status", + "--socket", + "/tmp/test.sock", + "--timeout-ms", + "0" + ]) + .is_err() + ); + } +} diff --git a/services/tilemaxsimd/src/cache.rs b/services/tilemaxsimd/src/cache.rs index bcfcf57a..6abcadc1 100644 --- a/services/tilemaxsimd/src/cache.rs +++ b/services/tilemaxsimd/src/cache.rs @@ -1,3 +1,5 @@ +// Copyright (c) 2026 HuXinjing + use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::hash::{Hash, Hasher}; diff --git a/services/tilemaxsimd/src/engine.rs b/services/tilemaxsimd/src/engine.rs index 7668d67d..d0365290 100644 --- a/services/tilemaxsimd/src/engine.rs +++ b/services/tilemaxsimd/src/engine.rs @@ -1,8 +1,12 @@ +// Copyright (c) 2026 HuXinjing + use crate::cache::{Admission, GpuCache}; use crate::gpu::Gpu; use crate::protocol::{Descriptor, Request}; use crate::shard::{ShardStore, cache_key}; use anyhow::{Result, anyhow, bail}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Instant; struct MissingTensor { candidate_index: usize, @@ -131,7 +135,22 @@ impl Engine { Ok(()) } - pub fn score(&mut self, request: &Request) -> Result> { + pub fn score_until( + &mut self, + request: &Request, + deadline: Instant, + canceled: &AtomicBool, + ) -> Result> { + self.score_controlled(request, Some(deadline), Some(canceled)) + } + + fn score_controlled( + &mut self, + request: &Request, + deadline: Option, + canceled: Option<&AtomicBool>, + ) -> Result> { + check_request_control(deadline, canceled)?; if request.candidates.is_empty() { return Ok(Vec::new()); } @@ -180,6 +199,7 @@ impl Engine { self.release_chunk(chunk)?; } + check_request_control(deadline, canceled)?; let payloads = self.store.resolve_many(&missing_descriptors)?; let mut pending = missing_indices .into_iter() @@ -194,6 +214,7 @@ impl Engine { .collect::>(); while !pending.is_empty() { + check_request_control(deadline, canceled)?; let mut chunks = (0..self.devices.len()) .map(|_| Vec::::new()) .collect::>(); @@ -293,6 +314,7 @@ impl Engine { for chunk in &chunks { self.release_chunk(chunk)?; } + check_request_control(deadline, canceled)?; pending.drain(..consumed); } @@ -426,6 +448,16 @@ impl Engine { } } +fn check_request_control(deadline: Option, canceled: Option<&AtomicBool>) -> Result<()> { + if canceled.is_some_and(|flag| flag.load(Ordering::Relaxed)) { + bail!("request canceled because the client disconnected"); + } + if deadline.is_some_and(|deadline| Instant::now() >= deadline) { + bail!("request deadline expired"); + } + Ok(()) +} + fn validate_entry(descriptor: &Descriptor, entry: &crate::cache::CacheEntry) -> Result<()> { let scalar_bytes = if descriptor.dtype == 1 { 4 } else { 2 }; let expected_bytes = descriptor.rows as usize * descriptor.dimension as usize * scalar_bytes; diff --git a/services/tilemaxsimd/src/gpu.rs b/services/tilemaxsimd/src/gpu.rs index 2c6c2df0..334707dc 100644 --- a/services/tilemaxsimd/src/gpu.rs +++ b/services/tilemaxsimd/src/gpu.rs @@ -1,3 +1,5 @@ +// Copyright (c) 2026 HuXinjing + use anyhow::{Result, anyhow, bail}; use std::ffi::{CStr, c_char, c_int, c_uchar, c_void}; use std::ptr::NonNull; diff --git a/services/tilemaxsimd/src/lifecycle.rs b/services/tilemaxsimd/src/lifecycle.rs new file mode 100644 index 00000000..50a18a03 --- /dev/null +++ b/services/tilemaxsimd/src/lifecycle.rs @@ -0,0 +1,188 @@ +// Copyright (c) 2026 HuXinjing + +use anyhow::{Context, Result, bail}; +use std::fs::{self, OpenOptions}; +use std::io::{ErrorKind, Write}; +use std::os::unix::fs::{MetadataExt, OpenOptionsExt}; +use std::path::{Path, PathBuf}; + +pub struct ManagedFile { + path: PathBuf, + device: u64, + inode: u64, + _file: fs::File, +} + +impl ManagedFile { + pub fn create_pid(path: &Path) -> Result { + remove_stale_pid_file(path)?; + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .mode(0o600) + .open(path) + .with_context(|| format!("cannot create PID file {}", path.display()))?; + writeln!(file, "{}", std::process::id())?; + file.sync_all()?; + let metadata = file.metadata()?; + Ok(Self { + path: path.to_owned(), + device: metadata.dev(), + inode: metadata.ino(), + _file: file, + }) + } + + pub fn create_ready(path: &Path, contents: &[u8]) -> Result { + if fs::symlink_metadata(path).is_ok() { + bail!("ready file already exists: {}", path.display()); + } + let parent = path.parent().filter(|value| !value.as_os_str().is_empty()); + let parent = parent.unwrap_or_else(|| Path::new(".")); + let name = path + .file_name() + .ok_or_else(|| anyhow::anyhow!("ready file path has no file name"))? + .to_string_lossy(); + let temporary = parent.join(format!(".{name}.{}.tmp", std::process::id())); + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .mode(0o600) + .open(&temporary) + .with_context(|| { + format!("cannot create temporary ready file {}", temporary.display()) + })?; + let result = (|| -> Result { + file.write_all(contents)?; + file.sync_all()?; + fs::hard_link(&temporary, path) + .with_context(|| format!("cannot publish ready file {}", path.display()))?; + let metadata = fs::symlink_metadata(path)?; + let guard = Self { + path: path.to_owned(), + device: metadata.dev(), + inode: metadata.ino(), + _file: file, + }; + fs::File::open(parent)?.sync_all()?; + Ok(guard) + })(); + let _ = fs::remove_file(&temporary); + result + } + + pub fn remove(mut self) -> Result<()> { + self.remove_if_owned() + } + + fn remove_if_owned(&mut self) -> Result<()> { + let metadata = match fs::symlink_metadata(&self.path) { + Ok(metadata) => metadata, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(error.into()), + }; + if metadata.dev() != self.device || metadata.ino() != self.inode { + bail!("refusing to remove replaced file {}", self.path.display()); + } + fs::remove_file(&self.path)?; + Ok(()) + } +} + +impl Drop for ManagedFile { + fn drop(&mut self) { + let _ = self.remove_if_owned(); + } +} + +fn remove_stale_pid_file(path: &Path) -> Result<()> { + let metadata = match fs::symlink_metadata(path) { + Ok(metadata) => metadata, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(error.into()), + }; + if !metadata.file_type().is_file() || metadata.uid() != unsafe { libc::geteuid() } { + bail!("refusing to replace unsafe PID file {}", path.display()); + } + let contents = fs::read_to_string(path) + .with_context(|| format!("cannot read existing PID file {}", path.display()))?; + let pid = contents + .trim() + .parse::() + .with_context(|| format!("invalid existing PID file {}", path.display()))?; + if pid <= 0 { + bail!("invalid existing PID file {}", path.display()); + } + // SAFETY: kill(pid, 0) performs existence/permission checking only. + if unsafe { libc::kill(pid, 0) } == 0 { + bail!("daemon recorded in {} is still running", path.display()); + } + let error = std::io::Error::last_os_error(); + if error.raw_os_error() != Some(libc::ESRCH) { + return Err(error).context("cannot verify existing daemon PID"); + } + let current = fs::symlink_metadata(path)?; + if current.dev() != metadata.dev() || current.ino() != metadata.ino() { + bail!("PID file changed while checking {}", path.display()); + } + fs::remove_file(path)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn directory(name: &str) -> PathBuf { + std::env::temp_dir().join(format!( + "vchord-tilemaxsim-lifecycle-{name}-{}", + std::process::id() + )) + } + + #[test] + fn managed_ready_file_is_removed() { + let directory = directory("ready-remove"); + fs::create_dir_all(&directory).unwrap(); + let path = directory.join("ready"); + let _ = fs::remove_file(&path); + let guard = ManagedFile::create_ready(&path, b"ready\n").unwrap(); + assert_eq!(fs::read(&path).unwrap(), b"ready\n"); + guard.remove().unwrap(); + assert!(!path.exists()); + fs::remove_dir(&directory).unwrap(); + } + + #[test] + fn dropping_guard_does_not_remove_replacement_file() { + let directory = directory("ready-replacement"); + fs::create_dir_all(&directory).unwrap(); + let path = directory.join("ready"); + let _ = fs::remove_file(&path); + let guard = ManagedFile::create_ready(&path, b"original").unwrap(); + fs::remove_file(&path).unwrap(); + fs::write(&path, b"replacement").unwrap(); + drop(guard); + assert_eq!(fs::read(&path).unwrap(), b"replacement"); + fs::remove_file(&path).unwrap(); + fs::remove_dir(&directory).unwrap(); + } + + #[test] + fn pid_file_refuses_live_process_and_replaces_stale_process() { + let directory = directory("pid"); + fs::create_dir_all(&directory).unwrap(); + let path = directory.join("daemon.pid"); + fs::write(&path, format!("{}\n", std::process::id())).unwrap(); + let error = ManagedFile::create_pid(&path).err().unwrap(); + assert!(format!("{error:#}").contains("still running")); + fs::write(&path, format!("{}\n", i32::MAX)).unwrap(); + let guard = ManagedFile::create_pid(&path).unwrap(); + assert_eq!( + fs::read_to_string(&path).unwrap().trim(), + std::process::id().to_string() + ); + guard.remove().unwrap(); + fs::remove_dir(&directory).unwrap(); + } +} diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs index b940a08b..983fb8ef 100644 --- a/services/tilemaxsimd/src/main.rs +++ b/services/tilemaxsimd/src/main.rs @@ -1,62 +1,279 @@ +// Copyright (c) 2026 HuXinjing + mod cache; mod engine; mod gpu; +mod lifecycle; mod protocol; +mod server; mod shard; use anyhow::{Context, Result, anyhow, bail}; -use clap::Parser; +use clap::{Parser, ValueEnum}; use engine::Engine; use gpu::Gpu; -use protocol::HEADER_BYTES; +use lifecycle::ManagedFile; use serde::Deserialize; +use server::ServerConfig; use shard::ShardStore; +use std::collections::HashSet; use std::fs; -use std::io::{BufRead, Read, Write}; -use std::os::unix::fs::{FileTypeExt, PermissionsExt}; -use std::os::unix::net::{UnixListener, UnixStream}; +use std::io::BufRead; use std::path::PathBuf; -use std::time::Instant; +use std::time::{Duration, Instant}; const GIB: usize = 1024 * 1024 * 1024; +const MIB: usize = 1024 * 1024; + +const AFTER_HELP: &str = r#"EXAMPLES: + Start an evictable 20 GB cache on GPU 1: + vchord-tilemaxsimd --socket /run/vectorchord/tilemaxsim.sock \ + --gpu-memory-gb 1=20 --gpu-workspace-gb 2 \ + --contract-root colqwen@1=/var/lib/vectorchord/colqwen + + Pin a complete tensor manifest before accepting requests: + vchord-tilemaxsimd --socket /run/vectorchord/tilemaxsim.sock \ + --gpu-memory-gb 1=20 --gpu-workspace-gb 2 \ + --contract-root colqwen@1=/var/lib/vectorchord/colqwen \ + --gpu-cache-mode resident \ + --resident-manifest colqwen@1=/var/lib/vectorchord/colqwen/descriptors.jsonl -#[derive(Parser)] -#[command(about = "Native Rust/CUDA TileMaxSim shard, cache, and scheduling daemon")] +Memory options named GB use GiB internally (1 GB = 1024^3 bytes). The daemon +allocates every configured GPU arena before creating the socket; startup fails +without leaving a ready socket if any allocation or prewarm step fails."#; + +#[derive(Debug, Parser)] +#[command( + name = "vchord-tilemaxsimd", + version, + about = "VectorChord TileMaxSim GPU cache and scoring service", + long_about = "Run the native VectorChord TileMaxSim service over a local Unix-domain socket.\n\ +The process owns the configured GPU memory, immutable tensor shards, cache\n\ +admission policy, and CUDA scoring scheduler.", + after_long_help = AFTER_HELP, + next_line_help = true +)] struct Args { - #[arg(long)] + /// Unix-domain socket used by the VectorChord PostgreSQL backend. + #[arg(long, value_name = "PATH", help_heading = "Connection options")] socket: PathBuf, - #[arg(long, required = true, value_parser = parse_gpu_memory)] + + /// Socket permissions as an octal mode. + #[arg( + long, + value_name = "MODE", + default_value = "600", + value_parser = parse_mode, + help_heading = "Connection options" + )] + socket_mode: u32, + + /// Maximum time for socket I/O, queueing, and scoring. + #[arg( + long, + value_name = "MILLISECONDS", + default_value_t = 2_000, + value_parser = parse_positive_usize, + help_heading = "Connection options" + )] + request_timeout_ms: usize, + + /// Number of clients that may be read or waiting for a response. + #[arg( + long, + value_name = "COUNT", + default_value_t = 8, + value_parser = parse_positive_usize, + help_heading = "Connection options" + )] + max_inflight: usize, + + /// Kernel listen backlog for the Unix-domain socket. + #[arg( + long, + value_name = "COUNT", + default_value_t = 64, + value_parser = parse_positive_usize, + help_heading = "Connection options" + )] + backlog: usize, + + /// Maximum parsed requests waiting for the GPU scheduler. + #[arg( + long, + value_name = "COUNT", + default_value_t = 64, + value_parser = parse_positive_usize, + help_heading = "Connection options" + )] + max_queued_requests: usize, + + /// Permit a client Unix user ID in addition to the daemon's own UID. + #[arg( + long, + value_name = "UID", + action = clap::ArgAction::Append, + help_heading = "Connection options" + )] + allow_peer_uid: Vec, + + /// Permit a client Unix group ID. + #[arg( + long, + value_name = "GID", + action = clap::ArgAction::Append, + help_heading = "Connection options" + )] + allow_peer_gid: Vec, + + /// Strict process-owned allocation in GPU=GB form; repeat for more GPUs. + #[arg( + long, + required = true, + value_name = "GPU=GB", + value_parser = parse_gpu_memory, + help_heading = "GPU and cache options" + )] gpu_memory_gb: Vec, - #[arg(long, default_value = "2", value_parser = parse_gb)] + + /// Per-GPU portion reserved for queries and scoring output. + #[arg( + long, + value_name = "GB", + default_value = "2", + value_parser = parse_gb, + help_heading = "GPU and cache options" + )] gpu_workspace_gb: usize, - #[arg(long, default_value = "8", value_parser = parse_gb)] + + /// Decoded host-memory tensor cache shared by shard readers. + #[arg( + long, + value_name = "GB", + default_value = "8", + value_parser = parse_gb, + help_heading = "GPU and cache options" + )] host_cache_gb: usize, - #[arg(long = "contract-root", required = true, value_parser = parse_contract_root)] - contract_roots: Vec<(String, PathBuf)>, - #[arg(long, default_value_t = 32)] + + /// GPU cache behavior: evict cold tensors or pin a complete manifest. + #[arg( + long, + value_name = "MODE", + default_value = "lru", + value_enum, + help_heading = "GPU and cache options" + )] + gpu_cache_mode: CacheMode, + + /// Base page size inside each preallocated GPU tensor arena. + #[arg( + long, + value_name = "KIB", + default_value_t = 32, + value_parser = parse_gpu_block_kib, + help_heading = "GPU and cache options" + )] gpu_block_kib: usize, - #[arg(long, default_value_t = 64 * 1024 * 1024)] - max_request_bytes: usize, - #[arg(long, default_value = "600", value_parser = parse_mode)] - socket_mode: u32, - #[arg(long)] - once: bool, - #[arg(long)] + + /// Immutable tensor shard root in MODEL_CONTRACT_ID=/absolute/path form. + #[arg( + long = "contract-root", + required = true, + value_name = "MODEL_CONTRACT_ID=PATH", + value_parser = parse_contract_root, + help_heading = "Tensor storage options" + )] + contract_roots: Vec<(String, PathBuf)>, + + /// Verify complete shard hashes lazily in addition to per-tensor hashes. + #[arg(long, help_heading = "Tensor storage options")] verify_full_shards: bool, - #[arg(long, default_value = "lru", value_parser = ["lru", "resident"])] - gpu_cache_mode: String, - #[arg(long = "resident-manifest", value_parser = parse_contract_root)] + + /// Descriptor manifest to pin before readiness; required in resident mode. + #[arg( + long = "resident-manifest", + value_name = "MODEL_CONTRACT_ID=PATH", + value_parser = parse_contract_root, + help_heading = "Tensor storage options" + )] resident_manifests: Vec<(String, PathBuf)>, - #[arg(long, default_value_t = 256)] + + /// Number of resident descriptors resolved and uploaded per prewarm batch. + #[arg( + long, + value_name = "COUNT", + default_value_t = 256, + value_parser = parse_positive_usize, + help_heading = "Tensor storage options" + )] prewarm_batch_size: usize, + + /// Maximum accepted request frame size. + #[arg( + long, + value_name = "MB", + default_value_t = 64, + value_parser = parse_positive_usize, + help_heading = "Resource limits" + )] + max_request_mb: usize, + + /// Maximum query plus candidate tensor tokens accepted per request. + #[arg( + long, + value_name = "COUNT", + default_value_t = 1_000_000, + value_parser = parse_positive_usize, + help_heading = "Resource limits" + )] + max_batch_tokens: usize, + + /// Maximum decoded query plus candidate tensor bytes per request. + #[arg( + long, + value_name = "MB", + default_value_t = 1024, + value_parser = parse_positive_usize, + help_heading = "Resource limits" + )] + max_batch_mb: usize, + + /// Atomically create this file after the service socket is ready. + #[arg(long, value_name = "PATH", help_heading = "Process control")] + ready_file: Option, + + /// Write the daemon PID to this file and remove it during clean shutdown. + #[arg(long, value_name = "PATH", help_heading = "Process control")] + pid_file: Option, + + /// Maximum time to drain accepted work after SIGINT or SIGTERM. + #[arg( + long, + value_name = "MILLISECONDS", + default_value_t = 30_000, + value_parser = parse_positive_usize, + help_heading = "Process control" + )] + shutdown_grace_ms: usize, + + /// Exit after one accepted request. Intended for tests and smoke checks. + #[arg(long, hide = true)] + once: bool, } -#[derive(Clone)] +#[derive(Clone, Debug)] struct GpuMemory { device: i32, bytes: usize, } +#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] +enum CacheMode { + Lru, + Resident, +} + fn parse_gpu_memory(value: &str) -> Result { let (device, gb) = value .strip_prefix("cuda:") @@ -109,8 +326,40 @@ fn parse_mode(value: &str) -> Result { Ok(mode) } +fn parse_positive_usize(value: &str) -> Result { + let parsed = value + .parse::() + .map_err(|_| "value must be a positive integer".to_owned())?; + if parsed == 0 { + return Err("value must be a positive integer".to_owned()); + } + Ok(parsed) +} + +fn parse_gpu_block_kib(value: &str) -> Result { + let kib = parse_positive_usize(value)?; + if !(4..=1024).contains(&kib) || !kib.is_power_of_two() { + return Err("GPU block size must be a power of two from 4 to 1024 KiB".to_owned()); + } + Ok(kib) +} + fn main() -> Result<()> { let args = Args::parse(); + validate_args(&args)?; + let _pid_file = args + .pid_file + .as_deref() + .map(ManagedFile::create_pid) + .transpose()?; + let max_request_bytes = args + .max_request_mb + .checked_mul(MIB) + .ok_or_else(|| anyhow!("maximum request size overflow"))?; + let max_batch_bytes = args + .max_batch_mb + .checked_mul(MIB) + .ok_or_else(|| anyhow!("maximum tensor batch size overflow"))?; let mut seen_devices = std::collections::HashSet::new(); for specification in &args.gpu_memory_gb { if args.gpu_workspace_gb >= specification.bytes { @@ -144,13 +393,13 @@ fn main() -> Result<()> { }) .collect::>>()?; let mut engine = Engine::new(gpus, block_bytes, store)?; - if args.gpu_cache_mode == "resident" && args.resident_manifests.is_empty() { + if args.gpu_cache_mode == CacheMode::Resident && args.resident_manifests.is_empty() { bail!("resident GPU cache mode requires at least one resident manifest"); } - if args.gpu_cache_mode == "lru" && !args.resident_manifests.is_empty() { + if args.gpu_cache_mode == CacheMode::Lru && !args.resident_manifests.is_empty() { bail!("resident manifests are valid only in resident GPU cache mode"); } - if args.gpu_cache_mode == "resident" { + if args.gpu_cache_mode == CacheMode::Resident { let prewarm_started = Instant::now(); let descriptors = load_resident_manifests(&args.resident_manifests)?; engine.prewarm(&descriptors, args.prewarm_batch_size)?; @@ -164,99 +413,99 @@ fn main() -> Result<()> { }) ); } - remove_stale_socket(&args.socket)?; - let listener = UnixListener::bind(&args.socket) - .with_context(|| format!("cannot bind {}", args.socket.display()))?; - fs::set_permissions(&args.socket, fs::Permissions::from_mode(args.socket_mode))?; - println!( - "{}", - serde_json::json!({ - "event": "tilemaxsim_rust_ready", - "socket": args.socket, - "devices": args.gpu_memory_gb.iter().map(|item| serde_json::json!({ - "device": item.device, - "allocated_bytes": item.bytes, - })).collect::>(), - "workspace_bytes": args.gpu_workspace_gb, - "cache": engine.status_json(), - }) - ); - let mut accepted = 0_usize; - for connection in listener.incoming() { - let mut connection = connection?; - let started = Instant::now(); - let response = match read_request(&mut connection, args.max_request_bytes) { - Ok(frame) => { - let request_id = header_request_id(&frame); - match protocol::parse(&frame).and_then(|request| { - let request_id = request.request_id; - engine - .score(&request) - .map(|results| protocol::success(request_id, &results)) - }) { - Ok(response) => response, - Err(error) => protocol::failure(request_id, 3, &format!("{error:#}")), - } - } - Err(error) => protocol::failure(0, 1, &format!("{error:#}")), - }; - connection.write_all(&response)?; - println!( - "{}", - serde_json::json!({ - "event": "tilemaxsim_rust_request", - "elapsed_ms": started.elapsed().as_secs_f64() * 1000.0, - "cache": engine.status_json(), - }) - ); - accepted += 1; - if args.once && accepted == 1 { - break; - } - } - drop(listener); - match fs::remove_file(&args.socket) { - Ok(()) => {} - Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} - Err(error) => return Err(error.into()), - } - Ok(()) + let mut allowed_uids = args.allow_peer_uid.iter().copied().collect::>(); + // SAFETY: geteuid has no preconditions and cannot fail. + allowed_uids.insert(unsafe { libc::geteuid() }); + let allowed_gids = args.allow_peer_gid.iter().copied().collect(); + let server_config = ServerConfig { + socket: args.socket.clone(), + socket_mode: args.socket_mode, + request_timeout: Duration::from_millis(args.request_timeout_ms as u64), + max_request_bytes, + max_batch_tokens: args.max_batch_tokens, + max_batch_bytes, + max_inflight: args.max_inflight, + backlog: args.backlog, + max_queued_requests: args.max_queued_requests, + allowed_uids, + allowed_gids, + shutdown_grace: Duration::from_millis(args.shutdown_grace_ms as u64), + once: args.once, + }; + server::install_signal_handlers()?; + let bound = server::bind(&server_config)?; + let ready = serde_json::json!({ + "event": "tilemaxsim_rust_ready", + "schema_version": 1, + "pid": std::process::id(), + "version": env!("CARGO_PKG_VERSION"), + "socket": bound.path(), + "devices": args.gpu_memory_gb.iter().map(|item| serde_json::json!({ + "device": item.device, + "allocated_bytes": item.bytes, + })).collect::>(), + "workspace_bytes": args.gpu_workspace_gb, + "limits": { + "request_bytes": max_request_bytes, + "batch_tokens": args.max_batch_tokens, + "batch_bytes": max_batch_bytes, + "max_inflight": args.max_inflight, + "connection_backlog": args.backlog, + "gpu_queue": args.max_queued_requests, + "request_timeout_ms": args.request_timeout_ms, + }, + "cache": engine.status_json(), + }); + let ready_contents = serde_json::to_vec(&ready)?; + let ready_file = args + .ready_file + .as_deref() + .map(|path| ManagedFile::create_ready(path, &ready_contents)) + .transpose()?; + println!("{ready}"); + server::serve(engine, bound, server_config, ready_file) } -fn read_request(connection: &mut UnixStream, maximum: usize) -> Result> { - let mut header = [0_u8; HEADER_BYTES]; - connection.read_exact(&mut header)?; - let body_bytes = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) - .context("request body does not fit this host")?; - let total = HEADER_BYTES - .checked_add(body_bytes) - .ok_or_else(|| anyhow!("request length overflow"))?; - if total > maximum { - bail!("request exceeds byte limit"); +fn validate_args(args: &Args) -> Result<()> { + if args.max_inflight > 1024 { + bail!("--max-inflight cannot exceed 1024"); } - let mut frame = Vec::with_capacity(total); - frame.extend_from_slice(&header); - frame.resize(total, 0); - connection.read_exact(&mut frame[HEADER_BYTES..])?; - Ok(frame) -} - -fn header_request_id(frame: &[u8]) -> u64 { - if frame.len() < HEADER_BYTES { - 0 - } else { - u64::from_le_bytes(frame[8..16].try_into().unwrap()) + if args.backlog > 65_535 { + bail!("--backlog cannot exceed 65535"); } -} - -fn remove_stale_socket(path: &PathBuf) -> Result<()> { - let Ok(metadata) = fs::symlink_metadata(path) else { - return Ok(()); - }; - if !metadata.file_type().is_socket() { - bail!("refusing to remove non-socket path {}", path.display()); + if args.max_queued_requests > 65_535 { + bail!("--max-queued-requests cannot exceed 65535"); + } + if args.max_request_mb > 1024 { + bail!("--max-request-mb cannot exceed 1024"); + } + if args.max_batch_tokens > 1_000_000 { + bail!("--max-batch-tokens cannot exceed 1000000"); + } + if args.max_batch_mb > 1024 { + bail!("--max-batch-mb cannot exceed 1024"); + } + if args.request_timeout_ms > u32::MAX as usize { + bail!("--request-timeout-ms is too large"); + } + if args.shutdown_grace_ms > u32::MAX as usize { + bail!("--shutdown-grace-ms is too large"); + } + let paths = [ + ("--socket", Some(&args.socket)), + ("--pid-file", args.pid_file.as_ref()), + ("--ready-file", args.ready_file.as_ref()), + ]; + for (index, (left_name, left_path)) in paths.iter().enumerate() { + let Some(left_path) = left_path else { + continue; + }; + for (right_name, right_path) in &paths[index + 1..] { + if right_path.is_some_and(|right_path| right_path == *left_path) { + bail!("{left_name} and {right_name} must use different paths"); + } + } } - fs::remove_file(path)?; Ok(()) } @@ -319,3 +568,59 @@ fn load_resident_manifests(values: &[(String, PathBuf)]) -> Result", + "--max-request-mb ", + "--max-batch-tokens ", + "--max-batch-mb ", + "--request-timeout-ms ", + "--allow-peer-uid ", + "--ready-file ", + "allocates every configured GPU arena before creating the socket", + "EXAMPLES:", + ] { + assert!(help.contains(expected), "long help is missing {expected:?}"); + } + assert!(!help.contains("--max-request-bytes")); + } + + #[test] + fn memory_is_configured_in_gigabytes_and_block_size_is_bounded() { + let memory = parse_gpu_memory("cuda:2=1.5").unwrap(); + assert_eq!(memory.device, 2); + assert_eq!(memory.bytes, GIB + GIB / 2); + assert_eq!(parse_gpu_block_kib("32").unwrap(), 32); + assert!(parse_gpu_block_kib("3").is_err()); + assert!(parse_gpu_block_kib("48").is_err()); + assert!(parse_gpu_block_kib("2048").is_err()); + } + + #[test] + fn process_control_paths_must_not_alias_the_socket() { + let args = Args::try_parse_from([ + "vchord-tilemaxsimd", + "--socket", + "/tmp/vchord-test.sock", + "--gpu-memory-gb", + "0=1", + "--contract-root", + "model@1=/tmp", + "--ready-file", + "/tmp/vchord-test.sock", + ]) + .unwrap(); + let error = validate_args(&args).unwrap_err(); + assert!(format!("{error:#}").contains("must use different paths")); + } +} diff --git a/services/tilemaxsimd/src/protocol.rs b/services/tilemaxsimd/src/protocol.rs index d02cae01..4a37fdd1 100644 --- a/services/tilemaxsimd/src/protocol.rs +++ b/services/tilemaxsimd/src/protocol.rs @@ -1,3 +1,5 @@ +// Copyright (c) 2026 HuXinjing + use anyhow::{Result, anyhow, bail}; use std::collections::HashSet; @@ -25,6 +27,8 @@ pub struct Request { pub dtype: u8, pub query: Vec, pub candidates: Vec, + pub tensor_tokens: usize, + pub tensor_bytes: usize, } struct Reader<'a> { @@ -201,6 +205,8 @@ pub fn parse(frame: &[u8]) -> Result { dtype, query, candidates, + tensor_tokens: total_tokens, + tensor_bytes: total_bytes, }) } @@ -236,3 +242,57 @@ pub fn failure(request_id: u64, status: u32, message: &str) -> Vec { frame.extend_from_slice(message); frame } + +#[cfg(test)] +mod tests { + use super::*; + + fn empty_request() -> Vec { + let contract = b"health"; + let body_bytes = 4 + 4 + 4 + 1 + 1 + 2 + 4 + contract.len() + 2; + let mut frame = Vec::with_capacity(HEADER_BYTES + body_bytes); + frame.extend_from_slice(MAGIC); + frame.extend_from_slice(&VERSION_EXTERNAL.to_le_bytes()); + frame.extend_from_slice(&REQUEST_KIND.to_le_bytes()); + frame.extend_from_slice(&7_u64.to_le_bytes()); + frame.extend_from_slice(&(body_bytes as u64).to_le_bytes()); + frame.extend_from_slice(&1_u32.to_le_bytes()); + frame.extend_from_slice(&1_u32.to_le_bytes()); + frame.extend_from_slice(&0_u32.to_le_bytes()); + frame.push(2); + frame.push(1); + frame.extend_from_slice(&0_u16.to_le_bytes()); + frame.extend_from_slice(&(contract.len() as u32).to_le_bytes()); + frame.extend_from_slice(contract); + frame.extend_from_slice(&0_u16.to_le_bytes()); + frame + } + + #[test] + fn zero_candidate_probe_is_a_valid_bounded_request() { + let request = parse(&empty_request()).unwrap(); + assert_eq!(request.request_id, 7); + assert_eq!(request.query_rows, 1); + assert_eq!(request.dimension, 1); + assert_eq!(request.dtype, 2); + assert!(request.candidates.is_empty()); + assert_eq!(request.tensor_tokens, 1); + assert_eq!(request.tensor_bytes, 2); + } + + #[test] + fn arbitrary_short_frames_fail_without_panicking() { + // Fixed seed and LCG constants make malformed-frame coverage reproducible. + let mut state = 0x9e37_79b9_7f4a_7c15_u64; + for length in 0..512 { + let mut frame = vec![0_u8; length]; + for byte in &mut frame { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + *byte = (state >> 32) as u8; + } + let _ = parse(&frame); + } + } +} diff --git a/services/tilemaxsimd/src/server.rs b/services/tilemaxsimd/src/server.rs new file mode 100644 index 00000000..12b5b62e --- /dev/null +++ b/services/tilemaxsimd/src/server.rs @@ -0,0 +1,796 @@ +// Copyright (c) 2026 HuXinjing + +use crate::engine::Engine; +use crate::lifecycle::ManagedFile; +use crate::protocol::{self, HEADER_BYTES}; +use anyhow::{Context, Result, anyhow, bail}; +use serde_json::json; +use std::collections::HashSet; +use std::fs; +use std::io::{ErrorKind, Read, Write}; +use std::os::fd::{AsRawFd, FromRawFd, OwnedFd}; +use std::os::unix::ffi::OsStrExt; +use std::os::unix::fs::{FileTypeExt, MetadataExt, PermissionsExt}; +use std::os::unix::net::{UnixListener, UnixStream}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::mpsc::{Receiver, RecvTimeoutError, SyncSender, TrySendError, sync_channel}; +use std::sync::{Arc, Mutex}; +use std::thread::{self, JoinHandle}; +use std::time::{Duration, Instant}; + +const STATUS_INVALID_REQUEST: u32 = 1; +const STATUS_RESOURCE_LIMIT: u32 = 2; +const STATUS_COMPUTE_ERROR: u32 = 3; +const WAIT_QUANTUM: Duration = Duration::from_millis(25); +static SHUTDOWN_REQUESTED: AtomicBool = AtomicBool::new(false); + +#[derive(Clone, Debug)] +pub struct ServerConfig { + pub socket: PathBuf, + pub socket_mode: u32, + pub request_timeout: Duration, + pub max_request_bytes: usize, + pub max_batch_tokens: usize, + pub max_batch_bytes: usize, + pub max_inflight: usize, + pub backlog: usize, + pub max_queued_requests: usize, + pub allowed_uids: HashSet, + pub allowed_gids: HashSet, + pub shutdown_grace: Duration, + pub once: bool, +} + +pub struct BoundServer { + socket: SocketGuard, + listener: UnixListener, +} + +impl BoundServer { + pub fn path(&self) -> &PathBuf { + &self.socket.path + } + + fn close(mut self) -> Result<()> { + self.socket.remove()?; + drop(self.listener); + Ok(()) + } +} + +struct SocketGuard { + path: PathBuf, + device: u64, + inode: u64, + removed: bool, +} + +impl SocketGuard { + fn remove(&mut self) -> Result<()> { + if self.removed { + return Ok(()); + } + let metadata = match fs::symlink_metadata(&self.path) { + Ok(metadata) => metadata, + Err(error) if error.kind() == ErrorKind::NotFound => { + self.removed = true; + return Ok(()); + } + Err(error) => return Err(error.into()), + }; + if metadata.dev() != self.device || metadata.ino() != self.inode { + bail!( + "refusing to remove replaced service socket {}", + self.path.display() + ); + } + fs::remove_file(&self.path)?; + self.removed = true; + Ok(()) + } +} + +impl Drop for SocketGuard { + fn drop(&mut self) { + let _ = self.remove(); + } +} + +#[derive(Clone, Copy, Debug)] +struct PeerCredentials { + pid: i32, + uid: u32, + gid: u32, +} + +struct ConnectionTask { + stream: UnixStream, + accepted_at: Instant, + peer: PeerCredentials, +} + +struct WorkItem { + request: protocol::Request, + accepted_at: Instant, + queued_at: Instant, + deadline: Instant, + peer: PeerCredentials, + canceled: Arc, + reply: SyncSender>, +} + +pub fn install_signal_handlers() -> Result<()> { + SHUTDOWN_REQUESTED.store(false, Ordering::Relaxed); + let mut action = unsafe { std::mem::zeroed::() }; + action.sa_sigaction = shutdown_signal_handler as *const () as usize; + // SAFETY: `action.sa_mask` is a valid signal set owned by this function. + unsafe { libc::sigemptyset(&mut action.sa_mask) }; + action.sa_flags = 0; + // SAFETY: the handler only performs an async-signal-safe atomic store. + if unsafe { libc::sigaction(libc::SIGINT, &action, std::ptr::null_mut()) } != 0 + || unsafe { libc::sigaction(libc::SIGTERM, &action, std::ptr::null_mut()) } != 0 + { + return Err(std::io::Error::last_os_error()).context("cannot install signal handlers"); + } + Ok(()) +} + +extern "C" fn shutdown_signal_handler(_signal: libc::c_int) { + SHUTDOWN_REQUESTED.store(true, Ordering::Relaxed); +} + +pub fn bind(config: &ServerConfig) -> Result { + remove_stale_socket(&config.socket)?; + let listener = UnixListener::bind(&config.socket) + .with_context(|| format!("cannot bind {}", config.socket.display()))?; + let metadata = fs::symlink_metadata(&config.socket)?; + let socket = SocketGuard { + path: config.socket.clone(), + device: metadata.dev(), + inode: metadata.ino(), + removed: false, + }; + fs::set_permissions( + &config.socket, + fs::Permissions::from_mode(config.socket_mode), + )?; + let backlog = i32::try_from(config.backlog).context("socket backlog is too large")?; + // SAFETY: the listener owns a valid listening Unix socket descriptor. + if unsafe { libc::listen(listener.as_raw_fd(), backlog) } != 0 { + return Err(std::io::Error::last_os_error()).context("cannot set socket backlog"); + } + listener.set_nonblocking(true)?; + Ok(BoundServer { socket, listener }) +} + +pub fn serve( + mut engine: Engine, + bound: BoundServer, + config: ServerConfig, + mut ready_file: Option, +) -> Result<()> { + let (connection_tx, connection_rx) = sync_channel::(config.backlog); + let connection_rx = Arc::new(Mutex::new(connection_rx)); + let (work_tx, work_rx) = sync_channel::(config.max_queued_requests); + let scheduler = thread::Builder::new() + .name("vchord-tilemaxsim-gpu".to_owned()) + .spawn(move || scheduler_loop(&mut engine, work_rx))?; + + let mut workers = Vec::with_capacity(config.max_inflight); + for index in 0..config.max_inflight { + let connections = Arc::clone(&connection_rx); + let work = work_tx.clone(); + let worker_config = config.clone(); + workers.push( + thread::Builder::new() + .name(format!("vchord-tilemaxsim-io-{index}")) + .spawn(move || connection_loop(connections, work, worker_config))?, + ); + } + drop(connection_rx); + drop(work_tx); + + let mut accepted = 0_u64; + let mut runtime_failure = None; + while !SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + if scheduler.is_finished() { + runtime_failure = Some("GPU scheduler stopped unexpectedly"); + break; + } + if workers.iter().any(JoinHandle::is_finished) { + runtime_failure = Some("connection worker stopped unexpectedly"); + break; + } + match bound.listener.accept() { + Ok((stream, _)) => { + let accepted_at = Instant::now(); + let peer = match peer_credentials(&stream) { + Ok(peer) => peer, + Err(error) => { + reject_connection( + stream, + config.request_timeout, + STATUS_INVALID_REQUEST, + "cannot read peer credentials", + ); + log_rejection(None, "peer_credentials", &format!("{error:#}")); + continue; + } + }; + if !peer_is_allowed(peer, &config) { + reject_connection( + stream, + config.request_timeout, + STATUS_INVALID_REQUEST, + "Unix peer credentials are not authorized", + ); + log_rejection(Some(peer), "unauthorized_peer", "peer rejected"); + continue; + } + let task = ConnectionTask { + stream, + accepted_at, + peer, + }; + match connection_tx.try_send(task) { + Ok(()) => { + accepted += 1; + if config.once && accepted == 1 { + break; + } + } + Err(TrySendError::Full(task)) => { + reject_connection( + task.stream, + config.request_timeout, + STATUS_RESOURCE_LIMIT, + "connection backlog is full", + ); + log_rejection(Some(task.peer), "connection_backlog_full", "busy"); + } + Err(TrySendError::Disconnected(_)) => { + bail!("connection worker pool stopped unexpectedly"); + } + } + } + Err(error) if error.kind() == ErrorKind::WouldBlock => { + thread::sleep(WAIT_QUANTUM); + } + Err(error) if error.kind() == ErrorKind::Interrupted => {} + Err(error) => return Err(error).context("service socket accept failed"), + } + } + + if let Some(file) = ready_file.take() { + file.remove().context("cannot remove readiness file")?; + } + let socket_path = bound.path().clone(); + bound.close()?; + drop(connection_tx); + let deadline = Instant::now() + config.shutdown_grace; + let workers_finished = join_until(workers, deadline, "connection workers"); + let scheduler_finished = join_one_until(scheduler, deadline, "GPU scheduler"); + let drained = workers_finished && scheduler_finished; + println!( + "{}", + json!({ + "event": "tilemaxsim_rust_stopped", + "schema_version": 1, + "socket": socket_path, + "accepted": accepted, + "drained": drained, + "runtime_failure": runtime_failure, + }) + ); + if let Some(message) = runtime_failure { + bail!(message); + } + if !drained { + bail!("shutdown grace period expired before all work drained"); + } + Ok(()) +} + +fn connection_loop( + connections: Arc>>, + work: SyncSender, + config: ServerConfig, +) { + loop { + let task = match connections.lock() { + Ok(receiver) => receiver.recv(), + Err(_) => return, + }; + let Ok(task) = task else { + return; + }; + handle_connection(task, &work, &config); + } +} + +fn handle_connection(task: ConnectionTask, work: &SyncSender, config: &ServerConfig) { + let deadline = task.accepted_at + config.request_timeout; + let mut stream = task.stream; + let result = (|| -> Result<()> { + set_stream_deadline(&stream, deadline)?; + let Some(frame) = read_request(&mut stream, config.max_request_bytes)? else { + log_io_result(task.peer, 0, "readiness_probe", task.accepted_at, None); + return Ok(()); + }; + let request_id = header_request_id(&frame); + let request = match protocol::parse(&frame) { + Ok(request) => request, + Err(error) => { + write_response( + &mut stream, + protocol::failure(request_id, STATUS_INVALID_REQUEST, &format!("{error:#}")), + deadline, + )?; + log_io_result( + task.peer, + request_id, + "invalid_request", + task.accepted_at, + None, + ); + return Ok(()); + } + }; + if request.tensor_tokens > config.max_batch_tokens + || request.tensor_bytes > config.max_batch_bytes + { + write_response( + &mut stream, + protocol::failure( + request.request_id, + STATUS_RESOURCE_LIMIT, + "request exceeds configured tensor batch limits", + ), + deadline, + )?; + log_io_result( + task.peer, + request.request_id, + "batch_limit", + task.accepted_at, + None, + ); + return Ok(()); + } + if Instant::now() >= deadline { + write_response( + &mut stream, + protocol::failure( + request.request_id, + STATUS_COMPUTE_ERROR, + "request deadline expired before queueing", + ), + deadline, + )?; + log_io_result( + task.peer, + request.request_id, + "timeout", + task.accepted_at, + None, + ); + return Ok(()); + } + let canceled = Arc::new(AtomicBool::new(false)); + let (reply_tx, reply_rx) = sync_channel(1); + let request_id = request.request_id; + let item = WorkItem { + request, + accepted_at: task.accepted_at, + queued_at: Instant::now(), + deadline, + peer: task.peer, + canceled: Arc::clone(&canceled), + reply: reply_tx, + }; + match work.try_send(item) { + Ok(()) => {} + Err(TrySendError::Full(_)) => { + write_response( + &mut stream, + protocol::failure( + request_id, + STATUS_RESOURCE_LIMIT, + "GPU request queue is full", + ), + deadline, + )?; + log_io_result( + task.peer, + request_id, + "gpu_queue_full", + task.accepted_at, + None, + ); + return Ok(()); + } + Err(TrySendError::Disconnected(_)) => { + bail!("GPU scheduler stopped unexpectedly"); + } + } + loop { + let now = Instant::now(); + if now >= deadline { + canceled.store(true, Ordering::Relaxed); + let response = + protocol::failure(request_id, STATUS_COMPUTE_ERROR, "request deadline expired"); + let _ = write_response(&mut stream, response, deadline); + return Ok(()); + } + let wait = (deadline - now).min(WAIT_QUANTUM); + match reply_rx.recv_timeout(wait) { + Ok(response) => { + write_response(&mut stream, response, deadline)?; + return Ok(()); + } + Err(RecvTimeoutError::Timeout) => { + if peer_disconnected(&stream)? { + canceled.store(true, Ordering::Relaxed); + return Ok(()); + } + } + Err(RecvTimeoutError::Disconnected) => { + bail!("GPU scheduler dropped a request without a response"); + } + } + } + })(); + if let Err(error) = result { + let request_id = 0; + let status = if is_timeout_error(&error) { + "timeout" + } else { + "io_error" + }; + let response = protocol::failure(request_id, STATUS_COMPUTE_ERROR, &format!("{error:#}")); + let _ = write_response(&mut stream, response, deadline); + log_io_result( + task.peer, + request_id, + status, + task.accepted_at, + Some(&format!("{error:#}")), + ); + } +} + +fn scheduler_loop(engine: &mut Engine, work: Receiver) { + for item in work { + let compute_started = Instant::now(); + let queue_ms = (compute_started - item.queued_at).as_secs_f64() * 1000.0; + let candidate_count = item.request.candidates.len(); + let request_id = item.request.request_id; + let result = if compute_started >= item.deadline { + Err(anyhow!("request deadline expired in the GPU queue")) + } else if item.canceled.load(Ordering::Relaxed) { + Err(anyhow!("request canceled because the client disconnected")) + } else { + engine.score_until(&item.request, item.deadline, &item.canceled) + }; + let compute_ms = compute_started.elapsed().as_secs_f64() * 1000.0; + let (status, error, response) = match result { + Ok(results) => ("ok", None, protocol::success(request_id, &results)), + Err(error) => { + let message = format!("{error:#}"); + let status = if message.contains("deadline") { + "timeout" + } else if message.contains("canceled") { + "canceled" + } else { + "compute_error" + }; + let response = protocol::failure(request_id, STATUS_COMPUTE_ERROR, &message); + (status, Some(message), response) + } + }; + let client_present = item.reply.try_send(response).is_ok(); + println!( + "{}", + json!({ + "event": "tilemaxsim_rust_request", + "schema_version": 1, + "request_id": request_id, + "status": status, + "error": error, + "peer_pid": item.peer.pid, + "peer_uid": item.peer.uid, + "peer_gid": item.peer.gid, + "candidates": candidate_count, + "queue_ms": queue_ms, + "compute_ms": compute_ms, + "total_ms": item.accepted_at.elapsed().as_secs_f64() * 1000.0, + "client_present": client_present, + "cache": engine.status_json(), + }) + ); + } +} + +fn peer_credentials(stream: &UnixStream) -> Result { + let mut credentials = libc::ucred { + pid: 0, + uid: 0, + gid: 0, + }; + let mut length = std::mem::size_of::() as libc::socklen_t; + // SAFETY: `credentials` and `length` point to writable objects of the + // exact type and size required by Linux SO_PEERCRED. + let status = unsafe { + libc::getsockopt( + stream.as_raw_fd(), + libc::SOL_SOCKET, + libc::SO_PEERCRED, + std::ptr::addr_of_mut!(credentials).cast(), + &mut length, + ) + }; + if status != 0 { + return Err(std::io::Error::last_os_error()).context("SO_PEERCRED failed"); + } + if length as usize != std::mem::size_of::() { + bail!("SO_PEERCRED returned an unexpected structure size"); + } + Ok(PeerCredentials { + pid: credentials.pid, + uid: credentials.uid, + gid: credentials.gid, + }) +} + +fn peer_is_allowed(peer: PeerCredentials, config: &ServerConfig) -> bool { + config.allowed_uids.contains(&peer.uid) || config.allowed_gids.contains(&peer.gid) +} + +fn peer_disconnected(stream: &UnixStream) -> Result { + let mut descriptor = libc::pollfd { + fd: stream.as_raw_fd(), + events: libc::POLLHUP | libc::POLLERR, + revents: 0, + }; + // SAFETY: exactly one valid pollfd is supplied for a nonblocking poll. + let result = unsafe { libc::poll(std::ptr::addr_of_mut!(descriptor), 1, 0) }; + if result < 0 { + let error = std::io::Error::last_os_error(); + if error.kind() == ErrorKind::Interrupted { + return Ok(false); + } + return Err(error).context("cannot inspect client connection"); + } + Ok(descriptor.revents & (libc::POLLHUP | libc::POLLERR | libc::POLLNVAL) != 0) +} + +fn set_stream_deadline(stream: &UnixStream, deadline: Instant) -> Result<()> { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + bail!("request deadline expired"); + } + stream.set_read_timeout(Some(remaining))?; + stream.set_write_timeout(Some(remaining))?; + Ok(()) +} + +fn write_response(stream: &mut UnixStream, response: Vec, deadline: Instant) -> Result<()> { + set_stream_deadline(stream, deadline)?; + stream.write_all(&response)?; + Ok(()) +} + +fn read_request(connection: &mut UnixStream, maximum: usize) -> Result>> { + let mut header = [0_u8; HEADER_BYTES]; + let mut read = 0; + while read < header.len() { + match connection.read(&mut header[read..])? { + 0 if read == 0 => return Ok(None), + 0 => bail!("client closed with an incomplete request header"), + bytes => read += bytes, + } + } + let body_bytes = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) + .context("request body does not fit this host")?; + let total = HEADER_BYTES + .checked_add(body_bytes) + .ok_or_else(|| anyhow!("request length overflow"))?; + if total > maximum { + bail!("request exceeds byte limit"); + } + let mut frame = Vec::with_capacity(total); + frame.extend_from_slice(&header); + frame.resize(total, 0); + connection.read_exact(&mut frame[HEADER_BYTES..])?; + Ok(Some(frame)) +} + +fn header_request_id(frame: &[u8]) -> u64 { + if frame.len() < HEADER_BYTES { + 0 + } else { + u64::from_le_bytes(frame[8..16].try_into().unwrap()) + } +} + +fn reject_connection(mut stream: UnixStream, timeout: Duration, status: u32, message: &str) { + let deadline = Instant::now() + timeout; + let _ = write_response(&mut stream, protocol::failure(0, status, message), deadline); +} + +fn log_rejection(peer: Option, status: &str, error: &str) { + println!( + "{}", + json!({ + "event": "tilemaxsim_rust_request", + "schema_version": 1, + "request_id": 0, + "status": status, + "error": error, + "peer_pid": peer.map(|item| item.pid), + "peer_uid": peer.map(|item| item.uid), + "peer_gid": peer.map(|item| item.gid), + }) + ); +} + +fn log_io_result( + peer: PeerCredentials, + request_id: u64, + status: &str, + accepted_at: Instant, + error: Option<&str>, +) { + println!( + "{}", + json!({ + "event": "tilemaxsim_rust_request", + "schema_version": 1, + "request_id": request_id, + "status": status, + "error": error, + "peer_pid": peer.pid, + "peer_uid": peer.uid, + "peer_gid": peer.gid, + "total_ms": accepted_at.elapsed().as_secs_f64() * 1000.0, + }) + ); +} + +fn is_timeout_error(error: &anyhow::Error) -> bool { + error.chain().any(|source| { + source + .downcast_ref::() + .is_some_and(|error| { + matches!(error.kind(), ErrorKind::TimedOut | ErrorKind::WouldBlock) + }) + }) || format!("{error:#}").contains("deadline") +} + +fn remove_stale_socket(path: &Path) -> Result<()> { + let metadata = match fs::symlink_metadata(path) { + Ok(metadata) => metadata, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(error.into()), + }; + if !metadata.file_type().is_socket() { + bail!("refusing to remove non-socket path {}", path.display()); + } + if socket_is_active(path)? { + bail!("another daemon is already accepting on {}", path.display()); + } + let current = fs::symlink_metadata(path)?; + if current.dev() != metadata.dev() || current.ino() != metadata.ino() { + bail!("service socket changed while checking {}", path.display()); + } + fs::remove_file(path)?; + Ok(()) +} + +fn socket_is_active(path: &Path) -> Result { + let bytes = path.as_os_str().as_bytes(); + let mut address = unsafe { std::mem::zeroed::() }; + if bytes.contains(&0) || bytes.len() >= address.sun_path.len() { + bail!("invalid Unix socket path {}", path.display()); + } + address.sun_family = libc::AF_UNIX as libc::sa_family_t; + for (target, source) in address.sun_path.iter_mut().zip(bytes.iter().copied()) { + *target = source as libc::c_char; + } + // SAFETY: socket has no pointer arguments; ownership moves to OwnedFd on success. + let raw = unsafe { + libc::socket( + libc::AF_UNIX, + libc::SOCK_STREAM | libc::SOCK_NONBLOCK | libc::SOCK_CLOEXEC, + 0, + ) + }; + if raw < 0 { + return Err(std::io::Error::last_os_error()).context("cannot probe service socket"); + } + // SAFETY: `raw` is a new and uniquely owned descriptor. + let descriptor = unsafe { OwnedFd::from_raw_fd(raw) }; + let address_length = + (std::mem::offset_of!(libc::sockaddr_un, sun_path) + bytes.len() + 1) as libc::socklen_t; + // SAFETY: the initialized AF_UNIX address and its exact length are supplied. + if unsafe { + libc::connect( + descriptor.as_raw_fd(), + std::ptr::addr_of!(address).cast(), + address_length, + ) + } == 0 + { + return Ok(true); + } + let error = std::io::Error::last_os_error(); + match error.raw_os_error() { + Some(libc::ECONNREFUSED | libc::ENOENT) => Ok(false), + Some(libc::EAGAIN | libc::EINPROGRESS | libc::EALREADY) => Ok(true), + _ => Err(error) + .with_context(|| format!("cannot prove that socket {} is stale", path.display())), + } +} + +fn join_until(handles: Vec>, deadline: Instant, name: &str) -> bool { + let mut pending = handles; + while !pending.is_empty() && Instant::now() < deadline { + let mut index = 0; + while index < pending.len() { + if pending[index].is_finished() { + let handle = pending.swap_remove(index); + if handle.join().is_err() { + eprintln!("{name} thread panicked"); + } + } else { + index += 1; + } + } + if !pending.is_empty() { + thread::sleep(WAIT_QUANTUM); + } + } + if !pending.is_empty() { + eprintln!("shutdown grace period expired while draining {name}"); + return false; + } + true +} + +fn join_one_until(handle: JoinHandle<()>, deadline: Instant, name: &str) -> bool { + while !handle.is_finished() && Instant::now() < deadline { + thread::sleep(WAIT_QUANTUM); + } + if !handle.is_finished() { + eprintln!("shutdown grace period expired while draining {name}"); + return false; + } + if handle.join().is_err() { + eprintln!("{name} thread panicked"); + return false; + } + true +} + +#[cfg(test)] +mod tests { + use super::*; + + fn socket_path(name: &str) -> PathBuf { + std::env::temp_dir().join(format!("vchord-tilemaxsim-{name}-{}", std::process::id())) + } + + #[test] + fn stale_socket_is_removed_but_live_socket_is_preserved() { + let path = socket_path("stale-socket"); + let _ = fs::remove_file(&path); + let listener = UnixListener::bind(&path).unwrap(); + let error = remove_stale_socket(&path).unwrap_err(); + assert!(format!("{error:#}").contains("already accepting")); + assert!(path.exists()); + drop(listener); + remove_stale_socket(&path).unwrap(); + assert!(!path.exists()); + } +} diff --git a/services/tilemaxsimd/src/shard.rs b/services/tilemaxsimd/src/shard.rs index 1a60e3cb..75d7e8b6 100644 --- a/services/tilemaxsimd/src/shard.rs +++ b/services/tilemaxsimd/src/shard.rs @@ -1,3 +1,5 @@ +// Copyright (c) 2026 HuXinjing + use crate::cache::TinyLfu; use crate::protocol::Descriptor; use anyhow::{Context, Result, anyhow, bail}; diff --git a/src/index/vchordrq/scanners/maxsim/candidate.rs b/src/index/vchordrq/scanners/maxsim/candidate.rs index af9b45b0..9463e76c 100644 --- a/src/index/vchordrq/scanners/maxsim/candidate.rs +++ b/src/index/vchordrq/scanners/maxsim/candidate.rs @@ -1,16 +1,4 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 HuXinjing use super::profile; use crate::index::fetcher::pointer_to_kv; diff --git a/src/index/vchordrq/scanners/maxsim/exact.rs b/src/index/vchordrq/scanners/maxsim/exact.rs index 020e8c5e..d602b59f 100644 --- a/src/index/vchordrq/scanners/maxsim/exact.rs +++ b/src/index/vchordrq/scanners/maxsim/exact.rs @@ -1,12 +1,4 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the ELv2, which has specific restrictions. -// -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 HuXinjing use super::candidate::{HeapKey, PageCandidate}; use super::external::{ diff --git a/src/index/vchordrq/scanners/maxsim/external.rs b/src/index/vchordrq/scanners/maxsim/external.rs index 8b726da5..fcf90282 100644 --- a/src/index/vchordrq/scanners/maxsim/external.rs +++ b/src/index/vchordrq/scanners/maxsim/external.rs @@ -1,12 +1,4 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): This software is also available under the ELv2, -// which has specific restrictions. -// -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 HuXinjing use super::candidate::PageCandidate; use super::rerank::RerankError; diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs index af8c5455..d59f9039 100644 --- a/src/index/vchordrq/scanners/maxsim/gpu.rs +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -1,16 +1,4 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 HuXinjing use super::candidate::{HeapKey, PageCandidate}; use super::external::{ diff --git a/src/index/vchordrq/scanners/maxsim/profile.rs b/src/index/vchordrq/scanners/maxsim/profile.rs index e8a6dffe..6db76a63 100644 --- a/src/index/vchordrq/scanners/maxsim/profile.rs +++ b/src/index/vchordrq/scanners/maxsim/profile.rs @@ -1,16 +1,4 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 HuXinjing use std::cell::RefCell; use std::time::{Duration, Instant}; diff --git a/src/index/vchordrq/scanners/maxsim/rerank.rs b/src/index/vchordrq/scanners/maxsim/rerank.rs index 151ae4ea..994831a7 100644 --- a/src/index/vchordrq/scanners/maxsim/rerank.rs +++ b/src/index/vchordrq/scanners/maxsim/rerank.rs @@ -1,16 +1,4 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 HuXinjing use super::candidate::{HeapKey, PageCandidate}; use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; diff --git a/src/index/vchordrq/scanners/maxsim/search.rs b/src/index/vchordrq/scanners/maxsim/search.rs index 41ea6e8f..baff76f3 100644 --- a/src/index/vchordrq/scanners/maxsim/search.rs +++ b/src/index/vchordrq/scanners/maxsim/search.rs @@ -1,26 +1,13 @@ -// This software is licensed under a dual license model: -// -// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the terms of the AGPLv3. -// -// Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the Elastic License v2, which has specific restrictions. -// -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025-2026 TensorChord Inc. - -use super::MaxsimBuilder; +// Copyright (c) 2026 HuXinjing + use super::candidate::{HeapKey, PageCandidate}; use super::external::{ CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorSourceBinding, ExternalTensorStorage, resolve_external_tensor_source, validate_descriptor, }; use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; -use super::profile; use super::rerank::RerankError; +use super::{MaxsimBuilder, profile}; use crate::index::fetcher::{ Fetcher, FilterableTuple, HeapFetcher, Tuple, TupleAttribute, ctid_to_key, }; From 9bb0f50d616e15ed1af176768d1d611c233d661d Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:05:08 +0800 Subject: [PATCH 307/324] fix(ci): address audit and cross-platform failures Upgrade vulnerable lockfile dependencies, document the validator macro advisory exception, install daemon lint components, use aligned Miri test storage, and resolve current beta/cross-platform warning failures. --- .github/workflows/check.yml | 1 + Cargo.lock | 8 +++---- crates/simd/src/lib.rs | 1 - crates/vchordrq/src/tuples.rs | 22 ++++++++++++++------ crates/xtask/src/main.rs | 17 +++++++-------- deny.toml | 4 ++++ services/tilemaxsimd/Cargo.lock | 4 ++-- src/index/vchordrq/scanners/maxsim/gpu.rs | 1 + src/index/vchordrq/scanners/maxsim/rerank.rs | 4 ++++ 9 files changed, 40 insertions(+), 22 deletions(-) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 6c28ef71..3bb9d7c1 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -170,6 +170,7 @@ jobs: rm -rf /var/lib/apt/lists/* curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ | sh -s -- -y --profile minimal --default-toolchain stable + "$HOME/.cargo/bin/rustup" component add rustfmt clippy echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" - name: Checkout diff --git a/Cargo.lock b/Cargo.lock index 80722f2f..062d9136 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -98,9 +98,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.102" +version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" [[package]] name = "arrayvec" @@ -395,9 +395,9 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 601da07d..8f6f87ef 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -14,7 +14,6 @@ #![allow(unsafe_code)] #![cfg_attr(feature = "nightly_f16", feature(f16))] -#![cfg_attr(target_arch = "s390x", feature(s390x_target_feature))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] #![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc_feature_detection))] #![cfg_attr(target_arch = "powerpc64", feature(powerpc_target_feature))] diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index 85dee786..89f20735 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -1583,6 +1583,12 @@ mod tests { } } + fn aligned_copy(bytes: &[u8]) -> Vec { + let mut storage = vec![0_u64; bytes.len().div_ceil(size_of::())]; + storage.as_mut_bytes()[..bytes.len()].copy_from_slice(bytes); + storage + } + #[test] fn meta_statistics_reuse_padding_without_moving_fields() { assert_eq!(size_of::(), 56); @@ -1608,9 +1614,11 @@ mod tests { #[test] fn meta_statistics_roundtrip_full_48_bit_range() { for expected in [0, 1, u32::MAX as u64 + 1, MAX_INDEXED_VECTORS] { - let bytes = meta(Some(expected)).serialize(); + let serialized = meta(Some(expected)).serialize(); + let storage = aligned_copy(&serialized); + let bytes = &storage.as_bytes()[..serialized.len()]; assert_eq!( - MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + MetaTuple::deserialize_ref(bytes).indexed_vectors(), Some(expected) ); } @@ -1618,12 +1626,14 @@ mod tests { #[test] fn old_meta_padding_reads_as_missing_statistics_and_can_be_upgraded() { - let mut bytes = meta(None).serialize(); - assert_eq!(MetaTuple::deserialize_ref(&bytes).indexed_vectors(), None); + let serialized = meta(None).serialize(); + let mut storage = aligned_copy(&serialized); + let bytes = &mut storage.as_mut_bytes()[..serialized.len()]; + assert_eq!(MetaTuple::deserialize_ref(bytes).indexed_vectors(), None); - MetaTuple::deserialize_mut(&mut bytes).set_indexed_vectors(42); + MetaTuple::deserialize_mut(bytes).set_indexed_vectors(42); assert_eq!( - MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + MetaTuple::deserialize_ref(bytes).indexed_vectors(), Some(42) ); } diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index b2b5c65d..70a00854 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -242,32 +242,31 @@ fn parse(rustc_cfg: &RustcCfg, obj: impl AsRef) -> Result, Box eprintln!("Reading {obj:?}"); let contents = std::fs::read(obj)?; let object = object::File::parse(contents.as_slice())?; - let exports; - if rustc_cfg.is_macos { - exports = object + let exports = if rustc_cfg.is_macos { + object .exports()? .into_iter() .flat_map(|x| std::str::from_utf8(x.name())) .flat_map(|x| x.strip_prefix("_")) .filter(|x| x.starts_with("__pgrx_internals")) .map(str::to_string) - .collect(); + .collect() } else if rustc_cfg.is_emscripten { - exports = object + object .symbols() .flat_map(|x| x.name().ok()) .filter(|x| x.starts_with("__pgrx_internals")) .map(str::to_string) - .collect(); + .collect() } else { - exports = object + object .exports()? .into_iter() .flat_map(|x| std::str::from_utf8(x.name())) .filter(|x| x.starts_with("__pgrx_internals")) .map(str::to_string) - .collect(); - } + .collect() + }; Ok(exports) } diff --git a/deny.toml b/deny.toml index 0fa41117..ae84eb04 100644 --- a/deny.toml +++ b/deny.toml @@ -5,6 +5,10 @@ all-features = true ignore = [ "RUSTSEC-2021-0127", # serde_cbor is unmaintained "RUSTSEC-2024-0436", # paste - no longer maintained + # Build-time-only transitive dependency of validator_derive 0.20.0. The + # advisory has no maintained compatible release yet; remove this exception + # as soon as validator publishes a replacement. + "RUSTSEC-2026-0173", ] [licenses] diff --git a/services/tilemaxsimd/Cargo.lock b/services/tilemaxsimd/Cargo.lock index d55f046d..c9c5a177 100644 --- a/services/tilemaxsimd/Cargo.lock +++ b/services/tilemaxsimd/Cargo.lock @@ -54,9 +54,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.102" +version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" [[package]] name = "block-buffer" diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs index d59f9039..4702bf33 100644 --- a/src/index/vchordrq/scanners/maxsim/gpu.rs +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -8,6 +8,7 @@ use super::rerank::{CandidateTensorSource, ExactMaxsimBackend, RerankError, Rera use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; +#[cfg(unix)] use std::mem::size_of_val; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; diff --git a/src/index/vchordrq/scanners/maxsim/rerank.rs b/src/index/vchordrq/scanners/maxsim/rerank.rs index 994831a7..ac733495 100644 --- a/src/index/vchordrq/scanners/maxsim/rerank.rs +++ b/src/index/vchordrq/scanners/maxsim/rerank.rs @@ -65,6 +65,10 @@ pub(super) trait ExactMaxsimBackend { #[derive(Debug)] pub(super) enum RerankError { TensorMismatch, + #[allow( + dead_code, + reason = "constructed only by the dynamically selected external descriptor source" + )] ModelContractMismatch, InvalidDescriptor(&'static str), Registry(String), From cd802eed369ecd74d3b9a6399bc5433278c8c063 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:21:38 +0800 Subject: [PATCH 308/324] fix(ci): stabilize planner sampling checks --- src/index/vchordg/am/am_build.rs | 18 ++++++++++-------- src/index/vchordrq/am/am_build.rs | 18 ++++++++++-------- tests/vchordrq/recall.slt | 8 +++++++- 3 files changed, 27 insertions(+), 17 deletions(-) diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index 819d79be..a5ab6361 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -418,15 +418,17 @@ pub unsafe extern "C-unwind" fn vchordg_parallel_build_main( .cast::() .cast_const() }; - let heap_lockmode; - let index_lockmode; - if unsafe { !(*vchordgshared).isconcurrent } { - heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; + let (heap_lockmode, index_lockmode) = if unsafe { !(*vchordgshared).isconcurrent } { + ( + pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) } else { - heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; - } + ( + pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) + }; let heap = unsafe { pgrx::pg_sys::table_open((*vchordgshared).heaprelid, heap_lockmode) }; let index = unsafe { pgrx::pg_sys::index_open((*vchordgshared).indexrelid, index_lockmode) }; let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 95b4bc45..15a5d982 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -826,15 +826,17 @@ pub unsafe extern "C-unwind" fn vchordrq_parallel_build_main( .cast::() .cast_const() }; - let heap_lockmode; - let index_lockmode; - if unsafe { !(*vchordrqshared).isconcurrent } { - heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; + let (heap_lockmode, index_lockmode) = if unsafe { !(*vchordrqshared).isconcurrent } { + ( + pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) } else { - heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; - } + ( + pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) + }; let heap = unsafe { pgrx::pg_sys::table_open((*vchordrqshared).heaprelid, heap_lockmode) }; let index = unsafe { pgrx::pg_sys::index_open((*vchordrqshared).indexrelid, index_lockmode) }; let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; diff --git a/tests/vchordrq/recall.slt b/tests/vchordrq/recall.slt index 475046fa..8246f31a 100644 --- a/tests/vchordrq/recall.slt +++ b/tests/vchordrq/recall.slt @@ -79,6 +79,9 @@ SHOW vchordrq.query_sampling_enable; ---- on +statement ok +SET enable_seqscan = off; + statement ok SELECT * from t ORDER BY val <-> '[0.50, 0.25, 1.00]'; @@ -163,6 +166,9 @@ SELECT COUNT(*) from public.vchordrq_sampled_queries; ---- 3 +statement ok +RESET enable_seqscan; + statement ok RESET search_path; @@ -179,4 +185,4 @@ statement ok SELECT pg_reload_conf(); statement ok -DROP TABLE t, t_dim4, t_expr; \ No newline at end of file +DROP TABLE t, t_dim4, t_expr; From 87ceb97d63e842bc6f4b3099b7bc9d6ccace5a96 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:33:48 +0800 Subject: [PATCH 309/324] fix(xtask): skip linker symbols in ppc schema helper --- crates/xtask/src/main.rs | 41 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index 70a00854..f7905b19 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -290,7 +290,7 @@ fn generate( .exports()? .into_iter() .flat_map(|x| std::str::from_utf8(x.name())) - .filter(|x| !["_start", "_IO_stdin_used", "main"].contains(x)) + .filter(|x| should_stub_postmaster_export(x)) .map(str::to_string) .collect::>() } else { @@ -350,6 +350,45 @@ fn generate( Ok(command_stdout) } +fn should_stub_postmaster_export(symbol: &str) -> bool { + // These symbols are supplied by the C runtime or linker when the schema + // helper executable is linked. Defining a PostgreSQL stub for any of them + // creates duplicate symbols on some targets (notably powerpc64le). + !matches!( + symbol, + "_start" + | "_IO_stdin_used" + | "main" + | "__data_start" + | "data_start" + | "__bss_start" + | "_edata" + | "_end" + ) +} + +#[cfg(test)] +mod tests { + use super::should_stub_postmaster_export; + + #[test] + fn schema_helper_does_not_stub_crt_or_linker_symbols() { + for symbol in [ + "_start", + "_IO_stdin_used", + "main", + "__data_start", + "data_start", + "__bss_start", + "_edata", + "_end", + ] { + assert!(!should_stub_postmaster_export(symbol), "{symbol}"); + } + assert!(should_stub_postmaster_export("PostgresMain")); + } +} + fn install_by_copying( src: impl AsRef, dst: impl AsRef, From 69aeb5b480f4cb6db87dd71da87c1f4061ad821d Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Mon, 13 Jul 2026 18:25:44 +0800 Subject: [PATCH 310/324] feat(vchordrq): add exact TileMaxSim reranking --- .github/workflows/check.yml | 9 + Cargo.lock | 1 + Cargo.toml | 3 + crates/vchordrq/src/build.rs | 1 + crates/vchordrq/src/bulkdelete.rs | 28 +- crates/vchordrq/src/cost.rs | 8 +- crates/vchordrq/src/lib.rs | 6 + crates/vchordrq/src/maxsim_cost.rs | 206 +++ crates/vchordrq/src/statistics.rs | 30 + crates/vchordrq/src/tuples.rs | 151 +- crates/vchordrq/src/types.rs | 34 + devtools/test_tilemaxsim_reference_sidecar.py | 383 +++++ devtools/tilemaxsim_reference_sidecar.py | 780 +++++++++ services/Dockerfile.tilemaxsim | 12 + services/benchmark_tilemaxsim_cuda.py | 152 ++ services/build_tilemaxsim_tensor_cache.py | 235 +++ services/test_tilemaxsim_cuda_sidecar.py | 313 ++++ services/tilemaxsim_cuda_sidecar.py | 802 +++++++++ src/datatype/mod.rs | 14 + src/datatype/operators_halfvec.rs | 8 + src/datatype/operators_rabitq4.rs | 8 + src/datatype/operators_rabitq8.rs | 8 + src/datatype/operators_vector.rs | 8 + src/index/fetcher.rs | 108 +- src/index/gucs.rs | 139 ++ src/index/vchordrq/am/am_build.rs | 31 +- src/index/vchordrq/am/mod.rs | 89 +- src/index/vchordrq/dispatch.rs | 27 +- src/index/vchordrq/opclass.rs | 8 + src/index/vchordrq/scanners/maxsim.rs | 1345 ++++++++------- .../vchordrq/scanners/maxsim/candidate.rs | 192 +++ .../vchordrq/scanners/maxsim/external.rs | 551 ++++++ src/index/vchordrq/scanners/maxsim/gpu.rs | 1495 +++++++++++++++++ src/index/vchordrq/scanners/maxsim/rerank.rs | 315 ++++ src/index/vchordrq/scanners/maxsim/search.rs | 578 +++++++ src/index/vchordrq/scanners/mod.rs | 8 + src/sql/finalize.sql | 821 +++++++++ tests/vchordrq/cost_estimator.slt | 171 ++ tests/vchordrq/maxsim_correctness.slt | 221 +++ tests/vchordrq/maxsim_source_registry.slt | 341 ++++ 40 files changed, 8942 insertions(+), 698 deletions(-) create mode 100644 crates/vchordrq/src/maxsim_cost.rs create mode 100644 crates/vchordrq/src/statistics.rs create mode 100644 devtools/test_tilemaxsim_reference_sidecar.py create mode 100644 devtools/tilemaxsim_reference_sidecar.py create mode 100644 services/Dockerfile.tilemaxsim create mode 100644 services/benchmark_tilemaxsim_cuda.py create mode 100644 services/build_tilemaxsim_tensor_cache.py create mode 100644 services/test_tilemaxsim_cuda_sidecar.py create mode 100644 services/tilemaxsim_cuda_sidecar.py create mode 100644 src/index/vchordrq/scanners/maxsim/candidate.rs create mode 100644 src/index/vchordrq/scanners/maxsim/external.rs create mode 100644 src/index/vchordrq/scanners/maxsim/gpu.rs create mode 100644 src/index/vchordrq/scanners/maxsim/rerank.rs create mode 100644 src/index/vchordrq/scanners/maxsim/search.rs create mode 100644 tests/vchordrq/maxsim_correctness.slt create mode 100644 tests/vchordrq/maxsim_source_registry.slt diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index df084841..f4eeb1c1 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -102,6 +102,15 @@ jobs: run: | ! grep -P '\t' -r ./sql + - name: TileMaxSim Reference Sidecar + run: | + python3 -m unittest -v devtools/test_tilemaxsim_reference_sidecar.py + python3 -m py_compile \ + services/tilemaxsim_cuda_sidecar.py \ + services/benchmark_tilemaxsim_cuda.py \ + services/build_tilemaxsim_tensor_cache.py \ + services/test_tilemaxsim_cuda_sidecar.py + miri: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +miri')) || diff --git a/Cargo.lock b/Cargo.lock index f324f59e..80722f2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1937,6 +1937,7 @@ dependencies = [ "index", "index_accessor", "k_means", + "libc", "mimalloc", "paste", "pgrx", diff --git a/Cargo.toml b/Cargo.toml index 5868e4d6..32654e1f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,6 +55,9 @@ zerocopy.workspace = true [target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos")))'.dependencies] mimalloc = { version = "0.1.49", features = ["local_dynamic_tls"] } +[target.'cfg(unix)'.dependencies] +libc = "0.2.182" + [lints] workspace = true diff --git a/crates/vchordrq/src/build.rs b/crates/vchordrq/src/build.rs index a7371353..d622fbc6 100644 --- a/crates/vchordrq/src/build.rs +++ b/crates/vchordrq/src/build.rs @@ -119,6 +119,7 @@ pub fn build( height_of_root: structures.len() as u32, is_residual, rerank_in_heap: vchordrq_options.rerank_in_table, + indexed_vectors: Some(0), centroids_first: centroids.first(), vectors_first: vectors, centroid_prefetch: pointer_of_centroids diff --git a/crates/vchordrq/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs index ae54ec54..0aae88fe 100644 --- a/crates/vchordrq/src/bulkdelete.rs +++ b/crates/vchordrq/src/bulkdelete.rs @@ -25,7 +25,8 @@ pub fn bulkdelete( index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, -) where +) -> u64 +where R::Page: Page, { let meta_guard = index.read(0); @@ -38,6 +39,8 @@ pub fn bulkdelete( drop(meta_guard); + let mut live = 0_u64; + let step = |state: State| { let mut results = Vec::new(); for first in state { @@ -64,19 +67,21 @@ pub fn bulkdelete( while current != u32::MAX { check(); let read = index.read(current); - let flag = 'flag: { + let (flag, page_live) = 'scan: { + let mut page_live = 0_u64; for i in 1..=read.len() { let bytes = read.get(i).expect("data corruption"); let tuple = FrozenTuple::deserialize_ref(bytes); if let FrozenTupleReader::_0(tuple) = tuple { for p in tuple.payload().iter() { if Some(true) == p.map(&callback) { - break 'flag true; + break 'scan (true, 0); } + page_live += u64::from(p.is_some()); } } } - false + (false, page_live) }; if flag { drop(read); @@ -89,9 +94,12 @@ pub fn bulkdelete( if Some(true) == p.map(&callback) { *p = None; } + live += u64::from(p.is_some()); } } } + } else { + live += page_live; } current = directory.next().unwrap_or(u32::MAX); } @@ -101,16 +109,18 @@ pub fn bulkdelete( while current != u32::MAX { check(); let read = index.read(current); - let flag = 'flag: { + let (flag, page_live) = 'scan: { + let mut page_live = 0_u64; for i in 1..=read.len() { let bytes = read.get(i).expect("data corruption"); let tuple = AppendableTuple::deserialize_ref(bytes); let p = tuple.payload(); if Some(true) == p.map(&callback) { - break 'flag true; + break 'scan (true, 0); } + page_live += u64::from(p.is_some()); } - false + (false, page_live) }; if flag { drop(read); @@ -122,14 +132,18 @@ pub fn bulkdelete( if Some(true) == p.map(&callback) { *p = None; } + live += u64::from(p.is_some()); } current = write.get_opaque().next; } else { + live += page_live; current = read.get_opaque().next; } } } } + + live } pub fn bulkdelete_vectors( diff --git a/crates/vchordrq/src/cost.rs b/crates/vchordrq/src/cost.rs index 60579401..c8e40252 100644 --- a/crates/vchordrq/src/cost.rs +++ b/crates/vchordrq/src/cost.rs @@ -18,6 +18,7 @@ use index::relation::{Page, RelationRead}; pub struct Cost { pub dim: u32, pub cells: Vec, + pub indexed_vectors: Option, } #[must_use] @@ -27,8 +28,13 @@ pub fn cost(index: &R) -> Cost { let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dim = meta_tuple.dim(); let cells = meta_tuple.cells().to_vec(); + let indexed_vectors = meta_tuple.indexed_vectors(); drop(meta_guard); - Cost { dim, cells } + Cost { + dim, + cells, + indexed_vectors, + } } diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index d0efef05..b60c2e3b 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -24,9 +24,11 @@ mod freepages; mod insert; mod linked_vec; mod maintain; +mod maxsim_cost; mod prewarm; mod rerank; mod search; +mod statistics; mod tape; mod tape_writer; mod tuples; @@ -43,9 +45,13 @@ pub use cost::cost; pub use fast_heap::FastHeap; pub use insert::{InsertChooser, insert, insert_vector}; pub use maintain::{MaintainChooser, maintain}; +pub use maxsim_cost::{ + MaxsimCostBackend, MaxsimCostEstimate, MaxsimCostInput, estimate_maxsim_cost, +}; pub use prewarm::prewarm; pub use rerank::{how, rerank_heap, rerank_index}; pub use search::{default_search, maxsim_search}; +pub use statistics::set_indexed_vectors; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; diff --git a/crates/vchordrq/src/maxsim_cost.rs b/crates/vchordrq/src/maxsim_cost.rs new file mode 100644 index 00000000..ca843d10 --- /dev/null +++ b/crates/vchordrq/src/maxsim_cost.rs @@ -0,0 +1,206 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the terms of the ELv2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[derive(Clone, Copy, Debug)] +pub enum MaxsimCostBackend { + CoarseOnly, + CpuExact, + Gpu, + Auto, +} + +#[derive(Clone, Copy, Debug)] +pub struct MaxsimCostInput { + pub heap_rows: f64, + pub index_tokens: f64, + pub token_nodes_per_query: f64, + pub base_index_pages: f64, + pub dimension: u32, + pub element_bits: u32, + pub query_tokens: u32, + pub limit_tuples: Option, + pub filter_selectivity: f64, + pub candidate_limit: Option, + pub backend: MaxsimCostBackend, +} + +#[derive(Clone, Copy, Debug)] +pub struct MaxsimCostEstimate { + pub startup_cost: f64, + pub total_cost: f64, + pub selectivity: f64, + pub index_pages: f64, +} + +pub fn estimate_maxsim_cost(input: MaxsimCostInput) -> MaxsimCostEstimate { + let heap_rows = input.heap_rows.max(1.0); + let index_tokens = input.index_tokens.max(heap_rows); + let query_tokens = f64::from(input.query_tokens.max(1)); + let average_document_tokens = (index_tokens / heap_rows).clamp(1.0, 65_536.0); + let token_visits = input.token_nodes_per_query.max(1.0) * query_tokens; + + // We do not yet persist page-level candidate statistics. Until then, use a + // conservative occupancy estimate: token visits are sampled from the token + // index, and a page is a candidate when at least one of its average tokens + // is visited. This is intentionally bounded by the heap row count. + let token_visit_fraction = (token_visits / index_tokens).clamp(0.0, 1.0); + let candidate_probability = 1.0 - (1.0 - token_visit_fraction).powf(average_document_tokens); + let generated_pages = (heap_rows * candidate_probability).clamp(1.0, heap_rows); + + let filtered_limit = input + .limit_tuples + .map(|limit| limit.max(1.0) / input.filter_selectivity.clamp(1e-9, 1.0)); + let exact_candidate_count = input + .candidate_limit + .map_or(generated_pages, |limit| { + f64::from(limit).min(generated_pages) + }) + .clamp(1.0, heap_rows); + let returned_pages = match input.backend { + MaxsimCostBackend::CoarseOnly => filtered_limit + .unwrap_or(generated_pages) + .min(generated_pages), + MaxsimCostBackend::CpuExact | MaxsimCostBackend::Gpu | MaxsimCostBackend::Auto => { + exact_candidate_count + } + }; + + // Candidate generation and aggregation are eager in the current scanner, + // so their full work belongs to startup cost even when SQL has a small + // LIMIT. The constants are deliberately conservative placeholders until + // committed corpus benchmarks replace them with fitted values. + let search_cost = 0.001 * token_visits; + let aggregation_cost = 0.01 * token_visits + 0.05 * generated_pages; + let exact_components = exact_candidate_count + * average_document_tokens + * query_tokens + * f64::from(input.dimension.max(1)); + let tensor_bytes = exact_candidate_count + * average_document_tokens + * f64::from(input.dimension.max(1)) + * f64::from(input.element_bits.max(1)) + / 8.0; + let cpu_exact_cost = exact_candidate_count + exact_components * 1e-6; + let gpu_exact_cost = 5.0 + tensor_bytes * 1e-7 + exact_components * 5e-8; + let backend_cost = match input.backend { + MaxsimCostBackend::CoarseOnly => 0.0, + MaxsimCostBackend::CpuExact => cpu_exact_cost, + MaxsimCostBackend::Gpu => gpu_exact_cost, + // Price a small but nonzero fallback risk. Runtime still performs a + // complete CPU rerank on every GPU failure. + MaxsimCostBackend::Auto => gpu_exact_cost + 0.05 * cpu_exact_cost, + }; + let startup_cost = search_cost + aggregation_cost + backend_cost; + let total_cost = startup_cost + returned_pages; + let selectivity = (returned_pages / heap_rows).clamp(1e-9, 1.0); + let index_pages = + input.base_index_pages.max(1.0) * (1.0 + 0.25 * (query_tokens - 1.0).max(0.0)); + + MaxsimCostEstimate { + startup_cost, + total_cost, + selectivity, + index_pages, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn input(backend: MaxsimCostBackend) -> MaxsimCostInput { + MaxsimCostInput { + heap_rows: 34_054.0, + index_tokens: 25_438_338.0, + token_nodes_per_query: 10_000.0, + base_index_pages: 20_000.0, + dimension: 320, + element_bits: 16, + query_tokens: 32, + limit_tuples: Some(20.0), + filter_selectivity: 1.0, + candidate_limit: Some(256), + backend, + } + } + + #[test] + fn maxsim_is_never_zero_cost() { + for backend in [ + MaxsimCostBackend::CoarseOnly, + MaxsimCostBackend::CpuExact, + MaxsimCostBackend::Gpu, + MaxsimCostBackend::Auto, + ] { + let estimate = estimate_maxsim_cost(input(backend)); + assert!(estimate.startup_cost > 0.0); + assert!(estimate.total_cost >= estimate.startup_cost); + assert!((1e-9..=1.0).contains(&estimate.selectivity)); + assert!(estimate.index_pages >= 1.0); + } + } + + #[test] + fn query_token_count_increases_eager_work() { + let one = estimate_maxsim_cost(MaxsimCostInput { + query_tokens: 1, + ..input(MaxsimCostBackend::CpuExact) + }); + let many = estimate_maxsim_cost(MaxsimCostInput { + query_tokens: 64, + ..input(MaxsimCostBackend::CpuExact) + }); + assert!(many.startup_cost > one.startup_cost); + assert!(many.index_pages > one.index_pages); + } + + #[test] + fn exact_candidate_limit_bounds_rows_and_cost() { + let small = estimate_maxsim_cost(MaxsimCostInput { + candidate_limit: Some(128), + ..input(MaxsimCostBackend::CpuExact) + }); + let large = estimate_maxsim_cost(MaxsimCostInput { + candidate_limit: Some(2048), + ..input(MaxsimCostBackend::CpuExact) + }); + assert!(small.selectivity < large.selectivity); + assert!(small.startup_cost < large.startup_cost); + } + + #[test] + fn auto_prices_more_than_gpu_for_fallback_risk() { + let gpu = estimate_maxsim_cost(input(MaxsimCostBackend::Gpu)); + let auto = estimate_maxsim_cost(input(MaxsimCostBackend::Auto)); + assert!(auto.startup_cost > gpu.startup_cost); + } + + #[test] + fn missing_stats_remain_finite() { + let estimate = estimate_maxsim_cost(MaxsimCostInput { + heap_rows: -1.0, + index_tokens: 0.0, + token_nodes_per_query: 0.0, + base_index_pages: 0.0, + filter_selectivity: 0.0, + limit_tuples: None, + candidate_limit: None, + ..input(MaxsimCostBackend::CoarseOnly) + }); + assert!(estimate.startup_cost.is_finite()); + assert!(estimate.total_cost.is_finite()); + assert!(estimate.selectivity.is_finite()); + assert!(estimate.index_pages.is_finite()); + } +} diff --git a/crates/vchordrq/src/statistics.rs b/crates/vchordrq/src/statistics.rs new file mode 100644 index 00000000..f651bec5 --- /dev/null +++ b/crates/vchordrq/src/statistics.rs @@ -0,0 +1,30 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the terms of the ELv2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::tuples::{MetaTuple, WithWriter}; +use index::relation::{Page, RelationWrite}; + +/// Store the number of live vector nodes observed by the latest complete +/// build or vacuum pass. +/// +/// This is deliberately refreshed in bulk instead of on every insert. A +/// per-insert update would serialize all writers on the metapage, which is a +/// poor tradeoff for a planner statistic. Like PostgreSQL's relation +/// statistics, the value may be stale between maintenance passes. +pub fn set_indexed_vectors(index: &R, indexed_vectors: u64) { + let mut meta_guard = index.write(0, false); + let meta_bytes = meta_guard.get_mut(1).expect("data corruption"); + let mut meta_tuple = MetaTuple::deserialize_mut(meta_bytes); + meta_tuple.set_indexed_vectors(indexed_vectors); +} diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index 8438c1ae..85dee786 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -21,6 +21,8 @@ pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); const VERSION: u64 = 1001; +const STATISTICS_VERSION: u16 = 1; +const MAX_INDEXED_VECTORS: u64 = (1_u64 << 48) - 1; #[inline(always)] fn tag(source: &[u8]) -> Tag { @@ -55,12 +57,13 @@ struct MetaTupleHeader { rerank_in_heap: Bool, cells_s: u16, cells_e: u16, - _padding_0: [Padding; 2], + statistics_version: u16, centroids_first: u32, vectors_first_s: u16, vectors_first_e: u16, freepages_first: u32, - _padding_1: [Padding; 6], + indexed_vectors_low: u32, + indexed_vectors_high: u16, // tree centroid_prefetch_s: u16, centroid_prefetch_e: u16, @@ -69,11 +72,24 @@ struct MetaTupleHeader { first: u32, } +// Statistics deliberately replace the old 2-byte and 6-byte padding regions. +// Keep these assertions in non-test builds: changing any offset would require +// an index format version bump and REINDEX instead of the compatibility path. +const _: () = { + assert!(size_of::() == 56); + assert!(std::mem::offset_of!(MetaTupleHeader, statistics_version) == 22); + assert!(std::mem::offset_of!(MetaTupleHeader, centroids_first) == 24); + assert!(std::mem::offset_of!(MetaTupleHeader, indexed_vectors_low) == 36); + assert!(std::mem::offset_of!(MetaTupleHeader, indexed_vectors_high) == 40); + assert!(std::mem::offset_of!(MetaTupleHeader, centroid_prefetch_s) == 42); +}; + pub struct MetaTuple { pub dim: u32, pub height_of_root: u32, pub is_residual: bool, pub rerank_in_heap: bool, + pub indexed_vectors: Option, pub cells: Vec, pub centroids_first: u32, pub vectors_first: Vec, @@ -94,6 +110,7 @@ impl Tuple for MetaTuple { height_of_root, is_residual, rerank_in_heap, + indexed_vectors, cells, centroids_first, vectors_first, @@ -103,6 +120,12 @@ impl Tuple for MetaTuple { centroid_norm, first, } => { + if let Some(indexed_vectors) = indexed_vectors { + assert!( + *indexed_vectors <= MAX_INDEXED_VECTORS, + "indexed vector count exceeds the on-disk 48-bit limit" + ); + } buffer.extend((MAGIC as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); // cells @@ -136,17 +159,20 @@ impl Tuple for MetaTuple { rerank_in_heap: (*rerank_in_heap).into(), cells_s, cells_e, + statistics_version: indexed_vectors + .map(|_| STATISTICS_VERSION) + .unwrap_or(0), centroids_first: *centroids_first, vectors_first_s, vectors_first_e, freepages_first: *freepages_first, + indexed_vectors_low: indexed_vectors.unwrap_or(0) as u32, + indexed_vectors_high: (indexed_vectors.unwrap_or(0) >> 32) as u16, centroid_prefetch_s, centroid_prefetch_e, centroid_head: *centroid_head, centroid_norm: *centroid_norm, first: *first, - _padding_0: Default::default(), - _padding_1: Default::default(), } .as_bytes(), ); @@ -186,6 +212,28 @@ impl WithReader for MetaTuple { } } +impl WithWriter for MetaTuple { + type Writer<'a> = MetaTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> MetaTupleWriter<'_> { + let tag = tag(source); + match tag { + MAGIC => { + let mut checker = MutChecker::new(source); + let header: &mut MetaTupleHeader = checker.prefix(size_of::()); + if VERSION != header.version { + panic!( + "deserialization: bad version number; {}", + "after upgrading VectorChord, please use REINDEX to rebuild the index." + ); + } + MetaTupleWriter { header } + } + _ => panic!("deserialization: bad magic number"), + } + } +} + #[derive(Debug, Clone, Copy)] pub struct MetaTupleReader<'a> { header: &'a MetaTupleHeader, @@ -207,6 +255,16 @@ impl<'a> MetaTupleReader<'a> { pub fn rerank_in_heap(self) -> bool { self.header.rerank_in_heap.into() } + pub fn indexed_vectors(self) -> Option { + match self.header.statistics_version { + 0 => None, + STATISTICS_VERSION => Some( + u64::from(self.header.indexed_vectors_low) + | (u64::from(self.header.indexed_vectors_high) << 32), + ), + _ => panic!("deserialization: unsupported statistics version"), + } + } pub fn cells(self) -> &'a [u32] { self.cells } @@ -233,6 +291,23 @@ impl<'a> MetaTupleReader<'a> { } } +#[derive(Debug)] +pub struct MetaTupleWriter<'a> { + header: &'a mut MetaTupleHeader, +} + +impl MetaTupleWriter<'_> { + pub fn set_indexed_vectors(&mut self, indexed_vectors: u64) { + assert!( + indexed_vectors <= MAX_INDEXED_VECTORS, + "indexed vector count exceeds the on-disk 48-bit limit" + ); + self.header.statistics_version = STATISTICS_VERSION; + self.header.indexed_vectors_low = indexed_vectors as u32; + self.header.indexed_vectors_high = (indexed_vectors >> 32) as u16; + } +} + #[repr(C, align(8))] #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FreepagesTupleHeader { @@ -1485,3 +1560,71 @@ impl AppendableTupleWriter<'_> { &mut self.header.payload } } + +#[cfg(test)] +mod tests { + use super::*; + + fn meta(indexed_vectors: Option) -> MetaTuple { + MetaTuple { + dim: 3, + height_of_root: 1, + is_residual: false, + rerank_in_heap: false, + indexed_vectors, + cells: vec![4], + centroids_first: 1, + vectors_first: vec![2], + freepages_first: 3, + centroid_prefetch: vec![4], + centroid_head: 0, + centroid_norm: 1.0, + first: 5, + } + } + + #[test] + fn meta_statistics_reuse_padding_without_moving_fields() { + assert_eq!(size_of::(), 56); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, statistics_version), + 22 + ); + assert_eq!(std::mem::offset_of!(MetaTupleHeader, centroids_first), 24); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, indexed_vectors_low), + 36 + ); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, indexed_vectors_high), + 40 + ); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, centroid_prefetch_s), + 42 + ); + } + + #[test] + fn meta_statistics_roundtrip_full_48_bit_range() { + for expected in [0, 1, u32::MAX as u64 + 1, MAX_INDEXED_VECTORS] { + let bytes = meta(Some(expected)).serialize(); + assert_eq!( + MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + Some(expected) + ); + } + } + + #[test] + fn old_meta_padding_reads_as_missing_statistics_and_can_be_upgraded() { + let mut bytes = meta(None).serialize(); + assert_eq!(MetaTuple::deserialize_ref(&bytes).indexed_vectors(), None); + + MetaTuple::deserialize_mut(&mut bytes).set_indexed_vectors(42); + assert_eq!( + MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + Some(42) + ); + } +} diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index c41ba501..d9b85cbc 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -12,12 +12,14 @@ // // Copyright (c) 2025-2026 TensorChord Inc. +use distance::Distance; use serde::{Deserialize, Serialize}; use simd::f16; use validator::{Validate, ValidationError}; use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; +use vector::{VectorBorrowed, VectorOwned}; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] @@ -61,6 +63,38 @@ pub enum OwnedVector { Rabitq4(Rabitq4Owned), } +impl OwnedVector { + pub fn dim(&self) -> u32 { + match self { + Self::Vecf32(vector) => vector.as_borrowed().dim(), + Self::Vecf16(vector) => vector.as_borrowed().dim(), + Self::Rabitq8(vector) => vector.as_borrowed().dim(), + Self::Rabitq4(vector) => vector.as_borrowed().dim(), + } + } + + pub fn operator_dot(&self, rhs: &Self) -> Option { + if self.dim() != rhs.dim() { + return None; + } + match (self, rhs) { + (Self::Vecf32(lhs), Self::Vecf32(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Vecf16(lhs), Self::Vecf16(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Rabitq8(lhs), Self::Rabitq8(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Rabitq4(lhs), Self::Rabitq4(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + _ => None, + } + } +} + #[derive(Debug, Clone, Copy)] pub enum BorrowedVector<'a> { Vecf32(VectBorrowed<'a, f32>), diff --git a/devtools/test_tilemaxsim_reference_sidecar.py b/devtools/test_tilemaxsim_reference_sidecar.py new file mode 100644 index 00000000..ee5ce666 --- /dev/null +++ b/devtools/test_tilemaxsim_reference_sidecar.py @@ -0,0 +1,383 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the terms of the ELv2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +from __future__ import annotations + +import hashlib +import os +import socket +import stat +import struct +import tempfile +import threading +import time +import unittest +from pathlib import Path + +try: + from . import tilemaxsim_reference_sidecar as sidecar +except ImportError: + import tilemaxsim_reference_sidecar as sidecar + + +def request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + candidates: list[tuple[int, list[list[float]]]], +) -> bytes: + dimension = len(query[0]) + code = "f" if dtype == sidecar.DTYPE_F32 else "e" + body = bytearray( + sidecar.REQUEST_FIXED.pack( + dimension, + len(query), + len(candidates), + dtype, + sidecar.SCORING_SUM_QUERY_MAX_DOCUMENT_DOT, + 0, + ) + ) + body.extend(struct.pack(f"<{len(query) * dimension}{code}", *sum(query, []))) + for candidate_id, tensor in candidates: + body.extend(sidecar.CANDIDATE_FIXED.pack(candidate_id, len(tensor))) + body.extend(struct.pack(f"<{len(tensor) * dimension}{code}", *sum(tensor, []))) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.VERSION, + sidecar.REQUEST_KIND, + request_id, + len(body), + ) + + body + ) + + +def external_request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + model_contract_id: str, + candidates: list[tuple[int, str, list[list[float]]]], +) -> tuple[bytes, dict[str, bytes]]: + dimension = len(query[0]) + code = "f" if dtype == sidecar.DTYPE_F32 else "e" + contract = model_contract_id.encode() + body = bytearray( + sidecar.EXTERNAL_REQUEST_FIXED.pack( + dimension, + len(query), + len(candidates), + dtype, + sidecar.SCORING_SUM_QUERY_MAX_DOCUMENT_DOT, + 0, + len(contract), + ) + ) + body.extend(contract) + body.extend(struct.pack(f"<{len(query) * dimension}{code}", *sum(query, []))) + objects = {} + for candidate_id, tensor_ref, tensor in candidates: + payload = struct.pack(f"<{len(tensor) * dimension}{code}", *sum(tensor, [])) + objects[tensor_ref] = payload + reference = tensor_ref.encode() + checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}".encode() + body.extend( + sidecar.EXTERNAL_CANDIDATE_FIXED.pack( + candidate_id, len(tensor), len(reference), len(checksum) + ) + ) + body.extend(reference) + body.extend(checksum) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.EXTERNAL_VERSION, + sidecar.REQUEST_KIND, + request_id, + len(body), + ) + + body, + objects, + ) + + +def decode_response(frame: bytes) -> tuple[int, int, list[tuple[int, float]] | str]: + magic, version, kind, request_id, body_len = sidecar.HEADER.unpack_from(frame) + assert magic == sidecar.MAGIC + assert version in (sidecar.VERSION, sidecar.EXTERNAL_VERSION) + assert kind == sidecar.RESPONSE_KIND + assert len(frame) == sidecar.HEADER.size + body_len + status, count_or_length = sidecar.RESPONSE_FIXED.unpack_from( + frame, sidecar.HEADER.size + ) + offset = sidecar.HEADER.size + sidecar.RESPONSE_FIXED.size + if status: + return request_id, status, frame[offset : offset + count_or_length].decode() + results = [] + for _ in range(count_or_length): + results.append(sidecar.RESULT.unpack_from(frame, offset)) + offset += sidecar.RESULT.size + assert offset == len(frame) + return request_id, status, results + + +class ReferenceSidecarTest(unittest.TestCase): + def test_f32_exact_scores_and_opaque_ids(self) -> None: + frame = request_frame( + 41, + sidecar.DTYPE_F32, + [[1.0, 0.0], [0.0, 1.0]], + [ + (17, [[1.0, 0.0], [0.0, 1.0]]), + (3, [[0.5, 0.5]]), + ], + ) + request_id, status, results = decode_response(sidecar.process_frame(frame)) + + self.assertEqual(request_id, 41) + self.assertEqual(status, 0) + self.assertEqual(results, [(17, 2.0), (3, 1.0)]) + + def test_f16_exact_scores(self) -> None: + frame = request_frame( + 42, + sidecar.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + [(0, [[0.75, 0.0], [0.0, 0.5]])], + ) + _, status, results = decode_response(sidecar.process_frame(frame)) + + self.assertEqual(status, 0) + self.assertEqual(results, [(0, 1.25)]) + + def test_shared_raw_decoder_preserves_inline_payloads(self) -> None: + frame = request_frame( + 420, + sidecar.DTYPE_F16, + [[1.0, 0.0]], + [(11, [[0.5, 0.25]])], + ) + request = sidecar.parse_request_frame(frame) + self.assertIsInstance(request, sidecar.InlineTensorRequest) + assert isinstance(request, sidecar.InlineTensorRequest) + self.assertEqual(request.request_id, 420) + self.assertEqual(request.query_rows, 1) + self.assertEqual(request.dimension, 2) + self.assertEqual([item.candidate_id for item in request.candidates], [11]) + self.assertEqual(len(request.query_payload), 4) + self.assertEqual(len(request.candidates[0].payload), 4) + + def test_duplicate_id_and_non_finite_input_fail_closed(self) -> None: + duplicate = request_frame( + 43, + sidecar.DTYPE_F32, + [[1.0]], + [(7, [[1.0]]), (7, [[2.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(duplicate)) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("duplicate candidate ID", message) + + non_finite = request_frame( + 44, + sidecar.DTYPE_F32, + [[float("nan")]], + [(0, [[1.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(non_finite)) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("non-finite", message) + + def test_truncated_and_oversized_frames_fail_closed(self) -> None: + valid = request_frame(45, sidecar.DTYPE_F32, [[1.0]], [(0, [[1.0]])]) + _, status, message = decode_response(sidecar.process_frame(valid[:-1])) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("length mismatch", message) + + _, status, message = decode_response( + sidecar.process_frame(valid, sidecar.Limits(max_request_bytes=32)) + ) + self.assertEqual(status, sidecar.STATUS_RESOURCE_LIMIT) + self.assertIn("byte limit", message) + + def test_header_reserved_trailing_and_token_limit_fail_closed(self) -> None: + valid = request_frame(47, sidecar.DTYPE_F32, [[1.0]], [(9, [[1.0]])]) + invalid_frames = [] + for offset, value in ((0, 0), (4, 3), (6, 2)): + invalid = bytearray(valid) + invalid[offset] = value + invalid_frames.append(bytes(invalid)) + + reserved = bytearray(valid) + reserved[sidecar.HEADER.size + 14] = 1 + invalid_frames.append(bytes(reserved)) + + trailing = bytearray(valid) + trailing.extend(b"x") + struct.pack_into(" None: + frame, objects = external_request_frame( + 48, + sidecar.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "colqwen@immutable-revision", + [ + (77, "object://fixture/page-1", [[1.0, 0.0], [0.0, 1.0]]), + (4, "object://fixture/page-2", [[0.5, 0.5]]), + ], + ) + seen = [] + + def resolver(request: sidecar.ExternalTensorRequest) -> bytes: + seen.append(request) + return objects[request.tensor_ref] + + request_id, status, results = decode_response( + sidecar.process_frame(frame, resolver=resolver) + ) + + self.assertEqual(request_id, 48) + self.assertEqual(status, 0) + self.assertEqual(results, [(77, 2.0), (4, 1.0)]) + self.assertEqual( + {request.model_contract_id for request in seen}, + {"colqwen@immutable-revision"}, + ) + self.assertEqual({request.tensor_ref for request in seen}, set(objects)) + + parsed = sidecar.parse_request_frame(frame) + self.assertIsInstance(parsed, sidecar.ParsedExternalTensorRequest) + assert isinstance(parsed, sidecar.ParsedExternalTensorRequest) + self.assertEqual(parsed.model_contract_id, "colqwen@immutable-revision") + self.assertEqual( + [candidate.candidate_id for candidate in parsed.candidates], [77, 4] + ) + + def test_external_v2_fails_closed_without_resolver_or_on_checksum_mismatch( + self, + ) -> None: + frame, objects = external_request_frame( + 49, + sidecar.DTYPE_F32, + [[1.0]], + "contract@1", + [(0, "object://immutable/page", [[2.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(frame)) + self.assertEqual(status, sidecar.STATUS_COMPUTE_ERROR) + self.assertIn("resolver is not configured", message) + + def corrupt_resolver(request: sidecar.ExternalTensorRequest) -> bytes: + return objects[request.tensor_ref] + b"x" + + _, status, message = decode_response( + sidecar.process_frame(frame, resolver=corrupt_resolver) + ) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("byte length", message) + + def test_external_v2_validates_complete_control_frame_before_resolution( + self, + ) -> None: + frame, objects = external_request_frame( + 50, + sidecar.DTYPE_F32, + [[1.0, 0.0]], + "contract@1", + [(0, "object://immutable/page", [[1.0, 0.0]])], + ) + invalid = bytearray(frame) + contract_length = len("contract@1") + candidate_offset = ( + sidecar.HEADER.size + + sidecar.EXTERNAL_REQUEST_FIXED.size + + contract_length + + 8 + ) + reference_offset = candidate_offset + sidecar.EXTERNAL_CANDIDATE_FIXED.size + invalid[reference_offset] = 0 + called = False + + def resolver(request: sidecar.ExternalTensorRequest) -> bytes: + nonlocal called + called = True + return objects[request.tensor_ref] + + _, status, message = decode_response( + sidecar.process_frame(bytes(invalid), resolver=resolver) + ) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("control characters", message) + self.assertFalse(called) + + def test_unix_socket_end_to_end(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "tilemaxsim.sock" + thread = threading.Thread( + target=sidecar.serve, + args=(path, sidecar.Limits()), + kwargs={"once": True}, + daemon=True, + ) + thread.start() + for _ in range(100): + if path.exists() and stat.S_IMODE(path.stat().st_mode) == 0o600: + break + time.sleep(0.01) + else: + self.fail("sidecar socket was not created with mode 0600") + self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600) + + frame = request_frame( + 46, + sidecar.DTYPE_F32, + [[1.0, 0.0], [0.0, 1.0]], + [(5, [[1.0, 0.0], [0.0, 1.0]])], + ) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(path)) + connection.sendall(frame) + header = connection.recv(sidecar.HEADER.size) + while len(header) < sidecar.HEADER.size: + header += connection.recv(sidecar.HEADER.size - len(header)) + body_len = sidecar.HEADER.unpack(header)[4] + body = b"" + while len(body) < body_len: + body += connection.recv(body_len - len(body)) + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + + request_id, status, results = decode_response(header + body) + self.assertEqual(request_id, 46) + self.assertEqual(status, 0) + self.assertEqual(results, [(5, 2.0)]) + self.assertFalse(path.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/devtools/tilemaxsim_reference_sidecar.py b/devtools/tilemaxsim_reference_sidecar.py new file mode 100644 index 00000000..8d9f7ee1 --- /dev/null +++ b/devtools/tilemaxsim_reference_sidecar.py @@ -0,0 +1,780 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the terms of the ELv2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""CPU reference implementation of the VectorChord TileMaxSim IPC sidecar. + +This executable is intentionally simple and single-threaded. It is a protocol +oracle and end-to-end development aid, not a production or GPU implementation. +The CLI serves inline v1 requests. External-descriptor v2 requests require an +explicit resolver injected by a test or embedding application and fail closed +when no resolver is configured. +""" + +from __future__ import annotations + +import argparse +import hashlib +import hmac +import math +import os +import signal +import socket +import stat +import struct +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterable + +MAGIC = b"VCTM" +VERSION = 1 +EXTERNAL_VERSION = 2 +REQUEST_KIND = 1 +RESPONSE_KIND = 2 +SCORING_SUM_QUERY_MAX_DOCUMENT_DOT = 1 +DTYPE_F32 = 1 +DTYPE_F16 = 2 + +HEADER = struct.Struct("<4sHHQQ") +REQUEST_FIXED = struct.Struct(" None: + super().__init__(message) + self.status = status + + +@dataclass(frozen=True) +class Limits: + max_request_bytes: int = 64 * 1024 * 1024 + max_batch_tokens: int = 1_000_000 + max_tensor_bytes: int = 1024 * 1024 * 1024 + max_candidates: int = 65_536 + + +@dataclass(frozen=True) +class ExternalTensorRequest: + model_contract_id: str + tensor_ref: str + rows: int + dimension: int + dtype: int + checksum: str + + +@dataclass(frozen=True) +class InlineTensorCandidate: + candidate_id: int + rows: int + payload: bytes + + +@dataclass(frozen=True) +class InlineTensorRequest: + request_id: int + dimension: int + query_rows: int + dtype: int + query_payload: bytes + candidates: tuple[InlineTensorCandidate, ...] + + +@dataclass(frozen=True) +class ExternalTensorCandidate: + candidate_id: int + descriptor: ExternalTensorRequest + + +@dataclass(frozen=True) +class ParsedExternalTensorRequest: + request_id: int + dimension: int + query_rows: int + dtype: int + model_contract_id: str + query_payload: bytes + candidates: tuple[ExternalTensorCandidate, ...] + + +ParsedRequest = InlineTensorRequest | ParsedExternalTensorRequest + + +# A reference resolver returns the exact row-major scalar bytes whose SHA-256 +# digest is stored in the descriptor. Production resolvers may adapt richer +# immutable object formats before returning this canonical tensor payload. +ExternalTensorResolver = Callable[[ExternalTensorRequest], bytes] + + +class Reader: + def __init__(self, payload: bytes) -> None: + self.payload = payload + self.offset = 0 + + def take(self, count: int) -> bytes: + end = self.offset + count + if count < 0 or end > len(self.payload): + raise SidecarError(STATUS_INVALID_REQUEST, "truncated request") + chunk = self.payload[self.offset : end] + self.offset = end + return chunk + + def unpack(self, layout: struct.Struct) -> tuple: + return layout.unpack(self.take(layout.size)) + + def finish(self) -> None: + if self.offset != len(self.payload): + raise SidecarError(STATUS_INVALID_REQUEST, "trailing request bytes") + + +def checked_elements(rows: int, dimension: int) -> int: + if rows <= 0 or dimension <= 0: + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor rows and dimension must be positive" + ) + elements = rows * dimension + if elements > (1 << 63) - 1: + raise SidecarError(STATUS_RESOURCE_LIMIT, "tensor shape is too large") + return elements + + +def dtype_size(dtype: int) -> int: + if dtype == DTYPE_F32: + return 4 + if dtype == DTYPE_F16: + return 2 + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported tensor dtype") + + +def checked_tensor_bytes(rows: int, dimension: int, dtype: int) -> int: + return checked_elements(rows, dimension) * dtype_size(dtype) + + +def validate_finite_tensor_payload( + payload: bytes, rows: int, dimension: int, dtype: int +) -> None: + expected = checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor byte length does not match its shape" + ) + code = "f" if dtype == DTYPE_F32 else "e" + try: + values = struct.iter_unpack(f"<{code}", payload) + if any(not math.isfinite(value[0]) for value in values): + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor contains non-finite value" + ) + except struct.error as error: + raise SidecarError(STATUS_INVALID_REQUEST, str(error)) from error + + +def read_text(reader: Reader, length: int, maximum: int, field: str) -> str: + if length <= 0 or length > maximum: + raise SidecarError(STATUS_RESOURCE_LIMIT, f"invalid {field} length") + try: + value = reader.take(length).decode("utf-8") + except UnicodeDecodeError as error: + raise SidecarError(STATUS_INVALID_REQUEST, f"{field} is not UTF-8") from error + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise SidecarError( + STATUS_INVALID_REQUEST, f"{field} contains control characters" + ) + return value + + +def read_tensor( + reader: Reader, rows: int, dimension: int, dtype: int +) -> list[tuple[float, ...]]: + elements = checked_elements(rows, dimension) + if dtype == DTYPE_F32: + code = "f" + element_size = 4 + elif dtype == DTYPE_F16: + code = "e" + element_size = 2 + else: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported tensor dtype") + raw = reader.take(elements * element_size) + try: + values = struct.unpack(f"<{elements}{code}", raw) + except struct.error as error: + raise SidecarError(STATUS_INVALID_REQUEST, str(error)) from error + if not all(math.isfinite(value) for value in values): + raise SidecarError(STATUS_INVALID_REQUEST, "tensor contains non-finite value") + return [ + tuple(values[offset : offset + dimension]) + for offset in range(0, elements, dimension) + ] + + +def decode_resolved_tensor( + payload: bytes, request: ExternalTensorRequest +) -> list[tuple[float, ...]]: + expected_bytes = checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + if len(payload) != expected_bytes: + raise SidecarError( + STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match descriptor", + ) + expected_checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}" + if not hmac.compare_digest(expected_checksum, request.checksum): + raise SidecarError(STATUS_INVALID_REQUEST, "resolved tensor checksum mismatch") + reader = Reader(payload) + tensor = read_tensor(reader, request.rows, request.dimension, request.dtype) + reader.finish() + return tensor + + +def tilemaxsim( + query: list[tuple[float, ...]], document: list[tuple[float, ...]] +) -> float: + if not query or not document: + raise SidecarError(STATUS_INVALID_REQUEST, "tensor must not be empty") + score = 0.0 + try: + for query_vector in query: + best = -math.inf + for document_vector in document: + dot = math.fsum( + left * right + for left, right in zip(query_vector, document_vector, strict=True) + ) + best = max(best, dot) + score += best + except (OverflowError, ValueError) as error: + raise SidecarError(STATUS_COMPUTE_ERROR, str(error)) from error + if not math.isfinite(score): + raise SidecarError(STATUS_COMPUTE_ERROR, "TileMaxSim result is non-finite") + return score + + +def success_response( + request_id: int, + results: Iterable[tuple[int, float]], + version: int = VERSION, +) -> bytes: + results = list(results) + body = bytearray(RESPONSE_FIXED.pack(0, len(results))) + for candidate_id, similarity in results: + body.extend(RESULT.pack(candidate_id, similarity)) + return HEADER.pack(MAGIC, version, RESPONSE_KIND, request_id, len(body)) + body + + +def error_response( + request_id: int, + status_code: int, + message: str, + version: int = VERSION, +) -> bytes: + encoded = message.encode("utf-8", errors="replace")[:MAX_ERROR_BYTES] + body = ERROR_FIXED.pack(status_code or STATUS_COMPUTE_ERROR, len(encoded)) + encoded + return HEADER.pack(MAGIC, version, RESPONSE_KIND, request_id, len(body)) + body + + +def validate_request_fixed( + dimension: int, + query_rows: int, + candidate_count: int, + dtype: int, + scoring: int, + reserved: int, + limits: Limits, +) -> None: + if dimension == 0 or dimension > 60_000: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid tensor dimension") + if query_rows == 0: + raise SidecarError(STATUS_INVALID_REQUEST, "query tensor is empty") + if candidate_count > limits.max_candidates: + raise SidecarError(STATUS_RESOURCE_LIMIT, "too many candidates") + dtype_size(dtype) + if scoring != SCORING_SUM_QUERY_MAX_DOCUMENT_DOT: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported scoring function") + if reserved != 0: + raise SidecarError(STATUS_INVALID_REQUEST, "reserved field must be zero") + if query_rows > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if checked_tensor_bytes(query_rows, dimension, dtype) > limits.max_tensor_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit") + + +def process_inline_request(reader: Reader, limits: Limits) -> list[tuple[int, float]]: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + ) = reader.unpack(REQUEST_FIXED) + validate_request_fixed( + dimension, query_rows, candidate_count, dtype, scoring, reserved, limits + ) + + query = read_tensor(reader, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = checked_tensor_bytes(query_rows, dimension, dtype) + candidates: list[tuple[int, list[tuple[float, ...]]]] = [] + candidate_ids: set[int] = set() + for _ in range(candidate_count): + candidate_id, rows = reader.unpack(CANDIDATE_FIXED) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + candidates.append((candidate_id, read_tensor(reader, rows, dimension, dtype))) + reader.finish() + return [ + (candidate_id, tilemaxsim(query, document)) + for candidate_id, document in candidates + ] + + +def process_external_request( + reader: Reader, + limits: Limits, + resolver: ExternalTensorResolver | None, +) -> list[tuple[int, float]]: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + validate_request_fixed( + dimension, query_rows, candidate_count, dtype, scoring, reserved, limits + ) + model_contract_id = read_text(reader, contract_length, 512, "model contract") + query = read_tensor(reader, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = checked_tensor_bytes(query_rows, dimension, dtype) + descriptors: list[tuple[int, ExternalTensorRequest]] = [] + candidate_ids: set[int] = set() + for _ in range(candidate_count): + candidate_id, rows, reference_length, checksum_length = reader.unpack( + EXTERNAL_CANDIDATE_FIXED + ) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + tensor_ref = read_text(reader, reference_length, 4096, "tensor reference") + checksum = read_text(reader, checksum_length, 512, "tensor checksum") + digest = checksum.removeprefix("sha256:") + if ( + not checksum.startswith("sha256:") + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise SidecarError( + STATUS_INVALID_REQUEST, + "tensor checksum must be a lowercase sha256 digest", + ) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + descriptors.append( + ( + candidate_id, + ExternalTensorRequest( + model_contract_id=model_contract_id, + tensor_ref=tensor_ref, + rows=rows, + dimension=dimension, + dtype=dtype, + checksum=checksum, + ), + ) + ) + reader.finish() + + # Validate the complete control frame before performing any external I/O. + if resolver is None: + raise SidecarError( + STATUS_COMPUTE_ERROR, "external tensor resolver is not configured" + ) + results = [] + for candidate_id, descriptor in descriptors: + payload = resolver(descriptor) + if not isinstance(payload, bytes): + raise SidecarError( + STATUS_COMPUTE_ERROR, + "external tensor resolver returned a non-bytes value", + ) + document = decode_resolved_tensor(payload, descriptor) + results.append((candidate_id, tilemaxsim(query, document))) + return results + + +def parse_request_frame( + frame: bytes, + limits: Limits = Limits(), + *, + validate_finite: bool = True, +) -> ParsedRequest: + """Decode and validate a request without resolving or computing tensors. + + The production CUDA sidecar uses this function so the protocol oracle and + deployable executor share one strict wire decoder. All control data and the + inline query are validated before a v2 resolver performs external I/O. + """ + + if len(frame) < HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "truncated frame header") + magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) + if magic != MAGIC: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") + if version not in (VERSION, EXTERNAL_VERSION): + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") + if kind != REQUEST_KIND: + raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") + if body_len != len(frame) - HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "request length mismatch") + if len(frame) > limits.max_request_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds byte limit") + + reader = Reader(frame[HEADER.size :]) + if version == VERSION: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + ) = reader.unpack(REQUEST_FIXED) + validate_request_fixed( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + limits, + ) + query_payload = reader.take(checked_tensor_bytes(query_rows, dimension, dtype)) + if validate_finite: + validate_finite_tensor_payload(query_payload, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = len(query_payload) + candidate_ids: set[int] = set() + candidates = [] + for _ in range(candidate_count): + candidate_id, rows = reader.unpack(CANDIDATE_FIXED) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + payload_bytes = checked_tensor_bytes(rows, dimension, dtype) + total_tokens += rows + total_tensor_bytes += payload_bytes + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + payload = reader.take(payload_bytes) + if validate_finite: + validate_finite_tensor_payload(payload, rows, dimension, dtype) + candidates.append(InlineTensorCandidate(candidate_id, rows, payload)) + reader.finish() + return InlineTensorRequest( + request_id, + dimension, + query_rows, + dtype, + query_payload, + tuple(candidates), + ) + + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + validate_request_fixed( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + limits, + ) + model_contract_id = read_text(reader, contract_length, 512, "model contract") + query_payload = reader.take(checked_tensor_bytes(query_rows, dimension, dtype)) + if validate_finite: + validate_finite_tensor_payload(query_payload, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = len(query_payload) + candidate_ids = set() + candidates = [] + for _ in range(candidate_count): + candidate_id, rows, reference_length, checksum_length = reader.unpack( + EXTERNAL_CANDIDATE_FIXED + ) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + tensor_ref = read_text(reader, reference_length, 4096, "tensor reference") + checksum = read_text(reader, checksum_length, 512, "tensor checksum") + digest = checksum.removeprefix("sha256:") + if ( + not checksum.startswith("sha256:") + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise SidecarError( + STATUS_INVALID_REQUEST, + "tensor checksum must be a lowercase sha256 digest", + ) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + candidates.append( + ExternalTensorCandidate( + candidate_id, + ExternalTensorRequest( + model_contract_id=model_contract_id, + tensor_ref=tensor_ref, + rows=rows, + dimension=dimension, + dtype=dtype, + checksum=checksum, + ), + ) + ) + reader.finish() + return ParsedExternalTensorRequest( + request_id, + dimension, + query_rows, + dtype, + model_contract_id, + query_payload, + tuple(candidates), + ) + + +def process_frame( + frame: bytes, + limits: Limits = Limits(), + resolver: ExternalTensorResolver | None = None, +) -> bytes: + request_id = 0 + response_version = VERSION + try: + if len(frame) < HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "truncated frame header") + magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) + if version in (VERSION, EXTERNAL_VERSION): + response_version = version + if magic != MAGIC: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") + if version not in (VERSION, EXTERNAL_VERSION): + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") + if kind != REQUEST_KIND: + raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") + if body_len != len(frame) - HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "request length mismatch") + if len(frame) > limits.max_request_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds byte limit") + + reader = Reader(frame[HEADER.size :]) + if version == VERSION: + results = process_inline_request(reader, limits) + else: + results = process_external_request(reader, limits, resolver) + return success_response(request_id, results, response_version) + except SidecarError as error: + return error_response(request_id, error.status, str(error), response_version) + except Exception as error: # Keep protocol failures inside the response boundary. + return error_response( + request_id, STATUS_COMPUTE_ERROR, str(error), response_version + ) + + +def receive_exact(connection: socket.socket, count: int) -> bytes: + chunks = bytearray() + while len(chunks) < count: + chunk = connection.recv(count - len(chunks)) + if not chunk: + raise SidecarError( + STATUS_INVALID_REQUEST, "connection closed during request" + ) + chunks.extend(chunk) + return bytes(chunks) + + +def handle_connection( + connection: socket.socket, + limits: Limits, + resolver: ExternalTensorResolver | None = None, +) -> None: + request_id = 0 + response_version = VERSION + try: + header = receive_exact(connection, HEADER.size) + _, version, _, request_id, body_len = HEADER.unpack(header) + if version in (VERSION, EXTERNAL_VERSION): + response_version = version + if body_len > limits.max_request_bytes - HEADER.size: + response = error_response( + request_id, + STATUS_RESOURCE_LIMIT, + "request exceeds byte limit", + response_version, + ) + else: + body = receive_exact(connection, body_len) + response = process_frame(header + body, limits, resolver) + except SidecarError as error: + response = error_response( + request_id, error.status, str(error), response_version + ) + except Exception as error: + response = error_response( + request_id, STATUS_COMPUTE_ERROR, str(error), response_version + ) + connection.sendall(response) + + +def remove_stale_socket(path: Path) -> None: + try: + mode = path.lstat().st_mode + except FileNotFoundError: + return + if not stat.S_ISSOCK(mode): + raise RuntimeError(f"refusing to replace non-socket path: {path}") + path.unlink() + + +def serve( + socket_path: Path, + limits: Limits, + socket_mode: int = 0o600, + once: bool = False, + stop: threading.Event | None = None, + resolver: ExternalTensorResolver | None = None, +) -> None: + stop = stop or threading.Event() + remove_stale_socket(socket_path) + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + listener.bind(os.fspath(socket_path)) + os.chmod(socket_path, socket_mode) + listener.listen(16) + listener.settimeout(0.25) + bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino + while not stop.is_set(): + try: + connection, _ = listener.accept() + except TimeoutError: + continue + with connection: + handle_connection(connection, limits, resolver) + if once: + break + finally: + listener.close() + try: + current = socket_path.lstat() + if (current.st_dev, current.st_ino) == bound_identity: + socket_path.unlink() + except (FileNotFoundError, UnboundLocalError): + pass + + +def parse_mode(value: str) -> int: + mode = int(value, 8) + if mode < 0 or mode > 0o777: + raise argparse.ArgumentTypeError("socket mode must be between 000 and 777") + return mode + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=parse_mode, default=0o600) + parser.add_argument("--max-request-bytes", type=int, default=64 * 1024 * 1024) + parser.add_argument("--max-batch-tokens", type=int, default=1_000_000) + parser.add_argument("--max-tensor-bytes", type=int, default=1024 * 1024 * 1024) + parser.add_argument("--max-candidates", type=int, default=65_536) + parser.add_argument("--once", action="store_true") + args = parser.parse_args() + limits = Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + if ( + min( + limits.max_request_bytes, + limits.max_batch_tokens, + limits.max_tensor_bytes, + limits.max_candidates, + ) + <= 0 + ): + parser.error("all limits must be positive") + + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + serve(args.socket, limits, args.socket_mode, args.once, stop) + + +if __name__ == "__main__": + main() diff --git a/services/Dockerfile.tilemaxsim b/services/Dockerfile.tilemaxsim new file mode 100644 index 00000000..9f58b88d --- /dev/null +++ b/services/Dockerfile.tilemaxsim @@ -0,0 +1,12 @@ +ARG BASE_IMAGE +FROM ${BASE_IMAGE} + +WORKDIR /opt/vectorchord + +COPY devtools/tilemaxsim_reference_sidecar.py devtools/tilemaxsim_reference_sidecar.py +COPY services/tilemaxsim_cuda_sidecar.py services/tilemaxsim_cuda_sidecar.py + +ENV PYTHONPATH=/opt/vectorchord \ + PYTHONUNBUFFERED=1 + +ENTRYPOINT ["python3", "-m", "services.tilemaxsim_cuda_sidecar"] diff --git a/services/benchmark_tilemaxsim_cuda.py b/services/benchmark_tilemaxsim_cuda.py new file mode 100644 index 00000000..698e4b84 --- /dev/null +++ b/services/benchmark_tilemaxsim_cuda.py @@ -0,0 +1,152 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Reproducible synthetic load probe for the CUDA TileMaxSim executor.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import time + +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import TorchTileMaxsimEngine, positive_int + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + index = max(0, math.ceil(fraction * len(ordered)) - 1) + return ordered[index] + + +def canonical_payload(tensor: torch.Tensor, dtype: int) -> bytes: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + return tensor.to(dtype=scalar_dtype).contiguous().numpy().tobytes() + + +def normalized_tensor( + shape: tuple[int, ...], generator: torch.Generator +) -> torch.Tensor: + tensor = torch.randn(shape, dtype=torch.float32, generator=generator) + return torch.nn.functional.normalize(tensor, dim=-1) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--dtype", choices=("f16", "f32"), default="f16") + parser.add_argument("--dimension", type=positive_int, default=320) + parser.add_argument("--query-rows", type=positive_int, default=32) + parser.add_argument("--document-rows", type=positive_int, default=747) + parser.add_argument("--candidates", type=positive_int, default=128) + parser.add_argument("--warmup", type=positive_int, default=3) + parser.add_argument("--iterations", type=positive_int, default=10) + parser.add_argument("--seed", type=int, default=20260713) + parser.add_argument( + "--max-device-bytes", type=positive_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument("--allow-tf32", action="store_true") + args = parser.parse_args() + + dtype = protocol.DTYPE_F32 if args.dtype == "f32" else protocol.DTYPE_F16 + generator = torch.Generator(device="cpu").manual_seed(args.seed) + query = normalized_tensor((args.query_rows, args.dimension), generator) + document_tensors = normalized_tensor( + (args.candidates, args.document_rows, args.dimension), generator + ) + query_payload = canonical_payload(query, dtype) + documents = [ + ( + candidate_id, + args.document_rows, + canonical_payload(document_tensors[candidate_id], dtype), + ) + for candidate_id in range(args.candidates) + ] + del document_tensors + + engine = TorchTileMaxsimEngine( + args.device, args.max_device_bytes, args.allow_tf32, 1 + ) + if engine.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(engine.device) + + total_samples: list[float] = [] + queue_samples: list[float] = [] + compute_samples: list[float] = [] + score_checksum = 0.0 + for iteration in range(args.warmup + args.iterations): + started = time.perf_counter() + results, queue_ms, compute_ms = engine.score( + query_payload, + args.query_rows, + args.dimension, + dtype, + documents, + time.monotonic() + 300, + lambda: False, + ) + total_ms = (time.perf_counter() - started) * 1000.0 + if iteration >= args.warmup: + total_samples.append(total_ms) + queue_samples.append(queue_ms) + compute_samples.append(compute_ms) + score_checksum = math.fsum(score for _, score in results) + + output = { + "benchmark": "tilemaxsim_cuda_synthetic_v1", + "device": str(engine.device), + "device_name": ( + torch.cuda.get_device_name(engine.device) + if engine.device.type == "cuda" + else "cpu" + ), + "torch_version": torch.__version__, + "dtype": args.dtype, + "dimension": args.dimension, + "query_rows": args.query_rows, + "document_rows": args.document_rows, + "candidates": args.candidates, + "candidate_tokens": args.candidates * args.document_rows, + "seed": args.seed, + "warmup": args.warmup, + "iterations": args.iterations, + "allow_tf32": args.allow_tf32, + "max_device_bytes": args.max_device_bytes, + "latency_ms": { + "mean": round(statistics.fmean(total_samples), 3), + "p50": round(percentile(total_samples, 0.50), 3), + "p95": round(percentile(total_samples, 0.95), 3), + "p99": round(percentile(total_samples, 0.99), 3), + "queue_mean": round(statistics.fmean(queue_samples), 3), + "compute_mean": round(statistics.fmean(compute_samples), 3), + }, + "score_checksum": score_checksum, + } + if engine.device.type == "cuda": + output["cuda_peak_allocated_bytes"] = torch.cuda.max_memory_allocated( + engine.device + ) + output["cuda_peak_reserved_bytes"] = torch.cuda.max_memory_reserved( + engine.device + ) + print(json.dumps(output, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/build_tilemaxsim_tensor_cache.py b/services/build_tilemaxsim_tensor_cache.py new file mode 100644 index 00000000..a48a1728 --- /dev/null +++ b/services/build_tilemaxsim_tensor_cache.py @@ -0,0 +1,235 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Publish NPY page tensors into the sidecar's immutable SHA-256 cache.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +import tempfile +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np + +from services.tilemaxsim_cuda_sidecar import positive_int + + +def canonical_tensor(path: Path, expected_rows: int, expected_dim: int) -> np.ndarray: + tensor = np.load(path, mmap_mode="r", allow_pickle=False) + if tensor.ndim != 2 or tensor.shape != (expected_rows, expected_dim): + raise ValueError( + f"{path}: expected shape {(expected_rows, expected_dim)}, got {tensor.shape}" + ) + if tensor.dtype == np.dtype("float16"): + little_dtype = np.dtype(" str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def write_payload(path: Path, payload: memoryview, digest: str, fsync: bool) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + if path.stat().st_size != len(payload): + raise ValueError(f"existing cache payload has wrong size: {path}") + if file_digest(path) != digest: + raise ValueError(f"existing cache payload has wrong checksum: {path}") + return + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".tmp", dir=path.parent + ) + try: + with os.fdopen(descriptor, "wb", closefd=True) as stream: + stream.write(payload) + stream.flush() + if fsync: + os.fsync(stream.fileno()) + try: + os.link(temporary_name, path) + except FileExistsError: + if path.stat().st_size != len(payload): + raise ValueError(f"concurrent cache payload has wrong size: {path}") + if file_digest(path) != digest: + raise ValueError(f"concurrent cache payload has wrong checksum: {path}") + if fsync: + directory_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + + +def process_record( + record: dict[str, object], + source_root: Path, + cache_root: Path, + fsync: bool, + dry_run: bool, +) -> dict[str, object]: + page_key = record.get("page_key") + relative = record.get("embedding_file") + rows = record.get("n_tokens") + dimension = record.get("dim") + if not isinstance(page_key, str) or not page_key: + raise ValueError("manifest record has no page_key") + if not isinstance(relative, str) or not relative: + raise ValueError(f"manifest page {page_key} has no embedding_file") + if not isinstance(rows, int) or rows <= 0: + raise ValueError(f"manifest page {page_key} has invalid n_tokens") + if not isinstance(dimension, int) or dimension <= 0: + raise ValueError(f"manifest page {page_key} has invalid dim") + source = (source_root / relative).resolve(strict=True) + try: + source.relative_to(source_root) + except ValueError as error: + raise ValueError( + f"embedding path escapes the source root: {relative}" + ) from error + tensor = canonical_tensor(source, rows, dimension) + payload = memoryview(tensor).cast("B") + digest = hashlib.sha256(payload).hexdigest() + destination = cache_root / digest[:2] / f"{digest}.bin" + if not dry_run: + write_payload(destination, payload, digest, fsync) + dtype_name = "float16" if tensor.dtype == np.dtype("float16") else "float32" + return { + "page_key": page_key, + "tensor_ref": f"sha256://{digest}", + "tensor_rows": rows, + "tensor_dim": dimension, + "tensor_dtype": dtype_name, + "tensor_checksum": f"sha256:{digest}", + "canonical_bytes": len(payload), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--cache-root", required=True, type=Path) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--workers", type=positive_int, default=4) + parser.add_argument("--no-fsync", action="store_true") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + + source_root = args.manifest.resolve(strict=True).parent + cache_root = args.cache_root.resolve() + if not cache_root.is_absolute(): + parser.error("--cache-root must be absolute") + records = [] + with args.manifest.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError( + f"invalid JSON at manifest line {line_number}" + ) from error + if not isinstance(record, dict): + raise ValueError(f"manifest line {line_number} is not an object") + records.append(record) + if not records: + raise ValueError("manifest is empty") + + cache_root.mkdir(parents=True, exist_ok=True) + descriptors = [] + with ThreadPoolExecutor(max_workers=args.workers) as workers: + results = workers.map( + lambda record: process_record( + record, + source_root, + cache_root, + not args.no_fsync, + args.dry_run, + ), + records, + ) + for completed, item in enumerate(results, 1): + descriptors.append(item) + if completed % 1000 == 0 or completed == len(records): + print( + json.dumps( + {"event": "tensor_cache_progress", "completed": completed}, + separators=(",", ":"), + ), + file=sys.stderr, + flush=True, + ) + + output_parent = args.descriptor_manifest.parent + output_parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{args.descriptor_manifest.name}.", + suffix=".tmp", + dir=output_parent, + text=True, + ) + try: + with os.fdopen(descriptor, "w", encoding="utf-8", closefd=True) as stream: + for item in descriptors: + stream.write(json.dumps(item, separators=(",", ":"), sort_keys=True)) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_name, args.descriptor_manifest) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + + total_bytes = sum(int(item["canonical_bytes"]) for item in descriptors) + print( + json.dumps( + { + "pages": len(descriptors), + "canonical_bytes": total_bytes, + "cache_root": os.fspath(cache_root), + "descriptor_manifest": os.fspath(args.descriptor_manifest), + "dry_run": args.dry_run, + }, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/services/test_tilemaxsim_cuda_sidecar.py b/services/test_tilemaxsim_cuda_sidecar.py new file mode 100644 index 00000000..093245e3 --- /dev/null +++ b/services/test_tilemaxsim_cuda_sidecar.py @@ -0,0 +1,313 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +from __future__ import annotations + +import hashlib +import os +import socket +import stat +import struct +import tempfile +import threading +import time +import unittest +from pathlib import Path + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from devtools.test_tilemaxsim_reference_sidecar import ( + decode_response, + external_request_frame, + request_frame, +) +from services import tilemaxsim_cuda_sidecar as cuda_sidecar +from services.build_tilemaxsim_tensor_cache import process_record + + +class CapturingMetrics(cuda_sidecar.JsonMetrics): + def __init__(self) -> None: + super().__init__() + self.events: list[dict[str, object]] = [] + + def emit(self, fields: dict[str, object]) -> None: + with self.lock: + self.events.append(fields.copy()) + + +def write_content_addressed(root: Path, payload: bytes) -> tuple[str, str]: + digest = hashlib.sha256(payload).hexdigest() + directory = root / digest[:2] + directory.mkdir(parents=True, exist_ok=True) + (directory / f"{digest}.bin").write_bytes(payload) + return f"sha256://{digest}", f"sha256:{digest}" + + +class CudaSidecarTest(unittest.TestCase): + def test_cache_builder_publishes_resolver_compatible_payload(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_root = root / "source" + cache_root = root / "cache" + source_root.mkdir() + tensor = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + cuda_sidecar.validate_finite_payload( + struct.pack("<2e", 1.0, 0.0), 1, 2, protocol.DTYPE_F16 + ) + with self.assertRaisesRegex(protocol.SidecarError, "non-finite"): + cuda_sidecar.validate_finite_payload( + struct.pack("<2f", 1.0, float("nan")), + 1, + 2, + protocol.DTYPE_F32, + ) + + def test_content_addressed_resolver_validates_and_caches(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 1024) + try: + request = protocol.ExternalTensorRequest( + "model@1", + tensor_ref, + 2, + 2, + protocol.DTYPE_F16, + checksum, + ) + first = resolver.resolve(request) + second = resolver.resolve(request) + self.assertEqual(first.payload, payload) + self.assertFalse(first.cache_hit) + self.assertTrue(second.cache_hit) + + bad = protocol.ExternalTensorRequest( + "model@1", + tensor_ref, + 2, + 2, + protocol.DTYPE_F16, + "sha256:" + "0" * 64, + ) + with self.assertRaisesRegex(protocol.SidecarError, "disagree"): + resolver.resolve(bad) + finally: + resolver.close() + + def test_content_addressed_resolver_rejects_symlink(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) / "root" + root.mkdir() + payload = struct.pack(" None: + query = [[1.0, 0.0], [0.0, 1.0]] + candidates = [ + (17, [[1.0, 0.0], [0.0, 1.0]]), + (3, [[0.5, 0.5], [0.25, 0.25]]), + ] + frame = request_frame(41, protocol.DTYPE_F32, query, candidates) + parsed = protocol.parse_request_frame(frame) + self.assertIsInstance(parsed, protocol.InlineTensorRequest) + assert isinstance(parsed, protocol.InlineTensorRequest) + documents = [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in parsed.candidates + ] + # 64 bytes fits one candidate but not both, exercising internal + # all-or-nothing device chunking. + engine = cuda_sidecar.TorchTileMaxsimEngine("cpu", 64, False, 1) + results, _, _ = engine.score( + parsed.query_payload, + parsed.query_rows, + parsed.dimension, + parsed.dtype, + documents, + time.monotonic() + 2, + lambda: False, + ) + _, status, oracle = decode_response(protocol.process_frame(frame)) + self.assertEqual(status, 0) + self.assertEqual(results, oracle) + + def test_compute_capacity_wait_uses_overall_deadline(self) -> None: + engine = cuda_sidecar.TorchTileMaxsimEngine("cpu", 1024, False, 1) + self.assertTrue(engine.compute_slots.acquire(blocking=False)) + try: + started = time.monotonic() + with self.assertRaisesRegex(protocol.SidecarError, "CUDA capacity"): + engine.score( + struct.pack(" None: + query = [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25]] + candidates = [ + (7, [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25]]), + (2, [[0.5, 0.5, 0.0], [-0.5, 0.25, 1.0]]), + ] + frame = request_frame(52, protocol.DTYPE_F16, query, candidates) + parsed = protocol.parse_request_frame(frame) + assert isinstance(parsed, protocol.InlineTensorRequest) + engine = cuda_sidecar.TorchTileMaxsimEngine("cuda:0", 1024 * 1024, False, 1) + results, _, _ = engine.score( + parsed.query_payload, + parsed.query_rows, + parsed.dimension, + parsed.dtype, + [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in parsed.candidates + ], + time.monotonic() + 5, + lambda: False, + ) + _, status, oracle = decode_response(protocol.process_frame(frame)) + self.assertEqual(status, 0) + assert isinstance(oracle, list) + self.assertEqual([item[0] for item in results], [item[0] for item in oracle]) + for (_, actual), (_, expected) in zip(results, oracle, strict=True): + self.assertAlmostEqual(actual, expected, places=5) + + def test_v2_unix_socket_end_to_end(self) -> None: + with tempfile.TemporaryDirectory() as directory: + directory_path = Path(directory) + root = directory_path / "tensors" + root.mkdir() + tensor = [[1.0, 0.0], [0.0, 1.0]] + payload = struct.pack("<4e", *sum(tensor, [])) + tensor_ref, _ = write_content_addressed(root, payload) + frame, _ = external_request_frame( + 61, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, tensor)], + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 1024) + metrics = CapturingMetrics() + service = cuda_sidecar.TileMaxsimService( + protocol.Limits(), + resolver, + cuda_sidecar.TorchTileMaxsimEngine("cpu", 1024 * 1024, False, 1), + 2000, + metrics, + ) + socket_path = directory_path / "tilemaxsim.sock" + stop = threading.Event() + thread = threading.Thread( + target=cuda_sidecar.serve, + args=(socket_path, 0o600, 4, 2, service, stop), + kwargs={"once": True}, + daemon=True, + ) + thread.start() + for _ in range(100): + if socket_path.exists(): + break + time.sleep(0.01) + else: + self.fail("CUDA sidecar socket was not created") + self.assertEqual(stat.S_IMODE(socket_path.stat().st_mode), 0o600) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_len = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_len) + thread.join(timeout=3) + resolver.close() + self.assertFalse(thread.is_alive()) + _, status, results = decode_response(response) + self.assertEqual(status, 0) + self.assertEqual(results, [(9, 2.0)]) + request_events = [ + event + for event in metrics.events + if event.get("event") == "tilemaxsim_request" + ] + self.assertEqual(len(request_events), 1) + self.assertEqual(request_events[0]["source"], "content_addressed") + self.assertEqual(request_events[0]["status"], "ok") + + +if __name__ == "__main__": + unittest.main() diff --git a/services/tilemaxsim_cuda_sidecar.py b/services/tilemaxsim_cuda_sidecar.py new file mode 100644 index 00000000..ce744119 --- /dev/null +++ b/services/tilemaxsim_cuda_sidecar.py @@ -0,0 +1,802 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Bounded CUDA executor for VectorChord TileMaxSim IPC v1 and v2. + +Version 1 consumes inline tensors. Version 2 resolves canonical tensor payloads +from an operations-configured, per-model content-addressed cache. Applications may +populate that cache from any object store; object-store credentials and routing +never enter PostgreSQL or this protocol. +""" + +from __future__ import annotations + +import argparse +import hashlib +import hmac +import json +import math +import os +import select +import signal +import socket +import stat +import struct +import sys +import threading +import time +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterable + +import numpy as np +import torch + +if __package__ in (None, ""): + sys.path.insert(0, os.fspath(Path(__file__).resolve().parents[1])) + +from devtools import tilemaxsim_reference_sidecar as protocol + + +def validate_finite_payload( + payload: bytes, rows: int, dimension: int, dtype: int +) -> None: + expected = protocol.checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor byte length does not match its shape", + ) + scalar_dtype = " None: + self.maximum_bytes = maximum_bytes + self.current_bytes = 0 + self.entries: OrderedDict[tuple[object, ...], bytes] = OrderedDict() + self.lock = threading.Lock() + + def get(self, key: tuple[object, ...]) -> bytes | None: + if self.maximum_bytes == 0: + return None + with self.lock: + payload = self.entries.get(key) + if payload is not None: + self.entries.move_to_end(key) + return payload + + def put(self, key: tuple[object, ...], payload: bytes) -> None: + if self.maximum_bytes == 0 or len(payload) > self.maximum_bytes: + return + with self.lock: + previous = self.entries.pop(key, None) + if previous is not None: + self.current_bytes -= len(previous) + self.entries[key] = payload + self.current_bytes += len(payload) + while self.current_bytes > self.maximum_bytes: + _, evicted = self.entries.popitem(last=False) + self.current_bytes -= len(evicted) + + +class ContentAddressedResolver: + """Resolve ``sha256://`` inside an allowlisted model cache root. + + A payload with digest ``abcdef...`` is stored as + ``/ab/abcdef....bin``. Directory and file symlinks are + rejected with ``openat(O_NOFOLLOW)``. The digest in the reference, the + registered checksum, the exact byte length, and the file content must all + agree before a payload is returned. + """ + + def __init__(self, roots: dict[str, Path], cache_bytes: int) -> None: + self.root_fds: dict[str, int] = {} + try: + for contract, path in roots.items(): + self.root_fds[contract] = os.open( + os.fspath(path), + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + ) + except Exception: + self.close() + raise + self.cache = PayloadCache(cache_bytes) + + def close(self) -> None: + for descriptor in self.root_fds.values(): + os.close(descriptor) + self.root_fds.clear() + + @staticmethod + def _digest(request: protocol.ExternalTensorRequest) -> str: + prefix = "sha256://" + if not request.tensor_ref.startswith(prefix): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "unsupported tensor reference; expected sha256://", + ) + digest = request.tensor_ref[len(prefix) :] + if len(digest) != 64 or any( + character not in "0123456789abcdef" for character in digest + ): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "invalid content-addressed tensor reference", + ) + if not hmac.compare_digest(request.checksum, f"sha256:{digest}"): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor reference and checksum disagree", + ) + return digest + + @staticmethod + def _read_exact_file(root_fd: int, digest: str, expected_bytes: int) -> bytes: + directory_fd = -1 + payload_fd = -1 + try: + directory_fd = os.open( + digest[:2], + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + payload_fd = os.open( + f"{digest}.bin", + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=directory_fd, + ) + metadata = os.fstat(payload_fd) + if not stat.S_ISREG(metadata.st_mode): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor is not a regular file", + ) + if metadata.st_size != expected_bytes: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match descriptor", + ) + chunks = bytearray() + while len(chunks) < expected_bytes: + chunk = os.read( + payload_fd, min(1024 * 1024, expected_bytes - len(chunks)) + ) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor file ended early", + ) + chunks.extend(chunk) + if os.read(payload_fd, 1): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor file grew during read", + ) + return bytes(chunks) + except FileNotFoundError as error: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "content-addressed tensor is missing" + ) from error + except OSError as error: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + f"content-addressed tensor read failed: {error.strerror}", + ) from error + finally: + if payload_fd >= 0: + os.close(payload_fd) + if directory_fd >= 0: + os.close(directory_fd) + + def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: + root_fd = self.root_fds.get(request.model_contract_id) + if root_fd is None: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "model contract has no configured tensor cache root", + ) + digest = self._digest(request) + expected_bytes = protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + key = ( + request.model_contract_id, + digest, + request.rows, + request.dimension, + request.dtype, + ) + cached = self.cache.get(key) + if cached is not None: + return ResolvedPayload(cached, True) + payload = self._read_exact_file(root_fd, digest, expected_bytes) + actual = hashlib.sha256(payload).hexdigest() + if not hmac.compare_digest(actual, digest): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, "resolved tensor checksum mismatch" + ) + validate_finite_payload(payload, request.rows, request.dimension, request.dtype) + self.cache.put(key, payload) + return ResolvedPayload(payload, False) + + +class TorchTileMaxsimEngine: + def __init__( + self, + device_name: str, + max_device_bytes: int, + allow_tf32: bool, + max_cuda_inflight: int, + ) -> None: + self.device = torch.device(device_name) + if self.device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested but torch.cuda.is_available() is false" + ) + if self.device.type not in ("cuda", "cpu"): + raise RuntimeError("device must be CUDA or CPU") + self.max_device_bytes = max_device_bytes + self.compute_slots = threading.BoundedSemaphore(max_cuda_inflight) + if self.device.type == "cuda": + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.backends.cudnn.allow_tf32 = allow_tf32 + with torch.inference_mode(): + left = torch.zeros((1, 1), dtype=torch.float32, device=self.device) + _ = left @ left + torch.cuda.synchronize(self.device) + + @staticmethod + def _cpu_tensor( + payload: bytes, rows: int, dimension: int, dtype: int + ) -> torch.Tensor: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + # bytearray gives torch a writable, owned buffer; clone detaches the + # resulting tensor before that temporary buffer leaves scope. + tensor = torch.frombuffer(bytearray(payload), dtype=scalar_dtype).reshape( + rows, dimension + ) + if scalar_dtype == torch.float32: + return tensor.clone() + return tensor.to(dtype=torch.float32) + + def _groups( + self, + query_rows: int, + dimension: int, + documents: list[tuple[int, int, bytes]], + ) -> Iterable[list[tuple[int, int, bytes]]]: + query_bytes = query_rows * dimension * 4 + group: list[tuple[int, int, bytes]] = [] + group_rows = 0 + for document in documents: + rows = document[1] + next_rows = group_rows + rows + # Device residency includes the f32 query, f32 documents, and the + # q-by-total-document-token similarity matrix. + required = ( + query_bytes + next_rows * dimension * 4 + query_rows * next_rows * 4 + ) + if required > self.max_device_bytes and group: + yield group + group = [] + group_rows = 0 + next_rows = rows + required = query_bytes + rows * dimension * 4 + query_rows * rows * 4 + if required > self.max_device_bytes: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one candidate exceeds the CUDA device-byte limit", + ) + group.append(document) + group_rows = next_rows + if group: + yield group + + def score( + self, + query_payload: bytes, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, int, bytes]], + deadline: float, + cancelled: Callable[[], bool], + ) -> tuple[list[tuple[int, float]], float, float]: + if not documents: + return [], 0.0, 0.0 + query_cpu = self._cpu_tensor(query_payload, query_rows, dimension, dtype) + results: list[tuple[int, float]] = [] + queue_started = time.monotonic() + remaining = deadline - time.monotonic() + if remaining <= 0 or not self.compute_slots.acquire(timeout=remaining): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired while waiting for CUDA capacity", + ) + queue_ms = (time.monotonic() - queue_started) * 1000.0 + compute_started = time.monotonic() + try: + with torch.inference_mode(): + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + query_device = query_cpu.to(self.device) + for group in self._groups(query_rows, dimension, documents): + if cancelled(): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request peer disconnected" + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + cpu_documents = [ + self._cpu_tensor(payload, rows, dimension, dtype) + for _, rows, payload in group + ] + document_device = torch.cat(cpu_documents).to(self.device) + similarities = query_device @ document_device.transpose(0, 1) + scores = [] + offset = 0 + for _, rows, _ in group: + scores.append( + similarities[:, offset : offset + rows] + .amax(dim=1) + .sum(dtype=torch.float32) + ) + offset += rows + host_scores = torch.stack(scores).to(device="cpu").tolist() + for (candidate_id, _, _), score in zip( + group, host_scores, strict=True + ): + if not math.isfinite(score): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "TileMaxSim result is non-finite", + ) + results.append((candidate_id, score)) + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + finally: + self.compute_slots.release() + return ( + results, + queue_ms, + (time.monotonic() - compute_started) * 1000.0, + ) + + +class JsonMetrics: + def __init__(self) -> None: + self.lock = threading.Lock() + + def emit(self, fields: dict[str, object]) -> None: + with self.lock: + print(json.dumps(fields, separators=(",", ":"), sort_keys=True), flush=True) + + +class TileMaxsimService: + def __init__( + self, + limits: protocol.Limits, + resolver: ContentAddressedResolver, + engine: TorchTileMaxsimEngine, + request_timeout_ms: int, + metrics: JsonMetrics, + ) -> None: + self.limits = limits + self.resolver = resolver + self.engine = engine + self.request_timeout_seconds = request_timeout_ms / 1000.0 + self.metrics = metrics + + @staticmethod + def _peer_disconnected(connection: socket.socket) -> bool: + poller = select.poll() + poller.register(connection, select.POLLHUP | select.POLLERR | select.POLLNVAL) + return bool(poller.poll(0)) + + @staticmethod + def _receive_exact_until( + connection: socket.socket, count: int, deadline: float + ) -> bytes: + chunks = bytearray() + while len(chunks) < count: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("request deadline expired during socket read") + connection.settimeout(remaining) + chunk = connection.recv(count - len(chunks)) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "connection closed during request", + ) + chunks.extend(chunk) + return bytes(chunks) + + def process_frame( + self, + frame: bytes, + connection: socket.socket, + deadline: float, + peer_credentials: tuple[int, int, int] | None, + ) -> bytes: + request_id = 0 + version = protocol.VERSION + started = time.monotonic() + metrics: dict[str, object] = {"event": "tilemaxsim_request"} + if peer_credentials is not None: + metrics["peer_pid"], metrics["peer_uid"], metrics["peer_gid"] = ( + peer_credentials + ) + try: + if len(frame) >= protocol.HEADER.size: + _, wire_version, _, request_id, _ = protocol.HEADER.unpack_from(frame) + if wire_version in (protocol.VERSION, protocol.EXTERNAL_VERSION): + version = wire_version + request = protocol.parse_request_frame( + frame, self.limits, validate_finite=False + ) + metrics.update( + request_id=request.request_id, + protocol_version=version, + query_rows=request.query_rows, + dimension=request.dimension, + candidate_count=len(request.candidates), + ) + validate_finite_payload( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + ) + resolve_started = time.monotonic() + cache_hits = 0 + documents: list[tuple[int, int, bytes]] = [] + if isinstance(request, protocol.InlineTensorRequest): + for candidate in request.candidates: + validate_finite_payload( + candidate.payload, + candidate.rows, + request.dimension, + request.dtype, + ) + documents = [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in request.candidates + ] + metrics["source"] = "inline" + else: + metrics["source"] = "content_addressed" + for candidate in request.candidates: + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during tensor resolution", + ) + resolved = self.resolver.resolve(candidate.descriptor) + cache_hits += int(resolved.cache_hit) + documents.append( + ( + candidate.candidate_id, + candidate.descriptor.rows, + resolved.payload, + ) + ) + metrics["cache_hits"] = cache_hits + metrics["resolve_ms"] = round( + (time.monotonic() - resolve_started) * 1000.0, 3 + ) + metrics["document_tokens"] = sum(rows for _, rows, _ in documents) + results, queue_ms, compute_ms = self.engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + metrics["queue_ms"] = round(queue_ms, 3) + metrics["compute_ms"] = round(compute_ms, 3) + metrics["status"] = "ok" + return protocol.success_response(request.request_id, results, version) + except protocol.SidecarError as error: + metrics.update(status="error", error_class=type(error).__name__) + return protocol.error_response( + request_id, error.status, str(error), version + ) + except torch.OutOfMemoryError: + if self.engine.device.type == "cuda": + torch.cuda.empty_cache() + metrics.update(status="error", error_class="CudaOutOfMemory") + return protocol.error_response( + request_id, + protocol.STATUS_RESOURCE_LIMIT, + "CUDA out of memory", + version, + ) + except Exception as error: + metrics.update(status="error", error_class=type(error).__name__) + return protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + f"TileMaxSim compute failed: {error}", + version, + ) + finally: + metrics["total_ms"] = round((time.monotonic() - started) * 1000.0, 3) + self.metrics.emit(metrics) + + def handle(self, connection: socket.socket) -> None: + deadline = time.monotonic() + self.request_timeout_seconds + request_id = 0 + version = protocol.VERSION + peer_credentials = None + if hasattr(socket, "SO_PEERCRED"): + try: + raw_credentials = connection.getsockopt( + socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i") + ) + peer_credentials = struct.unpack("3i", raw_credentials) + except OSError: + pass + try: + header = self._receive_exact_until( + connection, protocol.HEADER.size, deadline + ) + _, wire_version, _, request_id, body_len = protocol.HEADER.unpack(header) + if wire_version in (protocol.VERSION, protocol.EXTERNAL_VERSION): + version = wire_version + if body_len > self.limits.max_request_bytes - protocol.HEADER.size: + response = protocol.error_response( + request_id, + protocol.STATUS_RESOURCE_LIMIT, + "request exceeds byte limit", + version, + ) + else: + body = self._receive_exact_until(connection, body_len, deadline) + response = self.process_frame( + header + body, connection, deadline, peer_credentials + ) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("request deadline expired before socket write") + connection.settimeout(remaining) + connection.sendall(response) + except (BrokenPipeError, ConnectionResetError): + return + except (TimeoutError, socket.timeout): + try: + connection.sendall( + protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during socket I/O", + version, + ) + ) + except OSError: + pass + except Exception as error: + try: + connection.sendall( + protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + str(error), + version, + ) + ) + except OSError: + pass + + +def serve( + socket_path: Path, + socket_mode: int, + backlog: int, + max_inflight: int, + service: TileMaxsimService, + stop: threading.Event, + once: bool = False, +) -> None: + protocol.remove_stale_socket(socket_path) + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + slots = threading.BoundedSemaphore(max_inflight) + workers = ThreadPoolExecutor( + max_workers=max_inflight, thread_name_prefix="tilemaxsim" + ) + active = 0 + active_lock = threading.Lock() + + def handle(connection: socket.socket) -> None: + nonlocal active + try: + with connection: + service.handle(connection) + finally: + with active_lock: + active -= 1 + slots.release() + + try: + listener.bind(os.fspath(socket_path)) + os.chmod(socket_path, socket_mode) + listener.listen(backlog) + listener.settimeout(0.25) + bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino + service.metrics.emit( + { + "event": "tilemaxsim_ready", + "device": str(service.engine.device), + "max_inflight": max_inflight, + "socket": os.fspath(socket_path), + } + ) + accepted = 0 + while not stop.is_set(): + if not slots.acquire(timeout=0.25): + continue + try: + connection, _ = listener.accept() + except TimeoutError: + slots.release() + continue + with active_lock: + active += 1 + current_active = active + service.metrics.emit( + {"event": "tilemaxsim_accept", "inflight": current_active} + ) + workers.submit(handle, connection) + accepted += 1 + if once and accepted == 1: + break + finally: + listener.close() + workers.shutdown(wait=True, cancel_futures=False) + try: + current = socket_path.lstat() + if (current.st_dev, current.st_ino) == bound_identity: + socket_path.unlink() + except (FileNotFoundError, UnboundLocalError): + pass + + +def parse_mode(value: str) -> int: + mode = int(value, 8) + if mode < 0 or mode > 0o777: + raise argparse.ArgumentTypeError("socket mode must be between 000 and 777") + return mode + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def nonnegative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("value must be nonnegative") + return parsed + + +def contract_roots( + values: list[str], parser: argparse.ArgumentParser +) -> dict[str, Path]: + roots = {} + for value in values: + if "=" not in value: + parser.error("--contract-root must be MODEL_CONTRACT_ID=/absolute/path") + contract, raw_path = value.split("=", 1) + path = Path(raw_path) + if not contract or not path.is_absolute(): + parser.error("--contract-root must contain a nonempty ID and absolute path") + if contract in roots: + parser.error(f"duplicate --contract-root for {contract!r}") + roots[contract] = path + return roots + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=parse_mode, default=0o600) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--contract-root", action="append", default=[]) + parser.add_argument( + "--max-request-bytes", type=positive_int, default=64 * 1024 * 1024 + ) + parser.add_argument("--max-batch-tokens", type=positive_int, default=1_000_000) + parser.add_argument( + "--max-tensor-bytes", type=positive_int, default=1024 * 1024 * 1024 + ) + parser.add_argument("--max-candidates", type=positive_int, default=65_536) + parser.add_argument( + "--max-device-bytes", type=positive_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument( + "--cache-bytes", type=nonnegative_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument("--request-timeout-ms", type=positive_int, default=2000) + parser.add_argument("--max-inflight", type=positive_int, default=8) + parser.add_argument("--max-cuda-inflight", type=positive_int, default=1) + parser.add_argument("--backlog", type=positive_int, default=64) + parser.add_argument("--allow-tf32", action="store_true") + parser.add_argument("--once", action="store_true") + args = parser.parse_args() + + roots = contract_roots(args.contract_root, parser) + limits = protocol.Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + resolver = ContentAddressedResolver(roots, args.cache_bytes) + metrics = JsonMetrics() + try: + engine = TorchTileMaxsimEngine( + args.device, + args.max_device_bytes, + args.allow_tf32, + args.max_cuda_inflight, + ) + service = TileMaxsimService( + limits, resolver, engine, args.request_timeout_ms, metrics + ) + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + serve( + args.socket, + args.socket_mode, + args.backlog, + args.max_inflight, + service, + stop, + args.once, + ) + finally: + resolver.close() + + +if __name__ == "__main__": + main() diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 3725f050..ac550103 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -29,3 +29,17 @@ mod text_rabitq8; pub mod typmod; mod typmod_rabitq4; mod typmod_rabitq8; + +pub(crate) const MAX_MAXSIM_VECTORS: usize = u16::MAX as usize + 1; + +pub(crate) fn validate_maxsim_array_len(len: usize) { + if len == 0 { + pgrx::error!("MaxSim arrays must contain at least one vector"); + } + if len > MAX_MAXSIM_VECTORS { + pgrx::error!( + "MaxSim arrays cannot contain more than {} vectors", + MAX_MAXSIM_VECTORS + ); + } +} diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index 0ae307fe..f20ebc8d 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -95,12 +95,20 @@ fn _vchord_halfvec_operator_maxsim( lhs: Array<'_, HalfvecInput<'_>>, rhs: Array<'_, HalfvecInput<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/datatype/operators_rabitq4.rs b/src/datatype/operators_rabitq4.rs index bf0ff920..d1980fe9 100644 --- a/src/datatype/operators_rabitq4.rs +++ b/src/datatype/operators_rabitq4.rs @@ -125,12 +125,20 @@ fn _vchord_rabitq4_operator_maxsim( lhs: Array<'_, Rabitq4Input<'_>>, rhs: Array<'_, Rabitq4Input<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(Rabitq4Borrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/datatype/operators_rabitq8.rs b/src/datatype/operators_rabitq8.rs index d8d3cb72..7cac9aac 100644 --- a/src/datatype/operators_rabitq8.rs +++ b/src/datatype/operators_rabitq8.rs @@ -125,12 +125,20 @@ fn _vchord_rabitq8_operator_maxsim( lhs: Array<'_, Rabitq8Input<'_>>, rhs: Array<'_, Rabitq8Input<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(Rabitq8Borrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index 76b3bf8c..c125a18b 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -95,12 +95,20 @@ fn _vchord_vector_operator_maxsim( lhs: Array<'_, VectorInput<'_>>, rhs: Array<'_, VectorInput<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs index 704edfc8..fa836316 100644 --- a/src/index/fetcher.rs +++ b/src/index/fetcher.rs @@ -24,6 +24,21 @@ pub trait FilterableTuple: Tuple { pub trait Tuple { fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); + + /// Read one user attribute from the already fetched heap slot. + /// + /// Callers that may expose the value outside PostgreSQL must call + /// [`FilterableTuple::filter`] first. Attribute numbers are PostgreSQL + /// one-based `attnum` values, not zero-based Rust indexes. + #[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] + fn attribute(&mut self, attnum: i16) -> Option; +} + +#[derive(Clone, Copy)] +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +pub struct TupleAttribute { + pub datum: Datum, + pub is_null: bool, } pub trait Fetcher { @@ -52,6 +67,7 @@ pub struct HeapFetcher { heap_relation: pgrx::pg_sys::Relation, snapshot: pgrx::pg_sys::Snapshot, heapfetch: *mut pgrx::pg_sys::IndexFetchTableData, + owns_heapfetch: bool, slot: *mut pgrx::pg_sys::TupleTableSlot, values: [Datum; 32], is_nulls: [bool; 32], @@ -77,6 +93,7 @@ impl HeapFetcher { heap_relation, snapshot, heapfetch, + owns_heapfetch: false, slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), values: [Datum::null(); 32], is_nulls: [true; 32], @@ -84,6 +101,44 @@ impl HeapFetcher { } } } + + /// Create a heap fetch state that is not owned by an `IndexScanDesc`. + /// + /// This is used by the restricted external MaxSim executor to resolve the + /// root TIDs stored in the index through HOT chains before SQL-visible + /// descriptor projection. The table AM owns the rules for doing that; + /// looking up `ctid` directly does not follow a HOT chain. + pub unsafe fn new_standalone( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + ) -> Self { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + unsafe { + let table_am = (*heap_relation).rd_tableam; + if table_am.is_null() { + panic!("unknown heap access method"); + } + let index_fetch_begin = (*table_am) + .index_fetch_begin + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + let heapfetch = pg_guard_ffi_boundary(|| index_fetch_begin(heap_relation)); + if heapfetch.is_null() { + panic!("heap access method returned a null index fetch state"); + } + let mut fetcher = Self::new( + index_relation, + heap_relation, + snapshot, + heapfetch, + std::ptr::null_mut(), + ); + fetcher.owns_heapfetch = true; + fetcher + } + } } impl Drop for HeapFetcher { @@ -93,6 +148,16 @@ impl Drop for HeapFetcher { // free common resources pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); pgrx::pg_sys::FreeExecutorState(self.estate); + if self.owns_heapfetch { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + let table_am = (*self.heap_relation).rd_tableam; + let index_fetch_end = (*table_am) + .index_fetch_end + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| index_fetch_end(self.heapfetch)); + } } } } @@ -147,7 +212,14 @@ impl Fetcher for HeapFetcher { false }; if found { - Some(HeapTuple { this: self }) + // The heap table AM rewrites the requested root TID to the + // snapshot-visible HOT-chain member. The slot itself keeps + // the rewritten index TID as well, but carrying it explicitly + // avoids depending on slot representation details. + Some(HeapTuple { + this: self, + current_ctid: ctid, + }) } else { None } @@ -157,6 +229,16 @@ impl Fetcher for HeapFetcher { pub struct HeapTuple<'a> { this: &'a mut HeapFetcher, + current_ctid: ItemPointerData, +} + +impl HeapTuple<'_> { + /// Return the physical TID of the tuple version materialized in the slot. + /// This may differ from the root TID supplied to `Fetcher::fetch` after a + /// HOT update. + pub fn ctid(&self) -> ItemPointerData { + self.current_ctid + } } impl Tuple for HeapTuple<'_> { @@ -175,6 +257,30 @@ impl Tuple for HeapTuple<'_> { (&this.values, &this.is_nulls) } } + + fn attribute(&mut self, attnum: i16) -> Option { + unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + let slot = self.this.slot; + let tuple_descriptor = (*slot).tts_tupleDescriptor; + if attnum <= 0 + || tuple_descriptor.is_null() + || i32::from(attnum) > (*tuple_descriptor).natts + { + return None; + } + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + pgrx::pg_sys::slot_getsomeattrs_int(slot, i32::from(attnum)); + }); + let offset = usize::try_from(attnum - 1).ok()?; + Some(TupleAttribute { + datum: *(*slot).tts_values.add(offset), + is_null: *(*slot).tts_isnull.add(offset), + }) + } + } } impl FilterableTuple for HeapTuple<'_> { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index f2724593..26cde468 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -27,6 +27,18 @@ pub enum PostgresIo { ReadStream, } +#[derive(Debug, Clone, Copy, PostgresGucEnum)] +pub enum PostgresMaxsimBackend { + #[name = c"coarse_only"] + CoarseOnly, + #[name = c"cpu_exact"] + CpuExact, + #[name = c"gpu"] + Gpu, + #[name = c"auto"] + Auto, +} + static VCHORDRQ_QUERY_SAMPLING_ENABLE: GucSetting = GucSetting::::new(false); static VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS: GucSetting = GucSetting::::new(0); @@ -78,6 +90,24 @@ static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); +static VCHORDRQ_MAXSIM_CANDIDATE_LIMIT: GucSetting = GucSetting::::new(-1); + +static VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS: GucSetting = GucSetting::::new(32); + +static VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS: GucSetting = GucSetting::::new(256); + +static VCHORDRQ_MAXSIM_BACKEND: GucSetting = + GucSetting::::new(PostgresMaxsimBackend::CoarseOnly); + +static VCHORDRQ_MAXSIM_GPU_ENDPOINT: GucSetting> = + GucSetting::>::new(Some(c"")); + +static VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS: GucSetting = GucSetting::::new(2000); + +static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS: GucSetting = GucSetting::::new(1_000_000); + +static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES: GucSetting = GucSetting::::new(1_073_741_824); + static VCHORDRQ_PREFILTER: GucSetting = GucSetting::::new(false); static VCHORDRQ_IO_SEARCH: GucSetting = GucSetting::::new( @@ -151,6 +181,82 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_candidate_limit", + c"Maximum number of page candidates produced by MaxSim aggregation.", + c"Use -1 for no page-candidate limit.", + &VCHORDRQ_MAXSIM_CANDIDATE_LIMIT, + -1, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_planner_query_tokens", + c"Expected MaxSim query-token count used by the planner.", + c"Set this to the measured deployment average until expression statistics are available.", + &VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS, + 1, + 65_536, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_planner_document_tokens", + c"Fallback MaxSim document-token count used by the planner.", + c"Used for indexes that predate the native indexed-vector statistic; set it to the measured deployment average.", + &VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS, + 1, + 65_536, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordrq.maxsim_backend", + c"Page-level MaxSim rerank backend.", + c"GPU and auto modes require the native sidecar integration.", + &VCHORDRQ_MAXSIM_BACKEND, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_string_guc( + c"vchordrq.maxsim_gpu_endpoint", + c"Unix-socket endpoint for the TileMaxSim sidecar.", + c"An empty endpoint disables GPU transport.", + &VCHORDRQ_MAXSIM_GPU_ENDPOINT, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_timeout_ms", + c"Overall TileMaxSim sidecar deadline in milliseconds.", + c"The deadline covers connection, request write, and response read.", + &VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS, + 1, + 600_000, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_max_batch_tokens", + c"Maximum query plus document tokens in one TileMaxSim request.", + c"Requests exceeding the limit fail before connecting to the sidecar.", + &VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS, + 1, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_max_batch_bytes", + c"Maximum encoded TileMaxSim request size in bytes.", + c"Requests exceeding the limit fail before connecting to the sidecar.", + &VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES, + 1024, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); GucRegistry::define_bool_guc( c"vchordrq.prefilter", c"`prefilter` argument of vchordrq.", @@ -475,6 +581,39 @@ pub fn vchordrq_maxsim_threshold(index: pgrx::pg_sys::Relation) -> u32 { } } +pub fn vchordrq_maxsim_candidate_limit() -> Option { + let value = VCHORDRQ_MAXSIM_CANDIDATE_LIMIT.get(); + if value < 0 { None } else { Some(value as u32) } +} + +pub fn vchordrq_maxsim_planner_query_tokens() -> u32 { + VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_planner_document_tokens() -> u32 { + VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_backend() -> PostgresMaxsimBackend { + VCHORDRQ_MAXSIM_BACKEND.get() +} + +pub fn vchordrq_maxsim_gpu_endpoint() -> Option { + VCHORDRQ_MAXSIM_GPU_ENDPOINT.get() +} + +pub fn vchordrq_maxsim_gpu_timeout_ms() -> u32 { + VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS.get() as u32 +} + +pub fn vchordrq_maxsim_gpu_max_batch_tokens() -> u32 { + VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_gpu_max_batch_bytes() -> u32 { + VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES.get() as u32 +} + pub fn vchordrq_prefilter() -> bool { VCHORDRQ_PREFILTER.get() } diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 647d2478..95b4bc45 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -349,9 +349,10 @@ pub unsafe extern "C-unwind" fn ambuild( reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); order }, - |indtuples| { + |indtuples, indexed_vectors| { reporter.tuples_done(indtuples); reporter.tuples_total(indtuples); + vchordrq::set_indexed_vectors(&index, indexed_vectors); // enter the barrier let shared = leader.vchordrqshared; pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); @@ -429,9 +430,10 @@ pub unsafe extern "C-unwind" fn ambuild( || { reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); }, - |indtuples| { + |indtuples, indexed_vectors| { reporter.tuples_done(indtuples); reporter.tuples_total(indtuples); + vchordrq::set_indexed_vectors(&index, indexed_vectors); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); }, || {}, @@ -460,6 +462,7 @@ struct VchordrqShared { barrier_enter_0: i32, nparticipants: u32, indtuples: u64, + indexed_vectors: u64, barrier_leave_0: bool, barrier_enter_1: i32, barrier_leave_1: bool, @@ -721,6 +724,7 @@ impl VchordrqLeader { barrier_leave_2: false, mutex: std::mem::zeroed(), indtuples: 0, + indexed_vectors: 0, }); pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_enter_0); pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_leave_0); @@ -875,7 +879,7 @@ pub unsafe extern "C-unwind" fn vchordrq_parallel_build_main( pgrx::pg_sys::ConditionVariableCancelSleep(); order }, - |_| { + |_, _| { // enter the barrier let shared = vchordrqshared; pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); @@ -941,7 +945,7 @@ unsafe fn parallel_build( vchordrqcached: *const u8, mut callback: impl FnMut(u64), sync_0: impl FnOnce() -> u32, - sync_1: impl FnOnce(u64), + sync_1: impl FnOnce(u64, u64), sync_2: impl FnOnce(), ) { use vchordrq_cached::VchordrqCachedReader; @@ -989,6 +993,7 @@ unsafe fn parallel_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + let indexed_vectors = store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1010,6 +1015,7 @@ unsafe fn parallel_build( { pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); (*vchordrqshared).indtuples += 1; + (*vchordrqshared).indexed_vectors += indexed_vectors; indtuples = (*vchordrqshared).indtuples; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } @@ -1029,6 +1035,7 @@ unsafe fn parallel_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + let indexed_vectors = store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1050,6 +1057,7 @@ unsafe fn parallel_build( { pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); (*vchordrqshared).indtuples += 1; + (*vchordrqshared).indexed_vectors += indexed_vectors; indtuples = (*vchordrqshared).indtuples; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } @@ -1065,7 +1073,13 @@ unsafe fn parallel_build( drop(index); - sync_1(unsafe { (*vchordrqshared).indtuples }); + let (indtuples, indexed_vectors) = unsafe { + ( + (*vchordrqshared).indtuples, + (*vchordrqshared).indexed_vectors, + ) + }; + sync_1(indtuples, indexed_vectors); let index = unsafe { PostgresRelation::new(index_relation) }; @@ -1085,7 +1099,7 @@ unsafe fn sequential_build( vchordrqcached: &[u8], mut callback: impl FnMut(u64), sync_0: impl FnOnce(), - sync_1: impl FnOnce(u64), + sync_1: impl FnOnce(u64, u64), sync_2: impl FnOnce(), ) { use vchordrq_cached::VchordrqCachedReader; @@ -1125,6 +1139,7 @@ unsafe fn sequential_build( let index = unsafe { BufferedPostgresRelation::new(index_relation) }; let mut indtuples = 0; + let mut indexed_vectors = 0_u64; match cached { VchordrqCachedReader::_0(_) => { traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { @@ -1134,6 +1149,7 @@ unsafe fn sequential_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + indexed_vectors += store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1166,6 +1182,7 @@ unsafe fn sequential_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + indexed_vectors += store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1194,7 +1211,7 @@ unsafe fn sequential_build( drop(index); - sync_1(indtuples); + sync_1(indtuples, indexed_vectors); let index = unsafe { PostgresRelation::new(index_relation) }; diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index f995c600..976692b6 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -322,28 +322,13 @@ pub unsafe extern "C-unwind" fn amcostestimate( if !(*index_opt_info).hypothetical { let relation = Index::open((*index_opt_info).indexoid, pgrx::pg_sys::NoLock as _); let opfamily = opfamily(relation.raw()); - if !matches!( + let is_maxsim = matches!( opfamily, - Opfamily::HalfvecCosine - | Opfamily::HalfvecIp - | Opfamily::HalfvecL2 - | Opfamily::VectorCosine - | Opfamily::VectorIp - | Opfamily::VectorL2 - | Opfamily::Rabitq8Cosine - | Opfamily::Rabitq8Ip - | Opfamily::Rabitq8L2 - | Opfamily::Rabitq4Cosine - | Opfamily::Rabitq4Ip - | Opfamily::Rabitq4L2 - ) { - *index_startup_cost = 0.0; - *index_total_cost = 0.0; - *index_selectivity = 1.0; - *index_correlation = 0.0; - *index_pages = 1.0; - return; - } + Opfamily::VectorMaxsim + | Opfamily::HalfvecMaxsim + | Opfamily::Rabitq8Maxsim + | Opfamily::Rabitq4Maxsim + ); let index = PostgresRelation::::new(relation.raw()); let probes = gucs::vchordrq_probes(relation.raw()); let cost = vchordrq::cost(&index); @@ -354,15 +339,25 @@ pub unsafe extern "C-unwind" fn amcostestimate( probes.len() ); } + let estimated_index_vectors = if is_maxsim { + cost.indexed_vectors.map_or_else( + || { + (*index_opt_info).tuples.max(total_rows).max(1.0) + * f64::from(gucs::vchordrq_maxsim_planner_document_tokens()) + }, + |indexed_vectors| indexed_vectors as f64, + ) + } else { + (*index_opt_info).tuples.max(0.0) + }; let node_count = { - let tuples = (*index_opt_info).tuples as u32; let mut count = 0.0; - let r = cost.cells.iter().copied().rev(); - let numerator = std::iter::once(1).chain(probes.clone()); + let r = cost.cells.iter().copied().rev().map(f64::from); + let numerator = std::iter::once(1.0).chain(probes.iter().copied().map(f64::from)); let denumerator = r.clone(); - let scale = r.skip(1).chain(std::iter::once(tuples)); + let scale = r.skip(1).chain(std::iter::once(estimated_index_vectors)); for (scale, (numerator, denumerator)) in scale.zip(numerator.zip(denumerator)) { - count += (scale as f64) * 1.0f64.min((numerator as f64) / (denumerator as f64)); + count += scale * 1.0f64.min(numerator / denumerator); } count }; @@ -378,13 +373,41 @@ pub unsafe extern "C-unwind" fn amcostestimate( pages += cost.cells[0] as f64; pages }; + if is_maxsim { + let backend = match gucs::vchordrq_maxsim_backend() { + gucs::PostgresMaxsimBackend::CoarseOnly => { + vchordrq::MaxsimCostBackend::CoarseOnly + } + gucs::PostgresMaxsimBackend::CpuExact => vchordrq::MaxsimCostBackend::CpuExact, + gucs::PostgresMaxsimBackend::Gpu => vchordrq::MaxsimCostBackend::Gpu, + gucs::PostgresMaxsimBackend::Auto => vchordrq::MaxsimCostBackend::Auto, + }; + let estimate = vchordrq::estimate_maxsim_cost(vchordrq::MaxsimCostInput { + heap_rows: total_rows, + index_tokens: estimated_index_vectors, + token_nodes_per_query: node_count, + base_index_pages: page_count, + dimension: cost.dim, + element_bits: opfamily.vector_kind().number_of_bits_of_an_elements(), + query_tokens: gucs::vchordrq_maxsim_planner_query_tokens(), + limit_tuples: ((*root).limit_tuples > 0.0).then_some((*root).limit_tuples), + filter_selectivity, + candidate_limit: gucs::vchordrq_maxsim_candidate_limit(), + backend, + }); + *index_startup_cost = estimate.startup_cost; + *index_total_cost = estimate.total_cost; + *index_selectivity = estimate.selectivity; + *index_correlation = 0.0; + *index_pages = estimate.index_pages; + return; + } // `next_count` represents candidates we expect to process to // surface `limit_tuples` survivors after filter rejection. Clamp // by `node_count` so the estimate cannot exceed the candidates // the IVF visits at the configured probe count. let next_count = if (*root).limit_tuples > 0.0 { - ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)) - .min(node_count) + ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)).min(node_count) } else { node_count }; @@ -482,7 +505,9 @@ pub unsafe extern "C-unwind" fn ambulkdelete( pg_guard_ffi_boundary(|| callback(&mut ctid, callback_state)) } }; - crate::index::vchordrq::dispatch::bulkdelete(opfamily, &index, check, callback); + let indexed_vectors = + crate::index::vchordrq::dispatch::bulkdelete(opfamily, &index, check, callback); + vchordrq::set_indexed_vectors(&index, indexed_vectors); stats } @@ -538,6 +563,12 @@ pub unsafe extern "C-unwind" fn amrescan( max_scan_tuples: gucs::vchordrq_max_scan_tuples(), maxsim_refine: gucs::vchordrq_maxsim_refine((*scan).indexRelation), maxsim_threshold: gucs::vchordrq_maxsim_threshold((*scan).indexRelation), + maxsim_candidate_limit: gucs::vchordrq_maxsim_candidate_limit(), + maxsim_backend: gucs::vchordrq_maxsim_backend(), + maxsim_gpu_endpoint: gucs::vchordrq_maxsim_gpu_endpoint(), + maxsim_gpu_timeout_ms: gucs::vchordrq_maxsim_gpu_timeout_ms(), + maxsim_gpu_max_batch_tokens: gucs::vchordrq_maxsim_gpu_max_batch_tokens(), + maxsim_gpu_max_batch_bytes: gucs::vchordrq_maxsim_gpu_max_batch_bytes(), io_search: gucs::vchordrq_io_search(), io_rerank: gucs::vchordrq_io_rerank(), prefilter: gucs::vchordrq_prefilter(), diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs index 01d70a87..36baee63 100644 --- a/src/index/vchordrq/dispatch.rs +++ b/src/index/vchordrq/dispatch.rs @@ -71,42 +71,51 @@ pub fn bulkdelete( index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, -) where +) -> u64 +where R: RelationRead + RelationWrite, R::Page: Page, { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + live } (VectorKind::Vecf32, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + live } (VectorKind::Vecf16, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + live } (VectorKind::Vecf16, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + live } (VectorKind::Rabitq8, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } (VectorKind::Rabitq8, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } (VectorKind::Rabitq4, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } (VectorKind::Rabitq4, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } } } diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs index 7c51f5db..075b1075 100644 --- a/src/index/vchordrq/opclass.rs +++ b/src/index/vchordrq/opclass.rs @@ -92,6 +92,7 @@ impl Opfamily { Self::VectorMaxsim => { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -105,6 +106,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -118,6 +120,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -131,6 +134,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -203,6 +207,7 @@ impl Opfamily { Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Vecf32(vector.as_borrowed()))); @@ -213,6 +218,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Vecf16(vector.as_borrowed()))); @@ -223,6 +229,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Rabitq8(vector.as_borrowed()))); @@ -233,6 +240,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Rabitq4(vector.as_borrowed()))); diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 1c51c97f..238dbe2e 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -12,7 +12,17 @@ // // Copyright (c) 2025-2026 TensorChord Inc. +mod candidate; +mod external; +mod gpu; +mod rerank; +mod search; + +use self::candidate::{LegacyPageCandidateGenerator, PageCandidateGenerator}; +use self::gpu::{GpuTileMaxsimBackend, UnixSocketTransport, report_gpu_fallback}; +use self::rerank::{CpuExactMaxsimBackend, ExactMaxsimBackend, HeapArrayTensorSource}; use crate::index::fetcher::*; +use crate::index::gucs::PostgresMaxsimBackend; use crate::index::scanners::{Io, SearchBuilder}; use crate::index::vchordrq::dispatch::*; use crate::index::vchordrq::filter::filter; @@ -21,7 +31,6 @@ use crate::index::vchordrq::scanners::SearchOptions; use crate::recorder::Recorder; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; -use distance::Distance; use index::bump::Bump; use index::packed::PackedRefMut8; use index::prefetcher::*; @@ -29,8 +38,8 @@ use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadS use index_accessor::Dot; use simd::f16; use std::cmp::Reverse; -use std::collections::BinaryHeap; use std::num::NonZero; +use std::time::Duration; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; use vchordrq::{RerankMethod, how, maxsim_search, rerank_index}; use vector::VectorOwned; @@ -99,10 +108,33 @@ impl SearchBuilder for MaxsimBuilder { } let maxsim_refine = options.maxsim_refine; let maxsim_threshold = options.maxsim_threshold; + let maxsim_candidate_limit = options + .maxsim_candidate_limit + .map_or(usize::MAX, |value| value as usize); + let has_maxsim_candidate_limit = options.maxsim_candidate_limit.is_some(); + let maxsim_backend = options.maxsim_backend; + let maxsim_gpu_endpoint = options + .maxsim_gpu_endpoint + .as_deref() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let maxsim_gpu_timeout = Duration::from_millis(options.maxsim_gpu_timeout_ms as u64); + let maxsim_gpu_max_batch_tokens = options.maxsim_gpu_max_batch_tokens as usize; + let maxsim_gpu_max_batch_bytes = options.maxsim_gpu_max_batch_bytes as usize; + if !matches!(maxsim_backend, PostgresMaxsimBackend::CoarseOnly) + && (!has_maxsim_candidate_limit || maxsim_candidate_limit == 0) + { + pgrx::error!("exact MaxSim requires a positive vchordrq.maxsim_candidate_limit"); + } let opfamily = self.opfamily; let Some(vectors) = vectors else { return Box::new(std::iter::empty()) as Box>; }; + let expected_dim = vchordrq::cost(index).dim; + if vectors.iter().any(|vector| vector.dim() != expected_dim) { + pgrx::error!("dimension is not matched"); + } + let rerank_query = vectors.clone(); let method = how(index); if !matches!(method, RerankMethod::Index) { pgrx::error!("maxsim search with rerank_in_table is not supported"); @@ -124,669 +156,702 @@ impl SearchBuilder for MaxsimBuilder { _, AlwaysEqual, _, _)>>, )| (rough, payload); - let iter: Box> = match opfamily.vector_kind() { - VectorKind::Vecf32 => { - type Op = vchordrq::operator::Op, Dot>; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Vecf32(vector) = vector { - vector - } else { - unreachable!() - } - }) - .collect::>(); - let projected = unprojected - .iter() - .map(|vector| RandomProject::project(vector.as_borrowed())) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + let mut token_searches = { + let fetcher = &mut fetcher; + let iter: Box> = match opfamily.vector_kind() { + VectorKind::Vecf32 => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } - } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) - } - VectorKind::Vecf16 => { - type Op = vchordrq::operator::Op, Dot>; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Vecf16(vector) = vector { - vector } else { - unreachable!() + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - }) - .collect::>(); - let projected = unprojected - .iter() - .map(|vector| RandomProject::project(vector.as_borrowed())) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Vecf16 => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } - } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) - } - VectorKind::Rabitq8 => { - type Op = vchordrq::operator::Op; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Rabitq8(vector) = vector { - vector } else { - unreachable!() + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - }) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Rabitq8 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq8(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } - } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) - } - VectorKind::Rabitq4 => { - type Op = vchordrq::operator::Op; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Rabitq4(vector) = vector { - vector } else { - unreachable!() + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - }) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Rabitq4 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq4(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) + (accu_set, rough_set, estimation_by_threshold) + })) + } + }; + iter + }; + let mut candidates = + LegacyPageCandidateGenerator.generate(n, &mut token_searches, maxsim_candidate_limit); + drop(token_searches); + let iter: Box> = match maxsim_backend { + PostgresMaxsimBackend::CoarseOnly => Box::new(candidates.map(|candidate| { + let distance = candidate.approximate_distance.to_f32(); + let recheck = false; + (distance, candidate.heap_key, recheck) + })), + PostgresMaxsimBackend::CpuExact => { + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let exact = + CpuExactMaxsimBackend.rerank(&rerank_query, &mut candidates, &mut source); + let exact = exact.unwrap_or_else(|error| pgrx::error!("{error}")); + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) })) } - }; - let mut updates = Vec::new(); - let mut estimations = Vec::new(); - for (query_id, (accu_set, rough_set, estimation_by_threshold)) in iter.enumerate() { - updates.reserve(accu_set.len() + rough_set.len()); - let is_empty = accu_set.is_empty() && rough_set.is_empty(); - let mut estimation_by_scope = Distance::NEG_INFINITY; - for (distance, payload) in accu_set { - estimation_by_scope = std::cmp::max(estimation_by_scope, distance); - let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + PostgresMaxsimBackend::Gpu => { + let candidates = candidates.collect::>(); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let transport = UnixSocketTransport::new(maxsim_gpu_endpoint); + let mut backend = GpuTileMaxsimBackend::new( + transport, + maxsim_gpu_timeout, + maxsim_gpu_max_batch_tokens, + maxsim_gpu_max_batch_bytes, + ); + let exact = backend.rerank(&rerank_query, &mut candidate_iter, &mut source); + let exact = exact.unwrap_or_else(|error| pgrx::error!("{error}")); + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) + })) } - for (distance, payload) in rough_set { - let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + PostgresMaxsimBackend::Auto => { + let candidates = candidates.collect::>(); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let transport = UnixSocketTransport::new(maxsim_gpu_endpoint); + let mut backend = GpuTileMaxsimBackend::new( + transport, + maxsim_gpu_timeout, + maxsim_gpu_max_batch_tokens, + maxsim_gpu_max_batch_bytes, + ); + let exact = backend.rerank(&rerank_query, &mut candidate_iter, &mut source); + let exact = match exact { + Ok(exact) => exact, + Err(error) => { + report_gpu_fallback(&error); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + CpuExactMaxsimBackend + .rerank(&rerank_query, &mut candidate_iter, &mut source) + .unwrap_or_else(|error| pgrx::error!("{error}")) + } + }; + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) + })) } - estimations.push(if !is_empty { - std::cmp::max(estimation_by_scope, estimation_by_threshold) - } else { - Distance::ZERO - }); - } - updates.sort_unstable_by_key(|&(key, ..)| key); - let iter = updates - .chunk_by(|(kl, ..), (kr, ..)| kl == kr) - .map(|chunk| { - let key = chunk[0].0; - let mut value = vec![None; n]; - for &(_, query_id, distance) in chunk { - let this = value[query_id].get_or_insert(Distance::INFINITY); - *this = std::cmp::min(*this, distance); - } - let mut maxsim = 0.0f32; - for (query_id, distance) in value.into_iter().enumerate() { - let d = distance.unwrap_or(estimations[query_id]); - maxsim += Distance::to_f32(d); - } - (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) - }) - .collect::>() - .into_iter_sorted_polyfill() - .map(|(Reverse(distance), AlwaysEqual(key))| { - let distance = distance.to_f32(); - let recheck = false; - (distance, key, recheck) - }); - let iter: Box> = Box::new(iter); - let iter = if let Some(max_scan_tuples) = options.max_scan_tuples { - Box::new(iter.take(max_scan_tuples as _)) - } else { - iter }; #[allow(clippy::let_and_return)] iter } } -// Emulate unstable library feature `binary_heap_into_iter_sorted`. -// See https://github.com/rust-lang/rust/issues/59278. - -trait IntoIterSortedPolyfill { - fn into_iter_sorted_polyfill(self) -> IntoIterSorted; -} - -impl IntoIterSortedPolyfill for BinaryHeap { - fn into_iter_sorted_polyfill(self) -> IntoIterSorted { - IntoIterSorted { inner: self } - } -} - -#[derive(Clone, Debug)] -struct IntoIterSorted { - inner: BinaryHeap, -} - -impl Iterator for IntoIterSorted { - type Item = T; - - #[inline] - fn next(&mut self) -> Option { - self.inner.pop() - } - - #[inline] - fn size_hint(&self) -> (usize, Option) { - let exact = self.inner.len(); - (exact, Some(exact)) - } -} - -impl ExactSizeIterator for IntoIterSorted {} - -impl std::iter::FusedIterator for IntoIterSorted {} - #[inline(always)] pub fn id_0(f: F) -> F where diff --git a/src/index/vchordrq/scanners/maxsim/candidate.rs b/src/index/vchordrq/scanners/maxsim/candidate.rs new file mode 100644 index 00000000..9aa3e22f --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/candidate.rs @@ -0,0 +1,192 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::fetcher::pointer_to_kv; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::num::NonZero; + +pub(super) type HeapKey = [u16; 3]; +pub(super) type TokenCandidate = (Distance, NonZero); +pub(super) type TokenSearchResult = (Vec, Vec, Distance); + +#[derive(Clone, Copy, Debug)] +pub(super) struct PageCandidate { + pub approximate_distance: Distance, + pub heap_key: HeapKey, +} + +pub(super) trait PageCandidateGenerator { + type Candidates: Iterator; + + fn generate( + &mut self, + query_count: usize, + token_searches: &mut dyn Iterator, + candidate_limit: usize, + ) -> Self::Candidates; +} + +#[derive(Default)] +pub(super) struct LegacyPageCandidateGenerator; + +impl PageCandidateGenerator for LegacyPageCandidateGenerator { + type Candidates = LegacyPageCandidates; + + fn generate( + &mut self, + query_count: usize, + token_searches: &mut dyn Iterator, + candidate_limit: usize, + ) -> Self::Candidates { + let mut updates = Vec::new(); + let mut estimations = Vec::with_capacity(query_count); + for (query_id, (accu_set, rough_set, estimation_by_threshold)) in token_searches.enumerate() + { + updates.reserve(accu_set.len() + rough_set.len()); + let is_empty = accu_set.is_empty() && rough_set.is_empty(); + let mut estimation_by_scope = Distance::NEG_INFINITY; + for (distance, payload) in accu_set { + estimation_by_scope = std::cmp::max(estimation_by_scope, distance); + let (key, _) = pointer_to_kv(payload); + updates.push((key, query_id, distance)); + } + for (distance, payload) in rough_set { + let (key, _) = pointer_to_kv(payload); + updates.push((key, query_id, distance)); + } + estimations.push(if !is_empty { + std::cmp::max(estimation_by_scope, estimation_by_threshold) + } else { + Distance::ZERO + }); + } + debug_assert_eq!(estimations.len(), query_count); + updates.sort_unstable_by_key(|&(key, ..)| key); + let inner = updates + .chunk_by(|(kl, ..), (kr, ..)| kl == kr) + .map(|chunk| { + let key = chunk[0].0; + let mut value = vec![None; query_count]; + for &(_, query_id, distance) in chunk { + let this = value[query_id].get_or_insert(Distance::INFINITY); + *this = std::cmp::min(*this, distance); + } + let mut maxsim = 0.0f32; + for (query_id, distance) in value.into_iter().enumerate() { + let distance = distance.unwrap_or(estimations[query_id]); + maxsim += distance.to_f32(); + } + (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) + }) + .collect::>(); + LegacyPageCandidates { + inner, + remaining: candidate_limit, + } + } +} + +pub(super) struct LegacyPageCandidates { + inner: BinaryHeap<(Reverse, AlwaysEqual)>, + remaining: usize, +} + +impl Iterator for LegacyPageCandidates { + type Item = PageCandidate; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + return None; + } + let (Reverse(approximate_distance), AlwaysEqual(heap_key)) = self.inner.pop()?; + self.remaining -= 1; + Some(PageCandidate { + approximate_distance, + heap_key, + }) + } + + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len().min(self.remaining); + (exact, Some(exact)) + } +} + +impl ExactSizeIterator for LegacyPageCandidates {} +impl std::iter::FusedIterator for LegacyPageCandidates {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::index::fetcher::kv_to_pointer; + + fn distance(value: f32) -> Distance { + Distance::from_f32(value) + } + + #[test] + fn aggregates_tokens_by_page_and_preserves_distance_order() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let mut token_searches = vec![ + ( + vec![ + (distance(-1.0), kv_to_pointer((page_1, 0))), + (distance(-0.5), kv_to_pointer((page_2, 0))), + ], + vec![(distance(-0.4), kv_to_pointer((page_2, 1)))], + distance(-0.75), + ), + ( + vec![(distance(-2.0), kv_to_pointer((page_1, 1)))], + vec![], + distance(-1.0), + ), + ] + .into_iter(); + let candidates = LegacyPageCandidateGenerator + .generate(2, &mut token_searches, usize::MAX) + .collect::>(); + + assert_eq!(candidates.len(), 2); + assert_eq!(candidates[0].heap_key, page_1); + assert_eq!(candidates[0].approximate_distance.to_f32(), -3.0); + assert_eq!(candidates[1].heap_key, page_2); + assert_eq!(candidates[1].approximate_distance.to_f32(), -1.5); + } + + #[test] + fn candidate_limit_is_applied_after_global_ordering() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let mut token_searches = vec![( + vec![ + (distance(-1.0), kv_to_pointer((page_1, 0))), + (distance(-2.0), kv_to_pointer((page_2, 0))), + ], + vec![], + Distance::ZERO, + )] + .into_iter(); + let candidates = LegacyPageCandidateGenerator + .generate(1, &mut token_searches, 1) + .collect::>(); + + assert_eq!(candidates.len(), 1); + assert_eq!(candidates[0].heap_key, page_2); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/external.rs b/src/index/vchordrq/scanners/maxsim/external.rs new file mode 100644 index 00000000..0c8d1eb0 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/external.rs @@ -0,0 +1,551 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): This software is also available under the ELv2, +// which has specific restrictions. +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::PageCandidate; +use super::rerank::RerankError; +use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; +use pgrx::datum::{FromDatum, IntoDatum}; +use std::ffi::CString; + +const MAX_MODEL_CONTRACT_BYTES: usize = 512; +const MAX_TENSOR_REF_BYTES: usize = 4096; +const MAX_CHECKSUM_BYTES: usize = 512; +const MAX_TENSOR_ROWS: u32 = 65_536; +const MAX_TENSOR_DIMENSION: u32 = 60_000; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ExternalTensorDtype { + F32, + F16, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ExternalTensorStorage { + SameHeap, + DescriptorRelation, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) struct ExternalTensorColumns { + pub model_contract: i16, + pub public_id: i16, + pub tensor_ref: i16, + pub tensor_rows: i16, + pub tensor_dimension: i16, + pub tensor_dtype: i16, + pub tensor_checksum: i16, +} + +impl ExternalTensorColumns { + pub(super) fn validate(self) -> Result { + let columns = [ + self.model_contract, + self.public_id, + self.tensor_ref, + self.tensor_rows, + self.tensor_dimension, + self.tensor_dtype, + self.tensor_checksum, + ]; + if columns.iter().any(|column| *column <= 0) { + return Err(RerankError::InvalidDescriptor( + "registered attribute number is invalid", + )); + } + for (index, column) in columns.iter().enumerate() { + if columns[..index].contains(column) { + return Err(RerankError::InvalidDescriptor( + "registered descriptor columns are not distinct", + )); + } + } + Ok(self) + } +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorDescriptor { + pub candidate: PageCandidate, + /// Stable application identifier. It stays inside PostgreSQL and is never + /// encoded into the sidecar request. + #[allow( + dead_code, + reason = "returned by the score API through a separate visible-row map" + )] + pub public_id: i64, + pub tensor_ref: String, + pub rows: u32, + pub dimension: u32, + pub dtype: ExternalTensorDtype, + pub checksum: String, +} + +pub(super) trait CandidateTensorDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError>; +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +pub(super) struct HeapTensorRefSource<'a, F> { + fetcher: &'a mut F, + model_contract_id: String, + columns: ExternalTensorColumns, +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorSourceBinding { + pub index_oid: pgrx::pg_sys::Oid, + pub heap_oid: pgrx::pg_sys::Oid, + pub descriptor_oid: Option, + pub storage: ExternalTensorStorage, + pub model_contract_id: String, + #[allow( + dead_code, + reason = "physical attnums are consumed by the optional Phase 3C heap source" + )] + pub columns: Option, + pub column_names: ExternalTensorColumnNames, +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorColumnNames { + pub model_contract: String, + pub public_id: String, + pub descriptor_public_id: Option, + pub tensor_ref: String, + pub tensor_rows: String, + pub tensor_dimension: String, + pub tensor_dtype: String, + pub tensor_checksum: String, +} + +/// Resolve the privilege-aware SQL registry boundary. Same-heap bindings also +/// resolve the physical attribute numbers consumed by [`HeapTensorRefSource`]; +/// independent descriptor relations are projected later through SPI. +/// +/// The SECURITY DEFINER SQL function performs ownership/SELECT checks and +/// revalidates the live index, heap relation, opclass, column types, and NOT +/// NULL constraints. This Rust layer deliberately calls that function rather +/// than reading the private registry table with extension privileges. +pub(super) fn resolve_external_tensor_source( + index_oid: pgrx::pg_sys::Oid, +) -> Result { + pgrx::spi::Spi::connect(|client| { + let schema_rows = client + .select( + "SELECT n.nspname::text AS schema_name + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'", + Some(1), + &[], + ) + .map_err(registry_error)?; + if schema_rows.is_empty() { + return Err(RerankError::Registry( + "vchord extension schema is unavailable".into(), + )); + } + let schema_name = schema_rows + .first() + .get_by_name::("schema_name") + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry("vchord extension schema is NULL".into()))?; + let resolver = pgrx::spi::quote_qualified_identifier( + schema_name, + "vchordrq_maxsim_source_info".to_string(), + ); + let query = format!( + "SELECT registered_index::oid AS index_oid, + heap_relation::oid AS heap_oid, + descriptor_relation::oid AS descriptor_oid, + model_contract_id, + source_storage, + model_contract_column::text AS model_contract_column, + public_id_column::text AS public_id_column, + descriptor_public_id_column::text AS descriptor_public_id_column, + tensor_ref_column::text AS tensor_ref_column, + tensor_rows_column::text AS tensor_rows_column, + tensor_dim_column::text AS tensor_dim_column, + tensor_dtype_column::text AS tensor_dtype_column, + tensor_checksum_column::text AS tensor_checksum_column + FROM {resolver}($1::regclass)" + ); + let prepared = client + .prepare(query.as_str(), &pgrx::oids_of![pgrx::pg_sys::Oid]) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[index_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered MaxSim tensor source resolution returned no row".into(), + )); + } + let row = rows.first(); + let resolved_index_oid = required_column::(&row, "index_oid")?; + let heap_oid = required_column::(&row, "heap_oid")?; + let descriptor_oid = optional_column::(&row, "descriptor_oid")?; + let model_contract_id = required_column::(&row, "model_contract_id")?; + let storage = required_column::(&row, "source_storage")?; + if resolved_index_oid != index_oid { + return Err(RerankError::Registry( + "registered MaxSim tensor source resolved a different index".into(), + )); + } + let storage = match storage.as_str() { + "external_ref" if descriptor_oid.is_none() => ExternalTensorStorage::SameHeap, + "external_relation" if descriptor_oid.is_some() => { + ExternalTensorStorage::DescriptorRelation + } + "heap_array" => { + return Err(RerankError::Registry( + "registered MaxSim tensor source is not external".into(), + )); + } + _ => { + return Err(RerankError::Registry( + "registered external MaxSim tensor source is inconsistent".into(), + )); + } + }; + + let column_names = ExternalTensorColumnNames { + model_contract: required_column::(&row, "model_contract_column")?, + public_id: required_column::(&row, "public_id_column")?, + descriptor_public_id: optional_column::(&row, "descriptor_public_id_column")?, + tensor_ref: required_column::(&row, "tensor_ref_column")?, + tensor_rows: required_column::(&row, "tensor_rows_column")?, + tensor_dimension: required_column::(&row, "tensor_dim_column")?, + tensor_dtype: required_column::(&row, "tensor_dtype_column")?, + tensor_checksum: required_column::(&row, "tensor_checksum_column")?, + }; + let columns = match storage { + ExternalTensorStorage::SameHeap => Some( + ExternalTensorColumns { + model_contract: resolve_attnum(heap_oid, &column_names.model_contract)?, + public_id: resolve_attnum(heap_oid, &column_names.public_id)?, + tensor_ref: resolve_attnum(heap_oid, &column_names.tensor_ref)?, + tensor_rows: resolve_attnum(heap_oid, &column_names.tensor_rows)?, + tensor_dimension: resolve_attnum(heap_oid, &column_names.tensor_dimension)?, + tensor_dtype: resolve_attnum(heap_oid, &column_names.tensor_dtype)?, + tensor_checksum: resolve_attnum(heap_oid, &column_names.tensor_checksum)?, + } + .validate()?, + ), + ExternalTensorStorage::DescriptorRelation => { + if column_names.descriptor_public_id.is_none() { + return Err(RerankError::Registry( + "registered descriptor relation has no public ID column".into(), + )); + } + None + } + }; + validate_model_contract(&model_contract_id)?; + Ok(ExternalTensorSourceBinding { + index_oid, + heap_oid, + descriptor_oid, + storage, + model_contract_id, + columns, + column_names, + }) + }) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry(format!("registry resolver returned NULL {name}"))) +} + +fn optional_column( + row: &pgrx::spi::SpiTupleTable<'_>, + name: &str, +) -> Result, RerankError> +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name).map_err(registry_error) +} + +fn resolve_attnum(heap_oid: pgrx::pg_sys::Oid, column_name: &str) -> Result { + let column_name = CString::new(column_name) + .map_err(|_| RerankError::Registry("registry column contains a NUL byte".into()))?; + let attnum = unsafe { pgrx::pg_sys::get_attnum(heap_oid, column_name.as_ptr()) }; + if attnum <= 0 { + return Err(RerankError::Registry( + "registered descriptor column disappeared during resolution".into(), + )); + } + Ok(attnum) +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +impl<'a, F> HeapTensorRefSource<'a, F> { + pub(super) fn new( + fetcher: &'a mut F, + model_contract_id: String, + columns: ExternalTensorColumns, + ) -> Result { + validate_model_contract(&model_contract_id)?; + Ok(Self { + fetcher, + model_contract_id, + columns: columns.validate()?, + }) + } +} + +impl CandidateTensorDescriptorSource for HeapTensorRefSource<'_, F> { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + let Some(mut tuple) = self.fetcher.fetch(candidate.heap_key) else { + return Ok(None); + }; + + // This is the data-exfiltration boundary: no descriptor value may be + // read before PostgreSQL has accepted the active base-relation qual. + if !tuple.filter() { + return Ok(None); + } + + let model_contract = read_attribute::(&mut tuple, self.columns.model_contract)?; + if model_contract != self.model_contract_id { + return Err(RerankError::ModelContractMismatch); + } + let public_id = read_attribute::(&mut tuple, self.columns.public_id)?; + let tensor_ref = read_attribute::(&mut tuple, self.columns.tensor_ref)?; + let rows = read_attribute::(&mut tuple, self.columns.tensor_rows)?; + let dimension = read_attribute::(&mut tuple, self.columns.tensor_dimension)?; + let dtype = read_attribute::(&mut tuple, self.columns.tensor_dtype)?; + let checksum = read_attribute::(&mut tuple, self.columns.tensor_checksum)?; + + validate_descriptor( + candidate, public_id, tensor_ref, rows, dimension, dtype, checksum, + ) + .map(Some) + } +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +fn read_attribute(tuple: &mut impl Tuple, attnum: i16) -> Result { + let attribute = tuple + .attribute(attnum) + .ok_or(RerankError::InvalidDescriptor( + "registered attribute is unavailable", + ))?; + unsafe { T::from_datum(attribute.datum, attribute.is_null) }.ok_or( + RerankError::InvalidDescriptor("registered descriptor value is NULL or malformed"), + ) +} + +fn validate_model_contract(value: &str) -> Result<(), RerankError> { + if value.is_empty() || value.len() > MAX_MODEL_CONTRACT_BYTES || contains_control(value) { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn validate_descriptor( + candidate: PageCandidate, + public_id: i64, + tensor_ref: String, + rows: i32, + dimension: i32, + dtype: String, + checksum: String, +) -> Result { + if tensor_ref.is_empty() + || tensor_ref.len() > MAX_TENSOR_REF_BYTES + || contains_control(&tensor_ref) + { + return Err(RerankError::InvalidDescriptor( + "tensor reference is empty, oversized, or contains control characters", + )); + } + let rows = u32::try_from(rows) + .ok() + .filter(|rows| (1..=MAX_TENSOR_ROWS).contains(rows)) + .ok_or(RerankError::InvalidDescriptor( + "tensor row count is invalid", + ))?; + let dimension = u32::try_from(dimension) + .ok() + .filter(|dimension| (1..=MAX_TENSOR_DIMENSION).contains(dimension)) + .ok_or(RerankError::InvalidDescriptor( + "tensor dimension is invalid", + ))?; + let dtype = match dtype.as_str() { + "float32" => ExternalTensorDtype::F32, + "float16" => ExternalTensorDtype::F16, + _ => { + return Err(RerankError::InvalidDescriptor( + "tensor dtype must be float16 or float32", + )); + } + }; + if checksum.len() > MAX_CHECKSUM_BYTES || !is_sha256_checksum(&checksum) { + return Err(RerankError::InvalidDescriptor( + "tensor checksum must be a lowercase sha256 digest", + )); + } + Ok(ExternalTensorDescriptor { + candidate, + public_id, + tensor_ref, + rows, + dimension, + dtype, + checksum, + }) +} + +fn contains_control(value: &str) -> bool { + value.chars().any(char::is_control) +} + +fn is_sha256_checksum(value: &str) -> bool { + value.strip_prefix("sha256:").is_some_and(|digest| { + digest.len() == 64 + && digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::index::fetcher::TupleAttribute; + use distance::Distance; + + fn candidate() -> PageCandidate { + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + } + } + + fn columns() -> ExternalTensorColumns { + ExternalTensorColumns { + model_contract: 1, + public_id: 2, + tensor_ref: 3, + tensor_rows: 4, + tensor_dimension: 5, + tensor_dtype: 6, + tensor_checksum: 7, + } + } + + struct RejectingFetcher; + struct RejectingTuple; + + impl Tuple for RejectingTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + panic!("external descriptor sources do not build index expressions") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + panic!("a rejected tuple must not expose descriptor attributes") + } + } + + impl FilterableTuple for RejectingTuple { + fn filter(&mut self) -> bool { + false + } + } + + impl Fetcher for RejectingFetcher { + type Tuple<'a> = RejectingTuple; + + fn fetch(&mut self, _key: [u16; 3]) -> Option> { + Some(RejectingTuple) + } + } + + #[test] + fn source_filters_before_reading_descriptor_attributes() { + let mut fetcher = RejectingFetcher; + let mut source = + HeapTensorRefSource::new(&mut fetcher, "contract@1".into(), columns()).unwrap(); + assert!(source.fetch(candidate()).unwrap().is_none()); + } + + #[test] + fn descriptor_validation_accepts_bounded_immutable_metadata() { + let descriptor = validate_descriptor( + candidate(), + 42, + "s3://immutable-bucket/page-42.tensor".into(), + 747, + 320, + "float16".into(), + format!("sha256:{}", "a".repeat(64)), + ) + .unwrap(); + assert_eq!(descriptor.public_id, 42); + assert_eq!(descriptor.rows, 747); + assert_eq!(descriptor.dimension, 320); + assert_eq!(descriptor.dtype, ExternalTensorDtype::F16); + } + + #[test] + fn descriptor_validation_rejects_ambiguous_or_unbounded_metadata() { + assert!(matches!( + columns_with_duplicate().validate(), + Err(RerankError::InvalidDescriptor(_)) + )); + for (rows, dimension, dtype, checksum) in [ + (0, 320, "float16", format!("sha256:{}", "a".repeat(64))), + (1, 0, "float16", format!("sha256:{}", "a".repeat(64))), + (1, 320, "bf16", format!("sha256:{}", "a".repeat(64))), + (1, 320, "float16", "sha256:short".into()), + ] { + assert!(matches!( + validate_descriptor( + candidate(), + 1, + "s3://immutable/tensor".into(), + rows, + dimension, + dtype.into(), + checksum, + ), + Err(RerankError::InvalidDescriptor(_)) + )); + } + } + + fn columns_with_duplicate() -> ExternalTensorColumns { + ExternalTensorColumns { + public_id: 1, + ..columns() + } + } +} diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs new file mode 100644 index 00000000..af8c5455 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -0,0 +1,1495 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorDtype, +}; +use super::rerank::{CandidateTensorSource, ExactMaxsimBackend, RerankError, RerankResults}; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::mem::size_of_val; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; +use vchordrq::types::OwnedVector; + +const MAGIC: &[u8; 4] = b"VCTM"; +const VERSION: u16 = 1; +const EXTERNAL_VERSION: u16 = 2; +const REQUEST_KIND: u16 = 1; +const RESPONSE_KIND: u16 = 2; +const HEADER_LEN: usize = 24; +const MAX_REMOTE_ERROR_BYTES: usize = 64 * 1024; + +static NEXT_REQUEST_ID: AtomicU64 = AtomicU64::new(1); +static LAST_FALLBACK_WARNING_SECONDS: AtomicU64 = AtomicU64::new(0); + +pub(super) fn report_gpu_fallback(error: &RerankError) { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |duration| duration.as_secs()); + let should_warn = LAST_FALLBACK_WARNING_SECONDS + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |last| { + (last == 0 || now.saturating_sub(last) >= 60).then_some(now) + }) + .is_ok(); + if should_warn { + pgrx::warning!("GPU MaxSim failed; using cpu_exact: {error}"); + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum TensorDtype { + F32 = 1, + F16 = 2, +} + +pub(super) trait TileMaxsimTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + max_response_bytes: usize, + ) -> Result, RerankError>; +} + +pub(super) struct GpuTileMaxsimBackend { + transport: T, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, +} + +impl GpuTileMaxsimBackend { + pub fn new( + transport: T, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, + ) -> Self { + Self { + transport, + timeout, + max_batch_tokens, + max_batch_bytes, + } + } +} + +impl ExactMaxsimBackend for GpuTileMaxsimBackend { + type Results = RerankResults; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let encoded = encode_request( + request_id, + query, + candidates, + source, + self.max_batch_tokens, + self.max_batch_bytes, + )?; + if encoded.heap_keys.is_empty() { + return Ok(RerankResults { + inner: BinaryHeap::new(), + }); + } + let max_response_bytes = HEADER_LEN + .checked_add(8) + .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) + .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) + .ok_or(RerankError::RequestTooLarge)?; + let response = + self.transport + .round_trip(&encoded.frame, self.timeout, max_response_bytes)?; + decode_response(&response, request_id, &encoded.heap_keys) + } +} + +/// GPU backend for Phase 3B external tensor descriptors. +/// +/// This codec is intentionally separate from [`ExactMaxsimBackend`]: an +/// external full tensor may have different logical values from the indexed +/// sketch, so it cannot be substituted into ordinary `@#` execution. +pub(super) struct GpuExternalTileMaxsimBackend { + transport: T, + model_contract_id: String, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, +} + +impl GpuExternalTileMaxsimBackend { + pub(super) fn new( + transport: T, + model_contract_id: String, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, + ) -> Self { + Self { + transport, + model_contract_id, + timeout, + max_batch_tokens, + max_batch_bytes, + } + } +} + +impl GpuExternalTileMaxsimBackend { + pub(super) fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let encoded = encode_external_request( + request_id, + &self.model_contract_id, + query, + candidates, + source, + self.max_batch_tokens, + self.max_batch_bytes, + )?; + if encoded.heap_keys.is_empty() { + return Ok(RerankResults { + inner: BinaryHeap::new(), + }); + } + let max_response_bytes = HEADER_LEN + .checked_add(8) + .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) + .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) + .ok_or(RerankError::RequestTooLarge)?; + let response = + self.transport + .round_trip(&encoded.frame, self.timeout, max_response_bytes)?; + decode_response_for_version(&response, EXTERNAL_VERSION, request_id, &encoded.heap_keys) + } +} + +struct EncodedRequest { + frame: Vec, + heap_keys: Vec, +} + +fn encode_external_request( + request_id: u64, + model_contract_id: &str, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, +) -> Result { + const MAX_MODEL_CONTRACT_BYTES: usize = 512; + const MAX_TENSOR_REF_BYTES: usize = 4096; + const MAX_CHECKSUM_BYTES: usize = 512; + + if model_contract_id.is_empty() + || model_contract_id.len() > MAX_MODEL_CONTRACT_BYTES + || model_contract_id.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + let (dtype, dimension) = tensor_metadata(query)?; + let external_dtype = match dtype { + TensorDtype::F32 => ExternalTensorDtype::F32, + TensorDtype::F16 => ExternalTensorDtype::F16, + }; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let mut total_tokens = query.len(); + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + let mut declared_tensor_bytes = tensor_bytes(query_rows, dimension, dtype)?; + if declared_tensor_bytes > max_batch_bytes { + return Err(RerankError::RequestTooLarge); + } + + let mut writer = BoundedWriter::new(max_batch_bytes); + writer.zeros(HEADER_LEN)?; + writer.u32(dimension)?; + writer.u32(query_rows)?; + let candidate_count_offset = writer.len(); + writer.u32(0)?; + writer.u8(dtype as u8)?; + writer.u8(1)?; // sum_query_max_document_dot + writer.u16(0)?; + writer + .u32(u32::try_from(model_contract_id.len()).map_err(|_| RerankError::RequestTooLarge)?)?; + writer.bytes(model_contract_id.as_bytes())?; + encode_tensor_values(&mut writer, query, dtype)?; + + let mut heap_keys = Vec::new(); + for candidate in candidates { + let Some(descriptor) = source.fetch(candidate)? else { + continue; + }; + validate_external_for_request( + &descriptor, + dimension, + external_dtype, + MAX_TENSOR_REF_BYTES, + MAX_CHECKSUM_BYTES, + )?; + total_tokens = total_tokens + .checked_add(descriptor.rows as usize) + .ok_or(RerankError::RequestTooLarge)?; + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + declared_tensor_bytes = declared_tensor_bytes + .checked_add(tensor_bytes(descriptor.rows, dimension, dtype)?) + .ok_or(RerankError::RequestTooLarge)?; + if declared_tensor_bytes > max_batch_bytes { + return Err(RerankError::RequestTooLarge); + } + + let candidate_id = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.u32(candidate_id)?; + writer.u32(descriptor.rows)?; + writer.u32( + u32::try_from(descriptor.tensor_ref.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.u32( + u32::try_from(descriptor.checksum.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.bytes(descriptor.tensor_ref.as_bytes())?; + writer.bytes(descriptor.checksum.as_bytes())?; + heap_keys.push(descriptor.candidate.heap_key); + } + + let candidate_count = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_u32(candidate_count_offset, candidate_count); + let body_len = writer + .len() + .checked_sub(HEADER_LEN) + .ok_or_else(|| RerankError::Protocol("invalid request length".into()))?; + writer.patch_bytes(0, MAGIC); + writer.patch_u16(4, EXTERNAL_VERSION); + writer.patch_u16(6, REQUEST_KIND); + writer.patch_u64(8, request_id); + writer.patch_u64( + 16, + u64::try_from(body_len).map_err(|_| RerankError::RequestTooLarge)?, + ); + Ok(EncodedRequest { + frame: writer.finish(), + heap_keys, + }) +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: TensorDtype) -> Result { + let scalar_bytes = match dtype { + TensorDtype::F32 => 4usize, + TensorDtype::F16 => 2usize, + }; + usize::try_from(rows) + .ok() + .and_then(|rows| rows.checked_mul(dimension as usize)) + .and_then(|elements| elements.checked_mul(scalar_bytes)) + .ok_or(RerankError::RequestTooLarge) +} + +fn validate_external_for_request( + descriptor: &ExternalTensorDescriptor, + dimension: u32, + dtype: ExternalTensorDtype, + max_tensor_ref_bytes: usize, + max_checksum_bytes: usize, +) -> Result<(), RerankError> { + if descriptor.rows == 0 + || descriptor.rows > 65_536 + || dimension > 60_000 + || descriptor.dimension != dimension + || descriptor.dtype != dtype + { + return Err(RerankError::TensorMismatch); + } + if descriptor.tensor_ref.is_empty() + || descriptor.tensor_ref.len() > max_tensor_ref_bytes + || descriptor.tensor_ref.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "tensor reference is empty, oversized, or contains control characters", + )); + } + if descriptor.checksum.len() > max_checksum_bytes + || !descriptor + .checksum + .strip_prefix("sha256:") + .is_some_and(|digest| { + digest.len() == 64 + && digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + }) + { + return Err(RerankError::InvalidDescriptor( + "tensor checksum must be a lowercase sha256 digest", + )); + } + Ok(()) +} + +fn encode_request( + request_id: u64, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, +) -> Result { + let (dtype, dimension) = tensor_metadata(query)?; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let mut total_tokens = query.len(); + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + + let mut writer = BoundedWriter::new(max_batch_bytes); + writer.zeros(HEADER_LEN)?; + writer.u32(dimension)?; + writer.u32(query_rows)?; + let candidate_count_offset = writer.len(); + writer.u32(0)?; + writer.u8(dtype as u8)?; + writer.u8(1)?; // sum_query_max_document_dot + writer.u16(0)?; + encode_tensor_values(&mut writer, query, dtype)?; + + let mut heap_keys = Vec::new(); + for candidate in candidates { + let Some(tensor) = source.fetch(candidate)? else { + continue; + }; + let (candidate_dtype, candidate_dimension) = tensor_metadata(&tensor.vectors)?; + if candidate_dtype != dtype || candidate_dimension != dimension { + return Err(RerankError::TensorMismatch); + } + total_tokens = total_tokens + .checked_add(tensor.vectors.len()) + .ok_or(RerankError::RequestTooLarge)?; + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + let candidate_id = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + let rows = u32::try_from(tensor.vectors.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.u32(candidate_id)?; + writer.u32(rows)?; + encode_tensor_values(&mut writer, &tensor.vectors, dtype)?; + heap_keys.push(tensor.candidate.heap_key); + } + let candidate_count = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_u32(candidate_count_offset, candidate_count); + let body_len = writer + .len() + .checked_sub(HEADER_LEN) + .ok_or_else(|| RerankError::Protocol("invalid request length".into()))?; + let body_len = u64::try_from(body_len).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_bytes(0, MAGIC); + writer.patch_u16(4, VERSION); + writer.patch_u16(6, REQUEST_KIND); + writer.patch_u64(8, request_id); + writer.patch_u64(16, body_len); + Ok(EncodedRequest { + frame: writer.finish(), + heap_keys, + }) +} + +fn tensor_metadata(vectors: &[OwnedVector]) -> Result<(TensorDtype, u32), RerankError> { + let Some(first) = vectors.first() else { + return Err(RerankError::TensorMismatch); + }; + let dtype = match first { + OwnedVector::Vecf32(_) => TensorDtype::F32, + OwnedVector::Vecf16(_) => TensorDtype::F16, + OwnedVector::Rabitq8(_) | OwnedVector::Rabitq4(_) => { + return Err(RerankError::UnsupportedTensorKind); + } + }; + let dimension = first.dim(); + for vector in vectors { + let this_dtype = match vector { + OwnedVector::Vecf32(_) => TensorDtype::F32, + OwnedVector::Vecf16(_) => TensorDtype::F16, + OwnedVector::Rabitq8(_) | OwnedVector::Rabitq4(_) => { + return Err(RerankError::UnsupportedTensorKind); + } + }; + if this_dtype != dtype || vector.dim() != dimension { + return Err(RerankError::TensorMismatch); + } + } + Ok((dtype, dimension)) +} + +fn encode_tensor_values( + writer: &mut BoundedWriter, + vectors: &[OwnedVector], + dtype: TensorDtype, +) -> Result<(), RerankError> { + for vector in vectors { + match (dtype, vector) { + (TensorDtype::F32, OwnedVector::Vecf32(vector)) => { + for value in vector.slice() { + writer.bytes(&value.to_le_bytes())?; + } + } + (TensorDtype::F16, OwnedVector::Vecf16(vector)) => { + for value in vector.slice() { + writer.u16(value.to_bits())?; + } + } + _ => return Err(RerankError::TensorMismatch), + } + } + Ok(()) +} + +fn decode_response( + frame: &[u8], + request_id: u64, + heap_keys: &[HeapKey], +) -> Result { + decode_response_for_version(frame, VERSION, request_id, heap_keys) +} + +fn decode_response_for_version( + frame: &[u8], + expected_version: u16, + request_id: u64, + heap_keys: &[HeapKey], +) -> Result { + let mut cursor = Cursor::new(frame); + if cursor.bytes(4)? != MAGIC { + return Err(RerankError::Protocol("invalid magic".into())); + } + if cursor.u16()? != expected_version { + return Err(RerankError::Protocol("unsupported version".into())); + } + if cursor.u16()? != RESPONSE_KIND { + return Err(RerankError::Protocol("unexpected message kind".into())); + } + if cursor.u64()? != request_id { + return Err(RerankError::Protocol("request ID mismatch".into())); + } + let body_len = usize::try_from(cursor.u64()?) + .map_err(|_| RerankError::Protocol("response body is too large".into()))?; + if body_len != frame.len().saturating_sub(HEADER_LEN) { + return Err(RerankError::Protocol("response length mismatch".into())); + } + let status = cursor.u32()?; + if status != 0 { + let length = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("remote error is too large".into()))?; + if length > MAX_REMOTE_ERROR_BYTES { + return Err(RerankError::Protocol("remote error is too large".into())); + } + let message = std::str::from_utf8(cursor.bytes(length)?) + .map_err(|_| RerankError::Protocol("remote error is not UTF-8".into()))?; + cursor.finish()?; + return Err(RerankError::Remote(message.into())); + } + let result_count = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("result count is too large".into()))?; + if result_count != heap_keys.len() { + return Err(RerankError::Protocol("partial result set".into())); + } + let mut seen = vec![false; heap_keys.len()]; + let mut results = BinaryHeap::new(); + for _ in 0..result_count { + let candidate_id = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("candidate ID is too large".into()))?; + let Some(heap_key) = heap_keys.get(candidate_id).copied() else { + return Err(RerankError::Protocol("unknown candidate ID".into())); + }; + if std::mem::replace(&mut seen[candidate_id], true) { + return Err(RerankError::Protocol("duplicate candidate ID".into())); + } + let similarity = f32::from_bits(cursor.u32()?); + if !similarity.is_finite() { + return Err(RerankError::Protocol("non-finite similarity".into())); + } + let distance = Distance::from_f32(-similarity); + results.push((Reverse(distance), Reverse(heap_key))); + } + cursor.finish()?; + if seen.iter().any(|seen| !seen) { + return Err(RerankError::Protocol("partial result set".into())); + } + Ok(RerankResults { inner: results }) +} + +struct BoundedWriter { + bytes: Vec, + limit: usize, +} + +impl BoundedWriter { + fn new(limit: usize) -> Self { + Self { + bytes: Vec::new(), + limit, + } + } + + fn len(&self) -> usize { + self.bytes.len() + } + + fn ensure(&self, additional: usize) -> Result<(), RerankError> { + let size = self + .bytes + .len() + .checked_add(additional) + .ok_or(RerankError::RequestTooLarge)?; + if size > self.limit { + return Err(RerankError::RequestTooLarge); + } + Ok(()) + } + + fn zeros(&mut self, count: usize) -> Result<(), RerankError> { + self.ensure(count)?; + self.bytes.resize(self.bytes.len() + count, 0); + Ok(()) + } + + fn bytes(&mut self, bytes: &[u8]) -> Result<(), RerankError> { + self.ensure(bytes.len())?; + self.bytes.extend_from_slice(bytes); + Ok(()) + } + + fn u8(&mut self, value: u8) -> Result<(), RerankError> { + self.bytes(&[value]) + } + + fn u16(&mut self, value: u16) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + + fn u32(&mut self, value: u32) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + + fn patch_bytes(&mut self, offset: usize, bytes: &[u8]) { + self.bytes[offset..offset + bytes.len()].copy_from_slice(bytes); + } + + fn patch_u16(&mut self, offset: usize, value: u16) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn patch_u32(&mut self, offset: usize, value: u32) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn patch_u64(&mut self, offset: usize, value: u64) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn finish(self) -> Vec { + self.bytes + } +} + +struct Cursor<'a> { + bytes: &'a [u8], + offset: usize, +} + +impl<'a> Cursor<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, offset: 0 } + } + + fn bytes(&mut self, count: usize) -> Result<&'a [u8], RerankError> { + let end = self + .offset + .checked_add(count) + .ok_or_else(|| RerankError::Protocol("message offset overflow".into()))?; + let bytes = self + .bytes + .get(self.offset..end) + .ok_or_else(|| RerankError::Protocol("truncated message".into()))?; + self.offset = end; + Ok(bytes) + } + + fn u16(&mut self) -> Result { + Ok(u16::from_le_bytes(self.bytes(2)?.try_into().unwrap())) + } + + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.bytes(4)?.try_into().unwrap())) + } + + fn u64(&mut self) -> Result { + Ok(u64::from_le_bytes(self.bytes(8)?.try_into().unwrap())) + } + + fn finish(self) -> Result<(), RerankError> { + if self.offset != self.bytes.len() { + return Err(RerankError::Protocol("trailing response bytes".into())); + } + Ok(()) + } +} + +pub(super) struct UnixSocketTransport { + endpoint: String, +} + +impl UnixSocketTransport { + pub fn new(endpoint: String) -> Self { + Self { endpoint } + } +} + +#[cfg(unix)] +impl TileMaxsimTransport for UnixSocketTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + max_response_bytes: usize, + ) -> Result, RerankError> { + use std::time::Instant; + + if self.endpoint.is_empty() { + return Err(RerankError::Transport("endpoint is empty".into())); + } + let deadline = Instant::now() + .checked_add(timeout) + .ok_or_else(|| RerankError::Transport("invalid timeout".into()))?; + let mut stream = connect_interruptible(&self.endpoint, deadline)?; + let poll = remaining_until(deadline)?.min(Duration::from_millis(50)); + stream + .set_read_timeout(Some(poll)) + .map_err(|error| RerankError::Transport(error.to_string()))?; + stream + .set_write_timeout(Some(poll)) + .map_err(|error| RerankError::Transport(error.to_string()))?; + + write_interruptible(&mut stream, request, deadline)?; + let mut header = [0u8; HEADER_LEN]; + read_interruptible(&mut stream, &mut header, deadline)?; + let body_len = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) + .map_err(|_| RerankError::Protocol("response body is too large".into()))?; + let response_len = HEADER_LEN + .checked_add(body_len) + .ok_or_else(|| RerankError::Protocol("response length overflow".into()))?; + if response_len > max_response_bytes { + return Err(RerankError::Protocol( + "response exceeds configured limit".into(), + )); + } + let mut response = Vec::with_capacity(response_len); + response.extend_from_slice(&header); + response.resize(response_len, 0); + read_interruptible(&mut stream, &mut response[HEADER_LEN..], deadline)?; + Ok(response) + } +} + +#[cfg(unix)] +fn connect_interruptible( + endpoint: &str, + deadline: std::time::Instant, +) -> Result { + use std::os::fd::{AsRawFd, FromRawFd, OwnedFd}; + + let (address, address_len) = unix_socket_address(endpoint)?; + let raw_fd = unsafe { libc::socket(libc::AF_UNIX, libc::SOCK_STREAM, 0) }; + if raw_fd < 0 { + return Err(last_transport_error()); + } + let fd = unsafe { OwnedFd::from_raw_fd(raw_fd) }; + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFD, + libc::F_SETFD, + libc::FD_CLOEXEC, + true, + )?; + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFL, + libc::F_SETFL, + libc::O_NONBLOCK, + true, + )?; + + let connected = unsafe { + libc::connect( + fd.as_raw_fd(), + (&raw const address).cast::(), + address_len, + ) + } == 0; + if !connected { + let error = std::io::Error::last_os_error(); + let raw_error = error.raw_os_error(); + if raw_error != Some(libc::EINPROGRESS) + && raw_error != Some(libc::EAGAIN) + && raw_error != Some(libc::EWOULDBLOCK) + { + return Err(RerankError::Transport(error.to_string())); + } + wait_for_connect(fd.as_raw_fd(), deadline)?; + } + + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFL, + libc::F_SETFL, + libc::O_NONBLOCK, + false, + )?; + Ok(std::os::unix::net::UnixStream::from(fd)) +} + +#[cfg(unix)] +fn unix_socket_address( + endpoint: &str, +) -> Result<(libc::sockaddr_un, libc::socklen_t), RerankError> { + let path = endpoint.as_bytes(); + let mut address = unsafe { std::mem::zeroed::() }; + if path.contains(&0) { + return Err(RerankError::Transport( + "endpoint contains a NUL byte".into(), + )); + } + if path.len() >= address.sun_path.len() { + return Err(RerankError::Transport("endpoint path is too long".into())); + } + address.sun_family = libc::AF_UNIX as libc::sa_family_t; + unsafe { + std::ptr::copy_nonoverlapping( + path.as_ptr().cast::(), + address.sun_path.as_mut_ptr(), + path.len(), + ); + } + let length = std::mem::offset_of!(libc::sockaddr_un, sun_path) + .checked_add(path.len()) + .and_then(|length| length.checked_add(1)) + .and_then(|length| libc::socklen_t::try_from(length).ok()) + .ok_or_else(|| RerankError::Transport("endpoint path is too long".into()))?; + #[cfg(any( + target_os = "aix", + target_os = "dragonfly", + target_os = "freebsd", + target_os = "haiku", + target_os = "hurd", + target_os = "ios", + target_os = "macos", + target_os = "netbsd", + target_os = "openbsd", + target_os = "tvos", + target_os = "visionos", + target_os = "watchos" + ))] + { + address.sun_len = u8::try_from(length) + .map_err(|_| RerankError::Transport("endpoint path is too long".into()))?; + } + Ok((address, length)) +} + +#[cfg(unix)] +fn update_fd_flag( + fd: std::os::fd::RawFd, + get_command: libc::c_int, + set_command: libc::c_int, + flag: libc::c_int, + enabled: bool, +) -> Result<(), RerankError> { + let current = unsafe { libc::fcntl(fd, get_command) }; + if current < 0 { + return Err(last_transport_error()); + } + let updated = if enabled { + current | flag + } else { + current & !flag + }; + if unsafe { libc::fcntl(fd, set_command, updated) } < 0 { + return Err(last_transport_error()); + } + Ok(()) +} + +#[cfg(unix)] +fn wait_for_connect( + fd: std::os::fd::RawFd, + deadline: std::time::Instant, +) -> Result<(), RerankError> { + loop { + pgrx::check_for_interrupts!(); + let remaining = remaining_until(deadline)?; + let timeout_ms = remaining.min(Duration::from_millis(50)).as_millis().max(1) as libc::c_int; + let mut poll_fd = libc::pollfd { + fd, + events: libc::POLLOUT, + revents: 0, + }; + let result = unsafe { libc::poll(&mut poll_fd, 1, timeout_ms) }; + if result == 0 { + continue; + } + if result < 0 { + let error = std::io::Error::last_os_error(); + if error.kind() == std::io::ErrorKind::Interrupted { + continue; + } + return Err(RerankError::Transport(error.to_string())); + } + let mut socket_error = 0; + let mut socket_error_len = size_of_val(&socket_error) as libc::socklen_t; + if unsafe { + libc::getsockopt( + fd, + libc::SOL_SOCKET, + libc::SO_ERROR, + (&raw mut socket_error).cast(), + &raw mut socket_error_len, + ) + } < 0 + { + return Err(last_transport_error()); + } + if socket_error != 0 { + return Err(RerankError::Transport( + std::io::Error::from_raw_os_error(socket_error).to_string(), + )); + } + return Ok(()); + } +} + +#[cfg(unix)] +fn remaining_until(deadline: std::time::Instant) -> Result { + deadline + .checked_duration_since(std::time::Instant::now()) + .filter(|remaining| !remaining.is_zero()) + .ok_or_else(|| RerankError::Transport("request timed out".into())) +} + +#[cfg(unix)] +fn last_transport_error() -> RerankError { + RerankError::Transport(std::io::Error::last_os_error().to_string()) +} + +#[cfg(unix)] +fn write_interruptible( + stream: &mut std::os::unix::net::UnixStream, + mut bytes: &[u8], + deadline: std::time::Instant, +) -> Result<(), RerankError> { + use std::io::Write; + + while !bytes.is_empty() { + match stream.write(bytes) { + Ok(0) => return Err(RerankError::Transport("connection closed".into())), + Ok(count) => bytes = &bytes[count..], + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::Interrupted + | std::io::ErrorKind::WouldBlock + | std::io::ErrorKind::TimedOut + ) => {} + Err(error) => return Err(RerankError::Transport(error.to_string())), + } + pgrx::check_for_interrupts!(); + if std::time::Instant::now() >= deadline { + return Err(RerankError::Transport("request timed out".into())); + } + } + Ok(()) +} + +#[cfg(unix)] +fn read_interruptible( + stream: &mut std::os::unix::net::UnixStream, + mut bytes: &mut [u8], + deadline: std::time::Instant, +) -> Result<(), RerankError> { + use std::io::Read; + + while !bytes.is_empty() { + match stream.read(bytes) { + Ok(0) => return Err(RerankError::Transport("connection closed".into())), + Ok(count) => bytes = &mut bytes[count..], + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::Interrupted + | std::io::ErrorKind::WouldBlock + | std::io::ErrorKind::TimedOut + ) => {} + Err(error) => return Err(RerankError::Transport(error.to_string())), + } + pgrx::check_for_interrupts!(); + if std::time::Instant::now() >= deadline { + return Err(RerankError::Transport("request timed out".into())); + } + } + Ok(()) +} + +#[cfg(not(unix))] +impl TileMaxsimTransport for UnixSocketTransport { + fn round_trip( + &mut self, + _request: &[u8], + _timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + if self.endpoint.is_empty() { + return Err(RerankError::Transport("endpoint is empty".into())); + } + Err(RerankError::Transport( + "Unix sockets are not supported on this platform".into(), + )) + } +} + +#[cfg(test)] +mod tests { + use super::super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorDtype, + }; + use super::super::rerank::CandidateTensor; + use super::*; + use std::collections::BTreeMap; + use vector::vect::VectOwned; + + struct MockTensorSource(BTreeMap>); + + impl CandidateTensorSource for MockTensorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some(CandidateTensor { + candidate, + vectors: self + .0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + })) + } + } + + struct MockDescriptorSource(BTreeMap); + + impl CandidateTensorDescriptorSource for MockDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some( + self.0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + )) + } + } + + struct MockTransport { + similarities: Vec<(u32, f32)>, + } + + impl TileMaxsimTransport for MockTransport { + fn round_trip( + &mut self, + request: &[u8], + _timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + let request_id = u64::from_le_bytes(request[8..16].try_into().unwrap()); + let version = u16::from_le_bytes(request[4..6].try_into().unwrap()); + Ok(success_response_with_version( + version, + request_id, + &self.similarities, + )) + } + } + + fn vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf32(VectOwned::new(values.to_vec())) + } + + fn half_vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf16(VectOwned::new( + values.iter().copied().map(simd::f16::from_f32).collect(), + )) + } + + fn success_response(request_id: u64, similarities: &[(u32, f32)]) -> Vec { + success_response_with_version(VERSION, request_id, similarities) + } + + fn success_response_with_version( + version: u16, + request_id: u64, + similarities: &[(u32, f32)], + ) -> Vec { + let body_len = 8 + similarities.len() * 8; + let mut response = Vec::with_capacity(HEADER_LEN + body_len); + response.extend_from_slice(MAGIC); + response.extend_from_slice(&version.to_le_bytes()); + response.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + response.extend_from_slice(&request_id.to_le_bytes()); + response.extend_from_slice(&(body_len as u64).to_le_bytes()); + response.extend_from_slice(&0u32.to_le_bytes()); + response.extend_from_slice(&(similarities.len() as u32).to_le_bytes()); + for (candidate_id, similarity) in similarities { + response.extend_from_slice(&candidate_id.to_le_bytes()); + response.extend_from_slice(&similarity.to_bits().to_le_bytes()); + } + response + } + + fn external_descriptor( + candidate: PageCandidate, + public_id: i64, + tensor_ref: &str, + rows: u32, + dimension: u32, + dtype: ExternalTensorDtype, + ) -> ExternalTensorDescriptor { + ExternalTensorDescriptor { + candidate, + public_id, + tensor_ref: tensor_ref.into(), + rows, + dimension, + dtype, + checksum: format!("sha256:{}", "a".repeat(64)), + } + } + + fn error_response(request_id: u64, message: &str) -> Vec { + let body_len = 8 + message.len(); + let mut response = Vec::with_capacity(HEADER_LEN + body_len); + response.extend_from_slice(MAGIC); + response.extend_from_slice(&VERSION.to_le_bytes()); + response.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + response.extend_from_slice(&request_id.to_le_bytes()); + response.extend_from_slice(&(body_len as u64).to_le_bytes()); + response.extend_from_slice(&1u32.to_le_bytes()); + response.extend_from_slice(&(message.len() as u32).to_le_bytes()); + response.extend_from_slice(message.as_bytes()); + response + } + + #[test] + fn gpu_backend_maps_positive_similarity_to_ascending_distance() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page_1, + }, + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page_2, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0])]), + (page_2, vec![vector(&[0.5, 0.0])]), + ])); + let transport = MockTransport { + similarities: vec![(0, 1.0), (1, 2.0)], + }; + let results = GpuTileMaxsimBackend::new(transport, Duration::from_secs(1), 100, 4096) + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results[0].heap_key, page_2); + assert_eq!(results[0].distance.to_f32(), -2.0); + assert_eq!(results[1].heap_key, page_1); + assert_eq!(results[1].distance.to_f32(), -1.0); + } + + #[test] + fn response_rejects_partial_and_duplicate_results() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let partial = success_response(7, &[(0, 1.0)]); + assert!(matches!( + decode_response(&partial, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let duplicate = success_response(7, &[(0, 1.0), (0, 2.0)]); + assert!(matches!( + decode_response(&duplicate, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + + #[test] + fn response_ids_may_arrive_out_of_order() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let response = success_response(7, &[(1, 0.5), (0, 1.0)]); + let results = decode_response(&response, 7, &keys) + .unwrap() + .collect::>(); + + assert_eq!(results[0].heap_key, keys[0]); + assert_eq!(results[0].distance.to_f32(), -1.0); + assert_eq!(results[1].heap_key, keys[1]); + assert_eq!(results[1].distance.to_f32(), -0.5); + } + + #[test] + fn response_rejects_unknown_non_finite_and_trailing_results() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let unknown = success_response(7, &[(0, 1.0), (2, 2.0)]); + assert!(matches!( + decode_response(&unknown, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let non_finite = success_response(7, &[(0, f32::NAN), (1, 2.0)]); + assert!(matches!( + decode_response(&non_finite, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let mut trailing = success_response(7, &[(0, 1.0), (1, 2.0)]); + trailing.push(0); + assert!(matches!( + decode_response(&trailing, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + + #[test] + fn response_rejects_invalid_header_fields() { + let keys = [[0, 0, 1]]; + let valid = success_response(7, &[(0, 1.0)]); + + for (offset, replacement) in [(0, 0u8), (4, 2), (6, 1), (8, 8), (16, 0)] { + let mut invalid = valid.clone(); + invalid[offset] = replacement; + assert!(matches!( + decode_response(&invalid, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + } + + #[test] + fn response_surfaces_remote_error() { + let response = error_response(7, "CUDA queue is unavailable"); + assert!(matches!( + decode_response(&response, 7, &[]), + Err(RerankError::Remote(message)) if message == "CUDA queue is unavailable" + )); + } + + #[test] + fn request_frame_is_versioned_and_length_prefixed() { + let page = [0, 0, 1]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([(page, vec![vector(&[0.5, 0.0])])])); + let encoded = encode_request(9, &query, &mut candidates, &mut source, 100, 4096).unwrap(); + + assert_eq!(&encoded.frame[0..4], MAGIC); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + VERSION + ); + assert_eq!( + u16::from_le_bytes(encoded.frame[6..8].try_into().unwrap()), + REQUEST_KIND + ); + assert_eq!( + u64::from_le_bytes(encoded.frame[8..16].try_into().unwrap()), + 9 + ); + assert_eq!( + u64::from_le_bytes(encoded.frame[16..24].try_into().unwrap()) as usize, + encoded.frame.len() - HEADER_LEN + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[32..36].try_into().unwrap()), + 1 + ); + assert_eq!(encoded.heap_keys, vec![page]); + } + + #[test] + fn external_request_encodes_contract_and_opaque_descriptor_ids() { + let page = [0, 0, 7]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, -0.5])]; + let tensor_ref = "s3://immutable/page-9001.tensor"; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor(candidate, 9001, tensor_ref, 2, 2, ExternalTensorDtype::F16), + )])); + let contract = "colqwen@immutable-revision"; + let encoded = encode_external_request( + 19, + contract, + &query, + &mut candidates, + &mut source, + 100, + 4096, + ) + .unwrap(); + + assert_eq!(&encoded.frame[0..4], MAGIC); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + EXTERNAL_VERSION + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[32..36].try_into().unwrap()), + 1 + ); + let contract_len = u32::from_le_bytes(encoded.frame[40..44].try_into().unwrap()) as usize; + assert_eq!(&encoded.frame[44..44 + contract_len], contract.as_bytes()); + let candidate_offset = 44 + contract_len + 4; // one 2-D f16 query row + assert_eq!( + u32::from_le_bytes( + encoded.frame[candidate_offset..candidate_offset + 4] + .try_into() + .unwrap() + ), + 0 + ); + assert_eq!( + u32::from_le_bytes( + encoded.frame[candidate_offset + 4..candidate_offset + 8] + .try_into() + .unwrap() + ), + 2 + ); + let reference_len = u32::from_le_bytes( + encoded.frame[candidate_offset + 8..candidate_offset + 12] + .try_into() + .unwrap(), + ) as usize; + let reference_offset = candidate_offset + 16; + assert_eq!( + &encoded.frame[reference_offset..reference_offset + reference_len], + tensor_ref.as_bytes() + ); + assert_eq!(encoded.heap_keys, vec![page]); + } + + #[test] + fn external_backend_maps_scores_without_exposing_public_ids() { + let page = [0, 0, 8]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + i64::MAX, + "object://immutable/tensor", + 4, + 2, + ExternalTensorDtype::F32, + ), + )])); + let transport = MockTransport { + similarities: vec![(0, 3.5)], + }; + let results = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_secs(1), + 100, + 4096, + ) + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].heap_key, page); + assert_eq!(results[0].distance.to_f32(), -3.5); + } + + #[test] + fn external_request_rejects_shape_and_declared_payload_overflow() { + let page = [0, 0, 9]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, 0.0])]; + let make_source = |dimension, rows| { + MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + 9, + "object://immutable/tensor", + rows, + dimension, + ExternalTensorDtype::F16, + ), + )])) + }; + + let mut candidates = vec![candidate].into_iter(); + assert!(matches!( + encode_external_request( + 1, + "contract@1", + &query, + &mut candidates, + &mut make_source(3, 1), + 100, + 4096, + ), + Err(RerankError::TensorMismatch) + )); + + let mut candidates = vec![candidate].into_iter(); + assert!(matches!( + encode_external_request( + 1, + "contract@1", + &query, + &mut candidates, + &mut make_source(2, 100), + 1000, + 64, + ), + Err(RerankError::RequestTooLarge) + )); + } + + #[test] + fn request_frame_encodes_f16_tensor_bits() { + let page = [0, 0, 1]; + let query = vec![half_vector(&[1.0, -0.5])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = + MockTensorSource(BTreeMap::from([(page, vec![half_vector(&[0.25, 2.0])])])); + let encoded = encode_request(9, &query, &mut candidates, &mut source, 100, 4096).unwrap(); + + assert_eq!(encoded.frame[36], TensorDtype::F16 as u8); + assert_eq!( + u16::from_le_bytes(encoded.frame[40..42].try_into().unwrap()), + simd::f16::from_f32(1.0).to_bits() + ); + assert_eq!( + u16::from_le_bytes(encoded.frame[42..44].try_into().unwrap()), + simd::f16::from_f32(-0.5).to_bits() + ); + } + + #[test] + fn request_limits_are_enforced_before_transport() { + let page = [0, 0, 1]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([(page, vec![vector(&[1.0, 0.0])])])); + assert!(matches!( + encode_request(1, &query, &mut candidates, &mut source, 1, 4096), + Err(RerankError::RequestTooLarge) + )); + } + + #[cfg(unix)] + #[test] + fn unix_socket_address_is_length_bounded_and_nul_terminated() { + let (address, length) = unix_socket_address("/tmp/vectorchord.sock").unwrap(); + let path_offset = std::mem::offset_of!(libc::sockaddr_un, sun_path); + + assert_eq!(address.sun_family, libc::AF_UNIX as libc::sa_family_t); + assert_eq!( + length as usize, + path_offset + "/tmp/vectorchord.sock".len() + 1 + ); + assert_eq!( + &address.sun_path[.."/tmp/vectorchord.sock".len()], + "/tmp/vectorchord.sock" + .as_bytes() + .iter() + .map(|byte| *byte as libc::c_char) + .collect::>() + ); + assert_eq!(address.sun_path["/tmp/vectorchord.sock".len()], 0); + + let too_long = "x".repeat(address.sun_path.len()); + assert!(matches!( + unix_socket_address(&too_long), + Err(RerankError::Transport(message)) if message == "endpoint path is too long" + )); + assert!(matches!( + unix_socket_address("invalid\0path"), + Err(RerankError::Transport(message)) if message == "endpoint contains a NUL byte" + )); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/rerank.rs b/src/index/vchordrq/scanners/maxsim/rerank.rs new file mode 100644 index 00000000..151ae4ea --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/rerank.rs @@ -0,0 +1,315 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::{HeapKey, PageCandidate}; +use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; +use crate::index::vchordrq::opclass::Opfamily; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use vchordrq::types::OwnedVector; + +pub(super) struct CandidateTensor { + pub candidate: PageCandidate, + pub vectors: Vec, +} + +pub(super) trait CandidateTensorSource { + fn fetch(&mut self, candidate: PageCandidate) -> Result, RerankError>; +} + +pub(super) struct HeapArrayTensorSource<'a, F> { + fetcher: &'a mut F, + opfamily: Opfamily, +} + +impl<'a, F> HeapArrayTensorSource<'a, F> { + pub fn new(fetcher: &'a mut F, opfamily: Opfamily) -> Self { + Self { fetcher, opfamily } + } +} + +impl CandidateTensorSource for HeapArrayTensorSource<'_, F> { + fn fetch(&mut self, candidate: PageCandidate) -> Result, RerankError> { + let Some(mut tuple) = self.fetcher.fetch(candidate.heap_key) else { + return Ok(None); + }; + // Exact sources are the last boundary before a tensor may leave the + // executor process. Re-evaluate the active base-relation scan qual + // here even if token reranking already prefiltered some hits. This + // keeps same-relation quals ahead of CPU/GPU tensor access and + // also covers configurations with token-level refine disabled. + if !tuple.filter() { + return Ok(None); + } + let (values, is_nulls) = tuple.build(); + if is_nulls[0] { + return Err(RerankError::TensorMismatch); + } + let vectors = + unsafe { self.opfamily.input_vectors(values[0]) }.ok_or(RerankError::TensorMismatch)?; + Ok(Some(CandidateTensor { candidate, vectors })) + } +} + +pub(super) trait ExactMaxsimBackend { + type Results: Iterator; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result; +} + +#[derive(Debug)] +pub(super) enum RerankError { + TensorMismatch, + ModelContractMismatch, + InvalidDescriptor(&'static str), + Registry(String), + Configuration(&'static str), + UnsupportedTensorKind, + RequestTooLarge, + Transport(String), + Protocol(String), + Remote(String), +} + +impl std::fmt::Display for RerankError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::TensorMismatch => write!(f, "MaxSim tensor kind or dimension is not matched"), + Self::ModelContractMismatch => write!(f, "MaxSim model contract is not matched"), + Self::InvalidDescriptor(message) => { + write!(f, "invalid external MaxSim tensor descriptor: {message}") + } + Self::Registry(message) => { + write!(f, "MaxSim tensor-source registry error: {message}") + } + Self::Configuration(message) => write!(f, "MaxSim configuration error: {message}"), + Self::UnsupportedTensorKind => { + write!(f, "GPU MaxSim supports only vector and halfvec tensors") + } + Self::RequestTooLarge => write!(f, "GPU MaxSim request exceeds configured limits"), + Self::Transport(message) => write!(f, "GPU MaxSim transport error: {message}"), + Self::Protocol(message) => write!(f, "GPU MaxSim protocol error: {message}"), + Self::Remote(message) => write!(f, "GPU MaxSim sidecar error: {message}"), + } + } +} + +impl std::error::Error for RerankError {} + +#[derive(Default)] +pub(super) struct CpuExactMaxsimBackend; + +impl ExactMaxsimBackend for CpuExactMaxsimBackend { + type Results = RerankResults; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let mut results = BinaryHeap::new(); + for candidate in candidates { + let Some(tensor) = source.fetch(candidate)? else { + continue; + }; + let Some(distance) = exact_maxsim_distance(query, &tensor.vectors) else { + return Err(RerankError::TensorMismatch); + }; + results.push((Reverse(distance), Reverse(tensor.candidate.heap_key))); + } + Ok(RerankResults { inner: results }) + } +} + +fn exact_maxsim_distance(query: &[OwnedVector], document: &[OwnedVector]) -> Option { + if query.is_empty() || document.is_empty() { + return None; + } + let mut maxsim = 0.0f32; + for query_vector in query { + let mut best = Distance::INFINITY; + for document_vector in document { + let distance = document_vector.operator_dot(query_vector)?; + best = std::cmp::min(best, distance); + } + maxsim += best.to_f32(); + } + Some(Distance::from_f32(maxsim)) +} + +#[derive(Clone, Copy, Debug)] +pub(super) struct RerankedPage { + pub distance: Distance, + pub heap_key: HeapKey, +} + +pub(super) struct RerankResults { + pub(super) inner: BinaryHeap<(Reverse, Reverse)>, +} + +impl Iterator for RerankResults { + type Item = RerankedPage; + + fn next(&mut self) -> Option { + let (Reverse(distance), Reverse(heap_key)) = self.inner.pop()?; + Some(RerankedPage { distance, heap_key }) + } + + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len(); + (exact, Some(exact)) + } +} + +impl ExactSizeIterator for RerankResults {} +impl std::iter::FusedIterator for RerankResults {} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeMap; + use vector::vect::VectOwned; + + struct MockTensorSource(BTreeMap>); + + impl CandidateTensorSource for MockTensorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some(CandidateTensor { + candidate, + vectors: self + .0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + })) + } + } + + fn vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf32(VectOwned::new(values.to_vec())) + } + + struct RejectingFetcher; + + struct RejectingTuple; + + impl Tuple for RejectingTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + panic!("a rejected tuple must not materialize its tensor") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + panic!("a rejected tuple must not expose heap attributes") + } + } + + impl FilterableTuple for RejectingTuple { + fn filter(&mut self) -> bool { + false + } + } + + impl Fetcher for RejectingFetcher { + type Tuple<'a> = RejectingTuple; + + fn fetch(&mut self, _key: HeapKey) -> Option> { + Some(RejectingTuple) + } + } + + #[test] + fn heap_source_applies_scan_qual_before_materializing_tensor() { + let mut fetcher = RejectingFetcher; + let mut source = HeapArrayTensorSource::new(&mut fetcher, Opfamily::VectorMaxsim); + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + }; + + assert!(source.fetch(candidate).unwrap().is_none()); + } + + #[test] + fn cpu_backend_reorders_candidates_by_exact_page_maxsim() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::from_f32(-2.0), + heap_key: page_2, + }, + PageCandidate { + approximate_distance: Distance::from_f32(-1.0), + heap_key: page_1, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]), + (page_2, vec![vector(&[0.5, 0.5])]), + ])); + let results = CpuExactMaxsimBackend + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].heap_key, page_1); + assert_eq!(results[0].distance.to_f32(), -2.0); + assert_eq!(results[1].heap_key, page_2); + assert_eq!(results[1].distance.to_f32(), -1.0); + } + + #[test] + fn exact_ties_are_ordered_by_heap_key() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::from_f32(0.0), + heap_key: page_2, + }, + PageCandidate { + approximate_distance: Distance::from_f32(0.0), + heap_key: page_1, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0])]), + (page_2, vec![vector(&[1.0, 0.0])]), + ])); + + let results = CpuExactMaxsimBackend + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!( + results.iter().map(|page| page.heap_key).collect::>(), + vec![page_1, page_2] + ); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/search.rs b/src/index/vchordrq/scanners/maxsim/search.rs new file mode 100644 index 00000000..915fb87f --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/search.rs @@ -0,0 +1,578 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::MaxsimBuilder; +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorSourceBinding, + ExternalTensorStorage, resolve_external_tensor_source, validate_descriptor, +}; +use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; +use super::rerank::RerankError; +use crate::index::fetcher::{ + Fetcher, FilterableTuple, HeapFetcher, Tuple, TupleAttribute, ctid_to_key, +}; +use crate::index::gucs::{self, PostgresMaxsimBackend}; +use crate::index::scanners::SearchBuilder; +use crate::index::storage::PostgresRelation; +use crate::index::vchordrq::opclass::{Opfamily, opfamily}; +use crate::index::vchordrq::scanners::SearchOptions; +use crate::recorder::DefaultRecorder; +use distance::Distance; +use pgrx::datum::{DatumWithOid, FromDatum}; +use pgrx::iter::TableIterator; +use pgrx::{AnyArray, IntoDatum, name}; +use std::collections::{BTreeMap, BTreeSet}; +use std::time::Duration; + +const MAX_EXPLICIT_CANDIDATES: i32 = 65_536; + +/// Restricted Phase 3B search surface. Candidate generation reads only the +/// named index. Descriptor projection happens later through SPI, under the +/// caller's normal SELECT privileges and active MVCC snapshot. +#[pgrx::pg_extern(sql = "")] +fn _vchordrq_maxsim_search_external( + index_oid: pgrx::pg_sys::Oid, + query: AnyArray, + candidate_limit: i32, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = execute_external_search(index_oid, query, candidate_limit, top_k) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +fn execute_external_search( + index_oid: pgrx::pg_sys::Oid, + query: AnyArray, + candidate_limit: i32, + top_k: i32, +) -> Result, RerankError> { + validate_search_limits(candidate_limit, top_k)?; + if !matches!(gucs::vchordrq_maxsim_backend(), PostgresMaxsimBackend::Gpu) { + return Err(RerankError::Configuration( + "external MaxSim search currently requires vchordrq.maxsim_backend = 'gpu'", + )); + } + + let binding = resolve_external_tensor_source(index_oid)?; + let index_lock = RelationLock::open(index_oid, pgrx::pg_sys::AccessShareLock as _)?; + let heap_lock = RelationLock::open(binding.heap_oid, pgrx::pg_sys::AccessShareLock as _)?; + let descriptor_lock = binding + .descriptor_oid + .map(|oid| RelationLock::open(oid, pgrx::pg_sys::AccessShareLock as _)) + .transpose()?; + if binding.index_oid != index_lock.oid() || binding.heap_oid != heap_lock.oid() { + return Err(RerankError::Registry( + "registered MaxSim tensor source changed during execution".into(), + )); + } + if descriptor_lock.as_ref().map(RelationLock::oid) != binding.descriptor_oid { + return Err(RerankError::Registry( + "registered descriptor relation changed during execution".into(), + )); + } + + let opfamily = unsafe { opfamily(index_lock.raw()) }; + if !matches!(opfamily, Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim) { + return Err(RerankError::UnsupportedTensorKind); + } + let indexed_type = unsafe { pgrx::pg_sys::get_atttype(index_oid, 1) }; + if indexed_type != query.oid() { + return Err(RerankError::TensorMismatch); + } + let query_vectors = + unsafe { opfamily.input_vectors(query.datum()) }.ok_or(RerankError::TensorMismatch)?; + + // The registry resolution happens before the index read, and this + // privilege-only SELECT is planned/executed before any candidate CTID is + // generated. It fails early when the caller cannot project the registered + // descriptor columns. The same query shape is used for the actual fetch. + preflight_descriptor_access(&binding)?; + + let candidates = generate_candidates( + index_lock.raw(), + opfamily, + query.datum(), + candidate_limit as u32, + )?; + let resolved = + resolve_visible_candidates(index_lock.raw(), heap_lock.raw(), candidates.into_iter())?; + let (mut source, public_ids) = load_visible_descriptors(&binding, &resolved)?; + let visible_candidates = resolved + .into_iter() + .filter_map(|resolved| { + public_ids + .contains_key(&resolved.candidate.heap_key) + .then_some(resolved.candidate) + }) + .collect::>(); + let endpoint = gucs::vchordrq_maxsim_gpu_endpoint() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let transport = UnixSocketTransport::new(endpoint); + let mut backend = GpuExternalTileMaxsimBackend::new( + transport, + binding.model_contract_id, + Duration::from_millis(gucs::vchordrq_maxsim_gpu_timeout_ms() as u64), + gucs::vchordrq_maxsim_gpu_max_batch_tokens() as usize, + gucs::vchordrq_maxsim_gpu_max_batch_bytes() as usize, + ); + let mut candidate_iter = visible_candidates.into_iter(); + let exact = backend.rerank(&query_vectors, &mut candidate_iter, &mut source)?; + let mut rows = exact + .map(|result| { + let public_id = public_ids.get(&result.heap_key).copied().ok_or_else(|| { + RerankError::Protocol("sidecar result has no visible public ID".into()) + })?; + Ok((result.distance, public_id)) + }) + .collect::, RerankError>>()?; + rows.sort_unstable_by(|(left_distance, left_id), (right_distance, right_id)| { + left_distance + .cmp(right_distance) + .then_with(|| left_id.cmp(right_id)) + }); + rows.truncate(top_k as usize); + Ok(rows + .into_iter() + .map(|(distance, public_id)| (public_id, -distance.to_f32())) + .collect::>() + .into_iter()) +} + +#[derive(Clone, Copy)] +struct ResolvedCandidate { + candidate: PageCandidate, + current_ctid: pgrx::pg_sys::ItemPointerData, +} + +fn resolve_visible_candidates( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + candidates: impl Iterator, +) -> Result, RerankError> { + let snapshot = unsafe { pgrx::pg_sys::GetActiveSnapshot() }; + if snapshot.is_null() { + return Err(RerankError::Configuration( + "external MaxSim search requires an active MVCC snapshot", + )); + } + if unsafe { (*snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC { + return Err(RerankError::Configuration( + "external MaxSim search requires an MVCC snapshot", + )); + } + + let mut fetcher = + unsafe { HeapFetcher::new_standalone(index_relation, heap_relation, snapshot) }; + let mut resolved = Vec::new(); + for candidate in candidates { + pgrx::check_for_interrupts!(); + if let Some(tuple) = fetcher.fetch(candidate.heap_key) { + resolved.push(ResolvedCandidate { + candidate, + current_ctid: tuple.ctid(), + }); + } + } + Ok(resolved) +} + +fn validate_search_limits(candidate_limit: i32, top_k: i32) -> Result<(), RerankError> { + if !(1..=MAX_EXPLICIT_CANDIDATES).contains(&candidate_limit) { + return Err(RerankError::Configuration( + "candidate_limit must be between 1 and 65536", + )); + } + if top_k <= 0 || top_k > candidate_limit { + return Err(RerankError::Configuration( + "top_k must be positive and no greater than candidate_limit", + )); + } + Ok(()) +} + +fn generate_candidates( + index_relation: pgrx::pg_sys::Relation, + opfamily: Opfamily, + query: pgrx::pg_sys::Datum, + candidate_limit: u32, +) -> Result, RerankError> { + let index = unsafe { PostgresRelation::::new(index_relation) }; + let mut builder = MaxsimBuilder::new(opfamily); + unsafe { builder.add(3, Some(query)) }; + let options = SearchOptions { + epsilon: unsafe { gucs::vchordrq_epsilon(index_relation) }, + probes: unsafe { gucs::vchordrq_probes(index_relation) }, + max_scan_tuples: None, + maxsim_refine: gucs::vchordrq_maxsim_refine(index_relation), + maxsim_threshold: gucs::vchordrq_maxsim_threshold(index_relation), + maxsim_candidate_limit: Some(candidate_limit), + maxsim_backend: PostgresMaxsimBackend::CoarseOnly, + maxsim_gpu_endpoint: None, + maxsim_gpu_timeout_ms: 1, + maxsim_gpu_max_batch_tokens: 1, + maxsim_gpu_max_batch_bytes: 1, + io_search: gucs::vchordrq_io_search(), + io_rerank: gucs::vchordrq_io_rerank(), + // General same-relation quals would require the optional Phase 3C + // CustomScan. The restricted function applies PostgreSQL row + // visibility during the following SPI descriptor fetch. + prefilter: false, + }; + let bump = bumpalo::Bump::new(); + let recorder = DefaultRecorder { + enable: false, + rate: None, + max_records: 0, + index: unsafe { (*index_relation).rd_id.to_u32() }, + }; + let candidates = builder + .build(&index, options, NoHeapFetcher, &bump, recorder) + .map(|(distance, heap_key, _)| PageCandidate { + approximate_distance: Distance::from_f32(distance), + heap_key, + }) + .collect(); + Ok(candidates) +} + +fn preflight_descriptor_access(binding: &ExternalTensorSourceBinding) -> Result<(), RerankError> { + let query = descriptor_query(binding, true)?; + pgrx::spi::Spi::connect(|client| { + let privilege = client + .prepare( + "SELECT pg_catalog.has_table_privilege($1, 'SELECT') AS allowed", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + for relation_oid in std::iter::once(binding.heap_oid).chain(binding.descriptor_oid) { + let allowed = client + .select(&privilege, Some(1), &[relation_oid.into()]) + .map_err(registry_error)? + .first() + .get_by_name::("allowed") + .map_err(registry_error)? + .unwrap_or(false); + if !allowed { + return Err(RerankError::Registry( + "table-level SELECT privilege is required for every external MaxSim relation" + .into(), + )); + } + } + client + .select(query.as_str(), Some(1), &[]) + .map(|_| ()) + .map_err(registry_error) + }) +} + +fn load_visible_descriptors( + binding: &ExternalTensorSourceBinding, + candidates: &[ResolvedCandidate], +) -> Result<(MaterializedDescriptorSource, BTreeMap), RerankError> { + if candidates.is_empty() { + return Ok((MaterializedDescriptorSource::default(), BTreeMap::new())); + } + let mut candidates_by_key = BTreeMap::new(); + for resolved in candidates { + if candidates_by_key + .insert(ctid_to_key(resolved.current_ctid), resolved.candidate) + .is_some() + { + return Err(RerankError::Registry( + "multiple index candidates resolved to the same visible CTID".into(), + )); + } + } + let ctids = candidates + .iter() + .map(|resolved| resolved.current_ctid) + .collect::>(); + let query = descriptor_query(binding, false)?; + let (descriptors, public_ids) = pgrx::spi::Spi::connect(|client| { + let tid_array_oid = unsafe { pgrx::pg_sys::get_array_type(pgrx::pg_sys::TIDOID) }; + let prepared = client + .prepare( + query.as_str(), + &[ + pgrx::pg_sys::PgOid::from(tid_array_oid), + pgrx::pg_sys::PgOid::from(pgrx::pg_sys::TEXTOID), + ], + ) + .map_err(registry_error)?; + let args: [DatumWithOid<'_>; 2] = [ctids.into(), binding.model_contract_id.clone().into()]; + let rows = client + .select(&prepared, Some(candidates.len() as _), &args) + .map_err(registry_error)?; + let mut descriptors = BTreeMap::new(); + let mut public_ids = BTreeMap::new(); + let mut unique_public_ids = BTreeSet::new(); + for row in rows { + let ctid = required_heap_column::(&row, "heap_tid")?; + let heap_key = ctid_to_key(ctid); + let candidate = candidates_by_key.get(&heap_key).copied().ok_or_else(|| { + RerankError::Registry("descriptor query returned an unknown CTID".into()) + })?; + let public_id = required_heap_column::(&row, "public_id")?; + if !unique_public_ids.insert(public_id) { + return Err(RerankError::InvalidDescriptor( + "public IDs are not unique in the visible candidate batch", + )); + } + let descriptor = validate_descriptor( + candidate, + public_id, + required_heap_column::(&row, "tensor_ref")?, + required_heap_column::(&row, "tensor_rows")?, + required_heap_column::(&row, "tensor_dimension")?, + required_heap_column::(&row, "tensor_dtype")?, + required_heap_column::(&row, "tensor_checksum")?, + )?; + if descriptors.insert(candidate.heap_key, descriptor).is_some() { + return Err(RerankError::Registry( + "descriptor query returned a duplicate CTID".into(), + )); + } + public_ids.insert(candidate.heap_key, public_id); + } + Ok((descriptors, public_ids)) + })?; + Ok((MaterializedDescriptorSource(descriptors), public_ids)) +} + +fn descriptor_query( + binding: &ExternalTensorSourceBinding, + preflight: bool, +) -> Result { + let heap_relation = relation_name(binding.heap_oid)?; + let names = &binding.column_names; + let model_contract = pgrx::spi::quote_identifier(&names.model_contract); + let public_id = pgrx::spi::quote_identifier(&names.public_id); + let tensor_ref = pgrx::spi::quote_identifier(&names.tensor_ref); + let tensor_rows = pgrx::spi::quote_identifier(&names.tensor_rows); + let tensor_dimension = pgrx::spi::quote_identifier(&names.tensor_dimension); + let tensor_dtype = pgrx::spi::quote_identifier(&names.tensor_dtype); + let tensor_checksum = pgrx::spi::quote_identifier(&names.tensor_checksum); + let predicate = if preflight { + "false".to_string() + } else { + format!("h.ctid = ANY($1) AND h.{model_contract} = $2") + }; + match binding.storage { + ExternalTensorStorage::SameHeap => Ok(format!( + "SELECT h.ctid AS heap_tid, + h.{model_contract} AS model_contract, + h.{public_id} AS public_id, + h.{tensor_ref} AS tensor_ref, + h.{tensor_rows} AS tensor_rows, + h.{tensor_dimension} AS tensor_dimension, + h.{tensor_dtype} AS tensor_dtype, + h.{tensor_checksum} AS tensor_checksum + FROM ONLY {heap_relation} AS h + WHERE {predicate}" + )), + ExternalTensorStorage::DescriptorRelation => { + let descriptor_oid = binding.descriptor_oid.ok_or_else(|| { + RerankError::Registry("registered descriptor relation is missing".into()) + })?; + let descriptor_relation = relation_name(descriptor_oid)?; + let descriptor_public_id = pgrx::spi::quote_identifier( + names.descriptor_public_id.as_deref().ok_or_else(|| { + RerankError::Registry( + "registered descriptor public ID column is missing".into(), + ) + })?, + ); + Ok(format!( + "SELECT h.ctid AS heap_tid, + h.{model_contract} AS model_contract, + h.{public_id} AS public_id, + d.{tensor_ref} AS tensor_ref, + d.{tensor_rows} AS tensor_rows, + d.{tensor_dimension} AS tensor_dimension, + d.{tensor_dtype} AS tensor_dtype, + d.{tensor_checksum} AS tensor_checksum + FROM ONLY {heap_relation} AS h + LEFT JOIN ONLY {descriptor_relation} AS d + ON d.{descriptor_public_id} = h.{public_id} + WHERE {predicate}" + )) + } + } +} + +fn relation_name(relation_oid: pgrx::pg_sys::Oid) -> Result { + pgrx::spi::Spi::connect(|client| { + let prepared = client + .prepare( + "SELECT n.nspname::text AS schema_name, c.relname::text AS relation_name + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE c.oid = $1", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[relation_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered relation disappeared".into(), + )); + } + let row = rows.first(); + Ok(pgrx::spi::quote_qualified_identifier( + required_column::(&row, "schema_name")?, + required_column::(&row, "relation_name")?, + )) + }) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn required_heap_column( + row: &pgrx::spi::SpiHeapTupleData<'_>, + name: &str, +) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +#[derive(Default)] +struct MaterializedDescriptorSource(BTreeMap); + +impl CandidateTensorDescriptorSource for MaterializedDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(self.0.remove(&candidate.heap_key)) + } +} + +struct RelationLock { + raw: pgrx::pg_sys::Relation, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl RelationLock { + fn open(oid: pgrx::pg_sys::Oid, lockmode: pgrx::pg_sys::LOCKMODE) -> Result { + let raw = unsafe { pgrx::pg_sys::relation_open(oid, lockmode) }; + if raw.is_null() { + return Err(RerankError::Registry("relation open returned NULL".into())); + } + Ok(Self { raw, lockmode }) + } + + fn raw(&self) -> pgrx::pg_sys::Relation { + self.raw + } + + fn oid(&self) -> pgrx::pg_sys::Oid { + unsafe { (*self.raw).rd_id } + } +} + +impl Drop for RelationLock { + fn drop(&mut self) { + unsafe { pgrx::pg_sys::relation_close(self.raw, self.lockmode) }; + } +} + +struct NoHeapFetcher; +struct NoHeapTuple; + +impl Tuple for NoHeapTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + unreachable!("restricted external candidate generation must not read the heap") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +impl FilterableTuple for NoHeapTuple { + fn filter(&mut self) -> bool { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +impl Fetcher for NoHeapFetcher { + type Tuple<'a> = NoHeapTuple; + + fn fetch(&mut self, _key: HeapKey) -> Option> { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn explicit_search_limits_are_bounded() { + assert!(validate_search_limits(256, 10).is_ok()); + for (candidates, top_k) in [(0, 1), (65_537, 1), (1, 0), (10, 11)] { + assert!(matches!( + validate_search_limits(candidates, top_k), + Err(RerankError::Configuration(_)) + )); + } + } + + #[test] + fn materialized_source_only_returns_visible_descriptors_once() { + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + }; + let descriptor = validate_descriptor( + candidate, + 42, + "s3://immutable/tensor".into(), + 1, + 2, + "float16".into(), + format!("sha256:{}", "a".repeat(64)), + ) + .unwrap(); + let mut source = + MaterializedDescriptorSource(BTreeMap::from([(candidate.heap_key, descriptor)])); + assert_eq!(source.fetch(candidate).unwrap().unwrap().public_id, 42); + assert!(source.fetch(candidate).unwrap().is_none()); + } +} diff --git a/src/index/vchordrq/scanners/mod.rs b/src/index/vchordrq/scanners/mod.rs index b345da37..c5e392bc 100644 --- a/src/index/vchordrq/scanners/mod.rs +++ b/src/index/vchordrq/scanners/mod.rs @@ -15,7 +15,9 @@ mod default; mod maxsim; +use crate::index::gucs::PostgresMaxsimBackend; use crate::index::scanners::Io; +use std::ffi::CString; pub use default::DefaultBuilder; pub use maxsim::MaxsimBuilder; @@ -27,6 +29,12 @@ pub struct SearchOptions { pub max_scan_tuples: Option, pub maxsim_refine: u32, pub maxsim_threshold: u32, + pub maxsim_candidate_limit: Option, + pub maxsim_backend: PostgresMaxsimBackend, + pub maxsim_gpu_endpoint: Option, + pub maxsim_gpu_timeout_ms: u32, + pub maxsim_gpu_max_batch_tokens: u32, + pub maxsim_gpu_max_batch_bytes: u32, pub io_search: Io, pub io_rerank: Io, pub prefilter: bool, diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index a4f8f38a..a7e7af4e 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -610,3 +610,824 @@ FROM WHERE am.amname = 'vchordrq' ) AS index_oids CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; + +-- Phase 3B tensor-source bindings + +CREATE TABLE _vchordrq_maxsim_sources ( + index_oid oid PRIMARY KEY, + heap_oid oid NOT NULL, + model_contract_id text NOT NULL + CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['heap_array', 'external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint, + tensor_rows_attnum smallint, + tensor_dim_attnum smallint, + tensor_dtype_attnum smallint, + tensor_checksum_attnum smallint, + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'heap_array'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NULL + AND tensor_rows_attnum IS NULL + AND tensor_dim_attnum IS NULL + AND tensor_dtype_attnum IS NULL + AND tensor_checksum_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_maxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_maxsim_source( + index_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name DEFAULT NULL, + tensor_rows_column name DEFAULT NULL, + tensor_dim_column name DEFAULT NULL, + tensor_dtype_column name DEFAULT NULL, + tensor_checksum_column name DEFAULT NULL, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + heap_oid oid; + index_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + descriptor_id_is_unique boolean; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be heap_array, external_ref, or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + IF caller_oid IS NULL THEN + RAISE EXCEPTION 'could not resolve caller role'; + END IF; + + SELECT x.indrelid, i.relowner + INTO heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF heap_oid IS NULL THEN + RAISE EXCEPTION 'relation % is not a valid single-key vchordrq MaxSim index', + index_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may register its MaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a MaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO descriptor_id_is_unique; + IF NOT descriptor_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION '% sources must not specify a descriptor relation', normalized_storage; + END IF; + tensor_relation_oid := heap_oid; + END IF; + + IF normalized_storage = 'heap_array' THEN + IF tensor_ref_column IS NOT NULL + OR tensor_rows_column IS NOT NULL + OR tensor_dim_column IS NOT NULL + OR tensor_dtype_column IS NOT NULL + OR tensor_checksum_column IS NOT NULL THEN + RAISE EXCEPTION 'heap_array sources must not specify external tensor columns'; + END IF; + ELSE + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor sources require ref, rows, dim, dtype, and checksum columns'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', + tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', + tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', + tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', + tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'tensor source columns must be distinct'; + END IF; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_maxsim_sources ( + index_oid, heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) + ON CONFLICT (index_oid) DO UPDATE SET + heap_oid = EXCLUDED.heap_oid, + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + index_relation::oid, + heap_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_maxsim_source(index_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + index_owner oid; + removed_count bigint; +BEGIN + IF index_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO index_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = index_relation::oid AND c.relkind = 'i'; + IF index_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may unregister its MaxSim tensor source'; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources WHERE index_oid = $1', + ext_schema + ) USING index_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_maxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_maxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_source_info(index_relation regclass) +RETURNS TABLE( + registered_index regclass, + heap_relation regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + bound_heap_oid oid; + live_heap_oid oid; + index_owner oid; + bound_model_contract_id text; + bound_storage text; + bound_descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + expected_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + + EXECUTE pg_catalog.format( + 'SELECT heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_maxsim_sources + WHERE index_oid = $1', + ext_schema + ) INTO + bound_heap_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + bound_descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING index_relation::oid; + IF bound_heap_oid IS NULL THEN + RAISE EXCEPTION 'MaxSim tensor source is not registered for index %', + index_relation; + END IF; + + SELECT x.indrelid, i.relowner + INTO live_heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF live_heap_oid IS NULL OR live_heap_oid <> bound_heap_oid THEN + RAISE EXCEPTION 'registered MaxSim tensor source is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, live_heap_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered MaxSim tensor source'; + END IF; + + IF bound_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid storage'; + END IF; + + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, 'bigint'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = bound_heap_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap columns'; + END IF; + + IF bound_storage = 'heap_array' THEN + IF bound_descriptor_oid IS NOT NULL + OR descriptor_public_id_attnum IS NOT NULL + OR tensor_ref_attnum IS NOT NULL + OR tensor_rows_attnum IS NOT NULL + OR tensor_dim_attnum IS NOT NULL + OR tensor_dtype_attnum IS NOT NULL + OR tensor_checksum_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap_array binding'; + END IF; + tensor_relation_oid := NULL; + ELSIF bound_storage = 'external_ref' THEN + IF bound_descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_ref binding'; + END IF; + tensor_relation_oid := bound_heap_oid; + ELSE + IF bound_descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_relation binding'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = bound_descriptor_oid + AND c.relkind IN ('r', 'm'); + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor relation is stale or invalid'; + END IF; + IF NOT pg_catalog.has_table_privilege(caller_oid, bound_descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'SELECT privilege on the registered descriptor relation is required'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid = 'bigint'::regtype + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID column is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = bound_descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := bound_descriptor_oid; + END IF; + + expected_columns := CASE WHEN bound_storage = 'heap_array' THEN 0 ELSE 5 END; + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.attnum IS NOT NULL; + IF valid_columns <> expected_columns THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid descriptor columns'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + index_relation, + bound_heap_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + bound_descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_search( + index_relation regclass, + query anyarray, + candidate_limit integer, + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_maxsim_search_external_wrapper'; + +COMMENT ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) +IS 'Restricted Phase 3B external-tensor MaxSim search; returns exact similarity under caller MVCC, SELECT privileges, and PostgreSQL row visibility.'; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_maxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + registry regclass; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RETURN; + END IF; + registry := pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_maxsim_sources', ext_schema) + ); + IF registry IS NULL THEN + RETURN; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE d.objid = s.index_oid + OR ( + d.objid = s.heap_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.model_contract_attnum + OR d.objsubid = s.public_id_attnum + OR ( + s.storage = ''external_ref'' + AND d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.descriptor_public_id_attnum + OR d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); diff --git a/tests/vchordrq/cost_estimator.slt b/tests/vchordrq/cost_estimator.slt index 0f4ad266..87cf85e7 100644 --- a/tests/vchordrq/cost_estimator.slt +++ b/tests/vchordrq/cost_estimator.slt @@ -388,6 +388,174 @@ RESET vchordrq.enable_scan; statement ok RESET enable_seqscan; +# --------------------------------------------------------------------------- +# Case 12: MaxSim must have a nonzero, query-token-aware, backend-aware cost. +# The old special branch returned total_cost=0 and selectivity=1 for every +# MaxSim path, hiding all token expansion and exact rerank work. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_maxsim ( + id int PRIMARY KEY, + v vector(3)[] NOT NULL +); + +statement ok +INSERT INTO cost_test_maxsim +SELECT i, ARRAY[ + '[1,0,0]'::vector, + '[0,1,0]'::vector, + '[0,0,1]'::vector, + '[0.5,0.5,0]'::vector +] +FROM generate_series(1, 1000) i; + +statement ok +CREATE INDEX cost_test_maxsim_v +ON cost_test_maxsim +USING vchordrq (v vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [4] +$$); + +statement ok +ANALYZE cost_test_maxsim; + +statement ok +SET enable_seqscan = off; + +statement ok +SET vchordrq.probes = '4'; + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4; + +statement ok +CREATE TEMP TABLE maxsim_cost_observations ( + name text PRIMARY KEY, + value double precision NOT NULL +); + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 1; + +# A newly built index has a native indexed-vector count. The document-token +# GUC is only a compatibility fallback for pre-statistics indexes, so changing +# it must not change this index's cost. +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'native_document_4', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4096; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'native_document_4096', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'native_document_4') + = + (SELECT value FROM maxsim_cost_observations WHERE name = 'native_document_4096'); +---- +t + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'query_1', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 64; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'query_64', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'query_64') + > + (SELECT value FROM maxsim_cost_observations WHERE name = 'query_1'); +---- +t + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 32; + +statement ok +SET vchordrq.maxsim_candidate_limit = 16; + +statement ok +SET vchordrq.maxsim_backend = 'cpu_exact'; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'cpu_16', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_candidate_limit = 256; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'cpu_256', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'cpu_256') + > + (SELECT value FROM maxsim_cost_observations WHERE name = 'cpu_16'); +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); +---- +cost_test_maxsim_v + +statement ok +RESET vchordrq.maxsim_backend; + +statement ok +RESET vchordrq.maxsim_candidate_limit; + +statement ok +RESET vchordrq.maxsim_planner_query_tokens; + +statement ok +RESET vchordrq.maxsim_planner_document_tokens; + +statement ok +RESET vchordrq.probes; + # --------------------------------------------------------------------------- # Cleanup # --------------------------------------------------------------------------- @@ -404,6 +572,9 @@ DROP TABLE cost_test_cold; statement ok DROP TABLE cost_test_partial; +statement ok +DROP TABLE cost_test_maxsim; + statement ok DROP FUNCTION slow_true(int); diff --git a/tests/vchordrq/maxsim_correctness.slt b/tests/vchordrq/maxsim_correctness.slt new file mode 100644 index 00000000..a5d9652f --- /dev/null +++ b/tests/vchordrq/maxsim_correctness.slt @@ -0,0 +1,221 @@ +# Deterministic MaxSim semantics and input-boundary coverage. + +statement ok +SET enable_seqscan TO off; + +# @# is a distance: negative late-interaction similarity, ordered ascending. +query I +SELECT round(( + ARRAY['[1,0]'::vector, '[0,1]'::vector] + @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +)::numeric, 3); +---- +-2.000 + +# MaxSim is asymmetric: the right-hand array contains query vectors. +query I +SELECT round(( + ARRAY['[1,0]'::vector, '[0,1]'::vector] + @# ARRAY['[1,1]'::vector] +)::numeric, 3); +---- +-1.000 + +query I +SELECT round(( + ARRAY['[1,1]'::vector] + @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +)::numeric, 3); +---- +-2.000 + +# halfvec uses the same sign and orientation contract. +query I +SELECT round(( + ARRAY['[1,0]'::halfvec, '[0,1]'::halfvec] + @# ARRAY['[1,0]'::halfvec, '[0,1]'::halfvec] +)::numeric, 3); +---- +-2.000 + +statement error MaxSim arrays must contain at least one vector +SELECT ARRAY[]::vector[] @# ARRAY['[1,0]'::vector]; + +statement error MaxSim arrays must contain at least one vector +SELECT ARRAY['[1,0]'::vector] @# ARRAY[]::vector[]; + +statement error MaxSim arrays must not contain NULL vectors +SELECT ARRAY['[1,0]'::vector, NULL::vector] + @# ARRAY['[1,0]'::vector]; + +statement error dimension is not matched +SELECT ARRAY['[1,0]'::vector] @# ARRAY['[1,0,0]'::vector]; + +statement error MaxSim arrays cannot contain more than 65536 vectors +SELECT ARRAY['[1]'::vector] + @# ARRAY(SELECT '[1]'::vector FROM generate_series(1, 65537)); + +# Indexing an empty document array must fail instead of silently omitting it. +statement ok +CREATE TABLE maxsim_invalid_document (val vector(2)[]); + +statement ok +INSERT INTO maxsim_invalid_document VALUES (ARRAY[]::vector[]); + +statement error MaxSim arrays must contain at least one vector +CREATE INDEX ON maxsim_invalid_document +USING vchordrq (val vector_maxsim_ops); + +statement ok +TRUNCATE maxsim_invalid_document; + +statement ok +INSERT INTO maxsim_invalid_document +SELECT ARRAY( + SELECT '[1,0]'::vector + FROM generate_series(1, 65537) +); + +statement error MaxSim arrays cannot contain more than 65536 vectors +CREATE INDEX ON maxsim_invalid_document +USING vchordrq (val vector_maxsim_ops); + +statement ok +DROP TABLE maxsim_invalid_document; + +# Lock down current ranking with a non-flat index before the Phase 3 refactor. +statement ok +CREATE TABLE maxsim_deterministic ( + id integer primary key, + val vector(2)[] not null +); + +statement ok +INSERT INTO maxsim_deterministic VALUES + (1, ARRAY['[1,0]'::vector, '[0,1]'::vector]), + (2, ARRAY['[0.8,0]'::vector, '[0,0.8]'::vector]), + (3, ARRAY['[0.5,0.5]'::vector]), + (4, ARRAY['[-1,0]'::vector, '[0,-1]'::vector]); + +statement ok +CREATE INDEX maxsim_deterministic_idx +ON maxsim_deterministic +USING vchordrq (val vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [2] +$$); + +statement ok +SET vchordrq.probes = '2'; + +statement ok +SET vchordrq.maxsim_refine = 100; + +statement ok +SET vchordrq.maxsim_threshold = 0; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +statement ok +SET vchordrq.maxsim_candidate_limit = 2; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 + +statement ok +SET vchordrq.maxsim_candidate_limit = -1; + +statement ok +SET vchordrq.maxsim_backend = 'cpu_exact'; + +statement error exact MaxSim requires a positive vchordrq.maxsim_candidate_limit +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 1; + +statement ok +SET vchordrq.maxsim_candidate_limit = 4; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +query I +SELECT round((val @# ARRAY['[1,0]'::vector, '[0,1]'::vector])::numeric, 3) +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +-2.000 +-1.600 +-1.000 +0.000 + +statement ok +SET vchordrq.maxsim_backend = 'gpu'; + +statement error GPU MaxSim transport error: endpoint is empty +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 1; + +statement ok +SET vchordrq.maxsim_backend = 'auto'; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +statement ok +SET vchordrq.maxsim_candidate_limit = -1; + +# An empty query must also fail on the index scan path. +statement error MaxSim arrays must contain at least one vector +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY[]::vector[] +LIMIT 1; + +statement error dimension is not matched +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0,0]'::vector] +LIMIT 1; + +statement ok +DROP TABLE maxsim_deterministic; diff --git a/tests/vchordrq/maxsim_source_registry.slt b/tests/vchordrq/maxsim_source_registry.slt new file mode 100644 index 00000000..c408a469 --- /dev/null +++ b/tests/vchordrq/maxsim_source_registry.slt @@ -0,0 +1,341 @@ +# Phase 3B tensor-source registration must bind to a real MaxSim index by +# relation/attribute OID, reject incompatible descriptors, and fail closed when +# a bound index/table column is dropped. + +statement ok +CREATE TABLE maxsim_source_test ( + id bigint PRIMARY KEY, + model_contract text NOT NULL, + embedding vector(2)[] NOT NULL, + single_embedding vector(2) NOT NULL, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL, + application_note text +); + +statement ok +INSERT INTO maxsim_source_test VALUES ( + 1, + 'colqwen@test', + ARRAY['[1,0]'::vector, '[0,1]'::vector], + '[1,0]'::vector, + 'tensor://1', + 2, + 2, + 'float32', + 'sha256:test' +); + +statement ok +CREATE INDEX maxsim_source_test_idx +ON maxsim_source_test +USING vchordrq (embedding vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [1] +$$); + +statement ok +CREATE INDEX maxsim_source_wrong_idx +ON maxsim_source_test +USING vchordrq (single_embedding vector_l2_ops) +WITH (options = $$ +[build.internal] +lists = [1] +$$); + +statement error not a valid single-key vchordrq MaxSim index +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_wrong_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +statement error tensor rows column tensor_dtype must be a NOT NULL integer column +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_ref', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_dtype', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => ' colqwen@test ', + storage => 'EXTERNAL_REF', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +query TT +SELECT s.storage, s.model_contract_id +FROM _vchordrq_maxsim_sources AS s +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +external_ref colqwen@test + +query TTT +SELECT source_storage, tensor_ref_column::text, public_id_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +external_ref tensor_ref id + +# Production external tensors may live in a compact descriptor relation keyed +# by the stable application public ID, avoiding a rewrite of the indexed heap. +statement ok +CREATE TABLE maxsim_descriptor_no_unique ( + public_id bigint NOT NULL, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement error must have a non-partial single-key unique index +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_relation', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'maxsim_descriptor_no_unique'::regclass, + descriptor_public_id_column => 'public_id' +); + +statement ok +CREATE TABLE maxsim_descriptor_test ( + public_id bigint PRIMARY KEY, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement ok +INSERT INTO maxsim_descriptor_test VALUES ( + 1, 'tensor://1', 2, 2, 'float32', 'sha256:test' +); + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_relation', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'maxsim_descriptor_test'::regclass, + descriptor_public_id_column => 'public_id' +); + +query TTTT +SELECT source_storage, descriptor_relation::text, + descriptor_public_id_column::text, tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +external_relation maxsim_descriptor_test public_id tensor_ref + +statement ok +ALTER TABLE maxsim_descriptor_test RENAME tensor_ref TO tensor_location; + +query T +SELECT tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +tensor_location + +statement ok +ALTER TABLE maxsim_descriptor_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +0 + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_ref', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +# Unbound application columns are outside the registry contract. Dropping one +# must not invalidate an otherwise live source binding. +statement ok +ALTER TABLE maxsim_source_test DROP COLUMN application_note; + +query T +SELECT model_contract_id +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +colqwen@test + +# The explicit score surface validates bounded work and exact query/index type +# before any sidecar access. +statement error candidate_limit must be between 1 and 65536 +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::vector], + 0, + 1 +); + +statement error top_k must be positive and no greater than candidate_limit +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::vector], + 1, + 2 +); + +statement ok +SET vchordrq.maxsim_backend = 'gpu'; + +statement error MaxSim tensor kind or dimension is not matched +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::halfvec], + 1, + 1 +); + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +query T +SELECT a.attname +FROM _vchordrq_maxsim_sources AS s +JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = s.heap_oid AND a.attnum = s.tensor_ref_attnum +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +tensor_ref + +# Attribute numbers survive renames, so bindings do not depend on SQL text. +statement ok +ALTER TABLE maxsim_source_test RENAME tensor_ref TO tensor_location; + +query T +SELECT a.attname +FROM _vchordrq_maxsim_sources AS s +JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = s.heap_oid AND a.attnum = s.tensor_ref_attnum +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +tensor_location + +query T +SELECT tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +tensor_location + +# Type-altering DDL does not drop an attribute, so runtime resolution must +# revalidate the complete descriptor and fail closed. +statement ok +ALTER TABLE maxsim_source_test +ALTER COLUMN tensor_dtype TYPE varchar(16); + +statement error registered MaxSim tensor source has invalid descriptor columns +SELECT * FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); + +statement ok +ALTER TABLE maxsim_source_test +ALTER COLUMN tensor_dtype TYPE text; + +# Dropping a bound descriptor column invalidates the complete binding. +statement ok +ALTER TABLE maxsim_source_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +0 + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +query B +SELECT tensor_ref_attnum IS NULL +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +t + +query B +SELECT vchordrq_unregister_maxsim_source('maxsim_source_test_idx'::regclass); +---- +t + +query B +SELECT vchordrq_unregister_maxsim_source('maxsim_source_test_idx'::regclass); +---- +f + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +statement ok +DROP INDEX maxsim_source_test_idx; + +query I +SELECT count(*) FROM _vchordrq_maxsim_sources; +---- +0 + +statement ok +DROP TABLE maxsim_source_test; + +statement ok +DROP TABLE maxsim_descriptor_test, maxsim_descriptor_no_unique; From 03961017f015177b3137cfa6d3a97afa8a404795 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:09:38 +0800 Subject: [PATCH 311/324] feat(tilemaxsim): add bounded native GPU serving --- .gitignore | 4 + README.md | 219 +++ devtools/test_tilemaxsim_reference_sidecar.py | 66 +- devtools/tilemaxsim_reference_sidecar.py | 73 +- docs/MAXSIM_TENSOR_SOURCES.md | 162 ++ docs/TILEMAXSIM_CUDA_SIDECAR.md | 173 +++ docs/TILEMAXSIM_IPC_V1.md | 160 ++ docs/TILEMAXSIM_IPC_V2.md | 173 +++ services/Dockerfile.tilemaxsim | 3 + services/Dockerfile.tilemaxsimd | 22 + services/benchmark_full_corpus_tilemaxsim.py | 190 +++ services/benchmark_gbrain_phase3.py | 334 ++++ .../benchmark_gbrain_scoped_tilemaxsim.py | 252 +++ services/benchmark_npy_tilemaxsim_sidecar.py | 291 ++++ services/benchmark_postgres_single_vector.py | 336 ++++ services/benchmark_tilemaxsim_ablation.py | 765 ++++++++++ services/build_tilemaxsim_tensor_cache.py | 154 +- services/test_tilemaxsim_cuda_sidecar.py | 277 ++++ services/test_tilemaxsim_gpu_cache.py | 376 +++++ services/test_tilemaxsim_rust_daemon.py | 409 +++++ services/test_tilemaxsim_shard.py | 181 +++ services/tilemaxsim_cuda_sidecar.py | 1093 ++++++++++++- services/tilemaxsim_gpu_cache.py | 975 ++++++++++++ services/tilemaxsim_shard.py | 354 +++++ services/tilemaxsim_triton.py | 128 ++ services/tilemaxsimd/Cargo.lock | 416 +++++ services/tilemaxsimd/Cargo.toml | 19 + services/tilemaxsimd/build.rs | 12 + .../tilemaxsimd/native/tilemaxsim_cuda.cu | 291 ++++ services/tilemaxsimd/src/cache.rs | 762 ++++++++++ services/tilemaxsimd/src/engine.rs | 584 +++++++ services/tilemaxsimd/src/gpu.rs | 163 ++ services/tilemaxsimd/src/main.rs | 1349 +++++++++++++++++ services/tilemaxsimd/src/protocol.rs | 337 ++++ services/tilemaxsimd/src/scheduler.rs | 328 ++++ services/tilemaxsimd/src/shard.rs | 499 ++++++ src/index/gucs.rs | 48 + src/index/vchordrq/scanners/maxsim.rs | 74 +- .../vchordrq/scanners/maxsim/candidate.rs | 114 +- src/index/vchordrq/scanners/maxsim/exact.rs | 501 ++++++ .../vchordrq/scanners/maxsim/external.rs | 122 ++ src/index/vchordrq/scanners/maxsim/gpu.rs | 123 +- src/index/vchordrq/scanners/maxsim/profile.rs | 181 +++ src/index/vchordrq/scanners/maxsim/search.rs | 50 +- src/sql/finalize.sql | 753 +++++++++ tests/vchordrq/tilemaxsim_source_registry.slt | 91 ++ 46 files changed, 13830 insertions(+), 157 deletions(-) create mode 100644 docs/MAXSIM_TENSOR_SOURCES.md create mode 100644 docs/TILEMAXSIM_CUDA_SIDECAR.md create mode 100644 docs/TILEMAXSIM_IPC_V1.md create mode 100644 docs/TILEMAXSIM_IPC_V2.md create mode 100644 services/Dockerfile.tilemaxsimd create mode 100644 services/benchmark_full_corpus_tilemaxsim.py create mode 100644 services/benchmark_gbrain_phase3.py create mode 100644 services/benchmark_gbrain_scoped_tilemaxsim.py create mode 100644 services/benchmark_npy_tilemaxsim_sidecar.py create mode 100644 services/benchmark_postgres_single_vector.py create mode 100644 services/benchmark_tilemaxsim_ablation.py create mode 100644 services/test_tilemaxsim_gpu_cache.py create mode 100644 services/test_tilemaxsim_rust_daemon.py create mode 100644 services/test_tilemaxsim_shard.py create mode 100644 services/tilemaxsim_gpu_cache.py create mode 100644 services/tilemaxsim_shard.py create mode 100644 services/tilemaxsim_triton.py create mode 100644 services/tilemaxsimd/Cargo.lock create mode 100644 services/tilemaxsimd/Cargo.toml create mode 100644 services/tilemaxsimd/build.rs create mode 100644 services/tilemaxsimd/native/tilemaxsim_cuda.cu create mode 100644 services/tilemaxsimd/src/cache.rs create mode 100644 services/tilemaxsimd/src/engine.rs create mode 100644 services/tilemaxsimd/src/gpu.rs create mode 100644 services/tilemaxsimd/src/main.rs create mode 100644 services/tilemaxsimd/src/protocol.rs create mode 100644 services/tilemaxsimd/src/scheduler.rs create mode 100644 services/tilemaxsimd/src/shard.rs create mode 100644 src/index/vchordrq/scanners/maxsim/exact.rs create mode 100644 src/index/vchordrq/scanners/maxsim/profile.rs create mode 100644 tests/vchordrq/tilemaxsim_source_registry.slt diff --git a/.gitignore b/.gitignore index 85d1619b..d2f22b6e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,8 @@ /target +/services/tilemaxsimd/target +/benchmarks/ +**/__pycache__/ +.ruff_cache/ **/*.rs.bk .vscode .ignore diff --git a/README.md b/README.md index e305937c..6ae46489 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,224 @@
+# VectorChord TileMaxSim + +**An open-source VectorChord fork focused on exact multi-vector retrieval in PostgreSQL.** + +
+ +This fork extends VectorChord's `vchordrq` index with exact late-interaction +TileMaxSim retrieval. It is intended for applications that store one array of +token vectors per document and need PostgreSQL-native multi-vector search. + +## What this fork adds + +- Exact TileMaxSim reranking on CPU, plus an optional CUDA sidecar backend. +- Full caller-scoped tensor sets without an artificial candidate-count cap; + device-memory pressure is handled by the GPU page cache and request batching. +- External tensor-source registration for deployments that keep full-precision + token tensors outside the indexed PostgreSQL value. +- PostgreSQL-aware permission, MVCC, row-visibility, cancellation, and timeout + handling for the external-tensor search path. +- Planner statistics and cost estimation for multi-vector queries. +- Deterministic correctness, registry, sidecar-protocol, and planner-cost tests. + +## Performance ablation + +We measured each cache-path optimization independently on the same development +machine. The corpus contained 34,054 tensor descriptors (34,027 unique tensors, +16.28 GB of logical FP16 tensor data). The request-level tests sampled 100 +candidates containing 47.81 MB of tensor data. Absolute latency depends on the +storage and GPU, so the same-run comparisons are more useful than the raw +numbers. + +| Optimization | Baseline | Optimized | Result | +| --- | ---: | ---: | ---: | +| Immutable shards and batched reads | 333.38 ms, sequential files | 56.52 ms | 5.90x faster | +| Shards versus batched legacy files | 87.56 ms | 56.52 ms | 1.55x faster | +| Batched host-to-device transfer | 37.06 ms, 100 transfers | 14.19 ms, one transfer | 2.61x faster | +| TinyLFU/GDSF admission | 69.98% LRU hit rate | 76.18% hit rate | +6.20 percentage points | +| Rust/CUDA cold request | 855.34 ms, Python/Triton | 93.86 ms | 9.11x faster | +| Rust/CUDA warm request p50 | 14.26 ms, Python/Triton | 2.09 ms | 6.82x faster | +| Rust/CUDA warm request p95 | 14.84 ms, Python/Triton | 2.20 ms | 6.75x faster | + +A full resident-cache run assigned 20 GiB to one GPU: 18 GiB for tensors and +2 GiB for the TileMaxSim workspace. All 34,027 unique tensors were pinned before +the service became ready. Process-to-ready time was 23.07 seconds for the +Python/Triton sidecar and 14.86 seconds for the Rust/CUDA daemon, a 35.6% +reduction. This is a one-time prewarm cost; resident warm requests do not read +the tensors from disk. + +The GPU cache now suballocates exact contiguous page runs from one CUDA arena, +using best-fit size buckets and address-ordered coalescing. On the full corpus, +the 32 KiB default reduced allocated tensor space from 17.840 GB with the former +256 KiB power-of-two buddy allocator to 16.725 GB. It recovered 1.115 GB of GPU +space and reduced internal rounding waste from 8.82% to 2.74%. The default can +be overridden with `--gpu-block-kib`. + +In a deterministic 20,000-event churn trace, the former buddy allocator had 815 +failed cache-allocation attempts, while the segregated page-run allocator had +637; neither recorded an external-fragmentation failure on this workload. Exact +byte extents had 585 failures, all caused by external fragmentation. Page-run +metadata processing added about 1.4 microseconds per event versus the buddy +baseline. These are cache-admission attempts, not failed search requests: the +runtime evicts unpinned entries or streams oversized working sets in chunks. + +Rust/CUDA and Python/Triton produced identical top-10 results. The maximum +absolute score difference was 5.25e-6 and the mean difference was 3.45e-6. +The benchmark driver is +[`services/benchmark_tilemaxsim_ablation.py`](services/benchmark_tilemaxsim_ablation.py); +run it with `--help` for the corpus, cache-root, device, and output arguments. + +### Full-source retrieval and cache-scheduler stress + +We separately scored all 34,054 descriptors for each query to exercise a +working set larger than the cache. This is both a storage/cache stress test and +an upper bound for general semantic retrieval when no safe narrower hard scope +exists. GBrain still applies source, ACL, type, and other mandatory filters, but +lexical or graph recall must not hide semantic paraphrases. The traditional +embedding/HNSW path remains a separate `single_vector` mode and is not a +dependency of tensor retrieval. + +With an 18 GiB tensor arena, all 34,027 unique tensors were resident and native +daemon round trips averaged 1.08 seconds for a deliberately exhaustive scan. A +3 GiB tensor arena plus 1 GiB host cache produced the same exact top-K, but two +sequential scans caused 68,106 misses, 61,499 evictions, 32.53 GB of host-to- +device transfer, and only two cache hits. Native round trips averaged 18.47 +seconds. The slowdown is cyclic cache thrashing and repeated I/O, not TileMaxSim +compute scaling linearly with GPU-memory capacity. Candidate-scoped queries +whose hot working set fits the cache do not exhibit this full-scan behavior. + +The diagnostic driver is +[`services/benchmark_full_corpus_tilemaxsim.py`](services/benchmark_full_corpus_tilemaxsim.py). +A GBrain-scoped comparison against its original text-embedding/pgvector HNSW +path is reported separately; tensor-derived pooled vectors are not a valid +single-vector baseline. + +### Rejected lexical hard-gate ablation and the traditional embedding baseline + +We then used 40 real queries and their original 1,024-dimensional +`text-embedding-v4` vectors. The single-vector baseline searched 1,200 stored +chunk embeddings with PostgreSQL pgvector HNSW (`m=16`, `ef_search=40`). Tensor +mode did not read those vectors: a real lexical result list supplied every +returned chunk without a second truncation, the chunks were mapped to their +source-page tensors, and the native CUDA daemon ran exact TileMaxSim only over +that scope. This ablation tests whether lexical recall is safe as a mandatory +gate; it is not the accepted general semantic query plan. + +| Path | Mean latency | p95 | Quality | +| --- | ---: | ---: | ---: | +| Original text embedding + pgvector HNSW | 5.80 ms | 5.86 ms | document hit@1 1.000 | +| Candidate-scoped TileMaxSim, 18 GiB resident tensor cache, native round trips | 45.56 ms | 134.99 ms | document hit@1 0.475; hit@5 0.650 | +| Candidate-scoped TileMaxSim, resident Python driver end to end | 112.33 ms | 319.87 ms | same ranking | +| Candidate-scoped TileMaxSim, 3 GiB LRU tensor cache, native round trips | 573.43 ms | 1,812.70 ms | same ranking | +| Candidate-scoped TileMaxSim, small-cache Python driver end to end | 657.46 ms | 2,011.99 ms | same ranking | + +The lexical scope contained 19.3 chunks and 13.2 documents on average. Because +the source chunks span multiple PDF pages, that mapped to 1,459 page tensors on +average (p95 4,000); no whole-corpus tensor scan occurred. Candidate document +recall was 0.650, and TileMaxSim retained the relevant document in its top five +for all covered queries, so scope generation—not exact reranking—set the recall +ceiling in this run. A 35% loss before TileMaxSim is unacceptable, so lexical +and non-relational graph results neither exclude candidates nor reorder general +tensor retrieval. Only explicit relationship intent may use and fuse GBrain's +complete graph scope as a hard range. The HNSW baseline is an intentionally strong upper-bound +workload whose queries are excerpts from their gold chunks; page-tensor and +text-chunk granularity differ, so its quality number is not a claim that the two +rankers are interchangeable. + +The 3 GiB run completed every request with identical ranking, but incurred +24,432 GPU misses and 11.68 GB of host-to-device transfer across the query +sequence. Its 5.85x driver slowdown relative to resident mode is therefore a +cache-locality result that happens to resemble the 6x cache-capacity ratio; GPU +TileMaxSim arithmetic does not become six times slower when memory is smaller. + +The reproducible drivers are +[`services/benchmark_gbrain_scoped_tilemaxsim.py`](services/benchmark_gbrain_scoped_tilemaxsim.py) +and +[`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py). + +## Bounded multi-tenant GPU scheduling + +TileMaxSim remains opt-in. Ordinary single-vector VectorChord use does not +start the CUDA daemon and does not require a GPU-memory setting. When an +operator enables TileMaxSim, `tilemaxsimd` reserves exactly the configured +CUDA devices and GiB allocations at startup and fails closed if any allocation +cannot be obtained. + +The default `fair-priority` scheduler combines explicit request urgency with +weighted tenant fairness. Higher numeric priority is more urgent; requests in +the configured priority band are selected by normalized GPU service consumed, +not merely by request count. The default band spans the complete public priority +range, so urgency breaks ties and reorders work inside one tenant without letting +a high-priority noisy tenant starve an under-served tenant. Waiting requests age +up to the maximum priority. Operators that require global strict priority can +select `priority`; a narrower band is an explicit stronger-urgency trade-off. +This is a serving-design adaptation rather than wire compatibility with vLLM: +VectorChord intentionally uses higher-number-means-more-urgent throughout its +SQL, MCP, and IPC contracts. +`fair` and strict `priority` (priority then FCFS) are also available. This takes +the useful serving ideas from vLLM—bounded work budgets, continuous scheduling, +and resumable long work—while adding tenant isolation: a large request is split +at candidate/token quanta and re-enters the scheduler between CUDA launches. +It is cooperative preemption between kernels, not interruption of an executing +CUDA kernel. + +Admission is bounded both globally and per tenant before work enters the +scheduler. Client disconnects and end-to-end deadlines are checked between +quanta. GPU and host cache ownership also have per-tenant caps; optional GPU +reservations and tenant scheduling weights can be configured for differentiated +service. Tenant identifiers are accepted only as scheduling domains, never as +authorization evidence, and request logs expose only a stable tenant hash. + +Memory flags use GiB rather than byte counts. A representative launch is: + +```shell +tilemaxsimd \ + --socket /run/vectorchord/tilemaxsim.sock \ + --status-socket /run/vectorchord/tilemaxsim-status.sock \ + --gpu-memory-gb 0=20 \ + --gpu-workspace-gb 2 \ + --host-cache-gb 8 \ + --max-inflight-request-gb 1 \ + --contract-root MODEL_CONTRACT_ID=/srv/vectorchord/tensors \ + --scheduler-policy fair-priority \ + --max-queued-requests 128 \ + --max-tenant-queued-requests 16 \ + --tenant-weight foreground=2 \ + --tenant-cache-reservation foreground=4 +``` + +The optional status socket serves HTTP `GET /healthz` and Prometheus +`GET /metrics`. Metrics include readiness, scheduler depth, active CUDA work, +completed/error/timeout/disconnect outcomes, and global/per-tenant admission +rejections without exporting tenant identifiers. + +The in-flight request budget is also expressed in GiB. A reader must reserve +its complete declared frame after the fixed header is validated, and keeps that +permit through completion, timeout, or disconnect. This bounds aggregate query +and descriptor memory even when many clients submit maximum-size frames at once. + +PostgreSQL sends protocol v3 scheduling metadata only when +`vchordrq.maxsim_tenant` is set; otherwise it retains protocol v2 compatibility. +GBrain derives that tenant value from authenticated runtime context and may set +`vchordrq.maxsim_priority` in the range -100 through 100. Priority changes +latency ordering only and never bypasses PostgreSQL row visibility, ACL, source, +or other mandatory filters. + +The implementation is currently under active development. Its SQL interfaces +and deployment packaging may change before a stable release. This repository +contains only the public implementation and public-facing project information; +private planning and application documentation are intentionally excluded. + +This work is based on +[supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord). The +original VectorChord README is retained below for upstream installation, +licensing, and project information. + +--- + +
+ # VectorChord **Ready for the Billion-Scale Era. Host 100M vectors on a single i4i.xlarge ($247/mo) and [scale seamlessly to 1B+](https://blog.vectorchord.ai/scaling-vector-search-to-1-billion-on-postgresql).** diff --git a/devtools/test_tilemaxsim_reference_sidecar.py b/devtools/test_tilemaxsim_reference_sidecar.py index ee5ce666..cf3bc826 100644 --- a/devtools/test_tilemaxsim_reference_sidecar.py +++ b/devtools/test_tilemaxsim_reference_sidecar.py @@ -114,10 +114,46 @@ def external_request_frame( ) +def scheduled_external_request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + model_contract_id: str, + candidates: list[tuple[int, str, list[list[float]]]], + tenant: str, + priority: int, + timeout_ms: int, +) -> tuple[bytes, dict[str, bytes]]: + legacy, objects = external_request_frame( + request_id, dtype, query, model_contract_id, candidates + ) + body = legacy[sidecar.HEADER.size :] + fixed = sidecar.EXTERNAL_REQUEST_FIXED.unpack_from(body) + tenant_bytes = tenant.encode() + scheduled = bytearray( + sidecar.SCHEDULED_EXTERNAL_REQUEST_FIXED.pack( + *fixed, priority, timeout_ms, len(tenant_bytes) + ) + ) + contract_length = fixed[-1] + scheduled.extend(body[sidecar.EXTERNAL_REQUEST_FIXED.size :][:contract_length]) + scheduled.extend(tenant_bytes) + scheduled.extend(body[sidecar.EXTERNAL_REQUEST_FIXED.size + contract_length :]) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.SCHEDULED_EXTERNAL_VERSION, + sidecar.REQUEST_KIND, + request_id, + len(scheduled), + ) + + scheduled, + objects, + ) def decode_response(frame: bytes) -> tuple[int, int, list[tuple[int, float]] | str]: magic, version, kind, request_id, body_len = sidecar.HEADER.unpack_from(frame) assert magic == sidecar.MAGIC - assert version in (sidecar.VERSION, sidecar.EXTERNAL_VERSION) + assert version in sidecar.SUPPORTED_VERSIONS assert kind == sidecar.RESPONSE_KIND assert len(frame) == sidecar.HEADER.size + body_len status, count_or_length = sidecar.RESPONSE_FIXED.unpack_from( @@ -216,7 +252,7 @@ def test_truncated_and_oversized_frames_fail_closed(self) -> None: def test_header_reserved_trailing_and_token_limit_fail_closed(self) -> None: valid = request_frame(47, sidecar.DTYPE_F32, [[1.0]], [(9, [[1.0]])]) invalid_frames = [] - for offset, value in ((0, 0), (4, 3), (6, 2)): + for offset, value in ((0, 0), (4, 4), (6, 2)): invalid = bytearray(valid) invalid[offset] = value invalid_frames.append(bytes(invalid)) @@ -278,6 +314,32 @@ def resolver(request: sidecar.ExternalTensorRequest) -> bytes: [candidate.candidate_id for candidate in parsed.candidates], [77, 4] ) + def test_external_v3_preserves_scheduler_metadata_and_v2_scoring(self) -> None: + frame, objects = scheduled_external_request_frame( + 49, + sidecar.DTYPE_F16, + [[1.0, 0.0]], + "model@1", + [(7, "sha256://opaque", [[1.0, 0.0]])], + "tenant-a", + 17, + 4_000, + ) + parsed = sidecar.parse_request_frame(frame) + self.assertIsInstance(parsed, sidecar.ParsedExternalTensorRequest) + assert isinstance(parsed, sidecar.ParsedExternalTensorRequest) + self.assertEqual(parsed.scheduler_tenant, "tenant-a") + self.assertEqual(parsed.scheduler_priority, 17) + self.assertEqual(parsed.timeout_ms, 4_000) + request_id, status, results = decode_response( + sidecar.process_frame( + frame, resolver=lambda request: objects[request.tensor_ref] + ) + ) + self.assertEqual((request_id, status), (49, 0)) + assert isinstance(results, list) + self.assertAlmostEqual(results[0][1], 1.0) + def test_external_v2_fails_closed_without_resolver_or_on_checksum_mismatch( self, ) -> None: diff --git a/devtools/tilemaxsim_reference_sidecar.py b/devtools/tilemaxsim_reference_sidecar.py index 8d9f7ee1..2076d2c8 100644 --- a/devtools/tilemaxsim_reference_sidecar.py +++ b/devtools/tilemaxsim_reference_sidecar.py @@ -40,6 +40,8 @@ MAGIC = b"VCTM" VERSION = 1 EXTERNAL_VERSION = 2 +SCHEDULED_EXTERNAL_VERSION = 3 +SUPPORTED_VERSIONS = (VERSION, EXTERNAL_VERSION, SCHEDULED_EXTERNAL_VERSION) REQUEST_KIND = 1 RESPONSE_KIND = 2 SCORING_SUM_QUERY_MAX_DOCUMENT_DOT = 1 @@ -50,6 +52,7 @@ REQUEST_FIXED = struct.Struct(" list[tuple[int, float]]: - ( - dimension, - query_rows, - candidate_count, - dtype, - scoring, - reserved, - contract_length, - ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + if version == SCHEDULED_EXTERNAL_VERSION: + ( + dimension, query_rows, candidate_count, dtype, scoring, reserved, + contract_length, priority, timeout_ms, tenant_length, + ) = reader.unpack(SCHEDULED_EXTERNAL_REQUEST_FIXED) + if not -100 <= priority <= 100 or not 1 <= timeout_ms <= 600_000: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid scheduler metadata") + else: + ( + dimension, query_rows, candidate_count, dtype, scoring, reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + tenant_length = 0 validate_request_fixed( dimension, query_rows, candidate_count, dtype, scoring, reserved, limits ) model_contract_id = read_text(reader, contract_length, 512, "model contract") + if tenant_length: + read_text(reader, tenant_length, 256, "scheduler tenant") query = read_tensor(reader, query_rows, dimension, dtype) total_tokens = query_rows total_tensor_bytes = checked_tensor_bytes(query_rows, dimension, dtype) @@ -462,7 +475,7 @@ def parse_request_frame( magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) if magic != MAGIC: raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") - if version not in (VERSION, EXTERNAL_VERSION): + if version not in SUPPORTED_VERSIONS: raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") if kind != REQUEST_KIND: raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") @@ -525,15 +538,21 @@ def parse_request_frame( tuple(candidates), ) - ( - dimension, - query_rows, - candidate_count, - dtype, - scoring, - reserved, - contract_length, - ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + if version == SCHEDULED_EXTERNAL_VERSION: + ( + dimension, query_rows, candidate_count, dtype, scoring, reserved, + contract_length, scheduler_priority, timeout_ms, tenant_length, + ) = reader.unpack(SCHEDULED_EXTERNAL_REQUEST_FIXED) + if not -100 <= scheduler_priority <= 100 or not 1 <= timeout_ms <= 600_000: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid scheduler metadata") + else: + ( + dimension, query_rows, candidate_count, dtype, scoring, reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + scheduler_priority = 0 + timeout_ms = 0 + tenant_length = 0 validate_request_fixed( dimension, query_rows, @@ -544,6 +563,11 @@ def parse_request_frame( limits, ) model_contract_id = read_text(reader, contract_length, 512, "model contract") + scheduler_tenant = ( + read_text(reader, tenant_length, 256, "scheduler tenant") + if tenant_length + else "__default__" + ) query_payload = reader.take(checked_tensor_bytes(query_rows, dimension, dtype)) if validate_finite: validate_finite_tensor_payload(query_payload, query_rows, dimension, dtype) @@ -600,6 +624,9 @@ def parse_request_frame( model_contract_id, query_payload, tuple(candidates), + scheduler_tenant, + scheduler_priority, + timeout_ms, ) @@ -614,11 +641,11 @@ def process_frame( if len(frame) < HEADER.size: raise SidecarError(STATUS_INVALID_REQUEST, "truncated frame header") magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) - if version in (VERSION, EXTERNAL_VERSION): + if version in SUPPORTED_VERSIONS: response_version = version if magic != MAGIC: raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") - if version not in (VERSION, EXTERNAL_VERSION): + if version not in SUPPORTED_VERSIONS: raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") if kind != REQUEST_KIND: raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") @@ -631,7 +658,7 @@ def process_frame( if version == VERSION: results = process_inline_request(reader, limits) else: - results = process_external_request(reader, limits, resolver) + results = process_external_request(reader, limits, resolver, version) return success_response(request_id, results, response_version) except SidecarError as error: return error_response(request_id, error.status, str(error), response_version) @@ -663,7 +690,7 @@ def handle_connection( try: header = receive_exact(connection, HEADER.size) _, version, _, request_id, body_len = HEADER.unpack(header) - if version in (VERSION, EXTERNAL_VERSION): + if version in SUPPORTED_VERSIONS: response_version = version if body_len > limits.max_request_bytes - HEADER.size: response = error_response( diff --git a/docs/MAXSIM_TENSOR_SOURCES.md b/docs/MAXSIM_TENSOR_SOURCES.md new file mode 100644 index 00000000..1e0ae37f --- /dev/null +++ b/docs/MAXSIM_TENSOR_SOURCES.md @@ -0,0 +1,162 @@ +# MaxSim Tensor-Source Registry + +## Scope + +The registry binds one physical `vchordrq` MaxSim index to the tensor contract +that Phase 3B will use for exact reranking. It stores metadata only. Tensor +credentials and tensor contents never belong in this catalog. + +The registry lifecycle, internal descriptor reader, and restricted Phase 3B +score executor are implemented. The reader validates model contract, shape, +dtype, reference, and checksum before the descriptor v2 IPC encoder accepts a +candidate. The internal Rust resolver calls `vchordrq_maxsim_source_info` and +never reads the private registry table directly. Same-heap bindings resolve +current names back to physical attribute numbers; independent descriptors are +read with one snapshot-consistent SQL join after candidate visibility is +resolved. Registering a source does not change ordinary `@#` execution. + +Search an external full-tensor source with: + +```sql +SELECT public_id, similarity +FROM vchordrq_maxsim_search( + 'visual_page_embeddings_maxsim_idx'::regclass, + $1::halfvec[], + 1024, -- bounded coarse page candidates + 20 -- final exact rows +) +ORDER BY similarity DESC, public_id; +``` + +The current function is deliberately restricted: + +- the registered source must use `external_ref` or `external_relation` storage; +- the named index must use `vector_maxsim_ops` or `halfvec_maxsim_ops` and the + query array type must match it exactly; +- `vchordrq.maxsim_backend` must be `gpu`, with an operations-controlled Unix + socket endpoint; +- `candidate_limit` is between 1 and 65,536 and `top_k` cannot exceed it; +- the caller must have table-level `SELECT` on the indexed heap and, for + `external_relation`, the descriptor relation; projection uses the caller's + MVCC snapshot and PostgreSQL row visibility; +- the caller selects the physical relation/partition. Arbitrary SQL predicates + would require an optional future CustomScan path and are not accepted as + strings; +- HOT root TIDs from the index are resolved through the table AM before the + descriptor query, while backend IDs remain opaque root-candidate IDs. + +The result contains the registered stable `bigint` public ID and positive exact +similarity. Equal scores are ordered by public ID. GPU/sidecar failures fail the +whole statement; this surface does not silently fall back to coarse scores. + +The bundled CUDA sidecar uses `sha256://` tensor references backed by +an operations-configured, per-model local content-addressed cache. GBrain may +populate that cache from its own storage layer; storage credentials and remote +destinations are not registry fields. See +[`TILEMAXSIM_CUDA_SIDECAR`](TILEMAXSIM_CUDA_SIDECAR.md) for the file layout, +runtime limits, and deployment contract. + +## Registration + +Only the index owner, a member of the owning role, or a superuser can register +or replace a binding. An independently owned descriptor relation must also be +owned by the caller (or a role it belongs to). Descriptor columns are resolved +to attribute numbers at registration time. + +`external_relation` is the recommended production layout. It avoids widening +and rewriting an already indexed multi-vector heap: the heap keeps its stable +public ID, model contract, and coarse array, while a compact relation keeps one +external descriptor per public ID. + +```sql +CREATE TABLE visual_page_tensor_descriptors ( + document_page_id bigint PRIMARY KEY, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +SELECT vchordrq_register_maxsim_source( + index_relation => 'visual_page_embeddings_maxsim_idx'::regclass, + model_contract_id => 'colqwen3.5@revision+preprocessing-hash', + storage => 'external_relation', + model_contract_column => 'model_contract_id'::name, + public_id_column => 'document_page_id'::name, + tensor_ref_column => 'tensor_ref'::name, + tensor_rows_column => 'tensor_rows'::name, + tensor_dim_column => 'tensor_dim'::name, + tensor_dtype_column => 'tensor_dtype'::name, + tensor_checksum_column => 'tensor_checksum'::name, + descriptor_relation => 'visual_page_tensor_descriptors'::regclass, + descriptor_public_id_column => 'document_page_id'::name +); +``` + +The index must be a valid, ready, single-key `vchordrq` index using one of the +four MaxSim opclasses. The model-contract and descriptor string columns must be +`NOT NULL text`; the public ID must be `NOT NULL bigint`; rows and dimension +must be `NOT NULL integer`. For `external_relation`, its public ID must have a +valid, non-partial, single-key unique index. Its public ID and five descriptor +columns must be distinct. Missing descriptor rows fail the complete search. + +`external_ref` remains available when all five descriptor columns already live +on the indexed heap. In that layout the two heap metadata columns and five +descriptor columns must all be distinct; omit the two descriptor-relation +arguments. + +For `heap_array`, omit all five external descriptor columns. The model-contract +and public-ID columns are still required: + +```sql +SELECT vchordrq_register_maxsim_source( + index_relation => 'visual_page_embeddings_maxsim_idx'::regclass, + model_contract_id => 'colqwen3.5@revision+preprocessing-hash', + storage => 'heap_array', + model_contract_column => 'model_contract_id'::name, + public_id_column => 'document_page_id'::name +); +``` + +Unregister explicitly with: + +```sql +SELECT vchordrq_unregister_maxsim_source( + 'visual_page_embeddings_maxsim_idx'::regclass +); +``` + +Resolve and revalidate a binding without reading tensor values: + +```sql +SELECT * +FROM vchordrq_maxsim_source_info( + 'visual_page_embeddings_maxsim_idx'::regclass +); +``` + +This function resolves current column names from stored attribute numbers, +rechecks the live index/opclass and every column type, and requires either index +ownership or table `SELECT` privilege. It is the metadata boundary consumed by +the Phase 3B executor; it does not grant access to row values. + +## DDL and Security Semantics + +- Direct DML on `_vchordrq_maxsim_sources` is revoked from `PUBLIC`. +- Registration functions use definer rights with the fixed search path + `pg_catalog, pg_temp`; ownership is checked against the session user and its + role memberships. +- Attribute-number binding survives column renames. +- An `sql_drop` event trigger removes the binding when its index, heap relation, + descriptor relation, or any bound column is dropped. +- A concurrent index replacement gets a new OID and therefore requires explicit + re-registration. This is fail-closed behavior. +- Raw OID bindings are intentionally not portable pg_dump data. Restore tooling + must register sources after restoring tables and indexes. +- Runtime metadata lookup revalidates the relation, opclass, and descriptor + types. Type-altering DDL that does not drop an attribute therefore makes + resolution fail closed until the binding is corrected or re-registered. +- The search API uses the caller's snapshot and permissions and preserves + PostgreSQL row visibility. The registry itself grants no table or tensor + access and does not define an application authorization or routing model. diff --git a/docs/TILEMAXSIM_CUDA_SIDECAR.md b/docs/TILEMAXSIM_CUDA_SIDECAR.md new file mode 100644 index 00000000..c1fd2968 --- /dev/null +++ b/docs/TILEMAXSIM_CUDA_SIDECAR.md @@ -0,0 +1,173 @@ +# VectorChord CUDA TileMaxSim Sidecar + +## Scope + +`services/tilemaxsim_cuda_sidecar.py` is the deployable CUDA executor for IPC +v1 and v2. PostgreSQL remains responsible for candidate generation, MVCC row +visibility, descriptor validation, and final public-ID mapping. The sidecar +performs bounded tensor resolution and exact TileMaxSim only. + +The service does not fetch from S3, HTTP, or another network destination. +GBrain may populate a node-local immutable cache from its storage system. This +keeps storage credentials, application authorization, and routing outside the +VectorChord extension while retaining one VectorChord-owned retrieval call. + +Production acceptance still requires an application-level corpus benchmark. +The service implementation alone is not that acceptance result. + +## Runtime Requirements + +- Linux with a Unix-domain socket; +- PyTorch built for the installed NVIDIA driver/CUDA runtime; +- one CUDA device by default (`--device cuda:0`); +- a service account that can read configured cache roots and create the socket; +- the PostgreSQL service account must be able to connect to the socket. + +CPU mode exists for conformance tests only: + +```text +python3 -m services.tilemaxsim_cuda_sidecar --device cpu ... +``` + +The production command must select a CUDA device. Startup performs a CUDA +matrix multiply and synchronization, so an unusable device fails before the +socket is advertised. + +For container deployment, build the minimal service image on top of the +operations-approved PyTorch/CUDA runtime that matches the node driver: + +```text +docker build \ + --build-arg BASE_IMAGE= \ + -f services/Dockerfile.tilemaxsim \ + -t vectorchord-tilemaxsim:local \ + . +``` + +The Dockerfile intentionally has no default base image. CUDA/PyTorch/driver +compatibility is a deployment input and must not drift through an implicit +`latest` tag. + +## Content-Addressed Tensor Cache + +Each model contract is mapped explicitly to one absolute cache root. A v2 +reference has exactly this form: + +```text +sha256://0123456789abcdef...64-lowercase-hex-characters +``` + +For digest `abcdef...`, the canonical row-major scalar payload is stored at: + +```text +/ab/abcdef....bin +``` + +The registered checksum must be `sha256:`. The service verifies: + +- exact model-contract root mapping; +- reference/checksum agreement; +- regular file and exact descriptor byte length; +- SHA-256 content digest; +- finite f16/f32 values; +- descriptor shape and dtype. + +Directory and file symlinks below the configured root are rejected. Cache +publishers should write a temporary file, flush it, and atomically rename it to +the digest path only after the complete payload is available. Existing digest +paths are immutable. + +Example SQL descriptor values: + +```text +tensor_ref = sha256://0123...cdef +tensor_checksum = sha256:0123...cdef +tensor_rows = 747 +tensor_dim = 320 +tensor_dtype = float16 +``` + +## Start the Service + +```text +python3 -m services.tilemaxsim_cuda_sidecar \ + --socket /run/vectorchord/tilemaxsim.sock \ + --socket-mode 660 \ + --device cuda:0 \ + --contract-root 'colqwen3.5@revision+preprocessing-hash=/var/cache/gbrain/colqwen35' \ + --request-timeout-ms 2000 \ + --max-request-bytes 1073741824 \ + --max-tensor-bytes 1073741824 \ + --max-batch-tokens 1000000 \ + --max-device-bytes 8589934592 \ + --cache-bytes 8589934592 \ + --max-inflight 8 \ + --max-cuda-inflight 1 +``` + +Run the process under an operations-managed supervisor. The socket directory, +process user/group, cache roots, CUDA device visibility, and resource limits +belong in that supervisor configuration. `SIGINT` and `SIGTERM` stop new +accepts, drain accepted work, and remove only the socket inode created by the +process. + +The sidecar and PostgreSQL limits must agree. IPC v1 embeds tensor payloads, so +`--max-request-bytes` must accommodate the inline batch. IPC v2 carries only +descriptors but independently accounts for all declared canonical tensor bytes. +PostgreSQL's `vchordrq.maxsim_gpu_timeout_ms` and the service request timeout +both cover the complete request; operations should leave enough margin for +normal socket scheduling without allowing abandoned work to run indefinitely. + +## Batching, Backpressure, and Failure Semantics + +- The listener accepts at most `--max-inflight` connections. Further clients + remain in the bounded Unix-socket backlog and are governed by their + PostgreSQL-side deadline. +- `--max-cuda-inflight` bounds simultaneous work on one device. Waiting for a + slot consumes the same overall request deadline. +- A request may be split into device-memory-bounded candidate groups. Results + are returned only after every group succeeds. +- Peer disconnect and deadline checks occur before every device group. +- CUDA OOM, missing/corrupt tensors, unsupported references, timeout, and + non-finite scores fail the complete response. Partial candidate results are + never returned. +- f16 input scalars are promoted to f32 before matrix multiplication. TF32 is + disabled by default; `--allow-tf32` must be enabled only after corpus-level + score/ranking tolerance is accepted. + +The service emits one-line JSON events to standard output. Request events +contain request ID, protocol version, source kind, dimensions, candidate and +token counts, content-cache hits, resolution time, CUDA queue time, compute +time, total time, status, and the Unix peer PID/UID/GID where supported. The +peer PID plus request ID disambiguates PostgreSQL backend-local request +counters. Tensor references and public/application IDs are not logged. + +## Verification and Load Probe + +Protocol, resolver, deadline, device-chunking, live-socket, CPU-equivalence, +and optional CUDA-equivalence tests: + +```text +python3 -m unittest -v \ + devtools/test_tilemaxsim_reference_sidecar.py \ + services/test_tilemaxsim_cuda_sidecar.py +``` + +Synthetic ColQwen-shaped load probe: + +```text +python3 -m services.benchmark_tilemaxsim_cuda \ + --device cuda:0 \ + --dtype f16 \ + --dimension 320 \ + --query-rows 32 \ + --document-rows 747 \ + --candidates 256 \ + --warmup 3 \ + --iterations 20 +``` + +The command emits reproducibility metadata, latency percentiles, peak CUDA +allocation/reservation, and a deterministic score checksum as JSON. It is a +runtime smoke/load probe, not a substitute for recall and end-to-end latency on +the committed GBrain corpus. diff --git a/docs/TILEMAXSIM_IPC_V1.md b/docs/TILEMAXSIM_IPC_V1.md new file mode 100644 index 00000000..274258c2 --- /dev/null +++ b/docs/TILEMAXSIM_IPC_V1.md @@ -0,0 +1,160 @@ +# VectorChord TileMaxSim IPC Protocol v1 + +## Scope + +This protocol connects a PostgreSQL VectorChord backend to a local TileMaxSim +GPU sidecar. It carries full query/page tensors for Phase 3A. It is not a +network authentication protocol and must be exposed only through an +operations-controlled Unix socket. + +All integers and floating-point bit patterns are little-endian. The sidecar +must read and write exact frame lengths. Partial results are forbidden. + +## Common Frame Header + +Every frame starts with 24 bytes: + +| Offset | Size | Type | Meaning | +| --- | ---: | --- | --- | +| 0 | 4 | bytes | ASCII magic `VCTM` | +| 4 | 2 | `u16` | protocol version, currently `1` | +| 6 | 2 | `u16` | message kind: `1` request, `2` response | +| 8 | 8 | `u64` | request ID, echoed unchanged by the response | +| 16 | 8 | `u64` | body length, excluding the 24-byte header | + +Unknown versions, message kinds, request IDs, and inconsistent lengths are +fatal protocol errors. + +## Rerank Request Body + +The fixed portion is 16 bytes: + +| Offset in body | Size | Type | Meaning | +| --- | ---: | --- | --- | +| 0 | 4 | `u32` | tensor dimension | +| 4 | 4 | `u32` | query row/token count | +| 8 | 4 | `u32` | candidate page count | +| 12 | 1 | `u8` | dtype: `1` f32, `2` IEEE f16 | +| 13 | 1 | `u8` | scoring: `1` = sum of query-token max document dot products | +| 14 | 2 | `u16` | reserved, must be zero | + +The fixed portion is followed by the query tensor in row-major order. Each +candidate then has: + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | opaque request-local candidate ID | +| 4 | `u32` | page tensor row/token count | +| `rows * dimension * dtype_size` | bytes | row-major page tensor | + +Candidate IDs are assigned densely from zero in v1, but the sidecar must treat +them as opaque and echo them unchanged. Heap CTIDs and public GBrain page IDs +are not exposed to the sidecar. + +Every query and candidate tensor in one request has the same dtype and +dimension. v1 accepts f32 and f16 only. RaBitQ arrays must use CPU exact or a +future protocol version that defines their representation. + +## Successful Response Body + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | status `0` | +| 4 | `u32` | result count; must equal request candidate count | + +Each result then contains: + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | echoed candidate ID | +| 4 | `f32` bits | positive MaxSim similarity | + +Each requested candidate must appear exactly once. Unknown IDs, duplicate IDs, +missing IDs, non-finite scores, and trailing bytes fail the entire SQL query. +VectorChord converts public positive similarity to its SQL distance convention +by negating the value before ascending sort. Equal exact distances are ordered +deterministically by the candidate's internal heap key after response IDs have +been mapped back inside the backend. + +## Error Response Body + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | nonzero sidecar-defined status | +| 4 | `u32` | UTF-8 error-message length | +| `length` | bytes | UTF-8 diagnostic, maximum 64 KiB | + +The status namespace is reserved for the sidecar implementation in v1. Any +nonzero status fails `gpu` mode. In `auto` mode it triggers a complete CPU exact +rerank of the original candidate set; no partial GPU scores are reused. + +## Limits and Timeouts + +Before connecting, VectorChord enforces: + +- `vchordrq.maxsim_candidate_limit`; +- `vchordrq.maxsim_gpu_max_batch_tokens`, including query and page tokens; +- `vchordrq.maxsim_gpu_max_batch_bytes`, including the request frame. + +`vchordrq.maxsim_gpu_timeout_ms` is one overall connect-plus-write-plus-read +deadline. Connection and backend I/O poll PostgreSQL interrupts at least every +50 ms. The response is also bounded from the expected candidate count and the +64 KiB error limit. + +v1 sends one bounded request. Splitting a candidate set into multiple GPU +batches is future work and must preserve one overall SQL deadline and +all-or-nothing result semantics. + +External immutable tensor descriptors use the separate Phase 3B +[`TILEMAXSIM_IPC_V2`](TILEMAXSIM_IPC_V2.md) contract. They are never smuggled +into a v1 inline-tensor frame. + +## Reference Sidecar + +`devtools/tilemaxsim_reference_sidecar.py` is a dependency-free, single-threaded +CPU protocol oracle for v1 and resolver-injected v2 tests. It is intended for +codec and end-to-end development only; it is not a production or GPU executor. +The production-oriented CUDA implementation is documented in +[`TILEMAXSIM_CUDA_SIDECAR`](TILEMAXSIM_CUDA_SIDECAR.md). + +Start one request/response cycle with: + +```text +python3 devtools/tilemaxsim_reference_sidecar.py \ + --socket /tmp/vectorchord-tilemaxsim.sock \ + --once +``` + +Then configure the PostgreSQL session as an administrator and select the GPU +backend in the query session: + +```sql +SET vchordrq.maxsim_gpu_endpoint = '/tmp/vectorchord-tilemaxsim.sock'; +SET vchordrq.maxsim_candidate_limit = 256; +SET vchordrq.maxsim_backend = 'gpu'; +``` + +The endpoint setting is `SUSET`. The reference sidecar creates its socket with +mode `0600` by default and refuses to replace a non-socket filesystem path. +Its default request limit is intentionally 64 MiB, lower than the extension's +configurable production ceiling. + +Run its protocol and live-socket tests with: + +```text +python3 -m unittest -v devtools/test_tilemaxsim_reference_sidecar.py +``` + +## Sidecar Conformance Cases + +A conforming implementation must be tested for: + +- f32 and f16 score equivalence with CPU exact; +- reordered response IDs; +- duplicate, missing, and unknown IDs; +- invalid magic/version/kind/request ID/body length; +- truncated and trailing data; +- NaN and infinite similarity; +- explicit error response; +- connection close during request and response; +- queue/compute duration exceeding the overall deadline. diff --git a/docs/TILEMAXSIM_IPC_V2.md b/docs/TILEMAXSIM_IPC_V2.md new file mode 100644 index 00000000..d5e211df --- /dev/null +++ b/docs/TILEMAXSIM_IPC_V2.md @@ -0,0 +1,173 @@ +# VectorChord TileMaxSim IPC Protocol v2 and v3 + +## Scope + +Version 2 carries an inline query tensor plus immutable external tensor +descriptors. It is the Phase 3B transport for a coarse/sketch `vchordrq` index +whose final score comes from a different full tensor. It does not change the +meaning of ordinary SQL `@#` and is not enabled by source registration alone. + +Version 3 is the backward-compatible scheduled form. It adds an authenticated +upstream scheduling domain, priority, and end-to-end timeout. These fields +control resource scheduling only; they are not authorization claims and never +replace PostgreSQL row visibility or application ACL checks. + +The transport is an operations-controlled Unix socket. It is not a network +authentication protocol. All integers and floating-point bit patterns are +little-endian. The sidecar must read and write exact frame lengths; partial +results are forbidden. + +## Common Header + +The 24-byte header is the same shape as v1: + +| Offset | Size | Type | Meaning | +| --- | ---: | --- | --- | +| 0 | 4 | bytes | ASCII magic `VCTM` | +| 4 | 2 | `u16` | protocol version `2` | +| 6 | 2 | `u16` | message kind: `1` request, `2` response | +| 8 | 8 | `u64` | request ID, echoed unchanged | +| 16 | 8 | `u64` | body length, excluding the header | + +The response must echo the request version. Unknown versions, message kinds, +request IDs, or inconsistent lengths fail the complete request. + +## Version 3 Scheduling Extension + +For a version 3 request, the v2 fixed request body is immediately followed by +12 additional fixed bytes: + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `i32` | priority from -100 through 100; higher values are more urgent | +| 4 | `u32` | end-to-end timeout in milliseconds, from 1 through 600000 | +| 4 | `u32` | scheduling-tenant UTF-8 byte length | + +The variable body then contains model-contract bytes, scheduling-tenant bytes, +the query tensor, and the candidate descriptors, in that order. The tenant is +nonempty, limited to 256 bytes, and cannot contain control characters. + +The timeout begins when the daemon accepts the connection, not when GPU work +starts. A server-side timeout may shorten it. Priority affects queue order only; +it cannot expand the request deadline, cache quota, candidate scope, or access +rights. Protocol v2 requests enter the default scheduling domain at priority +zero and use the server timeout. + +## External Rerank Request + +The fixed request body is 20 bytes: + +| Offset | Size | Type | Meaning | +| --- | ---: | --- | --- | +| 0 | 4 | `u32` | tensor dimension | +| 4 | 4 | `u32` | query row/token count | +| 8 | 4 | `u32` | candidate count | +| 12 | 1 | `u8` | dtype: `1` f32, `2` IEEE f16 | +| 13 | 1 | `u8` | scoring: `1` = sum of query-token max document dot products | +| 14 | 2 | `u16` | reserved, must be zero | +| 16 | 4 | `u32` | UTF-8 model-contract byte length | + +The fixed body is followed, in order, by: + +1. the model-contract UTF-8 bytes; +2. the inline query tensor in row-major order; +3. exactly `candidate_count` descriptor entries. + +Each descriptor entry is: + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | opaque request-local candidate ID | +| 4 | `u32` | full tensor row/token count | +| 4 | `u32` | tensor-reference UTF-8 byte length | +| 4 | `u32` | checksum UTF-8 byte length | +| variable | bytes | tensor reference | +| variable | bytes | checksum | + +All tensors in one request have the fixed body's dtype and dimension. The +backend validates each registered row descriptor against them before IPC. The +model contract and every descriptor field are nonempty, bounded UTF-8 without +control characters. Current bounds are 512 bytes for the contract and +checksum, and 4096 bytes for the reference. + +The checksum encoding is `sha256:` followed by 64 lowercase hexadecimal +digits. It covers the canonical row-major scalar payload after any immutable +object envelope is decoded. This makes verification independent of an object +store's metadata and compression. A resolver for formats such as safetensors +must decode the selected tensor, validate its declared shape/dtype, produce the +canonical payload, and then verify this digest before compute. + +Candidate IDs are opaque. PostgreSQL heap CTIDs, registered public document +IDs, application routing identifiers, and credentials are never encoded. +VectorChord maps response IDs back to heap keys inside the backend. + +## Resolution and Security + +The sidecar owns descriptor resolution. A production resolver must: + +- allowlist schemes, buckets/namespaces, and immutable reference forms; +- obtain credentials from operations-managed sidecar configuration, never from + PostgreSQL rows or this protocol; +- reject redirects or network destinations outside its allowlist; +- enforce the model contract, dtype, shape, and checksum before GPU access; +- bound queue, fetch, decode, host-to-device, and compute work by the one SQL + request deadline. + +The complete control frame must be parsed and validated before external I/O. +An absent resolver fails closed. The dependency-free reference sidecar exposes +v2 resolution only through an injected callback; its CLI deliberately has no +filesystem, HTTP, or object-store resolver. + +## Limits + +Both peers enforce: + +- candidate count; +- total query plus declared document tokens; +- total declared canonical tensor bytes; +- control-frame bytes; +- one overall connect/write/read deadline. + +`vchordrq.maxsim_gpu_max_batch_tokens` bounds total declared tokens. +`vchordrq.maxsim_gpu_max_batch_bytes` bounds both the control frame and total +declared canonical tensor bytes independently. A small descriptor frame cannot +therefore authorize unbounded object loads. + +Version 2 currently sends one bounded batch. Future splitting must retain one +overall deadline and all-or-nothing result semantics. + +## Response + +The success and error bodies are identical to v1 except that the header version +echoes 2 or 3. A success body begins with status `0` and a result count equal to the +number of accepted descriptors, followed by `(candidate_id u32, similarity +f32)` pairs. Every candidate appears exactly once. Unknown, duplicate, missing, +non-finite, or trailing results fail the whole SQL operation. + +Similarity is positive public MaxSim. VectorChord negates it to its ascending +distance convention and uses the internal heap key as a deterministic tie +breaker for the query execution. + +## Current Integration State + +The Rust descriptor source, strict v2/v3 encoder/decoder, and reference-sidecar +protocol oracle are implemented and unit tested. The internal runtime resolver +also consumes the privilege-aware SQL registry boundary and produces physical +attribute bindings. The restricted `vchordrq_maxsim_search` score API now sends +validated, SQL-visible external descriptors over v2 and returns exact positive +similarity. Ordinary `@#` scans keep their existing stored-array semantics. + +The native Rust/CUDA daemon additionally provides bounded global and per-tenant +admission, fair/priority/fair-priority policies, request aging, candidate/token +quanta, deadline and disconnect cancellation between CUDA launches, per-tenant +cache caps, optional reservations, and immutable-shard reload on `SIGHUP`. +An optional operations-only Unix status socket exposes HTTP `/healthz` and +Prometheus `/metrics`; it is separate from this binary scoring protocol. + +The dependency-free reference sidecar remains a protocol oracle. The separate +[`TILEMAXSIM_CUDA_SIDECAR`](TILEMAXSIM_CUDA_SIDECAR.md) implementation executes +v1/v2 on CUDA and provides a per-model, allowlisted, content-addressed local +resolver with checksum, resource-limit, deadline, backpressure, and structured +metric enforcement. GBrain remains responsible for populating that immutable +node-local cache from its storage system. Production acceptance still requires +the committed GBrain corpus recall, latency, concurrency, and failure matrix. diff --git a/services/Dockerfile.tilemaxsim b/services/Dockerfile.tilemaxsim index 9f58b88d..2bf87252 100644 --- a/services/Dockerfile.tilemaxsim +++ b/services/Dockerfile.tilemaxsim @@ -5,6 +5,9 @@ WORKDIR /opt/vectorchord COPY devtools/tilemaxsim_reference_sidecar.py devtools/tilemaxsim_reference_sidecar.py COPY services/tilemaxsim_cuda_sidecar.py services/tilemaxsim_cuda_sidecar.py +COPY services/tilemaxsim_gpu_cache.py services/tilemaxsim_gpu_cache.py +COPY services/tilemaxsim_shard.py services/tilemaxsim_shard.py +COPY services/tilemaxsim_triton.py services/tilemaxsim_triton.py ENV PYTHONPATH=/opt/vectorchord \ PYTHONUNBUFFERED=1 diff --git a/services/Dockerfile.tilemaxsimd b/services/Dockerfile.tilemaxsimd new file mode 100644 index 00000000..d6dcb2f2 --- /dev/null +++ b/services/Dockerfile.tilemaxsimd @@ -0,0 +1,22 @@ +ARG CUDA_DEVEL_IMAGE=nvidia/cuda:12.6.3-devel-ubuntu24.04 +ARG CUDA_RUNTIME_IMAGE=nvidia/cuda:12.6.3-runtime-ubuntu24.04 + +FROM ${CUDA_DEVEL_IMAGE} AS build + +RUN apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + build-essential ca-certificates curl \ + && rm -rf /var/lib/apt/lists/* +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --profile minimal --default-toolchain 1.96.0 + +ENV PATH=/root/.cargo/bin:${PATH} +WORKDIR /build/tilemaxsimd +COPY services/tilemaxsimd/ . +RUN cargo build --release --locked + +FROM ${CUDA_RUNTIME_IMAGE} + +COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimd /usr/local/bin/tilemaxsimd + +ENTRYPOINT ["/usr/local/bin/tilemaxsimd"] diff --git a/services/benchmark_full_corpus_tilemaxsim.py b/services/benchmark_full_corpus_tilemaxsim.py new file mode 100644 index 00000000..42e5b0be --- /dev/null +++ b/services/benchmark_full_corpus_tilemaxsim.py @@ -0,0 +1,190 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Benchmark exact full-corpus TileMaxSim through the native CUDA daemon.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import time +from pathlib import Path + +import numpy as np + +from services.benchmark_tilemaxsim_ablation import encode_frame, request_round_trip + +MAX_BATCH_CANDIDATES = 65_536 +MAX_BATCH_TOKENS = 1_000_000 +MAX_BATCH_TENSOR_BYTES = 1024**3 + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + return ordered[max(0, math.ceil(len(ordered) * fraction) - 1)] + + +def summary(samples: list[float]) -> dict[str, float | int]: + return { + "count": len(samples), + "mean": statistics.fmean(samples), + "p50": percentile(samples, 0.50), + "p95": percentile(samples, 0.95), + "p99": percentile(samples, 0.99), + "max": max(samples), + } + + +def load_jsonl(path: Path) -> list[dict[str, object]]: + with path.open(encoding="utf-8") as stream: + records = [json.loads(line) for line in stream if line.strip()] + if not records: + raise ValueError(f"empty JSONL file: {path}") + return records + + +def descriptor_batches( + descriptors: list[dict[str, object]], query_rows: int +) -> list[tuple[int, list[dict[str, object]]]]: + batches: list[tuple[int, list[dict[str, object]]]] = [] + start = 0 + current: list[dict[str, object]] = [] + tokens = query_rows + tensor_bytes = query_rows * int(descriptors[0]["tensor_dim"]) * 2 + for descriptor in descriptors: + rows = int(descriptor["tensor_rows"]) + scalar_bytes = 2 if descriptor["tensor_dtype"] == "float16" else 4 + payload_bytes = rows * int(descriptor["tensor_dim"]) * scalar_bytes + would_overflow = current and ( + len(current) == MAX_BATCH_CANDIDATES + or tokens + rows > MAX_BATCH_TOKENS + or tensor_bytes + payload_bytes > MAX_BATCH_TENSOR_BYTES + ) + if would_overflow: + batches.append((start, current)) + start += len(current) + current = [] + tokens = query_rows + tensor_bytes = query_rows * int(descriptor["tensor_dim"]) * scalar_bytes + current.append(descriptor) + tokens += rows + tensor_bytes += payload_bytes + if current: + batches.append((start, current)) + return batches + + +def recall(expected: np.ndarray, actual: list[int], top_k: int) -> float: + return len(set(expected[:top_k].tolist()).intersection(actual[:top_k])) / top_k + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--query-dataset", required=True, type=Path) + parser.add_argument("--query-embeddings", required=True, type=Path) + parser.add_argument("--gold", required=True, type=Path) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--contract", required=True) + parser.add_argument("--query-limit", type=int, default=0) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + + descriptors = load_jsonl(args.descriptor_manifest) + queries = json.loads(args.query_dataset.read_text(encoding="utf-8")) + if not isinstance(queries, list) or not queries: + raise ValueError("query dataset must be a nonempty JSON array") + if args.query_limit < 0: + parser.error("--query-limit must be nonnegative") + queries = queries[: args.query_limit or None] + gold = np.load(args.gold, allow_pickle=False) + query_latencies: list[float] = [] + round_trip_latencies: list[float] = [] + batch_latencies: list[float] = [] + cases = [] + + for query_index, query_record in enumerate(queries): + query_id = query_record.get("query_id") + if not isinstance(query_id, str): + raise ValueError(f"query {query_index} has no query_id") + query = np.load(args.query_embeddings / f"{query_id}.npy", allow_pickle=False) + if query.dtype != np.dtype("float16"): + query = query.astype(" float: + ordered = sorted(samples) + index = max(0, math.ceil(fraction * len(ordered)) - 1) + return ordered[index] + + +def sql_halfvec_array(tensor: np.ndarray) -> str: + if tensor.ndim != 2 or tensor.shape[1] <= 0 or tensor.shape[0] <= 0: + raise ValueError("query tensor must have shape [rows, dimension]") + if tensor.dtype not in (np.dtype("float16"), np.dtype("float32")): + raise ValueError("query tensor must use float16 or float32") + if not np.isfinite(tensor).all(): + raise ValueError("query tensor contains non-finite values") + vectors = [] + for row in tensor: + value = "[" + ",".join(format(float(item), ".8g") for item in row) + "]" + vectors.append("'" + value + "'::halfvec") + return "ARRAY[" + ",".join(vectors) + "]" + + +def load_manifest(path: Path) -> list[str]: + page_keys = [] + with path.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + record = json.loads(line) + page_key = record.get("page_key") + if not isinstance(page_key, str) or not PAGE_KEY.fullmatch(page_key): + raise ValueError(f"invalid page_key at manifest line {line_number}") + page_keys.append(page_key) + if not page_keys: + raise ValueError("page manifest is empty") + return page_keys + + +def load_queries( + results_path: Path, embeddings: Path, limit: int +) -> list[tuple[str, Path]]: + results = json.loads(results_path.read_text(encoding="utf-8")) + cases = results.get("queries", {}).get("cases") + if not isinstance(cases, list) or not cases: + raise ValueError("results JSON has no queries.cases") + queries = [] + for case in cases[: limit or None]: + query_id = case.get("query_id") + if not isinstance(query_id, str) or not re.fullmatch(r"[0-9a-f]+", query_id): + raise ValueError("invalid query ID in results JSON") + path = embeddings / f"{query_id}.npy" + if not path.is_file(): + raise ValueError(f"missing query embedding: {path}") + queries.append((query_id, path)) + return queries + + +def recall(expected: list[str], actual: list[str], top_k: int) -> float: + expected_set = set(expected[:top_k]) + return len(expected_set.intersection(actual[:top_k])) / top_k + + +def candidate_recall(expected: list[str], candidates: list[str], top_k: int) -> float: + expected_set = set(expected[:top_k]) + return len(expected_set.intersection(candidates)) / top_k + + +def parse_profile(stderr: str) -> dict[str, int] | None: + profiles = [] + for line in stderr.splitlines(): + marker = line.find(PROFILE_PREFIX) + if marker < 0: + continue + payload = line[marker + len(PROFILE_PREFIX) :].strip() + profile = json.loads(payload) + if not isinstance(profile, dict) or profile.get("schema_version") != 1: + raise RuntimeError("unexpected MaxSim profile schema") + if any(not isinstance(value, int) for value in profile.values()): + raise RuntimeError("MaxSim profile values must be integers") + profiles.append(profile) + if len(profiles) > 1: + raise RuntimeError("query emitted multiple MaxSim profiles") + return profiles[0] if profiles else None + + +def execute_query( + psql_command: list[str], + table: str, + index: str | None, + query_sql: str, + endpoint: str, + probes: str, + refine: int, + candidate_limit: int, + timeout_ms: int, + profile: bool, +) -> tuple[list[str], float, dict[str, int] | None]: + if index is None: + search_sql = f""" +SELECT page_key +FROM {table} +ORDER BY embedding @# {query_sql} +LIMIT {candidate_limit}; +""" + else: + search_sql = f""" +SELECT p.page_key +FROM vchordrq_maxsim_search( + '{index}'::regclass, + {query_sql}, + {candidate_limit}, + {candidate_limit} + ) WITH ORDINALITY AS r(public_id, similarity, result_order) +JOIN {table} AS p ON p.id = r.public_id +ORDER BY r.result_order; +""" + sql = f""" +SET statement_timeout = '{timeout_ms}ms'; +SET enable_seqscan = off; +SET vchordrq.probes = '{probes}'; +SET vchordrq.maxsim_refine = {refine}; +SET vchordrq.maxsim_candidate_limit = {candidate_limit}; +SET vchordrq.maxsim_gpu_endpoint = '{endpoint}'; +SET vchordrq.maxsim_gpu_timeout_ms = {timeout_ms}; +SET vchordrq.maxsim_backend = 'gpu'; +SET vchordrq.maxsim_profile = {"on" if profile else "off"}; +{search_sql} +""" + started = time.perf_counter() + result = subprocess.run( + [*psql_command, "-X", "-q", "-A", "-t", "-v", "ON_ERROR_STOP=1"], + input=sql, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=max(5.0, timeout_ms / 1000.0 + 5.0), + check=False, + ) + latency_ms = (time.perf_counter() - started) * 1000.0 + if result.returncode != 0: + diagnostic = result.stderr.strip() or result.stdout.strip() + raise RuntimeError(f"psql failed: {diagnostic}") + page_keys = [line.strip() for line in result.stdout.splitlines() if line.strip()] + invalid = [page_key for page_key in page_keys if not PAGE_KEY.fullmatch(page_key)] + if invalid: + raise RuntimeError(f"psql returned non-page-key output: {invalid[0]!r}") + if len(page_keys) != len(set(page_keys)): + raise RuntimeError("Phase 3 query returned duplicate page keys") + query_profile = parse_profile(result.stderr) + if profile and query_profile is None: + raise RuntimeError("query did not emit a MaxSim profile") + return page_keys, latency_ms, query_profile + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--results-json", required=True, type=Path) + parser.add_argument("--query-embeddings", required=True, type=Path) + parser.add_argument("--gold", required=True, type=Path) + parser.add_argument("--psql-command", default="psql") + parser.add_argument("--table", default="kb_colqwen_pages") + parser.add_argument( + "--index", + help="use the Phase 3B external-tensor API with this MaxSim index", + ) + parser.add_argument("--endpoint", required=True) + parser.add_argument("--probes", default="8") + parser.add_argument("--refine", type=int, default=512) + parser.add_argument("--candidate-limit", type=positive_int, default=100) + parser.add_argument("--query-limit", type=int, default=0) + parser.add_argument("--timeout-ms", type=positive_int, default=60000) + parser.add_argument( + "--profile", + action="store_true", + help="enable vchordrq.maxsim_profile and include phase metrics", + ) + args = parser.parse_args() + + if not IDENTIFIER.fullmatch(args.table): + parser.error( + "--table must be an unquoted SQL identifier, optionally schema-qualified" + ) + if args.index is not None and not IDENTIFIER.fullmatch(args.index): + parser.error( + "--index must be an unquoted SQL identifier, optionally schema-qualified" + ) + if not args.endpoint.startswith("/") or "'" in args.endpoint: + parser.error("--endpoint must be an absolute Unix-socket path without quotes") + if args.refine < 0: + parser.error("--refine must be nonnegative") + if args.query_limit < 0: + parser.error("--query-limit must be nonnegative") + if args.profile and args.index is None: + parser.error("--profile requires --index") + if not re.fullmatch(r"[0-9]+(,[0-9]+)*", args.probes): + parser.error("--probes must be a comma-separated list of nonnegative integers") + psql_command = shlex.split(args.psql_command) + if not psql_command: + parser.error("--psql-command must not be empty") + + page_keys = load_manifest(args.manifest) + queries = load_queries(args.results_json, args.query_embeddings, args.query_limit) + gold = np.load(args.gold, allow_pickle=False) + per_query = [] + latencies = [] + for query_number, (query_id, query_path) in enumerate(queries): + index_key = f"q{query_number}_idx" + if index_key not in gold: + raise ValueError(f"gold archive is missing {index_key}") + gold_indices = gold[index_key].astype(np.int64, copy=False).tolist() + if any(index < 0 or index >= len(page_keys) for index in gold_indices): + raise ValueError(f"gold archive {index_key} contains an invalid page index") + expected = [page_keys[index] for index in gold_indices] + query = np.load(query_path, allow_pickle=False) + actual, latency_ms, query_profile = execute_query( + psql_command, + args.table, + args.index, + sql_halfvec_array(query), + args.endpoint, + args.probes, + args.refine, + args.candidate_limit, + args.timeout_ms, + args.profile, + ) + latencies.append(latency_ms) + per_query.append( + { + "query_id": query_id, + "query_rows": int(query.shape[0]), + "returned_candidates": len(actual), + "latency_ms": round(latency_ms, 3), + "recall_at_10": recall(expected, actual, 10), + "recall_at_20": recall(expected, actual, 20), + "candidate_recall_at_10": candidate_recall(expected, actual, 10), + "candidate_recall_at_20": candidate_recall(expected, actual, 20), + "top_10": actual[:10], + **({"profile": query_profile} if query_profile is not None else {}), + } + ) + + output = { + "benchmark": ( + "gbrain_vectorchord_phase3b_external_v1" + if args.index is not None + else "gbrain_vectorchord_phase3a_v1" + ), + "corpus_pages": len(page_keys), + "query_count": len(per_query), + "configuration": { + "table": args.table, + "index": args.index, + "probes": args.probes, + "maxsim_refine": args.refine, + "candidate_limit": args.candidate_limit, + }, + "metrics": { + "recall_at_10": sum(item["recall_at_10"] for item in per_query) + / len(per_query), + "recall_at_20": sum(item["recall_at_20"] for item in per_query) + / len(per_query), + "candidate_recall_at_10": sum( + item["candidate_recall_at_10"] for item in per_query + ) + / len(per_query), + "candidate_recall_at_20": sum( + item["candidate_recall_at_20"] for item in per_query + ) + / len(per_query), + }, + "latency_ms": { + "mean": sum(latencies) / len(latencies), + "p50": percentile(latencies, 0.50), + "p95": percentile(latencies, 0.95), + "p99": percentile(latencies, 0.99), + "max": max(latencies), + }, + "per_query": per_query, + } + profiles = [item["profile"] for item in per_query if "profile" in item] + if profiles: + profile_keys = sorted(set.intersection(*(set(profile) for profile in profiles))) + output["profile_summary"] = { + key: { + "mean": sum(profile[key] for profile in profiles) / len(profiles), + "p50": percentile([profile[key] for profile in profiles], 0.50), + "p95": percentile([profile[key] for profile in profiles], 0.95), + "max": max(profile[key] for profile in profiles), + } + for key in profile_keys + if key != "schema_version" + } + print(json.dumps(output, ensure_ascii=False, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/benchmark_gbrain_scoped_tilemaxsim.py b/services/benchmark_gbrain_scoped_tilemaxsim.py new file mode 100644 index 00000000..96c4d646 --- /dev/null +++ b/services/benchmark_gbrain_scoped_tilemaxsim.py @@ -0,0 +1,252 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Benchmark real CUDA TileMaxSim over a GBrain-style lexical/structured scope.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import statistics +import time +from pathlib import Path + +import numpy as np + +from services.benchmark_full_corpus_tilemaxsim import descriptor_batches +from services.benchmark_tilemaxsim_ablation import encode_frame, request_round_trip + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + return ordered[max(0, math.ceil(len(ordered) * fraction) - 1)] + + +def summary(samples: list[float]) -> dict[str, float | int]: + return { + "count": len(samples), + "mean": statistics.fmean(samples), + "p50": percentile(samples, 0.50), + "p95": percentile(samples, 0.95), + "p99": percentile(samples, 0.99), + "max": max(samples), + } + + +def load_jsonl(path: Path) -> list[dict[str, object]]: + with path.open(encoding="utf-8") as stream: + return [json.loads(line) for line in stream if line.strip()] + + +def select_method(report: dict[str, object], method_name: str) -> dict[str, object]: + methods = report.get("methods") + if not isinstance(methods, list): + raise ValueError("candidate report has no methods array") + for method in methods: + if isinstance(method, dict) and method.get("method") == method_name: + return method + raise ValueError(f"candidate method {method_name!r} was not found") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--page-manifest", required=True, type=Path) + parser.add_argument("--query-dataset", required=True, type=Path) + parser.add_argument("--corpus-json", required=True, type=Path) + parser.add_argument("--query-embeddings", required=True, type=Path) + parser.add_argument("--candidate-report", required=True, type=Path) + parser.add_argument("--candidate-method", default="es_bm25") + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--contract", required=True) + parser.add_argument("--query-limit", type=int, default=0) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + + raw_descriptors = load_jsonl(args.descriptor_manifest) + page_metadata = { + str(record["page_key"]): record for record in load_jsonl(args.page_manifest) + } + descriptors = [ + {**descriptor, **page_metadata.get(str(descriptor.get("page_key")), {})} + for descriptor in raw_descriptors + ] + queries = json.loads(args.query_dataset.read_text(encoding="utf-8")) + corpus_records = json.loads(args.corpus_json.read_text(encoding="utf-8")) + candidate_report = json.loads(args.candidate_report.read_text(encoding="utf-8")) + method = select_method(candidate_report, args.candidate_method) + per_query = method.get("per_query_results") + if not isinstance(queries, list) or not queries: + raise ValueError("query dataset must be a nonempty JSON array") + if not isinstance(corpus_records, list): + raise ValueError("corpus JSON must be an array") + if not isinstance(per_query, dict): + raise ValueError("candidate method has no per_query_results object") + if args.query_limit < 0: + parser.error("--query-limit must be nonnegative") + queries = queries[: args.query_limit or None] + + descriptors_by_doc: dict[str, list[dict[str, object]]] = {} + descriptors_by_doc_page: dict[tuple[str, int], list[dict[str, object]]] = {} + for descriptor in descriptors: + doc_name = descriptor.get("doc_name") + if isinstance(doc_name, str): + descriptors_by_doc.setdefault(doc_name, []).append(descriptor) + descriptors_by_doc_page.setdefault( + (doc_name, int(descriptor.get("page_no", 0))), [] + ).append(descriptor) + chunks_by_id = { + str(record["_id"]): record["_source"] for record in corpus_records + } + + cases: list[dict[str, object]] = [] + scope_latencies: list[float] = [] + cuda_latencies: list[float] = [] + end_to_end_latencies: list[float] = [] + for query_index, query_record in enumerate(queries): + query_id = query_record.get("query_id") + gold_doc = query_record.get("gold_doc_name") + if not isinstance(query_id, str) or not isinstance(gold_doc, str): + raise ValueError(f"query {query_index} lacks query_id or gold_doc_name") + raw_candidates = per_query.get(query_id, []) + if not isinstance(raw_candidates, list): + raise ValueError(f"candidate list for {query_id} is not an array") + + scope_started = time.perf_counter() + candidate_docs: list[str] = [] + seen_docs: set[str] = set() + for row in raw_candidates: + doc_name = row.get("doc_name") if isinstance(row, dict) else None + if isinstance(doc_name, str) and doc_name not in seen_docs: + seen_docs.add(doc_name) + candidate_docs.append(doc_name) + scoped_descriptors: list[dict[str, object]] = [] + seen_pages: set[str] = set() + for row in raw_candidates: + if not isinstance(row, dict): + continue + chunk_id = str(row.get("chunk_id", "")) + doc_name = row.get("doc_name") + source = chunks_by_id.get(chunk_id) + content = source.get("content_with_weight", "") if isinstance(source, dict) else "" + page_numbers = { + int(value) for value in re.findall(r"::(\d+)", str(content)) + } + mapped = [ + descriptor + for page_no in sorted(page_numbers) + for descriptor in descriptors_by_doc_page.get((str(doc_name), page_no), []) + ] + # Preserve recall when an imported chunk has no page markers. + if not mapped and isinstance(doc_name, str): + mapped = descriptors_by_doc.get(doc_name, []) + for descriptor in mapped: + page_key = str(descriptor.get("page_key", "")) + if page_key and page_key not in seen_pages: + seen_pages.add(page_key) + scoped_descriptors.append(descriptor) + scope_ms = (time.perf_counter() - scope_started) * 1000 + scope_latencies.append(scope_ms) + + query = np.load(args.query_embeddings / f"{query_id}.npy", allow_pickle=False) + if query.dtype != np.dtype("float16"): + query = query.astype(" None: + self.model_contract_id = model_contract_id + self.records = self._load_records(source_manifest, descriptor_manifest) + self.cache = sidecar.PayloadCache(cache_bytes) + + @staticmethod + def _json_lines(path: Path) -> list[dict[str, object]]: + records = [] + with path.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path}:{line_number}: invalid JSON") from error + if not isinstance(record, dict): + raise ValueError(f"{path}:{line_number}: expected a JSON object") + records.append(record) + if not records: + raise ValueError(f"{path}: manifest is empty") + return records + + @classmethod + def _load_records( + cls, source_manifest: Path, descriptor_manifest: Path + ) -> dict[str, TensorRecord]: + source_root = source_manifest.resolve(strict=True).parent + sources: dict[str, tuple[Path, int, int]] = {} + for record in cls._json_lines(source_manifest): + page_key = record.get("page_key") + relative = record.get("embedding_file") + rows = record.get("n_tokens") + dimension = record.get("dim") + if ( + not isinstance(page_key, str) + or not isinstance(relative, str) + or not isinstance(rows, int) + or rows <= 0 + or not isinstance(dimension, int) + or dimension <= 0 + ): + raise ValueError("source manifest contains an invalid tensor record") + path = (source_root / relative).resolve(strict=True) + try: + path.relative_to(source_root) + except ValueError as error: + raise ValueError( + f"embedding path escapes source root: {relative}" + ) from error + if page_key in sources: + raise ValueError(f"duplicate source page_key: {page_key}") + sources[page_key] = (path, rows, dimension) + + resolved: dict[str, TensorRecord] = {} + seen_pages = set() + for descriptor in cls._json_lines(descriptor_manifest): + page_key = descriptor.get("page_key") + tensor_ref = descriptor.get("tensor_ref") + rows = descriptor.get("tensor_rows") + dimension = descriptor.get("tensor_dim") + dtype_name = descriptor.get("tensor_dtype") + checksum = descriptor.get("tensor_checksum") + source = sources.get(page_key) if isinstance(page_key, str) else None + if source is None: + raise ValueError(f"descriptor has unknown page_key: {page_key!r}") + source_path, source_rows, source_dimension = source + if rows != source_rows or dimension != source_dimension: + raise ValueError(f"descriptor shape disagrees for page {page_key}") + dtype = { + "float16": protocol.DTYPE_F16, + "float32": protocol.DTYPE_F32, + }.get(dtype_name) + if dtype is None: + raise ValueError(f"descriptor has invalid dtype for page {page_key}") + if not isinstance(tensor_ref, str) or not tensor_ref.startswith( + "sha256://" + ): + raise ValueError( + f"descriptor has invalid tensor_ref for page {page_key}" + ) + digest = tensor_ref.removeprefix("sha256://") + if ( + len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + or checksum != f"sha256:{digest}" + ): + raise ValueError(f"descriptor has invalid SHA-256 for page {page_key}") + tensor_record = TensorRecord(source_path, rows, dimension, dtype, checksum) + previous = resolved.get(tensor_ref) + if previous is not None: + if ( + previous.rows != tensor_record.rows + or previous.dimension != tensor_record.dimension + or previous.dtype != tensor_record.dtype + or previous.checksum != tensor_record.checksum + ): + raise ValueError(f"conflicting duplicate tensor_ref: {tensor_ref}") + else: + resolved[tensor_ref] = tensor_record + seen_pages.add(page_key) + if seen_pages != set(sources): + raise ValueError("descriptor manifest does not cover the source manifest") + return resolved + + def resolve( + self, request: protocol.ExternalTensorRequest + ) -> sidecar.ResolvedPayload: + if request.model_contract_id != self.model_contract_id: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "model contract is not allowlisted by this benchmark resolver", + ) + record = self.records.get(request.tensor_ref) + if record is None: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor reference is not present in the benchmark manifest", + ) + if ( + request.rows != record.rows + or request.dimension != record.dimension + or request.dtype != record.dtype + or not hmac.compare_digest(request.checksum, record.checksum) + ): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor descriptor disagrees with the benchmark manifest", + ) + key = ( + request.tensor_ref, + request.rows, + request.dimension, + request.dtype, + ) + cached = self.cache.get(key) + if cached is not None: + return sidecar.ResolvedPayload(cached, True) + + dtype_name = "float16" if record.dtype == protocol.DTYPE_F16 else "float32" + tensor = canonical_tensor(record.path, record.rows, record.dimension) + if tensor.dtype != np.dtype(dtype_name): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "NPY dtype disagrees with the benchmark descriptor", + ) + payload = memoryview(tensor).cast("B").tobytes() + actual_checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}" + if not hmac.compare_digest(actual_checksum, record.checksum): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "NPY payload checksum disagrees with the benchmark descriptor", + ) + self.cache.put(key, payload) + return sidecar.ResolvedPayload(payload, False) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=sidecar.parse_mode, default=0o600) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--model-contract", required=True) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument( + "--max-request-bytes", type=sidecar.positive_int, default=64 * 1024 * 1024 + ) + parser.add_argument( + "--max-batch-tokens", type=sidecar.positive_int, default=1_000_000 + ) + parser.add_argument( + "--max-tensor-bytes", type=sidecar.positive_int, default=1024 * 1024 * 1024 + ) + parser.add_argument("--max-candidates", type=sidecar.positive_int, default=65_536) + parser.add_argument( + "--max-device-bytes", + type=sidecar.positive_int, + default=8 * 1024 * 1024 * 1024, + ) + parser.add_argument( + "--cache-bytes", type=sidecar.nonnegative_int, default=2 * 1024 * 1024 * 1024 + ) + parser.add_argument( + "--request-timeout-ms", type=sidecar.positive_int, default=60_000 + ) + parser.add_argument("--max-inflight", type=sidecar.positive_int, default=8) + parser.add_argument("--max-cuda-inflight", type=sidecar.positive_int, default=1) + parser.add_argument("--backlog", type=sidecar.positive_int, default=64) + parser.add_argument("--allow-tf32", action="store_true") + args = parser.parse_args() + + if not args.model_contract: + parser.error("--model-contract must not be empty") + limits = protocol.Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + resolver = NpyManifestResolver( + args.model_contract, + args.manifest, + args.descriptor_manifest, + args.cache_bytes, + ) + metrics = sidecar.JsonMetrics() + engine = sidecar.TorchTileMaxsimEngine( + args.device, + args.max_device_bytes, + args.allow_tf32, + args.max_cuda_inflight, + ) + service = sidecar.TileMaxsimService( + limits, resolver, engine, args.request_timeout_ms, metrics + ) + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + sidecar.serve( + args.socket, + args.socket_mode, + args.backlog, + args.max_inflight, + service, + stop, + ) + + +if __name__ == "__main__": + main() diff --git a/services/benchmark_postgres_single_vector.py b/services/benchmark_postgres_single_vector.py new file mode 100644 index 00000000..3041190d --- /dev/null +++ b/services/benchmark_postgres_single_vector.py @@ -0,0 +1,336 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Benchmark PostgreSQL/pgvector HNSW with original text-embedding vectors.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import shlex +import statistics +import subprocess +import time +from pathlib import Path + +import numpy as np + +IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*(?:\.[A-Za-z_][A-Za-z0-9_$]*)?$") + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + return ordered[max(0, math.ceil(len(ordered) * fraction) - 1)] + + +def summary(samples: list[float]) -> dict[str, float | int]: + return { + "count": len(samples), + "mean": statistics.fmean(samples), + "p50": percentile(samples, 0.50), + "p95": percentile(samples, 0.95), + "p99": percentile(samples, 0.99), + "max": max(samples), + } + + +def vector_text(vector: np.ndarray) -> str: + return "[" + ",".join(format(float(value), ".8g") for value in vector) + "]" + + +def run_psql(command: list[str], sql: str) -> str: + result = subprocess.run( + [*command, "-X", "-q", "-A", "-t", "-v", "ON_ERROR_STOP=1"], + input=sql, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + if result.returncode != 0: + diagnostic = result.stderr.strip() or result.stdout.strip() + raise RuntimeError(f"psql failed: {diagnostic}") + return result.stdout + + +def prepare_table( + command: list[str], table: str, vectors: np.ndarray, m: int, ef_construction: int +) -> float: + dimension = int(vectors.shape[1]) + process = subprocess.Popen( + [*command, "-X", "-q", "-v", "ON_ERROR_STOP=1"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + assert process.stdin is not None + started = time.perf_counter() + process.stdin.write( + f"DROP TABLE IF EXISTS {table};\n" + f"CREATE TABLE {table} (id integer PRIMARY KEY, embedding vector({dimension}) NOT NULL);\n" + f"COPY {table} (id, embedding) FROM STDIN;\n" + ) + for index, vector in enumerate(vectors): + process.stdin.write(f"{index}\t{vector_text(vector)}\n") + process.stdin.write( + "\\.\n" + f"CREATE INDEX ON {table} USING hnsw (embedding vector_ip_ops) " + f"WITH (m={m}, ef_construction={ef_construction});\n" + f"ANALYZE {table};\n" + ) + process.stdin.close() + stdout = process.stdout.read() if process.stdout is not None else "" + stderr = process.stderr.read() if process.stderr is not None else "" + return_code = process.wait() + if return_code != 0: + raise RuntimeError(f"psql prepare failed: {stderr.strip() or stdout.strip()}") + return (time.perf_counter() - started) * 1000 + + +def execute_queries( + command: list[str], table: str, queries: np.ndarray, ef_search: int, top_k: int +) -> list[tuple[int, float, list[int]]]: + dimension = int(queries.shape[1]) + query_values = ",\n".join( + f"({index}, '{vector_text(vector)}'::vector({dimension}))" + for index, vector in enumerate(queries) + ) + sql = f""" +SET hnsw.ef_search = {ef_search}; +CREATE TEMP TABLE single_vector_queries ( + query_index integer PRIMARY KEY, + embedding vector({dimension}) NOT NULL +) ON COMMIT PRESERVE ROWS; +INSERT INTO single_vector_queries VALUES +{query_values}; +CREATE TEMP TABLE single_vector_results ( + query_index integer PRIMARY KEY, + latency_ms double precision NOT NULL, + ids integer[] NOT NULL +) ON COMMIT PRESERVE ROWS; +DO $bench$ +DECLARE + q record; + started timestamptz; + result_ids integer[]; +BEGIN + -- Warm PostgreSQL buffers and the HNSW path before recording latency. + FOR q IN SELECT * FROM single_vector_queries ORDER BY query_index LOOP + EXECUTE 'SELECT array_agg(id ORDER BY distance, id) FROM ( + SELECT id, embedding <#> $1 AS distance + FROM {table} + ORDER BY embedding <#> $1 + LIMIT {top_k} + ) ranked' + INTO result_ids USING q.embedding; + END LOOP; + FOR q IN SELECT * FROM single_vector_queries ORDER BY query_index LOOP + started := clock_timestamp(); + EXECUTE 'SELECT array_agg(id ORDER BY distance, id) FROM ( + SELECT id, embedding <#> $1 AS distance + FROM {table} + ORDER BY embedding <#> $1 + LIMIT {top_k} + ) ranked' + INTO result_ids USING q.embedding; + INSERT INTO single_vector_results VALUES ( + q.query_index, + extract(epoch FROM clock_timestamp() - started) * 1000.0, + result_ids + ); + END LOOP; +END +$bench$; +SELECT query_index, latency_ms, array_to_string(ids, ',') + FROM single_vector_results + ORDER BY query_index; +""" + rows = [] + for line in run_psql(command, sql).splitlines(): + if not line.strip(): + continue + query_index, latency_ms, ids = line.split("|") + rows.append( + ( + int(query_index), + float(latency_ms), + [int(value) for value in ids.split(",")], + ) + ) + return rows + + +def recall(expected: np.ndarray, actual: list[int], top_k: int) -> float: + return len(set(expected[:top_k].tolist()).intersection(actual[:top_k])) / top_k + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--page-vectors", type=Path) + parser.add_argument("--query-vectors", type=Path) + parser.add_argument("--gold", type=Path) + parser.add_argument( + "--corpus-json", + type=Path, + help="RAG/GBrain corpus JSON whose _source.q_1024_vec is the stored text embedding", + ) + parser.add_argument( + "--query-dataset", + type=Path, + help="query JSON whose q_1024_vec is the original text query embedding", + ) + parser.add_argument("--psql-command", default="psql") + parser.add_argument("--table", default="public.tilemaxsim_single_vector_bench") + parser.add_argument("--prepare", action="store_true") + parser.add_argument("--m", type=int, default=16) + parser.add_argument("--ef-construction", type=int, default=64) + parser.add_argument("--ef-search", type=int, default=40) + parser.add_argument("--top-k", type=int, default=20) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + if not IDENTIFIER.fullmatch(args.table): + parser.error("--table must be an unquoted SQL identifier") + if min(args.m, args.ef_construction, args.ef_search, args.top_k) <= 0: + parser.error("HNSW parameters and top-k must be positive") + command = shlex.split(args.psql_command) + if not command: + parser.error("--psql-command must not be empty") + + corpus_chunk_ids: list[str] | None = None + corpus_doc_names: list[str] | None = None + query_records: list[dict[str, object]] | None = None + if args.corpus_json or args.query_dataset: + if not args.corpus_json or not args.query_dataset: + parser.error("--corpus-json and --query-dataset must be used together") + corpus_records = json.loads(args.corpus_json.read_text(encoding="utf-8")) + query_records = json.loads(args.query_dataset.read_text(encoding="utf-8")) + if not isinstance(corpus_records, list) or not isinstance(query_records, list): + raise ValueError("corpus and query JSON inputs must be arrays") + page_vectors = np.asarray( + [record["_source"]["q_1024_vec"] for record in corpus_records], + dtype=np.float32, + ) + query_vectors = np.asarray( + [record["q_1024_vec"] for record in query_records], + dtype=np.float32, + ) + corpus_chunk_ids = [str(record["_id"]) for record in corpus_records] + corpus_doc_names = [str(record["_source"]["docnm_kwd"]) for record in corpus_records] + else: + if not args.page_vectors or not args.query_vectors or not args.gold: + parser.error( + "use --corpus-json/--query-dataset or all of " + "--page-vectors/--query-vectors/--gold" + ) + page_vectors = np.load(args.page_vectors, mmap_mode="r", allow_pickle=False) + query_vectors = np.load(args.query_vectors, mmap_mode="r", allow_pickle=False) + if page_vectors.ndim != 2 or query_vectors.ndim != 2: + raise ValueError("page and query vectors must be rank-2 arrays") + if page_vectors.shape[1] != query_vectors.shape[1]: + raise ValueError("page and query dimensions differ") + prepare_ms = ( + prepare_table( + command, args.table, page_vectors, args.m, args.ef_construction + ) + if args.prepare + else None + ) + query_rows = execute_queries( + command, args.table, query_vectors, args.ef_search, args.top_k + ) + cases = [] + if query_records is not None and corpus_chunk_ids is not None and corpus_doc_names is not None: + for query_index, latency_ms, ids in query_rows: + query_record = query_records[query_index] + gold_chunk = str(query_record["gold_chunk_id"]) + gold_doc = str(query_record["gold_doc_name"]) + ranked_chunks = [corpus_chunk_ids[index] for index in ids] + ranked_docs = [corpus_doc_names[index] for index in ids] + cases.append( + { + "query_index": query_index, + "query_id": query_record.get("query_id"), + "latency_ms": latency_ms, + "chunk_hit_at_1": gold_chunk in ranked_chunks[:1], + "chunk_hit_at_10": gold_chunk in ranked_chunks[:10], + "chunk_hit_at_20": gold_chunk in ranked_chunks[:20], + "doc_hit_at_1": gold_doc in ranked_docs[:1], + "doc_hit_at_10": gold_doc in ranked_docs[:10], + "doc_hit_at_20": gold_doc in ranked_docs[:20], + "top_20_chunk_ids": ranked_chunks[:20], + "top_20_doc_names": ranked_docs[:20], + } + ) + else: + assert args.gold is not None + gold = np.load(args.gold, allow_pickle=False) + for query_index, latency_ms, ids in query_rows: + expected = gold[f"q{query_index}_idx"].astype(np.int64, copy=False) + cases.append( + { + "query_index": query_index, + "latency_ms": latency_ms, + "recall_at_10": recall(expected, ids, 10), + "recall_at_20": recall(expected, ids, 20), + "top_20": ids[:20], + } + ) + latencies = [item["latency_ms"] for item in cases] + report = { + "benchmark": "postgres_pgvector_hnsw_text_embedding_v2", + "corpus_vectors": int(page_vectors.shape[0]), + "query_vectors": int(query_vectors.shape[0]), + "dimension": int(page_vectors.shape[1]), + "configuration": { + "table": args.table, + "m": args.m, + "ef_construction": args.ef_construction, + "ef_search": args.ef_search, + "top_k": args.top_k, + }, + "prepare_ms": prepare_ms, + "latency_ms": summary(latencies), + "quality": ( + { + key: statistics.fmean(float(item[key]) for item in cases) + for key in ( + "chunk_hit_at_1", + "chunk_hit_at_10", + "chunk_hit_at_20", + "doc_hit_at_1", + "doc_hit_at_10", + "doc_hit_at_20", + ) + } + if query_records is not None + else { + "mean_recall_at_10": statistics.fmean( + item["recall_at_10"] for item in cases + ), + "mean_recall_at_20": statistics.fmean( + item["recall_at_20"] for item in cases + ), + } + ), + "cases": cases, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/benchmark_tilemaxsim_ablation.py b/services/benchmark_tilemaxsim_ablation.py new file mode 100644 index 00000000..d7ab0ab5 --- /dev/null +++ b/services/benchmark_tilemaxsim_ablation.py @@ -0,0 +1,765 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Run storage, cache, H2D, and native-daemon TileMaxSim ablations.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import random +import select +import socket +import statistics +import subprocess +import sys +import tempfile +import time +from collections import OrderedDict +from pathlib import Path + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from devtools.test_tilemaxsim_reference_sidecar import decode_response +from services.tilemaxsim_cuda_sidecar import ContentAddressedResolver, PayloadCache +from services.tilemaxsim_gpu_cache import ( + FixedBlockAllocator, + FreeExtentAllocator, + GpuArenaSpec, + GpuResourcePool, + GpuTensorCache, + GpuTensorLoad, +) + + +class _LegacyBuddyAllocator: + """Former power-of-two allocator retained only as an ablation baseline.""" + + def __init__(self, capacity: int, block_bytes: int = 256 * 1024) -> None: + self.block_bytes = block_bytes + self.block_count = capacity // block_bytes + self.capacity = self.block_count * block_bytes + self._free: dict[int, set[tuple[int, int, int]]] = {} + self._allocated: dict[int, tuple[int, int, int]] = {} + start = 0 + remaining = self.block_count + while remaining: + order = remaining.bit_length() - 1 + size = 1 << order + self._free.setdefault(order, set()).add((start, start, order)) + start += size + remaining -= size + + @property + def free_bytes(self) -> int: + return sum( + len(items) * (1 << order) * self.block_bytes + for order, items in self._free.items() + ) + + def allocation_bytes(self, payload_bytes: int) -> int: + raw = math.ceil(payload_bytes / self.block_bytes) + return (1 << (raw - 1).bit_length()) * self.block_bytes + + def allocate(self, payload_bytes: int) -> tuple[int, ...] | None: + required = self.allocation_bytes(payload_bytes) // self.block_bytes + order = required.bit_length() - 1 + available_order = next( + ( + candidate + for candidate in range(order, self.block_count.bit_length()) + if self._free.get(candidate) + ), + None, + ) + if available_order is None: + return None + start, root_start, root_order = self._free[available_order].pop() + while available_order > order: + available_order -= 1 + buddy = start + (1 << available_order) + self._free.setdefault(available_order, set()).add( + (buddy, root_start, root_order) + ) + self._allocated[start] = (order, root_start, root_order) + return tuple(range(start, start + required)) + + def release(self, blocks: tuple[int, ...]) -> None: + start = blocks[0] + order, root_start, root_order = self._allocated.pop(start) + while order < root_order: + buddy = root_start + ((start - root_start) ^ (1 << order)) + item = (buddy, root_start, root_order) + free = self._free.setdefault(order, set()) + if item not in free: + break + free.remove(item) + start = min(start, buddy) + order += 1 + self._free.setdefault(order, set()).add((start, root_start, root_order)) + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + return ordered[min(len(ordered) - 1, math.ceil(len(ordered) * fraction) - 1)] + + +def load_records(path: Path) -> list[dict[str, object]]: + with path.open(encoding="utf-8") as stream: + return [json.loads(line) for line in stream if line.strip()] + + +def request(record: dict[str, object], contract: str) -> protocol.ExternalTensorRequest: + dtype = ( + protocol.DTYPE_F16 + if record["tensor_dtype"] == "float16" + else protocol.DTYPE_F32 + ) + return protocol.ExternalTensorRequest( + contract, + str(record["tensor_ref"]), + int(record["tensor_rows"]), + int(record["tensor_dim"]), + dtype, + str(record["tensor_checksum"]), + ) + + +def evict_paths(paths: list[Path]) -> None: + if not hasattr(os, "posix_fadvise"): + return + for path in paths: + descriptor = os.open(path, os.O_RDONLY | os.O_CLOEXEC) + try: + os.posix_fadvise(descriptor, 0, 0, os.POSIX_FADV_DONTNEED) + finally: + os.close(descriptor) + + +def storage_ablation( + selected: list[dict[str, object]], + contract: str, + legacy_root: Path, + shard_root: Path, +) -> dict[str, object]: + requests = [request(record, contract) for record in selected] + legacy_paths = [ + legacy_root + / str(record["tensor_checksum"])[7:9] + / f"{str(record['tensor_checksum'])[7:]}.bin" + for record in selected + ] + shard_paths = sorted((shard_root / "shards").glob("*.vts")) + + evict_paths(legacy_paths) + resolver = ContentAddressedResolver({contract: legacy_root}, 0) + started = time.perf_counter() + try: + sequential = [resolver.resolve(item) for item in requests] + finally: + resolver.close() + sequential_ms = (time.perf_counter() - started) * 1000 + + evict_paths(legacy_paths) + resolver = ContentAddressedResolver({contract: legacy_root}, 0) + started = time.perf_counter() + try: + legacy_batch = resolver.resolve_many(requests) + finally: + resolver.close() + legacy_batch_ms = (time.perf_counter() - started) * 1000 + + evict_paths(shard_paths) + resolver = ContentAddressedResolver({contract: shard_root}, 0) + started = time.perf_counter() + try: + shard_batch = resolver.resolve_many(requests) + shard_status = resolver.status() + finally: + resolver.close() + shard_batch_ms = (time.perf_counter() - started) * 1000 + expected = [item.payload for item in sequential] + if expected != [item.payload for item in legacy_batch] or expected != [ + item.payload for item in shard_batch + ]: + raise RuntimeError("storage ablation payloads disagree") + return { + "candidates": len(selected), + "logical_bytes": sum(len(item) for item in expected), + "legacy_sequential_ms": sequential_ms, + "legacy_batch_ms": legacy_batch_ms, + "shard_batch_ms": shard_batch_ms, + "shard_batch_read_calls": shard_status["batch_read_calls"], + "shard_batch_read_bytes": shard_status["batch_read_bytes"], + } + + +def h2d_ablation( + selected: list[dict[str, object]], contract: str, shard_root: Path, device: int +) -> dict[str, object]: + requests = [request(record, contract) for record in selected] + resolver = ContentAddressedResolver({contract: shard_root}, 0) + try: + payloads = resolver.resolve_many(requests) + keys = [resolver.key(item) for item in requests] + finally: + resolver.close() + total_bytes = 768 * 1024**2 + workspace_bytes = 256 * 1024**2 + + pool = GpuResourcePool( + [GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + started = time.perf_counter() + handles = [] + for key, item, resolved in zip(keys, requests, payloads, strict=True): + handle, _ = cache.acquire( + key, + item.rows, + item.dimension, + item.dtype, + lambda payload=resolved.payload: payload, + ) + handles.append(handle) + torch.cuda.synchronize(device) + sequential_ms = (time.perf_counter() - started) * 1000 + sequential_status = pool.status()[0] + for handle in handles: + cache.release(handle) + finally: + pool.close() + + pool = GpuResourcePool( + [GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + loads = [ + GpuTensorLoad(key, item.rows, item.dimension, item.dtype, resolved.payload) + for key, item, resolved in zip(keys, requests, payloads, strict=True) + ] + started = time.perf_counter() + batch = cache.acquire_many(loads) + torch.cuda.synchronize(device) + batch_ms = (time.perf_counter() - started) * 1000 + batch_status = pool.status()[0] + if batch.bypassed or batch.deferred: + raise RuntimeError("H2D ablation cache is undersized") + for handle in batch.handles: + assert handle is not None + cache.release(handle) + finally: + pool.close() + return { + "candidates": len(selected), + "logical_bytes": sum(len(item.payload) for item in payloads), + "per_tensor_h2d_ms": sequential_ms, + "batch_h2d_ms": batch_ms, + "per_tensor_copy_batches": sequential_status["h2d_batches"], + "batch_copy_batches": batch_status["h2d_batches"], + "batch_copy_calls": batch_status["h2d_copy_calls"], + } + + +def allocator_ablation(records: list[dict[str, object]]) -> dict[str, object]: + sizes = [int(record["canonical_bytes"]) for record in records] + rng = random.Random(991) + capacity = 256 * 1024**2 + events: list[tuple[str, int, int]] = [] + abstract_live: list[int] = [] + next_identifier = 0 + for _ in range(20_000): + if abstract_live and rng.random() < 0.48: + index = rng.randrange(len(abstract_live)) + identifier = abstract_live.pop(index) + events.append(("release", identifier, 0)) + else: + size = rng.choice(sizes) + identifier = next_identifier + next_identifier += 1 + abstract_live.append(identifier) + events.append(("allocate", identifier, size)) + + def run_trace( + allocator: FreeExtentAllocator | FixedBlockAllocator | _LegacyBuddyAllocator, + ) -> dict[str, object]: + live: dict[int, tuple[int, ...] | tuple[int, int]] = {} + failures = 0 + fragmentation_failures = 0 + requested_bytes = 0 + allocated_bytes = 0 + started = time.perf_counter() + for operation, identifier, size in events: + if operation == "release": + allocation = live.pop(identifier, None) + if allocation is None: + continue + if isinstance(allocator, FreeExtentAllocator): + allocator.release(*allocation) + else: + allocator.release(allocation) + continue + required = allocator.allocation_bytes(size) + allocation = allocator.allocate(size) + if allocation is None: + failures += 1 + fragmentation_failures += int(allocator.free_bytes >= required) + else: + live[identifier] = allocation + requested_bytes += size + allocated_bytes += ( + allocation[1] + if isinstance(allocator, FreeExtentAllocator) + else len(allocation) * allocator.block_bytes + ) + return { + "allocation_failures": failures, + "fragmentation_failures": fragmentation_failures, + "internal_waste_ratio": ( + (allocated_bytes - requested_bytes) / allocated_bytes + ), + "trace_ms": (time.perf_counter() - started) * 1000, + } + + extent = FreeExtentAllocator(capacity) + legacy = _LegacyBuddyAllocator(capacity) + page_runs = FixedBlockAllocator(capacity) + legacy_full_bytes = sum(legacy.allocation_bytes(size) for size in sizes) + page_run_full_bytes = sum(page_runs.allocation_bytes(size) for size in sizes) + return { + "operations": 20_000, + "capacity_bytes": capacity, + "exact_byte_extents": run_trace(extent), + "legacy_power_of_two_buddy": { + "block_bytes": legacy.block_bytes, + "full_corpus_allocated_bytes": legacy_full_bytes, + **run_trace(legacy), + }, + "segregated_page_runs": { + "block_bytes": page_runs.block_bytes, + "full_corpus_allocated_bytes": page_run_full_bytes, + "full_corpus_space_saved_bytes": legacy_full_bytes - page_run_full_bytes, + **run_trace(page_runs), + }, + } + + +class LegacyLru: + def __init__(self, maximum_bytes: int) -> None: + self.maximum_bytes = maximum_bytes + self.bytes = 0 + self.entries: OrderedDict[str, bytes] = OrderedDict() + + def access(self, key: str, size: int) -> bool: + if key in self.entries: + self.entries.move_to_end(key) + return True + payload = bytes(size) + if size <= self.maximum_bytes: + self.entries[key] = payload + self.bytes += size + while self.bytes > self.maximum_bytes: + _, evicted = self.entries.popitem(last=False) + self.bytes -= len(evicted) + return False + + +def policy_ablation(records: list[dict[str, object]]) -> dict[str, object]: + rng = random.Random(77) + universe = records[:5000] + sizes = { + str(record["tensor_ref"]): int(record["canonical_bytes"]) for record in universe + } + hot = list(sizes)[:100] + cold = list(sizes)[100:] + # Warm every hot object before injecting scans so TinyLFU admission is + # measured directly rather than starting from an empty cache. + trace = hot * 5 + for cycle in range(30): + trace.extend(rng.choices(hot, weights=range(len(hot), 0, -1), k=1000)) + start = cycle * 100 % len(cold) + trace.extend(cold[start : start + 300]) + # The budget holds exactly the hot set but not the scan tail. + budget = sum(sizes[key] for key in hot) + lru = LegacyLru(budget) + gdsf = PayloadCache(budget) + lru_hits = 0 + gdsf_hits = 0 + for key in trace: + size = sizes[key] + lru_hits += int(lru.access(key, size)) + cached = gdsf.get((key,)) + gdsf_hits += int(cached is not None) + if cached is None: + gdsf.put((key,), bytes(size)) + return { + "accesses": len(trace), + "budget_bytes": budget, + "lru_hit_ratio": lru_hits / len(trace), + "tinylfu_gdsf_hit_ratio": gdsf_hits / len(trace), + "tinylfu_gdsf_status": gdsf.status(), + } + + +def encode_frame( + records: list[dict[str, object]], contract: str, query: np.ndarray, request_id: int +) -> bytes: + dtype = protocol.DTYPE_F16 if query.dtype == np.dtype(" tuple[float, list[tuple[int, float]]]: + started = time.perf_counter() + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_bytes = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_bytes) + elapsed = (time.perf_counter() - started) * 1000 + _, status, parsed = decode_response(response) + if status != 0 or not isinstance(parsed, list): + raise RuntimeError(f"daemon request failed: {parsed}") + if len(parsed) != len(protocol.parse_request_frame(frame).candidates): + raise RuntimeError("daemon returned an incomplete result") + return elapsed, parsed + + +def daemon_ablation( + records: list[dict[str, object]], + contract: str, + shard_root: Path, + rust_binary: Path, + device: int, + repeats: int, + gpu_block_kib: int = 32, +) -> dict[str, object]: + rng = np.random.default_rng(31) + query = rng.standard_normal((44, int(records[0]["tensor_dim"]))).astype(" dict[str, object]: + shard_paths = sorted((shard_root / "shards").glob("*.vts")) + results = {} + with tempfile.TemporaryDirectory() as directory: + commands = { + "python_triton": [ + sys.executable, + "-m", + "services.tilemaxsim_cuda_sidecar", + "--socket", + f"{directory}/python-resident.sock", + "--gpu-memory-gb", + f"{device}=20", + "--gpu-workspace-gb", + "2", + "--gpu-block-kib", + str(gpu_block_kib), + "--host-cache-gb", + "0.1", + "--contract-root", + f"{contract}={shard_root}", + "--gpu-cache-mode", + "resident", + "--resident-manifest", + f"{contract}={descriptor_manifest}", + "--prewarm-batch-size", + "256", + "--request-timeout-ms", + "20000", + ], + "rust_cuda": [ + os.fspath(rust_binary), + "--socket", + f"{directory}/rust-resident.sock", + "--gpu-memory-gb", + f"{device}=20", + "--gpu-workspace-gb", + "2", + "--gpu-block-kib", + str(gpu_block_kib), + "--host-cache-gb", + "0.1", + "--contract-root", + f"{contract}={shard_root}", + "--gpu-cache-mode", + "resident", + "--resident-manifest", + f"{contract}={descriptor_manifest}", + "--prewarm-batch-size", + "256", + ], + } + for name, command in commands.items(): + os.sync() + evict_paths(shard_paths) + started = time.perf_counter() + process = subprocess.Popen( + command, + cwd=Path(__file__).resolve().parents[1], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + lines = [] + prewarm_event = None + ready_event = None + deadline = time.monotonic() + 300 + try: + assert process.stdout is not None + while time.monotonic() < deadline: + readable, _, _ = select.select([process.stdout], [], [], 1.0) + if readable: + line = process.stdout.readline() + if line: + lines.append(line) + if line.startswith("{"): + event = json.loads(line) + if event.get("event") in ( + "tilemaxsim_prewarm_complete", + "tilemaxsim_rust_prewarm_complete", + ): + prewarm_event = event + if event.get("event") in ( + "tilemaxsim_ready", + "tilemaxsim_rust_ready", + ): + ready_event = event + break + if process.poll() is not None: + break + if ready_event is None: + remainder, _ = process.communicate(timeout=5) + raise RuntimeError( + f"{name} resident prewarm failed: {''.join(lines)}{remainder}" + ) + results[name] = { + "process_to_ready_ms": (time.perf_counter() - started) * 1000, + "prewarm_reported_ms": prewarm_event.get("elapsed_ms") + if prewarm_event + else None, + "cache": ready_event.get("gpu_cache", ready_event.get("cache")), + } + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=15) + process.communicate(timeout=5) + return results + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--legacy-root", required=True, type=Path) + parser.add_argument("--shard-root", required=True, type=Path) + parser.add_argument("--rust-binary", required=True, type=Path) + parser.add_argument("--contract", default="benchmark@1") + parser.add_argument("--device", required=True, type=int) + parser.add_argument("--sample-size", type=int, default=100) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--gpu-block-kib", type=int, default=32) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--full-prewarm", action="store_true") + args = parser.parse_args() + records = load_records(args.descriptor_manifest) + selected = random.Random(20260714).sample(records, args.sample_size) + report = { + "corpus": { + "records": len(records), + "logical_bytes": sum(int(record["canonical_bytes"]) for record in records), + "sample_size": len(selected), + }, + "storage": storage_ablation( + selected, args.contract, args.legacy_root, args.shard_root + ), + "h2d": h2d_ablation(selected, args.contract, args.shard_root, args.device), + "allocator": allocator_ablation(records), + "policy": policy_ablation(records), + "daemon": daemon_ablation( + selected, + args.contract, + args.shard_root, + args.rust_binary, + args.device, + args.repeats, + args.gpu_block_kib, + ), + } + if args.full_prewarm: + report["full_resident_prewarm"] = prewarm_ablation( + args.descriptor_manifest, + args.contract, + args.shard_root, + args.rust_binary, + args.device, + args.gpu_block_kib, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as stream: + json.dump(report, stream, indent=2, sort_keys=True) + stream.write("\n") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/build_tilemaxsim_tensor_cache.py b/services/build_tilemaxsim_tensor_cache.py index a48a1728..e6adf7ad 100644 --- a/services/build_tilemaxsim_tensor_cache.py +++ b/services/build_tilemaxsim_tensor_cache.py @@ -22,12 +22,15 @@ import os import sys import tempfile +from collections import deque from concurrent.futures import ThreadPoolExecutor +from decimal import Decimal, InvalidOperation from pathlib import Path import numpy as np from services.tilemaxsim_cuda_sidecar import positive_int +from services.tilemaxsim_shard import DEFAULT_SHARD_BYTES, ImmutableShardWriter def canonical_tensor(path: Path, expected_rows: int, expected_dim: int) -> np.ndarray: @@ -96,13 +99,10 @@ def write_payload(path: Path, payload: memoryview, digest: str, fsync: bool) -> pass -def process_record( +def prepare_record( record: dict[str, object], source_root: Path, - cache_root: Path, - fsync: bool, - dry_run: bool, -) -> dict[str, object]: +) -> tuple[dict[str, object], memoryview]: page_key = record.get("page_key") relative = record.get("embedding_file") rows = record.get("n_tokens") @@ -125,9 +125,6 @@ def process_record( tensor = canonical_tensor(source, rows, dimension) payload = memoryview(tensor).cast("B") digest = hashlib.sha256(payload).hexdigest() - destination = cache_root / digest[:2] / f"{digest}.bin" - if not dry_run: - write_payload(destination, payload, digest, fsync) dtype_name = "float16" if tensor.dtype == np.dtype("float16") else "float32" return { "page_key": page_key, @@ -137,7 +134,57 @@ def process_record( "tensor_dtype": dtype_name, "tensor_checksum": f"sha256:{digest}", "canonical_bytes": len(payload), - } + }, payload + + +def process_record( + record: dict[str, object], + source_root: Path, + cache_root: Path, + fsync: bool, + dry_run: bool, +) -> dict[str, object]: + """Publish one legacy per-tensor file. + + Kept for backwards compatibility and migration tests. New cache builds use + immutable shards by default. + """ + + descriptor, payload = prepare_record(record, source_root) + digest = str(descriptor["tensor_checksum"]).removeprefix("sha256:") + destination = cache_root / digest[:2] / f"{digest}.bin" + if not dry_run: + write_payload(destination, payload, digest, fsync) + return descriptor + + +def shard_size_gb(value: str) -> int: + try: + parsed = int(Decimal(value) * 1024**3) + except (InvalidOperation, ValueError) as error: + raise argparse.ArgumentTypeError("shard size must be a positive number of GB") from error + if parsed <= 0: + raise argparse.ArgumentTypeError("shard size must be a positive number of GB") + return parsed + + +def bounded_parallel_map(executor, function, items, maximum_pending: int): + """Preserve input order without retaining a corpus-sized future backlog.""" + + iterator = iter(items) + pending = deque() + for _ in range(maximum_pending): + try: + pending.append(executor.submit(function, next(iterator))) + except StopIteration: + break + while pending: + future = pending.popleft() + yield future.result() + try: + pending.append(executor.submit(function, next(iterator))) + except StopIteration: + pass def main() -> None: @@ -146,6 +193,18 @@ def main() -> None: parser.add_argument("--cache-root", required=True, type=Path) parser.add_argument("--descriptor-manifest", required=True, type=Path) parser.add_argument("--workers", type=positive_int, default=4) + parser.add_argument( + "--storage-format", + choices=("shards", "files"), + default="shards", + help="publish immutable shards (default) or legacy per-tensor files", + ) + parser.add_argument( + "--shard-size-gb", + type=shard_size_gb, + default=DEFAULT_SHARD_BYTES, + help="target immutable shard size in GB", + ) parser.add_argument("--no-fsync", action="store_true") parser.add_argument("--dry-run", action="store_true") args = parser.parse_args() @@ -171,28 +230,63 @@ def main() -> None: cache_root.mkdir(parents=True, exist_ok=True) descriptors = [] - with ThreadPoolExecutor(max_workers=args.workers) as workers: - results = workers.map( - lambda record: process_record( - record, - source_root, - cache_root, - not args.no_fsync, - args.dry_run, - ), - records, + shard_writer = None + if args.storage_format == "shards" and not args.dry_run: + shard_writer = ImmutableShardWriter( + cache_root, + target_bytes=args.shard_size_gb, + fsync=not args.no_fsync, ) - for completed, item in enumerate(results, 1): - descriptors.append(item) - if completed % 1000 == 0 or completed == len(records): - print( - json.dumps( - {"event": "tensor_cache_progress", "completed": completed}, - separators=(",", ":"), - ), - file=sys.stderr, - flush=True, + try: + with ThreadPoolExecutor(max_workers=args.workers) as workers: + if args.storage_format == "files": + results = ( + (item, None) + for item in bounded_parallel_map( + workers, + lambda record: process_record( + record, + source_root, + cache_root, + not args.no_fsync, + args.dry_run, + ), + records, + args.workers * 2, + ) ) + else: + results = bounded_parallel_map( + workers, + lambda record: prepare_record(record, source_root), + records, + args.workers * 2, + ) + for completed, (item, payload) in enumerate(results, 1): + descriptors.append(item) + if shard_writer is not None: + assert payload is not None + digest = str(item["tensor_checksum"]).removeprefix("sha256:") + shard_writer.add( + digest, + payload, + int(item["tensor_rows"]), + int(item["tensor_dim"]), + str(item["tensor_dtype"]), + ) + if completed % 1000 == 0 or completed == len(records): + print( + json.dumps( + {"event": "tensor_cache_progress", "completed": completed}, + separators=(",", ":"), + ), + file=sys.stderr, + flush=True, + ) + shard_index = shard_writer.finish() if shard_writer is not None else None + finally: + if shard_writer is not None: + shard_writer.close() output_parent = args.descriptor_manifest.parent output_parent.mkdir(parents=True, exist_ok=True) @@ -224,6 +318,8 @@ def main() -> None: "canonical_bytes": total_bytes, "cache_root": os.fspath(cache_root), "descriptor_manifest": os.fspath(args.descriptor_manifest), + "storage_format": args.storage_format, + "shard_index": os.fspath(shard_index) if shard_index else None, "dry_run": args.dry_run, }, sort_keys=True, diff --git a/services/test_tilemaxsim_cuda_sidecar.py b/services/test_tilemaxsim_cuda_sidecar.py index 093245e3..e015174d 100644 --- a/services/test_tilemaxsim_cuda_sidecar.py +++ b/services/test_tilemaxsim_cuda_sidecar.py @@ -15,10 +15,13 @@ from __future__ import annotations import hashlib +import json import os import socket import stat import struct +import subprocess +import sys import tempfile import threading import time @@ -57,6 +60,26 @@ def write_content_addressed(root: Path, payload: bytes) -> tuple[str, str]: class CudaSidecarTest(unittest.TestCase): + def test_cli_without_explicit_gpu_memory_keeps_tilemaxsim_disabled(self) -> None: + with tempfile.TemporaryDirectory() as directory: + socket_path = Path(directory) / "disabled.sock" + completed = subprocess.run( + [ + sys.executable, + "-m", + "services.tilemaxsim_cuda_sidecar", + "--socket", + os.fspath(socket_path), + ], + cwd=Path(__file__).resolve().parents[1], + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(completed.returncode, 2) + self.assertIn("TileMaxSim is disabled", completed.stderr) + self.assertFalse(socket_path.exists()) + def test_cache_builder_publishes_resolver_compatible_payload(self) -> None: with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -140,6 +163,14 @@ def test_content_addressed_resolver_validates_and_caches(self) -> None: finally: resolver.close() + def test_host_payload_cache_evicts_to_its_byte_budget(self) -> None: + cache = cuda_sidecar.PayloadCache(6) + cache.put(("first",), b"1234") + cache.put(("second",), b"5678") + self.assertIsNone(cache.get(("first",))) + self.assertEqual(cache.get(("second",)), b"5678") + self.assertEqual(cache.current_bytes, 4) + def test_content_addressed_resolver_rejects_symlink(self) -> None: with tempfile.TemporaryDirectory() as directory: root = Path(directory) / "root" @@ -246,6 +277,252 @@ def test_cuda_f16_matches_cpu_protocol_oracle(self) -> None: for (_, actual), (_, expected) in zip(results, oracle, strict=True): self.assertAlmostEqual(actual, expected, places=5) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_cli_gb_allocation_serves_tilemaxsim(self) -> None: + with tempfile.TemporaryDirectory() as directory: + directory_path = Path(directory) + root = directory_path / "tensors" + root.mkdir() + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, _ = write_content_addressed(root, payload) + frame, _ = external_request_frame( + 70, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, [[1.0, 0.0], [0.0, 1.0]])], + ) + device = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + socket_path = directory_path / "resident.sock" + process = subprocess.Popen( + [ + sys.executable, + "-m", + "services.tilemaxsim_cuda_sidecar", + "--socket", + os.fspath(socket_path), + "--gpu-memory-gb", + f"{device}=0.05", + "--gpu-workspace-gb", + "0.02", + "--host-cache-gb", + "0.01", + "--contract-root", + f"model@1={root}", + "--once", + ], + cwd=Path(__file__).resolve().parents[1], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + for _ in range(1000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + self.assertIsNone(process.poll()) + self.assertTrue(socket_path.exists()) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_len = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_len) + output, _ = process.communicate(timeout=10) + self.assertEqual(process.returncode, 0, output) + self.assertEqual(decode_response(response)[1:], (0, [(9, 2.0)])) + self.assertIn('"event":"tilemaxsim_ready"', output) + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=5) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_v2_gpu_resident_hit_does_not_resolve_payload_again(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + digest = checksum.removeprefix("sha256:") + frame, _ = external_request_frame( + 71, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, [[1.0, 0.0], [0.0, 1.0]])], + ) + device = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + pool = cuda_sidecar.GpuResourcePool( + [cuda_sidecar.GpuArenaSpec(f"cuda:{device}", 32 * 1024 * 1024)], + 16 * 1024 * 1024, + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 0) + try: + cache = cuda_sidecar.GpuTensorCache(pool, allow_eviction=False) + metrics = CapturingMetrics() + stream_engine = cuda_sidecar.TorchTileMaxsimEngine( + f"cuda:{device}", 16 * 1024 * 1024, False, 1 + ) + resident_engine = cuda_sidecar.ResidentTorchTileMaxsimEngine( + pool, 16 * 1024 * 1024, False, 1 + ) + service = cuda_sidecar.TileMaxsimService( + protocol.Limits(), + resolver, + stream_engine, + 2000, + metrics, + cache, + resident_engine, + pin_gpu_entries=True, + ) + client, server = socket.socketpair() + try: + first = service.process_frame( + frame, server, time.monotonic() + 2, None + ) + (root / digest[:2] / f"{digest}.bin").unlink() + second = service.process_frame( + frame, server, time.monotonic() + 2, None + ) + finally: + client.close() + server.close() + self.assertEqual(decode_response(first)[1:], (0, [(9, 2.0)])) + self.assertEqual(decode_response(second)[1:], (0, [(9, 2.0)])) + requests = [ + event + for event in metrics.events + if event.get("event") == "tilemaxsim_request" + ] + self.assertEqual(requests[0]["gpu_cache_misses"], 1) + self.assertEqual(requests[1]["gpu_cache_hits"], 1) + finally: + resolver.close() + pool.close() + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_lru_gpu_cache_streams_request_larger_than_its_arena(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + rows, dimension = 480, 320 + first_tensor = np.zeros((rows, dimension), dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + manifest = root / "descriptors.jsonl" + manifest.write_text( + json.dumps( + { + "page_key": "page-1", + "tensor_ref": tensor_ref, + "tensor_rows": 2, + "tensor_dim": 2, + "tensor_dtype": "float16", + "tensor_checksum": checksum, + "canonical_bytes": len(payload), + } + ) + + "\n", + encoding="utf-8", + ) + device = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + pool = cuda_sidecar.GpuResourcePool( + [cuda_sidecar.GpuArenaSpec(f"cuda:{device}", 32 * 1024 * 1024)], + 16 * 1024 * 1024, + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 0) + try: + cache = cuda_sidecar.GpuTensorCache(pool, allow_eviction=False) + metrics = CapturingMetrics() + cuda_sidecar.prewarm_resident_cache( + [("model@1", manifest)], resolver, cache, metrics + ) + status = cache.status() + self.assertEqual(status["entries"], 1) + self.assertEqual(status["pinned_entries"], 1) + self.assertEqual( + metrics.events[-1]["event"], "tilemaxsim_prewarm_complete" + ) + finally: + resolver.close() + pool.close() + def test_v2_unix_socket_end_to_end(self) -> None: with tempfile.TemporaryDirectory() as directory: directory_path = Path(directory) diff --git a/services/test_tilemaxsim_gpu_cache.py b/services/test_tilemaxsim_gpu_cache.py new file mode 100644 index 00000000..6394d7ff --- /dev/null +++ b/services/test_tilemaxsim_gpu_cache.py @@ -0,0 +1,376 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2026 Hu Xinjing + +from __future__ import annotations + +import time +import unittest + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import ResidentTorchTileMaxsimEngine +from services.tilemaxsim_gpu_cache import ( + FixedBlockAllocator, + FreeExtentAllocator, + GpuArenaSpec, + GpuResourcePool, + GpuTensorCache, + GpuTensorLoad, + parse_gpu_memory_gb, + parse_memory_gb, +) + + +def available_device() -> str: + index = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + return f"cuda:{index}" + + +class GpuCacheUnitTest(unittest.TestCase): + def test_public_memory_configuration_uses_gb(self) -> None: + self.assertEqual(parse_memory_gb("20"), 20 * 1024**3) + self.assertEqual(parse_memory_gb("0.5"), 512 * 1024**2) + self.assertEqual( + parse_gpu_memory_gb("2=12"), + GpuArenaSpec("cuda:2", 12 * 1024**3), + ) + self.assertEqual( + parse_gpu_memory_gb("cuda:2=12.5"), + GpuArenaSpec("cuda:2", int(12.5 * 1024**3)), + ) + with self.assertRaisesRegex(ValueError, "GPU=GB"): + parse_gpu_memory_gb("cuda:0") + with self.assertRaisesRegex(ValueError, "byte suffixes"): + parse_gpu_memory_gb("0=20GiB") + + def test_extent_allocator_coalesces_released_ranges(self) -> None: + allocator = FreeExtentAllocator(4096, alignment=256) + first = allocator.allocate(300) + second = allocator.allocate(700) + self.assertEqual(first, (0, 512)) + self.assertEqual(second, (512, 768)) + assert first is not None and second is not None + allocator.release(*first) + allocator.release(*second) + self.assertEqual(allocator.extents, [(0, 4096)]) + + def test_fixed_block_allocator_uses_exact_runs_and_coalesces(self) -> None: + allocator = FixedBlockAllocator(8 * 256, block_bytes=256) + first = allocator.allocate(300) + second = allocator.allocate(300) + third = allocator.allocate(700) + self.assertEqual(first, (0, 1)) + self.assertEqual(second, (2, 3)) + self.assertEqual(third, (4, 5, 6)) + self.assertEqual(allocator.free_bytes, 256) + assert first is not None and second is not None and third is not None + allocator.release(second) + allocator.release(first) + self.assertEqual(allocator.largest_free_extent, 4 * 256) + allocator.release(third) + self.assertEqual(allocator.largest_free_extent, 8 * 256) + + def test_fixed_block_allocator_reuses_best_fit_run(self) -> None: + allocator = FixedBlockAllocator(12 * 256, block_bytes=256) + first = allocator.allocate(2 * 256) + separator = allocator.allocate(256) + second = allocator.allocate(3 * 256) + tail = allocator.allocate(6 * 256) + assert first is not None and separator is not None + assert second is not None and tail is not None + allocator.release(first) + allocator.release(second) + reused = allocator.allocate(3 * 256) + self.assertEqual(reused, second) + self.assertEqual(allocator.largest_free_extent, 2 * 256) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_pool_rejects_budget_larger_than_currently_free_memory(self) -> None: + device = available_device() + free_bytes, _ = torch.cuda.mem_get_info(torch.device(device)) + with self.assertRaisesRegex(RuntimeError, "cannot acquire"): + GpuResourcePool( + [GpuArenaSpec(device, free_bytes + 1024 * 1024)], + 1024 * 1024, + ) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_gpu_cache_evicts_only_released_entries(self) -> None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 16 * 1024 * 1024)], 8 * 1024 * 1024 + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + rows, dimension = 8192, 320 + payload = np.ones((rows, dimension), dtype=" None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 32 * 1024 * 1024)], 16 * 1024 * 1024 + ) + try: + cache = GpuTensorCache(pool, allow_eviction=False) + query = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + device = available_device() + pool = GpuResourcePool([GpuArenaSpec(device, 4 * 1024 * 1024)], 2 * 1024 * 1024) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + first = np.ones((128, 320), dtype=" None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 768 * 1024)], + 512 * 1024, + block_bytes=256 * 1024, + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + tensor = np.ones((128, 320), dtype=" None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 64 * 1024 * 1024)], 32 * 1024 * 1024 + ) + try: + generator = np.random.default_rng(7) + query = generator.standard_normal((44, 320)).astype("= 2, + "two CUDA devices are unavailable", + ) + def test_resident_engine_scores_shards_on_multiple_gpus(self) -> None: + pool = GpuResourcePool( + [ + GpuArenaSpec("cuda:0", 32 * 1024 * 1024), + GpuArenaSpec("cuda:1", 32 * 1024 * 1024), + ], + 16 * 1024 * 1024, + ) + try: + cache = GpuTensorCache(pool, allow_eviction=False) + identity = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" Path: + return Path(__file__).parent / "tilemaxsimd" / "target" / "release" / "tilemaxsimd" + + def run_daemon( + self, + devices: list[int], + documents: list[np.ndarray] | None = None, + query: np.ndarray | None = None, + gpu_memory_gb: str = "0.05", + workspace_gb: str = "0.02", + resident: bool = False, + scheduled: bool = False, + ) -> tuple[str, list[tuple[int, float]]]: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shard_root = root / "cache" + shard_root.mkdir() + if documents is None: + documents = [ + np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + self.run_daemon([device]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_external_v3_scheduled_round_trip_hashes_tenant_and_preserves_priority(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + output, _ = self.run_daemon([device], scheduled=True) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + request = next( + event for event in events if event.get("event") == "tilemaxsim_rust_request" + ) + self.assertNotIn("tenant", request) + self.assertRegex(request["tenant_hash"], r"^[0-9a-f]{16}$") + self.assertEqual(request["priority"], 17) + + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.device_count() >= 2, + "two CUDA devices are unavailable", + ) + def test_multi_gpu_scheduler_uploads_and_scores_on_each_device(self) -> None: + output, _ = self.run_daemon([0, 1]) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + request_event = next( + event for event in events if event.get("event") == "tilemaxsim_rust_request" + ) + devices = request_event["cache"]["devices"] + self.assertEqual(len(devices), 2) + self.assertEqual([device["h2d_batches"] for device in devices], [1, 1]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_one_request_larger_than_gpu_cache_is_scored_in_chunks(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + rows, dimension = 480, 320 + first = np.zeros((rows, dimension), dtype=" None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + output, _ = self.run_daemon([device], resident=True) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + prewarm = next( + event + for event in events + if event.get("event") == "tilemaxsim_rust_prewarm_complete" + ) + self.assertEqual(prewarm["entries"], 2) + self.assertEqual(prewarm["cache"]["devices"][0]["gpu_pinned_entries"], 2) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_concurrent_readers_do_not_block_fair_priority_scheduler(self) -> None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shard_root = root / "cache" + shard_root.mkdir() + document = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" int: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.settimeout(10) + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_bytes = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_bytes) + request_id, status, _ = decode_response(response) + self.assertEqual(status, 0) + return request_id + + with ThreadPoolExecutor(max_workers=len(frames)) as executor: + completed = list(executor.map(call, frames)) + self.assertEqual(set(completed), set(range(1_000, 1_000 + len(frames)))) + slow.close() + slow = None + process.terminate() + output, _ = process.communicate(timeout=10) + self.assertEqual(process.returncode, 0, output) + events = [ + json.loads(line) + for line in output.splitlines() + if line.startswith("{") + ] + processed = [ + event["priority"] + for event in events + if event.get("event") == "tilemaxsim_rust_request" + ] + # All public priorities share the default fair-priority band. + # Urgency breaks the first tie, then equal-cost tenants + # alternate instead of one tenant draining its whole queue. + self.assertEqual(processed, [9, 3, 8, 1, 7, 0, 5, -2]) + finally: + if slow is not None: + slow.close() + if process.poll() is None: + process.terminate() + process.wait(timeout=10) + + +if __name__ == "__main__": + unittest.main() diff --git a/services/test_tilemaxsim_shard.py b/services/test_tilemaxsim_shard.py new file mode 100644 index 00000000..591e4954 --- /dev/null +++ b/services/test_tilemaxsim_shard.py @@ -0,0 +1,181 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +from __future__ import annotations + +import hashlib +import json +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import ContentAddressedResolver +from services.tilemaxsim_shard import ImmutableShardWriter, load_index + + +class ImmutableShardTest(unittest.TestCase): + def test_builder_defaults_to_shards_and_publishes_atomically(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source" + source.mkdir() + manifest = source / "pages.jsonl" + records = [] + for index in range(3): + tensor = np.full((2, 4), index + 1, dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + tensors = [ + np.arange(32, dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + tensor = np.eye(4, dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = np.ones((2, 4), dtype=" None: self.maximum_bytes = maximum_bytes self.current_bytes = 0 - self.entries: OrderedDict[tuple[object, ...], bytes] = OrderedDict() + self.entries: OrderedDict[tuple[object, ...], tuple[bytes, float]] = ( + OrderedDict() + ) self.lock = threading.Lock() + self.sketch = TinyLfuSketch() + self.inflation = 0.0 + self.hits = 0 + self.misses = 0 + self.evictions = 0 + self.admission_rejections = 0 def get(self, key: tuple[object, ...]) -> bytes | None: if self.maximum_bytes == 0: return None with self.lock: - payload = self.entries.get(key) - if payload is not None: + frequency = self.sketch.increment(key) + entry = self.entries.get(key) + if entry is not None: + payload, _ = entry + priority = self.inflation + frequency / len(payload) + self.entries[key] = (payload, priority) self.entries.move_to_end(key) - return payload + self.hits += 1 + return payload + self.misses += 1 + return None def put(self, key: tuple[object, ...], payload: bytes) -> None: if self.maximum_bytes == 0 or len(payload) > self.maximum_bytes: @@ -97,12 +140,39 @@ def put(self, key: tuple[object, ...], payload: bytes) -> None: with self.lock: previous = self.entries.pop(key, None) if previous is not None: - self.current_bytes -= len(previous) - self.entries[key] = payload + self.current_bytes -= len(previous[0]) + frequency = max(1, self.sketch.estimate(key)) + candidate_priority = self.inflation + frequency / len(payload) + while self.current_bytes + len(payload) > self.maximum_bytes: + victim_key, (victim_payload, victim_priority) = min( + self.entries.items(), key=lambda item: item[1][1] + ) + if candidate_priority < victim_priority: + if previous is not None: + self.entries[key] = previous + self.current_bytes += len(previous[0]) + self.admission_rejections += 1 + return + self.entries.pop(victim_key) + self.current_bytes -= len(victim_payload) + self.inflation = max(self.inflation, victim_priority) + candidate_priority = self.inflation + frequency / len(payload) + self.evictions += 1 + self.entries[key] = (payload, candidate_priority) self.current_bytes += len(payload) - while self.current_bytes > self.maximum_bytes: - _, evicted = self.entries.popitem(last=False) - self.current_bytes -= len(evicted) + + def status(self) -> dict[str, object]: + with self.lock: + return { + "entries": len(self.entries), + "bytes": self.current_bytes, + "maximum_bytes": self.maximum_bytes, + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "admission_rejections": self.admission_rejections, + "policy": "tinylfu-gdsf", + } class ContentAddressedResolver: @@ -115,20 +185,85 @@ class ContentAddressedResolver: agree before a payload is returned. """ - def __init__(self, roots: dict[str, Path], cache_bytes: int) -> None: + def __init__( + self, + roots: dict[str, Path], + cache_bytes: int, + verify_full_shards: bool = False, + ) -> None: self.root_fds: dict[str, int] = {} + self.shard_roots: dict[str, _OpenShardRoot] = {} + self.batch_read_calls = 0 + self.batch_read_bytes = 0 + self.batch_lock = threading.Lock() + self.verify_full_shards = verify_full_shards try: for contract, path in roots.items(): - self.root_fds[contract] = os.open( + root_fd = os.open( os.fspath(path), os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, ) + self.root_fds[contract] = root_fd + self._open_shard_root(contract, root_fd) except Exception: self.close() raise self.cache = PayloadCache(cache_bytes) + def _open_shard_root(self, contract: str, root_fd: int) -> None: + try: + index_fd = os.open( + INDEX_NAME, + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + except FileNotFoundError: + return + try: + with os.fdopen(index_fd, "r", encoding="utf-8", closefd=True) as stream: + index = parse_index(json.load(stream)) + except Exception: + raise + shard_directory_fd = -1 + shard_fds: dict[str, int] = {} + try: + shard_directory_fd = os.open( + "shards", + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + for relative, shard in index.shards.items(): + name = relative.removeprefix("shards/") + descriptor = os.open( + name, + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=shard_directory_fd, + ) + metadata = os.fstat(descriptor) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_size != shard.size: + os.close(descriptor) + raise ValueError(f"immutable shard metadata disagrees: {relative}") + shard_fds[relative] = descriptor + self.shard_roots[contract] = _OpenShardRoot( + index, + shard_directory_fd, + shard_fds, + {name: threading.Lock() for name in shard_fds}, + set(), + ) + except Exception: + for descriptor in shard_fds.values(): + os.close(descriptor) + if shard_directory_fd >= 0: + os.close(shard_directory_fd) + raise + def close(self) -> None: + for root in self.shard_roots.values(): + for descriptor in root.shard_fds.values(): + os.close(descriptor) + os.close(root.directory_fd) + self.shard_roots.clear() for descriptor in self.root_fds.values(): os.close(descriptor) self.root_fds.clear() @@ -214,7 +349,7 @@ def _read_exact_file(root_fd: int, digest: str, expected_bytes: int) -> bytes: if directory_fd >= 0: os.close(directory_fd) - def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: + def key(self, request: protocol.ExternalTensorRequest) -> tuple[object, ...]: root_fd = self.root_fds.get(request.model_contract_id) if root_fd is None: raise protocol.SidecarError( @@ -222,9 +357,6 @@ def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: "model contract has no configured tensor cache root", ) digest = self._digest(request) - expected_bytes = protocol.checked_tensor_bytes( - request.rows, request.dimension, request.dtype - ) key = ( request.model_contract_id, digest, @@ -232,18 +364,222 @@ def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: request.dimension, request.dtype, ) - cached = self.cache.get(key) - if cached is not None: - return ResolvedPayload(cached, True) - payload = self._read_exact_file(root_fd, digest, expected_bytes) - actual = hashlib.sha256(payload).hexdigest() - if not hmac.compare_digest(actual, digest): + return key + + @staticmethod + def _validate_shard_entry( + request: protocol.ExternalTensorRequest, digest: str, entry: ShardEntry + ) -> None: + dtype_name = "float32" if request.dtype == protocol.DTYPE_F32 else "float16" + if ( + entry.digest != digest + or entry.rows != request.rows + or entry.dimension != request.dimension + or entry.dtype != dtype_name + or entry.length + != protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + ): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "shard entry disagrees with the tensor descriptor", + ) + + @staticmethod + def _verify_shard(root: _OpenShardRoot, name: str) -> None: + if name in root.verified: + return + lock = root.verification_locks[name] + with lock: + if name in root.verified: + return + shard = root.index.shards[name] + expected = shard.checksum.removeprefix("sha256:") + digest = hashlib.sha256() + offset = 0 + descriptor = root.shard_fds[name] + while offset < shard.size: + chunk = os.pread( + descriptor, min(8 * 1024 * 1024, shard.size - offset), offset + ) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "immutable tensor shard ended early", + ) + digest.update(chunk) + offset += len(chunk) + if not hmac.compare_digest(digest.hexdigest(), expected): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "immutable tensor shard checksum mismatch", + ) + root.verified.add(name) + + @staticmethod + def _coalesced_ranges( + entries: list[tuple[tuple[object, ...], ShardEntry]], + maximum_gap: int = 64 * 1024, + maximum_span: int = 8 * 1024 * 1024, + ) -> list[tuple[int, int, list[tuple[tuple[object, ...], ShardEntry]]]]: + ranges = [] + current: list[tuple[tuple[object, ...], ShardEntry]] = [] + start = 0 + end = 0 + for item in sorted(entries, key=lambda value: value[1].offset): + entry = item[1] + entry_end = entry.offset + entry.length + if current and ( + entry.offset - end > maximum_gap or entry_end - start > maximum_span + ): + ranges.append((start, end, current)) + current = [] + if not current: + start = entry.offset + end = entry_end + else: + end = max(end, entry_end) + current.append(item) + if current: + ranges.append((start, end, current)) + return ranges + + def _read_shard_range( + self, + root: _OpenShardRoot, + shard_name: str, + start: int, + end: int, + entries: list[tuple[tuple[object, ...], ShardEntry]], + ) -> dict[tuple[object, ...], bytes]: + payload = os.pread(root.shard_fds[shard_name], end - start, start) + if len(payload) != end - start: raise protocol.SidecarError( - protocol.STATUS_INVALID_REQUEST, "resolved tensor checksum mismatch" + protocol.STATUS_COMPUTE_ERROR, "immutable tensor shard ended early" + ) + with self.batch_lock: + self.batch_read_calls += 1 + self.batch_read_bytes += len(payload) + result = {} + for key, entry in entries: + tensor = payload[entry.offset - start : entry.offset - start + entry.length] + actual = hashlib.sha256(tensor).hexdigest() + if not hmac.compare_digest(actual, entry.digest): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved shard tensor checksum mismatch", + ) + dtype = ( + protocol.DTYPE_F32 if entry.dtype == "float32" else protocol.DTYPE_F16 ) - validate_finite_payload(payload, request.rows, request.dimension, request.dtype) - self.cache.put(key, payload) - return ResolvedPayload(payload, False) + validate_finite_payload(tensor, entry.rows, entry.dimension, dtype) + result[key] = tensor + return result + + def resolve_many( + self, requests: list[protocol.ExternalTensorRequest] + ) -> list[ResolvedPayload]: + if not requests: + return [] + keys = [self.key(request) for request in requests] + payloads: dict[tuple[object, ...], bytes] = {} + hits: dict[tuple[object, ...], bool] = {} + missing: dict[tuple[object, ...], protocol.ExternalTensorRequest] = {} + for key, request in zip(keys, requests, strict=True): + cached = self.cache.get(key) + if cached is not None: + payloads[key] = cached + hits[key] = True + elif key not in missing: + missing[key] = request + hits[key] = False + + shard_groups: dict[ + tuple[str, str], list[tuple[tuple[object, ...], ShardEntry]] + ] = {} + legacy: list[tuple[tuple[object, ...], protocol.ExternalTensorRequest]] = [] + for key, request in missing.items(): + contract = request.model_contract_id + shard_root = self.shard_roots.get(contract) + if shard_root is None: + legacy.append((key, request)) + continue + digest = str(key[1]) + entry = shard_root.index.entries.get(digest) + if entry is None: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "content-addressed tensor is missing from the shard index", + ) + self._validate_shard_entry(request, digest, entry) + shard_groups.setdefault((contract, entry.shard), []).append((key, entry)) + + jobs: list[ + tuple[ + _OpenShardRoot, + str, + int, + int, + list[tuple[tuple[object, ...], ShardEntry]], + ] + ] = [] + for (contract, shard_name), entries in shard_groups.items(): + root = self.shard_roots[contract] + if self.verify_full_shards: + self._verify_shard(root, shard_name) + for start, end, grouped in self._coalesced_ranges(entries): + jobs.append((root, shard_name, start, end, grouped)) + if jobs: + with ThreadPoolExecutor(max_workers=min(8, len(jobs))) as workers: + for resolved in workers.map( + lambda job: self._read_shard_range(*job), jobs + ): + payloads.update(resolved) + + def read_legacy( + item: tuple[tuple[object, ...], protocol.ExternalTensorRequest], + ) -> tuple[tuple[object, ...], bytes]: + key, request = item + digest = str(key[1]) + expected_bytes = protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + payload = self._read_exact_file( + self.root_fds[request.model_contract_id], digest, expected_bytes + ) + actual = hashlib.sha256(payload).hexdigest() + if not hmac.compare_digest(actual, digest): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor checksum mismatch", + ) + validate_finite_payload( + payload, request.rows, request.dimension, request.dtype + ) + return key, payload + + if legacy: + with ThreadPoolExecutor(max_workers=min(8, len(legacy))) as workers: + for key, payload in workers.map(read_legacy, legacy): + payloads[key] = payload + for key in missing: + self.cache.put(key, payloads[key]) + return [ResolvedPayload(payloads[key], hits[key]) for key in keys] + + def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: + return self.resolve_many([request])[0] + + def status(self) -> dict[str, object]: + with self.batch_lock: + return { + "shard_contracts": len(self.shard_roots), + "verified_shards": sum( + len(root.verified) for root in self.shard_roots.values() + ), + "batch_read_calls": self.batch_read_calls, + "batch_read_bytes": self.batch_read_bytes, + } class TorchTileMaxsimEngine: @@ -393,6 +729,190 @@ def score( ) +class ResidentTorchTileMaxsimEngine: + """Score tensors already owned by one or more process GPU arenas.""" + + def __init__( + self, + pool: GpuResourcePool, + max_workspace_bytes: int, + allow_tf32: bool, + max_cuda_inflight: int, + ) -> None: + self.pool = pool + self.device = pool.primary_device + self.max_workspace_bytes = max_workspace_bytes + self.compute_slots = threading.BoundedSemaphore(max_cuda_inflight) + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.backends.cudnn.allow_tf32 = allow_tf32 + with torch.inference_mode(): + for arena in pool.arenas: + left = torch.zeros((1, 1), dtype=torch.float32, device=arena.device) + _ = left @ left + for arena in pool.arenas: + torch.cuda.synchronize(arena.device) + + @staticmethod + def _cpu_tensor( + payload: bytes, rows: int, dimension: int, dtype: int + ) -> torch.Tensor: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + return torch.frombuffer(bytearray(payload), dtype=scalar_dtype).reshape( + rows, dimension + ) + + def _groups( + self, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, GpuTensorHandle]], + ) -> Iterable[list[tuple[int, GpuTensorHandle]]]: + scalar_bytes = 4 if dtype == protocol.DTYPE_F32 else 2 + query_bytes = query_rows * dimension * scalar_bytes + group: list[tuple[int, GpuTensorHandle]] = [] + group_rows = 0 + for document in documents: + rows = document[1].rows + next_rows = group_rows + rows + # The resident document remains inside the arena. torch.cat makes + # one device-local contiguous scoring view; the other temporaries + # are the query and q-by-document-token similarity matrix. + required = ( + query_bytes + + next_rows * dimension * scalar_bytes + + query_rows * next_rows * scalar_bytes + ) + if required > self.max_workspace_bytes and group: + yield group + group = [] + group_rows = 0 + next_rows = rows + required = ( + query_bytes + + rows * dimension * scalar_bytes + + query_rows * rows * scalar_bytes + ) + if required > self.max_workspace_bytes: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one resident candidate exceeds the configured GPU workspace", + ) + group.append(document) + group_rows = next_rows + if group: + yield group + + def score( + self, + query_payload: bytes, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, GpuTensorHandle]], + deadline: float, + cancelled: Callable[[], bool], + ) -> tuple[list[tuple[int, float]], float, float]: + if not documents: + return [], 0.0, 0.0 + query_cpu = self._cpu_tensor(query_payload, query_rows, dimension, dtype) + by_device: dict[str, list[tuple[int, GpuTensorHandle]]] = {} + for document in documents: + handle = document[1] + if handle.dimension != dimension or handle.dtype != dtype: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resident tensor contract disagrees with the query", + ) + by_device.setdefault(str(handle.device), []).append(document) + + queue_started = time.monotonic() + remaining = deadline - time.monotonic() + if remaining <= 0 or not self.compute_slots.acquire(timeout=remaining): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired while waiting for CUDA capacity", + ) + queue_ms = (time.monotonic() - queue_started) * 1000.0 + compute_started = time.monotonic() + pending: list[tuple[list[tuple[int, GpuTensorHandle]], torch.Tensor]] = [] + try: + with torch.inference_mode(): + for arena in self.pool.arenas: + device_documents = by_device.get(str(arena.device), []) + if not device_documents: + continue + if cancelled(): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request peer disconnected", + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + query_device = query_cpu.to(arena.device) + if dtype == protocol.DTYPE_F16 and ragged_tilemaxsim_fp16: + offsets = torch.tensor( + [ + handle.offset_bytes // 2 + for _, handle in device_documents + ], + dtype=torch.int64, + device=arena.device, + ) + rows = torch.tensor( + [handle.rows for _, handle in device_documents], + dtype=torch.int32, + device=arena.device, + ) + assert arena.storage is not None + device_scores = ragged_tilemaxsim_fp16( + query_device, + arena.storage.view(torch.float16), + offsets, + rows, + max(handle.rows for _, handle in device_documents), + ) + pending.append((device_documents, device_scores)) + continue + for group in self._groups( + query_rows, dimension, dtype, device_documents + ): + document_device = torch.cat( + [handle.tensor() for _, handle in group] + ) + similarities = query_device @ document_device.transpose(0, 1) + scores = [] + offset = 0 + for _, handle in group: + scores.append( + similarities[:, offset : offset + handle.rows] + .amax(dim=1) + .sum(dtype=torch.float32) + ) + offset += handle.rows + pending.append((group, torch.stack(scores))) + + results: list[tuple[int, float]] = [] + for group, device_scores in pending: + host_scores = device_scores.to(device="cpu", dtype=torch.float32) + for (candidate_id, _), score in zip( + group, host_scores.tolist(), strict=True + ): + if not math.isfinite(score): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "TileMaxSim result is non-finite", + ) + results.append((candidate_id, score)) + for arena in self.pool.arenas: + torch.cuda.synchronize(arena.device) + finally: + self.compute_slots.release() + return results, queue_ms, (time.monotonic() - compute_started) * 1000.0 + + class JsonMetrics: def __init__(self) -> None: self.lock = threading.Lock() @@ -410,12 +930,22 @@ def __init__( engine: TorchTileMaxsimEngine, request_timeout_ms: int, metrics: JsonMetrics, + gpu_cache: GpuTensorCache | None = None, + resident_engine: ResidentTorchTileMaxsimEngine | None = None, + pin_gpu_entries: bool = False, ) -> None: self.limits = limits self.resolver = resolver self.engine = engine self.request_timeout_seconds = request_timeout_ms / 1000.0 self.metrics = metrics + self.gpu_cache = gpu_cache + self.resident_engine = resident_engine + self.pin_gpu_entries = pin_gpu_entries + if (gpu_cache is None) != (resident_engine is None): + raise ValueError( + "GPU cache and resident engine must be configured together" + ) @staticmethod def _peer_disconnected(connection: socket.socket) -> bool: @@ -453,6 +983,7 @@ def process_frame( version = protocol.VERSION started = time.monotonic() metrics: dict[str, object] = {"event": "tilemaxsim_request"} + resident_documents: list[tuple[int, GpuTensorHandle]] = [] if peer_credentials is not None: metrics["peer_pid"], metrics["peer_uid"], metrics["peer_gid"] = ( peer_credentials @@ -460,7 +991,7 @@ def process_frame( try: if len(frame) >= protocol.HEADER.size: _, wire_version, _, request_id, _ = protocol.HEADER.unpack_from(frame) - if wire_version in (protocol.VERSION, protocol.EXTERNAL_VERSION): + if wire_version in protocol.SUPPORTED_VERSIONS: version = wire_version request = protocol.parse_request_frame( frame, self.limits, validate_finite=False @@ -480,7 +1011,42 @@ def process_frame( ) resolve_started = time.monotonic() cache_hits = 0 + gpu_cache_hits = 0 + gpu_cache_misses = 0 + gpu_chunks = 0 + resident_results: list[tuple[int, float]] = [] + resident_queue_ms = 0.0 + resident_compute_ms = 0.0 + document_tokens = 0 documents: list[tuple[int, int, bytes]] = [] + + def flush_resident_documents() -> None: + nonlocal gpu_chunks, resident_queue_ms, resident_compute_ms + if not resident_documents: + return + assert self.resident_engine is not None + try: + batch_results, batch_queue_ms, batch_compute_ms = ( + self.resident_engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + resident_documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + ) + resident_results.extend(batch_results) + resident_queue_ms += batch_queue_ms + resident_compute_ms += batch_compute_ms + gpu_chunks += 1 + finally: + assert self.gpu_cache is not None + for _, resident_handle in resident_documents: + self.gpu_cache.release(resident_handle) + resident_documents.clear() + if isinstance(request, protocol.InlineTensorRequest): for candidate in request.candidates: validate_finite_payload( @@ -493,38 +1059,192 @@ def process_frame( (candidate.candidate_id, candidate.rows, candidate.payload) for candidate in request.candidates ] + document_tokens = sum( + candidate.rows for candidate in request.candidates + ) metrics["source"] = "inline" else: metrics["source"] = "content_addressed" - for candidate in request.candidates: - if time.monotonic() >= deadline: - raise protocol.SidecarError( - protocol.STATUS_COMPUTE_ERROR, - "request deadline expired during tensor resolution", + document_tokens = sum( + candidate.descriptor.rows for candidate in request.candidates + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during tensor resolution", + ) + if self.gpu_cache is None: + descriptors = [ + candidate.descriptor for candidate in request.candidates + ] + if hasattr(self.resolver, "resolve_many"): + resolved_payloads = self.resolver.resolve_many(descriptors) + else: + resolved_payloads = [ + self.resolver.resolve(descriptor) + for descriptor in descriptors + ] + for candidate, resolved in zip( + request.candidates, resolved_payloads, strict=True + ): + cache_hits += int(resolved.cache_hit) + documents.append( + ( + candidate.candidate_id, + candidate.descriptor.rows, + resolved.payload, + ) ) - resolved = self.resolver.resolve(candidate.descriptor) - cache_hits += int(resolved.cache_hit) - documents.append( + else: + keys = [ + self.resolver.key(candidate.descriptor) + for candidate in request.candidates + ] + probed, miss_indices = self.gpu_cache.probe_many(keys) + for candidate, handle in zip( + request.candidates, probed, strict=True + ): + if handle is not None: + resident_documents.append((candidate.candidate_id, handle)) + gpu_cache_hits = len(request.candidates) - len(miss_indices) + gpu_cache_misses = len(miss_indices) + missing_candidates = [ + request.candidates[index] for index in miss_indices + ] + missing_descriptors = [ + candidate.descriptor for candidate in missing_candidates + ] + if hasattr(self.resolver, "resolve_many"): + resolved_payloads = self.resolver.resolve_many( + missing_descriptors + ) + else: + resolved_payloads = [ + self.resolver.resolve(descriptor) + for descriptor in missing_descriptors + ] + cache_hits += sum(item.cache_hit for item in resolved_payloads) + pending = [ ( - candidate.candidate_id, - candidate.descriptor.rows, - resolved.payload, + candidate, + GpuTensorLoad( + key, + candidate.descriptor.rows, + candidate.descriptor.dimension, + candidate.descriptor.dtype, + resolved.payload, + self.pin_gpu_entries, + ), + False, ) - ) + for candidate, key, resolved in zip( + missing_candidates, + (keys[index] for index in miss_indices), + resolved_payloads, + strict=True, + ) + ] + admission_rejections = 0 + while pending: + batch = self.gpu_cache.acquire_many( + [item[1] for item in pending], + enforce_admission=( + not self.pin_gpu_entries + and not any(item[2] for item in pending) + ), + record_access=False, + count_stats=False, + ) + for (candidate, load, _force), handle in zip( + pending, batch.handles, strict=True + ): + if handle is not None: + resident_documents.append( + (candidate.candidate_id, handle) + ) + for index in batch.bypassed: + candidate, load, _force = pending[index] + stream_bytes = ( + request.query_rows * request.dimension * 4 + + load.rows * request.dimension * 4 + + request.query_rows * load.rows * 4 + ) + if stream_bytes <= self.engine.max_device_bytes: + documents.append( + (candidate.candidate_id, load.rows, load.payload) + ) + admission_rejections += len(batch.bypassed) + deferred = [pending[index] for index in batch.deferred] + forced = [ + (pending[index][0], pending[index][1], True) + for index in batch.bypassed + if ( + request.query_rows * request.dimension * 4 + + pending[index][1].rows * request.dimension * 4 + + request.query_rows * pending[index][1].rows * 4 + > self.engine.max_device_bytes + ) + ] + made_progress = len(deferred) < len(pending) + if resident_documents: + flush_resident_documents() + made_progress = True + if deferred and not made_progress: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "GPU block cache cannot make progress with its configured capacity", + ) + pending = forced + deferred + if resident_documents: + flush_resident_documents() + metrics["gpu_admission_rejections"] = admission_rejections metrics["cache_hits"] = cache_hits + metrics["host_cache_hits"] = cache_hits + metrics["gpu_cache_hits"] = gpu_cache_hits + metrics["gpu_cache_misses"] = gpu_cache_misses + metrics["gpu_chunks"] = gpu_chunks metrics["resolve_ms"] = round( - (time.monotonic() - resolve_started) * 1000.0, 3 - ) - metrics["document_tokens"] = sum(rows for _, rows, _ in documents) - results, queue_ms, compute_ms = self.engine.score( - request.query_payload, - request.query_rows, - request.dimension, - request.dtype, - documents, - deadline, - lambda: self._peer_disconnected(connection), + max( + 0.0, + (time.monotonic() - resolve_started) * 1000.0 + - resident_queue_ms + - resident_compute_ms, + ), + 3, ) + metrics["document_tokens"] = document_tokens + if self.gpu_cache is not None and isinstance( + request, protocol.ParsedExternalTensorRequest + ): + results = resident_results + queue_ms = resident_queue_ms + compute_ms = resident_compute_ms + if documents: + stream_results, stream_queue_ms, stream_compute_ms = ( + self.engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + ) + results.extend(stream_results) + queue_ms += stream_queue_ms + compute_ms += stream_compute_ms + metrics["streamed_candidates"] = len(documents) + else: + results, queue_ms, compute_ms = self.engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + documents, + deadline, + lambda: self._peer_disconnected(connection), + ) metrics["queue_ms"] = round(queue_ms, 3) metrics["compute_ms"] = round(compute_ms, 3) metrics["status"] = "ok" @@ -535,8 +1255,6 @@ def process_frame( request_id, error.status, str(error), version ) except torch.OutOfMemoryError: - if self.engine.device.type == "cuda": - torch.cuda.empty_cache() metrics.update(status="error", error_class="CudaOutOfMemory") return protocol.error_response( request_id, @@ -553,6 +1271,9 @@ def process_frame( version, ) finally: + if self.gpu_cache is not None: + for _, handle in resident_documents: + self.gpu_cache.release(handle) metrics["total_ms"] = round((time.monotonic() - started) * 1000.0, 3) self.metrics.emit(metrics) @@ -574,7 +1295,7 @@ def handle(self, connection: socket.socket) -> None: connection, protocol.HEADER.size, deadline ) _, wire_version, _, request_id, body_len = protocol.HEADER.unpack(header) - if wire_version in (protocol.VERSION, protocol.EXTERNAL_VERSION): + if wire_version in protocol.SUPPORTED_VERSIONS: version = wire_version if body_len > self.limits.max_request_bytes - protocol.HEADER.size: response = protocol.error_response( @@ -655,14 +1376,15 @@ def handle(connection: socket.socket) -> None: listener.listen(backlog) listener.settimeout(0.25) bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino - service.metrics.emit( - { - "event": "tilemaxsim_ready", - "device": str(service.engine.device), - "max_inflight": max_inflight, - "socket": os.fspath(socket_path), - } - ) + ready: dict[str, object] = { + "event": "tilemaxsim_ready", + "device": str(service.engine.device), + "max_inflight": max_inflight, + "socket": os.fspath(socket_path), + } + if service.gpu_cache is not None: + ready["gpu_cache"] = service.gpu_cache.status() + service.metrics.emit(ready) accepted = 0 while not stop.is_set(): if not slots.acquire(timeout=0.25): @@ -714,6 +1436,20 @@ def nonnegative_int(value: str) -> int: return parsed +def memory_gb(value: str) -> int: + try: + return parse_memory_gb(value) + except ValueError as error: + raise argparse.ArgumentTypeError(str(error)) from error + + +def gpu_memory_gb(value: str) -> GpuArenaSpec: + try: + return parse_gpu_memory_gb(value) + except (ValueError, RuntimeError) as error: + raise argparse.ArgumentTypeError(str(error)) from error + + def contract_roots( values: list[str], parser: argparse.ArgumentParser ) -> dict[str, Path]: @@ -731,12 +1467,180 @@ def contract_roots( return roots +def contract_manifests( + values: list[str], parser: argparse.ArgumentParser +) -> list[tuple[str, Path]]: + manifests = [] + for value in values: + if "=" not in value: + parser.error("--resident-manifest must be MODEL_CONTRACT_ID=/absolute/path") + contract, raw_path = value.split("=", 1) + path = Path(raw_path) + if not contract or not path.is_absolute(): + parser.error( + "--resident-manifest must contain a nonempty ID and absolute path" + ) + manifests.append((contract, path)) + return manifests + + +def prewarm_resident_cache( + manifests: list[tuple[str, Path]], + resolver: ContentAddressedResolver, + gpu_cache: GpuTensorCache, + metrics: JsonMetrics, + batch_size: int = 256, +) -> None: + completed = 0 + loaded_bytes = 0 + started = time.monotonic() + pending: list[tuple[protocol.ExternalTensorRequest, int]] = [] + + def flush() -> None: + nonlocal completed, loaded_bytes + if not pending: + return + keys = [resolver.key(request) for request, _ in pending] + probed, miss_indices = gpu_cache.probe_many(keys) + acquired = [handle for handle in probed if handle is not None] + try: + missing = [pending[index][0] for index in miss_indices] + resolved = resolver.resolve_many(missing) + loads = [ + GpuTensorLoad( + keys[index], + request.rows, + request.dimension, + request.dtype, + payload.payload, + True, + ) + for index, request, payload in zip( + miss_indices, missing, resolved, strict=True + ) + ] + batch = gpu_cache.acquire_many( + loads, + enforce_admission=False, + record_access=False, + count_stats=False, + ) + if ( + batch.bypassed + or batch.deferred + or any(handle is None for handle in batch.handles) + ): + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "resident manifest exceeds the configured GPU block arenas", + ) + acquired.extend(handle for handle in batch.handles if handle is not None) + finally: + for handle in acquired: + gpu_cache.release(handle) + completed += len(pending) + loaded_bytes += sum(size for _, size in pending) + if completed % 1000 < len(pending): + metrics.emit( + { + "event": "tilemaxsim_prewarm_progress", + "entries": completed, + "logical_bytes": loaded_bytes, + } + ) + pending.clear() + + for contract, path in manifests: + with path.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path}:{line_number}: invalid JSON") from error + if not isinstance(record, dict): + raise ValueError(f"{path}:{line_number}: record must be an object") + dtype_name = record.get("tensor_dtype") + if dtype_name == "float16": + dtype = protocol.DTYPE_F16 + elif dtype_name == "float32": + dtype = protocol.DTYPE_F32 + else: + raise ValueError(f"{path}:{line_number}: unsupported tensor_dtype") + tensor_ref = record.get("tensor_ref") + checksum = record.get("tensor_checksum") + rows = record.get("tensor_rows") + dimension = record.get("tensor_dim") + if not isinstance(tensor_ref, str) or not tensor_ref: + raise ValueError(f"{path}:{line_number}: invalid tensor_ref") + if not isinstance(checksum, str) or not checksum: + raise ValueError(f"{path}:{line_number}: invalid tensor_checksum") + if not isinstance(rows, int) or rows <= 0: + raise ValueError(f"{path}:{line_number}: invalid tensor_rows") + if not isinstance(dimension, int) or dimension <= 0: + raise ValueError(f"{path}:{line_number}: invalid tensor_dim") + request = protocol.ExternalTensorRequest( + contract, tensor_ref, rows, dimension, dtype, checksum + ) + expected_bytes = protocol.checked_tensor_bytes(rows, dimension, dtype) + declared_bytes = record.get("canonical_bytes") + if declared_bytes is not None and declared_bytes != expected_bytes: + raise ValueError( + f"{path}:{line_number}: canonical_bytes disagrees with shape" + ) + pending.append((request, expected_bytes)) + if len(pending) >= batch_size: + flush() + flush() + if completed == 0: + raise ValueError("resident manifests contain no tensor descriptors") + metrics.emit( + { + "event": "tilemaxsim_prewarm_complete", + "entries": completed, + "logical_bytes": loaded_bytes, + "elapsed_ms": round((time.monotonic() - started) * 1000.0, 3), + "gpu_cache": gpu_cache.status(), + } + ) + + def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--socket", required=True, type=Path) parser.add_argument("--socket-mode", type=parse_mode, default=0o600) - parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--gpu-cache-mode", + choices=("lru", "resident"), + default="lru", + help="use evictable GPU arenas or pin a full descriptor manifest", + ) + parser.add_argument( + "--gpu-memory-gb", + action="append", + type=gpu_memory_gb, + default=[], + metavar="GPU=GB", + help="repeatable strict allocation, for example 1=20; enables TileMaxSim", + ) + parser.add_argument( + "--gpu-workspace-gb", + type=memory_gb, + default=2 * 1024**3, + help="per-GPU portion of the configured GB reserved for scoring temporaries", + ) + parser.add_argument( + "--gpu-block-kib", + type=positive_int, + default=32, + help="base page size inside each preallocated GPU tensor arena", + ) parser.add_argument("--contract-root", action="append", default=[]) + parser.add_argument( + "--resident-manifest", + action="append", + default=[], + metavar="MODEL_CONTRACT_ID=/ABSOLUTE/PATH", + ) parser.add_argument( "--max-request-bytes", type=positive_int, default=64 * 1024 * 1024 ) @@ -746,37 +1650,84 @@ def main() -> None: ) parser.add_argument("--max-candidates", type=positive_int, default=65_536) parser.add_argument( - "--max-device-bytes", type=positive_int, default=8 * 1024 * 1024 * 1024 - ) - parser.add_argument( - "--cache-bytes", type=nonnegative_int, default=8 * 1024 * 1024 * 1024 + "--host-cache-gb", + type=memory_gb, + default=8 * 1024**3, + help="decoded host-memory tensor cache size in GB", ) parser.add_argument("--request-timeout-ms", type=positive_int, default=2000) parser.add_argument("--max-inflight", type=positive_int, default=8) parser.add_argument("--max-cuda-inflight", type=positive_int, default=1) + parser.add_argument("--prewarm-batch-size", type=positive_int, default=256) parser.add_argument("--backlog", type=positive_int, default=64) parser.add_argument("--allow-tf32", action="store_true") + parser.add_argument( + "--verify-full-shards", + action="store_true", + help="verify complete immutable shard digests lazily in addition to every tensor digest", + ) parser.add_argument("--once", action="store_true") args = parser.parse_args() roots = contract_roots(args.contract_root, parser) + manifests = contract_manifests(args.resident_manifest, parser) + if not args.gpu_memory_gb: + parser.error( + "TileMaxSim is disabled until at least one --gpu-memory-gb GPU=GB is configured" + ) + if args.gpu_cache_mode == "resident" and not manifests: + parser.error( + "--gpu-cache-mode resident requires at least one --resident-manifest" + ) + if args.gpu_cache_mode == "lru" and manifests: + parser.error("--resident-manifest is valid only in resident mode") limits = protocol.Limits( max_request_bytes=args.max_request_bytes, max_batch_tokens=args.max_batch_tokens, max_tensor_bytes=args.max_tensor_bytes, max_candidates=args.max_candidates, ) - resolver = ContentAddressedResolver(roots, args.cache_bytes) + resolver = ContentAddressedResolver( + roots, args.host_cache_gb, args.verify_full_shards + ) metrics = JsonMetrics() + pool: GpuResourcePool | None = None try: + pool = GpuResourcePool( + args.gpu_memory_gb, + args.gpu_workspace_gb, + args.gpu_block_kib * 1024, + ) engine = TorchTileMaxsimEngine( - args.device, - args.max_device_bytes, + str(pool.primary_device), + args.gpu_workspace_gb, args.allow_tf32, args.max_cuda_inflight, ) + gpu_cache = GpuTensorCache(pool, allow_eviction=args.gpu_cache_mode == "lru") + resident_engine = ResidentTorchTileMaxsimEngine( + pool, + args.gpu_workspace_gb, + args.allow_tf32, + args.max_cuda_inflight, + ) + if args.gpu_cache_mode == "resident": + prewarm_resident_cache( + manifests, + resolver, + gpu_cache, + metrics, + args.prewarm_batch_size, + ) service = TileMaxsimService( - limits, resolver, engine, args.request_timeout_ms, metrics + limits, + resolver, + engine, + args.request_timeout_ms, + metrics, + gpu_cache, + resident_engine, + pin_gpu_entries=args.gpu_cache_mode == "resident", ) stop = threading.Event() @@ -795,6 +1746,8 @@ def request_stop(_signum: int, _frame: object) -> None: args.once, ) finally: + if pool is not None: + pool.close() resolver.close() diff --git a/services/tilemaxsim_gpu_cache.py b/services/tilemaxsim_gpu_cache.py new file mode 100644 index 00000000..dd17e7a3 --- /dev/null +++ b/services/tilemaxsim_gpu_cache.py @@ -0,0 +1,975 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2026 Hu Xinjing + +"""Process-owned GPU arenas and a bounded tensor cache for TileMaxSim.""" + +from __future__ import annotations + +import re +import threading +from bisect import bisect_left, insort +from collections import OrderedDict +from dataclasses import dataclass +from decimal import Decimal, InvalidOperation +from hashlib import blake2b +from math import ceil +from typing import Callable + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol + + +_GPU_MEMORY_GB = re.compile(r"^(?:cuda:)?([0-9]+)=([0-9]+(?:\.[0-9]+)?)$") +_MEMORY_GB = re.compile(r"^[0-9]+(?:\.[0-9]+)?$") +GIB = 1024**3 +# A 32 KiB page keeps the allocator metadata small while avoiding the 8.82% +# rounding loss observed with the former 256 KiB default on 472--483 KiB +# document tensors. The native daemon exposes the same default via +# ``--gpu-block-kib``. +DEFAULT_BLOCK_BYTES = 32 * 1024 +DEFAULT_STAGING_BYTES = 64 * 1024 * 1024 + + +@dataclass(frozen=True) +class GpuArenaSpec: + device: str + total_bytes: int + + +def parse_gpu_memory_gb(value: str) -> GpuArenaSpec: + """Parse the public ``GPU=GB`` configuration into an internal byte budget.""" + + match = _GPU_MEMORY_GB.fullmatch(value.strip()) + if match is None: + raise ValueError( + "GPU memory must be GPU=GB, for example 1=20; byte suffixes are not accepted" + ) + raw_index, raw_gb = match.groups() + try: + total_bytes = int(Decimal(raw_gb) * GIB) + except InvalidOperation as error: + raise ValueError("GPU memory GB value is invalid") from error + if total_bytes <= 0: + raise ValueError("GPU memory GB value must be positive") + return GpuArenaSpec(f"cuda:{int(raw_index)}", total_bytes) + + +def parse_memory_gb(value: str) -> int: + if _MEMORY_GB.fullmatch(value.strip()) is None: + raise ValueError("memory size must be a positive number of GB") + try: + total_bytes = int(Decimal(value.strip()) * GIB) + except InvalidOperation as error: + raise ValueError("memory size must be a positive number of GB") from error + if total_bytes <= 0: + raise ValueError("memory size must be a positive number of GB") + return total_bytes + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +class FreeExtentAllocator: + """Best-fit allocator for contiguous slices of a preallocated byte arena.""" + + def __init__(self, capacity: int, alignment: int = 256) -> None: + if capacity <= 0: + raise ValueError("arena capacity must be positive") + self.capacity = capacity + self.alignment = alignment + self.extents: list[tuple[int, int]] = [(0, capacity)] + + @property + def free_bytes(self) -> int: + return sum(length for _, length in self.extents) + + @property + def largest_free_extent(self) -> int: + return max((length for _, length in self.extents), default=0) + + def allocation_bytes(self, payload_bytes: int) -> int: + return _align_up(payload_bytes, self.alignment) + + def allocate(self, payload_bytes: int) -> tuple[int, int] | None: + required = self.allocation_bytes(payload_bytes) + choices = [ + (length, index, start) + for index, (start, length) in enumerate(self.extents) + if length >= required + ] + if not choices: + return None + length, index, start = min(choices) + if length == required: + self.extents.pop(index) + else: + self.extents[index] = (start + required, length - required) + return start, required + + def release(self, start: int, length: int) -> None: + if start < 0 or length <= 0 or start + length > self.capacity: + raise ValueError("released extent is outside the arena") + self.extents.append((start, length)) + self.extents.sort() + merged: list[tuple[int, int]] = [] + for extent_start, extent_length in self.extents: + if merged and merged[-1][0] + merged[-1][1] == extent_start: + previous_start, previous_length = merged[-1] + merged[-1] = (previous_start, previous_length + extent_length) + else: + merged.append((extent_start, extent_length)) + self.extents = merged + + +class FixedBlockAllocator: + """Best-fit page-run allocator over one process-owned GPU arena. + + Tensors keep the dense layout required by the TileMaxSim kernel, but their + runs are rounded only to the base page instead of the next power of two. + Free runs are indexed by both address and exact size. Allocation therefore + finds the smallest suitable run, while release coalesces adjacent runs. + """ + + def __init__(self, capacity: int, block_bytes: int = DEFAULT_BLOCK_BYTES) -> None: + if capacity <= 0 or block_bytes <= 0: + raise ValueError("arena capacity and block size must be positive") + if block_bytes % 256: + raise ValueError("GPU block size must be 256-byte aligned") + self.block_bytes = block_bytes + self.block_count = capacity // block_bytes + if self.block_count == 0: + raise ValueError("arena must contain at least one GPU block") + self.capacity = self.block_count * block_bytes + self._free_by_start: dict[int, int] = {} + self._free_by_size: dict[int, set[int]] = {} + self._starts: list[int] = [] + self._sizes: list[int] = [] + self._allocated: dict[int, int] = {} + self._free_blocks = 0 + self._add_free_run(0, self.block_count) + + def _add_free_run(self, start: int, blocks: int) -> None: + if blocks <= 0 or start < 0 or start + blocks > self.block_count: + raise ValueError("free GPU run is outside the arena") + if start in self._free_by_start: + raise ValueError("duplicate free GPU run") + self._free_by_start[start] = blocks + insort(self._starts, start) + bucket = self._free_by_size.get(blocks) + if bucket is None: + bucket = set() + self._free_by_size[blocks] = bucket + insort(self._sizes, blocks) + bucket.add(start) + self._free_blocks += blocks + + def _remove_free_run(self, start: int, blocks: int) -> None: + if self._free_by_start.get(start) != blocks: + raise ValueError("free GPU run directory is inconsistent") + del self._free_by_start[start] + start_index = bisect_left(self._starts, start) + if start_index == len(self._starts) or self._starts[start_index] != start: + raise ValueError("free GPU address index is inconsistent") + self._starts.pop(start_index) + bucket = self._free_by_size[blocks] + bucket.remove(start) + if not bucket: + del self._free_by_size[blocks] + size_index = bisect_left(self._sizes, blocks) + if size_index == len(self._sizes) or self._sizes[size_index] != blocks: + raise ValueError("free GPU size index is inconsistent") + self._sizes.pop(size_index) + self._free_blocks -= blocks + + @property + def free_blocks(self) -> int: + return self._free_blocks + + @property + def free_bytes(self) -> int: + return self.free_blocks * self.block_bytes + + @property + def largest_free_extent(self) -> int: + return self._sizes[-1] * self.block_bytes if self._sizes else 0 + + def blocks_for(self, payload_bytes: int) -> int: + if payload_bytes <= 0: + raise ValueError("payload size must be positive") + return ceil(payload_bytes / self.block_bytes) + + def allocation_bytes(self, payload_bytes: int) -> int: + return self.blocks_for(payload_bytes) * self.block_bytes + + def allocate(self, payload_bytes: int) -> tuple[int, ...] | None: + required = self.blocks_for(payload_bytes) + size_index = bisect_left(self._sizes, required) + if size_index == len(self._sizes): + return None + available = self._sizes[size_index] + start = min(self._free_by_size[available]) + self._remove_free_run(start, available) + if available > required: + self._add_free_run(start + required, available - required) + self._allocated[start] = required + return tuple(range(start, start + required)) + + def release(self, blocks: tuple[int, ...]) -> None: + if ( + not blocks + or len(set(blocks)) != len(blocks) + or blocks != tuple(range(blocks[0], blocks[0] + len(blocks))) + ): + raise ValueError("released GPU blocks are invalid") + allocated_blocks = self._allocated.pop(blocks[0], None) + if allocated_blocks is None: + raise ValueError("GPU block was released more than once") + if len(blocks) != allocated_blocks: + raise ValueError("released GPU page run has the wrong size") + start = blocks[0] + length = allocated_blocks + insertion = bisect_left(self._starts, start) + if insertion: + previous_start = self._starts[insertion - 1] + previous_length = self._free_by_start[previous_start] + if previous_start + previous_length == start: + self._remove_free_run(previous_start, previous_length) + start = previous_start + length += previous_length + next_start = start + length + next_length = self._free_by_start.get(next_start) + if next_length is not None: + self._remove_free_run(next_start, next_length) + length += next_length + self._add_free_run(start, length) + + +class TinyLfuSketch: + """A small aging count-min sketch for cache admission and GDSF frequency.""" + + def __init__(self, width: int = 4096, depth: int = 4) -> None: + if width <= 0 or depth <= 0: + raise ValueError("TinyLFU dimensions must be positive") + self.width = width + self.depth = depth + self.tables = [[0] * width for _ in range(depth)] + self.samples = 0 + self.reset_at = width * 10 + + def _indices(self, key: tuple[object, ...]) -> tuple[int, ...]: + digest = blake2b(repr(key).encode("utf-8"), digest_size=16).digest() + return tuple( + int.from_bytes(digest[row * 4 : row * 4 + 4], "little") % self.width + for row in range(self.depth) + ) + + def increment(self, key: tuple[object, ...]) -> int: + indices = self._indices(key) + for row, index in enumerate(indices): + if self.tables[row][index] < 65535: + self.tables[row][index] += 1 + self.samples += 1 + estimate = min(self.tables[row][index] for row, index in enumerate(indices)) + if self.samples >= self.reset_at: + for table in self.tables: + for index, value in enumerate(table): + table[index] = value // 2 + self.samples //= 2 + return max(1, estimate) + + def estimate(self, key: tuple[object, ...]) -> int: + return min( + self.tables[row][index] for row, index in enumerate(self._indices(key)) + ) + + +class GpuArena: + """A CUDA byte buffer acquired atomically during process startup.""" + + def __init__( + self, + spec: GpuArenaSpec, + workspace_bytes: int, + block_bytes: int = DEFAULT_BLOCK_BYTES, + ) -> None: + self.device = torch.device(spec.device) + if not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested but torch.cuda.is_available() is false" + ) + if self.device.index is None or self.device.index >= torch.cuda.device_count(): + raise RuntimeError(f"configured CUDA device is unavailable: {spec.device}") + if workspace_bytes <= 0 or workspace_bytes >= spec.total_bytes: + raise RuntimeError( + f"{spec.device} allocation must exceed its TileMaxSim workspace" + ) + self.total_bytes = spec.total_bytes + self.workspace_bytes = workspace_bytes + raw_capacity = spec.total_bytes - workspace_bytes + self.allocator = FixedBlockAllocator(raw_capacity, block_bytes) + self.capacity = self.allocator.capacity + self.reserved_workspace_bytes = spec.total_bytes - self.capacity + if self.capacity <= 0: + raise RuntimeError(f"{spec.device} has no aligned tensor-cache capacity") + self.storage: torch.Tensor | None = None + self.host_staging: torch.Tensor | None = None + self.copy_stream: torch.cuda.Stream | None = None + self.h2d_batches = 0 + self.h2d_copy_calls = 0 + self.h2d_bytes = 0 + + with torch.cuda.device(self.device): + free_bytes, device_bytes = torch.cuda.mem_get_info(self.device) + self.device_bytes = device_bytes + if free_bytes < spec.total_bytes: + raise RuntimeError( + f"cannot acquire {spec.total_bytes} bytes on {spec.device}: " + f"only {free_bytes} bytes are free" + ) + try: + self.storage = torch.empty( + self.capacity, dtype=torch.uint8, device=self.device + ) + # Reserve the remaining configured budget in PyTorch's CUDA + # allocator. Releasing this temporary tensor leaves the block + # in the process-owned caching allocator for TileMaxSim + # workspaces instead of returning it to another process. + workspace = torch.empty( + self.reserved_workspace_bytes, + dtype=torch.uint8, + device=self.device, + ) + torch.cuda.synchronize(self.device) + del workspace + staging_bytes = min(self.capacity, DEFAULT_STAGING_BYTES) + staging_bytes = max( + self.allocator.block_bytes, + staging_bytes + // self.allocator.block_bytes + * self.allocator.block_bytes, + ) + self.host_staging = torch.empty( + staging_bytes, dtype=torch.uint8, pin_memory=True + ) + # Fault and pin every staging page during startup so the first + # cache-miss batch does not pay a request-path NUMA/page cost. + self.host_staging.zero_() + self.copy_stream = torch.cuda.Stream(device=self.device) + except Exception: + self.storage = None + self.host_staging = None + self.copy_stream = None + torch.cuda.empty_cache() + raise + + def tensor( + self, + blocks: tuple[int, ...], + payload_bytes: int, + rows: int, + dimension: int, + dtype: int, + ) -> torch.Tensor: + if ( + self.storage is None + or self.host_staging is None + or self.copy_stream is None + ): + raise RuntimeError("GPU arena is closed") + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + block_bytes = self.allocator.block_bytes + if all(right == left + 1 for left, right in zip(blocks, blocks[1:])): + raw = self.storage.narrow(0, blocks[0] * block_bytes, payload_bytes) + else: + raw = torch.cat( + [ + self.storage.narrow(0, block * block_bytes, block_bytes) + for block in blocks + ] + ).narrow(0, 0, payload_bytes) + return raw.view(scalar_dtype).reshape(rows, dimension) + + def copy_many_from_host(self, items: list[tuple[tuple[int, ...], bytes]]) -> None: + if self.storage is None: + raise RuntimeError("GPU arena is closed") + if not items: + return + block_bytes = self.allocator.block_bytes + runs: list[tuple[int, int, bytes]] = [] + for blocks, payload in items: + if ( + not blocks + or blocks != tuple(range(blocks[0], blocks[0] + len(blocks))) + or len(payload) > len(blocks) * block_bytes + ): + raise RuntimeError("GPU upload requires one valid contiguous page run") + runs.append((blocks[0], len(blocks), payload)) + runs.sort(key=lambda item: item[0]) + staging = self.host_staging + staging_array = staging.numpy() + staging_bytes = staging.numel() + copy_calls = 0 + with torch.cuda.device(self.device): + stream = self.copy_stream + run_index = 0 + while run_index < len(runs): + first_block, run_blocks, payload = runs[run_index] + allocation_bytes = run_blocks * block_bytes + if allocation_bytes > staging_bytes: + source_offset = 0 + while source_offset < len(payload): + length = min(staging_bytes, len(payload) - source_offset) + staging_array[:length] = np.frombuffer( + payload, + dtype=np.uint8, + count=length, + offset=source_offset, + ) + with torch.cuda.stream(stream): + self.storage.narrow( + 0, + first_block * block_bytes + source_offset, + length, + ).copy_(staging.narrow(0, 0, length), non_blocking=True) + stream.synchronize() + copy_calls += 1 + source_offset += length + run_index += 1 + continue + + batch_first_block = first_block + expected_block = first_block + staging_offset = 0 + while run_index < len(runs): + start_block, count, current_payload = runs[run_index] + current_allocation = count * block_bytes + if ( + start_block != expected_block + or staging_offset + current_allocation > staging_bytes + ): + break + payload_end = staging_offset + len(current_payload) + staging_array[staging_offset:payload_end] = np.frombuffer( + current_payload, dtype=np.uint8 + ) + allocation_end = staging_offset + current_allocation + staging_array[payload_end:allocation_end] = 0 + staging_offset = allocation_end + expected_block += count + run_index += 1 + with torch.cuda.stream(stream): + self.storage.narrow( + 0, batch_first_block * block_bytes, staging_offset + ).copy_( + staging.narrow(0, 0, staging_offset), + non_blocking=True, + ) + copy_calls += 1 + stream.synchronize() + self.h2d_batches += 1 + self.h2d_copy_calls += copy_calls + self.h2d_bytes += sum(len(payload) for _, payload in items) + + def copy_from_host(self, blocks: tuple[int, ...], payload: bytes) -> None: + self.copy_many_from_host([(blocks, payload)]) + + def status(self) -> dict[str, object]: + return { + "device": str(self.device), + "allocated_gb": round(self.total_bytes / GIB, 3), + "tensor_capacity_gb": round(self.capacity / GIB, 3), + "workspace_gb": round(self.workspace_bytes / GIB, 3), + "allocated_bytes": self.total_bytes, + "tensor_capacity_bytes": self.capacity, + "workspace_bytes": self.workspace_bytes, + "tensor_free_bytes": self.allocator.free_bytes, + "largest_free_extent_bytes": self.allocator.largest_free_extent, + "block_bytes": self.allocator.block_bytes, + "block_count": self.allocator.block_count, + "free_blocks": self.allocator.free_blocks, + "h2d_batches": self.h2d_batches, + "h2d_copy_calls": self.h2d_copy_calls, + "h2d_bytes": self.h2d_bytes, + "host_staging_bytes": self.host_staging.numel() + if self.host_staging is not None + else 0, + } + + def close(self) -> None: + if self.storage is None: + return + with torch.cuda.device(self.device): + self.storage = None + self.host_staging = None + self.copy_stream = None + torch.cuda.empty_cache() + + +class GpuResourcePool: + """Own all configured CUDA allocations or fail without a partial pool.""" + + def __init__( + self, + specs: list[GpuArenaSpec], + workspace_bytes: int, + block_bytes: int = DEFAULT_BLOCK_BYTES, + ) -> None: + if not specs: + raise RuntimeError("at least one GPU allocation is required") + devices = [spec.device for spec in specs] + if len(devices) != len(set(devices)): + raise RuntimeError("each CUDA device may be configured only once") + self.arenas: list[GpuArena] = [] + try: + for spec in specs: + self.arenas.append(GpuArena(spec, workspace_bytes, block_bytes)) + except Exception: + self.close() + raise + + @property + def primary_device(self) -> torch.device: + return self.arenas[0].device + + def status(self) -> list[dict[str, object]]: + return [arena.status() for arena in self.arenas] + + def close(self) -> None: + for arena in self.arenas: + arena.close() + self.arenas.clear() + + +@dataclass +class _GpuCacheEntry: + key: tuple[object, ...] + arena: GpuArena + blocks: tuple[int, ...] + allocated_bytes: int + payload_bytes: int + rows: int + dimension: int + dtype: int + references: int = 0 + pinned: bool = False + priority: float = 0.0 + + +@dataclass(frozen=True) +class GpuTensorHandle: + entry: _GpuCacheEntry + + @property + def arena(self) -> GpuArena: + return self.entry.arena + + @property + def device(self) -> torch.device: + return self.entry.arena.device + + @property + def rows(self) -> int: + return self.entry.rows + + @property + def dimension(self) -> int: + return self.entry.dimension + + @property + def dtype(self) -> int: + return self.entry.dtype + + @property + def payload_bytes(self) -> int: + return self.entry.payload_bytes + + @property + def offset_bytes(self) -> int: + return self.entry.blocks[0] * self.entry.arena.allocator.block_bytes + + @property + def block_ids(self) -> tuple[int, ...]: + return self.entry.blocks + + @property + def block_bytes(self) -> int: + return self.entry.arena.allocator.block_bytes + + def tensor(self) -> torch.Tensor: + return self.entry.arena.tensor( + self.entry.blocks, + self.entry.payload_bytes, + self.entry.rows, + self.entry.dimension, + self.entry.dtype, + ) + + +@dataclass(frozen=True) +class GpuTensorLoad: + key: tuple[object, ...] + rows: int + dimension: int + dtype: int + payload: bytes + pin: bool = False + + +@dataclass(frozen=True) +class GpuAcquireBatch: + handles: tuple[GpuTensorHandle | None, ...] + bypassed: tuple[int, ...] + deferred: tuple[int, ...] + hits: int + misses: int + admitted: int + + +class GpuTensorCache: + """Fixed-block GPU cache with TinyLFU admission and GDSF eviction.""" + + def __init__(self, pool: GpuResourcePool, allow_eviction: bool) -> None: + self.pool = pool + self.allow_eviction = allow_eviction + self.entries: OrderedDict[tuple[object, ...], _GpuCacheEntry] = OrderedDict() + self.lock = threading.Lock() + self.hits = 0 + self.misses = 0 + self.evictions = 0 + self.loaded_bytes = 0 + self.admission_rejections = 0 + self.inflation = 0.0 + self.sketch = TinyLfuSketch() + + def _find_arena(self, payload_bytes: int) -> GpuArena | None: + candidates = [ + arena + for arena in self.pool.arenas + if arena.allocator.largest_free_extent + >= arena.allocator.allocation_bytes(payload_bytes) + ] + if not candidates: + return None + return max(candidates, key=lambda arena: arena.allocator.free_bytes) + + def _victim(self, arena: GpuArena | None = None) -> _GpuCacheEntry | None: + candidates = [ + entry + for entry in self.entries.values() + if not entry.references + and not entry.pinned + and (arena is None or entry.arena is arena) + ] + return min(candidates, key=lambda entry: entry.priority, default=None) + + def _evict(self, entry: _GpuCacheEntry) -> None: + current = self.entries.pop(entry.key, None) + if current is not entry: + raise RuntimeError("GPU eviction directory is inconsistent") + entry.arena.allocator.release(entry.blocks) + self.inflation = max(self.inflation, entry.priority) + self.evictions += 1 + + @staticmethod + def _entry_cost(entry: _GpuCacheEntry) -> int: + return len(entry.blocks) + + def _priority(self, key: tuple[object, ...], blocks: int) -> float: + return self.inflation + self.sketch.estimate(key) / max(1, blocks) + + def _allocate( + self, + key: tuple[object, ...], + payload_bytes: int, + enforce_admission: bool, + ) -> tuple[GpuArena, tuple[int, ...], int] | None: + arena = self._find_arena(payload_bytes) + capable = [ + candidate + for candidate in self.pool.arenas + if candidate.allocator.capacity + >= candidate.allocator.allocation_bytes(payload_bytes) + ] + if not capable: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one tensor exceeds every configured GPU block arena", + ) + if arena is None and self.allow_eviction: + for candidate in sorted( + capable, key=lambda item: item.allocator.free_bytes, reverse=True + ): + required_blocks = candidate.allocator.blocks_for(payload_bytes) + required_bytes = candidate.allocator.allocation_bytes(payload_bytes) + while candidate.allocator.largest_free_extent < required_bytes: + victim = self._victim(candidate) + if victim is None: + break + candidate_priority = self._priority(key, required_blocks) + if enforce_admission and candidate_priority <= victim.priority: + self.admission_rejections += 1 + return None + self._evict(victim) + if candidate.allocator.largest_free_extent >= required_bytes: + arena = candidate + break + if arena is None: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "configured GPU tensor arenas have insufficient free blocks", + ) + blocks = arena.allocator.allocate(payload_bytes) + assert blocks is not None + return arena, blocks, len(blocks) * arena.allocator.block_bytes + + def acquire( + self, + key: tuple[object, ...], + rows: int, + dimension: int, + dtype: int, + loader: Callable[[], bytes], + *, + pin: bool = False, + ) -> tuple[GpuTensorHandle, bool]: + with self.lock: + frequency = self.sketch.increment(key) + cached = self.entries.get(key) + if cached is not None: + cached.references += 1 + cached.pinned = cached.pinned or pin + cached.priority = self.inflation + frequency / self._entry_cost(cached) + self.entries.move_to_end(key) + self.hits += 1 + return GpuTensorHandle(cached), True + + payload = loader() + expected = protocol.checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match its shape", + ) + + with self.lock: + cached = self.entries.get(key) + if cached is not None: + frequency = self.sketch.increment(key) + cached.references += 1 + cached.pinned = cached.pinned or pin + cached.priority = self.inflation + frequency / self._entry_cost(cached) + self.entries.move_to_end(key) + self.hits += 1 + return GpuTensorHandle(cached), True + allocated = self._allocate(key, len(payload), enforce_admission=False) + assert allocated is not None + arena, blocks, allocated_bytes = allocated + try: + arena.copy_from_host(blocks, payload) + except Exception: + arena.allocator.release(blocks) + raise + entry = _GpuCacheEntry( + key, + arena, + blocks, + allocated_bytes, + len(payload), + rows, + dimension, + dtype, + references=1, + pinned=pin, + priority=self._priority(key, len(blocks)), + ) + self.entries[key] = entry + self.misses += 1 + self.loaded_bytes += len(payload) + return GpuTensorHandle(entry), False + + def acquire_many( + self, + loads: list[GpuTensorLoad], + *, + enforce_admission: bool = True, + record_access: bool = True, + count_stats: bool = True, + ) -> GpuAcquireBatch: + """Acquire a request working set and upload all new slabs in batches. + + ``deferred`` items could not be allocated while earlier handles in the + same request are referenced. The caller scores/releases the returned + handles and retries those items. ``bypassed`` items lost TinyLFU + admission and should be scored through the bounded streaming engine. + """ + + handles: list[GpuTensorHandle | None] = [None] * len(loads) + bypassed: list[int] = [] + deferred: list[int] = [] + new_entries: list[tuple[int, _GpuCacheEntry, bytes]] = [] + hits = 0 + misses = 0 + admitted = 0 + with self.lock: + try: + for index, load in enumerate(loads): + expected = protocol.checked_tensor_bytes( + load.rows, load.dimension, load.dtype + ) + if len(load.payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match its shape", + ) + frequency = ( + self.sketch.increment(load.key) + if record_access + else max(1, self.sketch.estimate(load.key)) + ) + cached = self.entries.get(load.key) + if cached is not None: + cached.references += 1 + cached.pinned = cached.pinned or load.pin + cached.priority = self.inflation + frequency / self._entry_cost( + cached + ) + self.entries.move_to_end(load.key) + handles[index] = GpuTensorHandle(cached) + hits += 1 + continue + misses += 1 + try: + allocation = self._allocate( + load.key, + len(load.payload), + enforce_admission and not load.pin, + ) + except protocol.SidecarError as error: + if "one tensor exceeds" in str(error): + bypassed.append(index) + continue + deferred.append(index) + continue + if allocation is None: + bypassed.append(index) + continue + arena, blocks, allocated_bytes = allocation + entry = _GpuCacheEntry( + load.key, + arena, + blocks, + allocated_bytes, + len(load.payload), + load.rows, + load.dimension, + load.dtype, + references=1, + pinned=load.pin, + priority=self._priority(load.key, len(blocks)), + ) + self.entries[load.key] = entry + handles[index] = GpuTensorHandle(entry) + new_entries.append((index, entry, load.payload)) + admitted += 1 + + by_arena: dict[ + int, tuple[GpuArena, list[tuple[tuple[int, ...], bytes]]] + ] = {} + for _, entry, payload in new_entries: + bucket = by_arena.setdefault(id(entry.arena), (entry.arena, [])) + bucket[1].append((entry.blocks, payload)) + for arena, items in by_arena.values(): + arena.copy_many_from_host(items) + if count_stats: + self.hits += hits + self.misses += misses + self.loaded_bytes += sum(len(payload) for _, _, payload in new_entries) + except Exception: + new_ids = {id(entry) for _, entry, _ in new_entries} + for handle in handles: + if handle is None: + continue + entry = handle.entry + if id(entry) in new_ids: + if self.entries.pop(entry.key, None) is entry: + entry.arena.allocator.release(entry.blocks) + else: + entry.references -= 1 + raise + return GpuAcquireBatch( + tuple(handles), + tuple(bypassed), + tuple(deferred), + hits, + misses, + admitted, + ) + + def probe_many( + self, keys: list[tuple[object, ...]] + ) -> tuple[tuple[GpuTensorHandle | None, ...], tuple[int, ...]]: + """Acquire GPU hits without resolving the corresponding host payloads.""" + + handles: list[GpuTensorHandle | None] = [None] * len(keys) + misses = [] + with self.lock: + for index, key in enumerate(keys): + frequency = self.sketch.increment(key) + entry = self.entries.get(key) + if entry is None: + misses.append(index) + self.misses += 1 + continue + entry.references += 1 + entry.priority = self.inflation + frequency / self._entry_cost(entry) + self.entries.move_to_end(key) + handles[index] = GpuTensorHandle(entry) + self.hits += 1 + return tuple(handles), tuple(misses) + + def release(self, handle: GpuTensorHandle) -> None: + with self.lock: + entry = self.entries.get(handle.entry.key) + if entry is not handle.entry or entry.references <= 0: + raise RuntimeError("GPU tensor handle was released more than once") + entry.references -= 1 + + def status(self) -> dict[str, object]: + with self.lock: + allocated_bytes = sum( + entry.allocated_bytes for entry in self.entries.values() + ) + payload_bytes = sum(entry.payload_bytes for entry in self.entries.values()) + return { + "allocator": "segregated-page-runs", + "entries": len(self.entries), + "pinned_entries": sum(entry.pinned for entry in self.entries.values()), + "active_references": sum( + entry.references for entry in self.entries.values() + ), + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "admission_rejections": self.admission_rejections, + "policy": "tinylfu-gdsf", + "gdsf_inflation": self.inflation, + "loaded_bytes": self.loaded_bytes, + "allocated_bytes": allocated_bytes, + "payload_bytes": payload_bytes, + "internal_waste_bytes": allocated_bytes - payload_bytes, + "arenas": self.pool.status(), + } diff --git a/services/tilemaxsim_shard.py b/services/tilemaxsim_shard.py new file mode 100644 index 00000000..2837b4ed --- /dev/null +++ b/services/tilemaxsim_shard.py @@ -0,0 +1,354 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2026 Hu Xinjing + +"""Immutable, content-addressed TileMaxSim tensor shard format. + +The data files contain canonical tensor bytes and alignment padding only. A +generation index maps tensor SHA-256 digests to byte ranges. Data files are +published under their own SHA-256, so a writer never mutates a visible shard. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO, Iterable + + +FORMAT = "vectorchord.tilemaxsim.shards" +VERSION = 1 +INDEX_NAME = "tilemaxsim-shards-v1.json" +SHARD_DIRECTORY = "shards" +DEFAULT_ALIGNMENT = 4096 +DEFAULT_SHARD_BYTES = 2 * 1024**3 + + +def align_up(value: int, alignment: int) -> int: + if alignment <= 0 or alignment & (alignment - 1): + raise ValueError("alignment must be a positive power of two") + return (value + alignment - 1) // alignment * alignment + + +@dataclass(frozen=True) +class ShardEntry: + digest: str + shard: str + offset: int + length: int + rows: int + dimension: int + dtype: str + + +@dataclass(frozen=True) +class ShardFile: + name: str + size: int + checksum: str + + +@dataclass(frozen=True) +class ShardIndex: + alignment: int + shards: dict[str, ShardFile] + entries: dict[str, ShardEntry] + + +def _digest_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(8 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +class ImmutableShardWriter: + """Deterministically publish bounded immutable shards and one atomic index.""" + + def __init__( + self, + root: Path, + target_bytes: int = DEFAULT_SHARD_BYTES, + alignment: int = DEFAULT_ALIGNMENT, + fsync: bool = True, + ) -> None: + if target_bytes <= 0: + raise ValueError("target shard bytes must be positive") + if target_bytes < alignment: + raise ValueError("target shard bytes must be at least one alignment unit") + align_up(0, alignment) + self.root = root + self.shard_root = root / SHARD_DIRECTORY + self.target_bytes = target_bytes + self.alignment = alignment + self.fsync = fsync + self.entries: dict[str, ShardEntry] = {} + self.shards: list[ShardFile] = [] + self._stream: BinaryIO | None = None + self._temporary: Path | None = None + self._hasher: hashlib._Hash | None = None + self._offset = 0 + self._pending: list[tuple[str, int, int, int, int, str]] = [] + self.root.mkdir(parents=True, exist_ok=True) + self.shard_root.mkdir(parents=True, exist_ok=True) + + def _open(self) -> None: + descriptor, name = tempfile.mkstemp( + prefix=".tilemaxsim-shard-", suffix=".tmp", dir=self.shard_root + ) + self._stream = os.fdopen(descriptor, "wb", closefd=True) + self._temporary = Path(name) + self._hasher = hashlib.sha256() + self._offset = 0 + self._pending = [] + + def _write(self, payload: bytes | memoryview) -> None: + assert self._stream is not None and self._hasher is not None + self._stream.write(payload) + self._hasher.update(payload) + self._offset += len(payload) + + def add( + self, + digest: str, + payload: bytes | memoryview, + rows: int, + dimension: int, + dtype: str, + ) -> None: + if len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest): + raise ValueError("invalid tensor SHA-256 digest") + if digest in self.entries or any(item[0] == digest for item in self._pending): + return + if not payload: + raise ValueError("tensor payload must not be empty") + if rows <= 0 or dimension <= 0 or dtype not in ("float16", "float32"): + raise ValueError("invalid tensor metadata") + padded = align_up(len(payload), self.alignment) + if self._stream is not None and self._offset and self._offset + padded > self.target_bytes: + self._finish_shard() + if self._stream is None: + self._open() + offset = self._offset + self._write(payload) + padding = padded - len(payload) + if padding: + self._write(bytes(min(padding, self.alignment))) + remaining = padding - min(padding, self.alignment) + while remaining: + chunk = min(remaining, self.alignment) + self._write(bytes(chunk)) + remaining -= chunk + self._pending.append((digest, offset, len(payload), rows, dimension, dtype)) + + def _finish_shard(self) -> None: + if self._stream is None: + return + assert self._temporary is not None and self._hasher is not None + stream = self._stream + stream.flush() + if self.fsync: + os.fsync(stream.fileno()) + stream.close() + digest = self._hasher.hexdigest() + name = f"sha256-{digest}.vts" + destination = self.shard_root / name + size = self._offset + if destination.exists(): + if destination.stat().st_size != size or _digest_file(destination) != digest: + raise ValueError(f"existing immutable shard is corrupt: {destination}") + self._temporary.unlink() + else: + os.replace(self._temporary, destination) + if self.fsync: + descriptor = os.open(self.shard_root, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + relative = f"{SHARD_DIRECTORY}/{name}" + shard = ShardFile(relative, size, f"sha256:{digest}") + self.shards.append(shard) + for tensor_digest, offset, length, rows, dimension, dtype in self._pending: + self.entries[tensor_digest] = ShardEntry( + tensor_digest, + relative, + offset, + length, + rows, + dimension, + dtype, + ) + self._stream = None + self._temporary = None + self._hasher = None + self._offset = 0 + self._pending = [] + + def finish(self) -> Path: + self._finish_shard() + if not self.entries: + raise ValueError("cannot publish an empty shard generation") + document = { + "format": FORMAT, + "version": VERSION, + "alignment": self.alignment, + "shards": [ + {"name": shard.name, "bytes": shard.size, "checksum": shard.checksum} + for shard in self.shards + ], + "entries": [ + { + "digest": entry.digest, + "shard": entry.shard, + "offset": entry.offset, + "length": entry.length, + "rows": entry.rows, + "dimension": entry.dimension, + "dtype": entry.dtype, + } + for entry in self.entries.values() + ], + } + descriptor, temporary = tempfile.mkstemp( + prefix=f".{INDEX_NAME}.", suffix=".tmp", dir=self.root, text=True + ) + try: + with os.fdopen(descriptor, "w", encoding="utf-8", closefd=True) as stream: + json.dump(document, stream, separators=(",", ":"), sort_keys=True) + stream.write("\n") + stream.flush() + if self.fsync: + os.fsync(stream.fileno()) + destination = self.root / INDEX_NAME + os.replace(temporary, destination) + if self.fsync: + directory = os.open(self.root, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory) + finally: + os.close(directory) + return destination + finally: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + + def close(self) -> None: + if self._stream is not None: + self._stream.close() + if self._temporary is not None: + try: + self._temporary.unlink() + except FileNotFoundError: + pass + self._stream = None + self._temporary = None + + +def parse_index(document: object) -> ShardIndex: + if not isinstance(document, dict): + raise ValueError("shard index must be a JSON object") + if document.get("format") != FORMAT or document.get("version") != VERSION: + raise ValueError("unsupported TileMaxSim shard index") + alignment = document.get("alignment") + if not isinstance(alignment, int): + raise ValueError("shard index has invalid alignment") + align_up(0, alignment) + raw_shards = document.get("shards") + raw_entries = document.get("entries") + if not isinstance(raw_shards, list) or not isinstance(raw_entries, list): + raise ValueError("shard index has invalid arrays") + shards: dict[str, ShardFile] = {} + for raw in raw_shards: + if not isinstance(raw, dict): + raise ValueError("shard index contains an invalid shard") + name, size, checksum = raw.get("name"), raw.get("bytes"), raw.get("checksum") + if ( + not isinstance(name, str) + or not name.startswith(f"{SHARD_DIRECTORY}/sha256-") + or "/" in name[len(SHARD_DIRECTORY) + 1 :] + or not isinstance(size, int) + or size <= 0 + or not isinstance(checksum, str) + or not checksum.startswith("sha256:") + or name != f"{SHARD_DIRECTORY}/sha256-{checksum.removeprefix('sha256:')}.vts" + or name in shards + ): + raise ValueError("shard index contains invalid shard metadata") + shards[name] = ShardFile(name, size, checksum) + entries: dict[str, ShardEntry] = {} + intervals: dict[str, list[tuple[int, int]]] = {name: [] for name in shards} + for raw in raw_entries: + if not isinstance(raw, dict): + raise ValueError("shard index contains an invalid tensor entry") + digest = raw.get("digest") + shard = raw.get("shard") + offset = raw.get("offset") + length = raw.get("length") + rows = raw.get("rows") + dimension = raw.get("dimension") + dtype = raw.get("dtype") + if ( + not isinstance(digest, str) + or len(digest) != 64 + or any(c not in "0123456789abcdef" for c in digest) + or not isinstance(shard, str) + or shard not in shards + or not isinstance(offset, int) + or offset < 0 + or offset % alignment + or not isinstance(length, int) + or length <= 0 + or not isinstance(rows, int) + or rows <= 0 + or not isinstance(dimension, int) + or dimension <= 0 + or dtype not in ("float16", "float32") + or digest in entries + ): + raise ValueError("shard index contains invalid tensor metadata") + scalar_bytes = 2 if dtype == "float16" else 4 + if length != rows * dimension * scalar_bytes: + raise ValueError("shard tensor length disagrees with its shape") + end = offset + length + if end > shards[shard].size: + raise ValueError("shard tensor range is outside its data file") + intervals[shard].append((offset, align_up(end, alignment))) + entries[digest] = ShardEntry( + digest, shard, offset, length, rows, dimension, dtype + ) + if not entries: + raise ValueError("shard index contains no tensor entries") + for shard, ranges in intervals.items(): + previous_end = 0 + for start, end in sorted(ranges): + if start < previous_end: + raise ValueError(f"overlapping tensor ranges in {shard}") + previous_end = end + return ShardIndex(alignment, shards, entries) + + +def load_index(path: Path) -> ShardIndex: + with path.open(encoding="utf-8") as stream: + return parse_index(json.load(stream)) + + +def iter_index_entries(index: ShardIndex) -> Iterable[ShardEntry]: + return index.entries.values() diff --git a/services/tilemaxsim_triton.py b/services/tilemaxsim_triton.py new file mode 100644 index 00000000..ee0736cc --- /dev/null +++ b/services/tilemaxsim_triton.py @@ -0,0 +1,128 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2026 Hu Xinjing + +"""Fused ragged FP16 TileMaxSim over a process-owned GPU tensor arena.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _ragged_tilemaxsim_fp16_kernel( + query, + documents, + document_offsets, + document_rows, + scores, + query_rows, + dimension: tl.constexpr, + max_document_rows: tl.constexpr, + block_query: tl.constexpr, + block_document: tl.constexpr, + block_dimension: tl.constexpr, +): + document_index = tl.program_id(0).to(tl.int64) + query_block_index = tl.program_id(1).to(tl.int64) + query_indices = (query_block_index * block_query + tl.arange(0, block_query)).to( + tl.int64 + ) + query_mask = query_indices < query_rows + document_base = tl.load(document_offsets + document_index).to(tl.int64) + valid_document_rows = tl.load(document_rows + document_index) + running_max = tl.full([block_query], value=float("-inf"), dtype=tl.float32) + + for document_start in range(0, max_document_rows, block_document): + document_indices = (document_start + tl.arange(0, block_document)).to(tl.int64) + document_mask = document_indices < valid_document_rows + similarities = tl.zeros([block_query, block_document], dtype=tl.float32) + for dimension_start in range(0, dimension, block_dimension): + dimension_indices = (dimension_start + tl.arange(0, block_dimension)).to( + tl.int64 + ) + dimension_mask = dimension_indices < dimension + query_pointers = ( + query + query_indices[:, None] * dimension + dimension_indices[None, :] + ) + query_tile = tl.load( + query_pointers, + mask=query_mask[:, None] & dimension_mask[None, :], + other=0.0, + ) + document_pointers = ( + documents + + document_base + + document_indices[:, None] * dimension + + dimension_indices[None, :] + ) + document_tile = tl.load( + document_pointers, + mask=document_mask[:, None] & dimension_mask[None, :], + other=0.0, + ) + similarities += tl.dot(query_tile, tl.trans(document_tile)) + similarities = tl.where(document_mask[None, :], similarities, float("-inf")) + running_max = tl.maximum(running_max, tl.max(similarities, axis=1)) + + running_max = tl.where(query_mask, running_max, 0.0) + tl.atomic_add(scores + document_index, tl.sum(running_max, axis=0)) + + +def ragged_tilemaxsim_fp16( + query: torch.Tensor, + document_arena: torch.Tensor, + document_offsets: torch.Tensor, + document_rows: torch.Tensor, + maximum_document_rows: int, +) -> torch.Tensor: + if query.device.type != "cuda" or document_arena.device != query.device: + raise ValueError("query and document arena must be on the same CUDA device") + if query.dtype != torch.float16 or document_arena.dtype != torch.float16: + raise ValueError("ragged TileMaxSim currently requires FP16 tensors") + if query.ndim != 2 or document_arena.ndim != 1: + raise ValueError("invalid query or document arena shape") + if document_offsets.dtype != torch.int64 or document_rows.dtype != torch.int32: + raise ValueError("invalid ragged TileMaxSim metadata dtype") + if document_offsets.shape != document_rows.shape: + raise ValueError("document offsets and rows must have the same shape") + count = document_offsets.numel() + if count == 0: + return torch.empty(0, dtype=torch.float32, device=query.device) + if maximum_document_rows <= 0: + raise ValueError("maximum document rows must be positive") + block_query = 32 + block_document = 32 + block_dimension = 128 + padded_document_rows = ( + (maximum_document_rows + block_document - 1) // block_document * block_document + ) + query_blocks = (query.shape[0] + block_query - 1) // block_query + scores = torch.zeros(count, dtype=torch.float32, device=query.device) + with torch.cuda.device(query.device): + _ragged_tilemaxsim_fp16_kernel[(count, query_blocks)]( + query, + document_arena, + document_offsets, + document_rows, + scores, + query.shape[0], + dimension=query.shape[1], + max_document_rows=padded_document_rows, + block_query=block_query, + block_document=block_document, + block_dimension=block_dimension, + ) + return scores diff --git a/services/tilemaxsimd/Cargo.lock b/services/tilemaxsimd/Cargo.lock new file mode 100644 index 00000000..d55f046d --- /dev/null +++ b/services/tilemaxsimd/Cargo.lock @@ -0,0 +1,416 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cc" +version = "1.2.60" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom", + "libc", +] + +[[package]] +name = "libc" +version = "0.2.185" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f" + +[[package]] +name = "memchr" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tilemaxsimd" +version = "0.1.0" +dependencies = [ + "anyhow", + "cc", + "clap", + "hex", + "libc", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "typenum" +version = "1.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40ce102ab67701b8526c123c1bab5cbe42d7040ccfd0f64af1a385808d2f43de" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasip2" +version = "1.0.3+wasi-0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/services/tilemaxsimd/Cargo.toml b/services/tilemaxsimd/Cargo.toml new file mode 100644 index 00000000..30c5d82a --- /dev/null +++ b/services/tilemaxsimd/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "tilemaxsimd" +version = "0.1.0" +edition = "2024" +publish = false + +[dependencies] +anyhow = "1.0" +clap = { version = "4.5", features = ["derive"] } +hex = "0.4" +libc = "0.2" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +sha2 = "0.10" + +[build-dependencies] +cc = { version = "1.2", features = ["parallel"] } + +[workspace] diff --git a/services/tilemaxsimd/build.rs b/services/tilemaxsimd/build.rs new file mode 100644 index 00000000..15cbfae6 --- /dev/null +++ b/services/tilemaxsimd/build.rs @@ -0,0 +1,12 @@ +fn main() { + println!("cargo:rerun-if-changed=native/tilemaxsim_cuda.cu"); + cc::Build::new() + .cuda(true) + .flag("-O3") + .flag("--use_fast_math") + .flag("-lineinfo") + .file("native/tilemaxsim_cuda.cu") + .compile("tilemaxsim_cuda"); + println!("cargo:rustc-link-lib=cudart"); + println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64"); +} diff --git a/services/tilemaxsimd/native/tilemaxsim_cuda.cu b/services/tilemaxsimd/native/tilemaxsim_cuda.cu new file mode 100644 index 00000000..ff214e6e --- /dev/null +++ b/services/tilemaxsimd/native/tilemaxsim_cuda.cu @@ -0,0 +1,291 @@ +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +struct VctmGpu { + int device; + unsigned char *allocation; + size_t total_bytes; + size_t tensor_bytes; + size_t workspace_bytes; + unsigned char *host_staging; + size_t host_staging_bytes; + cudaStream_t upload_stream; + cudaStream_t compute_stream; +}; + +static int fail(char *error, size_t capacity, const char *message) { + if (error != nullptr && capacity != 0) { + std::snprintf(error, capacity, "%s", message); + } + return 1; +} + +static int cuda_fail(char *error, size_t capacity, const char *operation, + cudaError_t status) { + if (error != nullptr && capacity != 0) { + std::snprintf(error, capacity, "%s: %s", operation, + cudaGetErrorString(status)); + } + return 1; +} + +extern "C" int vctm_gpu_create(int device, size_t total_bytes, + size_t workspace_bytes, VctmGpu **output, + char *error, size_t error_capacity) { + if (output == nullptr || total_bytes == 0 || workspace_bytes == 0 || + workspace_bytes >= total_bytes) { + return fail(error, error_capacity, "invalid GPU arena configuration"); + } + cudaError_t status = cudaSetDevice(device); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaSetDevice", status); + } + size_t free_bytes = 0; + size_t device_bytes = 0; + status = cudaMemGetInfo(&free_bytes, &device_bytes); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaMemGetInfo", status); + } + if (free_bytes < total_bytes) { + return fail(error, error_capacity, + "configured GPU memory is not currently available"); + } + auto *gpu = new VctmGpu{}; + gpu->device = device; + gpu->total_bytes = total_bytes; + gpu->tensor_bytes = ((total_bytes - workspace_bytes) / 256) * 256; + gpu->workspace_bytes = total_bytes - gpu->tensor_bytes; + gpu->host_staging_bytes = + std::min(gpu->tensor_bytes, static_cast(64) * 1024 * 1024); + status = cudaMalloc(reinterpret_cast(&gpu->allocation), total_bytes); + if (status != cudaSuccess) { + delete gpu; + return cuda_fail(error, error_capacity, "cudaMalloc", status); + } + if ((status = cudaStreamCreateWithFlags(&gpu->upload_stream, + cudaStreamNonBlocking)) != cudaSuccess || + (status = cudaStreamCreateWithFlags(&gpu->compute_stream, + cudaStreamNonBlocking)) != cudaSuccess) { + if (gpu->upload_stream != nullptr) cudaStreamDestroy(gpu->upload_stream); + cudaFree(gpu->allocation); + delete gpu; + return cuda_fail(error, error_capacity, "cudaStreamCreate", status); + } + status = cudaHostAlloc(reinterpret_cast(&gpu->host_staging), + gpu->host_staging_bytes, cudaHostAllocPortable); + if (status != cudaSuccess) { + cudaStreamDestroy(gpu->upload_stream); + cudaStreamDestroy(gpu->compute_stream); + cudaFree(gpu->allocation); + delete gpu; + return cuda_fail(error, error_capacity, "cudaHostAlloc", status); + } + std::memset(gpu->host_staging, 0, gpu->host_staging_bytes); + *output = gpu; + return 0; +} + +extern "C" void vctm_gpu_destroy(VctmGpu *gpu) { + if (gpu == nullptr) return; + cudaSetDevice(gpu->device); + cudaStreamDestroy(gpu->upload_stream); + cudaStreamDestroy(gpu->compute_stream); + cudaFreeHost(gpu->host_staging); + cudaFree(gpu->allocation); + delete gpu; +} + +extern "C" size_t vctm_gpu_tensor_bytes(const VctmGpu *gpu) { + return gpu == nullptr ? 0 : gpu->tensor_bytes; +} + +extern "C" int vctm_gpu_upload_batch( + VctmGpu *gpu, const uint64_t *offsets, + const unsigned char *const *payloads, const size_t *lengths, size_t count, + char *error, size_t error_capacity) { + if (gpu == nullptr || offsets == nullptr || payloads == nullptr || + lengths == nullptr) { + return fail(error, error_capacity, "invalid upload batch"); + } + cudaError_t status = cudaSetDevice(gpu->device); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaSetDevice", status); + } + for (size_t i = 0; i < count; ++i) { + if (offsets[i] > gpu->tensor_bytes || + lengths[i] > gpu->tensor_bytes - offsets[i]) { + return fail(error, error_capacity, "upload is outside the tensor arena"); + } + } + size_t item = 0; + size_t item_offset = 0; + while (item < count) { + size_t staging_offset = 0; + while (item < count && staging_offset < gpu->host_staging_bytes) { + const size_t remaining = lengths[item] - item_offset; + const size_t chunk = + std::min(remaining, gpu->host_staging_bytes - staging_offset); + std::memcpy(gpu->host_staging + staging_offset, + payloads[item] + item_offset, chunk); + status = cudaMemcpyAsync(gpu->allocation + offsets[item] + item_offset, + gpu->host_staging + staging_offset, chunk, + cudaMemcpyHostToDevice, gpu->upload_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaMemcpyAsync(H2D)", status); + } + staging_offset += chunk; + item_offset += chunk; + if (item_offset == lengths[item]) { + item += 1; + item_offset = 0; + } + } + status = cudaStreamSynchronize(gpu->upload_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaStreamSynchronize(upload)", + status); + } + } + return 0; +} + +template +__device__ float scalar_to_float(Scalar value); + +template <> +__device__ float scalar_to_float(half value) { + return __half2float(value); +} + +template <> +__device__ float scalar_to_float(float value) { + return value; +} + +template +__global__ void tilemaxsim_kernel(const Scalar *query, uint32_t query_rows, + uint32_t dimension, + const unsigned char *documents, + const uint64_t *document_offsets, + const uint32_t *document_rows, float *scores) { + const uint32_t candidate = blockIdx.x; + const uint32_t query_row = blockIdx.y; + const uint32_t lane = threadIdx.x & 31; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t warps = blockDim.x >> 5; + const auto *document = reinterpret_cast( + documents + document_offsets[candidate]); + const Scalar *query_vector = query + static_cast(query_row) * dimension; + float best = -CUDART_INF_F; + for (uint32_t row = warp; row < document_rows[candidate]; row += warps) { + const Scalar *document_vector = + document + static_cast(row) * dimension; + float dot = 0.0f; + for (uint32_t index = lane; index < dimension; index += 32) { + dot = fmaf(scalar_to_float(query_vector[index]), + scalar_to_float(document_vector[index]), dot); + } + for (int delta = 16; delta != 0; delta >>= 1) { + dot += __shfl_down_sync(0xffffffff, dot, delta); + } + if (lane == 0) best = fmaxf(best, dot); + } + __shared__ float warp_best[8]; + if (lane == 0) warp_best[warp] = best; + __syncthreads(); + if (threadIdx.x == 0) { + float maximum = -CUDART_INF_F; + for (uint32_t index = 0; index < warps; ++index) { + maximum = fmaxf(maximum, warp_best[index]); + } + atomicAdd(scores + candidate, maximum); + } +} + +static size_t aligned(size_t value, size_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +extern "C" int vctm_gpu_score( + VctmGpu *gpu, const unsigned char *query, size_t query_bytes, + uint32_t query_rows, uint32_t dimension, uint8_t dtype, + const uint64_t *document_offsets, const uint32_t *document_rows, + size_t count, float *output, char *error, size_t error_capacity) { + if (gpu == nullptr || query == nullptr || document_offsets == nullptr || + document_rows == nullptr || output == nullptr || query_rows == 0 || + dimension == 0 || count == 0) { + return fail(error, error_capacity, "invalid TileMaxSim score request"); + } + cudaError_t status = cudaSetDevice(gpu->device); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaSetDevice", status); + } + unsigned char *workspace = gpu->allocation + gpu->tensor_bytes; + size_t cursor = 0; + const size_t query_offset = cursor; + cursor = aligned(cursor + query_bytes, 256); + const size_t offsets_offset = cursor; + cursor = aligned(cursor + count * sizeof(uint64_t), 256); + const size_t rows_offset = cursor; + cursor = aligned(cursor + count * sizeof(uint32_t), 256); + const size_t scores_offset = cursor; + cursor = aligned(cursor + count * sizeof(float), 256); + if (cursor > gpu->workspace_bytes) { + return fail(error, error_capacity, + "TileMaxSim request exceeds the configured GPU workspace"); + } + status = cudaMemcpyAsync(workspace + query_offset, query, query_bytes, + cudaMemcpyHostToDevice, gpu->compute_stream); + if (status == cudaSuccess) + status = cudaMemcpyAsync(workspace + offsets_offset, document_offsets, + count * sizeof(uint64_t), cudaMemcpyHostToDevice, + gpu->compute_stream); + if (status == cudaSuccess) + status = cudaMemcpyAsync(workspace + rows_offset, document_rows, + count * sizeof(uint32_t), cudaMemcpyHostToDevice, + gpu->compute_stream); + if (status == cudaSuccess) + status = cudaMemsetAsync(workspace + scores_offset, 0, + count * sizeof(float), gpu->compute_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "CUDA workspace initialization", + status); + } + dim3 grid(static_cast(count), query_rows); + dim3 block(256); + if (dtype == 2) { + tilemaxsim_kernel<<compute_stream>>>( + reinterpret_cast(workspace + query_offset), query_rows, + dimension, gpu->allocation, + reinterpret_cast(workspace + offsets_offset), + reinterpret_cast(workspace + rows_offset), + reinterpret_cast(workspace + scores_offset)); + } else if (dtype == 1) { + tilemaxsim_kernel<<compute_stream>>>( + reinterpret_cast(workspace + query_offset), query_rows, + dimension, gpu->allocation, + reinterpret_cast(workspace + offsets_offset), + reinterpret_cast(workspace + rows_offset), + reinterpret_cast(workspace + scores_offset)); + } else { + return fail(error, error_capacity, "unsupported tensor dtype"); + } + status = cudaGetLastError(); + if (status == cudaSuccess) + status = cudaMemcpyAsync(output, workspace + scores_offset, + count * sizeof(float), cudaMemcpyDeviceToHost, + gpu->compute_stream); + if (status == cudaSuccess) status = cudaStreamSynchronize(gpu->compute_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "TileMaxSim CUDA execution", status); + } + return 0; +} diff --git a/services/tilemaxsimd/src/cache.rs b/services/tilemaxsimd/src/cache.rs new file mode 100644 index 00000000..0ab1d5ac --- /dev/null +++ b/services/tilemaxsimd/src/cache.rs @@ -0,0 +1,762 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::hash::{Hash, Hasher}; + +#[derive(Debug)] +pub struct PageRunAllocator { + block_bytes: usize, + block_count: usize, + free_by_start: BTreeMap, + free_by_size: BTreeMap>, + allocated: HashMap, + free_blocks: usize, +} + +impl PageRunAllocator { + pub fn new(capacity: usize, block_bytes: usize) -> Result { + if capacity == 0 || block_bytes == 0 || !block_bytes.is_multiple_of(256) { + return Err("invalid fixed-block arena"); + } + let block_count = capacity / block_bytes; + if block_count == 0 { + return Err("fixed-block arena is too small"); + } + let mut allocator = Self { + block_bytes, + block_count, + free_by_start: BTreeMap::new(), + free_by_size: BTreeMap::new(), + allocated: HashMap::new(), + free_blocks: 0, + }; + allocator.add_free_run(0, block_count)?; + Ok(allocator) + } + + fn add_free_run(&mut self, start: usize, blocks: usize) -> Result<(), &'static str> { + if blocks == 0 + || start + .checked_add(blocks) + .is_none_or(|end| end > self.block_count) + { + return Err("free GPU page run is outside the arena"); + } + if self.free_by_start.insert(start, blocks).is_some() { + return Err("duplicate free GPU page run"); + } + self.free_by_size.entry(blocks).or_default().insert(start); + self.free_blocks = self + .free_blocks + .checked_add(blocks) + .ok_or("free GPU page count overflow")?; + Ok(()) + } + + fn remove_free_run(&mut self, start: usize, blocks: usize) -> Result<(), &'static str> { + if self.free_by_start.remove(&start) != Some(blocks) { + return Err("free GPU address index is inconsistent"); + } + let remove_size = { + let starts = self + .free_by_size + .get_mut(&blocks) + .ok_or("free GPU size index is inconsistent")?; + if !starts.remove(&start) { + return Err("free GPU size index is inconsistent"); + } + starts.is_empty() + }; + if remove_size { + self.free_by_size.remove(&blocks); + } + self.free_blocks = self + .free_blocks + .checked_sub(blocks) + .ok_or("free GPU page count underflow")?; + Ok(()) + } + + pub fn capacity(&self) -> usize { + self.block_count * self.block_bytes + } + + pub fn block_bytes(&self) -> usize { + self.block_bytes + } + + pub fn allocation_bytes(&self, payload_bytes: usize) -> Option { + if payload_bytes == 0 { + return None; + } + payload_bytes + .div_ceil(self.block_bytes) + .checked_mul(self.block_bytes) + } + + pub fn largest_free(&self) -> usize { + self.free_by_size + .last_key_value() + .map_or(0, |(blocks, _)| blocks * self.block_bytes) + } + + pub fn allocate(&mut self, payload_bytes: usize) -> Option<(usize, usize)> { + let allocation_bytes = self.allocation_bytes(payload_bytes)?; + let required = allocation_bytes / self.block_bytes; + let (&available, starts) = self.free_by_size.range(required..).next()?; + let start = *starts.first()?; + self.remove_free_run(start, available).ok()?; + if available > required { + self.add_free_run(start + required, available - required) + .ok()?; + } + self.allocated.insert(start, required); + Some((start * self.block_bytes, allocation_bytes)) + } + + pub fn release(&mut self, offset: usize) -> Result<(), &'static str> { + if !offset.is_multiple_of(self.block_bytes) { + return Err("unaligned fixed-block release"); + } + let mut start = offset / self.block_bytes; + let mut blocks = self + .allocated + .remove(&start) + .ok_or("fixed-block page run was released twice")?; + let previous = self + .free_by_start + .range(..start) + .next_back() + .map(|(&previous_start, &previous_blocks)| (previous_start, previous_blocks)); + if let Some((previous_start, previous_blocks)) = previous + && previous_start + previous_blocks == start + { + self.remove_free_run(previous_start, previous_blocks)?; + start = previous_start; + blocks += previous_blocks; + } + let next_start = start + blocks; + if let Some(&next_blocks) = self.free_by_start.get(&next_start) { + self.remove_free_run(next_start, next_blocks)?; + blocks += next_blocks; + } + self.add_free_run(start, blocks) + } + + pub fn free_bytes(&self) -> usize { + self.free_blocks * self.block_bytes + } + + pub fn allocated_bytes(&self) -> usize { + self.allocated + .values() + .map(|blocks| blocks * self.block_bytes) + .sum() + } + + #[cfg(test)] + pub fn validate(&self) -> Result<(), &'static str> { + let address_blocks = self.free_by_start.values().sum::(); + let size_blocks = self + .free_by_size + .iter() + .map(|(blocks, starts)| blocks * starts.len()) + .sum::(); + if address_blocks != self.free_blocks || size_blocks != self.free_blocks { + return Err("free GPU page accounting is inconsistent"); + } + for (&start, &blocks) in &self.free_by_start { + if !self + .free_by_size + .get(&blocks) + .is_some_and(|starts| starts.contains(&start)) + { + return Err("free GPU page indexes disagree"); + } + } + Ok(()) + } +} + +#[derive(Debug)] +pub struct TinyLfu { + width: usize, + tables: Vec>, + samples: usize, +} + +impl TinyLfu { + pub fn new(width: usize) -> Self { + Self { + width, + tables: vec![vec![0; width]; 4], + samples: 0, + } + } + + fn indices(&self, key: &T) -> [usize; 4] { + std::array::from_fn(|row| { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + row.hash(&mut hasher); + key.hash(&mut hasher); + hasher.finish() as usize % self.width + }) + } + + pub fn increment(&mut self, key: &T) -> u16 { + let indices = self.indices(key); + for (row, index) in indices.into_iter().enumerate() { + self.tables[row][index] = self.tables[row][index].saturating_add(1); + } + self.samples += 1; + let estimate = self.estimate(key).max(1); + if self.samples >= self.width * 10 { + for table in &mut self.tables { + for value in table { + *value /= 2; + } + } + self.samples /= 2; + } + estimate + } + + pub fn estimate(&self, key: &T) -> u16 { + self.indices(key) + .into_iter() + .enumerate() + .map(|(row, index)| self.tables[row][index]) + .min() + .unwrap_or(0) + } +} + +#[derive(Clone, Debug)] +pub struct CacheEntry { + pub offset: u64, + pub allocated_bytes: usize, + pub payload_bytes: usize, + pub rows: u32, + pub dimension: u32, + pub dtype: u8, + pub references: usize, + pub pinned: bool, + /// False while an H2D upload is in progress. Loading entries reserve pages + /// but are never visible as cache hits. + pub ready: bool, + owner_tenant: String, + priority: f64, +} + +#[derive(Debug, PartialEq)] +pub enum Admission { + Admitted { offset: u64, allocated_bytes: usize }, + Rejected, + Deferred, +} + +pub struct GpuCache { + allocator: PageRunAllocator, + entries: HashMap, + sketch: TinyLfu, + inflation: f64, + tenant_allocated: HashMap, + tenant_reservations: HashMap, + default_tenant_max_bytes: usize, + pinned_max_bytes: usize, + pinned_bytes: usize, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, +} + +impl GpuCache { + #[cfg(test)] + pub fn new(capacity: usize, block_bytes: usize) -> Result { + Self::new_with_limits(capacity, block_bytes, 100, 100, HashMap::new()) + } + + pub fn new_with_limits( + capacity: usize, + block_bytes: usize, + tenant_max_percent: u8, + pinned_max_percent: u8, + tenant_reservations: HashMap, + ) -> Result { + if tenant_max_percent == 0 + || tenant_max_percent > 100 + || pinned_max_percent > 100 + || tenant_reservations.values().sum::() > capacity + { + return Err("invalid GPU tenant cache policy"); + } + Ok(Self { + allocator: PageRunAllocator::new(capacity, block_bytes)?, + entries: HashMap::new(), + sketch: TinyLfu::new(4096), + inflation: 0.0, + tenant_allocated: HashMap::new(), + tenant_reservations, + default_tenant_max_bytes: capacity * tenant_max_percent as usize / 100, + pinned_max_bytes: capacity * pinned_max_percent as usize / 100, + pinned_bytes: 0, + hits: 0, + misses: 0, + evictions: 0, + admission_rejections: 0, + }) + } + + pub fn capacity(&self) -> usize { + self.allocator.capacity() + } + + pub fn block_bytes(&self) -> usize { + self.allocator.block_bytes() + } + + pub fn entry_count(&self) -> usize { + self.entries.len() + } + + pub fn pinned_entries(&self) -> usize { + self.entries.values().filter(|entry| entry.pinned).count() + } + + pub fn tenant_count(&self) -> usize { + self.tenant_allocated.len() + } + + pub fn pinned_bytes(&self) -> usize { + self.pinned_bytes + } + + pub fn free_bytes(&self) -> usize { + self.allocator.free_bytes() + } + + pub fn largest_free_extent(&self) -> usize { + self.allocator.largest_free() + } + + pub fn allocated_bytes(&self) -> usize { + self.allocator.allocated_bytes() + } + + pub fn payload_bytes(&self) -> usize { + self.entries.values().map(|entry| entry.payload_bytes).sum() + } + + pub fn get(&mut self, key: &str) -> Option { + let frequency = self.sketch.increment(&key); + let Some(entry) = self.entries.get_mut(key).filter(|entry| entry.ready) else { + self.misses += 1; + return None; + }; + entry.references += 1; + entry.priority = self.inflation + f64::from(frequency) / entry.allocated_bytes as f64; + self.hits += 1; + Some(entry.clone()) + } + + pub fn acquire_existing(&mut self, key: &str) -> Option { + let entry = self.entries.get_mut(key).filter(|entry| entry.ready)?; + entry.references += 1; + Some(entry.clone()) + } + + pub fn record_access_miss(&mut self, key: &str) { + self.sketch.increment(&key); + self.misses += 1; + } + + pub fn contains(&self, key: &str) -> bool { + self.entries.get(key).is_some_and(|entry| entry.ready) + } + + fn victim(&self) -> Option<(&String, &CacheEntry)> { + self.entries + .iter() + .filter(|(_, entry)| self.can_evict(entry)) + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + } + + fn tenant_victim(&self, tenant: &str) -> Option<(&String, &CacheEntry)> { + self.entries + .iter() + .filter(|(_, entry)| { + entry.owner_tenant == tenant && entry.references == 0 && !entry.pinned + }) + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + } + + fn can_evict(&self, entry: &CacheEntry) -> bool { + if entry.references != 0 || entry.pinned { + return false; + } + let usage = self + .tenant_allocated + .get(&entry.owner_tenant) + .copied() + .unwrap_or(0); + let reservation = self + .tenant_reservations + .get(&entry.owner_tenant) + .copied() + .unwrap_or(0); + usage.saturating_sub(entry.allocated_bytes) >= reservation + } + + fn remove_entry(&mut self, key: &str) -> Result { + let entry = self + .entries + .remove(key) + .ok_or("GPU cache entry disappeared")?; + self.allocator.release(entry.offset as usize)?; + if entry.pinned { + self.pinned_bytes = self.pinned_bytes.saturating_sub(entry.allocated_bytes); + } + let remove_tenant = if let Some(usage) = self.tenant_allocated.get_mut(&entry.owner_tenant) + { + *usage = usage.saturating_sub(entry.allocated_bytes); + *usage == 0 + } else { + false + }; + if remove_tenant { + self.tenant_allocated.remove(&entry.owner_tenant); + } + Ok(entry) + } + + fn evict_one(&mut self) -> bool { + let Some(key) = self.victim().map(|(key, _)| key.clone()) else { + return false; + }; + let entry = self.remove_entry(&key).expect("victim disappeared"); + self.inflation = self.inflation.max(entry.priority); + self.evictions += 1; + true + } + + #[allow(clippy::too_many_arguments)] + #[cfg(test)] + pub fn admit( + &mut self, + key: String, + payload_bytes: usize, + rows: u32, + dimension: u32, + dtype: u8, + pinned: bool, + force: bool, + ) -> Admission { + self.admit_for_tenant( + "__default__", + key, + payload_bytes, + rows, + dimension, + dtype, + pinned, + force, + ) + } + + #[allow(clippy::too_many_arguments)] + pub fn admit_for_tenant( + &mut self, + tenant: &str, + key: String, + payload_bytes: usize, + rows: u32, + dimension: u32, + dtype: u8, + pinned: bool, + force: bool, + ) -> Admission { + let Some(allocation_bytes) = self.allocator.allocation_bytes(payload_bytes) else { + return Admission::Deferred; + }; + if allocation_bytes > self.allocator.capacity() { + return Admission::Deferred; + } + if pinned && self.pinned_bytes.saturating_add(allocation_bytes) > self.pinned_max_bytes { + return Admission::Deferred; + } + if !pinned { + let tenant_max_bytes = self + .tenant_reservations + .get(tenant) + .copied() + .unwrap_or(0) + .max(self.default_tenant_max_bytes); + while self + .tenant_allocated + .get(tenant) + .copied() + .unwrap_or(0) + .saturating_add(allocation_bytes) + > tenant_max_bytes + { + let Some((victim_key, victim_priority)) = self + .tenant_victim(tenant) + .map(|(victim_key, victim)| (victim_key.clone(), victim.priority)) + else { + return Admission::Deferred; + }; + let candidate = self.inflation + + f64::from(self.sketch.estimate(&key).max(1)) / allocation_bytes as f64; + if !force && candidate <= victim_priority { + self.admission_rejections += 1; + return Admission::Rejected; + } + let entry = self + .remove_entry(&victim_key) + .expect("tenant victim disappeared"); + self.inflation = self.inflation.max(entry.priority); + self.evictions += 1; + } + } + while self.allocator.largest_free() < allocation_bytes { + let Some((_, victim)) = self.victim() else { + return Admission::Deferred; + }; + let candidate = self.inflation + + f64::from(self.sketch.estimate(&key).max(1)) / allocation_bytes as f64; + if !force && !pinned && candidate <= victim.priority { + self.admission_rejections += 1; + return Admission::Rejected; + } + if !self.evict_one() { + return Admission::Deferred; + } + } + let Some((offset, allocated_bytes)) = self.allocator.allocate(payload_bytes) else { + return Admission::Deferred; + }; + let priority = + self.inflation + f64::from(self.sketch.estimate(&key).max(1)) / allocated_bytes as f64; + self.entries.insert( + key, + CacheEntry { + offset: offset as u64, + allocated_bytes, + payload_bytes, + rows, + dimension, + dtype, + references: 1, + pinned, + ready: false, + owner_tenant: tenant.to_owned(), + priority, + }, + ); + *self.tenant_allocated.entry(tenant.to_owned()).or_default() += allocated_bytes; + if pinned { + self.pinned_bytes += allocated_bytes; + } + Admission::Admitted { + offset: offset as u64, + allocated_bytes, + } + } + + pub fn release(&mut self, key: &str) -> Result<(), &'static str> { + let entry = self + .entries + .get_mut(key) + .ok_or("GPU cache entry disappeared")?; + if entry.references == 0 { + return Err("GPU cache entry was released twice"); + } + entry.references -= 1; + Ok(()) + } + + pub fn mark_ready(&mut self, key: &str) -> Result<(), &'static str> { + let entry = self + .entries + .get_mut(key) + .ok_or("GPU cache entry disappeared before upload completed")?; + entry.ready = true; + Ok(()) + } + + #[cfg(test)] + pub fn entry(&self, key: &str) -> Option<&CacheEntry> { + self.entries.get(key) + } + + pub fn remove(&mut self, key: &str) -> Result<(), &'static str> { + self.remove_entry(key).map(|_| ()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn next_test_random(state: &mut u64) -> u64 { + // Deterministic xorshift64 sequence for allocator churn tests. This is + // test data, not a pointer, address, secret, or production RNG. + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + *state + } + + #[test] + fn page_runs_use_exact_blocks_and_coalesce() { + let mut allocator = PageRunAllocator::new(8 * 256, 256).unwrap(); + let first = allocator.allocate(300).unwrap(); + let second = allocator.allocate(300).unwrap(); + let third = allocator.allocate(700).unwrap(); + assert_eq!(first, (0, 512)); + assert_eq!(second, (512, 512)); + assert_eq!(third, (1024, 768)); + assert_eq!(allocator.free_bytes(), 256); + allocator.release(second.0).unwrap(); + allocator.release(first.0).unwrap(); + assert_eq!(allocator.largest_free(), 1024); + allocator.release(third.0).unwrap(); + assert_eq!(allocator.largest_free(), 2048); + allocator.validate().unwrap(); + } + + #[test] + fn page_run_indexes_stay_consistent_under_churn() { + let mut allocator = PageRunAllocator::new(64 * 256, 256).unwrap(); + let mut live = Vec::new(); + let mut state = 0x5eed_u64; + for _ in 0..20_000 { + let random = next_test_random(&mut state); + if !live.is_empty() && random & 3 == 0 { + let index = random as usize % live.len(); + let (offset, _) = live.swap_remove(index); + allocator.release(offset).unwrap(); + } else { + let payload = ((random >> 16) as usize % (12 * 256)) + 1; + if let Some(allocation) = allocator.allocate(payload) { + live.push(allocation); + } + } + allocator.validate().unwrap(); + assert_eq!( + allocator.free_bytes() + allocator.allocated_bytes(), + allocator.capacity() + ); + } + } + + #[test] + fn tinylfu_protects_hot_entry() { + let mut cache = GpuCache::new(256, 256).unwrap(); + cache.record_access_miss("hot"); + assert!(matches!( + cache.admit("hot".into(), 64, 1, 32, 2, false, false), + Admission::Admitted { .. } + )); + assert!( + cache.get("hot").is_none(), + "loading pages must not be visible" + ); + cache.mark_ready("hot").unwrap(); + cache.release("hot").unwrap(); + for _ in 0..3 { + cache.get("hot").unwrap(); + cache.release("hot").unwrap(); + } + cache.record_access_miss("cold"); + assert_eq!( + cache.admit("cold".into(), 64, 1, 32, 2, false, false), + Admission::Rejected + ); + assert!(cache.entry("hot").is_some()); + } + + #[test] + fn failed_loading_entry_can_be_rolled_back_without_leaking_pages() { + let mut cache = GpuCache::new(4 * 256, 256).unwrap(); + assert!(matches!( + cache.admit("loading".into(), 300, 1, 150, 2, false, true), + Admission::Admitted { .. } + )); + assert!(!cache.contains("loading")); + assert_eq!(cache.allocated_bytes(), 512); + cache.remove("loading").unwrap(); + assert_eq!(cache.allocated_bytes(), 0); + assert_eq!(cache.free_bytes(), cache.capacity()); + } + + #[test] + fn tenant_limit_prevents_one_tenant_from_owning_the_entire_arena() { + let mut cache = GpuCache::new_with_limits(4 * 256, 256, 50, 100, HashMap::new()).unwrap(); + for key in ["a-1", "a-2"] { + assert!(matches!( + cache.admit_for_tenant("tenant-a", key.into(), 200, 1, 100, 2, false, true), + Admission::Admitted { .. } + )); + cache.mark_ready(key).unwrap(); + cache.release(key).unwrap(); + } + assert!(matches!( + cache.admit_for_tenant("tenant-b", "b-1".into(), 200, 1, 100, 2, false, true,), + Admission::Admitted { .. } + )); + assert_eq!(cache.tenant_count(), 2); + } + + #[test] + fn pinned_budget_fails_closed_before_consuming_pages() { + let mut cache = GpuCache::new_with_limits(4 * 256, 256, 100, 25, HashMap::new()).unwrap(); + assert!(matches!( + cache.admit_for_tenant( + "__resident__", + "resident".into(), + 200, + 1, + 100, + 2, + true, + true, + ), + Admission::Admitted { .. } + )); + assert_eq!(cache.pinned_bytes(), 256); + assert_eq!( + cache.admit_for_tenant( + "__resident__", + "too-much".into(), + 200, + 1, + 100, + 2, + true, + true, + ), + Admission::Deferred + ); + assert_eq!(cache.allocated_bytes(), 256); + } + + #[test] + fn tenant_can_recycle_its_own_reserved_pages() { + let mut reservations = HashMap::new(); + reservations.insert("tenant-a".to_owned(), 2 * 256); + let mut cache = GpuCache::new_with_limits(2 * 256, 256, 50, 100, reservations).unwrap(); + for key in ["a-1", "a-2"] { + assert!(matches!( + cache.admit_for_tenant("tenant-a", key.into(), 200, 1, 100, 2, false, true), + Admission::Admitted { .. } + )); + cache.mark_ready(key).unwrap(); + cache.release(key).unwrap(); + } + assert!(matches!( + cache.admit_for_tenant("tenant-a", "a-3".into(), 200, 1, 100, 2, false, true), + Admission::Admitted { .. } + )); + assert_eq!(cache.entry_count(), 2); + } +} diff --git a/services/tilemaxsimd/src/engine.rs b/services/tilemaxsimd/src/engine.rs new file mode 100644 index 00000000..d953169d --- /dev/null +++ b/services/tilemaxsimd/src/engine.rs @@ -0,0 +1,584 @@ +use crate::cache::{Admission, GpuCache}; +use crate::gpu::Gpu; +use crate::protocol::{Descriptor, Request}; +use crate::shard::{ShardStore, cache_key}; +use anyhow::{Result, anyhow, bail}; +use std::collections::HashMap; +use std::sync::Arc; + +struct MissingTensor { + candidate_index: usize, + descriptor: Descriptor, + key: String, + payload: Arc<[u8]>, +} + +struct ResidentTensor { + candidate_index: usize, + device: usize, + key: String, + offset: u64, + rows: u32, + transient: bool, + newly_admitted: bool, +} + +struct DeviceState { + gpu: Gpu, + cache: GpuCache, + h2d_batches: u64, + h2d_bytes: u64, +} + +pub struct Engine { + devices: Vec, + store: ShardStore, + next_device: usize, +} + +impl Engine { + pub fn new( + gpus: Vec, + block_bytes: usize, + store: ShardStore, + tenant_cache_max_percent: u8, + pinned_cache_max_percent: u8, + tenant_reservations: &HashMap, + ) -> Result { + if gpus.is_empty() { + bail!("at least one GPU is required"); + } + let devices = gpus + .into_iter() + .map(|gpu| { + let cache = GpuCache::new_with_limits( + gpu.tensor_bytes(), + block_bytes, + tenant_cache_max_percent, + pinned_cache_max_percent, + tenant_reservations.clone(), + ) + .map_err(|message| anyhow!(message))?; + Ok(DeviceState { + gpu, + cache, + h2d_batches: 0, + h2d_bytes: 0, + }) + }) + .collect::>>()?; + Ok(Self { + devices, + store, + next_device: 0, + }) + } + + pub fn reload_shards(&mut self) -> Result<()> { + self.store.reload() + } + + pub fn prewarm(&mut self, descriptors: &[Descriptor], batch_size: usize) -> Result<()> { + if batch_size == 0 { + bail!("resident prewarm batch size must be positive"); + } + for batch in descriptors.chunks(batch_size) { + let payloads = self.store.resolve_many(batch, "__resident__")?; + let mut uploads = (0..self.devices.len()) + .map(|_| Vec::<(u64, &[u8])>::new()) + .collect::>(); + let mut acquired = Vec::<(usize, String, bool)>::new(); + for (descriptor, payload) in batch.iter().zip(&payloads) { + let key = cache_key(descriptor); + if let Some((device, _)) = + self.devices + .iter_mut() + .enumerate() + .find_map(|(index, device)| { + device + .cache + .acquire_existing(&key) + .map(|entry| (index, entry)) + }) + { + acquired.push((device, key, false)); + continue; + } + self.devices[self.next_device] + .cache + .record_access_miss(&key); + let mut admission = None; + for step in 0..self.devices.len() { + let device = (self.next_device + step) % self.devices.len(); + if let Admission::Admitted { offset, .. } = + self.devices[device].cache.admit_for_tenant( + "__resident__", + key.clone(), + payload.len(), + descriptor.rows, + descriptor.dimension, + descriptor.dtype, + true, + true, + ) + { + admission = Some((device, offset)); + self.next_device = (device + 1) % self.devices.len(); + break; + } + } + let Some((device, offset)) = admission else { + bail!("resident manifest exceeds the configured Rust GPU block caches"); + }; + uploads[device].push((offset, payload.as_ref())); + acquired.push((device, key, true)); + } + let mut upload_succeeded = vec![true; self.devices.len()]; + let mut upload_error = None; + for (device, items) in uploads.iter().enumerate() { + if items.is_empty() { + continue; + } + match self.devices[device].gpu.upload_batch(items) { + Ok(()) => { + self.devices[device].h2d_batches += 1; + self.devices[device].h2d_bytes += items + .iter() + .map(|(_, payload)| payload.len() as u64) + .sum::(); + } + Err(error) => { + upload_succeeded[device] = false; + upload_error.get_or_insert(error); + } + } + } + for (device, key, newly_admitted) in acquired { + if newly_admitted && upload_succeeded[device] { + self.devices[device] + .cache + .mark_ready(&key) + .map_err(|message| anyhow!(message))?; + } + if newly_admitted && !upload_succeeded[device] { + self.devices[device] + .cache + .remove(&key) + .map_err(|message| anyhow!(message))?; + } else { + self.devices[device] + .cache + .release(&key) + .map_err(|message| anyhow!(message))?; + } + } + if let Some(error) = upload_error { + return Err(error); + } + } + Ok(()) + } + + pub fn score(&mut self, request: &Request) -> Result> { + if request.candidates.is_empty() { + return Ok(Vec::new()); + } + let mut scores = vec![None; request.candidates.len()]; + let mut hit_chunks = (0..self.devices.len()) + .map(|_| Vec::::new()) + .collect::>(); + let mut missing_descriptors = Vec::new(); + let mut missing_indices = Vec::new(); + for (index, descriptor) in request.candidates.iter().enumerate() { + let key = cache_key(descriptor); + let hit_device = self + .devices + .iter() + .position(|device| device.cache.contains(&key)); + if let Some(device_index) = hit_device { + let entry = self.devices[device_index] + .cache + .get(&key) + .expect("cache hit disappeared"); + validate_entry(descriptor, &entry)?; + hit_chunks[device_index].push(ResidentTensor { + candidate_index: index, + device: device_index, + key, + offset: entry.offset, + rows: entry.rows, + transient: false, + newly_admitted: false, + }); + } else { + // Record one request-level miss on the device that will get the + // first admission opportunity. Other devices are not polluted. + self.devices[self.next_device] + .cache + .record_access_miss(&key); + missing_indices.push(index); + missing_descriptors.push(descriptor.clone()); + } + } + + let hit_result = self.score_devices(request, &hit_chunks, &mut scores); + let hit_cleanup = self.release_chunks(&hit_chunks); + hit_result?; + hit_cleanup?; + + let payloads = self + .store + .resolve_many(&missing_descriptors, &request.tenant)?; + let mut pending = missing_indices + .into_iter() + .zip(missing_descriptors) + .zip(payloads) + .map(|((candidate_index, descriptor), payload)| MissingTensor { + candidate_index, + key: cache_key(&descriptor), + descriptor, + payload, + }) + .collect::>(); + + while !pending.is_empty() { + let mut chunks = (0..self.devices.len()) + .map(|_| Vec::::new()) + .collect::>(); + let mut uploads = (0..self.devices.len()) + .map(|_| Vec::<(u64, &[u8])>::new()) + .collect::>(); + let mut consumed = 0; + for tensor in &pending { + if let Some((device_index, entry)) = + self.devices + .iter_mut() + .enumerate() + .find_map(|(index, device)| { + device + .cache + .acquire_existing(&tensor.key) + .map(|entry| (index, entry)) + }) + { + chunks[device_index].push(ResidentTensor { + candidate_index: tensor.candidate_index, + device: device_index, + key: tensor.key.clone(), + offset: entry.offset, + rows: entry.rows, + transient: false, + newly_admitted: false, + }); + consumed += 1; + continue; + } + + let mut admitted = None; + for step in 0..self.devices.len() { + let device_index = (self.next_device + step) % self.devices.len(); + let admission = self.devices[device_index].cache.admit_for_tenant( + &request.tenant, + tensor.key.clone(), + tensor.payload.len(), + tensor.descriptor.rows, + tensor.descriptor.dimension, + tensor.descriptor.dtype, + false, + false, + ); + if let Admission::Admitted { offset, .. } = admission { + admitted = Some((device_index, offset, false)); + self.next_device = (device_index + 1) % self.devices.len(); + break; + } + } + + if admitted.is_none() && chunks.iter().all(Vec::is_empty) { + // TinyLFU rejected the cold item on every device, but the + // request must still be computed. Use one transient slab + // and remove it after the chunk completes. + for step in 0..self.devices.len() { + let device_index = (self.next_device + step) % self.devices.len(); + if let Admission::Admitted { offset, .. } = + self.devices[device_index].cache.admit_for_tenant( + &request.tenant, + tensor.key.clone(), + tensor.payload.len(), + tensor.descriptor.rows, + tensor.descriptor.dimension, + tensor.descriptor.dtype, + false, + true, + ) + { + admitted = Some((device_index, offset, true)); + self.next_device = (device_index + 1) % self.devices.len(); + break; + } + } + } + + let Some((device_index, offset, transient)) = admitted else { + if chunks.iter().any(|chunk| !chunk.is_empty()) { + break; + } + bail!("one tensor cannot be scheduled in any configured Rust GPU block cache"); + }; + uploads[device_index].push((offset, tensor.payload.as_ref())); + chunks[device_index].push(ResidentTensor { + candidate_index: tensor.candidate_index, + device: device_index, + key: tensor.key.clone(), + offset, + rows: tensor.descriptor.rows, + transient, + newly_admitted: true, + }); + consumed += 1; + } + + if consumed == 0 { + bail!("Rust multi-GPU scheduler made no progress"); + } + let upload_succeeded = match self.upload_devices(&chunks, &uploads) { + Ok(upload_succeeded) => upload_succeeded, + Err((error, upload_succeeded)) => { + self.cleanup_after_upload_failure(&chunks, &upload_succeeded)?; + return Err(error); + } + }; + debug_assert!(upload_succeeded.iter().all(|succeeded| *succeeded)); + let score_result = self.score_devices(request, &chunks, &mut scores); + let cleanup_result = self.release_chunks(&chunks); + score_result?; + cleanup_result?; + pending.drain(..consumed); + } + + request + .candidates + .iter() + .enumerate() + .map(|(index, descriptor)| { + scores[index] + .map(|score| (descriptor.candidate_id, score)) + .ok_or_else(|| anyhow!("missing native TileMaxSim result")) + }) + .collect() + } + + fn upload_devices( + &mut self, + chunks: &[Vec], + uploads: &[Vec<(u64, &[u8])>], + ) -> Result, (anyhow::Error, Vec)> { + let results = std::thread::scope(|scope| { + let mut workers = Vec::new(); + for (device_index, ((device, chunk), upload)) in + self.devices.iter_mut().zip(chunks).zip(uploads).enumerate() + { + if upload.is_empty() { + continue; + } + workers.push(( + device_index, + scope.spawn(move || -> Result<()> { + device.gpu.upload_batch(upload)?; + device.h2d_batches += 1; + device.h2d_bytes += upload + .iter() + .map(|(_, payload)| payload.len() as u64) + .sum::(); + // Keep the chunk borrow in this worker so upload metadata and + // payload lifetimes remain tied to the scoped thread. + let _ = chunk; + Ok(()) + }), + )); + } + workers + .into_iter() + .map(|(device, worker)| { + let result = worker + .join() + .map_err(|_| anyhow!("GPU upload worker panicked")) + .and_then(|result| result); + (device, result) + }) + .collect::>() + }); + + let mut succeeded = vec![true; self.devices.len()]; + let mut first_error = None; + for (device, result) in results { + if let Err(error) = result { + succeeded[device] = false; + first_error.get_or_insert(error); + } + } + for (device, chunk) in chunks.iter().enumerate() { + if !succeeded[device] { + continue; + } + for tensor in chunk.iter().filter(|tensor| tensor.newly_admitted) { + if let Err(message) = self.devices[device].cache.mark_ready(&tensor.key) { + succeeded[device] = false; + first_error.get_or_insert_with(|| anyhow!(message)); + break; + } + } + } + if let Some(error) = first_error { + Err((error, succeeded)) + } else { + Ok(succeeded) + } + } + + fn score_devices( + &mut self, + request: &Request, + chunks: &[Vec], + scores: &mut [Option], + ) -> Result<()> { + let completed = std::thread::scope(|scope| -> Result>> { + let mut workers = Vec::new(); + for (device, chunk) in self.devices.iter_mut().zip(chunks) { + if chunk.is_empty() { + continue; + } + workers.push(scope.spawn(move || -> Result> { + let offsets = chunk.iter().map(|item| item.offset).collect::>(); + let rows = chunk.iter().map(|item| item.rows).collect::>(); + let computed = device.gpu.score( + &request.query, + request.query_rows, + request.dimension, + request.dtype, + &offsets, + &rows, + )?; + Ok(chunk + .iter() + .zip(computed) + .map(|(tensor, score)| (tensor.candidate_index, score)) + .collect()) + })); + } + let mut completed = Vec::new(); + for worker in workers { + completed.push( + worker + .join() + .map_err(|_| anyhow!("GPU worker panicked"))??, + ); + } + Ok(completed) + })?; + for device_scores in completed { + for (candidate_index, score) in device_scores { + scores[candidate_index] = Some(score); + } + } + Ok(()) + } + + fn release_chunks(&mut self, chunks: &[Vec]) -> Result<()> { + for chunk in chunks { + for tensor in chunk { + if tensor.transient { + self.devices[tensor.device] + .cache + .remove(&tensor.key) + .map_err(|message| anyhow!(message))?; + } else { + self.devices[tensor.device] + .cache + .release(&tensor.key) + .map_err(|message| anyhow!(message))?; + } + } + } + Ok(()) + } + + fn cleanup_after_upload_failure( + &mut self, + chunks: &[Vec], + upload_succeeded: &[bool], + ) -> Result<()> { + for chunk in chunks { + for tensor in chunk { + if tensor.newly_admitted && (!upload_succeeded[tensor.device] || tensor.transient) { + self.devices[tensor.device] + .cache + .remove(&tensor.key) + .map_err(|message| anyhow!(message))?; + } else { + self.devices[tensor.device] + .cache + .release(&tensor.key) + .map_err(|message| anyhow!(message))?; + } + } + } + Ok(()) + } + + pub fn status_json(&self) -> serde_json::Value { + let (host_hits, host_misses, host_evictions, host_rejections) = self.store.host_status(); + let devices = self + .devices + .iter() + .enumerate() + .map(|(index, device)| { + serde_json::json!({ + "index": index, + "gpu_allocator": "segregated-page-runs", + "gpu_tensor_bytes": device.cache.capacity(), + "gpu_block_bytes": device.cache.block_bytes(), + "gpu_free_bytes": device.cache.free_bytes(), + "gpu_largest_free_extent_bytes": device.cache.largest_free_extent(), + "gpu_allocated_bytes": device.cache.allocated_bytes(), + "gpu_payload_bytes": device.cache.payload_bytes(), + "gpu_internal_waste_bytes": device.cache.allocated_bytes() + - device.cache.payload_bytes(), + "gpu_entries": device.cache.entry_count(), + "gpu_pinned_entries": device.cache.pinned_entries(), + "gpu_pinned_bytes": device.cache.pinned_bytes(), + "gpu_tenant_count": device.cache.tenant_count(), + "gpu_hits": device.cache.hits, + "gpu_misses": device.cache.misses, + "gpu_evictions": device.cache.evictions, + "gpu_admission_rejections": device.cache.admission_rejections, + "h2d_batches": device.h2d_batches, + "h2d_bytes": device.h2d_bytes, + }) + }) + .collect::>(); + serde_json::json!({ + "devices": devices, + "host_hits": host_hits, + "host_misses": host_misses, + "host_evictions": host_evictions, + "host_admission_rejections": host_rejections, + "batch_read_calls": self.store.batch_read_calls, + "batch_read_bytes": self.store.batch_read_bytes, + }) + } +} + +fn validate_entry(descriptor: &Descriptor, entry: &crate::cache::CacheEntry) -> Result<()> { + let scalar_bytes = if descriptor.dtype == 1 { 4 } else { 2 }; + let expected_bytes = descriptor.rows as usize * descriptor.dimension as usize * scalar_bytes; + if entry.rows != descriptor.rows + || entry.dimension != descriptor.dimension + || entry.dtype != descriptor.dtype + || entry.payload_bytes != expected_bytes + || entry.allocated_bytes < expected_bytes + { + bail!("GPU cache metadata disagrees with the tensor descriptor"); + } + Ok(()) +} diff --git a/services/tilemaxsimd/src/gpu.rs b/services/tilemaxsimd/src/gpu.rs new file mode 100644 index 00000000..2c6c2df0 --- /dev/null +++ b/services/tilemaxsimd/src/gpu.rs @@ -0,0 +1,163 @@ +use anyhow::{Result, anyhow, bail}; +use std::ffi::{CStr, c_char, c_int, c_uchar, c_void}; +use std::ptr::NonNull; + +#[repr(C)] +struct NativeGpu(c_void); + +unsafe extern "C" { + fn vctm_gpu_create( + device: c_int, + total_bytes: usize, + workspace_bytes: usize, + output: *mut *mut NativeGpu, + error: *mut c_char, + error_capacity: usize, + ) -> c_int; + fn vctm_gpu_destroy(gpu: *mut NativeGpu); + fn vctm_gpu_tensor_bytes(gpu: *const NativeGpu) -> usize; + fn vctm_gpu_upload_batch( + gpu: *mut NativeGpu, + offsets: *const u64, + payloads: *const *const c_uchar, + lengths: *const usize, + count: usize, + error: *mut c_char, + error_capacity: usize, + ) -> c_int; + fn vctm_gpu_score( + gpu: *mut NativeGpu, + query: *const c_uchar, + query_bytes: usize, + query_rows: u32, + dimension: u32, + dtype: u8, + document_offsets: *const u64, + document_rows: *const u32, + count: usize, + output: *mut f32, + error: *mut c_char, + error_capacity: usize, + ) -> c_int; +} + +pub struct Gpu { + native: NonNull, + tensor_bytes: usize, +} + +// SAFETY: `Gpu` uniquely owns the native handle. It may move to a scoped +// worker, but no method exposes the pointer and all calls require `&mut self`. +unsafe impl Send for Gpu {} + +impl Gpu { + pub fn create(device: i32, total_bytes: usize, workspace_bytes: usize) -> Result { + let mut native = std::ptr::null_mut(); + let mut error = [0_i8; 512]; + // SAFETY: the C API writes one opaque pointer and a bounded error string. + let status = unsafe { + vctm_gpu_create( + device, + total_bytes, + workspace_bytes, + &mut native, + error.as_mut_ptr(), + error.len(), + ) + }; + if status != 0 { + bail!(native_error(&error)); + } + let native = NonNull::new(native).ok_or_else(|| anyhow!("CUDA returned a null arena"))?; + // SAFETY: `native` is live until Drop. + let tensor_bytes = unsafe { vctm_gpu_tensor_bytes(native.as_ptr()) }; + Ok(Self { + native, + tensor_bytes, + }) + } + + pub fn tensor_bytes(&self) -> usize { + self.tensor_bytes + } + + pub fn upload_batch(&mut self, items: &[(u64, &[u8])]) -> Result<()> { + if items.is_empty() { + return Ok(()); + } + let offsets = items.iter().map(|item| item.0).collect::>(); + let payloads = items.iter().map(|item| item.1.as_ptr()).collect::>(); + let lengths = items.iter().map(|item| item.1.len()).collect::>(); + let mut error = [0_i8; 512]; + // SAFETY: all slices remain alive for the synchronous native batch call. + let status = unsafe { + vctm_gpu_upload_batch( + self.native.as_ptr(), + offsets.as_ptr(), + payloads.as_ptr(), + lengths.as_ptr(), + items.len(), + error.as_mut_ptr(), + error.len(), + ) + }; + if status != 0 { + bail!(native_error(&error)); + } + Ok(()) + } + + pub fn score( + &mut self, + query: &[u8], + query_rows: u32, + dimension: u32, + dtype: u8, + document_offsets: &[u64], + document_rows: &[u32], + ) -> Result> { + if document_offsets.len() != document_rows.len() || document_offsets.is_empty() { + bail!("invalid native document metadata"); + } + let mut output = vec![0.0_f32; document_offsets.len()]; + let mut error = [0_i8; 512]; + // SAFETY: the native call is synchronous and receives valid slice pointers. + let status = unsafe { + vctm_gpu_score( + self.native.as_ptr(), + query.as_ptr(), + query.len(), + query_rows, + dimension, + dtype, + document_offsets.as_ptr(), + document_rows.as_ptr(), + document_offsets.len(), + output.as_mut_ptr(), + error.as_mut_ptr(), + error.len(), + ) + }; + if status != 0 { + bail!(native_error(&error)); + } + if output.iter().any(|score| !score.is_finite()) { + bail!("native TileMaxSim returned a non-finite score"); + } + Ok(output) + } +} + +impl Drop for Gpu { + fn drop(&mut self) { + // SAFETY: this is the unique owned native pointer. + unsafe { vctm_gpu_destroy(self.native.as_ptr()) }; + } +} + +fn native_error(buffer: &[c_char]) -> String { + // SAFETY: the native helper always NUL-terminates a nonempty error buffer. + unsafe { CStr::from_ptr(buffer.as_ptr()) } + .to_string_lossy() + .into_owned() +} diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs new file mode 100644 index 00000000..8a01793d --- /dev/null +++ b/services/tilemaxsimd/src/main.rs @@ -0,0 +1,1349 @@ +mod cache; +mod engine; +mod gpu; +mod protocol; +mod scheduler; +mod shard; + +use anyhow::{Context, Result, anyhow, bail}; +use clap::Parser; +use engine::Engine; +use gpu::Gpu; +use protocol::{HEADER_BYTES, VERSION_EXTERNAL, VERSION_SCHEDULED_EXTERNAL}; +use scheduler::{RequestQueue, Scheduled, SchedulerPolicy}; +use serde::Deserialize; +use sha2::{Digest, Sha256}; +use shard::ShardStore; +use std::collections::HashMap; +use std::fs::{self, OpenOptions}; +use std::io::{BufRead, Read, Write}; +use std::os::fd::AsRawFd; +use std::os::unix::fs::{FileTypeExt, PermissionsExt}; +use std::os::unix::net::{UnixListener, UnixStream}; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, mpsc}; +use std::thread; +use std::time::{Duration, Instant}; + +const GIB: usize = 1024 * 1024 * 1024; +static SHUTDOWN_REQUESTED: AtomicBool = AtomicBool::new(false); +static RELOAD_REQUESTED: AtomicBool = AtomicBool::new(false); + +extern "C" fn handle_signal(signal: libc::c_int) { + if signal == libc::SIGHUP { + RELOAD_REQUESTED.store(true, Ordering::Relaxed); + } else { + SHUTDOWN_REQUESTED.store(true, Ordering::Relaxed); + } +} + +fn install_signal_handlers() -> Result<()> { + for signal in [libc::SIGINT, libc::SIGTERM, libc::SIGHUP] { + if unsafe { libc::signal(signal, handle_signal as *const () as libc::sighandler_t) } + == libc::SIG_ERR + { + return Err(std::io::Error::last_os_error().into()); + } + } + Ok(()) +} + +#[derive(Parser)] +#[command(about = "Native Rust/CUDA TileMaxSim shard, cache, and scheduling daemon")] +struct Args { + #[arg(long)] + socket: PathBuf, + #[arg(long, required = true, value_parser = parse_gpu_memory)] + gpu_memory_gb: Vec, + #[arg(long, default_value = "2", value_parser = parse_gb)] + gpu_workspace_gb: usize, + #[arg(long, default_value = "8", value_parser = parse_gb)] + host_cache_gb: usize, + #[arg(long, default_value_t = 80, value_parser = clap::value_parser!(u8).range(1..=100))] + host_tenant_cache_max_percent: u8, + #[arg(long = "contract-root", required = true, value_parser = parse_contract_root)] + contract_roots: Vec<(String, PathBuf)>, + #[arg(long, default_value_t = 32)] + gpu_block_kib: usize, + #[arg(long, default_value_t = 64 * 1024 * 1024)] + max_request_bytes: usize, + #[arg(long, default_value = "1", value_parser = parse_gb)] + max_inflight_request_gb: usize, + #[arg(long, default_value = "600", value_parser = parse_mode)] + socket_mode: u32, + #[arg(long)] + status_socket: Option, + #[arg(long, default_value = "600", value_parser = parse_mode)] + status_socket_mode: u32, + #[arg(long)] + once: bool, + #[arg(long)] + verify_full_shards: bool, + #[arg(long, default_value = "lru", value_parser = ["lru", "resident"])] + gpu_cache_mode: String, + #[arg(long = "resident-manifest", value_parser = parse_contract_root)] + resident_manifests: Vec<(String, PathBuf)>, + #[arg(long, default_value_t = 256)] + prewarm_batch_size: usize, + #[arg(long, default_value_t = 80, value_parser = clap::value_parser!(u8).range(1..=100))] + tenant_cache_max_percent: u8, + #[arg(long, default_value_t = 100, value_parser = clap::value_parser!(u8).range(0..=100))] + pinned_cache_max_percent: u8, + #[arg(long = "tenant-cache-reservation", value_parser = parse_tenant_memory)] + tenant_cache_reservations: Vec<(String, usize)>, + #[arg(long, default_value_t = 256)] + max_connections: usize, + #[arg(long, default_value_t = 128)] + max_queued_requests: usize, + #[arg(long, default_value_t = 16)] + max_tenant_queued_requests: usize, + #[arg(long, default_value_t = 5000)] + socket_io_timeout_ms: u64, + #[arg(long, default_value_t = 8000)] + request_timeout_ms: u64, + #[arg(long, default_value = "fair-priority", value_parser = parse_scheduler_policy)] + scheduler_policy: SchedulerPolicy, + #[arg(long, default_value_t = 1000)] + priority_aging_ms: u64, + #[arg(long, default_value_t = 200)] + priority_band: i32, + #[arg(long, default_value_t = 2)] + scheduler_batch_window_ms: u64, + #[arg(long, default_value_t = 1024)] + scheduler_quantum_candidates: usize, + #[arg(long, default_value_t = 250_000)] + scheduler_quantum_tokens: u64, + #[arg(long = "tenant-weight", value_parser = parse_tenant_weight)] + tenant_weights: Vec<(String, f64)>, +} + +#[derive(Clone)] +struct GpuMemory { + device: i32, + bytes: usize, +} + +fn parse_gpu_memory(value: &str) -> Result { + let (device, gb) = value + .strip_prefix("cuda:") + .unwrap_or(value) + .split_once('=') + .ok_or_else(|| "GPU memory must be GPU=GB, for example 1=20".to_owned())?; + let device = device + .parse::() + .map_err(|_| "GPU index must be a nonnegative integer".to_owned())?; + if device < 0 { + return Err("GPU index must be nonnegative".to_owned()); + } + let bytes = parse_gb(gb)?; + Ok(GpuMemory { device, bytes }) +} + +fn parse_gb(value: &str) -> Result { + if value.is_empty() + || !value + .bytes() + .all(|byte| byte.is_ascii_digit() || byte == b'.') + { + return Err("memory size must be a positive number of GB".to_owned()); + } + let gb = value + .parse::() + .map_err(|_| "memory size must be a positive number of GB".to_owned())?; + if !gb.is_finite() || gb <= 0.0 || gb * GIB as f64 > usize::MAX as f64 { + return Err("memory size must be a positive number of GB".to_owned()); + } + Ok((gb * GIB as f64) as usize) +} + +fn kib_to_bytes(kib: usize) -> Result { + kib.checked_mul(1024).ok_or("GPU block size overflow") +} + +fn parse_contract_root(value: &str) -> Result<(String, PathBuf), String> { + let (contract, path) = value + .split_once('=') + .ok_or_else(|| "contract root must be MODEL_CONTRACT_ID=/absolute/path".to_owned())?; + let path = PathBuf::from(path); + if contract.is_empty() || !path.is_absolute() { + return Err("contract root must contain a nonempty ID and absolute path".to_owned()); + } + Ok((contract.to_owned(), path)) +} + +fn parse_mode(value: &str) -> Result { + let mode = u32::from_str_radix(value, 8).map_err(|_| "invalid octal socket mode".to_owned())?; + if mode > 0o777 { + return Err("socket mode must be between 000 and 777".to_owned()); + } + Ok(mode) +} + +fn parse_scheduler_policy(value: &str) -> Result { + SchedulerPolicy::parse(value) +} + +fn parse_tenant_weight(value: &str) -> Result<(String, f64), String> { + let (tenant, weight) = value + .split_once('=') + .ok_or_else(|| "tenant weight must be TENANT=WEIGHT".to_owned())?; + if tenant.is_empty() + || tenant.len() > 256 + || tenant.chars().any(|character| character.is_control()) + { + return Err("tenant weight has an invalid tenant name".to_owned()); + } + let weight = weight + .parse::() + .map_err(|_| "tenant weight must be a finite positive number".to_owned())?; + if !weight.is_finite() || weight <= 0.0 || weight > 1000.0 { + return Err("tenant weight must be between 0 and 1000".to_owned()); + } + Ok((tenant.to_owned(), weight)) +} + +fn parse_tenant_memory(value: &str) -> Result<(String, usize), String> { + let (tenant, gb) = value + .split_once('=') + .ok_or_else(|| "tenant cache reservation must be TENANT=GB".to_owned())?; + if tenant.is_empty() + || tenant.len() > 256 + || tenant.chars().any(|character| character.is_control()) + { + return Err("tenant cache reservation has an invalid tenant name".to_owned()); + } + Ok((tenant.to_owned(), parse_gb(gb)?)) +} + +fn main() -> Result<()> { + let args = Args::parse(); + if args.max_connections == 0 + || args.max_queued_requests == 0 + || args.max_tenant_queued_requests == 0 + || args.socket_io_timeout_ms == 0 + || args.request_timeout_ms == 0 + || args.priority_aging_ms == 0 + || args.scheduler_quantum_candidates == 0 + || args.scheduler_quantum_tokens == 0 + { + bail!("connection, queue, timeout, and priority-aging limits must be positive"); + } + let mut seen_devices = std::collections::HashSet::new(); + for specification in &args.gpu_memory_gb { + if args.gpu_workspace_gb >= specification.bytes { + bail!("every configured GPU allocation must exceed its workspace"); + } + if !seen_devices.insert(specification.device) { + bail!("each CUDA device may be configured only once"); + } + } + let block_bytes = kib_to_bytes(args.gpu_block_kib).map_err(|message| anyhow!(message))?; + if block_bytes == 0 || block_bytes % 256 != 0 { + bail!("GPU block size must be positive and 256-byte aligned"); + } + let store = ShardStore::open( + &args.contract_roots, + args.host_cache_gb, + args.host_tenant_cache_max_percent, + args.verify_full_shards, + )?; + let gpus = args + .gpu_memory_gb + .iter() + .map(|specification| { + Gpu::create( + specification.device, + specification.bytes, + args.gpu_workspace_gb, + ) + }) + .collect::>>()?; + let tenant_cache_reservations = args.tenant_cache_reservations.iter().cloned().collect(); + let mut engine = Engine::new( + gpus, + block_bytes, + store, + args.tenant_cache_max_percent, + args.pinned_cache_max_percent, + &tenant_cache_reservations, + )?; + if args.gpu_cache_mode == "resident" && args.resident_manifests.is_empty() { + bail!("resident GPU cache mode requires at least one resident manifest"); + } + if args.gpu_cache_mode == "lru" && !args.resident_manifests.is_empty() { + bail!("resident manifests are valid only in resident GPU cache mode"); + } + if args.gpu_cache_mode == "resident" { + let prewarm_started = Instant::now(); + let descriptors = load_resident_manifests(&args.resident_manifests)?; + engine.prewarm(&descriptors, args.prewarm_batch_size)?; + println!( + "{}", + serde_json::json!({ + "event": "tilemaxsim_rust_prewarm_complete", + "entries": descriptors.len(), + "elapsed_ms": prewarm_started.elapsed().as_secs_f64() * 1000.0, + "cache": engine.status_json(), + }) + ); + } + let _instance_lock = acquire_instance_lock(&args.socket)?; + remove_stale_socket(&args.socket)?; + let listener = UnixListener::bind(&args.socket) + .with_context(|| format!("cannot bind {}", args.socket.display()))?; + fs::set_permissions(&args.socket, fs::Permissions::from_mode(args.socket_mode))?; + listener.set_nonblocking(true)?; + if args.status_socket.as_ref() == Some(&args.socket) { + bail!("status socket must differ from the TileMaxSim protocol socket"); + } + let reload = Arc::new(AtomicBool::new(false)); + let metrics = Arc::new(RuntimeMetrics::default()); + install_signal_handlers()?; + let ready_cache = engine.status_json(); + + let (sender, receiver) = mpsc::sync_channel::(args.max_queued_requests); + let frame_admission = Arc::new(ByteAdmission::new(args.max_inflight_request_gb)); + let pending_admission = Arc::new(PendingAdmission::new( + args.max_queued_requests, + args.max_tenant_queued_requests, + )); + let tenant_weights = args.tenant_weights.iter().cloned().collect(); + let scheduler_config = SchedulerConfig { + policy: args.scheduler_policy, + priority_aging: Duration::from_millis(args.priority_aging_ms), + priority_band: args.priority_band, + batch_window: Duration::from_millis(args.scheduler_batch_window_ms), + quantum_candidates: args.scheduler_quantum_candidates, + quantum_tokens: args.scheduler_quantum_tokens, + socket_io_timeout: Duration::from_millis(args.socket_io_timeout_ms), + tenant_weights, + }; + let scheduler_reload = Arc::clone(&reload); + let scheduler_metrics = Arc::clone(&metrics); + let scheduler = thread::Builder::new() + .name("tilemaxsim-scheduler".to_owned()) + .spawn(move || { + run_scheduler( + engine, + receiver, + scheduler_config, + scheduler_reload, + scheduler_metrics, + ) + })?; + let status_server = if let Some(path) = args.status_socket.clone() { + remove_stale_socket(&path)?; + let status_listener = UnixListener::bind(&path) + .with_context(|| format!("cannot bind status socket {}", path.display()))?; + fs::set_permissions(&path, fs::Permissions::from_mode(args.status_socket_mode))?; + status_listener.set_nonblocking(true)?; + let status_metrics = Arc::clone(&metrics); + Some( + thread::Builder::new() + .name("tilemaxsim-status".to_owned()) + .spawn(move || run_status_server(status_listener, path, status_metrics))?, + ) + } else { + None + }; + metrics.ready.store(true, Ordering::Release); + println!( + "{}", + serde_json::json!({ + "event": "tilemaxsim_rust_ready", + "socket": args.socket, + "devices": args.gpu_memory_gb.iter().map(|item| serde_json::json!({ + "device": item.device, + "allocated_bytes": item.bytes, + })).collect::>(), + "workspace_bytes": args.gpu_workspace_gb, + "scheduler_policy": format!("{:?}", args.scheduler_policy), + "max_connections": args.max_connections, + "max_queued_requests": args.max_queued_requests, + "max_tenant_queued_requests": args.max_tenant_queued_requests, + "max_inflight_request_bytes": args.max_inflight_request_gb, + "status_socket": args.status_socket, + "cache": ready_cache, + }) + ); + + let live_readers = Arc::new(AtomicUsize::new(0)); + let mut readers = Vec::new(); + let mut accepted = 0_usize; + let mut status_failed = false; + while !SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + if status_server + .as_ref() + .is_some_and(thread::JoinHandle::is_finished) + { + status_failed = true; + metrics.ready.store(false, Ordering::Release); + break; + } + if RELOAD_REQUESTED.swap(false, Ordering::AcqRel) { + reload.store(true, Ordering::Release); + } + reap_readers(&mut readers); + match listener.accept() { + Ok((connection, _)) => { + if !try_acquire_reader(&live_readers, args.max_connections) { + // Closing immediately is deliberate: we have not read enough + // bytes to know whether the peer expects a v2 or v3 response. + metrics.rejected_global.fetch_add(1, Ordering::Relaxed); + drop(connection); + continue; + } + accepted += 1; + let reader_sender = sender.clone(); + let reader_count = Arc::clone(&live_readers); + let reader_admission = Arc::clone(&pending_admission); + let reader_metrics = Arc::clone(&metrics); + let reader_frame_admission = Arc::clone(&frame_admission); + let maximum = args.max_request_bytes; + let io_timeout = Duration::from_millis(args.socket_io_timeout_ms); + let request_timeout = Duration::from_millis(args.request_timeout_ms); + readers.push( + thread::Builder::new() + .name("tilemaxsim-reader".to_owned()) + .spawn(move || { + let _permit = ReaderPermit(reader_count); + read_and_enqueue( + connection, + &reader_sender, + maximum, + io_timeout, + request_timeout, + reader_admission, + reader_metrics, + reader_frame_admission, + ); + })?, + ); + if args.once && accepted == 1 { + SHUTDOWN_REQUESTED.store(true, Ordering::Relaxed); + } + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(5)); + } + Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {} + Err(error) => return Err(error.into()), + } + } + drop(listener); + for reader in readers { + if reader.join().is_err() { + eprintln!("tilemaxsim reader thread panicked during shutdown"); + } + } + drop(sender); + metrics.ready.store(false, Ordering::Release); + scheduler + .join() + .map_err(|_| anyhow!("TileMaxSim scheduler thread panicked"))??; + if let Some(status_server) = status_server { + if status_server.join().is_err() { + eprintln!("TileMaxSim status thread panicked during shutdown"); + } + } + if status_failed { + bail!("TileMaxSim status server exited unexpectedly"); + } + match fs::remove_file(&args.socket) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error.into()), + } + Ok(()) +} + +struct Work { + request: protocol::Request, + connection: UnixStream, + accepted_at: Instant, + deadline: Instant, + next_candidate: usize, + results: Vec<(u32, f32)>, + gpu_elapsed: Duration, + _pending_permit: PendingPermit, + _frame_permit: BytePermit, +} + +struct SchedulerConfig { + policy: SchedulerPolicy, + priority_aging: Duration, + priority_band: i32, + batch_window: Duration, + quantum_candidates: usize, + quantum_tokens: u64, + socket_io_timeout: Duration, + tenant_weights: std::collections::HashMap, +} + +#[derive(Default)] +struct RuntimeMetrics { + ready: AtomicBool, + scheduler_depth: AtomicUsize, + gpu_active: AtomicUsize, + completed: AtomicU64, + failed: AtomicU64, + timed_out: AtomicU64, + disconnected: AtomicU64, + rejected_global: AtomicU64, + rejected_tenant: AtomicU64, +} + +struct ReaderPermit(Arc); + +impl Drop for ReaderPermit { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::Release); + } +} + +struct PendingState { + total: usize, + tenants: HashMap, +} + +struct PendingAdmission { + state: Mutex, + max_total: usize, + max_tenant: usize, +} + +struct ByteAdmission { + used: AtomicUsize, + maximum: usize, +} + +struct BytePermit { + admission: Arc, + bytes: usize, +} + +impl ByteAdmission { + fn new(maximum: usize) -> Self { + Self { + used: AtomicUsize::new(0), + maximum, + } + } + + fn try_acquire(self: &Arc, bytes: usize) -> Option { + self.used + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + current + .checked_add(bytes) + .filter(|next| *next <= self.maximum) + }) + .ok()?; + Some(BytePermit { + admission: Arc::clone(self), + bytes, + }) + } +} + +impl Drop for BytePermit { + fn drop(&mut self) { + self.admission.used.fetch_sub(self.bytes, Ordering::Release); + } +} + +#[derive(Debug)] +enum AdmissionRejection { + Global, + Tenant, +} + +struct PendingPermit { + admission: Arc, + tenant: String, +} + +impl PendingAdmission { + fn new(max_total: usize, max_tenant: usize) -> Self { + Self { + state: Mutex::new(PendingState { + total: 0, + tenants: HashMap::new(), + }), + max_total, + max_tenant, + } + } + + fn try_acquire(self: &Arc, tenant: &str) -> Result { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.total >= self.max_total { + return Err(AdmissionRejection::Global); + } + if state.tenants.get(tenant).copied().unwrap_or(0) >= self.max_tenant { + return Err(AdmissionRejection::Tenant); + } + state.total += 1; + *state.tenants.entry(tenant.to_owned()).or_default() += 1; + Ok(PendingPermit { + admission: Arc::clone(self), + tenant: tenant.to_owned(), + }) + } +} + +impl Drop for PendingPermit { + fn drop(&mut self) { + let mut state = self + .admission + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + state.total = state.total.saturating_sub(1); + if let Some(count) = state.tenants.get_mut(&self.tenant) { + *count = count.saturating_sub(1); + if *count == 0 { + state.tenants.remove(&self.tenant); + } + } + } +} + +fn try_acquire_reader(counter: &AtomicUsize, maximum: usize) -> bool { + counter + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + (current < maximum).then_some(current + 1) + }) + .is_ok() +} + +fn reap_readers(readers: &mut Vec>) { + let mut index = 0; + while index < readers.len() { + if readers[index].is_finished() { + let reader = readers.swap_remove(index); + if reader.join().is_err() { + eprintln!("tilemaxsim reader thread panicked"); + } + } else { + index += 1; + } + } +} + +fn read_and_enqueue( + mut connection: UnixStream, + sender: &mpsc::SyncSender, + maximum: usize, + io_timeout: Duration, + server_timeout: Duration, + pending_admission: Arc, + metrics: Arc, + frame_admission: Arc, +) { + let accepted_at = Instant::now(); + if let Err(error) = connection.set_read_timeout(Some(io_timeout)) { + eprintln!("cannot configure TileMaxSim socket read timeout: {error}"); + return; + } + if let Err(error) = connection.set_write_timeout(Some(io_timeout)) { + eprintln!("cannot configure TileMaxSim socket write timeout: {error}"); + return; + } + let (frame, frame_permit) = match read_request(&mut connection, maximum, &frame_admission) { + Ok(frame) => frame, + Err(error) => { + metrics.failed.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure(VERSION_EXTERNAL, 0, 1, &format!("{error:#}")), + ); + return; + } + }; + let version = header_version(&frame); + let request_id = header_request_id(&frame); + let request = match protocol::parse(&frame) { + Ok(request) => request, + Err(error) => { + metrics.failed.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure(version, request_id, 1, &format!("{error:#}")), + ); + return; + } + }; + let client_timeout = if request.timeout_ms == 0 { + server_timeout + } else { + Duration::from_millis(u64::from(request.timeout_ms)).min(server_timeout) + }; + let deadline = accepted_at + .checked_add(client_timeout) + .unwrap_or(accepted_at); + if deadline <= Instant::now() { + metrics.timed_out.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure( + version, + request_id, + 2, + "request deadline expired before enqueue", + ), + ); + return; + } + let pending_permit = match pending_admission.try_acquire(&request.tenant) { + Ok(permit) => permit, + Err(AdmissionRejection::Global) => { + metrics.rejected_global.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure(version, request_id, 2, "TileMaxSim scheduler queue is full"), + ); + return; + } + Err(AdmissionRejection::Tenant) => { + metrics.rejected_tenant.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure( + version, + request_id, + 2, + "tenant TileMaxSim queue limit exceeded", + ), + ); + return; + } + }; + match sender.try_send(Work { + request, + connection, + accepted_at, + deadline, + next_candidate: 0, + results: Vec::new(), + gpu_elapsed: Duration::ZERO, + _pending_permit: pending_permit, + _frame_permit: frame_permit, + }) { + Ok(()) => {} + Err(mpsc::TrySendError::Full(mut work)) => { + metrics.rejected_global.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut work.connection, + &protocol::failure( + work.request.protocol_version, + work.request.request_id, + 2, + "TileMaxSim scheduler queue is full", + ), + ); + } + Err(mpsc::TrySendError::Disconnected(mut work)) => { + metrics.failed.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut work.connection, + &protocol::failure( + work.request.protocol_version, + work.request.request_id, + 3, + "TileMaxSim scheduler is unavailable", + ), + ); + } + } +} + +fn run_scheduler( + mut engine: Engine, + receiver: mpsc::Receiver, + config: SchedulerConfig, + reload: Arc, + metrics: Arc, +) -> Result<()> { + let mut queue = RequestQueue::new( + config.policy, + config.priority_aging, + config.priority_band, + config.tenant_weights.clone(), + ); + let mut channel_open = true; + while channel_open || queue.len() > 0 { + if reload.swap(false, Ordering::AcqRel) { + match engine.reload_shards() { + Ok(()) => println!( + "{}", + serde_json::json!({"event": "tilemaxsim_rust_shards_reloaded"}) + ), + Err(error) => eprintln!("TileMaxSim shard reload rejected: {error:#}"), + } + } + + if queue.len() == 0 && channel_open { + match receiver.recv_timeout(Duration::from_millis(50)) { + Ok(work) => enqueue_work(&mut queue, work, &config, &metrics), + Err(mpsc::RecvTimeoutError::Timeout) => continue, + Err(mpsc::RecvTimeoutError::Disconnected) => channel_open = false, + } + } + if channel_open { + let batching_deadline = Instant::now() + config.batch_window; + loop { + let remaining = batching_deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + break; + } + match receiver.recv_timeout(remaining) { + Ok(work) => enqueue_work(&mut queue, work, &config, &metrics), + Err(mpsc::RecvTimeoutError::Timeout) => break, + Err(mpsc::RecvTimeoutError::Disconnected) => { + channel_open = false; + break; + } + } + } + } + + let now = Instant::now(); + for mut expired in queue.drain_expired(now) { + metrics.timed_out.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut expired.payload.connection, + &protocol::failure( + expired.payload.request.protocol_version, + expired.payload.request.request_id, + 2, + "request deadline expired in scheduler queue", + ), + ); + } + metrics + .scheduler_depth + .store(queue.len(), Ordering::Relaxed); + let Some(scheduled) = queue.pop(Instant::now()) else { + continue; + }; + metrics + .scheduler_depth + .store(queue.len(), Ordering::Relaxed); + if peer_disconnected(&scheduled.payload.connection) { + metrics.disconnected.fetch_add(1, Ordering::Relaxed); + continue; + } + let started = Instant::now(); + let mut work = scheduled.payload; + let request_id = work.request.request_id; + let version = work.request.protocol_version; + let tenant = work.request.tenant.clone(); + let priority = work.request.priority; + let response = if work.deadline <= started { + metrics.timed_out.fetch_add(1, Ordering::Relaxed); + Some(protocol::failure( + version, + request_id, + 2, + "request deadline expired before execution", + )) + } else { + let end = next_quantum_end(&work, &config); + let quantum = protocol::Request { + protocol_version: work.request.protocol_version, + request_id: work.request.request_id, + tenant: work.request.tenant.clone(), + priority: work.request.priority, + timeout_ms: work.request.timeout_ms, + query_rows: work.request.query_rows, + dimension: work.request.dimension, + dtype: work.request.dtype, + query: work.request.query.clone(), + candidates: work.request.candidates[work.next_candidate..end].to_vec(), + }; + let quantum_started = Instant::now(); + metrics.gpu_active.store(1, Ordering::Relaxed); + match engine.score(&quantum) { + Ok(results) => { + metrics.gpu_active.store(0, Ordering::Relaxed); + work.gpu_elapsed += quantum_started.elapsed(); + work.results.extend(results); + work.next_candidate = end; + if work.deadline <= Instant::now() { + metrics.timed_out.fetch_add(1, Ordering::Relaxed); + Some(protocol::failure( + version, + request_id, + 2, + "request deadline expired during GPU execution", + )) + } else if end < work.request.candidates.len() { + let cost = estimated_next_work(&work, &config); + queue.push(Scheduled::new( + tenant.clone(), + priority, + cost, + work.accepted_at, + work.deadline, + work, + )); + metrics + .scheduler_depth + .store(queue.len(), Ordering::Relaxed); + continue; + } else { + metrics.completed.fetch_add(1, Ordering::Relaxed); + Some(protocol::success(version, request_id, &work.results)) + } + } + Err(error) => { + metrics.gpu_active.store(0, Ordering::Relaxed); + metrics.failed.fetch_add(1, Ordering::Relaxed); + work.gpu_elapsed += quantum_started.elapsed(); + Some(protocol::failure( + version, + request_id, + 3, + &format!("{error:#}"), + )) + } + } + }; + let Some(response) = response else { + continue; + }; + work.connection + .set_write_timeout(Some(config.socket_io_timeout)) + .ok(); + write_response_nonfatal(&mut work.connection, &response); + println!( + "{}", + serde_json::json!({ + "event": "tilemaxsim_rust_request", + "request_id": request_id, + "tenant_hash": tenant_hash(&tenant), + "priority": priority, + "total_ms": work.accepted_at.elapsed().as_secs_f64() * 1000.0, + "gpu_ms": work.gpu_elapsed.as_secs_f64() * 1000.0, + "queue_ms": work.accepted_at.elapsed().saturating_sub(work.gpu_elapsed).as_secs_f64() * 1000.0, + "queue_depth": queue.len(), + "cache": engine.status_json(), + }) + ); + } + Ok(()) +} + +fn enqueue_work( + queue: &mut RequestQueue, + work: Work, + config: &SchedulerConfig, + metrics: &RuntimeMetrics, +) { + let cost = estimated_next_work(&work, config); + queue.push(Scheduled::new( + work.request.tenant.clone(), + work.request.priority, + cost, + work.accepted_at, + work.deadline, + work, + )); + metrics + .scheduler_depth + .store(queue.len(), Ordering::Relaxed); +} + +fn tenant_hash(tenant: &str) -> String { + let digest = Sha256::digest(tenant.as_bytes()); + digest[..8] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn next_quantum_end(work: &Work, config: &SchedulerConfig) -> usize { + quantum_end( + &work.request.candidates, + work.next_candidate, + config.quantum_candidates, + config.quantum_tokens, + ) +} + +fn quantum_end( + candidates: &[protocol::Descriptor], + start: usize, + maximum_candidates: usize, + maximum_tokens: u64, +) -> usize { + let mut end = start; + let mut tokens = 0_u64; + while end < candidates.len() && end - start < maximum_candidates { + let rows = u64::from(candidates[end].rows); + if end > start && tokens.saturating_add(rows) > maximum_tokens { + break; + } + tokens = tokens.saturating_add(rows); + end += 1; + } + end +} + +fn estimated_next_work(work: &Work, config: &SchedulerConfig) -> u64 { + let end = next_quantum_end(work, config); + let document_rows = work.request.candidates[work.next_candidate..end] + .iter() + .map(|candidate| u64::from(candidate.rows)) + .sum::(); + u64::from(work.request.query_rows) + .saturating_mul(document_rows) + .max(1) +} + +fn peer_disconnected(connection: &UnixStream) -> bool { + let mut byte = 0_u8; + let result = unsafe { + libc::recv( + connection.as_raw_fd(), + (&raw mut byte).cast(), + 1, + libc::MSG_PEEK | libc::MSG_DONTWAIT, + ) + }; + if result == 0 { + return true; + } + if result < 0 { + let error = std::io::Error::last_os_error(); + return !matches!( + error.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::Interrupted + ); + } + false +} + +fn write_response_nonfatal(connection: &mut UnixStream, response: &[u8]) { + if let Err(error) = connection.write_all(response) { + eprintln!("TileMaxSim response write failed without stopping daemon: {error}"); + } +} + +fn run_status_server(listener: UnixListener, path: PathBuf, metrics: Arc) { + while !SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + match listener.accept() { + Ok((mut connection, _)) => handle_status_connection(&mut connection, &metrics), + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(20)); + } + Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {} + Err(error) => { + eprintln!("TileMaxSim status socket failed: {error}"); + break; + } + } + } + drop(listener); + match fs::remove_file(&path) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => eprintln!("cannot remove TileMaxSim status socket: {error}"), + } +} + +fn handle_status_connection(connection: &mut UnixStream, metrics: &RuntimeMetrics) { + connection + .set_read_timeout(Some(Duration::from_millis(250))) + .ok(); + connection + .set_write_timeout(Some(Duration::from_millis(250))) + .ok(); + let mut request = [0_u8; 1024]; + let Ok(count) = connection.read(&mut request) else { + return; + }; + let request = String::from_utf8_lossy(&request[..count]); + let (status, content_type, body) = if request.starts_with("GET /healthz ") { + let ready = metrics.ready.load(Ordering::Acquire); + ( + if ready { + "200 OK" + } else { + "503 Service Unavailable" + }, + "application/json", + serde_json::json!({"ready": ready}).to_string(), + ) + } else if request.starts_with("GET /metrics ") { + ( + "200 OK", + "text/plain; version=0.0.4", + render_metrics(metrics), + ) + } else { + ("404 Not Found", "text/plain", "not found\n".to_owned()) + }; + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + if let Err(error) = connection.write_all(response.as_bytes()) { + eprintln!("TileMaxSim status response failed: {error}"); + } +} + +fn render_metrics(metrics: &RuntimeMetrics) -> String { + format!( + concat!( + "# HELP tilemaxsim_ready Whether the daemon is ready to accept work.\n", + "# TYPE tilemaxsim_ready gauge\n", + "tilemaxsim_ready {}\n", + "# HELP tilemaxsim_scheduler_queue_depth Requests waiting for a GPU quantum.\n", + "# TYPE tilemaxsim_scheduler_queue_depth gauge\n", + "tilemaxsim_scheduler_queue_depth {}\n", + "# HELP tilemaxsim_gpu_active Whether a CUDA quantum is executing.\n", + "# TYPE tilemaxsim_gpu_active gauge\n", + "tilemaxsim_gpu_active {}\n", + "# HELP tilemaxsim_requests_total Completed request outcomes.\n", + "# TYPE tilemaxsim_requests_total counter\n", + "tilemaxsim_requests_total{{outcome=\"completed\"}} {}\n", + "tilemaxsim_requests_total{{outcome=\"failed\"}} {}\n", + "tilemaxsim_requests_total{{outcome=\"timeout\"}} {}\n", + "tilemaxsim_requests_total{{outcome=\"disconnected\"}} {}\n", + "# HELP tilemaxsim_admission_rejections_total Admission rejections.\n", + "# TYPE tilemaxsim_admission_rejections_total counter\n", + "tilemaxsim_admission_rejections_total{{reason=\"global\"}} {}\n", + "tilemaxsim_admission_rejections_total{{reason=\"tenant\"}} {}\n", + ), + usize::from(metrics.ready.load(Ordering::Relaxed)), + metrics.scheduler_depth.load(Ordering::Relaxed), + metrics.gpu_active.load(Ordering::Relaxed), + metrics.completed.load(Ordering::Relaxed), + metrics.failed.load(Ordering::Relaxed), + metrics.timed_out.load(Ordering::Relaxed), + metrics.disconnected.load(Ordering::Relaxed), + metrics.rejected_global.load(Ordering::Relaxed), + metrics.rejected_tenant.load(Ordering::Relaxed), + ) +} + +fn read_request( + connection: &mut UnixStream, + maximum: usize, + frame_admission: &Arc, +) -> Result<(Vec, BytePermit)> { + let mut header = [0_u8; HEADER_BYTES]; + connection.read_exact(&mut header)?; + let body_bytes = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) + .context("request body does not fit this host")?; + let total = HEADER_BYTES + .checked_add(body_bytes) + .ok_or_else(|| anyhow!("request length overflow"))?; + if total > maximum { + bail!("request exceeds byte limit"); + } + let permit = frame_admission + .try_acquire(total) + .ok_or_else(|| anyhow!("in-flight TileMaxSim request byte budget is exhausted"))?; + let mut frame = Vec::with_capacity(total); + frame.extend_from_slice(&header); + frame.resize(total, 0); + connection.read_exact(&mut frame[HEADER_BYTES..])?; + Ok((frame, permit)) +} + +fn header_request_id(frame: &[u8]) -> u64 { + if frame.len() < HEADER_BYTES { + 0 + } else { + u64::from_le_bytes(frame[8..16].try_into().unwrap()) + } +} + +fn header_version(frame: &[u8]) -> u16 { + if frame.len() < HEADER_BYTES { + VERSION_EXTERNAL + } else { + match u16::from_le_bytes(frame[4..6].try_into().unwrap()) { + VERSION_SCHEDULED_EXTERNAL => VERSION_SCHEDULED_EXTERNAL, + _ => VERSION_EXTERNAL, + } + } +} + +fn acquire_instance_lock(socket: &PathBuf) -> Result { + let lock_path = PathBuf::from(format!("{}.lock", socket.display())); + let mut lock = OpenOptions::new() + .create(true) + .read(true) + .write(true) + .open(&lock_path) + .with_context(|| format!("cannot open instance lock {}", lock_path.display()))?; + if unsafe { libc::flock(lock.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) } != 0 { + bail!( + "another TileMaxSim daemon holds the instance lock {}", + lock_path.display() + ); + } + lock.set_len(0)?; + writeln!(lock, "{}", std::process::id())?; + lock.flush()?; + Ok(lock) +} + +fn remove_stale_socket(path: &PathBuf) -> Result<()> { + let Ok(metadata) = fs::symlink_metadata(path) else { + return Ok(()); + }; + if !metadata.file_type().is_socket() { + bail!("refusing to remove non-socket path {}", path.display()); + } + if UnixStream::connect(path).is_ok() { + bail!( + "another TileMaxSim daemon is already listening on {}", + path.display() + ); + } + fs::remove_file(path)?; + Ok(()) +} + +#[derive(Deserialize)] +struct ResidentRecord { + tensor_ref: String, + tensor_rows: u32, + tensor_dim: u32, + tensor_dtype: String, + tensor_checksum: String, + canonical_bytes: Option, +} + +fn load_resident_manifests(values: &[(String, PathBuf)]) -> Result> { + let mut descriptors = Vec::new(); + for (contract, path) in values { + let file = fs::File::open(path) + .with_context(|| format!("cannot open resident manifest {}", path.display()))?; + for (line_number, line) in std::io::BufReader::new(file).lines().enumerate() { + let line = line?; + let record: ResidentRecord = serde_json::from_str(&line).with_context(|| { + format!( + "invalid resident manifest {}:{}", + path.display(), + line_number + 1 + ) + })?; + let digest = record + .tensor_ref + .strip_prefix("sha256://") + .ok_or_else(|| anyhow!("resident tensor reference is not SHA-256"))?; + if record.tensor_checksum != format!("sha256:{digest}") { + bail!("resident tensor reference and checksum disagree"); + } + let (dtype, scalar_bytes) = match record.tensor_dtype.as_str() { + "float16" => (2, 2), + "float32" => (1, 4), + _ => bail!("resident manifest has an unsupported tensor dtype"), + }; + let expected_bytes = + record.tensor_rows as usize * record.tensor_dim as usize * scalar_bytes; + if record + .canonical_bytes + .is_some_and(|bytes| bytes != expected_bytes) + { + bail!("resident canonical byte length disagrees with its shape"); + } + descriptors.push(protocol::Descriptor { + candidate_id: descriptors.len() as u32, + contract: contract.clone(), + digest: digest.to_owned(), + rows: record.tensor_rows, + dimension: record.tensor_dim, + dtype, + }); + } + } + if descriptors.is_empty() { + bail!("resident manifests contain no tensor descriptors"); + } + Ok(descriptors) +} + +#[cfg(test)] +mod tests { + use super::{ + ByteAdmission, PendingAdmission, RuntimeMetrics, kib_to_bytes, quantum_end, render_metrics, + tenant_hash, + }; + use crate::protocol::Descriptor; + use std::sync::Arc; + + #[test] + fn gpu_block_kib_is_converted_once() { + assert_eq!(kib_to_bytes(32), Ok(32 * 1024)); + assert!(kib_to_bytes(usize::MAX).is_err()); + } + + #[test] + fn scheduler_quantum_always_makes_progress_and_honours_both_limits() { + let descriptor = |rows| Descriptor { + candidate_id: rows, + contract: "model".to_owned(), + digest: "a".repeat(64), + rows, + dimension: 2, + dtype: 2, + }; + let candidates = vec![descriptor(60), descriptor(60), descriptor(60)]; + assert_eq!(quantum_end(&candidates, 0, 8, 100), 1); + assert_eq!(quantum_end(&candidates, 0, 2, 1_000), 2); + assert_eq!(quantum_end(&candidates, 2, 8, 10), 3); + } + + #[test] + fn pending_admission_is_globally_and_per_tenant_bounded() { + let admission = Arc::new(PendingAdmission::new(2, 1)); + let first = admission.try_acquire("a").unwrap(); + assert!(admission.try_acquire("a").is_err()); + let second = admission.try_acquire("b").unwrap(); + assert!(admission.try_acquire("c").is_err()); + drop(first); + assert!(admission.try_acquire("a").is_ok()); + drop(second); + } + + #[test] + fn in_flight_frame_bytes_are_bounded_until_the_permit_drops() { + let admission = Arc::new(ByteAdmission::new(100)); + let first = admission.try_acquire(60).unwrap(); + assert!(admission.try_acquire(41).is_none()); + let second = admission.try_acquire(40).unwrap(); + assert!(admission.try_acquire(1).is_none()); + drop(first); + assert!(admission.try_acquire(60).is_some()); + drop(second); + } + + #[test] + fn tenant_labels_are_stable_and_do_not_expose_raw_ids() { + assert_eq!(tenant_hash("tenant-a"), tenant_hash("tenant-a")); + assert_ne!(tenant_hash("tenant-a"), tenant_hash("tenant-b")); + assert!(!tenant_hash("tenant-a").contains("tenant")); + } + + #[test] + fn prometheus_status_contains_only_bounded_scheduler_counters() { + let metrics = RuntimeMetrics::default(); + metrics + .ready + .store(true, std::sync::atomic::Ordering::Relaxed); + metrics + .completed + .store(7, std::sync::atomic::Ordering::Relaxed); + let output = render_metrics(&metrics); + assert!(output.contains("tilemaxsim_ready 1")); + assert!(output.contains("outcome=\"completed\"} 7")); + assert!(!output.contains("tenant-a")); + } +} diff --git a/services/tilemaxsimd/src/protocol.rs b/services/tilemaxsimd/src/protocol.rs new file mode 100644 index 00000000..84ccccaa --- /dev/null +++ b/services/tilemaxsimd/src/protocol.rs @@ -0,0 +1,337 @@ +use anyhow::{Result, anyhow, bail}; +use std::collections::HashSet; + +pub const HEADER_BYTES: usize = 24; +pub const VERSION_EXTERNAL: u16 = 2; +pub const VERSION_SCHEDULED_EXTERNAL: u16 = 3; +const MAGIC: &[u8; 4] = b"VCTM"; +const REQUEST_KIND: u16 = 1; +const RESPONSE_KIND: u16 = 2; + +#[derive(Clone, Debug)] +pub struct Descriptor { + pub candidate_id: u32, + pub contract: String, + pub digest: String, + pub rows: u32, + pub dimension: u32, + pub dtype: u8, +} + +#[derive(Clone, Debug)] +pub struct Request { + pub protocol_version: u16, + pub request_id: u64, + /// Logical scheduling domain. It is never used for authorization or + /// tensor lookup; the caller must already have applied hard ACL filters. + pub tenant: String, + /// Higher values run first within the scheduler's fairness policy. + pub priority: i32, + /// Client-supplied end-to-end budget. Zero is used only by legacy v2. + pub timeout_ms: u32, + pub query_rows: u32, + pub dimension: u32, + pub dtype: u8, + pub query: Vec, + pub candidates: Vec, +} + +struct Reader<'a> { + bytes: &'a [u8], + offset: usize, +} + +impl<'a> Reader<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, offset: 0 } + } + + fn take(&mut self, count: usize) -> Result<&'a [u8]> { + let end = self + .offset + .checked_add(count) + .ok_or_else(|| anyhow!("request overflow"))?; + if end > self.bytes.len() { + bail!("truncated request"); + } + let result = &self.bytes[self.offset..end]; + self.offset = end; + Ok(result) + } + + fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + + fn u16(&mut self) -> Result { + Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap())) + } + + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + + fn i32(&mut self) -> Result { + Ok(i32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + + fn text(&mut self, count: usize, maximum: usize, name: &str) -> Result { + if count == 0 || count > maximum { + bail!("invalid {name} length"); + } + let value = std::str::from_utf8(self.take(count)?)?; + if value.chars().any(|character| character.is_control()) { + bail!("{name} contains control characters"); + } + Ok(value.to_owned()) + } + + fn finish(self) -> Result<()> { + if self.offset != self.bytes.len() { + bail!("trailing request bytes"); + } + Ok(()) + } +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: u8) -> Result { + let scalar = match dtype { + 1 => 4, + 2 => 2, + _ => bail!("unsupported tensor dtype"), + }; + if rows == 0 || dimension == 0 || dimension > 60_000 { + bail!("invalid tensor shape"); + } + (rows as usize) + .checked_mul(dimension as usize) + .and_then(|value| value.checked_mul(scalar)) + .ok_or_else(|| anyhow!("tensor shape is too large")) +} + +fn validate_finite(payload: &[u8], dtype: u8) -> Result<()> { + if dtype == 1 { + for scalar in payload.chunks_exact(4) { + if !f32::from_bits(u32::from_le_bytes(scalar.try_into().unwrap())).is_finite() { + bail!("query contains a non-finite value"); + } + } + } else { + for scalar in payload.chunks_exact(2) { + let bits = u16::from_le_bytes(scalar.try_into().unwrap()); + if bits & 0x7c00 == 0x7c00 { + bail!("query contains a non-finite value"); + } + } + } + Ok(()) +} + +pub fn parse(frame: &[u8]) -> Result { + if frame.len() < HEADER_BYTES { + bail!("truncated request header"); + } + if &frame[..4] != MAGIC { + bail!("invalid protocol magic"); + } + let version = u16::from_le_bytes(frame[4..6].try_into().unwrap()); + let kind = u16::from_le_bytes(frame[6..8].try_into().unwrap()); + let request_id = u64::from_le_bytes(frame[8..16].try_into().unwrap()); + let body_bytes = u64::from_le_bytes(frame[16..24].try_into().unwrap()); + if !matches!(version, VERSION_EXTERNAL | VERSION_SCHEDULED_EXTERNAL) || kind != REQUEST_KIND { + bail!("Rust daemon requires TileMaxSim external protocol v2 or v3"); + } + if usize::try_from(body_bytes).ok() != Some(frame.len() - HEADER_BYTES) { + bail!("request body length mismatch"); + } + let mut reader = Reader::new(&frame[HEADER_BYTES..]); + let dimension = reader.u32()?; + let query_rows = reader.u32()?; + let candidate_count = reader.u32()?; + let dtype = reader.u8()?; + let scoring = reader.u8()?; + let reserved = reader.u16()?; + let contract_bytes = reader.u32()? as usize; + if scoring != 1 || reserved != 0 { + bail!("unsupported scoring function or reserved bits"); + } + if candidate_count > 65_536 { + bail!("too many candidates"); + } + let (priority, timeout_ms, tenant_bytes) = if version == VERSION_SCHEDULED_EXTERNAL { + (reader.i32()?, reader.u32()?, reader.u32()? as usize) + } else { + (0, 0, 0) + }; + if !(-100..=100).contains(&priority) { + bail!("scheduler priority must be between -100 and 100"); + } + if version == VERSION_SCHEDULED_EXTERNAL && !(1..=600_000).contains(&timeout_ms) { + bail!("scheduler timeout must be between 1 and 600000 milliseconds"); + } + let contract = reader.text(contract_bytes, 512, "model contract")?; + let tenant = if version == VERSION_SCHEDULED_EXTERNAL { + reader.text(tenant_bytes, 256, "scheduler tenant")? + } else { + "__default__".to_owned() + }; + let query_bytes = tensor_bytes(query_rows, dimension, dtype)?; + let query = reader.take(query_bytes)?.to_vec(); + validate_finite(&query, dtype)?; + let mut total_tokens = query_rows as usize; + let mut total_bytes = query_bytes; + let mut candidate_ids = HashSet::new(); + let mut candidates = Vec::with_capacity(candidate_count as usize); + for _ in 0..candidate_count { + let candidate_id = reader.u32()?; + let rows = reader.u32()?; + let reference_bytes = reader.u32()? as usize; + let checksum_bytes = reader.u32()? as usize; + if !candidate_ids.insert(candidate_id) { + bail!("duplicate candidate ID"); + } + let tensor_ref = reader.text(reference_bytes, 4096, "tensor reference")?; + let checksum = reader.text(checksum_bytes, 512, "tensor checksum")?; + let digest = tensor_ref + .strip_prefix("sha256://") + .ok_or_else(|| anyhow!("unsupported tensor reference"))?; + if digest.len() != 64 + || !digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + || checksum != format!("sha256:{digest}") + { + bail!("invalid content-addressed tensor descriptor"); + } + let bytes = tensor_bytes(rows, dimension, dtype)?; + total_tokens = total_tokens + .checked_add(rows as usize) + .ok_or_else(|| anyhow!("token overflow"))?; + total_bytes = total_bytes + .checked_add(bytes) + .ok_or_else(|| anyhow!("byte overflow"))?; + if total_tokens > 1_000_000 || total_bytes > 1024 * 1024 * 1024 { + bail!("request exceeds tensor limits"); + } + candidates.push(Descriptor { + candidate_id, + contract: contract.clone(), + digest: digest.to_owned(), + rows, + dimension, + dtype, + }); + } + reader.finish()?; + Ok(Request { + protocol_version: version, + request_id, + tenant, + priority, + timeout_ms, + query_rows, + dimension, + dtype, + query, + candidates, + }) +} + +pub fn success(version: u16, request_id: u64, results: &[(u32, f32)]) -> Vec { + let body_bytes = 8 + results.len() * 8; + let mut frame = Vec::with_capacity(HEADER_BYTES + body_bytes); + frame.extend_from_slice(MAGIC); + frame.extend_from_slice(&version.to_le_bytes()); + frame.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + frame.extend_from_slice(&request_id.to_le_bytes()); + frame.extend_from_slice(&(body_bytes as u64).to_le_bytes()); + frame.extend_from_slice(&0_u32.to_le_bytes()); + frame.extend_from_slice(&(results.len() as u32).to_le_bytes()); + for (candidate_id, score) in results { + frame.extend_from_slice(&candidate_id.to_le_bytes()); + frame.extend_from_slice(&score.to_le_bytes()); + } + frame +} + +pub fn failure(version: u16, request_id: u64, status: u32, message: &str) -> Vec { + let message = message.as_bytes(); + let message = &message[..message.len().min(64 * 1024)]; + let body_bytes = 8 + message.len(); + let mut frame = Vec::with_capacity(HEADER_BYTES + body_bytes); + frame.extend_from_slice(MAGIC); + frame.extend_from_slice(&version.to_le_bytes()); + frame.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + frame.extend_from_slice(&request_id.to_le_bytes()); + frame.extend_from_slice(&(body_bytes as u64).to_le_bytes()); + frame.extend_from_slice(&status.max(1).to_le_bytes()); + frame.extend_from_slice(&(message.len() as u32).to_le_bytes()); + frame.extend_from_slice(message); + frame +} + +#[cfg(test)] +mod tests { + use super::*; + + fn scheduled_frame(priority: i32, timeout_ms: u32, tenant: &str) -> Vec { + let contract = "model@1"; + let digest = "a".repeat(64); + let reference = format!("sha256://{digest}"); + let checksum = format!("sha256:{digest}"); + let mut body = Vec::new(); + body.extend_from_slice(&2_u32.to_le_bytes()); + body.extend_from_slice(&1_u32.to_le_bytes()); + body.extend_from_slice(&1_u32.to_le_bytes()); + body.push(2); + body.push(1); + body.extend_from_slice(&0_u16.to_le_bytes()); + body.extend_from_slice(&(contract.len() as u32).to_le_bytes()); + body.extend_from_slice(&priority.to_le_bytes()); + body.extend_from_slice(&timeout_ms.to_le_bytes()); + body.extend_from_slice(&(tenant.len() as u32).to_le_bytes()); + body.extend_from_slice(contract.as_bytes()); + body.extend_from_slice(tenant.as_bytes()); + body.extend_from_slice(&0_u16.to_le_bytes()); + body.extend_from_slice(&0_u16.to_le_bytes()); + body.extend_from_slice(&7_u32.to_le_bytes()); + body.extend_from_slice(&1_u32.to_le_bytes()); + body.extend_from_slice(&(reference.len() as u32).to_le_bytes()); + body.extend_from_slice(&(checksum.len() as u32).to_le_bytes()); + body.extend_from_slice(reference.as_bytes()); + body.extend_from_slice(checksum.as_bytes()); + let mut frame = Vec::new(); + frame.extend_from_slice(MAGIC); + frame.extend_from_slice(&VERSION_SCHEDULED_EXTERNAL.to_le_bytes()); + frame.extend_from_slice(&REQUEST_KIND.to_le_bytes()); + frame.extend_from_slice(&42_u64.to_le_bytes()); + frame.extend_from_slice(&(body.len() as u64).to_le_bytes()); + frame.extend_from_slice(&body); + frame + } + + #[test] + fn scheduled_protocol_carries_tenant_priority_and_deadline() { + let request = parse(&scheduled_frame(17, 4_000, "tenant-a")).unwrap(); + assert_eq!(request.protocol_version, VERSION_SCHEDULED_EXTERNAL); + assert_eq!(request.request_id, 42); + assert_eq!(request.tenant, "tenant-a"); + assert_eq!(request.priority, 17); + assert_eq!(request.timeout_ms, 4_000); + assert_eq!(request.candidates[0].candidate_id, 7); + } + + #[test] + fn scheduled_protocol_rejects_priority_outside_the_public_contract() { + assert!(parse(&scheduled_frame(101, 4_000, "tenant-a")).is_err()); + } + + #[test] + fn response_uses_the_request_protocol_version() { + let response = success(VERSION_SCHEDULED_EXTERNAL, 9, &[(1, 0.5)]); + assert_eq!( + u16::from_le_bytes(response[4..6].try_into().unwrap()), + VERSION_SCHEDULED_EXTERNAL + ); + } +} diff --git a/services/tilemaxsimd/src/scheduler.rs b/services/tilemaxsimd/src/scheduler.rs new file mode 100644 index 00000000..95ec0055 --- /dev/null +++ b/services/tilemaxsimd/src/scheduler.rs @@ -0,0 +1,328 @@ +use std::collections::HashMap; +use std::time::{Duration, Instant}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SchedulerPolicy { + Fair, + Priority, + FairPriority, +} + +impl SchedulerPolicy { + pub fn parse(value: &str) -> Result { + match value { + "fair" => Ok(Self::Fair), + "priority" => Ok(Self::Priority), + "fair-priority" => Ok(Self::FairPriority), + _ => Err("scheduler policy must be fair, priority, or fair-priority".to_owned()), + } + } +} + +pub struct Scheduled { + pub tenant: String, + pub priority: i32, + pub cost: u64, + pub enqueued_at: Instant, + pub deadline: Instant, + pub payload: T, + sequence: u64, +} + +impl Scheduled { + pub fn new( + tenant: String, + priority: i32, + cost: u64, + enqueued_at: Instant, + deadline: Instant, + payload: T, + ) -> Self { + Self { + tenant, + priority, + cost: cost.max(1), + enqueued_at, + deadline, + payload, + sequence: 0, + } + } +} + +/// Hierarchical request scheduler inspired by vLLM's priority queue, with a +/// tenant-fair first level that vLLM does not need in its usual single-owner +/// deployment. `fair-priority` admits the highest aged priority band and then +/// selects the tenant with the lowest normalized service time. This preserves +/// explicit urgency while preventing equal-priority noisy neighbours from +/// monopolizing the GPU. +pub struct RequestQueue { + policy: SchedulerPolicy, + aging: Duration, + priority_band: i32, + pending: Vec>, + virtual_runtime: HashMap, + weights: HashMap, + next_sequence: u64, +} + +impl RequestQueue { + pub fn new( + policy: SchedulerPolicy, + aging: Duration, + priority_band: i32, + weights: HashMap, + ) -> Self { + Self { + policy, + aging, + priority_band: priority_band.max(0), + pending: Vec::new(), + virtual_runtime: HashMap::new(), + weights, + next_sequence: 0, + } + } + + pub fn len(&self) -> usize { + self.pending.len() + } + + pub fn push(&mut self, mut item: Scheduled) { + item.sequence = self.next_sequence; + self.next_sequence = self.next_sequence.wrapping_add(1); + if self.pending.is_empty() { + // A new busy period should not inherit unbounded floating-point + // service history from a previous burst. + self.virtual_runtime.clear(); + } + let tenant_is_active = self + .pending + .iter() + .any(|pending| pending.tenant == item.tenant); + let baseline = self + .pending + .iter() + .filter_map(|pending| self.virtual_runtime.get(&pending.tenant).copied()) + .reduce(f64::min) + .unwrap_or(0.0); + let runtime = self + .virtual_runtime + .entry(item.tenant.clone()) + .or_insert(baseline); + if !tenant_is_active { + *runtime = runtime.max(baseline); + } + self.pending.push(item); + } + + pub fn drain_expired(&mut self, now: Instant) -> Vec> { + let mut expired = Vec::new(); + let mut index = 0; + while index < self.pending.len() { + if self.pending[index].deadline <= now { + expired.push(self.pending.swap_remove(index)); + } else { + index += 1; + } + } + expired.sort_by_key(|item| item.sequence); + expired + } + + pub fn pop(&mut self, now: Instant) -> Option> { + if self.pending.is_empty() { + return None; + } + let effective = self + .pending + .iter() + .map(|item| self.effective_priority(item, now)) + .collect::>(); + let highest = effective.iter().copied().max().unwrap_or(0); + let mut best = 0; + for candidate in 1..self.pending.len() { + if self.better(candidate, best, &effective, highest) { + best = candidate; + } + } + let item = self.pending.swap_remove(best); + let weight = self.weights.get(&item.tenant).copied().unwrap_or(1.0); + let runtime = self.virtual_runtime.entry(item.tenant.clone()).or_default(); + *runtime += item.cost as f64 / weight; + Some(item) + } + + fn effective_priority(&self, item: &Scheduled, now: Instant) -> i32 { + if self.aging.is_zero() { + return item.priority; + } + let waited = now.saturating_duration_since(item.enqueued_at).as_millis(); + let steps = waited / self.aging.as_millis().max(1); + item.priority + .saturating_add(i32::try_from(steps.min(200)).unwrap_or(200)) + .min(100) + } + + fn better(&self, candidate: usize, current: usize, effective: &[i32], highest: i32) -> bool { + let left = &self.pending[candidate]; + let right = &self.pending[current]; + match self.policy { + SchedulerPolicy::Fair => self.fair_cmp(left, right), + SchedulerPolicy::Priority => { + effective[candidate] > effective[current] + || (effective[candidate] == effective[current] + && left.sequence < right.sequence) + } + SchedulerPolicy::FairPriority => { + let floor = highest.saturating_sub(self.priority_band); + let left_eligible = effective[candidate] >= floor; + let right_eligible = effective[current] >= floor; + if left_eligible != right_eligible { + return left_eligible; + } + if !left_eligible { + return effective[candidate] > effective[current] + || (effective[candidate] == effective[current] + && left.sequence < right.sequence); + } + if left.tenant == right.tenant { + return effective[candidate] > effective[current] + || (effective[candidate] == effective[current] + && left.sequence < right.sequence); + } + let left_runtime = self + .virtual_runtime + .get(&left.tenant) + .copied() + .unwrap_or(0.0); + let right_runtime = self + .virtual_runtime + .get(&right.tenant) + .copied() + .unwrap_or(0.0); + left_runtime < right_runtime + || (left_runtime == right_runtime + && (effective[candidate] > effective[current] + || (effective[candidate] == effective[current] + && left.sequence < right.sequence))) + } + } + } + + fn fair_cmp(&self, left: &Scheduled, right: &Scheduled) -> bool { + let left_runtime = self + .virtual_runtime + .get(&left.tenant) + .copied() + .unwrap_or(0.0); + let right_runtime = self + .virtual_runtime + .get(&right.tenant) + .copied() + .unwrap_or(0.0); + left_runtime < right_runtime + || (left_runtime == right_runtime && left.sequence < right.sequence) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn item(tenant: &str, priority: i32, cost: u64, now: Instant, id: u32) -> Scheduled { + Scheduled::new( + tenant.to_owned(), + priority, + cost, + now, + now + Duration::from_secs(60), + id, + ) + } + + #[test] + fn strict_priority_uses_higher_priority_then_fcfs() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::Priority, + Duration::from_secs(60), + 0, + HashMap::new(), + ); + queue.push(item("a", 0, 1, now, 1)); + queue.push(item("b", 10, 1, now, 2)); + queue.push(item("c", 10, 1, now, 3)); + assert_eq!(queue.pop(now).unwrap().payload, 2); + assert_eq!(queue.pop(now).unwrap().payload, 3); + assert_eq!(queue.pop(now).unwrap().payload, 1); + } + + #[test] + fn fair_priority_alternates_equal_priority_tenants() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::FairPriority, + Duration::from_secs(60), + 0, + HashMap::new(), + ); + queue.push(item("noisy", 0, 10, now, 1)); + queue.push(item("noisy", 0, 10, now, 2)); + queue.push(item("quiet", 0, 10, now, 3)); + assert_eq!(queue.pop(now).unwrap().payload, 1); + assert_eq!(queue.pop(now).unwrap().payload, 3); + assert_eq!(queue.pop(now).unwrap().payload, 2); + } + + #[test] + fn full_priority_band_keeps_priority_without_sacrificing_tenant_fairness() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::FairPriority, + Duration::from_secs(60), + 200, + HashMap::new(), + ); + queue.push(item("noisy", 100, 10, now, 1)); + queue.push(item("noisy", 100, 10, now, 2)); + queue.push(item("quiet", -100, 10, now, 3)); + assert_eq!(queue.pop(now).unwrap().payload, 1); + assert_eq!(queue.pop(now).unwrap().payload, 3); + assert_eq!(queue.pop(now).unwrap().payload, 2); + } + + #[test] + fn aging_eventually_promotes_waiting_work() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::Priority, + Duration::from_millis(10), + 0, + HashMap::new(), + ); + queue.push(item("old", 0, 1, now, 1)); + queue.push(item("new", 5, 1, now + Duration::from_millis(60), 2)); + assert_eq!( + queue.pop(now + Duration::from_millis(60)).unwrap().payload, + 1 + ); + } + + #[test] + fn expired_requests_leave_without_consuming_service_credit() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::FairPriority, + Duration::from_secs(1), + 0, + HashMap::new(), + ); + let mut expired = item("a", 0, 1, now, 1); + expired.deadline = now; + queue.push(expired); + queue.push(item("b", 0, 1, now, 2)); + assert_eq!(queue.drain_expired(now).len(), 1); + assert_eq!(queue.pop(now).unwrap().payload, 2); + } +} diff --git a/services/tilemaxsimd/src/shard.rs b/services/tilemaxsimd/src/shard.rs new file mode 100644 index 00000000..358c25a4 --- /dev/null +++ b/services/tilemaxsimd/src/shard.rs @@ -0,0 +1,499 @@ +use crate::cache::TinyLfu; +use crate::protocol::Descriptor; +use anyhow::{Context, Result, anyhow, bail}; +use serde::Deserialize; +use sha2::{Digest, Sha256}; +use std::collections::HashMap; +use std::fs::{File, OpenOptions}; +use std::io::Read; +use std::os::unix::fs::{FileExt, OpenOptionsExt}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +const INDEX_NAME: &str = "tilemaxsim-shards-v1.json"; + +#[derive(Deserialize)] +struct RawIndex { + format: String, + version: u32, + alignment: usize, + shards: Vec, + entries: Vec, +} + +#[derive(Deserialize)] +struct RawShard { + name: String, + #[serde(rename = "bytes")] + size: usize, + checksum: String, +} + +#[derive(Deserialize)] +struct RawEntry { + digest: String, + shard: String, + offset: u64, + length: usize, + rows: u32, + dimension: u32, + dtype: String, +} + +struct ShardFile { + file: File, + size: usize, + checksum: String, + verified: bool, +} + +#[derive(Clone)] +struct Entry { + shard: String, + offset: u64, + length: usize, + rows: u32, + dimension: u32, + dtype: u8, +} + +struct ContractStore { + shards: HashMap, + entries: HashMap, +} + +struct HostEntry { + payload: Arc<[u8]>, + priority: f64, + owner_tenant: String, +} + +struct HostCache { + maximum_bytes: usize, + current_bytes: usize, + entries: HashMap, + sketch: TinyLfu, + inflation: f64, + tenant_bytes: HashMap, + default_tenant_max_bytes: usize, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, +} + +impl HostCache { + fn new(maximum_bytes: usize, tenant_max_percent: u8) -> Self { + Self { + maximum_bytes, + current_bytes: 0, + entries: HashMap::new(), + sketch: TinyLfu::new(4096), + inflation: 0.0, + tenant_bytes: HashMap::new(), + default_tenant_max_bytes: maximum_bytes * tenant_max_percent as usize / 100, + hits: 0, + misses: 0, + evictions: 0, + admission_rejections: 0, + } + } + + fn get(&mut self, key: &str) -> Option> { + let frequency = self.sketch.increment(&key); + let entry = self.entries.get_mut(key); + if let Some(entry) = entry { + entry.priority = self.inflation + f64::from(frequency) / entry.payload.len() as f64; + self.hits += 1; + Some(entry.payload.clone()) + } else { + self.misses += 1; + None + } + } + + fn remove(&mut self, key: &str) -> Option { + let entry = self.entries.remove(key)?; + self.current_bytes = self.current_bytes.saturating_sub(entry.payload.len()); + let remove_tenant = if let Some(bytes) = self.tenant_bytes.get_mut(&entry.owner_tenant) { + *bytes = bytes.saturating_sub(entry.payload.len()); + *bytes == 0 + } else { + false + }; + if remove_tenant { + self.tenant_bytes.remove(&entry.owner_tenant); + } + Some(entry) + } + + fn put(&mut self, tenant: &str, key: String, payload: Arc<[u8]>) { + if self.maximum_bytes == 0 || payload.len() > self.maximum_bytes { + return; + } + self.remove(&key); + let frequency = self.sketch.estimate(&key).max(1); + let mut priority = self.inflation + f64::from(frequency) / payload.len() as f64; + while self + .tenant_bytes + .get(tenant) + .copied() + .unwrap_or(0) + .saturating_add(payload.len()) + > self.default_tenant_max_bytes + { + let Some((victim_key, victim_priority)) = self + .entries + .iter() + .filter(|(_, entry)| entry.owner_tenant == tenant) + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + .map(|(key, entry)| (key.clone(), entry.priority)) + else { + return; + }; + if priority <= victim_priority { + self.admission_rejections += 1; + return; + } + let victim = self + .remove(&victim_key) + .expect("host tenant victim disappeared"); + self.inflation = self.inflation.max(victim.priority); + priority = self.inflation + f64::from(frequency) / payload.len() as f64; + self.evictions += 1; + } + while self.current_bytes + payload.len() > self.maximum_bytes { + let Some((victim_key, victim_priority)) = self + .entries + .iter() + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + .map(|(key, entry)| (key.clone(), entry.priority)) + else { + return; + }; + if priority < victim_priority { + self.admission_rejections += 1; + return; + } + let victim = self.remove(&victim_key).expect("host victim disappeared"); + self.inflation = self.inflation.max(victim.priority); + priority = self.inflation + f64::from(frequency) / payload.len() as f64; + self.evictions += 1; + } + self.current_bytes += payload.len(); + *self.tenant_bytes.entry(tenant.to_owned()).or_default() += payload.len(); + self.entries.insert( + key, + HostEntry { + payload, + priority, + owner_tenant: tenant.to_owned(), + }, + ); + } +} + +pub struct ShardStore { + contracts: HashMap, + roots: Vec<(String, PathBuf)>, + host_cache: HostCache, + pub batch_read_calls: u64, + pub batch_read_bytes: u64, + verify_full_shards: bool, +} + +impl ShardStore { + pub fn open( + roots: &[(String, PathBuf)], + host_cache_bytes: usize, + host_tenant_max_percent: u8, + verify_full_shards: bool, + ) -> Result { + let mut contracts = HashMap::new(); + for (contract, root) in roots { + if contract.is_empty() || contracts.contains_key(contract) { + bail!("invalid or duplicate model contract {contract:?}"); + } + contracts.insert(contract.clone(), open_contract(root)?); + } + if contracts.is_empty() { + bail!("at least one immutable shard contract is required"); + } + Ok(Self { + contracts, + roots: roots.to_vec(), + host_cache: HostCache::new(host_cache_bytes, host_tenant_max_percent), + batch_read_calls: 0, + batch_read_bytes: 0, + verify_full_shards, + }) + } + + /// Re-open every immutable contract index and publish the new generation + /// only after all roots validate. Content-addressed host-cache entries stay + /// valid across the atomic metadata swap. + pub fn reload(&mut self) -> Result<()> { + let mut contracts = HashMap::new(); + for (contract, root) in &self.roots { + contracts.insert(contract.clone(), open_contract(root)?); + } + self.contracts = contracts; + Ok(()) + } + + pub fn resolve_many( + &mut self, + descriptors: &[Descriptor], + tenant: &str, + ) -> Result>> { + let mut output = vec![None; descriptors.len()]; + let mut groups = HashMap::<(String, String), Vec<(usize, Entry)>>::new(); + for (index, descriptor) in descriptors.iter().enumerate() { + let key = cache_key(descriptor); + if let Some(payload) = self.host_cache.get(&key) { + output[index] = Some(payload); + continue; + } + let contract = self + .contracts + .get(&descriptor.contract) + .ok_or_else(|| anyhow!("model contract has no immutable shard root"))?; + let entry = contract + .entries + .get(&descriptor.digest) + .ok_or_else(|| anyhow!("tensor is missing from the immutable shard index"))? + .clone(); + if entry.rows != descriptor.rows + || entry.dimension != descriptor.dimension + || entry.dtype != descriptor.dtype + || entry.length + != tensor_bytes(descriptor.rows, descriptor.dimension, descriptor.dtype)? + { + bail!("tensor descriptor disagrees with the immutable shard index"); + } + groups + .entry((descriptor.contract.clone(), entry.shard.clone())) + .or_default() + .push((index, entry)); + } + + for ((contract_name, shard_name), mut entries) in groups { + let contract = self.contracts.get_mut(&contract_name).unwrap(); + let shard = contract.shards.get_mut(&shard_name).unwrap(); + if self.verify_full_shards { + verify_shard(shard)?; + } + entries.sort_by_key(|(_, entry)| entry.offset); + let mut ranges = Vec::<(u64, u64, Vec<(usize, Entry)>)>::new(); + let mut cursor = 0; + while cursor < entries.len() { + let start = entries[cursor].1.offset; + let mut end = start + entries[cursor].1.length as u64; + let mut limit = cursor + 1; + while limit < entries.len() { + let candidate = &entries[limit].1; + let candidate_end = candidate.offset + candidate.length as u64; + if candidate.offset.saturating_sub(end) > 64 * 1024 + || candidate_end.saturating_sub(start) > 8 * 1024 * 1024 + { + break; + } + end = end.max(candidate_end); + limit += 1; + } + ranges.push((start, end, entries[cursor..limit].to_vec())); + cursor = limit; + } + let worker_count = ranges.len().min(8); + let ranges_per_worker = ranges.len().div_ceil(worker_count); + let resolved = std::thread::scope(|scope| -> Result> { + let mut workers = Vec::new(); + for ranges in ranges.chunks(ranges_per_worker) { + let file = shard.file.try_clone()?; + workers.push(scope.spawn(move || -> Result)>> { + let mut resolved = Vec::new(); + for (start, end, entries) in ranges { + let mut range = vec![0_u8; (end - start) as usize]; + file.read_exact_at(&mut range, *start)?; + for (output_index, entry) in entries { + let local = (entry.offset - start) as usize; + let payload = range[local..local + entry.length].to_vec(); + let actual = hex::encode(Sha256::digest(&payload)); + if actual != descriptors[*output_index].digest { + bail!("tensor checksum mismatch inside immutable shard"); + } + resolved.push((*output_index, payload)); + } + } + Ok(resolved) + })); + } + let mut resolved = Vec::new(); + for worker in workers { + resolved.extend( + worker + .join() + .map_err(|_| anyhow!("shard reader worker panicked"))??, + ); + } + Ok(resolved) + })?; + self.batch_read_calls += ranges.len() as u64; + self.batch_read_bytes += ranges + .iter() + .map(|(start, end, _)| end - start) + .sum::(); + for (output_index, payload) in resolved { + let payload: Arc<[u8]> = payload.into(); + output[output_index] = Some(Arc::clone(&payload)); + self.host_cache + .put(tenant, cache_key(&descriptors[output_index]), payload); + } + } + output + .into_iter() + .map(|payload| payload.ok_or_else(|| anyhow!("tensor resolution produced no payload"))) + .collect() + } + + pub fn host_status(&self) -> (u64, u64, u64, u64) { + ( + self.host_cache.hits, + self.host_cache.misses, + self.host_cache.evictions, + self.host_cache.admission_rejections, + ) + } +} + +fn open_contract(root: &Path) -> Result { + let index_path = root.join(INDEX_NAME); + let mut index_file = nofollow(&index_path)?; + let mut document = String::new(); + index_file.read_to_string(&mut document)?; + let raw: RawIndex = serde_json::from_str(&document).context("invalid immutable shard index")?; + if raw.format != "vectorchord.tilemaxsim.shards" || raw.version != 1 { + bail!("unsupported immutable shard index"); + } + if raw.alignment == 0 || !raw.alignment.is_power_of_two() { + bail!("invalid immutable shard alignment"); + } + let mut shards = HashMap::new(); + for shard in raw.shards { + let digest = shard + .checksum + .strip_prefix("sha256:") + .ok_or_else(|| anyhow!("invalid immutable shard checksum"))?; + if shard.name != format!("shards/sha256-{digest}.vts") + || !digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + { + bail!("invalid immutable shard name"); + } + let path = root.join(&shard.name); + let file = nofollow(&path)?; + if file.metadata()?.len() != shard.size as u64 { + bail!("immutable shard size mismatch"); + } + if shards + .insert( + shard.name, + ShardFile { + file, + size: shard.size, + checksum: digest.to_owned(), + verified: false, + }, + ) + .is_some() + { + bail!("duplicate immutable shard"); + } + } + let mut entries = HashMap::new(); + for entry in raw.entries { + let dtype = match entry.dtype.as_str() { + "float32" => 1, + "float16" => 2, + _ => bail!("invalid shard tensor dtype"), + }; + let shard = shards + .get(&entry.shard) + .ok_or_else(|| anyhow!("shard tensor references an unknown file"))?; + if !(entry.offset as usize).is_multiple_of(raw.alignment) + || entry.length != tensor_bytes(entry.rows, entry.dimension, dtype)? + || entry.offset + entry.length as u64 > shard.size as u64 + { + bail!("invalid shard tensor range"); + } + if entries + .insert( + entry.digest, + Entry { + shard: entry.shard, + offset: entry.offset, + length: entry.length, + rows: entry.rows, + dimension: entry.dimension, + dtype, + }, + ) + .is_some() + { + bail!("duplicate tensor digest in shard index"); + } + } + Ok(ContractStore { shards, entries }) +} + +fn nofollow(path: &Path) -> Result { + OpenOptions::new() + .read(true) + .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(path) + .with_context(|| format!("cannot open immutable shard path {}", path.display())) +} + +fn verify_shard(shard: &mut ShardFile) -> Result<()> { + if shard.verified { + return Ok(()); + } + let mut digest = Sha256::new(); + let mut offset = 0_u64; + let mut buffer = vec![0_u8; 8 * 1024 * 1024]; + while offset < shard.size as u64 { + let count = buffer.len().min(shard.size - offset as usize); + shard.file.read_exact_at(&mut buffer[..count], offset)?; + digest.update(&buffer[..count]); + offset += count as u64; + } + if hex::encode(digest.finalize()) != shard.checksum { + bail!("immutable shard checksum mismatch"); + } + shard.verified = true; + Ok(()) +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: u8) -> Result { + let scalar = match dtype { + 1 => 4, + 2 => 2, + _ => bail!("unsupported tensor dtype"), + }; + (rows as usize) + .checked_mul(dimension as usize) + .and_then(|value| value.checked_mul(scalar)) + .ok_or_else(|| anyhow!("tensor shape overflow")) +} + +pub fn cache_key(descriptor: &Descriptor) -> String { + format!( + "{}:{}:{}:{}:{}", + descriptor.contract, + descriptor.digest, + descriptor.rows, + descriptor.dimension, + descriptor.dtype + ) +} diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 26cde468..a0c831a4 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -92,6 +92,8 @@ static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut pgrx::pg_sys::config_generic = static VCHORDRQ_MAXSIM_CANDIDATE_LIMIT: GucSetting = GucSetting::::new(-1); +static VCHORDRQ_MAXSIM_PROFILE: GucSetting = GucSetting::::new(false); + static VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS: GucSetting = GucSetting::::new(32); static VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS: GucSetting = GucSetting::::new(256); @@ -104,6 +106,11 @@ static VCHORDRQ_MAXSIM_GPU_ENDPOINT: GucSetting> = static VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS: GucSetting = GucSetting::::new(2000); +static VCHORDRQ_MAXSIM_TENANT: GucSetting> = + GucSetting::>::new(Some(c"")); + +static VCHORDRQ_MAXSIM_PRIORITY: GucSetting = GucSetting::::new(0); + static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS: GucSetting = GucSetting::::new(1_000_000); static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES: GucSetting = GucSetting::::new(1_073_741_824); @@ -191,6 +198,14 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_bool_guc( + c"vchordrq.maxsim_profile", + c"Emit one client NOTICE with MaxSim phase timings and counters.", + c"Disabled by default; emitted profiles contain no tensor references or application IDs.", + &VCHORDRQ_MAXSIM_PROFILE, + GucContext::Userset, + GucFlags::default(), + ); GucRegistry::define_int_guc( c"vchordrq.maxsim_planner_query_tokens", c"Expected MaxSim query-token count used by the planner.", @@ -237,6 +252,24 @@ pub fn init() { GucContext::Suset, GucFlags::default(), ); + GucRegistry::define_string_guc( + c"vchordrq.maxsim_tenant", + c"Logical tenant used only by the TileMaxSim GPU scheduler.", + c"Authorization and source filtering must be completed before setting this request-local value.", + &VCHORDRQ_MAXSIM_TENANT, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_priority", + c"TileMaxSim request priority; larger values run first subject to tenant fairness and aging.", + c"The allowed request-local range is -100 through 100.", + &VCHORDRQ_MAXSIM_PRIORITY, + -100, + 100, + GucContext::Userset, + GucFlags::default(), + ); GucRegistry::define_int_guc( c"vchordrq.maxsim_gpu_max_batch_tokens", c"Maximum query plus document tokens in one TileMaxSim request.", @@ -586,6 +619,10 @@ pub fn vchordrq_maxsim_candidate_limit() -> Option { if value < 0 { None } else { Some(value as u32) } } +pub fn vchordrq_maxsim_profile() -> bool { + VCHORDRQ_MAXSIM_PROFILE.get() +} + pub fn vchordrq_maxsim_planner_query_tokens() -> u32 { VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS.get() as u32 } @@ -606,6 +643,17 @@ pub fn vchordrq_maxsim_gpu_timeout_ms() -> u32 { VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS.get() as u32 } +pub fn vchordrq_maxsim_tenant() -> Option { + VCHORDRQ_MAXSIM_TENANT.get().and_then(|tenant| { + let tenant = tenant.to_string_lossy().into_owned(); + (!tenant.is_empty()).then_some(tenant) + }) +} + +pub fn vchordrq_maxsim_priority() -> i32 { + VCHORDRQ_MAXSIM_PRIORITY.get() +} + pub fn vchordrq_maxsim_gpu_max_batch_tokens() -> u32 { VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS.get() as u32 } diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 238dbe2e..4f76c0db 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -13,12 +13,14 @@ // Copyright (c) 2025-2026 TensorChord Inc. mod candidate; +mod exact; mod external; mod gpu; +mod profile; mod rerank; mod search; -use self::candidate::{LegacyPageCandidateGenerator, PageCandidateGenerator}; +use self::candidate::{DensePageCandidateGenerator, PageCandidateGenerator}; use self::gpu::{GpuTileMaxsimBackend, UnixSocketTransport, report_gpu_fallback}; use self::rerank::{CpuExactMaxsimBackend, ExactMaxsimBackend, HeapArrayTensorSource}; use crate::index::fetcher::*; @@ -176,6 +178,7 @@ impl SearchBuilder for MaxsimBuilder { .map(|vector| RandomProject::project(vector.as_borrowed())) .collect::>(); Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); let (results, estimation_by_threshold) = match options.io_search { Io::Plain => maxsim_search::<_, Op>( index, @@ -208,6 +211,15 @@ impl SearchBuilder for MaxsimBuilder { make_h0_stream_prefetcher.clone(), ), }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { let sequence = Heap::from(results); @@ -312,6 +324,13 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); (accu_set, rough_set, estimation_by_threshold) })) } @@ -332,6 +351,7 @@ impl SearchBuilder for MaxsimBuilder { .map(|vector| RandomProject::project(vector.as_borrowed())) .collect::>(); Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); let (results, estimation_by_threshold) = match options.io_search { Io::Plain => maxsim_search::<_, Op>( index, @@ -364,6 +384,15 @@ impl SearchBuilder for MaxsimBuilder { make_h0_stream_prefetcher.clone(), ), }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { let sequence = Heap::from(results); @@ -468,6 +497,13 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); (accu_set, rough_set, estimation_by_threshold) })) } @@ -484,6 +520,7 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); let (results, estimation_by_threshold) = match options.io_search { Io::Plain => maxsim_search::<_, Op>( index, @@ -516,6 +553,15 @@ impl SearchBuilder for MaxsimBuilder { make_h0_stream_prefetcher.clone(), ), }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { let sequence = Heap::from(results); @@ -620,6 +666,13 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); (accu_set, rough_set, estimation_by_threshold) })) } @@ -636,6 +689,7 @@ impl SearchBuilder for MaxsimBuilder { }) .collect::>(); Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); let (results, estimation_by_threshold) = match options.io_search { Io::Plain => maxsim_search::<_, Op>( index, @@ -668,6 +722,15 @@ impl SearchBuilder for MaxsimBuilder { make_h0_stream_prefetcher.clone(), ), }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); if maxsim_refine != 0 && !results.is_empty() { let sequence = Heap::from(results); @@ -772,6 +835,13 @@ impl SearchBuilder for MaxsimBuilder { let rough_iter = results.into_iter(); rough_set.extend(rough_iter.map(rough_map)); } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); (accu_set, rough_set, estimation_by_threshold) })) } @@ -779,7 +849,7 @@ impl SearchBuilder for MaxsimBuilder { iter }; let mut candidates = - LegacyPageCandidateGenerator.generate(n, &mut token_searches, maxsim_candidate_limit); + DensePageCandidateGenerator.generate(n, &mut token_searches, maxsim_candidate_limit); drop(token_searches); let iter: Box> = match maxsim_backend { PostgresMaxsimBackend::CoarseOnly => Box::new(candidates.map(|candidate| { diff --git a/src/index/vchordrq/scanners/maxsim/candidate.rs b/src/index/vchordrq/scanners/maxsim/candidate.rs index 9aa3e22f..af9b45b0 100644 --- a/src/index/vchordrq/scanners/maxsim/candidate.rs +++ b/src/index/vchordrq/scanners/maxsim/candidate.rs @@ -12,11 +12,13 @@ // // Copyright (c) 2025-2026 TensorChord Inc. +use super::profile; use crate::index::fetcher::pointer_to_kv; use always_equal::AlwaysEqual; use distance::Distance; use std::cmp::Reverse; -use std::collections::BinaryHeap; +use std::collections::hash_map::Entry; +use std::collections::{BinaryHeap, HashMap}; use std::num::NonZero; pub(super) type HeapKey = [u16; 3]; @@ -41,10 +43,10 @@ pub(super) trait PageCandidateGenerator { } #[derive(Default)] -pub(super) struct LegacyPageCandidateGenerator; +pub(super) struct DensePageCandidateGenerator; -impl PageCandidateGenerator for LegacyPageCandidateGenerator { - type Candidates = LegacyPageCandidates; +impl PageCandidateGenerator for DensePageCandidateGenerator { + type Candidates = PageCandidates; fn generate( &mut self, @@ -52,60 +54,120 @@ impl PageCandidateGenerator for LegacyPageCandidateGenerator { token_searches: &mut dyn Iterator, candidate_limit: usize, ) -> Self::Candidates { - let mut updates = Vec::new(); + let mut page_lookup = HashMap::new(); + let mut page_keys = Vec::new(); + let mut best_by_page_query = Vec::new(); let mut estimations = Vec::with_capacity(query_count); + let mut hit_updates = 0u64; + let mut page_token_updates = 0u64; for (query_id, (accu_set, rough_set, estimation_by_threshold)) in token_searches.enumerate() { - updates.reserve(accu_set.len() + rough_set.len()); + debug_assert!(query_id < query_count); + let hit_collect_timer = profile::ProfileTimer::start(); + hit_updates += (accu_set.len() + rough_set.len()) as u64; let is_empty = accu_set.is_empty() && rough_set.is_empty(); let mut estimation_by_scope = Distance::NEG_INFINITY; for (distance, payload) in accu_set { estimation_by_scope = std::cmp::max(estimation_by_scope, distance); let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + page_token_updates += u64::from(record_page_token_distance( + &mut page_lookup, + &mut page_keys, + &mut best_by_page_query, + query_count, + query_id, + key, + distance, + )); } for (distance, payload) in rough_set { let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + page_token_updates += u64::from(record_page_token_distance( + &mut page_lookup, + &mut page_keys, + &mut best_by_page_query, + query_count, + query_id, + key, + distance, + )); } estimations.push(if !is_empty { std::cmp::max(estimation_by_scope, estimation_by_threshold) } else { Distance::ZERO }); + let hit_collect_elapsed = hit_collect_timer.elapsed(); + profile::update(|profile| { + profile.hit_collect_us += profile::duration_us(hit_collect_elapsed); + }); } debug_assert_eq!(estimations.len(), query_count); - updates.sort_unstable_by_key(|&(key, ..)| key); - let inner = updates - .chunk_by(|(kl, ..), (kr, ..)| kl == kr) - .map(|chunk| { - let key = chunk[0].0; - let mut value = vec![None; query_count]; - for &(_, query_id, distance) in chunk { - let this = value[query_id].get_or_insert(Distance::INFINITY); - *this = std::cmp::min(*this, distance); - } + profile::update(|profile| { + profile.hit_updates += hit_updates; + profile.page_token_updates += page_token_updates; + }); + let page_aggregate_timer = profile::ProfileTimer::start(); + let mut page_order = (0..page_keys.len()).collect::>(); + page_order.sort_unstable_by_key(|&page_index| page_keys[page_index]); + let inner = page_order + .into_iter() + .map(|page_index| { + let key = page_keys[page_index]; + let start = page_index * query_count; + let values = &best_by_page_query[start..start + query_count]; let mut maxsim = 0.0f32; - for (query_id, distance) in value.into_iter().enumerate() { + for (query_id, distance) in values.iter().copied().enumerate() { let distance = distance.unwrap_or(estimations[query_id]); maxsim += distance.to_f32(); } (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) }) .collect::>(); - LegacyPageCandidates { + let page_aggregate_elapsed = page_aggregate_timer.elapsed(); + profile::update(|profile| { + profile.page_aggregate_us += profile::duration_us(page_aggregate_elapsed); + profile.aggregated_pages += page_keys.len() as u64; + }); + PageCandidates { inner, remaining: candidate_limit, } } } -pub(super) struct LegacyPageCandidates { +fn record_page_token_distance( + page_lookup: &mut HashMap, + page_keys: &mut Vec, + best_by_page_query: &mut Vec>, + query_count: usize, + query_id: usize, + key: HeapKey, + distance: Distance, +) -> bool { + let page_index = match page_lookup.entry(key) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + let page_index = page_keys.len(); + entry.insert(page_index); + page_keys.push(key); + best_by_page_query.resize(best_by_page_query.len() + query_count, None); + page_index + } + }; + let slot = &mut best_by_page_query[page_index * query_count + query_id]; + let first_for_page_token = slot.is_none(); + let best = slot.get_or_insert(Distance::INFINITY); + *best = std::cmp::min(*best, distance); + first_for_page_token +} + +pub(super) struct PageCandidates { inner: BinaryHeap<(Reverse, AlwaysEqual)>, remaining: usize, } -impl Iterator for LegacyPageCandidates { +impl Iterator for PageCandidates { type Item = PageCandidate; fn next(&mut self) -> Option { @@ -126,8 +188,8 @@ impl Iterator for LegacyPageCandidates { } } -impl ExactSizeIterator for LegacyPageCandidates {} -impl std::iter::FusedIterator for LegacyPageCandidates {} +impl ExactSizeIterator for PageCandidates {} +impl std::iter::FusedIterator for PageCandidates {} #[cfg(test)] mod tests { @@ -158,7 +220,7 @@ mod tests { ), ] .into_iter(); - let candidates = LegacyPageCandidateGenerator + let candidates = DensePageCandidateGenerator .generate(2, &mut token_searches, usize::MAX) .collect::>(); @@ -182,7 +244,7 @@ mod tests { Distance::ZERO, )] .into_iter(); - let candidates = LegacyPageCandidateGenerator + let candidates = DensePageCandidateGenerator .generate(1, &mut token_searches, 1) .collect::>(); diff --git a/src/index/vchordrq/scanners/maxsim/exact.rs b/src/index/vchordrq/scanners/maxsim/exact.rs new file mode 100644 index 00000000..32d36ec4 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/exact.rs @@ -0,0 +1,501 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the ELv2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorStorage, + TileMaxsimSourceBinding, resolve_tilemaxsim_source, validate_descriptor, +}; +use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; +use super::profile; +use super::rerank::RerankError; +use crate::datatype::memory_halfvec::HalfvecInput; +use crate::datatype::memory_vector::VectorInput; +use crate::index::gucs::{self, PostgresMaxsimBackend}; +use distance::Distance; +use pgrx::datum::{Array, DatumWithOid, FromDatum, IntoDatum}; +use pgrx::iter::TableIterator; +use pgrx::{name, pg_extern}; +use std::collections::{BTreeMap, BTreeSet}; +use std::time::Duration; +use vchordrq::types::OwnedVector; +use vector::VectorBorrowed; + +/// Exact TileMaxSim for a caller-supplied candidate set of external tensors. +/// Candidate selection, tenancy, ACLs, graph traversal, and clustering are +/// intentionally outside this function. +#[pg_extern(sql = "")] +fn _vchordrq_tilemaxsim_rerank_vector( + source_oid: pgrx::pg_sys::Oid, + query: Array<'_, VectorInput<'_>>, + candidate_ids: Array<'_, i64>, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = collect_vector_query(query) + .and_then(|query| { + collect_candidate_ids(candidate_ids) + .and_then(|candidate_ids| execute_rerank(source_oid, query, candidate_ids, top_k)) + }) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +/// Half-precision overload of exact caller-scoped TileMaxSim. +#[pg_extern(sql = "")] +fn _vchordrq_tilemaxsim_rerank_halfvec( + source_oid: pgrx::pg_sys::Oid, + query: Array<'_, HalfvecInput<'_>>, + candidate_ids: Array<'_, i64>, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = collect_halfvec_query(query) + .and_then(|query| { + collect_candidate_ids(candidate_ids) + .and_then(|candidate_ids| execute_rerank(source_oid, query, candidate_ids, top_k)) + }) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +fn collect_vector_query( + query: Array<'_, VectorInput<'_>>, +) -> Result, RerankError> { + let mut vectors = Vec::with_capacity(query.len()); + for vector in query.iter() { + let vector = vector.ok_or(RerankError::TensorMismatch)?; + vectors.push(OwnedVector::Vecf32(vector.as_borrowed().own())); + } + validate_query(&vectors)?; + Ok(vectors) +} + +fn collect_halfvec_query( + query: Array<'_, HalfvecInput<'_>>, +) -> Result, RerankError> { + let mut vectors = Vec::with_capacity(query.len()); + for vector in query.iter() { + let vector = vector.ok_or(RerankError::TensorMismatch)?; + vectors.push(OwnedVector::Vecf16(vector.as_borrowed().own())); + } + validate_query(&vectors)?; + Ok(vectors) +} + +fn validate_query(query: &[OwnedVector]) -> Result<(), RerankError> { + if query.is_empty() { + return Err(RerankError::Configuration( + "query tensor must contain at least one token", + )); + } + Ok(()) +} + +fn collect_candidate_ids(candidate_ids: Array<'_, i64>) -> Result, RerankError> { + let mut result = Vec::with_capacity(candidate_ids.len()); + let mut unique = BTreeSet::new(); + for public_id in candidate_ids.iter() { + let public_id = public_id.ok_or(RerankError::Configuration( + "candidate_ids must not contain NULL", + ))?; + if !unique.insert(public_id) { + return Err(RerankError::Configuration( + "candidate_ids must not contain duplicates", + )); + } + result.push(public_id); + } + if result.is_empty() { + return Err(RerankError::Configuration( + "candidate_ids must contain at least one ID", + )); + } + Ok(result) +} + +fn execute_rerank( + source_oid: pgrx::pg_sys::Oid, + query: Vec, + candidate_ids: Vec, + top_k: i32, +) -> Result, RerankError> { + validate_top_k(candidate_ids.len(), top_k)?; + if !matches!(gucs::vchordrq_maxsim_backend(), PostgresMaxsimBackend::Gpu) { + return Err(RerankError::Configuration( + "external TileMaxSim rerank requires vchordrq.maxsim_backend = 'gpu'", + )); + } + require_mvcc_snapshot()?; + + let profile_guard = profile::ProfileGuard::start(gucs::vchordrq_maxsim_profile()); + let total_timer = profile::ProfileTimer::start(); + profile::update(|profile| { + profile.query_tokens = query.len() as u64; + profile.generated_candidates = candidate_ids.len() as u64; + }); + + let preflight_timer = profile::ProfileTimer::start(); + let binding = resolve_tilemaxsim_source(source_oid)?; + let source_lock = RelationLock::open(source_oid, pgrx::pg_sys::AccessShareLock as _)?; + let descriptor_lock = binding + .descriptor_oid + .map(|oid| RelationLock::open(oid, pgrx::pg_sys::AccessShareLock as _)) + .transpose()?; + if source_lock.oid() != binding.source_oid + || descriptor_lock.as_ref().map(RelationLock::oid) != binding.descriptor_oid + { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source changed during execution".into(), + )); + } + preflight_descriptor_access(&binding)?; + profile::update(|profile| { + profile.preflight_us += profile::duration_us(preflight_timer.elapsed()); + }); + + let descriptor_timer = profile::ProfileTimer::start(); + let (mut source, public_ids) = load_visible_descriptors(&binding, &candidate_ids)?; + let mut candidates = public_ids + .keys() + .copied() + .map(|heap_key| PageCandidate { + approximate_distance: Distance::ZERO, + heap_key, + }) + .collect::>() + .into_iter(); + profile::update(|profile| { + profile.descriptor_us += profile::duration_us(descriptor_timer.elapsed()); + profile.visible_candidates = public_ids.len() as u64; + profile.descriptors = public_ids.len() as u64; + }); + + let endpoint = gucs::vchordrq_maxsim_gpu_endpoint() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let mut backend = GpuExternalTileMaxsimBackend::new( + UnixSocketTransport::new(endpoint), + binding.model_contract_id, + Duration::from_millis(gucs::vchordrq_maxsim_gpu_timeout_ms() as u64), + gucs::vchordrq_maxsim_gpu_max_batch_tokens() as usize, + gucs::vchordrq_maxsim_gpu_max_batch_bytes() as usize, + ); + if let Some(tenant) = gucs::vchordrq_maxsim_tenant() { + backend = backend.with_scheduling(tenant, gucs::vchordrq_maxsim_priority()); + } + let sidecar_timer = profile::ProfileTimer::start(); + let exact = backend.rerank(&query, &mut candidates, &mut source)?; + let mut rows = exact + .map(|result| { + let public_id = public_ids.get(&result.heap_key).copied().ok_or_else(|| { + RerankError::Protocol("sidecar result has no visible public ID".into()) + })?; + Ok((result.distance, public_id)) + }) + .collect::, RerankError>>()?; + profile::update(|profile| { + profile.sidecar_us += profile::duration_us(sidecar_timer.elapsed()); + }); + + let result_timer = profile::ProfileTimer::start(); + rows.sort_unstable_by(|(left_distance, left_id), (right_distance, right_id)| { + left_distance + .cmp(right_distance) + .then_with(|| left_id.cmp(right_id)) + }); + rows.truncate(top_k as usize); + let output = rows + .into_iter() + .map(|(distance, public_id)| (public_id, -distance.to_f32())) + .collect::>(); + profile::update(|profile| { + profile.result_finalize_us += profile::duration_us(result_timer.elapsed()); + profile.returned_rows = output.len() as u64; + }); + profile_guard.finish(total_timer.elapsed()); + Ok(output.into_iter()) +} + +fn validate_top_k(candidate_count: usize, top_k: i32) -> Result<(), RerankError> { + if top_k <= 0 || top_k as usize > candidate_count { + return Err(RerankError::Configuration( + "top_k must be positive and no greater than candidate_ids length", + )); + } + Ok(()) +} + +fn require_mvcc_snapshot() -> Result<(), RerankError> { + let snapshot = unsafe { pgrx::pg_sys::GetActiveSnapshot() }; + if snapshot.is_null() + || unsafe { (*snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + { + return Err(RerankError::Configuration( + "external TileMaxSim rerank requires an active MVCC snapshot", + )); + } + Ok(()) +} + +fn preflight_descriptor_access(binding: &TileMaxsimSourceBinding) -> Result<(), RerankError> { + let query = descriptor_query(binding, true)?; + pgrx::spi::Spi::connect(|client| { + let privilege = client + .prepare( + "SELECT pg_catalog.has_table_privilege($1, 'SELECT') AS allowed", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + for relation_oid in std::iter::once(binding.source_oid).chain(binding.descriptor_oid) { + let allowed = client + .select(&privilege, Some(1), &[relation_oid.into()]) + .map_err(registry_error)? + .first() + .get_by_name::("allowed") + .map_err(registry_error)? + .unwrap_or(false); + if !allowed { + return Err(RerankError::Registry( + "table-level SELECT privilege is required for every TileMaxSim relation".into(), + )); + } + } + client + .select(query.as_str(), Some(1), &[]) + .map(|_| ()) + .map_err(registry_error) + }) +} + +fn load_visible_descriptors( + binding: &TileMaxsimSourceBinding, + candidate_ids: &[i64], +) -> Result<(MaterializedDescriptorSource, BTreeMap), RerankError> { + let mut candidates_by_id = BTreeMap::new(); + for (ordinal, public_id) in candidate_ids.iter().copied().enumerate() { + candidates_by_id.insert(public_id, candidate_for_ordinal(ordinal)?); + } + let query = descriptor_query(binding, false)?; + let (descriptors, public_ids) = pgrx::spi::Spi::connect(|client| { + let int8_array_oid = unsafe { pgrx::pg_sys::get_array_type(pgrx::pg_sys::INT8OID) }; + let prepared = client + .prepare( + query.as_str(), + &[ + pgrx::pg_sys::PgOid::from(int8_array_oid), + pgrx::pg_sys::PgOid::from(pgrx::pg_sys::TEXTOID), + ], + ) + .map_err(registry_error)?; + let args: [DatumWithOid<'_>; 2] = [ + candidate_ids.to_vec().into(), + binding.model_contract_id.clone().into(), + ]; + let rows = client + .select(&prepared, Some(candidate_ids.len() as _), &args) + .map_err(registry_error)?; + let mut descriptors = BTreeMap::new(); + let mut public_ids = BTreeMap::new(); + for row in rows { + let public_id = required_heap_column::(&row, "public_id")?; + let candidate = candidates_by_id.get(&public_id).copied().ok_or_else(|| { + RerankError::Registry("descriptor query returned an unknown public ID".into()) + })?; + let descriptor = validate_descriptor( + candidate, + public_id, + required_heap_column::(&row, "tensor_ref")?, + required_heap_column::(&row, "tensor_rows")?, + required_heap_column::(&row, "tensor_dimension")?, + required_heap_column::(&row, "tensor_dtype")?, + required_heap_column::(&row, "tensor_checksum")?, + )?; + if descriptors.insert(candidate.heap_key, descriptor).is_some() + || public_ids.insert(candidate.heap_key, public_id).is_some() + { + return Err(RerankError::Registry( + "descriptor query returned a duplicate public ID".into(), + )); + } + } + Ok((descriptors, public_ids)) + })?; + Ok((MaterializedDescriptorSource(descriptors), public_ids)) +} + +fn candidate_for_ordinal(ordinal: usize) -> Result { + let ordinal = u32::try_from(ordinal).map_err(|_| RerankError::RequestTooLarge)?; + Ok(PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, (ordinal >> 16) as u16, ordinal as u16], + }) +} + +fn descriptor_query( + binding: &TileMaxsimSourceBinding, + preflight: bool, +) -> Result { + let source_relation = relation_name(binding.source_oid)?; + let names = &binding.column_names; + let model_contract = pgrx::spi::quote_identifier(&names.model_contract); + let public_id = pgrx::spi::quote_identifier(&names.public_id); + let tensor_ref = pgrx::spi::quote_identifier(&names.tensor_ref); + let tensor_rows = pgrx::spi::quote_identifier(&names.tensor_rows); + let tensor_dimension = pgrx::spi::quote_identifier(&names.tensor_dimension); + let tensor_dtype = pgrx::spi::quote_identifier(&names.tensor_dtype); + let tensor_checksum = pgrx::spi::quote_identifier(&names.tensor_checksum); + let predicate = if preflight { + "false".to_string() + } else { + format!("h.{public_id}::bigint = ANY($1) AND h.{model_contract} = $2") + }; + match binding.storage { + ExternalTensorStorage::SameHeap => Ok(format!( + "SELECT h.{public_id}::bigint AS public_id, + h.{tensor_ref} AS tensor_ref, + h.{tensor_rows} AS tensor_rows, + h.{tensor_dimension} AS tensor_dimension, + h.{tensor_dtype} AS tensor_dtype, + h.{tensor_checksum} AS tensor_checksum + FROM ONLY {source_relation} AS h + WHERE {predicate}" + )), + ExternalTensorStorage::DescriptorRelation => { + let descriptor_oid = binding.descriptor_oid.ok_or_else(|| { + RerankError::Registry("registered descriptor relation is missing".into()) + })?; + let descriptor_relation = relation_name(descriptor_oid)?; + let descriptor_public_id = pgrx::spi::quote_identifier( + names.descriptor_public_id.as_deref().ok_or_else(|| { + RerankError::Registry( + "registered descriptor public ID column is missing".into(), + ) + })?, + ); + Ok(format!( + "SELECT h.{public_id}::bigint AS public_id, + d.{tensor_ref} AS tensor_ref, + d.{tensor_rows} AS tensor_rows, + d.{tensor_dimension} AS tensor_dimension, + d.{tensor_dtype} AS tensor_dtype, + d.{tensor_checksum} AS tensor_checksum + FROM ONLY {source_relation} AS h + JOIN ONLY {descriptor_relation} AS d + ON d.{descriptor_public_id}::bigint = h.{public_id}::bigint + WHERE {predicate}" + )) + } + } +} + +fn relation_name(relation_oid: pgrx::pg_sys::Oid) -> Result { + pgrx::spi::Spi::connect(|client| { + let prepared = client + .prepare( + "SELECT n.nspname::text AS schema_name, c.relname::text AS relation_name + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE c.oid = $1", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[relation_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered relation disappeared".into(), + )); + } + let row = rows.first(); + Ok(pgrx::spi::quote_qualified_identifier( + required_column::(&row, "schema_name")?, + required_column::(&row, "relation_name")?, + )) + }) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry(format!("query returned NULL {name}"))) +} + +fn required_heap_column( + row: &pgrx::spi::SpiHeapTupleData<'_>, + name: &str, +) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +#[derive(Default)] +struct MaterializedDescriptorSource(BTreeMap); + +impl CandidateTensorDescriptorSource for MaterializedDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(self.0.remove(&candidate.heap_key)) + } +} + +struct RelationLock { + raw: pgrx::pg_sys::Relation, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl RelationLock { + fn open(oid: pgrx::pg_sys::Oid, lockmode: pgrx::pg_sys::LOCKMODE) -> Result { + let raw = unsafe { pgrx::pg_sys::relation_open(oid, lockmode) }; + if raw.is_null() { + return Err(RerankError::Registry("relation open returned NULL".into())); + } + Ok(Self { raw, lockmode }) + } + + fn oid(&self) -> pgrx::pg_sys::Oid { + unsafe { (*self.raw).rd_id } + } +} + +impl Drop for RelationLock { + fn drop(&mut self) { + unsafe { pgrx::pg_sys::relation_close(self.raw, self.lockmode) }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn explicit_candidates_support_full_corpus_ordinals() { + assert!(validate_top_k(256, 10).is_ok()); + assert!(validate_top_k(10, 0).is_err()); + assert!(validate_top_k(10, 11).is_err()); + let first = candidate_for_ordinal(0).unwrap().heap_key; + let last = candidate_for_ordinal(1_000_000).unwrap().heap_key; + assert_ne!(first, last); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/external.rs b/src/index/vchordrq/scanners/maxsim/external.rs index 0c8d1eb0..8b726da5 100644 --- a/src/index/vchordrq/scanners/maxsim/external.rs +++ b/src/index/vchordrq/scanners/maxsim/external.rs @@ -116,6 +116,20 @@ pub(super) struct ExternalTensorSourceBinding { pub column_names: ExternalTensorColumnNames, } +/// External tensor metadata bound directly to an application source relation. +/// +/// Unlike [`ExternalTensorSourceBinding`], this binding has no ANN index: the +/// caller supplies an already-authorized candidate ID set and VectorChord only +/// performs exact TileMaxSim over those candidates. +#[derive(Clone, Debug)] +pub(super) struct TileMaxsimSourceBinding { + pub source_oid: pgrx::pg_sys::Oid, + pub descriptor_oid: Option, + pub storage: ExternalTensorStorage, + pub model_contract_id: String, + pub column_names: ExternalTensorColumnNames, +} + #[derive(Clone, Debug)] pub(super) struct ExternalTensorColumnNames { pub model_contract: String, @@ -264,6 +278,114 @@ pub(super) fn resolve_external_tensor_source( }) } +/// Resolve a source-relation binding for exact candidate-only TileMaxSim. +/// +/// The SECURITY DEFINER SQL resolver revalidates relation ownership, caller +/// SELECT privileges, column types, uniqueness, and descriptor metadata. This +/// Rust layer deliberately sees only the validated projection. +pub(super) fn resolve_tilemaxsim_source( + source_oid: pgrx::pg_sys::Oid, +) -> Result { + pgrx::spi::Spi::connect(|client| { + let schema_rows = client + .select( + "SELECT n.nspname::text AS schema_name + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'", + Some(1), + &[], + ) + .map_err(registry_error)?; + if schema_rows.is_empty() { + return Err(RerankError::Registry( + "vchord extension schema is unavailable".into(), + )); + } + let schema_name = schema_rows + .first() + .get_by_name::("schema_name") + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry("vchord extension schema is NULL".into()))?; + let resolver = pgrx::spi::quote_qualified_identifier( + schema_name, + "vchordrq_tilemaxsim_source_info".to_string(), + ); + let query = format!( + "SELECT registered_source::oid AS source_oid, + descriptor_relation::oid AS descriptor_oid, + model_contract_id, + source_storage, + model_contract_column::text AS model_contract_column, + public_id_column::text AS public_id_column, + descriptor_public_id_column::text AS descriptor_public_id_column, + tensor_ref_column::text AS tensor_ref_column, + tensor_rows_column::text AS tensor_rows_column, + tensor_dim_column::text AS tensor_dim_column, + tensor_dtype_column::text AS tensor_dtype_column, + tensor_checksum_column::text AS tensor_checksum_column + FROM {resolver}($1::regclass)" + ); + let prepared = client + .prepare(query.as_str(), &pgrx::oids_of![pgrx::pg_sys::Oid]) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[source_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source resolution returned no row".into(), + )); + } + let row = rows.first(); + let resolved_source_oid = required_column::(&row, "source_oid")?; + let descriptor_oid = optional_column::(&row, "descriptor_oid")?; + let model_contract_id = required_column::(&row, "model_contract_id")?; + let storage = required_column::(&row, "source_storage")?; + if resolved_source_oid != source_oid { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source resolved a different relation".into(), + )); + } + let storage = match storage.as_str() { + "external_ref" if descriptor_oid.is_none() => ExternalTensorStorage::SameHeap, + "external_relation" if descriptor_oid.is_some() => { + ExternalTensorStorage::DescriptorRelation + } + _ => { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source is inconsistent".into(), + )); + } + }; + let column_names = ExternalTensorColumnNames { + model_contract: required_column::(&row, "model_contract_column")?, + public_id: required_column::(&row, "public_id_column")?, + descriptor_public_id: optional_column::(&row, "descriptor_public_id_column")?, + tensor_ref: required_column::(&row, "tensor_ref_column")?, + tensor_rows: required_column::(&row, "tensor_rows_column")?, + tensor_dimension: required_column::(&row, "tensor_dim_column")?, + tensor_dtype: required_column::(&row, "tensor_dtype_column")?, + tensor_checksum: required_column::(&row, "tensor_checksum_column")?, + }; + if matches!(storage, ExternalTensorStorage::DescriptorRelation) + && column_names.descriptor_public_id.is_none() + { + return Err(RerankError::Registry( + "registered descriptor relation has no public ID column".into(), + )); + } + validate_model_contract(&model_contract_id)?; + Ok(TileMaxsimSourceBinding { + source_oid, + descriptor_oid, + storage, + model_contract_id, + column_names, + }) + }) +} + fn registry_error(error: impl std::fmt::Display) -> RerankError { RerankError::Registry(error.to_string()) } diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs index af8c5455..7e0928a6 100644 --- a/src/index/vchordrq/scanners/maxsim/gpu.rs +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -28,6 +28,7 @@ use vchordrq::types::OwnedVector; const MAGIC: &[u8; 4] = b"VCTM"; const VERSION: u16 = 1; const EXTERNAL_VERSION: u16 = 2; +const SCHEDULED_EXTERNAL_VERSION: u16 = 3; const REQUEST_KIND: u16 = 1; const RESPONSE_KIND: u16 = 2; const HEADER_LEN: usize = 24; @@ -134,6 +135,13 @@ pub(super) struct GpuExternalTileMaxsimBackend { timeout: Duration, max_batch_tokens: usize, max_batch_bytes: usize, + scheduling: Option, +} + +#[derive(Clone, Debug)] +pub(super) struct TileMaxsimScheduling { + pub tenant: String, + pub priority: i32, } impl GpuExternalTileMaxsimBackend { @@ -150,8 +158,14 @@ impl GpuExternalTileMaxsimBackend { timeout, max_batch_tokens, max_batch_bytes, + scheduling: None, } } + + pub(super) fn with_scheduling(mut self, tenant: String, priority: i32) -> Self { + self.scheduling = Some(TileMaxsimScheduling { tenant, priority }); + self + } } impl GpuExternalTileMaxsimBackend { @@ -170,6 +184,8 @@ impl GpuExternalTileMaxsimBackend { source, self.max_batch_tokens, self.max_batch_bytes, + self.scheduling.as_ref(), + self.timeout, )?; if encoded.heap_keys.is_empty() { return Ok(RerankResults { @@ -184,13 +200,14 @@ impl GpuExternalTileMaxsimBackend { let response = self.transport .round_trip(&encoded.frame, self.timeout, max_response_bytes)?; - decode_response_for_version(&response, EXTERNAL_VERSION, request_id, &encoded.heap_keys) + decode_response_for_version(&response, encoded.version, request_id, &encoded.heap_keys) } } struct EncodedRequest { frame: Vec, heap_keys: Vec, + version: u16, } fn encode_external_request( @@ -201,10 +218,13 @@ fn encode_external_request( source: &mut S, max_batch_tokens: usize, max_batch_bytes: usize, + scheduling: Option<&TileMaxsimScheduling>, + timeout: Duration, ) -> Result { const MAX_MODEL_CONTRACT_BYTES: usize = 512; const MAX_TENSOR_REF_BYTES: usize = 4096; const MAX_CHECKSUM_BYTES: usize = 512; + const MAX_TENANT_BYTES: usize = 256; if model_contract_id.is_empty() || model_contract_id.len() > MAX_MODEL_CONTRACT_BYTES @@ -214,6 +234,17 @@ fn encode_external_request( "model contract is empty, oversized, or contains control characters", )); } + if let Some(scheduling) = scheduling { + if scheduling.tenant.is_empty() + || scheduling.tenant.len() > MAX_TENANT_BYTES + || scheduling.tenant.chars().any(char::is_control) + || !(-100..=100).contains(&scheduling.priority) + { + return Err(RerankError::Configuration( + "TileMaxSim scheduler tenant or priority is invalid", + )); + } + } let (dtype, dimension) = tensor_metadata(query)?; let external_dtype = match dtype { TensorDtype::F32 => ExternalTensorDtype::F32, @@ -240,7 +271,23 @@ fn encode_external_request( writer.u16(0)?; writer .u32(u32::try_from(model_contract_id.len()).map_err(|_| RerankError::RequestTooLarge)?)?; + let version = if let Some(scheduling) = scheduling { + writer.i32(scheduling.priority)?; + writer.u32( + u32::try_from(timeout.as_millis().clamp(1, 600_000)) + .map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.u32( + u32::try_from(scheduling.tenant.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + SCHEDULED_EXTERNAL_VERSION + } else { + EXTERNAL_VERSION + }; writer.bytes(model_contract_id.as_bytes())?; + if let Some(scheduling) = scheduling { + writer.bytes(scheduling.tenant.as_bytes())?; + } encode_tensor_values(&mut writer, query, dtype)?; let mut heap_keys = Vec::new(); @@ -291,7 +338,7 @@ fn encode_external_request( .checked_sub(HEADER_LEN) .ok_or_else(|| RerankError::Protocol("invalid request length".into()))?; writer.patch_bytes(0, MAGIC); - writer.patch_u16(4, EXTERNAL_VERSION); + writer.patch_u16(4, version); writer.patch_u16(6, REQUEST_KIND); writer.patch_u64(8, request_id); writer.patch_u64( @@ -301,6 +348,7 @@ fn encode_external_request( Ok(EncodedRequest { frame: writer.finish(), heap_keys, + version, }) } @@ -422,6 +470,7 @@ fn encode_request( Ok(EncodedRequest { frame: writer.finish(), heap_keys, + version: VERSION, }) } @@ -602,6 +651,10 @@ impl BoundedWriter { self.bytes(&value.to_le_bytes()) } + fn i32(&mut self, value: i32) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + fn patch_bytes(&mut self, offset: usize, bytes: &[u8]) { self.bytes[offset..offset + bytes.len()].copy_from_slice(bytes); } @@ -1287,6 +1340,8 @@ mod tests { &mut source, 100, 4096, + None, + Duration::from_secs(2), ) .unwrap(); @@ -1331,6 +1386,66 @@ mod tests { assert_eq!(encoded.heap_keys, vec![page]); } + #[test] + fn scheduled_external_request_encodes_tenant_priority_and_timeout() { + let page = [0, 0, 8]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, -0.5])]; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + 9002, + "sha256://opaque", + 2, + 2, + ExternalTensorDtype::F16, + ), + )])); + let scheduling = TileMaxsimScheduling { + tenant: "tenant-a".to_owned(), + priority: 17, + }; + let encoded = encode_external_request( + 20, + "contract@1", + &query, + &mut candidates, + &mut source, + 100, + 4096, + Some(&scheduling), + Duration::from_millis(4_000), + ) + .unwrap(); + + assert_eq!(encoded.version, SCHEDULED_EXTERNAL_VERSION); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + SCHEDULED_EXTERNAL_VERSION + ); + assert_eq!( + i32::from_le_bytes(encoded.frame[44..48].try_into().unwrap()), + 17 + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[48..52].try_into().unwrap()), + 4_000 + ); + let tenant_len = + u32::from_le_bytes(encoded.frame[52..56].try_into().unwrap()) as usize; + let contract_len = + u32::from_le_bytes(encoded.frame[40..44].try_into().unwrap()) as usize; + assert_eq!( + &encoded.frame[56 + contract_len..56 + contract_len + tenant_len], + b"tenant-a" + ); + } + #[test] fn external_backend_maps_scores_without_exposing_public_ids() { let page = [0, 0, 8]; @@ -1402,6 +1517,8 @@ mod tests { &mut make_source(3, 1), 100, 4096, + None, + Duration::from_secs(2), ), Err(RerankError::TensorMismatch) )); @@ -1416,6 +1533,8 @@ mod tests { &mut make_source(2, 100), 1000, 64, + None, + Duration::from_secs(2), ), Err(RerankError::RequestTooLarge) )); diff --git a/src/index/vchordrq/scanners/maxsim/profile.rs b/src/index/vchordrq/scanners/maxsim/profile.rs new file mode 100644 index 00000000..d2590a9f --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/profile.rs @@ -0,0 +1,181 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2026 Hu Xinjing + +use std::cell::RefCell; +use std::time::{Duration, Instant}; + +#[derive(Default)] +pub(super) struct MaxsimProfile { + pub query_tokens: u64, + pub token_search_calls: u64, + pub token_search_us: u64, + pub token_search_results: u64, + pub token_refine_calls: u64, + pub token_refine_us: u64, + pub accurate_token_hits: u64, + pub rough_token_hits: u64, + pub hit_collect_us: u64, + pub hit_updates: u64, + pub page_token_updates: u64, + pub hit_sort_us: u64, + pub page_aggregate_us: u64, + pub aggregated_pages: u64, + pub preflight_us: u64, + pub candidate_generation_us: u64, + pub generated_candidates: u64, + pub visibility_us: u64, + pub visible_candidates: u64, + pub descriptor_us: u64, + pub descriptors: u64, + pub sidecar_us: u64, + pub result_finalize_us: u64, + pub returned_rows: u64, + pub total_us: u64, +} + +impl MaxsimProfile { + fn to_json(&self) -> String { + format!( + concat!( + "{{", + "\"schema_version\":1,", + "\"query_tokens\":{},", + "\"token_search_calls\":{},", + "\"token_search_us\":{},", + "\"token_search_results\":{},", + "\"token_refine_calls\":{},", + "\"token_refine_us\":{},", + "\"accurate_token_hits\":{},", + "\"rough_token_hits\":{},", + "\"hit_collect_us\":{},", + "\"hit_updates\":{},", + "\"page_token_updates\":{},", + "\"hit_sort_us\":{},", + "\"page_aggregate_us\":{},", + "\"aggregated_pages\":{},", + "\"preflight_us\":{},", + "\"candidate_generation_us\":{},", + "\"generated_candidates\":{},", + "\"visibility_us\":{},", + "\"visible_candidates\":{},", + "\"descriptor_us\":{},", + "\"descriptors\":{},", + "\"sidecar_us\":{},", + "\"result_finalize_us\":{},", + "\"returned_rows\":{},", + "\"total_us\":{}", + "}}" + ), + self.query_tokens, + self.token_search_calls, + self.token_search_us, + self.token_search_results, + self.token_refine_calls, + self.token_refine_us, + self.accurate_token_hits, + self.rough_token_hits, + self.hit_collect_us, + self.hit_updates, + self.page_token_updates, + self.hit_sort_us, + self.page_aggregate_us, + self.aggregated_pages, + self.preflight_us, + self.candidate_generation_us, + self.generated_candidates, + self.visibility_us, + self.visible_candidates, + self.descriptor_us, + self.descriptors, + self.sidecar_us, + self.result_finalize_us, + self.returned_rows, + self.total_us, + ) + } +} + +thread_local! { + static ACTIVE_PROFILE: RefCell> = const { RefCell::new(None) }; +} + +pub(super) struct ProfileGuard { + enabled: bool, + finished: bool, +} + +impl ProfileGuard { + pub fn start(enabled: bool) -> Self { + if enabled { + ACTIVE_PROFILE.with(|slot| { + *slot.borrow_mut() = Some(MaxsimProfile::default()); + }); + } + Self { + enabled, + finished: false, + } + } + + pub fn finish(mut self, total: Duration) { + if self.enabled { + let profile = ACTIVE_PROFILE.with(|slot| { + let mut profile = slot.borrow_mut().take(); + if let Some(profile) = profile.as_mut() { + profile.total_us = duration_us(total); + } + profile + }); + if let Some(profile) = profile { + pgrx::notice!("vchordrq_maxsim_profile {}", profile.to_json()); + } + } + self.finished = true; + } +} + +impl Drop for ProfileGuard { + fn drop(&mut self) { + if self.enabled && !self.finished { + ACTIVE_PROFILE.with(|slot| { + slot.borrow_mut().take(); + }); + } + } +} + +pub(super) struct ProfileTimer(Option); + +impl ProfileTimer { + pub fn start() -> Self { + let enabled = ACTIVE_PROFILE.with(|slot| slot.borrow().is_some()); + Self(enabled.then(Instant::now)) + } + + pub fn elapsed(self) -> Duration { + self.0.map_or(Duration::ZERO, |started| started.elapsed()) + } +} + +pub(super) fn update(f: impl FnOnce(&mut MaxsimProfile)) { + ACTIVE_PROFILE.with(|slot| { + if let Some(profile) = slot.borrow_mut().as_mut() { + f(profile); + } + }); +} + +pub(super) fn duration_us(duration: Duration) -> u64 { + duration.as_micros().min(u128::from(u64::MAX)) as u64 +} diff --git a/src/index/vchordrq/scanners/maxsim/search.rs b/src/index/vchordrq/scanners/maxsim/search.rs index 915fb87f..65443bc0 100644 --- a/src/index/vchordrq/scanners/maxsim/search.rs +++ b/src/index/vchordrq/scanners/maxsim/search.rs @@ -19,6 +19,7 @@ use super::external::{ ExternalTensorStorage, resolve_external_tensor_source, validate_descriptor, }; use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; +use super::profile; use super::rerank::RerankError; use crate::index::fetcher::{ Fetcher, FilterableTuple, HeapFetcher, Tuple, TupleAttribute, ctid_to_key, @@ -59,6 +60,8 @@ fn execute_external_search( candidate_limit: i32, top_k: i32, ) -> Result, RerankError> { + let profile_guard = profile::ProfileGuard::start(gucs::vchordrq_maxsim_profile()); + let total_timer = profile::ProfileTimer::start(); validate_search_limits(candidate_limit, top_k)?; if !matches!(gucs::vchordrq_maxsim_backend(), PostgresMaxsimBackend::Gpu) { return Err(RerankError::Configuration( @@ -66,6 +69,7 @@ fn execute_external_search( )); } + let preflight_timer = profile::ProfileTimer::start(); let binding = resolve_external_tensor_source(index_oid)?; let index_lock = RelationLock::open(index_oid, pgrx::pg_sys::AccessShareLock as _)?; let heap_lock = RelationLock::open(binding.heap_oid, pgrx::pg_sys::AccessShareLock as _)?; @@ -94,21 +98,41 @@ fn execute_external_search( } let query_vectors = unsafe { opfamily.input_vectors(query.datum()) }.ok_or(RerankError::TensorMismatch)?; + profile::update(|profile| { + profile.query_tokens = query_vectors.len() as u64; + }); // The registry resolution happens before the index read, and this // privilege-only SELECT is planned/executed before any candidate CTID is // generated. It fails early when the caller cannot project the registered // descriptor columns. The same query shape is used for the actual fetch. preflight_descriptor_access(&binding)?; + let preflight_elapsed = preflight_timer.elapsed(); + profile::update(|profile| { + profile.preflight_us += profile::duration_us(preflight_elapsed); + }); + let candidate_timer = profile::ProfileTimer::start(); let candidates = generate_candidates( index_lock.raw(), opfamily, query.datum(), candidate_limit as u32, )?; + let candidate_elapsed = candidate_timer.elapsed(); + profile::update(|profile| { + profile.candidate_generation_us += profile::duration_us(candidate_elapsed); + profile.generated_candidates = candidates.len() as u64; + }); + let visibility_timer = profile::ProfileTimer::start(); let resolved = resolve_visible_candidates(index_lock.raw(), heap_lock.raw(), candidates.into_iter())?; + let visibility_elapsed = visibility_timer.elapsed(); + profile::update(|profile| { + profile.visibility_us += profile::duration_us(visibility_elapsed); + profile.visible_candidates = resolved.len() as u64; + }); + let descriptor_timer = profile::ProfileTimer::start(); let (mut source, public_ids) = load_visible_descriptors(&binding, &resolved)?; let visible_candidates = resolved .into_iter() @@ -118,6 +142,11 @@ fn execute_external_search( .then_some(resolved.candidate) }) .collect::>(); + let descriptor_elapsed = descriptor_timer.elapsed(); + profile::update(|profile| { + profile.descriptor_us += profile::duration_us(descriptor_elapsed); + profile.descriptors = public_ids.len() as u64; + }); let endpoint = gucs::vchordrq_maxsim_gpu_endpoint() .map(|endpoint| endpoint.to_string_lossy().into_owned()) .unwrap_or_default(); @@ -129,7 +158,11 @@ fn execute_external_search( gucs::vchordrq_maxsim_gpu_max_batch_tokens() as usize, gucs::vchordrq_maxsim_gpu_max_batch_bytes() as usize, ); + if let Some(tenant) = gucs::vchordrq_maxsim_tenant() { + backend = backend.with_scheduling(tenant, gucs::vchordrq_maxsim_priority()); + } let mut candidate_iter = visible_candidates.into_iter(); + let sidecar_timer = profile::ProfileTimer::start(); let exact = backend.rerank(&query_vectors, &mut candidate_iter, &mut source)?; let mut rows = exact .map(|result| { @@ -139,17 +172,28 @@ fn execute_external_search( Ok((result.distance, public_id)) }) .collect::, RerankError>>()?; + let sidecar_elapsed = sidecar_timer.elapsed(); + profile::update(|profile| { + profile.sidecar_us += profile::duration_us(sidecar_elapsed); + }); + let result_finalize_timer = profile::ProfileTimer::start(); rows.sort_unstable_by(|(left_distance, left_id), (right_distance, right_id)| { left_distance .cmp(right_distance) .then_with(|| left_id.cmp(right_id)) }); rows.truncate(top_k as usize); - Ok(rows + let output = rows .into_iter() .map(|(distance, public_id)| (public_id, -distance.to_f32())) - .collect::>() - .into_iter()) + .collect::>(); + let result_finalize_elapsed = result_finalize_timer.elapsed(); + profile::update(|profile| { + profile.result_finalize_us += profile::duration_us(result_finalize_elapsed); + profile.returned_rows = output.len() as u64; + }); + profile_guard.finish(total_timer.elapsed()); + Ok(output.into_iter()) } #[derive(Clone, Copy)] diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index a7e7af4e..932660c2 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -1431,3 +1431,756 @@ REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop ON sql_drop EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); + +-- Exact TileMaxSim over a caller-scoped tensor set. This source registry is +-- deliberately relation-keyed rather than index-keyed: the caller owns ACL, +-- graph, and application filtering decisions. VectorChord accepts the full +-- visible ID set without an artificial count cap, loads its tensor pages, and +-- applies the configured GPU cache policy. + +CREATE TABLE _vchordrq_tilemaxsim_sources ( + source_oid oid PRIMARY KEY, + model_contract_id text NOT NULL CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint NOT NULL CHECK ( + tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_rows_attnum smallint NOT NULL CHECK ( + tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dim_attnum smallint NOT NULL CHECK ( + tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dtype_attnum smallint NOT NULL CHECK ( + tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_checksum_attnum smallint NOT NULL CHECK ( + tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_tilemaxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_tilemaxsim_source( + source_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + public_id_is_unique boolean; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor descriptor columns must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be external_ref or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.oid, c.relowner + INTO source_oid, source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid + AND c.relkind IN ('r', 'm'); + IF source_oid IS NULL THEN + RAISE EXCEPTION 'source_relation % must be a table or materialized view', source_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may register a TileMaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL integer or bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'public ID column % must have a non-partial single-key unique index', + public_id_column; + END IF; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a TileMaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL integer or bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION 'external_ref sources must not specify a descriptor relation'; + END IF; + tensor_relation_oid := source_oid; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'TileMaxSim tensor source columns must be distinct'; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_tilemaxsim_sources ( + source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) + ON CONFLICT (source_oid) DO UPDATE SET + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + source_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_tilemaxsim_source(source_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_owner oid; + removed_count bigint; +BEGIN + IF source_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid AND c.relkind IN ('r', 'm'); + IF ext_schema IS NULL + OR source_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may unregister its TileMaxSim source'; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources WHERE source_oid = $1', + ext_schema + ) USING source_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_source_info(source_relation regclass) +RETURNS TABLE( + registered_source regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + bound_model_contract_id text; + bound_storage text; + descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + EXECUTE pg_catalog.format( + 'SELECT source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_tilemaxsim_sources + WHERE source_oid = $1', + ext_schema + ) INTO + source_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING source_relation::oid; + IF source_oid IS NULL THEN + RAISE EXCEPTION 'TileMaxSim tensor source is not registered for relation %', source_relation; + END IF; + + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_oid AND c.relkind IN ('r', 'm'); + IF source_owner IS NULL THEN + RAISE EXCEPTION 'registered TileMaxSim source relation is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, source_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered TileMaxSim source'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, NULL::oid) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = source_oid + AND a.attnum = expected.attnum + AND (expected.atttypid IS NULL OR a.atttypid = expected.atttypid) + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.atttypid IS NOT NULL + OR a.atttypid IN ('integer'::regtype, 'bigint'::regtype); + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered TileMaxSim source columns are invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim source public ID is no longer unique'; + END IF; + + IF bound_storage = 'external_ref' THEN + IF descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered external_ref TileMaxSim binding is invalid'; + END IF; + tensor_relation_oid := source_oid; + ELSIF bound_storage = 'external_relation' THEN + IF descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered external_relation TileMaxSim binding is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_oid AND c.relkind IN ('r', 'm'); + IF NOT FOUND OR NOT pg_catalog.has_table_privilege(caller_oid, descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor relation is unavailable'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid IN ('integer'::regtype, 'bigint'::regtype) + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + RAISE EXCEPTION 'registered TileMaxSim storage is invalid'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 5 THEN + RAISE EXCEPTION 'registered TileMaxSim tensor descriptor columns are invalid'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + source_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query vector[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_vector_wrapper'; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query halfvec[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_halfvec_wrapper'; + +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) FROM PUBLIC; +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) TO PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_tilemaxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL OR pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_tilemaxsim_sources', ext_schema) + ) IS NULL THEN + RETURN; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE ( + d.objid = s.source_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.model_contract_attnum, + s.public_id_attnum, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_ref_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_rows_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dim_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dtype_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_checksum_attnum END + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.descriptor_public_id_attnum, + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_tilemaxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_tilemaxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_tilemaxsim_source_sql_drop(); diff --git a/tests/vchordrq/tilemaxsim_source_registry.slt b/tests/vchordrq/tilemaxsim_source_registry.slt new file mode 100644 index 00000000..340bfd7d --- /dev/null +++ b/tests/vchordrq/tilemaxsim_source_registry.slt @@ -0,0 +1,91 @@ +# Caller-scoped TileMaxSim binds external tensor descriptors directly to an +# application relation. It must not require or inspect a VectorChord ANN index. + +statement ok +CREATE TABLE tilemaxsim_source_test ( + id integer PRIMARY KEY, + model_contract text NOT NULL, + embedding vector(2) +); + +statement ok +CREATE TABLE tilemaxsim_descriptor_test ( + source_id integer PRIMARY KEY, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement ok +INSERT INTO tilemaxsim_source_test VALUES + (1, 'colqwen@test', '[1,0]'), + (2, 'colqwen@test', '[0,1]'); + +statement ok +INSERT INTO tilemaxsim_descriptor_test VALUES + (1, 'tensor://1', 2, 2, 'float32', 'sha256:test1'), + (2, 'tensor://2', 2, 2, 'float32', 'sha256:test2'); + +statement ok +SELECT vchordrq_register_tilemaxsim_source( + source_relation => 'tilemaxsim_source_test'::regclass, + model_contract_id => ' colqwen@test ', + storage => 'EXTERNAL_RELATION', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'tilemaxsim_descriptor_test'::regclass, + descriptor_public_id_column => 'source_id' +); + +query TTTT +SELECT registered_source::text, source_storage, + descriptor_relation::text, public_id_column::text +FROM vchordrq_tilemaxsim_source_info('tilemaxsim_source_test'::regclass); +---- +tilemaxsim_source_test external_relation tilemaxsim_descriptor_test id + +# Candidate validation occurs before external tensor or sidecar access. +statement error candidate_ids must not contain duplicates +SELECT * FROM vchordrq_tilemaxsim_rerank( + 'tilemaxsim_source_test'::regclass, + ARRAY['[1,0]'::vector], + ARRAY[1,1]::bigint[], + 1 +); + +statement error top_k must be positive and no greater than candidate_ids length +SELECT * FROM vchordrq_tilemaxsim_rerank( + 'tilemaxsim_source_test'::regclass, + ARRAY['[1,0]'::halfvec], + ARRAY[1,2]::bigint[], + 3 +); + +# Attribute OIDs survive renames and the resolver returns the live name. +statement ok +ALTER TABLE tilemaxsim_descriptor_test RENAME tensor_ref TO tensor_location; + +query T +SELECT tensor_ref_column::text +FROM vchordrq_tilemaxsim_source_info('tilemaxsim_source_test'::regclass); +---- +tensor_location + +# Dropping a bound descriptor column removes the binding through sql_drop. +statement ok +ALTER TABLE tilemaxsim_descriptor_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) FROM _vchordrq_tilemaxsim_sources; +---- +0 + +statement ok +DROP TABLE tilemaxsim_source_test, tilemaxsim_descriptor_test; From 8e37405b8b7b162d41c02454ca7ba0a07dea08a6 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:25:30 +0800 Subject: [PATCH 312/324] fix(tilemaxsim): batch external rerank requests --- src/index/vchordrq/scanners/maxsim.rs | 29 ++ src/index/vchordrq/scanners/maxsim/exact.rs | 18 +- src/index/vchordrq/scanners/maxsim/gpu.rs | 496 +++++++++++++++++-- src/index/vchordrq/scanners/maxsim/search.rs | 18 +- 4 files changed, 494 insertions(+), 67 deletions(-) diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 4f76c0db..1e0fe486 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -49,6 +49,18 @@ use vector::rabitq4::Rabitq4Owned; use vector::rabitq8::Rabitq8Owned; use vector::vect::VectOwned; +fn retain_top_k(best: &mut std::collections::BinaryHeap, limit: usize, item: T) { + if limit == 0 { + return; + } + if best.len() < limit { + best.push(item); + } else if best.peek().is_some_and(|worst| item < *worst) { + best.pop(); + best.push(item); + } +} + pub struct MaxsimBuilder { opfamily: Opfamily, orderbys: Vec>>, @@ -929,3 +941,20 @@ where { f } + +#[cfg(test)] +mod tests { + use super::retain_top_k; + use std::collections::BinaryHeap; + + #[test] + fn bounded_top_k_keeps_global_order_and_ties() { + let mut best = BinaryHeap::new(); + for item in [(3, 30), (1, 20), (1, 10), (2, 40), (0, 50)] { + retain_top_k(&mut best, 3, item); + } + let mut rows = best.into_vec(); + rows.sort_unstable(); + assert_eq!(rows, vec![(0, 50), (1, 10), (1, 20)]); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/exact.rs b/src/index/vchordrq/scanners/maxsim/exact.rs index 32d36ec4..d43bf9f7 100644 --- a/src/index/vchordrq/scanners/maxsim/exact.rs +++ b/src/index/vchordrq/scanners/maxsim/exact.rs @@ -23,7 +23,7 @@ use distance::Distance; use pgrx::datum::{Array, DatumWithOid, FromDatum, IntoDatum}; use pgrx::iter::TableIterator; use pgrx::{name, pg_extern}; -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::{BTreeMap, BTreeSet, BinaryHeap}; use std::time::Duration; use vchordrq::types::OwnedVector; use vector::VectorBorrowed; @@ -190,15 +190,19 @@ fn execute_rerank( backend = backend.with_scheduling(tenant, gucs::vchordrq_maxsim_priority()); } let sidecar_timer = profile::ProfileTimer::start(); - let exact = backend.rerank(&query, &mut candidates, &mut source)?; - let mut rows = exact - .map(|result| { + let result_limit = top_k as usize; + let mut best = BinaryHeap::with_capacity(result_limit.saturating_add(1)); + backend.rerank_batches(&query, &mut candidates, &mut source, |batch| { + for result in batch { let public_id = public_ids.get(&result.heap_key).copied().ok_or_else(|| { RerankError::Protocol("sidecar result has no visible public ID".into()) })?; - Ok((result.distance, public_id)) - }) - .collect::, RerankError>>()?; + let row = (result.distance, public_id); + super::retain_top_k(&mut best, result_limit, row); + } + Ok(()) + })?; + let mut rows = best.into_vec(); profile::update(|profile| { profile.sidecar_us += profile::duration_us(sidecar_timer.elapsed()); }); diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs index 7e0928a6..a204566e 100644 --- a/src/index/vchordrq/scanners/maxsim/gpu.rs +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -22,7 +22,7 @@ use std::cmp::Reverse; use std::collections::BinaryHeap; use std::mem::size_of_val; use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; +use std::time::{Duration, Instant}; use vchordrq::types::OwnedVector; const MAGIC: &[u8; 4] = b"VCTM"; @@ -33,6 +33,11 @@ const REQUEST_KIND: u16 = 1; const RESPONSE_KIND: u16 = 2; const HEADER_LEN: usize = 24; const MAX_REMOTE_ERROR_BYTES: usize = 64 * 1024; +const MAX_EXTERNAL_CANDIDATES_PER_BATCH: usize = 65_536; +const MAX_MODEL_CONTRACT_BYTES: usize = 512; +const MAX_TENSOR_REF_BYTES: usize = 4096; +const MAX_CHECKSUM_BYTES: usize = 512; +const MAX_TENANT_BYTES: usize = 256; static NEXT_REQUEST_ID: AtomicU64 = AtomicU64::new(1); static LAST_FALLBACK_WARNING_SECONDS: AtomicU64 = AtomicU64::new(0); @@ -169,47 +174,104 @@ impl GpuExternalTileMaxsimBackend { } impl GpuExternalTileMaxsimBackend { + #[cfg(test)] pub(super) fn rerank( &mut self, query: &[OwnedVector], candidates: &mut dyn Iterator, source: &mut S, ) -> Result { - let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); - let encoded = encode_external_request( - request_id, - &self.model_contract_id, - query, - candidates, - source, - self.max_batch_tokens, - self.max_batch_bytes, - self.scheduling.as_ref(), - self.timeout, - )?; - if encoded.heap_keys.is_empty() { - return Ok(RerankResults { - inner: BinaryHeap::new(), - }); + let mut inner = BinaryHeap::new(); + self.rerank_batches(query, candidates, source, |batch| { + inner.extend(batch.inner); + Ok(()) + })?; + Ok(RerankResults { inner }) + } + + /// Execute one logical rerank as bounded sidecar requests. + /// + /// The callback is invoked after every complete batch so callers that only + /// need a global top-k do not have to retain every score. `self.timeout` is + /// a deadline for the whole logical query, rather than a fresh timeout for + /// each IPC request. + pub(super) fn rerank_batches( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + mut consume: F, + ) -> Result<(), RerankError> + where + S: CandidateTensorDescriptorSource, + F: FnMut(RerankResults) -> Result<(), RerankError>, + { + let deadline = Instant::now() + .checked_add(self.timeout) + .ok_or_else(|| RerankError::Transport("logical request deadline overflow".into()))?; + let mut pending = None; + + loop { + let descriptors = collect_external_batch( + &self.model_contract_id, + query, + candidates, + source, + self.max_batch_tokens, + self.max_batch_bytes, + self.scheduling.as_ref(), + &mut pending, + )?; + if descriptors.is_empty() { + return Ok(()); + } + + let remaining = remaining_logical_timeout(deadline)?; + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let encoded = encode_external_descriptors( + request_id, + &self.model_contract_id, + query, + &descriptors, + self.max_batch_tokens, + self.max_batch_bytes, + self.scheduling.as_ref(), + remaining, + )?; + let max_response_bytes = HEADER_LEN + .checked_add(8) + .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) + .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) + .ok_or(RerankError::RequestTooLarge)?; + let response = self.transport.round_trip( + &encoded.frame, + remaining_logical_timeout(deadline)?, + max_response_bytes, + )?; + consume(decode_response_for_version( + &response, + encoded.version, + request_id, + &encoded.heap_keys, + )?)?; } - let max_response_bytes = HEADER_LEN - .checked_add(8) - .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) - .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) - .ok_or(RerankError::RequestTooLarge)?; - let response = - self.transport - .round_trip(&encoded.frame, self.timeout, max_response_bytes)?; - decode_response_for_version(&response, encoded.version, request_id, &encoded.heap_keys) } } +fn remaining_logical_timeout(deadline: Instant) -> Result { + deadline + .checked_duration_since(Instant::now()) + .filter(|remaining| !remaining.is_zero()) + .ok_or_else(|| RerankError::Transport("logical request timed out".into())) +} + struct EncodedRequest { frame: Vec, heap_keys: Vec, version: u16, } +#[cfg(test)] fn encode_external_request( request_id: u64, model_contract_id: &str, @@ -221,11 +283,34 @@ fn encode_external_request( scheduling: Option<&TileMaxsimScheduling>, timeout: Duration, ) -> Result { - const MAX_MODEL_CONTRACT_BYTES: usize = 512; - const MAX_TENSOR_REF_BYTES: usize = 4096; - const MAX_CHECKSUM_BYTES: usize = 512; - const MAX_TENANT_BYTES: usize = 256; + let mut descriptors = Vec::new(); + for candidate in candidates { + if let Some(descriptor) = source.fetch(candidate)? { + descriptors.push(descriptor); + } + } + encode_external_descriptors( + request_id, + model_contract_id, + query, + &descriptors, + max_batch_tokens, + max_batch_bytes, + scheduling, + timeout, + ) +} +fn encode_external_descriptors( + request_id: u64, + model_contract_id: &str, + query: &[OwnedVector], + descriptors: &[ExternalTensorDescriptor], + max_batch_tokens: usize, + max_batch_bytes: usize, + scheduling: Option<&TileMaxsimScheduling>, + timeout: Duration, +) -> Result { if model_contract_id.is_empty() || model_contract_id.len() > MAX_MODEL_CONTRACT_BYTES || model_contract_id.chars().any(char::is_control) @@ -290,13 +375,13 @@ fn encode_external_request( } encode_tensor_values(&mut writer, query, dtype)?; - let mut heap_keys = Vec::new(); - for candidate in candidates { - let Some(descriptor) = source.fetch(candidate)? else { - continue; - }; + if descriptors.len() > MAX_EXTERNAL_CANDIDATES_PER_BATCH { + return Err(RerankError::RequestTooLarge); + } + let mut heap_keys = Vec::with_capacity(descriptors.len()); + for descriptor in descriptors { validate_external_for_request( - &descriptor, + descriptor, dimension, external_dtype, MAX_TENSOR_REF_BYTES, @@ -352,6 +437,137 @@ fn encode_external_request( }) } +#[allow(clippy::too_many_arguments)] +fn collect_external_batch( + model_contract_id: &str, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, + scheduling: Option<&TileMaxsimScheduling>, + pending: &mut Option, +) -> Result, RerankError> { + validate_external_request_identity(model_contract_id, scheduling)?; + let (dtype, dimension) = tensor_metadata(query)?; + let external_dtype = match dtype { + TensorDtype::F32 => ExternalTensorDtype::F32, + TensorDtype::F16 => ExternalTensorDtype::F16, + }; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let query_bytes = tensor_bytes(query_rows, dimension, dtype)?; + let mut total_tokens = query.len(); + let mut declared_tensor_bytes = query_bytes; + let mut frame_bytes = external_request_base_bytes(model_contract_id, scheduling, query_bytes)?; + if total_tokens > max_batch_tokens + || declared_tensor_bytes > max_batch_bytes + || frame_bytes > max_batch_bytes + { + return Err(RerankError::RequestTooLarge); + } + + let mut batch = Vec::new(); + while batch.len() < MAX_EXTERNAL_CANDIDATES_PER_BATCH { + let descriptor = if let Some(descriptor) = pending.take() { + descriptor + } else { + let mut fetched = None; + for candidate in &mut *candidates { + if let Some(descriptor) = source.fetch(candidate)? { + fetched = Some(descriptor); + break; + } + } + let Some(descriptor) = fetched else { + break; + }; + descriptor + }; + + validate_external_for_request( + &descriptor, + dimension, + external_dtype, + MAX_TENSOR_REF_BYTES, + MAX_CHECKSUM_BYTES, + )?; + let next_tokens = total_tokens + .checked_add(descriptor.rows as usize) + .ok_or(RerankError::RequestTooLarge)?; + let next_declared_bytes = declared_tensor_bytes + .checked_add(tensor_bytes(descriptor.rows, dimension, dtype)?) + .ok_or(RerankError::RequestTooLarge)?; + let descriptor_frame_bytes = 16usize + .checked_add(descriptor.tensor_ref.len()) + .and_then(|size| size.checked_add(descriptor.checksum.len())) + .ok_or(RerankError::RequestTooLarge)?; + let next_frame_bytes = frame_bytes + .checked_add(descriptor_frame_bytes) + .ok_or(RerankError::RequestTooLarge)?; + if next_tokens > max_batch_tokens + || next_declared_bytes > max_batch_bytes + || next_frame_bytes > max_batch_bytes + { + if batch.is_empty() { + return Err(RerankError::RequestTooLarge); + } + *pending = Some(descriptor); + break; + } + + total_tokens = next_tokens; + declared_tensor_bytes = next_declared_bytes; + frame_bytes = next_frame_bytes; + batch.push(descriptor); + } + Ok(batch) +} + +fn validate_external_request_identity( + model_contract_id: &str, + scheduling: Option<&TileMaxsimScheduling>, +) -> Result<(), RerankError> { + if model_contract_id.is_empty() + || model_contract_id.len() > MAX_MODEL_CONTRACT_BYTES + || model_contract_id.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + if let Some(scheduling) = scheduling { + if scheduling.tenant.is_empty() + || scheduling.tenant.len() > MAX_TENANT_BYTES + || scheduling.tenant.chars().any(char::is_control) + || !(-100..=100).contains(&scheduling.priority) + { + return Err(RerankError::Configuration( + "TileMaxSim scheduler tenant or priority is invalid", + )); + } + } + Ok(()) +} + +fn external_request_base_bytes( + model_contract_id: &str, + scheduling: Option<&TileMaxsimScheduling>, + query_bytes: usize, +) -> Result { + let fixed = if scheduling.is_some() { + 56usize + } else { + 44usize + }; + fixed + .checked_add(model_contract_id.len()) + .and_then(|size| { + size.checked_add(scheduling.map_or(0, |scheduling| scheduling.tenant.len())) + }) + .and_then(|size| size.checked_add(query_bytes)) + .ok_or(RerankError::RequestTooLarge) +} + fn tensor_bytes(rows: u32, dimension: u32, dtype: TensorDtype) -> Result { let scalar_bytes = match dtype { TensorDtype::F32 => 4usize, @@ -737,8 +953,6 @@ impl TileMaxsimTransport for UnixSocketTransport { timeout: Duration, max_response_bytes: usize, ) -> Result, RerankError> { - use std::time::Instant; - if self.endpoint.is_empty() { return Err(RerankError::Transport("endpoint is empty".into())); } @@ -778,7 +992,7 @@ impl TileMaxsimTransport for UnixSocketTransport { #[cfg(unix)] fn connect_interruptible( endpoint: &str, - deadline: std::time::Instant, + deadline: Instant, ) -> Result { use std::os::fd::{AsRawFd, FromRawFd, OwnedFd}; @@ -904,10 +1118,7 @@ fn update_fd_flag( } #[cfg(unix)] -fn wait_for_connect( - fd: std::os::fd::RawFd, - deadline: std::time::Instant, -) -> Result<(), RerankError> { +fn wait_for_connect(fd: std::os::fd::RawFd, deadline: Instant) -> Result<(), RerankError> { loop { pgrx::check_for_interrupts!(); let remaining = remaining_until(deadline)?; @@ -952,9 +1163,9 @@ fn wait_for_connect( } #[cfg(unix)] -fn remaining_until(deadline: std::time::Instant) -> Result { +fn remaining_until(deadline: Instant) -> Result { deadline - .checked_duration_since(std::time::Instant::now()) + .checked_duration_since(Instant::now()) .filter(|remaining| !remaining.is_zero()) .ok_or_else(|| RerankError::Transport("request timed out".into())) } @@ -968,7 +1179,7 @@ fn last_transport_error() -> RerankError { fn write_interruptible( stream: &mut std::os::unix::net::UnixStream, mut bytes: &[u8], - deadline: std::time::Instant, + deadline: Instant, ) -> Result<(), RerankError> { use std::io::Write; @@ -986,7 +1197,7 @@ fn write_interruptible( Err(error) => return Err(RerankError::Transport(error.to_string())), } pgrx::check_for_interrupts!(); - if std::time::Instant::now() >= deadline { + if Instant::now() >= deadline { return Err(RerankError::Transport("request timed out".into())); } } @@ -997,7 +1208,7 @@ fn write_interruptible( fn read_interruptible( stream: &mut std::os::unix::net::UnixStream, mut bytes: &mut [u8], - deadline: std::time::Instant, + deadline: Instant, ) -> Result<(), RerankError> { use std::io::Read; @@ -1015,7 +1226,7 @@ fn read_interruptible( Err(error) => return Err(RerankError::Transport(error.to_string())), } pgrx::check_for_interrupts!(); - if std::time::Instant::now() >= deadline { + if Instant::now() >= deadline { return Err(RerankError::Transport("request timed out".into())); } } @@ -1046,7 +1257,9 @@ mod tests { }; use super::super::rerank::CandidateTensor; use super::*; + use std::cell::RefCell; use std::collections::BTreeMap; + use std::rc::Rc; use vector::vect::VectOwned; struct MockTensorSource(BTreeMap>); @@ -1102,6 +1315,52 @@ mod tests { } } + #[derive(Default)] + struct BatchObservations { + candidate_counts: Vec, + transport_timeouts: Vec, + scheduled_timeouts_ms: Vec, + } + + struct BatchingTransport { + observations: Rc>, + delay: Duration, + } + + impl TileMaxsimTransport for BatchingTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + let request_id = u64::from_le_bytes(request[8..16].try_into().unwrap()); + let version = u16::from_le_bytes(request[4..6].try_into().unwrap()); + let candidate_count = u32::from_le_bytes(request[32..36].try_into().unwrap()); + let call = { + let mut observations = self.observations.borrow_mut(); + let call = observations.candidate_counts.len(); + observations.candidate_counts.push(candidate_count); + observations.transport_timeouts.push(timeout); + if version == SCHEDULED_EXTERNAL_VERSION { + observations + .scheduled_timeouts_ms + .push(u32::from_le_bytes(request[48..52].try_into().unwrap())); + } + call + }; + std::thread::sleep(self.delay); + let similarities = (0..candidate_count) + .map(|candidate_id| (candidate_id, call as f32 * 10.0 + candidate_id as f32)) + .collect::>(); + Ok(success_response_with_version( + version, + request_id, + &similarities, + )) + } + } + fn vector(values: &[f32]) -> OwnedVector { OwnedVector::Vecf32(VectOwned::new(values.to_vec())) } @@ -1436,10 +1695,8 @@ mod tests { u32::from_le_bytes(encoded.frame[48..52].try_into().unwrap()), 4_000 ); - let tenant_len = - u32::from_le_bytes(encoded.frame[52..56].try_into().unwrap()) as usize; - let contract_len = - u32::from_le_bytes(encoded.frame[40..44].try_into().unwrap()) as usize; + let tenant_len = u32::from_le_bytes(encoded.frame[52..56].try_into().unwrap()) as usize; + let contract_len = u32::from_le_bytes(encoded.frame[40..44].try_into().unwrap()) as usize; assert_eq!( &encoded.frame[56 + contract_len..56 + contract_len + tenant_len], b"tenant-a" @@ -1485,6 +1742,139 @@ mod tests { assert_eq!(results[0].distance.to_f32(), -3.5); } + #[test] + fn external_backend_splits_one_logical_query_and_shares_its_deadline() { + let pages = [[0, 0, 1], [0, 0, 2], [0, 0, 3]]; + let candidates = pages.map(|heap_key| PageCandidate { + approximate_distance: Distance::ZERO, + heap_key, + }); + let query = vec![vector(&[1.0, 0.0])]; + let mut candidate_iter = candidates.into_iter(); + let mut source = + MockDescriptorSource(BTreeMap::from_iter(candidates.into_iter().enumerate().map( + |(index, candidate)| { + ( + candidate.heap_key, + external_descriptor( + candidate, + index as i64, + &format!("object://immutable/tensor-{index}"), + 2, + 2, + ExternalTensorDtype::F32, + ), + ) + }, + ))); + let observations = Rc::new(RefCell::new(BatchObservations::default())); + let transport = BatchingTransport { + observations: Rc::clone(&observations), + delay: Duration::from_millis(5), + }; + let results = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_millis(100), + 5, + 4096, + ) + .with_scheduling("tenant-a".into(), 9) + .rerank(&query, &mut candidate_iter, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 3); + let observations = observations.borrow(); + assert_eq!(observations.candidate_counts, vec![2, 1]); + assert!(observations.transport_timeouts[1] < observations.transport_timeouts[0]); + assert!(observations.scheduled_timeouts_ms[1] < observations.scheduled_timeouts_ms[0]); + } + + #[test] + fn external_backend_rejects_a_single_unsplittable_tensor() { + let page = [0, 0, 1]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + 1, + "object://immutable/tensor", + 5, + 2, + ExternalTensorDtype::F32, + ), + )])); + let observations = Rc::new(RefCell::new(BatchObservations::default())); + let transport = BatchingTransport { + observations: Rc::clone(&observations), + delay: Duration::ZERO, + }; + let result = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_millis(100), + 5, + 4096, + ) + .rerank(&query, &mut candidates, &mut source); + + assert!(matches!(result, Err(RerankError::RequestTooLarge))); + assert!(observations.borrow().candidate_counts.is_empty()); + } + + #[test] + fn external_backend_does_not_refresh_timeout_for_later_batches() { + let pages = [[0, 0, 1], [0, 0, 2], [0, 0, 3]]; + let candidates = pages.map(|heap_key| PageCandidate { + approximate_distance: Distance::ZERO, + heap_key, + }); + let query = vec![vector(&[1.0, 0.0])]; + let mut candidate_iter = candidates.into_iter(); + let mut source = + MockDescriptorSource(BTreeMap::from_iter(candidates.into_iter().enumerate().map( + |(index, candidate)| { + ( + candidate.heap_key, + external_descriptor( + candidate, + index as i64, + &format!("object://immutable/tensor-{index}"), + 2, + 2, + ExternalTensorDtype::F32, + ), + ) + }, + ))); + let observations = Rc::new(RefCell::new(BatchObservations::default())); + let transport = BatchingTransport { + observations: Rc::clone(&observations), + delay: Duration::from_millis(20), + }; + let result = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_millis(5), + 5, + 4096, + ) + .rerank(&query, &mut candidate_iter, &mut source); + + assert!(matches!( + result, + Err(RerankError::Transport(message)) if message == "logical request timed out" + )); + assert_eq!(observations.borrow().candidate_counts, vec![2]); + } + #[test] fn external_request_rejects_shape_and_declared_payload_overflow() { let page = [0, 0, 9]; diff --git a/src/index/vchordrq/scanners/maxsim/search.rs b/src/index/vchordrq/scanners/maxsim/search.rs index 65443bc0..04baba34 100644 --- a/src/index/vchordrq/scanners/maxsim/search.rs +++ b/src/index/vchordrq/scanners/maxsim/search.rs @@ -34,7 +34,7 @@ use distance::Distance; use pgrx::datum::{DatumWithOid, FromDatum}; use pgrx::iter::TableIterator; use pgrx::{AnyArray, IntoDatum, name}; -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::{BTreeMap, BTreeSet, BinaryHeap}; use std::time::Duration; const MAX_EXPLICIT_CANDIDATES: i32 = 65_536; @@ -163,15 +163,19 @@ fn execute_external_search( } let mut candidate_iter = visible_candidates.into_iter(); let sidecar_timer = profile::ProfileTimer::start(); - let exact = backend.rerank(&query_vectors, &mut candidate_iter, &mut source)?; - let mut rows = exact - .map(|result| { + let result_limit = top_k as usize; + let mut best = BinaryHeap::with_capacity(result_limit.saturating_add(1)); + backend.rerank_batches(&query_vectors, &mut candidate_iter, &mut source, |batch| { + for result in batch { let public_id = public_ids.get(&result.heap_key).copied().ok_or_else(|| { RerankError::Protocol("sidecar result has no visible public ID".into()) })?; - Ok((result.distance, public_id)) - }) - .collect::, RerankError>>()?; + let row = (result.distance, public_id); + super::retain_top_k(&mut best, result_limit, row); + } + Ok(()) + })?; + let mut rows = best.into_vec(); let sidecar_elapsed = sidecar_timer.elapsed(); profile::update(|profile| { profile.sidecar_us += profile::duration_us(sidecar_elapsed); From 2e19e4f031b6a5b10b8a24fec8046e60417f1e0b Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:33:41 +0800 Subject: [PATCH 313/324] feat(tilemaxsimd): add production supervision and probes --- README.md | 5 + deploy/systemd/tilemaxsimd.env.example | 3 + deploy/systemd/tilemaxsimd.service | 34 +++ docs/TILEMAXSIM_CUDA_SIDECAR.md | 249 ++++++++---------- services/Dockerfile.tilemaxsimd | 5 + services/test_tilemaxsim_rust_daemon.py | 82 ++++++ services/tilemaxsimd/src/bin/tilemaxsimctl.rs | 82 ++++++ services/tilemaxsimd/src/main.rs | 126 ++++++--- 8 files changed, 406 insertions(+), 180 deletions(-) create mode 100644 deploy/systemd/tilemaxsimd.env.example create mode 100644 deploy/systemd/tilemaxsimd.service create mode 100644 services/tilemaxsimd/src/bin/tilemaxsimctl.rs diff --git a/README.md b/README.md index 6ae46489..1ae883d6 100644 --- a/README.md +++ b/README.md @@ -193,6 +193,11 @@ The optional status socket serves HTTP `GET /healthz` and Prometheus completed/error/timeout/disconnect outcomes, and global/per-tenant admission rejections without exporting tenant identifiers. +`GET /livez` reports process liveness separately from readiness. The packaged +`tilemaxsimctl` probe can wait on the status socket without curl or a TCP port. +An opt-in hardened systemd unit and environment example are provided under +`deploy/systemd`; the CUDA container uses the same probe for its health check. + The in-flight request budget is also expressed in GiB. A reader must reserve its complete declared frame after the fixed header is validated, and keeps that permit through completion, timeout, or disconnect. This bounds aggregate query diff --git a/deploy/systemd/tilemaxsimd.env.example b/deploy/systemd/tilemaxsimd.env.example new file mode 100644 index 00000000..08834ee2 --- /dev/null +++ b/deploy/systemd/tilemaxsimd.env.example @@ -0,0 +1,3 @@ +# This unit is opt-in. Do not enable it on nodes where TileMaxSim is disabled. +# Memory values are GiB; startup fails if the requested CUDA arena is unavailable. +TILEMAXSIMD_ARGS="--socket /run/vectorchord/tilemaxsim.sock --status-socket /run/vectorchord/tilemaxsim-status.sock --socket-mode 660 --status-socket-mode 660 --gpu-memory-gb 0=20 --gpu-workspace-gb 2 --host-cache-gb 8 --max-inflight-request-gb 1 --contract-root MODEL_CONTRACT_ID=/var/lib/vectorchord/tensors --scheduler-policy fair-priority --max-connections 256 --max-queued-requests 128 --max-tenant-queued-requests 16" diff --git a/deploy/systemd/tilemaxsimd.service b/deploy/systemd/tilemaxsimd.service new file mode 100644 index 00000000..e13b94ef --- /dev/null +++ b/deploy/systemd/tilemaxsimd.service @@ -0,0 +1,34 @@ +[Unit] +Description=VectorChord native CUDA TileMaxSim daemon +Documentation=https://github.com/HuXinjing/VectorChord +After=local-fs.target +StartLimitIntervalSec=60 +StartLimitBurst=3 + +[Service] +Type=exec +User=vectorchord +Group=vectorchord +RuntimeDirectory=vectorchord +RuntimeDirectoryMode=0750 +StateDirectory=vectorchord +StateDirectoryMode=0750 +EnvironmentFile=/etc/vectorchord/tilemaxsimd.env +ExecStart=/usr/local/bin/tilemaxsimd $TILEMAXSIMD_ARGS +ExecStartPost=/usr/local/bin/tilemaxsimctl --socket /run/vectorchord/tilemaxsim-status.sock --wait-timeout-ms 300000 +ExecReload=/bin/kill -HUP $MAINPID +Restart=on-failure +RestartSec=2 +TimeoutStopSec=120 +KillSignal=SIGTERM +UMask=0027 +LimitMEMLOCK=infinity +LimitNOFILE=1048576 +NoNewPrivileges=true +PrivateTmp=true +ProtectSystem=strict +ProtectHome=true +ReadWritePaths=/run/vectorchord /var/lib/vectorchord + +[Install] +WantedBy=multi-user.target diff --git a/docs/TILEMAXSIM_CUDA_SIDECAR.md b/docs/TILEMAXSIM_CUDA_SIDECAR.md index c1fd2968..f5986c25 100644 --- a/docs/TILEMAXSIM_CUDA_SIDECAR.md +++ b/docs/TILEMAXSIM_CUDA_SIDECAR.md @@ -1,173 +1,146 @@ -# VectorChord CUDA TileMaxSim Sidecar +# VectorChord Native CUDA TileMaxSim Daemon ## Scope -`services/tilemaxsim_cuda_sidecar.py` is the deployable CUDA executor for IPC -v1 and v2. PostgreSQL remains responsible for candidate generation, MVCC row -visibility, descriptor validation, and final public-ID mapping. The sidecar -performs bounded tensor resolution and exact TileMaxSim only. +`services/tilemaxsimd` is VectorChord's production CUDA execution service for +external TileMaxSim protocol v2/v3. PostgreSQL remains responsible for MVCC +visibility, hard filters, descriptor validation, and public-ID mapping. The +daemon owns only bounded tensor loading, host/GPU caches, scheduling, and exact +TileMaxSim execution. -The service does not fetch from S3, HTTP, or another network destination. -GBrain may populate a node-local immutable cache from its storage system. This -keeps storage credentials, application authorization, and routing outside the -VectorChord extension while retaining one VectorChord-owned retrieval call. +The daemon does not fetch S3 or HTTP objects and does not make authorization, +graph, Fact, community, or application-routing decisions. An application such +as GBrain publishes immutable tensors into a configured node-local shard root. -Production acceptance still requires an application-level corpus benchmark. -The service implementation alone is not that acceptance result. +The Python sidecar remains a reference and conformance implementation. New +production deployments should use the Rust/CUDA daemon. -## Runtime Requirements +## Fail-closed resource contract -- Linux with a Unix-domain socket; -- PyTorch built for the installed NVIDIA driver/CUDA runtime; -- one CUDA device by default (`--device cuda:0`); -- a service account that can read configured cache roots and create the socket; -- the PostgreSQL service account must be able to connect to the socket. +TileMaxSim is opt-in. Do not start this service when no GPU memory has been +explicitly assigned. Every `--gpu-memory-gb GPU=GB` value is required, uses GiB +rather than bytes, and is reserved with one CUDA allocation during startup. +An invalid device, insufficient free memory, failed pinned-host allocation, or +failed CUDA stream creation terminates the process before either socket becomes +ready. -CPU mode exists for conformance tests only: +The configured allocation is split into a persistent tensor arena and +`--gpu-workspace-gb`. The workspace must be smaller than every GPU allocation. +Host cache and aggregate in-flight request budgets are also configured in GiB. -```text -python3 -m services.tilemaxsim_cuda_sidecar --device cpu ... -``` +## Build the production image -The production command must select a CUDA device. Startup performs a CUDA -matrix multiply and synchronization, so an unusable device fails before the -socket is advertised. +The build and runtime CUDA images are explicit deployment inputs: -For container deployment, build the minimal service image on top of the -operations-approved PyTorch/CUDA runtime that matches the node driver: - -```text +```shell docker build \ - --build-arg BASE_IMAGE= \ - -f services/Dockerfile.tilemaxsim \ - -t vectorchord-tilemaxsim:local \ + --build-arg CUDA_DEVEL_IMAGE=nvidia/cuda:12.6.3-devel-ubuntu24.04 \ + --build-arg CUDA_RUNTIME_IMAGE=nvidia/cuda:12.6.3-runtime-ubuntu24.04 \ + -f services/Dockerfile.tilemaxsimd \ + -t vectorchord-tilemaxsimd:local \ . ``` -The Dockerfile intentionally has no default base image. CUDA/PyTorch/driver -compatibility is a deployment input and must not drift through an implicit -`latest` tag. +Pin both images by digest in a production build. The container health check +expects the status socket at +`/run/vectorchord/tilemaxsim-status.sock`; mount the socket directory so the +PostgreSQL process can reach the protocol socket. -## Content-Addressed Tensor Cache +## Start the daemon -Each model contract is mapped explicitly to one absolute cache root. A v2 -reference has exactly this form: - -```text -sha256://0123456789abcdef...64-lowercase-hex-characters +```shell +tilemaxsimd \ + --socket /run/vectorchord/tilemaxsim.sock \ + --status-socket /run/vectorchord/tilemaxsim-status.sock \ + --socket-mode 660 \ + --status-socket-mode 660 \ + --gpu-memory-gb 0=20 \ + --gpu-workspace-gb 2 \ + --host-cache-gb 8 \ + --max-inflight-request-gb 1 \ + --contract-root 'MODEL_CONTRACT_ID=/var/lib/vectorchord/tensors' \ + --scheduler-policy fair-priority \ + --max-connections 256 \ + --max-queued-requests 128 \ + --max-tenant-queued-requests 16 ``` -For digest `abcdef...`, the canonical row-major scalar payload is stored at: +Multiple GPU assignments may be supplied. Device selection and the requested +allocation are never inferred from all currently free VRAM. -```text -/ab/abcdef....bin -``` +The example systemd unit and environment file are in `deploy/systemd`. The unit +is deliberately opt-in: enabling it is the operator's explicit TileMaxSim/GPU +configuration. It waits for readiness, reloads immutable shard indexes with +SIGHUP, restarts on failure, and gives SIGTERM shutdown time to stop accepts and +drain admitted work. -The registered checksum must be `sha256:`. The service verifies: +## Health and metrics -- exact model-contract root mapping; -- reference/checksum agreement; -- regular file and exact descriptor byte length; -- SHA-256 content digest; -- finite f16/f32 values; -- descriptor shape and dtype. +The status Unix socket serves: -Directory and file symlinks below the configured root are rejected. Cache -publishers should write a temporary file, flush it, and atomically rename it to -the digest path only after the complete payload is available. Existing digest -paths are immutable. +- `GET /livez`: the status server is alive; +- `GET /healthz`: the protocol listener and scheduler are ready; +- `GET /metrics`: bounded Prometheus scheduler and outcome counters. -Example SQL descriptor values: +Use the image's dependency-free probe from systemd, Docker, or Kubernetes: -```text -tensor_ref = sha256://0123...cdef -tensor_checksum = sha256:0123...cdef -tensor_rows = 747 -tensor_dim = 320 -tensor_dtype = float16 +```shell +tilemaxsimctl \ + --socket /run/vectorchord/tilemaxsim-status.sock \ + --wait-timeout-ms 30000 ``` -## Start the Service +Readiness is removed before shutdown. Unexpected scheduler or status-thread +exit makes the whole daemon fail, so a supervisor cannot keep routing requests +to a process whose CUDA worker has disappeared. -```text -python3 -m services.tilemaxsim_cuda_sidecar \ - --socket /run/vectorchord/tilemaxsim.sock \ - --socket-mode 660 \ - --device cuda:0 \ - --contract-root 'colqwen3.5@revision+preprocessing-hash=/var/cache/gbrain/colqwen35' \ - --request-timeout-ms 2000 \ - --max-request-bytes 1073741824 \ - --max-tensor-bytes 1073741824 \ - --max-batch-tokens 1000000 \ - --max-device-bytes 8589934592 \ - --cache-bytes 8589934592 \ - --max-inflight 8 \ - --max-cuda-inflight 1 -``` +## Immutable tensor storage -Run the process under an operations-managed supervisor. The socket directory, -process user/group, cache roots, CUDA device visibility, and resource limits -belong in that supervisor configuration. `SIGINT` and `SIGTERM` stop new -accepts, drain accepted work, and remove only the socket inode created by the -process. - -The sidecar and PostgreSQL limits must agree. IPC v1 embeds tensor payloads, so -`--max-request-bytes` must accommodate the inline batch. IPC v2 carries only -descriptors but independently accounts for all declared canonical tensor bytes. -PostgreSQL's `vchordrq.maxsim_gpu_timeout_ms` and the service request timeout -both cover the complete request; operations should leave enough margin for -normal socket scheduling without allowing abandoned work to run indefinitely. - -## Batching, Backpressure, and Failure Semantics - -- The listener accepts at most `--max-inflight` connections. Further clients - remain in the bounded Unix-socket backlog and are governed by their - PostgreSQL-side deadline. -- `--max-cuda-inflight` bounds simultaneous work on one device. Waiting for a - slot consumes the same overall request deadline. -- A request may be split into device-memory-bounded candidate groups. Results - are returned only after every group succeeds. -- Peer disconnect and deadline checks occur before every device group. -- CUDA OOM, missing/corrupt tensors, unsupported references, timeout, and - non-finite scores fail the complete response. Partial candidate results are - never returned. -- f16 input scalars are promoted to f32 before matrix multiplication. TF32 is - disabled by default; `--allow-tf32` must be enabled only after corpus-level - score/ranking tolerance is accepted. - -The service emits one-line JSON events to standard output. Request events -contain request ID, protocol version, source kind, dimensions, candidate and -token counts, content-cache hits, resolution time, CUDA queue time, compute -time, total time, status, and the Unix peer PID/UID/GID where supported. The -peer PID plus request ID disambiguates PostgreSQL backend-local request -counters. Tensor references and public/application IDs are not logged. - -## Verification and Load Probe - -Protocol, resolver, deadline, device-chunking, live-socket, CPU-equivalence, -and optional CUDA-equivalence tests: +Each model contract maps to one absolute shard root. A protocol descriptor uses +a content-addressed reference and matching checksum: ```text -python3 -m unittest -v \ - devtools/test_tilemaxsim_reference_sidecar.py \ - services/test_tilemaxsim_cuda_sidecar.py +tensor_ref = sha256://0123...cdef +tensor_checksum = sha256:0123...cdef +tensor_rows = 747 +tensor_dim = 320 +tensor_dtype = float16 ``` -Synthetic ColQwen-shaped load probe: - -```text -python3 -m services.benchmark_tilemaxsim_cuda \ - --device cuda:0 \ - --dtype f16 \ - --dimension 320 \ - --query-rows 32 \ - --document-rows 747 \ - --candidates 256 \ - --warmup 3 \ - --iterations 20 +The shard publisher must write complete immutable records and atomically +publish its index. SIGHUP reloads shard indexes but never changes the configured +contract-to-root mapping. The daemon validates contract, digest, shape, dtype, +length, and finite values before GPU admission. + +## Batching, scheduling, and failure semantics + +- PostgreSQL splits one logical candidate set into protocol-bounded batches and + merges a deterministic global top-k. +- All batches share the caller's one logical deadline; a new IPC round trip does + not refresh it. +- Admission is bounded by connections, queued requests, per-scheduling-domain + queue depth, and aggregate frame bytes. +- Long batches re-enter the selected scheduler between candidate/token quanta. + This is cooperative preemption between CUDA kernels; a running kernel cannot + be interrupted safely. +- `fair-priority` preserves weighted fairness within the configured priority + band. Strict global priority is available with `--scheduler-policy priority`. +- Disconnects and deadlines are checked between quanta. CUDA, shard, checksum, + workspace, or timeout failure fails the logical request; PostgreSQL does not + expose partial results. + +Scheduling-domain strings affect latency ordering and cache quotas only. They +are not authorization evidence. + +## Verification + +```shell +cargo test --manifest-path services/tilemaxsimd/Cargo.toml --locked +python3 -m unittest -v services.test_tilemaxsim_rust_daemon ``` -The command emits reproducibility metadata, latency percentiles, peak CUDA -allocation/reservation, and a deterministic score checksum as JSON. It is a -runtime smoke/load probe, not a substitute for recall and end-to-end latency on -the committed GBrain corpus. +The Python integration suite exercises real Unix sockets and, when CUDA is +available, multi-GPU, resident-cache, oversized-working-set, concurrent-reader, +health, and score-equivalence paths. Application-level corpus recall and +end-to-end latency remain a release acceptance gate rather than an inference +from synthetic service tests. diff --git a/services/Dockerfile.tilemaxsimd b/services/Dockerfile.tilemaxsimd index d6dcb2f2..328b7f7a 100644 --- a/services/Dockerfile.tilemaxsimd +++ b/services/Dockerfile.tilemaxsimd @@ -18,5 +18,10 @@ RUN cargo build --release --locked FROM ${CUDA_RUNTIME_IMAGE} COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimd /usr/local/bin/tilemaxsimd +COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimctl /usr/local/bin/tilemaxsimctl + +STOPSIGNAL SIGTERM +HEALTHCHECK --interval=10s --timeout=2s --start-period=5m --retries=3 \ + CMD ["/usr/local/bin/tilemaxsimctl", "--socket", "/run/vectorchord/tilemaxsim-status.sock"] ENTRYPOINT ["/usr/local/bin/tilemaxsimd"] diff --git a/services/test_tilemaxsim_rust_daemon.py b/services/test_tilemaxsim_rust_daemon.py index 523210c4..2fb21746 100644 --- a/services/test_tilemaxsim_rust_daemon.py +++ b/services/test_tilemaxsim_rust_daemon.py @@ -174,6 +174,7 @@ def run_daemon( self.assertTrue(socket_path.exists()) self.assertTrue(status_socket_path.exists()) for path, expected in ( + ("/livez", b'200 OK'), ("/healthz", b'200 OK'), ("/metrics", b"tilemaxsim_ready 1"), ): @@ -186,6 +187,18 @@ def run_daemon( while part := status.recv(4096): response_parts.append(part) self.assertIn(expected, b"".join(response_parts)) + probe = subprocess.run( + [ + os.fspath(binary.with_name("tilemaxsimctl")), + "--socket", + os.fspath(status_socket_path), + ], + capture_output=True, + text=True, + timeout=5, + check=False, + ) + self.assertEqual(probe.returncode, 0, probe.stderr) with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: connection.connect(os.fspath(socket_path)) connection.sendall(frame) @@ -215,6 +228,75 @@ def run_daemon( process.terminate() process.wait(timeout=5) + def test_gpu_assignment_is_required_before_startup(self) -> None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + socket_path = root / "tilemaxsimd.sock" + completed = subprocess.run( + [ + os.fspath(binary), + "--socket", + os.fspath(socket_path), + "--contract-root", + f"model@1={root}", + ], + capture_output=True, + text=True, + timeout=5, + check=False, + ) + self.assertNotEqual(completed.returncode, 0) + self.assertIn("--gpu-memory-gb", completed.stderr) + self.assertFalse(socket_path.exists()) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_unavailable_configured_gpu_fails_before_socket_ready(self) -> None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + writer = ImmutableShardWriter( + root, target_bytes=4096, alignment=256, fsync=False + ) + payload = np.asarray([[1.0, 0.0]], dtype=" None: device = max( diff --git a/services/tilemaxsimd/src/bin/tilemaxsimctl.rs b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs new file mode 100644 index 00000000..9fca086d --- /dev/null +++ b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs @@ -0,0 +1,82 @@ +use anyhow::{Context, Result, anyhow, bail}; +use clap::Parser; +use std::io::{Read, Write}; +use std::os::unix::net::UnixStream; +use std::path::PathBuf; +use std::thread; +use std::time::{Duration, Instant}; + +#[derive(Parser)] +#[command(about = "Probe the native TileMaxSim daemon readiness socket")] +struct Args { + #[arg( + long, + default_value = "/run/vectorchord/tilemaxsim-status.sock" + )] + socket: PathBuf, + #[arg(long, default_value_t = 500)] + io_timeout_ms: u64, + #[arg(long, default_value_t = 0)] + wait_timeout_ms: u64, +} + +fn main() -> Result<()> { + let args = Args::parse(); + if args.io_timeout_ms == 0 { + bail!("I/O timeout must be positive"); + } + let io_timeout = Duration::from_millis(args.io_timeout_ms); + let wait_timeout = Duration::from_millis(args.wait_timeout_ms); + let deadline = Instant::now() + .checked_add(wait_timeout) + .ok_or_else(|| anyhow!("readiness deadline overflow"))?; + + loop { + let error = match probe(&args.socket, io_timeout) { + Ok(()) => return Ok(()), + Err(error) => error, + }; + if wait_timeout.is_zero() || Instant::now() >= deadline { + return Err(error); + } + thread::sleep(Duration::from_millis(50)); + } +} + +fn probe(socket: &PathBuf, timeout: Duration) -> Result<()> { + let mut stream = UnixStream::connect(socket) + .with_context(|| format!("cannot connect to status socket {}", socket.display()))?; + stream.set_read_timeout(Some(timeout))?; + stream.set_write_timeout(Some(timeout))?; + stream.write_all(b"GET /healthz HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")?; + let mut response = Vec::new(); + stream.read_to_end(&mut response)?; + if !is_ready_response(&response) { + bail!("TileMaxSim daemon is not ready"); + } + Ok(()) +} + +fn is_ready_response(response: &[u8]) -> bool { + response.starts_with(b"HTTP/1.1 200 ") + && response + .windows(b"\r\n\r\n".len()) + .any(|window| window == b"\r\n\r\n") + && response.ends_with(b"{\"ready\":true}") +} + +#[cfg(test)] +mod tests { + use super::is_ready_response; + + #[test] + fn readiness_requires_success_status_and_true_body() { + assert!(is_ready_response( + b"HTTP/1.1 200 OK\r\nContent-Length: 14\r\n\r\n{\"ready\":true}" + )); + assert!(!is_ready_response( + b"HTTP/1.1 503 Service Unavailable\r\n\r\n{\"ready\":false}" + )); + assert!(!is_ready_response(b"HTTP/1.1 200 OK\r\n\r\n")); + } +} diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs index 8a01793d..4e159c31 100644 --- a/services/tilemaxsimd/src/main.rs +++ b/services/tilemaxsimd/src/main.rs @@ -20,7 +20,7 @@ use std::io::{BufRead, Read, Write}; use std::os::fd::AsRawFd; use std::os::unix::fs::{FileTypeExt, PermissionsExt}; use std::os::unix::net::{UnixListener, UnixStream}; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, mpsc}; use std::thread; @@ -372,13 +372,22 @@ fn main() -> Result<()> { let mut readers = Vec::new(); let mut accepted = 0_usize; let mut status_failed = false; + let mut scheduler_failed = false; + let mut fatal_error = None; while !SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + if scheduler.is_finished() { + scheduler_failed = true; + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + break; + } if status_server .as_ref() .is_some_and(thread::JoinHandle::is_finished) { status_failed = true; metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); break; } if RELOAD_REQUESTED.swap(false, Ordering::AcqRel) { @@ -400,26 +409,33 @@ fn main() -> Result<()> { let reader_admission = Arc::clone(&pending_admission); let reader_metrics = Arc::clone(&metrics); let reader_frame_admission = Arc::clone(&frame_admission); - let maximum = args.max_request_bytes; - let io_timeout = Duration::from_millis(args.socket_io_timeout_ms); - let request_timeout = Duration::from_millis(args.request_timeout_ms); - readers.push( - thread::Builder::new() - .name("tilemaxsim-reader".to_owned()) - .spawn(move || { - let _permit = ReaderPermit(reader_count); - read_and_enqueue( - connection, - &reader_sender, - maximum, - io_timeout, - request_timeout, - reader_admission, - reader_metrics, - reader_frame_admission, - ); - })?, - ); + let reader_config = ReaderConfig { + maximum: args.max_request_bytes, + io_timeout: Duration::from_millis(args.socket_io_timeout_ms), + server_timeout: Duration::from_millis(args.request_timeout_ms), + }; + match thread::Builder::new() + .name("tilemaxsim-reader".to_owned()) + .spawn(move || { + let _permit = ReaderPermit(reader_count); + read_and_enqueue( + connection, + &reader_sender, + reader_config, + reader_admission, + reader_metrics, + reader_frame_admission, + ); + }) { + Ok(reader) => readers.push(reader), + Err(error) => { + live_readers.fetch_sub(1, Ordering::Release); + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + fatal_error = Some(error.into()); + break; + } + } if args.once && accepted == 1 { SHUTDOWN_REQUESTED.store(true, Ordering::Relaxed); } @@ -428,9 +444,15 @@ fn main() -> Result<()> { thread::sleep(Duration::from_millis(5)); } Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {} - Err(error) => return Err(error.into()), + Err(error) => { + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + fatal_error = Some(error.into()); + break; + } } } + SHUTDOWN_REQUESTED.store(true, Ordering::Release); drop(listener); for reader in readers { if reader.join().is_err() { @@ -439,22 +461,29 @@ fn main() -> Result<()> { } drop(sender); metrics.ready.store(false, Ordering::Release); - scheduler + let scheduler_result = scheduler .join() - .map_err(|_| anyhow!("TileMaxSim scheduler thread panicked"))??; - if let Some(status_server) = status_server { - if status_server.join().is_err() { - eprintln!("TileMaxSim status thread panicked during shutdown"); - } - } - if status_failed { - bail!("TileMaxSim status server exited unexpectedly"); + .map_err(|_| anyhow!("TileMaxSim scheduler thread panicked")) + .and_then(|result| result); + if status_server.is_some_and(|status_server| status_server.join().is_err()) { + eprintln!("TileMaxSim status thread panicked during shutdown"); } match fs::remove_file(&args.socket) { Ok(()) => {} Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} Err(error) => return Err(error.into()), } + if status_failed { + bail!("TileMaxSim status server exited unexpectedly"); + } + if scheduler_failed { + scheduler_result?; + bail!("TileMaxSim scheduler exited unexpectedly"); + } + scheduler_result?; + if let Some(error) = fatal_error { + return Err(error); + } Ok(()) } @@ -631,26 +660,32 @@ fn reap_readers(readers: &mut Vec>) { } } -fn read_and_enqueue( - mut connection: UnixStream, - sender: &mpsc::SyncSender, +#[derive(Clone, Copy)] +struct ReaderConfig { maximum: usize, io_timeout: Duration, server_timeout: Duration, +} + +fn read_and_enqueue( + mut connection: UnixStream, + sender: &mpsc::SyncSender, + config: ReaderConfig, pending_admission: Arc, metrics: Arc, frame_admission: Arc, ) { let accepted_at = Instant::now(); - if let Err(error) = connection.set_read_timeout(Some(io_timeout)) { + if let Err(error) = connection.set_read_timeout(Some(config.io_timeout)) { eprintln!("cannot configure TileMaxSim socket read timeout: {error}"); return; } - if let Err(error) = connection.set_write_timeout(Some(io_timeout)) { + if let Err(error) = connection.set_write_timeout(Some(config.io_timeout)) { eprintln!("cannot configure TileMaxSim socket write timeout: {error}"); return; } - let (frame, frame_permit) = match read_request(&mut connection, maximum, &frame_admission) { + let (frame, frame_permit) = + match read_request(&mut connection, config.maximum, &frame_admission) { Ok(frame) => frame, Err(error) => { metrics.failed.fetch_add(1, Ordering::Relaxed); @@ -660,7 +695,7 @@ fn read_and_enqueue( ); return; } - }; + }; let version = header_version(&frame); let request_id = header_request_id(&frame); let request = match protocol::parse(&frame) { @@ -675,9 +710,9 @@ fn read_and_enqueue( } }; let client_timeout = if request.timeout_ms == 0 { - server_timeout + config.server_timeout } else { - Duration::from_millis(u64::from(request.timeout_ms)).min(server_timeout) + Duration::from_millis(u64::from(request.timeout_ms)).min(config.server_timeout) }; let deadline = accepted_at .checked_add(client_timeout) @@ -1065,7 +1100,13 @@ fn handle_status_connection(connection: &mut UnixStream, metrics: &RuntimeMetric return; }; let request = String::from_utf8_lossy(&request[..count]); - let (status, content_type, body) = if request.starts_with("GET /healthz ") { + let (status, content_type, body) = if request.starts_with("GET /livez ") { + ( + "200 OK", + "application/json", + serde_json::json!({"live": true}).to_string(), + ) + } else if request.starts_with("GET /healthz ") { let ready = metrics.ready.load(Ordering::Acquire); ( if ready { @@ -1173,10 +1214,11 @@ fn header_version(frame: &[u8]) -> u16 { } } -fn acquire_instance_lock(socket: &PathBuf) -> Result { +fn acquire_instance_lock(socket: &Path) -> Result { let lock_path = PathBuf::from(format!("{}.lock", socket.display())); let mut lock = OpenOptions::new() .create(true) + .truncate(false) .read(true) .write(true) .open(&lock_path) From 73a65abc4c4b84c4ad7a1fcedec01af03c19cff1 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:56:26 +0800 Subject: [PATCH 314/324] feat(tilemaxsimd): publish online tensor objects --- services/test_tilemaxsim_rust_daemon.py | 91 ++++++++ services/tilemaxsimd/src/bin/tilemaxsimctl.rs | 210 +++++++++++++++++- services/tilemaxsimd/src/shard.rs | 104 ++++++++- 3 files changed, 392 insertions(+), 13 deletions(-) diff --git a/services/test_tilemaxsim_rust_daemon.py b/services/test_tilemaxsim_rust_daemon.py index 2fb21746..b2e38fc7 100644 --- a/services/test_tilemaxsim_rust_daemon.py +++ b/services/test_tilemaxsim_rust_daemon.py @@ -305,6 +305,97 @@ def test_external_v2_shard_round_trip_matches_protocol_oracle(self) -> None: ) self.run_daemon([device]) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_published_object_is_queryable_without_daemon_restart(self) -> None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + object_root = root / "objects-root" + object_root.mkdir() + socket_path = root / "tilemaxsimd.sock" + process = subprocess.Popen( + [ + os.fspath(binary), + "--socket", os.fspath(socket_path), + "--gpu-memory-gb", f"{device}=0.05", + "--gpu-workspace-gb", "0.02", + "--host-cache-gb", "0.01", + "--contract-root", f"model@1={object_root}", + "--once", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + for _ in range(1000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + self.assertIsNone(process.poll()) + document = np.asarray( + [[1.0, 0.0], [0.0, 1.0]], dtype=" None: device = max( diff --git a/services/tilemaxsimd/src/bin/tilemaxsimctl.rs b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs index 9fca086d..0f1c55b0 100644 --- a/services/tilemaxsimd/src/bin/tilemaxsimctl.rs +++ b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs @@ -1,18 +1,20 @@ use anyhow::{Context, Result, anyhow, bail}; -use clap::Parser; +use clap::{Parser, Subcommand}; +use sha2::{Digest, Sha256}; +use std::fs::{self, File, OpenOptions}; use std::io::{Read, Write}; +use std::os::unix::fs::{OpenOptionsExt, PermissionsExt}; use std::os::unix::net::UnixStream; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::thread; use std::time::{Duration, Instant}; #[derive(Parser)] #[command(about = "Probe the native TileMaxSim daemon readiness socket")] struct Args { - #[arg( - long, - default_value = "/run/vectorchord/tilemaxsim-status.sock" - )] + #[command(subcommand)] + command: Option, + #[arg(long, default_value = "/run/vectorchord/tilemaxsim-status.sock")] socket: PathBuf, #[arg(long, default_value_t = 500)] io_timeout_ms: u64, @@ -20,8 +22,55 @@ struct Args { wait_timeout_ms: u64, } +#[derive(Subcommand)] +enum Command { + /// Publish one canonical tensor from stdin into the immutable object store. + PublishObject { + #[arg(long)] + root: PathBuf, + #[arg(long)] + rows: u32, + #[arg(long)] + dimension: u32, + #[arg(long, value_parser = ["float16", "float32"])] + dtype: String, + #[arg(long)] + expected_sha256: Option, + }, +} + fn main() -> Result<()> { let args = Args::parse(); + if let Some(Command::PublishObject { + root, + rows, + dimension, + dtype, + expected_sha256, + }) = args.command + { + let expected_bytes = tensor_bytes(rows, dimension, &dtype)?; + let mut payload = Vec::with_capacity(expected_bytes); + std::io::stdin() + .take(expected_bytes.saturating_add(1) as u64) + .read_to_end(&mut payload)?; + if payload.len() != expected_bytes { + bail!( + "stdin tensor length {} does not match declared length {expected_bytes}", + payload.len() + ); + } + let descriptor = publish_object( + &root, + rows, + dimension, + &dtype, + &payload, + expected_sha256.as_deref(), + )?; + println!("{}", serde_json::to_string(&descriptor)?); + return Ok(()); + } if args.io_timeout_ms == 0 { bail!("I/O timeout must be positive"); } @@ -43,6 +92,131 @@ fn main() -> Result<()> { } } +#[derive(serde::Serialize)] +struct PublishedDescriptor { + tensor_ref: String, + tensor_rows: u32, + tensor_dim: u32, + tensor_dtype: String, + tensor_checksum: String, +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: &str) -> Result { + if rows == 0 || rows > 65_536 || dimension == 0 || dimension > 60_000 { + bail!("invalid tensor shape"); + } + let scalar_bytes = match dtype { + "float16" => 2usize, + "float32" => 4usize, + _ => bail!("unsupported tensor dtype"), + }; + (rows as usize) + .checked_mul(dimension as usize) + .and_then(|elements| elements.checked_mul(scalar_bytes)) + .ok_or_else(|| anyhow!("tensor shape overflow")) +} + +fn publish_object( + root: &Path, + rows: u32, + dimension: u32, + dtype: &str, + payload: &[u8], + expected_sha256: Option<&str>, +) -> Result { + if payload.len() != tensor_bytes(rows, dimension, dtype)? { + bail!("tensor payload length disagrees with its shape"); + } + let digest = hex::encode(Sha256::digest(payload)); + if expected_sha256.is_some_and(|expected| expected != digest) { + bail!("tensor payload checksum disagrees with --expected-sha256"); + } + fs::create_dir_all(root)?; + let root_metadata = fs::symlink_metadata(root)?; + if root_metadata.file_type().is_symlink() || !root_metadata.is_dir() { + bail!("tensor root must be a real directory"); + } + let objects = root.join("objects"); + fs::create_dir_all(&objects)?; + let objects_metadata = fs::symlink_metadata(&objects)?; + if objects_metadata.file_type().is_symlink() || !objects_metadata.is_dir() { + bail!("tensor objects path must be a real directory"); + } + let directory = objects.join(&digest[..2]); + fs::create_dir_all(&directory)?; + let directory_metadata = fs::symlink_metadata(&directory)?; + if directory_metadata.file_type().is_symlink() || !directory_metadata.is_dir() { + bail!("tensor object directory must be a real directory"); + } + let destination = directory.join(format!("{digest}.tensor")); + if destination.try_exists()? { + verify_existing_object(&destination, payload.len(), &digest)?; + } else { + let temporary = directory.join(format!(".{digest}.{}.tmp", std::process::id())); + let mut file = OpenOptions::new() + .create_new(true) + .write(true) + .mode(0o640) + .open(&temporary) + .with_context(|| format!("cannot create {}", temporary.display()))?; + let publish = (|| -> Result<()> { + file.write_all(payload)?; + file.sync_all()?; + fs::set_permissions(&temporary, fs::Permissions::from_mode(0o440))?; + match fs::hard_link(&temporary, &destination) { + Ok(()) => File::open(&directory)?.sync_all()?, + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + verify_existing_object(&destination, payload.len(), &digest)?; + } + Err(error) => return Err(error.into()), + } + Ok(()) + })(); + drop(file); + let remove_result = fs::remove_file(&temporary); + if let Err(error) = remove_result + && error.kind() != std::io::ErrorKind::NotFound + { + return Err(error.into()); + } + publish?; + } + Ok(PublishedDescriptor { + tensor_ref: format!("sha256://{digest}"), + tensor_rows: rows, + tensor_dim: dimension, + tensor_dtype: dtype.to_owned(), + tensor_checksum: format!("sha256:{digest}"), + }) +} + +fn verify_existing_object(path: &Path, expected_bytes: usize, expected_digest: &str) -> Result<()> { + let metadata = fs::symlink_metadata(path)?; + if metadata.file_type().is_symlink() + || !metadata.is_file() + || metadata.len() != expected_bytes as u64 + { + bail!("existing immutable tensor object is invalid"); + } + let mut file = OpenOptions::new() + .read(true) + .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(path)?; + let mut digest = Sha256::new(); + let mut buffer = [0_u8; 1024 * 1024]; + loop { + let count = file.read(&mut buffer)?; + if count == 0 { + break; + } + digest.update(&buffer[..count]); + } + if hex::encode(digest.finalize()) != expected_digest { + bail!("existing immutable tensor object checksum mismatch"); + } + Ok(()) +} + fn probe(socket: &PathBuf, timeout: Duration) -> Result<()> { let mut stream = UnixStream::connect(socket) .with_context(|| format!("cannot connect to status socket {}", socket.display()))?; @@ -67,7 +241,8 @@ fn is_ready_response(response: &[u8]) -> bool { #[cfg(test)] mod tests { - use super::is_ready_response; + use super::{is_ready_response, publish_object}; + use std::fs; #[test] fn readiness_requires_success_status_and_true_body() { @@ -79,4 +254,25 @@ mod tests { )); assert!(!is_ready_response(b"HTTP/1.1 200 OK\r\n\r\n")); } + + #[test] + fn publish_object_is_content_addressed_and_idempotent() { + let root = std::env::temp_dir().join(format!( + "tilemaxsimctl-publish-{}-{}", + std::process::id(), + std::thread::current().name().unwrap_or("test") + )); + let _ = fs::remove_dir_all(&root); + let payload = [0_u8, 60, 0, 0]; + let first = publish_object(&root, 1, 2, "float16", &payload, None).unwrap(); + let second = publish_object(&root, 1, 2, "float16", &payload, None).unwrap(); + assert_eq!(first.tensor_ref, second.tensor_ref); + assert_eq!(first.tensor_checksum, second.tensor_checksum); + assert!( + root.join("objects") + .join(&first.tensor_checksum[7..9]) + .exists() + ); + fs::remove_dir_all(root).unwrap(); + } } diff --git a/services/tilemaxsimd/src/shard.rs b/services/tilemaxsimd/src/shard.rs index 358c25a4..044b27fa 100644 --- a/services/tilemaxsimd/src/shard.rs +++ b/services/tilemaxsimd/src/shard.rs @@ -58,6 +58,7 @@ struct Entry { } struct ContractStore { + root: PathBuf, shards: HashMap, entries: HashMap, } @@ -258,11 +259,26 @@ impl ShardStore { .contracts .get(&descriptor.contract) .ok_or_else(|| anyhow!("model contract has no immutable shard root"))?; - let entry = contract - .entries - .get(&descriptor.digest) - .ok_or_else(|| anyhow!("tensor is missing from the immutable shard index"))? - .clone(); + let Some(entry) = contract.entries.get(&descriptor.digest).cloned() else { + let expected = + tensor_bytes(descriptor.rows, descriptor.dimension, descriptor.dtype)?; + let mut file = open_object(&contract.root, &descriptor.digest) + .with_context(|| "tensor is missing from immutable shards and objects")?; + if file.metadata()?.len() != expected as u64 { + bail!("immutable tensor object length disagrees with its descriptor"); + } + let mut payload = Vec::with_capacity(expected); + file.read_to_end(&mut payload)?; + if hex::encode(Sha256::digest(&payload)) != descriptor.digest { + bail!("immutable tensor object checksum mismatch"); + } + self.batch_read_calls += 1; + self.batch_read_bytes += payload.len() as u64; + let payload: Arc<[u8]> = payload.into(); + output[index] = Some(Arc::clone(&payload)); + self.host_cache.put(tenant, key, payload); + continue; + }; if entry.rows != descriptor.rows || entry.dimension != descriptor.dimension || entry.dtype != descriptor.dtype @@ -368,6 +384,18 @@ impl ShardStore { fn open_contract(root: &Path) -> Result { let index_path = root.join(INDEX_NAME); + let root_metadata = std::fs::symlink_metadata(root) + .with_context(|| format!("cannot inspect immutable tensor root {}", root.display()))?; + if root_metadata.file_type().is_symlink() || !root_metadata.is_dir() { + bail!("immutable tensor root must be a real directory"); + } + if !index_path.try_exists()? { + return Ok(ContractStore { + root: root.to_path_buf(), + shards: HashMap::new(), + entries: HashMap::new(), + }); + } let mut index_file = nofollow(&index_path)?; let mut document = String::new(); index_file.read_to_string(&mut document)?; @@ -444,7 +472,38 @@ fn open_contract(root: &Path) -> Result { bail!("duplicate tensor digest in shard index"); } } - Ok(ContractStore { shards, entries }) + Ok(ContractStore { + root: root.to_path_buf(), + shards, + entries, + }) +} + +fn object_path(root: &Path, digest: &str) -> Result { + if digest.len() != 64 + || !digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + { + bail!("invalid tensor object digest"); + } + Ok(root + .join("objects") + .join(&digest[..2]) + .join(format!("{digest}.tensor"))) +} + +fn open_object(root: &Path, digest: &str) -> Result { + let path = object_path(root, digest)?; + let objects = root.join("objects"); + let prefix = objects.join(&digest[..2]); + for directory in [&objects, &prefix] { + let metadata = std::fs::symlink_metadata(directory)?; + if metadata.file_type().is_symlink() || !metadata.is_dir() { + bail!("immutable tensor object path contains a non-directory component"); + } + } + nofollow(&path) } fn nofollow(path: &Path) -> Result { @@ -497,3 +556,36 @@ pub fn cache_key(descriptor: &Descriptor) -> String { descriptor.dtype ) } + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + + #[test] + fn content_addressed_objects_are_visible_without_index_reload() { + let root = + std::env::temp_dir().join(format!("tilemaxsim-object-store-{}", std::process::id())); + let _ = fs::remove_dir_all(&root); + fs::create_dir_all(&root).unwrap(); + let mut store = + ShardStore::open(&[("model@1".to_owned(), root.clone())], 1024, 100, false).unwrap(); + let payload = [0_u8, 60, 0, 0]; + let digest = hex::encode(Sha256::digest(payload)); + let directory = root.join("objects").join(&digest[..2]); + fs::create_dir_all(&directory).unwrap(); + fs::write(directory.join(format!("{digest}.tensor")), payload).unwrap(); + let descriptor = Descriptor { + candidate_id: 1, + contract: "model@1".to_owned(), + digest, + rows: 1, + dimension: 2, + dtype: 2, + }; + + let resolved = store.resolve_many(&[descriptor], "tenant-a").unwrap(); + assert_eq!(&*resolved[0], &payload); + fs::remove_dir_all(root).unwrap(); + } +} From 89a48e6687657e9d06b2d921c4b8b61ad5667bb7 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 18:08:07 +0800 Subject: [PATCH 315/324] feat(tilemaxsimd): expose production cache metrics --- docs/TILEMAXSIM_CUDA_SIDECAR.md | 53 +- services/test_tilemaxsim_rust_daemon.py | 6 + services/tilemaxsimd/src/engine.rs | 126 +++- services/tilemaxsimd/src/gpu.rs | 6 + services/tilemaxsimd/src/main.rs | 858 +++++++++++++++++++++--- services/tilemaxsimd/src/shard.rs | 30 +- 6 files changed, 959 insertions(+), 120 deletions(-) diff --git a/docs/TILEMAXSIM_CUDA_SIDECAR.md b/docs/TILEMAXSIM_CUDA_SIDECAR.md index f5986c25..3723935c 100644 --- a/docs/TILEMAXSIM_CUDA_SIDECAR.md +++ b/docs/TILEMAXSIM_CUDA_SIDECAR.md @@ -80,7 +80,10 @@ The status Unix socket serves: - `GET /livez`: the status server is alive; - `GET /healthz`: the protocol listener and scheduler are ready; -- `GET /metrics`: bounded Prometheus scheduler and outcome counters. +- `GET /metrics`: bounded Prometheus resource, scheduler, cache, transfer, + storage, latency, timeout, and outcome metrics. GPU labels contain only the + configured numeric device/slot; application scheduling-domain names are + never exported. Use the image's dependency-free probe from systemd, Docker, or Kubernetes: @@ -94,6 +97,30 @@ Readiness is removed before shutdown. Unexpected scheduler or status-thread exit makes the whole daemon fail, so a supervisor cannot keep routing requests to a process whose CUDA worker has disappeared. +The most useful production signals are: + +- `tilemaxsim_pending_requests` and `tilemaxsim_scheduler_queue_depth` for + saturation; +- `tilemaxsim_admission_rejections_total` and `tilemaxsim_timeouts_total` for + overload or an undersized deadline; +- `tilemaxsim_gpu_cache_events_total`, `tilemaxsim_gpu_h2d_bytes_total`, and + `tilemaxsim_storage_read_bytes_total` for cache churn and cold-load traffic; +- `tilemaxsim_gpu_cache_bytes` for free space, largest free extent, payload, + allocator waste, and pinned capacity; +- `tilemaxsim_host_cache_bytes` and `tilemaxsim_host_cache_events_total` for + the L1 cache; +- `tilemaxsim_request_duration_seconds`, `tilemaxsim_queue_duration_seconds`, + and `tilemaxsim_gpu_duration_seconds` for rate-derived mean latency; +- `tilemaxsim_requests_admitted_total{priority_class=...}` and + `tilemaxsim_scheduler_requeues_total` for priority/cooperative scheduling. + +At minimum, alert when readiness is zero, any admission-rejection rate is +nonzero under expected load, timeout ratio exceeds the application SLO, queue +depth remains above 80% of its limit, or cache admission rejections rise while +`largest_free_extent / free` is small. A rising miss/eviction/H2D rate with +stable traffic means the assigned GPU cache is too small or the access set has +poor locality; it is not a PostgreSQL HNSW symptom. + ## Immutable tensor storage Each model contract maps to one absolute shard root. A protocol descriptor uses @@ -107,10 +134,26 @@ tensor_dim = 320 tensor_dtype = float16 ``` -The shard publisher must write complete immutable records and atomically -publish its index. SIGHUP reloads shard indexes but never changes the configured -contract-to-root mapping. The daemon validates contract, digest, shape, dtype, -length, and finite values before GPU admission. +Online writers publish canonical tensor bytes through VectorChord's publisher, +so application code never owns or duplicates the disk layout: + +```shell +install -d -m 0750 /var/lib/vectorchord/tensors +tilemaxsimctl publish-object \ + --root /var/lib/vectorchord/tensors \ + --rows 747 \ + --dimension 320 \ + --dtype float16 \ + --expected-sha256 0123...cdef < document.tensor.f16 +``` + +The command fsyncs a temporary object, publishes it by immutable hard link, and +returns the JSON descriptor. Repeating the same publication is idempotent. A +daemon that already has the contract root open can read the new content address +without restart or SIGHUP. SIGHUP is needed only for a newly published bulk +shard index. It never changes the configured contract-to-root mapping. The +daemon validates contract, digest, shape, dtype, length, and checksum before GPU +admission. ## Batching, scheduling, and failure semantics diff --git a/services/test_tilemaxsim_rust_daemon.py b/services/test_tilemaxsim_rust_daemon.py index b2e38fc7..94cf978f 100644 --- a/services/test_tilemaxsim_rust_daemon.py +++ b/services/test_tilemaxsim_rust_daemon.py @@ -187,6 +187,12 @@ def run_daemon( while part := status.recv(4096): response_parts.append(part) self.assertIn(expected, b"".join(response_parts)) + if path == "/metrics": + metrics_body = b"".join(response_parts) + self.assertIn(b"tilemaxsim_gpu_cache_bytes", metrics_body) + self.assertIn(b"tilemaxsim_host_cache_bytes", metrics_body) + self.assertIn(b"tilemaxsim_storage_read_bytes_total", metrics_body) + self.assertNotIn(b"tenant-a", metrics_body) probe = subprocess.run( [ os.fspath(binary.with_name("tilemaxsimctl")), diff --git a/services/tilemaxsimd/src/engine.rs b/services/tilemaxsimd/src/engine.rs index d953169d..5f6d946b 100644 --- a/services/tilemaxsimd/src/engine.rs +++ b/services/tilemaxsimd/src/engine.rs @@ -1,7 +1,7 @@ use crate::cache::{Admission, GpuCache}; use crate::gpu::Gpu; use crate::protocol::{Descriptor, Request}; -use crate::shard::{ShardStore, cache_key}; +use crate::shard::{HostCacheStatus, ShardStore, cache_key}; use anyhow::{Result, anyhow, bail}; use std::collections::HashMap; use std::sync::Arc; @@ -36,6 +36,37 @@ pub struct Engine { next_device: usize, } +#[derive(Clone, Debug, Default)] +pub struct DeviceStatus { + pub slot: usize, + pub device: i32, + pub capacity_bytes: usize, + pub block_bytes: usize, + pub free_bytes: usize, + pub largest_free_extent_bytes: usize, + pub allocated_bytes: usize, + pub payload_bytes: usize, + pub internal_waste_bytes: usize, + pub entries: usize, + pub pinned_entries: usize, + pub pinned_bytes: usize, + pub tenants: usize, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, + pub h2d_batches: u64, + pub h2d_bytes: u64, +} + +#[derive(Clone, Debug, Default)] +pub struct EngineStatus { + pub devices: Vec, + pub host: HostCacheStatus, + pub batch_read_calls: u64, + pub batch_read_bytes: u64, +} + impl Engine { pub fn new( gpus: Vec, @@ -526,32 +557,69 @@ impl Engine { Ok(()) } - pub fn status_json(&self) -> serde_json::Value { - let (host_hits, host_misses, host_evictions, host_rejections) = self.store.host_status(); + pub fn status_snapshot(&self) -> EngineStatus { let devices = self .devices .iter() .enumerate() - .map(|(index, device)| { + .map(|(slot, device)| DeviceStatus { + slot, + device: device.gpu.device(), + capacity_bytes: device.cache.capacity(), + block_bytes: device.cache.block_bytes(), + free_bytes: device.cache.free_bytes(), + largest_free_extent_bytes: device.cache.largest_free_extent(), + allocated_bytes: device.cache.allocated_bytes(), + payload_bytes: device.cache.payload_bytes(), + internal_waste_bytes: device + .cache + .allocated_bytes() + .saturating_sub(device.cache.payload_bytes()), + entries: device.cache.entry_count(), + pinned_entries: device.cache.pinned_entries(), + pinned_bytes: device.cache.pinned_bytes(), + tenants: device.cache.tenant_count(), + hits: device.cache.hits, + misses: device.cache.misses, + evictions: device.cache.evictions, + admission_rejections: device.cache.admission_rejections, + h2d_batches: device.h2d_batches, + h2d_bytes: device.h2d_bytes, + }) + .collect(); + EngineStatus { + devices, + host: self.store.host_status(), + batch_read_calls: self.store.batch_read_calls, + batch_read_bytes: self.store.batch_read_bytes, + } + } + + pub fn status_json(&self) -> serde_json::Value { + let status = self.status_snapshot(); + let devices = status + .devices + .iter() + .map(|device| { serde_json::json!({ - "index": index, + "index": device.slot, + "device": device.device, "gpu_allocator": "segregated-page-runs", - "gpu_tensor_bytes": device.cache.capacity(), - "gpu_block_bytes": device.cache.block_bytes(), - "gpu_free_bytes": device.cache.free_bytes(), - "gpu_largest_free_extent_bytes": device.cache.largest_free_extent(), - "gpu_allocated_bytes": device.cache.allocated_bytes(), - "gpu_payload_bytes": device.cache.payload_bytes(), - "gpu_internal_waste_bytes": device.cache.allocated_bytes() - - device.cache.payload_bytes(), - "gpu_entries": device.cache.entry_count(), - "gpu_pinned_entries": device.cache.pinned_entries(), - "gpu_pinned_bytes": device.cache.pinned_bytes(), - "gpu_tenant_count": device.cache.tenant_count(), - "gpu_hits": device.cache.hits, - "gpu_misses": device.cache.misses, - "gpu_evictions": device.cache.evictions, - "gpu_admission_rejections": device.cache.admission_rejections, + "gpu_tensor_bytes": device.capacity_bytes, + "gpu_block_bytes": device.block_bytes, + "gpu_free_bytes": device.free_bytes, + "gpu_largest_free_extent_bytes": device.largest_free_extent_bytes, + "gpu_allocated_bytes": device.allocated_bytes, + "gpu_payload_bytes": device.payload_bytes, + "gpu_internal_waste_bytes": device.internal_waste_bytes, + "gpu_entries": device.entries, + "gpu_pinned_entries": device.pinned_entries, + "gpu_pinned_bytes": device.pinned_bytes, + "gpu_tenant_count": device.tenants, + "gpu_hits": device.hits, + "gpu_misses": device.misses, + "gpu_evictions": device.evictions, + "gpu_admission_rejections": device.admission_rejections, "h2d_batches": device.h2d_batches, "h2d_bytes": device.h2d_bytes, }) @@ -559,12 +627,16 @@ impl Engine { .collect::>(); serde_json::json!({ "devices": devices, - "host_hits": host_hits, - "host_misses": host_misses, - "host_evictions": host_evictions, - "host_admission_rejections": host_rejections, - "batch_read_calls": self.store.batch_read_calls, - "batch_read_bytes": self.store.batch_read_bytes, + "host_capacity_bytes": status.host.capacity_bytes, + "host_used_bytes": status.host.used_bytes, + "host_entries": status.host.entries, + "host_tenant_count": status.host.tenants, + "host_hits": status.host.hits, + "host_misses": status.host.misses, + "host_evictions": status.host.evictions, + "host_admission_rejections": status.host.admission_rejections, + "batch_read_calls": status.batch_read_calls, + "batch_read_bytes": status.batch_read_bytes, }) } } diff --git a/services/tilemaxsimd/src/gpu.rs b/services/tilemaxsimd/src/gpu.rs index 2c6c2df0..ed00749d 100644 --- a/services/tilemaxsimd/src/gpu.rs +++ b/services/tilemaxsimd/src/gpu.rs @@ -43,6 +43,7 @@ unsafe extern "C" { pub struct Gpu { native: NonNull, + device: i32, tensor_bytes: usize, } @@ -73,10 +74,15 @@ impl Gpu { let tensor_bytes = unsafe { vctm_gpu_tensor_bytes(native.as_ptr()) }; Ok(Self { native, + device, tensor_bytes, }) } + pub fn device(&self) -> i32 { + self.device + } + pub fn tensor_bytes(&self) -> usize { self.tensor_bytes } diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs index 4e159c31..4b050776 100644 --- a/services/tilemaxsimd/src/main.rs +++ b/services/tilemaxsimd/src/main.rs @@ -7,7 +7,7 @@ mod shard; use anyhow::{Context, Result, anyhow, bail}; use clap::Parser; -use engine::Engine; +use engine::{Engine, EngineStatus}; use gpu::Gpu; use protocol::{HEADER_BYTES, VERSION_EXTERNAL, VERSION_SCHEDULED_EXTERNAL}; use scheduler::{RequestQueue, Scheduled, SchedulerPolicy}; @@ -15,6 +15,7 @@ use serde::Deserialize; use sha2::{Digest, Sha256}; use shard::ShardStore; use std::collections::HashMap; +use std::fmt::Write as FmtWrite; use std::fs::{self, OpenOptions}; use std::io::{BufRead, Read, Write}; use std::os::fd::AsRawFd; @@ -298,15 +299,25 @@ fn main() -> Result<()> { bail!("status socket must differ from the TileMaxSim protocol socket"); } let reload = Arc::new(AtomicBool::new(false)); - let metrics = Arc::new(RuntimeMetrics::default()); + let metrics = Arc::new(RuntimeMetrics::new( + args.max_connections, + args.max_queued_requests, + args.max_tenant_queued_requests, + args.max_inflight_request_gb, + )); + metrics.update_engine(engine.status_snapshot()); install_signal_handlers()?; let ready_cache = engine.status_json(); let (sender, receiver) = mpsc::sync_channel::(args.max_queued_requests); - let frame_admission = Arc::new(ByteAdmission::new(args.max_inflight_request_gb)); + let frame_admission = Arc::new(ByteAdmission::new( + args.max_inflight_request_gb, + Arc::clone(&metrics), + )); let pending_admission = Arc::new(PendingAdmission::new( args.max_queued_requests, args.max_tenant_queued_requests, + Arc::clone(&metrics), )); let tenant_weights = args.tenant_weights.iter().cloned().collect(); let scheduler_config = SchedulerConfig { @@ -368,7 +379,6 @@ fn main() -> Result<()> { }) ); - let live_readers = Arc::new(AtomicUsize::new(0)); let mut readers = Vec::new(); let mut accepted = 0_usize; let mut status_failed = false; @@ -396,18 +406,18 @@ fn main() -> Result<()> { reap_readers(&mut readers); match listener.accept() { Ok((connection, _)) => { - if !try_acquire_reader(&live_readers, args.max_connections) { + if !try_acquire_reader(&metrics.active_connections, args.max_connections) { // Closing immediately is deliberate: we have not read enough // bytes to know whether the peer expects a v2 or v3 response. - metrics.rejected_global.fetch_add(1, Ordering::Relaxed); + metrics.rejected_connections.fetch_add(1, Ordering::Relaxed); drop(connection); continue; } accepted += 1; let reader_sender = sender.clone(); - let reader_count = Arc::clone(&live_readers); - let reader_admission = Arc::clone(&pending_admission); let reader_metrics = Arc::clone(&metrics); + let reader_admission = Arc::clone(&pending_admission); + let request_metrics = Arc::clone(&metrics); let reader_frame_admission = Arc::clone(&frame_admission); let reader_config = ReaderConfig { maximum: args.max_request_bytes, @@ -417,19 +427,19 @@ fn main() -> Result<()> { match thread::Builder::new() .name("tilemaxsim-reader".to_owned()) .spawn(move || { - let _permit = ReaderPermit(reader_count); + let _permit = ReaderPermit(reader_metrics); read_and_enqueue( connection, &reader_sender, reader_config, reader_admission, - reader_metrics, + request_metrics, reader_frame_admission, ); }) { Ok(reader) => readers.push(reader), Err(error) => { - live_readers.fetch_sub(1, Ordering::Release); + metrics.active_connections.fetch_sub(1, Ordering::Release); metrics.ready.store(false, Ordering::Release); SHUTDOWN_REQUESTED.store(true, Ordering::Release); fatal_error = Some(error.into()); @@ -513,21 +523,101 @@ struct SchedulerConfig { #[derive(Default)] struct RuntimeMetrics { ready: AtomicBool, + max_connections: usize, + max_pending_requests: usize, + max_tenant_pending_requests: usize, + max_inflight_request_bytes: usize, + active_connections: AtomicUsize, + inflight_request_bytes: AtomicUsize, + pending_requests: AtomicUsize, + pending_tenants: AtomicUsize, scheduler_depth: AtomicUsize, + scheduler_depth_high_water: AtomicUsize, gpu_active: AtomicUsize, completed: AtomicU64, failed: AtomicU64, timed_out: AtomicU64, disconnected: AtomicU64, - rejected_global: AtomicU64, + rejected_connections: AtomicU64, + rejected_frame_bytes: AtomicU64, + rejected_queue_global: AtomicU64, rejected_tenant: AtomicU64, + frame_read_failures: AtomicU64, + invalid_requests: AtomicU64, + gpu_failures: AtomicU64, + scheduler_failures: AtomicU64, + reload_succeeded: AtomicU64, + reload_failed: AtomicU64, + timeout_before_enqueue: AtomicU64, + timeout_in_queue: AtomicU64, + timeout_before_execution: AtomicU64, + timeout_during_execution: AtomicU64, + gpu_quantums: AtomicU64, + scheduler_requeues: AtomicU64, + admitted_priority_negative: AtomicU64, + admitted_priority_zero: AtomicU64, + admitted_priority_positive: AtomicU64, + candidates_scored: AtomicU64, + document_rows_scored: AtomicU64, + latency_observations: AtomicU64, + total_latency_us: AtomicU64, + gpu_latency_us: AtomicU64, + engine: Mutex, +} + +impl RuntimeMetrics { + fn new( + max_connections: usize, + max_pending_requests: usize, + max_tenant_pending_requests: usize, + max_inflight_request_bytes: usize, + ) -> Self { + Self { + max_connections, + max_pending_requests, + max_tenant_pending_requests, + max_inflight_request_bytes, + ..Self::default() + } + } + + fn update_engine(&self, status: EngineStatus) { + *self + .engine + .lock() + .unwrap_or_else(|error| error.into_inner()) = status; + } + + fn update_scheduler_depth(&self, depth: usize) { + self.scheduler_depth.store(depth, Ordering::Relaxed); + self.scheduler_depth_high_water + .fetch_max(depth, Ordering::Relaxed); + } + + fn observe_latency(&self, total: Duration, gpu: Duration) { + self.latency_observations.fetch_add(1, Ordering::Relaxed); + saturating_atomic_add(&self.total_latency_us, duration_micros(total)); + saturating_atomic_add(&self.gpu_latency_us, duration_micros(gpu)); + } } -struct ReaderPermit(Arc); +fn duration_micros(duration: Duration) -> u64 { + u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) +} + +fn saturating_atomic_add(counter: &AtomicU64, value: u64) { + counter + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.saturating_add(value)) + }) + .ok(); +} + +struct ReaderPermit(Arc); impl Drop for ReaderPermit { fn drop(&mut self) { - self.0.fetch_sub(1, Ordering::Release); + self.0.active_connections.fetch_sub(1, Ordering::Release); } } @@ -540,11 +630,13 @@ struct PendingAdmission { state: Mutex, max_total: usize, max_tenant: usize, + metrics: Arc, } struct ByteAdmission { used: AtomicUsize, maximum: usize, + metrics: Arc, } struct BytePermit { @@ -553,21 +645,32 @@ struct BytePermit { } impl ByteAdmission { - fn new(maximum: usize) -> Self { + fn new(maximum: usize, metrics: Arc) -> Self { Self { used: AtomicUsize::new(0), maximum, + metrics, } } fn try_acquire(self: &Arc, bytes: usize) -> Option { - self.used + let result = self + .used .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { current .checked_add(bytes) .filter(|next| *next <= self.maximum) }) - .ok()?; + .ok(); + if result.is_none() { + self.metrics + .rejected_frame_bytes + .fetch_add(1, Ordering::Relaxed); + return None; + } + self.metrics + .inflight_request_bytes + .fetch_add(bytes, Ordering::Relaxed); Some(BytePermit { admission: Arc::clone(self), bytes, @@ -578,6 +681,10 @@ impl ByteAdmission { impl Drop for BytePermit { fn drop(&mut self) { self.admission.used.fetch_sub(self.bytes, Ordering::Release); + self.admission + .metrics + .inflight_request_bytes + .fetch_sub(self.bytes, Ordering::Release); } } @@ -593,7 +700,7 @@ struct PendingPermit { } impl PendingAdmission { - fn new(max_total: usize, max_tenant: usize) -> Self { + fn new(max_total: usize, max_tenant: usize, metrics: Arc) -> Self { Self { state: Mutex::new(PendingState { total: 0, @@ -601,6 +708,7 @@ impl PendingAdmission { }), max_total, max_tenant, + metrics, } } @@ -614,6 +722,12 @@ impl PendingAdmission { } state.total += 1; *state.tenants.entry(tenant.to_owned()).or_default() += 1; + self.metrics + .pending_requests + .store(state.total, Ordering::Relaxed); + self.metrics + .pending_tenants + .store(state.tenants.len(), Ordering::Relaxed); Ok(PendingPermit { admission: Arc::clone(self), tenant: tenant.to_owned(), @@ -635,6 +749,14 @@ impl Drop for PendingPermit { state.tenants.remove(&self.tenant); } } + self.admission + .metrics + .pending_requests + .store(state.total, Ordering::Relaxed); + self.admission + .metrics + .pending_tenants + .store(state.tenants.len(), Ordering::Relaxed); } } @@ -686,15 +808,16 @@ fn read_and_enqueue( } let (frame, frame_permit) = match read_request(&mut connection, config.maximum, &frame_admission) { - Ok(frame) => frame, - Err(error) => { - metrics.failed.fetch_add(1, Ordering::Relaxed); - write_response_nonfatal( - &mut connection, - &protocol::failure(VERSION_EXTERNAL, 0, 1, &format!("{error:#}")), - ); - return; - } + Ok(frame) => frame, + Err(error) => { + metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.frame_read_failures.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure(VERSION_EXTERNAL, 0, 1, &format!("{error:#}")), + ); + return; + } }; let version = header_version(&frame); let request_id = header_request_id(&frame); @@ -702,6 +825,7 @@ fn read_and_enqueue( Ok(request) => request, Err(error) => { metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.invalid_requests.fetch_add(1, Ordering::Relaxed); write_response_nonfatal( &mut connection, &protocol::failure(version, request_id, 1, &format!("{error:#}")), @@ -719,6 +843,9 @@ fn read_and_enqueue( .unwrap_or(accepted_at); if deadline <= Instant::now() { metrics.timed_out.fetch_add(1, Ordering::Relaxed); + metrics + .timeout_before_enqueue + .fetch_add(1, Ordering::Relaxed); write_response_nonfatal( &mut connection, &protocol::failure( @@ -733,7 +860,9 @@ fn read_and_enqueue( let pending_permit = match pending_admission.try_acquire(&request.tenant) { Ok(permit) => permit, Err(AdmissionRejection::Global) => { - metrics.rejected_global.fetch_add(1, Ordering::Relaxed); + metrics + .rejected_queue_global + .fetch_add(1, Ordering::Relaxed); write_response_nonfatal( &mut connection, &protocol::failure(version, request_id, 2, "TileMaxSim scheduler queue is full"), @@ -767,7 +896,9 @@ fn read_and_enqueue( }) { Ok(()) => {} Err(mpsc::TrySendError::Full(mut work)) => { - metrics.rejected_global.fetch_add(1, Ordering::Relaxed); + metrics + .rejected_queue_global + .fetch_add(1, Ordering::Relaxed); write_response_nonfatal( &mut work.connection, &protocol::failure( @@ -780,6 +911,7 @@ fn read_and_enqueue( } Err(mpsc::TrySendError::Disconnected(mut work)) => { metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.scheduler_failures.fetch_add(1, Ordering::Relaxed); write_response_nonfatal( &mut work.connection, &protocol::failure( @@ -810,11 +942,18 @@ fn run_scheduler( while channel_open || queue.len() > 0 { if reload.swap(false, Ordering::AcqRel) { match engine.reload_shards() { - Ok(()) => println!( - "{}", - serde_json::json!({"event": "tilemaxsim_rust_shards_reloaded"}) - ), - Err(error) => eprintln!("TileMaxSim shard reload rejected: {error:#}"), + Ok(()) => { + metrics.reload_succeeded.fetch_add(1, Ordering::Relaxed); + metrics.update_engine(engine.status_snapshot()); + println!( + "{}", + serde_json::json!({"event": "tilemaxsim_rust_shards_reloaded"}) + ); + } + Err(error) => { + metrics.reload_failed.fetch_add(1, Ordering::Relaxed); + eprintln!("TileMaxSim shard reload rejected: {error:#}"); + } } } @@ -846,6 +985,8 @@ fn run_scheduler( let now = Instant::now(); for mut expired in queue.drain_expired(now) { metrics.timed_out.fetch_add(1, Ordering::Relaxed); + metrics.timeout_in_queue.fetch_add(1, Ordering::Relaxed); + metrics.observe_latency(expired.payload.accepted_at.elapsed(), Duration::ZERO); write_response_nonfatal( &mut expired.payload.connection, &protocol::failure( @@ -856,17 +997,14 @@ fn run_scheduler( ), ); } - metrics - .scheduler_depth - .store(queue.len(), Ordering::Relaxed); + metrics.update_scheduler_depth(queue.len()); let Some(scheduled) = queue.pop(Instant::now()) else { continue; }; - metrics - .scheduler_depth - .store(queue.len(), Ordering::Relaxed); + metrics.update_scheduler_depth(queue.len()); if peer_disconnected(&scheduled.payload.connection) { metrics.disconnected.fetch_add(1, Ordering::Relaxed); + metrics.observe_latency(scheduled.payload.accepted_at.elapsed(), Duration::ZERO); continue; } let started = Instant::now(); @@ -877,6 +1015,9 @@ fn run_scheduler( let priority = work.request.priority; let response = if work.deadline <= started { metrics.timed_out.fetch_add(1, Ordering::Relaxed); + metrics + .timeout_before_execution + .fetch_add(1, Ordering::Relaxed); Some(protocol::failure( version, request_id, @@ -898,15 +1039,33 @@ fn run_scheduler( candidates: work.request.candidates[work.next_candidate..end].to_vec(), }; let quantum_started = Instant::now(); + metrics.gpu_quantums.fetch_add(1, Ordering::Relaxed); + saturating_atomic_add( + &metrics.candidates_scored, + u64::try_from(end - work.next_candidate).unwrap_or(u64::MAX), + ); + saturating_atomic_add( + &metrics.document_rows_scored, + quantum + .candidates + .iter() + .map(|candidate| u64::from(candidate.rows)) + .sum(), + ); metrics.gpu_active.store(1, Ordering::Relaxed); - match engine.score(&quantum) { + let score_result = engine.score(&quantum); + metrics.gpu_active.store(0, Ordering::Relaxed); + metrics.update_engine(engine.status_snapshot()); + match score_result { Ok(results) => { - metrics.gpu_active.store(0, Ordering::Relaxed); work.gpu_elapsed += quantum_started.elapsed(); work.results.extend(results); work.next_candidate = end; if work.deadline <= Instant::now() { metrics.timed_out.fetch_add(1, Ordering::Relaxed); + metrics + .timeout_during_execution + .fetch_add(1, Ordering::Relaxed); Some(protocol::failure( version, request_id, @@ -914,6 +1073,7 @@ fn run_scheduler( "request deadline expired during GPU execution", )) } else if end < work.request.candidates.len() { + metrics.scheduler_requeues.fetch_add(1, Ordering::Relaxed); let cost = estimated_next_work(&work, &config); queue.push(Scheduled::new( tenant.clone(), @@ -923,9 +1083,7 @@ fn run_scheduler( work.deadline, work, )); - metrics - .scheduler_depth - .store(queue.len(), Ordering::Relaxed); + metrics.update_scheduler_depth(queue.len()); continue; } else { metrics.completed.fetch_add(1, Ordering::Relaxed); @@ -933,8 +1091,8 @@ fn run_scheduler( } } Err(error) => { - metrics.gpu_active.store(0, Ordering::Relaxed); metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.gpu_failures.fetch_add(1, Ordering::Relaxed); work.gpu_elapsed += quantum_started.elapsed(); Some(protocol::failure( version, @@ -952,6 +1110,7 @@ fn run_scheduler( .set_write_timeout(Some(config.socket_io_timeout)) .ok(); write_response_nonfatal(&mut work.connection, &response); + metrics.observe_latency(work.accepted_at.elapsed(), work.gpu_elapsed); println!( "{}", serde_json::json!({ @@ -976,6 +1135,17 @@ fn enqueue_work( config: &SchedulerConfig, metrics: &RuntimeMetrics, ) { + match work.request.priority.cmp(&0) { + std::cmp::Ordering::Less => metrics + .admitted_priority_negative + .fetch_add(1, Ordering::Relaxed), + std::cmp::Ordering::Equal => metrics + .admitted_priority_zero + .fetch_add(1, Ordering::Relaxed), + std::cmp::Ordering::Greater => metrics + .admitted_priority_positive + .fetch_add(1, Ordering::Relaxed), + }; let cost = estimated_next_work(&work, config); queue.push(Scheduled::new( work.request.tenant.clone(), @@ -985,9 +1155,7 @@ fn enqueue_work( work.deadline, work, )); - metrics - .scheduler_depth - .store(queue.len(), Ordering::Relaxed); + metrics.update_scheduler_depth(queue.len()); } fn tenant_hash(tenant: &str) -> String { @@ -1136,38 +1304,540 @@ fn handle_status_connection(connection: &mut UnixStream, metrics: &RuntimeMetric } fn render_metrics(metrics: &RuntimeMetrics) -> String { - format!( - concat!( - "# HELP tilemaxsim_ready Whether the daemon is ready to accept work.\n", - "# TYPE tilemaxsim_ready gauge\n", - "tilemaxsim_ready {}\n", - "# HELP tilemaxsim_scheduler_queue_depth Requests waiting for a GPU quantum.\n", - "# TYPE tilemaxsim_scheduler_queue_depth gauge\n", - "tilemaxsim_scheduler_queue_depth {}\n", - "# HELP tilemaxsim_gpu_active Whether a CUDA quantum is executing.\n", - "# TYPE tilemaxsim_gpu_active gauge\n", - "tilemaxsim_gpu_active {}\n", - "# HELP tilemaxsim_requests_total Completed request outcomes.\n", - "# TYPE tilemaxsim_requests_total counter\n", - "tilemaxsim_requests_total{{outcome=\"completed\"}} {}\n", - "tilemaxsim_requests_total{{outcome=\"failed\"}} {}\n", - "tilemaxsim_requests_total{{outcome=\"timeout\"}} {}\n", - "tilemaxsim_requests_total{{outcome=\"disconnected\"}} {}\n", - "# HELP tilemaxsim_admission_rejections_total Admission rejections.\n", - "# TYPE tilemaxsim_admission_rejections_total counter\n", - "tilemaxsim_admission_rejections_total{{reason=\"global\"}} {}\n", - "tilemaxsim_admission_rejections_total{{reason=\"tenant\"}} {}\n", + let mut output = String::with_capacity(8 * 1024); + writeln!( + output, + "# HELP tilemaxsim_ready Whether the daemon is ready to accept work." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_ready gauge").unwrap(); + writeln!( + output, + "tilemaxsim_ready {}", + usize::from(metrics.ready.load(Ordering::Relaxed)) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_connections Active reader connections and configured limit." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_connections gauge").unwrap(); + writeln!( + output, + "tilemaxsim_connections{{kind=\"active\"}} {}", + metrics.active_connections.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_connections{{kind=\"limit\"}} {}", + metrics.max_connections + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_pending_requests Requests admitted but not yet completed." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_pending_requests gauge").unwrap(); + writeln!( + output, + "tilemaxsim_pending_requests{{kind=\"current\"}} {}", + metrics.pending_requests.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_pending_requests{{kind=\"global_limit\"}} {}", + metrics.max_pending_requests + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_pending_requests{{kind=\"per_tenant_limit\"}} {}", + metrics.max_tenant_pending_requests + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_pending_tenants Tenants with admitted requests." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_pending_tenants gauge").unwrap(); + writeln!( + output, + "tilemaxsim_pending_tenants {}", + metrics.pending_tenants.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_inflight_request_bytes Encoded request frames retained in memory." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_inflight_request_bytes gauge").unwrap(); + writeln!( + output, + "tilemaxsim_inflight_request_bytes{{kind=\"current\"}} {}", + metrics.inflight_request_bytes.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_inflight_request_bytes{{kind=\"limit\"}} {}", + metrics.max_inflight_request_bytes + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_scheduler_queue_depth Requests waiting for a GPU quantum." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_scheduler_queue_depth gauge").unwrap(); + writeln!( + output, + "tilemaxsim_scheduler_queue_depth{{kind=\"current\"}} {}", + metrics.scheduler_depth.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_scheduler_queue_depth{{kind=\"high_water\"}} {}", + metrics.scheduler_depth_high_water.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_active Whether a CUDA quantum is executing." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_active gauge").unwrap(); + writeln!( + output, + "tilemaxsim_gpu_active {}", + metrics.gpu_active.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_requests_total Terminal request outcomes." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_requests_total counter").unwrap(); + for (outcome, value) in [ + ("completed", metrics.completed.load(Ordering::Relaxed)), + ("failed", metrics.failed.load(Ordering::Relaxed)), + ("timeout", metrics.timed_out.load(Ordering::Relaxed)), + ("disconnected", metrics.disconnected.load(Ordering::Relaxed)), + ] { + writeln!( + output, + "tilemaxsim_requests_total{{outcome=\"{outcome}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_admission_rejections_total Bounded admission rejection reasons." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_admission_rejections_total counter" + ) + .unwrap(); + for (reason, value) in [ + ( + "connections", + metrics.rejected_connections.load(Ordering::Relaxed), + ), + ( + "frame_bytes", + metrics.rejected_frame_bytes.load(Ordering::Relaxed), ), - usize::from(metrics.ready.load(Ordering::Relaxed)), - metrics.scheduler_depth.load(Ordering::Relaxed), - metrics.gpu_active.load(Ordering::Relaxed), - metrics.completed.load(Ordering::Relaxed), - metrics.failed.load(Ordering::Relaxed), - metrics.timed_out.load(Ordering::Relaxed), - metrics.disconnected.load(Ordering::Relaxed), - metrics.rejected_global.load(Ordering::Relaxed), - metrics.rejected_tenant.load(Ordering::Relaxed), + ( + "queue_global", + metrics.rejected_queue_global.load(Ordering::Relaxed), + ), + ( + "queue_tenant", + metrics.rejected_tenant.load(Ordering::Relaxed), + ), + ] { + writeln!( + output, + "tilemaxsim_admission_rejections_total{{reason=\"{reason}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_failures_total Internal failure categories." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_failures_total counter").unwrap(); + for (reason, value) in [ + ( + "frame_io_or_validation", + metrics.frame_read_failures.load(Ordering::Relaxed), + ), + ("request", metrics.invalid_requests.load(Ordering::Relaxed)), + ("gpu", metrics.gpu_failures.load(Ordering::Relaxed)), + ( + "scheduler", + metrics.scheduler_failures.load(Ordering::Relaxed), + ), + ] { + writeln!( + output, + "tilemaxsim_failures_total{{reason=\"{reason}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_timeouts_total Request timeout phase." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_timeouts_total counter").unwrap(); + for (phase, value) in [ + ( + "before_enqueue", + metrics.timeout_before_enqueue.load(Ordering::Relaxed), + ), + ("queue", metrics.timeout_in_queue.load(Ordering::Relaxed)), + ( + "before_execution", + metrics.timeout_before_execution.load(Ordering::Relaxed), + ), + ( + "gpu", + metrics.timeout_during_execution.load(Ordering::Relaxed), + ), + ] { + writeln!( + output, + "tilemaxsim_timeouts_total{{phase=\"{phase}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_requests_admitted_total Requests admitted by bounded priority class." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_requests_admitted_total counter").unwrap(); + for (priority_class, value) in [ + ( + "negative", + metrics.admitted_priority_negative.load(Ordering::Relaxed), + ), + ( + "zero", + metrics.admitted_priority_zero.load(Ordering::Relaxed), + ), + ( + "positive", + metrics.admitted_priority_positive.load(Ordering::Relaxed), + ), + ] { + writeln!( + output, + "tilemaxsim_requests_admitted_total{{priority_class=\"{priority_class}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_scheduler_quantums_total CUDA scheduling quanta executed." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_scheduler_quantums_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_scheduler_quantums_total {}", + metrics.gpu_quantums.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_scheduler_requeues_total Cooperative quantum yields requeued." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_scheduler_requeues_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_scheduler_requeues_total {}", + metrics.scheduler_requeues.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_candidates_scored_total Candidate tensors submitted to CUDA." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_candidates_scored_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_candidates_scored_total {}", + metrics.candidates_scored.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_document_rows_scored_total Document token rows submitted to CUDA." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_document_rows_scored_total counter" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_document_rows_scored_total {}", + metrics.document_rows_scored.load(Ordering::Relaxed) ) + .unwrap(); + let observations = metrics.latency_observations.load(Ordering::Relaxed); + let total_us = metrics.total_latency_us.load(Ordering::Relaxed); + let gpu_us = metrics.gpu_latency_us.load(Ordering::Relaxed); + writeln!( + output, + "# HELP tilemaxsim_request_duration_seconds Request wall-clock duration summary." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_request_duration_seconds summary").unwrap(); + writeln!( + output, + "tilemaxsim_request_duration_seconds_count {observations}" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_request_duration_seconds_sum {}", + total_us as f64 / 1_000_000.0 + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_duration_seconds CUDA execution duration summary." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_duration_seconds summary").unwrap(); + writeln!( + output, + "tilemaxsim_gpu_duration_seconds_count {observations}" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_gpu_duration_seconds_sum {}", + gpu_us as f64 / 1_000_000.0 + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_queue_duration_seconds Non-CUDA request duration summary." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_queue_duration_seconds summary").unwrap(); + writeln!( + output, + "tilemaxsim_queue_duration_seconds_count {observations}" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_queue_duration_seconds_sum {}", + total_us.saturating_sub(gpu_us) as f64 / 1_000_000.0 + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_shard_reloads_total Immutable shard metadata reload outcomes." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_shard_reloads_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_shard_reloads_total{{outcome=\"success\"}} {}", + metrics.reload_succeeded.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_shard_reloads_total{{outcome=\"failed\"}} {}", + metrics.reload_failed.load(Ordering::Relaxed) + ) + .unwrap(); + + let engine = metrics + .engine + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone(); + writeln!( + output, + "# HELP tilemaxsim_gpu_cache_bytes GPU tensor-cache byte accounting." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_cache_bytes gauge").unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_cache_entries GPU tensor-cache entry accounting." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_cache_entries gauge").unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_cache_events_total GPU tensor-cache cumulative events." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_cache_events_total counter").unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_h2d_batches_total Host-to-device transfer batches." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_h2d_batches_total counter").unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_h2d_bytes_total Host-to-device transfer bytes." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_h2d_bytes_total counter").unwrap(); + for device in &engine.devices { + for (kind, value) in [ + ("capacity", device.capacity_bytes), + ("free", device.free_bytes), + ("largest_free_extent", device.largest_free_extent_bytes), + ("allocated", device.allocated_bytes), + ("payload", device.payload_bytes), + ("internal_waste", device.internal_waste_bytes), + ("pinned", device.pinned_bytes), + ("block", device.block_bytes), + ] { + writeln!( + output, + "tilemaxsim_gpu_cache_bytes{{slot=\"{}\",device=\"{}\",kind=\"{kind}\"}} {value}", + device.slot, device.device + ) + .unwrap(); + } + for (kind, value) in [ + ("total", device.entries), + ("pinned", device.pinned_entries), + ("tenants", device.tenants), + ] { + writeln!( + output, + "tilemaxsim_gpu_cache_entries{{slot=\"{}\",device=\"{}\",kind=\"{kind}\"}} {value}", + device.slot, device.device + ) + .unwrap(); + } + for (event, value) in [ + ("hit", device.hits), + ("miss", device.misses), + ("eviction", device.evictions), + ("admission_rejection", device.admission_rejections), + ] { + writeln!(output, "tilemaxsim_gpu_cache_events_total{{slot=\"{}\",device=\"{}\",event=\"{event}\"}} {value}", device.slot, device.device).unwrap(); + } + writeln!( + output, + "tilemaxsim_gpu_h2d_batches_total{{slot=\"{}\",device=\"{}\"}} {}", + device.slot, device.device, device.h2d_batches + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_gpu_h2d_bytes_total{{slot=\"{}\",device=\"{}\"}} {}", + device.slot, device.device, device.h2d_bytes + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_host_cache_bytes Host tensor-cache byte accounting." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_host_cache_bytes gauge").unwrap(); + writeln!( + output, + "tilemaxsim_host_cache_bytes{{kind=\"capacity\"}} {}", + engine.host.capacity_bytes + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_host_cache_bytes{{kind=\"used\"}} {}", + engine.host.used_bytes + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_host_cache_entries Host tensor-cache entries and active tenants." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_host_cache_entries gauge").unwrap(); + writeln!( + output, + "tilemaxsim_host_cache_entries{{kind=\"entries\"}} {}", + engine.host.entries + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_host_cache_entries{{kind=\"tenants\"}} {}", + engine.host.tenants + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_host_cache_events_total Host tensor-cache cumulative events." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_host_cache_events_total counter").unwrap(); + for (event, value) in [ + ("hit", engine.host.hits), + ("miss", engine.host.misses), + ("eviction", engine.host.evictions), + ("admission_rejection", engine.host.admission_rejections), + ] { + writeln!( + output, + "tilemaxsim_host_cache_events_total{{event=\"{event}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_storage_read_calls_total Immutable tensor storage read calls." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_storage_read_calls_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_storage_read_calls_total {}", + engine.batch_read_calls + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_storage_read_bytes_total Immutable tensor storage read bytes." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_storage_read_bytes_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_storage_read_bytes_total {}", + engine.batch_read_bytes + ) + .unwrap(); + output } fn read_request( @@ -1318,7 +1988,9 @@ mod tests { ByteAdmission, PendingAdmission, RuntimeMetrics, kib_to_bytes, quantum_end, render_metrics, tenant_hash, }; + use crate::engine::{DeviceStatus, EngineStatus}; use crate::protocol::Descriptor; + use crate::shard::HostCacheStatus; use std::sync::Arc; #[test] @@ -1345,7 +2017,8 @@ mod tests { #[test] fn pending_admission_is_globally_and_per_tenant_bounded() { - let admission = Arc::new(PendingAdmission::new(2, 1)); + let metrics = Arc::new(RuntimeMetrics::default()); + let admission = Arc::new(PendingAdmission::new(2, 1, metrics)); let first = admission.try_acquire("a").unwrap(); assert!(admission.try_acquire("a").is_err()); let second = admission.try_acquire("b").unwrap(); @@ -1357,7 +2030,8 @@ mod tests { #[test] fn in_flight_frame_bytes_are_bounded_until_the_permit_drops() { - let admission = Arc::new(ByteAdmission::new(100)); + let metrics = Arc::new(RuntimeMetrics::default()); + let admission = Arc::new(ByteAdmission::new(100, metrics)); let first = admission.try_acquire(60).unwrap(); assert!(admission.try_acquire(41).is_none()); let second = admission.try_acquire(40).unwrap(); @@ -1383,9 +2057,31 @@ mod tests { metrics .completed .store(7, std::sync::atomic::Ordering::Relaxed); + metrics.update_engine(EngineStatus { + devices: vec![DeviceStatus { + slot: 0, + device: 3, + capacity_bytes: 1_024, + free_bytes: 512, + largest_free_extent_bytes: 256, + ..DeviceStatus::default() + }], + host: HostCacheStatus { + capacity_bytes: 2_048, + used_bytes: 128, + ..HostCacheStatus::default() + }, + batch_read_calls: 4, + batch_read_bytes: 256, + }); let output = render_metrics(&metrics); assert!(output.contains("tilemaxsim_ready 1")); assert!(output.contains("outcome=\"completed\"} 7")); + assert!(output.contains( + "tilemaxsim_gpu_cache_bytes{slot=\"0\",device=\"3\",kind=\"capacity\"} 1024" + )); + assert!(output.contains("tilemaxsim_host_cache_bytes{kind=\"used\"} 128")); + assert!(output.contains("tilemaxsim_storage_read_bytes_total 256")); assert!(!output.contains("tenant-a")); } } diff --git a/services/tilemaxsimd/src/shard.rs b/services/tilemaxsimd/src/shard.rs index 044b27fa..ce015210 100644 --- a/services/tilemaxsimd/src/shard.rs +++ b/services/tilemaxsimd/src/shard.rs @@ -203,6 +203,18 @@ pub struct ShardStore { verify_full_shards: bool, } +#[derive(Clone, Debug, Default)] +pub struct HostCacheStatus { + pub capacity_bytes: usize, + pub used_bytes: usize, + pub entries: usize, + pub tenants: usize, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, +} + impl ShardStore { pub fn open( roots: &[(String, PathBuf)], @@ -372,13 +384,17 @@ impl ShardStore { .collect() } - pub fn host_status(&self) -> (u64, u64, u64, u64) { - ( - self.host_cache.hits, - self.host_cache.misses, - self.host_cache.evictions, - self.host_cache.admission_rejections, - ) + pub fn host_status(&self) -> HostCacheStatus { + HostCacheStatus { + capacity_bytes: self.host_cache.maximum_bytes, + used_bytes: self.host_cache.current_bytes, + entries: self.host_cache.entries.len(), + tenants: self.host_cache.tenant_bytes.len(), + hits: self.host_cache.hits, + misses: self.host_cache.misses, + evictions: self.host_cache.evictions, + admission_rejections: self.host_cache.admission_rejections, + } } } From 947a779a117e4c50533f3758de5319f475716ce3 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 19:12:28 +0800 Subject: [PATCH 316/324] fix(tilemaxsim): close production release blockers --- .dockerignore | 16 + .github/workflows/check.yml | 103 +- .github/workflows/cla.yml | 43 - .github/workflows/release.yml | 40 +- META.json.in | 13 +- README.md | 25 +- crates/vchordrq/src/maxsim_cost.rs | 2 +- crates/vchordrq/src/statistics.rs | 2 +- devtools/test_tilemaxsim_reference_sidecar.py | 8 +- devtools/tilemaxsim_reference_sidecar.py | 8 +- services/Dockerfile.tilemaxsimd | 13 + services/benchmark_gbrain_phase3.py | 4 - services/benchmark_npy_tilemaxsim_sidecar.py | 4 - services/benchmark_tilemaxsim_cuda.py | 6 +- services/build_tilemaxsim_tensor_cache.py | 6 +- services/test_tilemaxsim_cuda_sidecar.py | 6 +- services/test_tilemaxsim_gpu_cache.py | 4 - services/test_tilemaxsim_rust_daemon.py | 193 ++ services/tilemaxsim_cuda_sidecar.py | 6 +- services/tilemaxsim_gpu_cache.py | 4 - services/tilemaxsim_shard.py | 4 - services/tilemaxsim_triton.py | 4 - services/tilemaxsimd/build.rs | 11 +- .../tilemaxsimd/native/tilemaxsim_cuda.cu | 81 +- services/tilemaxsimd/src/bin/tilemaxsimctl.rs | 10 + services/tilemaxsimd/src/cache.rs | 10 + services/tilemaxsimd/src/engine.rs | 10 + services/tilemaxsimd/src/gpu.rs | 10 + services/tilemaxsimd/src/main.rs | 63 +- services/tilemaxsimd/src/protocol.rs | 10 + services/tilemaxsimd/src/scheduler.rs | 10 + services/tilemaxsimd/src/shard.rs | 10 + sql/install/vchord--1.2.0.sql | 2912 +++++++++++++++++ sql/upgrade/vchord--1.1.1--1.2.0.sql | 1587 +++++++++ src/index/vchordrq/scanners/maxsim/exact.rs | 4 +- .../vchordrq/scanners/maxsim/external.rs | 6 +- src/index/vchordrq/scanners/maxsim/gpu.rs | 6 +- src/index/vchordrq/scanners/maxsim/profile.rs | 4 - src/index/vchordrq/scanners/maxsim/search.rs | 3 +- 39 files changed, 5086 insertions(+), 175 deletions(-) create mode 100644 .dockerignore delete mode 100644 .github/workflows/cla.yml create mode 100644 sql/install/vchord--1.2.0.sql create mode 100644 sql/upgrade/vchord--1.1.1--1.2.0.sql diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..774769fa --- /dev/null +++ b/.dockerignore @@ -0,0 +1,16 @@ +.git +.github +build +target +**/target +TODO* +docs +sql +tests +*.md +*.zip +*.deb +__pycache__ +**/__pycache__ +.pytest_cache +**/.pytest_cache diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index f4eeb1c1..d4e1f591 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -32,7 +32,7 @@ jobs: uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Typos - uses: crate-ci/typos@master + uses: crate-ci/typos@3bc303c295c081add5df4a4d52a1e117f2fb2dce - name: Taplo run: taplo fmt --check @@ -45,7 +45,7 @@ jobs: - name: License Header run: | - HEADER=$(cat < ARRAY[1, 0, 0]::vector + LIMIT 1; + SQL + test "$(psql -At -d vchord_upgrade_test -c \ + "SELECT extversion FROM pg_extension WHERE extname = 'vchord'")" = "1.2.0" + test "$(psql -At -d vchord_upgrade_test -c \ + "SELECT to_regprocedure('vchordrq_tilemaxsim_rerank(regclass,vector[],bigint[],integer)') IS NOT NULL")" = "t" + + dropdb --if-exists vchord_fresh_install_test + createdb vchord_fresh_install_test + psql -v ON_ERROR_STOP=1 -d vchord_fresh_install_test <<'SQL' + CREATE EXTENSION vector; + CREATE EXTENSION vchord VERSION '1.2.0'; + SQL + test "$(psql -At -d vchord_fresh_install_test -c \ + "SELECT extversion FROM pg_extension WHERE extname = 'vchord'")" = "1.2.0" + test "$(psql -At -d vchord_fresh_install_test -c \ + "SELECT to_regprocedure('vchordrq_tilemaxsim_rerank(regclass,vector[],bigint[],integer)') IS NOT NULL")" = "t" + - name: Logging if: always() run: | diff --git a/.github/workflows/cla.yml b/.github/workflows/cla.yml deleted file mode 100644 index 97e090b8..00000000 --- a/.github/workflows/cla.yml +++ /dev/null @@ -1,43 +0,0 @@ -name: "CLA Assistant" -on: - issue_comment: - types: [created] - pull_request_target: - types: [opened,closed,synchronize] - -# explicitly configure permissions, in case your GITHUB_TOKEN workflow permissions are set to read-only in repository settings -permissions: - actions: write - contents: write # this can be 'read' if the signatures are in remote repository - pull-requests: write - statuses: write - -jobs: - CLAAssistant: - runs-on: ubuntu-latest - steps: - - name: "CLA Assistant" - if: (github.event.comment.body == 'recheck' || github.event.comment.body == 'I have read the CLA Document and I hereby sign the CLA') || github.event_name == 'pull_request_target' - uses: contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08 # v2.6.1 - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - # the below token should have repo scope and must be manually added by you in the repository's secret - # This token is required only if you have configured to store the signatures in a remote repository/organization - PERSONAL_ACCESS_TOKEN: ${{ secrets.CLA_GITHUB_TOKEN }} - with: - path-to-signatures: 'signatures/version1/cla.json' - path-to-document: 'https://gist.githubusercontent.com/gaocegege/92886a860e3648842e07651feb836719/raw/10663798a5e459707042a3d37bd6f89ac784dada/CLA.md' # e.g. a CLA or a DCO document - # The branch is protected, but we use this: https://github.com/contributor-assistant/github-action/issues/150 - branch: 'main' - allowlist: user1,bot* - - # the followings are the optional inputs - If the optional inputs are not given, then default values will be taken - remote-organization-name: tensorchord - remote-repository-name: CLA - #create-file-commit-message: 'For example: Creating file for storing CLA Signatures' - #signed-commit-message: 'For example: $contributorName has signed the CLA in $owner/$repo#$pullRequestNo' - #custom-notsigned-prcomment: 'pull request comment with Introductory message to ask new contributors to sign' - #custom-pr-sign-comment: 'The signature to be committed in order to sign the CLA' - #custom-allsigned-prcomment: 'pull request comment when all contributors has signed, defaults to **CLA Assistant Lite bot** All Contributors have signed the CLA.' - lock-pullrequest-aftermerge: false - if you don't want this bot to automatically lock the pull request after merging (default - true) - #use-dco-flag: true - If you are using DCO instead of CLA diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 2c426c2a..c469997c 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -14,6 +14,10 @@ concurrency: group: ${{ github.ref }}-${{ github.workflow }} cancel-in-progress: true +permissions: + contents: write + packages: write + jobs: semver: runs-on: "ubuntu-latest" @@ -101,10 +105,11 @@ jobs: mkdir -p ~/debian/usr/share/doc/postgresql-${VERSION}-vchord/ echo "Format: https://www.debian.org/doc/packaging-manuals/copyright-format/1.0/ Upstream-Name: VectorChord - Source: https://github.com/tensorchord/VectorChord + Source: https://github.com/HuXinjing/VectorChord Files: * Copyright: Copyright (c) 2025-2026 TensorChord Inc. + Copyright (c) 2026 Hu Xinjing License: AGPL-3.0-only OR Elastic-2.0 License: AGPL-3.0-only @@ -123,18 +128,44 @@ jobs: Section: database Priority: optional Architecture: ${PLATFORM} - Maintainer: Tensorchord + Maintainer: Hu Xinjing Description: Scalable, fast, and disk-friendly vector search in Postgres VectorChord is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search. It supports L2 distance, inner product, cosine distance and maxsim. - Homepage: https://vectorchord.ai/" \ + Homepage: https://github.com/HuXinjing/VectorChord" \ > ~/debian/DEBIAN/control (cd ~/debian && find usr -type f -print0 | xargs -0 md5sum) > ~/debian/DEBIAN/md5sums dpkg-deb --root-owner-group -Zxz --build ~/debian ~/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb gh release upload --clobber $SEMVER ~/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb + tilemaxsimd: + runs-on: ubuntu-24.04 + needs: ["semver"] + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Free space for the CUDA build image + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /usr/local/.ghcup + docker builder prune --all --force + + - name: Log in to the fork container registry + env: + GH_TOKEN: ${{ github.token }} + run: echo "$GH_TOKEN" | docker login ghcr.io -u "${{ github.actor }}" --password-stdin + + - name: Build and publish native TileMaxSim image + env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + run: | + IMAGE=$(echo "ghcr.io/${GITHUB_REPOSITORY_OWNER}/vectorchord-tilemaxsimd" | tr '[:upper:]' '[:lower:]') + docker build -f services/Dockerfile.tilemaxsimd -t "${IMAGE}:${SEMVER}" . + docker push "${IMAGE}:${SEMVER}" + pgxn: + if: ${{ vars.PUBLISH_PGXN == 'true' }} runs-on: "ubuntu-latest" needs: ["semver", "build"] @@ -145,6 +176,7 @@ jobs: - name: Upload env: SEMVER: ${{ needs.semver.outputs.SEMVER }} + PGXN_USERNAME: ${{ secrets.PGXN_USERNAME }} PGXN_PASSWORD: ${{ secrets.PGXN_PASSWORD }} run: | mkdir -p ./build @@ -153,7 +185,7 @@ jobs: git archive --format zip --prefix "vchord-${SEMVER}/" --add-file META.json -o "./build/vchord--${SEMVER}.zip" HEAD curl --fail -sS \ - --user "tensorchord:${PGXN_PASSWORD}" \ + --user "${PGXN_USERNAME}:${PGXN_PASSWORD}" \ -F "submit=Release It!" \ -F "archive=@./build/vchord--${SEMVER}.zip" \ -H "X-Requested-With: XMLHttpRequest" \ diff --git a/META.json.in b/META.json.in index ef79c686..19a2ffde 100644 --- a/META.json.in +++ b/META.json.in @@ -4,8 +4,7 @@ "description": "VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search, it supports L2 distance, inner product, cosine distance and maxsim", "version": "@DISTVERSION@", "maintainer": [ - "TensorChord support@tensorchord.ai", - "usamoi usamoi@tensorchord.ai" + "Hu Xinjing HuXinjing@users.noreply.github.com" ], "license": "agpl_3", "prereqs": { @@ -24,17 +23,17 @@ } }, "resources": { - "homepage": "https://github.com/tensorchord/VectorChord", + "homepage": "https://github.com/HuXinjing/VectorChord", "bugtracker": { - "web": "https://github.com/tensorchord/VectorChord/issues" + "web": "https://github.com/HuXinjing/VectorChord/issues" }, "repository": { - "url": "https://github.com/tensorchord/VectorChord.git", - "web": "https://github.com/tensorchord/VectorChord", + "url": "https://github.com/HuXinjing/VectorChord.git", + "web": "https://github.com/HuXinjing/VectorChord", "type": "git" } }, - "generated_by": "TensorChord", + "generated_by": "HuXinjing/VectorChord", "meta-spec": { "version": "1.0.0", "url": "http://pgxn.org/meta/spec.txt" diff --git a/README.md b/README.md index 1ae883d6..9415b69a 100644 --- a/README.md +++ b/README.md @@ -347,36 +347,39 @@ This software is licensed under a dual license model: 2. **Elastic License v2 (ELv2)**: You may also use, modify, and distribute this software under the Elastic License v2, which has specific restrictions. -You may choose either license based on your needs. We welcome any commercial collaboration or support, so please email us with any questions or requests regarding the licenses. +You may choose either license based on its terms. The original VectorChord code +retains its upstream copyright notices. TileMaxSim fork additions are Copyright +(c) 2026 Hu Xinjing; use this fork's +[issue tracker](https://github.com/HuXinjing/VectorChord/issues) for support. [cost-estimation]: https://github.com/user-attachments/assets/168fe550-6465-4eee-a224-8c848c301e3d [image-compare]: https://github.com/user-attachments/assets/2d985f1e-7093-4c3a-8bf3-9f0b92c0e7e7 -[license-1-link]: https://github.com/tensorchord/VectorChord#license +[license-1-link]: https://github.com/HuXinjing/VectorChord#license [license-1-shield]: https://img.shields.io/badge/License-AGPLv3-green?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCIgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiBmaWxsPSIjZmZmZmZmIj48cGF0aCBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik0xMi43NSAyLjc1YS43NS43NSAwIDAwLTEuNSAwVjQuNUg5LjI3NmExLjc1IDEuNzUgMCAwMC0uOTg1LjMwM0w2LjU5NiA1Ljk1N0EuMjUuMjUgMCAwMTYuNDU1IDZIMi4zNTNhLjc1Ljc1IDAgMTAwIDEuNUgzLjkzTC41NjMgMTUuMThhLjc2Mi43NjIgMCAwMC4yMS44OGMuMDguMDY0LjE2MS4xMjUuMzA5LjIyMS4xODYuMTIxLjQ1Mi4yNzguNzkyLjQzMy42OC4zMTEgMS42NjIuNjIgMi44NzYuNjJhNi45MTkgNi45MTkgMCAwMDIuODc2LS42MmMuMzQtLjE1NS42MDYtLjMxMi43OTItLjQzMy4xNS0uMDk3LjIzLS4xNTguMzEtLjIyM2EuNzUuNzUgMCAwMC4yMDktLjg3OEw1LjU2OSA3LjVoLjg4NmMuMzUxIDAgLjY5NC0uMTA2Ljk4NC0uMzAzbDEuNjk2LTEuMTU0QS4yNS4yNSAwIDAxOS4yNzUgNmgxLjk3NXYxNC41SDYuNzYzYS43NS43NSAwIDAwMCAxLjVoMTAuNDc0YS43NS43NSAwIDAwMC0xLjVIMTIuNzVWNmgxLjk3NGMuMDUgMCAuMS4wMTUuMTQuMDQzbDEuNjk3IDEuMTU0Yy4yOS4xOTcuNjMzLjMwMy45ODQuMzAzaC44ODZsLTMuMzY4IDcuNjhhLjc1Ljc1IDAgMDAuMjMuODk2Yy4wMTIuMDA5IDAgMCAuMDAyIDBhMy4xNTQgMy4xNTQgMCAwMC4zMS4yMDZjLjE4NS4xMTIuNDUuMjU2Ljc5LjRhNy4zNDMgNy4zNDMgMCAwMDIuODU1LjU2OCA3LjM0MyA3LjM0MyAwIDAwMi44NTYtLjU2OWMuMzM4LS4xNDMuNjA0LS4yODcuNzktLjM5OWEzLjUgMy41IDAgMDAuMzEtLjIwNi43NS43NSAwIDAwLjIzLS44OTZMMjAuMDcgNy41aDEuNTc4YS43NS43NSAwIDAwMC0xLjVoLTQuMTAyYS4yNS4yNSAwIDAxLS4xNC0uMDQzbC0xLjY5Ny0xLjE1NGExLjc1IDEuNzUgMCAwMC0uOTg0LS4zMDNIMTIuNzVWMi43NXpNMi4xOTMgMTUuMTk4YTUuNDE4IDUuNDE4IDAgMDAyLjU1Ny42MzUgNS40MTggNS40MTggMCAwMDIuNTU3LS42MzVMNC43NSA5LjM2OGwtMi41NTcgNS44M3ptMTQuNTEtLjAyNGMuMDgyLjA0LjE3NC4wODMuMjc1LjEyNi41My4yMjMgMS4zMDUuNDUgMi4yNzIuNDVhNS44NDYgNS44NDYgMCAwMDIuNTQ3LS41NzZMMTkuMjUgOS4zNjdsLTIuNTQ3IDUuODA3eiI+PC9wYXRoPjwvc3ZnPgo= -[license-2-link]: https://github.com/tensorchord/VectorChord#license +[license-2-link]: https://github.com/HuXinjing/VectorChord#license [license-2-shield]: https://img.shields.io/badge/License-ELv2-green?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCIgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiBmaWxsPSIjZmZmZmZmIj48cGF0aCBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik0xMi43NSAyLjc1YS43NS43NSAwIDAwLTEuNSAwVjQuNUg5LjI3NmExLjc1IDEuNzUgMCAwMC0uOTg1LjMwM0w2LjU5NiA1Ljk1N0EuMjUuMjUgMCAwMTYuNDU1IDZIMi4zNTNhLjc1Ljc1IDAgMTAwIDEuNUgzLjkzTC41NjMgMTUuMThhLjc2Mi43NjIgMCAwMC4yMS44OGMuMDguMDY0LjE2MS4xMjUuMzA5LjIyMS4xODYuMTIxLjQ1Mi4yNzguNzkyLjQzMy42OC4zMTEgMS42NjIuNjIgMi44NzYuNjJhNi45MTkgNi45MTkgMCAwMDIuODc2LS42MmMuMzQtLjE1NS42MDYtLjMxMi43OTItLjQzMy4xNS0uMDk3LjIzLS4xNTguMzEtLjIyM2EuNzUuNzUgMCAwMC4yMDktLjg3OEw1LjU2OSA3LjVoLjg4NmMuMzUxIDAgLjY5NC0uMTA2Ljk4NC0uMzAzbDEuNjk2LTEuMTU0QS4yNS4yNSAwIDAxOS4yNzUgNmgxLjk3NXYxNC41SDYuNzYzYS43NS43NSAwIDAwMCAxLjVoMTAuNDc0YS43NS43NSAwIDAwMC0xLjVIMTIuNzVWNmgxLjk3NGMuMDUgMCAuMS4wMTUuMTQuMDQzbDEuNjk3IDEuMTU0Yy4yOS4xOTcuNjMzLjMwMy45ODQuMzAzaC44ODZsLTMuMzY4IDcuNjhhLjc1Ljc1IDAgMDAuMjMuODk2Yy4wMTIuMDA5IDAgMCAuMDAyIDBhMy4xNTQgMy4xNTQgMCAwMC4zMS4yMDZjLjE4NS4xMTIuNDUuMjU2Ljc5LjRhNy4zNDMgNy4zNDMgMCAwMDIuODU1LjU2OCA3LjM0MyA3LjM0MyAwIDAwMi44NTYtLjU2OWMuMzM4LS4xNDMuNjA0LS4yODcuNzktLjM5OWEzLjUgMy41IDAgMDAuMzEtLjIwNi43NS43NSAwIDAwLjIzLS44OTZMMjAuMDcgNy41aDEuNTc4YS43NS43NSAwIDAwMC0xLjVoLTQuMTAyYS4yNS4yNSAwIDAxLS4xNC0uMDQzbC0xLjY5Ny0xLjE1NGExLjc1IDEuNzUgMCAwMC0uOTg0LS4zMDNIMTIuNzVWMi43NXpNMi4xOTMgMTUuMTk4YTUuNDE4IDUuNDE4IDAgMDAyLjU1Ny42MzUgNS40MTggNS40MTggMCAwMDIuNTU3LS42MzVMNC43NSA5LjM2OGwtMi41NTcgNS44M3ptMTQuNTEtLjAyNGMuMDgyLjA0LjE3NC4wODMuMjc1LjEyNi41My4yMjMgMS4zMDUuNDUgMi4yNzIuNDVhNS44NDYgNS44NDYgMCAwMDIuNTQ3LS41NzZMMTkuMjUgOS4zNjdsLTIuNTQ3IDUuODA3eiI+PC9wYXRoPjwvc3ZnPgo= [docker-release-link]: https://hub.docker.com/r/tensorchord/vchord-postgres [docker-release-shield]: https://img.shields.io/docker/v/tensorchord/vchord-postgres?color=369eff&label=docker&labelColor=black&logo=docker&logoColor=white&style=flat -[github-release-link]: https://github.com/tensorchord/VectorChord/releases -[github-release-shield]: https://img.shields.io/github/v/release/tensorchord/VectorChord?color=369eff&labelColor=black&logo=github&style=flat -[ghcr-release-link]: https://github.com/orgs/tensorchord/packages/container/package/vchord-postgres +[github-release-link]: https://github.com/HuXinjing/VectorChord/releases +[github-release-shield]: https://img.shields.io/github/v/release/HuXinjing/VectorChord?color=369eff&labelColor=black&logo=github&style=flat +[ghcr-release-link]: https://github.com/users/HuXinjing/packages/container/package/vectorchord-tilemaxsimd [ghcr-release-shield]: https://img.shields.io/docker/v/tensorchord/vchord-postgres?color=369eff&label=GHCR&labelColor=black&logo=github&logoColor=white&style=flat [docker-pulls-link]: https://hub.docker.com/r/tensorchord/vchord-postgres [docker-pulls-shield]: https://img.shields.io/docker/pulls/tensorchord/vchord-postgres?color=45cc11&labelColor=black&style=flat&sort=semver [previous-docker-pulls-link]: https://hub.docker.com/r/tensorchord/pgvecto-rs [previous-docker-pulls-shield]: https://img.shields.io/docker/pulls/tensorchord/pgvecto-rs?color=45cc11&labelColor=black&style=flat&sort=semver -[github-downloads-link]: https://github.com/tensorchord/VectorChord/releases -[github-downloads-shield]: https://img.shields.io/github/downloads/tensorchord/VectorChord/total?color=45cc11&labelColor=black&style=flat&sort=semver +[github-downloads-link]: https://github.com/HuXinjing/VectorChord/releases +[github-downloads-shield]: https://img.shields.io/github/downloads/HuXinjing/VectorChord/total?color=45cc11&labelColor=black&style=flat&sort=semver [discord-link]: https://discord.gg/KqswhpVgdU [discord-shield]: https://img.shields.io/discord/974584200327991326?&logoColor=white&color=5865F2&style=flat&logo=discord&cacheSeconds=60 [X-link]: https://x.com/TensorChord [X-shield]: https://img.shields.io/badge/follow-%40tensorchord-1DA1F2?logo=x&style=flat&logoColor=white&color=1da1f2 -[deepwiki-link]: https://deepwiki.com/tensorchord/VectorChord +[deepwiki-link]: https://deepwiki.com/HuXinjing/VectorChord [deepwiki-shield]: https://deepwiki.com/badge.svg [blog-link]: https://blog.vectorchord.ai/ [official-site-link]: https://vectorchord.ai/ -[github-issues-link]: https://github.com/tensorchord/VectorChord/issues -[email-link]: mailto:support@tensorchord.ai +[github-issues-link]: https://github.com/HuXinjing/VectorChord/issues +[email-link]: https://github.com/HuXinjing/VectorChord/issues [docs-link]: https://docs.vectorchord.ai/vectorchord/ diff --git a/crates/vchordrq/src/maxsim_cost.rs b/crates/vchordrq/src/maxsim_cost.rs index ca843d10..c63966d7 100644 --- a/crates/vchordrq/src/maxsim_cost.rs +++ b/crates/vchordrq/src/maxsim_cost.rs @@ -4,7 +4,7 @@ // distribute this software under the terms of the AGPLv3. // // Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the terms of the ELv2, which has specific restrictions. +// software under the Elastic License v2, which has specific restrictions. // // We welcome any commercial collaboration or support. For inquiries // regarding the licenses, please contact us at: diff --git a/crates/vchordrq/src/statistics.rs b/crates/vchordrq/src/statistics.rs index f651bec5..80d850e2 100644 --- a/crates/vchordrq/src/statistics.rs +++ b/crates/vchordrq/src/statistics.rs @@ -4,7 +4,7 @@ // distribute this software under the terms of the AGPLv3. // // Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the terms of the ELv2, which has specific restrictions. +// software under the Elastic License v2, which has specific restrictions. // // We welcome any commercial collaboration or support. For inquiries // regarding the licenses, please contact us at: diff --git a/devtools/test_tilemaxsim_reference_sidecar.py b/devtools/test_tilemaxsim_reference_sidecar.py index cf3bc826..19f0f01a 100644 --- a/devtools/test_tilemaxsim_reference_sidecar.py +++ b/devtools/test_tilemaxsim_reference_sidecar.py @@ -4,13 +4,9 @@ # distribute this software under the terms of the AGPLv3. # # Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the terms of the ELv2, which has specific restrictions. +# software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 Hu Xinjing from __future__ import annotations diff --git a/devtools/tilemaxsim_reference_sidecar.py b/devtools/tilemaxsim_reference_sidecar.py index 2076d2c8..267a77fc 100644 --- a/devtools/tilemaxsim_reference_sidecar.py +++ b/devtools/tilemaxsim_reference_sidecar.py @@ -4,13 +4,9 @@ # distribute this software under the terms of the AGPLv3. # # Elastic License v2 (ELv2): You may also use, modify, and distribute this -# software under the terms of the ELv2, which has specific restrictions. +# software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 Hu Xinjing """CPU reference implementation of the VectorChord TileMaxSim IPC sidecar. diff --git a/services/Dockerfile.tilemaxsimd b/services/Dockerfile.tilemaxsimd index 328b7f7a..b15b7088 100644 --- a/services/Dockerfile.tilemaxsimd +++ b/services/Dockerfile.tilemaxsimd @@ -15,10 +15,23 @@ WORKDIR /build/tilemaxsimd COPY services/tilemaxsimd/ . RUN cargo build --release --locked +FROM build AS test +RUN rustup component add clippy rustfmt \ + && cargo fmt --check \ + && cargo clippy --all-targets --locked -- -D warnings \ + && cargo test --locked + FROM ${CUDA_RUNTIME_IMAGE} COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimd /usr/local/bin/tilemaxsimd COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimctl /usr/local/bin/tilemaxsimctl +RUN groupadd --system vectorchord \ + && useradd --system --gid vectorchord --home-dir /nonexistent \ + --shell /usr/sbin/nologin vectorchord \ + && install -d -o vectorchord -g vectorchord -m 0750 \ + /run/vectorchord /var/lib/vectorchord + +USER vectorchord:vectorchord STOPSIGNAL SIGTERM HEALTHCHECK --interval=10s --timeout=2s --start-period=5m --retries=3 \ diff --git a/services/benchmark_gbrain_phase3.py b/services/benchmark_gbrain_phase3.py index 6359da41..d3955ca2 100644 --- a/services/benchmark_gbrain_phase3.py +++ b/services/benchmark_gbrain_phase3.py @@ -6,10 +6,6 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# # Copyright (c) 2026 Hu Xinjing """Measure native Phase 3A ranking against a committed GBrain TileMaxSim gold.""" diff --git a/services/benchmark_npy_tilemaxsim_sidecar.py b/services/benchmark_npy_tilemaxsim_sidecar.py index 212bde54..6b3f4b26 100644 --- a/services/benchmark_npy_tilemaxsim_sidecar.py +++ b/services/benchmark_npy_tilemaxsim_sidecar.py @@ -6,10 +6,6 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# # Copyright (c) 2026 Hu Xinjing """Benchmark CUDA sidecar that resolves verified tensors from an NPY corpus. diff --git a/services/benchmark_tilemaxsim_cuda.py b/services/benchmark_tilemaxsim_cuda.py index 698e4b84..31799d26 100644 --- a/services/benchmark_tilemaxsim_cuda.py +++ b/services/benchmark_tilemaxsim_cuda.py @@ -6,11 +6,7 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 Hu Xinjing """Reproducible synthetic load probe for the CUDA TileMaxSim executor.""" diff --git a/services/build_tilemaxsim_tensor_cache.py b/services/build_tilemaxsim_tensor_cache.py index e6adf7ad..e42e0da9 100644 --- a/services/build_tilemaxsim_tensor_cache.py +++ b/services/build_tilemaxsim_tensor_cache.py @@ -6,11 +6,7 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 Hu Xinjing """Publish NPY page tensors into the sidecar's immutable SHA-256 cache.""" diff --git a/services/test_tilemaxsim_cuda_sidecar.py b/services/test_tilemaxsim_cuda_sidecar.py index e015174d..fb4b88e7 100644 --- a/services/test_tilemaxsim_cuda_sidecar.py +++ b/services/test_tilemaxsim_cuda_sidecar.py @@ -6,11 +6,7 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 Hu Xinjing from __future__ import annotations diff --git a/services/test_tilemaxsim_gpu_cache.py b/services/test_tilemaxsim_gpu_cache.py index 6394d7ff..2fd82cc4 100644 --- a/services/test_tilemaxsim_gpu_cache.py +++ b/services/test_tilemaxsim_gpu_cache.py @@ -6,10 +6,6 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# # Copyright (c) 2026 Hu Xinjing from __future__ import annotations diff --git a/services/test_tilemaxsim_rust_daemon.py b/services/test_tilemaxsim_rust_daemon.py index 94cf978f..9b9dd71c 100644 --- a/services/test_tilemaxsim_rust_daemon.py +++ b/services/test_tilemaxsim_rust_daemon.py @@ -311,6 +311,21 @@ def test_external_v2_shard_round_trip_matches_protocol_oracle(self) -> None: ) self.run_daemon([device]) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_exact_scores_are_repeatable_across_cuda_contexts(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + generator = np.random.default_rng(20260715) + document = generator.normal(size=(79, 64)).astype(" None: binary = self._release_binary() @@ -402,6 +417,95 @@ def test_published_object_is_queryable_without_daemon_restart(self) -> None: process.terminate() process.wait(timeout=5) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_sigkill_restart_recovers_stale_socket_and_cache_root(self) -> None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shard_root = root / "cache" + shard_root.mkdir() + document = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: device = max( @@ -583,6 +687,95 @@ def call(frame: bytes) -> int: process.terminate() process.wait(timeout=10) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_overload_rejects_one_tenant_without_crashing_daemon(self) -> None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shard_root = root / "cache" + shard_root.mkdir() + document = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" tuple[int, int]: + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_bytes = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_bytes) + request_id, status, _ = decode_response(response) + return request_id, status + + first = None + try: + for _ in range(1000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + self.assertIsNone(process.poll()) + first = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + first.settimeout(10) + first.connect(os.fspath(socket_path)) + first.sendall(frames[0]) + time.sleep(0.05) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as second: + second.settimeout(10) + second.connect(os.fspath(socket_path)) + second.sendall(frames[1]) + self.assertEqual(receive(second), (2_002, 2)) + self.assertEqual(receive(first), (2_001, 0)) + self.assertIsNone(process.poll()) + finally: + if first is not None: + first.close() + if process.poll() is None: + process.terminate() + output, _ = process.communicate(timeout=10) + self.assertEqual(process.returncode, 0, output) + if __name__ == "__main__": unittest.main() diff --git a/services/tilemaxsim_cuda_sidecar.py b/services/tilemaxsim_cuda_sidecar.py index a84c33b0..fbbfd24a 100644 --- a/services/tilemaxsim_cuda_sidecar.py +++ b/services/tilemaxsim_cuda_sidecar.py @@ -6,11 +6,7 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# -# Copyright (c) 2025-2026 TensorChord Inc. +# Copyright (c) 2026 Hu Xinjing """Bounded CUDA executor for VectorChord TileMaxSim IPC v1 and v2. diff --git a/services/tilemaxsim_gpu_cache.py b/services/tilemaxsim_gpu_cache.py index dd17e7a3..0c9b5ef2 100644 --- a/services/tilemaxsim_gpu_cache.py +++ b/services/tilemaxsim_gpu_cache.py @@ -6,10 +6,6 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# # Copyright (c) 2026 Hu Xinjing """Process-owned GPU arenas and a bounded tensor cache for TileMaxSim.""" diff --git a/services/tilemaxsim_shard.py b/services/tilemaxsim_shard.py index 2837b4ed..453d58c9 100644 --- a/services/tilemaxsim_shard.py +++ b/services/tilemaxsim_shard.py @@ -6,10 +6,6 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# # Copyright (c) 2026 Hu Xinjing """Immutable, content-addressed TileMaxSim tensor shard format. diff --git a/services/tilemaxsim_triton.py b/services/tilemaxsim_triton.py index ee0736cc..8e9afce4 100644 --- a/services/tilemaxsim_triton.py +++ b/services/tilemaxsim_triton.py @@ -6,10 +6,6 @@ # Elastic License v2 (ELv2): You may also use, modify, and distribute this # software under the Elastic License v2, which has specific restrictions. # -# We welcome any commercial collaboration or support. For inquiries -# regarding the licenses, please contact us at: -# vectorchord-inquiry@tensorchord.ai -# # Copyright (c) 2026 Hu Xinjing """Fused ragged FP16 TileMaxSim over a process-owned GPU tensor arena.""" diff --git a/services/tilemaxsimd/build.rs b/services/tilemaxsimd/build.rs index 15cbfae6..a8cd9fff 100644 --- a/services/tilemaxsimd/build.rs +++ b/services/tilemaxsimd/build.rs @@ -1,9 +1,18 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + fn main() { println!("cargo:rerun-if-changed=native/tilemaxsim_cuda.cu"); cc::Build::new() .cuda(true) .flag("-O3") - .flag("--use_fast_math") .flag("-lineinfo") .file("native/tilemaxsim_cuda.cu") .compile("tilemaxsim_cuda"); diff --git a/services/tilemaxsimd/native/tilemaxsim_cuda.cu b/services/tilemaxsimd/native/tilemaxsim_cuda.cu index ff214e6e..6c6fc527 100644 --- a/services/tilemaxsimd/native/tilemaxsim_cuda.cu +++ b/services/tilemaxsimd/native/tilemaxsim_cuda.cu @@ -1,3 +1,13 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + #include #include #include @@ -175,7 +185,7 @@ __global__ void tilemaxsim_kernel(const Scalar *query, uint32_t query_rows, uint32_t dimension, const unsigned char *documents, const uint64_t *document_offsets, - const uint32_t *document_rows, float *scores) { + const uint32_t *document_rows, float *maxima) { const uint32_t candidate = blockIdx.x; const uint32_t query_row = blockIdx.y; const uint32_t lane = threadIdx.x & 31; @@ -206,14 +216,36 @@ __global__ void tilemaxsim_kernel(const Scalar *query, uint32_t query_rows, for (uint32_t index = 0; index < warps; ++index) { maximum = fmaxf(maximum, warp_best[index]); } - atomicAdd(scores + candidate, maximum); + maxima[static_cast(candidate) * query_rows + query_row] = maximum; + } +} + +__global__ void tilemaxsim_sum_kernel(const float *maxima, uint32_t query_rows, + size_t count, float *scores) { + const size_t candidate = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (candidate >= count) return; + float score = 0.0f; + for (uint32_t query_row = 0; query_row < query_rows; ++query_row) { + score += maxima[candidate * query_rows + query_row]; } + scores[candidate] = score; } static size_t aligned(size_t value, size_t alignment) { return (value + alignment - 1) / alignment * alignment; } +static bool reserve_aligned(size_t *cursor, size_t bytes, size_t *offset) { + constexpr size_t alignment = 256; + *offset = *cursor; + if (bytes > std::numeric_limits::max() - *cursor) return false; + const size_t end = *cursor + bytes; + if (end > std::numeric_limits::max() - (alignment - 1)) return false; + *cursor = aligned(end, alignment); + return true; +} + extern "C" int vctm_gpu_score( VctmGpu *gpu, const unsigned char *query, size_t query_bytes, uint32_t query_rows, uint32_t dimension, uint8_t dtype, @@ -229,15 +261,29 @@ extern "C" int vctm_gpu_score( return cuda_fail(error, error_capacity, "cudaSetDevice", status); } unsigned char *workspace = gpu->allocation + gpu->tensor_bytes; + if (count > std::numeric_limits::max() / query_rows) { + return fail(error, error_capacity, "TileMaxSim maxima count overflow"); + } + const size_t maxima_count = count * query_rows; + if (count > std::numeric_limits::max() / sizeof(uint64_t) || + count > std::numeric_limits::max() / sizeof(uint32_t) || + count > std::numeric_limits::max() / sizeof(float) || + maxima_count > std::numeric_limits::max() / sizeof(float)) { + return fail(error, error_capacity, "TileMaxSim workspace size overflow"); + } size_t cursor = 0; - const size_t query_offset = cursor; - cursor = aligned(cursor + query_bytes, 256); - const size_t offsets_offset = cursor; - cursor = aligned(cursor + count * sizeof(uint64_t), 256); - const size_t rows_offset = cursor; - cursor = aligned(cursor + count * sizeof(uint32_t), 256); - const size_t scores_offset = cursor; - cursor = aligned(cursor + count * sizeof(float), 256); + size_t query_offset = 0; + size_t offsets_offset = 0; + size_t rows_offset = 0; + size_t maxima_offset = 0; + size_t scores_offset = 0; + if (!reserve_aligned(&cursor, query_bytes, &query_offset) || + !reserve_aligned(&cursor, count * sizeof(uint64_t), &offsets_offset) || + !reserve_aligned(&cursor, count * sizeof(uint32_t), &rows_offset) || + !reserve_aligned(&cursor, maxima_count * sizeof(float), &maxima_offset) || + !reserve_aligned(&cursor, count * sizeof(float), &scores_offset)) { + return fail(error, error_capacity, "TileMaxSim workspace size overflow"); + } if (cursor > gpu->workspace_bytes) { return fail(error, error_capacity, "TileMaxSim request exceeds the configured GPU workspace"); @@ -252,9 +298,6 @@ extern "C" int vctm_gpu_score( status = cudaMemcpyAsync(workspace + rows_offset, document_rows, count * sizeof(uint32_t), cudaMemcpyHostToDevice, gpu->compute_stream); - if (status == cudaSuccess) - status = cudaMemsetAsync(workspace + scores_offset, 0, - count * sizeof(float), gpu->compute_stream); if (status != cudaSuccess) { return cuda_fail(error, error_capacity, "CUDA workspace initialization", status); @@ -267,18 +310,26 @@ extern "C" int vctm_gpu_score( dimension, gpu->allocation, reinterpret_cast(workspace + offsets_offset), reinterpret_cast(workspace + rows_offset), - reinterpret_cast(workspace + scores_offset)); + reinterpret_cast(workspace + maxima_offset)); } else if (dtype == 1) { tilemaxsim_kernel<<compute_stream>>>( reinterpret_cast(workspace + query_offset), query_rows, dimension, gpu->allocation, reinterpret_cast(workspace + offsets_offset), reinterpret_cast(workspace + rows_offset), - reinterpret_cast(workspace + scores_offset)); + reinterpret_cast(workspace + maxima_offset)); } else { return fail(error, error_capacity, "unsupported tensor dtype"); } status = cudaGetLastError(); + if (status == cudaSuccess) { + constexpr unsigned int threads = 256; + const auto blocks = static_cast((count + threads - 1) / threads); + tilemaxsim_sum_kernel<<compute_stream>>>( + reinterpret_cast(workspace + maxima_offset), query_rows, + count, reinterpret_cast(workspace + scores_offset)); + status = cudaGetLastError(); + } if (status == cudaSuccess) status = cudaMemcpyAsync(output, workspace + scores_offset, count * sizeof(float), cudaMemcpyDeviceToHost, diff --git a/services/tilemaxsimd/src/bin/tilemaxsimctl.rs b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs index 0f1c55b0..8c6a9ed7 100644 --- a/services/tilemaxsimd/src/bin/tilemaxsimctl.rs +++ b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs @@ -1,3 +1,13 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + use anyhow::{Context, Result, anyhow, bail}; use clap::{Parser, Subcommand}; use sha2::{Digest, Sha256}; diff --git a/services/tilemaxsimd/src/cache.rs b/services/tilemaxsimd/src/cache.rs index 0ab1d5ac..e53f1b49 100644 --- a/services/tilemaxsimd/src/cache.rs +++ b/services/tilemaxsimd/src/cache.rs @@ -1,3 +1,13 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::hash::{Hash, Hasher}; diff --git a/services/tilemaxsimd/src/engine.rs b/services/tilemaxsimd/src/engine.rs index 5f6d946b..29fe9d74 100644 --- a/services/tilemaxsimd/src/engine.rs +++ b/services/tilemaxsimd/src/engine.rs @@ -1,3 +1,13 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + use crate::cache::{Admission, GpuCache}; use crate::gpu::Gpu; use crate::protocol::{Descriptor, Request}; diff --git a/services/tilemaxsimd/src/gpu.rs b/services/tilemaxsimd/src/gpu.rs index ed00749d..de37e093 100644 --- a/services/tilemaxsimd/src/gpu.rs +++ b/services/tilemaxsimd/src/gpu.rs @@ -1,3 +1,13 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + use anyhow::{Result, anyhow, bail}; use std::ffi::{CStr, c_char, c_int, c_uchar, c_void}; use std::ptr::NonNull; diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs index 4b050776..569dd6cc 100644 --- a/services/tilemaxsimd/src/main.rs +++ b/services/tilemaxsimd/src/main.rs @@ -1,3 +1,13 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + mod cache; mod engine; mod gpu; @@ -1094,12 +1104,19 @@ fn run_scheduler( metrics.failed.fetch_add(1, Ordering::Relaxed); metrics.gpu_failures.fetch_add(1, Ordering::Relaxed); work.gpu_elapsed += quantum_started.elapsed(); - Some(protocol::failure( - version, - request_id, - 3, - &format!("{error:#}"), - )) + let diagnostic = format!("{error:#}"); + let failure = protocol::failure(version, request_id, 3, &diagnostic); + if is_fatal_cuda_diagnostic(&diagnostic) { + // CUDA execution/context failures are not ordinary bad + // requests. Stop advertising readiness and terminate + // the scheduler after best-effort notification so the + // service supervisor can recreate the CUDA context. + metrics.ready.store(false, Ordering::Release); + write_response_nonfatal(&mut work.connection, &failure); + metrics.observe_latency(work.accepted_at.elapsed(), work.gpu_elapsed); + return Err(anyhow!("fatal CUDA failure: {diagnostic}")); + } + Some(failure) } } }; @@ -1166,6 +1183,20 @@ fn tenant_hash(tenant: &str) -> String { .collect() } +fn is_fatal_cuda_diagnostic(diagnostic: &str) -> bool { + [ + "cudaSetDevice", + "cudaMemcpy", + "cudaStream", + "CUDA workspace initialization", + "TileMaxSim CUDA execution", + "GPU upload worker panicked", + "GPU worker panicked", + ] + .iter() + .any(|marker| diagnostic.contains(marker)) +} + fn next_quantum_end(work: &Work, config: &SchedulerConfig) -> usize { quantum_end( &work.request.candidates, @@ -1985,8 +2016,8 @@ fn load_resident_manifests(values: &[(String, PathBuf)]) -> Result */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:47 +-- bootstrap +-- List of shell types + +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE rabitq8; +CREATE TYPE sphere_rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq4; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ + "rhs" halfvec[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:41 +-- vchord::datatype::functions_rabitq4::_vchord_halfvec_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:41 +-- vchord::datatype::functions_rabitq8::_vchord_halfvec_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:72 +-- vchord::datatype::functions_rabitq4::_vchord_rabitq4_dequantize_to_halfvec +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:59 +-- vchord::datatype::functions_rabitq4::_vchord_rabitq4_dequantize_to_vector +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:20 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_in +CREATE FUNCTION "_vchord_rabitq4_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:41 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_cosine +CREATE FUNCTION "_vchord_rabitq4_operator_cosine"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:31 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_ip +CREATE FUNCTION "_vchord_rabitq4_operator_ip"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:21 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_l2 +CREATE FUNCTION "_vchord_rabitq4_operator_l2"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:123 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:140 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_out +CREATE FUNCTION "_vchord_rabitq4_out"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:36 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_recv +CREATE FUNCTION "_vchord_rabitq4_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:21 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_send +CREATE FUNCTION "_vchord_rabitq4_send"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:99 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq4_sphere_cosine_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:75 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq4_sphere_ip_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:51 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq4_sphere_l2_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq4.rs:17 +-- vchord::datatype::typmod_rabitq4::_vchord_rabitq4_typmod_in +CREATE FUNCTION "_vchord_rabitq4_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:72 +-- vchord::datatype::functions_rabitq8::_vchord_rabitq8_dequantize_to_halfvec +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:59 +-- vchord::datatype::functions_rabitq8::_vchord_rabitq8_dequantize_to_vector +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:20 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_in +CREATE FUNCTION "_vchord_rabitq8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:41 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_cosine +CREATE FUNCTION "_vchord_rabitq8_operator_cosine"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:31 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_ip +CREATE FUNCTION "_vchord_rabitq8_operator_ip"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:21 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_l2 +CREATE FUNCTION "_vchord_rabitq8_operator_l2"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:123 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:140 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_out +CREATE FUNCTION "_vchord_rabitq8_out"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:36 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_recv +CREATE FUNCTION "_vchord_rabitq8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:21 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_send +CREATE FUNCTION "_vchord_rabitq8_send"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:99 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq8_sphere_cosine_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:75 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq8_sphere_ip_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:51 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq8_sphere_l2_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq8.rs:17 +-- vchord::datatype::typmod_rabitq8::_vchord_rabitq8_typmod_in +CREATE FUNCTION "_vchord_rabitq8_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ + "rhs" vector[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:23 +-- vchord::datatype::functions_rabitq4::_vchord_vector_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:23 +-- vchord::datatype::functions_rabitq8::_vchord_vector_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:110 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordg_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordg_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordg_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordg_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordg_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordg_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:192 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/vchordrq/scanners/maxsim/search.rs:45 +-- vchord::index::vchordrq::scanners::maxsim::search::_vchordrq_maxsim_search_external +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:100 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:95 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:90 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:140 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:130 +-- vchord::index::opclass::_vchordrq_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:125 +-- vchord::index::opclass::_vchordrq_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:120 +-- vchord::index::opclass::_vchordrq_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:150 +-- vchord::index::opclass::_vchordrq_support_rabitq4_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:115 +-- vchord::index::opclass::_vchordrq_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:110 +-- vchord::index::opclass::_vchordrq_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:105 +-- vchord::index::opclass::_vchordrq_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:145 +-- vchord::index::opclass::_vchordrq_support_rabitq8_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:85 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:135 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/scanners/maxsim/exact.rs:51 +-- vchord::index::vchordrq::scanners::maxsim::exact::_vchordrq_tilemaxsim_rerank_halfvec +/* */ + +/* */ +-- src/index/vchordrq/scanners/maxsim/exact.rs:34 +-- vchord::index::vchordrq::scanners::maxsim::exact::_vchordrq_tilemaxsim_rerank_vector +/* */ + +/* */ +-- src/lib.rs:48 +-- finalize +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external +); + +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, + radius REAL +); + +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_vector'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_halfvec'; + +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq8'; + +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq4'; + +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq8) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq8) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq4) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq4) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; + +-- Phase 3B tensor-source bindings + +CREATE TABLE _vchordrq_maxsim_sources ( + index_oid oid PRIMARY KEY, + heap_oid oid NOT NULL, + model_contract_id text NOT NULL + CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['heap_array', 'external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint, + tensor_rows_attnum smallint, + tensor_dim_attnum smallint, + tensor_dtype_attnum smallint, + tensor_checksum_attnum smallint, + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'heap_array'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NULL + AND tensor_rows_attnum IS NULL + AND tensor_dim_attnum IS NULL + AND tensor_dtype_attnum IS NULL + AND tensor_checksum_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_maxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_maxsim_source( + index_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name DEFAULT NULL, + tensor_rows_column name DEFAULT NULL, + tensor_dim_column name DEFAULT NULL, + tensor_dtype_column name DEFAULT NULL, + tensor_checksum_column name DEFAULT NULL, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + heap_oid oid; + index_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + descriptor_id_is_unique boolean; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be heap_array, external_ref, or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + IF caller_oid IS NULL THEN + RAISE EXCEPTION 'could not resolve caller role'; + END IF; + + SELECT x.indrelid, i.relowner + INTO heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF heap_oid IS NULL THEN + RAISE EXCEPTION 'relation % is not a valid single-key vchordrq MaxSim index', + index_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may register its MaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a MaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO descriptor_id_is_unique; + IF NOT descriptor_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION '% sources must not specify a descriptor relation', normalized_storage; + END IF; + tensor_relation_oid := heap_oid; + END IF; + + IF normalized_storage = 'heap_array' THEN + IF tensor_ref_column IS NOT NULL + OR tensor_rows_column IS NOT NULL + OR tensor_dim_column IS NOT NULL + OR tensor_dtype_column IS NOT NULL + OR tensor_checksum_column IS NOT NULL THEN + RAISE EXCEPTION 'heap_array sources must not specify external tensor columns'; + END IF; + ELSE + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor sources require ref, rows, dim, dtype, and checksum columns'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', + tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', + tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', + tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', + tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'tensor source columns must be distinct'; + END IF; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_maxsim_sources ( + index_oid, heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) + ON CONFLICT (index_oid) DO UPDATE SET + heap_oid = EXCLUDED.heap_oid, + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + index_relation::oid, + heap_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_maxsim_source(index_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + index_owner oid; + removed_count bigint; +BEGIN + IF index_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO index_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = index_relation::oid AND c.relkind = 'i'; + IF index_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may unregister its MaxSim tensor source'; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources WHERE index_oid = $1', + ext_schema + ) USING index_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_maxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_maxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_source_info(index_relation regclass) +RETURNS TABLE( + registered_index regclass, + heap_relation regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + bound_heap_oid oid; + live_heap_oid oid; + index_owner oid; + bound_model_contract_id text; + bound_storage text; + bound_descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + expected_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + + EXECUTE pg_catalog.format( + 'SELECT heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_maxsim_sources + WHERE index_oid = $1', + ext_schema + ) INTO + bound_heap_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + bound_descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING index_relation::oid; + IF bound_heap_oid IS NULL THEN + RAISE EXCEPTION 'MaxSim tensor source is not registered for index %', + index_relation; + END IF; + + SELECT x.indrelid, i.relowner + INTO live_heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF live_heap_oid IS NULL OR live_heap_oid <> bound_heap_oid THEN + RAISE EXCEPTION 'registered MaxSim tensor source is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, live_heap_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered MaxSim tensor source'; + END IF; + + IF bound_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid storage'; + END IF; + + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, 'bigint'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = bound_heap_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap columns'; + END IF; + + IF bound_storage = 'heap_array' THEN + IF bound_descriptor_oid IS NOT NULL + OR descriptor_public_id_attnum IS NOT NULL + OR tensor_ref_attnum IS NOT NULL + OR tensor_rows_attnum IS NOT NULL + OR tensor_dim_attnum IS NOT NULL + OR tensor_dtype_attnum IS NOT NULL + OR tensor_checksum_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap_array binding'; + END IF; + tensor_relation_oid := NULL; + ELSIF bound_storage = 'external_ref' THEN + IF bound_descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_ref binding'; + END IF; + tensor_relation_oid := bound_heap_oid; + ELSE + IF bound_descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_relation binding'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = bound_descriptor_oid + AND c.relkind IN ('r', 'm'); + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor relation is stale or invalid'; + END IF; + IF NOT pg_catalog.has_table_privilege(caller_oid, bound_descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'SELECT privilege on the registered descriptor relation is required'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid = 'bigint'::regtype + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID column is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = bound_descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := bound_descriptor_oid; + END IF; + + expected_columns := CASE WHEN bound_storage = 'heap_array' THEN 0 ELSE 5 END; + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.attnum IS NOT NULL; + IF valid_columns <> expected_columns THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid descriptor columns'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + index_relation, + bound_heap_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + bound_descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_search( + index_relation regclass, + query anyarray, + candidate_limit integer, + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_maxsim_search_external_wrapper'; + +COMMENT ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) +IS 'Restricted Phase 3B external-tensor MaxSim search; returns exact similarity under caller MVCC, SELECT privileges, and PostgreSQL row visibility.'; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_maxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + registry regclass; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RETURN; + END IF; + registry := pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_maxsim_sources', ext_schema) + ); + IF registry IS NULL THEN + RETURN; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE d.objid = s.index_oid + OR ( + d.objid = s.heap_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.model_contract_attnum + OR d.objsubid = s.public_id_attnum + OR ( + s.storage = ''external_ref'' + AND d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.descriptor_public_id_attnum + OR d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); + +-- Exact TileMaxSim over a caller-scoped tensor set. This source registry is +-- deliberately relation-keyed rather than index-keyed: the caller owns ACL, +-- graph, and application filtering decisions. VectorChord accepts the full +-- visible ID set without an artificial count cap, loads its tensor pages, and +-- applies the configured GPU cache policy. + +CREATE TABLE _vchordrq_tilemaxsim_sources ( + source_oid oid PRIMARY KEY, + model_contract_id text NOT NULL CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint NOT NULL CHECK ( + tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_rows_attnum smallint NOT NULL CHECK ( + tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dim_attnum smallint NOT NULL CHECK ( + tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dtype_attnum smallint NOT NULL CHECK ( + tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_checksum_attnum smallint NOT NULL CHECK ( + tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_tilemaxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_tilemaxsim_source( + source_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + public_id_is_unique boolean; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor descriptor columns must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be external_ref or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.oid, c.relowner + INTO source_oid, source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid + AND c.relkind IN ('r', 'm'); + IF source_oid IS NULL THEN + RAISE EXCEPTION 'source_relation % must be a table or materialized view', source_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may register a TileMaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL integer or bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'public ID column % must have a non-partial single-key unique index', + public_id_column; + END IF; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a TileMaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL integer or bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION 'external_ref sources must not specify a descriptor relation'; + END IF; + tensor_relation_oid := source_oid; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'TileMaxSim tensor source columns must be distinct'; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_tilemaxsim_sources ( + source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) + ON CONFLICT (source_oid) DO UPDATE SET + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + source_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_tilemaxsim_source(source_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_owner oid; + removed_count bigint; +BEGIN + IF source_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid AND c.relkind IN ('r', 'm'); + IF ext_schema IS NULL + OR source_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may unregister its TileMaxSim source'; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources WHERE source_oid = $1', + ext_schema + ) USING source_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_source_info(source_relation regclass) +RETURNS TABLE( + registered_source regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + bound_model_contract_id text; + bound_storage text; + descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + EXECUTE pg_catalog.format( + 'SELECT source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_tilemaxsim_sources + WHERE source_oid = $1', + ext_schema + ) INTO + source_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING source_relation::oid; + IF source_oid IS NULL THEN + RAISE EXCEPTION 'TileMaxSim tensor source is not registered for relation %', source_relation; + END IF; + + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_oid AND c.relkind IN ('r', 'm'); + IF source_owner IS NULL THEN + RAISE EXCEPTION 'registered TileMaxSim source relation is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, source_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered TileMaxSim source'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, NULL::oid) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = source_oid + AND a.attnum = expected.attnum + AND (expected.atttypid IS NULL OR a.atttypid = expected.atttypid) + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.atttypid IS NOT NULL + OR a.atttypid IN ('integer'::regtype, 'bigint'::regtype); + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered TileMaxSim source columns are invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim source public ID is no longer unique'; + END IF; + + IF bound_storage = 'external_ref' THEN + IF descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered external_ref TileMaxSim binding is invalid'; + END IF; + tensor_relation_oid := source_oid; + ELSIF bound_storage = 'external_relation' THEN + IF descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered external_relation TileMaxSim binding is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_oid AND c.relkind IN ('r', 'm'); + IF NOT FOUND OR NOT pg_catalog.has_table_privilege(caller_oid, descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor relation is unavailable'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid IN ('integer'::regtype, 'bigint'::regtype) + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + RAISE EXCEPTION 'registered TileMaxSim storage is invalid'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 5 THEN + RAISE EXCEPTION 'registered TileMaxSim tensor descriptor columns are invalid'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + source_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query vector[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_vector_wrapper'; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query halfvec[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_halfvec_wrapper'; + +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) FROM PUBLIC; +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) TO PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_tilemaxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL OR pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_tilemaxsim_sources', ext_schema) + ) IS NULL THEN + RETURN; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE ( + d.objid = s.source_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.model_contract_attnum, + s.public_id_attnum, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_ref_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_rows_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dim_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dtype_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_checksum_attnum END + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.descriptor_public_id_attnum, + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_tilemaxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_tilemaxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_tilemaxsim_source_sql_drop(); +/* */ + diff --git a/sql/upgrade/vchord--1.1.1--1.2.0.sql b/sql/upgrade/vchord--1.1.1--1.2.0.sql new file mode 100644 index 00000000..b0bbb785 --- /dev/null +++ b/sql/upgrade/vchord--1.1.1--1.2.0.sql @@ -0,0 +1,1587 @@ +-- Preserve the post-1.1.1 composite-type compatibility fix when upgrading. + +CREATE OR REPLACE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_vector'; + +CREATE OR REPLACE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_halfvec'; + +CREATE OR REPLACE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq8'; + +CREATE OR REPLACE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq4'; + +-- Phase 3B tensor-source bindings + +CREATE TABLE _vchordrq_maxsim_sources ( + index_oid oid PRIMARY KEY, + heap_oid oid NOT NULL, + model_contract_id text NOT NULL + CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['heap_array', 'external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint, + tensor_rows_attnum smallint, + tensor_dim_attnum smallint, + tensor_dtype_attnum smallint, + tensor_checksum_attnum smallint, + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'heap_array'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NULL + AND tensor_rows_attnum IS NULL + AND tensor_dim_attnum IS NULL + AND tensor_dtype_attnum IS NULL + AND tensor_checksum_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_maxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_maxsim_source( + index_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name DEFAULT NULL, + tensor_rows_column name DEFAULT NULL, + tensor_dim_column name DEFAULT NULL, + tensor_dtype_column name DEFAULT NULL, + tensor_checksum_column name DEFAULT NULL, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + heap_oid oid; + index_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + descriptor_id_is_unique boolean; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be heap_array, external_ref, or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + IF caller_oid IS NULL THEN + RAISE EXCEPTION 'could not resolve caller role'; + END IF; + + SELECT x.indrelid, i.relowner + INTO heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF heap_oid IS NULL THEN + RAISE EXCEPTION 'relation % is not a valid single-key vchordrq MaxSim index', + index_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may register its MaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a MaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO descriptor_id_is_unique; + IF NOT descriptor_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION '% sources must not specify a descriptor relation', normalized_storage; + END IF; + tensor_relation_oid := heap_oid; + END IF; + + IF normalized_storage = 'heap_array' THEN + IF tensor_ref_column IS NOT NULL + OR tensor_rows_column IS NOT NULL + OR tensor_dim_column IS NOT NULL + OR tensor_dtype_column IS NOT NULL + OR tensor_checksum_column IS NOT NULL THEN + RAISE EXCEPTION 'heap_array sources must not specify external tensor columns'; + END IF; + ELSE + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor sources require ref, rows, dim, dtype, and checksum columns'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', + tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', + tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', + tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', + tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'tensor source columns must be distinct'; + END IF; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_maxsim_sources ( + index_oid, heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) + ON CONFLICT (index_oid) DO UPDATE SET + heap_oid = EXCLUDED.heap_oid, + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + index_relation::oid, + heap_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_maxsim_source(index_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + index_owner oid; + removed_count bigint; +BEGIN + IF index_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO index_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = index_relation::oid AND c.relkind = 'i'; + IF index_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may unregister its MaxSim tensor source'; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources WHERE index_oid = $1', + ext_schema + ) USING index_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_maxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_maxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_source_info(index_relation regclass) +RETURNS TABLE( + registered_index regclass, + heap_relation regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + bound_heap_oid oid; + live_heap_oid oid; + index_owner oid; + bound_model_contract_id text; + bound_storage text; + bound_descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + expected_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + + EXECUTE pg_catalog.format( + 'SELECT heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_maxsim_sources + WHERE index_oid = $1', + ext_schema + ) INTO + bound_heap_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + bound_descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING index_relation::oid; + IF bound_heap_oid IS NULL THEN + RAISE EXCEPTION 'MaxSim tensor source is not registered for index %', + index_relation; + END IF; + + SELECT x.indrelid, i.relowner + INTO live_heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF live_heap_oid IS NULL OR live_heap_oid <> bound_heap_oid THEN + RAISE EXCEPTION 'registered MaxSim tensor source is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, live_heap_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered MaxSim tensor source'; + END IF; + + IF bound_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid storage'; + END IF; + + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, 'bigint'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = bound_heap_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap columns'; + END IF; + + IF bound_storage = 'heap_array' THEN + IF bound_descriptor_oid IS NOT NULL + OR descriptor_public_id_attnum IS NOT NULL + OR tensor_ref_attnum IS NOT NULL + OR tensor_rows_attnum IS NOT NULL + OR tensor_dim_attnum IS NOT NULL + OR tensor_dtype_attnum IS NOT NULL + OR tensor_checksum_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap_array binding'; + END IF; + tensor_relation_oid := NULL; + ELSIF bound_storage = 'external_ref' THEN + IF bound_descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_ref binding'; + END IF; + tensor_relation_oid := bound_heap_oid; + ELSE + IF bound_descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_relation binding'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = bound_descriptor_oid + AND c.relkind IN ('r', 'm'); + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor relation is stale or invalid'; + END IF; + IF NOT pg_catalog.has_table_privilege(caller_oid, bound_descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'SELECT privilege on the registered descriptor relation is required'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid = 'bigint'::regtype + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID column is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = bound_descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := bound_descriptor_oid; + END IF; + + expected_columns := CASE WHEN bound_storage = 'heap_array' THEN 0 ELSE 5 END; + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.attnum IS NOT NULL; + IF valid_columns <> expected_columns THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid descriptor columns'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + index_relation, + bound_heap_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + bound_descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_search( + index_relation regclass, + query anyarray, + candidate_limit integer, + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_maxsim_search_external_wrapper'; + +COMMENT ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) +IS 'Restricted Phase 3B external-tensor MaxSim search; returns exact similarity under caller MVCC, SELECT privileges, and PostgreSQL row visibility.'; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_maxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + registry regclass; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RETURN; + END IF; + registry := pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_maxsim_sources', ext_schema) + ); + IF registry IS NULL THEN + RETURN; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE d.objid = s.index_oid + OR ( + d.objid = s.heap_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.model_contract_attnum + OR d.objsubid = s.public_id_attnum + OR ( + s.storage = ''external_ref'' + AND d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.descriptor_public_id_attnum + OR d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); + +-- Exact TileMaxSim over a caller-scoped tensor set. This source registry is +-- deliberately relation-keyed rather than index-keyed: the caller owns ACL, +-- graph, and application filtering decisions. VectorChord accepts the full +-- visible ID set without an artificial count cap, loads its tensor pages, and +-- applies the configured GPU cache policy. + +CREATE TABLE _vchordrq_tilemaxsim_sources ( + source_oid oid PRIMARY KEY, + model_contract_id text NOT NULL CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint NOT NULL CHECK ( + tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_rows_attnum smallint NOT NULL CHECK ( + tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dim_attnum smallint NOT NULL CHECK ( + tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dtype_attnum smallint NOT NULL CHECK ( + tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_checksum_attnum smallint NOT NULL CHECK ( + tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_tilemaxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_tilemaxsim_source( + source_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + public_id_is_unique boolean; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor descriptor columns must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be external_ref or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.oid, c.relowner + INTO source_oid, source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid + AND c.relkind IN ('r', 'm'); + IF source_oid IS NULL THEN + RAISE EXCEPTION 'source_relation % must be a table or materialized view', source_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may register a TileMaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL integer or bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'public ID column % must have a non-partial single-key unique index', + public_id_column; + END IF; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a TileMaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL integer or bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION 'external_ref sources must not specify a descriptor relation'; + END IF; + tensor_relation_oid := source_oid; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'TileMaxSim tensor source columns must be distinct'; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_tilemaxsim_sources ( + source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) + ON CONFLICT (source_oid) DO UPDATE SET + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + source_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_tilemaxsim_source(source_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_owner oid; + removed_count bigint; +BEGIN + IF source_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid AND c.relkind IN ('r', 'm'); + IF ext_schema IS NULL + OR source_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may unregister its TileMaxSim source'; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources WHERE source_oid = $1', + ext_schema + ) USING source_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_source_info(source_relation regclass) +RETURNS TABLE( + registered_source regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + bound_model_contract_id text; + bound_storage text; + descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + EXECUTE pg_catalog.format( + 'SELECT source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_tilemaxsim_sources + WHERE source_oid = $1', + ext_schema + ) INTO + source_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING source_relation::oid; + IF source_oid IS NULL THEN + RAISE EXCEPTION 'TileMaxSim tensor source is not registered for relation %', source_relation; + END IF; + + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_oid AND c.relkind IN ('r', 'm'); + IF source_owner IS NULL THEN + RAISE EXCEPTION 'registered TileMaxSim source relation is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, source_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered TileMaxSim source'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, NULL::oid) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = source_oid + AND a.attnum = expected.attnum + AND (expected.atttypid IS NULL OR a.atttypid = expected.atttypid) + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.atttypid IS NOT NULL + OR a.atttypid IN ('integer'::regtype, 'bigint'::regtype); + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered TileMaxSim source columns are invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim source public ID is no longer unique'; + END IF; + + IF bound_storage = 'external_ref' THEN + IF descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered external_ref TileMaxSim binding is invalid'; + END IF; + tensor_relation_oid := source_oid; + ELSIF bound_storage = 'external_relation' THEN + IF descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered external_relation TileMaxSim binding is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_oid AND c.relkind IN ('r', 'm'); + IF NOT FOUND OR NOT pg_catalog.has_table_privilege(caller_oid, descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor relation is unavailable'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid IN ('integer'::regtype, 'bigint'::regtype) + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + RAISE EXCEPTION 'registered TileMaxSim storage is invalid'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 5 THEN + RAISE EXCEPTION 'registered TileMaxSim tensor descriptor columns are invalid'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + source_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query vector[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_vector_wrapper'; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query halfvec[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_halfvec_wrapper'; + +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) FROM PUBLIC; +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) TO PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_tilemaxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL OR pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_tilemaxsim_sources', ext_schema) + ) IS NULL THEN + RETURN; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE ( + d.objid = s.source_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.model_contract_attnum, + s.public_id_attnum, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_ref_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_rows_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dim_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dtype_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_checksum_attnum END + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.descriptor_public_id_attnum, + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_tilemaxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_tilemaxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_tilemaxsim_source_sql_drop(); diff --git a/src/index/vchordrq/scanners/maxsim/exact.rs b/src/index/vchordrq/scanners/maxsim/exact.rs index d43bf9f7..2d7e750d 100644 --- a/src/index/vchordrq/scanners/maxsim/exact.rs +++ b/src/index/vchordrq/scanners/maxsim/exact.rs @@ -1,10 +1,10 @@ // This software is licensed under a dual license model: // // GNU Affero General Public License v3 (AGPLv3): You may use, modify, and -// distribute this software under the AGPLv3. +// distribute this software under the terms of the AGPLv3. // // Elastic License v2 (ELv2): You may also use, modify, and distribute this -// software under the ELv2, which has specific restrictions. +// software under the Elastic License v2, which has specific restrictions. // // Copyright (c) 2026 Hu Xinjing diff --git a/src/index/vchordrq/scanners/maxsim/external.rs b/src/index/vchordrq/scanners/maxsim/external.rs index 8b726da5..713fdd7f 100644 --- a/src/index/vchordrq/scanners/maxsim/external.rs +++ b/src/index/vchordrq/scanners/maxsim/external.rs @@ -3,10 +3,10 @@ // GNU Affero General Public License v3 (AGPLv3): You may use, modify, and // distribute this software under the terms of the AGPLv3. // -// Elastic License v2 (ELv2): This software is also available under the ELv2, -// which has specific restrictions. +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. // -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 Hu Xinjing use super::candidate::PageCandidate; use super::rerank::RerankError; diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs index a204566e..16d382e8 100644 --- a/src/index/vchordrq/scanners/maxsim/gpu.rs +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -6,11 +6,7 @@ // Elastic License v2 (ELv2): You may also use, modify, and distribute this // software under the Elastic License v2, which has specific restrictions. // -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// -// Copyright (c) 2025-2026 TensorChord Inc. +// Copyright (c) 2026 Hu Xinjing use super::candidate::{HeapKey, PageCandidate}; use super::external::{ diff --git a/src/index/vchordrq/scanners/maxsim/profile.rs b/src/index/vchordrq/scanners/maxsim/profile.rs index d2590a9f..da6bfd7c 100644 --- a/src/index/vchordrq/scanners/maxsim/profile.rs +++ b/src/index/vchordrq/scanners/maxsim/profile.rs @@ -6,10 +6,6 @@ // Elastic License v2 (ELv2): You may also use, modify, and distribute this // software under the Elastic License v2, which has specific restrictions. // -// We welcome any commercial collaboration or support. For inquiries -// regarding the licenses, please contact us at: -// vectorchord-inquiry@tensorchord.ai -// // Copyright (c) 2026 Hu Xinjing use std::cell::RefCell; diff --git a/src/index/vchordrq/scanners/maxsim/search.rs b/src/index/vchordrq/scanners/maxsim/search.rs index 04baba34..da9aca34 100644 --- a/src/index/vchordrq/scanners/maxsim/search.rs +++ b/src/index/vchordrq/scanners/maxsim/search.rs @@ -12,15 +12,14 @@ // // Copyright (c) 2025-2026 TensorChord Inc. -use super::MaxsimBuilder; use super::candidate::{HeapKey, PageCandidate}; use super::external::{ CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorSourceBinding, ExternalTensorStorage, resolve_external_tensor_source, validate_descriptor, }; use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; -use super::profile; use super::rerank::RerankError; +use super::{MaxsimBuilder, profile}; use crate::index::fetcher::{ Fetcher, FilterableTuple, HeapFetcher, Tuple, TupleAttribute, ctid_to_key, }; From e54544d7cdd297669d128c90739e553de7d86cf3 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 19:24:48 +0800 Subject: [PATCH 317/324] fix(tilemaxsimd): support tall query tensors --- services/test_tilemaxsim_rust_daemon.py | 14 +++- .../tilemaxsimd/native/tilemaxsim_cuda.cu | 71 +++++++++++-------- 2 files changed, 54 insertions(+), 31 deletions(-) diff --git a/services/test_tilemaxsim_rust_daemon.py b/services/test_tilemaxsim_rust_daemon.py index 9b9dd71c..09e7864e 100644 --- a/services/test_tilemaxsim_rust_daemon.py +++ b/services/test_tilemaxsim_rust_daemon.py @@ -326,6 +326,18 @@ def test_exact_scores_are_repeatable_across_cuda_contexts(self) -> None: ] self.assertEqual(observed, [observed[0]] * len(observed)) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_query_rows_can_cross_the_cuda_grid_y_limit(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + document = np.asarray([[1.0, 0.0]], dtype=" None: binary = self._release_binary() @@ -463,7 +475,7 @@ def test_sigkill_restart_recovers_stale_socket_and_cache_root(self) -> None: time.sleep(0.01) self.assertIsNone(crashed.poll()) crashed.kill() - crashed.wait(timeout=5) + crashed.communicate(timeout=5) self.assertTrue(socket_path.exists()) status_socket_path = root / "restarted-status.sock" diff --git a/services/tilemaxsimd/native/tilemaxsim_cuda.cu b/services/tilemaxsimd/native/tilemaxsim_cuda.cu index 6c6fc527..6ebf470c 100644 --- a/services/tilemaxsimd/native/tilemaxsim_cuda.cu +++ b/services/tilemaxsimd/native/tilemaxsim_cuda.cu @@ -185,38 +185,43 @@ __global__ void tilemaxsim_kernel(const Scalar *query, uint32_t query_rows, uint32_t dimension, const unsigned char *documents, const uint64_t *document_offsets, - const uint32_t *document_rows, float *maxima) { - const uint32_t candidate = blockIdx.x; - const uint32_t query_row = blockIdx.y; + const uint32_t *document_rows, + size_t task_count, float *maxima) { const uint32_t lane = threadIdx.x & 31; const uint32_t warp = threadIdx.x >> 5; const uint32_t warps = blockDim.x >> 5; - const auto *document = reinterpret_cast( - documents + document_offsets[candidate]); - const Scalar *query_vector = query + static_cast(query_row) * dimension; - float best = -CUDART_INF_F; - for (uint32_t row = warp; row < document_rows[candidate]; row += warps) { - const Scalar *document_vector = - document + static_cast(row) * dimension; - float dot = 0.0f; - for (uint32_t index = lane; index < dimension; index += 32) { - dot = fmaf(scalar_to_float(query_vector[index]), - scalar_to_float(document_vector[index]), dot); - } - for (int delta = 16; delta != 0; delta >>= 1) { - dot += __shfl_down_sync(0xffffffff, dot, delta); - } - if (lane == 0) best = fmaxf(best, dot); - } __shared__ float warp_best[8]; - if (lane == 0) warp_best[warp] = best; - __syncthreads(); - if (threadIdx.x == 0) { - float maximum = -CUDART_INF_F; - for (uint32_t index = 0; index < warps; ++index) { - maximum = fmaxf(maximum, warp_best[index]); + for (size_t task = blockIdx.x; task < task_count; task += gridDim.x) { + const size_t candidate = task / query_rows; + const uint32_t query_row = static_cast(task % query_rows); + const auto *document = reinterpret_cast( + documents + document_offsets[candidate]); + const Scalar *query_vector = + query + static_cast(query_row) * dimension; + float best = -CUDART_INF_F; + for (uint32_t row = warp; row < document_rows[candidate]; row += warps) { + const Scalar *document_vector = + document + static_cast(row) * dimension; + float dot = 0.0f; + for (uint32_t index = lane; index < dimension; index += 32) { + dot = fmaf(scalar_to_float(query_vector[index]), + scalar_to_float(document_vector[index]), dot); + } + for (int delta = 16; delta != 0; delta >>= 1) { + dot += __shfl_down_sync(0xffffffff, dot, delta); + } + if (lane == 0) best = fmaxf(best, dot); + } + if (lane == 0) warp_best[warp] = best; + __syncthreads(); + if (threadIdx.x == 0) { + float maximum = -CUDART_INF_F; + for (uint32_t index = 0; index < warps; ++index) { + maximum = fmaxf(maximum, warp_best[index]); + } + maxima[task] = maximum; } - maxima[static_cast(candidate) * query_rows + query_row] = maximum; + __syncthreads(); } } @@ -302,7 +307,13 @@ extern "C" int vctm_gpu_score( return cuda_fail(error, error_capacity, "CUDA workspace initialization", status); } - dim3 grid(static_cast(count), query_rows); + // A CUDA grid's Y dimension is limited to 65,535 even on modern devices. + // Flatten candidate/query-row work into X and let blocks stride when the + // task count is larger. This covers every protocol-valid query shape without + // constructing an excessive launch grid. + constexpr size_t maximum_kernel_blocks = 65'535; + const size_t kernel_blocks = std::min(maxima_count, maximum_kernel_blocks); + dim3 grid(static_cast(kernel_blocks)); dim3 block(256); if (dtype == 2) { tilemaxsim_kernel<<compute_stream>>>( @@ -310,14 +321,14 @@ extern "C" int vctm_gpu_score( dimension, gpu->allocation, reinterpret_cast(workspace + offsets_offset), reinterpret_cast(workspace + rows_offset), - reinterpret_cast(workspace + maxima_offset)); + maxima_count, reinterpret_cast(workspace + maxima_offset)); } else if (dtype == 1) { tilemaxsim_kernel<<compute_stream>>>( reinterpret_cast(workspace + query_offset), query_rows, dimension, gpu->allocation, reinterpret_cast(workspace + offsets_offset), reinterpret_cast(workspace + rows_offset), - reinterpret_cast(workspace + maxima_offset)); + maxima_count, reinterpret_cast(workspace + maxima_offset)); } else { return fail(error, error_capacity, "unsupported tensor dtype"); } From c3c0d400b37a1f38b3ede40a0a1095e9469a7a2b Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 19:29:34 +0800 Subject: [PATCH 318/324] fix(tilemaxsimd): bound non-preemptible GPU work --- README.md | 4 ++ deploy/systemd/tilemaxsimd.env.example | 2 +- services/tilemaxsimd/src/main.rs | 71 ++++++++++++++++++++++---- 3 files changed, 65 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 9415b69a..c572cd34 100644 --- a/README.md +++ b/README.md @@ -184,6 +184,7 @@ tilemaxsimd \ --scheduler-policy fair-priority \ --max-queued-requests 128 \ --max-tenant-queued-requests 16 \ + --scheduler-quantum-fmas 4000000000 \ --tenant-weight foreground=2 \ --tenant-cache-reservation foreground=4 ``` @@ -202,6 +203,9 @@ The in-flight request budget is also expressed in GiB. A reader must reserve its complete declared frame after the fixed header is validated, and keeps that permit through completion, timeout, or disconnect. This bounds aggregate query and descriptor memory even when many clients submit maximum-size frames at once. +The FMA quantum additionally bounds one non-preemptible CUDA launch by actual +`query rows × document rows × dimension` work. A single candidate above that +operator-configured limit is rejected instead of monopolizing a shared GPU. PostgreSQL sends protocol v3 scheduling metadata only when `vchordrq.maxsim_tenant` is set; otherwise it retains protocol v2 compatibility. diff --git a/deploy/systemd/tilemaxsimd.env.example b/deploy/systemd/tilemaxsimd.env.example index 08834ee2..6f3b6267 100644 --- a/deploy/systemd/tilemaxsimd.env.example +++ b/deploy/systemd/tilemaxsimd.env.example @@ -1,3 +1,3 @@ # This unit is opt-in. Do not enable it on nodes where TileMaxSim is disabled. # Memory values are GiB; startup fails if the requested CUDA arena is unavailable. -TILEMAXSIMD_ARGS="--socket /run/vectorchord/tilemaxsim.sock --status-socket /run/vectorchord/tilemaxsim-status.sock --socket-mode 660 --status-socket-mode 660 --gpu-memory-gb 0=20 --gpu-workspace-gb 2 --host-cache-gb 8 --max-inflight-request-gb 1 --contract-root MODEL_CONTRACT_ID=/var/lib/vectorchord/tensors --scheduler-policy fair-priority --max-connections 256 --max-queued-requests 128 --max-tenant-queued-requests 16" +TILEMAXSIMD_ARGS="--socket /run/vectorchord/tilemaxsim.sock --status-socket /run/vectorchord/tilemaxsim-status.sock --socket-mode 660 --status-socket-mode 660 --gpu-memory-gb 0=20 --gpu-workspace-gb 2 --host-cache-gb 8 --max-inflight-request-gb 1 --contract-root MODEL_CONTRACT_ID=/var/lib/vectorchord/tensors --scheduler-policy fair-priority --scheduler-quantum-fmas 4000000000 --max-connections 256 --max-queued-requests 128 --max-tenant-queued-requests 16" diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs index 569dd6cc..48bc715e 100644 --- a/services/tilemaxsimd/src/main.rs +++ b/services/tilemaxsimd/src/main.rs @@ -125,6 +125,8 @@ struct Args { scheduler_quantum_candidates: usize, #[arg(long, default_value_t = 250_000)] scheduler_quantum_tokens: u64, + #[arg(long, default_value_t = 4_000_000_000)] + scheduler_quantum_fmas: u64, #[arg(long = "tenant-weight", value_parser = parse_tenant_weight)] tenant_weights: Vec<(String, f64)>, } @@ -237,6 +239,7 @@ fn main() -> Result<()> { || args.priority_aging_ms == 0 || args.scheduler_quantum_candidates == 0 || args.scheduler_quantum_tokens == 0 + || args.scheduler_quantum_fmas == 0 { bail!("connection, queue, timeout, and priority-aging limits must be positive"); } @@ -337,6 +340,7 @@ fn main() -> Result<()> { batch_window: Duration::from_millis(args.scheduler_batch_window_ms), quantum_candidates: args.scheduler_quantum_candidates, quantum_tokens: args.scheduler_quantum_tokens, + quantum_fmas: args.scheduler_quantum_fmas, socket_io_timeout: Duration::from_millis(args.socket_io_timeout_ms), tenant_weights, }; @@ -384,6 +388,7 @@ fn main() -> Result<()> { "max_queued_requests": args.max_queued_requests, "max_tenant_queued_requests": args.max_tenant_queued_requests, "max_inflight_request_bytes": args.max_inflight_request_gb, + "scheduler_quantum_fmas": args.scheduler_quantum_fmas, "status_socket": args.status_socket, "cache": ready_cache, }) @@ -433,6 +438,7 @@ fn main() -> Result<()> { maximum: args.max_request_bytes, io_timeout: Duration::from_millis(args.socket_io_timeout_ms), server_timeout: Duration::from_millis(args.request_timeout_ms), + maximum_candidate_fmas: args.scheduler_quantum_fmas, }; match thread::Builder::new() .name("tilemaxsim-reader".to_owned()) @@ -526,6 +532,7 @@ struct SchedulerConfig { batch_window: Duration, quantum_candidates: usize, quantum_tokens: u64, + quantum_fmas: u64, socket_io_timeout: Duration, tenant_weights: std::collections::HashMap, } @@ -797,6 +804,7 @@ struct ReaderConfig { maximum: usize, io_timeout: Duration, server_timeout: Duration, + maximum_candidate_fmas: u64, } fn read_and_enqueue( @@ -843,6 +851,23 @@ fn read_and_enqueue( return; } }; + if request.candidates.iter().any(|candidate| { + candidate_fmas(request.query_rows, request.dimension, candidate.rows) + > config.maximum_candidate_fmas + }) { + metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.invalid_requests.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure( + version, + request_id, + 1, + "one candidate exceeds the configured CUDA kernel work limit", + ), + ); + return; + } let client_timeout = if request.timeout_ms == 0 { config.server_timeout } else { @@ -1203,6 +1228,9 @@ fn next_quantum_end(work: &Work, config: &SchedulerConfig) -> usize { work.next_candidate, config.quantum_candidates, config.quantum_tokens, + config.quantum_fmas, + work.request.query_rows, + work.request.dimension, ) } @@ -1211,28 +1239,47 @@ fn quantum_end( start: usize, maximum_candidates: usize, maximum_tokens: u64, + maximum_fmas: u64, + query_rows: u32, + dimension: u32, ) -> usize { let mut end = start; let mut tokens = 0_u64; + let mut fmas = 0_u64; while end < candidates.len() && end - start < maximum_candidates { let rows = u64::from(candidates[end].rows); - if end > start && tokens.saturating_add(rows) > maximum_tokens { + let candidate_fmas = candidate_fmas(query_rows, dimension, candidates[end].rows); + if end > start + && (tokens.saturating_add(rows) > maximum_tokens + || fmas.saturating_add(candidate_fmas) > maximum_fmas) + { break; } tokens = tokens.saturating_add(rows); + fmas = fmas.saturating_add(candidate_fmas); end += 1; } end } +fn candidate_fmas(query_rows: u32, dimension: u32, document_rows: u32) -> u64 { + u64::from(query_rows) + .saturating_mul(u64::from(document_rows)) + .saturating_mul(u64::from(dimension)) +} + fn estimated_next_work(work: &Work, config: &SchedulerConfig) -> u64 { let end = next_quantum_end(work, config); - let document_rows = work.request.candidates[work.next_candidate..end] + work.request.candidates[work.next_candidate..end] .iter() - .map(|candidate| u64::from(candidate.rows)) - .sum::(); - u64::from(work.request.query_rows) - .saturating_mul(document_rows) + .map(|candidate| { + candidate_fmas( + work.request.query_rows, + work.request.dimension, + candidate.rows, + ) + }) + .fold(0_u64, u64::saturating_add) .max(1) } @@ -2016,8 +2063,8 @@ fn load_resident_manifests(values: &[(String, PathBuf)]) -> Result Date: Wed, 15 Jul 2026 23:14:07 +0800 Subject: [PATCH 319/324] feat: harden TileMaxSim commercial runtime --- .dockerignore | 1 - .github/workflows/check.yml | 9 +- .github/workflows/release.yml | 51 ++- Cargo.lock | 48 ++- README.md | 20 +- services/Dockerfile.postgres | 36 ++ services/Dockerfile.tilemaxsimd | 20 +- services/test_tilemaxsim_rust_daemon.py | 26 +- services/tilemaxsimd/Cargo.lock | 4 +- services/tilemaxsimd/src/bin/tilemaxsimctl.rs | 310 ++++++++++++++++-- services/tilemaxsimd/src/main.rs | 269 +++++++++++---- src/index/vchordrq/scanners/maxsim/gpu.rs | 108 +++++- vchord.control | 2 +- 13 files changed, 770 insertions(+), 134 deletions(-) create mode 100644 services/Dockerfile.postgres diff --git a/.dockerignore b/.dockerignore index 774769fa..4623fd97 100644 --- a/.dockerignore +++ b/.dockerignore @@ -5,7 +5,6 @@ target **/target TODO* docs -sql tests *.md *.zip diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index d4e1f591..d8c794a4 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -41,7 +41,9 @@ jobs: run: cargo fmt --check -- --config-path /dev/null --config imports_granularity=Module - name: Deny - run: cargo deny check + run: | + cargo deny check + cargo deny --manifest-path services/tilemaxsimd/Cargo.toml check - name: License Header run: | @@ -149,6 +151,9 @@ jobs: - name: Build, lint, and test native TileMaxSim run: docker build --target test -f services/Dockerfile.tilemaxsimd . + - name: Validate VectorChord PostgreSQL Dockerfile + run: docker build --check -f services/Dockerfile.postgres . + miri: if: | (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +miri')) || @@ -368,7 +373,7 @@ jobs: if: matrix.version == '14' && matrix.arch == 'x86_64' run: | EXTENSION_DIR="$(pg_config --sharedir)/extension" - cmp sql/install/vchord--1.2.0.sql build/sharedir/extension/vchord--0.0.0.sql + cmp sql/install/vchord--1.2.0.sql build/sharedir/extension/vchord--1.2.0.sql sudo install -m 0644 sql/install/vchord--1.1.1.sql "$EXTENSION_DIR/" sudo install -m 0644 sql/install/vchord--1.2.0.sql "$EXTENSION_DIR/" sudo install -m 0644 sql/upgrade/vchord--1.1.1--1.2.0.sql "$EXTENSION_DIR/" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c469997c..f701f7b1 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -141,7 +141,7 @@ jobs: tilemaxsimd: runs-on: ubuntu-24.04 - needs: ["semver"] + needs: ["semver", "build"] steps: - name: Checkout uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -161,9 +161,58 @@ jobs: SEMVER: ${{ needs.semver.outputs.SEMVER }} run: | IMAGE=$(echo "ghcr.io/${GITHUB_REPOSITORY_OWNER}/vectorchord-tilemaxsimd" | tr '[:upper:]' '[:lower:]') + docker build --target test -f services/Dockerfile.tilemaxsimd . docker build -f services/Dockerfile.tilemaxsimd -t "${IMAGE}:${SEMVER}" . + docker run --rm -v /var/run/docker.sock:/var/run/docker.sock \ + aquasec/trivy:0.72.0@sha256:cffe3f5161a47a6823fbd23d985795b3ed72a4c806da4c4df16266c02accdd6f \ + image --scanners vuln --severity HIGH,CRITICAL --ignore-unfixed \ + --exit-code 1 "${IMAGE}:${SEMVER}" docker push "${IMAGE}:${SEMVER}" + postgres-image: + runs-on: ubuntu-24.04 + needs: ["semver", "build"] + strategy: + matrix: + version: ["16"] + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Log in to the fork container registry + env: + GH_TOKEN: ${{ github.token }} + run: echo "$GH_TOKEN" | docker login ghcr.io -u "${{ github.actor }}" --password-stdin + + - name: Build and publish VectorChord PostgreSQL image + env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + run: | + IMAGE=$(echo "ghcr.io/${GITHUB_REPOSITORY_OWNER}/vectorchord-postgres" | tr '[:upper:]' '[:lower:]') + docker build \ + --build-arg POSTGRES_VERSION=${{ matrix.version }} \ + -f services/Dockerfile.postgres \ + -t "${IMAGE}:pg${{ matrix.version }}-${SEMVER}" \ + -t "${IMAGE}:${{ matrix.version }}-${SEMVER}" . + CID=$(docker run -d \ + -e POSTGRES_HOST_AUTH_METHOD=trust \ + "${IMAGE}:pg${{ matrix.version }}-${SEMVER}" \ + -c shared_preload_libraries=vchord) + trap 'docker rm -f "$CID" >/dev/null 2>&1 || true' EXIT + for attempt in $(seq 1 60); do + if docker exec "$CID" pg_isready -U postgres >/dev/null 2>&1; then break; fi + sleep 1 + done + docker exec "$CID" psql -v ON_ERROR_STOP=1 -U postgres -d postgres \ + -c 'CREATE EXTENSION vchord CASCADE' \ + -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord'" + docker run --rm -v /var/run/docker.sock:/var/run/docker.sock \ + aquasec/trivy:0.72.0@sha256:cffe3f5161a47a6823fbd23d985795b3ed72a4c806da4c4df16266c02accdd6f \ + image --scanners vuln --severity HIGH,CRITICAL --ignore-unfixed \ + --exit-code 1 "${IMAGE}:pg${{ matrix.version }}-${SEMVER}" + docker push "${IMAGE}:pg${{ matrix.version }}-${SEMVER}" + docker push "${IMAGE}:${{ matrix.version }}-${SEMVER}" + pgxn: if: ${{ vars.PUBLISH_PGXN == 'true' }} runs-on: "ubuntu-latest" diff --git a/Cargo.lock b/Cargo.lock index 80722f2f..a2e604b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -98,9 +98,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.102" +version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" [[package]] name = "arrayvec" @@ -395,9 +395,9 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] @@ -416,9 +416,9 @@ checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "darling" -version = "0.20.11" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" dependencies = [ "darling_core", "darling_macro", @@ -426,11 +426,10 @@ dependencies = [ [[package]] name = "darling_core" -version = "0.20.11" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" dependencies = [ - "fnv", "ident_case", "proc-macro2", "quote", @@ -440,9 +439,9 @@ dependencies = [ [[package]] name = "darling_macro" -version = "0.20.11" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ "darling_core", "quote", @@ -582,12 +581,6 @@ dependencies = [ "miniz_oxide", ] -[[package]] -name = "fnv" -version = "1.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" - [[package]] name = "foldhash" version = "0.1.5" @@ -1303,22 +1296,22 @@ dependencies = [ ] [[package]] -name = "proc-macro-error-attr2" -version = "2.0.0" +name = "proc-macro-error-attr3" +version = "3.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" +checksum = "34e4dd828515431dd6c4a030d26f7eaed7dd4778226e9d2bb968d65ca4ec3d4d" dependencies = [ "proc-macro2", "quote", ] [[package]] -name = "proc-macro-error2" -version = "2.0.1" +name = "proc-macro-error3" +version = "3.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" +checksum = "5ee475e440453418ff1335189eddf7101ba502cd818ab7ae04209bc83aa925aa" dependencies = [ - "proc-macro-error-attr2", + "proc-macro-error-attr3", "proc-macro2", "quote", "syn", @@ -1910,13 +1903,12 @@ dependencies = [ [[package]] name = "validator_derive" -version = "0.20.0" +version = "0.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7df16e474ef958526d1205f6dda359fdfab79d9aa6d54bafcb92dcd07673dca" +checksum = "240e4b81c20a1d6d50d1d7265c658dfbd204e8b9ac4d80f3c931f39462196335" dependencies = [ "darling", - "once_cell", - "proc-macro-error2", + "proc-macro-error3", "proc-macro2", "quote", "syn", diff --git a/README.md b/README.md index c572cd34..f1b9bcba 100644 --- a/README.md +++ b/README.md @@ -172,9 +172,15 @@ authorization evidence, and request logs expose only a stable tenant hash. Memory flags use GiB rather than byte counts. A representative launch is: +The published daemon is built with CUDA 12.6 and requires the NVIDIA container +runtime plus a host driver compatible with CUDA 12.6. Its final image keeps the +statically linked CUDA application runtime but intentionally omits the unused +CUDA toolkit; the container runtime injects the host driver libraries/devices. + ```shell tilemaxsimd \ --socket /run/vectorchord/tilemaxsim.sock \ + --listen 0.0.0.0:9191 \ --status-socket /run/vectorchord/tilemaxsim-status.sock \ --gpu-memory-gb 0=20 \ --gpu-workspace-gb 2 \ @@ -189,13 +195,23 @@ tilemaxsimd \ --tenant-cache-reservation foreground=4 ``` -The optional status socket serves HTTP `GET /healthz` and Prometheus +PostgreSQL may use either a local Unix path or a `tcp://HOST:PORT` value for +`vchordrq.maxsim_gpu_endpoint`. The TCP listener lets a long-lived GPU service +remain independent from PostgreSQL pod failover. It intentionally has no +application authentication layer, so it must be exposed only on a private +network and restricted to PostgreSQL clients with a firewall or Kubernetes +NetworkPolicy. Tenant and priority fields remain scheduling hints, not trust +boundaries. + +The optional status socket, and an optional TCP status listener for +network-isolated orchestrator probes, serve HTTP `GET /healthz` and Prometheus `GET /metrics`. Metrics include readiness, scheduler depth, active CUDA work, completed/error/timeout/disconnect outcomes, and global/per-tenant admission rejections without exporting tenant identifiers. `GET /livez` reports process liveness separately from readiness. The packaged -`tilemaxsimctl` probe can wait on the status socket without curl or a TCP port. +`tilemaxsimctl --probe live|ready` probe can wait on the status socket without +curl or a TCP port. An opt-in hardened systemd unit and environment example are provided under `deploy/systemd`; the CUDA container uses the same probe for its health check. diff --git a/services/Dockerfile.postgres b/services/Dockerfile.postgres new file mode 100644 index 00000000..42a0db71 --- /dev/null +++ b/services/Dockerfile.postgres @@ -0,0 +1,36 @@ +ARG POSTGRES_VERSION=16 +ARG POSTGRES_RUNTIME_IMAGE=pgvector/pgvector:pg16@sha256:1d533553fefe4f12e5d80c7b80622ba0c382abb5758856f52983d8789179f0fb + +FROM rust:1.96.0-bookworm@sha256:5e2214abe154fe26e39f64488952e5c991eeed1d6d6da7cc8381ae83927f0cfc AS build + +ARG POSTGRES_VERSION +RUN apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + ca-certificates clang curl make postgresql-common \ + && /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + postgresql-server-dev-${POSTGRES_VERSION} \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /build/vectorchord +COPY . . +ENV CC=clang +RUN make PG_CONFIG=/usr/lib/postgresql/${POSTGRES_VERSION}/bin/pg_config build \ + && make PG_CONFIG=/usr/lib/postgresql/${POSTGRES_VERSION}/bin/pg_config \ + DESTDIR=/opt/vectorchord install + +FROM ${POSTGRES_RUNTIME_IMAGE} + +ARG POSTGRES_VERSION +COPY --from=build /opt/vectorchord/ / + +# The upstream PostgreSQL entrypoint only uses gosu to become the postgres +# account. Replace that call with Debian's maintained setpriv and remove the +# statically linked Go binary, whose embedded standard library otherwise leaves +# scanners unable to distinguish unreachable Go CVEs from exploitable code. +RUN postgres --version | grep -Eq "PostgreSQL\\) ${POSTGRES_VERSION}([.[:space:]]|$)" \ + && sed -ri \ + 's!exec gosu postgres "\$BASH_SOURCE" "\$@"!exec setpriv --reuid=postgres --regid=postgres --init-groups "\$BASH_SOURCE" "\$@"!' \ + /usr/local/bin/docker-entrypoint.sh \ + && ! grep -q 'exec gosu postgres' /usr/local/bin/docker-entrypoint.sh \ + && rm -f /usr/local/bin/gosu diff --git a/services/Dockerfile.tilemaxsimd b/services/Dockerfile.tilemaxsimd index b15b7088..6665366e 100644 --- a/services/Dockerfile.tilemaxsimd +++ b/services/Dockerfile.tilemaxsimd @@ -1,5 +1,5 @@ -ARG CUDA_DEVEL_IMAGE=nvidia/cuda:12.6.3-devel-ubuntu24.04 -ARG CUDA_RUNTIME_IMAGE=nvidia/cuda:12.6.3-runtime-ubuntu24.04 +ARG CUDA_DEVEL_IMAGE=nvidia/cuda:12.6.3-devel-ubuntu24.04@sha256:392c0df7b577ecae17a17f6ba7f2009c217bb4422f8431c053ae9af61a8c148a +ARG RUNTIME_IMAGE=ubuntu:24.04@sha256:4fbb8e6a8395de5a7550b33509421a2bafbc0aab6c06ba2cef9ebffbc7092d90 FROM ${CUDA_DEVEL_IMAGE} AS build @@ -21,11 +21,23 @@ RUN rustup component add clippy rustfmt \ && cargo clippy --all-targets --locked -- -D warnings \ && cargo test --locked -FROM ${CUDA_RUNTIME_IMAGE} +FROM ${RUNTIME_IMAGE} AS runtime-base + +# nvcc links cudart statically for this daemon. The release binary needs only +# glibc/libstdc++ from Ubuntu; libcuda and the devices are injected by the +# NVIDIA container runtime. Keeping CUDA toolkits out of the final image avoids +# shipping an unused 1.4 GiB runtime and materially reduces its attack surface. +RUN apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get upgrade -y \ + && rm -rf /var/lib/apt/lists/* + +FROM runtime-base COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimd /usr/local/bin/tilemaxsimd COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimctl /usr/local/bin/tilemaxsimctl -RUN groupadd --system vectorchord \ +RUN /usr/local/bin/tilemaxsimd --help >/dev/null \ + && /usr/local/bin/tilemaxsimctl --help >/dev/null \ + && groupadd --system vectorchord \ && useradd --system --gid vectorchord --home-dir /nonexistent \ --shell /usr/sbin/nologin vectorchord \ && install -d -o vectorchord -g vectorchord -m 0750 \ diff --git a/services/test_tilemaxsim_rust_daemon.py b/services/test_tilemaxsim_rust_daemon.py index 09e7864e..777fe45f 100644 --- a/services/test_tilemaxsim_rust_daemon.py +++ b/services/test_tilemaxsim_rust_daemon.py @@ -47,6 +47,7 @@ def run_daemon( workspace_gb: str = "0.02", resident: bool = False, scheduled: bool = False, + tcp: bool = False, ) -> tuple[str, list[tuple[int, float]]]: binary = self._release_binary() if not binary.exists(): @@ -108,6 +109,11 @@ def run_daemon( ) socket_path = root / "tilemaxsimd.sock" status_socket_path = root / "tilemaxsimd-status.sock" + tcp_port = None + if tcp: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as reservation: + reservation.bind(("127.0.0.1", 0)) + tcp_port = reservation.getsockname()[1] command = [ os.fspath(binary), "--socket", @@ -115,6 +121,8 @@ def run_daemon( "--status-socket", os.fspath(status_socket_path), ] + if tcp_port is not None: + command.extend(("--listen", f"127.0.0.1:{tcp_port}")) for device in devices: command.extend(("--gpu-memory-gb", f"{device}={gpu_memory_gb}")) command.extend( @@ -205,8 +213,12 @@ def run_daemon( check=False, ) self.assertEqual(probe.returncode, 0, probe.stderr) - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: - connection.connect(os.fspath(socket_path)) + family = socket.AF_INET if tcp_port is not None else socket.AF_UNIX + with socket.socket(family, socket.SOCK_STREAM) as connection: + if tcp_port is None: + connection.connect(os.fspath(socket_path)) + else: + connection.connect(("127.0.0.1", tcp_port)) connection.sendall(frame) header = protocol.receive_exact(connection, protocol.HEADER.size) body_bytes = protocol.HEADER.unpack(header)[4] @@ -311,6 +323,16 @@ def test_external_v2_shard_round_trip_matches_protocol_oracle(self) -> None: ) self.run_daemon([device]) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_tcp_scoring_round_trip_matches_protocol_oracle(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + output, results = self.run_daemon([device], tcp=True, scheduled=True) + self.assertIn('"listen":"127.0.0.1:', output) + self.assertEqual([candidate for candidate, _ in results], [11, 12]) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") def test_exact_scores_are_repeatable_across_cuda_contexts(self) -> None: device = max( diff --git a/services/tilemaxsimd/Cargo.lock b/services/tilemaxsimd/Cargo.lock index d55f046d..c9c5a177 100644 --- a/services/tilemaxsimd/Cargo.lock +++ b/services/tilemaxsimd/Cargo.lock @@ -54,9 +54,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.102" +version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" [[package]] name = "block-buffer" diff --git a/services/tilemaxsimd/src/bin/tilemaxsimctl.rs b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs index 8c6a9ed7..cb8d4140 100644 --- a/services/tilemaxsimd/src/bin/tilemaxsimctl.rs +++ b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs @@ -11,13 +11,15 @@ use anyhow::{Context, Result, anyhow, bail}; use clap::{Parser, Subcommand}; use sha2::{Digest, Sha256}; +use std::collections::HashSet; use std::fs::{self, File, OpenOptions}; -use std::io::{Read, Write}; +use std::io::{BufRead, BufReader, Read, Write}; +use std::os::fd::AsRawFd; use std::os::unix::fs::{OpenOptionsExt, PermissionsExt}; use std::os::unix::net::UnixStream; use std::path::{Path, PathBuf}; use std::thread; -use std::time::{Duration, Instant}; +use std::time::{Duration, Instant, SystemTime}; #[derive(Parser)] #[command(about = "Probe the native TileMaxSim daemon readiness socket")] @@ -30,6 +32,8 @@ struct Args { io_timeout_ms: u64, #[arg(long, default_value_t = 0)] wait_timeout_ms: u64, + #[arg(long, default_value = "ready", value_parser = ["ready", "live"])] + probe: String, } #[derive(Subcommand)] @@ -47,6 +51,16 @@ enum Command { #[arg(long)] expected_sha256: Option, }, + /// Remove old immutable objects not named by the sha256:// refs on stdin. + GcObjects { + #[arg(long)] + root: PathBuf, + #[arg(long, default_value_t = 86_400)] + grace_seconds: u64, + /// Actually remove eligible objects. The default is a dry run. + #[arg(long)] + delete: bool, + }, } fn main() -> Result<()> { @@ -57,9 +71,9 @@ fn main() -> Result<()> { dimension, dtype, expected_sha256, - }) = args.command + }) = args.command.as_ref() { - let expected_bytes = tensor_bytes(rows, dimension, &dtype)?; + let expected_bytes = tensor_bytes(*rows, *dimension, dtype)?; let mut payload = Vec::with_capacity(expected_bytes); std::io::stdin() .take(expected_bytes.saturating_add(1) as u64) @@ -71,16 +85,27 @@ fn main() -> Result<()> { ); } let descriptor = publish_object( - &root, - rows, - dimension, - &dtype, + root, + *rows, + *dimension, + dtype, &payload, expected_sha256.as_deref(), )?; println!("{}", serde_json::to_string(&descriptor)?); return Ok(()); } + if let Some(Command::GcObjects { + root, + grace_seconds, + delete, + }) = args.command.as_ref() + { + let live = read_live_refs(BufReader::new(std::io::stdin().lock()))?; + let outcome = gc_objects(root, &live, *grace_seconds, *delete)?; + println!("{}", serde_json::to_string(&outcome)?); + return Ok(()); + } if args.io_timeout_ms == 0 { bail!("I/O timeout must be positive"); } @@ -91,7 +116,7 @@ fn main() -> Result<()> { .ok_or_else(|| anyhow!("readiness deadline overflow"))?; loop { - let error = match probe(&args.socket, io_timeout) { + let error = match probe(&args.socket, io_timeout, &args.probe) { Ok(()) => return Ok(()), Err(error) => error, }; @@ -111,6 +136,175 @@ struct PublishedDescriptor { tensor_checksum: String, } +#[derive(serde::Serialize)] +struct GcOutcome { + live_refs: usize, + scanned_objects: u64, + eligible_objects: u64, + eligible_bytes: u64, + deleted_objects: u64, + deleted_bytes: u64, + skipped_entries: u64, + dry_run: bool, +} + +struct ObjectStoreLock { + _file: File, +} + +fn prepare_store_root(root: &Path) -> Result<()> { + fs::create_dir_all(root)?; + let metadata = fs::symlink_metadata(root)?; + if metadata.file_type().is_symlink() || !metadata.is_dir() { + bail!("tensor root must be a real directory"); + } + Ok(()) +} + +fn lock_object_store(root: &Path, exclusive: bool) -> Result { + prepare_store_root(root)?; + let path = root.join(".tilemaxsim-objects.lock"); + let file = OpenOptions::new() + .create(true) + .read(true) + .write(true) + .mode(0o640) + .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(&path) + .with_context(|| format!("cannot open object-store lock {}", path.display()))?; + let operation = if exclusive { + libc::LOCK_EX + } else { + libc::LOCK_SH + }; + loop { + if unsafe { libc::flock(file.as_raw_fd(), operation) } == 0 { + break; + } + let error = std::io::Error::last_os_error(); + if error.kind() != std::io::ErrorKind::Interrupted { + return Err(error) + .with_context(|| format!("cannot lock tensor object store {}", root.display())); + } + } + Ok(ObjectStoreLock { _file: file }) +} + +fn read_live_refs(reader: impl BufRead) -> Result> { + let mut live = HashSet::new(); + for (index, line) in reader.lines().enumerate() { + let line = line?; + let value = line.trim(); + if value.is_empty() { + continue; + } + let Some(digest) = value.strip_prefix("sha256://") else { + bail!("live ref on line {} is not a sha256:// URI", index + 1); + }; + if !is_lower_hex_digest(digest) { + bail!("live ref on line {} has an invalid digest", index + 1); + } + live.insert(digest.to_owned()); + } + Ok(live) +} + +fn is_lower_hex_digest(value: &str) -> bool { + value.len() == 64 + && value + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} + +fn gc_objects( + root: &Path, + live: &HashSet, + grace_seconds: u64, + delete: bool, +) -> Result { + let _lock = lock_object_store(root, true)?; + let mut outcome = GcOutcome { + live_refs: live.len(), + scanned_objects: 0, + eligible_objects: 0, + eligible_bytes: 0, + deleted_objects: 0, + deleted_bytes: 0, + skipped_entries: 0, + dry_run: !delete, + }; + let objects = root.join("objects"); + if !objects.try_exists()? { + return Ok(outcome); + } + let objects_metadata = fs::symlink_metadata(&objects)?; + if objects_metadata.file_type().is_symlink() || !objects_metadata.is_dir() { + bail!("tensor objects path must be a real directory"); + } + + let now = SystemTime::now(); + let grace = Duration::from_secs(grace_seconds); + for directory in fs::read_dir(&objects)? { + let directory = directory?; + let directory_metadata = fs::symlink_metadata(directory.path())?; + if directory_metadata.file_type().is_symlink() { + bail!("tensor objects path contains a symlink"); + } + let Some(prefix) = directory.file_name().to_str().map(str::to_owned) else { + outcome.skipped_entries += 1; + continue; + }; + if !directory_metadata.is_dir() + || prefix.len() != 2 + || !prefix + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + { + outcome.skipped_entries += 1; + continue; + } + for object in fs::read_dir(directory.path())? { + let object = object?; + let path = object.path(); + let metadata = fs::symlink_metadata(&path)?; + if metadata.file_type().is_symlink() { + bail!("tensor object path contains a symlink"); + } + let Some(filename) = object.file_name().to_str().map(str::to_owned) else { + outcome.skipped_entries += 1; + continue; + }; + let Some(digest) = filename.strip_suffix(".tensor") else { + outcome.skipped_entries += 1; + continue; + }; + if !metadata.is_file() || !is_lower_hex_digest(digest) || !digest.starts_with(&prefix) { + outcome.skipped_entries += 1; + continue; + } + outcome.scanned_objects += 1; + if live.contains(digest) { + continue; + } + let modified = metadata.modified().with_context(|| { + format!("cannot read tensor object timestamp {}", path.display()) + })?; + if now.duration_since(modified).unwrap_or_default() < grace { + continue; + } + outcome.eligible_objects += 1; + outcome.eligible_bytes = outcome.eligible_bytes.saturating_add(metadata.len()); + if delete { + fs::remove_file(&path) + .with_context(|| format!("cannot remove tensor object {}", path.display()))?; + outcome.deleted_objects += 1; + outcome.deleted_bytes = outcome.deleted_bytes.saturating_add(metadata.len()); + } + } + } + Ok(outcome) +} + fn tensor_bytes(rows: u32, dimension: u32, dtype: &str) -> Result { if rows == 0 || rows > 65_536 || dimension == 0 || dimension > 60_000 { bail!("invalid tensor shape"); @@ -141,11 +335,7 @@ fn publish_object( if expected_sha256.is_some_and(|expected| expected != digest) { bail!("tensor payload checksum disagrees with --expected-sha256"); } - fs::create_dir_all(root)?; - let root_metadata = fs::symlink_metadata(root)?; - if root_metadata.file_type().is_symlink() || !root_metadata.is_dir() { - bail!("tensor root must be a real directory"); - } + let _lock = lock_object_store(root, false)?; let objects = root.join("objects"); fs::create_dir_all(&objects)?; let objects_metadata = fs::symlink_metadata(&objects)?; @@ -227,42 +417,67 @@ fn verify_existing_object(path: &Path, expected_bytes: usize, expected_digest: & Ok(()) } -fn probe(socket: &PathBuf, timeout: Duration) -> Result<()> { +fn probe(socket: &PathBuf, timeout: Duration, kind: &str) -> Result<()> { let mut stream = UnixStream::connect(socket) .with_context(|| format!("cannot connect to status socket {}", socket.display()))?; stream.set_read_timeout(Some(timeout))?; stream.set_write_timeout(Some(timeout))?; - stream.write_all(b"GET /healthz HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")?; + let path = if kind == "live" { "/livez" } else { "/healthz" }; + write!( + stream, + "GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n" + )?; let mut response = Vec::new(); stream.read_to_end(&mut response)?; - if !is_ready_response(&response) { - bail!("TileMaxSim daemon is not ready"); + if !is_probe_response(&response, kind) { + bail!("TileMaxSim daemon failed its {kind} probe"); } Ok(()) } -fn is_ready_response(response: &[u8]) -> bool { +fn is_probe_response(response: &[u8], kind: &str) -> bool { + let expected_body = if kind == "live" { + b"{\"live\":true}".as_slice() + } else { + b"{\"ready\":true}".as_slice() + }; response.starts_with(b"HTTP/1.1 200 ") && response .windows(b"\r\n\r\n".len()) .any(|window| window == b"\r\n\r\n") - && response.ends_with(b"{\"ready\":true}") + && response.ends_with(expected_body) } #[cfg(test)] mod tests { - use super::{is_ready_response, publish_object}; + use super::{gc_objects, is_probe_response, publish_object, read_live_refs}; + use std::collections::HashSet; use std::fs; + use std::io::Cursor; #[test] fn readiness_requires_success_status_and_true_body() { - assert!(is_ready_response( - b"HTTP/1.1 200 OK\r\nContent-Length: 14\r\n\r\n{\"ready\":true}" + assert!(is_probe_response( + b"HTTP/1.1 200 OK\r\nContent-Length: 14\r\n\r\n{\"ready\":true}", + "ready" )); - assert!(!is_ready_response( - b"HTTP/1.1 503 Service Unavailable\r\n\r\n{\"ready\":false}" + assert!(!is_probe_response( + b"HTTP/1.1 503 Service Unavailable\r\n\r\n{\"ready\":false}", + "ready" + )); + assert!(!is_probe_response(b"HTTP/1.1 200 OK\r\n\r\n", "ready")); + } + + #[test] + fn liveness_is_distinct_from_readiness() { + assert!(is_probe_response( + b"HTTP/1.1 200 OK\r\nContent-Length: 13\r\n\r\n{\"live\":true}", + "live", + )); + assert!(!is_probe_response( + b"HTTP/1.1 200 OK\r\nContent-Length: 14\r\n\r\n{\"ready\":true}", + "live", )); - assert!(!is_ready_response(b"HTTP/1.1 200 OK\r\n\r\n")); } #[test] @@ -285,4 +500,47 @@ mod tests { ); fs::remove_dir_all(root).unwrap(); } + + #[test] + fn gc_requires_valid_refs_and_preserves_live_objects() { + let root = std::env::temp_dir().join(format!( + "tilemaxsimctl-gc-{}-{}", + std::process::id(), + std::thread::current().name().unwrap_or("test") + )); + let _ = fs::remove_dir_all(&root); + let first = publish_object(&root, 1, 2, "float16", &[0, 60, 0, 0], None).unwrap(); + let second = publish_object(&root, 1, 2, "float16", &[0, 56, 0, 0], None).unwrap(); + let live = read_live_refs(Cursor::new(format!("{}\n", first.tensor_ref))).unwrap(); + let dry_run = gc_objects(&root, &live, 0, false).unwrap(); + assert_eq!(dry_run.scanned_objects, 2); + assert_eq!(dry_run.eligible_objects, 1); + assert_eq!(dry_run.deleted_objects, 0); + + let deleted = gc_objects(&root, &live, 0, true).unwrap(); + assert_eq!(deleted.deleted_objects, 1); + let first_digest = &first.tensor_checksum[7..]; + let second_digest = &second.tensor_checksum[7..]; + assert!( + root.join("objects") + .join(&first_digest[..2]) + .join(format!("{first_digest}.tensor")) + .exists() + ); + assert!( + !root + .join("objects") + .join(&second_digest[..2]) + .join(format!("{second_digest}.tensor")) + .exists() + ); + assert!(read_live_refs(Cursor::new("not-a-ref\n")).is_err()); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn empty_live_manifest_is_explicitly_supported_for_an_empty_database() { + let live = read_live_refs(Cursor::new("\n")).unwrap(); + assert_eq!(live, HashSet::new()); + } } diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs index 48bc715e..c12e69f9 100644 --- a/services/tilemaxsimd/src/main.rs +++ b/services/tilemaxsimd/src/main.rs @@ -28,6 +28,7 @@ use std::collections::HashMap; use std::fmt::Write as FmtWrite; use std::fs::{self, OpenOptions}; use std::io::{BufRead, Read, Write}; +use std::net::{SocketAddr, TcpListener, TcpStream}; use std::os::fd::AsRawFd; use std::os::unix::fs::{FileTypeExt, PermissionsExt}; use std::os::unix::net::{UnixListener, UnixStream}; @@ -65,7 +66,10 @@ fn install_signal_handlers() -> Result<()> { struct Args { #[arg(long)] socket: PathBuf, - #[arg(long, required = true, value_parser = parse_gpu_memory)] + /// Optional TCP scoring listener for a database running in another pod. + #[arg(long)] + listen: Option, + #[arg(long, required = true, value_delimiter = ',', value_parser = parse_gpu_memory)] gpu_memory_gb: Vec, #[arg(long, default_value = "2", value_parser = parse_gb)] gpu_workspace_gb: usize, @@ -85,6 +89,8 @@ struct Args { socket_mode: u32, #[arg(long)] status_socket: Option, + #[arg(long)] + status_listen: Option, #[arg(long, default_value = "600", value_parser = parse_mode)] status_socket_mode: u32, #[arg(long)] @@ -308,6 +314,17 @@ fn main() -> Result<()> { .with_context(|| format!("cannot bind {}", args.socket.display()))?; fs::set_permissions(&args.socket, fs::Permissions::from_mode(args.socket_mode))?; listener.set_nonblocking(true)?; + if args.listen.is_some() && args.listen == args.status_listen { + bail!("scoring and status TCP listeners must use different addresses"); + } + let tcp_listener = if let Some(address) = args.listen { + let listener = TcpListener::bind(address) + .with_context(|| format!("cannot bind scoring listener {address}"))?; + listener.set_nonblocking(true)?; + Some(listener) + } else { + None + }; if args.status_socket.as_ref() == Some(&args.socket) { bail!("status socket must differ from the TileMaxSim protocol socket"); } @@ -372,12 +389,26 @@ fn main() -> Result<()> { } else { None }; + let status_tcp_server = if let Some(address) = args.status_listen { + let status_listener = TcpListener::bind(address) + .with_context(|| format!("cannot bind status listener {address}"))?; + status_listener.set_nonblocking(true)?; + let status_metrics = Arc::clone(&metrics); + Some( + thread::Builder::new() + .name("tilemaxsim-status-tcp".to_owned()) + .spawn(move || run_tcp_status_server(status_listener, status_metrics))?, + ) + } else { + None + }; metrics.ready.store(true, Ordering::Release); println!( "{}", serde_json::json!({ "event": "tilemaxsim_rust_ready", "socket": args.socket, + "listen": args.listen, "devices": args.gpu_memory_gb.iter().map(|item| serde_json::json!({ "device": item.device, "allocated_bytes": item.bytes, @@ -390,6 +421,7 @@ fn main() -> Result<()> { "max_inflight_request_bytes": args.max_inflight_request_gb, "scheduler_quantum_fmas": args.scheduler_quantum_fmas, "status_socket": args.status_socket, + "status_listen": args.status_listen, "cache": ready_cache, }) ); @@ -409,6 +441,9 @@ fn main() -> Result<()> { if status_server .as_ref() .is_some_and(thread::JoinHandle::is_finished) + || status_tcp_server + .as_ref() + .is_some_and(thread::JoinHandle::is_finished) { status_failed = true; metrics.ready.store(false, Ordering::Release); @@ -419,57 +454,14 @@ fn main() -> Result<()> { reload.store(true, Ordering::Release); } reap_readers(&mut readers); + let mut connections = Vec::with_capacity(2); match listener.accept() { - Ok((connection, _)) => { - if !try_acquire_reader(&metrics.active_connections, args.max_connections) { - // Closing immediately is deliberate: we have not read enough - // bytes to know whether the peer expects a v2 or v3 response. - metrics.rejected_connections.fetch_add(1, Ordering::Relaxed); - drop(connection); - continue; - } - accepted += 1; - let reader_sender = sender.clone(); - let reader_metrics = Arc::clone(&metrics); - let reader_admission = Arc::clone(&pending_admission); - let request_metrics = Arc::clone(&metrics); - let reader_frame_admission = Arc::clone(&frame_admission); - let reader_config = ReaderConfig { - maximum: args.max_request_bytes, - io_timeout: Duration::from_millis(args.socket_io_timeout_ms), - server_timeout: Duration::from_millis(args.request_timeout_ms), - maximum_candidate_fmas: args.scheduler_quantum_fmas, - }; - match thread::Builder::new() - .name("tilemaxsim-reader".to_owned()) - .spawn(move || { - let _permit = ReaderPermit(reader_metrics); - read_and_enqueue( - connection, - &reader_sender, - reader_config, - reader_admission, - request_metrics, - reader_frame_admission, - ); - }) { - Ok(reader) => readers.push(reader), - Err(error) => { - metrics.active_connections.fetch_sub(1, Ordering::Release); - metrics.ready.store(false, Ordering::Release); - SHUTDOWN_REQUESTED.store(true, Ordering::Release); - fatal_error = Some(error.into()); - break; - } - } - if args.once && accepted == 1 { - SHUTDOWN_REQUESTED.store(true, Ordering::Relaxed); - } - } - Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { - thread::sleep(Duration::from_millis(5)); - } - Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {} + Ok((connection, _)) => connections.push(ClientStream::Unix(connection)), + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::Interrupted + ) => {} Err(error) => { metrics.ready.store(false, Ordering::Release); SHUTDOWN_REQUESTED.store(true, Ordering::Release); @@ -477,9 +469,82 @@ fn main() -> Result<()> { break; } } + if let Some(tcp_listener) = &tcp_listener { + match tcp_listener.accept() { + Ok((connection, _)) => { + connection.set_nodelay(true).ok(); + connections.push(ClientStream::Tcp(connection)); + } + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::Interrupted + ) => {} + Err(error) => { + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + fatal_error = Some(error.into()); + break; + } + } + } + if connections.is_empty() { + thread::sleep(Duration::from_millis(5)); + continue; + } + for connection in connections { + if SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + break; + } + if !try_acquire_reader(&metrics.active_connections, args.max_connections) { + // Closing immediately is deliberate: we have not read enough + // bytes to know whether the peer expects a v2 or v3 response. + metrics.rejected_connections.fetch_add(1, Ordering::Relaxed); + drop(connection); + continue; + } + accepted += 1; + let reader_sender = sender.clone(); + let reader_metrics = Arc::clone(&metrics); + let reader_admission = Arc::clone(&pending_admission); + let request_metrics = Arc::clone(&metrics); + let reader_frame_admission = Arc::clone(&frame_admission); + let reader_config = ReaderConfig { + maximum: args.max_request_bytes, + io_timeout: Duration::from_millis(args.socket_io_timeout_ms), + server_timeout: Duration::from_millis(args.request_timeout_ms), + maximum_candidate_fmas: args.scheduler_quantum_fmas, + }; + match thread::Builder::new() + .name("tilemaxsim-reader".to_owned()) + .spawn(move || { + let _permit = ReaderPermit(reader_metrics); + read_and_enqueue( + connection, + &reader_sender, + reader_config, + reader_admission, + request_metrics, + reader_frame_admission, + ); + }) { + Ok(reader) => readers.push(reader), + Err(error) => { + metrics.active_connections.fetch_sub(1, Ordering::Release); + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + fatal_error = Some(error.into()); + break; + } + } + if args.once && accepted == 1 { + SHUTDOWN_REQUESTED.store(true, Ordering::Relaxed); + } + } } SHUTDOWN_REQUESTED.store(true, Ordering::Release); drop(listener); + drop(tcp_listener); for reader in readers { if reader.join().is_err() { eprintln!("tilemaxsim reader thread panicked during shutdown"); @@ -494,6 +559,9 @@ fn main() -> Result<()> { if status_server.is_some_and(|status_server| status_server.join().is_err()) { eprintln!("TileMaxSim status thread panicked during shutdown"); } + if status_tcp_server.is_some_and(|status_server| status_server.join().is_err()) { + eprintln!("TileMaxSim TCP status thread panicked during shutdown"); + } match fs::remove_file(&args.socket) { Ok(()) => {} Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} @@ -513,9 +581,64 @@ fn main() -> Result<()> { Ok(()) } +enum ClientStream { + Unix(UnixStream), + Tcp(TcpStream), +} + +impl ClientStream { + fn set_read_timeout(&self, timeout: Option) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.set_read_timeout(timeout), + Self::Tcp(stream) => stream.set_read_timeout(timeout), + } + } + + fn set_write_timeout(&self, timeout: Option) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.set_write_timeout(timeout), + Self::Tcp(stream) => stream.set_write_timeout(timeout), + } + } +} + +impl Read for ClientStream { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + match self { + Self::Unix(stream) => stream.read(buffer), + Self::Tcp(stream) => stream.read(buffer), + } + } +} + +impl Write for ClientStream { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + match self { + Self::Unix(stream) => stream.write(buffer), + Self::Tcp(stream) => stream.write(buffer), + } + } + + fn flush(&mut self) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.flush(), + Self::Tcp(stream) => stream.flush(), + } + } +} + +impl AsRawFd for ClientStream { + fn as_raw_fd(&self) -> std::os::fd::RawFd { + match self { + Self::Unix(stream) => stream.as_raw_fd(), + Self::Tcp(stream) => stream.as_raw_fd(), + } + } +} + struct Work { request: protocol::Request, - connection: UnixStream, + connection: ClientStream, accepted_at: Instant, deadline: Instant, next_candidate: usize, @@ -808,7 +931,7 @@ struct ReaderConfig { } fn read_and_enqueue( - mut connection: UnixStream, + mut connection: ClientStream, sender: &mpsc::SyncSender, config: ReaderConfig, pending_admission: Arc, @@ -1283,7 +1406,7 @@ fn estimated_next_work(work: &Work, config: &SchedulerConfig) -> u64 { .max(1) } -fn peer_disconnected(connection: &UnixStream) -> bool { +fn peer_disconnected(connection: &ClientStream) -> bool { let mut byte = 0_u8; let result = unsafe { libc::recv( @@ -1306,7 +1429,7 @@ fn peer_disconnected(connection: &UnixStream) -> bool { false } -fn write_response_nonfatal(connection: &mut UnixStream, response: &[u8]) { +fn write_response_nonfatal(connection: &mut ClientStream, response: &[u8]) { if let Err(error) = connection.write_all(response) { eprintln!("TileMaxSim response write failed without stopping daemon: {error}"); } @@ -1315,7 +1438,15 @@ fn write_response_nonfatal(connection: &mut UnixStream, response: &[u8]) { fn run_status_server(listener: UnixListener, path: PathBuf, metrics: Arc) { while !SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { match listener.accept() { - Ok((mut connection, _)) => handle_status_connection(&mut connection, &metrics), + Ok((mut connection, _)) => { + connection + .set_read_timeout(Some(Duration::from_millis(250))) + .ok(); + connection + .set_write_timeout(Some(Duration::from_millis(250))) + .ok(); + handle_status_connection(&mut connection, &metrics); + } Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { thread::sleep(Duration::from_millis(20)); } @@ -1334,13 +1465,35 @@ fn run_status_server(listener: UnixListener, path: PathBuf, metrics: Arc) { + while !SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + match listener.accept() { + Ok((mut connection, _)) => { + configure_tcp_status_connection(&connection); + handle_status_connection(&mut connection, &metrics); + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(20)); + } + Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {} + Err(error) => { + eprintln!("TileMaxSim TCP status listener failed: {error}"); + break; + } + } + } +} + +fn configure_tcp_status_connection(connection: &TcpStream) { connection .set_read_timeout(Some(Duration::from_millis(250))) .ok(); connection .set_write_timeout(Some(Duration::from_millis(250))) .ok(); +} + +fn handle_status_connection(connection: &mut (impl Read + Write), metrics: &RuntimeMetrics) { let mut request = [0_u8; 1024]; let Ok(count) = connection.read(&mut request) else { return; @@ -1919,7 +2072,7 @@ fn render_metrics(metrics: &RuntimeMetrics) -> String { } fn read_request( - connection: &mut UnixStream, + connection: &mut ClientStream, maximum: usize, frame_admission: &Arc, ) -> Result<(Vec, BytePermit)> { diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs index 16d382e8..f530ca50 100644 --- a/src/index/vchordrq/scanners/maxsim/gpu.rs +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -16,7 +16,9 @@ use super::rerank::{CandidateTensorSource, ExactMaxsimBackend, RerankError, Rera use distance::Distance; use std::cmp::Reverse; use std::collections::BinaryHeap; +use std::io::{Read, Write}; use std::mem::size_of_val; +use std::net::{TcpStream, ToSocketAddrs}; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; use vchordrq::types::OwnedVector; @@ -955,7 +957,7 @@ impl TileMaxsimTransport for UnixSocketTransport { let deadline = Instant::now() .checked_add(timeout) .ok_or_else(|| RerankError::Transport("invalid timeout".into()))?; - let mut stream = connect_interruptible(&self.endpoint, deadline)?; + let mut stream = connect_endpoint_interruptible(&self.endpoint, deadline)?; let poll = remaining_until(deadline)?.min(Duration::from_millis(50)); stream .set_read_timeout(Some(poll)) @@ -985,6 +987,102 @@ impl TileMaxsimTransport for UnixSocketTransport { } } +#[cfg(unix)] +enum TransportStream { + Unix(std::os::unix::net::UnixStream), + Tcp(TcpStream), +} + +#[cfg(unix)] +impl TransportStream { + fn set_read_timeout(&self, timeout: Option) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.set_read_timeout(timeout), + Self::Tcp(stream) => stream.set_read_timeout(timeout), + } + } + + fn set_write_timeout(&self, timeout: Option) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.set_write_timeout(timeout), + Self::Tcp(stream) => stream.set_write_timeout(timeout), + } + } +} + +#[cfg(unix)] +impl Read for TransportStream { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + match self { + Self::Unix(stream) => stream.read(buffer), + Self::Tcp(stream) => stream.read(buffer), + } + } +} + +#[cfg(unix)] +impl Write for TransportStream { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + match self { + Self::Unix(stream) => stream.write(buffer), + Self::Tcp(stream) => stream.write(buffer), + } + } + + fn flush(&mut self) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.flush(), + Self::Tcp(stream) => stream.flush(), + } + } +} + +#[cfg(unix)] +fn connect_endpoint_interruptible( + endpoint: &str, + deadline: Instant, +) -> Result { + let Some(authority) = endpoint.strip_prefix("tcp://") else { + return connect_interruptible(endpoint, deadline).map(TransportStream::Unix); + }; + if authority.is_empty() + || authority.contains('/') + || authority.contains('@') + || !authority.contains(':') + { + return Err(RerankError::Transport( + "TCP endpoint must be tcp://HOST:PORT without credentials or a path".into(), + )); + } + let remaining = remaining_until(deadline)?; + let addresses = authority + .to_socket_addrs() + .map_err(|error| RerankError::Transport(error.to_string()))? + .collect::>(); + if addresses.is_empty() { + return Err(RerankError::Transport( + "TCP endpoint resolved to no addresses".into(), + )); + } + let mut last_error = None; + for address in addresses { + pgrx::check_for_interrupts!(); + let attempt = remaining_until(deadline)?.min(remaining); + match TcpStream::connect_timeout(&address, attempt) { + Ok(stream) => { + stream.set_nodelay(true).ok(); + return Ok(TransportStream::Tcp(stream)); + } + Err(error) => last_error = Some(error), + } + } + Err(RerankError::Transport( + last_error + .map(|error| error.to_string()) + .unwrap_or_else(|| "TCP connection failed".to_owned()), + )) +} + #[cfg(unix)] fn connect_interruptible( endpoint: &str, @@ -1173,12 +1271,10 @@ fn last_transport_error() -> RerankError { #[cfg(unix)] fn write_interruptible( - stream: &mut std::os::unix::net::UnixStream, + stream: &mut TransportStream, mut bytes: &[u8], deadline: Instant, ) -> Result<(), RerankError> { - use std::io::Write; - while !bytes.is_empty() { match stream.write(bytes) { Ok(0) => return Err(RerankError::Transport("connection closed".into())), @@ -1202,12 +1298,10 @@ fn write_interruptible( #[cfg(unix)] fn read_interruptible( - stream: &mut std::os::unix::net::UnixStream, + stream: &mut TransportStream, mut bytes: &mut [u8], deadline: Instant, ) -> Result<(), RerankError> { - use std::io::Read; - while !bytes.is_empty() { match stream.read(bytes) { Ok(0) => return Err(RerankError::Transport("connection closed".into())), diff --git a/vchord.control b/vchord.control index 603cc3ff..a08ebfe7 100644 --- a/vchord.control +++ b/vchord.control @@ -1,5 +1,5 @@ comment = 'vchord: Vector database plugin for Postgres, written in Rust, specifically designed for LLM' -default_version = '0.0.0' +default_version = '1.2.0' module_pathname = 'vchord' relocatable = true superuser = true From b28efd9490045b0b439df1f3f779da25a2cc0dc0 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 23:33:22 +0800 Subject: [PATCH 320/324] docs: document TileMaxSim limitations in English and Chinese --- README.md | 72 +++++++++++++++++++ README.zh-CN.md | 184 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+) create mode 100644 README.zh-CN.md diff --git a/README.md b/README.md index f1b9bcba..03112b37 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,8 @@ **An open-source VectorChord fork focused on exact multi-vector retrieval in PostgreSQL.** +English | [简体中文](README.zh-CN.md) +
This fork extends VectorChord's `vchordrq` index with exact late-interaction @@ -137,6 +139,76 @@ The reproducible drivers are and [`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py). +## Limitations + +This fork is suitable for controlled production pilots, but the following +limits must be addressed or explicitly accepted before a general-availability +deployment with strict latency, availability, or multi-tenant SLOs. + +### Retrieval scalability + +- TileMaxSim is an exact scorer, not a tensor-native approximate-nearest- + neighbour index. Mandatory source, ACL, document-type, and other caller- + authorized filters may safely reduce its input, but lexical or graph matches + are not high-recall semantic gates for ordinary queries. +- Keeping all tensors resident on a GPU removes storage transfer but not the + exact scoring arithmetic. On the development corpus, exhaustive scoring of + 34,054 descriptors with an 18 GiB resident tensor arena averaged 1.08 seconds; + a high-recall application scope averaged 45.56 ms. Larger low-latency + deployments need a tensor-native ANN or another measured high-recall + candidate generator before exact TileMaxSim. +- A cache smaller than the active working set remains correct, but can thrash. + The 3 GiB tensor-arena stress run returned the same exact top-K while averaging + 18.47 seconds for exhaustive scans. Capacity planning must therefore cover + the hot working set, or the application must preserve a high-recall scope. + +### Scheduling and multi-user operation + +- Fairness, priority, admission limits, and cache quotas are enforced by each + daemon independently. Multiple GPU replicas do not yet share a global queue, + global entitlement ledger, or work-aware load balancer, so skew between + replicas is possible. +- Preemption is cooperative between CUDA kernel launches. The daemon cannot + interrupt a kernel that is already executing; the configured FMA quantum + bounds that non-preemptible interval. +- GPU failover preserves correctness through durable tensor storage, but a new + replica starts with a cold cache. Availability and warm latency are separate + SLOs and should be tested independently. +- Tenant identifiers and priorities are scheduling hints only. VectorChord does + not implement application tenancy, identity, ACL policy, graph governance, or + privileged-user membership. The authenticated caller must enforce those + policies and pass only an authorized hard scope. + +### Deployment, security, and disaster recovery + +- The TileMaxSim TCP protocol currently has no built-in TLS or client + authentication. Keep it on a private network and restrict callers with a + firewall or Kubernetes NetworkPolicy; use a mutually authenticated proxy if + traffic crosses a trust boundary. +- Container images, health probes, metrics, PostgreSQL integration, and HA- + friendly TCP serving are provided, but manifests alone do not prove a + production installation. Operators must rehearse PostgreSQL switchover, + GPU-pod loss, rolling upgrades, restore, and migration with their actual + storage and Kubernetes environment. +- This repository does not configure continuous PostgreSQL WAL archiving, + point-in-time recovery, or cross-site disaster recovery. Those remain + database-platform responsibilities and require a verified restore test. +- Metrics endpoints are available, but ready-made SLO dashboards and alert + rules are not included. Production deployments need alerts for queueing, + rejection, timeout, cache churn, disk capacity, GPU health, and cold-cache + failover latency. +- Tensor-object garbage collection and very large registries require explicit + operational testing and maintenance planning. Filesystem scans and database + locking can become visible at sufficiently large object counts. + +### Release status + +The SQL surface, external-tensor protocol, image tags, and deployment packaging +remain under active development and may change before a stable release. The +real-GPU and fault-path tests in this repository validate bounded execution and +correctness; they are not a substitute for a multi-hour, many-user soak and +chaos test on the target hardware. + ## Bounded multi-tenant GPU scheduling TileMaxSim remains opt-in. Ordinary single-vector VectorChord use does not diff --git a/README.zh-CN.md b/README.zh-CN.md new file mode 100644 index 00000000..932c0578 --- /dev/null +++ b/README.zh-CN.md @@ -0,0 +1,184 @@ +
+ +# VectorChord TileMaxSim + +**面向 PostgreSQL 精确多向量检索的开源 VectorChord 分支。** + +[English](README.md) | 简体中文 + +
+ +本分支为 VectorChord 的 `vchordrq` 索引增加精确 late-interaction +TileMaxSim 检索，适用于“每个文档保存一个 token 向量张量，并在 +PostgreSQL 内完成多向量搜索”的应用。 + +VectorChord 在这里负责向量/张量存储、精确 TileMaxSim、GPU 缓存和调度。 +实体、事件、Fact、关系图谱、社区、意图路由、用户身份与 ACL 等应用治理能力 +应由 GBrain 等上层系统负责，不进入 VectorChord。 + +## 本分支增加的能力 + +- CPU 精确 TileMaxSim 重排，以及可选的原生 Rust/CUDA 后端。 +- 不人为限制调用方授权范围内的张量候选数量；显存压力由 GPU 页面缓存、 + 分块计算和有界请求调度处理。 +- 外部张量源注册，可将完整精度的 token 张量存放在 PostgreSQL 行外。 +- 外部张量查询路径遵循 PostgreSQL 权限、MVCC、行可见性、取消和超时语义。 +- 多向量查询的规划器统计与成本估算。 +- GPU/内存/磁盘三级存储路径、常驻预热、缓存换入换出和指标接口。 +- 按 daemon 实例执行的公平调度、优先级、准入限制和缓存配额。 +- 正确性、注册表、sidecar 协议、规划成本和真实 GPU 故障路径测试。 + +TileMaxSim 是可选功能。没有显式配置 GPU 显存时不会启动 CUDA daemon， +传统单向量 VectorChord 使用不需要 GPU。启用时，daemon 会在启动阶段预占用户 +指定设备上的 GiB 级显存；任何设备或容量无法获得都会直接失败退出。 + +## 性能实验 + +开发机语料包含 34,054 个张量描述符、34,027 个唯一张量，逻辑 FP16 数据量 +为 16.28 GB。绝对延迟会随存储、CPU 和 GPU 改变，因此同机对照比单个数字 +更有参考价值。 + +| 优化项 | 基线 | 优化后 | 结果 | +| --- | ---: | ---: | ---: | +| 不可变分片与批量读取 | 333.38 ms，顺序文件 | 56.52 ms | 快 5.90 倍 | +| 分片相对批量旧文件 | 87.56 ms | 56.52 ms | 快 1.55 倍 | +| 批量 Host→Device 传输 | 37.06 ms，100 次传输 | 14.19 ms，1 次传输 | 快 2.61 倍 | +| TinyLFU/GDSF 准入 | LRU 命中率 69.98% | 命中率 76.18% | +6.20 个百分点 | +| Rust/CUDA 冷请求 | Python/Triton 855.34 ms | 93.86 ms | 快 9.11 倍 | +| Rust/CUDA 热请求 p50 | Python/Triton 14.26 ms | 2.09 ms | 快 6.82 倍 | +| Rust/CUDA 热请求 p95 | Python/Triton 14.84 ms | 2.20 ms | 快 6.75 倍 | + +20 GiB 显存实验把 18 GiB 用于张量、2 GiB 用于 TileMaxSim 工作区，服务就绪前 +预热全部 34,027 个唯一张量。Python/Triton 从进程启动到就绪耗时 23.07 秒， +Rust/CUDA daemon 为 14.86 秒；预热后的常驻请求不再从磁盘读取张量。 + +GPU 缓存从一块 CUDA arena 中分配连续 page run，并使用 best-fit 尺寸桶和按 +地址合并。默认 32 KiB 页面将语料的分配空间从原 256 KiB buddy allocator 的 +17.840 GB 降至 16.725 GB，内部取整浪费从 8.82% 降至 2.74%。页面大小可以用 +`--gpu-block-kib` 调整。 + +### 全范围检索与小缓存压力 + +18 GiB 张量 arena 中全部张量常驻时，34,054 个描述符的刻意全量精确扫描平均 +为 1.08 秒。3 GiB 张量 arena 加 1 GiB host cache 返回相同的精确 top-K，但 +连续全量扫描产生 68,106 次 miss、61,499 次 eviction 和 32.53 GB Host→Device +传输，平均延迟达到 18.47 秒。 + +这说明小缓存下结果仍然正确，但循环抖动会显著增加延迟；它不是“显存缩小几倍， +GPU 算术就线性变慢”。高召回范围内的热工作集能够放入缓存时不会出现这种全库 +扫描行为。 + +在 40 个真实查询上的传统 `text-embedding-v4`/pgvector HNSW 对照如下： + +| 路径 | 平均延迟 | p95 | 质量 | +| --- | ---: | ---: | ---: | +| 原始文本 embedding + pgvector HNSW | 5.80 ms | 5.86 ms | 文档 hit@1 1.000 | +| 18 GiB 常驻缓存、候选范围内精确 TileMaxSim | 45.56 ms | 134.99 ms | 文档 hit@1 0.475；hit@5 0.650 | +| 3 GiB 缓存、候选范围内精确 TileMaxSim | 573.43 ms | 1,812.70 ms | 排序相同 | + +该实验中的词法范围召回率只有 0.650，因此词法或普通图谱结果不能作为一般语义 +检索的唯一硬门。它们可以加权、解释或预热缓存；关系类意图可以使用完整图谱范围， +而 source、ACL、文档类型等授权约束仍必须硬过滤。 + +复现实验入口： + +- [`services/benchmark_tilemaxsim_ablation.py`](services/benchmark_tilemaxsim_ablation.py) +- [`services/benchmark_full_corpus_tilemaxsim.py`](services/benchmark_full_corpus_tilemaxsim.py) +- [`services/benchmark_gbrain_scoped_tilemaxsim.py`](services/benchmark_gbrain_scoped_tilemaxsim.py) +- [`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py) + +## GPU 缓存与调度 + +默认 `fair-priority` 调度器组合显式优先级和加权公平性。大请求会按候选/token +计算量切分，并在 CUDA kernel 之间重新进入调度器。队列、进行中请求内存、GPU +和 host cache 都可以设置全局及调度域上限。 + +内存配置使用 GiB，不使用字节数。示例： + +```shell +tilemaxsimd \ + --socket /run/vectorchord/tilemaxsim.sock \ + --listen 0.0.0.0:9191 \ + --status-socket /run/vectorchord/tilemaxsim-status.sock \ + --gpu-memory-gb 0=20 \ + --gpu-workspace-gb 2 \ + --host-cache-gb 8 \ + --max-inflight-request-gb 1 \ + --contract-root MODEL_CONTRACT_ID=/srv/vectorchord/tensors \ + --scheduler-policy fair-priority \ + --max-queued-requests 128 \ + --max-tenant-queued-requests 16 \ + --scheduler-quantum-fmas 4000000000 \ + --tenant-weight foreground=2 \ + --tenant-cache-reservation foreground=4 +``` + +发布镜像使用 CUDA 12.6 构建，需要 NVIDIA Container Runtime，以及兼容 +CUDA 12.6 的宿主机驱动。最终镜像保留静态链接的 CUDA 应用运行时，但不包含 +不需要的完整 CUDA toolkit；驱动库和 GPU 设备由 NVIDIA 运行时注入。 + +PostgreSQL 的 `vchordrq.maxsim_gpu_endpoint` 可以是本地 Unix socket，也可以是 +`tcp://HOST:PORT`。`GET /livez`、`GET /healthz` 和 `GET /metrics` 分别用于 +存活、就绪和 Prometheus 指标检查。 + +## 已知限制（Limitations） + +本分支可以用于受控生产试点，但在严格低延迟、高可用或大规模多用户场景正式 +GA 前，必须解决或明确接受以下限制。 + +### 检索规模 + +- TileMaxSim 是精确打分器，不是张量原生 ANN。source、ACL、文档类型等授权硬 + 过滤可以安全缩小范围，但普通词法或图谱命中不能充当一般语义查询的高召回硬门。 +- 张量全部常驻 GPU 只能消除存储传输，不能消除精确计算量。本次实测全量 34,054 + 个描述符平均 1.08 秒，而高召回应用范围平均 45.56 ms。大知识库的严格低延迟 + 部署仍需要张量原生 ANN，或经过召回率验证的候选生成器，再做精确 TileMaxSim。 +- 缓存小于活跃工作集时不会返回错误结果，但可能频繁换入换出。前述 3 GiB 压力 + 实验的全量扫描平均 18.47 秒，生产容量必须覆盖热工作集或保留高召回范围。 + +### 多用户与调度 + +- 公平、priority、准入上限和缓存配额目前由每个 daemon 独立执行。多个 GPU + 副本之间还没有全局队列、全局 entitlement 账本或感知工作量的负载均衡器， + 因此可能出现副本负载倾斜。 +- 抢占发生在 CUDA kernel 之间，无法中断已经开始执行的 kernel；FMA quantum + 负责限制单次不可抢占的最长计算量。 +- GPU 故障转移可依靠持久张量存储保持正确性，但新副本是冷缓存。高可用和热缓存 + 延迟是两个独立 SLO，需要分别验证。 +- tenant 标识和 priority 只是调度提示，不是授权凭据。VectorChord 不负责应用 + 多租户、身份、ACL、图谱治理或特权名单；上层认证系统必须只传入已授权硬范围。 + +### 部署、安全与灾备 + +- TileMaxSim TCP 协议当前没有内置 TLS 和客户端认证，只能放在私网并通过防火墙 + 或 Kubernetes NetworkPolicy 限制 PostgreSQL 调用；跨信任域时应增加双向认证 + 代理。 +- 容器、探针、指标、PostgreSQL 集成和适配 HA 的 TCP 服务已经提供，但 manifest + 不等于生产验证。必须在目标环境演练 PostgreSQL 主从切换、GPU Pod 丢失、滚动 + 升级、备份恢复和数据迁移。 +- 本仓库不负责持续 WAL 归档、PITR 和跨站灾备；这些属于数据库平台职责，而且 + 必须做真实恢复测试。 +- 已提供指标接口，但没有开箱即用的 SLO dashboard 和告警规则。生产环境需要监控 + 排队、拒绝、超时、缓存抖动、磁盘容量、GPU 健康和故障后的冷缓存延迟。 +- 张量对象数量很大时，GC 文件扫描和数据库锁可能变得可见，需要按实际规模压测并 + 规划维护窗口或后续分代/墓碑式 GC。 + +### 发布状态 + +SQL 接口、外部张量协议、镜像 tag 和部署封装仍在持续开发，稳定版之前可能变化。 +仓库中的真实 GPU 与故障路径测试验证了有界执行和正确性，但不能替代目标硬件上的 +多小时、多用户 soak/chaos 压测。 + +## 安全边界 + +PostgreSQL 只有在设置 `vchordrq.maxsim_tenant` 时才发送 v3 调度元数据，否则 +保持 v2 协议兼容。上层系统可以设置 -100 到 100 的 priority，但优先级只影响 +延迟顺序，不能绕过 PostgreSQL 行可见性、ACL、source 或其他硬过滤。 + +本仓库只包含公开实现与公开项目信息，不包含私有规划或上层应用文档。 + +## 项目关系 + +本项目基于 [supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord)。 +上游安装方式、SQL 用法、兼容性与许可证信息请参阅[英文 README](README.md) 中保留的 +原始 VectorChord 文档。 From 3c3673fcb66a823cb67d78b417f9d6b0a6fff00e Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Wed, 15 Jul 2026 23:35:55 +0800 Subject: [PATCH 321/324] fix(ci): preserve public branch portability fixes --- crates/simd/src/lib.rs | 1 - crates/xtask/src/main.rs | 41 ++++++++++++++++++++++++++++++- deny.toml | 4 +++ src/index/vchordg/am/am_build.rs | 18 ++++++++------ src/index/vchordrq/am/am_build.rs | 18 ++++++++------ tests/vchordrq/recall.slt | 8 +++++- 6 files changed, 71 insertions(+), 19 deletions(-) diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs index 601da07d..8f6f87ef 100644 --- a/crates/simd/src/lib.rs +++ b/crates/simd/src/lib.rs @@ -14,7 +14,6 @@ #![allow(unsafe_code)] #![cfg_attr(feature = "nightly_f16", feature(f16))] -#![cfg_attr(target_arch = "s390x", feature(s390x_target_feature))] #![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] #![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc_feature_detection))] #![cfg_attr(target_arch = "powerpc64", feature(powerpc_target_feature))] diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index b2b5c65d..66afc023 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -291,7 +291,7 @@ fn generate( .exports()? .into_iter() .flat_map(|x| std::str::from_utf8(x.name())) - .filter(|x| !["_start", "_IO_stdin_used", "main"].contains(x)) + .filter(|x| should_stub_postmaster_export(x)) .map(str::to_string) .collect::>() } else { @@ -351,6 +351,45 @@ fn generate( Ok(command_stdout) } +fn should_stub_postmaster_export(symbol: &str) -> bool { + // These symbols are supplied by the C runtime or linker when the schema + // helper executable is linked. Defining a PostgreSQL stub for any of them + // creates duplicate symbols on some targets (notably powerpc64le). + !matches!( + symbol, + "_start" + | "_IO_stdin_used" + | "main" + | "__data_start" + | "data_start" + | "__bss_start" + | "_edata" + | "_end" + ) +} + +#[cfg(test)] +mod tests { + use super::should_stub_postmaster_export; + + #[test] + fn schema_helper_does_not_stub_crt_or_linker_symbols() { + for symbol in [ + "_start", + "_IO_stdin_used", + "main", + "__data_start", + "data_start", + "__bss_start", + "_edata", + "_end", + ] { + assert!(!should_stub_postmaster_export(symbol), "{symbol}"); + } + assert!(should_stub_postmaster_export("PostgresMain")); + } +} + fn install_by_copying( src: impl AsRef, dst: impl AsRef, diff --git a/deny.toml b/deny.toml index 0fa41117..ae84eb04 100644 --- a/deny.toml +++ b/deny.toml @@ -5,6 +5,10 @@ all-features = true ignore = [ "RUSTSEC-2021-0127", # serde_cbor is unmaintained "RUSTSEC-2024-0436", # paste - no longer maintained + # Build-time-only transitive dependency of validator_derive 0.20.0. The + # advisory has no maintained compatible release yet; remove this exception + # as soon as validator publishes a replacement. + "RUSTSEC-2026-0173", ] [licenses] diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs index 819d79be..a5ab6361 100644 --- a/src/index/vchordg/am/am_build.rs +++ b/src/index/vchordg/am/am_build.rs @@ -418,15 +418,17 @@ pub unsafe extern "C-unwind" fn vchordg_parallel_build_main( .cast::() .cast_const() }; - let heap_lockmode; - let index_lockmode; - if unsafe { !(*vchordgshared).isconcurrent } { - heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; + let (heap_lockmode, index_lockmode) = if unsafe { !(*vchordgshared).isconcurrent } { + ( + pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) } else { - heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; - } + ( + pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) + }; let heap = unsafe { pgrx::pg_sys::table_open((*vchordgshared).heaprelid, heap_lockmode) }; let index = unsafe { pgrx::pg_sys::index_open((*vchordgshared).indexrelid, index_lockmode) }; let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 95b4bc45..15a5d982 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -826,15 +826,17 @@ pub unsafe extern "C-unwind" fn vchordrq_parallel_build_main( .cast::() .cast_const() }; - let heap_lockmode; - let index_lockmode; - if unsafe { !(*vchordrqshared).isconcurrent } { - heap_lockmode = pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE; + let (heap_lockmode, index_lockmode) = if unsafe { !(*vchordrqshared).isconcurrent } { + ( + pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) } else { - heap_lockmode = pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE; - index_lockmode = pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE; - } + ( + pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) + }; let heap = unsafe { pgrx::pg_sys::table_open((*vchordrqshared).heaprelid, heap_lockmode) }; let index = unsafe { pgrx::pg_sys::index_open((*vchordrqshared).indexrelid, index_lockmode) }; let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; diff --git a/tests/vchordrq/recall.slt b/tests/vchordrq/recall.slt index 475046fa..8246f31a 100644 --- a/tests/vchordrq/recall.slt +++ b/tests/vchordrq/recall.slt @@ -79,6 +79,9 @@ SHOW vchordrq.query_sampling_enable; ---- on +statement ok +SET enable_seqscan = off; + statement ok SELECT * from t ORDER BY val <-> '[0.50, 0.25, 1.00]'; @@ -163,6 +166,9 @@ SELECT COUNT(*) from public.vchordrq_sampled_queries; ---- 3 +statement ok +RESET enable_seqscan; + statement ok RESET search_path; @@ -179,4 +185,4 @@ statement ok SELECT pg_reload_conf(); statement ok -DROP TABLE t, t_dim4, t_expr; \ No newline at end of file +DROP TABLE t, t_dim4, t_expr; From 02ab2367adac3b1bcecc48a48821140b1796e379 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Thu, 16 Jul 2026 15:36:57 +0800 Subject: [PATCH 322/324] fix(tilemaxsimd): deduplicate content-addressed tensors --- services/tilemaxsimd/src/engine.rs | 64 +++++++++++++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/services/tilemaxsimd/src/engine.rs b/services/tilemaxsimd/src/engine.rs index 29fe9d74..ee08515c 100644 --- a/services/tilemaxsimd/src/engine.rs +++ b/services/tilemaxsimd/src/engine.rs @@ -13,7 +13,7 @@ use crate::gpu::Gpu; use crate::protocol::{Descriptor, Request}; use crate::shard::{HostCacheStatus, ShardStore, cache_key}; use anyhow::{Result, anyhow, bail}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; struct MissingTensor { @@ -123,6 +123,12 @@ impl Engine { if batch_size == 0 { bail!("resident prewarm batch size must be positive"); } + // A content-addressed manifest may legitimately reference the same tensor + // from more than one logical candidate. Upload each cache key once. Without + // this normalization, two duplicates in one batch both try to admit the + // same not-yet-ready entry and the second insert replaces the first arena + // allocation. + let descriptors = unique_descriptors(descriptors); for batch in descriptors.chunks(batch_size) { let payloads = self.store.resolve_many(batch, "__resident__")?; let mut uploads = (0..self.devices.len()) @@ -230,8 +236,15 @@ impl Engine { .collect::>(); let mut missing_descriptors = Vec::new(); let mut missing_indices = Vec::new(); + let mut first_candidate_by_key = HashMap::::new(); + let mut duplicate_candidates = Vec::<(usize, usize)>::new(); for (index, descriptor) in request.candidates.iter().enumerate() { let key = cache_key(descriptor); + if let Some(first_index) = first_candidate_by_key.get(&key) { + duplicate_candidates.push((index, *first_index)); + continue; + } + first_candidate_by_key.insert(key.clone(), index); let hit_device = self .devices .iter() @@ -397,6 +410,13 @@ impl Engine { pending.drain(..consumed); } + // Equal content-addressed tensors have equal TileMaxSim scores. Preserve + // every logical candidate id while avoiding duplicate cache acquisitions, + // uploads, and kernel work inside one request. + for (duplicate_index, first_index) in duplicate_candidates { + scores[duplicate_index] = scores[first_index]; + } + request .candidates .iter() @@ -651,6 +671,15 @@ impl Engine { } } +fn unique_descriptors(descriptors: &[Descriptor]) -> Vec { + let mut seen = HashSet::with_capacity(descriptors.len()); + descriptors + .iter() + .filter(|descriptor| seen.insert(cache_key(descriptor))) + .cloned() + .collect() +} + fn validate_entry(descriptor: &Descriptor, entry: &crate::cache::CacheEntry) -> Result<()> { let scalar_bytes = if descriptor.dtype == 1 { 4 } else { 2 }; let expected_bytes = descriptor.rows as usize * descriptor.dimension as usize * scalar_bytes; @@ -664,3 +693,36 @@ fn validate_entry(descriptor: &Descriptor, entry: &crate::cache::CacheEntry) -> } Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn descriptor(candidate_id: u32, digest: &str, rows: u32) -> Descriptor { + Descriptor { + candidate_id, + contract: "colqwen35@test".to_owned(), + digest: digest.to_owned(), + rows, + dimension: 320, + dtype: 2, + } + } + + #[test] + fn resident_prewarm_deduplicates_content_references() { + let descriptors = vec![ + descriptor(1, "a", 100), + descriptor(2, "a", 100), + descriptor(3, "b", 120), + // Shape is part of the cache identity and must not be collapsed. + descriptor(4, "a", 101), + ]; + + let unique = unique_descriptors(&descriptors); + assert_eq!(unique.len(), 3); + assert_eq!(unique[0].candidate_id, 1); + assert_eq!(unique[1].candidate_id, 3); + assert_eq!(unique[2].candidate_id, 4); + } +} From 6f2317b12d434a1d04ace6c4634fea941d66b15c Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Thu, 16 Jul 2026 15:36:57 +0800 Subject: [PATCH 323/324] docs: present the TileMaxSim retrieval project --- README.md | 734 ++++++++++++++++++++---------------------------- README.zh-CN.md | 374 ++++++++++++++++-------- 2 files changed, 563 insertions(+), 545 deletions(-) diff --git a/README.md b/README.md index 03112b37..b25c1944 100644 --- a/README.md +++ b/README.md @@ -2,252 +2,110 @@ # VectorChord TileMaxSim -**An open-source VectorChord fork focused on exact multi-vector retrieval in PostgreSQL.** +**PostgreSQL-native vector and tensor retrieval with exact TileMaxSim and a persistent GPU cache.** English | [简体中文](README.zh-CN.md)
-This fork extends VectorChord's `vchordrq` index with exact late-interaction -TileMaxSim retrieval. It is intended for applications that store one array of -token vectors per document and need PostgreSQL-native multi-vector search. - -## What this fork adds - -- Exact TileMaxSim reranking on CPU, plus an optional CUDA sidecar backend. -- Full caller-scoped tensor sets without an artificial candidate-count cap; - device-memory pressure is handled by the GPU page cache and request batching. -- External tensor-source registration for deployments that keep full-precision - token tensors outside the indexed PostgreSQL value. -- PostgreSQL-aware permission, MVCC, row-visibility, cancellation, and timeout - handling for the external-tensor search path. -- Planner statistics and cost estimation for multi-vector queries. -- Deterministic correctness, registry, sidecar-protocol, and planner-cost tests. - -## Performance ablation - -We measured each cache-path optimization independently on the same development -machine. The corpus contained 34,054 tensor descriptors (34,027 unique tensors, -16.28 GB of logical FP16 tensor data). The request-level tests sampled 100 -candidates containing 47.81 MB of tensor data. Absolute latency depends on the -storage and GPU, so the same-run comparisons are more useful than the raw -numbers. - -| Optimization | Baseline | Optimized | Result | -| --- | ---: | ---: | ---: | -| Immutable shards and batched reads | 333.38 ms, sequential files | 56.52 ms | 5.90x faster | -| Shards versus batched legacy files | 87.56 ms | 56.52 ms | 1.55x faster | -| Batched host-to-device transfer | 37.06 ms, 100 transfers | 14.19 ms, one transfer | 2.61x faster | -| TinyLFU/GDSF admission | 69.98% LRU hit rate | 76.18% hit rate | +6.20 percentage points | -| Rust/CUDA cold request | 855.34 ms, Python/Triton | 93.86 ms | 9.11x faster | -| Rust/CUDA warm request p50 | 14.26 ms, Python/Triton | 2.09 ms | 6.82x faster | -| Rust/CUDA warm request p95 | 14.84 ms, Python/Triton | 2.20 ms | 6.75x faster | - -A full resident-cache run assigned 20 GiB to one GPU: 18 GiB for tensors and -2 GiB for the TileMaxSim workspace. All 34,027 unique tensors were pinned before -the service became ready. Process-to-ready time was 23.07 seconds for the -Python/Triton sidecar and 14.86 seconds for the Rust/CUDA daemon, a 35.6% -reduction. This is a one-time prewarm cost; resident warm requests do not read -the tensors from disk. - -The GPU cache now suballocates exact contiguous page runs from one CUDA arena, -using best-fit size buckets and address-ordered coalescing. On the full corpus, -the 32 KiB default reduced allocated tensor space from 17.840 GB with the former -256 KiB power-of-two buddy allocator to 16.725 GB. It recovered 1.115 GB of GPU -space and reduced internal rounding waste from 8.82% to 2.74%. The default can -be overridden with `--gpu-block-kib`. - -In a deterministic 20,000-event churn trace, the former buddy allocator had 815 -failed cache-allocation attempts, while the segregated page-run allocator had -637; neither recorded an external-fragmentation failure on this workload. Exact -byte extents had 585 failures, all caused by external fragmentation. Page-run -metadata processing added about 1.4 microseconds per event versus the buddy -baseline. These are cache-admission attempts, not failed search requests: the -runtime evicts unpinned entries or streams oversized working sets in chunks. - -Rust/CUDA and Python/Triton produced identical top-10 results. The maximum -absolute score difference was 5.25e-6 and the mean difference was 3.45e-6. -The benchmark driver is -[`services/benchmark_tilemaxsim_ablation.py`](services/benchmark_tilemaxsim_ablation.py); -run it with `--help` for the corpus, cache-root, device, and output arguments. - -### Full-source retrieval and cache-scheduler stress - -We separately scored all 34,054 descriptors for each query to exercise a -working set larger than the cache. This is both a storage/cache stress test and -an upper bound for general semantic retrieval when no safe narrower hard scope -exists. GBrain still applies source, ACL, type, and other mandatory filters, but -lexical or graph recall must not hide semantic paraphrases. The traditional -embedding/HNSW path remains a separate `single_vector` mode and is not a -dependency of tensor retrieval. - -With an 18 GiB tensor arena, all 34,027 unique tensors were resident and native -daemon round trips averaged 1.08 seconds for a deliberately exhaustive scan. A -3 GiB tensor arena plus 1 GiB host cache produced the same exact top-K, but two -sequential scans caused 68,106 misses, 61,499 evictions, 32.53 GB of host-to- -device transfer, and only two cache hits. Native round trips averaged 18.47 -seconds. The slowdown is cyclic cache thrashing and repeated I/O, not TileMaxSim -compute scaling linearly with GPU-memory capacity. Candidate-scoped queries -whose hot working set fits the cache do not exhibit this full-scan behavior. - -The diagnostic driver is -[`services/benchmark_full_corpus_tilemaxsim.py`](services/benchmark_full_corpus_tilemaxsim.py). -A GBrain-scoped comparison against its original text-embedding/pgvector HNSW -path is reported separately; tensor-derived pooled vectors are not a valid -single-vector baseline. - -### Rejected lexical hard-gate ablation and the traditional embedding baseline - -We then used 40 real queries and their original 1,024-dimensional -`text-embedding-v4` vectors. The single-vector baseline searched 1,200 stored -chunk embeddings with PostgreSQL pgvector HNSW (`m=16`, `ef_search=40`). Tensor -mode did not read those vectors: a real lexical result list supplied every -returned chunk without a second truncation, the chunks were mapped to their -source-page tensors, and the native CUDA daemon ran exact TileMaxSim only over -that scope. This ablation tests whether lexical recall is safe as a mandatory -gate; it is not the accepted general semantic query plan. +VectorChord TileMaxSim is an open-source retrieval engine built on PostgreSQL +and derived from the upstream +[VectorChord](https://github.com/supervc-stack/VectorChord) codebase. It preserves +the traditional single-vector path while adding exact late-interaction retrieval +for documents represented as arrays of token vectors. + +This repository is maintained as its own project. Its public README describes +this project's implementation, measured results, limitations, and roadmap; it +does not reproduce the upstream project's README or present upstream product +claims as our own. + +## Project scope + +VectorChord TileMaxSim provides infrastructure primitives: + +- vector, tensor-descriptor, and retrieval-metadata storage in PostgreSQL; +- traditional single-vector retrieval and exact multi-vector TileMaxSim; +- PostgreSQL permission, MVCC, row-visibility, cancellation, and timeout + semantics around retrieval; +- a persistent GPU tensor arena, host-memory cache, and immutable disk shards; +- bounded admission, fair/priority scheduling, cache quotas, health probes, and + metrics for the optional GPU service. + +It does not implement application identity, authentication, ACL policy, entity +registries, facts, events, relationship graphs, communities, or query-intent +routing. Those belong to an authenticated application or knowledge-governance +layer. VectorChord accepts only the already-authorized candidate scope and +opaque scheduling hints supplied by that caller. + +## Core capabilities + +- Exact TileMaxSim reranking on CPU and through the native Rust/CUDA + `tilemaxsimd` backend. +- No artificial cap on a caller-authorized tensor scope. Oversized working sets + are processed in bounded chunks. +- External tensor-source registration, so full token tensors can remain outside + the indexed PostgreSQL value while stable public IDs and descriptors remain + queryable. +- A three-level tensor path: + - L0: persistent GPU page cache; + - L1: bounded host-memory cache; + - L2: immutable tensor shards or content-addressed objects on durable storage. +- Batched shard reads, batched host-to-device transfers, content-addressed + deduplication, TinyLFU/GDSF admission, and a coalescing page-run allocator. +- Optional resident prewarming that completes before the daemon reports ready. +- Per-daemon fair, strict-priority, or fair-priority scheduling, request aging, + tenant weights, deadlines, disconnect cancellation, and bounded CUDA work + quanta. +- PostgreSQL planner statistics and cost estimation for multi-vector queries. +- Prometheus metrics and separate liveness/readiness probes. +- Correctness, registry, protocol, planner, real-GPU, and fault-path tests. + +TileMaxSim is opt-in. Ordinary single-vector operation does not start the CUDA +daemon and does not require a GPU-memory setting. When enabled, `tilemaxsimd` +reserves the configured devices and GiB allocations during startup and exits if +any requested device or allocation cannot be obtained. + +## Architecture + +```text +Authenticated application / retrieval planner + │ authorized IDs + scheduling hints + ▼ + PostgreSQL + VectorChord + │ tensor descriptors + ▼ + tilemaxsimd + ┌─────────┼──────────┐ + │ │ │ + L0 GPU L1 host L2 durable + cache cache shards + │ + └── exact CUDA TileMaxSim +``` -| Path | Mean latency | p95 | Quality | -| --- | ---: | ---: | ---: | -| Original text embedding + pgvector HNSW | 5.80 ms | 5.86 ms | document hit@1 1.000 | -| Candidate-scoped TileMaxSim, 18 GiB resident tensor cache, native round trips | 45.56 ms | 134.99 ms | document hit@1 0.475; hit@5 0.650 | -| Candidate-scoped TileMaxSim, resident Python driver end to end | 112.33 ms | 319.87 ms | same ranking | -| Candidate-scoped TileMaxSim, 3 GiB LRU tensor cache, native round trips | 573.43 ms | 1,812.70 ms | same ranking | -| Candidate-scoped TileMaxSim, small-cache Python driver end to end | 657.46 ms | 2,011.99 ms | same ranking | - -The lexical scope contained 19.3 chunks and 13.2 documents on average. Because -the source chunks span multiple PDF pages, that mapped to 1,459 page tensors on -average (p95 4,000); no whole-corpus tensor scan occurred. Candidate document -recall was 0.650, and TileMaxSim retained the relevant document in its top five -for all covered queries, so scope generation—not exact reranking—set the recall -ceiling in this run. A 35% loss before TileMaxSim is unacceptable, so lexical -and non-relational graph results neither exclude candidates nor reorder general -tensor retrieval. Only explicit relationship intent may use and fuse GBrain's -complete graph scope as a hard range. The HNSW baseline is an intentionally strong upper-bound -workload whose queries are excerpts from their gold chunks; page-tensor and -text-chunk granularity differ, so its quality number is not a claim that the two -rankers are interchangeable. - -The 3 GiB run completed every request with identical ranking, but incurred -24,432 GPU misses and 11.68 GB of host-to-device transfer across the query -sequence. Its 5.85x driver slowdown relative to resident mode is therefore a -cache-locality result that happens to resemble the 6x cache-capacity ratio; GPU -TileMaxSim arithmetic does not become six times slower when memory is smaller. - -The reproducible drivers are -[`services/benchmark_gbrain_scoped_tilemaxsim.py`](services/benchmark_gbrain_scoped_tilemaxsim.py) -and -[`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py). +PostgreSQL remains the source of row visibility and hard filtering. The daemon +does not infer authorization from tenant or priority fields. It receives an +explicit tensor set, resolves the corresponding immutable objects through the +cache hierarchy, and returns exact scores. -## Limitations +## GPU cache and scheduling -This fork is suitable for controlled production pilots, but the following -limits must be addressed or explicitly accepted before a general-availability -deployment with strict latency, availability, or multi-tenant SLOs. - -### Retrieval scalability - -- TileMaxSim is an exact scorer, not a tensor-native approximate-nearest- - neighbour index. Mandatory source, ACL, document-type, and other caller- - authorized filters may safely reduce its input, but lexical or graph matches - are not high-recall semantic gates for ordinary queries. -- Keeping all tensors resident on a GPU removes storage transfer but not the - exact scoring arithmetic. On the development corpus, exhaustive scoring of - 34,054 descriptors with an 18 GiB resident tensor arena averaged 1.08 seconds; - a high-recall application scope averaged 45.56 ms. Larger low-latency - deployments need a tensor-native ANN or another measured high-recall - candidate generator before exact TileMaxSim. -- A cache smaller than the active working set remains correct, but can thrash. - The 3 GiB tensor-arena stress run returned the same exact top-K while averaging - 18.47 seconds for exhaustive scans. Capacity planning must therefore cover - the hot working set, or the application must preserve a high-recall scope. - -### Scheduling and multi-user operation - -- Fairness, priority, admission limits, and cache quotas are enforced by each - daemon independently. Multiple GPU replicas do not yet share a global queue, - global entitlement ledger, or work-aware load balancer, so skew between - replicas is possible. -- Preemption is cooperative between CUDA kernel launches. The daemon cannot - interrupt a kernel that is already executing; the configured FMA quantum - bounds that non-preemptible interval. -- GPU failover preserves correctness through durable tensor storage, but a new - replica starts with a cold cache. Availability and warm latency are separate - SLOs and should be tested independently. -- Tenant identifiers and priorities are scheduling hints only. VectorChord does - not implement application tenancy, identity, ACL policy, graph governance, or - privileged-user membership. The authenticated caller must enforce those - policies and pass only an authorized hard scope. - -### Deployment, security, and disaster recovery - -- The TileMaxSim TCP protocol currently has no built-in TLS or client - authentication. Keep it on a private network and restrict callers with a - firewall or Kubernetes NetworkPolicy; use a mutually authenticated proxy if - traffic crosses a trust boundary. -- Container images, health probes, metrics, PostgreSQL integration, and HA- - friendly TCP serving are provided, but manifests alone do not prove a - production installation. Operators must rehearse PostgreSQL switchover, - GPU-pod loss, rolling upgrades, restore, and migration with their actual - storage and Kubernetes environment. -- This repository does not configure continuous PostgreSQL WAL archiving, - point-in-time recovery, or cross-site disaster recovery. Those remain - database-platform responsibilities and require a verified restore test. -- Metrics endpoints are available, but ready-made SLO dashboards and alert - rules are not included. Production deployments need alerts for queueing, - rejection, timeout, cache churn, disk capacity, GPU health, and cold-cache - failover latency. -- Tensor-object garbage collection and very large registries require explicit - operational testing and maintenance planning. Filesystem scans and database - locking can become visible at sufficiently large object counts. +The default `fair-priority` scheduler combines explicit urgency with weighted +fairness. Large requests are split by candidate count, document tokens, and +actual `query rows × document rows × dimension` work. A request re-enters the +scheduler between CUDA launches, allowing another scheduling domain to run. +An executing CUDA kernel cannot be interrupted, so the FMA quantum bounds the +cooperative, non-preemptible interval. -### Release status +GPU and host caches have per-scheduling-domain maximums. An optional +`--tenant-cache-reservation TENANT=GB` protects already-warmed GPU pages from +cross-domain eviction below the configured allocation. This is currently an +aggregate byte reservation, not a permanent partition for one particular +knowledge base: it does not prewarm data, does not provide an L1 host-cache +minimum, and does not distinguish multiple sources owned by one tenant. These +gaps are listed explicitly in the roadmap. -The SQL surface, external-tensor protocol, image tags, and deployment packaging -remain under active development and may change before a stable release. The -real-GPU and fault-path tests in this repository validate bounded execution and -correctness; they are not a substitute for a multi-hour, many-user soak and -chaos test on the target hardware. - -## Bounded multi-tenant GPU scheduling - -TileMaxSim remains opt-in. Ordinary single-vector VectorChord use does not -start the CUDA daemon and does not require a GPU-memory setting. When an -operator enables TileMaxSim, `tilemaxsimd` reserves exactly the configured -CUDA devices and GiB allocations at startup and fails closed if any allocation -cannot be obtained. - -The default `fair-priority` scheduler combines explicit request urgency with -weighted tenant fairness. Higher numeric priority is more urgent; requests in -the configured priority band are selected by normalized GPU service consumed, -not merely by request count. The default band spans the complete public priority -range, so urgency breaks ties and reorders work inside one tenant without letting -a high-priority noisy tenant starve an under-served tenant. Waiting requests age -up to the maximum priority. Operators that require global strict priority can -select `priority`; a narrower band is an explicit stronger-urgency trade-off. -This is a serving-design adaptation rather than wire compatibility with vLLM: -VectorChord intentionally uses higher-number-means-more-urgent throughout its -SQL, MCP, and IPC contracts. -`fair` and strict `priority` (priority then FCFS) are also available. This takes -the useful serving ideas from vLLM—bounded work budgets, continuous scheduling, -and resumable long work—while adding tenant isolation: a large request is split -at candidate/token quanta and re-enters the scheduler between CUDA launches. -It is cooperative preemption between kernels, not interruption of an executing -CUDA kernel. - -Admission is bounded both globally and per tenant before work enters the -scheduler. Client disconnects and end-to-end deadlines are checked between -quanta. GPU and host cache ownership also have per-tenant caps; optional GPU -reservations and tenant scheduling weights can be configured for differentiated -service. Tenant identifiers are accepted only as scheduling domains, never as -authorization evidence, and request logs expose only a stable tenant hash. - -Memory flags use GiB rather than byte counts. A representative launch is: - -The published daemon is built with CUDA 12.6 and requires the NVIDIA container -runtime plus a host driver compatible with CUDA 12.6. Its final image keeps the -statically linked CUDA application runtime but intentionally omits the unused -CUDA toolkit; the container runtime injects the host driver libraries/devices. +Memory flags use GiB rather than raw byte counts. A representative launch is: ```shell tilemaxsimd \ @@ -267,211 +125,243 @@ tilemaxsimd \ --tenant-cache-reservation foreground=4 ``` -PostgreSQL may use either a local Unix path or a `tcp://HOST:PORT` value for -`vchordrq.maxsim_gpu_endpoint`. The TCP listener lets a long-lived GPU service -remain independent from PostgreSQL pod failover. It intentionally has no -application authentication layer, so it must be exposed only on a private -network and restricted to PostgreSQL clients with a firewall or Kubernetes -NetworkPolicy. Tenant and priority fields remain scheduling hints, not trust -boundaries. - -The optional status socket, and an optional TCP status listener for -network-isolated orchestrator probes, serve HTTP `GET /healthz` and Prometheus -`GET /metrics`. Metrics include readiness, scheduler depth, active CUDA work, -completed/error/timeout/disconnect outcomes, and global/per-tenant admission -rejections without exporting tenant identifiers. - -`GET /livez` reports process liveness separately from readiness. The packaged -`tilemaxsimctl --probe live|ready` probe can wait on the status socket without -curl or a TCP port. -An opt-in hardened systemd unit and environment example are provided under -`deploy/systemd`; the CUDA container uses the same probe for its health check. - -The in-flight request budget is also expressed in GiB. A reader must reserve -its complete declared frame after the fixed header is validated, and keeps that -permit through completion, timeout, or disconnect. This bounds aggregate query -and descriptor memory even when many clients submit maximum-size frames at once. -The FMA quantum additionally bounds one non-preemptible CUDA launch by actual -`query rows × document rows × dimension` work. A single candidate above that -operator-configured limit is rejected instead of monopolizing a shared GPU. - -PostgreSQL sends protocol v3 scheduling metadata only when -`vchordrq.maxsim_tenant` is set; otherwise it retains protocol v2 compatibility. -GBrain derives that tenant value from authenticated runtime context and may set -`vchordrq.maxsim_priority` in the range -100 through 100. Priority changes -latency ordering only and never bypasses PostgreSQL row visibility, ACL, source, -or other mandatory filters. - -The implementation is currently under active development. Its SQL interfaces -and deployment packaging may change before a stable release. This repository -contains only the public implementation and public-facing project information; -private planning and application documentation are intentionally excluded. - -This work is based on -[supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord). The -original VectorChord README is retained below for upstream installation, -licensing, and project information. - ---- - -
- -# VectorChord - -**Ready for the Billion-Scale Era. Host 100M vectors on a single i4i.xlarge ($247/mo) and [scale seamlessly to 1B+](https://blog.vectorchord.ai/scaling-vector-search-to-1-billion-on-postgresql).** - -[Official Site][official-site-link] · [Blog][blog-link] · [Docs][docs-link] · [Feedback][github-issues-link] · [Contact Us][email-link] - -[![][github-release-shield]][github-release-link] -[![][docker-release-shield]][docker-release-link] -[![][docker-pulls-shield]][docker-pulls-link] -[![][ghcr-release-shield]][ghcr-release-link] -[![][github-downloads-shield]][github-downloads-link] -[![][discord-shield]][discord-link] -[![][X-shield]][X-link] -[![][deepwiki-shield]][deepwiki-link] -[![][license-1-shield]][license-1-link] -[![][license-2-shield]][license-2-link] - -
- -VectorChord (vchord) is a PostgreSQL extension engineered for scalable, high-performance, and cost-effective vector search. +PostgreSQL can connect through a local Unix socket or a `tcp://HOST:PORT` +endpoint. `GET /livez`, `GET /healthz`, and `GET /metrics` expose process +liveness, readiness, and bounded operational metrics. See +[`docs/TILEMAXSIM_CUDA_SIDECAR.md`](docs/TILEMAXSIM_CUDA_SIDECAR.md) and +[`docs/TILEMAXSIM_IPC_V2.md`](docs/TILEMAXSIM_IPC_V2.md) for deployment and +protocol details. -To efficiently store vectors while preserving search quality, VectorChord applies RaBitQ[^1] compression together with autonomous reranking. With VectorChord, you can store 400,000 vectors for just $1, enabling significant savings: 6x more vectors compared to Pinecone's optimized storage and 26x more than pgvector/pgvecto.rs for the same price. +## Measured performance -[^1]: Gao, Jianyang, and Cheng Long. "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search." Proceedings of the ACM on Management of Data 2.3 (2024): 1-27. - -![][image-compare] - -## Features - -VectorChord introduces remarkable enhancements over pgvecto.rs and pgvector: - -**💰 Affordable Vector Search**: Host 100M × 768-dimensional vectors → AWS i4i.xlarge ($247/month)[^2], host 1B × 96-dimensional vectors → i7ie.6xlarge ($2246/month)[^3], helping you keep infrastructure costs down while maintaining competitive search quality. - -[^2]: Please check out our [blog post](https://blog.vectorchord.ai/vectorchord-store-400k-vectors-for-1-in-postgresql) for more details. -[^3]: Please check out our [blog post](https://blog.vectorchord.ai/scaling-vector-search-to-1-billion-on-postgresql) for more details. - -**⚡ Accelerated Index Build**: Index 100 million vectors in just 20 minutes. Powered by hierarchical K-means and highly optimized disk operations, VectorChord eliminates the bottleneck of vector indexing on a single machine with limited hardware resources. - -[^4]: Please check out our [blog post](https://blog.vectorchord.ai/how-we-made-100m-vector-indexing-in-20-minutes-possible-on-postgresql#heading-hierarchical-k-means) for more technique details and [document](https://docs.vectorchord.ai/vectorchord/usage/partitioning-tuning.html#hierarchical-k-means) for usages. - -**📈 Smoothly Scale Up**: Scale with confidence as your data grows. Through dimensionality reduction and sampling[^5], VectorChord effectively controls memory growth, enabling 1B-vector indexes to be built on machines with 128GB of memory in practice. - -[^5]: Please check out our [blog post](https://blog.vectorchord.ai/how-we-made-100m-vector-indexing-in-20-minutes-possible-on-postgresql#heading-dimensionality-reduction) for more technique details and [document](https://docs.vectorchord.ai/vectorchord/usage/partitioning-tuning.html#reduce-sampling-factor) for usages. +The development corpus contains 34,054 descriptors, 34,027 unique tensors, and +16.28 GB of logical FP16 tensor data. Absolute latency depends on storage, CPU, +and GPU hardware; same-machine comparisons are more meaningful than isolated +numbers. -**🔌 Seamless Integration**: Fully compatible with pgvector data types and syntax while providing optimal defaults out of the box - no complex parameter tuning needed. Just drop in VectorChord for enhanced experience. +| Optimization | Baseline | Optimized | Result | +| --- | ---: | ---: | ---: | +| Immutable shards and batched reads | 333.38 ms, sequential files | 56.52 ms | 5.90x faster | +| Shards versus batched legacy files | 87.56 ms | 56.52 ms | 1.55x faster | +| Batched host-to-device transfer | 37.06 ms, 100 transfers | 14.19 ms, one transfer | 2.61x faster | +| TinyLFU/GDSF admission | 69.98% LRU hit rate | 76.18% hit rate | +6.20 percentage points | +| Rust/CUDA cold request | 855.34 ms, Python/Triton | 93.86 ms | 9.11x faster | +| Rust/CUDA warm request p50 | 14.26 ms, Python/Triton | 2.09 ms | 6.82x faster | +| Rust/CUDA warm request p95 | 14.84 ms, Python/Triton | 2.20 ms | 6.75x faster | -**💾 Efficient Storage with Low-Bit Data type**: Drastically reduce storage costs with our [native 4-bit (RaBitQ4) and 8-bit (RaBitQ8) vector types](https://docs.vectorchord.ai/vectorchord/usage/quantization-types.html). Achieve massive space savings without compromising search quality—RaBitQ8 maintains high precision with <1% recall loss. +A 20 GiB run assigned 18 GiB to tensors and 2 GiB to workspace. Prewarming all +34,027 unique tensors took 14.86 seconds in the Rust/CUDA daemon versus 23.07 +seconds in the Python/Triton implementation. The default 32 KiB page-run +allocator reduced allocated space from 17.840 GB with the former 256 KiB buddy +allocator to 16.725 GB, reducing internal rounding waste from 8.82% to 2.74%. -## Quick Start +### Full-range and small-cache stress -For new users, we recommend using the Docker image to get started quickly. If you do not prefer Docker, please read [installation guide](https://docs.vectorchord.ai/vectorchord/getting-started/installation.html) for other installation methods. +With all tensors resident, an intentionally exhaustive exact scan of all 34,054 +descriptors averaged 1.08 seconds. A 3 GiB tensor arena plus 1 GiB host cache +returned the same exact top-K, but sequential full scans caused 68,106 misses, +61,499 evictions, and 32.53 GB of host-to-device transfer; native round trips +averaged 18.47 seconds. -```bash -docker run \ - --name vectorchord-demo \ - -e POSTGRES_PASSWORD=mysecretpassword \ - -p 5432:5432 \ - -d ghcr.io/tensorchord/vchord-postgres:pg18-v1.1.1 -``` +The slowdown is cyclic cache thrashing and repeated I/O, not TileMaxSim +arithmetic becoming linearly slower as memory shrinks. Candidate-scoped queries +whose hot working set fits the cache do not exhibit this full-scan behavior. -> [!NOTE] -> In addition to the base image with the VectorChord extension, we provide an all-in-one image, `tensorchord/vchord-suite:pg17-latest`. This comprehensive image includes all official TensorChord extensions, including `VectorChord`, `VectorChord-bm25` and `pg_tokenizer.rs` . Developers should select an image tag that is compatible with their extension's version, as indicated in [the support matrix](https://github.com/tensorchord/VectorChord-images?tab=readme-ov-file#support-matrix). +### Candidate-scope ablation -Then you can connect to the database using the `psql` command line tool. The default username is `postgres`, and the default password is `mysecretpassword`. +Forty real queries compared the original 1,024-dimensional text embedding with +pgvector HNSW against exact TileMaxSim over a lexical candidate scope: -```bash -psql -h localhost -p 5432 -U postgres -``` +| Path | Mean latency | p95 | Quality | +| --- | ---: | ---: | ---: | +| Original text embedding + pgvector HNSW | 5.80 ms | 5.86 ms | document hit@1 1.000 | +| Candidate-scoped TileMaxSim, 18 GiB resident cache | 45.56 ms | 134.99 ms | hit@1 0.475; hit@5 0.650 | +| Candidate-scoped TileMaxSim, 3 GiB cache | 573.43 ms | 1,812.70 ms | identical TileMaxSim ranking | -Now you can play with VectorChord! +The lexical scope covered the relevant document for only 65% of queries. +TileMaxSim kept the relevant document in its top five for every covered query, +so candidate generation—not exact scoring—set the recall ceiling. Lexical and +ordinary graph matches must therefore not be the only hard gate for general +semantic retrieval. Authorization filters remain hard; relational graph scopes +may be used for explicitly relational intent when the application can construct +a complete range. -VectorChord depends on pgvector, including the vector representation. Since you can use them directly, your application can be easily migrated without pain! +Reproducible drivers: -```sql -CREATE EXTENSION IF NOT EXISTS vchord CASCADE; -``` +- [`services/benchmark_tilemaxsim_ablation.py`](services/benchmark_tilemaxsim_ablation.py) +- [`services/benchmark_full_corpus_tilemaxsim.py`](services/benchmark_full_corpus_tilemaxsim.py) +- [`services/benchmark_gbrain_scoped_tilemaxsim.py`](services/benchmark_gbrain_scoped_tilemaxsim.py) +- [`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py) -Similar to pgvector, you can create a table with vector column and insert some rows to it. +## Limitations -```sql -CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3)); -INSERT INTO items (embedding) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); -``` +This project is suitable for controlled production pilots. Strict low-latency, +high-availability, or large multi-user deployments must address or explicitly +accept the following limitations before general availability. + +### Retrieval scale + +- Exact TileMaxSim is a scorer, not a tensor-native approximate-nearest- + neighbour index. Full residency removes transfer time but not exact scoring + arithmetic. +- Large low-latency corpora still need a measured high-recall candidate + generator or a future tensor-native ANN before exact TileMaxSim. +- A cache smaller than the active working set remains correct but can thrash; + production capacity must cover the hot set or preserve a safe high-recall + range. + +### Cache and multi-user isolation + +- Fairness, priority, admission, and reservations are enforced independently by + each daemon. Replicas do not share a global queue, entitlement ledger, or + residency-aware router. +- A GPU reservation protects aggregate bytes for a scheduling tenant, not a + named `source_id` or knowledge base. There is no dynamic cache-domain API, + source-level prewarm contract, or host-cache minimum. +- Reservations are configured at daemon startup. On multi-GPU daemons the + current value is applied to each device, rather than being expressed as one + explicit cluster-wide total. +- Globally resident manifests are pinned under a service-owned domain. If + pinned pages consume the entire arena, unrelated cold tensors cannot evict + them and the request fails instead of waiting for pinned data to leave. +- Preemption occurs only between CUDA launches. A running kernel is not + interruptible. +- Tenant and priority values are scheduling hints, not authorization evidence. + +### Operations and security + +- The TileMaxSim TCP protocol has no built-in TLS or client authentication. It + must remain on a private network or behind a mutually authenticated proxy. +- A replacement GPU replica starts cold. Durable storage preserves correctness, + but availability and warm-cache latency are separate SLOs. +- PostgreSQL WAL archiving, PITR, cross-site disaster recovery, and verified + restore procedures remain database-platform responsibilities. +- Metrics exist, but ready-made SLO dashboards and alert rules are not yet + included. +- Very large tensor registries and garbage collection require scale-specific + operational testing. -With VectorChord, you can create `vchordrq` indexes. +### Release status -```SQL -CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops); +The SQL surface, external-tensor protocol, image tags, and deployment packaging +remain under active development. Real-GPU and fault-path tests do not replace +multi-hour, many-user soak and chaos testing on target hardware. + +## Outlook: an AI infrastructure retrieval plane + +VectorChord TileMaxSim is intended to become the retrieval plane in a broader, +composable AI infrastructure rather than absorb application or model-serving +responsibilities: + +```text +Agent / application layer + │ + ▼ +AI gateway: identity, quota, priority, deadline, trace + │ + ├── Knowledge layer ── PostgreSQL / VectorChord ── tilemaxsimd + │ + └── Ray Serve ── vLLM and other model engines + +Kubernetes / KubeRay: placement, GPU nodes, scaling, recovery +Object storage: model weights, tensor shards, ingestion artifacts +Prometheus + OpenTelemetry: end-to-end observability ``` -And then perform a vector search using `SELECT ... ORDER BY ... LIMIT ...`. +In this design: + +- [Ray Serve LLM](https://docs.ray.io/en/latest/serve/llm/) coordinates + distributed model replicas, autoscaling, and model-aware routing. +- [vLLM](https://docs.vllm.ai/) owns generation batching, model parallelism, + KV-cache management, and prefix-cache reuse. +- VectorChord owns PostgreSQL retrieval semantics and vector/tensor metadata; + `tilemaxsimd` owns exact TileMaxSim, persistent tensor residency, and bounded + retrieval scheduling. +- The authenticated application owns identity, ACL, graph/fact/event + governance, query intent, and the final RAG pipeline. + +The components should not silently compete for the same GPU memory. Initial +deployments should place vLLM and TileMaxSim in separate GPU pools. Kubernetes +or another resource manager must account for the memory reserved by the +long-lived TileMaxSim daemon before scheduling model workers; explicit MIG or +other hardware partitioning can be evaluated later. + +The next infrastructure milestones are: + +1. Separate `scheduler_tenant` from an opaque `cache_domain`, typically derived + by the authorized caller from tenant and source identity. +2. Add GPU and host `min_resident_gb`, `max_gb`, and prewarm contracts per cache + domain, with unambiguous per-device and fleet-wide semantics. +3. Build a bounded global admission layer and a residency-aware router across + TileMaxSim replicas, while leaving exact local quantum scheduling in each + daemon. +4. Propagate one request envelope—request ID, tenant, source, priority, + deadline, trace ID, and cache domain—through retrieval and model serving. +5. Validate cold-start recovery, rolling upgrades, PostgreSQL failover, GPU-pod + loss, multi-tenant isolation, and long-running soak/chaos behavior. +6. Add a tensor-native ANN or another recall-measured first stage for corpora + where exact full-range TileMaxSim cannot meet the latency SLO. + +Ray is an orchestration layer, not durable storage. PostgreSQL and immutable +object storage remain authoritative. `tilemaxsimd` should remain a supervised, +stateful GPU service rather than a short-lived generic Ray task, so its reserved +arena and cache lifetime remain predictable. + +## Security boundary + +PostgreSQL sends scheduled protocol metadata only when +`vchordrq.maxsim_tenant` is set. Priority affects latency ordering only; it +cannot bypass row visibility, ACL, source, or other mandatory filters. Request +logs expose a stable hash rather than raw tenant identifiers. + +This public repository contains implementation and public-facing project +information only. Private application and planning documents are intentionally +excluded. + +## Build and test + +Useful entry points include: -```SQL -SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5; +```shell +make build +cargo test --locked --workspace --exclude vchord --no-fail-fast +cargo test --manifest-path services/tilemaxsimd/Cargo.toml --locked ``` -For more usage, please read: - -- [Indexing](https://docs.vectorchord.ai/vectorchord/usage/indexing.html) -- [Multi-Vector Retrieval](https://docs.vectorchord.ai/vectorchord/usage/indexing-with-maxsim-operators.html) -- [Quantization Types](https://docs.vectorchord.ai/vectorchord/usage/quantization-types.html) -- [Graph Index](https://docs.vectorchord.ai/vectorchord/usage/graph-index.html) -- [Quantization Types](https://docs.vectorchord.ai/vectorchord/usage/quantization-types.html) -- [Similarity Filter](https://docs.vectorchord.ai/vectorchord/usage/range-query.html) -- [PostgreSQL Tuning](https://docs.vectorchord.ai/vectorchord/usage/performance-tuning.html) -- [Monitoring](https://docs.vectorchord.ai/vectorchord/usage/monitoring.html) -- [Fallback Parameters](https://docs.vectorchord.ai/vectorchord/usage/fallback-parameters.html) -- [Measure Recall](https://docs.vectorchord.ai/vectorchord/usage/measure-recall.html) -- [Prewarm](https://docs.vectorchord.ai/vectorchord/usage/prewarm.html) -- [Prefilter](https://docs.vectorchord.ai/vectorchord/usage/prefilter.html) -- [Prefetch](https://docs.vectorchord.ai/vectorchord/usage/prefetch.html) -- [Rerank in Table](https://docs.vectorchord.ai/vectorchord/usage/rerank-in-table.html) -- [Partitioning Tuning](https://docs.vectorchord.ai/vectorchord/usage/partitioning-tuning.html) -- [External Build](https://docs.vectorchord.ai/vectorchord/usage/external-index-precomputation.html) - -## License - -This software is licensed under a dual license model: - -1. **GNU Affero General Public License v3 (AGPLv3)**: You may use, modify, and distribute this software under the terms of the AGPLv3. - -2. **Elastic License v2 (ELv2)**: You may also use, modify, and distribute this software under the Elastic License v2, which has specific restrictions. - -You may choose either license based on its terms. The original VectorChord code -retains its upstream copyright notices. TileMaxSim fork additions are Copyright -(c) 2026 Hu Xinjing; use this fork's -[issue tracker](https://github.com/HuXinjing/VectorChord/issues) for support. - -[cost-estimation]: https://github.com/user-attachments/assets/168fe550-6465-4eee-a224-8c848c301e3d -[image-compare]: https://github.com/user-attachments/assets/2d985f1e-7093-4c3a-8bf3-9f0b92c0e7e7 -[license-1-link]: https://github.com/HuXinjing/VectorChord#license -[license-1-shield]: https://img.shields.io/badge/License-AGPLv3-green?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCIgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiBmaWxsPSIjZmZmZmZmIj48cGF0aCBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik0xMi43NSAyLjc1YS43NS43NSAwIDAwLTEuNSAwVjQuNUg5LjI3NmExLjc1IDEuNzUgMCAwMC0uOTg1LjMwM0w2LjU5NiA1Ljk1N0EuMjUuMjUgMCAwMTYuNDU1IDZIMi4zNTNhLjc1Ljc1IDAgMTAwIDEuNUgzLjkzTC41NjMgMTUuMThhLjc2Mi43NjIgMCAwMC4yMS44OGMuMDguMDY0LjE2MS4xMjUuMzA5LjIyMS4xODYuMTIxLjQ1Mi4yNzguNzkyLjQzMy42OC4zMTEgMS42NjIuNjIgMi44NzYuNjJhNi45MTkgNi45MTkgMCAwMDIuODc2LS42MmMuMzQtLjE1NS42MDYtLjMxMi43OTItLjQzMy4xNS0uMDk3LjIzLS4xNTguMzEtLjIyM2EuNzUuNzUgMCAwMC4yMDktLjg3OEw1LjU2OSA3LjVoLjg4NmMuMzUxIDAgLjY5NC0uMTA2Ljk4NC0uMzAzbDEuNjk2LTEuMTU0QS4yNS4yNSAwIDAxOS4yNzUgNmgxLjk3NXYxNC41SDYuNzYzYS43NS43NSAwIDAwMCAxLjVoMTAuNDc0YS43NS43NSAwIDAwMC0xLjVIMTIuNzVWNmgxLjk3NGMuMDUgMCAuMS4wMTUuMTQuMDQzbDEuNjk3IDEuMTU0Yy4yOS4xOTcuNjMzLjMwMy45ODQuMzAzaC44ODZsLTMuMzY4IDcuNjhhLjc1Ljc1IDAgMDAuMjMuODk2Yy4wMTIuMDA5IDAgMCAuMDAyIDBhMy4xNTQgMy4xNTQgMCAwMC4zMS4yMDZjLjE4NS4xMTIuNDUuMjU2Ljc5LjRhNy4zNDMgNy4zNDMgMCAwMDIuODU1LjU2OCA3LjM0MyA3LjM0MyAwIDAwMi44NTYtLjU2OWMuMzM4LS4xNDMuNjA0LS4yODcuNzktLjM5OWEzLjUgMy41IDAgMDAuMzEtLjIwNi43NS43NSAwIDAwLjIzLS44OTZMMjAuMDcgNy41aDEuNTc4YS43NS43NSAwIDAwMC0xLjVoLTQuMTAyYS4yNS4yNSAwIDAxLS4xNC0uMDQzbC0xLjY5Ny0xLjE1NGExLjc1IDEuNzUgMCAwMC0uOTg0LS4zMDNIMTIuNzVWMi43NXpNMi4xOTMgMTUuMTk4YTUuNDE4IDUuNDE4IDAgMDAyLjU1Ny42MzUgNS40MTggNS40MTggMCAwMDIuNTU3LS42MzVMNC43NSA5LjM2OGwtMi41NTcgNS44M3ptMTQuNTEtLjAyNGMuMDgyLjA0LjE3NC4wODMuMjc1LjEyNi41My4yMjMgMS4zMDUuNDUgMi4yNzIuNDVhNS44NDYgNS44NDYgMCAwMDIuNTQ3LS41NzZMMTkuMjUgOS4zNjdsLTIuNTQ3IDUuODA3eiI+PC9wYXRoPjwvc3ZnPgo= -[license-2-link]: https://github.com/HuXinjing/VectorChord#license -[license-2-shield]: https://img.shields.io/badge/License-ELv2-green?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCIgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiBmaWxsPSIjZmZmZmZmIj48cGF0aCBmaWxsLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik0xMi43NSAyLjc1YS43NS43NSAwIDAwLTEuNSAwVjQuNUg5LjI3NmExLjc1IDEuNzUgMCAwMC0uOTg1LjMwM0w2LjU5NiA1Ljk1N0EuMjUuMjUgMCAwMTYuNDU1IDZIMi4zNTNhLjc1Ljc1IDAgMTAwIDEuNUgzLjkzTC41NjMgMTUuMThhLjc2Mi43NjIgMCAwMC4yMS44OGMuMDguMDY0LjE2MS4xMjUuMzA5LjIyMS4xODYuMTIxLjQ1Mi4yNzguNzkyLjQzMy42OC4zMTEgMS42NjIuNjIgMi44NzYuNjJhNi45MTkgNi45MTkgMCAwMDIuODc2LS42MmMuMzQtLjE1NS42MDYtLjMxMi43OTItLjQzMy4xNS0uMDk3LjIzLS4xNTguMzEtLjIyM2EuNzUuNzUgMCAwMC4yMDktLjg3OEw1LjU2OSA3LjVoLjg4NmMuMzUxIDAgLjY5NC0uMTA2Ljk4NC0uMzAzbDEuNjk2LTEuMTU0QS4yNS4yNSAwIDAxOS4yNzUgNmgxLjk3NXYxNC41SDYuNzYzYS43NS43NSAwIDAwMCAxLjVoMTAuNDc0YS43NS43NSAwIDAwMC0xLjVIMTIuNzVWNmgxLjk3NGMuMDUgMCAuMS4wMTUuMTQuMDQzbDEuNjk3IDEuMTU0Yy4yOS4xOTcuNjMzLjMwMy45ODQuMzAzaC44ODZsLTMuMzY4IDcuNjhhLjc1Ljc1IDAgMDAuMjMuODk2Yy4wMTIuMDA5IDAgMCAuMDAyIDBhMy4xNTQgMy4xNTQgMCAwMC4zMS4yMDZjLjE4NS4xMTIuNDUuMjU2Ljc5LjRhNy4zNDMgNy4zNDMgMCAwMDIuODU1LjU2OCA3LjM0MyA3LjM0MyAwIDAwMi44NTYtLjU2OWMuMzM4LS4xNDMuNjA0LS4yODcuNzktLjM5OWEzLjUgMy41IDAgMDAuMzEtLjIwNi43NS43NSAwIDAwLjIzLS44OTZMMjAuMDcgNy41aDEuNTc4YS43NS43NSAwIDAwMC0xLjVoLTQuMTAyYS4yNS4yNSAwIDAxLS4xNC0uMDQzbC0xLjY5Ny0xLjE1NGExLjc1IDEuNzUgMCAwMC0uOTg0LS4zMDNIMTIuNzVWMi43NXpNMi4xOTMgMTUuMTk4YTUuNDE4IDUuNDE4IDAgMDAyLjU1Ny42MzUgNS40MTggNS40MTggMCAwMDIuNTU3LS42MzVMNC43NSA5LjM2OGwtMi41NTcgNS44M3ptMTQuNTEtLjAyNGMuMDgyLjA0LjE3NC4wODMuMjc1LjEyNi41My4yMjMgMS4zMDUuNDUgMi4yNzIuNDVhNS44NDYgNS44NDYgMCAwMDIuNTQ3LS41NzZMMTkuMjUgOS4zNjdsLTIuNTQ3IDUuODA3eiI+PC9wYXRoPjwvc3ZnPgo= - -[docker-release-link]: https://hub.docker.com/r/tensorchord/vchord-postgres -[docker-release-shield]: https://img.shields.io/docker/v/tensorchord/vchord-postgres?color=369eff&label=docker&labelColor=black&logo=docker&logoColor=white&style=flat -[github-release-link]: https://github.com/HuXinjing/VectorChord/releases -[github-release-shield]: https://img.shields.io/github/v/release/HuXinjing/VectorChord?color=369eff&labelColor=black&logo=github&style=flat -[ghcr-release-link]: https://github.com/users/HuXinjing/packages/container/package/vectorchord-tilemaxsimd - -[ghcr-release-shield]: https://img.shields.io/docker/v/tensorchord/vchord-postgres?color=369eff&label=GHCR&labelColor=black&logo=github&logoColor=white&style=flat -[docker-pulls-link]: https://hub.docker.com/r/tensorchord/vchord-postgres -[docker-pulls-shield]: https://img.shields.io/docker/pulls/tensorchord/vchord-postgres?color=45cc11&labelColor=black&style=flat&sort=semver -[previous-docker-pulls-link]: https://hub.docker.com/r/tensorchord/pgvecto-rs -[previous-docker-pulls-shield]: https://img.shields.io/docker/pulls/tensorchord/pgvecto-rs?color=45cc11&labelColor=black&style=flat&sort=semver -[github-downloads-link]: https://github.com/HuXinjing/VectorChord/releases -[github-downloads-shield]: https://img.shields.io/github/downloads/HuXinjing/VectorChord/total?color=45cc11&labelColor=black&style=flat&sort=semver -[discord-link]: https://discord.gg/KqswhpVgdU -[discord-shield]: https://img.shields.io/discord/974584200327991326?&logoColor=white&color=5865F2&style=flat&logo=discord&cacheSeconds=60 -[X-link]: https://x.com/TensorChord -[X-shield]: https://img.shields.io/badge/follow-%40tensorchord-1DA1F2?logo=x&style=flat&logoColor=white&color=1da1f2 -[deepwiki-link]: https://deepwiki.com/HuXinjing/VectorChord -[deepwiki-shield]: https://deepwiki.com/badge.svg -[blog-link]: https://blog.vectorchord.ai/ -[official-site-link]: https://vectorchord.ai/ -[github-issues-link]: https://github.com/HuXinjing/VectorChord/issues -[email-link]: https://github.com/HuXinjing/VectorChord/issues -[docs-link]: https://docs.vectorchord.ai/vectorchord/ +The CUDA container and its test stage are defined in +[`services/Dockerfile.tilemaxsimd`](services/Dockerfile.tilemaxsimd). PostgreSQL +tensor-source integration is documented in +[`docs/MAXSIM_TENSOR_SOURCES.md`](docs/MAXSIM_TENSOR_SOURCES.md). + +## Acknowledgements + +We thank the maintainers and contributors of the upstream +[supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord) +project. Its PostgreSQL extension, vector types, `vchordrq` index, and existing +MaxSim support provide the foundation on which this project builds. + +The project's IO-aware GPU MaxSim direction and the `TileMaxSim` name are +informed by Ashutosh Sharma's paper, +[“TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization”](https://arxiv.org/abs/2606.26439), +arXiv:2606.26439 (2026), and its accompanying open-source +[ashutoshuiuc/tilemaxsim](https://github.com/ashutoshuiuc/tilemaxsim) +Triton implementation. We gratefully acknowledge that work and its published +analysis of fused, tiled MaxSim scoring. + +This repository is a separate PostgreSQL/Rust/CUDA systems integration. Its +IPC protocols, persistent three-level tensor cache, allocator, multi-user +scheduler, database bindings, and operational packaging are maintained here; +the acknowledgement does not imply endorsement or shared maintenance by the +upstream projects or paper author. + +## Project provenance and license + +This project is derived from +[supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord). +Original source files retain their upstream copyright and license notices. +TileMaxSim project additions are Copyright (c) 2026 Hu Xinjing. + +The repository is offered under the license terms in [`LICENSE`](LICENSE), +including AGPLv3 and Elastic License v2 options where applicable. Use this +project's [issue tracker](https://github.com/HuXinjing/VectorChord/issues) for +support and project-specific reports. diff --git a/README.zh-CN.md b/README.zh-CN.md index 932c0578..dbc8d10e 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -2,41 +2,122 @@ # VectorChord TileMaxSim -**面向 PostgreSQL 精确多向量检索的开源 VectorChord 分支。** +**基于 PostgreSQL 的向量/张量检索、精确 TileMaxSim 与常驻 GPU 缓存。** [English](README.md) | 简体中文 -本分支为 VectorChord 的 `vchordrq` 索引增加精确 late-interaction -TileMaxSim 检索，适用于“每个文档保存一个 token 向量张量，并在 -PostgreSQL 内完成多向量搜索”的应用。 +VectorChord TileMaxSim 是一个构建在 PostgreSQL 上的开源检索引擎，代码源自 +上游 [VectorChord](https://github.com/supervc-stack/VectorChord)。项目保留传统 +单向量检索通路，并增加面向“每个文档由一组 token 向量表示”的精确 +late-interaction 检索。 + +本仓库作为独立项目维护。README 只介绍本项目的实现、实测结果、局限和路线图， +不再附带上游官方 README，也不会把上游项目的宣传、镜像或产品指标表述成本项目 +自己的能力。 + +## 项目定位 + +VectorChord TileMaxSim 提供基础设施能力： + +- 在 PostgreSQL 中存储向量、张量描述符和检索元数据； +- 兼容传统单向量检索，并提供精确多向量 TileMaxSim； +- 让检索遵循 PostgreSQL 权限、MVCC、行可见性、取消和超时语义； +- 提供常驻 GPU 张量 arena、host 内存缓存和不可变磁盘分片； +- 为可选 GPU 服务提供有界准入、公平/优先级调度、缓存配额、健康探针和指标。 + +本项目不负责应用身份、认证、ACL 策略、实体注册表、Fact、事件账本、关系图谱、 +社区或查询意图路由。这些属于经过认证的应用层或知识治理层。VectorChord 只接收 +上层已经完成授权的候选范围，以及不具有授权含义的调度提示。 + +## 核心能力 + +- CPU 精确 TileMaxSim，以及原生 Rust/CUDA `tilemaxsimd` 后端。 +- 不人为限制调用方授权范围内的张量候选数量，超出显存的工作集按有界批次处理。 +- 外部张量源注册：PostgreSQL 保存稳定公开 ID 和描述符，完整 token 张量可存放在 + 被注册的行外存储中。 +- GPU/内存/磁盘三级张量路径： + - L0：常驻 GPU 页面缓存； + - L1：有界 host 内存缓存； + - L2：持久存储上的不可变张量分片或内容寻址对象。 +- 批量分片读取、批量 Host→Device 传输、内容寻址去重、TinyLFU/GDSF 准入，以及 + 可合并的 page-run 分配器。 +- 可选 resident 预热；预热完成前 daemon 不会报告 ready。 +- daemon 内的公平、严格 priority 或 fair-priority 调度，请求 aging、租户权重、 + deadline、断连取消和有界 CUDA quantum。 +- 多向量查询的 PostgreSQL 规划器统计与成本估算。 +- Prometheus 指标，以及相互独立的存活和就绪探针。 +- 正确性、注册表、协议、规划器、真实 GPU 和故障路径测试。 + +TileMaxSim 是可选功能。传统单向量通路不会启动 CUDA daemon，也不要求配置显存。 +启用时，`tilemaxsimd` 会在启动阶段一次性预占用户指定设备上的 GiB 级显存；任何 +设备或容量无法获得都会直接退出。 + +## 架构 + +```text +认证后的应用 / 检索规划器 + │ 已授权 ID + 调度提示 + ▼ + PostgreSQL + VectorChord + │ 张量描述符 + ▼ + tilemaxsimd + ┌─────────┼──────────┐ + │ │ │ +L0 GPU L1 内存 L2 持久存储 +缓存 缓存 分片 + │ + └── CUDA 精确 TileMaxSim +``` + +PostgreSQL 是行可见性和硬过滤的权威来源。daemon 不会根据 tenant 或 priority 推断 +权限，只会对调用方明确传入的张量集合进行缓存解析和精确打分。 + +## GPU 缓存与调度 -VectorChord 在这里负责向量/张量存储、精确 TileMaxSim、GPU 缓存和调度。 -实体、事件、Fact、关系图谱、社区、意图路由、用户身份与 ACL 等应用治理能力 -应由 GBrain 等上层系统负责，不进入 VectorChord。 +默认 `fair-priority` 调度器组合显式优先级和加权公平性。大请求按候选数、文档 +token 数和实际 `query rows × document rows × dimension` 计算量切分。每个 CUDA +launch 结束后，请求重新进入调度器，因此其他调度域可以插入执行。已经开始的 +CUDA kernel 不能被中断，FMA quantum 用来限制这段协作式不可抢占时间。 -## 本分支增加的能力 +GPU 和 host cache 都有调度域最大占用限制。可选的 +`--tenant-cache-reservation TENANT=GB` 会保护已经预热的 GPU 页面，使其他调度域 +不能把它淘汰到低于指定容量。不过它当前只是“聚合字节保底”，不是某个知识库的 +永久显存分区：它不会自动预热，不为 L1 内存提供最低保留量，也不能区分同一租户 +拥有的多个 source。这些缺口已明确列入展望。 -- CPU 精确 TileMaxSim 重排，以及可选的原生 Rust/CUDA 后端。 -- 不人为限制调用方授权范围内的张量候选数量；显存压力由 GPU 页面缓存、 - 分块计算和有界请求调度处理。 -- 外部张量源注册，可将完整精度的 token 张量存放在 PostgreSQL 行外。 -- 外部张量查询路径遵循 PostgreSQL 权限、MVCC、行可见性、取消和超时语义。 -- 多向量查询的规划器统计与成本估算。 -- GPU/内存/磁盘三级存储路径、常驻预热、缓存换入换出和指标接口。 -- 按 daemon 实例执行的公平调度、优先级、准入限制和缓存配额。 -- 正确性、注册表、sidecar 协议、规划成本和真实 GPU 故障路径测试。 +所有内存参数使用 GiB，不使用裸字节数。示例： + +```shell +tilemaxsimd \ + --socket /run/vectorchord/tilemaxsim.sock \ + --listen 0.0.0.0:9191 \ + --status-socket /run/vectorchord/tilemaxsim-status.sock \ + --gpu-memory-gb 0=20 \ + --gpu-workspace-gb 2 \ + --host-cache-gb 8 \ + --max-inflight-request-gb 1 \ + --contract-root MODEL_CONTRACT_ID=/srv/vectorchord/tensors \ + --scheduler-policy fair-priority \ + --max-queued-requests 128 \ + --max-tenant-queued-requests 16 \ + --scheduler-quantum-fmas 4000000000 \ + --tenant-weight foreground=2 \ + --tenant-cache-reservation foreground=4 +``` -TileMaxSim 是可选功能。没有显式配置 GPU 显存时不会启动 CUDA daemon， -传统单向量 VectorChord 使用不需要 GPU。启用时，daemon 会在启动阶段预占用户 -指定设备上的 GiB 级显存；任何设备或容量无法获得都会直接失败退出。 +PostgreSQL 可以通过本地 Unix socket 或 `tcp://HOST:PORT` 连接 daemon。 +`GET /livez`、`GET /healthz` 和 `GET /metrics` 分别提供存活、就绪和有界运行指标。 +部署和协议细节见 +[`docs/TILEMAXSIM_CUDA_SIDECAR.md`](docs/TILEMAXSIM_CUDA_SIDECAR.md) 与 +[`docs/TILEMAXSIM_IPC_V2.md`](docs/TILEMAXSIM_IPC_V2.md)。 -## 性能实验 +## 性能实测 -开发机语料包含 34,054 个张量描述符、34,027 个唯一张量，逻辑 FP16 数据量 -为 16.28 GB。绝对延迟会随存储、CPU 和 GPU 改变，因此同机对照比单个数字 -更有参考价值。 +开发机语料包含 34,054 个张量描述符、34,027 个唯一张量，逻辑 FP16 数据量为 +16.28 GB。绝对延迟会随存储、CPU 和 GPU 改变，因此同机对照比孤立数字更有意义。 | 优化项 | 基线 | 优化后 | 结果 | | --- | ---: | ---: | ---: | @@ -48,37 +129,36 @@ TileMaxSim 是可选功能。没有显式配置 GPU 显存时不会启动 CUDA d | Rust/CUDA 热请求 p50 | Python/Triton 14.26 ms | 2.09 ms | 快 6.82 倍 | | Rust/CUDA 热请求 p95 | Python/Triton 14.84 ms | 2.20 ms | 快 6.75 倍 | -20 GiB 显存实验把 18 GiB 用于张量、2 GiB 用于 TileMaxSim 工作区，服务就绪前 -预热全部 34,027 个唯一张量。Python/Triton 从进程启动到就绪耗时 23.07 秒， -Rust/CUDA daemon 为 14.86 秒；预热后的常驻请求不再从磁盘读取张量。 +20 GiB 实验将 18 GiB 分给张量、2 GiB 分给工作区。Rust/CUDA daemon 预热全部 +34,027 个唯一张量耗时 14.86 秒，Python/Triton 实现为 23.07 秒。默认 32 KiB +page-run 分配器将原 256 KiB buddy allocator 的 17.840 GB 分配量降至 +16.725 GB，内部取整浪费从 8.82% 降至 2.74%。 -GPU 缓存从一块 CUDA arena 中分配连续 page run，并使用 best-fit 尺寸桶和按 -地址合并。默认 32 KiB 页面将语料的分配空间从原 256 KiB buddy allocator 的 -17.840 GB 降至 16.725 GB，内部取整浪费从 8.82% 降至 2.74%。页面大小可以用 -`--gpu-block-kib` 调整。 +### 全范围与小缓存压力 -### 全范围检索与小缓存压力 +全部张量常驻时，刻意对 34,054 个描述符做全量精确扫描平均为 1.08 秒。3 GiB +张量 arena 加 1 GiB host cache 返回相同的精确 top-K，但连续全量扫描产生 +68,106 次 miss、61,499 次 eviction 和 32.53 GB Host→Device 传输；原生 +round trip 平均达到 18.47 秒。 -18 GiB 张量 arena 中全部张量常驻时，34,054 个描述符的刻意全量精确扫描平均 -为 1.08 秒。3 GiB 张量 arena 加 1 GiB host cache 返回相同的精确 top-K，但 -连续全量扫描产生 68,106 次 miss、61,499 次 eviction 和 32.53 GB Host→Device -传输，平均延迟达到 18.47 秒。 +该差距来自循环缓存抖动和重复 I/O，不是显存缩小几倍后 TileMaxSim 算术就线性 +变慢。高召回范围内的热工作集能放进缓存时，不会出现这种全库扫描行为。 -这说明小缓存下结果仍然正确，但循环抖动会显著增加延迟；它不是“显存缩小几倍， -GPU 算术就线性变慢”。高召回范围内的热工作集能够放入缓存时不会出现这种全库 -扫描行为。 +### 候选范围消融 -在 40 个真实查询上的传统 `text-embedding-v4`/pgvector HNSW 对照如下： +40 个真实查询使用原始 1,024 维文本 embedding/pgvector HNSW，与词法范围内的 +精确 TileMaxSim 对照： | 路径 | 平均延迟 | p95 | 质量 | | --- | ---: | ---: | ---: | | 原始文本 embedding + pgvector HNSW | 5.80 ms | 5.86 ms | 文档 hit@1 1.000 | -| 18 GiB 常驻缓存、候选范围内精确 TileMaxSim | 45.56 ms | 134.99 ms | 文档 hit@1 0.475；hit@5 0.650 | -| 3 GiB 缓存、候选范围内精确 TileMaxSim | 573.43 ms | 1,812.70 ms | 排序相同 | +| 18 GiB 常驻缓存、候选范围内 TileMaxSim | 45.56 ms | 134.99 ms | hit@1 0.475；hit@5 0.650 | +| 3 GiB 缓存、候选范围内 TileMaxSim | 573.43 ms | 1,812.70 ms | TileMaxSim 排序相同 | -该实验中的词法范围召回率只有 0.650，因此词法或普通图谱结果不能作为一般语义 -检索的唯一硬门。它们可以加权、解释或预热缓存；关系类意图可以使用完整图谱范围， -而 source、ACL、文档类型等授权约束仍必须硬过滤。 +词法范围只覆盖了 65% 查询的相关文档；对于已经覆盖的查询，TileMaxSim 全部把相关 +文档保留在 top 5。因此召回上限由候选生成而不是精确打分决定。普通词法和图谱命中 +不能作为一般语义查询的唯一硬门；授权过滤仍然必须硬执行，明确的关系类意图可以在 +上层能构造完整关系范围时使用图谱范围。 复现实验入口： @@ -87,98 +167,146 @@ GPU 算术就线性变慢”。高召回范围内的热工作集能够放入缓 - [`services/benchmark_gbrain_scoped_tilemaxsim.py`](services/benchmark_gbrain_scoped_tilemaxsim.py) - [`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py) -## GPU 缓存与调度 +## 已知局限 -默认 `fair-priority` 调度器组合显式优先级和加权公平性。大请求会按候选/token -计算量切分，并在 CUDA kernel 之间重新进入调度器。队列、进行中请求内存、GPU -和 host cache 都可以设置全局及调度域上限。 +本项目适合受控生产试点。严格低延迟、高可用或大规模多用户场景在正式 GA 前，必须 +解决或明确接受以下限制。 -内存配置使用 GiB，不使用字节数。示例： +### 检索规模 -```shell -tilemaxsimd \ - --socket /run/vectorchord/tilemaxsim.sock \ - --listen 0.0.0.0:9191 \ - --status-socket /run/vectorchord/tilemaxsim-status.sock \ - --gpu-memory-gb 0=20 \ - --gpu-workspace-gb 2 \ - --host-cache-gb 8 \ - --max-inflight-request-gb 1 \ - --contract-root MODEL_CONTRACT_ID=/srv/vectorchord/tensors \ - --scheduler-policy fair-priority \ - --max-queued-requests 128 \ - --max-tenant-queued-requests 16 \ - --scheduler-quantum-fmas 4000000000 \ - --tenant-weight foreground=2 \ - --tenant-cache-reservation foreground=4 +- 精确 TileMaxSim 是打分器，不是张量原生 ANN。全部常驻只能消除传输，不能消除 + 精确计算量。 +- 大规模低延迟知识库仍需要经过实测的高召回候选生成器，或未来的张量原生 ANN， + 再执行精确 TileMaxSim。 +- 缓存小于活跃工作集时结果仍然正确，但可能发生抖动；生产容量必须覆盖热工作集， + 或保留经过验证的高召回范围。 + +### 缓存与多用户隔离 + +- 公平、priority、准入和 reservation 由每个 daemon 独立执行。多个副本之间没有 + 全局队列、全局 entitlement 账本或感知缓存驻留的路由器。 +- GPU reservation 保护的是某个调度 tenant 的聚合字节数，不是命名的 + `source_id` 或知识库；当前没有动态 cache-domain API、source 级预热契约或 + host-cache 最低保留量。 +- reservation 在 daemon 启动时静态配置。多 GPU daemon 当前会把该数值应用到 + 每张卡，而不是一个语义明确的集群总量。 +- 全局 resident manifest 使用服务内部域固定。如果 pinned 页面占满 arena， + 无关的冷张量不能淘汰它们，请求会失败，而不是等待 pinned 数据离开。 +- 抢占只能发生在 CUDA launch 之间，正在执行的 kernel 不可中断。 +- tenant 和 priority 只是调度提示，不是授权证据。 + +### 运维与安全 + +- TileMaxSim TCP 协议没有内置 TLS 和客户端认证，必须放在私网或双向认证代理后。 +- GPU 副本替换后是冷缓存。持久存储保证正确性，但可用性和热缓存延迟是两个独立 + SLO。 +- PostgreSQL WAL 归档、PITR、跨站灾备和真实恢复验证仍属于数据库平台职责。 +- 已提供指标，但还没有开箱即用的 SLO dashboard 和告警规则。 +- 超大张量注册表与垃圾回收需要按实际规模做运维压测。 + +### 发布状态 + +SQL 接口、外部张量协议、镜像 tag 和部署封装仍在持续演进。真实 GPU 和故障路径 +测试不能替代目标硬件上的长时间、多用户 soak/chaos 测试。 + +## 展望：AI Infra 的检索平面 + +VectorChord TileMaxSim 的目标是在可组合 AI Infra 中承担检索平面，而不是把应用 +治理和模型服务职责都塞进自身： + +```text +Agent / 应用层 + │ + ▼ +AI Gateway：身份、配额、priority、deadline、Trace + │ + ├── 知识层 ── PostgreSQL / VectorChord ── tilemaxsimd + │ + └── Ray Serve ── vLLM 与其他模型引擎 + +Kubernetes / KubeRay：节点放置、GPU、扩缩容、故障恢复 +对象存储：模型权重、张量分片、注入产物 +Prometheus + OpenTelemetry：端到端可观测性 ``` -发布镜像使用 CUDA 12.6 构建，需要 NVIDIA Container Runtime，以及兼容 -CUDA 12.6 的宿主机驱动。最终镜像保留静态链接的 CUDA 应用运行时，但不包含 -不需要的完整 CUDA toolkit；驱动库和 GPU 设备由 NVIDIA 运行时注入。 +在这套分工中： + +- [Ray Serve LLM](https://docs.ray.io/en/latest/serve/llm/) 负责分布式模型副本、 + 自动扩缩容和模型感知路由； +- [vLLM](https://docs.vllm.ai/) 负责生成模型的连续批处理、模型并行、KV cache 和 + prefix cache； +- VectorChord 负责 PostgreSQL 检索语义以及向量/张量元数据，`tilemaxsimd` 负责 + 精确 TileMaxSim、张量驻留和有界检索调度； +- 经过认证的应用层负责身份、ACL、图谱/Fact/事件治理、查询意图和最终 RAG 链路。 + +各组件不能在未感知对方的情况下竞争同一块显存。初期部署应让 vLLM 与 TileMaxSim +使用独立 GPU 池；Kubernetes 或其他资源管理器必须知道常驻 TileMaxSim daemon +已经预占的显存，之后再评估 MIG 或其他硬件分区方案。 + +下一阶段基础设施里程碑： + +1. 将 `scheduler_tenant` 与不透明的 `cache_domain` 分离；后者通常由已认证上层根据 + tenant 和 source 派生，但 VectorChord 不解释其中的业务含义。 +2. 为每个 cache domain 增加 GPU 和 host 的 `min_resident_gb`、`max_gb` 与预热 + 契约，并明确区分单卡和全局语义。 +3. 建立有界的全局准入层和感知张量驻留的多副本路由，同时保留 daemon 内精确的 + quantum 调度。 +4. 让 request ID、tenant、source、priority、deadline、trace ID 和 cache domain + 组成统一请求信封，贯穿检索与模型服务。 +5. 验证冷启动恢复、滚动升级、PostgreSQL 切换、GPU Pod 丢失、多租户隔离和长时间 + soak/chaos 行为。 +6. 为无法用全范围精确 TileMaxSim 满足延迟 SLO 的大知识库增加张量原生 ANN，或 + 其他经过召回率验证的第一阶段。 + +Ray 是编排层，不是持久存储。PostgreSQL 和不可变对象存储仍是权威数据源。 +`tilemaxsimd` 应继续作为受监督的有状态 GPU 服务，而不是短生命周期的普通 Ray +任务，这样才能保证预占 arena 和缓存生命周期可预测。 -PostgreSQL 的 `vchordrq.maxsim_gpu_endpoint` 可以是本地 Unix socket，也可以是 -`tcp://HOST:PORT`。`GET /livez`、`GET /healthz` 和 `GET /metrics` 分别用于 -存活、就绪和 Prometheus 指标检查。 +## 安全边界 -## 已知限制（Limitations） +只有设置 `vchordrq.maxsim_tenant` 时，PostgreSQL 才发送带调度信息的协议。 +priority 只改变延迟顺序，不能绕过行可见性、ACL、source 或其他硬过滤。日志只 +暴露稳定哈希，不输出原始 tenant 标识。 -本分支可以用于受控生产试点，但在严格低延迟、高可用或大规模多用户场景正式 -GA 前，必须解决或明确接受以下限制。 +本公开仓库只包含实现和对外项目信息，私有应用文档与规划文档不会提交到本仓库。 -### 检索规模 +## 构建与测试 -- TileMaxSim 是精确打分器，不是张量原生 ANN。source、ACL、文档类型等授权硬 - 过滤可以安全缩小范围，但普通词法或图谱命中不能充当一般语义查询的高召回硬门。 -- 张量全部常驻 GPU 只能消除存储传输，不能消除精确计算量。本次实测全量 34,054 - 个描述符平均 1.08 秒，而高召回应用范围平均 45.56 ms。大知识库的严格低延迟 - 部署仍需要张量原生 ANN，或经过召回率验证的候选生成器，再做精确 TileMaxSim。 -- 缓存小于活跃工作集时不会返回错误结果，但可能频繁换入换出。前述 3 GiB 压力 - 实验的全量扫描平均 18.47 秒，生产容量必须覆盖热工作集或保留高召回范围。 - -### 多用户与调度 - -- 公平、priority、准入上限和缓存配额目前由每个 daemon 独立执行。多个 GPU - 副本之间还没有全局队列、全局 entitlement 账本或感知工作量的负载均衡器， - 因此可能出现副本负载倾斜。 -- 抢占发生在 CUDA kernel 之间，无法中断已经开始执行的 kernel；FMA quantum - 负责限制单次不可抢占的最长计算量。 -- GPU 故障转移可依靠持久张量存储保持正确性，但新副本是冷缓存。高可用和热缓存 - 延迟是两个独立 SLO，需要分别验证。 -- tenant 标识和 priority 只是调度提示，不是授权凭据。VectorChord 不负责应用 - 多租户、身份、ACL、图谱治理或特权名单；上层认证系统必须只传入已授权硬范围。 - -### 部署、安全与灾备 - -- TileMaxSim TCP 协议当前没有内置 TLS 和客户端认证，只能放在私网并通过防火墙 - 或 Kubernetes NetworkPolicy 限制 PostgreSQL 调用；跨信任域时应增加双向认证 - 代理。 -- 容器、探针、指标、PostgreSQL 集成和适配 HA 的 TCP 服务已经提供，但 manifest - 不等于生产验证。必须在目标环境演练 PostgreSQL 主从切换、GPU Pod 丢失、滚动 - 升级、备份恢复和数据迁移。 -- 本仓库不负责持续 WAL 归档、PITR 和跨站灾备；这些属于数据库平台职责，而且 - 必须做真实恢复测试。 -- 已提供指标接口，但没有开箱即用的 SLO dashboard 和告警规则。生产环境需要监控 - 排队、拒绝、超时、缓存抖动、磁盘容量、GPU 健康和故障后的冷缓存延迟。 -- 张量对象数量很大时，GC 文件扫描和数据库锁可能变得可见，需要按实际规模压测并 - 规划维护窗口或后续分代/墓碑式 GC。 +常用入口： -### 发布状态 +```shell +make build +cargo test --locked --workspace --exclude vchord --no-fail-fast +cargo test --manifest-path services/tilemaxsimd/Cargo.toml --locked +``` -SQL 接口、外部张量协议、镜像 tag 和部署封装仍在持续开发，稳定版之前可能变化。 -仓库中的真实 GPU 与故障路径测试验证了有界执行和正确性，但不能替代目标硬件上的 -多小时、多用户 soak/chaos 压测。 +CUDA 容器及测试阶段定义在 +[`services/Dockerfile.tilemaxsimd`](services/Dockerfile.tilemaxsimd)，PostgreSQL +张量源集成见 [`docs/MAXSIM_TENSOR_SOURCES.md`](docs/MAXSIM_TENSOR_SOURCES.md)。 -## 安全边界 +## 致谢 + +感谢上游 [supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord) +项目的维护者与贡献者。其 PostgreSQL 扩展、向量类型、`vchordrq` 索引和已有 +MaxSim 能力构成了本项目继续开发的基础。 + +本项目的 I/O 感知 GPU MaxSim 方向和 `TileMaxSim` 名称受到 Ashutosh Sharma +论文 +[《TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization》](https://arxiv.org/abs/2606.26439) +（arXiv:2606.26439，2026）及其配套开源 +[ashutoshuiuc/tilemaxsim](https://github.com/ashutoshuiuc/tilemaxsim) Triton 实现 +的启发。感谢该工作公开的融合、分块 MaxSim 设计与性能分析。 -PostgreSQL 只有在设置 `vchordrq.maxsim_tenant` 时才发送 v3 调度元数据，否则 -保持 v2 协议兼容。上层系统可以设置 -100 到 100 的 priority，但优先级只影响 -延迟顺序，不能绕过 PostgreSQL 行可见性、ACL、source 或其他硬过滤。 +本仓库是独立的 PostgreSQL/Rust/CUDA 系统集成，IPC 协议、GPU/内存/磁盘三级 +张量缓存、分配器、多用户调度、数据库绑定和部署运维由本项目维护。上述致谢不表示 +上游项目或论文作者对本项目背书，也不表示双方共同维护。 -本仓库只包含公开实现与公开项目信息，不包含私有规划或上层应用文档。 +## 项目来源与许可 -## 项目关系 +本项目源自 [supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord)， +原始源文件保留其上游版权和许可声明。TileMaxSim 项目新增部分 Copyright (c) 2026 +Hu Xinjing。 -本项目基于 [supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord)。 -上游安装方式、SQL 用法、兼容性与许可证信息请参阅[英文 README](README.md) 中保留的 -原始 VectorChord 文档。 +仓库按 [`LICENSE`](LICENSE) 中的许可条款提供，在适用范围内包括 AGPLv3 和 +Elastic License v2 选项。项目问题与支持请求请提交到本项目的 +[Issue Tracker](https://github.com/HuXinjing/VectorChord/issues)。 From 794e76325fef29f30bf5c35b5ea19b380c1349c0 Mon Sep 17 00:00:00 2001 From: HuXinjing <49928618+HuXinjing@users.noreply.github.com> Date: Thu, 16 Jul 2026 16:11:23 +0800 Subject: [PATCH 324/324] feat(tilemaxsimd): overlap tensor IO with CUDA scoring --- README.md | 25 + README.zh-CN.md | 20 + docs/TILEMAXSIM_CUDA_SIDECAR.md | 40 + services/benchmark_tilemaxsim_io_pipeline.py | 241 ++++++ services/test_tilemaxsim_rust_daemon.py | 43 + services/tilemaxsimd/src/engine.rs | 845 ++++++++++++++----- services/tilemaxsimd/src/gpu.rs | 12 +- services/tilemaxsimd/src/main.rs | 185 +++- 8 files changed, 1189 insertions(+), 222 deletions(-) create mode 100644 services/benchmark_tilemaxsim_io_pipeline.py diff --git a/README.md b/README.md index b25c1944..da55600d 100644 --- a/README.md +++ b/README.md @@ -115,12 +115,15 @@ tilemaxsimd \ --gpu-memory-gb 0=20 \ --gpu-workspace-gb 2 \ --host-cache-gb 8 \ + --io-pipeline overlap \ + --io-batch-gb 0.05 \ --max-inflight-request-gb 1 \ --contract-root MODEL_CONTRACT_ID=/srv/vectorchord/tensors \ --scheduler-policy fair-priority \ --max-queued-requests 128 \ --max-tenant-queued-requests 16 \ --scheduler-quantum-fmas 4000000000 \ + --scheduler-quantum-io-gb 1 \ --tenant-weight foreground=2 \ --tenant-cache-reservation foreground=4 ``` @@ -148,6 +151,15 @@ numbers. | Rust/CUDA cold request | 855.34 ms, Python/Triton | 93.86 ms | 9.11x faster | | Rust/CUDA warm request p50 | 14.26 ms, Python/Triton | 2.09 ms | 6.82x faster | | Rust/CUDA warm request p95 | 14.84 ms, Python/Triton | 2.20 ms | 6.75x faster | +| Small-VRAM cold I/O pipeline | 757.96 ms, serial | 685.85 ms, overlap | 1.11x faster; identical scores | + +The I/O-pipeline row used five cold trials over 1,000 real tensors +(478,016,000 logical bytes) with only 0.5 GiB assigned to the daemon: 0.4 GiB +for tensor pages and 0.1 GiB for CUDA workspace. The 0.05 GiB bounded resolver +stage overlaps SSD/L1 resolution of batch N+1, H2D upload of batch N, and exact +TileMaxSim for batch N-1. The maximum score delta between serial and overlap +was zero. Results are hardware- and batch-size-sensitive, so production must +measure the exposed pipeline metrics rather than assume this exact ratio. A 20 GiB run assigned 18 GiB to tensors and 2 GiB to workspace. Prewarming all 34,027 unique tensors took 14.86 seconds in the Rust/CUDA daemon versus 23.07 @@ -192,6 +204,7 @@ Reproducible drivers: - [`services/benchmark_full_corpus_tilemaxsim.py`](services/benchmark_full_corpus_tilemaxsim.py) - [`services/benchmark_gbrain_scoped_tilemaxsim.py`](services/benchmark_gbrain_scoped_tilemaxsim.py) - [`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py) +- [`services/benchmark_tilemaxsim_io_pipeline.py`](services/benchmark_tilemaxsim_io_pipeline.py) ## Limitations @@ -209,6 +222,11 @@ accept the following limitations before general availability. - A cache smaller than the active working set remains correct but can thrash; production capacity must cover the hot set or preserve a safe high-recall range. +- The current Tutti-inspired pipeline still uses CPU shard reads and pinned + host-to-device copies; it is not GPU io_uring or a GPUDirect Storage backend. + SSD, PCIe, and H2D contention make `--io-batch-gb` hardware-sensitive. Use + `--io-pipeline serial` as the ablation/fallback mode and inspect + `tilemaxsim_io_pipeline_*` before enabling overlap for an SLO. ### Cache and multi-user isolation @@ -348,6 +366,13 @@ arXiv:2606.26439 (2026), and its accompanying open-source Triton implementation. We gratefully acknowledge that work and its published analysis of fused, tiled MaxSim scoring. +The bounded asynchronous tensor pipeline is informed by +[“Tutti: Making SSD-Backed KV Cache Practical for Long-Context LLM Serving”](https://arxiv.org/abs/2605.03375), +particularly its GPU-native object, asynchronous I/O, and slack-aware +scheduling principles. This implementation currently adopts the scheduling +and pipelining ideas only; it does not include Tutti's GPU io_uring storage +stack. + This repository is a separate PostgreSQL/Rust/CUDA systems integration. Its IPC protocols, persistent three-level tensor cache, allocator, multi-user scheduler, database bindings, and operational packaging are maintained here; diff --git a/README.zh-CN.md b/README.zh-CN.md index dbc8d10e..cbf9bb20 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -98,12 +98,15 @@ tilemaxsimd \ --gpu-memory-gb 0=20 \ --gpu-workspace-gb 2 \ --host-cache-gb 8 \ + --io-pipeline overlap \ + --io-batch-gb 0.05 \ --max-inflight-request-gb 1 \ --contract-root MODEL_CONTRACT_ID=/srv/vectorchord/tensors \ --scheduler-policy fair-priority \ --max-queued-requests 128 \ --max-tenant-queued-requests 16 \ --scheduler-quantum-fmas 4000000000 \ + --scheduler-quantum-io-gb 1 \ --tenant-weight foreground=2 \ --tenant-cache-reservation foreground=4 ``` @@ -128,6 +131,13 @@ PostgreSQL 可以通过本地 Unix socket 或 `tcp://HOST:PORT` 连接 daemon。 | Rust/CUDA 冷请求 | Python/Triton 855.34 ms | 93.86 ms | 快 9.11 倍 | | Rust/CUDA 热请求 p50 | Python/Triton 14.26 ms | 2.09 ms | 快 6.82 倍 | | Rust/CUDA 热请求 p95 | Python/Triton 14.84 ms | 2.20 ms | 快 6.75 倍 | +| 小显存冷 I/O 流水线 | 串行 757.96 ms | 重叠 685.85 ms | 快 1.11 倍；分数一致 | + +I/O 流水线一行使用 1,000 个真实张量（逻辑 478,016,000 字节）做了 5 次冷请求， +daemon 只分配 0.5 GiB：0.4 GiB 张量页面加 0.1 GiB CUDA 工作区。0.05 GiB +有界解析批次让 N+1 批 SSD/L1 解析、N 批 H2D 上传与 N-1 批精确 TileMaxSim +重叠执行。串行与重叠路径的最大分数差为 0。实际收益对硬件和批次大小敏感，生产 +环境应以流水线指标实测，不能直接套用该倍数。 20 GiB 实验将 18 GiB 分给张量、2 GiB 分给工作区。Rust/CUDA daemon 预热全部 34,027 个唯一张量耗时 14.86 秒，Python/Triton 实现为 23.07 秒。默认 32 KiB @@ -166,6 +176,7 @@ round trip 平均达到 18.47 秒。 - [`services/benchmark_full_corpus_tilemaxsim.py`](services/benchmark_full_corpus_tilemaxsim.py) - [`services/benchmark_gbrain_scoped_tilemaxsim.py`](services/benchmark_gbrain_scoped_tilemaxsim.py) - [`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py) +- [`services/benchmark_tilemaxsim_io_pipeline.py`](services/benchmark_tilemaxsim_io_pipeline.py) ## 已知局限 @@ -180,6 +191,10 @@ round trip 平均达到 18.47 秒。 再执行精确 TileMaxSim。 - 缓存小于活跃工作集时结果仍然正确，但可能发生抖动；生产容量必须覆盖热工作集， 或保留经过验证的高召回范围。 +- 当前受 Tutti 启发的流水线仍通过 CPU 读取分片并经 pinned host memory 上传， + 还不是 GPU io_uring 或 GPUDirect Storage 后端。SSD、PCIe 与 H2D 竞争使 + `--io-batch-gb` 对硬件敏感；应使用 `--io-pipeline serial` 做消融/回退，并在 + SLO 上线前检查 `tilemaxsim_io_pipeline_*` 指标。 ### 缓存与多用户隔离 @@ -297,6 +312,11 @@ MaxSim 能力构成了本项目继续开发的基础。 [ashutoshuiuc/tilemaxsim](https://github.com/ashutoshuiuc/tilemaxsim) Triton 实现 的启发。感谢该工作公开的融合、分块 MaxSim 设计与性能分析。 +有界异步张量流水线也受到 +[《Tutti: Making SSD-Backed KV Cache Practical for Long-Context LLM Serving》](https://arxiv.org/abs/2605.03375) +中 GPU 原生对象、异步 I/O 与 slack-aware 调度原则的启发。当前实现只引入其调度 +与流水思想，不包含 Tutti 的 GPU io_uring 存储栈。 + 本仓库是独立的 PostgreSQL/Rust/CUDA 系统集成，IPC 协议、GPU/内存/磁盘三级 张量缓存、分配器、多用户调度、数据库绑定和部署运维由本项目维护。上述致谢不表示 上游项目或论文作者对本项目背书，也不表示双方共同维护。 diff --git a/docs/TILEMAXSIM_CUDA_SIDECAR.md b/docs/TILEMAXSIM_CUDA_SIDECAR.md index 3723935c..c1951caa 100644 --- a/docs/TILEMAXSIM_CUDA_SIDECAR.md +++ b/docs/TILEMAXSIM_CUDA_SIDECAR.md @@ -57,9 +57,12 @@ tilemaxsimd \ --gpu-memory-gb 0=20 \ --gpu-workspace-gb 2 \ --host-cache-gb 8 \ + --io-pipeline overlap \ + --io-batch-gb 0.05 \ --max-inflight-request-gb 1 \ --contract-root 'MODEL_CONTRACT_ID=/var/lib/vectorchord/tensors' \ --scheduler-policy fair-priority \ + --scheduler-quantum-io-gb 1 \ --max-connections 256 \ --max-queued-requests 128 \ --max-tenant-queued-requests 16 @@ -68,6 +71,39 @@ tilemaxsimd \ Multiple GPU assignments may be supplied. Device selection and the requested allocation are never inferred from all currently free VRAM. +## Tutti-inspired tensor I/O pipeline + +`--io-pipeline overlap` is the default for the opt-in CUDA daemon. It uses a +bounded one-batch resolver channel and separate CUDA upload/compute streams to +pipeline three stages: + +```text +immutable shard or L1 resolve N+1 + │ + ├── H2D upload N + │ + └── exact TileMaxSim N-1 +``` + +`--io-batch-gb` bounds the logical payload held by each resolver stage; the +default is 0.05 GiB. The channel has capacity one, so speculative prefetch +cannot grow without bound. GPU cache references keep the current compute pages +non-evictable, and a next batch that cannot acquire pages while computation is +active is deferred and retried after the current batch releases them. The +request therefore remains correct even when its working set exceeds L0. + +The cooperative scheduler also caps the L0-miss payload in one quantum with +`--scheduler-quantum-io-gb` in addition to candidate, token, and FMA limits. +This shortens the period for which a cold request can occupy the I/O/GPU path +before a higher-priority or under-served scheduling domain can run. + +Use `--io-pipeline serial` for a controlled fallback or an A/B comparison. +The current implementation retains ordinary CPU shard reads and pinned H2D +copies. It adopts Tutti's bounded asynchronous pipeline and I/O-aware +scheduling ideas, but it is not yet a GPU io_uring or GPUDirect Storage data +path. SSD and GPU traffic may share a PCIe root complex, so tune the batch on +target hardware and reject configurations that regress the application SLO. + The example systemd unit and environment file are in `deploy/systemd`. The unit is deliberately opt-in: enabling it is the operator's explicit TileMaxSim/GPU configuration. It waits for readiness, reloads immutable shard indexes with @@ -105,6 +141,10 @@ The most useful production signals are: overload or an undersized deadline; - `tilemaxsim_gpu_cache_events_total`, `tilemaxsim_gpu_h2d_bytes_total`, and `tilemaxsim_storage_read_bytes_total` for cache churn and cold-load traffic; +- `tilemaxsim_io_pipeline_batches_total`, + `tilemaxsim_io_pipeline_resolve_bytes_total`, and + `tilemaxsim_io_pipeline_seconds_total` for resolve/upload/compute time and + time hidden by overlap; - `tilemaxsim_gpu_cache_bytes` for free space, largest free extent, payload, allocator waste, and pinned capacity; - `tilemaxsim_host_cache_bytes` and `tilemaxsim_host_cache_events_total` for diff --git a/services/benchmark_tilemaxsim_io_pipeline.py b/services/benchmark_tilemaxsim_io_pipeline.py new file mode 100644 index 00000000..df2eeaac --- /dev/null +++ b/services/benchmark_tilemaxsim_io_pipeline.py @@ -0,0 +1,241 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Compare serial and overlapped tensor I/O on the native CUDA daemon.""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import tempfile +import time +from pathlib import Path + +import numpy as np + +from services.benchmark_tilemaxsim_ablation import ( + encode_frame, + evict_paths, + request_round_trip, +) + + +def load_records(path: Path, maximum_candidates: int) -> list[dict[str, object]]: + records: list[dict[str, object]] = [] + tokens = 0 + logical_bytes = 0 + with path.open(encoding="utf-8") as stream: + for line in stream: + if not line.strip(): + continue + record = json.loads(line) + rows = int(record["tensor_rows"]) + payload_bytes = int(record["canonical_bytes"]) + if records and ( + len(records) >= maximum_candidates + or tokens + rows > 1_000_000 + or logical_bytes + payload_bytes > 1024**3 + ): + break + records.append(record) + tokens += rows + logical_bytes += payload_bytes + if len(records) < 2: + raise ValueError("the benchmark needs at least two protocol-valid tensors") + return records + + +def run_mode( + *, + mode: str, + trial: int, + binary: Path, + shard_root: Path, + contract: str, + device: int, + gpu_memory_gb: str, + workspace_gb: str, + host_cache_gb: str, + io_batch_gb: str, + frame: bytes, + shard_paths: list[Path], + directory: Path, +) -> tuple[float, list[tuple[int, float]], dict[str, object]]: + evict_paths(shard_paths) + socket_path = directory / f"{mode}-{trial}.sock" + command = [ + os.fspath(binary), + "--socket", + os.fspath(socket_path), + "--gpu-memory-gb", + f"{device}={gpu_memory_gb}", + "--gpu-workspace-gb", + workspace_gb, + "--host-cache-gb", + host_cache_gb, + "--io-pipeline", + mode, + "--io-batch-gb", + io_batch_gb, + "--contract-root", + f"{contract}={shard_root}", + "--request-timeout-ms", + "60000", + "--socket-io-timeout-ms", + "60000", + "--once", + ] + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + for _ in range(3000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + if process.poll() is not None or not socket_path.exists(): + output, _ = process.communicate(timeout=5) + raise RuntimeError(f"{mode} daemon failed to start: {output}") + latency_ms, scores = request_round_trip(socket_path, frame) + output, _ = process.communicate(timeout=70) + if process.returncode != 0: + raise RuntimeError(f"{mode} daemon failed: {output}") + events = [ + json.loads(line) for line in output.splitlines() if line.startswith("{") + ] + request_event = next( + event for event in events if event.get("event") == "tilemaxsim_rust_request" + ) + return latency_ms, scores, request_event["cache"]["io_pipeline"] + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=10) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--query-embedding", required=True, type=Path) + parser.add_argument("--shard-root", required=True, type=Path) + parser.add_argument("--rust-binary", required=True, type=Path) + parser.add_argument("--contract", required=True) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--candidates", type=int, default=1000) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--gpu-memory-gb", default="1") + parser.add_argument("--workspace-gb", default="0.2") + parser.add_argument("--host-cache-gb", default="0.1") + parser.add_argument("--io-batch-gb", default="0.05") + parser.add_argument("--output", type=Path) + args = parser.parse_args() + if args.candidates < 2 or args.trials < 1: + parser.error("candidates must be >= 2 and trials must be positive") + + records = load_records(args.descriptor_manifest, args.candidates) + query = np.load(args.query_embedding, allow_pickle=False) + if query.dtype != np.dtype("float16"): + query = query.astype(" tuple[str, list[tuple[int, float]]]: binary = self._release_binary() if not binary.exists(): @@ -131,6 +133,10 @@ def run_daemon( workspace_gb, "--host-cache-gb", "0.01", + "--io-pipeline", + io_pipeline, + "--io-batch-gb", + io_batch_gb, "--contract-root", f"model@1={shard_root}", "--once", @@ -200,6 +206,7 @@ def run_daemon( self.assertIn(b"tilemaxsim_gpu_cache_bytes", metrics_body) self.assertIn(b"tilemaxsim_host_cache_bytes", metrics_body) self.assertIn(b"tilemaxsim_storage_read_bytes_total", metrics_body) + self.assertIn(b"tilemaxsim_io_pipeline_enabled", metrics_body) self.assertNotIn(b"tenant-a", metrics_body) probe = subprocess.run( [ @@ -323,6 +330,42 @@ def test_external_v2_shard_round_trip_matches_protocol_oracle(self) -> None: ) self.run_daemon([device]) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_io_overlap_matches_serial_scores_and_reports_pipeline_cycles(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + serial_output, serial_results = self.run_daemon( + [device], io_pipeline="serial", io_batch_gb="0.000000004" + ) + overlap_output, overlap_results = self.run_daemon( + [device], io_pipeline="overlap", io_batch_gb="0.000000004" + ) + self.assertEqual(overlap_results, serial_results) + serial_events = [ + json.loads(line) for line in serial_output.splitlines() if line.startswith("{") + ] + overlap_events = [ + json.loads(line) + for line in overlap_output.splitlines() + if line.startswith("{") + ] + serial_cache = next( + event["cache"] + for event in serial_events + if event.get("event") == "tilemaxsim_rust_request" + ) + overlap_cache = next( + event["cache"] + for event in overlap_events + if event.get("event") == "tilemaxsim_rust_request" + ) + self.assertEqual(serial_cache["io_pipeline"]["mode"], "serial") + self.assertEqual(overlap_cache["io_pipeline"]["mode"], "overlap") + self.assertEqual(serial_cache["io_pipeline"]["overlap_cycles"], 0) + self.assertGreaterEqual(overlap_cache["io_pipeline"]["overlap_cycles"], 1) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") def test_tcp_scoring_round_trip_matches_protocol_oracle(self) -> None: device = max( diff --git a/services/tilemaxsimd/src/engine.rs b/services/tilemaxsimd/src/engine.rs index ee08515c..99480882 100644 --- a/services/tilemaxsimd/src/engine.rs +++ b/services/tilemaxsimd/src/engine.rs @@ -13,8 +13,37 @@ use crate::gpu::Gpu; use crate::protocol::{Descriptor, Request}; use crate::shard::{HostCacheStatus, ShardStore, cache_key}; use anyhow::{Result, anyhow, bail}; -use std::collections::{HashMap, HashSet}; -use std::sync::Arc; +use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::{Arc, Mutex, mpsc}; +use std::time::{Duration, Instant}; + +#[derive(Clone, Copy, Debug)] +pub struct IoPipelineConfig { + pub overlap: bool, + pub batch_bytes: usize, +} + +impl IoPipelineConfig { + pub fn validate(self) -> Result { + if self.batch_bytes == 0 { + bail!("I/O pipeline batch size must be positive"); + } + Ok(self) + } +} + +#[derive(Clone, Debug, Default)] +pub struct IoPipelineStatus { + pub overlap_enabled: bool, + pub batch_bytes: usize, + pub resolve_batches: u64, + pub resolve_bytes: u64, + pub resolve_microseconds: u64, + pub upload_microseconds: u64, + pub compute_microseconds: u64, + pub overlap_cycles: u64, + pub overlap_microseconds: u64, +} struct MissingTensor { candidate_index: usize, @@ -23,6 +52,36 @@ struct MissingTensor { payload: Arc<[u8]>, } +struct UnresolvedBatch { + indices: Vec, + descriptors: Vec, + payload_bytes: usize, +} + +struct ResolvedBatch { + tensors: VecDeque, + payload_bytes: usize, + elapsed: Duration, +} + +struct PreparedBatch { + chunks: Vec>, + uploads: Vec)>>, +} + +impl PreparedBatch { + fn empty(device_count: usize) -> Self { + Self { + chunks: (0..device_count).map(|_| Vec::new()).collect(), + uploads: (0..device_count).map(|_| Vec::new()).collect(), + } + } + + fn is_empty(&self) -> bool { + self.chunks.iter().all(Vec::is_empty) + } +} + struct ResidentTensor { candidate_index: usize, device: usize, @@ -34,7 +93,7 @@ struct ResidentTensor { } struct DeviceState { - gpu: Gpu, + gpu: Arc, cache: GpuCache, h2d_batches: u64, h2d_bytes: u64, @@ -42,8 +101,10 @@ struct DeviceState { pub struct Engine { devices: Vec, - store: ShardStore, + store: Arc>, next_device: usize, + io_pipeline: IoPipelineConfig, + io_status: IoPipelineStatus, } #[derive(Clone, Debug, Default)] @@ -75,6 +136,7 @@ pub struct EngineStatus { pub host: HostCacheStatus, pub batch_read_calls: u64, pub batch_read_bytes: u64, + pub io_pipeline: IoPipelineStatus, } impl Engine { @@ -85,10 +147,12 @@ impl Engine { tenant_cache_max_percent: u8, pinned_cache_max_percent: u8, tenant_reservations: &HashMap, + io_pipeline: IoPipelineConfig, ) -> Result { if gpus.is_empty() { bail!("at least one GPU is required"); } + let io_pipeline = io_pipeline.validate()?; let devices = gpus .into_iter() .map(|gpu| { @@ -101,7 +165,7 @@ impl Engine { ) .map_err(|message| anyhow!(message))?; Ok(DeviceState { - gpu, + gpu: Arc::new(gpu), cache, h2d_batches: 0, h2d_bytes: 0, @@ -110,13 +174,40 @@ impl Engine { .collect::>>()?; Ok(Self { devices, - store, + store: Arc::new(Mutex::new(store)), next_device: 0, + io_pipeline, + io_status: IoPipelineStatus { + overlap_enabled: io_pipeline.overlap, + batch_bytes: io_pipeline.batch_bytes, + ..IoPipelineStatus::default() + }, }) } pub fn reload_shards(&mut self) -> Result<()> { - self.store.reload() + self.store + .lock() + .map_err(|_| anyhow!("immutable shard store lock is poisoned"))? + .reload() + } + + /// Return the H2D payload expected for a descriptor at the instant the + /// scheduler builds its next cooperative quantum. Host-cache residency does + /// not make this zero: every L0 miss still consumes upload bandwidth. + pub fn gpu_miss_bytes(&self, descriptor: &Descriptor) -> u64 { + let key = cache_key(descriptor); + if self + .devices + .iter() + .any(|device| device.cache.contains(&key)) + { + return 0; + } + descriptor_payload_bytes(descriptor) + .ok() + .and_then(|bytes| u64::try_from(bytes).ok()) + .unwrap_or(u64::MAX) } pub fn prewarm(&mut self, descriptors: &[Descriptor], batch_size: usize) -> Result<()> { @@ -130,7 +221,11 @@ impl Engine { // allocation. let descriptors = unique_descriptors(descriptors); for batch in descriptors.chunks(batch_size) { - let payloads = self.store.resolve_many(batch, "__resident__")?; + let payloads = self + .store + .lock() + .map_err(|_| anyhow!("immutable shard store lock is poisoned"))? + .resolve_many(batch, "__resident__")?; let mut uploads = (0..self.devices.len()) .map(|_| Vec::<(u64, &[u8])>::new()) .collect::>(); @@ -275,186 +370,392 @@ impl Engine { } } - let hit_result = self.score_devices(request, &hit_chunks, &mut scores); - let hit_cleanup = self.release_chunks(&hit_chunks); - hit_result?; - hit_cleanup?; + let unresolved = split_unresolved_batches( + missing_indices, + missing_descriptors, + self.io_pipeline.batch_bytes, + )?; + let hits = PreparedBatch { + chunks: hit_chunks, + uploads: (0..self.devices.len()).map(|_| Vec::new()).collect(), + }; + if self.io_pipeline.overlap { + let stage_before = self.pipeline_stage_microseconds(); + let pipeline_started = Instant::now(); + self.score_overlapped(request, hits, unresolved, &mut scores)?; + let stage_delta = self + .pipeline_stage_microseconds() + .saturating_sub(stage_before); + let hidden = + stage_delta.saturating_sub(duration_microseconds(pipeline_started.elapsed())); + self.io_status.overlap_microseconds = + self.io_status.overlap_microseconds.saturating_add(hidden); + } else { + self.score_serial(request, hits, unresolved, &mut scores)?; + } - let payloads = self - .store - .resolve_many(&missing_descriptors, &request.tenant)?; - let mut pending = missing_indices - .into_iter() - .zip(missing_descriptors) - .zip(payloads) - .map(|((candidate_index, descriptor), payload)| MissingTensor { - candidate_index, - key: cache_key(&descriptor), - descriptor, - payload, + // Equal content-addressed tensors have equal TileMaxSim scores. Preserve + // every logical candidate id while avoiding duplicate cache acquisitions, + // uploads, and kernel work inside one request. + for (duplicate_index, first_index) in duplicate_candidates { + scores[duplicate_index] = scores[first_index]; + } + + request + .candidates + .iter() + .enumerate() + .map(|(index, descriptor)| { + scores[index] + .map(|score| (descriptor.candidate_id, score)) + .ok_or_else(|| anyhow!("missing native TileMaxSim result")) }) - .collect::>(); + .collect() + } - while !pending.is_empty() { - let mut chunks = (0..self.devices.len()) - .map(|_| Vec::::new()) - .collect::>(); - let mut uploads = (0..self.devices.len()) - .map(|_| Vec::<(u64, &[u8])>::new()) - .collect::>(); - let mut consumed = 0; - for tensor in &pending { - if let Some((device_index, entry)) = - self.devices - .iter_mut() - .enumerate() - .find_map(|(index, device)| { - device - .cache - .acquire_existing(&tensor.key) - .map(|entry| (index, entry)) - }) - { - chunks[device_index].push(ResidentTensor { - candidate_index: tensor.candidate_index, - device: device_index, - key: tensor.key.clone(), - offset: entry.offset, - rows: entry.rows, - transient: false, - newly_admitted: false, - }); - consumed += 1; - continue; - } + fn score_serial( + &mut self, + request: &Request, + mut current: PreparedBatch, + mut unresolved: VecDeque, + scores: &mut [Option], + ) -> Result<()> { + let mut pending = VecDeque::new(); + loop { + if !current.is_empty() { + let result = score_chunks(&self.gpus(), request, ¤t.chunks); + let cleanup = self.release_chunks(¤t.chunks); + let (computed, elapsed) = result?; + self.io_status.compute_microseconds = self + .io_status + .compute_microseconds + .saturating_add(duration_microseconds(elapsed)); + apply_scores(scores, computed); + cleanup?; + } + if pending.is_empty() && unresolved.is_empty() { + return Ok(()); + } + current = + self.fill_next_batch(&request.tenant, &mut unresolved, &mut pending, false)?; + } + } - let mut admitted = None; - for step in 0..self.devices.len() { - let device_index = (self.next_device + step) % self.devices.len(); - let admission = self.devices[device_index].cache.admit_for_tenant( - &request.tenant, - tensor.key.clone(), - tensor.payload.len(), - tensor.descriptor.rows, - tensor.descriptor.dimension, - tensor.descriptor.dtype, - false, - false, - ); - if let Admission::Admitted { offset, .. } = admission { - admitted = Some((device_index, offset, false)); - self.next_device = (device_index + 1) % self.devices.len(); + fn score_overlapped( + &mut self, + request: &Request, + current: PreparedBatch, + unresolved: VecDeque, + scores: &mut [Option], + ) -> Result<()> { + let (sender, receiver) = mpsc::sync_channel::>(1); + let store = Arc::clone(&self.store); + let tenant = request.tenant.clone(); + std::thread::scope(|scope| { + let resolver = scope.spawn(move || { + for batch in unresolved { + let result = resolve_batch(&store, batch, &tenant); + let failed = result.is_err(); + if sender.send(result).is_err() || failed { break; } } + }); + let result = self.score_overlapped_inner(request, current, &receiver, scores); + // An early CUDA or cache error must unblock a resolver waiting on the + // bounded one-batch channel before the scoped worker can join. + drop(receiver); + let resolver_result = resolver + .join() + .map_err(|_| anyhow!("tensor resolver coordinator panicked")); + result.and(resolver_result) + }) + } + + fn score_overlapped_inner( + &mut self, + request: &Request, + mut current: PreparedBatch, + receiver: &mpsc::Receiver>, + scores: &mut [Option], + ) -> Result<()> { + let mut pending = VecDeque::new(); + let mut resolver_done = false; + loop { + if current.is_empty() { + current = self.fill_next_resolved_batch( + &request.tenant, + receiver, + &mut pending, + &mut resolver_done, + false, + )?; + if current.is_empty() && resolver_done { + return Ok(()); + } + continue; + } + + if pending.is_empty() && resolver_done { + let result = score_chunks(&self.gpus(), request, ¤t.chunks); + let cleanup = self.release_chunks(¤t.chunks); + let (computed, elapsed) = result?; + self.io_status.compute_microseconds = self + .io_status + .compute_microseconds + .saturating_add(duration_microseconds(elapsed)); + apply_scores(scores, computed); + cleanup?; + return Ok(()); + } - if admitted.is_none() && chunks.iter().all(Vec::is_empty) { - // TinyLFU rejected the cold item on every device, but the - // request must still be computed. Use one transient slab - // and remove it after the chunk completes. - for step in 0..self.devices.len() { - let device_index = (self.next_device + step) % self.devices.len(); - if let Admission::Admitted { offset, .. } = - self.devices[device_index].cache.admit_for_tenant( - &request.tenant, - tensor.key.clone(), - tensor.payload.len(), - tensor.descriptor.rows, - tensor.descriptor.dimension, - tensor.descriptor.dtype, - false, - true, - ) - { - admitted = Some((device_index, offset, true)); - self.next_device = (device_index + 1) % self.devices.len(); - break; - } + let gpus = self.gpus(); + let (score_result, next_result) = std::thread::scope(|scope| { + let score_worker = scope.spawn(|| score_chunks(&gpus, request, ¤t.chunks)); + let next = self.fill_next_resolved_batch( + &request.tenant, + receiver, + &mut pending, + &mut resolver_done, + true, + ); + let score = score_worker + .join() + .map_err(|_| anyhow!("GPU score coordinator panicked")) + .and_then(|result| result); + (score, next) + }); + let cleanup_current = self.release_chunks(¤t.chunks); + + let (computed, compute_elapsed) = match score_result { + Ok(result) => result, + Err(error) => { + if let Ok(next) = next_result { + self.release_chunks(&next.chunks)?; } + cleanup_current?; + return Err(error); } + }; + self.io_status.compute_microseconds = self + .io_status + .compute_microseconds + .saturating_add(duration_microseconds(compute_elapsed)); + apply_scores(scores, computed); + cleanup_current?; - let Some((device_index, offset, transient)) = admitted else { - if chunks.iter().any(|chunk| !chunk.is_empty()) { - break; - } - bail!("one tensor cannot be scheduled in any configured Rust GPU block cache"); - }; - uploads[device_index].push((offset, tensor.payload.as_ref())); - chunks[device_index].push(ResidentTensor { + let next = next_result?; + self.io_status.overlap_cycles = self.io_status.overlap_cycles.saturating_add(1); + current = next; + } + } + + fn fill_next_resolved_batch( + &mut self, + tenant: &str, + receiver: &mpsc::Receiver>, + pending: &mut VecDeque, + resolver_done: &mut bool, + allow_deferred: bool, + ) -> Result { + if pending.is_empty() && !*resolver_done { + match receiver.recv() { + Ok(result) => { + let resolved = result?; + self.record_resolution(&resolved); + *pending = resolved.tensors; + } + Err(_) => *resolver_done = true, + } + } + if pending.is_empty() { + return Ok(PreparedBatch::empty(self.devices.len())); + } + let mut prepared = self.prepare_pending(pending, tenant, allow_deferred)?; + if prepared.is_empty() { + return Ok(prepared); + } + let upload_elapsed = self.upload_prepared(&mut prepared)?; + self.io_status.upload_microseconds = self + .io_status + .upload_microseconds + .saturating_add(duration_microseconds(upload_elapsed)); + Ok(prepared) + } + + fn fill_next_batch( + &mut self, + tenant: &str, + unresolved: &mut VecDeque, + pending: &mut VecDeque, + allow_deferred: bool, + ) -> Result { + if pending.is_empty() + && let Some(batch) = unresolved.pop_front() + { + let resolved = resolve_batch(&self.store, batch, tenant)?; + self.record_resolution(&resolved); + *pending = resolved.tensors; + } + if pending.is_empty() { + return Ok(PreparedBatch::empty(self.devices.len())); + } + let mut prepared = self.prepare_pending(pending, tenant, allow_deferred)?; + if prepared.is_empty() { + return Ok(prepared); + } + let upload_elapsed = self.upload_prepared(&mut prepared)?; + self.io_status.upload_microseconds = self + .io_status + .upload_microseconds + .saturating_add(duration_microseconds(upload_elapsed)); + Ok(prepared) + } + + fn record_resolution(&mut self, resolved: &ResolvedBatch) { + self.io_status.resolve_batches = self.io_status.resolve_batches.saturating_add(1); + self.io_status.resolve_bytes = self + .io_status + .resolve_bytes + .saturating_add(u64::try_from(resolved.payload_bytes).unwrap_or(u64::MAX)); + self.io_status.resolve_microseconds = self + .io_status + .resolve_microseconds + .saturating_add(duration_microseconds(resolved.elapsed)); + } + + fn pipeline_stage_microseconds(&self) -> u64 { + self.io_status + .resolve_microseconds + .saturating_add(self.io_status.upload_microseconds) + .saturating_add(self.io_status.compute_microseconds) + } + + fn prepare_pending( + &mut self, + pending: &mut VecDeque, + tenant: &str, + allow_deferred: bool, + ) -> Result { + let mut prepared = PreparedBatch::empty(self.devices.len()); + let mut consumed = 0; + for tensor in pending.iter() { + if let Some((device_index, entry)) = + self.devices + .iter_mut() + .enumerate() + .find_map(|(index, device)| { + device + .cache + .acquire_existing(&tensor.key) + .map(|entry| (index, entry)) + }) + { + prepared.chunks[device_index].push(ResidentTensor { candidate_index: tensor.candidate_index, device: device_index, key: tensor.key.clone(), - offset, - rows: tensor.descriptor.rows, - transient, - newly_admitted: true, + offset: entry.offset, + rows: entry.rows, + transient: false, + newly_admitted: false, }); consumed += 1; + continue; } - if consumed == 0 { - bail!("Rust multi-GPU scheduler made no progress"); + let mut admitted = None; + for step in 0..self.devices.len() { + let device_index = (self.next_device + step) % self.devices.len(); + let admission = self.devices[device_index].cache.admit_for_tenant( + tenant, + tensor.key.clone(), + tensor.payload.len(), + tensor.descriptor.rows, + tensor.descriptor.dimension, + tensor.descriptor.dtype, + false, + false, + ); + if let Admission::Admitted { offset, .. } = admission { + admitted = Some((device_index, offset, false)); + self.next_device = (device_index + 1) % self.devices.len(); + break; + } } - let upload_succeeded = match self.upload_devices(&chunks, &uploads) { - Ok(upload_succeeded) => upload_succeeded, - Err((error, upload_succeeded)) => { - self.cleanup_after_upload_failure(&chunks, &upload_succeeded)?; - return Err(error); + + if admitted.is_none() && prepared.is_empty() { + // TinyLFU may reject a cold item, but every requested tensor must + // still be scored. A transient page run is removed after this + // chunk completes. During overlap it may be temporarily deferred + // because the current compute batch still owns all freeable pages. + for step in 0..self.devices.len() { + let device_index = (self.next_device + step) % self.devices.len(); + if let Admission::Admitted { offset, .. } = + self.devices[device_index].cache.admit_for_tenant( + tenant, + tensor.key.clone(), + tensor.payload.len(), + tensor.descriptor.rows, + tensor.descriptor.dimension, + tensor.descriptor.dtype, + false, + true, + ) + { + admitted = Some((device_index, offset, true)); + self.next_device = (device_index + 1) % self.devices.len(); + break; + } } + } + + let Some((device_index, offset, transient)) = admitted else { + if !prepared.is_empty() || allow_deferred { + break; + } + bail!("one tensor cannot be scheduled in any configured Rust GPU block cache"); }; - debug_assert!(upload_succeeded.iter().all(|succeeded| *succeeded)); - let score_result = self.score_devices(request, &chunks, &mut scores); - let cleanup_result = self.release_chunks(&chunks); - score_result?; - cleanup_result?; - pending.drain(..consumed); + prepared.uploads[device_index].push((offset, Arc::clone(&tensor.payload))); + prepared.chunks[device_index].push(ResidentTensor { + candidate_index: tensor.candidate_index, + device: device_index, + key: tensor.key.clone(), + offset, + rows: tensor.descriptor.rows, + transient, + newly_admitted: true, + }); + consumed += 1; } - // Equal content-addressed tensors have equal TileMaxSim scores. Preserve - // every logical candidate id while avoiding duplicate cache acquisitions, - // uploads, and kernel work inside one request. - for (duplicate_index, first_index) in duplicate_candidates { - scores[duplicate_index] = scores[first_index]; + if consumed == 0 { + if allow_deferred { + return Ok(prepared); + } + bail!("Rust multi-GPU scheduler made no progress"); } - - request - .candidates - .iter() - .enumerate() - .map(|(index, descriptor)| { - scores[index] - .map(|score| (descriptor.candidate_id, score)) - .ok_or_else(|| anyhow!("missing native TileMaxSim result")) - }) - .collect() + pending.drain(..consumed); + Ok(prepared) } - fn upload_devices( - &mut self, - chunks: &[Vec], - uploads: &[Vec<(u64, &[u8])>], - ) -> Result, (anyhow::Error, Vec)> { + fn upload_prepared(&mut self, prepared: &mut PreparedBatch) -> Result { + let started = Instant::now(); + let gpus = self.gpus(); let results = std::thread::scope(|scope| { let mut workers = Vec::new(); - for (device_index, ((device, chunk), upload)) in - self.devices.iter_mut().zip(chunks).zip(uploads).enumerate() - { + for (device_index, (gpu, upload)) in gpus.iter().zip(&prepared.uploads).enumerate() { if upload.is_empty() { continue; } workers.push(( device_index, scope.spawn(move || -> Result<()> { - device.gpu.upload_batch(upload)?; - device.h2d_batches += 1; - device.h2d_bytes += upload + let items = upload .iter() - .map(|(_, payload)| payload.len() as u64) - .sum::(); - // Keep the chunk borrow in this worker so upload metadata and - // payload lifetimes remain tied to the scoped thread. - let _ = chunk; - Ok(()) + .map(|(offset, payload)| (*offset, payload.as_ref())) + .collect::>(); + gpu.upload_batch(&items) }), )); } @@ -476,9 +777,17 @@ impl Engine { if let Err(error) = result { succeeded[device] = false; first_error.get_or_insert(error); + continue; } + self.devices[device].h2d_batches = self.devices[device].h2d_batches.saturating_add(1); + self.devices[device].h2d_bytes = self.devices[device].h2d_bytes.saturating_add( + prepared.uploads[device] + .iter() + .map(|(_, payload)| payload.len() as u64) + .sum(), + ); } - for (device, chunk) in chunks.iter().enumerate() { + for (device, chunk) in prepared.chunks.iter().enumerate() { if !succeeded[device] { continue; } @@ -490,59 +799,20 @@ impl Engine { } } } + prepared.uploads.iter_mut().for_each(Vec::clear); if let Some(error) = first_error { - Err((error, succeeded)) + self.cleanup_after_upload_failure(&prepared.chunks, &succeeded)?; + Err(error) } else { - Ok(succeeded) + Ok(started.elapsed()) } } - fn score_devices( - &mut self, - request: &Request, - chunks: &[Vec], - scores: &mut [Option], - ) -> Result<()> { - let completed = std::thread::scope(|scope| -> Result>> { - let mut workers = Vec::new(); - for (device, chunk) in self.devices.iter_mut().zip(chunks) { - if chunk.is_empty() { - continue; - } - workers.push(scope.spawn(move || -> Result> { - let offsets = chunk.iter().map(|item| item.offset).collect::>(); - let rows = chunk.iter().map(|item| item.rows).collect::>(); - let computed = device.gpu.score( - &request.query, - request.query_rows, - request.dimension, - request.dtype, - &offsets, - &rows, - )?; - Ok(chunk - .iter() - .zip(computed) - .map(|(tensor, score)| (tensor.candidate_index, score)) - .collect()) - })); - } - let mut completed = Vec::new(); - for worker in workers { - completed.push( - worker - .join() - .map_err(|_| anyhow!("GPU worker panicked"))??, - ); - } - Ok(completed) - })?; - for device_scores in completed { - for (candidate_index, score) in device_scores { - scores[candidate_index] = Some(score); - } - } - Ok(()) + fn gpus(&self) -> Vec> { + self.devices + .iter() + .map(|device| Arc::clone(&device.gpu)) + .collect() } fn release_chunks(&mut self, chunks: &[Vec]) -> Result<()> { @@ -617,11 +887,13 @@ impl Engine { h2d_bytes: device.h2d_bytes, }) .collect(); + let store = self.store.lock().unwrap_or_else(|error| error.into_inner()); EngineStatus { devices, - host: self.store.host_status(), - batch_read_calls: self.store.batch_read_calls, - batch_read_bytes: self.store.batch_read_bytes, + host: store.host_status(), + batch_read_calls: store.batch_read_calls, + batch_read_bytes: store.batch_read_bytes, + io_pipeline: self.io_status.clone(), } } @@ -667,10 +939,156 @@ impl Engine { "host_admission_rejections": status.host.admission_rejections, "batch_read_calls": status.batch_read_calls, "batch_read_bytes": status.batch_read_bytes, + "io_pipeline": { + "mode": if status.io_pipeline.overlap_enabled { "overlap" } else { "serial" }, + "batch_bytes": status.io_pipeline.batch_bytes, + "resolve_batches": status.io_pipeline.resolve_batches, + "resolve_bytes": status.io_pipeline.resolve_bytes, + "resolve_microseconds": status.io_pipeline.resolve_microseconds, + "upload_microseconds": status.io_pipeline.upload_microseconds, + "compute_microseconds": status.io_pipeline.compute_microseconds, + "overlap_cycles": status.io_pipeline.overlap_cycles, + "overlap_microseconds": status.io_pipeline.overlap_microseconds, + }, + }) + } +} + +fn resolve_batch( + store: &Arc>, + batch: UnresolvedBatch, + tenant: &str, +) -> Result { + let started = Instant::now(); + let payloads = store + .lock() + .map_err(|_| anyhow!("immutable shard store lock is poisoned"))? + .resolve_many(&batch.descriptors, tenant)?; + let elapsed = started.elapsed(); + let tensors = batch + .indices + .into_iter() + .zip(batch.descriptors) + .zip(payloads) + .map(|((candidate_index, descriptor), payload)| MissingTensor { + candidate_index, + key: cache_key(&descriptor), + descriptor, + payload, }) + .collect(); + Ok(ResolvedBatch { + tensors, + payload_bytes: batch.payload_bytes, + elapsed, + }) +} + +fn split_unresolved_batches( + indices: Vec, + descriptors: Vec, + maximum_bytes: usize, +) -> Result> { + if indices.len() != descriptors.len() { + bail!("internal missing tensor metadata disagrees"); + } + let mut output = VecDeque::new(); + let mut batch_indices = Vec::new(); + let mut batch_descriptors = Vec::new(); + let mut batch_bytes = 0_usize; + for (index, descriptor) in indices.into_iter().zip(descriptors) { + let payload_bytes = descriptor_payload_bytes(&descriptor)?; + if !batch_descriptors.is_empty() + && batch_bytes.saturating_add(payload_bytes) > maximum_bytes + { + output.push_back(UnresolvedBatch { + indices: std::mem::take(&mut batch_indices), + descriptors: std::mem::take(&mut batch_descriptors), + payload_bytes: batch_bytes, + }); + batch_bytes = 0; + } + batch_indices.push(index); + batch_descriptors.push(descriptor); + batch_bytes = batch_bytes + .checked_add(payload_bytes) + .ok_or_else(|| anyhow!("I/O pipeline batch byte count overflow"))?; + } + if !batch_descriptors.is_empty() { + output.push_back(UnresolvedBatch { + indices: batch_indices, + descriptors: batch_descriptors, + payload_bytes: batch_bytes, + }); + } + Ok(output) +} + +fn descriptor_payload_bytes(descriptor: &Descriptor) -> Result { + let scalar_bytes = match descriptor.dtype { + 1 => 4_usize, + 2 => 2_usize, + _ => bail!("unsupported tensor dtype"), + }; + (descriptor.rows as usize) + .checked_mul(descriptor.dimension as usize) + .and_then(|value| value.checked_mul(scalar_bytes)) + .ok_or_else(|| anyhow!("tensor payload byte count overflow")) +} + +fn score_chunks( + gpus: &[Arc], + request: &Request, + chunks: &[Vec], +) -> Result<(Vec<(usize, f32)>, Duration)> { + let started = Instant::now(); + let completed = std::thread::scope(|scope| -> Result>> { + let mut workers = Vec::new(); + for (gpu, chunk) in gpus.iter().zip(chunks) { + if chunk.is_empty() { + continue; + } + workers.push(scope.spawn(move || -> Result> { + let offsets = chunk.iter().map(|item| item.offset).collect::>(); + let rows = chunk.iter().map(|item| item.rows).collect::>(); + let computed = gpu.score( + &request.query, + request.query_rows, + request.dimension, + request.dtype, + &offsets, + &rows, + )?; + Ok(chunk + .iter() + .zip(computed) + .map(|(tensor, score)| (tensor.candidate_index, score)) + .collect()) + })); + } + let mut completed = Vec::new(); + for worker in workers { + completed.push( + worker + .join() + .map_err(|_| anyhow!("GPU worker panicked"))??, + ); + } + Ok(completed) + })?; + Ok((completed.into_iter().flatten().collect(), started.elapsed())) +} + +fn apply_scores(scores: &mut [Option], computed: Vec<(usize, f32)>) { + for (candidate_index, score) in computed { + scores[candidate_index] = Some(score); } } +fn duration_microseconds(duration: Duration) -> u64 { + u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) +} + fn unique_descriptors(descriptors: &[Descriptor]) -> Vec { let mut seen = HashSet::with_capacity(descriptors.len()); descriptors @@ -725,4 +1143,21 @@ mod tests { assert_eq!(unique[1].candidate_id, 3); assert_eq!(unique[2].candidate_id, 4); } + + #[test] + fn io_batches_are_payload_bounded_and_keep_one_oversized_tensor() { + let descriptors = vec![ + descriptor(1, "a", 1), + descriptor(2, "b", 1), + descriptor(3, "c", 3), + ]; + // One row is 640 bytes. The first two fit exactly; the final tensor is + // larger than the target but must remain a schedulable one-item batch. + let batches = split_unresolved_batches(vec![0, 1, 2], descriptors, 1_280).unwrap(); + assert_eq!(batches.len(), 2); + assert_eq!(batches[0].indices, vec![0, 1]); + assert_eq!(batches[0].payload_bytes, 1_280); + assert_eq!(batches[1].indices, vec![2]); + assert_eq!(batches[1].payload_bytes, 1_920); + } } diff --git a/services/tilemaxsimd/src/gpu.rs b/services/tilemaxsimd/src/gpu.rs index de37e093..2059b161 100644 --- a/services/tilemaxsimd/src/gpu.rs +++ b/services/tilemaxsimd/src/gpu.rs @@ -57,9 +57,13 @@ pub struct Gpu { tensor_bytes: usize, } -// SAFETY: `Gpu` uniquely owns the native handle. It may move to a scoped -// worker, but no method exposes the pointer and all calls require `&mut self`. +// SAFETY: `Gpu` uniquely owns the native handle. The engine permits at most one +// upload and one score call at a time. The native implementation assigns those +// operations distinct CUDA streams and disjoint host/device workspaces; tensor +// cache references prevent an upload from overwriting pages used by scoring. +// Destruction happens only after every scoped worker has joined. unsafe impl Send for Gpu {} +unsafe impl Sync for Gpu {} impl Gpu { pub fn create(device: i32, total_bytes: usize, workspace_bytes: usize) -> Result { @@ -97,7 +101,7 @@ impl Gpu { self.tensor_bytes } - pub fn upload_batch(&mut self, items: &[(u64, &[u8])]) -> Result<()> { + pub fn upload_batch(&self, items: &[(u64, &[u8])]) -> Result<()> { if items.is_empty() { return Ok(()); } @@ -124,7 +128,7 @@ impl Gpu { } pub fn score( - &mut self, + &self, query: &[u8], query_rows: u32, dimension: u32, diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs index c12e69f9..b35beda8 100644 --- a/services/tilemaxsimd/src/main.rs +++ b/services/tilemaxsimd/src/main.rs @@ -17,7 +17,7 @@ mod shard; use anyhow::{Context, Result, anyhow, bail}; use clap::Parser; -use engine::{Engine, EngineStatus}; +use engine::{Engine, EngineStatus, IoPipelineConfig}; use gpu::Gpu; use protocol::{HEADER_BYTES, VERSION_EXTERNAL, VERSION_SCHEDULED_EXTERNAL}; use scheduler::{RequestQueue, Scheduled, SchedulerPolicy}; @@ -75,6 +75,12 @@ struct Args { gpu_workspace_gb: usize, #[arg(long, default_value = "8", value_parser = parse_gb)] host_cache_gb: usize, + /// Overlap tensor resolution and H2D upload with the current CUDA batch. + #[arg(long, default_value = "overlap", value_parser = ["serial", "overlap"])] + io_pipeline: String, + /// Maximum logical tensor payload resolved by one pipeline stage, in GB. + #[arg(long, default_value = "0.05", value_parser = parse_gb)] + io_batch_gb: usize, #[arg(long, default_value_t = 80, value_parser = clap::value_parser!(u8).range(1..=100))] host_tenant_cache_max_percent: u8, #[arg(long = "contract-root", required = true, value_parser = parse_contract_root)] @@ -133,6 +139,9 @@ struct Args { scheduler_quantum_tokens: u64, #[arg(long, default_value_t = 4_000_000_000)] scheduler_quantum_fmas: u64, + /// Maximum L0-miss payload admitted to one cooperative quantum, in GB. + #[arg(long, default_value = "1", value_parser = parse_gb)] + scheduler_quantum_io_gb: usize, #[arg(long = "tenant-weight", value_parser = parse_tenant_weight)] tenant_weights: Vec<(String, f64)>, } @@ -287,6 +296,10 @@ fn main() -> Result<()> { args.tenant_cache_max_percent, args.pinned_cache_max_percent, &tenant_cache_reservations, + IoPipelineConfig { + overlap: args.io_pipeline == "overlap", + batch_bytes: args.io_batch_gb, + }, )?; if args.gpu_cache_mode == "resident" && args.resident_manifests.is_empty() { bail!("resident GPU cache mode requires at least one resident manifest"); @@ -358,6 +371,7 @@ fn main() -> Result<()> { quantum_candidates: args.scheduler_quantum_candidates, quantum_tokens: args.scheduler_quantum_tokens, quantum_fmas: args.scheduler_quantum_fmas, + quantum_io_bytes: args.scheduler_quantum_io_gb, socket_io_timeout: Duration::from_millis(args.socket_io_timeout_ms), tenant_weights, }; @@ -420,6 +434,9 @@ fn main() -> Result<()> { "max_tenant_queued_requests": args.max_tenant_queued_requests, "max_inflight_request_bytes": args.max_inflight_request_gb, "scheduler_quantum_fmas": args.scheduler_quantum_fmas, + "scheduler_quantum_io_bytes": args.scheduler_quantum_io_gb, + "io_pipeline": args.io_pipeline, + "io_batch_bytes": args.io_batch_gb, "status_socket": args.status_socket, "status_listen": args.status_listen, "cache": ready_cache, @@ -656,6 +673,7 @@ struct SchedulerConfig { quantum_candidates: usize, quantum_tokens: u64, quantum_fmas: u64, + quantum_io_bytes: usize, socket_io_timeout: Duration, tenant_weights: std::collections::HashMap, } @@ -1117,7 +1135,7 @@ fn run_scheduler( if queue.len() == 0 && channel_open { match receiver.recv_timeout(Duration::from_millis(50)) { - Ok(work) => enqueue_work(&mut queue, work, &config, &metrics), + Ok(work) => enqueue_work(&mut queue, work, &config, &engine, &metrics), Err(mpsc::RecvTimeoutError::Timeout) => continue, Err(mpsc::RecvTimeoutError::Disconnected) => channel_open = false, } @@ -1130,7 +1148,7 @@ fn run_scheduler( break; } match receiver.recv_timeout(remaining) { - Ok(work) => enqueue_work(&mut queue, work, &config, &metrics), + Ok(work) => enqueue_work(&mut queue, work, &config, &engine, &metrics), Err(mpsc::RecvTimeoutError::Timeout) => break, Err(mpsc::RecvTimeoutError::Disconnected) => { channel_open = false; @@ -1183,7 +1201,7 @@ fn run_scheduler( "request deadline expired before execution", )) } else { - let end = next_quantum_end(&work, &config); + let end = next_quantum_end(&work, &config, &engine); let quantum = protocol::Request { protocol_version: work.request.protocol_version, request_id: work.request.request_id, @@ -1232,7 +1250,7 @@ fn run_scheduler( )) } else if end < work.request.candidates.len() { metrics.scheduler_requeues.fetch_add(1, Ordering::Relaxed); - let cost = estimated_next_work(&work, &config); + let cost = estimated_next_work(&work, &config, &engine); queue.push(Scheduled::new( tenant.clone(), priority, @@ -1298,6 +1316,7 @@ fn enqueue_work( queue: &mut RequestQueue, work: Work, config: &SchedulerConfig, + engine: &Engine, metrics: &RuntimeMetrics, ) { match work.request.priority.cmp(&0) { @@ -1311,7 +1330,7 @@ fn enqueue_work( .admitted_priority_positive .fetch_add(1, Ordering::Relaxed), }; - let cost = estimated_next_work(&work, config); + let cost = estimated_next_work(&work, config, engine); queue.push(Scheduled::new( work.request.tenant.clone(), work.request.priority, @@ -1345,8 +1364,8 @@ fn is_fatal_cuda_diagnostic(diagnostic: &str) -> bool { .any(|marker| diagnostic.contains(marker)) } -fn next_quantum_end(work: &Work, config: &SchedulerConfig) -> usize { - quantum_end( +fn next_quantum_end(work: &Work, config: &SchedulerConfig, engine: &Engine) -> usize { + let compute_end = quantum_end( &work.request.candidates, work.next_candidate, config.quantum_candidates, @@ -1354,9 +1373,37 @@ fn next_quantum_end(work: &Work, config: &SchedulerConfig) -> usize { config.quantum_fmas, work.request.query_rows, work.request.dimension, + ); + io_quantum_end( + &work.request.candidates, + work.next_candidate, + compute_end, + config.quantum_io_bytes, + |candidate| engine.gpu_miss_bytes(candidate), ) } +fn io_quantum_end( + candidates: &[protocol::Descriptor], + start: usize, + compute_end: usize, + maximum_io_bytes: usize, + mut miss_bytes: impl FnMut(&protocol::Descriptor) -> u64, +) -> usize { + let maximum_io_bytes = u64::try_from(maximum_io_bytes).unwrap_or(u64::MAX); + let mut end = start; + let mut bytes = 0_u64; + while end < compute_end { + let candidate_bytes = miss_bytes(&candidates[end]); + if end > start && bytes.saturating_add(candidate_bytes) > maximum_io_bytes { + break; + } + bytes = bytes.saturating_add(candidate_bytes); + end += 1; + } + end.max((start + 1).min(compute_end)) +} + fn quantum_end( candidates: &[protocol::Descriptor], start: usize, @@ -1391,8 +1438,8 @@ fn candidate_fmas(query_rows: u32, dimension: u32, document_rows: u32) -> u64 { .saturating_mul(u64::from(dimension)) } -fn estimated_next_work(work: &Work, config: &SchedulerConfig) -> u64 { - let end = next_quantum_end(work, config); +fn estimated_next_work(work: &Work, config: &SchedulerConfig, engine: &Engine) -> u64 { + let end = next_quantum_end(work, config, engine); work.request.candidates[work.next_candidate..end] .iter() .map(|candidate| { @@ -2068,6 +2115,89 @@ fn render_metrics(metrics: &RuntimeMetrics) -> String { engine.batch_read_bytes ) .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_enabled Whether tensor I/O overlaps CUDA scoring." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_io_pipeline_enabled gauge").unwrap(); + writeln!( + output, + "tilemaxsim_io_pipeline_enabled {}", + usize::from(engine.io_pipeline.overlap_enabled) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_batch_bytes Configured logical tensor bytes per I/O stage." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_io_pipeline_batch_bytes gauge").unwrap(); + writeln!( + output, + "tilemaxsim_io_pipeline_batch_bytes {}", + engine.io_pipeline.batch_bytes + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_batches_total Tensor pipeline stage executions." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_io_pipeline_batches_total counter" + ) + .unwrap(); + for (stage, value) in [ + ("resolve", engine.io_pipeline.resolve_batches), + ("overlap", engine.io_pipeline.overlap_cycles), + ] { + writeln!( + output, + "tilemaxsim_io_pipeline_batches_total{{stage=\"{stage}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_resolve_bytes_total Logical tensor payload passed through resolve stages." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_io_pipeline_resolve_bytes_total counter" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_io_pipeline_resolve_bytes_total {}", + engine.io_pipeline.resolve_bytes + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_seconds_total Cumulative tensor pipeline stage and hidden time." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_io_pipeline_seconds_total counter" + ) + .unwrap(); + for (stage, microseconds) in [ + ("resolve", engine.io_pipeline.resolve_microseconds), + ("upload", engine.io_pipeline.upload_microseconds), + ("compute", engine.io_pipeline.compute_microseconds), + ("overlap", engine.io_pipeline.overlap_microseconds), + ] { + writeln!( + output, + "tilemaxsim_io_pipeline_seconds_total{{stage=\"{stage}\"}} {}", + microseconds as f64 / 1_000_000.0 + ) + .unwrap(); + } output } @@ -2216,10 +2346,10 @@ fn load_resident_manifests(values: &[(String, PathBuf)]) -> Result